%% Atm S 547 HW4: 1st order closure model for stable boundary layer
%  Chris Bretherton
%
%%  Input parameters and theta profile (thetai):
%  Use evenly spaced flux levels from 0-ztop m with spacing dz, with 
%  initial theta given at half-levels zh. lambda is the asymptotic Blackadar 
%  lengthscale, and Hs0 is the surface sensible heat flux (W/m2)

  % User-specifiable time and space grid and output parameters
  
  tf = 3*3600;
  dtout = 600;  % Time between outputs for z-t plots
  dtprof = 3600; % Time between outputs for z-profiles
  dz = 10;
  ztop = 150;

  % User-specifiable internal model parameters
  lambda = 5;
  Hs0 = -25;
  ustar = 0.2;
  ug = 10;
  vg = 0;
  theta0 = 290;
  dthdz = 0.01;
  
  % Physical constants
  g=9.8;
  f = 1e-4;
  rhoref = 1.2;
  cp = 1000;
  k = 0.4;
  
  % Define grid and auxiliary parameters
  
  tspan = 0:dtout:tf;
  tprof = 0:dtprof:tf;
  nprof = tf/dtprof;
  noutperprof = dtprof/dtout;
  
  zh = (dz/2:dz:(ztop-dz/2))';
  zf = (0:dz:ztop)';
  nz = length(zh);

  % Initial theta, u and v profiles on the half-levels
  thetai = theta0 + dthdz*zh;
  ui = ug*ones(nz,1);
  vi = zeros(nz,1);

  %  Specify Blackadar master lengthscale for interior flux (full) levels
  %  for use in ODE solver.

  l = lambda./(1 + lambda./(k*zf(2:nz)));

  % Compute/print surface buoyancy flux and Obukhov length (prob. 1)  
  Fth0 = Hs0/(rhoref*cp);
  B0 = g*Fth0/theta0
  Lob = -ustar^3/(k*B0)

%% Use Matlab ODE solver ode45 to solve the equation. 

  % Define initial solution vector 
  yi = [thetai;ui;vi];
  
  % Call the ODE solver...Uses dydt.m to specify time derivatives
  [t,y] = ode45('dydt',tspan,yi,[],zf,zh,l,Fth0,ustar,ug,vg);

%% Separate solution y into theta;u;v
  theta = y(:,1:nz);
  u = y(:,nz+1:2*nz);
  v = y(:,2*nz+1:3*nz);
  
  theta = theta';  % Now each column of theta corresponds to a time in tspan.
  u=u';
  v=v';
  
  figure(1)
  clf
  subplot(2,2,1)
  colormap(gray)
  pcolor(t/3600,zh,theta)
  shading interp
  hold on
  contour(t/3600,zh,theta,10,'k')
  colorbar
  xlabel('time (hr)')
  ylabel('z (m)')
  title('Time-height section of \theta')
  hold off

  subplot(2,2,2)
  colormap(gray)
  pcolor(t/3600,zh,u)
  shading interp
  hold on
  contour(t/3600,zh,u,10,'k')
  colorbar
  xlabel('time (hr)')
  ylabel('z (m)')
  title('Time-height section of u')
  hold off

  subplot(2,2,3)
  colormap(gray)
  pcolor(t/3600,zh,v)
  shading interp
  hold on
  contour(t/3600,zh,v,10,'k')
  colorbar
  xlabel('time (hr)')
  ylabel('z (m)')
  title('Time-height section of v')
  hold off
  
  subplot(2,2,4)
  plot([0; l],zf(1:nz))
  xlabel('[m]')
  ylabel('z [m]')
  title(['Blackadar lengthscale, \lambda =' num2str(lambda) ' m'])

%%
  figure(2)
  clf

  % Plot hourly theta profiles
  
  subplot(2,2,1)
  iprof = (0:nprof)*noutperprof + 1;
  plot(theta(:,iprof),zh)
  xlabel('\theta (K)')
  ylabel('z (m)')
  legend('0 hr','1 hr','2 hr','3 hr','Location','West')
  title('Profiles of \theta')
  text(0.5 + theta(1,end), 0.8*ztop, ['\lambda = ' num2str(lambda) ' m'])
  text(0.5 + theta(1,end), 0.9*ztop, ['H_s(0) = ' num2str(Hs0) ' W m^{-2}'])

  % Generate and plot hourly profiles of heat flux, Ri and K

  Kprof = zeros(nz+1,nprof+1);
  Fthprof = zeros(nz+1,nprof+1);
  Riprof = zeros(nz+1,nprof+1);
  
  for it = 1:(nprof+1)
    ip = iprof(it);
    [tend, Kh, Fth, Ri] = dydt(t(ip),y(ip,:)',flag,zf,zh,l,Fth0,ustar,ug,vg);
    Kprof(:,it) = [0; Kh; 0];
    Fthprof(:,it) = Fth;
    Riprof(:,it) = [0; Ri; 0];
  end
  
  subplot(2,2,2)
  plot(Riprof,zf)
  xlabel('Ri')
  ylabel('z (m)')
  title('Ri profiles')
  legend('0 hr','1 hr','2 hr','3 hr','Location','West')
  xlim([0 0.25])
  text(0.02, 0.8*ztop, ['\lambda = ' num2str(lambda) ' m'])
  text(0.02, 0.9*ztop, ['H_s(0) = ' num2str(Hs0) ' W m^{-2}'])

  subplot(2,2,3)
  plot(rhoref*cp*Fthprof,zf)
  xlabel('Sensible heat flux (W m^{-2})')
  ylabel('z (m)')
  title('Heat flux profiles')
  legend('0 hr','1 hr','2 hr','3 hr','Location','NorthWest')

  subplot(2,2,4)
  plot(Kprof,zf)
  xlabel('K_h (m s^{-2})')
  ylabel('z (m)')
  title('K_h profiles')
  legend('0 hr','1 hr','2 hr','3 hr')



  
