V OLUMETRIC O PTIMAL T RANSPORTATION BY F AST F OURIER T RANSFORM
A BSTRACT
The optimal transportation map finds the most economical way to transport one probability measure to another, and it has been applied in a broad range of applications in machine learning and computer vision.
By the Brenier theory, computing the optimal transport map is equivalent to solving a Monge-Amp ` ere equation, which is highly non-linear.
Therefore, the computation of optimal transportation maps is intrinsically challenging.
In this work, we propose a novel and powerful method, the FFT-OT (fast Fourier transform-optimal transport), to compute the 3-dimensional OT problems.
The method is based on several key ideas: first, the
Monge-Amp ` ere equation is linearized to a sequence of linear elliptic PDEs with spacial and temporal variant coefficients; second, the obliqueness property of optimal transportation maps is reformulated as a Neumann boundary condition; and third, the variant coefficient elliptic PDEs are approximated by constant coefficient elliptic PDEs and solved by FFT on GPUs.
We also prove that the algorithm converges linearly.
Experimental results show that the FFT-OT algorithm is more than a hundred times faster than the conventional methods based on the convex geometry.
Furthermore, the method can be directly applied for sampling from complex 3D density functions in machine learning and magnifying the volumetric data in medical imaging.
I NTRODUCTION
Optimal transportation (OT) transports one probability measure to another in the most economical way, and it plays a fundamental role in areas like machine learning Courty et al.
; Altschuler et al.
computer vision Arjovsky et al.
; Tolstikhin
2020), and computer graphics Solomon et al.
; Nader & Guennebaud (2018).
Given a Riemannian manifold X , all the probability distributions on X form an infinite dimensional space P ( X ) .
Given any two distributions µ, ν ∈ P ( X ) , the optimal transportation map defines a distance between them, and the McCann interpolation McCann (1997) defines the geodesic connecting them.
Hence optimal transportation equips P ( X ) with a Riemannian metric and defines its covariant differentiation, which provides a variational calculus framework for optimization in it.
As the optimal transportation problem is highly non-linear, it is quite challenging to compute the OT maps.
Recently, researchers have developed many algorithms.
The geometric variational approach
Aurenhammer et al.
(1998); Gu et al.
; Levy (2015) based on the Brenier theorem Brenier (1991) is capable of achieving high accuracy for low dimensional problems, but it requires complicated geometric data structure and the storage complexity grows exponentially as the dimension increases.
The Sinkhorn method Cuturi (2013) based on the Kantorovich theorem adds an entropic regularizer to the primal problem and can handle high dimensional tasks, but it suffers from the intrinsic approximation error.
We propose a novel method to tackle this challenging problem through Fast Fourier Transformation (FFT).
According to the Brenier theorem Brenier (1991), under the quadratic distance cost, the optimal transportation map is the gradient of the Brenier potential, which satisfies the MongeAmp`ere equation.
With the continuity method Delano¨e
(1991), the Monge-Amp`ere equation can be linearized as a sequence of elliptic partial differential equations (PDEs) with spacial and temporal variant coefficients.
By iteratively solving the linearized Monge-Amp`ere equations, we can obtain the OT map.
Specifically, we propose to approximate the linearized Monge-Amp`ere equation by constant coefficient elliptic PDEs and solve them using the FFT on GPUs.
Our proposed FFT-OT method has many merits: (i) it is generalizable for arbitrary dimension; (ii) it has a linear convergence rate, namely the approximation error decays exponentially fast; (iii) in each iteration, the computational complexity of FFT is O ( n log n ) , thus our algorithm can solve large scale OT problems; and (iv) it is highly parallelable and can be efficiently implemented on GPUs.
We demonstrate the efficiency of the FFT-OT algorithm by solving the volumetric OT problems for machine learning and medical imaging applications including sampling from given 3D density functions and volumetric magnifier.
The algorithm also has its own limitations: (i) although it can be generalized to any dimensions, the storage complexity increase exponentially with respect to the dimension, so its power is limited by the memory size of the GPUs; (ii) Since the algorithm uses FFT, the current version of the method only works well for continuous density functions. (iii)
In this work, we mainly focus on the computation of the OT map from the uniform distribution to another arbitrary continuous distribution. To extend the method to find the OT map between any two continuous measures, we can compute two OT maps from the uniform distribution to the both continuous measures, then combine them together. The combination will give a reasonable approximation of the OT map Nader & Guennebaud (2018).
Though Lei and Gu Lei & Gu (2021) also uses FFT to solve the 2-dimensional OT problem, our method differs their works in the following two aspects: (i) Lei and Gu’s method uses the fixed point method to compute the 2D OT problems, ours is based on the linearization of the Monge-Amp ` ere operator to solve the 3D OT problems, these are two different methodologies in PDE theory; (ii) In our paper, we also provide the theoretical convergence analysis of the proposed method.
For more detailed analysis and related work, please refer to the appendix A.
O PTIMAL T RANSPORTATION T HEORY
In this section, we review the fundamental concepts and theorems of the OT problem and the MongeAmper ` e equation, more details can be found in Villani (2008).
Optimal Transportation Map and the Monge-Amp`ere equation Suppose the source domain Ω is an open set in R d with the probability measure µ , the target domain Σ is with the probability measure ν .
Both µ and ν have density functions
dµ ( x ) =
f
( x ) dx and dν ( y )
=
g ( y ) dy , respectively, with the equal total mass: (cid:82) Ω f ( x ) dx =
(cid:82) Σ g
( y ) dy, which is called the balance condition .
Suppose T : Ω → Σ is a measurable map.
The mapping T is called measure preserving and denoted as T # µ =
ν if the following relation
[Equation]
for every Borel subset A ⊂ Σ .
A cost function c : Ω × Σ → R measures the transportation cost for transporting the unit mass from x ∈ Ω to y ∈ Σ .
Problem 1 (Monge) .
The optimal transportation problem finds the measure preserving map with the minimal total transportation cost,
[Equation]
The solution to the Monge’s problem is called the optimal transport map between µ and ν .
The existence, uniqueness and regularity of OT maps depend on the boundedness and the continuity of the density functions, the convexity of the supporting domains, the continuity of their boundaries, and the cost function.
In our current work, we focus on the similar situation in Saumier et al.
• The cost function is quadratic Euclidean distance c ( x, y ) =
∥ x − y ∥ 2 / 2 ; • The supports of the source and the target measures are the canonical cube Ω
=
[ − 1 , 1] 3 , which is uniformly convex; • The source and the target measures µ, ν are absolutely continuous with respect to the Lebesgue measure, their densities
f, g are positive and bounded away from zero;
[Equation]
m
[Equation] f, g [Equation] M, and f, g are of class C α (Ω) , • The boundary condition is second boundary condition (OT boundary condition), T (Ω) =
Ω .
Then according to (Villani (2003) Theorem 14.4, Saumier et al.
Theorem 2.1), the OT maps
T : Ω → Ω exists and is unique and invertible ( µ a.e), and the Brenier potential is of class C 2 ,β (Ω) form some 0 < β < α .
Theorem 2.
Assume that Ω , µ, ν, f and g are defined as above.
Then there exists a convex function u : Ω → R , u ∈ C 2 ,β (Ω) for some 0 < β < α , such that ∇ u pushes µ forward to ν , ( ∇ u ) # µ = ν .
Moreover, ∇ u is unique and invertible ( µ a.e), and its inverse ∇ v satisfies ( ∇ v ) # ν = µ .
We call such a convex function u the Brenier potential , it satisfies the Monge-Amp ` ere equation,
[Equation]
with the boundary condition ∇ u (Ω) = Σ .
Then finding the optimal transportation map is equivalent to solving the corresponding Monge-Amp ` ere equation.
In the current work, the target measure is always the Lebesgue measure, and the source density f is of class C 2 ,α (Ω) .
Amp` ere
Operator The Monge-Amp ` ere operator is defined as
[Equation]
] = det D 2 u, which is highly non-linear.
It can be linearized as following:
[Equation]
where Adj ( A ) is the adjoint (co-factor) matrix of A , Adj ( A ) :
= det ( A ) A − T .
Therefore the linearized Monge-Amp`ere operator is defined as
[Equation]
where ( u pq )
=
Adj ( D 2 u ) is the adjoint matrix of the Hessian of u , and ∂ p ∂ q :=
[Equation]
Continuity Method For simplicity, we assume the source domain coincides with the target domain, that is Ω = Σ , and the target density is g ( x ) ≡ 1 .
The Monge-Amp ` ere equation Eqn.
(2) is simplified as det D 2 u ( x ) =
f
( x ) .
Define a flow of density as
[Equation]
The corresponding flow of the Brenier potentials is u ( x, t ) :
Ω × [0 , 1] → R ,
[Equation]
where D 2 x u ( x, t ) is the Hessian of u ( x, t ) with respect to x , and u ( x, 1) is the solution to the initial
Monge-Amp`ere equation Eqn.
Take the derivative w.r.t.
time t on both sides of the linearized
Monge-Amp`ere operator Eqn.
(4), we obtain an elliptic PDE with the spacial and temporal variant coefficients of the unknown v ( x, t ) :
= ˙ u ( x, t ) , namely the “velocity” of the Brenier potential, [Equation]
[Equation]
At time t = 0 , the initial Brenier potential is known as u ( x, 0)
=
∥ x ∥ 2 .
Suppose at time t , we have obtained u ( x, t ) already, then we can compute the adjoint matrix u pq ( x, t ) of the Hessian D 2 x u ( x, t ) , and solve Eqn.
(6) to get the velocity v ( x, t ) = ˙
u ( x, t ) .
In turn, we move forward to time t + δt , and update u ( x, t + δt ) by u ( x, t ) + ˙ u ( x, t ) δt .
By repeating this procedure, eventually we reach time t = 1 and obtain the solution u ( x ) :
= u ( x, 1) to the initial Monge-Amp`ere Eqn.
Obliqueness Boundary Condition Suppose the boundary of Ω is C 1 almost everywhere, therefore at a C 1 point x ∈
∂ Ω , the outer normal n ( x ) is well defined.
For almost every boundary point x ∈ ∂ Ω , the obliqueness condition is represented as
[Equation]
Suppose Ω is a cuboid and has 6 faces, if a boundary point x ∈ ∂ Ω is on a face, by the cyclic monotonicity of the map and the strict convexity of u Villani (2008), its image ∇ u ( x ) must be on the same face of x , namely,
[Equation]
We can rewrite the Brenier potential as u ( x 1 , x 2 , . . .
, x d ) = 12 (cid:80) di
=1 x 2 i + v ( x 1 , · · · , x d ) , then ∇ u ( x ) − x = ∇ v ( x ) .
By Eqn.
(8), v ( x ) satisfies the Neumann boundary condition,
[Equation]
Similarly, the velocity of the (modified)
Brenier potential v in Eqn.
(6) also satisfies the Neumann boundary condition.
The analysis about the existence and regularity of the solutions to Eqn.
(6) with boundary condition Eqn.
(9) can be found in the supplementary material.
C OMPUTATIONAL A LGORITHM
Here we introduce the 3-dimensional FFT-OT algorithm, which can be generalized to any dimensions.
We approximate the Monge-Amp ` ere equation by a sequence of constant coefficient elliptic
PDEs, and solve them by FFT on GPUs.
More detailed analysis about the solution of the discretized
Monge-Amp ` ere equation, and the proofs of the lemmas and theorems are given by Appendix B.
3. C ONTINUITY M ETHOD FOR S OLVING THE M ONGE -A MP ` ERE E QUATION
By using the continuity method, we can solve the Monge-Amp ` ere equation iteratively.
For simplicity, we assume the target measure is the Lebesgue’s measure with g ≡ 1 .
At the n -th iteration, the
Brenier potential is represented as 12 ∥ x ∥ 2 + u n ( x ) , its Hessian matrix is H n ( x ) :
= I + D 2 u n ( x ) , the corresponding density function is defined as the determinant of the Hessian ρ n = det ( H n ) , and the velocity of the Brenier potential is v n ( x ) .
In the beginning, the Brenier potential u 0 ( x ) is zero, the Hessian is H
= I and the density is ρ 0 = 1 .
At the n -th step, we compute the adjoint matrix [ H pqn ( x )] of the Hessian matrix H n ( x ) for any x ∈ Ω .
According to Eqn.
(3), the velocity v n ( x ) satisfies the variant coefficient elliptic PDE induced by the linearized Monge-Amp`ere operator,
[Equation]
Note that the right hand side of Eqn.
(6) is the difference between the initial and the target densities, whereas here it is replaced by the difference between the initial and the current densities.
The step length parameter τ ≥ 1 can be chosen to guarantee the convergence Loeper & Rapetti (2005).
The elliptic PDE Eqn. (10) is with spatially variant coefficients.
Although the traditional finite element method (FEM) can solve it using the GMRES algorithm Saad (2003), this algorithm can not be directly accelerated by GPUs.
To overcome this difficulty, we approximate Eqn.
(10) by a much simpler elliptic PDE with constant coefficients, which can be directly solved using the following
FFT-OT Algorithm Alg.
1 on GPUs.
At the n -th iteration, after obtaining the adjoint matrix [ H pqn ( x )] , x ∈ Ω , we compute the mean adjoint matrix [ ¯ H pqn ( x )]
[Equation]
and replace the elliptic PDE Eqn.(10) with variant coefficients by the elliptic PDE with constant coefficients,
[Equation]
where DMA is called the mean linearized Monge-Amp ` ere operator.
Then we solve the constant coefficient elliptic PDE Eqn.
(12) by FFT Algorithm Alg.
Although the original variant coefficient PDE Eqn.
(10) is replaced by its constant coefficient approximation
Eqn.
(12), the algorithm still converges to the solution with a linear convergence rate.
This replacement allows the whole algorithm to be solved by FFT on GPUs, which greatly improves the computational efficiency.
Theorem 3 (main) .
Given a domain Ω ⊂ R d , which is a canonical cuboid Ω =
[ − 1 , 1] d , and a positive density function f : Ω → R with the balance condition (cid:82) Ω f ( x ) dx = (cid:82) Ω dx, suppose the mirror reflection extension Eqn.
(14) of f to the flat torus ˜ f
: T n → R is C α , α ∈ (0 , 1) , then the
Monge-Amp`ere equation,
[Equation]
can be solved using the FFT-OT Algorithm Alg.
In particular, one can choose the step length parameter τ , such that there is a constant 0 < γ < 1 that the approximation error satisfies
[Equation]
namely the algorithm has a linear convergence rate.
3. FFT S OLVER FOR C ONSTANT C OEFFICIENT E LLIPTIC PDE S
To solve the constant coefficient elliptic PDE Eqn.
(12), we first extend the PDE to the flat torus by mirror reflection, then discretize the domain and compute the differential operators by central difference scheme.
Finally the PDE is converted to algebraic equations in the frequency domain by
FFT and can be efficiently solved on GPUs.
Extension by Mirror Reflection Suppose Ω = [0 , 1] 3 and f : Ω → R are given, we extend Ω to ˜Ω =
[ − 1 , 1] 3 and f to ˜ f : ˜Ω → R by mirror reflection
[Equation]
By definition, ˜ f satisfies the periodic boundary condition and can be treated as a function defined on the flat torus T 3 .
˜Ω is one of the fundamental domain of T 3 .
The constant coefficients a p,q keep unchanged.
Then we solve the following constant coefficient elliptic PDE Eqn.
(18) L [˜ u ] = ˜
f
with the periodic boundary condition.
Finally, the restriction of ˜ u on Ω gives the initial solution u to
L [ u ] = f with Neumann boundary condition.
In the following, to avoid using overly complicated symbols, we use ( u, f, Ω) to represent (˜ u, ˜ f, ˜Ω) for simplicity.
Tessellation Suppose Ω =
[ − 1 , 1] 3 is the canonical cube (a fundamental domain of a flat torus), we tessellate it to the regular cells, and the centers of the cells form a grid M × N × L .
The Brenier potential u : Ω → R is discretized to a tensor u i,j,k with { i, j, k } ∈ { 0 , . . .
, M − 1 } × { 0 , . . .
, N − 1 } × { 0 , . . .
, L − 1 } .
The spacial step lengths are ( h x , h y , h z ) =
(2 /M, 2 /N, 2 /L ) .
The coordinate of each sample point ( x i , y j , z k ) is ( x i , y j , z k ) =
( − 1 + h x
( i + 1 / 2) , − 1 + h y
( j + 1 / 2) , − 1 + h z ( k + 1 / 2)) .
The periodic boundary condition is then formulated as
[Equation]
Finite Difference Differential Operator We use the standard central differences to compute the differential operators.
The first order derivative D x is approximated by
[Equation]
where the index i + 1 means i + 1 modulus M .
The operators D y , D z are defined in a similar way.
The second order derivative operator D xx and D xy are approximated by
[Equation]
The other operators D
yy , D zz , D yz and D xz are defined similarly.
Discrete Fourier Transformation
The discrete Fourier transformation (DFT) of u i,j,k is given by
[Equation]
L − 1 (cid:88) k
= u i,j,k ˆ ω mnl (16) u i,j,k =
MNL
M − 1 (cid:88) m =
N − 1 (cid:88) n = L − 1 (cid:88) l = ˆ u m,n,l ω mnl (17) where ˆ ω mnl =
e − ι 2 πmiM e − ι 2 πnjN e − ι 2 πlkL , ω mnl = e ι 2 πmiM e ι 2 πnjN e ι 2 πlkL and ι = √− 1 , { m, n, l } are the indices of the frequency coefficients.
By using DFT, the differential operators are converted to algebraic operators in the frequency domain.
Lemma 4.
Suppose the discrete function is u i,j,k , with the discrete Fourier transformation Eqn.
(16) and Eqn.
(17), by using the central difference scheme, the first order differential operator is given by
[Equation]
the second order differential operators are represented by
[Equation]
The other differential operators D
[Equation]
yy , D zz , D yz and D xz are also represented accordingly.
The detailed proofs can be found in the supplementary material.
FFT Solver Suppose we want to solve an elliptic PDE with constant coefficients on Ω ⊂ R 3 ,
[Equation]
with the periodic boundary condition, where a p,q , b r , c are constants, the matrix ( a p,q ) is positive definite, namely the PDE is uniformly elliptic.
By the discrete Fourier transformation F , we convert the differential equation to an algebraic equation in the frequency domain,
[Equation]
By applying Lemma 4 and defining
[Equation]
We have the algebraic equations in frequency domain,
[Equation]
With ˆ u m,n,l ’s, we can easily obtain u i,j,k ’s by the Inverse Discrete Fourier Transform (IDFT), which means solving the constant coefficient elliptic equation.
The algorithm is described in Alg. 2. Algorithm 2:
Input: Domain Ω =
[ − 1 , 1] 3 , M, N, L , { a pq } , b r , c , function f with the periodic boundary condition
Output: Solution u to the elliptic PDE Eqn. (18) Discretize the domain Ω to a M × N × L grid; Sample the function f to f i,j,k ; Compute FFT using Eqn.
(16), { ˆ f m,n,l } ← FFT ( { f i,j,k } ) ; for ( m, n, l ) ∈ [0 , M − 1] × [0 , N − 1] ×
[0 , L − 1] do
Compute the factor λ m,n,l using Eqn.
(19); if λ m,n,l is 0 then ˆ u m,n,l ← 0 ; else ˆ u m,n,l ← ˆ f m
The FFT for solving the constant coefficient elliptic PDE can be efficiently computed with GPUs.
Moreover, the algorithm Alg.
2 solves the constant coefficient elliptic PDEs with a periodic boundary condition , which can be generalized to solving the same type of PDEs with Neumann boundary condition by extending the PDE to the flat torus T 3 using mirror reflection Eqn.
E XPERIMENTAL R ESULTS
In this section, we firstly show that the our proposed FFT-OT algorithm converges linearly and runs 100 × faster than the conventional convex geometry based solver Levy (2015), then demonstrate the method in two applications: 3D adaptive sampling and Volume Magnifier.
All the algorithms are developed using generic C++ with CUDA Toolkit.
All the experiments are conducted on a Windows laptop with Intel Core i7-7700HQ CPU with 16 GB memory and NVIDIA GeForce GTX
Graphics Cards.
More experiments can be found in Appendix C.
4. R UNNING T IME AND C ONVERGENCE A NALYSIS
To show the performance of the proposed method, we experiment on the density functions defined by the Gaussian mixture models.
To be specific, the domain is a cube Ω
=
[0 , 1] 3 , the 3-dimensional density function defined on Ω is set to be f ( x ) = (cid:80) 30 i =1 p i N ( µ i , Σ i ) , where N ( µ i , Σ i ) represents
Gaussian distribution with mean µ i and variance Σ i = diag (
σ 2 i 0 , σ 2 i 1 , σ 2 i 2 ) .
µ i ∈ R 3 is uniformly sampled from [0 , 1] 3 , σ ij is uniformly sampled from [0 , 0 . 5] , p i ∈ R is uniformly sampled from [0 . 2 , 1] and normalized such that (cid:82) Ω f ( x ) dx = 1 .
Thus the source distribution µ is a complicated
Gaussian mixture distribution restricted on Ω .
Then by mirror reflection in Sec. 3.2, we obtain the complex density function which is defined on [ − 1 , 1] 3 and satisfies the periodic boundary condition.
We directly use the FFT-OT algorithm Alg.
1 to solve the linearized Monge-Amp`ere equation.
With the approximation error threshold ε = 1 .
× 10 − 6 and the resolution 256 × 256 × 256 , the running time for our FFT-OT algorithm with double precision on GPU is less thanseconds.
The conventional convex geometry based algorithm for 3D optimal transportation Levy (2015) can neither handle such large data sets nor be implemented on GPUs.
It can only compute OT map with resolution no greater than 100 × 100 × 100 on our system, which takes about 2700 seconds.
When handling problem with 128 × 128 × 128 resolution, our FFT-OT consumes about 20 .
3 seconds, which is 130 × faster than the power diagram based method Levy (2015).
Fig.
1 shows the approximation error for the above Gaussian mixture density with respect to iterations, namely log ∥ f − ρ n ∥ 22 .
Our algorithm does converge linearly and the result is consistent with the prediction Eqn.
(13) in Thm. 3.
Therefore, this experiment validates the theorem.
4. 3D A DAPTIVE S AMPLING
Generating random samples matching a given density function plays an essential role in the applications like Monte-Carlo integration or stippling.
Efficiently obtaining high quality samples is still an on-going research topic Bauer et al.
; Perrier et al.
And optimal transportation has been successfully applied for generating high quality 2D samples de Goes et al.
; Nader & Guennebaud (2018).
Most of the current research focuses on generating 2D samples fitting the given density function.
Here we apply the proposed 3D FFT-OT method to generate high quality 3D samples according to the given complex density functions.
To the best of our knowledge, it is the first work that uses OT to sample from 3D density functions.
Suppose the source probability distribution dµ ( x ) =
f ( x ) dx is defined on Ω =
[0 , 1] 3 with µ (Ω) = 1 .
The target distribution dν ( y ) = dy is the uniform distribution.
We use the FFT-OT algorithm
Alg. 1 to compute the OT map T : Ω → Ω , T # µ = ν .
The domain is tessellated to a 256 × 256 ×grid.
For each x ijk , i, j, k ∈ { 0 , 1 , . . .
, 255 } , the image T ( x ijk ) can be obtained.
We use { T ( x ijk ) } as vertices to compute the Delaunay triangulation of Ω .
Then representing the OT map T : (Ω , µ ) → (Ω , ν ) as a piecewise linear map, the restriction of T on each tetrahedron is a linear map.
Then the inverse OT map T − 1 : (Ω , ν ) → (Ω , µ ) is also a piecewise linear map.
Namely, given a grid point y mnl , we can find a tetrahedron containing it.
Suppose the vertices of the tetrahedron are { T ( x i ) , T ( x j ) , T ( x k ) , T ( x l ) } , then y mnl is computed as
[Equation]
where the non-negative barycenter coordinates satisfy λ i + λ j + λ k + λ l = 1 .
Then the image of the inverse OT map is given by
[Equation]
We generate random samples { y k } according to the uniform distribution ν on Ω , then their images { T − 1 ( y k ) } are the desired random samples following the distribution µ .
In our experiment, we use the same Gaussian mixture settings of the density function as Sec. 4.1.
Fig.
2 visualizes the generated samples.
We randomly pick the k -th slice along the z -direction from the discretized volume, draw the source density function on this slice, and use pixel intensity to represent the density in Fig. 2(a).
(i)
We uniformly generate 100 k random samples { y k } ⊂ Ω , and obtain the desired random samples by applying the inverse OT map { T − 1 ( y k ) } .
(ii) We also set { y k } as the grid centers of Ω and obtain the corresponding samples of the desired distribution µ .
The samples around the k -th slice of both sampling strategies are plotted in Fig. 2(e) and Fig. 2(f).
By visual comparison, it is obvious that the distributions of Fig. 2(e) and Fig. 2(f) are consistent with the density function in Fig. 2(a).
The consistency of the boundary of Fig. 2(e) and (f) and Fig.
2(a) also verifies the obliqueness boundary condition of the Monge-Amp`ere equation.
To further show the performance of the proposed method, we compare it with the classical sampling methods, namely rejection sampling, the Metropolis-Hastings algorithm Bishop (2006) and the slice sampling [Figure]
Neal (2003), shown in Fig. 2(b), Fig. 2(c) and Fig. 2(d).
To quantitatively compare the sampling results, we use the Chi-square goodness-of-fit test, which firstly groups the data and then computes the L 2 norm of the difference between the actual number of observations in each group and the expected number of observations.
In our experiment, we set the group number to 64 × 64 × 64 and use 500K samples to make the comparison.
The corresponding L 2 norm of each method is shown in the top-right of the corresponding figure.
We can see that the both sampling strategies of our method give smaller scores than the classical ones.
4. V OLUMETRIC M AGNIFIER
In reality, physical magnifiers can only magnify planar images.
In medical image processing, it is highly desirable to magnify certain regions of the 3D MRIs or CT images.
Our algorithm can address such requests with the user prescribed region of interest (ROI) and magnifying factor.
Suppose the ROI is a symmetric region with the center (¯ x, ¯ y, ¯ z ) ∈ Ω and the radius σ x , σ y , σ z in different directions.
The density function f of the source measure µ is defined as
[Equation]
We compute OT map T : (Ω , µ ) →
(Ω , ν ) , where ν is the uniform distribution.
Similar to the method in 3D adaptive sampling, we compute the Delaunay triangulation of the images { T ( x ijk ) } , then the OT map T is represented as a piecewise linear map.
The inverse optimal transportation map
T − 1 : (Ω , ν ) → (Ω , µ ) is also piecewise linear.
For each grid point y mnl ∈ Ω , we use Eqn.
(20) to find its pre-image.
Similarly, its corresponding intensity I mnl is computed by linear interpolation.
Then we obtain the new volumetric data { I mnl } with the magnified ROI and visualize the result with Voreen Meyer-Spradow et al.
Fig.
3 demonstrates our volumetric magnifier by magnifying an aneurysm on blood vessel Hansen &
Johnson (2004).
We choose the aneurysm region as the ROI.
The first column gives the snapshot of the blood vessel, and the yellow circle denotes the location of the aneurysm.
The last three columns show the magnified aneurysm with different magnifying ratio from the same viewpoints.
Moreover, we show the magnified volumetric knee from different viewpoints with different ROIs denoted by the blue boxes in Fig.
Our method only magnifies the ROIs and keeps other regions unchanged.
Compared with the traditional method requiring tedious zoom in/out, our method only magnifies the
ROI region and keeps the whole subject in the field of view, which enables doctors to visualize the overall anatomy while scrutinize detailed anatomical structure at the same time.
C ONCLUSION
In this paper, we propose the FFT-OT method to solve the optimal transportation problem.
According to the Brenier theory, under the quadratic distance cost, finding the solution to the OT problem is equivalent to solving the Monge-Amp ` ere equation, which can be linearized as a sequence of variant coefficient elliptic PDEs.
Later, the variant coefficient PDEs are approximated by constant coefficient PDEs and solved by Fast Fourier Transformation.
We also prove that the proposed method converges linearly.
Experiments on volumetric data show that the FFT-OT can be used to sample from complex 3D density functions and magnify the volumetric data in medical images. A CKNOWLEDGEMENT
This research was partially supported by National Key R&D Program of China 2021YFAand the NSFC No.
61936002, T2225012.
This work was also partially supported by NIH 3R01LM012434-05S1, 1R21EB029733-01A1, NSF FAIN-2115095 and NSF CMMI-1762287.
A R ELATED W ORK
There is a huge literature about optimal transportation.
Here we will only briefly review the most related works.
For detailed reviews, we refer readers to Santambrogio (2015); Peyr´e & Cuturi (2019).
The first type of algorithms is based on the Kantorovich theory.
When both the input and output domains are Dirac masses, the Kantorovich problem can be treated as a standard linear programming (LP) task.
In order to tackle large data sets, Cuturi (2013) adds an entropic regularizer to the original LP problem and the regularized problem can be quickly solved by the Sinkhorn algorithm.
Recently, various algorithms have been proposed to further accelerate the computation by improving the efficiency of matrix-vector multiplications, including the Greenkhorn Altschuler et al.
(2017), Screenkhorn Alaya et al.
and the NYS-SINK Altschuler et al.
(2019) algorithms.
Dvurechensky et al. Dvurechensky et al.
(2018) also propose the adaptive primal-dual accelerated gradient descent algorithm (APDAGD) to solve the discrete OT problem.
An et al. An et al.
(2022) compute the approximate OT plan by smoothing the dual Kantorovich problem and solving it with the FISTA method.
This kind of methods have limitations: (i) they only give transport plans and cannot produce the bijective transportation maps; and (ii) the computational complexity is too high to apply them in the scenarios with huge number of samples.
The second type of algorithms is based on the Brenier theory Brenier (1987) and its intrinsic connection with convex geometry Gu et al.
The semi-discrete OT algorithm proposed in Aurenhammer et al.
(1998) finds the transport map between a continuous distribution and a discrete measure via a variational approach by dynamically constructing the power diagrams.
Its efficiency can be further improved Levy (2015); Merigot (2011) by the multi-resolution strategy.
The algorithms proposed in Kitagawa et al.
Su et al.
(2017) also improve the efficiency by applying the Newton’s method.
When both the source and target measures are continuous, some interpolation methods are necessary Schwartzburg et al.
The major drawback of this type of algorithms is the high computational complexity of constructing the dynamic power diagram, which prevents them from handling high dimensional tasks.
For example, for the 3D OT problems, these algorithms usually run very slow.
The third type of algorithms is based on computational fluid dynamics Benamou et al.
; Papadakis et al.
These methods aim at finding a special temporal-spacial flow field that transports the initial source density to the target density with the minimal total kinetic energy.
Then the diffeomorphism induced by the flow gives the optimal transport map under the quadratic Euclidean distance cost.
However, this kind of algorithms are difficult to extend to high dimensional space.
The fourth type of algorithms directly solve the Monge-Amp ` ere equation using numerical methods.
Loeper and Rapetti Loeper & Rapetti (2005) propose to solve the linearized Monge-Amp ` ere equation defined on a flat torus in each iteration.
Its corresponding variant coefficient elliptic PDE is converted to a positive definite linear system using the finite-difference scheme, which can be solved by the BiCG algorithm Endre (2020).
Benamou et al. Benamou et al.
(2014) propose to solve the linearized Monge-Amp ` ere on more general domains using Newton’s method.
Nader and
Guennebaud Nader & Guennebaud (2018) apply the similar discretization strategy and solve the
Monge-Amp`ere equation by conjugate gradient method.
Saumier et al. Saumier et al.
(2013) propose to solve the linearized Monge-Amp`ere equation using FFT.
In each iteration the elliptic PDE with spacial and temporal variant coefficients is converted to a group of linear equations in the frequency domain, which is solved by the GMRES algorithm.
Although the GMRES algorithm can be implemented on GPUs Aliaga et al.
(2019), there is no available open source code.
The work in
Saumier et al.
(2013) focuses on periodic boundary condition, but this our proposed work focuses on general second boundary condition; the work in Saumier et al.
concerns planar OT maps, ours emphasizes on volumetric OT maps, which has higher complexity.
The work in Saumier et al.
(2013) can handle more general target measures, the proposed work currently only deals with the
Lebesgue target measure.
Nevertheless, the current work can be directly generalized to handle general target measures as well.
Lei and Gu Lei & Gu (2021) use the fixed point method to compute the 2-dimensional OT problem based on FFT, but it cannot be extended to solve the 3-dimensional problems.
In this work, we combine the idea of linearizing the Monge-Amp`ere equation Loeper & Rapetti (2005) and the idea of FFT Saumier et al.
The key novelty of our proposed method is to use the mean linearized Monge-Amp`ere operator Eqn.
(12) to replace the conventional linearized
Monge-Ampere operator Eqn.
This replacement allows the algorithm to be implemented on
GPUs and makes the algorithm hundreds of times faster.
In the following, we compute the 3dimensional optimal transport problem by applying the proposed algorithm.
Our method also runs more than 100 × faster than the convex geometry based method Levy (2015).
B A PPENDIX T HEORY
In the section, we give the detailed proofs for several lemmas and theorems.
Some of them are well known in the Monge-Amp`ere PDE field and the applied mathematics field, we include them for the completeness.
B. E XISTENCE OF THE S OLUTION TO THE T IME D EPENDENT M ONGE -A MP ` ERE E QNUATION
Let T n = R n / Z n be the n -dimensional flat torus.
Below we sometimes identify it with Ω = [0 , 1] n and assume all data are periodic.
The existence and regularity of solutions to the Monge-Amp`ere equation are given by the following theorem,
Theorem 5.
Suppose a positive density function f : Ω → R is defined on Ω = [0 , 1] n , such that (cid:82) Ω f
( x ) dx = 1 , and f ∈ C α (Ω) , then the solution u : Ω × [0 , 1] to the time-dependent MongeAmp`ere equation
[Equation]
[Equation]
exists and is unique up to a constant.
Furthermore, there exist constants 0 < λ < Λ , such that
[Equation]
We refer readers to Cordero-Erausquin (1999) for detailed proof.
Weak Solution In practice, we compute the weak solution of the linearized Monge-Amp`ere
Eqn.
(6) using numerical methods.
We first rewrite the differential operator to a divergence form, then define a bi-linear form.
Since ( u pq ( x, t )) is the adjoint matrix of D 2 x u ( x, t ) , by direct computation, we obtain n (cid:88)
[Equation]
(23) so Eqn.
(6) can be converted into the divergence form:
[Equation]
u
[Equation]
, we obtain n (cid:88)
[Equation]
(24) with Neumann boundary condition ∂v ( x
[Equation]
t )
[Equation]
n =
[Equation]
For any w ∈ H 1 (Ω) , by differentiation of product, we obtain n (cid:88)
[Equation]
by integrating both sides, and from the fact that v satisfies the Neumann boundary condition, we deduce (cid:90)
[Equation]
For any fixed time t ∈ [0 , 1] , by the divergence form, we can construct a bilinear form a : H 1 (Ω) ×
H 1 (Ω) and a linear form l : H 1 (Ω) → R ,
[Equation]
A weak solution to Eqn.
(24) is a function v ∈ H 1 (Ω) , such that
[Equation]
By the uniform ellipticity Eqn.
(22), the Lax-Milgram theorem Endre (2020) shows the existence of the weak solution.
B. D ISCRETE L INEARIZED M ONGE -A MP ` ERE E QUATION S OLVABILITY
Galerkin Method In practice, we construct a triangulation T of Ω , such that the ratio between the diameter and inscribe-sphere radius of each simplex is bounded, and variation of the diameters of all the simplexes is small.
We call such kind of T a quasi-uniform triangulation, and denote the largest diameter as h .
For each vertex v i ∈ T , we construct a piecewise linear base function φ i , such that φ i is linear on each triangle, φ i ( v j ) is δ ij .
We define a finite dimensional subspace V h ⊂ H 1 (Ω) ,
V
[Equation]
[Equation]
∈
[Equation]
Given a function u ∈ H 1 (Ω) , we use u h ∈ V h to denote its approximation in V h .
Furthermore, u h = (cid:80)
i λ i
φ i , we also use u h to represent the coefficient vector ( λ 1 , λ 2 , . . .
, λ k ) T depending on the context.
The weak solution Eqn.
(28) to the Monge-Amp ` ere equation (6) is equivalent to find a v ∈ H 1 (Ω) , such that a ( v, w )
=
l ( w ) for all w ∈ H 1 (Ω) .
In discrete cases, we want to find v h ∈ V h , such that
[Equation]
Eqn. (29) is equivalent to the linear system,
[Equation]
From the weak solution to the linearized Monge-Amp ` ere equation (10), we obtain the linear system
Eqn. (30).
We denote the stiffness matrix A
= ( a ( φ i , φ j )) .
By the uniform ellipticity Eqn.
(22), and
V h
⊂ H 1 (Ω)
[Equation]
Assume (cid:82) Ω vdx = 0 , by Poincar´e inequality,
[Equation]
where the constant C 1 (Ω) depends on Ω .
Combine the above two inequalities, we obtain
[Equation]
Similarly, By the uniform ellipticity Eqn. 22, and V h ⊂ H 1 (Ω)
[Equation]
For linear finite element and quasi-uniform triangulation, we have the inverse Poincar´e inequality,
[Equation]
where h is the diameter of each element.
Combine the above two inequalities, we obtain
[Equation]
[Equation]
By combining the inequalities Eqn.
(31) and Eqn.
(32), we obtain
[Equation]
where C 3 > 1 is a constant.
Suppose v h = (cid:80) ni
=1 ξ i φ i , then
[Equation]
where ξ = ( ξ i ) and the matrix Φ = (cid:0)(cid:82) Ω φ i φ j (cid:1) is positive definite.
Therefore,
[Equation]
By a ( v h , v h ) =
ξ T Aξ , combing inequalities Eqn.
(33) and Eqn.
(34), we obtain
[Equation]
where C 3 C 4 > 1 .
This proves the following lemma,
Lemma 6.
By using Galerkin method using linear elements to numerically approximate the weak solution Eqn.
(28) to the linearized Monge-Amp`ere Eqn.
(6), if the uniform ellipticity Eqn.
(22) holds, and the triangulation T is quasi-uniform, then the stiffness matrix of the linear system Eqn.
(30) is positive definite on the space (cid:80) ni [Equation]
[Equation]
where C 3 C 4 > 1 .
Since the uniform ellipticity Eqn.
(22) holds for any time t ∈ [0 , 1] , then we obtain
Corollary 7.
By using Galerkin method with linear elements on quasi-uniform triangulations, the linearized Monge-Amp`ere equation in the continuity method Eqn.
(6) always has a solution v h ∈ V h for any t ∈ [0 , 1] .
Please note that the central differential scheme can be treated as Galerkin’s method on a special uniform triangulation.
Therefore, the above estimates still hold.
B. C ONVERGENCE R ATE
Theorem 8 (main) .
Given a domain Ω ⊂ R n , which is a canonical cuboid Ω = [ − 1 , 1] n , and a positive density function f : Ω → R with the balance condition
[Equation]
suppose the mirror reflection extension Eqn.
(14) of f to the flat torus ˜ f
: T n → R is C α , α ∈ (0 , 1) , then Monge-Amp`ere equation,
[Equation]
can be solved using FFT-OT Algorithm Alg.
In particular, one can choose the step length parameter τ , such that there is a constant 0 < γ < 1 , the approximation error satisfies
[Equation]
namely the algorithm has a linear convergence rate.
Proof.
Suppose at the k + 1 -th iteration, ρ k +1 = det
( I
[Equation] [Equation]
where L k [ v k ] = (cid:80) pq u pqk ∂ p
∂ q v k .
Hence by integration by parts Eqn.
[Equation]
where a k is the bilinear form in Eqn.(27).
In the discrete case, all functions are in V h , we denote
[Equation]
by the inequality Eqn.
(34) and Eqn. 35,
[Equation]
Therefore
[Equation]
where A k is the stiffness matrix in Eqn.(30), and ¯ A k is the mean stiffness matrix.
( By the uniform ellipticity Eqn.
(22), the eigen values of the adjoint matrix ( u pq )( x, t ) is uniformly bounded away from zero in the space H :
= { ξ ∈ R n | (cid:80) i ξ i = 0 } , so the eigen value of the mean adjoint matrix ¯ u pq ( t ) is bounded away from zero in H . After discretization, the eigen values of ¯ A k is strictly positive in H , hence ¯ A k is invertible in H . In the following discussion, the term o ( τ − 1 ) will be ignored.) Remark that the following displayed equation is a scalar
[Equation]
Since A k and ¯ A k are symmetric, positive definite on the space (cid:80) i ξ i = 0 ,
∥ A k ∥ 2 ≤ C 3 C 4 and ∥ ¯ A k ∥ 2 ≤ C 3 C 4 , so are their inverses.
Since A n and ¯ A n are symmetric, positive definite on the space orthogonal to (1 , 1 , . .
, 1) T , by Eqn.
(35) and ∥ A k ¯ A − 1 k ∥ ≤ ∥
A k ∥∥ ¯ A − 1 k ∥ , we have
[Equation]
Plug into Eqn.
(37), we have
[Equation]
We can choose the step-length τ − 1 , such that γ ∈ (0 , 1) , where
[Equation]
Therefore
[Equation]
B. D IFFERENTIAL O PERATOR U SING FFT
By using the Discrete Fourier Transformation, the differential operators can be converted to algebraic operators in the frequency domain.
Lemma 9.
Suppose the discrete function is u i,j,k , with discrete Fourier transformation
[Equation]
then the differential operator using central difference ∂ i ∂ i u i,j,k is given by
[Equation]
where ι = √− 1 , and ∂ i ∂ j u i,j,k is given by,
[Equation]
Proof.
By equations
[Equation]
and sin
[Equation]
we obtain
[Equation]
by direct computation
[Equation]
we have
[Equation]
Similarly, by equations
[Equation] + β ) =2 cos( A + B ) cos( α + β ) − 2 cos( A + B ) cos( α − β )
=2 cos( A + B )(cos( α + β ) − [-1 (cid:1) h 
-]cos( α − β ))
=2 cos( A + B )(cos α cos β − sin α sin β − cos α cos β − sin α − sin β )
= − 4 cos( A + B ) sin α sin β and sin( A + α + B + β ) + sin( A − α + B − β ) − sin
( A + α + B − β ) − sin
A − α + B + β )
=
sin( A + B + α + β ) + sin( A + B − α − β ) − sin( A + B + α − β ) − sin( A + B − α + β ) =2 sin( A + B ) cos( α + β ) − 2 sin( A + B ) cos( α − β ) =2 sin( A + B )(cos( α + β ) − cos( α − β )) =2 sin( A + B )(cos α cos β − sin α sin β − cos α cos β − sin α − sin β ) = − 4 sin( A + B ) sin α sin β we deduce the following equation,
[Equation]
Similarly
[Equation]
we have the representations of other differential operators in the frequency domain,
[Equation]
i,j +1 ,k − 1 − u i,j − 1 ,k +1 ) =
MNL
M − 1 (cid:88) m = N − 1 (cid:88) n = L − 1 (cid:88) l = ˆ u m,n,l − sin 2 πnN sin
2 πlL h y h
z e ι 2 πmiM e ι 2 πnjN e ι 2 πlkL ∂ k ∂ i u i,j,k = 4 h
{+z h x ( u i +1 ,j,k +
 + u i − 1 ,j,k − 1 − u i +1 ,j,k − 1 − u+}
i − 1 ,j,k +1 ) =
MNL M − 1 (cid:88) m = N − 1 (cid:88) n = L − 1 (cid:88) l
= ˆ u m,n,l − sin 2 πlL sin 2 πmM h z h x e ι 2 πmiM e ι 2 πnjN e ι 2 πlkL
C A PPENDIX E XPERIMENTS
In this section, as a compensation of the experiments in the main paper, we give more results on the 3D adaptive sampling and volumetric magnifier.
C. M ORE RESULTS ON 3D A DAPTIVE S AMPLING
In the experiments, we set the density function f ( x ) = (cid:80) 30 i =1 p i N ( µ i , Σ i ) , where N ( µ i , Σ i ) represents Gaussian distribution with mean µ i and variance Σ
i = diag ( σ 2 i 0 , σ 2 i 1 , σ 2 i 2 ) .
µ i ∈ R 3 is uniformly sampled from [0 , 1] 3 , σ ij is uniformly sampled from [0 , 0 . 5] , p i ∈ R is uniformly sampled from [0 . 2 , 1] and normalized such that (cid:82) Ω f ( x ) dx = 1 .
Thus the source distribution µ is a complicated Gaussian mixture distribution restricted on Ω
= [0 , 1] 3 .
After computing the OT map
T from µ to the uniform distribution ν defined on [ − 1 , 1] 3 , we conduct two groups of experiments: (i) we map the cell centers of the grid { y k } of [ − 1 , 1] 3 back to [ − 1 , 1] 3 through the inverse OT map
T − 1 ( y k ) defined by Eqn.
(20); (ii) we randomly sample 100 k samples { y k } from the Uniform distribution defined in [ − 1 , 1] 3 , then map them back to [ − 1 , 1] 3 through the inverse OT map T − 1 ( y k ) .
In order to keep the consistency with the mirror reflection process in the FFT-OT algorithm, we also reflect the the generated samples back to Ω .
To visualize the results of the k th slice, we plot the samples whose z coordinates satisfy the inequality,
[Equation]
[Figure]
In Fig.
5 and Fig.
6, we give more sampling results of different slices correspond to the two models used in Fig.
2 in the main paper.
Fig.
5 visualize the density function restricted on the 40th, 56th, 72th and 80th slices for different methods of the model displayed in the first row of 2.
Fig.
6 visualize the density function restricted on the 56th, 64th, 80th and 88th slices for different methods of the model displayed in the second row of 2.
Compared with the classical methods, the both sampling strategies of our method give decent sampling results that fit the prescribed density function well.
Moreover, the number of generated samples for different slices of the same 3D model fits the density functions restricted to the corresponding slices well, namely more samples are generated in the brighter regions for different slices.
C. M ORE RESULTS ON V OLUMETRIC M AGNIFIER
In this experiment, we magnify the volumetric MRI image of the aneurysm by different amplification factors.
In Fig.
7, we show the original aneurysm viewed from difference angles in the first column.
The last three columns give the magnified results with different amplification factors from the viewpoints same as those in the first column.
We can see that the aneurysm region is successfully magnified by different factors and the rest parts of the volume nearly keeps the same. [Figure] [Figure]