Vector and Matrix Optimal Mass Transport: Theory, Algorithm, and Applications

In many applications such as color image processing, data has more than one piece of information associated with each spatial coordinate, and in such cases the classical optimal mass transport (OMT) must be generalized to handle vector-valued or matrix-valued densities. In this paper, we discuss the vector and matrix optimal mass transport and present three contributions. We first present a rigorous mathematical formulation for these setups and provide analytical results including existence of solutions and strong duality. Next, we present a simple, scalable, and parallelizable methods to solve the vector and matrix-OMT problems. Finally, we implement the proposed methods on a CUDA GPU and present experiments and applications.

However, in many applications such as color image processing, there is more than one piece of information associated with each spatial coordinate, and such data can be interpreted as vector-valued or matrix-valued densities. As the classical optimal mass transport works with scalar probability densities, such applications require a new notion of mass transport. For this purpose, Chen et al. [18,19,14,12] recently developed a framework for vector-valued and matrix-valued optimal mass transport. See also [41,40,42,25,45,55] for other different frameworks. For vector-valued OMT, potential applications include color image processing, multi-modality medical imaging, and image processing involving textures. For matrix-valued OMT, we have diffusion tensor imaging, multivariate spectral analysis, and stress tensor analysis.
Several mathematical aspects of vector and matrix-valued OMT were not addressed in the previous work [18,19,14]. As the first contribution of this paper, we present duality and existence results of the continuous vector and matrix-valued OMT problems along with rigorous problem formulations.
Although the classical theory of OMT is very rich, only recently has there been much attention to numerical methods to compute the OMT. Several recent work proposed algorithms to solve the L 2 OMT [3,22,7,28,6,13,17,27] and the L 1 OMT [34,50]. As the second contribution of this paper, we present first-order primal-dual methods to solve the vector and matrix-valued OMT problems. The methods simultaneously solve for both the primal and dual solutions (hence a primal-dual method) and are scalable as they are first-order methods. We also discuss the convergence of the methods.
As the third contribution of this paper, we implement the proposed method on a CUDA GPU and present several applications. The proposed algorithms' simple structure allows us to effectively utilize the computing capability of the CUDA architecture, and we demonstrate this through our experiments. We furthermore release the code for scientific reproducibility.
The rest of the paper is structured as follows. In Section 2 we give a quick review of the classic OMT theory, which allows us to present the later sections in an analogous manner and thereby outline the similarities and differences. In Section 3 and Section 4 we present the vector and matrix-valued OMT problems and state a few theoretical results. In Section 5, we present and prove the analytical results. In Section 7, we present the algorithm. In Section 8, we present the experiments and applications.

Scalar optimal mass transport
The theory of optimal transport [26,53,54] remarkably points out that (2) is equivalent to the following flux minimization problem: where u = (u 1 , . . . , u d ) : Ω → R d is the optimization variable and div x denote the (spatial) divergence operator. Although (3) and (2) are mathematically equivalent, (3) is much more computationally effective as its optimization variable u is much smaller when discretized. It is worth mentioning that OMT in formulation (3) is very close to the problems in compressed sensing. Its objective function is homogeneous degree one and the constraint is linear. It can be observed that for characterizing the OMT, the gradient operator in (3) and divergence operator in (3) play the key roles. Later on, we extend the definition of L 1 OMT problem by extending these differential operators into a general meaning.
The optimization problem (3) has the following dual problem: where the optimization and φ : Ω → R is a function. We write · * for the dual norm of · . It is well-known that, strong duality holds between (3) and (4) in the sense that the minimized and maximized objective values are equal [53]. Therefore, we take either (3) or (4) as the definition of S.
Rigorous definitions of the optimization problems (3) or (4) are somewhat technical. We skip this discussion, as scalar optimal mass transport is standard. In Section 5, we rigorously discuss the vector-OMT problems, so any rigorous discussion of the scalar-OMT problems can be inferred as a special case. To make c j the cost of traversing edge j, the edge weight is defined to be 1/c 2 j for j = 1, . . . , .

Theoretical properties
The algorithm we present in Section 7 is a primal-dual algorithm and, as such, finds solutions to both Problems (3) and (4). This is well-defined as both the primal and dual problems have solutions [53]. Write R + for the set of nonnegative real numbers. Write P(Ω, R) for the space of nonnegative densities supported on Ω with unit mass. We can use S(λ 0 , λ 1 ) as a distance measure between λ 0 , λ 1 ∈ P(Ω, R). The value S : P(Ω, R) × P(Ω, R) → R + defines a metric on P(Ω, R) [53].

Vector optimal mass transport
Next we discuss the vector-valued optimal transport, proposed recently in [19]. The basic idea is to combine scalar optimal mass transport with network flow problems [1].

Gradient and divergence on graphs
Consider a connected, positively weighted, undirected graph G with k nodes and edges. To define an incidence matrix for G, we say an edge {i, j} points from i to j, i.e., i → j, if i < j. This choice is arbitrary and does not affect the final result. With this edge orientation, the incidence matrix D ∈ R k× is For example, the incidence matrix of the graph of Figure 1 is Write ∆ G = −D diag{1/c 2 1 , · · · , 1/c 2 }D T for the (negative) graph Laplacian, where 1/c 2 1 , · · · , 1/c 2 are the edge weights. The edge weights are defined so that and c j represents the cost of traversing edge j for j = 1, . . . , .
We define the gradient operator on G as ∇ G x = diag{1/c 1 , · · · , 1/c }D T x and the divergence operator as div G y = −D diag{1/c 1 , · · · , 1/c }y. So the Laplacian can be rewritten as ∆ G = div G ∇ G . Note that div G = −∇ * G , where ∇ * G is the adjoint of ∇ G . This is in analogy with the usual spatial gradient and divergence operators [36,20,21,19].

Vector optimal mass transport
We say λ : Ω → R k + is a nonnegative vector-valued density with unit mass if Assume λ 0 and λ 1 are nonnegative vector-valued densities supported on Ω with unit mass. We define the optimal mass transport between vector-valued densities as where u : Ω → R k×d and w : Ω → R are the optimization variables, α > 0 is a parameter, and · u is a norm on R k×d and · w is a norm on R . The parameter α represents the relative importance of the two flux terms u and w. We write We call div x the spatial divergence operator. The zero-flux boundary condition is for i = 1, . . . , k. Note that w has no boundary conditions. The optimization problem (5) has the following dual problem: where the optimization variable φ : Ω → R k is a function. We write · u * and · w * for the dual norms of · u and · w , respectively. As stated in Theorem 1, strong duality holds between (5) and (6) in the sense that the minimized and maximized objective values are equal. Therefore, we take either (5) or (6) as the definition of V . In Section 5, we rigorously define the primal and dual problems and prove Theorem 1.

Theoretical properties
The algorithm we present in Section 7 is a primal-dual algorithm and, as such, finds solutions to both Problems (5) and (6). This is well-defined as both the primal and dual problems have solutions. Write P(Ω, R k ) for the space of nonnegative vector-valued densities supported on Ω with unit mass. We can use V ( λ 0 , λ 1 ) as a distance measure between λ 0 , λ 1 ∈ P(Ω, R k ). Theorem 2. V : P(Ω, R k ) × P(Ω, R k ) → R + defines a metric on P(Ω, R k ).

Quantum gradient operator and matrix optimal mass transport
We closely follow the treatment in [18]. In particular, we define a notion of gradient on the space of Hermitian matrices and its dual, i.e., the (negative) divergence.
Some applications of matrix-OMT, such as diffusion tensor imaging, have real-valued data while some applications, such as multivariate spectral analysis, have complex-valued data [51]. To accommodate the wide range of applications, we develop the matrix-OMT with complex-valued matrices.
Write C, H, and S for the set of k × k complex, Hermitian, and skew-Hermitian matrices respectively. We write H + for the set of k × k positive semidefinite Hermitian matrices, i.e., M ∈ H + if v * M v ≥ 0 for all v ∈ C k . Write tr for the trace, i.e. for any M ∈ H, we have tr(M ) = k i=1 M ii . Write C N for the block-column concatenation of N matrices in C, i.e., Z ∈ C N if Z = [Z * 1 · · · Z * N ] * and Z 1 , . . . , Z N ∈ C. Define H N and S N likewise. For X, Y ∈ C, we use the Hilbert-Schmidt inner product (This is the standard inner product when we view C as the real vector space R 2n 2 .) For X ∈ C, we use the norm

Quantum gradient and divergence operators
We define the gradient operator, given a L = [L 1 , · · · , L ] * ∈ H , as Define the divergence operator as This is in analogy with the usual spatial gradient and divergence operators. Write ∆ L = div L ∇ L .
This notion of gradient and divergence operators is motivated by the Lindblad equation in Quantum mechanics [18]. The choice of L affects ∇ L . There is no standard way of choosing L. A standing assumption throughout, is that the null space of ∇ L , denoted by ker(∇ L ), contains only scalar multiples of the identity matrix I.
Assume Λ 0 and Λ 1 are nonnegative matrix-valued densities supported on Ω with unit mass.
We define the optimal mass transport between matrix-valued densities as where U : Ω → H d and W : Ω → S are the optimization variables, α > 0 is a parameter, and · u is a norm on H d and · w is a norm on S . The parameter α represents the relative importance of the two flux terms U and W. We write We define the spatial divergence as The zero-flux boundary condition is for i, j = 1, . . . , k. Note that W has no boundary conditions [14]. The optimization problem (7) has the following dual problem: where the optimization variable Φ : Ω → H is a function. We write · u * and · w * for the dual norms of · u and · w , respectively. As stated in Theorem 3, strong duality holds between (7) and (8) in the sense that the minimized and maximized objective values are equal. Therefore we take either (7) or (8) as the definition of M .
To avoid repeating the same argument with different notation, we simply point out that the rigorous definitions of the matrix optimal mass transport problems are analogous to those of the vector setup. So the precise definitions of (7) or (8) can be inferred from the discussion of Section 5.

Theoretical properties
The algorithm we present in Section 7 is a primal-dual algorithm and, as such, finds solutions to both Problems (7) and (8). This is well-defined as both the primal and dual problems have solutions.
Theorem 3. The (infinite dimensional) primal and dual optimization problems (7) and (8) have solutions, and their optimal values are the same, i.e., strong duality holds.

Duality proof
In this section, we establish the theoretical results. For notational simplicity, we only prove the results for the vector-OMT primal and dual problems (5) and (6). Analogous results for the matrix-OMT primal and dual problems (7) and (8) follow from the same logic.
Although the classical scalar-OMT literature is very rich, standard techniques for proving scalar-OMT duality do not simply apply to our setup. For example, Villani's proof of strong duality, presented as Theorem 1.3 of [53], relies on and works with the linear optimization formulation (2). However, our vector and matrix-OMT formulations directly generalize the flux formulation (3) and do not have formulations analogous to (2). We need a direct approach to analyze duality between the flux formulation and the dual (with one function variable), and we provide this in this section.
We further assume Ω ⊂ R d has a piecewise smooth boundary. Write Ω • and ∂Ω for the interior and boundary of Ω. For simplicity, assume Ω as full affine dimensions, i.e., Ω • = Ω.
The rigorous form the dual problem (6) is where W 1,∞ (Ω, R k ) is the standard Sobolev space of functions from Ω to R k with bounded weak gradients. That (9) has a solution directly follows from the Arzelà-Ascoli Theorem.
To rigorously define the primal problem (5) requires more definitions, and we do so later as (11).

Fenchel-Rockafellar duality
Let L : X → Y be a continuous linear map between locally convex topological vector spaces X and Y , and let f : In this framework of Fenchel-Rockafellar duality, d ≤ p , i.e., weak duality, holds unconditionally, and this is not difficult to prove. To establish d = p , i.e., strong duality, requires additional assumptions and is more difficult to prove. The following theorem does this with a condition we can use.

Theorem 5.
[Theorem 17 and 18 [48]] If there is an x ∈ X such that f (x) < ∞ and g is bounded above in a neighborhood of Lx,

Spaces
Throughout this section, · 1 , · 2 , . . . denote unspecified finite dimensional norms. As all finite dimensional norms are equivalent, we do not bother to precisely specify which norms they are. Define Then C 1 (Ω, R k ) is a Banach space equipped with the norm · ∞,∞ . Write M(Ω, R k ) for the space of R k -valued signed finite Borel measures on Ω, and define M(Ω, R k×d ) likewise. Write (C(Ω, R k×d )) * , (C 1 (Ω, R k )) * for the topological dual of C(Ω, R k×d ), C 1 (Ω, R k ), respectively. The standard Riesz-Markov theorem tells us that (C(Ω, R k×d )) * = M(Ω, R k×d ).
Fully characterizing (C 1 (Ω, R k )) * is hard, but we do not need to do so. Instead, we only use the following simple fact. Any g ∈ M(Ω, R k ) defines the bounded linear map φ → Ω φ(x), g(dx) for any φ ∈ C 1 (Ω, R k ). In other words, M(Ω, R k×d ) ⊂ (C 1 (Ω, R k )) * with the appropriate identification.

Operators
We redefine ∇ x : C 1 (Ω, R k ) → C(Ω, R k×d ) so that ∇ x φ is the continuous exntension of the usual ∇ x φ to all of Ω. This makes ∇ x a bounded linear operator. Define the dual (adjoint) operator ∇ * for any φ ∈ C 1 (Ω, R k ) and u ∈ M(Ω, R k×d ). Define the ∇ G (which is simply a multiplication by a R k× matrix) as ∇ G : C 1 (Ω, R k ) → C(Ω, R k ). Since C 1 (Ω, R k ) ⊂ C(Ω, R k ), there is nothing wrong with defining the range of ∇ G to be C(Ω, R k ), and this still makes ∇ G a bounded linear operator. Define the dual (adjoint) operator ∇ * G : M(Ω, R k ) → (C 1 (Ω, R k )) * by identifying ∇ * G with the transpose of the matrix that defines ∇ G . Since ∇ * G is simply multiplication by a matrix, we can further say

Zero-flux boundary condition
Let m : Ω → R k×d a smooth function. Then integration by parts tells us that holds for all smooth ϕ : Ω → R k if and only if m(x) satisfies the zero-flux boundary condition, i.e., m(x)n(x) = 0 for all x ∈ ∂Ω, where n(x) denotes the normal vector at x. Here div x denotes the usual (spatial) divergence. To put it differently, ∇ * x = −div x holds when the zero-flux boundary condition holds. We generalize this notion to measures. We say u ∈ M(Ω, R k×d ) satisfies the zero-flux boundary condition in the weak sense if there is a g ∈ M(Ω, R k ) ⊂ (C 1 (Ω, R k )) * such that holds for all φ ∈ C 1 (Ω, R k ). In other words, u satisfies the zero-flux boundary condition if ∇ * x u ∈ M(Ω, R k ) ⊂ (C 1 (Ω, R k )) * . In this case, we write div x ( u) = g and div x ( u) = −∇ * x ( u). This definition is often used in elasticity theory.

Duality
To establish duality, we view the dual problem (6) as the primal problem and obtain the primal problem (5) as the dual of the dual. We do this because the dual of C(Ω, R k×d ) is known, while the dual of M(Ω, R k×d ) is difficult to characterize.
Consider the problem maximize which is equivalent to (9). Define Rewrite (10) as and consider its Fenchel-Rockafellar dual , u satisfies the zero-flux boundary condition.
We now state the rigorous form of the primal problem (5) minimize The point φ = 0 satisfies the assumption of Theorem 5. Furthermore, it is easy to verify that the optimal value of the dual problem (6) is bounded. This implies strong duality, (11) is feasible, and (11) has a solution.

Algorithmic preliminaries
Consider the Lagrangian for the vector optimal transport problems (5) and its dual (6) which is convex with respect to u and w and concave with respect to φ. Finding a saddle point of (12) is equivalent to solving (5) and (6), when the primal problem (5) has a solution, the dual problem (6) has a solution, and the optimal values of (5) and (6) [48] for further discussion on this point.
To solve the optimal transport problems, we discretize the continuous problems and apply PDHG method to solve the discretized convex-concave saddle point problem.

PDHG method
Consider the convex-concave saddle function where f , g, and h are (closed and proper) convex functions and x ∈ R n , y ∈ R m , z ∈ R l , A ∈ R l×n , and B ∈ R l×m . Note L is convex in x and y and concave in z. Assume L has a saddle point and step sizes µ, ν, τ > 0 satisfy 1 > τ µλ max (A T A) + τ νλ max (B T B).
Write · 2 for the standard Euclidean norm. Then the method converges to a saddle point. This method is called the the Primal-Dual Hybrid Gradient (PDHG) method or the (preconditioned) Chambolle-Pock method [23,10,46]. PDHG can be interpreted as a proximal point method under a certain metric [30]. The quantity is the fixed-point residual of the non-expansive mapping defined by the proximal point method. Therefore R k = 0 if and only if (x k , y k , z k ) is a saddle point of L, and R k decreases monotonically to 0, cf., review paper [49]. We can use R k as a measure of progress and as a termination criterion.

Shrink operators
As the subproblems of PDHG (13) are optimization problems themselves, PDHG is most effective when these subproblems have closed-form solutions. The problem definitions of scalar, vector, and matrix-OMT involve norms. For some, but not all, choices of norms, the "shrink" operators shrink(x 0 ; µ) = argmin for u ∈ R k×d and The shrink operators of these norms have closed-form solutions. For the matrix-OMT, we focus on norms for U ∈ H d and · 2 , · 1 , and · 1,nuc for W ∈ S , which are defined likewise. The nuclear norm · nuc is the sum of the singular values. The shrink operators of these norms have closed-form solutions. We provide further information and details on shrink operators in the appendix.

Algorithms
We now present simple and parallelizable algorithms for the OMT problems. These algorithms are, in particular, very well-suited for GPU computing. In Section 7 and 8 we deal with discretized optimization variables that approximate solutions to the continuous problems. For simplicity of notation, we use the same symbol to denote the discretizations and their continuous counterparts. Whether we are referring to the continuous variable or its discretization should be clear from context.
As mentioned in Section 6, these methods are the PDHG method applied to discretizations of the continuous problems. In the implementation, it is important to get the discretization at the boundary correct in order to respect the zero-flux boundary conditions. For interested readers, the details are provided in the appendix.
Instead of detailing the somewhat repetitive derivations of the algorithms in full, we simply show the key steps and arguments for the u update of vector-OMT. The other steps follow from similar logic.
When we discretize the primal and dual vector-OMT problems and apply PDHG to the discretized Lagrangian form of (12), we get Since the minimization splits over the i, j indices, we write At the boundary, these manipulations need special care. When we incorporate ghost cells in our discretization, these seemingly cavalier manipulations are also correct on the boundary. We further explain the ghost cells and discretization in the appendix.

Scalar-OMT algorithm
The scalar-OMT algorithm can be viewed as a special case of vector-OMT or matrix-OMT algorithms. This scalar-OMT algorithm was presented in [34], but we restate it here for completeness.

Matrix-OMT algorithm
Write shrink u and shrink w for the shrink operators with respect to · u and · w . The matrix-OMT algorithm is as follows:

Parallelization
For the vector-OMT, the computation for the u, w, and φ updates splits over the indices (i, j), i.e., the computation splits pixel-by-pixel. Furthermore, the u and w updates can be done in parallel. Parallel processors handing jobs split over (i, j) must be synchronized before and after the φ update. The same is true for the scalar-OMT and matrix-OMT. This highly parallel and regular algorithmic structure makes the proposed algorithms very well suited for CUDA GPUs. We demonstrate this through our experiments.

Examples
In this section, we provide example applications of vector and matrix-OMT with numerical tests to demonstrate the effectiveness of our algorithms. As mentioned in the introduction, potential applications of vector and matrix-OMT are broad. Here, we discuss two of the simplest applications.
We implemented the algorithm on C++ CUDA and ran it on a Nvidia GPU. For convenience, we MEXed this code, i.e., the code is made available as a Matlab function. For scientific reproducibility, we release this code.

Color images
Color images in RGB format is one of the more immediate examples of vector-valued densities. At each spatial position of a 2D color image, the color is represented as a combination of the three basic colors red (R), green (G), and blue (B). We allow any basic color to change to another basic color with cost c 1 for R to G, c 2 for R to B, and c 3 for G to B. So the graph G as described in Section 3.1 has 3 nodes and 3 edges.
Consider the two color images on the domain Ω = [0, 1] × [0, 1] with 256 × 256 discretization shown in Figures 2a and 2b. The initial and target densities λ 0 and λ 1 both have three disks at the same location, but with different colors. The optimal flux depends on the choice of norms. Figures 2c and 2d, show fluxes optimal with respect to different norms.
Whether it is optimal to spatially transport the colors or to change the colors depends on the parameters c 1 , c 2 , c 3 , α as well as the norms · u and · w . With the parameters of Figure 3a it is optimal to spatially transport the colors, while with the parameters of Figure 3b it is optimal to change the colors.
Finally, we test our algorithm on the setup of Figure 2c with grid sizes 32 × 32, 64 × 64, 128 × 128, and 256 × 256. Table 1 shows the number of iterations and runtime tested on a Nvidia Titan Xp GPU required to achieve a 10 −3 precision, measured as the ratio between the duality gap and primal objective value.

Diffusion tensor imaging
Diffusion tensor imaging (DTI) is a technique used in magnetic resonance imaging. DTI captures orientations and intensities of brain fibers at each spatial position as ellipsoids and gives us a matrix-valued density. Therefore, the metric defined by the matrix-OMT provides a natural way to compare the differences between brain diffusion images. In Figure 4 we visualize diffusion tensor images by using colors to indicate different orientations of the tensors at each voxel. In this paper, we present simple 2D examples as a proof of concept and leave actual 3D imaging for a topic of future work.   Table 2: Distances between the three images Λ 0 , Λ 1 and Λ 2 .
Consider three synthetic matrix-valued densities Λ 0 , Λ 1 , and Λ 2 in Figure 5. The densities Λ 0 and Λ 1 have mass at the same spatial location, but the ellipsoids have different shapes. The densities Λ 0 and Λ 2 have the same ellipsoids at different spatial locations.
Again, we test our algorithm on the setup of Figure 5 with α = 1 on grid sizes 32 × 32, 64 × 64, 128 × 128, and 256 × 256. Table 3 shows the number of iterations and runtime tested on a Nvidia Titan Xp GPU required to achieve a precision of 10 −3 , measured as the ratio between the duality gap and primal objective value.

Conclusions
In this paper, we studied the extensions of Wasserstein-1 optimal transport to vector and matrix-valued densities. This extension, as a tool of applied mathematics, is interesting if the mathematics is sound, if the  numerical methods are good, and if the applications are interesting. In this paper, we investigated all three concerns. From a practical viewpoint, that we can solve vector and matrix-OMT problems of realistic sizes in modest time with GPU computing is the most valuable observation. Applying our algorithms to tackle real world problems in signal/imaging processing, medical imaging, and machine learning would be interesting directions of future research.
Another interesting direction of study is quadratic regularization. In general, the solutions to vector and matrix-OMT problems are not unique. However, the regularized version of (5) is strictly convex and therefore has a unique solution. A similar regularization can be done for matrix-OMT. As discussed in [34,50], this form of regularization is particularly useful as a slight modification to the proposed method solves the regularized problem.

Acknowledgments
We would like to thank Wilfrid Gangbo for many fruitful and inspirational discussions on the related topics. The Titan Xp's used for this research were donated by the NVIDIA Corporation. This work was funded by ONR grants N000141410683, N000141210838, DOE grant de-sc00183838 and a startup funding from Iowa State University.

A Discretization
Here, we discribe the discretization for the vect-OMT problem (5) and (6). The discretization for the matrix-OMT problem is similar. We consider the 2D square domain, but (with more complicated notation) our approach immediately generalizes to more general domains. Again, we use the same symbol for the discretization and its continuous counterpart. Consider a n × n discretization of Ω with finite difference ∆x in both x and y directions. Write the x and y coordinates of the points as x 1 , . . . , x n and y 1 , . . . , y n . So we are approximating the domain Ω with {x 1 , . . . , x n } × {y 1 , . . . , y n }. Write C(x, y) for the ∆x × ∆x cube centered at (x, y), i.e., We use a finite volume approximation for λ 0 , λ 1 , w and φ. Specifically, we write λ 0 ∈ R n×n×k with λ 0 ij ≈ C(xi,yj ) λ 0 (x, y) dxdy, for i, j = 1, . . . , n. The discretizations λ 1 , φ ∈ R n×n×k and w ∈ R n×n× are defined the same way. Write u = ( u x , u y ) for both the continuous variables and their discretizations. To be clear, the subscripts of u x and u y do not denote differentiation. We use the discretization u x ∈ R (n−1)×n×k and u y ∈ R n×(n−1)×k . For i = 1, . . . , n − 1 and j = 1, . . . , n u x,ij ≈ C(xi+∆x/2,yj ) u x (x, y) dxdy, and for i = 1, . . . , n and j = 1, . . . , n − 1 u y,ij ≈ C(xi,yj +∆x/2) u y (x, y) dxdy.
In defining u x and u y , the center points are placed between the n × n grid points to make the finite difference operator symmetric.
So (∇ φ) x ∈ R (n−1)×n×k and (∇ φ) y ∈ R n×(n−1)×k , and the ∇ x is the transpose (as a matrix) of −div x . Ghost cells are convenient for both describing and implementing the method. This approach is similar to that of [34]. We redefine the variable u = ( u x , u y ) so that u x,ij = u x,ij for i < n 0 for i = n u y,ij = u y,ij for j < n 0 for j = n , for i, j = 1, . . . , n, and u x , u y ∈ R n×n×k . We also redefine ∇ φ = ((∇ φ) x , (∇ φ) y ) so that for i, j = 1, . . . , n, and (∇ φ) x , (∇ φ) y ∈ R n×n×k . With some abuse of notation, we write w ij w .

B Shrink operators with closed-form solutions
All shrink operators we consider in this paper can be built from these shrink operators. These ideas are well-known to the compressed sensing and proximal methods community [44]. However, there is a subtlety we must address when applying the shrink operators to matrix-OMT: the U update is defined as the minimization over H, not all of C, and the W update is defined as the minimization over S, not all of C. Fortunately, this is not a problem when we use unitarily invariant norms (such as the nuclear norm) thanks to the following lemma. Lemma 1. Let M ∈ H and N ∈ S. Let · be a unitarily invariant norm, i.e., U AV = A for any A ∈ C, and U, V ∈ C unitary. Then is Skew-Hermitian.
Proof. The minimum over H ⊂ C is the same as the minimum over all of C, if the minimum over all of C is in H. Likewise, the minimum over S ⊂ C is the same as the minimum over all of C, if the minimum over all of C is in S.
Write M = U ΣU T for M 's eigenvalue decomposition, which is also its singular value decomposition. Define σ + = argmin s {f (s) + (1/2) s − σ(M ) 2 2 }, which implies 0 ∈ ∂f (σ + ) + σ + − σ(M ). With this, we can verify that 0 ∈ U diag(∂f (σ + ))U T + U diag(σ + )U T − U diag(σ(M ))U T , i.e., U diag(σ + )U T satisfies the optimality conditions of the optimization problem that defines M + . So M + = U diag(σ + )U T , which is Hermitian. Likewise, write N = U ΛU T for N 's eigenvalue decomposition. N has orthonormal eigenvectors and its eigenvalues are purely imaginary. We separate out the magnitude and phase of Λ and write N = U diag(σ(N ))P U T . More precisely, σ(N ) = | diag(Λ)| and P is a diagonal matrix with diagonal components ±i. Then the SVD of N is N = U diag(σ(N ))V T where P U T = V T . With the same argument as before, we conclude that N + = U diag(σ + )V T for some σ + ∈ R n . So N + = U Λ + U T where Λ + = diag(σ + )P is diagonal and purely imaginary, and we conclude N + is skew-Hermitian.
The norm · 1 as described in Section 6.2, is not unitarily invariant. However, its shrink operator acts element-by-element, so it is easy to arrive at a conclusion similar to that of Lemma 1.