============================================================== The Back-And-Forth Method (BFM) for Wasserstein Gradient Flows ============================================================== This page provides documentation for the source code used in the paper `The back-and-forth method for Wasserstein gradient flows `_ [1]. A MATLAB wrapper to our C++ code is provided here: https://github.com/wonjunee/wgfBFMcodes. Introduction ============ We are interested in simulating Darcy's law (also known as the generalized porous medium equation) .. math:: \partial_t \rho - \nabla \cdot (\rho \nabla \phi) = 0\\ \phi = \delta U(\rho) on a domain :math:`\Omega\subset\mathbb{R}^2`. This equation describes the evolution of a mass density :math:`\rho` flowing along a pressure gradient :math:`\nabla \phi` generated by an energy functional :math:`U`. The evolution of the density can be approximated by a discrete-in-time scheme known as the **JKO scheme** [3]. Given an initial density :math:`\rho^{(0)}`, this scheme constructs approximate solutions by iterating .. math:: :label: eq:jko \rho^{(n+1)} := \underset{\rho}{\text{argmin}} \,\, U(\rho) + \frac{1}{2\tau} W_2 (\rho, \rho^{(n)})^2. Here, :math:`\tau` plays the role of a time step and :math:`W_2(\cdot,\cdot)` is the 2-Wasserstein distance. We extend the `back-and-forth method `_ (BFM) [2] for optimal transport to efficiently solve (1). The **BFM** solves (1) via the dual formulation .. math:: :label: eq:dual \sup_{\phi,\psi} \int_\Omega \psi(x) \,\rho^{(n)}(x) dx - U^*(\phi), where the supremum runs over functions :math:`\phi,\psi\colon \Omega \rightarrow \mathbb{R}` satisfying .. math:: \psi(x) - \phi(y) \leq \frac{|x-y|^2}{2\tau}. Installation ============ .. highlight:: matlab Requirements: FFTW (`download here `_), MATLAB. Clone or download the `repository `_: .. code:: bash git clone https://github.com/wonjunee/wgfBFMcodes cd wgfBFMcodes/ Compile: in a MATLAB session run :: mex -O CFLAGS="\$CFLAGS -std=c++11" src/wgfslow.cpp -lfftw3 mex -O CFLAGS="\$CFLAGS -std=c++11" src/wgfinc.cpp -lfftw3 This will produce two MEX functions that you can use in MATLAB: :code:`wgfslow` for slow diffusions and :code:`wgfinc` for incompressible energies. You may need to use :code:`-I` and :code:`-L` to link to the FFTW3 library, e.g. :code:`mex -O CFLAGS="\$CFLAGS -std=c++11" src/wgfslow.cpp -lfftw3 -I/usr/local/include`. See `this page `_ for more information on how to compile MEX files. Slow diffusion ============== Usage ----- In a MATLAB session, the command :: rhoFinal = wgfslow(rhoInitial, V, obstacle, m, gamma, ... maxIters, TOL, tau, nt, folder, verbose); computes the JKO iterates :math:`\rho^{(n)}` from an initial density :math:`\rho^{(0)}` with the internal energy .. math:: :label: eq:slowenergy U(\rho) = \frac{\gamma}{m-1} \int_\Omega \rho^m(x) dx + \int_\Omega V(x)\,\rho(x)dx. The iterates :math:`\rho^{(n)}` will be stored as binary files :file:`rho-0000.dat`, :file:`rho-0001.dat`, etc in a user-specified folder. Input: * :code:`rhoInitial` : The nonnegative initial density :math:`\rho^{(0)}`. * :code:`V` : The potential function :math:`V(x)`. * :code:`obstacle` : An obstacle in :math:`\Omega`. Any positive values indicate the obstacle. The densities :math:`\rho^{(n)}` will always be :math:`0` on the obstacle. * :code:`m` : The constant :math:`m` in the internal energy :math:`U`. * :code:`gamma` : The constant :math:`\gamma` in the internal energy :math:`U`. * :code:`maxIters` : The maximum number of inner loop iterations :math:`k` of the BFM in each outer iteration of the JKO scheme. * :code:`TOL` : The tolerance for the BFM. Once the error :math:`\| T_{\phi_k\#} \rho^{(n)} - \delta U^*(\phi_k) \|_{L^1}` becomes less than the tolerance, the algorithm stops the BFM and proceeds to the next outer iteration :math:`n+1`. * :code:`tau` : The time step :math:`\tau` in (1). * :code:`nt` : The number of outer iterations :math:`n` of the JKO scheme. * :code:`folder` : Path to the folder where the successive iterates :math:`\rho^{(n)}` will be saved. Don't include a trailing slash. * :code:`verbose` : :code:`1` - print out process in :code:`c++` code. :code:`0` - be silent. Output: * :code:`rhoFinal` : The final density :math:`\rho^{(nt)}` after :code:`nt` outer iterations. Example ------- We compute the Wasserstein gradient flow of a slow diffusion energy (3) on the 2d domain :math:`[0,1]^2`. We discretize the problem on a :math:`512\times 512` grid and plot the densities. :: % Parameters n = 512; % The size of the n x n grid maxIters = 200; % Maximum number of BFM iterations TOL = 1e-2; % Tolerance for BFM nt = 60; % Number of outer iterations tau = 0.005; % Time step in the JKO scheme m = 2; % m in the internal energy gamma = 0.05; % gamma in the internal energy folder = 'data'; % Output directory verbose = 1; % pPrint out logs [x,y] = meshgrid(linspace(0,1,n)); % Initial density rhoInitial = zeros(n); idx = (x-0.5).^2 + (y-0.5).^2 < 0.1^2; rhoInitial(idx) = 1; rhoInitial = rhoInitial / sum(rhoInitial(:)) * n^2; % Potential V = sin(3*pi*x) .* sin(3*pi*y); % No obstacle obstacle = zeros(n); % Plots subplot(1,2,1) contourf(x, y, rhoInitial) title('Initial density') axis square subplot(1,2,2) contourf(x, y, V); title('Potential V') axis square .. image:: example-slowdiff-init.png :width: 100% :alt: Initial density and potential V Here :math:`\rho^{(0)}` is renormalized so that :math:`\int_\Omega\rho^{(0)}(x)dx=1`. Now run the algorithm :: rhoFinal = wgfslow(rhoInitial, V, obstacle, m, gamma, ... maxIters, TOL, nt, tau, folder, verbose); Output: .. code-block:: text XXX Slow Diffusion XXX n : 512 nt : 60 tau : 0.005 m : 2 gamma: 0.05 Max Iteration : 200 Tolerance : 0.01 CPU time for FFT: 0.005817 seconds. The total mass : 1 Iter : 1 50 Dual value: 1.2282 L1 error: 0.0262 100 Dual value: 1.2280 L1 error: 0.0193 150 Dual value: 1.2279 L1 error: 0.0172 ... Iter : 60 1 Dual value: -0.2241 L1 error: 0.0024 CPU time for GF: 70.906929 seconds. Make movie :: fig = figure; movieName = 'movie.gif'; for i = 0:nt file = fopen(sprintf("%s/rho-%04d.dat", folder, i), 'r'); rho = fread(file, [n n], 'double'); imagesc(rho) axis xy square set(gca,'XTickLabel',[], 'YTickLabel',[]) frame = getframe(fig); im = frame2im(frame); [X,cmap] = rgb2ind(im, 256); % Write to the GIF file if i == 0 imwrite(X, cmap, movieName, 'gif', 'Loopcount',inf, 'DelayTime',0.03); else imwrite(X, cmap, movieName, 'gif', 'WriteMode','append', 'DelayTime',0.03); end end The file :file:`movie.gif` .. image:: slowdiff.gif :width: 100% :alt: evolution of a density :align: center Incompressible energy ===================== Usage ----- In a MATLAB session, the command :: rhoFinal = wgfinc(rhoInitial, V, obstacle, maxIters, TOL, ... tau, nt, folder, verbose); computes the JKO iterates :math:`\rho^{(n)}` from an initial density :math:`\rho^{(0)}` with the internal energy .. math:: :label: eq:inc U(\rho) = \int_\Omega u_\infty(\rho(x)) dx + \int_\Omega V(x)\,\rho(x)dx. Here :math:`u_\infty(\rho)=0` if :math:`0\le\rho\le 1` and :math:`+\infty` otherwise. The iterates :math:`\rho^{(n)}` will be stored as binary files :file:`rho-0000.dat`, :file:`rho-0001.dat`, etc in a user-specified folder. Input: * :code:`rhoInitial` : The nonnegative initial density :math:`\rho^{(0)}`. * :code:`V` : The potential function :math:`V(x)`. * :code:`obstacle` : An obstacle in :math:`\Omega`. Any positive values indicate the obstacle. The densities :math:`\rho^{(n)}` will always be :math:`0` on the obstacle. * :code:`maxIters` : The maximum number of inner loop iterations :math:`k` of the BFM in each outer iteration of the JKO scheme. * :code:`TOL` : The tolerance for the BFM. Once the error :math:`\| T_{\phi_k\#} \rho^{(n)} - \delta U^*(\phi_k) \|_{L^1}` becomes less than the tolerance, the algorithm stops the BFM and proceeds to the next outer iteration :math:`n+1`. * :code:`tau` : The time step :math:`\tau` in (1). * :code:`nt` : The number of outer iterations :math:`n` of the JKO scheme. * :code:`folder` : Path to the folder where the successive iterates :math:`\rho^{(n)}` will be saved. Don't include a trailing slash. * :code:`verbose` : :code:`1` - print out process in C++ code. :code:`0` - be silent. Output: * :code:`rhoFinal` : The final density :math:`\rho^{(nt)}` after :code:`nt` outer iterations. Example ------- We compute the Wasserstein gradient flow of an incompressible energy (4) on the 2d domain :math:`[0,1]^2`. We discretize the problem on a :math:`512\times 512` grid and plot the densities. :: n = 512; % The size of the n x n grid maxIters = 300; % Maximum number of BFM iterations TOL = 3e-2; % Tolerance for BFM nt = 180; % Number of outer iterations tau = 0.005; % Time step in the JKO scheme folder = 'data'; % Output directory verbose = 1; % Print out logs [x,y] = meshgrid(linspace(0,1,n)); % Initial density rhoInitial = zeros(n); idx = (x-0.25).^2 + (y-0.75).^2 < 0.15^2; rhoInitial(idx) = 1; % Potential V = 3 * ((x-0.7).^2 + (y-0.3).^2); % Add an obstacle obstacle = zeros(n); idx = (x-0.5).^2 + (y-0.5).^2 < 0.1^2; obstacle(idx) = 1; % Plots subplot(1,3,1) contourf(x, y, rhoInitial) title('Initial density') axis('square') subplot(1,3,2) contourf(x, y, V); title('Potential V') axis('square') subplot(1,3,3) contourf(x, y, obstacle) colormap(gca, 'copper') title('Obstacle') axis('square') .. image:: example-inc-init.png :width: 100% :alt: Initial and final densities Run the algorithm :: rhoFinal = wgfinc(rhoInitial, V, obstacle, maxIters, TOL, nt, tau, ... folder, verbose); Output: .. code-block:: text XXX Incompressible Flow XXX n : 512 nt : 180 tau : 0.005 Max Iteration : 300 Tolerance : 0.03 CPU time for FFT: 0.295334 seconds. Total mass : 0.0703964 Iter : 1 10 dual value: 0.085350 L1 error: 0.1135 20 dual value: 0.085378 L1 error: 0.0807 30 dual value: 0.085386 L1 error: 0.0628 40 dual value: 0.085390 L1 error: 0.0525 50 dual value: 0.085391 L1 error: 0.0446 56 dual value: 0.085392 L1 error: 0.0407 ... Iter : 180 7 dual value: 0.002379 L1 error: 0.0329 CPU time for GF: 616.0 seconds. Note that here there is a ceiling contraint :math:`\rho(x)\le 1` so the total mass *must* be less than 1. Make movie :: fig = figure; movieName = 'movie.gif'; for i = 0:nt file = fopen(sprintf("%s/rho-%04d.dat", folder, i), 'r'); rho = fread(file, [n n], 'double'); imagesc(rho) axis xy square set(gca,'XTickLabel',[], 'YTickLabel',[]) frame = getframe(fig); im = frame2im(frame); [X,cmap] = rgb2ind(im, 256); % Write to the GIF file if i == 0 imwrite(X, cmap, movieName, 'gif', 'Loopcount',inf, 'DelayTime',0.001); else imwrite(X, cmap, movieName, 'gif', 'WriteMode','append', 'DelayTime',0.001); end end The file :file:`movie.gif` .. image:: incompressible.gif :width: 80% :alt: evolution of a density :align: center More Videos =========== Here are other cool things you can do with the BFM. - Slow diffusion flow with :math:`m=4`, a quadratic potential and a star-shaped obstacle. (512 x 512 pixels) .. raw:: html
- Slow diffusion flow with potential :math:`V(x)=-\sin(5\pi x_1)\sin(3\pi x_2)`. Top: :math:`m=2`. Bottom: :math:`m=4`. (512 x 512 pixels) .. raw:: html
- Aggregation-diffusion: slow diffusion + a concave interaction kernel. (512 x 512 pixels) .. raw:: html
- Incompressible flow. (1024 x 1024 pixels) .. raw:: html
.. raw:: html
- Moving obstacles and moving potential. (512 x 512 pixels) .. raw:: html
References ========== [1] Matt Jacobs, Wonjun Lee, and Flavien Léger. `The back-and-forth method for Wasserstein gradient flows `_. Submitted. 2020. [2] Matt Jacobs and Flavien Léger. `A fast approach to optimal transport: the back-and-forth method `_. `Numerische Mathematik` (2020): 1-32. [3] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker–Planck equation. `SIAM journal on mathematical analysis` (1998): 29(1), 1–17.