Plot 3d in Python

Question

I wish to know how to plot a surface of type $(t,x,u=u(t,x))$ in Python. More precisely, $t,x$ are vectors and $u$ is a matrix that are initialized as np.zeros(), while the function plot does not draw the surface as I desire. Could someone help? The code is as follow:

 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))
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)

Which format should be used such that plt.plot(time,space,sol) works?

PS : I do research in maths and I code rarely. Sorry if my statement is not clear.

Answer 1

You can plot that function like so:

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

# ... your original code here ...

def plot_surface_from_arrays(X, Y, Z, rotate=0):
    assert Y.shape + X.shape == Z.shape, "X and Y shapes don't match Z"
    X_mesh, Y_mesh = np.meshgrid(X, Y)
    fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
    ax.view_init(elev=30, azim=-60 + rotate)
    surf = ax.plot_surface(X_mesh, Y_mesh, Z, cmap=cm.coolwarm,
                           linewidth=0, antialiased=False)
plot_surface_from_arrays(space, time, sol, rotate=90)

Result:

Code adapted from this documentation example .