Change to:
ax.plot_wireframe(T, z, abs(U), cstride=1000)
and call:
drawPropagation(1.0, 1.0, numpy.linspace(-2, 2, 10))
If you need the curve been filled with white color:
import numpyfrom mpl_toolkits.mplot3d import Axes3Dfrom matplotlib import pyplotfrom matplotlib.collections import PolyCollectiondef drawPropagation(beta2, C, z): """ beta2 in ps / km C is chirp z is an array of z positions """ T = numpy.linspace(-10, 10, 100) sx = T.size sy = z.size T = numpy.tile(T, (sy, 1)) z = numpy.tile(z, (sx, 1)).T U = 1 / numpy.sqrt(1 - 1j*beta2*z * (1 + 1j * C)) * numpy.exp(- 0.5 * (1 + 1j * C) * T * T / (1 - 1j*beta2*z*(1 + 1j*C))) fig = pyplot.figure() ax = fig.add_subplot(1,1,1, projection='3d') U = numpy.abs(U) verts = [] for i in xrange(T.shape[0]): verts.append(zip(T[i, :], U[i, :])) poly = PolyCollection(verts, facecolors=(1,1,1,1), edgecolors=(0,0,1,1)) ax.add_collection3d(poly, zs=z[:, 0], zdir='y') ax.set_xlim3d(numpy.min(T), numpy.max(T)) ax.set_ylim3d(numpy.min(z), numpy.max(z)) ax.set_zlim3d(numpy.min(U), numpy.max(U))drawPropagation(1.0, 1.0, numpy.linspace(-2, 2, 10))pyplot.show()
