Numerical Optimal Transport from 1D to 2D using a Non-local Monge-Amp\`ere Equation

We consider the numerical solution of the optimal transport problem between densities that are supported on sets of unequal dimension. Recent work by McCann and Pass reformulates this problem into a non-local Monge-Amp\`ere type equation. We provide a new level set framework for interpreting this non-linear PDE. We also propose a novel discretisation that combines carefully constructed monotone finite difference schemes with a variable-support discrete version of the Dirac delta function. The resulting method is consistent and monotone. These new techniques are described and implemented in the setting of 1D to 2D transport, but can easily be generalised to higher dimensions. Several challenging computational tests validate the new numerical method.


Introduction
The problem of optimal transport [27] involves finding a mapping T : X → Y that rearranges a density f supported on X into a density g supported on Y while minimising a given cost function (1) inf This problem has been widely studied, both theoretically and numerically, in the setting where the sets X and Y lie in same ambient space (eg: X, Y ⊂ R n ).However, in many important problems this simplifying assumption is not valid.In this article, we consider the problem of optimal transport between densities that are supported on sets of different dimension (X ⊂ R n , Y ⊂ R m , n < m).This setting is critical to applications such as economics and game theory (m ≥ n ≥ 1) [7,15,21], simulation of atmosphere and ocean dynamics (2 ≤ n ≤ 3 = m) [8], and Generalized Adversarial Networks (m = n) [19].
Despite the importance of this problem, almost nothing is known about its numerical solution except in the case of a one-dimensional target (n = 1) combined with significant extra structure on the problem data [7].Utilizing the linear programming structure of the optimal transport problem is theoretically possible but computationally infeasible as it requires the discretization of an (m+n)-dimensional domain.
In this article, we propose the first PDE-based approach to the problem of optimal transportation between unequal dimensions.The new method we propose relies on a reformulation in terms of a non-local Monge-Ampère type equation that was recently proposed by McCann and Pass [20]: (2) Here m ≥ n ≥ 1, H k denotes the k-dimensional Haussdorf measure, ∂ c u is the c-subgradient of u, and the matrix-valued function A and scalar-valued function ψ encode information about the cost and density functions.
While the numerical solution of the Monge-Ampère equation has received a great deal of attention in recent years [2,4,5,9,[11][12][13][14]17,22,24,25], nothing is known about non-local extensions.We propose a novel numerical method for (2) that is inspired by recent developments in the numerical analysis of fully nonlinear elliptic equations, but requires significant new ideas in order to couple the unequal dimensions.Because the techniques we utilise are motivated by existing convergence theory, there is a great reason to hope that a proof of convergence of this method will soon be developed.
As a proof-of-concept, the new method is implemented in the 1D to 2D case.However, we emphasise that this method employs techniques that are designed to easily generalise to higher dimensions.A range of computational examples demonstrate that this is an effective approach for solving a very challenging problem.

Numerical Solution of Monge-Ampère equations
While nothing is known about the numerical solution of general non-local Monge-Ampère equations of the form (2), there has been a great deal of progress in the simpler case resulting from optimal transport between equal dimensions (m = n) with quadratic cost (c(x, y) = 1 2 |x − y| 2 ).In this setting, the PDE reduces to the second boundary value problem for the Monge-Ampère equation: (3) The Monge-Ampère equation is an example of a second-order fully nonlinear elliptic PDE.

Definition 1 (Degenerate elliptic). The operator
Much of the recent progress in the design of provably convergent methods for the Monge-Ampère equation is inspired by the convergence framework of Barles and Souganidis [1].This foundational work demonstrated that a consistent, monotone, stable numerical method will converge to the weak solution of the PDE, provided the underlying PDE satisfies a strong form of comparison principle (subsolutions always lie below supersolutions).

Definition 2 (Consistency). The scheme
Definition 3 (Monotonicity).The scheme Definition 4 (Stability).The scheme A nice property of monotone schemes is that it is generally very easy to ensure the existence of a solution by including, at most, a small perturbation to make the scheme proper.

Definition 5 (Proper). The scheme
Any monotone scheme F h can be perturbed to a proper scheme via where h (x) → 0 as h → 0. A typical choice, which lends itself to nice stability properties, is to let h scale like the truncation error of F h .Theorem 6 (Existence and uniqueness [23,Theorem 8]).Let F h be a continuous, proper, monotone scheme.Then F h (x, u(x), u(x) − u(•)) = 0 has a unique solution.
The Barles-Souganidis convergence framework has inspired many different monotone discretisations of the Monge-Ampère equation [2,5,12,14,17,22,24]. Typically, these discretisations involve a PDE operator that incorporates the convexity constraint (D 2 u(x) ≥ 0) through the use of a modified determinant satisfying (4) det Monotone discretisations can then be constructed for the modified Monge-Ampère equation ( 5) The Barles-Souganidis convergence framework does not immediately apply the second boundary value problem for Monge-Ampère (3), which does not satisfy the required comparison principle.However, this powerful framework has recently been extended to include this important PDE [3,16].By adding mild additional structure to the numerical methods, these works establish the convergence of consistent, monotone methods.Moreover, these works demonstrate that as long as the interior problem is set up carefully, there is a great deal of flexibility in the choice of numerical boundary conditions.Indeed, even boundary conditions that are inconsistent with (3) can nevertheless yield a numerical solution that converges to the true solution in L ∞ (Ω) for any compact subset Ω ⊂ X.
A similar convergence result was recently established by one of the authors for a class of more general Monge-Ampère type equations of the form (6) ψ(x, ∇u(x)) det(D 2 u(x) posed on the 2-sphere S 2 [18].In particular, this work showed that consistent, monotone methods converge with the addition of mild additional structure (which is easily incorporated into the numerical method).With these results in mind, a reasonable starting point for the numerical solution of non-local Monge-Ampère type equations ( 2) is to design consistent, monotone methods.Given recent progress in the numerical analysis of other Monge-Ampère type equations, there is great reason to hope that it will soon be possible to produce a proof of convergence for such methods.

PDE Formulation
The starting point of our numerical method for optimal transport between unequal dimensions is the PDE formulation (2) proposed by McCann and Pass.We begin this section with a formal overview of the derivation of the PDE, referring to [20] for complete details.We then propose a local level set formulation that highlights the elliptic structure of the equation and will lead into the numerical method introduced in section 4.
To formulate the optimal transport problem, we begin with a density function f supported on a set X ⊂ R n , a density g supported on a set Y ⊂ R m , and a cost function c : X × Y → R. We require the data to satisfy the mass balance constraint We assume here that n < m and seek a mapping T * : Y → X that minimises the cost (8) T * = argmin T #g=f Y c(T (y), y)g(y) dy.
We remark here that, under appropriate conditions on the data, we can hope for a well-defined mapping from the higher-dimensional space Y to the lower-dimensional space.However, we expect that any "mapping" from Y to X would have to be multi-valued as points in the lower-dimensional space must "spread" to fill the higher-dimensional space.
3.1.Feasibility.In a traditional Optimal Transport problem (n = m), the pushforward condition T #g = f can be expressed via the change of variables theorem as (9) f (T (y)) det(∇T (y)) = g(y).
We now seek a similar respresentation of the pushforward condition for the unequal dimension setting (n < m).
We begin by expressing the condition T #g = f in the form Next we make the particular choice We notice that if y ∈ T −1 (x) then T (y) = x, which allows us to rewrite the final integral as for every Φ ∈ L 1 (X).From this we conclude that almost everywhere, ( 13) We note that in the special case of equal dimensions (m = n), this coincides with the usual result of the change of variables formula (9).

Optimality.
Having characterised the feasibility condition T #g = f , we now seek a representation of the optimal mapping.In the simplest case of equal dimensions (m = n) and a quadratic cost (c(x, y) = 1 2 |x − y| 2 ), optimality is achieved by a mapping that can be expressed as the gradient of a convex function.
Our starting point in the current setting is the usual dual formulation of the Optimal Transport problem [27].This leads us to seek a c-convex dual pair u(x), v(y) satisfying u(x) + v(y) + c(x, y) ≥ 0 with equality when x = T (y).Equivalently, we require the inverse optimal mapping (which may be set-valued) to be contained in the c-subgradient of the function u, which we write as ( 14) From here, we obtain the following first-and second-order optimality conditions: Differentiating the first-order optimality condition ( 16) leads to the expression xx c(T (y), y)∇T (y) + D 2 xy c(T (y), y) = 0.Under a slightly stronger version of the second-order optimality condition (16) (assuming positive definiteness instead of merely semi-positive definiteness), we can obtain the following expression for the Jacobian of the optimal mapping: We can now combine this with the feasibility condition (13).We first compute the relevant Jacobian determinant (12): .
Finally, we substitute this into the pushforward condition ( 13) and simplify to obtain the non-local Monge-Ampère type equation We recall that this equation is coupled to the c-convexity constraint ( 16), , which plays the role of the convexity constraint (D 2 u(x) ≥ 0) in the traditional Monge-Ampère equation.

Local formulation.
As written, equations ( 18)-( 19) are non-local since the domain of integration involves the c-subgradient of u, which relies on global information (14).However, in many settings it is possible to replace this global definition with a local construction derived from the first-and second-order optimality conditions [20].
To accomplish this, we first define sets on which the first-and second-order optimality conditions ( 15)-( 16) hold: These are necessary conditions for optimality in (14), which leads to the containment However, when the problem has sufficient regularity (i.e.C 2 solutions), this containment becomes an equality up to H m−n -negligible sets [20].In this setting, (18) simplifies to (23) subject to the c-convexity condition (19).This simplification of the domain of integration reduces the problem to a local PDE.However, given the dependence of the domain of integration on second-order information, it is not at all obvious that this equation is actually elliptic and that the arguments developed in section 2 are reasonable.
To re-express this in a way that ellucidates the elliptic structure of this equation, we recall that the c-convexity constraint (19) This allows us to absorb the c-convexity constraint into the PDE and reformulate ( 19) and ( 23) into the single equation Following the arguments of [16, Lemma 2.6], we emphasise that solutions of this equation are automatically c-convex on the domain X since the density function f is positive and f and g are required to satisfy the mass balance constraint (7).This observation makes a further simplification possible.We notice that for any x ∈ X and y ∈ Thus we can replace the domain of integration in (24) with the set of points satisfying the first-order optimality conditions to obtain (25) f Importantly, equation ( 25) is now a degenerate elliptic Monge-Ampère type equation since the only second-order term is det + D 2 u(x) + D 2 xx c(x, y) , while all coefficients of this term are non-negative.
3.4.Level set formulation.While Monge-Ampère type equations (and their numerical solution) have received a great deal of attention in recent years, it is less clear how to deal with their integration over a manifold Y 1 (x, ∇u(x)) that itself depends on the solution of the equation.Here we propose a level-set formulation of ( 25) that transfers all of the u-dependence into the integrand.
First, we note that we can easily enlarge the domain of integration by introducing for each x ∈ X a Dirac delta function δ supported on the m−n dimensional manifold Y 1 (x, ∇u(x)).To make this precise, we let d Y1(x,∇u(x)) (y) denote the Euclidean distance between y and the manifold Y 1 (x, u (x)).This allows us to rewrite (25) as ( 26) Next, we seek an easier representation of the delta function since the distance function may not be easy to compute in practice.To facilitate this, for each x ∈ X we define the n level-set functions ( 27) We notice immediately that the domain of integration in ( 25) is equivalent to the intersection of the zero level-sets of the φ x i .That is, (28) . ., n}.Following [6], we can rewrite equation (26) in terms of these level set functions, though the specific form is dependent on the smaller dimension n.When n = 1, the equation is equivalent to When n = 2, the PDE becomes (30) where P v w denotes a projection of the vector w onto the subspace orthogonal to the vector v.
For higher values of the dimension n, this same procedure can be continued using additional orthogonal projections.
Finally, we note that the gradients of these level set functions are given by (31) which do not depend on the potential function u.

Numerical Method
In this section, we introduce a numerical method for solving the optimal transport problem between unequal dimensions, starting from the level-set PDE formulation introduced in section 3.For simplicity of the exposition, we will describe the method in the low-dimensional setting (n = 1, m = 2).However, all of the techniques we propose can be extended to higher dimensions.
Our goal here is to produce a consistent, monotone discretisation of the ODE operator in (29), which now takes the form Here we let (33) ψ(x, y) = g(y) ∇ y ∂c ∂x (x, y) which is a known function with no dependence on the potential u.
Since to date little is known about this equation except in the classical setting, we will focus our numerical analysis on the setting in which C 2 solutions exist.However, we expect that with ongoing development of the theory of weak solutions, we will be able to extend the numerical analysis to the non-smooth setting using the techniques of [16].
In particular, we make the following assumptions.
(H2) The mixed partial derivatives of the cost are bounded away from zero.(H3) The positive density functions f ∈ C 0 ( X), g ∈ C 2 ( Ȳ ) are bounded away from zero and infinity.(H4) The solution of (32 4.1.Discretisation of domains.We begin by discretising the sets X ⊂ R and Y ⊂ R 2 .Since Y need only contain the support of the density g, we may assume that Y is a square. To discretise X = [x 0 , x max ], we choose an integer N ∈ N and introduce the grid spacing h X = (x max − x 0 )/N .Then grid points in the domain are given by x i = x 0 + ih X for i = 0, . . ., N .
To discretise Y = [y 0 , y max ] 2 , we choose an integer M ∈ N and introduce the grid spacing h Y = (y max − y 0 )/M .Then grid points in the domain are given by y jk = (y 0 + jh Y , y 0 + kh Y ) for j, k = 0, . . ., M .
We will require that asymptotically, M N .That is, the one-dimensional set X needs to be more highly resolved than the two-dimensional set Y .
Throughout, we use the convention that the values of a function u defined on the one-dimensional grid can be represented as u i = u(x i ).

Discretisation of ODE.
We seek a finite difference discretisation of (32) of the form for i = 1, . . ., N − 1.For F h X ,h Y to be monotone, it must be a non-decreasing function of all of its arguments.This, in turn, is guaranteed if each term appearing in the summation is non-negative and a non-increasing function of the differences u i − u i−1 and u i − u i+1 .
The first term in the summation is easily approximated using standard centered differences via This trivially has the required monotonicity properties, which are preserved in the operations of addition with a given grid function and restriction to the positive part.
Discretisation of the delta function, which is supported on a curve in R 2 , is much more delicate.Indeed, naive extensions of typical discretisations of one-dimensional delta functions can lead to O(1) errors in the approximation [26].
We take as a starting point an approximate one-dimensional delta function defined in terms of the linear hat function via Typically, the support of this discrete delta function is chosen to be proportional to the underlying grid spacing, that is ∼ h Y .
It is tempting at this point to define the discrete delta function in (34) by where D h X x u i denotes some finite difference approximation of u (x i ) and m ∈ N is some fixed positive integer.However, there are two problems with this: (1) fixing the value of in relation to the grid parameter h Y can lead to O(1) errors [26] and (2) because of the nonlinearity, we cannot expect the overall approximation to be monotone even if D h X x u i is itself monotone.Fortunately, the first problem can be resolved by allowing to vary in space [10].We begin by discretising the delta function in the y variable only, allowing x to be continuous.Then we can introduce the approximate delta function where the support parameter in the linear hat approximation to the delta function is given by (38) We note that since the cost c is a given function, this spatially varying parameter can be computed a priori.
To achieve a monotone discretisation of this, we need to exploit the structure of the linear hat function.To this end, we notice that This is automatically non-negative.We can also produce a fully discrete approximation with the required monotonicity properties by utilising a careful combination of forward and backward differences.This leads us to propose with allowed to vary according to (38).
4.3.Boundary conditions.While Optimal Transport PDEs such as ( 25) typically do not require a traditional boundary conditions [16], the corresponding finite difference approximation (34) will require some additional conditions to be supplied at the boundary grid points x 0 and x N .
We have a great deal of freedom in the choice of condition to enforce at the remaining boundary point.For the usual second boundary value problem for the Monge-Ampère equation, it was shown in [16] that even enforcing a possibly inconsistent boundary condition such as u(x) = 0, x ∈ ∂X would lead to a convergent numerical method.However, this type of artificial boundary condition will result in a boundary layer in the computed solution.
To find a more natural boundary condition, we recall that the integrand in ( 29) is supported on the set of points y for which u (x) + ∂c ∂x (x, y) = 0.
In, particular, at the boundary point x i there must be some The cost function c (and its partial derivatives) are all known functions.However, the relevant values of y are not obvious.We choose to proceed using a monotonicity assumption and enforce the Neumann boundary conditions These can easily be dicretised in a consistent, monotone way by defining the approximations 4.4.Consistency and monotonicity.We now verify the consistency and monotonicity of the approximation scheme described by ( 34)-( 40).First, we note that monotonicity was built into the approximation scheme by design.
Next, we verify consistency of the scheme.Because of the interplay between the finite difference discretisations and the approximation of the Dirac measure, this is a more delicate calculation that requires a careful balance between the discretisation parameters h X and h Y .
Lemma 9 (Consistency).Under the assumptions of Hypothesis 7, consider the approximation scheme (34)-( 40) where the grid parameters are chosen so that h X /h Y → 0 as h Y → 0. Then the scheme is consistent with the ODE (32).Moreover, the formal truncation error is given by O(h 2 Y + h X /h Y ).Before proving this lemma, we introduce some notation designed to simplify the representation of the semi-and fully-discrete versions of the approximation.
Denote the coefficient of the delta distribution in (32) by be the discretised version.
We recall that the level set function used to locate the curve Y 1 (x, u (x)) is given by while its discretised version satisfies Using this notation, the integral appearing in equation (32) takes the form (45) while its semi-and fully-discrete approximations are (46) Proof of Lemma 9. We begin by noting that the width of the support of the approximate delta function can be bounded in terms of the spatial discretisation parameter h Y .Let Then (h X , ∇φ x (y)) ≤ Kh Y for every x, y.
We also notice that many of the terms in the sum of (46) are zero.To identify the non-zero terms, we define Since φ x (y) is uniformly continuous in both x and y and its zero level-set defines a curve in Y , the cardinality of this set will be We can also try to identify the terms that are non-zero in the fully discrete sums in (47).To this end, we set Any point (j, k) in this set will satisfy for some y ∈ I h Y (x i ).This means that for each such y, we need to consider possible neighbouring grid points that satisfy the conditions in (49).Adding up all possible points y ∈ I h Y (x i ) leads to the bound Using the results of [10], we know that the discretisation of the Dirac delta function satisfies Finally, we can consider the approximation accuracy of the fully discrete version of the integral (47).
In the last step, we notice that the error term will only be non-zero at indices where the summand is non-zero in either the semi-or fully-discrete approximation of ( 46)-(47).
We can further simplify this to Finally, we recall that h Y ≤ = O(h Y ) so that By hypothesis, h X h Y so that the approximation of the integral is consistent with a formal truncation error of O(h 2 Y +h X /h Y ).Since the remaining terms in the approximation (34) share the same consistency error, this completes the proof.
We also make the remark here that the formal consistency error in the scheme can achieve an optimal value of O(h In particular, this means that the two-dimensional set Y can (and should) have a much lower resolution than the one-dimensional set X. We also note that typical convergence analysis for Monge-Ampère equations in optimal transport involves compactness arguments without error bounds; actual computational error need not coincide with the formal truncation error.4.5.Normalisation.Finally, we notice that solutions of the ODE (32) are only determined up to additive constants.The particular constant is not typically important to applications, as the information about the transport mapping is encoded in u (x).
For the purposes of this article, we will seek to approximate the mean-zero solution of (32).There is no particular reason that the approximate solution we obtain from (34)-(40) should have mean zero, or even that the mean need be the same given different values of the discretisation parameters h X , h Y .This is typical of Monge-Ampère equations in optimal transport, and an easy solution is to simply shift the approximate solution [18].
Thus we propose the following two-step procedure: (1) Let v h X ,h Y be the unique solution of for all i = 0, . . ., N .(2) Normalise to

Computational Results
Finally, we validate our proposed numerical method on several computational examples.
In each example, we take as our domain X = [0, 1], which is discretised using N grid points.The computational target set is taken to be the smallest possible square that encloses the actual target Y .
The grid spacing parameters are chosen to satisfy h Y ≈ h X , with slight deviations from equality used to ensure that this leads to an integer M number of grid points along each dimension of Y .
The resulting systems of nonlinear equations were solved using Newton's method.
5.1.Rectangular target.The first example we consider involves a rectangular domain Note that this domain is rotated, so it will not line up with the grid used to discretise the y variables.
We solve (1) using the density functions See Figure 1 for a visual of the problem data.We can easily verify that the exact mean-zero solution of the resulting Monge-Ampère equation is u(x) = e x − e + 1.The resulting error is displayed in Figure 2. We notice that for this particular example, the error decreases, but not monotonically.It is well known that monotone finite difference schemes do not have to yield monotonic convergence.In this case, it appears to be the simplicity of the geometry that leads to the oscillations in error.This is because for some choices of discretisation parameters M, h Y , the grid is by chance quite well-aligned with the relevant curve Y 1 (x, u (x)), the contours of g, and the boundaries of the target Y .This, in turn, can lead to artificially low error for certain values of the discretisation parameters.
and the quadratic cost as in the previous example.This time, we use density functions that are allowed to vanish to zero at the boundaries of X and Y .This setting is known to lead to a loss of uniform ellipticity in the Monge-Ampère equation, which is often a challenge for numerical methods.We choose the particular density functions See Figure 3 for a visual of the problem data.In this case, the exact solution is the constant function Although the exact solution is constant, we cannot expect to recover this exactly (which would occur with a typical finite difference discretisation of a standard Monge-Ampère equation).This is because the approximation scheme will still involve (approximate) integration of a spatially varying function over a non-trivial curve, which will introduce numerical errors.
The error obtained for this example is displayed in Figure 4.In this example, the error appears to decrease monotonically, likely because the contours of g are no longer straight lines that can align with the grid.Despite the degeneracy of this example, there is no difficult in computing this solution.
We solve (1) using the density functions otherwise .
These densities are also allowed to vanish along part of the boundary.In this example, we use the cost function See Figure 5 for a visual of the problem data.The resulting error is displayed in Figure 6.Once again, we observe clear monotonic convergence.The more complicated geometry does not any difficulty for the numerical method.

Conclusions
In this article, we have proposed a numerical method for a Monge-Ampère type equation arising in the problem of optimal transport between unequal dimensions.In to the strong nonlinearity of this problem, the equation involves an extra layer of complexity because it requires integrating the equation over a manifold that itself depends on the solution of the Monge-Ampère equation.
We have proposed a level-set formulation of this PDE, which allows the integral to be posed over a much simpler domain.However, the resulting equation still involves (1) a fully nonlinear elliptic operator of Monge-Ampère type and (2) integration against a Dirac delta distribution whose support is dependent on the solution gradient.We have proposed a very novel method for numerically approximating the resulting equation.For simplicity, we focus on the setting of optimal transport from 1D to 2D.However, all of the techniques we proposed can be easily extended to higher dimensions.The numerical method utilises monotone finite difference approximations of the Monge-Ampère type operator.These are combined with a very careful discrete version of the Dirac delta function, which is allowed to have a spatially varying support width in order to preserve consistency.The level set function that is input into this discrete delta function is discretised using a careful combination of forward and backward differences in order to preserve monotonicity.
The resulting numerical method is both consistent and monotone.For more wellunderstood Monge-Ampère type equations in optimal transport, consistency and monotonicity have proven to be the key ingredients in proofs of convergence [3,16,18].With ongoing development of the theory for the non-standard Monge-Ampère equation considered in this article, there is great reason to hope that a convergence proof can be constructed for the method proposed here.This will be an avenue of continued research.
We have performed computational tests using several challenging examples, which validate the performance of the new method.
A natural avenue for future work is to extend this numerical method to higher dimensions.An advantage of the current method is that the higher-dimensional set Y can be discretised using a fairly coarse grid, which is important for keeping the computational cost reasonable.Future work will further improve the computational cost by devising better strategies to estimate the support of the discrete delta function, allowing for a great reduction in the cost of the numerical quadrature.We expect the resulting method to be easy to implement, while having a computational cost equivalent to integrating over an m − n-dimensional manifold (instead of the full m-dimensional target set) at each point in the lower-dimensional domain X.

Figure 1 .
Figure 1.The density functions (a) f and (b) g for the rectangular target.

Figure 3 .
Figure 3.The density functions (a) f and (b) g for the rectangular target.

Figure 5 .
Figure 5.The density functions (a) f and (b) g for the rectangular target.The exact mean-zero solution of the resulting Monge-Ampère equation is u(x) = − cos πx 2 + 2 π .

4 Figure 6 .
Figure 6.Maximum error in u for the curved target.