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.

• drum: stackoverflow.com/help/how-to-ask
• Nick ODell: I don't follow the mathematics of this question. You want to plot the function u(t, x), where t and x are both vectors, which means that the plot is at least 4D? (2 or more dimensions for first vector, and 2 or more for second vector.) How does the plot look?
• Philo18: @NickODell Indeed, time=(i/m: 0<=i<=m), space=(j/n: -n<=j<=n) and sol=(u_i,j: 0<=i<=m, -n<=j<=n). I wish to plot sol as a function of time/space
• Nick ODell: So say that time is X coordinate, space is Y coordinate, and sol is Z coordinate. sol is a 2D array, right? So it can't be used directly. How should that be turned into a 1D array?
• Philo18: @NickODell My function u is not explicite, while I have its values at the chosen grid

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:

surface picture

Code adapted from this documentation example.

• Philo18: It is still not working. Indeed, even my matrix sol is created as np.zeros(m+1,2*n+1). The system still treats it as an one-dimensional array/vector, instead of a matrix or 2d array
• Nick ODell: Can you clarify? On my system, sol is created as a 2D array, eg. sol.shape shows (2001, 201).
• Philo18: The same for me. I've uploaded the code which yields an error
• Nick ODell: You swapped the order of time and space. If you want to do that you have to swap the order of axes for sol too, using e.g. np.swapaxes().
• Philo18: I don't undestand... Do you mean ax.plot_surface(x_mesh, t_mesh, sol) instead of ax.plot_surface(t_mesh, x_mesh, sol)?