Drawing Surfaces in Python

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

ax = plt.axes(projection="3d")

x = np.arange(0,50,0.5)
y = np.arange(0,50,0.5)
X,Y = np.meshgrid(x,y)
Z = X*X*Y

max_z = np.max(Z)
min_z = np.min(Z)

ax.plot_surface(X,Y,Z)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_title('Function z = x^2 * y')
ax.legend()
plt.show()