Edit notation mu->rhoInitial

example_incompressible.m

@@ -1,75 +1,56 @@
 %% Example of the Back-And-Forth method
+% 
+% For more explanation see also the documentation here:
+% https://wasserstein-gradient-flows.netlify.app/
+% 
+
+
+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
 
-% Parameters
-n = 512;         % The size of the grid
-maxIters = 300;  % Maximum iterations for the Back-and-Forth method
-TOL = 3e-2;      % Tolerance for the Back-and-Forth method
-nt  = 200;        % The number of outer iterations
-tau = 0.005;     % Time step for JKO scheme
-verbose  = 1;    % print out logs 
-folder = 'data'; % Set the filename
 
 [x,y] = meshgrid(linspace(0,1,n));
 
-% Define initial density
-mu = zeros(n);
-idx = (x-0.25).^2 + (y-0.25).^2 < 0.15^2;
-mu(idx) = 1;
+% Initial density
+rhoInitial = zeros(n);
+idx = (x-0.25).^2 + (y-0.75).^2 < 0.15^2;
+rhoInitial(idx) = 1;
 
-% Define a quadratic potential
-V = 3 * ((x-0.7).^2 + (y-0.7).^2);
+% Potential
+V = 3 * ((x-0.7).^2 + (y-0.3).^2);
 
-% Define an obstacle
+% Add an obstacle
 obstacle = zeros(n);
 idx = (x-0.5).^2 + (y-0.5).^2 < 0.1^2;
 obstacle(idx) = 1;
 
-
-%% Run BFM!
-
-result = wgfinc(mu, V, obstacle, maxIters, TOL, nt, tau, folder, verbose);
-
-
-%% Plot
-
-subplot(1,2,1)
-imagesc(mu);
-title("initial density")
-% colormap bone
-axis square
-
-hold on
-subplot(1,2,2)
-imagesc(result);
-title("final density")
-% colormap bone
-axis square
-hold off
-
+% Plots
 subplot(1,3,1)
-contourf(x,y,mu);
-title("initial density")
-% colormap bone
-axis square
+contourf(x, y, rhoInitial)
+title('Initial density')
+axis('square')
 
-hold on
 subplot(1,3,2)
-contourf(x,y,V);
-title("V")
-% colormap bone
-axis square
+contourf(x, y, V);
+title('Potential V')
+axis('square')
 
-hold on
 subplot(1,3,3)
-contourf(x,y,obstacle);
-title("Obstacle")
+contourf(x, y, obstacle)
 colormap(gca, 'copper')
-axis square
+title('Obstacle')
+axis('square')
 
-hold off
+%% Run BFM!
+rhoFinal = wgfinc(rhoInitial, V, obstacle, maxIters, TOL, nt, tau, folder, verbose);
 
-saveas(gcf, "./figures/initial-final.png");
 
+%% Make movie
 
 fig = figure;
 movieName = 'movie.gif';
@@ -78,7 +59,7 @@
     file = fopen(sprintf("%s/rho-%04d.dat", folder, i), 'r');
     rho = fread(file, [n n], 'double');
     imagesc(rho)
-    axis square
+    axis xy square
     set(gca,'XTickLabel',[], 'YTickLabel',[])
 
     frame = getframe(fig);
@@ -87,8 +68,8 @@
 
     % Write to the GIF file
     if i == 0
-        imwrite(X, cmap, movieName, 'gif', 'Loopcount',inf, 'DelayTime',0.03);
+        imwrite(X, cmap, movieName, 'gif', 'Loopcount',inf, 'DelayTime',0.001);
     else
-        imwrite(X, cmap, movieName, 'gif', 'WriteMode','append', 'DelayTime',0.03);
+        imwrite(X, cmap, movieName, 'gif', 'WriteMode','append', 'DelayTime',0.001);
     end
 end