title: Colored mesh plot of a function in Python
body: <p>I am trying to draw a colored mesh plot of a travelling wave like this (just ignore the grey background):</p>
  
  <p><img src="https://i.stack.imgur.com/e1vsl.png" alt="enter image description here"></p>
  
  <p>I am looking at the example shown in <a href="https://stackoverflow.com/questions/24909256/how-to-obtain-3d-colored-surface-via-python/24958192">How to obtain 3D colored surface via Python?</a>, but I'm not sure how to adapt it for my function <code>z</code>:</p>
  
  <pre><code>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()
  </code></pre>