Python matplotlib — set_data и set_3d_properties, похоже, не обновляют мой график

#python #arrays #python-3.x #numpy #matplotlib

#python #массивы #python-3.x #numpy #matplotlib

Вопрос:

В настоящее время я работаю над скриптом Yee Solver для uni, но когда я пытаюсь анимировать свой 3D-график, график не соответствует ожиданиям. Это работает для 2D-графика, но, похоже, я не могу перевести это в 3D. Насколько я понимаю, set_data и set_3d_properties для работы требуется одномерный массив, который я ввожу.

 import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from matplotlib.widgets import Slider

# Program Variables
wv_lgth_num = 10
graph_type  = '3d'
t           = 0

# Physical Constants
c           = 3e8
mu_r        = 1
eps_r       = 1

# Source Constants
f           = 2e9
omega       = 2*f*(np.pi)
amp         = 1.0
wv_lgth     = c/f
period      = 1/f

# Step size
dz          = wv_lgth/20
dt          = ((c/f)/20)/c
#dt          = ((1/f)/20)/1

# Axis Grids
z_grid      = np.arange(0,wv_lgth_num*wv_lgth,dz)
t_grid      = np.arange(0,10*period,dt)

# Number of steps
num_z       = z_grid.size
num_t       = t_grid.size

# Coefficients
coe_E = c*dt/(eps_r*dz)
coe_H = c*dt/(mu_r*dz)

# E and H Matricies
E_mat = np.zeros((num_z,num_t))
H_mat = np.zeros((num_z,num_t))

# Generating Values for E and H
for time in range(0,num_t-1):
    for pos in range(0,num_z-1):
    
        # Source Wave
        if pos == 0:
            H_mat[0,time] = amp*np.sin(omega*t_grid[time])

        # All cases of Yee Solver
        if pos == 1:
            if time == 0:
                H_mat[1,0] = 0
                E_mat[0,0] = 0
            else:
                H_mat[1,time] = H_mat[1,time-1]   coe_H*(E_mat[1,time-1] - E_mat[0,time-1])
                E_mat[0,time] = E_mat[0,time-1]   coe_E*(H_mat[1,time] - H_mat[0,time])
        if pos > 1 and pos != num_z-1:
            if time == 0:
                H_mat[pos,0] = 0
                E_mat[pos-1,0] = 0
            if time > 0:
                H_mat[pos,time] = H_mat[pos,time-1]   coe_H*(E_mat[pos,time-1] - E_mat[pos-1,time-1])
                E_mat[pos-1,time] = E_mat[pos-1,time-1]   coe_E*(H_mat[pos,time] - H_mat[pos-1,time])
        if pos == num_z-1:
            if time == 0:
                H_mat[num_z-1,0] = 0
                E_mat[num_z-2,0] = 0
                E_mat[num_z-1,0] = 0
            if time > 0:
                H_mat[num_z-1,time] = H_mat[num_z-1,time-1]   coe_H*(E_mat[num_z-1,time-1] - E_mat[num_z-2,time-1])
                E_mat[num_z-2,time] = E_mat[num_z-2,time-1]   coe_E*(H_mat[num_z-1,time] - H_mat[num_z-2,time])
                E_mat[num_z-1,time] = E_mat[num_z-2,time]     

def update(val):
    t = slider_time.val
    if graph_type == '2d':
        a.set_ydata(E_mat[:,t])
        b.set_ydata(H_mat[:,t])
    if graph_type == '3d':
        a.set_3d_properties(E_mat[:,t])
        a.set_data(z_grid,np.zeros((num_z,num_t))[:,t])
        b.set_3d_properties(np.zeros((num_z,num_t))[:,t])
        b.set_data(z_grid,H_mat[:t])
    fig.canvas.draw_idle()

print(H_mat)
print(H_mat[:,t].size)
print(z_grid)
print(np.zeros((num_z,num_t))[:,t].size)

# Creating plot
if graph_type == '3d':
    fig, ax = plt.subplots()
    ax = plt.axes(projection='3d')
    b, = ax.plot3D(z_grid,H_mat[:,t],np.zeros((num_z,num_t))[:,t], label='H')
    a, = ax.plot3D(z_grid,np.zeros((num_z,num_t))[:,t],E_mat[:,t], label='E')
    plt.title('Light Wave')
    ax.set_xlabel('z')
    ax.set_ylabel('x')
    ax.set_zlabel('y')
    plt.legend()

    ax_time = plt.axes([0.25,0.1,0.65,0.03])
    slider_time = Slider(ax_time,'Time',0,num_t-2,valinit=0,valstep=1)
    slider_time.on_changed(update)
    
    plt.show()

if graph_type == '2d':
    fig, ax = plt.subplots()
    plt.subplots_adjust(left=0.25, bottom=0.25)
    a, = plt.plot(z_grid,E_mat[:,t], label='E (yz plane)')
    b, = plt.plot(z_grid,H_mat[:,t], label='H (xz plane)')
    plt.title('Light Wave')
    plt.xlabel('z')
    plt.ylabel('x')
    plt.legend()

    ax_time = plt.axes([0.25,0.1,0.65,0.03])
    slider_time = Slider(ax_time,'Time',0,num_t-2,valinit=0,valstep=1)
    slider_time.on_changed(update)
    
    plt.show()
  

Буду признателен за любую помощь. Средний цикл for просто генерирует мои функции, используя метод Yee.