% Example script for Assignment 2
% MATLAB

%% Open the data file and see what its contents are

% Put the path to your data file here
data_file = '../NCEP2_time_mean.nc';

% Display the contents of the data file
ncdisp(data_file)

%% Read the data we require

% Read the coordinate variables
lat = ncread(data_file,'lat');
lon = ncread(data_file,'lon');
p = ncread(data_file,'p');

% Read the variables you want (Load the temperature and relative humidity for now)
T = ncread(data_file,'T');
RH = ncread(data_file,'RH');

% Display the size of T to ensure everything is okay
disp(' ')
disp(' ')
disp('Just read the temperature. it has dimensions:')
disp(num2str(size(T)))
disp(' ')
disp(' ')

%% Some useful constants
Rd = 287.04; % Dry gas constant (J/K/kg)

%% Calculate any derived variables

% Matlab is not very good at loops, but good at multiplying entire arrays.
% In order to make use of this functionality, it helps to create arrays
% containing the coordinate variables that are the same size as the 
% dynamical variables

% First, create a pressure matrix that is the same size as the zonal-mean temperature matrix
p_matrix = repmat(p,[1 length(lat) length(lon)]); % Use repmat to create copies of the pressure
p_matrix = permute(p_matrix,[3 2 1]); % Now switch the order of the dimensions using permute

% Do the same for latitude
lat_matrix = repmat(lat,[1 length(p) length(lon)]); % Use repmat to create copies of the pressure
lat_matrix = permute(lat_matrix,[3 1 2]); % Now switch the order of the dimensions using permute

% Matlab is very good at element-wise arithmetic. Just remember to include
% a dot "." before multiplication or division symbols to ensure you are
% doing element-wise operations and not matrix operations.

% For example, we can calculate the density
rho = p_matrix./(Rd.*T);

% Note that here we use ".*" and ".^" to refer to element-wise operations.
% Simply using "*" causes an error, because MATLAB expects that you want to
% take use matrix multiplication.
% Also, we used "p_matrix", since we have constructed it to be the same
% size as T.



%% Calculate the zonal-mean temperature & relative humidity

% Calculate a mean along the first dimension
T_zonal_mean = mean(T,1);
RH_zonal_mean = mean(RH,1);

% "squeeze" removes extraneous dimensions from variables
T_zonal_mean = squeeze(T_zonal_mean);
RH_zonal_mean = squeeze(RH_zonal_mean);



%% Plot the relative humidity as a function of latitude and pressure

% Create a figure
fig1 = figure;

% Create axes
ax = axes;

% Transpose the dimensions, because that is how the plotting functions want the data
plot_var = permute(RH_zonal_mean,[2 1]);

% Pick the levels you want to put contours at
contour_levels = [0 5 10 15 20 25 30 35 40 45 50 55 60 65 70 75 80 85 90 95 100];

% Use either "contour" or "contourf"
contourf(lat,p,plot_var,contour_levels);
%contour(lat,p,plot_var,contour_levels);

% Make the pressure decrease upwards
set(ax,'ydir','reverse')

% Set the axis labels
xlabel(ax,'latitude (deg)')
ylabel(ax,'pressure (hPa)')

% Create a title
title(ax,'Zonal- and time-mean relative humidity')

% Make a colour bar
c = colorbar;
ylabel(c,'relative humidity (%)')

% Set the colorbar limits (from 0 to 100 %)
set(ax,'clim',[0 100])

% Save the figure as a pdf
print('-dpdf','./zonal_mean_RH_contour_plot.pdf')


%% Plot the zonal-mean temperature at 500 hPa and 850 hPa

% Create a figure
fig2 = figure;

% Create axes
ax = axes;

% 500 hPa corresponds to the 6th element of the pressure array
Ilev = 6;

% Use "plot" to make a line plot
plot(lat,T_zonal_mean(:,Ilev));

% This allows you to make another plot on top of this one
hold on

% 850 hPa corresponds to the 3rd element of the pressure array
Ilev = 3;

% Use "plot" to make a line plot
plot(lat,T_zonal_mean(:,Ilev));

% Set the axis labels
xlabel(ax,'latitude (deg)')
ylabel(ax,'temperature (K)')

% Set the axis limits
set(ax,'xlim',[-90 90])
set(ax,'ylim',[230 295])

% Create a title
title(ax,'Zonal- and time-mean temperature')

% create a legend
legend('500 hPa','850 hPa')

% Save the figure as a pdf
print('-dpdf','./zonal_mean_T_line_plot.pdf')
