Matplotlib cmap没有按预期工作

时间:2017-04-29 04:03:31

标签: python python-3.x numpy matplotlib

我有以下Python 3代码,它会随着时间的推移生成波函数并将结果绘制成3D。请注意,schroedinger1D(...)函数返回两个numpy数组,每个数组都是shape(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轴上变化,但在这里我不知道测量的是什么:Plot of wave function over time

1 个答案:

答案 0 :(得分:1)

在减少网格点数量时,可能更容易理解发生了什么

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

并且还使用了1的rstride和cstride。

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

enter image description here

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

唯一可行的解​​决方案是使网格比振荡频率更密集。即尝试可视化更平滑变化的波浪,或仅显示波函数的一部分但具有更密集的网格。