Skip to main content

Visualizing differential geometry in Jupyter notebooks

(you can download this post as a notebook here)

Visualizing differential geometry in Jupyter notebooks

In [19]:
 
Out[19]:
In [30]:
 
Out[30]:

I taught a senior seminar on differential geometry last year. I'll be honest: it was a selfish course. Markus Deserno writes all of these cool papers about membranes, and I didn't know enough math to follow them. So, the goal of our course was basically to learn enough differential geometry to read several of his papers. It was fantastic. We used several of his papers, as well as Kreyszig's textbook, as our core materials.

Here's the thing: a lot of this was pretty foreign to my physics students. In particular, the discussion of surfaces and mappings was new. So, we wrote some tools in Jupyter Notebooks to help us visualize and solve problems. I particularly like the stuff we wrote to visualize a mapping, and I don't know of a comparable resource elsewhere.

As a teaser, the above movies show you how to interpolate between a surface and a curve in the first case, and how the Monge gauge works for a membrane in the second case.

Let's jump in.

Surfaces

We get our introduction to surfaces in Kreyszig Ch. 3. Specifically, $\S3.24$ In most of our classes, we've defined surfaces like

$$z = f(x,y) $$

and we've been good at visualizing them. Here, we run into the definition of a surface in terms of a mapping that sends

$$ (u_1,u_2) \to (x(u_1,u_2), y(u_1,u_2), z(u_1,u_2)) $$

and Kreyszig warns us that something like

$$ (u_1+u_2, (u_1+u_2)^2, (u_1+u_2)^3) $$

is problematic because, since it can be parameterized by a single variable, it's actually a curve rather than a surface. We also run into the Jacobian and some linear algebra. Our goal right now is to work with concrete curves and surfaces in $\mathbb{R}^3$ So, let's figure out how to visualize these sorts of mappings. First, some gunk to set up our python environment. For the experts, note that we could use %matplotlib notebook to make things prettier, but it seems really slow.

In [1]:
import numpy as np, scipy as sp, pandas as pd, seaborn as sns, matplotlib.pyplot as plt, matplotlib as mpl
import itertools
from mpl_toolkits.mplot3d import Axes3D, proj3d
import matplotlib.animation as animation
from matplotlib.patches import FancyArrowPatch
from IPython.display import HTML
from ipywidgets import interact, interactive, fixed

In class, we used interact to visualize things. However, I'm making this available on the web using the fact that Jake Vanderplas' awesome JSAnimation package is built in to matplotlib 2.1. I like it way better than just an html5 video, because I like to be able to step through the frames. I still strongly recommend interact for playing around on your own server.

In [2]:
# You can change this to %matplotlib notebook if you want. It looks better, but is much slower for me.
#%matplotlib inline
In [3]:
# mpl.rcParams['animation.html'] = 'jshtml'
# Doesn't seem to work for me, so I make the HTML explicitly below

Now, let's just try one of these out. One of the examples we want to look at in the book is

$$ (u_1+u_2, \sin(u_1), \cos(u_1+u_2)) $$

so let's make a grid of points where $u_1$ and $u_2$ vary from -2 to 2, then plot the corresponding surface. We'll plot $(u_1, u_2, 0)$ in grey and the surface in blue.

In [4]:
u1 = np.linspace(-2,2,100)
u2 = np.linspace(-2,2,100)
U1, U2 = np.meshgrid(u1, u2)
X = U1 + U2
Y = np.sin(U1)
Z = np.cos(U1+U2)

ax = plt.gca(projection='3d')
surf2 = ax.plot_surface(U1, U2, np.zeros_like(U1), color='grey', alpha=0.5)
surf = ax.scatter(X, Y, Z)
plt.show()

You can play around with the above, using more or fewer points, etc. You can clearly see that the surface extends vertically and horizontally beyond the patch from $u_1$ and $u_2$, although it's kind of cheating to plot them both at the same time. So let's wrap that up in a function and try it with different maps.

To make it cuter, we'll draw red arrows from some specific points to see which points get mapped where. And we'll plot a surface not a collection of points, by default, because it's way faster. Most impoartantly, we'll wrap all of the plotting stuff up in a function of its own so that we just need to define a map, and then ask to plot it.

In [5]:
class Arrow3D(FancyArrowPatch):
    #http://stackoverflow.com/questions/22867620/putting-arrowheads-on-vectors-in-matplotlibs-3d-plot
    def __init__(self, xs, ys, zs, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        FancyArrowPatch.draw(self, renderer)

We played around with things quite a bit in class, so the function below has options for changing from surface to points and for changing urange (the range included in the $u_1,u_2$ patch)

In [6]:
def plot_surface(f,plotmode='surface',urange=2*np.pi):
    u1 = np.linspace(-urange,urange,100)
    u2 = np.linspace(-urange,urange,100)
    ax = plt.gca(projection='3d')
    U1, U2 = np.meshgrid(u1, u2)
    X,Y,Z = f(U1,U2)
    surf2 = ax.plot_surface(U1, U2, np.zeros_like(U1), color='grey', alpha=0.5)
    if plotmode == 'surface':
        surf = ax.plot_surface(X,Y,Z)
    elif plotmode == 'points':
        surf = ax.scatter(X,Y,Z)
    else:
        raise Exception('Unknown plot mode {p}'.format(p=plotmode))
    for (w1,w2) in itertools.permutations([-urange,-urange/2,0,urange/2,urange],2):
        x,y,z = f(w1,w2)
        a = Arrow3D([w1,x],[w2,y],[0,z], mutation_scale=20, 
                    lw=1, arrowstyle="-|>", color="r",
                   alpha=0.7,zorder=-1)
        ax.add_artist(a)
    plt.show()

And now let's plot that same surface $(u_1,u_2) \to u_1 + u_2, \sin(u_1), \cos(u_1+u_2)$

In [7]:
def f(U1,U2):
    return U1+U2, np.sin(U1), np.cos(U1+U2)
plot_surface(f)

Changing the $(u1,u2)$ patch

We spent a bit of time thinking about what exactly $u1$ and $u2$ were doing. If you make their range bigger or smaller, you change the size of the surface. That's built into the function defined above. By default, they each go from -urange to urange. As you can see in the figure above, urange defaults to $2\pi$. We can look at a smaller patch to see what's going on:

In [8]:
plot_surface(f,urange=2)

Playing around

At this point, the class just played around with the results. Everybody defined their own functions, and looked at the resulting maps. You can do that too if you run the notebook interactively! Here's one that the students particularly liked, despite the attempted bad math.

In [10]:
def f(U1,U2):
    return U1,np.log(U2),3*U1+5*U2
plot_surface(f)
/Users/mglerner/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: invalid value encountered in log
  
/Users/mglerner/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log
  
/Users/mglerner/anaconda3/lib/python3.6/site-packages/mpl_toolkits/mplot3d/proj3d.py:141: RuntimeWarning: invalid value encountered in true_divide
  txs, tys, tzs = vecw[0]/w, vecw[1]/w, vecw[2]/w

The curve

OK. A large portion of why we did this was because it was really not obvious to the class why $(u_1+u_2, (u_1+u_2)^2, (u_1+u_2)^3)$ mapped to a curve. We could all follow the mathematical arguments, but it didn't really stick without a good visualization. So, let's look:

so, what does the curve look like?

In [11]:
def f(U1,U2):
    return U1+U2, (U1+U2)**2, (U1+U2)**3
plot_surface(f, plotmode='points')

That's a great start. We can see where representative points match, and it's definitely useful to visualize things and "prove" to ourselves that you do indeed get a curve. It's much more convincing when we make an interactive version (below). Before we do that, let's also quickly check out a cylinder.

In [12]:
def f(U1,U2):
    return 2*np.cos(U1),2*np.sin(U1),U2
plot_surface(f, urange=np.pi)

In class, we played around with urange and the scale factors to make sure we understood what was going on. Speaking of which, Markus Deserno was kind enough to look at a preliminary version of this post, and he noticed that I was originally using urange=2*np.pi which got a "double covering" of the cylinder (i.e. the $u_1,u_2$ patch mapped two points to every point on the cylinder. Check it out:

In [13]:
plot_surface(f, urange=np.pi/2)
plot_surface(f, urange=np.pi)
plot_surface(f, urange=2*np.pi)

An interactive one

That was OK, and it convinced us that the maps really did go from, e.g., a plane to a cylinder. But (1) it still wasn't quite so obvious exactly how/why it worked (2) it's obvious that playing around with the input parameters can give a lot of insight.

Then I had the idea to make it interactive: you get a slider so you can watch the green surface interpolate linearly between the grey patch $(u_1,u_2)$ and the surface defined by $(x,y,z)$. This is worth running interactively, and playing with, as it really cleared things up. If you run it locally and interactively, you get sliders that let you rotate the surface around in 3D, so you can look at it from all angles. If you're runnign it locally, you wnat the interp step to be more like 0.01 and the degree step to be more like 1.

For notebooks to be viewable online, I'm JSAnimations with a coarser interp step and I'm skipping the controls for elevation and azimuth altogether.

If you just want the results, skip the code here and look for the next pretty picture!

In [14]:
def plot_surface_interpolate(f=f,plotmode='surface',urange=2*np.pi,interp=0.1,
                             elev=None,azim=None):
    u1 = np.linspace(-urange,urange,100)
    u2 = np.linspace(-urange,urange,100)
    #fig = plt.figure()
    ax = plt.gca(projection='3d')
    U1, U2 = np.meshgrid(u1, u2)
    X,Y,Z = f(U1,U2)
    U3 = np.zeros_like(Z)
    A = (1-interp)*U1 + interp*X
    B = (1-interp)*U2 + interp*Y
    C = (1-interp)*U3 + interp*Z
    surf2 = ax.plot_surface(U1, U2, np.zeros_like(U1), color='grey', alpha=0.5)
    if plotmode == 'surface':
        surf = ax.plot_surface(X,Y,Z,alpha=0.4,color='blue')
        surf3 = ax.plot_surface(A,B,C,color='green',zorder=-2,alpha=0.9)
    elif plotmode == 'points':
        surf = ax.scatter(X,Y,Z,alpha=0.4,color='blue')
        surf3 = ax.scatter(A,B,C,color='green',zorder=-2,alpha=0.9)
    else:
        raise Exception('Unknown plot mode {p}'.format(p=plotmode))
    
    for (w1,w2) in itertools.permutations([-urange,-urange/2,0,urange/2,urange],2):
        x,y,z = f(w1,w2)
        a = Arrow3D([w1,x],[w2,y],[0,z], mutation_scale=20, 
                    lw=1, arrowstyle="-|>", color="r",
                   alpha=0.7,zorder=-1)
        ax.add_artist(a)
    plt.title(f'interp = {interp}')        
    ax.view_init(elev=elev,azim=azim)
    plt.show()
def do_interact(f,interp_step=0.01,degree_step=1,urange=np.pi,azim=30,elev=30,):   
    def plot_this(f,interp=0.1,azim=azim,elev=elev,plotmode='surface',urange=np.pi):
        plot_surface_interpolate(f,interp=interp,azim=azim,elev=elev,plotmode=plotmode,urange=urange)
    interact(plot_this,interp=(0,1,interp_step),
             azim=(0,360,degree_step),
             elev=(0,360,degree_step),
             plotmode=['points','surface'],
             urange=fixed(urange),
             f=fixed(f));

def update_interp(i,f,surfs,ax,
                  U1,U2,U3,
                  X,Y,Z,
                  interpsteps,
                  elev,azim,urange):
    surf,surf2,surf3 = surfs
    ax.clear()
    interp = i/interpsteps
    A = (1-interp)*U1 + interp*X
    B = (1-interp)*U2 + interp*Y
    C = (1-interp)*U3 + interp*Z

    surf2 = ax.plot_surface(U1, U2, U3, color='grey', alpha=0.4)
    surf = ax.plot_surface(X,Y,Z,alpha=0.4,color='blue')
    surf3 = ax.plot_surface(A,B,C,color='green',zorder=-2,alpha=0.9)

    for (w1,w2) in itertools.permutations([-urange,-urange/2,0,urange/2,urange],2):
        x,y,z = f(w1,w2)
        a = Arrow3D([w1,x],[w2,y],[0,z], mutation_scale=20, 
                    lw=1, arrowstyle="-|>", color="r",
                   alpha=0.7,zorder=-1)
        ax.add_artist(a)

    plt.title(f'interp = {interp}')

    ax.view_init(elev=elev,azim=azim)
    return (surf,surf2,surf3)

def get_ani(f,urange=2*np.pi,interpsteps=10,elev=30,azim=30):
    u1 = np.linspace(-urange,urange,100)
    u2 = np.linspace(-urange,urange,100)
    U1, U2 = np.meshgrid(u1, u2)
    U3 = np.zeros_like(U1)

    X,Y,Z = f(U1,U2)

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

    surf2 = ax.plot_surface(U1,U2,U3, color='grey', alpha=0.5)
    surf = ax.plot_surface(U1,U2,U3,alpha=0.9,color='blue')
    surf3 = ax.plot_surface(U1,U2,U3,color='green',zorder=-2,alpha=0.9)
    surfs = (surf,surf2,surf3)

    ani = animation.FuncAnimation(fig, update_interp, interpsteps, fargs=(f,surfs,ax,U1,U2,U3,X,Y,Z,interpsteps,
                                                                         elev,azim,urange),
                                 interval=250, blit=False)
    return ani    

We can make this interactive in two ways. For live notebooks, we use interact and for the web we use JSAnimation. In both cases, once the boilerplate is done, we just need to define a function f and call either do_interact or get_ani, as seen below. To be clear, we expect a function f to be defined before each new call. f maps u1 and u2 into x, y, and z. See below for examples.

Now play with the interp slider to see how you'd interpolate linearly from $(u_1, u_2)$ to $(x,y,z)$

We just looked at mapping the plane to a cylinder, so we'll start with it. Good pedagogical tip: make the students write down (with pictures) what they think will happen first.

In [15]:
def f(X,Y):
    return 2*np.cos(X),2*np.sin(X),Y
do_interact(f)

Or, if you're looking at the blog post and want an animation where the slider controls the interpolation ...

In [16]:
def f(X,Y):
    return 2*np.cos(X),2*np.sin(X),Y
HTML(get_ani(f,interpsteps=10,urange=np.pi).to_jshtml())
Out[16]:


Once Loop Reflect