def get_hx(fil1, fil2, **initializer):
bc = dict(
v=['mirror', 'circsymh', 'zeros', 'circsymq', 'circsymh', 'neumann'])
with pym6.Dataset(fil2, **initializer) as ds, pym6.Dataset(
fil1, **initializer) as ds1:
hx = ds1.h_Cv.bc_type(bc).xsm().xep().read().dbyd(3).move_to(
'v').nanmean(axis=(0, 2)).compute()
h = ds1.h_Cv.read().nanmean(axis=(0, 2)).compute()
hx.values[h.values < initializer['htol']] = 0
return h, hx
def get_eprimesq(fil1, fil2, **initializer):
#initializer['final_loc'] = 'vl'
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
esq = ds1.esq.read().nanmean(axis=0)
e = ds2.e.zep().read().nanmean(axis=0).move_to('l')**2
e = e.compute(check_loc=False)
edlsqm = esq - e
edlsqm = edlsqm.compute()
return edlsqm
def get_edlsqmx(fil1, fil2, **initializer):
initializer['final_loc'] = 'ul'
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
esq = ds1.esq.xep().read().nanmean(axis=0)
e = ds2.e.xep().zep().read().nanmean(axis=0).move_to('l')**2
e = e.compute(check_loc=False)
edlsqmx = esq - e
edlsqmx = edlsqmx.dbyd(3).compute()
return edlsqmx
def get_edlsqmy(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
esq = ds1.esq.yep().read().nanmean(axis=0)
e = ds2.e.yep().zep().read().nanmean(axis=0).move_to('l')**2
e = e.compute(check_loc=False)
edlsqmy = esq - e
edlsqmy = edlsqmy.dbyd(2).multiply_by('dxCv').reduce_(
np.nansum, axis=3).compute()
return edlsqmy
def diapycnal(fil1, fil2, **initializer):
with pym6.Dataset(fil2,
**initializer) as ds, pym6.Dataset(fil1,
**initializer) as ds1:
dwd = ds.wd.final_loc('ql').xep().yep().zep().read().nanmean(
axis=0).np_ops(np.diff, axis=1,
sets_vloc='l').move_to('u').move_to('q').compute()
term7 = (-dwd).compute()
term7.math = r'$(\bar{\sigma}\hat{\varpi})_{\tilde{b}}$'
term7.name = r"Diapycnal term"
return term7
def ZoverturnSFstanding(expt):
with pym6.Dataset(expt.fil2, geometry=expt.geometry) as ds, pym6.Dataset(
expt.fil1) as ds1:
u = ds.u.read().nanmean(axis=0).compute()
u.values = u.values - np.mean(u.values, axis=2, keepdims=True)
u = u.multiply_by('dyCu')
h = ds1.h_Cu.read().nanmean(axis=0).compute()
h.values = h.values - np.mean(h.values, axis=2, keepdims=True)
zoc = (u * h).reduce_(np.sum, axis=2).compute()
zoc.array = np.cumsum(zoc.values, axis=1)
zoc = zoc.tob(axis=1).to_DataArray().squeeze()
return zoc / 1e6
def MoverturnSFstanding(expt):
with pym6.Dataset(expt.fil2, geometry=expt.geometry) as ds, pym6.Dataset(
expt.fil1) as ds1:
v = ds.v.read().nanmean(axis=0).compute()
v.values = v.values - np.mean(v.values, axis=3, keepdims=True)
v = v.multiply_by('dxCv')
h = ds1.h_Cv.read().nanmean(axis=0).compute()
h.values = h.values - np.mean(h.values, axis=3, keepdims=True)
moc = (v * h).reduce_(np.sum, axis=3).compute()
moc.array = np.cumsum(moc.values, axis=1)
moc = moc.tob(axis=1).to_DataArray().squeeze()
return moc / 1e6
def extract_budget(fil1, fil2, fil3=None, **initializer):
blist, PV = extract_twapv_terms(fil1, fil2, **initializer)
meanax = initializer.get('meanax', 2)
initializer.pop('meanax')
z = initializer.get('z', None)
if 'z' in initializer:
initializer.pop('z')
initializer['final_loc'] = 'qi'
with pym6.Dataset(fil2, **initializer) as ds:
e = ds.e.xep().zep().yep().read().move_to('u').move_to('q').nanmean(
axis=(0, 2)).compute()
initializer['final_loc'] = 'ql'
if fil3 is not None:
with pym6.Dataset(fil3) as ds:
islaydeepmax = ds.islayerdeep.read().compute(check_loc=False).array
islaydeepmax = islaydeepmax[0, 0, 0, 0]
with pym6.Dataset(fil3, fillvalue=np.nan, **initializer) as ds:
swash = ds.islayerdeep
swash = swash.read().nanmean(axis=(0, 2)).compute()
swash = ((-swash + islaydeepmax) / islaydeepmax * 100).compute()
if z is not None:
swash = swash.toz(z, e)
swash = swash.tokm(3).to_DataArray()
else:
swash = None
PV = PV.nanmean(axis=meanax)
if z is not None:
PV = PV.toz(z, e)
PV = PV.tokm(3).to_DataArray(check_loc=False)
namelist = []
for i, b in enumerate(blist):
namelist.append(b.name)
b = b.nanmean(axis=meanax)
if z is not None:
b = b.toz(z, e).compute(check_loc=False)
blist[i] = b.tokm(3).to_DataArray(check_loc=False)
blist_concat = xr.concat(blist, dim=pd.Index(namelist, name='Term'))
blist_concat.name = 'TWA PV budget'
e = e.tokm(3).to_DataArray()
withPV = initializer.get('withPV', True)
if withPV:
return_dict = dict(blist_concat=blist_concat,
blist=blist,
e=e,
PV=PV,
swash=swash)
else:
return_dict = dict(blist_concat=blist_concat,
blist=blist,
e=e,
swash=swash)
return return_dict
def eape(fil1, fil2, **kwargs):
htol = kwargs.get('htol', 1e-10)
# meanape = mape(fil1, fil2, **kwargs)
with pym6.Dataset(fil1, **kwargs) as ds1, pym6.Dataset(fil2,
**kwargs) as ds2:
db = np.diff(ds1.zl)[0] * 9.8 / 1000
esqm = ds1.esq.read().nanmean(axis=0).compute()
emsq = ds2.e.final_loc('hl').zep().read().nanmean(
axis=0).move_to('l')**2
emsq = emsq.compute()
eape = ((esqm - emsq) * db * 0.5).compute()
eape.name = r'EAPE (m$^3$s$^{-2}$)'
return eape
def beta_term(fil1, fil2, **initializer):
with pym6.Dataset(fil2,
**initializer) as ds, pym6.Dataset(fil1,
**initializer) as ds1:
vh = ds.vh.final_loc('ql').xep().read().divide_by('dxCv').nanmean(
axis=0).move_to('q').compute()
y = (initializer['south_lat'] + initializer['north_lat']) / 2
f_slice = vh.get_slice_2D()._slice_2D
f = initializer['geometry'].f[f_slice]
bsv = ((-vh).compute() * beta(y) / f).compute()
bsv.math = r'$ -\frac{\beta\bar{\sigma}\hat{v}}{f}$'
bsv.name = r"beta term"
return bsv
def sigmyvhat(fil1, fil2, **kwargs):
htol = kwargs.get('htol', 1e-10)
with pym6.Dataset(
fil1, final_loc='hl', **kwargs) as ds1, pym6.Dataset(
fil2, final_loc='hl', **kwargs) as ds2:
h = ds1.h_Cv.ysm().read().nanmean(axis=0).move_to('h').compute()
vr = ds2.vh.ysm().read().nanmean(
axis=0).divide_by('dxCv').move_to('h').compute()
vr = (vr / h).compute()
vr.values[h.values < htol] = 0
hpfv = (-1 * ds1.twa_hpfv.ysm().read().nanmean(axis=0).move_to('h')
).compute()
return (vr * hpfv).compute()
def sigmxuhat(fil1, fil2, **kwargs):
htol = kwargs.get('htol', 1e-10)
with pym6.Dataset(
fil1, final_loc='hl', **kwargs) as ds1, pym6.Dataset(
fil2, final_loc='hl', **kwargs) as ds2:
h = ds1.h_Cu.xsm().read().nanmean(axis=0).move_to('h').compute()
ur = ds2.uh.xsm().read().nanmean(
axis=0).divide_by('dyCu').move_to('h').compute()
ur = (ur / h).compute()
ur.values[h.values < htol] = 0
hpfu = (-1 * ds1.twa_hpfu.xsm().read().nanmean(axis=0).move_to('h')
).compute()
return (ur * hpfu).compute()
def mape(fil1, fil2):
htol = kwargs.get('htol', 1e-10)
with pym6.Dataset(fil1) as ds1, pym6.Dataset(fil2) as ds2:
db = np.diff(ds1.zl)[0] * 9.8 / 1000
e = ds2.e.final_loc('hl').zep().read().nanmean(
axis=0).move_to('h').compute()
espatialmean = ds2.e.final_loc('hl').zep().read().nanmean(
axis=(0, 2, 3)).movw_to('h').compute()
mape = ((e * e).compute().nanmean(axis=(2, 3)).compute() -
(espatialmean**2).compute()).compute()
mape = (mape * db * 0.5).compute()
mape.name = r'MAPE (m$^3$s$^{-2}$)'
return mape
def get_corpfum(fil1, fil2, **initializer):
initializer['final_loc'] = 'ul'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
hfvm = ds1.twa_hfv.read().nanmean(axis=0).compute()
pfum = ds2.PFu.read().nanmean(axis=0).compute()
pfum = (h * pfum).compute()
corpfum = (hfvm + pfum).multiply_by('dyCu').reduce_(np.nansum,
axis=2).compute()
corpfum.name = 'Ageostrophic term'
corpfum.math = r'$f\bar{\sigma}\hat{v}-\bar{\sigma}\bar{m}_{\tilde{x}}$'
corpfum.units = r'm$^2$s$^{-2}$'
return corpfum
def get_corpfvm(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil2, **initializer) as ds2, pym6.Dataset(
fil1, **initializer) as ds1:
hmfum = ds1.twa_hmfu.read().nanmean(axis=0).compute()
pfvm = ds2.PFv.read().nanmean(axis=0).compute()
pfvm = (h * pfvm).compute()
corpfvm = (hmfum + pfvm).multiply_by('dxCv').reduce_(
np.nansum, axis=3).compute()
corpfvm.name = 'Ageostrophic term'
corpfvm.math = r'$-f\bar{\sigma}\hat{u}-\bar{\sigma}\bar{m}_{\tilde{y}}$'
corpfvm.units = r'm$^2$s$^{-2}$'
return corpfvm
def get_bdivep1(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1:
hvwbm = ds1.twa_hvwb.read().nanmean(axis=0).compute()
bdivep1 = (hvwbm / h).compute()
return bdivep1
def get_sigFD(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
edlsqmy = ty.get_edlsqmy(fil1, fil2, **initializer)
h = ty.get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
hpfv = ds1.twa_hpfv.read().nanmean(axis=0).compute()
pfvm = ds2.PFv.read().nanmean(axis=0).compute()
db = np.diff(ds1.zl)[0] * 9.8 / 1000
bdivep4 = -((-hpfv).compute() + (h * pfvm).compute() -
(edlsqmy * db * 0.5).compute()).compute()
bdivep4.name = 'Form drag'
bdivep4.math = (r"""$-(\overline{\zeta ^\prime m_{\tilde{y}}^\prime})"""
r"""_{\tilde{b}}$""")
bdivep4.units = r'ms$^{-2}$'
return bdivep4
def get_vr(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil2, **initializer) as ds2:
vr = ds2.vh.read().nanmean(axis=0).divide_by('dxCv') / h
vr = vr.compute()
return vr
def wparamsurf(self, eddy, x, y, rangex, rangey):
idx = self.nearest(eddy)
with pym6.Dataset(
self.file_,
east_lon=x + rangex,
west_lon=x - rangex,
south_lat=y - rangey,
north_lat=y + rangey) as ds:
ux = (ds.u.final_loc('ql').geometry(self.geometry).isel(
z=slice(0, 1), t=slice(eddy.time - 2,
eddy.time - 1)).xsm().xep().yep().
read().dbyd(3).move_to('u').move_to('q').to_DataArray())
vy = (ds.v.final_loc('ql').geometry(self.geometry).isel(
z=slice(0, 1), t=slice(eddy.time - 2,
eddy.time - 1)).ysm().yep().xep().
read().dbyd(2).move_to('v').move_to('q').to_DataArray())
uy = (ds.u.final_loc('ql').geometry(self.geometry).isel(
z=slice(0, 1),
t=slice(eddy.time - 2,
eddy.time - 1)).yep().read().dbyd(2).compute())
vx = (ds.v.final_loc('ql').geometry(self.geometry).isel(
z=slice(0, 1),
t=slice(eddy.time - 2,
eddy.time - 1)).xep().read().dbyd(3).compute())
vxuy4 = ((vx * uy).to_DataArray() * 4).squeeze()
W = ((ux - vy)**2 + vxuy4)
return W
def MeridSectionofMOCcomponents(expt, **initializer):
z = initializer.get('z', None)
if 'z' in initializer:
initializer.pop('z')
initializer['final_loc'] = 'ui'
with pym6.Dataset(expt.fil2, **initializer) as ds:
e = ds.e.xep().zep().read().move_to('u').nanmean(axis=(0, 3)).compute()
initializer['final_loc'] = 'ul'
blist = mb.extract_twamomx_terms(expt.fil1, expt.fil2, **initializer)
namelist = []
for i, b in enumerate(blist):
namelist.append(b.name)
#b = b.reduce_(np.sum, axis=2)
if z is not None:
b = b.toz(z, e)
else:
b.values = np.cumsum(b.values, axis=1)
b = b.tob(axis=1)
blist[i] = (-b).to_DataArray()
blist_concat = xr.concat(blist, dim=pd.Index(namelist, name='Term'))
blist_concat.name = 'Merid OC budget'
if z is None:
e = e.tob(axis=1)
e = e.to_DataArray()
return_dict = dict(blist_concat=blist_concat, blist=blist, e=e, swash=None)
return return_dict
def get_epe_flux(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
db = np.diff(ds1.zl)[0] * 9.8 / 1000
esq = ds1.esq.yep().read().nanmean(axis=0)
e = ds2.e.yep().zep().read().nanmean(axis=0).move_to('l')**2
e = e.compute(check_loc=False)
edlsqm = esq - e
edlsqm = edlsqm.move_to('v').compute()
epeflux = (edlsqm / h * 0.5 * db).compute()
ex = (ds2.e.final_loc('vl').zep().yep().xsm().xep().read().nanmean(
axis=0).dbyd(3).move_to('h').move_to('v').move_to('l').compute())
epeflux_ex = (ex * epeflux).compute()
return epeflux, epeflux_ex
def get_KE_budget(fil1, fil2, **initializer):
initializer['final_loc'] = 'hl'
with pym6.Dataset(fil2, **initializer) as ds2:
pe_to_ke = ds2.PE_to_KE.read().nanmean(axis=0).compute()
pe_to_ke.name = 'PE to KE'
pe_to_ke.math = (r'$-(\overline{\sigma u m_{\tilde{x}}}+'
r'\overline{\sigma v m_{\tilde{y}}})$')
ke_adv = ds2.KE_adv.read().nanmean(axis=0).compute()
ke_adv.name = 'KE_adv'
ke_adv.math = (r'$-((\overline{\sigma u KE})_{\tilde{x}}+'
r'(\overline{\sigma v KE})_{\tilde{y}})$')
ke_dia = ds2.KE_dia.read().nanmean(axis=0).compute()
ke_dia.name = 'KE_dia'
ke_dia.math = r'$-(\overline{\sigma \varpi KE})_{\tilde{b}}$'
ke_visc = ds2.KE_visc.read().nanmean(axis=0).compute()
ke_visc.name = 'KE_visc'
ke_visc.math = '$\overline{\sigma u X^V} + \overline{\sigma v Y^V}$'
ke_horvisc = ds2.KE_horvisc.read().nanmean(axis=0).compute()
ke_horvisc.name = 'KE_horvisc'
ke_horvisc.math = '$\overline{\sigma u X^H} + \overline{\sigma v Y^H}$'
KE_list = [pe_to_ke, ke_adv, ke_horvisc, ke_visc, ke_dia]
only = initializer.get('only', range(len(KE_list)))
return_list = [KE_list[i] for i in only]
return return_list
def get_ur(fil1, fil2, **initializer):
initializer['final_loc'] = 'ul'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil2, **initializer) as ds2:
ur = ds2.uh.read().nanmean(axis=0).divide_by('dyCu') / h
ur = ur.compute()
return ur
def extract_momx_terms(fil1, fil2, **initializer):
with pym6.Dataset(fil2, **initializer) as ds2:
cau = ds2.CAu.read().nanmean(axis=0).compute()
cau.name = 'Coriolis term'
cau.math = r'$-f\bar{v}$'
pfu = ds2.PFu.read().nanmean(axis=0).compute()
pfu.name = 'Grad Mont Pot'
pfu.math = r'$-\bar{m_{\tilde{x}}}$'
dudt_dia = ds2.dudt_dia.read().nanmean(axis=0).compute()
dudt_dia.name = 'Diapycnal advection'
dudt_dia.math = r'$-\bar{\varpi} \bar{u}_{\tilde{b}}$'
dudt_visc = ds2.du_dt_visc.read().nanmean(axis=0).compute()
diffu = ds2.diffu.read().nanmean(axis=0).compute()
diffu = (diffu + dudt_visc).compute()
diffu.name = 'Friction terms'
diffu.math = r'$\bar{X}$'
conventional_list = [cau, pfu, dudt_dia, diffu]
only = initializer.get('only', range(len(conventional_list)))
if 'only' in initializer:
return_list = []
initializer.pop('only')
for i, term in enumerate(conventional_list):
if i in only:
return_list.append(term)
else:
return_list = conventional_list
return return_list
def get_bdivep1(fil1, fil2, **initializer):
initializer['final_loc'] = 'ul'
with pym6.Dataset(fil1, **initializer) as ds1:
huwbm = ds1.twa_huwb.read().nanmean(axis=0).compute()
bdivep1 = (huwbm.multiply_by('dyCu').reduce_(np.nansum,
axis=2)).compute()
return bdivep1
def form_drag(fil1, fil2, **initializer):
bc_type = dict(
v=['mirror', 'circsymh', 'zeros', 'circsymq', 'circsymh', 'neumann'])
with pym6.Dataset(fil2,
**initializer) as ds, pym6.Dataset(fil1,
**initializer) as ds1:
h = ds1.h_Cv.bc_type(bc_type).final_loc('ql').xep().read().nanmean(
axis=0).compute(check_loc=False)
bdivep = ty.get_bdivep(fil1, fil2, **initializer).xep().bc_type(
bc_type).implement_BC_if_necessary().multiply_by('dyCv').compute()
term5 = (bdivep * h).dbyd(3, weights='area').compute(check_loc=False)
f_slice = term5.get_slice_2D()._slice_2D
f = initializer['geometry'].f[f_slice]
term5 = (term5 / f).compute(check_loc=False)
term5.math = r'$-\frac{(\overline{\zeta^{\prime}m_{\tilde{y}}^{\prime}})_{\tilde{x}\tilde{b}}}{f}$'
term5.name = r"Form drag"
return term5
def get_xdivep1(fil1, fil2, **initializer):
initializer['final_loc'] = 'ul'
with pym6.Dataset(fil1, **initializer) as ds1:
huuxm = ds1.huu_Cu.xsm().xep().read().dbyd(
3, weights='area').nanmean(axis=0).move_to('u').compute()
xdivep1 = (
-huuxm.multiply_by('dyCu').reduce_(np.nansum, axis=2)).compute()
return xdivep1
def get_ydivep1(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
with pym6.Dataset(fil1, **initializer) as ds1:
hvvym = ds1.hvv_Cv.ysm().yep().read().dbyd(
2, weights='area').nanmean(axis=0).move_to('v').compute()
ydivep1 = (-hvvym).multiply_by('dxCv').reduce_(
np.nansum, axis=3).compute()
return ydivep1
def get_advy(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
htol = initializer.get('htol', 1e-3)
vr = get_vr(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1, pym6.Dataset(
fil2, **initializer) as ds2:
h_fory = ds1.h_Cv.ysm().yep().read().nanmean(axis=0).compute()
a = h_fory.values
a[a < htol] = np.nan
h_fory.values = a
vry = ds2.vh.ysm().yep().read().nanmean(axis=0) / h_fory
vry = vry.dbyd(2, weights='area').move_to('v').compute()
advy = (-vr * vry).compute()
advy.name = 'Meridional advection'
advy.math = r'$-\hat{v}\hat{v}_{\tilde{y}}$'
advy.units = r'ms$^{-2}$'
return advy
def get_ydivep1(fil1, fil2, **initializer):
initializer['final_loc'] = 'vl'
h = get_h(fil1, fil2, **initializer)
with pym6.Dataset(fil1, **initializer) as ds1:
hvvym = ds1.hvv_Cv.ysm().yep().read().dbyd(
2, weights='area').nanmean(axis=0).move_to('v').compute()
ydivep1 = (-hvvym / h).compute()
return ydivep1
