import numpy as np
import sys

from matplotlib import pyplot as plt
from matplotlib import cm
from pylab import meshgrid
from mpl_toolkits.mplot3d.axes3d import Axes3D

if len(sys.argv) != 1:
  print('Usage: ', sys.argv[0])
  sys.exit(1)

def peaks(x, y):
  return 3 * (1 - x) * (1 - x) * np.exp(- x * x - (y + 1) * (y + 1)) \
    - 10 * (x / 5 - x**3 - y**5) * np.exp(- x*x - y*y) \
    - 1/3 * np.exp(-(x+1)*(x+1) - y*y) 

x = np.arange(-3.0, 3.0, 0.05)
y = np.arange(-3.0, 3.0, 0.05)
X, Y = meshgrid(x, y)
Z = peaks(X, Y)

fig = plt.figure(figsize = (18, 8))
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1, 
                       cmap = cm.seismic, linewidth = 0, antialiased = False)
ax.view_init(21, -155)

fig.colorbar(surf, shrink = 0.5, aspect = 5)
#fig.savefig('peaks.png')

plt.show()