clear all
close all
clc

load('Figure4.mat');
% each .mat file contains 
% a vector t, which contains the timepoints in seconds from the full simulation
% a matrix mov2D_psi which is a set of 2D wavefunctions, sampled at regular intervals over the full simulation. 
% The time between sampled wave functions is [max value of t]/[size of third dimension of mov2D_psi - 1].

%% Constants
%
a = 100.44*5.29e-11;           % approx scattering length for Rb87 in meters.
m = 1.5e-25;                % m = mass of Rb87 in kg
hbar = 1.05457e-34;         % Planck's constant in J*s

% need to calculate the local speed of sound from the BEC parameters and
% the atomic density.
% the local speed of sound is used to scale the velocity profile in column 3.
% a global speed of sound is used to scale the velocity profile in column 4.

% start with 3D experimental parameters
%
omegar = 2*pi*8;    % radial trap freq is 8 Hz.
omegaz = 2*pi*90;           % axial trap freq is 90 Hz.
omega = (omegar^2*omegaz)^(1/3);
%
lr = sqrt(hbar/m/omegar);  % H.O. lengths
lz = sqrt(hbar/m/omegaz);
l = (lr^2*lz)^(1/3);
%
u = (15*N*a/l)^0.4 * 0.5 * hbar * omega;  % 3D chemical potential
%
g = 4*pi*hbar^2*a/m;       % 3D  interaction prefactor
%
c0 = sqrt(u/m);             % bulk speed of sound.

% % 2D effective parameters for numerics
N2D = (15*N*a/l)^0.6/16 * sqrt(pi)/a * sqrt(2 + sqrt(2)) * sqrt(hbar*omega^3/m/omegar^4)
omegar2D = 2^(-.25) * omegar
lz2D = sqrt(pi*hbar/m/omega) * sqrt(2)*sqrt(2+sqrt(2)) * (15*N*a/l)^(-0.2)/sqrt(2*pi)  
%
g2D = g/lz2D/sqrt(2*pi);
%
G2D = g2D*N2D;             % nonlinear prefactor in 2D.

%%   Set up grids and time steps
%   All units are in mks.
%

%% space
xmax    = 120e-6;             % x grid size in meters
ymax    = xmax;
nxy      = 512;             % number of grid points
nmid  = floor(nxy/2);
v = 0:1:nxy-1;
dx = xmax/(nxy-1);
dy = ymax/(nxy-1);
[x,y] = meshgrid(v,v);
x = -xmax/2 + x*dx;
y = -ymax/2 + y*dy;
xar = -xmax/2 + v*dx;
yar = -ymax/2 + v*dy;

%%  time
tspiral = ts;                %time for spiral
dt = t(2)-t(1)              % seconds
nt = length(t);              % total number of time points in the simulation
ncycle = floor(nt/500);    % Saved 500 frames at regular intervals. 
dt_movie = ncycle*dt;       % time betweeen movie frames.

%% set up figure

[sy sx sz] = size(mov2D_psi);

clims1 = [0.2 Inf];
clims2 = clims1;

str_ar = {'(a)', '(b)', '(c)', '(d)', '(e)', '(f)', '(g)', '(h)'};

% timepoints to plot
ttt = [0,.2,.4,.6, .65,.67,1, 1.8];   %time / tspiral

% center location for zoomed in speed profile plots (column 4)
x1_ar = 1.0e-04 * [0.4    0.15   -0.2   -0.25   -0.1   0.07   -.01   -0.02];
y1_ar = 1.0e-04 * [-.01    0.25   -0.25    0.1   -0.27   -.27  -.02   -.02];

qctr = 0;

% set up plots
psize = 90;   % size of each plot. make smaller if they don't all fit.
dpsize = 2;   % distances between plots
Bshift = 55; % pads length to give room for column labels.  Add extra padding as needed.

% for the column label placement
A2 = psize - 20;   % length of colorbar
A3 = 8;   % height of color bar
Ashift = Bshift - dpsize-A3;   % vertical placement of colorbar
A1 = dpsize + (psize - A2)/2;  % horizontal placement of colorbar

%fontsizes
f1 = 14;  % fontsize for timing labels.  
f2 = 12;  % fontsize for colormap numbers.   
f3 = 16;  % fontsize for colormap labels.   

% % set up plots  - original sizing requires a large screen.
% psize = 150;   % size of each plot
% dpsize = 5;   % distances between plots
% Bshift = 55+dpsize*4; % pads length to give room for column labels (I think)
% % for the column label placement
% A2 = psize - 25;   % length of colorbar
% A3 = 10;   % height of color bar
% Ashift = Bshift -A3;   % vertical placement of colorbar
% A1 = (psize - A2 + dpsize)/2;  % horizontal placement of colorbar
% 
% %larger fontsizes
% f1 = 22;  % fontsize for timing labels.  % was 22 
% f2 = 18;  % fontsize for colormap numbers.   % was 18
% f3 = 24;  % fontsize for colormap labels.   % was 24

h3 = figure(103)
set(h3, 'position', [1        1    psize*4+dpsize psize*length(ttt)+Bshift]);
set(h3,'color','w');

%% make plots
for qq = round(ttt*ts/dt/ncycle)+1
    qctr = qctr+1;
    
    % atomic density profile at time ttt*ts
    n =  abs(mov2D_psi(:,:,qq)).^2;
    nmax = max(max(n));
    
    % set up masks to set regions nearly zero density to zero. 
    % this just makes the plots less messy.
    phmask1 = ones(size(n)).*(n/nmax > 1e-6);
    phmask2 = ones(size(n)).*(n/nmax > 1e-3);
    vmask = ones(size(n)).*(n/nmax > 5e-4).*(x.^2 + y.^2 < (1*RrTF)^2);
    vmask1 = phmask2;
    %   
    n = n.*phmask2;
    %
    ph = (angle(mov2D_psi(:,:,qq))+pi()).*phmask1;
    ph1 = ph.*phmask1;
    
    % calculate velocity profile from the phase.
    [gx, gy] = gradient(ph1,dx,dx);
    vx = gx*hbar/m.*phmask2;
    vy = gy*hbar/m.*phmask2; 
    vx1 = vx;
    vy1 = vy;
    vv = sqrt(vx1.^2 + vy1.^2);   % speed profile
    
    % mask the phase a bit more.
    ph2 = ph1.*phmask2;
    
    % find area that will be set to black in each plot.
    n(n == 0) = NaN;
    ph2(ph2 == 0) = NaN;
    vmask1(vmask1 == 0) = NaN;
          
    imAlpha=ones(size(n));
    imAlpha(isnan(n))=0;

    imAlpha1=ones(size(ph2));
    imAlpha1(isnan(ph2))=0;
    imAlpha2=ones(size(vmask1));
    imAlpha2(isnan(vmask1))=0;
    
    %% smoothing algorithm
    % method is for each point where v_x or v_y > 0.001, take the average
    % of the legit data points around that point and replace the value with
    % that data point.  note: have a copy so that I don't end up
    % accidentally making streaks.  
    
    clear cy
    clear cx
    [cy, cx] = find(abs(vy1) > 0.001);   % look at vy
    
    for cc = 1:length(cx)
        if cx(cc) == 1
            vxr = cx(cc):cx(cc)+1;
        else
            if cx(cc) == 256
                vxr = cx(cc)-1:cx(cc);
            else
                vxr = cx(cc)-1:cx(cc)+1;
            end
        end
        %
        if cy(cc) == 1
            vyr = cy(cc):cy(cc)+1;
        else
            if cy(cc) == 256
                vyr = cy(cc)-1:cy(cc);
            else
                vyr = cy(cc)-1:cy(cc)+1;
            end
        end
        
        vytest = (vy1(vyr,vxr))*1e6;   % smoothing vy
        vylegit = find(abs(vytest) ~= 0 & abs(vytest) < 1e3);
        vynew = mean(vytest(vylegit)) ;   % in um/s
        vy(cy(cc),cx(cc)) = vynew*1e-6;   % in m/s
    end
    
    
    clear cy
    clear cx
    [cy, cx] = find(abs(vx1) > 0.001);  % look at vx
    
    for cc = 1:length(cx)
        if cx(cc) == 1
            vxr = cx(cc):cx(cc)+1;
        else
            if cx(cc) == 256
                vxr = cx(cc)-1:cx(cc);
            else
                vxr = cx(cc)-1:cx(cc)+1;
            end
        end
        %
        if cy(cc) == 1
            vyr = cy(cc):cy(cc)+1;
        else
            if cy(cc) == 256
                vyr = cy(cc)-1:cy(cc);
            else
                vyr = cy(cc)-1:cy(cc)+1;
            end
        end
        
        vxtest = (vx1(vyr,vxr))*1e6;    % smoothing vx
        vxlegit = find(abs(vxtest) ~= 0 & abs(vxtest) < 1e3);
        vxnew = mean(vxtest(vxlegit));   % in um/s
        vx(cy(cc),cx(cc)) = vxnew*1e-6;   % in m/s
    end
    
    %% velocity field requires both magnitude and direction.
    v = sqrt(vx.^2 + vy.^2);        % speed profile = magnitude of the velocity.
    
    % gets rid of long arrows on quiver plot.
    vx = vx.*vmask.*(v<.3*c0);
    vy = vy.*vmask.*(v<.3*c0);   % unscaled
    
    % downsample for quiver plot.
    s = 12*nxy/256;
    %
    nn = downsample(downsample(n,s)',s)';
    xx = downsample(downsample(x,s)',s)';
    yy = downsample(downsample(y,s)',s)';
    vxx = downsample(downsample(vx,s)',s)';
    vyy = downsample(downsample(vy,s)',s)';
     
    %% make figure
    figure(103)
    % plot density profile
    axes('Units', 'pixels', 'position', [dpsize, psize*(length(ttt)-qctr)+dpsize+Bshift, psize - dpsize, psize - dpsize])
    imagesc(xar,yar, n/nmax,'AlphaData',imAlpha)
    set(gca,'color',0*[1 1 1]);
    set(gca, 'CLim', [0 1])
    axis image
    axis xy
    set(gca,'XTickLabel','')
    set(gca,'YTickLabel','')
    hold on
    if ttt(qctr) <= 1
        str = ['\bf{$t/\tau_\mathrm{s} = ' num2str(ttt(qctr)) '$}'];
    else
        str = ['{Beam off}'];
    end
    
    text(-xmax/2+dx*10, xmax/2 - dx*50,str, 'Interpreter', 'latex', 'fontsize', f1, 'color', 'white')
    
    hold off
    
    % add label / colorbar at the bottom of each column.
    if qctr == 1
        % column 1
        cax = colorbar('SouthOutside');
        set(cax, 'Units', 'pixels', 'position', [A1, Ashift, A2, A3], 'fontsize', f2);
        cax.Label.Interpreter = 'latex';
        cax.TickLabelInterpreter = 'latex';
        cax.Label.String = '{$n/n_\mathrm{max}$}';
        cax.Label.FontSize = f3;    
    end
    
    % plot phase
    axes('Units', 'pixels', 'position', [psize+dpsize, psize*(length(ttt)-qctr)+dpsize+Bshift, psize - dpsize, psize - dpsize])
    imagesc(ph2/2/pi,'AlphaData',imAlpha)
    set(gca, 'CLim', [0 1])
    axis xy
    axis image
    set(gca,'color',0*[1 1 1]);
    set(gca,'XTickLabel','')
    set(gca,'YTickLabel','')
    
    % add label / colorbar at the bottom of each column.
    if qctr == 1 
        % column 2
        cax = colorbar('SouthOutside');
        set(cax, 'Units', 'pixels', 'position', [psize+A1, Ashift, A2, A3], 'fontsize', f2);
        cax.Label.Interpreter = 'latex';
        cax.TickLabelInterpreter = 'latex';
        cax.Label.String = '{$\phi/2\pi$}';
        cax.Label.FontSize = f3;
    end
    
    % plot speed profile
    axes('Units', 'pixels', 'position', [psize*2+dpsize, psize*(length(ttt)-qctr)+dpsize+Bshift, psize - dpsize, psize - dpsize])
    imagesc(xar, yar, v./sqrt(g2D*n*N2D*sqrt(2)/m),'AlphaData',imAlpha)  
    set(gca, 'CLim', [0 3])
    axis xy
    axis image
    set(gca,'color',0*[1 1 1]);
    set(gca,'XTickLabel','')
    set(gca,'YTickLabel','')
    
    % add label / colorbar at the bottom of each column.
    if qctr == 1
        % column 3.
        cax = colorbar('SouthOutside');
        set(cax, 'Units', 'pixels', 'position', [psize*2+A1, Ashift, A2, A3], 'fontsize', f2);
        cax.Label.Interpreter = 'latex';
        cax.TickLabelInterpreter = 'latex';
        cax.Label.String = '{$v/c_\mathrm{local}$}';
        cax.Label.FontSize = f3;
    end
    
    x1 = x1_ar(qctr);
    y1 = y1_ar(qctr);

    % plot zoomed in speed profile
    axes('Units', 'pixels', 'position', [psize*3+dpsize, psize*(length(ttt)-qctr)+dpsize+Bshift, psize - dpsize, psize - dpsize]);
    imagesc(xar-x1, yar-y1, v/c0.*(v/c0<1),'AlphaData',imAlpha)
    axis xy
    axis image
    set(gca, 'CLim', [0 .5]) 
    set(gca,'color',0*[1 1 1]);
    set(gca,'XTickLabel','')
    set(gca,'YTickLabel','')

    
    
    % now add quiver plot to show the direction of flow
    scalefactor = .00003; % scales the length of the arrows.
    quiver(xx-x1, yy-y1, vxx/c0*scalefactor, vyy/c0*scalefactor, 'AutoScale', 'off');
    hax = gca; %get the axis handle
    axis xy;
    axis image %plot the quiver to see the dimensions of the plot
    imagesc(hax.XLim,hax.YLim,v/c0.*(v/c0<1), 'AlphaData',imAlpha); %plot the image within the axis limits
    hold on; % enable plotting overwrite
    quiver(xx-x1, yy-y1, vxx/c0*scalefactor, vyy/c0*scalefactor, 'color', 'white', 'linewidth', 1, 'AutoScale', 'off')
    axis xy
    axis image
    hold off
    
    xmax1 = 2.5e-5;
    xlim([-xmax1 xmax1])
    ylim([-xmax1 xmax1])
    set(gca, 'CLim', [0 .5]) 
    set(gca,'color',0*[1 1 1]);
    set(gca,'XTickLabel','')
    set(gca,'YTickLabel','')
    
    
    % add label / colorbar at the bottom of each column.
    if qctr == 1
        % column 4
        cax = colorbar('SouthOutside');
        set(cax, 'Units', 'pixels', 'position', [psize*3+A1, Ashift, A2, A3], 'fontsize', f2);
        cax.Label.Interpreter = 'latex';
        cax.TickLabelInterpreter = 'latex';
        cax.Label.String = '{$v/c_0$}';
        cax.Label.FontSize = f3;
    end
    
end
%%
% print(h3,'-dpng', ['Figure_SimsImages_v0'])
% print(h3,'-dsvg', ['Figure_SimsImages_v0'])
% print(h3,'-dpdf', ['Figure_SimsImages_v0'])