% MATLAB file:  ocean_basin_mode_staggered
% based on forced_equatorial_mode2.m from Chapter 11 of Holton
% Revised 12/4/08
% Finite difference model for linear equatorial shallow water model.
% Given a transient thermal forcing computes wave response.
% Method is finite differences in time and spa.
% Domain is equatorial beta-plane channel periodic in x.
clear all
%close all
Lx = 4000;  Ly = 3000;             % channel dimensions in km
Nx = 81; Ny = 61;                   % number of grid points in each direction
Nxl = Nx-1; Nyl = Ny-1;
xx=linspace(-Lx/2,Lx/2,Nx);         % Nx gridpoints in x   
yy=linspace( -Ly/2,Ly/2,Ny);        % Ny gridpoints in y
%convert distances to meters for calculations
[x,y] = meshgrid(xx*1000,yy*1000);  % Sets matrix for grid system    
dx = Lx/(Nx-1)*1000;                % grid distance in x  
dy = Ly/(Ny-1)*1000;                % grid distance in y
%define constants and parameters
y00 = 0;                            % y coordinate of heating maximum (m)
rad = 6.37e6;                       % radius of earth
omega = 7.292e-5;                   % earth angular velocity
lat0 = input('specify latitude at the center of the beta plane lat0 =   ');
f0=2*omega*sin(lat0*pi/180);        % coriolis parameter in the center
beta = 2*omega/rad*cos(lat0*pi/180);                 % beta effect
A = 2.e-3;                          % forcing amplitude
alpha = 1.e-7;                      % damping rate
CB2 = 2.^2;                         % shallow wave speed squared
Rx = 2.e5;  Ry = 2.e5;              % forcing scales in x and y

%set initial fields to zero
u0 = zeros(size(x)); u = u0;  un = u0;
v0 = u0; v = v0;  vn = v0;
phi0 = u0; phi = phi0;  phin = phi0;
div = u0;
eps = u0;
dphix = u0; dphiy = u0;
%%  time integration parameters:
nsv = 0;                            % counter for output figures
t = 0;
time_end = input('specify ending time of integration in days time_end =   ');
time_endsec = time_end*3600*24   ;     % ending time in seconds
dt = 1800;                          % time increment in seconds
tau = 15*24*3600;                    % forcing period
dt2 = 2*dt;  
adt = 1-alpha*dt;
adt2 = 1-alpha*dt2;
tsave = 3600*24*1/4;     % interval for saving to output --and for forward stepping
ntime = fix(time_endsec/dt);        % total number of time steps
nsave = fix(time_endsec/tsave);     % number of times saved to output

% take forward step as first step
% Begin time looping using leapfrog method
        u(:,end)=0;
        v(end,:)=0;
for s=1:ntime
    % define the transient thermal forcing function
    heat=zeros(size(x)); 
    if t<tau
        heat = A*exp(-((x+Lx*1000/8).^2/Rx^2+(y-y00).^2/Ry^2))...
            *sin(pi*t/tau);          % thermal forcing
    end
    [dphix, dphiy] = grad_wall_staggered(phi,dx,dy); % gradient of geopotential
    t = t + dt;
    tt = t/tsave;
    if rem(tt,1) == 0;          % forward step to damp computational mode
        tmpv=([zeros(1,Nx-1); v(1:end-1,1:end-1)]+v(1:end,1:end-1)+[zeros(1,Nx-1); v(1:end-1,2:end)]+v(1:end,2:end))/4;
        tmpy=(y(1:Ny,1:Nx-1)+y(1:Ny,2:Nx))/2;
        un(1:Ny,1:Nx-1) = adt*u(1:end,1:end-1)+dt*((f0+beta*tmpy).*tmpv(1:Ny,1:Nx-1)-dphix(1:Ny,1:Nx-1));
        tmpu=([zeros(Ny-1,1) u(1:end-1,1:end-1)]+u(1:end-1,1:end)+[zeros(Ny-1,1) u(2:end,1:end-1)]+u(2:end,1:end))/4;
        tmpy=(y(1:Ny-1,1:Nx)+y(2:Ny,1:Nx))/2;
        vn(1:Ny-1,1:Nx) = adt*v(1:Ny-1,1:end)+dt*(-(f0+beta*tmpy).*tmpu(1:Ny-1,1:Nx)-dphiy(1:Ny-1,1:Nx));
        div = divergence(u,v,dx,dy);
        phin = adt*phi + dt*(-CB2.*div + heat);
    else                        % leapfrog step
        tmpv=([zeros(1,Nx-1); v(1:end-1,1:end-1)]+v(1:end,1:end-1)+[zeros(1,Nx-1); v(1:end-1,2:end)]+v(1:end,2:end))/4;
        tmpy=(y(1:Ny,1:Nx-1)+y(1:Ny,2:Nx))/2;
        un(1:Ny,1:Nx-1) = adt2*u0(1:end,1:end-1)+dt2*((f0+beta*tmpy).*tmpv(1:Ny,1:Nx-1)-dphix(1:Ny,1:Nx-1));
        tmpu=([zeros(Ny-1,1) u(1:end-1,1:end-1)]+u(1:end-1,1:end)+[zeros(Ny-1,1) u(2:end,1:end-1)]+u(2:end,1:end))/4;
        tmpy=(y(1:Ny-1,1:Nx)+y(2:Ny,1:Nx))/2;
        vn(1:Ny-1,1:Nx) = adt2*v0(1:Ny-1,1:end)+dt2*(-(f0+beta*tmpy).*tmpu(1:Ny-1,1:Nx)-dphiy(1:Ny-1,1:Nx));
        div = divergence(u,v,dx,dy);
        phin = adt2*phi0 + dt2*(-CB2.*div+heat);  
    end
    %update all fields
    u0 = u;  v0 = v;  phi0 = phi;
    u = un;  v = vn;  phi = phin;
    if rem(tt,1) == 0;
        
        time = tt*tsave/3600;   % show output  every tsave/3600 hours
        max(max(u));
        
        figure(1)
        clf
        subplot(2,1,1)
        pcolor(x/1.e6,y/1.e6,phi/9.8)  % x, y, converted to km for graph 
        caxis([-5 5])
        shading interp
        colorbar('h')
        hold on
        V = 0.03e-3;
        contour(x/1.e6,y/1.e6,heat,[V V],'k') 
        ylabel('y  (km \times 10^3)')
        title( ['perturbaton height (color), forcing (contour) Time=' num2str(time/24) 'day'],'Fontsize',12)
        hold off
        subplot(2,1,2)
        % Reduces number of grid points for plotting velocity vectors
        for i = 1:Nyl/4
            for j = 1:Nxl/4
                xd(i,j)=x(4*i,4*j);
                yd(i,j)=y(4*i,4*j);
                uxd(i,j)=u(4*i,4*j);
                vyd(i,j)=v(4*i,4*j);
            end
        end   
        quiver(xd/1.e6,yd/1.e6,uxd,vyd)
        axis([-Lx/2000 Lx/2000 -Ly/2000  Ly/2000  ])
        title('scaled velocity field','Fontsize',12)
        xlabel('x  (km \times 10^3)'), ylabel('y  (km) \times 10^3')
        hold off
    end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



