% McDermott
% 1-15-2007
% driven_cavity_div_RK2.m
%
% Simple staggered primitive variable treatment for the driven
% cavity problem using an RK2 time integrator.

close all
clear all

Lx = 1; % cavity dimension in x
Ly = 0.25; % cavity dimension in y
nx = 64; % number of pcells in x
ny = 16; % number of pcells in y
dx = Lx/nx; % uniform grid spacing in x
dy = Ly/ny; % uniform grid spacing in y
twodx = 2*dx;
twody = 2*dy;
dxdx = dx^2;
dydy = dy^2;
U0 = 1; % lid velocity
nu = 1/1000; % kinematic viscosity (1/Re)
Fo = 0.5; % Fourier number
CFL = 0.5; % Courant number
dt_dif = Fo/(nu*(1/dxdx+1/dydy)) % time step based on Fo
dt_adv = CFL*min(dx,dy)/U0 % time step based on CFL
dt = min([dt_dif,dt_adv]) % time step used in Forward Euler integrator
T = 1000*dt; % total simulation time

onethird = 1/3;
eightthirds = 8/3;

% STAGGERED GRID ARRANGEMENT:
%
% Let this represent the bottom left pcell control volume.
% Velocities are stored on their respective faces.  Normal stresses
% and normal advective fluxes are stored at the pcell center.
% Off-diagonal stresses and advective fluxes are stored at vertices
% marked by the Xs.  The bottom left most element is prescribed the
% indices (1,1).  Because of this, some of the differencing
% below looks confusing, but this seems somewhat unavoidable.
%
%       uv(1,2)                                  uv(2,2)
%       tau12(1,2)          ^ v(1,2)             tau12(2,2)
%       X-------------------|--------------------X
%       |                                        |
%       |                                        |
%       |                                        |
%       |                                        |
%       |                                        |
%       |                                        |
%       |                                        |
%       |                                        |
%      ---> u(1,1)          O                   ---> u(2,1)
%       |                  phi(1,1)              |
%       |                  uu(1,1)               |
%       |                  vv(1,1)               |
%       |                  tau11(1,1)            |
%       |                  tau22(1,1)            |
%       |                                        |
%       |                                        |
%       |                   ^ v(1,1)             |
%       X-------------------|--------------------X
%       uv(1,1)                                  uv(2,1)
%       tau12(1,1)                               tau12(2,1)
%

% initialization
u = zeros(nx+1,ny);
v = zeros(nx,ny+1);
u_star = zeros(nx+1,ny);
v_star = zeros(nx,ny+1);
uu = zeros(nx,ny);
vv = zeros(nx,ny);
tau11 = zeros(nx,ny);
tau22 = zeros(nx,ny);
uv = zeros(nx+1,ny+1);
tau12 = zeros(nx+1,ny+1);
b_vec = zeros(nx*ny,1);

% grid
x = linspace(dx/2,Lx-dx/2,nx);
y = linspace(dy/2,Ly-dy/2,ny);
for j = 1:ny
    X(:,j) = x;
end
for i = 1:nx
    Y(i,:) = y;
end

% Build A
A = build_sparse_matrix_neumann(nx,ny,dx,dy);

for t = dt:dt:T
    
    % STAGE 1
    u0 = u;
    v0 = v;

    % compute normal stresses
    for j = 1:ny
        for i = 1:nx     
            uu(i,j) = ( 0.5*( u(i+1,j)+u(i,j) ) )^2;
            vv(i,j) = ( 0.5*( v(i,j+1)+v(i,j) ) )^2;
            tau11(i,j) = -2*nu*( u(i+1,j)-u(i,j) )/dx;
            tau22(i,j) = -2*nu*( v(i,j+1)-v(i,j) )/dy;
        end
    end
    
    % compute interior off-diagonal stresses
    for j = 2:ny
        for i = 2:nx
            uv(i,j) = 0.25*( u(i,j)+u(i,j-1) )*( v(i,j)+v(i-1,j) );
            tau12(i,j) = -nu*( (u(i,j)-u(i,j-1))/dy + (v(i,j)-v(i-1,j))/dx );
        end
    end
    
    % bottom wall stress
    j = 1;
    for i = 2:nx
        dudy_wall = ( 0 + 3*u(i,j) - onethird*u(i,j+1) )/dy;
        tau12(i,j) = -nu*dudy_wall;
    end
    % top wall stress
    j = ny+1;
    for i = 2:nx
        dudy_wall = ( eightthirds*U0 - 3*u(i,ny) + onethird*u(i,ny-1) )/dy;
        tau12(i,j) = -nu*dudy_wall;
    end
    % left wall stress
    i = 1;
    for j = 2:ny
        dvdx_wall = ( 0 + 3*v(i,j) - onethird*v(i+1,j) )/dx;
        tau12(i,j) = -nu*dvdx_wall;
    end
    % right wall stress
    i = nx+1;
    for j = 2:ny
        dvdx_wall = ( 0 - 3*v(nx,j) + onethird*v(nx-1,j) )/dx;
        tau12(i,j) = -nu*dvdx_wall;
    end

    % ucell momentum
    for j = 1:ny
        for i = 2:nx-1
            Fu_adv = -( (uu(i,j)-uu(i-1,j))/dx + (uv(i,j+1)-uv(i,j))/dy );
            Fu_dif = -( (tau11(i,j)-tau11(i-1,j))/dx + (tau12(i,j+1)-tau12(i,j))/dy );
            u_star(i,j) = u(i,j) + dt*( Fu_adv + Fu_dif );
        end
    end
    
    % vcell momentum
    for j = 2:ny-1
        for i = 1:nx
            Fv_adv = -( (uv(i+1,j)-uv(i,j))/dx + (vv(i,j)-vv(i,j-1))/dy );
            Fv_dif = -( (tau12(i+1,j)-tau12(i,j))/dx + (tau22(i,j)-tau22(i,j-1))/dy );
            v_star(i,j) = v(i,j) + dt*( Fv_adv + Fv_dif );
        end
    end

    % build source
    for j = 1:ny
        for i = 1:nx
            b(i,j) = ( u_star(i+1,j)-u_star(i,j) )/dx + ( v_star(i,j+1)-v_star(i,j) )/dy;
        end
    end
    b = b/dt;

    % map b to source vector
    for j = 1:ny
        for i = 1:nx
            p = (j-1)*nx+i;
            b_vec(p) = b(i,j);
        end
    end
    % subtract mean for discrete compatibility condition
    b_vec = b_vec - mean(b_vec);

    % solve for phi vector
    phi_vec = A\b_vec;
    % subtract mean for arbitrary solution
    phi_vec = phi_vec - mean(phi_vec);
    phimin = min(phi_vec);
    phimax = max(phi_vec);

    % map phi_vec to grid
    for j = 1:ny
        for i = 1:nx
            p = (j-1)*nx+i;
            phi(i,j) = phi_vec(p);
        end
    end

    % project velocities
    for j = 1:ny
        for i = 2:nx
            u(i,j) = u_star(i,j) - dt*( phi(i,j)-phi(i-1,j) )/dx;
        end
    end
    for j = 2:ny
        for i = 1:nx
            v(i,j) = v_star(i,j) - dt*( phi(i,j)-phi(i,j-1) )/dy;
        end
    end

    % check divergence
    for j = 1:ny
        for i = 1:nx
            b(i,j) = ( u(i+1,j)-u(i,j) )/dx + ( v(i,j+1)-v(i,j) )/dy;
        end
    end
    display(['max stage 1 velocity divergence = ',num2str( max(max(abs(b))) )])
    
    % STAGE 2
    
    % compute normal stresses
    for j = 1:ny
        for i = 1:nx     
            uu(i,j) = ( 0.5*( u(i+1,j)+u(i,j) ) )^2;
            vv(i,j) = ( 0.5*( v(i,j+1)+v(i,j) ) )^2;
            tau11(i,j) = -2*nu*( u(i+1,j)-u(i,j) )/dx;
            tau22(i,j) = -2*nu*( v(i,j+1)-v(i,j) )/dy;
        end
    end
    
    % compute interior off-diagonal stresses
    for j = 2:ny
        for i = 2:nx
            uv(i,j) = 0.25*( u(i,j)+u(i,j-1) )*( v(i,j)+v(i-1,j) );
            tau12(i,j) = -nu*( (u(i,j)-u(i,j-1))/dy + (v(i,j)-v(i-1,j))/dx );
        end
    end
    
    % bottom wall stress
    j = 1;
    for i = 2:nx
        dudy_wall = ( 0 + 3*u(i,j) - onethird*u(i,j+1) )/dy;
        tau12(i,j) = -nu*dudy_wall;
    end
    % top wall stress
    j = ny+1;
    for i = 2:nx
        dudy_wall = ( eightthirds*U0 - 3*u(i,ny) + onethird*u(i,ny-1) )/dy;
        tau12(i,j) = -nu*dudy_wall;
    end
    % left wall stress
    i = 1;
    for j = 2:ny
        dvdx_wall = ( 0 + 3*v(i,j) - onethird*v(i+1,j) )/dx;
        tau12(i,j) = -nu*dvdx_wall;
    end
    % right wall stress
    i = nx+1;
    for j = 2:ny
        dvdx_wall = ( 0 - 3*v(nx,j) + onethird*v(nx-1,j) )/dx;
        tau12(i,j) = -nu*dvdx_wall;
    end
    
    % ucell momentum
    for j = 1:ny
        for i = 2:nx-1
            Fu_adv = -( (uu(i,j)-uu(i-1,j))/dx + (uv(i,j+1)-uv(i,j))/dy );
            Fu_dif = -( (tau11(i,j)-tau11(i-1,j))/dx + (tau12(i,j+1)-tau12(i,j))/dy );
            u_star(i,j) = 0.5*( u0(i,j) + (u(i,j) + dt*( Fu_adv + Fu_dif )) );
        end
    end
    
    % vcell momentum
    for j = 2:ny-1
        for i = 1:nx
            Fv_adv = -( (uv(i+1,j)-uv(i,j))/dx + (vv(i,j)-vv(i,j-1))/dy );
            Fv_dif = -( (tau12(i+1,j)-tau12(i,j))/dx + (tau22(i,j)-tau22(i,j-1))/dy );
            v_star(i,j) = 0.5*( v0(i,j) + (v(i,j) + dt*( Fv_adv + Fv_dif )) );
        end
    end

    % build source
    for j = 1:ny
        for i = 1:nx
            b(i,j) = ( u_star(i+1,j)-u_star(i,j) )/dx + ( v_star(i,j+1)-v_star(i,j) )/dy;
        end
    end
    b = b/dt;

    % map b to source vector
    for j = 1:ny
        for i = 1:nx
            p = (j-1)*nx+i;
            b_vec(p) = b(i,j);
        end
    end
    % subtract mean for discrete compatibility condition
    b_vec = b_vec - mean(b_vec);

    % solve for phi vector
    phi_vec = A\b_vec;
    % subtract mean for arbitrary solution
    phi_vec = phi_vec - mean(phi_vec);
    phimin = min(phi_vec);
    phimax = max(phi_vec);

    % map phi_vec to grid
    for j = 1:ny
        for i = 1:nx
            p = (j-1)*nx+i;
            phi(i,j) = phi_vec(p);
        end
    end

    % project velocities
    for j = 1:ny
        for i = 2:nx
            u(i,j) = u_star(i,j) - dt*( phi(i,j)-phi(i-1,j) )/dx;
        end
    end
    for j = 2:ny
        for i = 1:nx
            v(i,j) = v_star(i,j) - dt*( phi(i,j)-phi(i,j-1) )/dy;
        end
    end

    % check divergence
    for j = 1:ny
        for i = 1:nx
            b(i,j) = ( u(i+1,j)-u(i,j) )/dx + ( v(i,j+1)-v(i,j) )/dy;
        end
    end
    display(['max stage 1 velocity divergence = ',num2str( max(max(abs(b))) )])
    
    % interpolate velocities to cell centers
    for j = 1:ny
        for i = 1:nx
            up(i,j) = 0.5*( u(i+1,j)+u(i,j) );
            vp(i,j) = 0.5*( v(i,j+1)+v(i,j) );
        end
    end
    
    SUBPLOT(2,1,1)
    velmag = sqrt( up.*up + vp.*vp );
    V = linspace(0,1,20);
    %contourf(X,Y,velmag,V); hold on
    %colorbar
    quiver(X,Y,up,vp); hold off
    axis([0 Lx 0 Ly])
    set(gca,'PlotBoxAspectRatio',[Lx Ly 1])
    
    SUBPLOT(2,1,2)
    V = linspace(phimin,phimax,20);
    contourf(X,Y,phi,V)
    colorbar
    set(gca,'PlotBoxAspectRatio',[Lx Ly 1])
    
    pause(0.001)

end