%  geoadjPE; adjustment under gravity and rotation in 1-D
%  A 1-D numerical ocean model of a single
%  shallow layer of rotating, incompressible fluid.
%  Useful for showing the numerical version of geostrophic
%  adjustment. Uses an A-grid (velocity and height at the
%  same grid points) and upwind differencing on advection
%  terms to preserve a shock. Time-stepping is by the  
%  leap-frog/trapezoidal time method.  This may be about the 
%  simplest PE model that works reasonably well for non-linear
%  cases, i.e., you can recognize the major features of the 
%  solution, including a fairly reasonable shock (or bore) in
%  the case of low latitude adjustment.  For zero rotation, see 
%  instead the script bore.m, which differs mainly in the form
%  of non-dimensionalization. 
%
%  Jim Price,  January, 2001.

clear
close all
format compact

set(0,'DefaultLineLineWidth',1.5)
set(0,'DefaultTextFontSize',12)
set(0,'DefaultAxesLineWidth',1.4)
set(0,'DefaultAxesFontSize',12)

%  display the text string str2
str2(1) =  {'geoadjPE:' };
str2(2) =  {'  '};
str2(3) =  {'Geostrophic adjustment in one dimension.  '};
str2(4) =  {'Solved by a primitive equation numerical '};
str2(5) =  {'model representing one (shallow) layer.'};
str2(6) =  {'The initial thickness represents a double front'}; 
str2(7) =  {'having zero initial current.  Parameters'}; 
str2(8) =  {'to enter are the latitude (2 < lat < 90), and '}; 
str2(9) =  {'the time of integration in inertial periods. '};
str2(10) = {'This is public domain for educational purposes.'};
str2(11) = {' '};
str2(12) = {'Hit any key to continue.'};
str2(13) = {''};
str2(14) = {'Jim Price, January, 2001.'};
hf3 = figure(1);
clf
set(hf3,'Position',[50 200 500 400])
set(gca,'Visible','off')
text(0, 0.50, str2,'FontSize',12,'HorizontalAlignment','left')
pause
close(1) 

H = 100.;                %  the nominal layer thickness, m
g = 9.8;
delr = 1.e-3;            %  fractional density change          

% ICtype = menu('What sort of IC?', ...
%   'initial interface displacement', 'initial current')
ICtype = 1;   %  hardwire this to avoid the many hangups caused by menu


lat = input(' Enter the latitude (degrees, 0, or 2 < lat < 90) '); 
f = 2*7.29e-5*sin(lat*pi/180);
IP = 2*pi/(f + 1.e-10);

%  set a flag to indicate if rotation is 'on'
rota = 1;
if abs(lat) <= 0.1
   rota = 0;
end


fwidth = input(' Enter the width of the ridge (km, 50 - 500)  ');

delH = input(' Enter the amplitude of the surface bump (m, 1 - 50) ');

c = sqrt(delr*g*(H + delH));    %  shallow water phase speed
disp(' the non-disersive phase speed, m/s is ')
c

Rd = 1.e10;
if rota == 1
Rd = c/(f + 1.e-10);        %  Rossby radius of deformation
disp(' the Rossby radius of deformation, km is ')
Rdkm = Rd/1000.
end


Ugeo = g*delr*delH/(f*Rd);
Usc = Ugeo;                %  this U scale is good for rotation

if rota == 0               %  this U scale is OK if no rotation
Usc = (delH/H)*c;
end
disp(' the velocity scale, m/s is ')
Usc

%  define the grid
dx = Rd/5;               %  grid interval, units of Rd
L = 40.*Rd;              %  domain half-width, Rd  (the larger, the better)

%  change this to a fixed interval (not scaled with Rd)
dx = 10.e3;
L = 200*dx;

xu = -L:dx:L;            %  grid points for u,v,h 

ndays = input(' Enter the time to integrate (days, 1 - 20) ');

CFL = 0.10;                 %  the CFL number, used to set dt
dt = CFL*dx/c;              %  time step, seconds
nstep = round(ndays*8.64e4/dt);  %  number of time steps to take
  
K = 1.e2;                  %  horizontal diffusion (keep this small)
disp('  the diffusion length scale, km is  ') 
Kd = sqrt(K*ndays*8.64e4)/1000.  
                         
udecay = (1./(5.*8.64e4));    %  Rayleigh drag  (not used at present)
udecay = 0.;

nonlinu = 1.;            %  set these both to 0 to get a linear model
nonlinh = 1.;            %   (for advection of momentum and thickness)

[m nu] = size(xu);
xh = ones(1,nu-1);
xh = cumsum(xh)*dx;
nh = nu - 1;

mov = 1;
if mov == 1 
M = moviein(round(nstep/20));
end

%  pre-allocate some arrays 
u = zeros(1,nu); unew = u; uold = u; uold2 = u;  ut = u;
v = u; vnew = u; vold = u; vold2 = u; vt = u;
h = u; ht = u; hnew = u; holda = u; hold2 = u; 
na = 0; ha = u; ua = u; va = u;

%  pre-allocate some arrays for storage of diagncostics
uc = zeros(1,nstep);
hc = uc; vc = uc; he = uc; ue = uc; ve = vc; ts = uc;
ns9 = floor(nstep/20);    % for data that are decimated in time
ts9 = zeros(1, ns9);
hs = zeros(nu, ns9); 
vs = hs; 
ns = 0;    %  a counter used to keep track of how many stored  

%  Specify the IC, first on h, the interface displacement.

Hscale = 10.e3;   %  scale over which h changes, 10-30 km is OK
xl = fwidth*1000.;      %  convert xl to m 
for m = 1:nh
   xp = xh(m) - L;
   if xp <= -xl; h(m) = delH*(0.5 + 0.5*tanh((xp + xl)/Hscale)); end;
   if xp > -xl; h(m) = delH*(0.5  + 0.5*tanh(-(xp - xl)/Hscale)); end;
end

%  or switch this to the current if you chose 
if ICtype == 2
   v = h*Usc/delH;
   h = 0;
end


%  initialize time levels
holda = h;
hold2 = h;
h0 = h;
vold = v;
vold2 = v;
v0 = v;

%  store the initial h and PV

hi = h;
PVi = f./(H + h);


for m = 1:nstep   

%  time step this baby; use a leap-frog method with an occasional 
%   trapezoidal correction to supress time splitting

if m == 1

%  actually, the first time step is Euler forward
   du = 0.5*(diff([0 u]) + diff([u 0]));
   dv = 0.5*(diff([0 v]) + diff([v 0]));
   dh = 0.5*(diff([0 h]) + diff([h 0]));

   du(1) = 0.; du(nu) = 0.; 
   dv(1) = 0.; dv(nu) = 0.; 
   dh(1) = 0.; dh(nu) = 0.;

   ut = f*v - (delr*g*dh/dx + udecay*u + nonlinu*u.*du/dx);
   vt = -(f*u + udecay*v + nonlinu*u.*dv/dx);
   ht = -(H*du/dx + nonlinh*h.*du/dx + nonlinh*u.*dh/dx);
  
   uold = u;
   vold = v;
   holda = h;
  
   u = uold + dt*ut;
   v = vold + dt*vt;
   h = holda + dt*ht;   

else  
   
   mtype = mod(m,100);  %  decide whether to do a trapezoidal step
   
   if mtype >= 1       %  come here for a simple leap frog 
          
   du = 0.5*(diff([0 u]) + diff([u 0]));
   dv = 0.5*(diff([0 v]) + diff([v 0]));
   dh = 0.5*(diff([0 h]) + diff([h 0]));
   du(1) = 0.; du(nu) = 0.; 
   dv(1) = 0.; dv(nu) = 0.; 
   dh(1) = 0.; dh(nu) = 0.;

%  velocity and derivatives needed for an upwind difference

   ul = 0.5*(u + abs(u));
   ur = 0.5*(u - abs(u));

   dhl = diff([0 h]); dhl(1) = 0.; dhl(nu) = 0.;
   dhr = diff([h 0]); dhr(1) = 0.; dhr(nu) = 0.;
   dul = diff([0 u]); dul(1) = 0.; dul(nu) = 0.;
   dur = diff([u 0]); dur(1) = 0.; dur(nu) = 0.;
   dvl = diff([0 v]); dvl(1) = 0.; dvl(nu) = 0.;
   dvr = diff([v 0]); dvr(1) = 0.; dvr(nu) = 0.;


   du2 = [0 diff(u,2) 0]/dx^2;
   dv2 = [0 diff(v,2) 0]/dx^2;
   dh2 = [0 diff(h,2) 0]/dx^2;
      
   ut = f*v + K*du2 - (delr*g*dh/dx + udecay*u + ...
           nonlinu*(ur.*dur/dx + ul.*dul/dx));

   vt = K*dv2 - (f*u + udecay*v + ...
           nonlinu*(ur.*dvr/dx + ul.*dvl/dx));

   ht = K*dh2 - (H*du/dx + nonlinh*h.*du/dx + ...
           nonlinh*(ur.*dhr/dx + ul.*dhl/dx));
   
   uold2 = uold;
   vold2 = vold;
   hold2 = holda;
   
   uold = u;
   vold = v;
   holda = h;
  
   u = uold2 + 2*dt*ut;
   v = vold2 + 2*dt*vt;
   h = hold2 + 2*dt*ht;   
   
   else    %  do a leap-frog/trapezoidal time step
    
   du = 0.5*(diff([0 u]) + diff([u 0]));
   dv = 0.5*(diff([0 v]) + diff([v 0]));
   dh = 0.5*(diff([0 h]) + diff([h 0])); 
   du(1) = 0.; du(nu) = 0.; 
   dv(1) = 0.; dv(nu) = 0.; 
   dh(1) = 0.; dh(nu) = 0.;

%  velocity and derivatives needed for an upwind difference

   ul = 0.5*(u + abs(u));
   ur = 0.5*(u - abs(u));

   dhl = diff([0 h]); dhl(1) = 0.; dhl(nu) = 0.;
   dhr = diff([h 0]); dhr(1) = 0.; dhr(nu) = 0.;
   dul = diff([0 u]); dul(1) = 0.; dul(nu) = 0.;
   dur = diff([u 0]); dur(1) = 0.; dur(nu) = 0.;
   dvl = diff([0 v]); dvl(1) = 0.; dvl(nu) = 0.;
   dvr = diff([v 0]); dvr(1) = 0.; dvr(nu) = 0.;

   du2 = [0 diff(u,2) 0]/dx^2;
   dv2 = [0 diff(v,2) 0]/dx^2;
   dh2 = [0 diff(h,2) 0]/dx^2;
      
   ut = f*v + K*du2 - (delr*g*dh/dx + udecay*u + ...
           nonlinu*(ur.*dur/dx + ul.*dul/dx));

   vt = K*dv2 - (f*u + udecay*v + ...
           nonlinu*(ur.*dvr/dx + ul.*dvl/dx));

   ht = K*dh2 - (H*du/dx + nonlinh*h.*du/dx + ...
           nonlinh*(ur.*dhr/dx + ul.*dhl/dx));
 
   u1 = u + 0.5*dt*ut;
   v1 = v + 0.5*dt*vt;
   h1 = h + 0.5*dt*ht; 
  
   
   du = 0.5*(diff([0 u1]) + diff([u1 0]));
   dv = 0.5*(diff([0 v1]) + diff([v1 0]));
   dh = 0.5*(diff([0 h1]) + diff([h1 0]));
   du(1) = 0.; du(nu) = 0.; 
   dv(1) = 0.; dv(nu) = 0.; 
   dh(1) = 0.; dh(nu) = 0.;

%  velocity and derivatives needed for an upwind difference

   ul = 0.5*(u1 + abs(u1));
   ur = 0.5*(u1 - abs(u1));

   dhl = diff([0 h1]); dhl(1) = 0.; dhl(nu) = 0.;
   dhr = diff([h1 0]); dhr(1) = 0.; dhr(nu) = 0.;
   dul = diff([0 u1]); dul(1) = 0.; dul(nu) = 0.;
   dur = diff([u1 0]); dur(1) = 0.; dur(nu) = 0.;
   dvl = diff([0 v1]); dvl(1) = 0.; dvl(nu) = 0.;
   dvr = diff([v1 0]); dvr(1) = 0.; dvr(nu) = 0.;
 
   du2 = [0 diff(u1,2) 0]/dx^2;
   dv2 = [0 diff(v1,2) 0]/dx^2;
   dh2 = [0 diff(h1,2) 0]/dx^2;
      
   ut = f*v + K*du2 - (delr*g*dh/dx + udecay*u + ...
           nonlinu*(ur.*dur/dx + ul.*dul/dx));

   vt = K*dv2 - (f*u + udecay*v + ...
           nonlinu*(ur.*dvr/dx + ul.*dvl/dx));

   ht = K*dh2 - (H*du/dx + nonlinh*h.*du/dx + ...
           nonlinh*(ur.*dhr/dx + ul.*dhl/dx));
     
  
   u = uold + dt*ut;
   v = vold + dt*vt;
   h = holda + dt*ht;
    
   uold = u;
   vold = v;
   holda = h;
 
   ut = 0.*ut;
   vt = 0.*vt;
   ht = 0.*ht;
   
end  %  if on mtype


end    %  if on m ==1

 
%  Apply a radiation BC to the grid ends.  This is only moderately 
%   succesful because the phase speed changes in time.
   
   rbc = c*dt/dx;
   
   u(1) = u(1) + rbc*(u(2) - u(1));
   v(1) = v(1) + rbc*(v(2) - v(1));
   h(1) = h(1) + rbc*(h(2) - h(1));
   u(nu) = u(nu) + rbc*(u(nu-1) - u(nu));
   v(nu) = v(nu) + rbc*(v(nu-1) - v(nu));
   h(nu) = h(nu) + rbc*(h(nu-1) - h(nu));
   
% store some data for plotting later
   
 
   ts(m) = m*dt/8.64e4;       %  the time (days)
   np = length(u);
   nf = np - round(xl/dx);
   hf(m) = h(nf);    %  far field h,u,v
   uf(m) = u(nf);
   vf(m) = v(nf);
   
   ne = round(np/2 + 500.e3/dx);
   he(m) = h(ne);  %  h,u,v at a point near the edge of the front
   ue(m) = u(ne);
   ve(m) = v(ne);

   na = na + 1;               %  sum for averages
   ha = ha + h;
   ua = ua + u;
   va = va + v;

%  evaluate the energy balance

ke(m) = sum(0.5*(H + h).*(u.^2 + v.^2));
pe(m) = sum(0.5*g*delr*(h.^2));


%  some diagnostic data are decimated in time
if mod(m,20) == 1    %  save the entire grid occasionally
   ns = ns+1;
   hs(:,ns) = h(:);
   vs(:,ns) = sqrt(v(:).^2 + u(:).^2);
   ts9(ns) = ts(m);

%  make a movie frame, occasionally

if mov == 1
hf1 = figure(1);
clf reset

for q=1:3:nu
   hold on
   xa = xu(q)/1000.; 
   ya = 2.5;
   xb = xa + (2*L/(1000*4))*u(q)/Usc;
   yb = ya + v(q)/Usc;
if rota == 0            %  switch the components if there is no rotation
   xb = xa + (2*L/(1000*4))*v(q)/Usc;
   yb = ya + u(q)/Usc;
end
   plot([xa xb], [ya yb], 'LineWidth', 1.0)
end
axis([-L/1000 L/1000 -0.5 3.5])  %  may need to change these for each case
title('1-D geostrophic adjustment')

hold on
plot(xu/1000., h/delH); 
fill([xu(1)/1000. xu/1000. max(xu)/1000.], [-9 h/delH -9], [0.6 0.6 0.8]);
ylabel('displacement, \eta/\deltah  and  current, U/Usc')
if rota == 0
   ylabel('displacement, \eta/\deltah  and U current, U/Usc')
end
xlabel('distance, km ')
timed = (floor(10*m*dt/8.64e4))/10.;
text(15., 1.7, ['time, days = ', num2str(timed)])
if rota == 1
   tIP = timed*8.64e4/IP;
   timeIP = (floor(10*tIP))/10.;
   text(15., 1.4, ['time, IP = ', num2str(timeIP)])
end


M(:,ns) = getframe;


end    %  if on movie or not

end    %  if on plot this time step ot not
   
end    %  end of the time stepping loop

%  compute time averages

ha = ha/na;
ua = ua/na;
va = va/na;

%  plot some things

figure(2)
clf reset

tip = ts;
subplot(2,1,1)
plot(ts, uf/Usc, 'b', ts, vf/Usc, '--r')
ylabel('U far/Usc')
legend('U','V')
title('Currents at the far field and edge')

subplot(2,1,2)
plot(ts, ue/Usc, 'b', ts, ve/Usc, '--r')
xlabel('time, days'); ylabel('Uedge/Usc')
legend('U','V')


figure(3)
clf reset
mesh(ts9, xu/1000., hs/delH)
xlabel('time, days')
ylabel('distance, km')
zlabel('height/\deltah')
axis([0 ndays -L/1000 L/1000 -0.5 1.5])
title('Geostrophic/Gravitational adjustment; height') 

if rota == 1
%  compute the final PV

dv = 0.5*(diff([0 va]) + diff([va 0]));
dv(1) = 0.; dv(nu) = 0.; 
rvort = dv/dx;
PVf = (f + rvort)./(H + h);
PVflin = (f + 0*rvort)./(H + h);

figure(4)
clf reset
subplot(2,1,1)
plot(xu/1000., hi/delH, 'g', xu/1000., h/delH,'r'); 
%axis([-10 10 -30 30]);
ylabel('height/\delta h')
legend('initial', 'final')
title('Initial and final thickness and PV')

subplot(2,1,2)
plot(xu/1000., PVi,'g', xu/1000., PVf, 'r', xu/1000., PVflin, 'b'); 
%axis([-10 10 2.e-7 3.e-7]); 
xlabel('distance, km')
ylabel('Potential Vorticity')
legend('initial', 'final','final, linear')

figure(5)
clf reset
subplot(2,1,1)
plot(xu/1000., ua/Usc, xu/1000., va/Usc, '--r')
xlabel('distance, km')
ylabel('U, V avg speeds/Ugeo')
legend('Uavg', 'Vavg')
title('Time average currents, etc.')

subplot(2,1,2)
plot(xu/1000., rvort/f)
xlabel('distance, km')
ylabel('rel vorticity/f')
end


figure(8)
clf reset
Esc = g*delr*delH^2*2*xl/dx;
plot(ts, ke/Esc, ts, pe/Esc, ts, (ke+pe)/Esc)
legend('KE', 'PE', 'KE+PE')
xlabel('time, days')
ylabel('energy/PE(t=0)')
title('Energy budget')

% shuffle the figures
if rota == 1
figure(5)
figure(4)
end
figure(3)
figure(2)

% replay the movie
if mov == 1
figure(1)
movie(M,1)
end