Primal Dual Methods for Wasserstein Gradient Flows

Abstract

Combining the classical theory of optimal transport with modern operator splitting techniques, we develop a new numerical method for nonlinear, nonlocal partial differential equations, arising in models of porous media, materials science, and biological swarming. Our method proceeds as follows: first, we discretize in time, either via the classical JKO scheme or via a novel Crank–Nicolson-type method we introduce. Next, we use the Benamou–Brenier dynamical characterization of the Wasserstein distance to reduce computing the solution of the discrete time equations to solving fully discrete minimization problems, with strictly convex objective functions and linear constraints. Third, we compute the minimizers by applying a recently introduced, provably convergent primal dual splitting scheme for three operators (Yan in J Sci Comput 1–20, 2018). By leveraging the PDEs’ underlying variational structure, our method overcomes stability issues present in previous numerical work built on explicit time discretizations, which suffer due to the equations’ strong nonlinearities and degeneracies. Our method is also naturally positivity and mass preserving and, in the case of the JKO scheme, energy decreasing. We prove that minimizers of the fully discrete problem converge to minimizers of the spatially continuous, discrete time problem as the spatial discretization is refined. We conclude with simulations of nonlinear PDEs and Wasserstein geodesics in one and two dimensions that illustrate the key properties of our approach, including higher-order convergence our novel Crank–Nicolson-type method, when compared to the classical JKO method.

A variational finite volume scheme for Wasserstein gradient flows

“…The approximation of the system ( 12) is a natural strategy to approximate the solution to (1). This approach was for instance at the basis of the works [7,21]. These methods require a sub-time stepping to solve system (12) on each interval (t n−1 , t n ), yielding a possibly important computational cost.…”

“…set on , complemented with homogeneous Neumann boundary condition ∇φ n τ •n = 0 on ∂ . We can interpret (21) as the one step resolvent of the mean-field game (12). Both the forward in time continuity equation and the backward in time HJ equation are discretized thanks to one step of backward Euler scheme.…”

“…Indeed, it merely demands to approximate two functions ρ n τ , φ n τ rather than time depending trajectories in function space as for the JKO scheme (12). This allows in particular to avoid inner time stepping as in [7,21], making our approach much more tractable to solve complex problems.…”

“…Moreover, they naturally transpose to the discrete setting the monotonicity properties of the continuous operators. Monotonicity was crucial in the derivation of the optimality conditions (21), as it will also be the case in the fully discrete framework later on. This led us to use upstream mobilities in the definition of the discrete counterpart of the squaredḢ 1 ρ norm.…”

Computation of optimal transport with finite volumes

“…Later in [26] Papadakis, Peyré and Oudet introduced a finite difference discretization using staggered grids, which are better suited for the discretization of the continuity equation. Similar finite differences approaches have been used also in more recent works [9,23]. Note that the use of regular grids can be beneficial for the efficient solution of the scheme, but is not adapated to complex domains.…”