%%%  Set up model  %%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Begin by defining constants and model structure
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%  Model geometry
Ly = 5;                  %  North / South Boundaries (non-dim)
%Lx = 120;                %  Wavelength (degrees)
dy = 0.2;               %  Grid spacing (non-dim)

%  Define constants
g = 9.8;                 %  Gravity
OMEGA = 2*pi/(3600*24);  %  Earth's angular velocity
RADUS = 6375000;         %  Earth's radius
beta = 2*OMEGA/RADUS;    %  d/dy of the Coriolis Parameter
ij = sqrt(-1);

%  Set up Atmospheric Constants
He = 200;                %  Equivalent Depth
ru = 1/(60*60*24*2);     %  Momentum damping (s^-1)
rh = 1/(60*60*24*2);     %  Thermal damping (geopotential)
Kq = 7e-3;             %  Heating (per degree C); Hirst, 1986
%Kq = 1.2e-2;             %  Heating (per degree C); Xie, 1997

%  Set up Ocean Contants
%wesh = 15;               %  WES feedback W/m^-^2 per m/s
rho = 1e3;               %  Density of ocean water
co = 4.2e3;              %  Specific Heat
Ho = 50;                 %  Depth of mixed layer (in meters)
rt = 1/(3600*24*300);     %  Linear temperature damping
gamma = 5e4;             %  Harmonic Damping
To = 1;                   %  Non-dim. temperature (this doesn't matter)

%%%  Now, scale constants so we have non-dimensional values  %%%
c = sqrt(g*He);          %  Gravity wave phase speed
ae = sqrt(c/(beta));   %  Equatorial Rossby Radius
ru = ru / sqrt(c*beta);  
rh = rh / sqrt(c*beta);
rt = rt / (sqrt(c*beta)); %(rho*co*Ho*sqrt(c*beta));
wesh = wesh * c/(sqrt(c*beta)*rho*co*Ho*To);
Kq = Kq * To * ae/c^3; %/(c^2*sqrt(c*beta));
gamma = gamma/(c*ae);    %  Note:  DO NOT divide by rho*co*Ho - this is
                         %  already done in defining gamma originally.
k = 2*pi/((RADUS*pi*Lx/180)/ae);  %  Non-dim wavenumber

%%%  Set up grid geometry  %%%
y = [-Ly:dy:Ly];
n = length(y);
n4 = n*4;
n3 = n*3;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Give meridional structure to Kq and wes feedbacks
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%  Kq structure  %%%

%  Symmetric heating
lat0p5 = 20;    %  Latitude at which Kq = e^(-1)
lat0p5 = (lat0p5*RADUS*pi/180)/ae;  %  Non-dim.
fx = exp(-((y-0)/lat0p5).^4);
%fx = ones(size(y));

%  Asymmetric heating
%  lat0p5 = 5;    %  Latitude at which Kq = e^(-1)
%  lat0p5 = (lat0p5*RADUS*pi/180)/ae;  %  Non-dim.
%  fx = exp(-((y-0.7)/lat0p5).^2);

Kq = Kq * fx;

%%%  WES structure %%%
lat0p5 = 35;
lat0p5 = (lat0p5*RADUS*pi/180)/ae;  %  Non-dim.
fx = exp(-((y/lat0p5)).^4);

wes = wesh*fx;


%%%  Details to stabilize numerical modes  %%%
%  Set up sponge layer for damping - ramp to 1e1 times the damping at the
%  edges.  

rstruct = exp(-1*(y/Ly).^20);   %  This applies a sponge layer to
                                %  approximately -Ly to -4*Ly/5, and from
                                %  4*Ly/5 to Ly. 
rstruct = 10 - 9 * (rstruct - min(rstruct))/(1 - min(rstruct));

%  Now, set latitudinal dependencies for damping
ru = ru*rstruct;
rh = rh*rstruct;
rt = rt*rstruct;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%  Define the linear system matrix
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%  Note:  each set of four rows (a block) contains the equations for one
%  y-value.  Elements in each column will operate on the corresponding
%  element of the state vector.  Thus, the matrix is set up as follows:
%
%  M(4*(i-1)+[1:4],:) = | -ru   beta*y/2   -ij*k   0 |
%                       | -beta*y/2  -rv   -d/dy   0 |
%                       | -ij*k  -d/dy     -rh    Kq |
%                       |  wes     0        0    -rt |
%
%  Note that y-derivatives are discretized, so they will involve terms in
%  the previous, or following y-block.
%
%  The state vector, then, is:  [u1, v1, h1, T1, u2, v2, h2, T2, ... ,
%  un, vn, hn, Tn]'.  It has n blocks, and 4*n elements.
%
%  A separate Mu is set up for the uncoupled case.  Note that this is
%  really easy to do, by eliminating the last row and column from each
%  block (thus, Mu = M(sort([1:4:end 2:4:end 3:4:end]), sort([1:4:end
%  2:4:end 3:4:end]);.
%

M = zeros(n4, n4);

%  Iterate through, except for the boundaries

for i = 2:(n-1);
  
  %%%  Zonal momentum equation
  ii = 4*(i-1)+1;
  
  %  Damping (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -1*ru(i);
  
  %  Coriolis (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) + y(i);
  
  %  PGF (h):
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) + -1*ij*k;
  
  %%%  Meridional momentum equation
  ii = 4*(i-1)+2;
  
  %  Coriolis (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -1*y(i);
  
  %  Damping (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) - ru(i);
  
  %  PGF
  jj = 4*(i-1)+3;  %  jj index for h in the current block
  jjn = jj + 4;    %  jj index for h to the north
  jjs = jj - 4;    %  jj index for h to the south
  M(ii, jjn) = M(ii, jjn) - 1/(2*dy);  %  Centered Difference
  M(ii, jjs) = M(ii, jjs) + 1/(2*dy);  %  for -d/dy
  
  %%%  Continuity equation
  ii = 4*(i-1)+3;
  
  %  Convergence (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - ij*k;
  
  %  Convergence (v):
  jj = 4*(i-1)+2;  %  jj index for h in the current block
  jjn = jj + 4;    %  jj index for h to the north
  jjs = jj - 4;    %  jj index for h to the south
  M(ii, jjn) = M(ii, jjn) - 1/(2*dy);  %  Centered Difference
  M(ii, jjs) = M(ii, jjs) + 1/(2*dy);  %  for -d/dy
  
  %  Damping
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) - rh(i);
  
  %  Heating
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) + Kq(i);
  
  %%% Temperature equation
  ii = 4*(i-1)+4;
  
  %  WES feedback
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - wes(i)/2;  %  Note:  the atmosphere is a two
                                     %  layer model, so u = utop-ubot.
                                     %  Thus, we use -wes/2.
  
  %  Linear damping
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) - rt(i);
  
  %  Harmonic damping
  jj = 4*(i-1)+4;  %  jj index for h in the current block
  jjn = jj + 4;    %  jj index for h to the north
  jjs = jj - 4;    %  jj index for h to the south
  M(ii, jjn) = M(ii, jjn) + gamma/(dy^2);  %  Centered Difference
  M(ii, jjs) = M(ii, jjs) + gamma/(dy^2);  %  for d^2/(dy^2)
  M(ii, jj) = M(ii, jj) - 2*gamma/(dy^2);  %
  
end  %  End interior

%%%  Now, do boundary blocks:

i = 1;  %  Southern boundary

  %%%  Zonal momentum equation
  ii = 4*(i-1)+1;
  
  %  Damping (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -1*ru(i);
  
  %  Coriolis (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) + y(i);
  
  %  PGF (h):
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) + -1*ij*k;
  
  %%%  Meridional momentum equation
  ii = 4*(i-1)+2;
  
  %  Coriolis (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -y(i);
  
  %  Damping (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) - ru(i);
  
  %  PGF
  jj = 4*(i-1)+3;  %  jj index for h in the current block
  jjn = jj + 4;    %  jj index for h to the north
  M(ii, jjn) = M(ii, jjn) - 1/dy;      %  Forward Difference
  M(ii, jj ) = M(ii, jj ) + 1/dy;      %  for -d/dy
  
  %%%  Continuity equation
  ii = 4*(i-1)+3;
  
  %  Convergence (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - ij*k;
  
  %  Convergence (v):
  jj = 4*(i-1)+2;  %  jj index for h in the current block
  jjn = jj + 4;    %  jj index for h to the north
  M(ii, jjn) = M(ii, jjn) - 1/dy;      %  Forward Difference
  M(ii, jj ) = M(ii, jj ) + 1/dy;      %  for -d/dy
  
  %  Damping
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) - rh(i);
  
  %  Heating
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) + Kq(i);
  
  %%% Temperature equation
  ii = 4*(i-1)+4;
  
  %  WES feedback
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - wes(i)/2;
  
  %  Linear damping
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) - rt(i);
  

i = n;  %  Northern boundary

  %%%  Zonal momentum equation
  ii = 4*(i-1)+1;
  
  %  Damping (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -1*ru(i);
  
  %  Coriolis (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) + y(i);
  
  %  PGF (h):
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) + -1*ij*k;
  
  %%%  Meridional momentum equation
  ii = 4*(i-1)+2;
  
  %  Coriolis (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) + -y(i);
  
  %  Damping (v):
  jj = 4*(i-1)+2;
  M(ii, jj) = M(ii, jj) - ru(i);
  
  %  PGF
  jj = 4*(i-1)+3;  %  jj index for h in the current block
  jjs = jj - 4;    %  jj index for h to the south
  M(ii, jj ) = M(ii, jj ) - 1/dy;      %  Backward Difference
  M(ii, jjs) = M(ii, jjs) + 1/dy;      %  for -d/dy
  
  %%%  Continuity equation
  ii = 4*(i-1)+3;
  
  %  Convergence (u):
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - ij*k;
  
  %  Convergence (v):
  jj = 4*(i-1)+2;  %  jj index for v in the current block
  jjs = jj - 4;    %  jj index for v to the south
  M(ii, jj ) = M(ii, jj ) - 1/dy;      %  Backward Difference
  M(ii, jjs) = M(ii, jjs) + 1/dy;      %  for -d/dy
  
  %  Damping
  jj = 4*(i-1)+3;
  M(ii, jj) = M(ii, jj) - rh(i);
  
  %  Heating
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) + Kq(i);
  
  %%% Temperature equation
  ii = 4*(i-1)+4;
  
  %  WES feedback
  jj = 4*(i-1)+1;
  M(ii, jj) = M(ii, jj) - wes(i)/2;
  
  %  Linear damping
  jj = 4*(i-1)+4;
  M(ii, jj) = M(ii, jj) - rt(i);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Done!!!
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%