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:{\mathbb{R}}^d\to [0,\mathrm{\infty }]$ is a potential function.
• (The incompressible limit) Congestion energies $$U(\rho )={\int }_{{\mathbb{R}}^d}{u}_{\mathrm{\infty }}(\rho (x))+V(x)dx,$$ where $u_{\mathrm{\infty }}(\rho )=0$ if $0\le \rho \le 1$ and $\mathrm{\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}|x-y{|}^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 (-\mathrm{\nabla }V))=\gamma \mathrm{\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}\Vert x-y{\Vert }^2\rho (x)\rho (y)dxdy.$ The second term encourages aggregation and is concave with respect to $\rho$.