在python中使用matplotlib創建3D曲面圖

Question

我試圖繪制一個3D表面，但我遇到了一些麻煩，因為matplotlib的文檔似乎並不是非常徹底，並且在示例中缺乏。 無論如何，我編寫的程序是通過有限差分方法以數值方式求解熱方程。 這是我的代碼：

    ## This program is to implement a Finite Difference method approximation
## to solve the Heat Equation, u_t = k * u_xx,
## in 1D w/out sources & on a finite interval 0 < x < L. The PDE
## is subject to B.C: u(0,t) = u(L,t) = 0,
## and the I.C: u(x,0) = f(x).
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

# Parameters    
L = 1 # length of the rod
T = 10 # terminal time
N = 40 # spatial values
M = 1600 # time values/hops; (M ~ N^2)
s = 0.25 # s := k * ( (dt) / (dx)^2 )

# uniform mesh
x_init = 0
x_end = L
dx = float(x_end - x_init) / N

x = np.arange(x_init, x_end, dx)
x[0] = x_init

# time discretization
t_init = 0
t_end = T
dt = float(t_end - t_init) / M

t = np.arange(t_init, t_end, dt)
t[0] = t_init

# time-vector
for m in xrange(0, M):
    t[m] = m * dt

# spatial-vector
for j in xrange(0, N):
    x[j] = j * dx

# definition of the solution u(x,t) to u_t = k * u_xx
u = np.zeros((N, M+1)) # array to store values of the solution

# Finite Difference Scheme:

u[:,0] = x * (x - 1) #initial condition

for m in xrange(0, M):
    for j in xrange(1, N-1):
        if j == 1:
            u[j-1,m] = 0 # Boundary condition
        elif j == N-1:
            u[j+1,m] = 0 # Boundary Condition
        else:
            u[j,m+1] = u[j,m] + s * ( u[j+1,m] - 
            2 * u[j,m] + u[j-1,m] )

這就是我為編寫3D曲面圖而編寫的內容：

    # for 3D graph
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(x, t, u, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()

當我運行代碼繪制圖形時，我收到此錯誤：“ValueError：shape mismatch：兩個或多個數組在軸1上具有不兼容的尺寸。”

請，任何和所有的幫助是非常有用的。 我認為錯誤出現是因為我將u定義為Nx(M+1)矩陣，但是必須使原始程序運行。 我不確定如何糾正這一點，因此圖表正確繪制。 謝謝！

Answer 1

打印出變量x ， t和u的形狀是有幫助的：

x.shape == (40,)
t.shape == (1600,)
u.shape == (40, 1601)

所以這里有兩個問題。 第一個是x和t是1維的，即使它們需要是2維的。 第二個是u在第二維中有一個元素而不是t 。 您可以通過運行來修復它們

t, x = np.meshgrid(t, x)
u = u[:,:-1]

在創建3d圖之前。

Answer 2

使用此代碼（查看注釋）：

# plot 3d surface
# create a meshgrid of (x,t) points
# T and X are 2-d arrays
T, X = np.meshgrid(t,x)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# Use X and T arrays to plot u 
# shape of X, T and u must to be the same
# but shape of u is [40,1601] and I will skip last row while plotting
surf = ax.plot_surface(X, T, u[:,:1600], rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()

結果：

因為matplotlib的文檔看起來不是很透徹，並且在示例中缺乏

http://matplotlib.org/examples/mplot3d/index.html