% Script Ek.m
% Input parameters:
dt = 108;                               % Time Step (s)
tf = 8*86400;                           % Final Time (s)
dz = 50;                                % Spatial Resolution (m)
h = 10000;                              % Maximum height (top of model) (m)
n = floor(h/dz);                        % Number of vertical grid points
m = floor(tf/dt);                       % Number of time steps
Km = 12.5;                              % Eddy viscosity (m^2 s^{-1})
f = 1e-4;                               % Coriolis parameter (s^{-1})
ug = 10*ones([n 1]) ;%+ 40*[1:n]'./n;   % Geostrophic Wind: u (m s^{-1})
vg = 0*ones([n 1]);                     % Geostrophic Wind: v (m s^{-1})

% Initialize variables

u = zeros([n m]);                       % U velocity (m/s)
v = zeros([n m]);                       % V velovity (m/s)
%u0 = u(:,1);   
%v0 = v(:,1);
u0 = ug;                                % Set initial wind to geostrophic wind    
v0 = vg;
u1 = u(:,1);
v1 = v(:,1);

% Set up diffusion matrix, with boundary conditions (no-slip at surface,
% free slip at top
Au = -diag(ones([n-1 1])*Km/dz/dz,-1)+diag(ones([n 1])*(1/dt +2*Km/dz/dz),0) + ...
    -diag(ones([n-1 1])*Km/dz/dz,1);
Av = -diag(ones([n-1 1])*Km/dz/dz,-1)+diag(ones([n 1])*(1/dt +2*Km/dz/dz),0) + ...
    -diag(ones([n-1 1])*Km/dz/dz,1);
Au(n,n) = 1/dt+Km/dz/dz;
Av(n,n) = 1/dt+Km/dz/dz;


for i = 1:m;

%  Adjust geostrophic wind (if desired)
%    if (i<m/2)
%        ug = 10 + 40*[1:n]'./n*i/(m/2);
%    else
    ug = 10*ones([n 1]) ;%+ 40.0*[1:n]'./n;
%    end
%  Integrate wind to account for coriolis accelerations 
    u1 = u0 + f*(v0-vg)*dt;
    v1 = v0 - f*(u0-ug)*dt;
%  Integrate wind to account for eddy viscosity, using fully implicit
%  scheme
    u1 = (u1'/Au)'/dt;
    v1 = (v1'/Av)'/dt;
% Apply newtonian relaxation to geostrophic wind in top model layers (not
% strictly necessary....)
    u1(n-9:n) = u1(n-9:n) + (ug(n-9:n)-u1(n-9:n))*dt/800;
    v1(n-9:n) = v1(n-9:n) + (vg(n-9:n)-v1(n-9:n))*dt/800;
% Store results
    u(:,i) = u1;
    v(:,i) = v1;
% Increment initial wind
    u0 = u1;
    v0 = v1;
end
z = [dz:dz:h];
a = (f/(2*Km)).^0.5;
ut = ug(1)*(1-exp(-a*z).*cos(a*z));
vt = ug(1)*exp(-a*z).*sin(a*z);