This notebooks is used to develope and test a function that will calulate the depth average of a model variable.
Finding the depth might be tricky. We could use the bathymetry but I think there are differences between the bathymetry file and the actual depth of the grid cell. I can think of two approaches
My plan is to use the first approach.
%matplotlib inline
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import os
def depth_average(var,depths,depth_axis):
"""Average over depth using the trapezoid rule. The variable should be masked to apply this function.
The depth is calcluated based on masking.
If the variable is not masked then the maximum depth of the entire domain is used.
:arg var: variable to average
:type var: masked numpy array
:arg depths: the depths associated with var
:type depths: numpy array
:arg depth_axis: The axis in var associated with depth
:type depth_axis: int
:returns: avg, the depth averaged var.
"""
# Make sure depths is an array and not an netcf variable.
de=np.array(depths)
# Integrate, the easy part
integral = np.trapz(var,x=de,axis=depth_axis)
# Find depth for averaging
# Need to expand the depths array to same shape as the variable. This is really awkward..
for n in np.arange(var.ndim-1):
de=de[:,np.newaxis]
roll=np.rollaxis(var,depth_axis)
expanded_depths= de +np.zeros(roll.shape)
expanded_depths=np.rollaxis(expanded_depths,0,depth_axis+1)
# Apply variable mask to depth masks
mask = np.ma.getmask(var)
depth_masked = np.ma.array(expanded_depths,mask=mask )
# Look up maximum depth and surface depth to calculate total depth of water column
max_depths = np.ma.max(depth_masked, axis=depth_axis)
surface_depths = depth_masked.take(0,axis=depth_axis)
total_depth = max_depths-surface_depths
# Divide integral by total depth
average = integral/total_depth
return average
#depths (linear)
depths = np.arange(0,40)
print depths.shape
(40,)
Zero array
#1D zeros array
var = np.zeros((depths.shape[0],1))
print depth_average(var,depths,depth_axis=0)
print var.shape
[0.0] (40, 1)
Testing size. What happens if the shape of the array is (N,) instead of (N,1)
#1D zeros array
var = np.zeros((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape
0.0 (40,)
Constant array
#1D ones array
var = np.ones((depths.shape[0],1))
print depth_average(var,depths,depth_axis=0)
print var.shape
[1.0] (40, 1)
First axis time, second axis depth
#1D ones array
#num output times
ts = 10;
var = np.ones((ts,depths.shape[0]))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
[1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] (10,) (10, 40)
First axis time, second axis depth, thrid axis x
#1D ones array
#num output times
ts = 10; xs=11
var = np.ones((ts,depths.shape[0],xs))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
[[1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0]] (10, 11) (10, 40, 11)
First axis time, second axis depth, third axis x, fourth axis y
#1D ones array
#num output times
ts = 1; xs=2; ys=3
var = np.ones((ts,depths.shape[0],xs,ys))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
[[[1.0 1.0 1.0] [1.0 1.0 1.0]]] (1, 2, 3) (1, 40, 2, 3)
Move depth axis
#1D ones array
#num output times
ts = 1; xs=2;
var = np.ones((ts,xs,depths.shape[0]))
avg= depth_average(var,depths,depth_axis=2)
print avg
print avg.shape
print var.shape
[[1.0 1.0]] (1, 2) (1, 2, 40)
Different values along time axis
ts = 2; xs=3;
var = np.ones((ts,depths.shape[0],xs))
var[0,...] = 2*var[0,...]
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
print 'First time'
print avg[0,:]
print 'Second time'
print avg[1,:]
[[2.0 2.0 2.0] [1.0 1.0 1.0]] (2, 3) (2, 40, 3) First time [2.0 2.0 2.0] Second time [1.0 1.0 1.0]
Different values along x axis
ts = 3; xs=2;
var = np.ones((ts,depths.shape[0],xs))
var[...,1] = 2*var[...,1]
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
print 'First x'
print avg[...,0]
print 'Second x'
print avg[...,1]
[[1.0 2.0] [1.0 2.0] [1.0 2.0]] (3, 2) (3, 40, 2) First x [1.0 1.0 1.0] Second x [2.0 2.0 2.0]
# Load depths from model data.
run='dwr_notsmooth'
base='/data/nsoontie/MEOPAR/SalishSea/results/stratification/'
path = os.path.join(base,'{}/SalishSea_1d_20030819_20030927_grid_T.nc'.format(run))
f = nc.Dataset(path,'r');
depths = f.variables['deptht']
print depths.shape
(40,)
Zeros
var = np.zeros((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape
0.0 (40,)
Ones
var = np.ones((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape
1.00000000014 (40,)
Some small error here. But why...?
print depths[-1] - depths[0]
440.966
diffs = depths[1:]-depths[0:-1]
s = np.sum(diffs)
print s
440.966
inte = np.trapz(var, x=depths[:],axis=0)
print inte
print inte-s
440.966094732 3.02791595459e-05
In this case, the trapezoid rule should just give us the sum of the grid spacings. There may be some a mismatch in the floating point precision, causing the error. But the error is small and I won't worry about it.
Entire array masked, 1D
var = np.zeros((depths.shape[0]))
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=0)
print var.shape
print var
-- (40,) [-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --]
Entire array masked, multiple dimensions
var = np.zeros((3,depths.shape[0]))
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var
[-- -- --] (3, 40) [[-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --] [-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --] [-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --]]
Part of array is masked over entire water column
var = np.ones((4,depths.shape[0]))
var[0,:]=0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var
[-- 1.0000000001351683 1.0000000001351683 1.0000000001351683] (4, 40) [[-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0] [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0]]
Mask a bunch of the values
var = np.ones((depths.shape[0]))
var[10:] = 0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=0)
print var.shape
print var
1.00000000662 (40,) [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --]
Multidimensional masking
ts=2; xs=3
var = np.ones((ts,depths.shape[0],xs))
var[0,10:,0] = 0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var[0,:,0]
[[1.000000006620949 1.0000000001351683 1.0000000001351683] [1.0000000001351683 1.0000000001351683 1.0000000001351683]] (2, 40, 3) [1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --]
I'm pretty satisified with how this is working. There are some small rouding errors. And it must be used with maskd data.
Next I will add this function to analyze module.