Graphs in Python by MatPlotLib and NumPy

contour.py

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np

X = np.arange(0, 1, .1)
Y = np.arange(0, 1, .1)
X, Y = np.meshgrid(X, Y)
Z = X**2 + Y**2

levels = MaxNLocator(nbins=15).tick_values(Z.min(), Z.max())
cmap = plt.get_cmap('PiYG')

plt.contourf(X, Y, Z, levels=levels, cmap=cmap)
plt.colorbar()
plt.title('U(X, Y)')
plt.xlabel('X')
plt.ylabel('Y')
plt.show()

surface.py

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X = np.arange(0, 1, .1)
Y = np.arange(0, 1, .1)
X, Y = np.meshgrid(X, Y)
Z = X**2 + Y**2

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
                        linewidth=0, antialiased=False)

ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)
plt.title('U(X, Y)')
plt.xlabel('X')
plt.ylabel('Y')
plt.show()