Python中绘制(t,x,u=u(t,x))类型曲面的问题求助

问题背景 #

作为数学研究者，日常很少编程，想在Python中绘制$(t,x,u=u(t,x))$类型的3D曲面。t、x是用np.zeros()初始化的向量，u是同维度的矩阵，但使用plot_surface函数未能得到预期效果，代码如下：

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import math

eps=0.1
m=2000
n=100
dt=1.0/m
dx=1.0/(n*n)
time=np.zeros(m+1)
for i in range(m+1):
  time[i]=i*dt
space=np.zeros(2*n+1)
for j in range(2*n+1):
  space[j]=(j-n)*dx*n
sol=np.zeros((m+1,2*n+1))
# 这里建议添加初始条件，比如设置最后一行（t=1时）的解
# sol[m, :] = np.exp(-(space)**2/(0.1**2))  # 示例高斯初始条件
for i in range(m):
  index_i=m-1-i
  for j in range(1,2*n):
    sol[index_i, j] =sol[index_i+1, j]-0.5*dt*math.log(eps+abs(sol[index_i+1, j+1]+sol[index_i+1, j-1]-2*sol[index_i+1, j])/dx)
t_mesh, x_mesh = np.meshgrid(time, space)
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
surf = ax.plot_surface(t_mesh, x_mesh, sol, cmap=cm.coolwarm,
                           linewidth=0, antialiased=False)
plt.show()

核心问题与解决方法 #

你遇到的主要问题是网格矩阵与解矩阵的维度不匹配，以及初始条件缺失导致解全为0的问题，具体修正方式如下：

1. 维度匹配修正

np.meshgrid(time, space)默认生成的t_mesh和x_mesh形状是(2n+1, m+1)，而你的sol矩阵形状是(m+1, 2n+1)（行对应时间，列对应空间），二者维度不兼容，导致plot_surface无法正确映射。有两种修正方式：

• 方式一：转置解矩阵
  将plot_surface中的sol改为sol.T，让其形状与网格矩阵一致：

  surf = ax.plot_surface(t_mesh, x_mesh, sol.T, cmap=cm.coolwarm, linewidth=0, antialiased=False)
  

• 方式二：指定meshgrid的索引模式
  使用indexing='ij'参数生成与sol维度匹配的网格，这样无需转置：

  t_mesh, x_mesh = np.meshgrid(time, space, indexing='ij')
  surf = ax.plot_surface(t_mesh, x_mesh, sol, cmap=cm.coolwarm, linewidth=0, antialiased=False)
  

2. 添加初始条件

你的sol初始全为0，迭代计算时从最后一行（对应t=1的时刻）往前推，但最后一行没有设置初始值，导致整个解矩阵始终为0，绘制出来是平面。需要给sol[m, :]（最后一行）设置初始条件，比如：

# 示例：高斯分布初始条件
sol[m, :] = np.exp(-(space)**2/(0.1**2))

或者根据你的偏微分方程需求设置对应的初始边界值。

3. 关于plt.plot(time, space, sol)

plt.plot()是用于绘制2D曲线的函数，无法直接绘制3D曲面。要绘制3D曲面必须使用ax.plot_surface()（基于3D坐标轴），上面的修正就能让它正常生效。

内容的提问来源于stack exchange，提问作者Philo18