A Vectorial Approach to Unbalanced Optimal Mass Transport

Abstract

Unbalanced optimal mass transport (OMT) seeks to remove the conservation of mass constraint by adding a source term to the standard continuity equation in the Benamou-Brenier formulation of OMT. In this study, we show how the unbalanced case fits into the vector-valued OMT framework simply by adding an auxiliary source layer and taking the flow between the source layer and the original layer(s) as the source term. This allows for unbalanced models both in the scalar and vector-valued density settings. The results are demonstrated on a number of synthetic and real vector-valued data sets.

I. INTRODUCTION

OPTIMAL mass transport (OMT) is a very important subject in mathematics, originating with the French civil engineer and mathematician Gaspard Monge in 1781 [11, 18, 20, 21]. Recently, the theory has undergone a massive growth, with numerous applications to various areas including signal processing, machine learning, computer vision, meteorology, statistical physics, quantum mechanics, and network theory [1, 3, 12, 15, 16, 19]. Several works deal with extensions of the theory to the unbalanced and vector-valued cases; see [7, 9, 14] and the references therein. In standard treatments, vector-valued and unbalanced extensions are treated separately. In the present work, we show that unbalanced OMT can fit into the model of vector OMT by taking a special set of weight parameters, which gives us a fresh way to treat the problem.

In the L² setting, both extensions arise from the computational fluid dynamics (CFD) approach to OMT introduced by Benamou and Brenier [2], which was a major development in OMT theory. For “square of distance” based cost functionals, the Benamou and Brenier framework is equivalent to the original one but also gives an explicit interpolation between two mass distributions regarded as a geodesic on the space of probability distributions [17]. The CFD method views the OMT as the minimizer of a kinetic energy functional subject to a continuity constraint. By modifying the energy functional and the constraint, unbalanced OMT and vector-valued OMT versions were developed [7, 9, 14]. More specifically, by adding a source term, unbalanced OMT can handle transport problems when the mass is not conserved. That is, mass can be created or destroyed during the process. By introducing the divergence of the flow between channels (layers), vector OMT extends the density from scalar to vector-valued and even matrix-valued [4, 5, 6].

In this study, we indicate how to reformulate unbalanced OMT into a vector-valued OMT framework, and thus may be implemented through any existing code for vector OMT. All that is required is the addition of a new weight matrix that does not change the main structure of the code.

In this study, we focus on the connection between the unbalanced and vector-valued extensions of OMT. We will first give some background on these two extensions. After that, we show how to reformulate unbalanced OMT as a vector OMT problem. We also look into reformulating unbalanced vector OMT as a common vector OMT problem and conclude with some illustrative numerical results.

II. BENAMOU AND BRENIER’S APPROACH

In this section, we formalize the OMT problem with particular regard to the Benamou-Brenier approach.

The original formulation of OMT due to Monge [20, 21] may be expressed as follows:

$$\underset{T}{\inf }\{{\int }_{E}c(x,T(x)){\rho }_0(x)dx\mid T_{\#}{\rho }_0={\rho }_1\},$$ (1)

where c(x, y) is the cost of moving unit mass from x to y, T is the transport map, and ρ₀, ρ₁ are two probability distributions defined on E, a subdomain of ${ℝ}^n$. T_# denotes the push forward of T.

As pioneered by Leonid Kantorovich [13], the Monge formulation of OMT may be relaxed by replacing the transport map T with couplings π:

$$\underset{\pi \in \mathrm{\Pi }\left({\rho }_0,{\rho }_1\right)}{\inf }{\int }_{E}c(x,y)\pi (dx,dy),$$ (2)

where Π(ρ₀, ρ₁) denotes the set of all the couplings between ρ₀ and ρ₁ (joint distributions whose two marginal distributions are ρ₀ and ρ₁). One may show that for c(x, y) = ∥x − y∥² (square of distance function), the Kantorovich and Monge formulations are equivalent; see [20, 21] and the references therein.

Moreover for c(x, y) = ∥x − y∥², the specific infimum is called the Wasserstein 2 distance $({\mathcal{W}}_2)$. Benamou and Brenier [2] pointed out that this can be written in an equivalent computational fluid dynamic formulation:

$${\mathcal{W}}_2{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,v}{\inf }{\int }_0^1{\int }_E\rho (t,x)\Vert v(t,x){\Vert }^{2}dxdt$$ (3a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot (\rho v)=0$$ (3b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1.$$ (3c)

Here, v is the space-time velocity and ρ denotes the time-interpolant between the given distributions ρ₀, ρ₁ over the normalized time interval t ∈ [0, 1]. Thus, the latter reformulates OMT as one of minimizing a kinetic energy functional subject to a continuity constraint. The continuity constraint (3b) states that the change of density at each point over time is due to the flux of its neighborhood.

The Benamou-Brenier optimal solution may be related to the original one of Monge. Indeed, if we have the optimal transport plan T in (1), the interpolation can be expressed as $\rho (t,\cdot )={(t\cdot T+(1-t)\cdot id)}_{\#}{\rho }_0$. For computational purposes, the Benamou-Brenier model is usually written as a convex optimization problem in momentum form (ρ = ρv):

$${\mathcal{W}}_2{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p}{\inf }{\int }_0^1{\int }_E\frac{p{(t,x)}^2}{\rho (t,x)}dxdt$$ (4a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p=0$$ (4b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1.$$ (4c)

III. UNBALANCED SCALAR OMT

In the present work, we treat the unbalanced case in which a source term is added under the Fisher-Rao smooth formulation [9]. Namely,

$${\mathcal{W}}_{FR}{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p}{\inf }{\int }_0^1{\int }_E\frac{p{(t,x)}^2}{\rho (t,x)}+\gamma \frac{s{(t,x)}^2}{\rho (t,x)}dxdt$$ (5a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p=s$$ (5b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1.$$ (5c)

The extra term in the continuity equation originally formulated by Benamou and Brenier [2] is the source term s, defined on time and space [0, 1] × E. The change of density at each point over time is no longer only due to the flux of its neighborhood, but also a source. Now mass can be created or destroyed at any point and time. The parameter γ is a weight to control how much source one wants to use.

IV. VECTOR-VALUED OMT

In this section, we outline the approach to vector-valued OMT from [7].

A. GRAPH LAPLACIAN, GRADIENT, AND DIVERGENCE

We will need some standard definitions of weighted graphs. See [10] for all the details. Accordingly, let $\mathcal{F}$ denote a connected, positively weighted, undirected graph with n nodes labeled as i, with 1 ≤ i ≤ n, and m edges. $\mathbb{F}\in {ℝ}^{n\times m}$ is the incident matrix of the graph $\mathcal{F}$. The weight associated with node i is denoted by w_i ≥ 0 and W = diag(w₁, w₂, …, w_n) is the diagonal weight matrix. We define the graph Laplacian as

$${\Delta }_{\mathcal{F}}:=-{\nabla }_{\mathcal{F}}^{*}{\nabla }_{\mathcal{F}},$$

where

$${\nabla }_{\mathcal{F}}:{ℝ}^n\to {ℝ}^m,x\mapsto W^{1/2}{\mathbb{F}}^{T}x$$

denotes the gradient operator and

$${\nabla }_{\mathcal{F}}^{*}:{ℝ}^m\to {ℝ}^n,y\mapsto \mathbb{F}{W}^{1/2}y$$

denotes its dual. ${\nabla }_{\mathcal{F}}$ and ${\nabla }_{\mathcal{F}}^{*}$ are discrete analogies of general gradient (∇) and divergence (∇∙) in ${ℝ}^N$. The latter will be essential in formulating the vector-valued continuity equation [7] given below.

B. VECTOR-VALUED DENSITIES

A vector-valued density $\rho ={[{\rho }^1,{\rho }^2,\cdots ,{\rho }^M]}^T$ on ${ℝ}^N$ may represent a physical entity that can mutate or transit between alternative manifestations, e.g., power reflected off a surface at different frequencies or polarizations. More formally, an M-channel vector-valued density ρ on ${ℝ}^N$ is a map from ${ℝ}^N$ to ${ℝ}_{+}^M$ whose total mass is defined as ${\sum }_{i=1}^M{\int }_{{ℝ}^N}{\rho }^i(x)dx$.

We denote the set of all vector-valued densities and its interior by $\mathcal{D}$ and ${\mathcal{D}}_{+}$, respectively. We consider the components ρⁱ of ρ to represent the density or mass of particles in each of the various states that can mutate between one another. In the standard formulation, ρ₀, ρ₁ need to be balanced, i.e., have the same total mass. We may enforce this by simply normalizing the distributions to be probability distributions, i.e., ${\sum }_{i=1}^M{\int }_{{ℝ}^N}{\rho }_0^i(x)dx=1={\sum }_{i=1}^M{\int }_{{ℝ}^N}{\rho }_1^i(x)dx$. In the unbalanced case, total mass is not required to be preserved. This is very important for a number of applications in which mass may be created or destroyed.

C. BALANCED VECTOR-VALUED OMT

We now outline the balanced version of vector-valued OMT from [7]. With more than one channel, the change of mass is not only due to the flux of its spatial neighborhood but also the redistribution within the given channels. Transport between channels in a vector-valued density ρ is accomplished by treating ρ graphically. Specifically, each channel of ρ is represented by a node and an edge between two nodes allows for direct transport between the corresponding channels. We therefore use node and channel interchangeably. Namely, we have

$$\begin{gathered} {\mathcal{W}}_V{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,u}{\inf }{\int }_0^1{\int }_E{p}^T\mathrm{diag}{(\rho )}^{-1}p \\ +\gamma u^T[\mathrm{diag}{\left({\mathbb{F}}_2^T\rho \right)}^{-1}+\mathrm{diag}{\left({\mathbb{F}}_1^T\rho \right)}^{-1}]udxdt \end{gathered}$$ (6a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p-{\nabla }_{\mathcal{F}}^{*}u=0$$ (6b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1.$$ (6c)

As mentioned above, there are two kinds of flows for the vector-valued density ρ, reflected respectively by last two terms in the continuity function (6b). The first is the spatial flux p defined on a subdomain $E\subset {ℝ}^N$ whose corresponding divergence is denoted as ∇_x∙ and the second is the flux u between channels, whose discrete divergence is denoted by ${\nabla }_{\mathcal{F}}^{*}$. As above, $\mathcal{F}$ denotes a connected positively weighted graph. In the above expression, p and u are both vectors. We note that ${\mathbb{F}}_1$ is the all-1 part of the incident matrix $\mathbb{F}$ (sources of all the edges) and ${\mathbb{F}}_2={\mathbb{F}}_1-\mathbb{F}$ (sinks of all the edges). As u is defined on every edge and the density is only defined on nodes, ${[\mathrm{diag}{\left({\mathbb{F}}_2^T\rho \right)}^{-1}+\mathrm{diag}{\left({\mathbb{F}}_1^T\rho \right)}^{-1}]}^{-1}$ assigns a density to each edge using the density of two end points so that we can compute the kinetic energy for the flows on edges.

Vector-valued OMT can be considered as a general OMT problem on $E\times \mathcal{F}$, and so we can rewrite the energy functional in the following equivalent form:

$$\begin{gathered} {\mathcal{W}}_V{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,u}{\inf }{\int }_0^1{\int }_E\sum _{c\in Nodes(\mathcal{F})}\frac{p{(t,x,c)}^2}{\rho (t,x,c)}dxdt \\ +\gamma {\int }_0^1{\int }_E\sum _{e\in Edges(\mathcal{F})}\frac{u{(t,x,e)}^2}{\tilde{\rho }(t,x,e)}dxdt, \end{gathered}$$ (7)

where

$$\tilde{\rho }(t,x,e)=1/(\frac{1}{\rho (t,x,start(e))}+\frac{1}{\rho (t,x,end(e))})$$

is the harmonic average of the density at the endpoints, denoted start(e) and end(e), of edge e, which assigns a density to each edge.

D. UNBALANCED VECTOR-VALUED OMT

In this section, we write down the unbalanced vector-valued version of OMT, directly generalizing [9]. Namely,

$$\begin{gathered} {\mathcal{W}}_{VS}{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,u}{\inf }{\int }_0^1{\int }_E{p}^T\mathrm{diag}{(\rho )}^{-1}p \\ +\gamma u^T\left[\mathrm{diag}{\left({\mathbb{F}}_2^T\rho \right)}^{-1}+\mathrm{diag}{\left({\mathbb{F}}_1^T\rho \right)}^{-1}\right]u \\ +\eta s^T\mathrm{diag}{(\rho )}^{-1}sdxdt \end{gathered}$$ (8a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p-{\nabla }_{\mathcal{F}}^{*}u=s$$ (8b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1.$$ (8c)

Note that the source term is a space-time varying vector-valued map $s:{ℝ}^N\to {ℝ}^M$, and thus we need to write the energy functional for this vector term. Regarding the continuity equation, there are three possibilities for the change of mass over time: the flux of its neighborhood, flows between channels, and the source.

V. REFORMULATION OF UNBALANCED SCALAR OMT AS VECTOR-VALUED OMT

When considering the form of vector OMT in (6), we should note the similarity with unbalanced OMT (5). The flow between channels may be utilized as the source. The source term describes the creation or vanishing of mass. This can be modeled via an extra layer: the created mass originates from the added extra layer and the vanishing mass goes to that layer.

Here the input starting and target are regarded as two mass distributions on the given subdomain E. Total mass may or may not be preserved. We add an extra source layer which is parallel to the original space. For each point in E, there is an edge connecting it to the source layer. As we now have an extra layer, we can put the net difference of mass into the new layer so that the total mass of the new 2-vector structure is preserved. The net difference of mass can simply be distributed uniformly over the source layer or given a preferred spatial distribution of interest for further specific applications.

Now we view the flow between channels as the source term. Consider the continuity equations of (5b) and (6b):

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p=s$$

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p={\nabla }_{\mathcal{F}}^{*}u.$$

We can take ${\nabla }_{\mathcal{F}}^{*}u$ to act as the source term s and because of the special graph structure (at each point there is only one edge), we have:

$${\nabla }_{\mathcal{F}}^{*}u=u=s.$$ (9)

Now consider the second term in (7). Since there is only one edge, we can just omit the summation:

$$\begin{gathered} {\int }_0^1{\int }_E\sum _{e\in Edges(F)}\frac{u{(t,x,e)}^2}{\tilde{\rho }(t,x,e)}dxdt={\int }_0^1{\int }_E\frac{u{(t,x)}^2}{\tilde{\rho }(t,x)}dxdt \\ ={\int }_0^1{\int }_E\frac{u{(t,x)}^2}{\rho \left(t,x,c_1\right)}+\frac{u{(t,x)}^2}{\rho \left(t,x,c_2\right)}dxdt, \end{gathered}$$ (10)

where c₁ denotes the original layer and c₂ denotes the source layer; where the two end points are located (i.e., e = (c₁, c₂)). Taking the density of each edge as $\tilde{\rho }(t,x)=\rho (t,x)$, this term is exactly the same as the second term in the energy functional in the original unbalanced OMT setting (5).

We can still preserve the vector OMT setting. Note that we only use the source layer to compensate the mass difference. If we add a very large amount of mass g to both the starting and target source layers to make the density of the source layer large at all times, the second term in (10) will be very small (close to zero). Then $\tilde{\rho }(t,x)\approx \rho (t,x)$, which makes this term almost the same as the second term in the energy functional in unbalanced OMT (5).

Next we consider the first term of the energy functional (7):

$$\begin{gathered} {\int }_0^1{\int }_E\sum _{c\in Nodes(F)}\frac{p{(t,x,c)}^2}{\rho (t,x,c)}dxdt \\ ={\int }_0^1{\int }_E\frac{p{\left(t,x,c_1\right)}^2}{\rho \left(t,x,c_1\right)}+\frac{p{\left(t,x,c_2\right)}^2}{\rho \left(t,x,c_2\right)}dxdt. \end{gathered}$$ (11)

There are only two channels. Notice that the integral term over c₁ is exactly the first term of (5). Although we added a very large amount of mass to the source layer, the integral term over c₂ is not necessarily small. Due to the high density, mass may move freely in the given layer so the spatial flux p(t, x, c₂) can be very large. In order to get rid of the c₂ integral term, we can simply introduce a small weight parameter for that layer.

A. A NAIVE GENERALIZATION OF VECTOR-VALUED OMT: ADDITION OF WEIGHT

In the original setting [7], one treats the kinetic energy associated with each layer in an identical manner. However, by simply introducing a weight, we can treat each layer differently:

$$\begin{gathered} {\mathcal{W}}_{V^{\prime }}{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,u}{\inf }{\int }_0^1{\int }_E\sum _{c\in \mathrm{Nodes}(\mathcal{F})}w(c)\frac{p{(t,x,c)}^2}{\rho (t,x,c)}dxdt \\ +\gamma {\int }_0^1{\int }_E\sum _{e\in Edges(\mathcal{F})}\frac{u{(t,x,e)}^2}{\tilde{\rho }(t,x,e)}dxdt, \end{gathered}$$ (12)

where w(c) is the set of weighting parameters. We can thus express vector OMT with different chosen weights for each channel collected in the vector w in the following manner:

$$\begin{gathered} {\mathcal{W}}_{V^{\prime }}{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,u}{\inf }{\int }_0^1{\int }_E{p}^T\mathrm{diag}(w)\mathrm{diag}{(\rho )}^{-1}p \\ +\gamma u^T\left[\mathrm{diag}{\left({\mathbb{F}}_2^T\rho \right)}^{-1}+\mathrm{diag}{\left({\mathbb{F}}_1^T\rho \right)}^{-1}\right]udxdt. \end{gathered}$$ (13)

In this setting, if we take w(c₁) = 1 and w(c₂) = ε very small, then the integral in the source layer will also be very small in the energy functional. Clearly, the standard formulation of unbalanced OMT (5) is almost equivalent to the extra source layer vector OMT with the latter set of parameters.

B. IMPLEMENTATION: SCALAR CASE

By introducing the source layer, there is a new edge between the two channels, which utilizes the flow on the edge as the source, and hence there is an extra integral defined on the new layer. We add a very large amount of mass and use a very small weight parameter so that the two energy functionals are almost identical. See Figure 1.

FIGURE 1.

The top layer (black) is the original scalar distribution and the bottom layer (gray) is the appended source layer. The density of the source layer is very high and the weight for the flow within it is very small. The flow on the edge between two layers is our source.

In our setup, it is quite straightforward to implement unbalanced OMT from the vector-valued OMT algorithm [8]. In fact, we only need to alter a few lines of the code in the numerical implementation of the algorithm described in [8] to make it work for the unbalanced case:

1. Initialization: From the input, first construct two extended structures for vector OMT. The first layer is the original input and the second layer contains the mass difference.

2. Set density for source layer: If the starting density distribution has more total mass, then the mass difference is added to the second layer of target density distribution. If the starting density distribution has smaller total mass, then the mass difference is added to the second layer of the starting density distribution itself. In addition to the mass difference, we add a large number g to both starting and target distributions in the second layer.

3. Set weighting parameters: Employing the algorithm from [8] for the computation of the energy functional, add weighting parameters w(c₁) or w(c₂) to the corresponding layers. See Figures 3 and 4 for code snippets added to the original code from [8] to carry out the proposed unbalanced method.

FIGURE 3.

Code snippet 1: Construction of extended structures for starting and target distributions. d is the total mass difference. When d > 0, the mass difference is uniformly added to the source layer of target extended structure (rho1_ext). When d < 0, the mass difference is uniformly added to the source layer of starting extended structure (rho0_ext). A large number g is added to both the starting and target source layers.

FIGURE 4.

Code snippet 2: Addition of weight matrices. w1 is the diagonal weight matrix for different layers and w2 is another weight matrix containing the parameters γ and η. The weight matrices are added to the cost function as well as derivatives for the numerical algorithm.

VI. REFORMULATION OF UNBALANCED VECTOR OMT TO VECTOR-VALUED OMT

The reformulation of the vector-valued case is very similar to the scalar case. The same simple idea is to add a new source layer which connects to each of the original layers. The flows on those added edges are our sources.

A. SOURCE LAYER: VECTOR-VALUED CASE

In the vector-valued case, the input source and target are n-vector valued distributions. (See Section IV-B for the definition.) Total mass may or may not be preserved. We add an extra source layer which is connected to each of the existing layers. As before, we can distribute the net difference of mass over the new layer so that the total mass of the (n + 1)-vector new structure is preserved.

We denote the graph of the newly added edges by $\mathcal{G}$, and the expanded system by $\tilde{\mathcal{F}}$. Thus we have:

$${\nabla }_{\tilde{\mathcal{F}}}^{*}\left[\begin{array}{c}u\\ v\end{array}\right]=\left[{\nabla }_{\mathcal{F}}^{*}{\nabla }_{\mathcal{G}}^{*}\right]\left[\begin{array}{c}u\\ v\end{array}\right]={\nabla }_{\mathcal{F}}^{*}u+{\nabla }_{\mathcal{G}}^{*}v,$$ (14)

where u is the flow within the existing edges and v is the flow within the newly added edges.

Next we consider the continuity equations of unbalanced vector OMT and balanced vector OMT with the new layer:

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p-{\nabla }_{\mathcal{F}}^{*}u=s$$

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p-{\nabla }_{\mathcal{F}}^{*}u={\nabla }_{\mathcal{G}}^{*}v.$$

We can use ${\nabla }_{\mathcal{G}}^{*}v$ as the source s, and similarly to the scalar case:

$${\nabla }_{\mathcal{G}}^{*}v=v=s.$$ (15)

Thus we define:

$$\begin{gathered} {\mathcal{W}}_{V^{\prime \prime }}{\left({\rho }_0,{\rho }_1\right)}^2:=\underset{\rho ,p,u,v}{\inf }{\int }_0^1{\int }_E\sum _{c\in \mathit{Nodes}(\tilde{\mathcal{F}})}\frac{p{(t,x,c)}^2}{\rho (t,x,c)}dxdt \\ +\gamma {\int }_0^1{\int }_E\sum _{e\in Edges(\mathcal{F})}\frac{u{(t,x,e)}^2}{\tilde{\rho }(t,x,e)}dxdt \\ +\eta {\int }_0^1{\int }_E\sum _{e\in Edges(\mathcal{G})}\frac{v{(t,x,e)}^2}{\tilde{\rho }(t,x,e)}dxdt. \end{gathered}$$ (16)

Exactly as in the scalar case, if we add a very large amount of mass g to both starting and target source layers, $\tilde{\rho }(t,x,e)\approx \rho (t,x,start(e))$, where start(e) is the vertex of the edge e located in the original distribution space. Taken together with the above relationship (15), it is easy to see that the last term fits the original source energy term (8a).

Again there is a new component in the kinetic energy term due to the newly added source layer. Specifically, p and ρ are now n + 1-vectors, and the same goes for $\tilde{\rho },u$ and v. We can introduce weight parameters w₁(c_i) for i = 1, …, M + 1 which is 1 for original layers (i = 1, …, M) and a small value ε for the new kinetic term (i = M + 1) to make this extra term negligible. Indeed, using the identical technique described above in the scalar case of adding weights for different layers, we can incorporate γ and η into a weight matrix for different edges in order to make the final form more concise:

$$\begin{gathered} {\mathcal{W}}_{V^{\prime \prime }}{\left({\rho }_0,{\rho }_1\right)}^2=\underset{\rho ,p,\tilde{u}}{\inf }{\int }_0^1{\int }_E{p}^T\mathrm{diag}\left(w_1\right)\mathrm{diag}{(\rho )}^{-1}p \\ +{\tilde{u}}^T\mathrm{diag}\left(w_2\right)\mathrm{diag}{(\tilde{\rho })}^{-1}\tilde{u}dxdt \end{gathered}$$ (17a)

$$\frac{\partial \rho }{\partial t}+{\nabla }_x\cdot p-{\nabla }_{\tilde{\mathcal{F}}}^{*}\tilde{u}=0$$ (17b)

$$\rho (0,\cdot )={\rho }_0,\rho (1,\cdot )={\rho }_1,$$ (17c)

where diag(w₁) denotes the weight matrix for different layers, and diag(w₂) denotes the weight matrix for the added edges (weight= η) and existing edges (weight= γ). The flow between channels within the new large graph $\tilde{\mathcal{F}}$ is denoted by $\tilde{u}=\left[\begin{array}{l}u\hfill \\ v\hfill \end{array}\right]$ and $\tilde{\rho }$ is the density assigned to each edge.

B. IMPLEMENTATION: VECTOR-VALUED CASE

The addition of a source layer introduces edges between the source layer and each existing channel. We use the flows on those edges as sources, see Figure 2. As before, we add very small weight parameters for the source layer so that the two energy functionals are almost identical.

FIGURE 2.

Illustration of a vector-valued distribution with source. The red, green and blue channels of a color image constitute a vector-valued distribution. Each channel is depicted as a layer and the bottom layer (gray) is the source layer. The density of the source layer is very high and the weight for the flow within it is very small. The flow on the edges between the source layer and other existing layers is our source.

As above in the scalar case, it is straightforward to implement unbalanced vector OMT from the vector OMT algorithm described in [8]:

1. Initialization: From the input, first construct the extended structures for vector OMT. Add an extra layer and connect that layer with each of the existing layers.

2. Set density for source layer: Put a large number g plus the difference of total mass into that newly added layer.

3. Set weight parameters: Employing the algorithm proposed in [8] for the computation of the energy functional, add weighting parameters w₁(c) and w₂(c) to corresponding layers and corresponding edges (existing edges and newly added edges). See Figures 3 and 4 for code snippets added to the original code from [8] to carry out the proposed unbalanced method in the present work.

With the simple change in interpretation of the original vector-valued OMT, we can use the same code for unbalanced vector OMT case without even touching the main structure of the original code based on the numerical method of [8].

VII. NUMERICAL RESULTS

We tested our new formulation on several images using the numerical algorithm from [8]. Gray scale images are general mass distributions on a rectangular area while color images are vector valued distributions.

A. CHOICE OF PARAMETERS

As described above, we added two parameters as inputs to the vector OMT code: the first is a large number g added to the starting and target source layers and the second is a small weight ε for the new kinetic energy term due to the newly added source layer. After a number of tests, we found that results do not seem to be very sensitive to the choice of parameters, so we can choose them in a wide range.

More specifically, for general images in the uint8 format (gray scale or color), we used the original scale of [0, 255]. In numerical experiments, we found that the large number g for the source term may be chosen from 10³ to 10⁶, and the small weight ε for the kinetic energy may be chosen from 10⁻² to 10⁻⁴, giving similar results in the simulations.

B. UNBALANCED OMT

We tested an example of moving two Gaussian distributions. Though this example has preserved total mass, the optimal solution of OMT still uses the source term. We can look at the effect of tuning the weight parameter γ. The starting and target images are shown in Figure 5.

FIGURE 5.

Balanced Gaussians: starting and target distributions

Using the two distributions of Figure 5 as inputs to our algorithm, we obtained a series of interpolating distributions, which represent the geodesic path connecting starting and target distributions for the vector-valued metric with large and small γ. We show the density of the original space as well as the source layer (g is subtracted before plotting) at 5 intermediate time points. See Figures 6, 7, 8, and 9.

FIGURE 6.

Balanced Gaussians with large γ: density over time

FIGURE 7.

Balanced Gaussians with large γ: source over time

FIGURE 8.

Balanced Gaussians with small γ: density over time

FIGURE 9.

Balanced Gaussians with small γ: source over time

With large γ, it is expensive to use the source term, so the source layer is almost zero at all times, as shown in Figures 6 and 7.

With small γ, it is cheap to use the source term, thus much of the mass goes through the source layer, see Figures 8 and 9.

We also tested two images with a very large total mass difference, as shown in Figure 10. The corresponding interpolated density and source images are shown respectively in Figures 11 and 12.

FIGURE 10.

Unbalanced Gaussians: starting and target distributions

FIGURE 11.

Unbalanced Gaussians: density over time

FIGURE 12.

Unbalanced Gaussians: source over time

C. UNBALANCED VECTOR OMT

We tested our unbalanced vector-valued approach on color image data, where the vector values are taken as the red, green and blue color channels. While the interpolations of density are color, the source layer is still gray scale (same as the scalar case, the large number g is subtracted before plotting). Specifically, we demonstrate our method on color images of a solar flare (Figures 13–15) and a supernova (Figures 16–18).

FIGURE 13.

Solar flare: starting and target distributions

FIGURE 15.

Solar flare: source over time

FIGURE 16.

Supernova: starting and target distributions

FIGURE 18.

Supernova: source over time

The interpolated images of two frames of solar flare give an evolution path that looks quite natural. Clearly the total mass is not preserved because of the complicated energy changes during this diffusive process. By absorbing the extra density in the source layer, unbalanced OMT gives a reasonable displacement interpolation.

The major mass difference in the latter example is due to the outbreak of the supernova in the image, in which the background only slightly moves. Our algorithm captures this process by using the source in the vicinity of the supernova region. See Figures 17 and 18.

FIGURE 17.

Supernova: density over time

VIII. CONCLUSION

Vector-valued OMT is a very powerful model. In this study, we reformulated unbalanced OMT and unbalanced vector-valued OMT by adding the source as a new layer. We gave a new way to include a source term for both scalar and vector-valued densities and we proposed a very simple way to implement unbalanced OMT from the general vector-valued OMT code. In the present work, we only considered one kind of unbalanced formulation, namely Fisher-Rao. However, our model is adaptable to fit other flow-based unbalanced OMT settings as well. We believe our generalization of vector-valued OMT has not reached its full potential. Indeed, we only give different weight parameters to different layers but could also consider giving different weight parameters to different areas, different nodes or different edges. Different parameters may be used according to specific applications. This will be the subject of future research. More specifically, we plan to employ our methodology to analyze medical data, in particular radiation dose distribution problems.

FIGURE 14.

Solar flare: density over time

Biography

JIENING ZHU Jiening Zhu is a graduate student in Applied Mathematics and Statistics at Stony Brook University. His research is in bioinformatics and medical imaging with an emphasis to extending vector-valued optimal mass transport for studying cancer biology.

RENA ELKIN Dr. Rena Elkin received her Ph.D. in December 2019. She is presently a postdoctoral fellow in Medical Physics at Memorial Sloan Kettering Cancer Center. Her research is concerned with applying optimal mass transport for studying tumor vasculature and geometric graphical methods for cancer networks.

JUNG HUN OH Dr. Jung Hun Oh is an assistant attending physicist in Medical Physics at Memorial Sloan Kettering Cancer Center. He uses cutting-edge computational and statistical methods, informed by bioinformatics and machine learning techniques, to identify novel diagnostic biomarkers and to build models that predict radiation treatment outcomes.

JOSEPH O. DEASY Dr. Joseph O. Deasy is the Chair of the Department of Medical Physics and Enid A. Haupt Chair in Medical Physics at Memorial Sloan Kettering Cancer Center. He and his group apply statistical modeling to the analysis of large, complex datasets in order to understand the relationship between treatment, patient, and disease characteristics and the probability of local control and normal tissue toxicity.

ALLEN TANNENBAUM Dr. Allen Tannenbaum is Distinguished Professor of Computer Science and Applied Mathematics and Statistics at Stony Brook University. His research is in biological networks, medical imaging, computer vision, and control theory.