In other posts, I've touched on the idea of sampling from probability measures defined over spaces of functions. We might be inferring the density of material within the earth from seismic surveys, the elevation of the earth's surface from lidar measurements, or the conductivity of an aquifer from well head measurements. For the sake of definiteness, we'll assume throughout that the function we're looking for is an element of the space $X = L^2(\Omega)$ for some domain $\Omega$. We could consider Sobolev or other spaces but the idea is the same. I'll also tacitly assume that $\Omega$ is regular enough for PDE theory to apply. That usually means that the boundary of $\Omega$ has to be piecewise smooth and have no cusps.

You can think of a random variable with values in $X$, i.e. a function selected at random according to some statistical law, as a stochastic process. I'll use the two terms interchangeably. In the stochastic process literature, it's common to focus on processes that have certain useful statistical properties. For example, we might think of $\Omega$ as some time interval $[0, T]$ and a random function $z(t)$ as a function of time $t$. We'd then like to have some notion of causality. The values of $z(t)$ should depend on the values of $z(s)$ for $s$ coming before $t$ but not after. The Markov property gives us this idea of causality. If $\Omega$ is instead a spatial domain, we might instead be interested in random fields where we specify analytically the covariance between $z(x)$ and $z(y)$ as a function of $x - y$. For example, it's common to assume that the conductivity of an aquifer is a random field whose logarithm has the spatial correlation structure $$\mathbb E[(z(x) - m)(z(y) - m)] = \text{const}\times\exp(-|x - y|^2/\lambda^2).$$ This is all the province of geostatistics. But in the following I'll be focused more on general principles without making any assumptions about the dependence structure of the process.

The normal distribution has a kind of universal character in finite-dimensional statistical inference. We want to generalize the idea of a normal distribution to function spaces -- a Gaussian process. We'll have to think a bit about measure theory. This will lead us by turns to ask about the relation between two distinct normal random variables with values in a function space. This relationship will be summed up in the Feldman-Hájek theorem. It's a doozy. I'll do some numerical experiments to illustrate the theorem.

Normal random variables¶

A normal random variable $Z$ has a density $$\rho(z) = \frac{\exp\left(-\frac{1}{2}(z - m)^*C^{-1}(z - m)\right)}{\sqrt{(2\pi)^n|\det C|}}$$ where $m$ is the mean, $C$ the covariance matrix. To talk about normal random variables on function spaces, we'll have to think about things that we take for granted. What does it mean to say that a random variable $Z$ has a density at all? In finite dimensions, this means that we can evaluate expectations of any function of $Z$ by evaluating an integral: $$\mathbb E[f(Z)] = \int f(z)\rho(z)\mathrm dz.$$ Fine, but what is $\mathrm dz$? It's a symbolic shorthand for integration with respect to Lebesgue measure. I'll write Lebesgue measure as $\lambda$ for the time being. Let $\mu$ be the distribution for $Z$, i.e. $$\mu(Z) = \mathbb P[Z \in A].$$ When we say that $Z$ has a density, we really mean that, for any set $A$ such that $\lambda(A) = 0$, we must also have that $\mu(A) = 0$ as well. The fancy term for this is that $\mu$ is absolutely continuous with respect to $\lambda$, usually written as $\mu \ll \lambda$. The Radon-Nikodym theorem then tells us that there is a unique positive function $\rho$ such that, for all $f$, $$\int f\,\mathrm d\mu = \int f\,\rho\,\mathrm d\lambda.$$ At a formal level we might say that $\mathrm d\mu = \rho\,\mathrm d\lambda$, hence the notation $$\rho = \frac{\mathrm d\mu}{\mathrm d\lambda}.$$ The usual Lebesgue measure provides a kind of background, reference measure to which other probability measures can be absolutely continuous. Lebesgue measure is not itself a probability measure because the volume of the entire space is infinite. Discrete random variables don't have densities. The Lebesgue measure of a single point is zero. Discrete random variables instead have densities with respect to a different measure, for example the counting measure on $\mathbb Z$. The counting measure, which is again not itself a probability measure, fulfills the role of the reference measure to which discrete random variables can be absolutely continuous.

In infinite dimensions, there is no spoon Lebesgue measure. To be more specific, there's no measure on an infinite-dimensional vector space that is non-trivial and translation-invariant. There is no background reference measure "$\mathrm dz$" to integrate against. We can't write down a normal density -- a density is defined w.r.t. a background measure and we don't have one. There are still well-defined normal random variables and probability measures.

Not having a measure $\mathrm dz$ or a density $\rho(z)$ to work with doesn't stop physicists from writing functional integrals. I won't use this notation going forward, but I do find it appealing.

To define a normal random variable on a function space, we can instead work with arbitrary finite-dimensional projections. Let $\{g_1, \ldots, g_N\}$ be a finite collection of elements of the dual space $X^*$. We can then define a linear operator $G : X \to \mathbb R^N$ by $$G = \sum_k e_k\otimes g_k$$ where $e_k$ is the $k$-th standard basis vector in $\mathbb R^N$. Given a random variable $Z$ with values in $X$, $$GZ = \left[\begin{matrix}\langle g_1, Z\rangle \\ \vdots \\ \langle g_N, Z\rangle\end{matrix}\right]$$ is a random vector taking values in $\mathbb R^N$. We say that $Z$ is a normal random variable if $GZ$ is normal for all finite choices of $N$ and $G$. The mean of $Z$ is the vector $m$ in $X$ such that $GZ$ has the mean $Gm$ for all linear mappings $G$. Likewise, the covariance $C$ of $Z$ is an element of $X\otimes X$ such that $GZ$ has the covariance matrix $GCG^*$ for all $G$.

Not just any operator can be the covariance of a normal random variable on a function space. It ought to be familiar from the finite-dimensional theory that $C$ has to be symmetric and positive-definite. You might imagine writing down a decomposition for $C$ in terms of its eigenfunctions $\{\psi_n\}$. We know that the eigenfunctions are orthonormal because $C$ is symmetric. Since the eigenvalues are real and positive, we can write them as the squares $\sigma_n^2$ of a sequence $\{\sigma_n\}$ of real numbers as a mnemonic. I'll do this throughout because it becomes convenient in a number of places. The action of $C$ on a function $\phi$ is then $$C\phi = \sum_n\sigma_n^2\langle\psi_n, \phi\rangle\,\psi_n.$$ In the function space setting, we need that the sum of all the eigenvalues is finite: $$\sum_n\sigma_n^2 < \infty.$$ The fancy term for this is that $C$ has to be of trace class. In finite dimensions we don't have to worry about this condition at all. In function spaces, it's easy to try and use a covariance operator with $\sigma_n^2 \sim n^{-1}$ by mistake. The trace then diverges like the harmonic series.

The Karhunen-Loève expansion¶

The spectral theorem assures us that the covariance operator of a Gaussian process has to have an eigenvalue factorization. The Karhunen-Loève theorem then shows how we can use that decomposition to understand the stochastic process. It states that, if $\{\xi_k\}$ are an infinite sequence of i.i.d. standard normal random variables, then the random variable $$Z = \sum_n\sigma_n\,\xi_n\,\psi_n$$ has the distribution $N(0, C)$. In other words, if we know the eigendecomposition, we can generate samples from the process. We'll use this property in just a moment to simulate Brownian bridges. The eigenfunctions for the covariance operator of a Brownian bridge are trigonometric functions. We're simulating the stochastic process by Fourier synthesis.

Rather than start with a covariance operator and then compute its eigenvalue decomposition, the KL expansion also gives us a way to synthesize random variables with useful properties by picking the eigenvalues and eigenfunctions. The covariance operator is then $$C = \sum_n\sigma_n^2\,\psi_n\otimes\psi_n.$$ We can compute a closed form expression for $C$ as some integral operator $$C\phi = \int_\Omega \mathscr C(x, y)\phi(y)dy.$$ in some cases. But we don't need an explicit expression for the kernel for the KL expansion to be useful.

The code below generates some random functions on the domain $\Omega = [0, 1]$. Here we use the basis functions $$\psi_n = \sin(\pi n x).$$ We can then try eigenvalue sequences $\sigma_n = \alpha n^{-\gamma}$ for some power $\gamma$. Taking $\gamma \ge 1$ guarantees that $\sum_n\sigma_n^2$ is finite.

The plot below compares two function-valued random variables with different eigenvalue sequences. The plot on the top shows a classic Brownian bridge: $\sigma_n^2 = 2 / \pi^2n^2$. The plot on the bottom shows $\sigma_n^2 = 2 / \pi^2n^{2\gamma}$ where $\gamma = 1.6$. The second stochastic process is much smoother, which we expect -- the Fourier coefficients decay faster.

The Feldman-Hájek theorem¶

Now consider two $X$-valued normal random variables $W$ and $Z$. To simplify things, we'll assume that both of them have zero mean. The plots above might suggest that these two stochastic processes behave differently in some essential way that we can quantify. For example, the first process is a Brownian bridge. A Brownian bridge is continuous with probability 1, but it is nowhere differentiable. The second process, on the other hand, is almost surely differentiable. We can be a little more formal about this and say that $$\mathbb P[W \in H^1(\Omega)] = 0, \quad \mathbb P[Z \in H^1(\Omega)] = 1$$ where $H^1(\Omega)$ is the space of functions with square-integrable derivatives. The two random variables are not mutually absolutely continuous. That might not be too surprising. We picked the eigenvalue distributions of each process to have different decay rates: $$\sigma_n^2(W) \sim n^{-2}, \quad \sigma_n^2(Z) \sim n^{-2\gamma}.$$ Just to really hammer this point home, let's evaluate an approximation to the mean-square derivative of each group of samples.

If you run this notebook with an even finer discretization of the unit interval, you'll see that $\int|\nabla w|^2\mathrm dx$ keeps increasing while the value of $\int|\nabla z|^2\mathrm dx$ stabilizes.

What conditions guarantee that two function-value random variables are absolutely continuous? In other words, we want to know when it is impossible for there to be some subset $A$ of $X$ such that there is positive probability that $W$ is in $A$ but zero probability that $Z$ is in $A$ or vice versa. What needs to be true about the covariance operators of the two random variables to fulfill this condition? This is the content of the Feldman-Hájek theorem.

The calculation above suggests that two normal distributions will concentrate on different sets if the ratios $\sigma_n^2(W)/\sigma_n^2(Z)$ of the eigenvalues of their covariance operators can go to either zero or infinity. The example above used two eigenvalue sequences that look like $n^{-\alpha}$ for different values of $\alpha$. The gap in the decay rates then lets us construct a quadratic functional for which the process $Z$ takes finite values and $W$ takes infinite values. The fact of this functional taking on finite or infinite values on one process or the other then implies the concentration of the two probability measures on different sets. The fact that the mean square gradient of $W$ is infinite implies that $W$ is not differentiable with probability 1, while $Z$ is differentiable with probability 1. Put another way, the probability distribution for $Z$ is concentrated in the set of differentiable functions while the probability distribution for $W$ is not.

We can always construct an operator like this when the decay rates differ, although the effect might be subtler. For example, the distinction might be in having a higher-order derivative or not.

What if the ratio of the eigenvalue sequences doesn't go to zero or infinity? Are the laws of the two processes equivalent, or can they still concentrate on different sets? We can make it even simpler by supposing that $W \sim N(0, C)$ and $Z \sim N(0, \alpha C)$ for some non-zero $\alpha$. We'll show that even these random variables concentrate on different sets. I got this example from MathOverflow. Let $\{\sigma_n^2\}$ be the eigenvalues of $C$ (note the square) and $\{\psi_n\}$ the eigenfunctions. Define the sequence of random variables $$U_n = \langle \psi_n, W\rangle / \sigma_n.$$ We know from the KL expansion that the $U_n$ are independent, identically-distributed standard normal random variables. If we also define $$V_n = \langle\psi_n, Z\rangle / \sigma_n$$ then these random variables are i.i.d. normals with mean zero and variance $\alpha$.

Now remember that $W$ and $Z$ take values in a function space and so there are infinitely many modes to work with. The strong law of large numbers then implies that $$\frac{1}{n}\sum_nU_n^2 \to 1$$ with probability 1. This would not be possible in a finite-dimensional vector space -- the sum would have to terminate at some $N$. The average of $U_n^2$ would be a random variable with some non-trivial spread around its mean. It can take on any positive value with non-zero probability. That probability could be minute, but it is still positive. Likewise, we also find that $$\frac{1}{n}\sum_nV_n^2 \to \alpha$$ with probability 1. But we've now shown that the two probability measures concentrate on different sets as long as $\alpha$ is not equal to 1. The weighted average of squared expansion coefficients takes distinct values, each with probability 1, for the distributions of $W$ and $Z$. So the two distributions are mutually singular.

Just to illustrate things a bit, the plot below shows samples obtained with a spectrum of $\{\sigma_n\}$ generated as before with $\gamma = 1.5$ and with $2\sigma_n$ respectively.

What this example shows is that it's not enough that $\sigma_n(W)$ and $\sigma_n(Z)$ have the same decay rate. They both decay like some constant $\times$ $n^{-\gamma}$ for the same value of $\gamma$. The two distributions are mutually singular even if the rates are the same but the constants are different.

At this point you might wonder if it's enough to have $$\sigma_n(W)/\sigma_n(Z) \to 1.$$ Even this much tighter condition is not enough!

The Feldman-Hájek theorem states that, in order for $W$ and $Z$ to not concentrate on different sets, we need that $$\sum_n\left(\frac{\sigma_n(W)}{\sigma_n(Z)} - 1\right)^2 < \infty.$$ Not only do we need that the eigenvalue ratio goes to 1, we need it to go to 1 fast enough to be summable. For example, suppose we had a pair of $X$-valued random variables such that $$\frac{\sigma_n(W)}{\sigma_n(Z)} - 1 \sim \frac{\text{const}}{\log n}.$$ The relative error between the two eigenvalue sequences is decreasing to zero. But it isn't decreasing fast enough, so the two random variables would concentrate on different sets.

Why does this matter? I stumbled on the Feldman-Hájek theorem while trying to understand two papers, Beskos et al. (2008) and Cotter et al. (2013). Both of these papers are about how to do Markov-chain Monte Carlo sampling of probability distributions defined over function spaces. The main ingredient of MCMC is a proposal mechanism, which generates a new candidate sample $z$ from the current sample $w$. We can then also define a reverse proposal -- what is the probability of proposing $w$ if we instead started from $z$? In order for an MCMC method to converge, Tierney (1998) showed that the forward and reverse proposals need to be mutually absolutely continuous. They cannot concentrate probability on different sets. In finite dimensions, it's so hard to violate this condition that it often goes unstated. In function spaces, on the other hand, you're almost guaranteed to violate it unless you're careful. Even for the nicest class of distributions (normal random variables), the spectra of the covariance operators have to have nearly identical asymptotic behavior. In light of that, it's not surprising that MCMC sampling in function spaces is so challenging.

In this post we'll look at a fourth-order equation describing the vibrations of clamped plates. Fourth-order problems are much harder to discretize than second-order. Here we'll get around that by using its dual or mixed form. This will use some of the same techniques that we've already seen with the Stokes equations. But it won't be a perfectly conforming discretization, and as a consequence we'll need some jump terms that ought to be familiar from the discontinuous Galerkin method. Finally we'll compute some eigenfunctions of plate operator. This will look similar to what we did with Weyl's law but with an extra mathematical subtlety.

Our starting point is linear elasticity. Here we want to solve for the 3D displacement $u$ of the medium. The first equation is momentum conservation: $$\rho\ddot u = \nabla\cdot\sigma + f$$ where $\sigma$ is the stress tensor and $f$ the body forces. To close the system, we need to supply a constitutive equation relating the stress tensor and the strain tensor $$\varepsilon = \frac{1}{2}\left(\nabla u + \nabla u^*\right).$$ The most general linear constitutive equation we can write down is $$\sigma = \mathscr C\,\varepsilon$$ where $\mathscr C$ is the rank-4 elasticity tensor. We need that $\mathscr C$ maps symmetric tensors to symmetric tensors and that $\varepsilon:\mathscr C\varepsilon$ is always positive. The equilibrium displacement of the medium is the minimizer of the energy functional $$J(u) = \int_\Omega\left(\frac{1}{2}\sigma:\varepsilon + f\cdot u\right)\mathrm dx + \ldots$$ where the ellipses stand for various boundary forcings that we'll ignore.

For a medium that is homogeneous and isotropic, the elasticity tensor has to have the form $$\mathscr C\,\varepsilon = 2\,\mu\,\varepsilon + \lambda\,\text{tr}(\varepsilon)I$$ where $\mu$ and $\lambda$ are the Lamé parameters. As an aside, there are a mess of alternate forms of the elasticity equations. The wiki page has a conversion table at the bottom. Now take this with a grain of salt because I do fluid mechanics. But if I were a solid mechanics kinda guy, this would embarrass me. Get it together folks.

Plate theory is what you get when you assume the medium is thin along the vertical dimension and that this restricts the form that the displacements can take. The first and most widely agreed-upon simplification is that the vertical displacement is some function $w$ of the horizontal coordinates $x$, $y$: $$u_z = w(x, y).$$ From here we have to make additional assumptions about the horizontal displacements $u_x$ and $u_y$. Different theories make different sets of assumptions. The classical theory is the Kirchoff-Love plate. The Kirchoff theory assumes that any straight line that's perpendicular to the middle of the plate remains straight and perpendicular after deformation. This theory has some deficiences, see for example this paper. But it's a good starting point for experimentation. These assumptions let us write down the other components of the deformation in terms of $w$: $$u_x = -z\frac{\partial w}{\partial x}, \quad u_y = -z\frac{\partial w}{\partial y}.$$ It's possible that $u_x$ and $u_y$ are also offset by in-plane displacements. I'll assume that the boundary conditions make those equal to zero. We can then work out what the displacement gradient is: $$\nabla u = \left[\begin{matrix}-z\nabla^2 w & -\nabla w \\ +\nabla w^* & 0 \end{matrix}\right]$$ I'm being a little loose in my use of the gradient operator -- it's 3D on the left-hand side of this equation and 2D on the right. Notice how the antisymmetric part of the displacement gradient is all in the $x-z$ and $y-z$ components. When we symmetrize the displacement gradient, we get the strain tensor: $$\varepsilon = -z\left[\begin{matrix}\nabla^2 w & 0 \\ 0 & 0\end{matrix}\right].$$ Now remember that the medium is a thin plate. We can express the spatial domain as the product of a 2D footprint domain $\Phi$ and the interval $[-h / 2, +h / 2]$ where $h$ is the plate thickness. The strain energy is then $$\begin{align} J(w) & = \int_\Phi\int_{-h/2}^{+h/2}\left(\mu z^2|\nabla^2 w|^2 + \frac{\lambda}{2} z^2|\Delta w|^2 + fw\right)\mathrm dz\;\mathrm dx\\ & = \int_\Phi\left\{\frac{h^3}{24}\left(2\mu|\nabla ^2 w|^2 + \lambda|\Delta w|^2\right) + hfw\right\}\mathrm dx \end{align}.$$ Here I've used the fact that $\text{tr}(\nabla^2w) = \Delta w$ where $\Delta$ is the Laplace operator. In a gross abuse of notation I'll write this as $$J(w) = \int_\Omega\left(\frac{1}{2}\mathscr C\nabla^2w :\nabla^2 w + fw\right)h\,\mathrm dx$$ where $\mathscr C$ is an elasticity tensor. We'll need the explicit form $$\mathscr C\kappa = \frac{h^2}{12}\left(2\,\mu\,\kappa + \lambda\,\text{tr}(\kappa)\,I\right)$$ in a moment.

We can't discretize this problem right away using a conventional finite element basis. The usual piecewise polynomial basis functions are differentiable and piecewise continuous across cell edges. A conforming basis for a minimization problem involving 2nd-order derivatives would need to instead by continuously differentiable. It's much harder to come up with $C^1$ bases.

We have a few ways out.

1. Use a $C^1$ basis like Argyris or Hsieh-Clough-Tocher. This approach makes forming the problem easy, but applying boundary conditions hard. Kirby and Mitchell (2018) implemented the Argyris element in Firedrake but had to use Nitsche's method to enforce the boundary conditions.
2. Use $C^0$ elements and discretize the second derivatives using an interior penalty formulation of the problem. This approach is analogous to discretizing a 2nd-order elliptic problem using DG elements. Forming the problem is harder but applying the boundary conditions is easier. Bringmann et al. (2023) work out the exact form the interior penalty parameters necessary to make the discrete form work right.
3. Use the mixed or dual form of the problem, which introduces the moment tensor explicitly as an unknown. Discretize it with the Hellan-Herrmann-Johnson or HHJ element. Arnold and Walker (2020) describe this formulation of the problem in more detail.

In the following, I'll take the third approach. The dual form of the problem introduces the moment tensor, which in another gross abuse of notation I'll write as $\sigma$, explicitly as an unknown. We then add the constraint that $\sigma = \mathscr C\nabla^2 w$. But switching to the dual form of the problem inverts constitutive relations, so we instead enforce $$\nabla^2w = \mathscr A\sigma$$ where $\mathscr A$ is the inverse to $\mathscr C$. If $\mathscr C$ is like the elasticity tensor in 3D, then $\mathscr A$ is like the compliance tensor. Because we'll need it later, the explicit form of the compliance tensor is $$\mathscr A\sigma = \frac{6}{h^2\mu}\left(\sigma - \frac{\lambda}{2(\mu + \lambda)}\text{tr}(\sigma)I\right).$$ The full Lagrangian for the dual form is $$L(w, \sigma) = \int_\Omega\left(\frac{1}{2}\mathscr A\sigma : \sigma - \sigma : \nabla^2 w - fw\right)h\,\mathrm dx.$$ Here the displacement $w$ acts like a multiplier enforcing the constraint that $\nabla\cdot\nabla\cdot\sigma + f = 0$.

To get a conforming discretization of the dual problem, we'd still need the basis functions for the displacements to have continuous derivatives. That's exactly what we were trying to avoid by using the dual form in the first place. At the same time, we don't assume any continuity properties of the moments. We can work around this problem by relaxing the continuity requirements for the displacements while strengthening them for the moments. The Hellann-Herrmann-Johnson element discretizes the space of symmetric tensor fields with just enough additional regularity to make this idea work. The HHJ element has normal-normal continuity: the quantity $n\cdot\sigma n$ is continuous across cell boundaries. The tangential and mixed components are unconstrained. This extra continuity is enough to let us use the conventional Lagrange finite elements for the displacements as long as we add a correction term to the Lagrangian: $$\ldots + \sum_\Gamma\int_\Gamma (n\cdot\sigma n)\,\left[\!\!\left[\frac{\partial w}{\partial n}\right]\!\!\right]h\,\mathrm dS$$ where $\Gamma$ are all the edges of the mesh and the double square bracket denotes the jump of a quantity across a cell boundary. What is especially nice about the HHJ element is that this correction is all we need to add; there are no mesh-dependent penalty factors.

Experiments¶

First let's try and solve a plate problem using simple input data. We'll set all the physical constants equal to 1 but you can look these up for the material of your fancy.

The code below creates the dual form of the problem using the formula I wrote down above for the explicit form of the compliance tensor. Lord Jesus I hope I did all that algebra correct.

The right degrees are $p + 1$ for the displacements and $p$ for the moments.

Here we'll work on the unit square. An interesting feature of plate problems is that they don't become much easier to solve analytically on the unit square by separation of variables because of the mixed derivatives. The eigenfunctions of the biharmonic operator (we'll get to them below) are easier to derive on the circle than the square.

For plate problems, we almost always have some Dirichlet boundary condition $$w|_{\partial\Omega} = w_0.$$ Since the problem is fourth-order, we need to supply more boundary conditions than we do for, say, the diffusion equation. There are two kinds. The first is the clamped boundary condition: $$\frac{\partial w}{\partial n} = 0.$$ The second is the simply-supported boundary condition: $$n\cdot\sigma n = 0.$$ Before I started writing this I was dreading how I was going to deal with either one. I was shocked at how easy it was to get both using the HHJ element. First, let's see what happens if we only enforce zero displacement at the boundary.

So it looks like the clamped boundary condition is natural with HHJ elements. We can confirm that by looking at a 3D plot.

Now let's see what happens if we apply a Dirichlet boundary condition to the moment part of the solution.

Now the boundary stresses are zero and so the boundary slopes are unconstrained.

The reason why this works is that the boundary degrees of freedom for the HHJ element are the normal-normal stress components. So if we set all of those to zero, we get simply-supported boundary conditions. This is analogous to the simpler setting of $H(\text{div})$ elements for vector fields having normal continuity. We can easily enforce boundary conditions like $u\cdot n = 0$ by setting the right degrees of freedom to zero.

We have to work a lot harder to get the boundary conditions right using other finite element bases. For example, the Argyris element has both value and derivative degrees of freedom at the mesh vertices. We can't set the boundary values alone in the same way. In order to set the boundary values with Argyris elements, Kirby and Mitchell had to use Nitsche's method. The Hsieh-Clough-Tocher (HCT) element makes it easy to use clamped boundary conditions but simply-supported is harder. The fact that both types are relatively easy with HHJ is a definite advantage, although it has more total degrees of freedom than HCT.

Eigenfunctions¶

The eigenfunctions of the plate problem, also known as Chladni patterns, have a fascinating history. This review is good reading. Here I'll show the first few eigenfunctions in a square with simply-supported boundary conditions.

Note the $M$ matrix that we supplied above. We've looked at generalized eigenvalue problems before, like in the post on Weyl's law. In that setting we were solving eigenvalue problems of the form $$A\phi = \lambda M\phi$$ where $M$ was some symmetric positive-definite matrix instead of the identity. In those cases we always used the mass matrix. For saddle-point eigenvalue problems, we're often using something that is no longer definite. Here the form we want is instead $$\left[\begin{matrix} A & B^* \\ B & 0\end{matrix}\right]\left[\begin{matrix}\sigma \\ w\end{matrix}\right] = \lambda\left[\begin{matrix} 0 & 0 \\ 0 & M\end{matrix}\right]\left[\begin{matrix}\sigma \\ w\end{matrix}\right]$$ where $M$ is the mass matrix for the displacement block.

The plot below shows the nodal lines of the eigenfunctions. Many of them come in pairs due to the axis-flipping symmetry. These are the patterns that so fascinated Chladni and others.

Mantle convection¶

This code is taken from a chapter in The FEniCS Book, which is available for free online. That chapter is in turn an implementation of a model setup from van Keken et al (1997). I took the thermomechanical parts and removed the chemistry. This paper and the code from the FEniCS book all use a non-dimensional form of the equations, which I've adopted here.

In this notebook, we'll see the Stokes equations again, but we'll couple them to the evolution of temperature through an advection-diffusion equation. The extensive quantity is not the temperature itself but the internal energy density $E = \rho c_p T$ where $\rho$, $c_p$ are the mass density and the specific heat at constant pressure. The flux of heat is $$F = \rho c_p T u - k\nabla T.$$ We'll assume there are no heat sources but, but the real physics would include sources from decay of radioactive elements, strain heating, and chemical reactions. The variational form of the heat equation is: $$\int_\Omega\left\{\partial_t(\rho c_p T)\,\phi - (\rho c_p Tu - k\nabla T)\cdot\nabla\phi - Q\phi\right\}dx = 0$$ for all test functions $\phi$. I've written this in such a way that it still makes sense when the density and heat capacity aren't constant. We'll assume they are (mostly) constant in the following. The variational form of the momentum balance equation is $$\int_\Omega\left\{2\mu\varepsilon(u): \dot\varepsilon(v) - p\nabla\cdot v - q\nabla\cdot u - \rho g\cdot v\right\}dx = 0$$ for all velocity and pressure test functions $v$, $q$.

The characteristic that makes this problem interesting is that we assume the flow is just barely compressible. This is called the Boussinesq approximation). In practice, it means we do two apparently contradictory things. First, we assume that the flow is incompressible: $$\nabla\cdot u = 0.$$ Second, we assume that the density that we use on the right-hand side of the momentum balance equation actually is compressible, and that the degree of compression is linear in the temperature: $$\rho = \rho_0(1 - \beta(T - T_0))$$ where $\rho_0$ is a reference density and $\beta$ is thermal expansion coefficient. These assumptions aren't consistent with each other. We're also not consistent in using this everywhere. For example, we only use this expression on the right-hand side of the momentum balance equation, but we'll still take the internal energy density to be equal to $\rho_0c_pT$ in the energy equation. In other words, the material is compressible as far as momentum balance is concerned but not energy balance. I find that uncomfortable. But I would find writing a full compressible solver to be blinding agony. Anyway it lets us make pretty pictures.

I won't go into a full non-dimensionalization of the problem, which is detailed in the van Keken paper. The most important dimensionless numbers for this system are the Prandtl number and the Rayleigh number. The Prandtl number quantifies the ratio of momentum diffusivity to thermal diffusivity: $$\text{Pr} = \frac{\mu / \rho}{k / \rho c_p} = \frac{\nu}{\kappa}$$ where $\nu$ and $\kappa$ are respectively the kinematic viscosity and thermal diffusivity. For mantle flows, the Prandtl number is on the order of $10^{25}$. We can get a sense for a characteristic velocity scale from the momentum balance equation. First, we can take a factor of $\rho_0g$ into the pressure. If we take $U$, $L$, and $\Delta T$ to be the characteristic velocity, length, and temperature difference scales, non-dimensionalizing the momentum balance equation gives $$\frac{\mu U}{L^2} \approx \rho_0g\beta\Delta T$$ which then implies that $U = \rho_0g\beta L^2\Delta T / \mu$. Now we can compute the thermal Péclet number $$\text{Pe} = \frac{UL}{\kappa}.$$ When the Péclet number is large, we're in the convection-dominated regime; when it's small, we're in the diffusion-dominated regime. Taking our estimate of velocity scale from before, we get $$\text{Pe} = \frac{\rho_0 g \beta L^3 \Delta T}{\kappa\mu}.$$ This is identical to the Rayleigh number, which we'll write as $\text{Ra}$. Lord Rayleigh famously used linear stability analysis to show that a fluid that is heated from below and cooled from above is unstable when the Rayleigh number exceeds 600 or so. Here we'll take the Rayleigh number to be equal to $10^6$ as in the van Keken paper. It's a worthwhile exercise in numerical analysis to see what it takes to make a solver go for higher values of the Rayleigh number.

Here we're using a rectangular domain. Again, the spatial scales have been non-dimensionalized.

The initial condition for temperature involves loads of annoyingly long expressions, but it's all in Appendix A of the van Keken paper.

The important point is to make the fluid hotter below and on one side and cooler above and on the other side, shown in the plot below.

Next we need to make some function spaces for the fluid velocity and pressure. Note how the degree of the velocity space is one higher than that of the pressure -- we're using the Taylor-Hood element again.

Once we've created a function in the mixed space, we can then pull out the two parts with the split method.

The code below creates the variational form of the Stokes equations $$\nabla\cdot\tau - \nabla p + \text{Ra}\;T\; g = 0$$ with temperature-dependent gravitational forcing.

We can use the .sub method to pull parts out of mixed spaces, which we need in order to create the right boundary conditions.

A bit of magic in order to tell the linear solver that the Stokes equations have a null space we need to project out.

Here I'm creating solver object since we'll need to repeatedly solve the same system many times.

Here we're setting up the temperature solver, which includes both convection and diffusion. I've set the timestep based on the mesh size and the maximum speed that we just found above from the velocity solution.

And the timestepping loop. Note that the final time is on a non-dimensional scale again, in physical time it works out to be on the order of a hundred million years. Here we're using an operator splitting approach. We first update the temperature, and then we compute a new velocity and pressure. The splitting error goes like $\mathscr{O}(\delta t)$ as the timestep is reduced. So if we use a first-order integration scheme like backward Euler then the splitting error is asymptotically the same as that of the discretization itself. We could get $\mathscr{O}(\delta t^3)$ or higher convergence by using, say, the Radau-IIA method, but the total error would be dominated by splitting, so there would be little point in trying harder.

The movie below shows the temperature evolution. The rising and sinking plumes of hot and cold fluid correspond to the mantle plumes that produce surface volcanism.

What next¶

There are a few features in the van Keken paper and otherwise that I didn't attempt here. The original van Keken paper adds chemistry. Solving the associated species transport equation with as little numerical diffusion as possible is hard. You could then make the fluid buoyancy depend on chemical composition and add temperature-dependent chemical reactions.

Real mantle fluid also has a temperature-dependent viscosity. When I tried to add this using the splitting scheme above, the solver quickly diverges. Getting that to work might require a fully coupled time-integration scheme, which I did try. Rather than split up the temperature and the velocity/pressure solves, a fully coupled scheme would solve for all three fields at once. If you could make this work, it would open up the possibility of using higher-order methods like Radau-IIA. Whether it would go or not seemed to depend on the machine and the day of the week. I'll explore that more in a future post.

In previous posts, I looked at how to discretize the incompressible Stokes equations. The Stokes equations are a good approximation when the fluid speed is small enough that inertial effects are negligible. The relevant dimensionless number is the ratio $$\text{Re} = \frac{\rho UL}{\mu},$$ the Reynolds number. Stokes flow applies when the Reynolds number is substantially less than 1. The incompressibility constraint adds a new difficulty: we have to make good choices of finite element bases for the velocity and pressure. If we fail to do that, the resulting linear systems can have no solution or infinitely many.

Here I'll look at how we can discretize the full Navier-Stokes equations: $$\frac{\partial}{\partial t}\rho u + \nabla\cdot\rho u\otimes u = -\nabla p + \nabla\cdot \tau$$ where the deviatoric stress tensor is $\tau = 2\mu\dot\varepsilon$. The inertial terms are nonlinear, which makes this problem more difficult yet than the Stokes equations.

The goal here will be to simulate the famous von Kármán vortex street.

Making the initial geometry¶

First, we'll make a domain consisting of a circle punched out of a box. The fluid flow in the wake of the circle will produce vortices.

Initial velocity¶

We'll use a fixed inflow velocity $$u_x = 4 \frac{y}{L_y}\left(1 - \frac{y}{L_y}\right) u_{\text{in}}.$$ We'll take as characteristic length scale the radius of the disc $L_y / 8$. In order to see the effect we want, we'll need a Reynolds number that's on the order of 100 or greater.

I've used Taylor-Hood elements for this demonstration. It would be a good exercise to repeat this using, say, Crouzeix-Raviart or other elements.

Here we can use the fact that the Stokes problem has a minimization form. The Navier-Stokes equations do not because of the convective term.

A bit of a philosophical point is in order here. We picked the inflow velocity and the viscosity to produce a high Reynolds number (about 200). The Stokes equations, on the other hand, are physically realistic only when the Reynolds number is $\ll 1$. But we can still solve the Stokes equations in the high-Reynolds number limit. In other words, the mathematical model remains well-posed even in regimes where it is not applicable. Here we're only using the result to initialize a simulation using the "right" model. But it's a mistake -- and one I've made -- to lull yourself into a false sense of correctness merely because the model gave you an answer.

Solution method¶

There are a host of methods for solving the Navier-Stokes equations by breaking them up into two simpler problems. These are known as projection methods). I started out trying those but it required effort, which I find disgusting. So I threw backward Euler at it and it worked.

There are some non-trivial decisions to make about both the variational form and the boundary conditions. I started writing this wanting to use pressure boundary conditions at both the inflow and outflow. (Of course I had to be reminded that you can't prescribe pressures but rather tractions.) This went badly. If a wave reflects off of the obstacle, back upstream, and out the inflow boundary, the simulation will crash. So I had to dial back the challenge and use a fixed inflow velocity and a traction boundary condition at the outflow.

Almost any writing you see about the Navier-Stokes equations will express the problem in differential form and will use incompressibility to apply what might appear to be simplifications. For example, if the fluid is incompressible and the viscosity is constant, then you can rewrite the viscous term like so: $$\nabla\cdot \mu(\nabla u + \nabla u^\top) = \mu\nabla^2u.$$ You'll see this form in almost all numerical methods or engineering textbooks. I don't like it for two reasons. First, the apparent simplification only gets in your way as soon as you want to consider fluids with variable viscosity. Mantle convection is one obvious case -- the temperature and chemistry of mantle rock can change the viscosity by several orders of magnitude. Second, it gives the wrong boundary conditions when you try to discretize the problem (see Limache et al. (2007)). I've retained the symmetric gradients of the velocity and test function in the form below.

The second apparent simplification uses incompressibility to rewrite the convection term: $$\nabla\cdot \rho u\otimes u = u\cdot\nabla \rho u.$$ This form is ubiquitous and reflects a preference for thinking about fluid flow in a Lagrangian reference frame. I prefer to avoid it although both are correct, unlike the Laplacian form of the viscosity. Given any extensive density $\phi$, regardless of its tensor rank, the flux will include a term $\phi\cdot u$. The conversion of this flux from the conservation to the variational form is then $$-\int_\Omega \phi u\cdot\nabla\psi\,dx$$ and this is true of mass, momentum, energy, whatever you like. Taking something that was a divergence and making it not a divergence obfuscates the original conservation principle. It also stops you from pushing the differential operator over onto a test function. So I've instead coded up the convection term as a discretization of $$-\int_\Omega u\otimes u\cdot\dot\varepsilon(v)\,dx + \int_{\partial\Omega\cap\{u\cdot n \ge 0\}}(u\cdot n)(v\cdot n)ds.$$ In the first term, I've used the symmetric gradient of $v$ because the contraction of a symmetric and an anti-symmetric tensor is zero.

All together now, the variational form that I'm using is $$\begin{align} 0 & = \int_\Omega\left\{\rho\,\partial_tu\cdot v - \rho u\otimes u:\varepsilon(v) - p\nabla\cdot v - q\nabla\cdot u + 2\mu\,\dot\varepsilon(u):\dot\varepsilon(v)\right\}dx \\ & \qquad\qquad + \int_\Gamma (\rho u\cdot n)(u \cdot v)ds. \end{align}$$ for all test functions $v$ and $q$.

We'll need to make some choice about the timestep. Here I've computed the CFL time for the mesh and the initial velocity that we computed above. This choice might not be good enough. If we initialized the velocity by solving the Stokes equations, the fluid could evolve to a much higher speed. We might then find that this timestep is inadequate. A principled solution would be to use an adaptive scheme.

I've added a bit of code to show some diagnostic information in the progress bar. First I have it printing out the number of Newton iterations that were required to compute each timestep. If you see this going much above 20 then something is off. Second, I had it print out the maximum pressure. Both of these were useful when I was debugging this code.

Finally, we'll make an animated quiver plot because it looks pretty.

There's an empirical formula for the frequency of vortex shedding. A fun follow-up to this would be to compute the shedding frequency from the simulation using, say, a windowed Fourier transform, and comparing the result to the empirical formula. Next on the docket is comparing the results using different spatial finite elements.

In this post, we'll look at overland flow -- how rainwater drains across a landscape. The equations of motion are pretty rowdy and have some fascinating effects. To derive them, we'll start from the shallow water or Saint Venant equations for the water layer thickness $h$ and velocity $u$:

$$\begin{align} \frac{\partial h}{\partial t} + \nabla\cdot hu & = \dot a \\ \frac{\partial}{\partial t}hu + \nabla\cdot hu\otimes u & = -gh\nabla (b + h) - k|u|u \end{align}$$

The final term in the momentum equation represents frictional dissipation and $k$ is a (dimensionless) friction coefficient. Using the shallow water equations for predicting overland flow is challenging because the thickness can go to zero.

For many thin open channel flows, however, the fluid velocity can be expressed purely in terms of the surface slope and other factors. You could arrive at one such simplification by assuming that the inertial terms in the momentum equation are zero:

$$k|u|u + gh\nabla(b + h) = 0.$$

This approximation is known as the Darcy-Weisbach equation. We'll use it in the following because it's simple and it illustrates all the major difficulties. For serious work, you'd probably want to use the Manning formula, as it has some theoretical justification for turbulent open channel flows. The overall form of the equation and the resulting numerical challenges are the same in each case.

Putting together the Darcy-Weisbach equation for the velocity with the mass conservation equation gives a single PDE for the water layer thickness:

$$\frac{\partial h}{\partial t} - \nabla\cdot\left(\sqrt{\frac{gh^3}{k}}\frac{\nabla(b + h)}{\sqrt{|\nabla(b + h)|}}\right) = \dot a.$$

This looks like a parabolic equation, but there's a catch! The diffusion coefficient is proportional to $h^{3/2}$, so it can go to zero when $h = 0$; all the theory for elliptic and parabolic equations assumes that the diffusion coefficient is bounded below. For a non-degenerate parabolic PDE, disturbances propagate with infinite speed. For the degenerate problem we're considering, that's no longer true -- the $h = 0$ contour travels with finite speed! While we're using the Darcy-Weisbach equation to set the velocity here, we still get finite propagation speed if we use the Manning equation instead. What's important is that the velocity is propertional to some power of the thickness and surface slope.

Eliminating the velocity entirely from the problem is convenient for analysis, but not necessarily the best way to go numerically. We'll retain the velocity as an unknown, which gives the resulting variational form much of the same character as the mixed discretization of the heat equation.

As our model problem, we'll use the dam break test case from Santillana and Dawson (2009). They discretized the overland flow equations using the local discontinuous Galerkin or LDG method, which extends DG for hyperbolic systems to mixed advection-diffusion problems. We'll use different numerics because Firedrake has all the hipster elements. I'm eyeballing the shape of the domain from their figures.

The bed profile consists of an upper, elevated basin, together with a ramp down to a lower basin.

As I alluded to before, rather than eliminate the velocity entirely, we'll keep it as a field to be solved for explicitly. The problem we're solving, while degenerate, is pretty similar to the mixed form of the heat equation. This suggests that we should use element pairs that are stable for mixed Poisson. Here I'm using the MINI element: continuous linear basis functions for the thickness, and continuous linear enriched with cubic bubbles for the velocity. We could also have used a more proper $H(\text{div})$-conforming pair, like discontinuous Galerkin for the thickness and Raviart-Thomas or Brezzi-Douglas-Marini for the velocity.

The dam break problem specifies that the thickness is equal to 1 in the upper basin and 0 elsewhere. I've done a bit of extra work below because the expression for $h$ is discontinuous, and interpolating it directly gives some obvious mesh artifacts. Instead, I've chosen to project the expression and clamp it above and below.

The test case in the Santillana and Dawson paper uses a variable friction coefficient in order to simulate the effect of increased drag when flowing over vegetation.

The code below defines the variational form of the overland flow equations.

We'll run into trouble if we try and use a Newton-type method on the true variational form. Notice how the $q$-$q$ block of the derivative will go to zero whenever $q = 0$. This will happen whenever the thickness is zero too. The usual hack is to put a fudge factor $\varepsilon$ into the variational form, as shown below.

The disadvantage of is that we're then solving a slightly different physics problem. We don't have a great idea ahead of time of what $\varepsilon$ should be either. If we choose it too large, the deviation from the true problem is large enough that we can't believe the results. But if we choose it too small, the derivative will fail to be invertible.

We can take a middle course by instead just using the perturbed variational form just to define the derivative in Newton's method, but keep the true variational form as the quantity to find a root for. To do this, we'll pass the derivative of $F_\varepsilon$ as the Jacobian or J argument to the nonlinear variational problem object. Choosing $\varepsilon$ too small will still cause the solver to crash. Taking it to be too large, instead of causing us to solve a completely different problem, will now only make the solver go slower instead. We still want to make $\varepsilon$ as small as possible, but to my mind, getting the right answer slowly is a more tolerable failure mode than getting the wrong answer.

We'll have one final difficulty to overcome -- what happens if the thickness inadvertently becomes negative? There's a blunt solution that everyone uses, which is to clamp the thickness to 0 from below at every step. Clamping can work in many cases. But if you're using a Runge-Kutta method, it only assures positivity at the end of each step and not in any of the intermediate stages. We can instead formulate the whole problem, including the non-negativity constraint, as a variational inequality. Much like how some but not all variational problems arise from minimization principles, some variational inequalities arise from minimization principles with inequality constraints, like the obstacle problem. But variational inequalities are a more general class of problem than inequality-constrained minimization. Formulating overland flow as as a variational inequality is a bit of overkill for the time discretization that we're using. Nonetheless, I'll show how to do that in the following just for illustrative purposes. We first need two functions representing the upper and lower bounds for the solution. In this case, the upper bound is infinity.

The thickness is bounded below by 0, but there's no lower bound at all on the flux, so we'll set only the flux entries to negative infinity.

When we want to solve variational inequalities, we can't use the usual Newton solvers in PETSc -- we have a choice between a semi-smooth Newton (vinewtonssls) and an active set solver (vinewtonrsls). I couldn't get the semi-smooth Newton solver to work and I have no idea why.

Finally, we'll run the timestepping loop. Here we pass the bounds explicitly on each call to solve.

Movie time as always.

As some a posteriori sanity checking, we can evaluate how much the total water volume deviates.

Where a truly conservative scheme would exactly preserve the volume up to some small multiple of machine precision, we can only get global conservation up to the mesh resolution with our scheme. Instead, there are spurious "sources" at the free boundary. Likewise, there can be spurious sinks in the presence of ablation, so the sign error can go either way. This topic is covered in depth in this paper.

We can examine the fluxes after the fact in order to see where the value of $\varepsilon$ that we picked sits.

The fudge flux is 1/25 that of the average. This is quite a bit smaller, but not so much so that we should feel comfortable with this large a perturbation to the physics equations themselves. The ability to use it only in the derivative and not in the residual is a huge help.

To wrap things up, the overland flow equations are a perfect demonstration of how trivially equivalent forms of the same physical problem can yield vastly different discretizations. Writing the system as a single parabolic PDE might seem simplest, but there are several potential zeros in the denominator that require some form of regularization. By contrast, using a mixed form introduces more unknowns and a nonlinear equation for the flux, but there's wiggle room within that nonlinear equation. This makes it much easier to come up with a robust solution procedure, even if it includes a few uncomfortable hacks like using a different Jacobian from that of the true problem. Finally, while our discretization still works ok with no positivity constraint, PETSc has variational inequality solvers that make it possible to enforce positivity.

In the previous post, I showed how to integrate Hamiltonian systems

$$\begin{align} \dot q & = +\frac{\partial H}{\partial p} \\ \dot p & = -\frac{\partial H}{\partial q} \end{align}$$

using methods that approximately preserve the energy. Here I'd like to look at what happens when there are non-trivial constraints

$$g(q) = 0$$

on the configuration of the system. The simplest example is the pendulum problem, where the position $x$ of the pendulum is constrained to lie on the circle of radius $L$ centered at the origin. These constraints are easy to eliminate by instead working with the angle $\theta$. A more complicated example is a problem with rotational degrees of freedom, where the angular configuration $Q$ is a 3 $\times$ 3 matrix. The constraint comes from the fact that this matrix has to be orthogonal:

$$Q^\top Q = I.$$

We could play similar tricks to the case of the pendulum and use Euler angles, but these introduce other problems when used for numerics. For this or other more complex problems, we'll instead enforce the constraints using a Lagrange multiplier $\lambda$, and working with the constrained Hamiltonian

$$H' = H - \lambda\cdot g(q).$$

We're then left with a differential-algebraic equation:

$$\begin{align} \dot q & = +\frac{\partial H}{\partial p} \\ \dot p & = -\frac{\partial H}{\partial q} + \lambda\cdot\nabla g \\ 0 & = g(q). \end{align}$$

If you feel like I pulled this multiplier trick out of a hat, you might find it more illuminating to think back to the Lagrangian formulation of mechanics, which corresponds more directly with optimization via the stationary action principle. Alternatively, you can view the Hamiltonian above as the limit of

$$H_\epsilon' = H + \frac{|p_\lambda|^2}{2\epsilon} - \lambda\cdot g(q)$$

as $\epsilon \to 0$, where $p_\lambda$ is a momentum variable conjugate to $\lambda$. This zero-mass limit is a singular perturbation, so actually building a practical algorithm from this formulation is pretty awful, but it can be pretty helpful conceptually.

For now we'll assume that the Hamiltonian has the form

$$H = \frac{1}{2}p^*M^{-1}p + U(q)$$

for some mass matrix $M$ and potential energy $U$. The 2nd-order splitting scheme to solve Hamilton's equations of motion in the absence of any constraints are

$$\begin{align} p_{n + \frac{1}{2}} & = p_n - \frac{\delta t}{2}\nabla U(q_n) \\ q_{n + 1} & = q_n + \delta t\cdot M^{-1}p_{n + \frac{1}{2}} \\ p_{n + 1} & = p_{n + \frac{1}{2}} - \frac{\delta t}{2}\nabla U(q_{n + 1}). \end{align}$$

To enforce the constraints, we'll add some extra steps where we project back onto the surface or, in the case of the momenta, onto its cotangent space. In the first stage, we solve the system

$$\begin{align} p_{n + \frac{1}{2}} & = p_n - \frac{\delta t}{2}\left\{\nabla U(q_n) - \lambda_{n + 1}\cdot \nabla g(q_n)\right\} \\ q_{n + 1} & = q_n - \delta t\cdot M^{-1}p_{n + \frac{1}{2}} \\ 0 & = g(q_{n + 1}). \end{align}$$

If we substitute the formula for $p_{n + 1/2}$ into the second equation and then substitute the resulting formula for $q_{n + 1}$ into the constraint $0 = g(q_{n + 1})$, we get a nonlinear system of equations for the new Lagrange multiplier $\lambda_{n + 1}$ purely in terms of the current positions and momenta. Having solved this nonlinear system, we can then substitute the value of $\lambda_{n + 1}$ to obtain the values of $p_{n + 1/2}$ and $q_{n + 1}$. Next, we compute the momentum at step $n + 1$, but subject to the constraint that it has to lie in the cotangent space of the surface:

$$\begin{align} p_{n + 1} & = p_{n + \frac{1}{2}} - \frac{\delta t}{2}\left\{\nabla U(q_{n + 1}) - \mu_{n + 1}\cdot \nabla g(q_{n + 1})\right\} \\ 0 & = \nabla g(q_{n + 1})\cdot M^{-1}p_{n + 1}. \end{align}$$

Once again, we can substitute the first equation into the second to obtain a linear system for the momentum-space multiplier $\mu$. Having solved for $\mu$, we can then back-substitute into the first equation to get $p_{n + 1}$. This is the RATTLE algorithm. (I'm pulling heavily from chapter 7 of Leimkuhler and Reich here if you want to see a comparison with other methods and proofs that it's symplectic.)

Surfaces¶

Next we have to pick an example problem to work on. To start out, we'll assume that the potential energy for the problem is 0 and focus solely on the free motion of a particle on some interesting surface. The simplest surface we could look at is the sphere:

$$g(x, y, z) = x^2 + y^2 + z^2 - R^2$$

or the torus:

$$g(x, y, z) = \left(\sqrt{x^2 + y^2} - R\right)^2 + z^2 - r^2.$$

Just for kicks, I'd like to instead look at motion on surfaces of genus 2 or higher. There are simple parametric equations for tracing out spheres and tori in terms of the trigonometric functions, so the machinery of explicitly enforcing constraints isn't really necessary. There is no such direct parameterization for higher-genus surfaces, so we'll actually need to be clever in defining the surface and in simulating motion on it. As an added bonus, the ability to trace out curves on the surface will give us a nice way of visualizing it.

To come up with an implicit equation for a higher-genus surface, we'll start with an implicit equation for a 2D curve and inflate it into 3D. For example, the equation for the torus that we defined above is obtained by inflating the implicit equation $\sqrt{x^2 + y^2} - R = 0$ for the circle in 2D. What we want to generate higher-genus surfaces is a lemniscate. An ellipse is defined as the set of points such that the sum of the distances to two foci is constant. Likewise, a lemniscate is defined as the set of points such that the product of the distances to two or more foci is constant. The Bernoulli leminscate is one such example, which traces out a figure-8 in 2D. The Bernoulli leminscate is the zero set of the polynomial

$$f(x, y) = (x^2 + y^2)^2 - a^2(x^2 - y^2)$$

and it also has the parametric equation

$$\begin{align} x & = a\frac{\sin t}{1 + \cos^2t} \\ y & = a\frac{\sin t\cdot\cos t}{1 + \cos^2t} \end{align}$$

which gives us a simple way to visualize what we're starting with.

We've loosely referred to the idea of inflating the zero-contour of a function $f(x, y)$ into 3D. The 3D function defining the desired implicit surface is

$$g(x, y, z) = f(x, y)^2 + z^2 - r^2,$$

where $r$ is a free parameter that we'll have to tune. I'm going to guess that $r < \frac{a}{\sqrt{2}}$ but it could be much less; beyond that we'll have to figure out what $r$ is by trial and error.

The code below uses the sympy software package to create a symbolic representation of the function $g$ defining our surface. Having a symbolic expression for $g$ allows us to evaluate it and its derivatives, but to actually visualize the surface we'll have to sample points on it somehow.

Symbolically evaluating $g$ every time is expensive, so the code below uses the lambdify function from sympy to convert our symbolic expression into an ordinary Python function. I've added some additional wrappers so that we can pass in a numpy array of coordinates rather than the $x$, $y$, and $z$ values as separate arguments.

One of the first algorithms for constrained mechanical systems was called SHAKE, so naturally some clever bastard had to make one called RATTLE and there's probably a ROLL out there too. The code below implements the RATTLE algorithm. You can view this as analogous to the Stormer-Verlet method, which does a half-step of the momentum solve, a full step of the position solve, and finally another half-step of the momentum solve again. In the RATTLE algorithm, we have to exercise a bit of foresight in the initial momentum half-step and position full-step in order to calculate a Lagrange multiplier to project an arbitrary position back onto the zero-contour of $g$. Solving for the position multiplier is a true nonlinear equation, whereas the final momentum half-step is just a linear equation for the momentum and its multiplier, which we've written here as $\mu$. Here we only have one constraint, so each multiplier is a scalar, which is a convenient simplification.

I'll add that this algorithm was exceedingly fiddly to implement and I had to debug about 5 or 6 times before I got it right. The sanity checking shown below was essential to making sure it was right.

As a sanity check, we'll evaluate the change in energy throughout the course of the simulation relative to the mean kinetic energy. The relative differences are on the order of 1%, which suggests that the method is doing a pretty good job. I re-ran this notebook with half the timestep and the energy deviation is cut by a factor of four, indicative of second-order convergence.

Finally, let's make a movie of the results.

My Riemannian geometry kung fu is weak is but I think that the geodesic flow on this surface is ergodic (see these notes).

It's also interesting to have a look at what the respective Lagrange multipliers for position and velocity are doing.

Note how the Lagrange multipliers aren't smooth -- they have pretty sharp transitions. If you think of the Lagrange multipliers as fictitious "forces" that push the trajectories back onto the constraint manifold, then their amplitude is probably some kind of indicator of the local curvature of the constraint surface.

More interesting now¶

This all worked well enough for a single particle on the surface. Now let's see what happens if we put several particles on the surface and make them interact. I'd like to find some potential that's repulsive at distances shorter than equilibrium, attractive at longer distances, and falls off to zero at infinity. We could use the Lennard-Jones potential shown in the last demo but the singularity at the origin is going to create more difficulty than necessary. Instead, I'll use a variant of the Ricker wavelet, which is plotted below.