python - 使用 matplotlib 制作 3d 冲浪动画

Question

我正在用 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]), np.zeros([no_nodes, 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/6*kappa/dx*dy-1/3*kappa*dx/dy, 1/3*kappa/dx*dy+1/3*kappa*dx/dy, -1/3*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/6*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/18*dy*dx, 1/36*dy*dx],
        [1/36*dy*dx, 1/18*dy*dx, 1/9*dy*dx, 1/18*dy*dx],
        [1/18*dy*dx, 1/36*dy*dx, 1/18*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], [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, no_nodes_y, 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, no_nodes_y, 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, no_nodes_y, 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, y, zarray[:,:,frame_number], 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, Y, 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

编辑添加了整个代码。

标签： pythonmatplotlibanimationplot3d