Gradient flow methods are classical techniques for the analysis and numerical simulation of partial differential equations. Historically, such methods were exclusively based on gradient flows arising from a Hilbert space structure, particularly L 2 (R d ), but since the work of Jordan, Kinderlehrer, and Otto in the late 90's [75,93,94], interest has emerged in a range of nonlinear, nonlocal partial differential equations that are gradient flows in the Wasserstein metric,
When Ω = R d , we consider no-flux boundary conditions. Equations of this form arise in a number of physical and biological applications, including models in granular media [12,45,46,102], material science [71], and biological swarming [6,39,77]. Furthermore, many well-known equations may be written in this way: when V = W = 0 and α = 1, Eq. ( 1) reduces to the heat equation (m = 1), porous medium equation (m > 1), and fast diffusion equation (m < 1) [103]. In the presence of a drift potential V , it becomes a Fokker-Planck equation (m = 1) or nonlinear Fokker-Planck equation (m > 1), as used in models of tumor growth [96,100]. When the interaction potential W is given by a repulsive-attractive Morse or power-law potential, W (x) = -C a e -|x|/l a + C r e -|x|/l r , C r /C a < (l r /l a ) -d , 0 < l r < l a , 0 < C a < C r ,
we recover a range of nonlocal interaction models, which are repulsive at short length scales and attractive at long length scales [4,5,34,101]. When W = (Δ) -1 , the Newtonian potential, we have the Keller-Segel equation and its nonlinear diffusion variants [17,19,25,26,32,41,76]. Finally, as the diffusion exponent m → +∞, we recover congested aggregation and drift equations arising in models of pedestrian crowd dynamics and shape optimization problems [23,58,67,84,90,91].
In order to describe the gradient flow structure of equation ( 1), we begin by rewriting it as a continuity equation in ρ(x, t) for a velocity field v(x, t),
In this form, two key properties of the equation become evident: it is positivity preserving and conserves mass. In what follows, we will always consider nonnegative initial data, and we will typically renormalize so that the mass of the initial data equals one, i.e., ρ 0 ∈ P ac (Ω), where P ac (Ω) is the set of probability measures on Ω that are absolutely continuous with respect to Lebesgue measure. Furthermore, as our objective is to develop a numerical method for these equations, we will exclusively consider the case when Ω is a bounded domain. Throughout, we commit a mild abuse of notation and identify all such probability measures with their densities, dρ(x) = ρ(x)dx.
As discovered by Otto [93], given an energy E : P ac (Ω) → R ∪ {+∞}, we may formally define its gradient with respect to the Wasserstein metric d W using the formula
(See Sect. 2.1 for the definition of the Wasserstein metric d W .) In this way, gradient flows of E, ∂ t ρ = -∇ d W E(ρ), correspond to solutions of the continuity equation with velocity v = -∇ δE δρ . In particular, Eq. ( 3) is the gradient flow of the energy
Differentiating the energy (4) along solutions of (3), one formally obtains that the energy is decreasing along the gradient flow
which coincides with the theoretical interpretation of gradient flows as solutions that evolve in the direction of steepest descent of an energy, where the notion of steepest descent is induced by the Wasserstein metric structure.
A key feature of equations of the form (3) is the competition between repulsive and attractive effects. For repulsive-attractive interaction kernels W , as in equation ( 2), these effects can arise purely through nonlocal interactions, leading to rich structure of the steady states [4,13,14,34,65]. For purely attractive interaction kernels W , as in the Keller-Segel equation, the competition instead arises from the combination of nonlocal interaction with diffusion. In this case, different choices of interaction kernel W , diffusion exponent m, and initial data ρ 0 can lead to widely different behaviorfrom bounded solutions being globally well posed to smooth solutions blowing up in finite time [17,19,25,26,32,41].
The goal of the present work is to develop new numerical approach for partial differential equations of the form (1) that combine gradient flow methods with modern operator splitting techniques. Our approach applies to equations of this form with any combination of diffusion αU m (ρ) (α ≥ 0), drift V , or interaction W * ρ terms-in particular, it is not necessary for diffusion to be present in order for our scheme to converge. The main idea of our approach is to discretize the PDE/Wasserstein gradient flow at two levels. First, we consider a time discretization of the gradient flow with time step τ (see Fig. 1b), either given by the classical JKO scheme (Eq. ( 6) below) or a new Crank-Nicolson inspired variant (Eq. (7) below). This reduces computation of the gradient flow to solving a sequence of infinite-dimensional minimization problems. Then, we consider a dynamical reformulation of these minimization problems, stemming from Benamou and Brenier's dynamic characterization of the Wasserstein metric, by which the problem becomes the minimization of a strictly convex integral functional subject to a linear PDE constraint (see Fig. 1c). At this level, the problem remains continuous in space and time. We conclude by considering a further discretization of the problem, with inner time step (Δt) and spatial discretization (Δx), by taking piecewise constant approximations of the functions and using a finite difference approximation of the PDE constraint (see Fig. 1d). In this final, fully discrete form, we then compute the minimizer using modern operator splitting techniques, applying Yan's recent extension of the classical primal dual algorithm for minimizing sums of three convex functions [106].
Our paper is organized as follows. In Sect. 1.2, we discuss the relationship between our numerical approach and previous work. In Sect. 1.3, we summarize our contribution. In Sect. 2, we describe the details of our numerical method. Along with numerically simulating Wasserstein gradient flows, our method also provides, as a special case, a new method for computing Wasserstein geodesics and the Wasserstein distance between probability densities; see Remark 1. In Sect. 3, we prove that, provided a smooth, positive solution of the continuum JKO scheme exists and the energy corresponding to the PDE is sufficiently regular, then minimizers of the fully discrete problem exist (Theorem 1), the objective functions of the discrete problems Γconverge to the objective function of the continuum problem (Theorem 2), and thus, solutions of the fully discrete scheme converge, up to a subsequence, to a solution of the continuum scheme (Theorem 3). As a special case, we also recover convergence of a numerical method for computing Wasserstein geodesics, similar to that introduced by Papadakis, Péyre, and Oudet [95]. Finally, in Sect. 4, we provide several numerical simulations illustrating our approach in both one and two dimensions, computing Wasserstein geodesics, nonlinear Fokker-Planck equations, aggregation diffusion equations, and other related equations.
We now compare our approach to existing numerical methods. Perhaps the most common numerical approach for equations of the form (1) is to consider the equation as an advection-diffusion equation and apply classical finite difference, finite volume, or Galerkin discretizations [3,29,54,66,85]. However, when such methods are based on explicit time discretizations, they suffer from stability constraints due either to the degeneracy of the diffusion (when m > 1) or the nonlocality from the interaction potential W . (See for instance the mesa problem [83].) Implicit time discretizations, on the other hand, are computationally intensive, due to the difficulty of matrix inversion, even when the implicit steps are solved by smart iterative methods to avoid the high computation cost of convolution [3].
Another common approach is to leverage structural similarities between (3) and equations from fluid dynamics to develop particle methods [14,27,30,36,43,48,57,60,88,92]. Until recently, the key limitation of such methods has been developing approaches to incorporate diffusion. Following the analogy with the Navier-Stokes equations, stochastic particle methods have been proposed in the case of linear diffusion (m = 1) [72][73][74]86]. More recently the first two authors and Patacchini developed a deterministic blob method for linear and nonlinear diffusion (m ≥ 1) [31]. On the one hand, particle methods naturally conserve mass and positivity, and they can also be designed to respect the underlying gradient flow structure of the equation, including the energy dissipation property (5). On the other hand, a large number of particles are often required to resolve finer properties of solutions.
In contrast with such classical methods, our method introduces an auxiliary momentum variable m and an additional inner layer of time discretization, which enlarges the dimension of the problem. However, as later pointed out in [80], the inner layer of time can be discretized with just one step without violating the overall first-order accuracy, there completely eliminating the additional cost introduced by the inner layer. Another major advantage of our approach is that, by reforming the PDE problem into an optimization problem, we obtain unconditional stability (for the JKO discretization, see Eq. ( 6) below) while avoiding the inversion of a full matrix in the general implicit setting, which is extremely expensive, especially in higher dimensions; see for instance [3]. Finally, compared to other implicit methods, such as the backward Euler method, the suboptimization problems can be solved independently at each gridpoint, and therefore are massively parallelizable and suitable for high-dimensional problems.
Compared to the classical numerical PDE approaches described in the previous section, a more modern class of numerical methods leverages the gradient flow structure of (1) to approximate solutions of the PDE by solving a sequence of minimization problems. This is the approach we take in the present work. Originally introduced by Jordan, Kinderlehrer, and Otto as a technique for computing solutions of the Fokker-Planck equation (Eq. ( 1), W = 0, m = 1) [75], this scheme approximates the solution ρ(x, t) at time t by solving the following sequence of n minimization problems with time step τ = t/n,
The JKO scheme is precisely the analogue of the implicit Euler method in the infinitedimensional Wasserstein space. The constraint ρ ∈ P ac (Ω) ensures that the method is positivity and mass preserving, and the fact that d 2 W (ρ, ρ n ) ≥ 0 ensures the energy decreasing along the scheme, E(ρ n+1 τ ) ≤ E(ρ n τ ). Under sufficient assumptions on the underlying domain Ω, drift potential V , interaction potential W , and initial data ρ 0 (see Sect. 2.1), the solution of the JKO scheme ρ n τ converges to the solution ρ(x, t) of the partial differential equation ( 1), with a first-order rate in terms of the time step τ = t/n [2, Theorem 4.0.4],
In our numerical simulations, we observe that this discretization error dominates other errors in our numerical method; see Sects. 4.2.1 and 4.2.2. Consequently, we also introduce a new time discretization, in analogy with the Crank-Nicolson method
The connection between the above scheme and the classical Crank-Nicolson discretization can be seen by considering the optimality conditions for (7):
Like the JKO scheme, our Crank-Nicolson inspired method is also positivity and masspreserving, though it is not energy decreasing. In Figs. 7, 8, and 10 of our numerics section, we conduct a preliminary analysis of the rate of convergence of this method, which verifies that it is indeed higher order than the JKO scheme. As the goal of the present work is primarily the development of fully discrete numerical schemes, we leave a thorough analysis of the rate of convergence of our Crank-Nicolson inspired method as τ → 0 to future work. On the one hand, our Crank-Nicolson inspired method (7) is not the first higher-order method proposed for metric space gradient flows: Matthes and Plazotta developed a provably second-order scheme for general metric space gradient flows by generalizing the backward differentiation formula [89]. The Matthes-Plazotta method, however, requires two evaluations of the Wasserstein distance at each outer time step and thus is less practical for our purpose of numerically computing gradient flows in higher dimensions. Another method was introduced by Legendre and Turinici [79] based on the midpoint method. This method can be reformulated as the classical JKO step with half time step followed by an extrapolation. This extrapolation step could be implemented by solving the corresponding continuity equation either explicitly or implicitly; however, solving the equation explicitly could potentially violate conservation of positivity, while solving it implicitly would require an additional matrix inversion. Another higher-order variational method was also proposed in [78], which resembles explicit Runge-Kutta methods and, again, require two or more evaluations of the Wasserstein distance at each outer time step.
To use either the classical JKO scheme (6) or our new Crank-Nicolson inspired scheme (7) as a basis for numerical simulations, one must first develop a fully discrete approximation of the minimization problem at each step of the scheme. Here, the main numerical difficulty arises in approximating the Wassserstein distance, and there are several different approaches for dealing with this term. First, one can reformulate the Wasserstein distance in terms of a Monge-Ampére equation with nonstandard boundary conditions [11,68], though difficulties arise due to the lack of a comparison principle [70]. Second, one can reframe the problem as a classical L 2 (R d ) gradient flow at the level of diffeomorphisms [16,37,47,49,69], but to pursue this approach, one has to overcome complications arising from the underlying geometry and the structure of the PDE system for the diffeomorphisms. Third, one can discretize the Wasserstein distance term as a finite-dimensional linear program, overcoming the lack of strict convexity of the objective function by adding a small amount of entropic regularization [8,55,61]. (For a detailed survey of computational optimal transport, we refer the reader to the recent book by Péyre and Cuturi for [97].)
A fourth approach for computing the Wasserstein distance, and the one which we develop in the present work, is to consider a dynamic formulation due to Benamou and Brenier [7]. This reframes the problem as a strictly convex optimization problem with linear PDE constraints, which can be discretized using Benamou and Brenier's original augmented Lagrangian method ALG2 or, more generally, a range of modern proximal splitting methods, as shown by Papadakis, Peyre, and Oudet [95]. (See also [21,22] for related work on mean field games.) Adding an additional Fisher information term in this dynamic formulation (in analogy with entropic regularization) has also been explored in [82].
Only recently have these above approaches for computing the Wasserstein distance been integrated with the JKO scheme (6) in order to simulate partial differential equations of the form (1). The Monge-Ampére approach extends naturally, though the presence of a diffusion term αU m (ρ) for α > 0 is required to enforce convexity constraints at the discrete level [10]. Similarly, entropic regularization (or the addition of a Fisher information term) vastly accelerates the computation of gradient flows, but at the level of the partial differential equation, this corresponds to introducing numerical diffusion, which may disrupt the delicate balance between aggregation and diffusion inherent in PDEs of this type [28,55,82]. Finally, Benamou and Brenier's dynamic reformulation of the Wasserstein distance has also been adopted in recent work to approximate gradient flows [9]. A key benefit of this latter approach when compared to entropic regularization is that it leads to an optimization problem in N d
x × N t variables, where N x and N t are the number of spatial and temporal gridpoints, whereas the latter leads to an optimization problem in N 2d
x variables. In the present work, we further develop this last approach, using Benamou and Brenier's dynamic reformulation of the Wasserstein distance to simulate Wasserstein gradient flows, via both the classical JKO scheme (6) and our new Crank-Nicolson inspired scheme (7). This leads to a sequence of minimization problems (Fig. 1C), which we discretize (Fig. 1D) and then solve using a modern primal dual three operator splitting scheme due to Yan [106], instead of the classical ALG2 method. See Sect. 2 for a detailed description of our approach.
Due to the fact that we use operator splitting methods to compute the minimizer in Benamou and Brenier's dynamic formulation of the Wasserstein distance, our work can be seen as an extension of previous work by Papadakis, Peyre, and Oudet [95], which applied similar two operator splitting schemes to simulate the Wasserstein distance. However, there are a few key differences between our approach and previous work. First, we are able to implement the primal dual splitting scheme in a manner that does not require matrix inversion of the finite difference operator, which reduces the computational cost. Second, we succeed in obtaining the exact expression for the proximal operator, which allows our method to be truly positivity preserving, while other similar methods are only positivity preserving in the limit as Δx, Δt → 0; see Remark 5. Third, instead of imposing the linear PDE constraint in Benamou and Brenier's dynamic reformulation exactly, via a finite difference approximation, we allow the linear PDE constraint to hold up to an error of order δ > 0, which can be tuned according to the spatial discretization (Δx), the inner temporal discretization (Δt), and the outer time step τ to respect the order of accuracy of the finite difference approximation; see Remark 3. Numerically, this allows our method to converge in fewer iterations, without any reduction in accuracy, as demonstrated in Fig. 3. From a theoretical perspective, the fact that we only require the PDE constraint to hold up to an error of order δ > 0 makes it possible to prove convergence of minimizers of the fully discrete problem to minimizers of the JKO scheme (6), since minimizers of the fully discrete problem always exist for δ > 0, which is not the case when the PDE constraint is enforced exactly (δ = 0); see Remark 8 and Theorem 1.
The main components of our numerical method for computing solutions to (1) are: (a) an outer time discretization, of either JKO (6) or Crank-Nicolson type (7) (Fig. 1B) (b) a dynamic interpretation of the Wasserstein distance (Fig. 1C), which when discretized via finite difference approximations leads to a sequence of constrained optimization problems (Fig. 1D) (c) an application of modern three operator splitting schemes for solving these optimization problems.
Our main contributions are:
-Unlike classical explicit methods, our JKO-type method is unconditionally stable. Unlike classical implicit methods, it achieves this stability without an expensive matrix inversion. -In practice, we observe that our Crank-Nicolson-type method performs even better than our JKO-type method, in terms of rate of convergence with respect to the outer time step (see Figs. 7,8, and 10). We leave a thorough analysis of the rate of this convergence of this method to future work. -By formulating our optimization problem with a linear inequality constraint instead of a linear equality constraint, our algorithm converges in fewer iterations when compared to related algorithms for Wasserstein geodesics; see Remark 3 and Fig. 3. -We prove convergence of our fully discrete method (Fig. 1D) to the JKO scheme (Fig. 1B,C) as the spatial discretization and inner time discretization go to zero.
As described in the previous section, our numerical method for computing the JKO scheme is based on the following dynamic reformulation of the Wasserstein distance due to Benamou and Brenier [7]:
where (ρ, v) ∈ AC(0, 1; P(Ω)) × L 1 (0, 1; L 2 (ρ)) belongs to the constraint set C 0 provided that
where ν is the outer unit normal on the boundary of the domain Ω. A curve ρ in P(Ω) is absolutely continuous in time, denoted ρ ∈ AC(0, 1; P(Ω)), if there exists
t 0 w(s)ds for all 0 < t 0 ≤ t 1 < 1. The PDE constraint (9 and 10) holds in the duality with smooth test functions on
This dynamic reformulation reduces the problem of finding the Wasserstein distance between any two measures to identifying the curve in P(Ω) that connects them with minimal kinetic energy. However, the objective function ( 8) is not strictly convex, and the PDE constraint ( 9) is nonlinear. For these reasons, in Benamou and Brenier's original work, they restrict their attention to the case ρ(•, t) ∈ P ac (Ω) and introduce the momentum variables m = vρ, in order to rewrite (8) as
where
and (ρ, m) ∈ AC(0, 1; P ac (Ω)) × L 1 (0, 1; L 2 (ρ -1 )) belong to the constraint set C 1 provided that
After this reformulation, the integral functional
is strictly convex along linear interpolations and lower semicontinuous with respect to weak-* convergence [1, Example 2.36], and the PDE constraint is linear. As an immediate consequence, one can conclude that minimizers are unique. Furthermore, for any ρ 0 , ρ 1 ∈ P ac (Ω), a direct computation shows that the minimizer ( ρ, m) is given by the Wasserstein geodesic from ρ 0 to ρ 1 ,
where T is the optimal transport map from ρ 0 to ρ 1 . (See [2,98,105] for further background on optimal transport.) Consequently, given any minimizer ( ρ, m) of ( 12), we can recover the optimal transport plan T via the following formula:
Building upon Benamou and Brenier's dynamic reformulation of the Wasserstein distance, one can also consider a dynamic reformulation of the JKO scheme (6). In particular, substituting (12) in (6) leads to the following dynamic JKO scheme: Problem 1 (Dynamic JKO) Given τ > 0, E, and ρ 0 , solve the constrained optimization problem,
where (ρ, m) ∈ AC(0, 1;
)) belong to the constraint set C provided that
We emphasize that the requirement ρ(x, t) ∈ P ac (Ω) for all t ∈ [0, 1] ensures that ρ(x, t) ≥ 0.
Remark 1 (Wasserstein geodesics) Note that for any ρ 1 ∈ P ac (Ω), we may take
in which case Problem 1 reduces to the Benamou-Brenier formulation of the Wasserstein distance (12). Consequently, the numerical method we develop for Problem 1 offers, as a particular case, a provably convergent numerical method for computing the Wasserstein geodesic and Wasserstein distance between ρ 0 and ρ 1 . On the one hand, there are many alternative methods for computing Wasserstein geodesics in Euclidean space. Indeed, the many algorithms described in the introduction for computing the Wasserstein distance also provide an optimal transport plan, which can be linearly interpolated to give the Wasserstein geodesic [8,11,55,61,68,97]. On the other hand, our method is distinguished because it could be more naturally extended to variants of the Wasserstein metric built on the Benamou-Brenier formulation [33,64,87], as well as to Wasserstein geodesics on non-Euclidean manifolds, where the geodesic equations on the underlying manifold may no longer be explicit, so that one cannot pass directly from the optimal transport plan to the Wasserstein geodesic.
Remark 2 (existence and uniqueness of minimizers) If the underling domain Ω is convex and the energy E is proper, lower semicontinuous, coercive, and λ-convex along generalized geodesics, and also satisfies {μ : E(μ) < +∞} ⊆ P ac (Ω), then, for τ > 0 sufficiently small, there exists a unique solution to Problem 1 [2, Theorem 4.0.4, Theorem 8.3.1]. In particular, these assumptions are satisfied by the energy G ρ 1 (18), as well as by the drift-diffusion interaction energy from the introduction (4), for U as in Eq. ( 3), V , W ∈ C 2 (Ω). (See, for example, [2, Section 9.3] or [56] for more general conditions on U , V , W .)
Thus, if we denote by ( ρ, m) the minimizer of Problem 1, then for τ > 0 sufficiently small, the proximal map, J τ (ρ 0 ) := ρ(•, 1) , is well defined for all ρ 0 ∈ D(E). Furthermore, the energy decreases under the proximal map,
which can be seen by comparing the value of the objective function at the minimizer (ρ, m) to the value of the objective function at (ρ(x, 0), 0) ∈ C and using that Φ(ρ, m) ≥ 0. Given ρ 0 ∈ D(E), if we recursively define the discrete time gradient flow sequence
then, taking τ = t/n, ρ n τ converges to ρ(x, t), the gradient flow of the energy E with initial data ρ 0 at time t, and under mild regularity assumptions on ρ 0 , we have
In this way, the classical JKO scheme provides a first-order approximation of the gradient flow [2, Theorem 4.0.4]. In our numerical simulations, we observe that this discretization error dominates other errors in our numerical method; see Sects. 4.2.1 and 4.2.2. For this reason, we introduce the following new scheme, inspired by the Crank-Nicolson method.
Problem 2 (Crank-Nicolson Inspired Dynamic JKO) Given τ > 0, E, and ρ 0 , solve the constrained optimization problem,
In Sect. 4.2.2, we provide numerical examples comparing the above method to the classical JKO scheme from Problem 1, illustrating that it achieves a higher-order rate of convergence in practice (see Figs. 7, 8, and 10), in spite of the fact that that it lacks the energy decay property of Problem 1. Under what conditions a higher-order analogue of inequality ( 21) holds for the new scheme is an interesting open question that we leave to future work, as the main goal of the present work is the development of fully discrete numerical methods for computing minimizers of Problem 1 and 2. By iterating either of these minimization problems, as in Eq. ( 20), we obtain a numerical method for simulating Wasserstein gradient flows.
We now turn to the discretization of the dynamic JKO scheme, Problem 1, and the Crank-Nicolson inspired scheme, Problem 2. We begin by noting that the Crank-Nicolson inspired Problem 2 can be rewritten in the same form as Problem 1 by considering the energy
Using this observation, we will now describe our discretization of both problems simultaneously.
where in the lth direction, we suppose there are N l intervals of spacing (z) l = (b la l )/N l :
Piecewise constant functions with respect to this discretization are given by
For any i ∈ N d+1 , write i = ( j, k), for the spatial index j ∈ N d and the temporal index k ∈ N. We let N x ∈ N denote the number of intervals in each spatial direction and N t ∈ N denote the number of intervals in the temporal direction. Take z = (x, Δt) for (x) l = (Δx) > 0 for all l = 1, . . . , d and Δt > 0.
We consider piecewise constant approximations (ρ h , m h ) of the functions (ρ, m), with coefficients denoted by (ρ j,k , m j,k ). For any (ρ, m) ∈ C(Ω × [0, 1]), one such approximation is the pointwise piecewise approximation ( ρh , mh ), obtained by defining the coefficients ( ρ j,k , m j,k ) to be the value of (ρ, m) on a regular grid of spacing (Δx) × (Δt):
where
), we have that ( ρh , mh ) converges to (ρ, m) uniformly.
Next, we approximate the energy functionals by discrete energies E h , beginning with energies of the form (4). Given a piecewise constant function ρ h with coefficients ρ j ,
where
Likewise, for energies of the form (4), we consider the following discretization of the energy H ρ 0 from Eq. ( 22) for the Crank-Nicolson inspired scheme, Problem 2,
Finally, to compute Wasserstein geodesics between two measures ρ 0 , ρ 1 ∈ P ac (Ω), we consider a discretization of the energy G ρ 1 from Eq. (18). Given a piecewise constant approximation ρ h 1 of ρ 1 and δ ≥ 0, define
Let D h t ρ h and D h x m h denote the discrete time derivative and spatial divergence on Ω × [0, 1] and let ν h denote the discrete outer unit normal of Ω. (See Hypothesis 3 for the precise requirements we impose on each of these discretizations). For example, in one dimension we may choose a centered difference in space and a forward Euler method in time,
or a Crank-Nicolson method,
We compute these discretizations of the derivatives at the boundary by extending m j,k to be zero in the direction of the outer unit normal vector. As we can only expect these approximations of the temporal and spatial derivatives to hold up to an error term, we relax the equality constraints from (17) in the following discrete dynamic JKO scheme.
The discretizations described in the previous sections lead to a fully discrete dynamic JKO problem:
where (ρ j,k , m j,k ) belong to the constraint set C h provided that for all j, k, j,k
The inequalities (30) enforce the PDE constraint and the boundary condition; the inequalities (31) enforce the mass constraint and the initial conditions. Recall that, by definition of Φ in Eq. ( 13), Φ(ρ j,k , m j,k ) < +∞ only if ρ j,k is nonnegative. Consequently, if a minimizer ρ j,k exists, it must be nonnegative.
A key element of our numerical method is that we relax the equality constraint (17) at the fully discrete level. This reflects the fact that even an exact solution of the continuum PDE will only satisfy the discrete constraints (30)(31) up to an error term depending on the order of the finite difference operators.
We allow the choice of δ i to vary for each of the above constraints. However, when the desired exact solution is sufficiently smooth, the optimal choice of δ i for a second-order discretization of the spatial and temporal derivatives is
where τ > 0 is the size of the timestep in the outer time discretization; see equations (6)(7). As we will demonstrate in Fig. 3 of our numerics section, relaxing the PDE constraint accelerates convergence to a minimizer of the fully discrete Problem 1 j,k without any loss of accuracy with respect to the exact continuum solution.
Finally, note that while the discrete PDE constraint (30) automatically enforces the mass constraint up to order δ 2 1 +δ 2 2 , we choose to impose the mass constraint separately via the first Eq. in (31). This leads to better performance in examples where the exact solution is not smooth enough to satisfy the discrete PDE constraint up to a high order of accuracy but imposing a stricter mass constraint leads to a higher quality numerical solution; see Fig. 4.
Under sufficient hypotheses on the discrete energy E h and the initial data ρ h 0 , minimizers of Problem 1 j,k exist; see Theorem 1. Furthermore, this discrete dynamic JKO scheme preserves the energy decreasing property of the original JKO scheme. To see this, note that, given an energy E h , time step τ > 0, and initial data (ρ h 0 ) j we may define the fully discrete proximal map by
where (ρ j,k , m j,k ) is any minimizer of Problem 1 j,k . Independently of which minimizer is chosen, we have
which can be seen by comparing the value of the objective function at the minimizer (ρ j,k , m j,k ) to the value of the objective function at (ρ j,k , m j,k ) = ((ρ 0 ) j , 0) ∈ C and using the fact that Φ ≥ 0. Furthermore, by iterating the fully discrete proximal map, we may construct a fully discrete gradient flow sequence
In analogy with the continuum case, we will use this fully discrete JKO scheme to simulate gradient flows. (See Algorithm 3.)
In order to find minimizers of Problem 1 j,k , we apply a primal dual operator splitting method. Since the constraints in Problem 1 j,k are linear inequality constraints, we may rewrite them in the form Ãi u -bi 2 ≤ δ i for i = 1, 2, 3, 4, where u = (ae, m), and ae and m are vector representations of the matrices ρ j,k and m j,k . (See the Appendix A for explicit formulas for Ãi and bi , in one spatial dimension). Similarly, we may rewrite the first term of the objective function (29) in terms of u, defining
We consider two cases for the energy term in the objective function. When the energy is of the form G h ρ 1 , as in Eq. ( 26), we reframe the problem by removing the energy from the objective function and adding
to the constraints (30) and (31), denoting A i ub i 2 ≤ δ i , for i = 1, 2, 3, 4, 5, as the modified constraints. On the other hand, when the energy is of the form ( 24) or ( 25), we rewrite it in terms of u as
In particular, if we let S be the selection matrix
, where F h and H h ρ 0 are defined in ( 24) and (25), respectively. This leads to the following two optimization problems:
To compute the Wasserstein distance, we solve Problem 2.3, and to compute the gradient flow of an energy, we iterate Problem 2
Primal-dual methods for solving optimization problems in which the objective function is the sum of two convex functions, as in Problem 2.3, are widely available [52]. However, analogous methods for optimizations problems in which the objective function is the sum of three convex functions, as in Problem 2.3, have only recently emerged [62,106]. In particular, in Algorithm 1, for Problem 2.3, we use Chambolle and Pock's well-known primal dual algorithm, and in Algorithm 2, for Problem 2.3, we use Yan's recent extension of this algorithm to objective functions with three convex terms. Both algorithms offer an extended range of primal and dual step sizes λ and σ and low per-iteration complexity, due to the sparseness of S, A, and Ã. Note specifically that the success of Algorithm 1 depends on the ease of computing the proximal operators related to φ and i δ , and therefore if we simply group the additional energy term in Problem 2.3 to either φ or i δ , it would violate such property. Instead, we shall consider E(u) as a separate term and take advantage of its smoothness, as shown in Algorithm 2. Finally, in Algorithm 3, we describe how Algorithm 2 can be iterated to approximate the full JKO sequence and, consequently, solutions of a range of nonlinear partial differential equations of Wasserstein gradient flow type.
until stopping criteria is achieved;
To initialize both algorithms, we choose φ 0 and m 0 to be zero vectors, and for ρ 0 , we let its components at the initial time (i.e., k = 0) be ρ 0 (x) evaluated on an equally spaced grid of width Δx, and other times to be zero. The stopping criteria consist of checking the PDE constraint ( 30)- (31) along with the convergence monitors:
The proximal operator, which appears in Algorithms 1 and 2, is defined by
For both h = σ i * δ and h = λΦ, there are explicit formulas for the proximal operators. By Moreau's identity, we may write Prox σ i * δ (x) in terms of projections onto balls of radius δ i centered at b i for the ith portion of vector x:
For the proximal operator of Φ, as shown by Peyré, Papadakis, and Oudet [95, Proposition 1],
where ρ * is the largest real root of the cubic polynomial equation P(x) := (x -ρ)(x + λ) 2 -λ 2 |m| 2 = 0, and m * can be obtained by m * = ρ * m/(ρ * + λ). By computing the proximal operator exactly, our primal dual method is positivity preserving, respecting a key property of the original Problems 1 and 1 j,k .
As the computations of both proximal operators (35), (36) are component-wise, they can easily be parallelized. Likewise, the computation of the gradient ∇ E is also component-wise:
Remark 4 (discrete convolution) As written, the above functionals involves a computation of the convolutions l W j,l ρ l,N t and l W j,l ρ l,0 , which can be achieved efficiently using the fast Fourier transform. Note that since the product of the discrete Fourier transforms of two vectors is the Fourier transform of the circular convolution and the interaction potential W j-k = W (x jx k ) is not a periodic function, we need zero-padding for computing the convolution. For the 1D case, we can first use the fast Fourier transform to compute the circular convolution of
and (ae, (0) N x -2 ), and then extract the last N x -1 elements, which are the desired
Embedding Algorithm 2 to the JKO iteration, we have the following algorithm for Wasserstein gradient flows. Note that line 6 in Algorithm 3 is to construct a better Algorithm 3: Primal-Dual for JKO sequence Input: ρ(x, t 0 ), Iter max , λ, σ, τ, n > 0 Output: ρ(x, t k ) for 0 ≤ k ≤ n and the corresponding energy E(ρ(x, t k ))
initial guess for ρ at each JKO iteration by applying an extrapolation.
Remark 5 (Comparison of our numerical method to previous work) Our definition of the indicator function in Problems 3(a) and 3 (b) differs from previous work, and as a result, our primal-dual algorithm does not require the inversion of the matrix AA T [7,95], which makes it quite efficient in high dimensions thanks to the sparsity of A. A similar approach is taken in a recent preprint [81] to compute the earth mover's distance W 1 , though, in this context, the earth mover's distance is dissimilar from the Wasserstein distance, since it does not require an extra time dimension and is thus a lower-dimensional problem.
A second difference between our method and the approach in previous works is that, since P(x) has at most one strictly positive root, it can be obtained by the general solution formula for cubic polynomials with real coefficients. Therefore, in our numerical simulations, we may compute the proximal operator Prox λΦ (u) by using this general solution formula, rather than via Newton iteration [95]. As a consequence, our method is truly positivity preserving, as opposed to positivity preserving in the limit as Δx, Δt → 0.
We close this section by recalling sufficient conditions on the primal and dual step sizes σ and λ that ensure Algorithms 1 and 2 converge to minimizers of Problems 2.3 and 2.3.
Proposition 1 (Convergence of Algorithm 1, c.f. [52]) Suppose σ λ < 1/λ max (AA t ) and a minimizer of Problem 2.3 exists. Then, as Iter max → +∞, and 1 , 2 → 0 in the stopping criteria (33) (34), the output u * of Algorithm 1 converges to a minimizer of Problem 2.3.
Proposition 2 (Convergence of Algorithm 2, c.f. [106]) Suppose that the discrete energy E(u) defined in Eq. ( 32) is proper, lower semi-continuous, convex, and there exists
Suppose further that σ λ < 1/λ max ( Ã Ãt ), λ < 2β, and a minimizer of Problem 2.3 exists. Then, as Iter max → +∞ and 1 , 2 → 0, the output u * converges to a minimizer of Problem 2.3.
Note here that the co-coercivity requirement on ∇ E in the above proposition is equivalent to require the Lipschitz continuity of ∇ E, i.e.,
For the energy of the form (4), this requirement reduces to the boundedness of U (ρ) and W , which can be satisfied independent of the numerical resolution if we consider bounded solution (no finite time blow up in ρ) and nonsingular interaction kernel. In the case when W is singular, for example when W is a Newtonian interaction potential, we approximate W by a continuous function via convolution with a mollifier; see Remark 7.
We now prove the convergence of solutions of the fully discrete JKO scheme, Problem 1 j,k , to a solution of the continuum JKO scheme, Problem 1. We begin, in Sect. 3.1, by describing the hypotheses we place on the underlying domain Ω, the energy E, the initial data ρ 0 , and the discretization operators. Then, in Sect. 3.2, we show that minimizers of Problem 1 j,k exist, provided the discretization is sufficiently refined. Finally, in Sect. 3.3, we prove that any sequence of minimizers of Problem 1 j,k has a subsequence that converges to a minimizer of Problem 1. In order for our finite difference approximation to converge, we assume throughout that a smooth, positive minimizer of the continuum JKO scheme Problem 1 exists. See hypothesis (H6) and Remark 9 for further discussion of this assumption.
We impose the following hypotheses on the underlying domain, energy, and discretization operators.
We assume that the spacing of the spatial discretization (Δx) > 0 and the temporal discretization (Δt) > 0 are both functions of h satisfying lim h→0 (Δx) = lim h→0 (Δt) = 0. (H2) For any piecewise constant function ρ h on Ω, the discrete energy functional E h has one of the following forms, as described in Sect. 2.2.2:
+∞ otherwise.
We place the following assumptions on U , V , and W and the target measure ρ 1 :
1 is a pointwise piecewise constant approximation of ρ 1 .
(H3) D h t and D h x are finite difference approximations of the time derivative and spatial divergence. We assume that D h t is a forward Euler method in time, whereas D h x can be given by an explicit or implicit scheme of first or higher order. We denote by D -h t and D -h x the dual operators with respect to the 2 inner product, and we assume the following integration by parts formulas hold for all piecewise constant functions ρ h ,
where ν h : Ω → R d is the discrete outer unit normal of Ω. Finally, we assume there exists C > 0 depending on the domain Ω × [0, 1], so that, for any
(See Sect. 2.2.3 for finite difference approximations satisfying these hypotheses.) (H4) The constraint relaxation parameters δ 1 , δ 2 , δ 3 , δ 4 ≥ 0 are functions of h with lim h→0 δ i = 0, for all i. If the energy is of the form (H2c), we require that δ 5 is a function of h satisfying lim h→0 δ 5 = 0 and lim h→0 (Δx + Δt) /δ 5 = 0. (H5) The initial data of the continuum problem satisfy ρ 0 ∈ C 1 (Ω) and ρ h 0 is a pointwise piecewise constant approximation of ρ 0 . (H6) Given the domain, energy, and initial data described in the previous hypotheses, there exists a minimizer (ρ, m) of the continuum Problem 1 satisfying
To ease notation in the following convergence proof, we observe that Problem 1 j,k may be rewritten as follows in terms of (ρ h , m h ), the piecewise constant functions on Ω × [0, 1] corresponding to the coefficients (ρ j,k , m j,k ). Problem 1 h (Discrete Dynamic JKO) Fix τ, δ 1 , δ 2 , δ 3 , δ 4 > 0, E h , and ρ h 0 . Solve the constrained optimization problem,
where (ρ h , m h ) belong to the constraint set C h provided that they are piecewise constant functions on Ω × [0, 1] and the following inequalities hold
Similarly, we may rewrite the definition of the discrete energies in hypothesis (H2) in terms of a piecewise constant functions ρ h on Ω corresponding to ρ j ,
Recall that, by definition of Φ in equation (13), Φ(ρ h , m h ) < +∞ only if ρ h is nonnegative. Consequently, if a minimizer ρ exists, it must be nonnegative.
We conclude this section with several remarks on the sharpness of the preceding hypotheses.
In hypothesis (H1), we assume that Ω is an ndimensional hyperrectangle. We impose this assumption for simplicity, as it provides an natural interpretation of the discretized outer unit normal ν h , which is essential in imposing the boundary conditions for our PDE constraint at the discrete level. More generally, our convergence result can be extended to any Lipschitz domain, as long as sufficient care is taken to define the discrete outer unit normal and the corresponding no flux boundary conditions. Remark 7 (assumption on energy) As described in hypothesis (H2), our convergence result applies to internal U , drift V , and interaction W potentials that are sufficiently regular on Ω. Our assumptions on U are classical and ensure that the internal energy is lower semicontinuous with respect to weak-* convergence [2, Remark 9.3.8]. Our assumptions on V and W , on the other hand, are somewhat stronger, and in practice, one often encounters partial differential equations for which the corresponding choices of V and W are not continuous. However, there are robust methods for approximating these potentials by continuous functions that ensure convergence of the gradient flows.
For example, the second author and Topaloglu provide sufficient conditions on discontinuous interaction potentials W for which gradient flows of the regularized interaction potential, W ε := W * ϕ ε for a smooth mollifier ϕ ε , converge to gradient flows of the original interaction potential W , as well as conditions that ensure minimizers of W ε converge to minimizers of W [59]. (The convergence of general stationary points of W that are not global minimizers to stationary point of W remains open.) Remark 8 (assumption on δ 5 ) In hypothesis (H4), it is essential that δ 5 not vanish too quickly with respect to other parameters in the discretization. A simple illustration of this fact arises in the case that δ 1 ≡ δ 2 ≡ δ 3 ≡ δ 4 ≡ 0. In this case, we cannot choose δ 5 ≡ 0, since our pointwise piecewise approximation of the initial data ρ h 0 will not generally have the same mass as our pointwise piecewise approximation of the target measure ρ h 1 , and if they do not have the same mass, minimizers of the discrete problem do not exist. Consequently, it would be impossible to prove that minimizers of the fully discrete problem converge to minimizers of the continuum problem. On the one hand, this does not greatly impact the performance of our numerical method, as can be seen by considering previous work by Papadakis, Péyre, and Oudet, which numerically implements this approach [95]. On the other hand, our numerical simulation in Fig. 3 indicates that poor choice of the relaxation parameters can cause the method to iterate longer than necessary, without any improvement in accuracy.
Our requirement that lim h→0 (Δx + Δt)/δ 5 = 0 is sufficient to fix this problem and ensure convergence of the method, and this requirement is nearly sharp. To see this, note that, for an arbitrary pointwise piecewise approximation ρ h 0 of a continuous function ρ 0 , we cannot in general achieve accuracy of | Ω ρ h 0 -Ω ρ 0 | better than O(Δx). If either δ 1 and δ 3 , the parameters for the PDE constraint and the mass constraint, are chosen arbitrarily small, then | Ω ρ h (•, 1) -Ω ρ h 0 | can likewise be made arbitrarily small. Thus, since ρ 0 , ρ 1 ∈ P ac (Ω),
so we much have δ 5 ≥ O(Δx). While a CFL-type condition is not necessary for the stability of our discretization of the PDE constraint, since ρ and m indeed become coupled in the continuum limit (see Eqs. ( 8) and ( 12)), one should expect (Δt) ≤ O(Δx) to give the best balance between computational accuracy and cost, and we indeed observe this numerically. Combining these facts shows that enforcing that δ 5 cannot decay faster than O(Δx + Δt) by assuming lim h→0 (Δx + Δt)/δ 5 = 0 is nearly optimal.
Remark 9 (assumption on existence of smooth, positive minimizer) In hypothesis (H6), we suppose that there exists a sufficiently regular minimizer (ρ, m), ρ > 0, of the continuum problem. Our proof of the existence of minimizers of the fully discrete problem and our proof that minimizers of the discrete problems converge to a minimizer of the continuum problem as h → 0 strongly rely on this assumption.
In particular, the smoothness assumption allows us to use convergence of the finite difference operators, described in hypothesis (H3), to construct an element of C h in Proposition 3. The positivity assumption allows us to conclude that ∇ ρ,m Φ is uniformly bounded on the range of ρ, which we use to prove the lim sup inequality for the recovery sequence in Theorem 2(b).
From the perspective of approximating gradient flows, which are solutions of diffusive partial differential equations (3), such regularity and positivity can be guaranteed as long as the initial data are smooth and positive and either the diffusion is sufficiently strong or the drift and interaction terms do not cause loss of regularity. On the other hand, developing conditions on the energy and initial data that ensure such regularity and positivity holds at the level of the JKO scheme, for minimizers of Problem 1, remains largely open: results on the propagation of L p (R d ) or BV bounds along the scheme have only recently emerged [17,50,63].
From the perspective of approximating Wasserstein geodesics, the now classical regularity theory developed by Caffarelli and Urbas ensures that if the source and target measures ρ 0 and ρ 1 are smooth and strictly positive, then the minimizer of Problem 1 ( ρ, m) is also smooth and strictly positive. (See, for example, [105, Section 4.3] and [2, Section 8.3].)
Along with this analytical justification for our smoothness and positivity assumptions, our numerical results also indicate that such assumptions are in general necessary. For example in Fig. 4, we observe that if the source and target measure of a Wasserstein geodesic are not sufficiently smooth, the numerical solution introduces artificial regularity. Likewise, even in Fig. 6, we observe that the numerical simulation is strictly positive (though very close to zero in places), while the exact solution is identically zero outside of its support. Still, in spite of the fact that our theoretical convergence result requires smoothness and positivity assumptions, in practice our numerical method still performs well on nonsmooth or nonpositive problems, provided that the spatial and temporal discretization are taken to be sufficiently small; see Figs. 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, and 18. Finally, these types of smoothness and positivity assumptions are typically needed in convergence proofs for numerical methods based on the JKO scheme. For example, in a method based on the Monge Ampére approximation of the Wasserstein distance, the exact solution is required to be uniformly bounded above and below [10]. Likewise, while rigorous convergence results for fully discrete numerical methods based on entropic or Fisher information regularization remain open, since these methods correspond to introducing numerical diffusion at the level of the PDE, they automatically enforce smoothness and positivity [28,55,82].
We now show that, under the hypotheses described in the previous section, minimizers of the fully discrete JKO scheme, Problem 1 h , exist for all h > 0 sufficiently small. We begin with the following proposition, which constructs a specific element in the constraint set C h , which we will use both in our proof of existence of minimizers and in our Γ -convergence results in the next section.
Then for h > 0 sufficiently small, there exists ( ρh , mh
If, in addition, the energy satisfies hypothesis (H2c) and E(ρ(•, 1)) < +∞, then we have
for all h > 0 sufficiently small.
Let mh be a pointwise piecewise constant approximation of m; see Eq. (23). Recall that ν h is the discrete outer unit normal vector. We define mh : Ω × [0, 1] → R d component-wise to respect the no flux boundary conditions, letting ( mh ) l denote the lth component of the vector for l = 1, . . . , d. If x ∈ ∂Ω, then we define
Otherwise, we take mh (x, t) = mh (x, t). Define ρh :
x mh (x, t) ≡ 0. We begin by showing that ( ρh , mh ) ∈ C h . By construction, for all h > 0,
Taking f h ≡ 1 in Hypothesis (H3) and applying the PDE constraint ensures that, for all s ∈ [0, 1] and k ∈ N so that k(Δt) ≤ s < (k + 1)Δt,
Thus, we also obtain
This concludes the proof that ( ρh , mh ) ∈ C h . We now show that ( ρh , mh ) → (ρ, m) uniformly on Ω × [0, 1] as h → 0. We begin by proving convergence of mh to m. Due to hypothesis (H1) on our domain Ω, whenever e i • ν h (x) = 0, there exists y ∈ ∂Ω so that |y -x| ≤ 2 √ d(Δx) and ν(y) = e i . Thus, whenever e i • ν h (x) = 0, the continuum boundary condition m(y, t) • ν(y) = 0 ensures that for all t ∈ [0, 1],
We also have that, for all (x, t)
Therefore, for all (x, t) ∈ Ω × [0, 1], there exists
We now prove the convergence of ρh to ρ. Since (ρ, m) is a classical solution of the PDE constraint and ρh : Ω × [0, 1] → R is defined by the conditions that ρh (x, 0) = ρh 0 and
Since ρh → ρ uniformly and ρ > 0, we immediately obtain (39). Finally, suppose the energy satisfies (H2c). Since E(ρ(•, 1)) = G ρ 1 (ρ(•, 1)) < +∞, we have ρ(•, 1) = ρ 1 . By inequality (41) and the fact that ρ h 1 is a pointwise piecewise approximation of ρ(•, 1),
where
→ 0. Thus, for h sufficiently small,
which completes the proof.
Theorem 1 (minimizers of discrete dynamic JKO exist) Suppose that hypotheses (H1)-(H6) hold. Then for all h > 0 sufficiently small, a minimizer of Problem 1 h exists.
Proof First, we note that Proposition 3 ensures that, for h > 0 sufficiently small, the constraint set C h is nonempty and contains some (ρ h , m h ) satisfying ρ h > 0. If the energy satisfies (H2a) or (H2b), then we immediately obtain E h (ρ h (•, 1)) < +∞.
Similarly, if the energy satisfies (H2c), then inequality (40) in Proposition 3 again ensures that E h (ρ h (•, 1)) < +∞.
Since Φ(ρ h , m h ) < +∞ whenever ρ h ≥ 0, this ensures that value of the objective function in the discrete minimization problem 1 h is not identically +∞ on the constraint set. Therefore, inf
and we may choose a minimizing sequence (ρ h n , m h n ) ∈ C h that converges to the infimum. We may assume, without loss of generality, that sup
, and there exist ρ ∈ P(Ω × [0, 1]) and μ ∈ P(Ω) so that, up to a subsequence, ρ h * ρ and ρ h (•, 1) * μ. Furthermore, for any piecewise constant function
Proof By hypothesis (H6), ρ h 0 → ρ 0 uniformly on Ω. Likewise, the constraint on the initial data ( 38) and (H4) ensure lim h→0 ρ h (•, 0) -
We now turn to Eq. (47). By the PDE constraint and boundary conditions ( 37) and summation by parts, via hypotheses (H3),
where, in the last line, we use that (H4) ensures δ 2 , δ 4 → 0 and the fact that
Next, we show that there exist ρ ∈ P(Ω × [0, 1]) and μ ∈ P(Ω) so that, up to a subsequence, ρ h * ρ and ρ h (•, 1)
* μ. By Hölder's inequality and the mass constraint (38),
where, in the last line, we use that (H4) ensures δ 3 → 0. Since hypothesis (H6) ensures ρ h 0 → ρ 0 uniformly and Ω ρ 0 = 1, we obtain,
Furthermore, since 1 0 Ω Φ(ρ h , m h ) < +∞ for each h > 0, we must have ρ h ≥ 0 on Ω × [0, 1], and the above equation ensures sup h>0 ρ h L 1 (Ω×[0,1]) < +∞. Thus, classical functional analysis results ensure there exists a subsequence that converges to some ρ ∈ P(Ω × [0, 1]) in the weak-* topology (see, e.g., [20,Section 3]).
Finally, taking f h ≡ 1 in Eq. ( 47) gives,
Arguing as above, we obtain that, up to a further subsequence,
We now prove that the discrete energies E h are lower semicontinuous along weak-* convergent sequences.
Proof First, suppose the energy satisfies (H2a). Since the piecewise constant approximations V h and Ŵ h converge to V and W uniformly, for any sequence ρ h * ρ,
Furthermore, our assumptions on U guarantee that the internal energy term is lower semicontinuous with respect to weak-* convergence [2, Remark 9.3.8], so lim inf h→0 Ω U (ρ h (x))dx ≥ Ω U (ρ(x))dx. Combining this with equations (48)(49) gives the result. Next, suppose the energy satisfies (H2b). Since ρ 0 > 0 on the compact set Ω and U is uniformly continuous on ρ 0 (Ω) ⊂ (0, +∞), the fact that hypothesis (H6) ensures ρh 0 → ρ 0 uniformly ensures U ( ρh 0 ) → U (ρ 0 ) uniformly. Therefore,
Likewise, since V h and Ŵ h converge to V and W uniformly, we also have
Combining these limits with the lim inf inequality for energies of the form (H2a) gives the result. Finally, suppose the energy satisfies (H2c). Without loss of generality, we may assume that lim inf h→0 G h ρ 1 (ρ h ) < +∞, so that up to a subsequence,
We now apply Proposition 4 to prove the Γ -convergence of Problem 1 h to Problem 1.
, there exists a sequence ( ρh , mh ) ∈ C h so that ( ρh , mh ) → (ρ, m) uniformly and
Proof We first prove part (a). Suppose (ρ h , m h ) ∈ C h , with ρ h * ρ and m h * m. We begin by showing that the limit (ρ, m) belongs to C. Fix f ∈ C ∞ (Ω × [0, 1]) and let f h be a pointwise piecewise constant approximation of f . (See Eq. ( 23).) By Lemma 1 and hypothesis (H3),
We conclude that (ρ, m) satisfies the PDE constraint in the sense of distributions (17), which gives ρ ∈ AC([0, 1], P(Ω)) [2, Lemma 8.1.2]. In particular, since ρ is continuous in time, we have that the μ defined in Lemma 1 satisfies μ = ρ(•, 1). We now consider the inequality in part (a). Since the integral functional (ρ, m) → 1 0 Ω Φ(ρ, m) is lower semicontinuous with respect to weak-* convergence of measures [1, Example 2.36], we immediately obtain lim inf
which completes the proof of part (a). We now turn to part (b). Let ( ρh , mh ) ∈ C h be the sequence constructed in Proposition 3, so ( ρh , mh ) → (ρ, m) uniformly. By inequality (39), there exists c > 0 so that ρ h (x, t) ≥ c for h sufficiently small. Therefore,
It remains to show that lim sup
First, suppose the energy satisfies either (H2a) or (H2b). By Eqs. ( 48)-( 50), which hold for any weak-* convergent sequence, and the fact that
) uniformly, which gives the result. Finally, suppose the energy satisfies (H2c). Without loss of generality, suppose (40) ensures that, for h sufficiently small,
which gives the result.
We conclude this section by applying the Γ -convergence proof from Theorem 2 to prove that any sequence of minimizers of the discrete Problem 1 h converges, up to a subsequence, to a minimizer of the continuum Problem 1.
Then, for any sequence of minimizers (ρ h , m h ) of Problem 1 h , we have, up to a subsequence, ρ h * ρ and m h * m, where (ρ, m) is a minimizer of Problem 1.
Note that, if the minimizer of the continuum Problem 1 is unique, then this theorem ensures that any sequence of minimizers of the discrete Problem 1 j,k has a further subsequence that converges to this minimizer. Therefore, the sequence itself must converge to the unique minimizer of the continuum problem. (See Remark 2 for sufficient conditions that ensure the minimizer of the continuum problem is unique.) Proof of Theorem 3 First, note that Lemma 1 ensures that there exist ρ ∈ P(Ω ×[0, 1]) and μ ∈ P(Ω) so that, up to a subsequence, ρ h * ρ and ρ h (•, 1)
* μ. In order to prove an analogous weak-* compactness result for m h we first prove that, up to a subsequence,
By (H6), there exists a minimizer ( ρ, m) of the continuum Problem
Furthermore, Proposition 4 ensures that lim inf
which is bounded below by some constant, since hypothesis (H2c) ensures E ≥ 0 and hypotheses (H2a) or (H2b) ensures E(μ) > -∞, since U , V , and W are bounded below and U is bounded below on the range of the strictly positive density ρ 0 . Therefore, up to a subsequence, we obtain (51). We now deduce weak-* convergence of m h . By Hölder's inequality, the fact that ρ h * ρ, and the definition of Φ, we have
Thus, up to another subsequence, m h * m on Ω × [0, 1]. It remains to show that the limit (ρ, m) of (ρ h , m h ) is a minimizer of Problem 1. By Theorem 2, part (a), we have (ρ, m) ∈ C and
Combining this with inequality (52) above, we conclude that (ρ, m) ∈ C is also a minimizer of Problem 1, which completes the proof.
In this section, we provide several examples demonstrating the efficiency and accuracy of our algorithms. We begin by using Algorithm 1 to compute Wasserstein geodesics between given source and target measures, and we then turn to Algorithm 3 to compute solutions of nonlinear gradient flows. In the following simulations, we take our computational domain Ω to be a square, imposing the no flux boundary conditions on m dimension by dimension. In practice, unless otherwise specified, we always impose the discrete PDE constraint via the Crank-Nicolson finite difference operators ( 28), and we choose 1 = 2 = in the stopping criteria to be 10 -5 unless otherwise specified. For the relaxation of the constraints in ( 30) and ( 31), we choose δ 1 = δ 2 = δ 4 = δ 5 = δ, and δ 3 differently, as specified in each example.
As described in Remark 1, a particular case of our numerical scheme provides a method for computing the Wasserstein geodesic between two probability densities. We begin by computing the Wasserstein geodesic between rescaled Gaussians in one dimension: ). We choose σ = 0.1 and σ λ = 1.5/λ max (A A t ) and compute 10 5 iterations. Left: evolution of geodesic from time t = 0 to t = 1. Right: rate of convergence of numerical solution to exact solution, as a function of the number of iterations in Algorithm 1
The target measure is simply a translation and dilation of the initial measure, ρ 0 (x) = (0.5)g μ 0 ,θ 0 (x) and ρ 1 (x) = (0.5)g μ 1 ,θ 1 (x). The optimal transport map T (x) from ρ 0 (x) to ρ 1 (x) is given explicitly byfoot_0
Rewriting Eq. ( 15) for the geodesic ρ(x, t) and velocity v(x, t) induced by this transport map, via the definition of the push forward, we obtain
In Fig. 2, we apply Algorithm 1 to compute the Wasserstein geodesic ρ(x, t) between the initial and target densities (53), with means and variances μ 0 = -1.5, θ 0 = 0.3, μ 1 = 1.5, and θ 1 = 0.6. On the left, we plot the evolution of the geodesic at various times. On the right, we plot the 1 error of the densities, momenta, and Wasserstein distance as a function of the number of iterations, l, observing a rate of convergence of order O(1/l) (dashed black line). Here, the error is defined as
Fig. 3 Analysis of how the scaling relationship between the relaxation parameter δ and the spatial discretization (Δx) affects the accuracy of the numerical method and the number of iterations required to converge. We contrast the choices δ = (Δx) 2 , δ = (Δx) 3 and δ = 10 -8 for the example of the Wasserstein distance between geodesics, illustrated in Fig. 2. We take N t = 30, N x = 300, σ = 1, σ λ = 0.99/λ max (A A t ) and δ 3 = δ
In Fig. 3, we illustrate how choosing the optimal scaling relationship between the relaxation parameter δ and the spatial and temporal discretizations (Δx), (Δt) allows the method to converge in fewer iterations. We contrast the choices δ = (Δx) 2 , δ = (Δx) 3 , and δ = 10 -8 , for the example of the Wasserstein distance between geodesics, illustrated in Fig. 2, where the outer time step τ = 1, (Δx) ∼ (Δt), and δ 3 = δ. Based on the order of accuracy of our Crank-Nicolson approximation of the PDE constraint, we expect that δ = (Δx) 2 should give the optimal balance between accuracy and computational efficiency. (See Remark 3.)
In the plot on the left, we observe that for all choices of δ, the error between the numerical solution ρ (l) and the exact solution ρ * is identical, with the error saturating after 10 5 iterations. Thus, all three choices of δ provide the same level of accuracy, and the best way to distinguish between them is to identify which choice of δ causes the stopping criteria (33 and 34) to be satisfied in the least number of excess iterations after 10 5 . The behavior of two key stopping criteria is shown in the plot on the rightthe PDE constraint Au (l)b and the convergence monitor for the relative error of the dual variables φ (l) -φ (l-1) / φ (l) . Of the four stopping criteria we consider (PDE constraint and three convergence monitors), these two are the last to be satisfied in all of the numerical simulations contained in this manuscript, hence these determine when our method terminates its iterations.
For the case of δ = (Δx) 2 (red lines), we indeed observe that the PDE constraint (solid line) satisfies its stopping criteria (dashed line) by 10 4 iterations and the dual variables (starred line) satisfy their stopping criteria of 10 -5 by 10 5 iterations. On the other hand, for the cases of δ = (Δx) 3 (blue lines) and δ = 10 -8 (green lines), we see that while the dual variables (starred lines) have satisfied their stopping criteria of 10 -5 by 10 4 iterations, the PDE constraints (solid lines) do not satisfy their stopping criteria (dashed lines) until later-it takes more than 10 5 iterations for δ = (Δx) 3 and more than 10 7 iterations for δ = 10 -8 . This example shows that choosing δ without respecting the order of accuracy of the finite difference approximation in the PDE . Here, σ = 0.1, σ λ = 0.99/λ max (A A t ) and then λ = 0.9727, δ = 10 -5 , and δ 3 = 10 -8 constraint, one wastes computational effort without improving the accuracy of the numerical solution.
Next, we compute Wasserstein geodesics between initial and target measures when neither are smooth nor strictly positive. In Fig. 4, we compute the geodesic between a profile of the British Parliament and its translation. We do not observe convergence to the exact geodesic, which would be a constant speed translation, and instead observe degradation of the parliamentary building at intermediate times, due to numerical smoothing. Similarly, in Fig. 5, we compute the geodesic between Pac-Man and a ghost, visualized as characteristic functions on sets in two dimensions. Again, we observe numerical smoothing around the edges of discontinuity. Both of these examples offer a numerical justification for the smoothness assumption we impose in our main convergence Theorem 3. In the absence of such smoothness, it appears that the method does not converge. Similar smoothness assumptions are required in the other numerical methods for Wasserstein geodesics for which rigorous convergence has been analyzed, including Monge Ampére-type methods [11,68].
In this and the next section, we consider several examples of Wasserstein gradient flows, including some which have appeared in previous numerical studies [3,29,49,99], to demonstrate the performance of our method for simulating solutions of nonlinear partial differential equations.
The porous medium equation
is the Wasserstein gradient flow of the energy (4), with U (ρ) = 1 m-1 ρ m and V = W = 0 . A well-known family of exact solutions is given by Barenblatt profiles (c.f. [104]), which are densities of the form
We now apply Algorithm 3 to simulate solutions of the m = 2 porous medium equation with Barenblatt initial data, t 0 = 10 -3 and C = (3/16) 1/3 . Here, the Euler discretization ( 27) is used. In Fig. 6, we plot the evolution of the numerical solution over time, and we observe good agreement with the exact solution of the form (56), which is displayed in dashed curve.
In Fig. 7, we analyze how the rate of convergence depends on the inner time step Δt, the spatial discretization Δx, and outer time step of the JKO scheme τ . We compute the error between the exact solution and the numerical solution in the 1 norm, i.e.,
In the plot on the left of Fig. 7, we consider two fixed values of τ and examine how the error depends on N t and N x = 10N t . In both cases, the error quickly saturates, indicating that the outer time step τ dominates the error. In the plot on the right, we fix N t = 20 and N x = 200 and consider how the error depends on τ . We observe slightly less than first-order convergence in τ for the classical JKO scheme (E h = F h ) and higher-order convergence for the Crank-Nicolson inspired scheme (E h = H h ). We believe these slower rates of convergence are due to the lower regularity of solutions to the porous medium equation with compactly supported initial data, which are merely Hölder continuous. In Fig. 8, we consider the case of smooth, strictly positive initial data, given by a Gaussian with mean μ = 0 and variance θ = 0.2 (53), in which case solutions of the PDE remain smooth over time. On the left, we show the evolution of the solutions over time, and on the right, we illustrate that the classical JKO scheme indeed attains first-order accuracy, though the Crank-Nicolson inspired scheme is still less than second-order accurate.
We now consider a nonlinear variant of the Fokker-Planck equation, Right: The rate of convergence for various choices of τ , contrasting the traditional first-order JKO scheme with the new Crank-Nicolson inspired scheme. For each choice of τ in our computation of the higher-order method, we choose our stopping criteria = 10 -4 * 2 -0.01/τ inspired by the porous medium equation described in the previous section (55). When V is a confining drift potential, all solutions approach the unique steady state
where C > 0 depends on the mass of the initial data, so that ρ ∞ dx = ρ 0 dx, see [44,51].
In Fig. 9, we simulate the evolution of solutions to the nonlinear Fokker-Planck equation with V (x) = x 2 , m = 2, and initial data given by a Gaussian with mean μ = 0 and variance θ = 0.2 (53). On the left, we plot the evolution of the density Fig. 9 Evolution of the solution ρ(x, t) to the one-dimensional nonlinear Fokker-Planck equation, with m = 2 and V (x) = x 2 . We choose τ = 0.05, Δx = 0.04, Δt = 0.1, λ = 0.1641, σ = 1, δ = 10 -5 , and δ 3 = 10 -5 . Left: evolution of density ρ(x, t) toward equilibrium ρ ∞ (x). Right: Rate of decay of corresponding energy with respect to time Fig. 10 Analysis of rate of convergence for a solution of the nonlinear Fokker-Planck equation, as in Fig. 9. We choose Δt = 0.1, Δx = 0.04 and consider the error (57) for various choices of τ , contrasting the traditional first-order JKO scheme with the new Crank-Nicolson inspired scheme ρ(x, t) toward the steady state ρ ∞ (x). On the right, we compute the rate of decay of the corresponding energy (4) as a function of time, observing exponential decay as the solution approaches equilibrium. In this way, our method recovers analytic results on convergence to equilibrium of Carrillo, DiFrancesco, and Toscani [35,51].
In Fig. 10, we analyze how the rate of convergence depends on the outer time step τ of the scheme, for sufficiently small inner time step Δt = 0.1 and spatial discretization Δx = 0.04. We compute the error We observe slightly faster than first-order convergence for the traditional JKO scheme (E h = F h ) and higher-order convergence for the new Crank-Nicolson inspired scheme (E h = H h ). We believe this improvement in the rate of convergence as compared to our previous example for the porous medium equation, Fig. 7, is due to the rapid convergence to the steady state ρ ∞ .
In this section, we consider a nonlocal partial differential equation of Wasserstein gradient flow type, known as the aggregation equation
In recent years, there has been significant interest in interaction kernels W that are repulsive at short length scales and attractive at longer distances, such as the kernel with logarithmic repulsion and quadratic attraction
For this particular choice of W , there exists a unique equilibrium profile [38], given by
In Fig. 11, we simulate the solution to the aggregation equation with Gaussian initial data (53) with mean μ = 0 and variance θ = 1, analyzing convergence to equilibrium. On the left, we plot the evolution of the density ρ(x, t) at varying times, observing convergence to the equilibrium profile ρ ∞ (x). On the right, we compute the rate of the decay of the energy as a function of time, observing exponential decay as obtained by Carrillo, Ferreira, and Precioso [38] with a slightly slower numerical rate.
As the interaction potential W defined in Eq. ( 59) is not continuous, we make the following modifications to our discretization of the JKO scheme. To avoid evaluation of W (x) at x = 0, we set W (0) to equal the average value of W on the cell of width 2h centered at 0, i.e., W (0) = 1 2h h -h W (x)dx, where we apply Gauss-Legendre quadrature rule with four grid points to evaluate the integral. In addition to modifying the interaction kernel in this way, we also introduce an artificial diffusion term of the form ∂ x (ρ∂ x ρ) with = 1.6 × (Δx) 2 to the right-hand side of (58), to avoid the possible overshoot at the boundary. (See also [29] for a similar treatment.)
In the following, we consider a few gradient flows in two dimensions. Here, the constraint relaxation parameters are always chosen as δ = δ 3 = 10 -6 .
We now continue our study of the aggregation equation ( 58) with repulsive-attractive interaction potentials in two dimensions, with interaction kernels of the form
using the convention that |x| 0 0 = ln(|x|). It is well known that the repulsion near the origin of the potential determines the dimension of the support of the steady state measure, see [4,34]. In the following simulations, we take the initial data to be a gaussian (53) with mean μ = 0 and variance θ = 0.25. In Fig. 12, we simulate the evolution of solutions to the aggregation equation, with a = 4 and b = 2 in the interaction potential W , defined in Eq. ( 60). We observe that the solution concentrates on a Dirac ring with radius 0.5 centered at the origin, recovering analytical results on the existence of a stable Dirac ring equilibrium for these values of a and b [5,13].
In Fig. 13, we simulate the evolution of solutions to the aggregation equation, with a = 2 and b = 0. We observe that the solution converges to a characteristic function on the disk of radius 1, centered at the origin, recovering analytic results on solutions of the aggregation equation with Newtonian repulsion [14,34,65]. We follow the same strategy described in Sect. 4.2.3 with = 1.6 × (Δx 2 + Δy 2 ) to overcome the singularity of the interaction potential at x = 0 and potential overshooting.
Next, we compute solutions of aggregation-drift equations
where W (x) = |x| 2 /2-ln(|x|) and V (x) = -α β ln(|x|). As shown in several analytical and numerical results [29,42,53], the steady state is a characteristic function on a torus or "milling profile", with inner and outer radius given by
In Fig. 14, we simulate the long time behavior of a solution of the aggregation-drift equation with α = 1 and β = 4 and Gaussian initial data (53), μ = 0, θ = 0.25, as well as the rate of the decay of the entropy as the solution converges to equilibrium. In Fig. (15), we plot the evolution of the density from a nonradially symmetric initial data, given by five Gaussians to the same equilibrium profile. We follow the same strategy described in Sect. 4.2.3 to overcome the singularity of the interaction potential at x = 0 In recent years, there has been significant activity studying equations of this form, both analytically and numerically. When the interaction kernel W is attractive, the competition between the nonlocal ∇ • (ρ∇W * ρ) and nonlinear diffusion νΔρ m causes solutions to blow up in certain regimes and exist globally in time in others, see for example [18,19,25,26,41] and the survey [32]. With fixed m, and in the presence of nonlocal interaction, the equation has a unique steady state which is radially decreasing up to a translation [15,40]. In Fig. 16, we simulate a solution of the aggregation-diffusion equation with W (x) = -e -|x| 2 /π , ν = 0.1, and m = 3, and initial data given by a rescaled characteristic function on the square,
Diffusion dominates both the short and long ranges, and the medium range aggregation leads to the formation of four bumps, which ultimately approach a single bump equilibrium. (See also [29].) In Fig. 17, we simulate solutions of the Keller-Segel equation, which is an aggregation-diffusion equation (61) with a Newtonian interaction kernel, i.e., W (x) = 1 2π ln(|x|) in two dimensions for ν = 1 and both m = 1 and m = 2, illustrating the role of the diffusion exponent in blowup or global existence of solutions. We choose the initial data to be given by a rescaled gaussian, obtained by multiplying equation ( 53) by a mass M = 9π , with mean μ = 0 and variance θ = 0.5. On the left, we In Fig. 18, we again simulate solutions of the Keller-Segel equation with m = 2, but in this case we take the initial data to be given by three localized bumps (Gaussian rings, i.e., the radial cut of the ring is a Gaussian with a center on the circle.) We observe a two-stage evolution in which the each of the bumps converges to a localized quasistationary state, and then interact and merge into one single bump in the long time limit. This is a manifestation of the typical metastability phenomena, which is likely present in the majority of the diffusion dominated Keller-Segel models [24,29,32].
In this section, we provide explicit formulas for the matrix à and the vectors u and b introduced in Problem 3(b) in Sect. 2.3, which play a key role in the implementation of Algorithms 1, 2, and 3. For simplicity, we consider the case of one space dimension, and the discretization takes the form (28). The constructions of A and b in Problem 3(a) are very similar except a slightly different treatment of ρ at final time. From now on, for simplicity of notation, we will drop the tildes for the matrix à and vector b.
Define N = (N x +1)(N t +1). Let ⊗ denote the Kronecker tensor product, I N x +1 the identity matrix of size N x + 1, and (x) M the column vector in R M with all components equal to x. Then we define u = (ae .,k ) N t k=0 ; (m .,k ) N t k=0 ∈ R N , ae .,k = (ρ j,k ) N x j=0 , m .,k = (m j,k ) N x j=0 and the matrix A ∈ R M×2N takes the form
Here, A ρ ∈ R N ×N reads
t ⊗ I (1) x + D
(2) t ⊗ I (1) x := A (1) ρ + A (2) ρ ,
where D
(1)
t ∈ R (N t +1)×(N t +1) , and I
(1)
x ∈ R (N x +1)×(N x +1) are
t = 1 0 .
Here, D
t and D
(2)
t correspond to the temporal discretization and initial condition for ρ. Likewise,
t ⊗ D (1) x + I N t +1 ⊗ D (2) x := A (1) m + A (2) m , where D (1) x , D (2) x ∈ R (N x +1) , and B
(1) t ∈ R (N t +1)×(N t +1) :
For mass conservation, let S ρ = (x) t N x +1 , then A mass = I N t +1 ⊗ S ρ . In sum, different A i can be written as
One way to see that this is the unique optimal transport map from ρ 0 to ρ 1 is to note that T #ρ 0 = ρ 1 and T (x) is the gradient of a convex function; see, for example, [2, Section 6.2.3].