PaperlayerView exact revision on arXiv

Source-linked AI summary

Primal dual methods for Wasserstein gradient flows

Jose A. Carrillo, Katy Craig, Li Wang, Chaozhen Wei

arXiv:1901.08081v2math.NAcs.NA

TL;DR

Existing numerical approaches to nonlinear, nonlocal PDEs face stability constraints or costly implicit solves. This paper combines Wasserstein time stepping, a dynamic transport formulation, and primal-dual splitting; it proves convergence of fully discrete minimizers to the continuum discrete-time scheme and illustrates the approach in simulations.

  • Problem

    Explicit PDE discretizations face stability constraints from degenerate diffusion or nonlocal interactions, while implicit methods can be computationally intensive.

    Source: §1.2.1 · p. 4

  • Method

    The method uses a dynamic Wasserstein formulation and primal-dual splitting to solve fully discrete minimization problems.

    Source: §1.2.3 · p. 7, §2.3 · p. 13

  • Results

    Under the stated assumptions, fully discrete minimizers converge, up to a subsequence, to a solution of the continuum JKO scheme.

    Source: §1.1 · p. 4

  • Takeaways & Limitations

    Allowing the PDE constraint to hold up to an error δ > 0 permits fewer iterations without reducing accuracy and makes discrete minimizers exist.

    Source: §1.2.3 · p. 7

    With exact PDE-constraint enforcement, minimizers may fail to exist when the pointwise approximation and target measure have different masses.

    Source: §3.1 · p. 18

Abstract

from arXiv · show

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 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.

1 Introduction

The paper develops a Wasserstein-gradient-flow method for nonlinear, nonlocal PDEs, combining time discretization, a dynamic optimal-transport formulation, and primal-dual optimization. It situates the method against existing numerical approaches and reports convergence of fully discrete minimizers to the continuum JKO scheme under stated assumptions.

1 Introduction

  • Wasserstein gradient flows extend gradient-flow modeling beyond Hilbert-space L2 methods to nonlinear, nonlocal PDEs whose dynamics vary with interaction and diffusion choices.The equations include settings with behavior ranging from globally well-posed bounded solutions to finite-time blow-up.Source: §1 · pp. 2–3

1.1 Summary of numerical approach

  • The method discretizes time by JKO or a Crank–Nicolson-inspired scheme, reformulates the Wasserstein step dynamically, then discretizes space and time and solves the constrained problem by three-operator splitting.It applies with diffusion, drift, or interaction terms, including cases without diffusion.Source: §1.1 · p. 4
  • The convergence result concerns fully discrete minimizers approaching the continuum discrete-time scheme, while simulations cover Wasserstein geodesics and nonlinear PDEs in one and two dimensions.The convergence statement is conditional on the stated regularity assumptions.Source: §1.1 · p. 4

1.2.1 Classical numerical PDE methods

  • Explicit PDE discretizations face stability constraints from degenerate diffusion or nonlocal interactions, while implicit methods can require costly matrix inversion.The proposed optimization formulation is unconditionally stable for JKO and avoids full-matrix inversion; its gridpoint subproblems are independently solvable and parallelizable.Source: §1.2.1 · pp. 4–5

1.2.2 Variational methods

  • JKO provides Wasserstein time stepping, while the new Crank–Nicolson-inspired scheme shows preliminary higher-order convergence but is not energy decreasing.The authors leave a thorough analysis of its convergence rate as the outer time step tends to zero to future work.Source: §1.2.2 · p. 5

1.2.3 Numerical methods for the Wasserstein distance

  • The method uses Benamou–Brenier’s dynamic formulation and a primal-dual splitting scheme to solve fully discrete Wasserstein minimization problems.Its relaxed linear PDE constraint permits fewer iterations without reducing accuracy and supports the convergence proof; unlike exact enforcement, the relaxed problem has minimizers for δ > 0.Source: §1.2.3 · pp. 6–7

1.3 Contribution

  • The JKO-based method is unconditionally stable without expensive matrix inversion, and the Crank–Nicolson variant performs better in observed outer-step convergence rates.Using a linear inequality constraint rather than equality also reduces iterations relative to related Wasserstein-geodesic algorithms; fully discrete solutions converge to the JKO scheme as spatial and inner-time discretizations vanish.Source: §1.2.1 · p. 5, §1.3 · p. 7

2 Numerical Method

This section reformulates JKO and Wasserstein-geodesic computation as constrained dynamic minimization problems, then discretizes and solves the resulting convex formulations using primal-dual splitting. It specifies solvers for fully discrete JKO problems, including exact proximal operators and convergence guarantees under stated conditions.

2.1 Dynamic formulation of JKO scheme

  • The Benamou–Brenier reformulation yields a strictly convex functional with a linear PDE constraint, unique minimizers, and Wasserstein geodesics [7].The resulting method also computes geodesics and distances, with potential extensions to other Benamou–Brenier metrics and non-Euclidean manifolds.Source: §2.1 · pp. 7–9
  • The classical JKO scheme is first-order, whereas the Crank–Nicolson-inspired scheme shows higher-order convergence in practice but lacks JKO’s energy-decay property.The authors introduce it because JKO discretization error dominates other numerical errors in their simulations.Source: §2.1 · pp. 9–10

2.2 Fully discrete JKO

  • The Crank–Nicolson-inspired problem is rewritten in the same form as the dynamic JKO problem using a modified energy, allowing both schemes to be discretized together.The passage establishes the shared formulation but does not specify the modified energy’s form.Source: §2.2 · p. 10

2.2.1 Discretization of functions and domain

  • Piecewise-constant approximations on a regular space-time grid converge uniformly to the underlying functions when those functions are continuous.The discretization uses spatial and temporal grid spacings.Source: §2.2.1 · p. 10

2.2.2 Discretization of energy functionals

  • The energy discretizations use piecewise-constant approximations of the potential and symmetric interaction terms, and also discretize energies for the Crank–Nicolson scheme and Wasserstein geodesics.The passages identify these discretization targets but do not provide their full formulas.Source: §2.2.2 · p. 11

2.2.3 Discretization of derivative operators

  • Finite-difference derivative approximations can use centered spatial differences with forward Euler or Crank–Nicolson time stepping, with zero normal flux imposed at boundaries.Because these derivatives are approximate, the fully discrete scheme relaxes the continuum equality constraints.Source: §2.2.3 · p. 11

2.2.4 Discrete dynamic JKO

  • The fully discrete JKO scheme preserves nonnegativity and energy decrease; relaxing PDE constraints accelerates minimization without accuracy loss, while separate mass constraints can improve results for rough solutions.Minimizers exist under sufficient hypotheses on the discrete energy and initial data.Source: §2.2.4 · p. 12

2.3 Primal dual algorithms for fully discrete JKO

  • Chambolle–Pock splitting handles two-term objectives, while Yan’s extension handles three-term objectives by keeping the smooth energy separate.Sparse operators give low per-iteration complexity, and the method can be iterated to approximate JKO sequences and Wasserstein gradient flows [63,107].Source: §2.3 · p. 13
  • Exact proximal-operator computation makes the primal-dual method positivity preserving, rather than positivity preserving only as Δx, Δt → 0.The positive root needed for the operator can be computed using the general cubic solution formula instead of Newton iteration [96].Source: §2.3 · p. 15
  • The method avoids inverting AAT, enabling efficiency in high dimensions through the sparsity of A, unlike prior approaches [7,96].Its indicator-function definition differs from previous work; the cited earth mover’s distance approach [82] omits the Wasserstein time dimension.Source: §2.3 · p. 15
  • Algorithms 1 and 2 converge to minimizers as iterations increase and stopping tolerances vanish, under their respective step-size and problem assumptions.Algorithm 1 requires σλ < 1/λmax(AAt) and an existing minimizer [53]; Algorithm 2 additionally requires a proper, lower semi-continuous, convex energy, co-coercivity of its gradient, σλ < 1/λmax(ÃÃt), λ < 2β, and an existing minimizer [107].Source: §2.3 · pp. 15–16
  • For the stated energy form, Algorithm 2’s co-coercivity condition reduces to bounded U′′(ρ) and W, achievable with bounded solutions and a nonsingular interaction kernel.For singular kernels such as the Newtonian potential, the paper approximates W by convolution with a mollifier.Source: §2.3 · p. 16

3 Convergence

This section establishes convergence of fully discrete JKO minimizers to continuum minimizers, stating the required hypotheses and proving existence for sufficiently fine discretizations. Under (H1)–(H6), the discrete dynamic JKO problems Γ-converge to the continuum problem, and subsequences of discrete minimizers converge to continuum minimizers; the convergence proof assumes a smooth, positive continuum minimizer.

3 Convergence

  • The section aims to show that fully discrete JKO minimizers have subsequences converging to continuum JKO minimizers, under the stated smoothness and positivity assumption.It first establishes discrete minimizer existence before proving subsequential convergence.Source: §3 · p. 16

3.1 Hypotheses

  • The result assumes a hyperrectangle and sufficiently regular continuous potentials, though the domain result extends to Lipschitz domains with carefully defined discrete normals and no-flux conditions.Discontinuous interaction potentials can be approximated by regularizations under conditions that ensure convergence of gradient flows and minimizers.Source: §3.1 · p. 18
  • The relaxation condition lim_(h→0)(∆x + ∆t)/δ5 = 0 is nearly optimal: δ5 cannot decay faster than O(∆x + ∆t), or discrete minimizers may fail to exist when approximate masses differ.Poor choices of relaxation parameters can also lengthen iterations without improving accuracy.Source: §3.1 · p. 18
  • The convergence proof relies on a smooth, strictly positive continuum minimizer; conditions ensuring this regularity and positivity at the JKO-scheme level remain largely open.In practice, the method performs well on nonsmooth or nonpositive problems when spatial and temporal discretizations are sufficiently small.Source: §3.1 · pp. 18–19

3.2 Existence of minimizers

  • Under (H1)–(H6), minimizers of the fully discrete JKO problem exist for all sufficiently small h, and any minimizer is nonnegative.The objective is finite only for nonnegative densities.Source: §3.1 · p. 17, §3.2 · p. 21

3.3 Convergence of minimizers

  • Bounded finite-action discrete admissible pairs have subsequences whose densities converge weak-* to a probability measure satisfying the limiting PDE constraint and initial data.The discrete initial densities converge to ρ0 in L2(Ω), and the limiting terminal measure is the endpoint density.Source: §3.3 · pp. 22, 24
  • For weak-* convergent piecewise-constant densities, the discrete energies satisfy lim inf as h → 0 of Eh(ρh) ≥ E(ρ).This lower-semicontinuity result holds under (H1)–(H6).Source: §3.3 · p. 23
  • The discrete problems Γ-converge to the continuum problem under (H1)–(H6), with a recovery sequence for smooth, positive admissible pairs.The recovery-sequence result assumes ρ ∈ C2([0, 1]; C1(Ω)), ρ > 0, and m ∈ C([0, 1]; C2(Ω)).Source: §3.3 · p. 23
  • Under (H1)–(H6), every sequence of discrete minimizers has a subsequence converging weak-* to a minimizer of the continuum problem.If the continuum minimizer is unique, the full sequence converges to it.Source: §3.3 · pp. 24–25

4 Numerical results

The numerical section tests the algorithms on Wasserstein geodesics and nonlinear gradient flows, using a square computational domain with no-flux boundary conditions.

4 Numerical results

  • The examples apply algorithm 1 to Wasserstein geodesics and algorithm 3 to nonlinear gradient flows.Unless specified otherwise, they use Crank-Nicolson finite differences, stopping tolerances of 10^-5, and specified relaxation parameters.Source: §4 · p. 25
  • The computational domain is a square, with no-flux boundary conditions imposed dimension by dimension on m.Source: §4 · p. 25
  • Unless otherwise specified, the constraint relaxation uses δ1 = δ2 = δ4 = δ5 = δ, with δ3 set separately for each example.Source: §4 · p. 25

4.1 Wasserstein geodesics

The geodesic examples show that relaxation scaling affects iteration cost without changing accuracy in the tested Gaussian case, while nonsmooth endpoints exhibit numerical smoothing. The examples also illustrate geodesic computation and convergence behavior across one- and two-dimensional settings.

4.1 Wasserstein geodesics

  • The Gaussian example computes a one-dimensional Wasserstein geodesic between rescaled Gaussian measures and plots its evolution and numerical error against iteration count.Source: §4.1 · p. 26, Figure 2 · p. 26
  • The Gaussian geodesic example reports an error-convergence rate of O(1/l) for densities, momenta, and Wasserstein distance.Source: §4.1 · pp. 26–27
  • The relaxation-scaling experiment contrasts δ = (∆x)^2, δ = (∆x)^3, and δ = 10^-8, with δ = (∆x)^2 expected to balance accuracy and computational efficiency for the Crank-Nicolson constraint approximation.Source: §4.1 · p. 27, Figure 3 · p. 27
  • For δ = (∆x)^2, the PDE constraint meets its stopping criterion by 10^4 iterations and dual variables by 10^5; δ = (∆x)^3 takes over 10^5 iterations and δ = 10^-8 over 10^7.All three choices reach identical error, so the slower convergence adds iterations without improving accuracy.Source: §4.1 · p. 27
  • For nonsmooth, non-strictly-positive measures, the parliamentary geodesic exhibits intermediate-time degradation from numerical smoothing rather than the exact constant-speed translation.The Pac-Man and ghost example likewise shows smoothing around discontinuous edges; the authors report that the method appears not to converge without sufficient smoothness.Source: §4.1 · p. 28, Figure 4 · p. 28

4.2 Wasserstein gradient flows: one dimension

The one-dimensional simulations test the method on porous-medium, nonlinear Fokker–Planck, and aggregation equations. They show higher-order convergence for the Crank–Nicolson-inspired scheme and recovery of reported equilibrium behavior.

4.2.1 Porous medium equation

  • For the porous-medium equation, the Crank–Nicolson-inspired scheme converges at higher order than JKO on compactly supported data, while smooth positive data give first-order JKO accuracy and less-than-second-order accuracy for the new scheme.The authors attribute the slower rates for compactly supported data to the solutions' lower, Hölder regularity.Source: §4.2.1 · p. 30, Figure 7 · p. 30, Figure 8 · p. 31

4.2.2 Nonlinear Fokker-Planck equation

  • With quadratic confinement and m = 2, the nonlinear Fokker–Planck density approaches equilibrium with exponential energy decay; the new scheme has higher-order convergence, versus slightly faster than first order for JKO.The authors attribute the improved rate over the porous-medium example to rapid convergence to equilibrium.Source: Figure 9 · p. 31, §4.2.2 · p. 32, Figure 10 · p. 32

4.2.3 Aggregation equation

  • For the aggregation equation, the simulated density approaches its equilibrium profile with exponential energy decay, at a slightly slower numerical rate than the reported analytical result.For the discontinuous interaction potential, the discretization averages the kernel at the origin and adds artificial diffusion to avoid possible boundary overshoot.Source: §4.2.3 · p. 33, Figure 11 · p. 33

4.3 Wasserstein gradient flows: two dimensions

Two-dimensional simulations apply the method to aggregation, aggregation-drift, and aggregation-diffusion equations. They recover several reported steady-state profiles and illustrate how interaction, diffusion, and the diffusion exponent shape the observed evolution.

4.3.1 Aggregation equation

  • For the two-dimensional aggregation equation, one interaction potential yields a stable Dirac ring of radius 0.5, while Newtonian repulsion yields a characteristic function on the unit disk.Both simulations recover previously reported analytical steady-state results.Source: §4.3.1 · pp. 33–34, Figure 12 · p. 34, Figure 13 · p. 34

4.3.3 Aggregation diffusion equation

  • In aggregation-diffusion, medium-range aggregation forms four bumps that ultimately approach a single-bump equilibrium; for Keller–Segel, m = 2 gives a single bump while m = 1 produces a sharply peaked blowup profile by t = 2.A separate m = 2 Keller–Segel simulation from three localized bumps shows quasi-stationary stages before merging into one bump in the long-time limit.Source: §4.3.3 · pp. 35–36, Figure 16 · p. 37, Figure 17 · p. 37

4.3.2 Aggregation drift equation

A Further details of numerical implementation

This section specifies the one-dimensional discrete system used by Algorithms 1–3, including its matrices and vectors. It identifies how the constraint vector collects initial and boundary conditions.

A Further details of numerical implementation

  • The section gives explicit formulas for A, u, and b used in Algorithms 1–3 for the one-dimensional discretization (28).The corresponding tilde A and tilde b use a slightly different treatment of ρ at final time.Source: A Further details of numerical… · p. 37
  • N=(Nx+1)(Nt+1), and the system matrices use Kronecker products, including Amass=I_(Nt+1)⊗Sρ.The section defines the identity matrix and the all-equal column vector notation used in these constructions.Source: A Further details of numerical… · pp. 38–39
  • The vector b collects the initial conditions for ρ and the boundary conditions for m.Its displayed construction includes initial-density and boundary-value terms; b1=0.Source: A Further details of numerical… · p. 39
Loading 1901.08081v2…