Plotting 3D sets in Jupyter

%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
S = 50

def f(x,y):
    return np.exp(- x**2 - y**2)

def plot(x, y, f):
    xx, yy = np.meshgrid(X,Y)
    Z = map(lambda x: map(lambda y: f(x,y), Y), X)
    p = ax.plot_surface(xx, yy, Z, rstride=1, cstride=1, cmap=cm.jet,
                        linewidth=0.01, antialiased=True, shade=True)

X = np.linspace(-2,2,S)
Y = np.linspace(-2,2,S)
plot(X, Y, f)