The back-and-forth method for Wasserstein gradient flows

This page illustrates the method proposed in (Jacobs et al., 2021). The back-and-forth method was initially introduced to compute optimal transport maps (Jacobs & Léger, 2020), and it can be generalized to compute Wasserstein gradient flows, via the JKO scheme (Jordan et al., 1998). This allows to efficiently solve PDEs of the type

for a large class of interesting energies $U(\rho)$. We consider the following energies:

• (The porous medium equation) Slow diffusion energies $$ U(\rho) = \int_{\mathbb{R}^d} u_m(\rho(x)) + V(x)\,dx, $$ where $u_m(\rho) = \gamma \rho(x)^m$ for $m>1$ and $\gamma>0$. Here $V\colon\mathbb{R}^d\to[0,\infty]$ is a potential function.
• (The incompressible limit) Congestion energies $$ U(\rho) = \int_{\mathbb{R}^d} u_{\infty}(\rho(x)) + V(x)\,dx, $$ where $u_\infty(\rho)=0$ if $0\le \rho \le 1$ and $\infty$ otherwise.
• (An aggregation-diffusion equation) Convex-concave energies $$ U(\rho) = \int_{\mathbb{R}^d} u_m(\rho(x))\,dx + \iint_{\mathbb{R}^d\times \mathbb{R}^d} \lvert x-y\rvert^2\rho(x)\rho(y)\,dxdy. $$ Note that the first term in convex in $\rho$ while the second term is concave in $\rho$.

Code

The code used in the paper is available on GitHub. The documentation is available here.

Porous medium equation

The Wasserstein gradient flow of $U(\rho)=\int_{\mathbb{R}^d} u_m(\rho(x)) + V(x)\,dx$ is the advection slow diffusion equation \(\partial_t\rho +\mathrm{div}(\rho (-\nabla V))= \gamma\Delta \rho^m,\) together with Neumann boundary conditions and an initial condition $\rho(0,\cdot) = \rho_0$.

Incompressible limit

Aggregation-diffusion equations

Consider the energy \(U(\rho) = \int_{\mathbb{R}^d} u_m(\rho(x))\,dx + \iint_{\mathbb{R}^d\times \mathbb{R}^d} \lVert x-y\rVert^2\rho(x)\rho(y)\,dxdy.\) The second term encourages aggregation and is concave with respect to $\rho$.