<jupyter_start><jupyter_text>True Zonal MeanCalculate the *true zonal mean* of a scalar quantity regardless of the horizontal mesh. Specifically, we calculate the volume weighted mean along all grid cells whose centres fall within finite latitude intervals rather than the arithmetic mean of cells along the model's curvilinear grid. The method presented can also be used to re-grid models onto the same latitudinal grid and the general principles can be used to define any multidimensional sum or average using the `xhistogram` package.    **Requirements:** Select the `conda/analysis3-20.01` (or later) kernel.This code should work for just about any MOM5 configuration since all we are grabbing is temeprature and standard grid information. You can swap temperature with any other scalar variable. You can also in principle swap latitude with another scalar.Last updated by <NAME> 8-4-2020<jupyter_code>%matplotlib inline
%config InlineBackend.figure_format='retina'

import cosima_cookbook as cc
import matplotlib.pyplot as plt
import xarray as xr
import numpy as np
from dask.distributed import Client
from xhistogram.xarray import histogram #This package has only recently been added to conda/analysis3-20.01
client = Client(n_workers=4)
client
session = cc.database.create_session()<jupyter_output><empty_output><jupyter_text>Choose the run and the variable we want to average. I am using temperature but you can choose any scalar.<jupyter_code>expt =  '025deg_jra55v13_ryf8485_gmredi6' #Choose any run
variable = 'temp' #any scalar variable where volume weighted averaging makes sense
variable_to_average = cc.querying.getvar(expt, variable, session, ncfile='ocean.nc')
variable_to_average<jupyter_output><empty_output><jupyter_text>First we will show the standard approach, which is to take the arithmetic mean of all grid cells along the quasi-longitudinal coordinate. For MOM5's tri-polar grid this approach is in principle "okay" for the southern hemisphere where grid cell areas are constant at fixed latitude. It doesn't though, take into account partial cells. The `xarray`'s method `.mean(dim='dimension')` applies `numpy.mean()` across that dimension. This is simply the arithmetic mean.For some scalar $T$ the arithemtic mean is given by$$ \left_{j,k} = \frac{1}{I}\sum_{i=1}^{I} T_{i,j,k},$$where $i$, $j$ and $k$ are the indicies in the $x$, $y$ and $z$ directions respectively of the curvilinear grid and $I$ is the number of indicies along the $x$ axis.<jupyter_code>x_arith_mean = variable_to_average.groupby('time.year').mean(dim='time').mean(dim='xt_ocean')

plt.figure(figsize=(10, 5))
x_arith_mean.sel(year=2000).plot(yincrease=False, vmin=273, vmax=300, cmap='Oranges')
plt.title('x-coordinate arithmetic mean');<jupyter_output><empty_output><jupyter_text>The main issue with this average is the 'latitude' coordinate may be meaningless near the north pole, particulalrly when comparing to observational analysis or other models which will have either a regular grid or a different curvilinear grid. Even different versions of MOM might have different grids. Let us consider what the true zonal average looks like. That is consider a set of latitude 'edges' $\{\phi'_{1/2},\phi'_{1+1/2},...,\phi'_{\ell-1/2},\phi'_{\ell+1/2},...,\phi'_{L+1/2}\}$ between which we want to compute an average of $T$ at $\{\phi'_{1},\phi'_{2},...,\phi'_{\ell},...,\phi'_{L}\}$ such that$$ \overline{T}(\phi'_\ell,\sigma) = \dfrac{\iint_{\phi'_{\ell-1/2} < \phi \leq \phi'_{\ell+1/2}} T(\phi,\lambda,\sigma)\frac{\partial z}{\partial \sigma}(\phi,\lambda,\sigma)\,\mathrm{d}A}{\iint_{\phi'_{\ell-1/2} < \phi \leq \phi'_{\ell+1/2}}\frac{\partial z}{\partial \sigma}(\phi,\lambda,\sigma)\,\mathrm{d}A},$$where $\lambda$ is longitude and $\sigma$ is an arbitrary vertical coordinate. In discrete form this average is$$\overline{T}_{\ell,k} = \frac{\sum_{i=1}^{I}\sum_{j=1}^{J}\delta_{i,j}T_{i,j,k}\Delta Z_{i,j,k}\Delta Area_{i,j}}{\sum_{i=1}^{I}\sum_{j=1}^{J}\delta_{i,j,k}\Delta Z_{i,j,k}\Delta Area_{i,j}},$$where $\delta_{i,j} = 1$ if $\phi'_{\ell-1/2}<\phi_{i,j}\leq \phi'_{\ell+1/2}$ and $\delta_{i,j} = 0$ elsewhere, $\Delta Z$ is the grid cell vertical thickness and $\Delta Area$ is the grid cell horizontal area.For our purposes we will use the edges of the models `xt_ocean` coordinate to define $\phi'_{\ell+1/2}$ so the number of 'bins' $L$ will be the same as the length of the quasi-latitude coordinate ($J$). Fortunately, as you can see below, the two sums are weighted histograms (one for $T$ times volume and the other for just volume) and these can be rapidly computed using `xhistogram`.First lets load the scalar variable (latitude) we want to use as our coordinate then define the bin edges.<jupyter_code>coord = 'geolat_t' #can be any scalar (2D, 3D, eulerian, lagrangian etc)
variable_as_coord = cc.querying.getvar(expt, coord, session, n=1) #might need ncfile='ocean.nc' if a tracer

#Define the coordinate bins as the latitude edges of the T-cells
yu_ocean = cc.querying.getvar(expt, 'yu_ocean', session, n=1)
#make numpy array (using .values) and add 1st edge at -90
bins =  np.insert(yu_ocean.values, 0, np.array(-90), axis=0) 

#Alternatively we could just use some regular grid like this 
#coordbins =  np.linspace(-80, 90, 50)
#or use a grid from a different (coarser) model.

dzt = cc.querying.getvar(expt, 'dzt', session, ncfile='ocean.nc') #thickness of cells
area_t = cc.querying.getvar(expt, 'area_t', session, n=1) #area of cells
dVol = dzt*area_t #Volume of cells<jupyter_output><empty_output><jupyter_text>Now let's compute the numerator and denominator of the equation above using `xhistogram`, then the time mean and then the zonal mean.<jupyter_code>histVolCoordDepth = histogram(variable_as_coord.broadcast_like(dVol).where(~np.isnan(dVol)), bins=[bins], weights=dVol, dim=['yt_ocean', 'xt_ocean'])
histTVolCoordDepth = histogram(variable_as_coord.broadcast_like(dVol).where(~np.isnan(dVol)), bins=[bins], weights=dVol*variable_to_average, dim=['yt_ocean', 'xt_ocean'])
coord_mean = (histTVolCoordDepth/histVolCoordDepth).groupby('time.year').mean(dim='time')<jupyter_output><empty_output><jupyter_text>Now we can plot the results which thankfully retain all the data-array info on variables and axis etc.<jupyter_code>plt.figure(figsize=(10, 5))
coord_mean.sel(year=2000).plot(yincrease=False, vmin=273, vmax=300, cmap='Oranges')
plt.title('True zonal mean');<jupyter_output><empty_output><jupyter_text>Since we used the same bin edges as the standard `yt_ocean` coordinate we can take a difference between the arithmetic mean along the model's x-axis and our mean along grid cells within latitude bands. The main differences are near the North Pole where the grid is furthest for being regular. There are also differences near the Antacrtic Shelf suggesting partial cells also matter.<jupyter_code>zonal_minus_x_mean = coord_mean.sel(year=2000)-x_arith_mean.sel(year=2000).values

plt.figure(figsize=(10, 5))
zonal_minus_x_mean.plot(yincrease=False, vmin=-0.8, vmax=0.8, cmap='RdBu_r')
plt.title('True zonal minus x-coordinate arithmetic mean');<jupyter_output><empty_output><jupyter_text>`xarray` has a new `weighted` functionality which allows it to do weighted means instead of arithmetic mean.Let's see how that works out... We chose `dVol` as the weights and we only do the comparisson for year 2000.<jupyter_code>variable_to_average_weighted = variable_to_average.copy().sel(time='2000').mean(dim='time')
variable_to_average_weighted = variable_to_average_weighted.weighted(dVol.sel(time='2000').fillna(0))
meanweighted_y2000 = variable_to_average_weighted.mean(dim='xt_ocean').groupby('time.year').mean(dim='time').sel(year=2000)
zonal_minus_x_mean = coord_mean.sel(year=2000) - meanweighted_y2000.values

plt.figure(figsize=(10, 5))
zonal_minus_x_mean.plot(yincrease=False, vmin=-0.8, vmax=0.8, cmap='RdBu_r')
plt.title("True zonal minus xarray''s x-coordinate weighted mean");<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12330487192365645, 'fraction_numerical': 0.02021597187343044, 'mean_word_length': 6.048672566371682, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '462275', 'license': 'Apache-2.0', 'n_tokens_mistral': 2707, 'n_tokens_neox': 2495, 'n_words': 838, 'path': 'DocumentedExamples/True_Zonal_Mean.ipynb', 'repo_name': 'hrsdawson/cosima-recipes', 'size': '7966'}