Fast Fourier Transformation for Optimal Transport Map (FFT-OT)

buddaha
input
input
Alex
The input image: Lenna The means and the deviations of Guassian densities The optimal transport map The Gaussian mixture

Optimal Transport and the Monge-Ampere Equation

Given a pair of measures $(\Omega,\mu)$ and $(\Omega^*,\nu)$ with densities \[ d\mu(x) = f(x)dx,\quad d\nu(y) = g(y)dy, \] such that $\mu(\Omega)=\nu(\Omega^*)$. A mapping $T:\Omega\to\Omega^*$ is measure-preserving, if for any Borel set $B\subset \Omega^*$, we have \[ \int_{T^{-1}(B)} d\mu = \int_B d\nu. \] A meseaure preserving map is denoted as $T_\#\mu=\nu$. The cost function $c:\Omega\times \Omega^*\to\mathbb{R}$ defines the cost for transporting a unit mass from $x$ to $y$, the total transportation cost for a transportation map $T$ is given by \[ \int_\Omega c(x,T(x))d\mu(x). \] The Monge's problem asks to find the measure-preserving map $T$ that minimizes the total transportation cost, \[ \min_{T_\#\mu = \nu }\int_\Omega c(x,T(x)) d\mu(x), \] such a map is called the optimal transportation map. Brenier's theorem shows that for $c(x,y)=\frac{1}{2}|x-y|^2$, the optimal transportation map is the gradient map of a convex function $u:\Omega\to\mathbb{R}$, the so-called Brenier potential. The Brenier potential satisfies the Monge-Ampere equation: \[ det(D^2 u)(x) = \frac{f(x)}{g\circ \nabla u(x)}, \] with natural boundary condition $\nabla u(\Omega)=\Omega^*$.

The Poisson Equation

In 2D cases, the Monge-Ampere equation can be rewritten as \[ u_{xx}u_{yy}-u_{xy}^2 = \frac{f}{g\circ Du}, \] hence \[ (u_{xx}+u_{yy})^2 = u_{xx}^2 + u_{yy}^2 + 2u_{xy}^2 + 2 \frac{f}{g\circ Du}, \] convert to Poisson equation \[ \Delta u = \sqrt{u_{xx}^2 + u_{yy}^2 + 2u_{xy}^2 + 2 f/g\circ Du}. \]

Fixed Point

Let's define an operator $\mathcal{T}:H^2(\Omega)\to H^2(\Omega)$, \[ \mathcal{T}[u] = \Delta^{-1}\left\{\sqrt{u_{xx}^2 + u_{yy}^2 + 2u_{xy}^2 + 2 f/g\circ Du}\right\}, \] then the Brenier potential is the fixed point of $\mathcal{T}$. Suppose $\Omega$ is a rectangular domain, we define \[ u(x,y) = \varphi(x,y) + \frac{1}{2}(x^2+y^2),\quad Du=D\varphi + Id, \quad \Delta u=\Delta \varphi+2, \] we convert the operator $\mathcal{T}$ to $\mathcal{P}$, \[ \mathcal{P}(\varphi):= \Delta^{-1}\left\{ \sqrt{(\varphi_{xx}+1)^2+(\varphi_{yy}+1)^2+2\varphi^2_{xy}+2f/g\circ(Id+D\varphi)}-2 \right\} \] Suppose $\varphi^*$ is the fixed point of $\mathcal{P}$, $\mathcal{P}(\varphi^*)=\varphi^*$, then the optimal transportation map is given by $D\varphi+Id$. The fixed point $\varphi^*$ can be obtained by iterations \[ \varphi^{(n+1)}\leftarrow \mathcal{P}\left(\varphi^{(n)}\right),\quad \lim_{n\to\infty}\varphi^{(n)}\to \varphi^* \]

Poisson Operator

We sample $\Omega$ by regular grid points with horizontal and vertical step lengths $h_x$ and $h_y$ respectively, the finite difference operators are defined as \[ \begin{split} \mathcal{D}_{xx}^2 u_{ij} &= \frac{1}{h_x^2} (u_{i+1,j}+u_{i-1,j}-2u_{i,j}) \\ \mathcal{D}_{yy}^2 u_{ij} &= \frac{1}{h_y^2} (u_{i,j+1}+u_{i,j-1}-2u_{i,j}) \\ \mathcal{D}_{xy}^2 u_{ij} &= \frac{1}{4h_xh_y}(u_{i+1,j+1}+u_{i-1,j-1}-u_{i-1,j+1}-u_{i+1,j-1}) \end{split} \] and obtain the discrete Poisson equation \[ u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1} - 4u_{i,j} = \rho_{i,j}, \] the obliqueness boundary condition is converted to Neumann boundary condition for the Poisson equation.

Discrete Cosine Transformation

Given two dimensional array $u(i,j)$, the 2D discrete cosine transformation is defined as: \[ \tilde{u}(m,n) = c(m,n) \sum_{i,j} u(i,j)\cos \frac{(2i+1)m\pi}{2M} \cos \frac{(2j+1)n\pi}{2N}, \] the inverse discrete Cosine Transformation is given by: \[ u(i,j)=\sum_{m,n} c(m,n)\tilde{u}(m,n)\cos\frac{(2i+1)m\pi}{2M}\cos\frac{2(2j+1)n\pi}{2N}, \] where $m,i=0,\dots,M-1$ and $n,j=0,\dots,N-1$, \[ c(m,n)=\left\{ \begin{array}{ll} \frac{\sqrt{2}}{\sqrt{MN}}& m=0,n=0\\ \frac{2}{\sqrt{MN}}& otherwise\\ \end{array} \right. \]

Poisson Solver using FFT

Given a discrete Poisson equation with the Neumann boundary condition: \[ \Delta u = \rho,\quad \frac{\partial u}{\partial n} = 0. \] Let $\tilde{\rho} = DCT(\rho)$, $\tilde{u} = DCT(u)$, then we have \[ \tilde{u}(m,n) = \frac{\tilde{\rho}(m,n)}{2[\cos\frac{m\pi}{M}+\cos\frac{n\pi}{N}-2]}. \] Different solutions differ by a constant. Let $\tilde{u}(0,0)=0$, we get the unique solution. DCT and IDCT can be computed using FFT on GPUs.
buddaha
input
input
Alex
The input image: Mona Lisa The means and the deviations of Guassian densities The optimal transport map The Gaussian mixture
buddaha
input
input
Alex
The input image: Sophie The means and the deviations of Guassian densities The optimal transport map The Gaussian mixture

Comannd Line

FFT_OT.exe -UI < image_name >

Examples: FFT_OT.exe -UI lenna.png

Hot Keys

Hold the Ctrl key, click the left mouse button to select the mean $\mu$ of the Gauss distribution; Hold the Ctrl key, click the right mouse button to select the standard deviation $\sigma$ of the Gauss distribution; Press 'c' to visualize the support of the Gaussian distribution \[ \mu = \begin{pmatrix} \mu_x\\ \mu_y \end{pmatrix},\quad \Sigma = \begin{pmatrix} \sigma^2 &0\\ 0& \sigma^2 \end{pmatrix}, \quad f(x,y) = \frac{1}{2\pi \sigma^2} \exp \left(-\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu)\right). \] One can select multiple Gaussian distributions using the same procedure to define a Gaussian mixture, which is the target measure. The source measure is the standard Lebesgue measure.
FFT_OT.exe hot keys the following options are supported:
'c' show the supports of the Gaussian distrubitions $(\mu,\sigma)$'s
'o' Compute the optimal transportation map
's' Continue to display the next OT image
'Esc' Terminate the program

Download and Instructions

Environment and Dependencies

  1. The FFT_OT mapper is developed using generic C++ language on Windows Visual Studio.
  2. You can use OpenCV for FFT and image processing.

References

For any further information, comments and criticisms, please contact David Gu.