Wonjun Lee

Ph.D. Candidate
Department of Mathematics
University of California, Los Angeles
Office : Virtual
profile photo

Wasserstein Gradient Flows

PDF Code

This page briefly describes the paper, The back-and-forth method for Wasserstein gradient flows [1]. To see the details, check the PDF and Code links above.

In this project, we proposed a new algorithm to compute the following class of PDEs (also known as Darcy's law):

$$\begin{aligned} \partial_t \rho - \nabla \cdot (\rho \nabla \phi)=0&\\ \phi = \delta U(\rho)& \end{aligned}$$

where \(\rho\) is a mass density, \(\phi\) is a pressure, and \(\delta U(\rho)\) is the Frechet derivative of \(U\) at \(\rho\). This class of PDEs models various physical phenomena such as fluid flow, heat transfer, aggregation-diffusion, and crowd motion. For example, if

$$U(\rho) = \frac{1}{m-1} \int \rho^m \, dx,$$

then the PDE becomes the porous medium equation,

$$\partial_t \rho - \Delta (\rho^m) = 0.$$

We solve this PDE via the JKO scheme [2], a discrete-in-time variational formulation of the PDE. Given \(\rho^{(0)}\), the scheme iterates

$$\rho^{(n+1)} = \text{arg}\min_{\rho}\, U(\rho) + \frac{1}{2\tau} W_2(\rho, \rho^{(n)})^2$$

where \(\tau\) is a time step size, \(W_2\) is a 2-Wasserstein distance, and \(\rho^{(n)}\) is an approximate solution \(\rho(n\tau, \cdot)\) of the PDE. We solve developed an algorithm to solve the JKO scheme expanding upon the back-and-forth method (BFM) [3], a fast numerical method for optimal transport problems.

Below are some of our results solving difficult PDE problems including porous medium equations, incompressible flows, and aggregation-diffusion equations.


Porous medium equations with an obstacle and potential

$$ \begin{aligned} U(\rho) &= \int \frac{1}{3} \rho^4(x) + V(x)\,dx \\ V(x) &= \|x - a\|^2,\quad a=(0.9,0.9) . \end{aligned} $$

Aggregation-diffusion equations

$$ U(\rho) = \int \rho^2(x)\,dx + \int \int \rho(x) \|x-y\|^2 \rho(y) \, dx\,dy. $$

Incompressible flows (crowd motion) with an obstacle and potential

$$ \begin{aligned} U(\rho) &= \int u_\infty (\rho(x)) + V(x)\,dx \\ V(x) &= \|x - a\|^2,\quad a=(0.9,0.9) \\ \end{aligned} $$

$$u_\infty(t) = \begin{cases} 0 & \text{if } 0 \leq t \leq 1\\ \infty & \text{otherwise.} \end{cases}$$


[1] Matt Jacobs, Wonjun Lee and Flavien Léger. The back-and-forth method for Wasserstein gradient flows. 2020.

[2] Richard Jordan and David Kinderlehrer and Felix Otto. The variational formulation of the Fokker-Planck equation. SIAM journal on mathematical analysis 29.1 (1998): 1–17.

[3] Matt Jacobs and Flavien Léger. A fast approach to optimal transport: The back-and-forth method. Numerische Mathematik (2020): 1-32.