An Efficient Algorithm for Matrix-Valued and Vector-Valued Optimal Mass Transport

Abstract

We present an efficient algorithm for recent generalizations of optimal mass transport theory to matrix-valued and vector-valued densities. These generalizations lead to several applications including diffusion tensor imaging, color image processing, and multi-modality imaging. The algorithm is based on sequential quadratic programming. By approximating the Hessian of the cost and solving each iteration in an inexact manner, we are able to solve each iteration with relatively low cost while still maintaining a fast convergence rate. The core of the algorithm is solving a weighted Poisson equation, where different efficient preconditioners may be employed. We utilize incomplete Cholesky factorization, which yields an efficient and straightforward solver for our problem. Several illustrative examples are presented for both the matrix and vector-valued cases.

Appendix

As in the matrix-valued case, for simplicity of exposition, we consider the discretization in 1D case, and then describe the 2D case in Sect. 7.4 below. Thus, we take $E=[0,1]$, and as before our technique extends almost verbatim to the higher dimensional setting. We should note that the algorithm presented here in the vector-valued case is very similar to the matrix optimal transport just described in the preceding sections. Nevertheless, for the sake of completeness, we present the key details here.

We discretize the space-time domain $[0,\,1]\times [0,\,1]$ into $n_x\times n_t$ rectangular cells. Denote $\Omega _{ij}, 1\le i\le n_x, 1\le j\le n_t$ as the (i, j) box. We use staggered grid to discretize p and $\rho$. The variable u is, however, valued at the centers of the cells $\{\Omega _{ij}\}$. More specifically,

$$\begin{aligned} p=\left( p_{i+\frac{1}{2},j}\right) , \quad 0\le i\le n_x, ~1\le j\le n_t \end{aligned}$$

$$\begin{aligned} \rho =\left( \rho _{i,j+\frac{1}{2}}\right) , \quad 1\le i\le n_x,~0\le j\le n_t \end{aligned}$$

$$\begin{aligned} u=\left( u_{i,j}\right) , \quad 1\le i \le n_x, ~1\le j\le n_t. \end{aligned}$$

Note that the boundary values are

$$\begin{aligned} p_{\frac{1}{2},j}=0, \quad p_{n_x+\frac{1}{2},j}=0, ~1\le j\le n_t \end{aligned}$$

and

$$\begin{aligned} \rho _{i,\frac{1}{2}}=\rho ^0_i, \quad \rho _{i,n_t+\frac{1}{2}}=\rho ^1_i, ~1\le i\le n_x. \end{aligned}$$

We exclude the boundary values from the variables and denote

$$\begin{aligned} p=\left( p_{i+\frac{1}{2},j}\right) , \quad 1\le i\le n_x-1, ~1\le j\le n_t \end{aligned}$$

$$\begin{aligned} \rho =\left( \rho _{i,j+\frac{1}{2}}\right) , \quad 1\le i\le n_x,~1\le j\le n_t-1. \end{aligned}$$

1.1 Continuity Equation: Vector-Valued Case

We use the preceding discretizing scheme, together with the boundary conditions to rewrite the continuity Eq. (13b) as

$$\begin{aligned} D_1 p+D_2 \rho +D_3 u=b. \end{aligned}$$

(14)

Here the linear operators $D_1, D_2, D_3$ are defined as

$$\begin{aligned}&\left( D_1 p\right) _{i,j}= {\left\{ \begin{array}{ll} \left( p_{i+\frac{1}{2},j} -p_{i-\frac{1}{2},j}\right) /h_x, &{} \quad \text{ if } ~2\le i \le n_x-1, \\ p_{\frac{3}{2},j}/h_x, &{} \quad \text{ if } ~i=1, \\ -p_{n_x-\frac{1}{2},j}/h_x, &{} \quad \text{ if } ~i=n_x, \end{array}\right. }\\&(D_2 \rho )_{i,j}= {\left\{ \begin{array}{ll} \left( \rho _{i,j+\frac{1}{2}}-\rho _{i,j-\frac{1}{2}}\right) /h_t, &{} \quad \text{ if } ~2\le j \le n_t-1, \\ \rho _{i,\frac{3}{2}}/h_t, &{} \quad \text{ if } ~j=1, \\ - \rho _{i,n_t-\frac{1}{2}}/h_t, &{} \quad \text{ if } ~j=n_t, \end{array}\right. }\\&\left( D_3 u\right) _{i,j}=-\nabla _{{\mathcal {F}}}^* u_{i,j}, \quad 1\le i \le n_x, ~1\le j\le n_t. \end{aligned}$$

The parameter b carries the information of the boundary values $\rho ^0$ and $\rho ^1$. More specifically,

$$\begin{aligned} b_{i,j}= {\left\{ \begin{array}{ll} \rho ^0_i/h_t &{} \quad \text{ if }~j=1,\\ -\rho ^1_i/h_t &{} \quad \text{ if }~j=n_t,\\ 0 &{} \quad \text{ otherwise }. \end{array}\right. } \end{aligned}$$

1.2 Discretization of the Cost Function: Vector-Valued Case

Let $A_1$ be the averaging operator over the spatial domain and $A_2$ be the averaging operator over the time domain (as before one needs to be careful about the boundaries). Then the cost function (13a) may be approximated by

$$\begin{aligned} \left\langle A_1 \left( p^2\right) , A_2 (1/\rho )+a\right\rangle h_xh_t+ \left\langle u^2, A_2 \left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) +c\right\rangle \gamma h_xh_t, \end{aligned}$$

(15)

where $a\ge 0$ depends only on the boundary values $\rho ^0$ and $\rho ^1$. The inverse operator and multiplication operators are applied block-wise. The expressions for $A_1, A_2, a$ are

$$\begin{aligned}&\left( A_1 \left( p^2\right) \right) _{i,j}= {\left\{ \begin{array}{ll} \frac{1}{2}\left( p^2_{i-\frac{1}{2},j} +p^2_{i+\frac{1}{2},j}\right) , &{} \quad \text{ if } ~2\le i \le n_x-1, \\ \frac{1}{2}p^2_{\frac{3}{2},j}, &{} \quad \text{ if } ~i=1, \\ \frac{1}{2}p^2_{n_x-\frac{1}{2},j}, &{} \quad \text{ if } ~i=n_x, \end{array}\right. }\\&\left( A_2 \left( 1/\rho \right) \right) _{i,j}= {\left\{ \begin{array}{ll} \frac{1}{2}\left( 1/\rho _{i,j-\frac{1}{2}} +1/\rho _{i,j+\frac{1}{2}}\right) , &{}\quad \text{ if } ~2\le j \le n_t-1, \\ 1/\rho _{i,\frac{3}{2}}/2, &{} \quad \text{ if } ~j=1, \\ 1/\rho _{i,n_t-\frac{1}{2}}/2, &{} \quad \text{ if } ~j=n_t, \end{array}\right. }\\&a_{i,j}= {\left\{ \begin{array}{ll} 1/\rho ^{0}_i/2&{} \quad \text{ if }~j=1,\\ 1/\rho ^{1}_i/2 &{} \quad \text{ if }~j=n_t,\\ 0 &{} \quad \text{ otherwise }, \end{array}\right. }\\&c_{i,j}= {\left\{ \begin{array}{ll} 1/{{\mathbb {D}}}_2^T\rho ^{0}_i/2+1/{{\mathbb {D}}}_1^T\rho ^{0}_i/2&{} \quad \text{ if }~j=1,\\ 1/{{\mathbb {D}}}_2^T\rho ^{1}_i/2+1/{{\mathbb {D}}}_1^T\rho ^{1}_i/2 &{} \quad \text{ if }~j=n_t,\\ 0 &{} \quad \text{ otherwise }. \end{array}\right. } \end{aligned}$$

1.3 Sequential Quadratic Programming (SQP)

From the above discussion, we obtain the discrete convex optimization problem

$$\begin{aligned}&\min \quad f\left( p, \rho , u\right) =\left\langle A_1 \left( p^2\right) , A_2 \left( 1/\rho \right) \!+\!a\right\rangle h_xh_t \nonumber \\&\qquad \qquad \qquad \ \qquad \qquad + \left\langle u^2, A_2 \left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) \!+\!c\right\rangle \gamma h_xh_t \end{aligned}$$

(16a)

$$\begin{aligned}&\mathrm{subject to} \quad D_1 p+D_2 \rho +D_3 u=b. \end{aligned}$$

(16b)

The Lagrangian of this problem is

$$\begin{aligned} {{\mathcal {L}}}\left( p, \rho , u\right) = f\left( p, \rho , u\right) /\left( h_xh_t\right) +\left\langle \lambda , D_1 p+D_2 \rho +D_3 u-b\right\rangle . \end{aligned}$$

It follows that the KKT conditions are given by

$$\begin{aligned} \nabla _p {{\mathcal {L}}}= & {} D_1^T \lambda +2p\circ A_1^T\left( A_2\left( 1/\rho \right) +a\right) =0 \end{aligned}$$

(17a)

$$\begin{aligned} \nabla _\rho {{\mathcal {L}}}= & {} D_2^T \lambda -A_2^TA_1\left( p^2\right) / \rho ^{2} -\gamma {{\mathbb {D}}}_2\left( A_2^T\left( u^2\right) /\left( {{\mathbb {D}}}_2^T\rho \right) ^2\right) \nonumber \\&-\gamma {{\mathbb {D}}}_1\left( A_2^T\left( u^2\right) /\left( {{\mathbb {D}}}_1^T\rho \right) ^2\right) =0 \end{aligned}$$

(17b)

$$\begin{aligned} \nabla _u {{\mathcal {L}}}= & {} D_3^T\lambda +2\gamma u\circ \left( A_2\left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) +c\right) =0 \end{aligned}$$

(17c)

$$\begin{aligned} \nabla _\lambda {{\mathcal {L}}}= & {} D_1 p+D_2 \rho +D_3 u-b=0, \end{aligned}$$

(17d)

with $\circ$ denoting block-wise multiplication.

Let $w=(p,\rho , u).$ Then at each SQP iteration, we solve the system

$$\begin{aligned} \left( \begin{matrix}{\hat{A}} &{}\quad D^T \\ D &{}\quad 0\end{matrix}\right) \left( \begin{array}{c}\delta w\\ \delta \lambda \end{array}\right) = -\left( \begin{array}{c}\nabla _w{{\mathcal {L}}}\\ \nabla _\lambda {{\mathcal {L}}}\end{array}\right) , \end{aligned}$$

(18)

where $D=(D_1, D_2, D_3)$. Again, the matrix ${\hat{A}}$ is an approximation of the Hessian of the objective function

$$\begin{aligned} {\hat{A}}= \left( \begin{matrix} 2 \mathrm{diag} \left( A_1^T \left( A_2\left( 1/\rho \right) +a\right) \right) &{}\quad 0 &{} \quad 0 \\ 0 &{} \quad \mathrm{diag} \left( g\left( p,\rho ,u\right) \right) &{} \quad 0 \\ 0 &{} \quad 0 &{} \quad 2 \gamma \mathrm{diag} \left( A_2\left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) +c\right) \end{matrix}\right) . \end{aligned}$$

The operator $g(p, \rho , u)$ is the Hessian of f over $\rho$ with $g_{i,j+\frac{1}{2}}$ being the map

$$\begin{aligned} g_{i,j+\frac{1}{2}}(X)= & {} 2\left( A_2^TA_1\left( p^2\right) \right) _{i,j+\frac{1}{2}} /\rho ^{3}_{i,j+\frac{1}{2}}X\\&+ 2\gamma {{\mathbb {D}}}_2\left[ \left( A_2^T\left( u^2\right) \right) _{i,j+\frac{1}{2}} /\left( {{\mathbb {D}}}_2^T\rho \right) ^{3}_{i,j+\frac{1}{2}}{{\mathbb {D}}}_2^TX\right] \\&+2\gamma {{\mathbb {D}}}_1\left[ \left( A_2^T\left( u^2\right) \right) _{i,j+\frac{1}{2}} /\left( {{\mathbb {D}}}_1^T\rho \right) ^{3}_{i,j+\frac{1}{2}}{{\mathbb {D}}}_1^TX\right] . \end{aligned}$$

1.4 Higher Dimensions: Vector-Valued Case

We concretely work out the 2D case in this section. The higher dimensional cases are very similar, but naturally involve additional indices. We should note that for all of the applications we have in mind, the 2D case is sufficient.

We have the discrete convex optimization problem

$$\begin{aligned}&\min \quad f\left( p, \rho , u\right) =\left\langle A_{1x} \left( p_x^2\right) +A_{1y}\left( p_y^2\right) , A_2 \left( 1/\rho \right) + a\right\rangle h_xh_yh_t \\&\qquad \qquad \qquad \ \qquad \qquad + \left\langle u^2, A_2 \left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) +c\right\rangle \gamma h_xh_yh_t\\&\mathrm{subject~ to} \quad D_{1x} p_x +D_{1y} p_y+D_2 \rho +D_3 u=b. \end{aligned}$$

The Lagrangian of this problem is

$$\begin{aligned} {{\mathcal {L}}}\left( p, \rho , u\right) = f\left( p, \rho , u\right) /\left( h_xh_yh_t\right) +\left\langle \lambda , D_{1x} p_x+D_{1y} p_y+D_2 \rho +D_3 u-b\right\rangle . \end{aligned}$$

In the above,

$$\begin{aligned} a_{i,j,k}= {\left\{ \begin{array}{ll} 1/\rho ^{0}_{i,j}/2&{} \quad \text{ if }~k=1,\\ 1/\rho ^{1}_{i,j}/2 &{} \quad \text{ if }~k=n_t,\\ 0 &{} \quad \text{ otherwise }, \end{array}\right. } \end{aligned}$$

and

$$\begin{aligned}&b_{i,j,k}= {\left\{ \begin{array}{ll} \rho ^0_{i,j}/h_t &{} \quad \text{ if }~k=1,\\ -\rho ^1_{i,j}/h_t &{} \quad \text{ if }~k=n_t,\\ 0 &{} \quad \text{ otherwise }, \end{array}\right. }\\&c_{i,j,k}= {\left\{ \begin{array}{ll} 1/{{\mathbb {D}}}_2^T\rho ^{0}_{i,j}/2+1/{{\mathbb {D}}}_1^T\rho ^{0}_{i,j}/2&{} \quad \text{ if }~k=1,\\ 1/{{\mathbb {D}}}_2^T\rho ^{1}_{i,j}/2+1/{{\mathbb {D}}}_1^T\rho ^{1}_{i,j}/2 &{} \quad \text{ if }~k=n_t,\\ 0 &{} \quad \text{ otherwise }. \end{array}\right. } \end{aligned}$$

The KKT conditions now are

$$\begin{aligned} \nabla _{p_x} {{\mathcal {L}}}= & {} D_{1x}^T \lambda +2p_x\circ A_{1x}^T\left( A_2\left( 1/\rho \right) +a\right) =0\\ \nabla _{p_y} {{\mathcal {L}}}= & {} D_{1y}^T \lambda +2p_y\circ A_{1y}^T\left( A_2\left( 1/\rho \right) +a\right) =0\\ \nabla _\rho {{\mathcal {L}}}= & {} D_2^T \lambda -A_2^T\left( A_{1x}\left( p_x^2\right) +A_{1y}\left( p_y^2\right) \right) /\rho ^2 -\gamma {{\mathbb {D}}}_2\left( A_2^T\left( u^2\right) /\left( {{\mathbb {D}}}_2^T\rho \right) ^2\right) \\&-\gamma {{\mathbb {D}}}_1\left( A_2^T\left( u^2\right) /\left( {{\mathbb {D}}}_1^T\rho \right) ^2\right) =0\\ \nabla _u {{\mathcal {L}}}= & {} D_3^T\lambda +2\gamma u\circ \left( A_2\left( 1/\left( {{\mathbb {D}}}_2^T\rho \right) +1/\left( {{\mathbb {D}}}_1^T\rho \right) \right) +c\right) =0\\ \nabla _\lambda {{\mathcal {L}}}= & {} D_1 p+D_2 \rho +D_3 u-b=0, \end{aligned}$$

with $\circ$ denoting block-wise multiplication.

Let $w=(p_x,p_y,\rho , u)$, then at each SQP iteration we solve the system

$$\begin{aligned} \left( \begin{matrix}{\hat{A}} &{} \quad D^* \\ D &{} \quad 0\end{matrix}\right) \left( \begin{array}{c}\delta w\\ \delta \lambda \end{array}\right) = -\left( \begin{array}{c}\nabla _w{{\mathcal {L}}}\\ \nabla _\lambda {{\mathcal {L}}}\end{array}\right) , \end{aligned}$$

(19)

where $D=(D_{1x},D_{1y}, D_2, D_3)$. The matrix ${\hat{A}}$ is an approximation of the Hessian of the objective function

The operator $g(p, \rho , u)$ is the Hessian of f over $\rho$ with $g_{i,j,k+\frac{1}{2}}$ being the map

$$\begin{aligned} g_{i,j+\frac{1}{2}}(X)= & {} 2\left( A_2^T\left( A_{1x}\left( p_x^2\right) +A_{1y}\left( p_y^2\right) \right) \right) _{i,j+\frac{1}{2}} /\rho ^{3}_{i,j+\frac{1}{2}}X\\&+ 2\gamma {{\mathbb {D}}}_2\left[ \left( A_2^T\left( u^2\right) \right) _{i,j+\frac{1}{2}} /\left( {{\mathbb {D}}}_2^T\rho \right) ^{3}_{i,j+\frac{1}{2}}{{\mathbb {D}}}_2^TX\right] \\&+2\gamma {{\mathbb {D}}}_1\left[ \left( A_2^T\left( u^2\right) \right) _{i,j+\frac{1}{2}} /\left( {{\mathbb {D}}}_1^T\rho \right) ^{3}_{i,j+\frac{1}{2}}{{\mathbb {D}}}_1^TX\right] . \end{aligned}$$