matplotlib surface plot extends past axis limits

Question

How do I make a nice paraboloid in Matplotlib that looks like

All I can get is this,

where the top is not "cut off". I've tried just dropping all values of the Z array outside of the radius of the parabola at the top, but that gives very jagged edges. Can someone help me?

Here is my code:

from matplotlib import *
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from pylab import *
import math
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

X = np.arange(-5, 5, 0.1)
Y = np.arange(-5, 5, 0.1)
X, Y = np.meshgrid(X, Y)
Z = (X**2 + Y**2)

ax.set_zlim(-10, 20)

ax.plot_surface(X, Y, Z,  alpha=0.9, rstride=4, cstride=4, linewidth=0.5, cmap=cm.summer)

plt.show()

Answer 1

For future reference, I had a thought to parametrize the surface in cylindrical coordinates, and it looks exactly how I want it:

from matplotlib import *
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from pylab import *
import math
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

r = T = np.arange(0, 2*pi, 0.01)
r, T = np.meshgrid(r, T)
#Parametrise it
X = r*np.cos(T)
Y = r*np.sin(T)
Z = r**2

ax.plot_surface(X, Y, Z,  alpha=0.9, rstride=10, cstride=10, linewidth=0.5, cmap=cm.summer)

plt.show()

I guess it makes sense: when working with a cylindrical object, use cylindrical coordinates!