clear all

%---------------------------------------------------------
% Calculate and plot w vs t, z vs t, and w vs z given B(z)

% Use Euler Forward or Heun scheme

% Parameters to change: dt, alt_buoy, Heun 

dt = 0.001 % time step (s)  converged (true solution)
%dt = 10 % time step (s)  Use dt from 10 to 0.001 s 
alt_buoy = 0
Heun = 0

% alt_buoy = 1 changes the buoyancy profile for the parcel below its EL
% after it first reaches its maximum height. The buoyancy profile is 
% similar to that for a parcel undergoing dry adiabatic descent.

% ----Do not change any parameters below----

% Use Euler Forward scheme and a quadratic B(z) function,
% B(z) = a * z^2 + b * z + c  (J/kg, z in meters)
a = -2.3438e-08
b = 2.3437e-04
c = -0.2109

q = 9000; % EL (m) (needed for alt_buoy = T)

% Stopping conditions
zmax = 20000 % maximum height allowed (m)
zmin = -1;   % minimum height allowed (m)
tmax = 2000; % maximum time allowed (s)

% Initial values
t = 0;          % time (s)
w0 = 14.2548    % m/s    14.2548 just barely reaches LFC
% w0 = 14
w0 = 14.3
z0 = 0;         % m

i = 0;
irrev = 0;

tic

while z0 < zmax & z0 > zmin & t < tmax
% save into arrays for plotting
i = i + 1; % increment time step counter 
w(i) = w0;
z(i) = z0;
time(i) = t;
% calculate buoyancy at current height, z0  
B = a * z0^2 + b * z0 + c;
if alt_buoy
    if ~irrev
        if B > 0 & w0 < 0
            irrev = 1
        end
    else
% Irreversible buoyancy profile differs below EL = q:
% B' = dB/dz(EL) * (z - q) = (2 * a * q + b) * (z - q)
        if z0 < q
            B = (2 * a * q + b) * (z0 - q);  % larger only below EL
        end
    end
end
% update w and z from w0 to w1 and z0 to z1
w1 = w0 + B * dt;   % dw/dt = B
z1 = z0 + w0 * dt;  % dz/dt = w

B_plot(i) = B;

if Heun
    wstar = w1;
    zstar = z1;
    % calculate buoyancy at current height, zstar
    Bstar = a * zstar^2 + b * zstar + c;
    if alt_buoy
        if  zstar < q
            Bstar = (2 * a * q + b) * (zstar - q);
        end
    end
    % Heun scheme
    w1 = w0 + 0.5 * (B + Bstar) * dt;   % dw/dt = B
    z1 = z0 + 0.5 * (w0 + wstar) * dt;  % dz/dt = w
end
    
% update the time
t = t + dt;
% copy w1 to w0 and z1 to z0 (to use in the next time step)
w0 = w1;
z0 = z1;
end


toc
%------------------------------------------------------------------
% PLOTS