问题
I am trying to draw a colored mesh plot of a travelling wave like this (just ignore the grey background):
I am looking at the example shown in How to obtain 3D colored surface via Python?, but I'm not sure how to adapt it for my function z
:
from mpl_toolkits.mplot3d import axes3d, art3d
import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X = T = np.arange(0,10,0.1)
z = np.array([np.cos(2*np.pi*(x/3.0-t/6.0)) for x,t in zip(np.ravel(X), np.ravel(T))])
Z = z.reshape(X.shape)
#X, Y, Z = axes3d.get_test_data(0.05)
wire = ax.plot_wireframe(X, T, Z, rstride=10, cstride=10)
# Retrive data from internal storage of plot_wireframe, then delete it
nx, nt, _ = np.shape(wire._segments3d)
wire_x = np.array(wire._segments3d)[:, :, 0].ravel()
wire_t = np.array(wire._segments3d)[:, :, 1].ravel()
wire_z = np.array(wire._segments3d)[:, :, 2].ravel()
wire.remove()
# create data for a LineCollection
wire_x1 = np.vstack([wire_x, np.roll(wire_x, 1)])
wire_t1 = np.vstack([wire_t, np.roll(wire_t, 1)])
wire_z1 = np.vstack([wire_z, np.roll(wire_z, 1)])
to_delete = np.arange(0, nx*ny, ny)
wire_x1 = np.delete(wire_x1, to_delete, axis=1)
wire_t1 = np.delete(wire_t1, to_delete, axis=1)
wire_z1 = np.delete(wire_z1, to_delete, axis=1)
scalars = np.delete(wire_z, to_delete)
segs = [list(zip(xl, tl, zl)) for xl, tl, zl in \
zip(wire_x1.T, wire_t1.T, wire_z1.T)]
# Plots the wireframe by a a line3DCollection
my_wire = art3d.Line3DCollection(segs, cmap="hsv")
my_wire.set_array(scalars)
ax.add_collection(my_wire)
ax.set_title(r'Travelling wave , $z=f(x,t)=\cos(2\pi(x/3-t/6))$')
ax.set_xlabel(r'$x$', fontsize=16)
ax.set_ylabel(r'$t_0$', fontsize=16)
ax.set_zlabel(r'$z$', fontsize=16)
plt.show()
回答1:
np.meshgrid()
is your friend:
from mpl_toolkits.mplot3d import axes3d, art3d
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
x, t = np.arange(0,10,0.1), np.arange(0,10,0.1)
X, T = np.meshgrid(x, t) # generate two 100x100 arrays
Z = np.cos(2*np.pi*(X/3.0 - T/6.0))
wire = ax.plot_wireframe(X, T, Z, rstride=10, cstride=10)
# Retrive data from internal storage of plot_wireframe, then delete it
nx, ny, _ = np.shape(wire._segments3d)
wire_x = np.array(wire._segments3d)[:, :, 0].ravel()
wire_y = np.array(wire._segments3d)[:, :, 1].ravel()
wire_z = np.array(wire._segments3d)[:, :, 2].ravel()
wire.remove()
# create data for a LineCollection
wire_x1 = np.vstack([wire_x, np.roll(wire_x, 1)])
wire_y1 = np.vstack([wire_y, np.roll(wire_y, 1)])
wire_z1 = np.vstack([wire_z, np.roll(wire_z, 1)])
to_delete = np.arange(0, nx*ny, ny)
wire_x1 = np.delete(wire_x1, to_delete, axis=1)
wire_y1 = np.delete(wire_y1, to_delete, axis=1)
wire_z1 = np.delete(wire_z1, to_delete, axis=1)
scalars = np.delete(wire_z, to_delete)
segs = [list(zip(xl, yl, zl)) for xl, yl, zl in \
zip(wire_x1.T, wire_y1.T, wire_z1.T)]
# Plots the wireframe by a a line3DCollection
my_wire = art3d.Line3DCollection(segs, cmap="hsv")
my_wire.set_array(scalars)
ax.add_collection(my_wire)
plt.colorbar(my_wire)
#cl1 = fig.colorbar(pl1)
#cl1.set_label("$z$-axis")
fig.suptitle(r'Travelling wave $z=f(x,t)=\cos(2\pi(x/3-t/6))$')
ax.set_xlabel(r'$x$', fontsize=16)
ax.set_ylabel(r'$t_0$', fontsize=16)
ax.set_zlabel(r'$z$', fontsize=16)
fig.canvas.draw()
plt.show()
The result is
.来源:https://stackoverflow.com/questions/28751535/colored-mesh-plot-of-a-function-in-python