使用 matplotlib 制作 3d 冲浪动画

如何解决使用 matplotlib 制作 3d 冲浪动画

我正在使用 matplotlib 创建 3d 冲浪动画。我基于 matlab 代码。

完整代码如下：

from __future__ import division
import numpy as np
from pylab import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation


lx = 35
ly = 53
divide = [35,53]
dx = lx/divide[0]
dy = ly/divide[1]
no_nodes_x = lx + 1
no_nodes_y = ly + 1
no_nodes = no_nodes_x * no_nodes_y
no_el = (no_nodes_x - 1) * (no_nodes_y - 1)
no_nodes_el = 4
timesteps = 30
dt = 1
# dx = 1
# dy = 1

#physical parameters
kappa = 1
q = 0

K,K_trans = np.zeros([no_nodes,no_nodes]),np.zeros([no_nodes,no_nodes])
T = np.zeros([no_nodes_x,no_nodes_y])
F,R = np.zeros([no_nodes,1]),1])
#initial condition
T[round(no_nodes_x/2) - 5:round(no_nodes_x/2) + 4,round(no_nodes_y/2) - 5:round(no_nodes_y/2) + 4,] = 30

#setup elemnt matrix
#conductivity terms matrix
A = np.array([[-2/3*kappa/dx*dy-2/3*kappa*dx/dy+kappa*(1/(dx**2)+1/(dy**2))*dx*dy,-1/3*kappa/dx*dy+1/6*kappa*dx/dy,-1/6*kappa/dx*dy-1/6*kappa*dx/dy,1/6*kappa/dx*dy-1/3*kappa*dx/dy],[-1/3*kappa/dx*dy+1/6*kappa*dx/dy,1/3*kappa/dx*dy+1/3*kappa*dx/dy,1/6*kappa/dx*dy-1/3*kappa*dx/dy,-1/6*kappa/dx*dy-1/6*kappa*dx/dy],[-1/6*kappa/dx*dy-1/6*kappa*dx/dy,-1/3*kappa/dx*dy+1/6*kappa*dx/dy],[1/6*kappa/dx*dy-1/3*kappa*dx/dy,1/3*kappa/dx*dy+1/3*kappa*dx/dy]])

#transient terms matrix
A_trans = np.array([[1/9*dy*dx,1/18*dy*dx,1/36*dy*dx,1/18*dy*dx],[1/18*dy*dx,1/9*dy*dx,1/36*dy*dx],[1/36*dy*dx,1/9*dy*dx]])

#right hand side
R_el = np.array([[1/4*q*dx*dy],[1/4*q*dx*dy],[1/4*q*dx*dy]])


nodes = np.zeros([1855,4],dtype=int)
for j in range(1,no_nodes_y):
    for i in range(1,no_nodes_x):
        iel = i + (j-1) * (no_nodes_x - 1) - 1
        nodes[iel,0] = i + (j-1) * no_nodes_x
        nodes[iel,1] = i + 1 + (j-1) * no_nodes_x
        nodes[iel,2] = i + 1 + j * no_nodes_x
        nodes[iel,3] = i + j * no_nodes_x




#creates vector from temperature matrix
np.set_printoptions(threshold=np.inf)

T = np.array([T.flatten('F')]).T

for iel in range(no_el):
    for i in range(no_nodes_el):
        k = int(nodes[iel,i]) - 1
        for j in range(no_nodes_el):
            p = int(nodes[iel,j]) - 1
            K[k,p] += A[i,j]
            K_trans[k,p] += A_trans[i,j]

        R[k] += R_el[i]


timesteps = 10
T2d = np.zeros([no_nodes_x,no_nodes_y,timesteps+1])
T_step = T.reshape(no_nodes_x,order='F').copy()
T2d[:,:,0] = T_step
for t in range(timesteps):
    F = (R * dt) + np.dot(K_trans,T)

    K_tot = K_trans + K * dt

    T = np.linalg.solve(K_tot,F)

    T_step = T.reshape(no_nodes_x,order='F').copy()
    T2d[:,t+1] = T_step


x = np.arange(0,lx + dx,dx)
y = np.arange(0,ly + dy,dy)
x,y = np.meshgrid(x,y)
fps = 10  # frame per sec
frn = timesteps + 1
fig = plt.figure()
ax = fig.add_subplot(111,projection='3d')
plot = [ax.plot_surface(no_nodes_x,T2d[:,0],color='0.75',rstride=1,cstride=1)]
# ax.set_zlim(0,1.1)

def update_plot(frame_number,zarray,plot):
    plot[0].remove()
    plot[0] = ax.plot_surface(x,y,zarray[:,frame_number],cmap="magma")

ani = animation.FuncAnimation(fig,update_plot,frn,fargs=(T2d,plot),interval=1000 / fps)
plt.show()

我收到此错误：

Traceback (most recent call last):
  File "D:\FEM\venv\lib\site-packages\matplotlib\cbook\__init__.py",line 270,in process
    func(*args,**kwargs)
  File "D:\FEM\venv\lib\site-packages\matplotlib\animation.py",line 991,in _start
    self._init_draw()
  File "D:\FEM\venv\lib\site-packages\matplotlib\animation.py",line 1753,in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
  File "D:\FEM\venv\lib\site-packages\matplotlib\animation.py",line 1776,in _draw_frame
    self._drawn_artists = self._func(framedata,*self._args)
  File "D:\FEM\fem2d_v2.py",line 104,in update_plot
    plot[0] = ax.plot_surface(x,cmap="magma")
  File "D:\FEM\venv\lib\site-packages\matplotlib\_api\deprecation.py",line 431,in wrapper
    return func(*inner_args,**inner_kwargs)
  File "D:\FEM\venv\lib\site-packages\mpl_toolkits\mplot3d\axes3d.py",line 1665,in plot_surface
    X,Y,Z = np.broadcast_arrays(X,Z)
  File "<__array_function__ internals>",line 5,in broadcast_arrays
  File "D:\FEM\venv\lib\site-packages\numpy\lib\stride_tricks.py",line 538,in broadcast_arrays
    shape = _broadcast_shape(*args)
  File "D:\FEM\venv\lib\site-packages\numpy\lib\stride_tricks.py",line 420,in _broadcast_shape
    b = np.broadcast(*args[:32])
ValueError: shape mismatch: objects cannot be broadcast to a single shape

Process finished with exit code 0

我知道形状是错误的，但真的不知道如何正确修复它。 这是基于 matlab 代码，x、y 和 T2d 的形状与那里相同，如果 lx 和 ly 相同，则它有效.我不知道我应该如何以正确的方式将数据保存到 T2d。

这里也是我参考的matlab代码：https://www.files.ethz.ch/structuralgeology/sms/NumModRocDef/ML_2D_diffusion_numint.txt

编辑 添加了完整代码。