% Script to calculate eta level values from z height values to specify WRF 
% model levels.

% 28 Jan 2014
% John D McMillen
% john.mcmillen@utah.edu

% This script will calculate eta values at w grid points for WRF model 
% levels when given the model top pressure and the vertical distribution 
% of model level w grid points in height, z.  The eta levels will be 
% calculated using the WRF base state atmopshere and assuming the vertical 
% distribution of model level heights is given at sea level.  The 
% formulation of eta values is such that vertical distribution in z at sea 
% level will result in the same vertical ditribution of levels in AGL over 
% terrain elsewhere.

% The variable eta is an array of eta values that can be used in the WRF
% namelist saved as the variable eta_levels.

% This code is adapted from the WRF subroutine compute_eta

clear all;
close all;

% User Supplied Variables

p_top_requested = 5000; % Model top presure.  Specified in the WRF namelist.
                        % If not specified assumed to be 5000 Pa.
                        
z_levs = [0:250:1000 1500:500:6000 6500:1000:18000];    
                        % Run Section 1 first to determine the model top in z (m).
                        % This will help you set up arrays of model levels
                        % in z.  You can specify the levels by setting
                        % z_levels = surface:dz:modeltop 
                        % Note regardless of what z levels you chose start
                        % the array at the surface, 0 and end the array with the
                        % calculated model top.  If you do not do this you
                        % could end up with eta values that are not
                        % monotonic.
                        
                        % If you prefer, you can concatenate arrays with
                        % different intervals into a single array
                        % z_levels = [surface:dz:subtop subbottom:dz:modeltop] 
                        % Note: Do not allow subtop to equal subbottom.
                        % this will cause two duplicate eta values to be 
                        % calculated which will not work in WRF code.
                        
% WRF Base State Atmosphere Variables

p00 = 100000;       % Surface pressure, Pa
t00 = 290;          % Sea level temperature, K
t0  = 300;          % Base state potential temperature, K
a   = 50;           % Base state lapse rate 1000 - 300 hPa, K

% Constants
r_d = 287;          % Dry gas constant
cp = 1004;          % Specific heat of dry air at const P.
cvpm= - 717/1004;   % -cv/cp 
g   = 9.81;         % gravity

% WRF Standard Eta Levels
znw = [1.000000;0.993000;0.983000;0.970000;0.954000;0.934000;0.909000;...
    0.880000;0.844022;0.808045;0.772067;0.736089;0.671281;0.610883;...
    0.554646;0.502330;0.453708;0.408567;0.366699;0.327910;0.292015;...
    0.258837;0.228210;0.199974;0.173979;0.150082;0.128148;0.108049;...
    0.089663;0.072875;0.057576;0.043663;0.031039;0.019611;0.009292;...
    0.000000];

%% 
% Section 1.  Find the model top

mub = p00 - p_top_requested;

% This loop calculates the base state atmosphere at the WRF standard
% eta levels for use in the hydrostatic equation below.

% znu       is the eta value on the mass grid points
% dnw       is the change in eta between eta levels
% p         is the base state pressure on the mass grid points
% t         is the base state temperature on the mass grid points
% t_init    is the base state potential temperature perturbation on the mass grid points
% alb       is the inverse base state density on the mass grid points

for k = 1:(size(znw) - 1)
    znu(k) = (znw(k) + znw(k+1)).*0.5;
    dnw(k) = znw(k+1) - znw(k);
    p(k) = znu(k).*mub + p_top_requested;
    t(k) = t00 + a.*log(p(k)./p00);
    t_init(k) = t(k).*((p00./p(k)).^(r_d./cp)) - t0;
    alb(k) = (r_d./p00).*(t_init(k) + t0).*((p(k)./p00).^cvpm);
end

% This loops solves the hydrostatic equation using base state
% geopotential to find the model top geopotential.  

% phb is the base state geopotential on the w grid points

phb(1) = 0;

for k = 2:size(znw)
    phb(k) = phb(k-1) - dnw(k-1).*mub.*alb(k-1);
end

% Solve for model top on the w grid points in z coords, m.

ztop = phb(length(phb))./g
%% 


%% 
% Section 2.  Compute eta values from z_levs.

% The surface eta value is always equal to 1 and the top is 0.

eta(1) = 1.000;

% This loop calculates the base state atmosphere and the new eta levels 
% starting from the surface and integrating the geopotential/eta form of 
% the hydrostatic relationship to the model top.

% dz         is thickness between z levels
% pw         is the base state pressure on the w grid points
% tw         is the base state temperature on the w grid points
% t_initw    is the base state potential temperature perturbation on the w grid points
% albw       is the inverse base state density on the w grid points

for k = 1:(length(z_levs) - 1)
    dz(k) = z_levs(k+1) - z_levs(k); 
    pw(k) = eta(k).*mub + p_top_requested;
    tw(k) = t00 + a.*log(pw(k)./p00);
    t_initw(k) = tw(k).*((p00./pw(k)).^(r_d./cp)) - t0;
    albw(k) = (r_d./p00).*(t_initw(k) + t0).*((pw(k)./p00).^cvpm);
    eta(k+1) = eta(k) - ((dz(k).*g)./(mub.*albw(k)));    
end

eta(length(z_levs)) = 0.000;

% Print number of levels to screen.

disp('The number of eta levels specified in the WRF namelist:')
e_vert = length(eta)

% Print wrf namelist input to screen.

disp('eta_levels =')
dlmwrite('eta.dat',eta,'delimiter',',','precision','%.4f');
type eta.dat
%% 

%% 
% Section 3.  Calculate model levels on z coords from our new eta levs as a
% visual test.  If the distribution looks good use the eta values in the
% WRF namelist.

% This loop calculates the base state atmosphere at the calculated
% eta levels for use in the hydrostatic equation below.

% znue       is the eta value on the mass grid points
% dnwe       is the change in eta between eta levels
% pe         is the base state pressure on the mass grid points
% te         is the base state temperature on the mass grid points
% t_inite    is the base state potential temperature perturbation on the mass grid points
% albe       is the inverse base state density on the mass grid points

for k = 1:(length(eta) - 1)
    znue(k) = (eta(k) + eta(k+1)).*0.5;
    dnwe(k) = eta(k+1) - eta(k);
    pe(k) = znue(k).*mub + p_top_requested;
    te(k) = t00 + a.*log(pe(k)./p00);
    t_inite(k) = te(k).*((p00./pe(k)).^(r_d./cp)) - t0;
    albe(k) = (r_d./p00).*(t_inite(k) + t0).*((pe(k)./p00).^cvpm);
end

% This loops solves the hydrostatic equation using base state
% geopotential to find the model to in z coords.

% phbe is the base state geopotential on the w grid points from the
% calculated eta values

phbe(1) = 0;

for k = 2:length(eta)
    phbe(k) = phbe(k-1) - dnwe(k-1).*mub.*albe(k-1);
end

% Convert geopotential to heights in z, m.

heights = phbe/g;

% Plot heights of every preposed model level. 

figure(1)
for i = 1:length(heights)
line([0 1],[heights(i) heights(i)]);
hold on
end
axis([0 1 0 20000]);
title('WRF Model Levels')
ylabel('AGL Height (m)')
set(gca,'Fontname','Times','xticklabel','','xtick',[])
set(findall(gcf,'type','text'),'Fontname','Times')
set(gca,'FontSize',16)
set(findall(gcf,'type','text'),'FontSize',16)
hold off;

% Optional code to save image of the vertical height plot.

%print -dpng -r200 modellevels.png;

% Plot heights of every proposed model level for a subset of the total depth. 

figure(2)
for i = 1:length(heights)
line([0 1],[heights(i) heights(i)]);
hold on
end

% Setting the axis plots the portion of the column of interest.
axis([0 1 0 7000]);
title('WRF Model Levels')
ylabel('AGL Height (m)')
set(gca,'Fontname','Times','xticklabel','','xtick',[])
set(findall(gcf,'type','text'),'Fontname','Times')
set(gca,'FontSize',16)
set(findall(gcf,'type','text'),'FontSize',16)
hold off;

% Optional code to save image of the vertical height plot.

%print -dpng -r200 modellevels.png;

%% 
    
