python - Matplotlib cmap 未按预期工作

Question

我有以下 Python 3 代码，它随时间生成波函数并以 3D 形式绘制结果。请注意，schroedinger1D(...) 函数返回两个 numpy 数组，每个数组的形状为 (36,1000)。

import numpy as np
from scipy.integrate import fixed_quad
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# Initial conditions, outside for plotting.
f_re = lambda x: np.exp(-(x-xc)**2.0/s)*np.cos(2.0*np.pi*(x-xc)/wl)
f_im = lambda x: np.exp(-(x-xc)**2.0/s)*np.sin(2.0*np.pi*(x-xc)/wl)

def schroedinger1D(xl, xr, yb, yt, M, N, xc, wl, s):
    """
        Schrödinger Equation Simulation (no potential)
    """
    f = lambda x: f_re(x)**2 + f_im(x)**2
    area = fixed_quad(f, xl, xr, n=5)[0]
    f_real = lambda x: f_re(x)/area
    f_imag = lambda x: f_im(x)/area
    # Boundary conditions for all t
    l = lambda t: 0*t
    r = lambda t: 0*t
    # "Diffusion coefficient"
    D = 1
    # Step sizes and sigma constant
    h, k = (xr-xl)/M, (yt-yb)/N
    m, n = M-1, N
    sigma = D*k/(h**2)
    print("Sigma=%f" % sigma)
    print("k=%f" % k)
    print("h=%f" % h)
    # Finite differences matrix
    A_real = np.diag(2*sigma*np.ones(m)) + np.diag(-sigma*np.ones(m-1),1) + np.diag(-sigma*np.ones(m-1),-1)
    A_imag = -A_real
    # Left boundary condition u(xl,t) from time yb
    lside = l(yb+np.arange(0,n)*k)
    # Right boundary condition u(xr,t) from time tb
    rside = r(yb+np.arange(0,n)*k)
    # Initial conditions
    W_real = np.zeros((m, n))
    W_imag = np.zeros((m, n))
    W_real[:,0] = f_real(xl + np.arange(0,m)*h)
    W_imag[:,0] = f_imag(xl + np.arange(0,m)*h)
    # Solving for imaginary and real part
    for j in range(0,n-1):
        b_cond = np.concatenate(([lside[j]], np.zeros(m-2),[rside[j]]))
        W_real[:,j+1] =  W_real[:,j] + A_real.dot(W_imag[:,j]) - sigma*b_cond
        W_imag[:,j+1] =  W_imag[:,j] + A_imag.dot(W_real[:,j]) + sigma*b_cond
    return np.vstack([lside, W_real, rside]), np.vstack([lside, W_imag, rside])


xl, xr, yb, yt, M, N, xc, wl, s = (-9, 5, 0, 4, 35, 1000, -5, 4.0, 3.0)
W_real, W_imag = schroedinger1D(xl, xr, yb, yt, M, N, xc, wl, s)
[X, T] = np.meshgrid(np.linspace(xl, xr, M+1), np.linspace(yb, yt,N))
# Plot results
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel("$x$", fontsize=20)
ax.set_ylabel("$t$", fontsize=20)
ax.set_zlabel("$\Psi(x,t)$", fontsize=20)
print(X.shape)
print(T.shape)
print(W_real.T.shape)
surface = ax.plot_surface(X, T, W_real.T, cmap=cm.jet, linewidth=0, 
                          antialiased=True, rstride=10, cstride=10)
fig.colorbar(surface)
plt.tight_layout()
plt.show()

输出具有正确的 X、T 和 W_real.T (1000,36) 形状，但据我所知，表面指定的颜色是错误的。我原以为颜色在 Z 轴上会有所不同，但在这里我无法分辨正在测量的是什么:

Answer 1

减少网格点的数量可能更容易理解发生的情况

xl, xr, yb, yt, M, N, xc, wl, s = (-9, 5, 0, 4, 10, 10, -5, 4.0, 3.0)

并且还使用 rstride 和 cstride 为 1。

在这种情况下，情节可能看起来像这样

现在很容易发现问题:波函数中的振荡频率大于网格的分辨率。这意味着曲面图的单个面片可能从非常低的值开始，然后上升到非常高的值。在这种情况下，它的颜色可以是任何颜色，因为它是由 block 的单个边缘的值决定的。 (如果色 block 从高值开始并下降到低值，则它比从低值开始并上升到高值时更红。)

唯一可能的解决方案是使网格的密度大于振荡频率。 IE。尝试可视化更平滑变化的波函数，或者仅可视化波函数的一部分但具有更密集的网格。