title: Matplotlib - Assign specific colors to a surface
body: <p>This snippet of code is going to generate the following two pictures, which represents a complex function. Is it possible to apply the colors of the first image to the surface? If so, how?</p>
  <p><a href="https://i.stack.imgur.com/isZ13.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/isZ13.png" alt="domain coloring" /></a></p>
  <p><a href="https://i.stack.imgur.com/YLfeA.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/YLfeA.png" alt="3d magnitude" /></a></p>
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  import matplotlib.pyplot as plt
  from matplotlib.colors import hsv_to_rgb
  
  def saw_func(x, dx, a, b):
      x = x / dx - np.floor(x / dx)
      return a + (b - a) * x
  
  def domain_coloring(mag, arg, phaseres=20):
      arg[arg &lt; 0] += 2 * np.pi
      arg /= (2 * np.pi)
      blackp = saw_func(arg, 1 / phaseres, 0.75, 1)
      blackm = saw_func(np.log(mag), 2 * np.pi / phaseres, 0.75, 1)
      black = blackp * blackm
      H = arg
      S, V = np.ones_like(H), black
      return (hsv_to_rgb(np.dstack([H, S, V])) * 255).astype(np.uint8)
  
  x = y = np.linspace(-2, 2, 500)
  x, y = np.meshgrid(x, y)
  z = x + 1j * y
  f = (z - 1) / (z**2 + z + 1)
  mag, arg = np.absolute(f), np.angle(f)
  img = domain_coloring(mag, arg)
  
  fig1, ax1 = plt.subplots()
  ax1.imshow(
      img,
      extent = [np.amin(x), np.amax(x), np.amin(y), np.amax(y)], 
      interpolation = &quot;nearest&quot;,
      origin = &quot;lower&quot;,
  )
  plt.show()
  
  fig2 = plt.figure()
  ax2 = fig2.add_subplot(1, 1, 1, projection=&quot;3d&quot;)
  ax2.plot_surface(x, y, mag)
  ax2.set_zlim([0, 10])
  plt.show()
  </code></pre>