def plot_wf_spectrum_vort(initializer, perc=99):
domain = Domain.Domain(initializer)
with mfdset(initializer.fil) as fh:
e = gv('e', domain, 'hi', fh, plot_loc='qi').xsm().ysm().read_array(
tmean=False,
extend_kwargs=dict(method='mirror')).move_to('ui').move_to('qi')
uy = gv('u', domain, 'ul', fh,
plot_loc='ql').yep().read_array(
tmean=False,
filled=0,
extend_kwargs=dict(method='vorticity')).ddx(2)
vx = gv('v', domain, 'vl', fh,
plot_loc='ql').xep().read_array(
tmean=False,
filled=0,
extend_kwargs=dict(method='vorticity')).ddx(3)
vort = (vx - uy).toz([-1], e=e)
vort.name = ''
fig, ax = plt.subplots(1, 2)
vort.plot('nanmean',
initializer.meanax,
perc=98,
contour=False,
cbar=True,
plot_kwargs=dict(cmap='RdBu_r'),
ax=ax[0])
initializer.ls = 0
initializer.le = None
domain1 = Domain.Domain(initializer)
h = gv('e', domain1, 'hi', fh).read_array().o1diff(1).values
omega = 2 * np.pi * (1 / 24 / 3600 + 1 / 365 / 24 / 3600)
lr = np.sum(
np.sqrt(h * domain.db) / np.pi /
(2 * omega *
np.sin(np.radians(0.5 * (initializer.slat + initializer.nlat)))),
axis=1,
keepdims=True)
lr = np.mean(lr, axis=(0, 1, 2, 3)).squeeze()
sig = vort.mean(axis=2).values.squeeze()
fx, fx1, ft, Rsq = cross_spectrum(sig, 5 * 24 * 3600, 4000)
vmax = np.percentile(Rsq, 99.9)
im = ax[1].pcolormesh(fx1, ft, Rsq, cmap='Reds', vmin=0, vmax=vmax)
fig.colorbar(im, ax=ax[1])
beta = 2 * omega * np.cos(
np.radians(0.5 *
(initializer.slat + initializer.nlat))) / domain.R_earth
w = -beta * fx / (fx**2 + lr**-2)
ax[1].plot(-fx, -w, 'k')
return fig, ax
def plot_wf_spectrum(initializer, perc=99):
domain = Domain.Domain(initializer)
with mfdset(initializer.fil) as fh:
e = gv('e', domain, 'hi', fh).read_array(tmean='anom')
fig, ax = plt.subplots(1, 2)
e.plot('nanmean',
initializer.meanax,
ax=ax[0],
cbar=True,
contour=False,
plot_kwargs=dict(cmap='RdBu_r'))
initializer.ls = 0
initializer.le = None
domain1 = Domain.Domain(initializer)
h = gv('e', domain1, 'hi', fh).read_array().o1diff(1).values
omega = 2 * np.pi * (1 / 24 / 3600 + 1 / 365 / 24 / 3600)
lr = np.sum(
np.sqrt(h * domain.db) / np.pi /
(2 * omega *
np.sin(np.radians(0.5 * (initializer.slat + initializer.nlat)))),
axis=1,
keepdims=True)
lr = np.mean(lr, axis=(0, 1, 2, 3)).squeeze()
sig = e.mean(axis=2).values.squeeze()
fx, fx1, ft, Rsq = cross_spectrum(sig, 5 * 24 * 3600, 4000)
vmax = np.percentile(Rsq, 99.9)
im = ax[1].pcolormesh(fx1, ft, Rsq, cmap='Reds', vmin=0, vmax=vmax)
fig.colorbar(im, ax=ax[1])
beta = 2 * omega * np.cos(
np.radians(0.5 *
(initializer.slat + initializer.nlat))) / domain.R_earth
w = -beta * fx / (fx**2 + lr**-2)
ax[1].plot(-fx, -w, 'k')
return fig, ax
def extract_cb_terms_pym6(initializer):
domain = Domain.Domain(initializer)
with mfdset(initializer.fil) as fh, mfdset(initializer.fil2) as fh2:
vhy = gv('vh',domain,'vl',fh2,fh,plot_loc='hl').ysm().read_array(filled=0).ddx(2,div_by_area=True)
uhx = gv('uh',domain,'ul',fh2,fh,plot_loc='hl').xsm().read_array(filled=0).ddx(3,div_by_area=True)
wd = gv('wd',domain,'hi',fh2).read_array().o1diff(1)
# budgetlist = [-uhx*(1/domain.dyT[uhx._slice[2:]]),-vhy*(1/domain.dxT[vhy._slice[2:]]),-wd]
budgetlist = [-uhx,-vhy,-wd]
lab = [ r'$-(\bar{h}\hat{u})_{\tilde{x}}$',
r'$-(\bar{h}\hat{v})_{\tilde{y}}$',
r'$-(\bar{h}\hat{\varpi})_{\tilde{b}}$']
for i,var in enumerate(budgetlist):
var.name = lab[i]
return budgetlist
def plot_sst_vort(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
ls,
le,
fil3=None,
z=np.linspace(-10, -9, 4),
htol=1e-3,
whichterms=None):
domain = Domain.Domain2(geofil,
vgeofil,
xstart,
xend,
ystart,
yend,
ls=ls,
le=le,
ts=450,
te=451)
with mfdset(fil2) as fh2:
rho = gv('e', domain, 'hi', fh2).read_array(tmean=False).toz(z[0],
rho=True)
T = (-rho + 1031) * (1 / 0.2)
T.name = ''
fig, ax = plt.subplots(1, 2, sharex=True, sharey=True)
T.plot('nanmean', (0, 1),
contour=False,
cbar=True,
plot_kwargs=dict(cmap='RdYlBu_r', vmax=20, vmin=10),
ax=ax[0])
rho = gv('e', domain, 'hi', fh2).read_array(tmean=False).toz(z[1],
rho=True)
T = (-rho + 1031) * (1 / 0.2)
T.name = ''
T.plot('nanmean', (0, 1),
contour=False,
cbar=True,
plot_kwargs=dict(cmap='RdYlBu_r', vmax=10, vmin=6),
ax=ax[1])
# e = gv('e',domain,'hi',fh2,plot_loc='qi').xsm().ysm().read_array(tmean=False,
# extend_kwargs=dict(method='mirror')).move_to('ui').move_to('qi')
# uy = gv('u',domain,'ul',fh2,plot_loc='ql').yep().read_array(tmean=False,
# extend_kwargs=dict(method='vorticity')).ddx(2)
# vx = gv('v',domain,'vl',fh2,plot_loc='ql').xep().read_array(tmean=False,
# extend_kwargs=dict(method='vorticity')).ddx(3)
# vort = (vx - uy)
# vort.name = ''
# vort.plot('nanmean',(0,1),perc=98,contour=False,cbar=True,plot_kwargs=dict(cmap='RdBu_r'),ax=ax[1])
# ax[1].set_ylabel('')
#
# vortz = vort.toz([-1],e=e)
# vortz.name = ''
# vortz.plot('nanmean',(0,1),perc=98,contour=False,cbar=True,plot_kwargs=dict(cmap='RdBu_r'),ax=ax[2])
# vmax = np.nanpercentile(np.fabs(vortz.values),98)
# xx,yy = np.meshgrid(vortz.dom.lonq[vortz._plot_slice[3,0]:vortz._plot_slice[3,1]],
# vortz.dom.latq[vortz._plot_slice[2,0]:vortz._plot_slice[2,1]])
# print(xx.shape, yy.shape, vortz.values[0,0,:,:].shape)
# im = ax[2].pcolormesh(xx,yy,
# vortz.values[0,0,:,:],cmap='RdBu_r',vmin=-vmax,vmax=vmax)
# print(e.values.shape)
# im = ax[2].pcolormesh(e.dom.lonq[e._plot_slice[3,0]:e._plot_slice[3,1]],
# e.dom.latq[e._plot_slice[2,0]:e._plot_slice[2,1]],
# e.values[0,0,:,:],cmap='RdBu_r')
# cbar = fig.colorbar(im,ax=ax[2])
# cbar.formatter.set_powerlimits((-2,2))
# cbar.update_ticks()
# ax[2].set_ylabel('')
domain = Domain.Domain2(geofil,
vgeofil,
xstart,
xend,
ystart,
yend,
ls=0,
le=1)
with mfdset(fil) as fh:
e = gv('e', domain, 'hi', fh).read_array()
cs = ax[0].contour(e.dom.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
e.dom.lath[e._plot_slice[2, 0]:e._plot_slice[2, 1]],
e.values.squeeze(),
levels=[-0.4, -0.3, -0.15, 0, 0.15, 0.3, 0.4],
colors='k')
cs.clabel(inline=True)
return fig
def plot_twamomy_pym6(initializer, **kwargs):
budgetlist = extract_twamomy_terms_pym6(initializer)
terms = kwargs.get('terms', None)
if terms:
budgetlist = [budgetlist[i] for i in terms]
with mfdset(initializer.fil) as fh, mfdset(initializer.fil2) as fh2:
e = (gv('e', budgetlist[0].dom, 'hi', fh2, fh,
plot_loc='vl').yep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('vi'))
plot_kwargs = dict(cmap='RdBu_r')
plotter_kwargs = dict(xtokm=True,
zcoord=True,
z=initializer.z,
e=e,
isop_mean=True)
ax, fig = Plotter.budget_plot(budgetlist,
initializer.meanax,
plot_kwargs=plot_kwargs,
plotter_kwargs=plotter_kwargs,
individual_cbars=False,
**kwargs)
swash = kwargs.get('swash', True)
plotter_kwargs.pop('ax')
if swash:
initializer_for_swash = Domain.Initializer(geofil=initializer.geofil,
vgeofil=initializer.vgeofil,
wlon=-0.5,
elon=0,
slat=10,
nlat=11)
domain_for_swash0 = Domain.Domain(initializer_for_swash)
with mfdset(initializer.fil3) as fh:
islayerdeep = (gv('islayerdeep',
budgetlist[0].dom,
'ql',
fh,
plot_loc='vl').xsm().read_array(tmean=False,
extend_kwargs={
'method':
'mirror'
}).move_to('vl'))
islayerdeep0 = (gv('islayerdeep',
domain_for_swash0,
'ql',
fh,
plot_loc='vl').xsm().read_array(
tmean=False,
extend_kwargs={
'method': 'mirror'
}).move_to('vl'))
swashperc = (-islayerdeep / islayerdeep0.values[-1, 0, 0, 0] +
1) * 100
plot_kwargs = kwargs.get('swash_plot_kwargs',
dict(colors='grey', linewidths=4))
for axc in ax.ravel():
axc, _ = swashperc.plot('nanmean',
initializer.meanax,
only_contour=True,
clevs=np.array([1]),
annotate='None',
ax=axc,
plot_kwargs=plot_kwargs,
**plotter_kwargs)
return ax, fig
def extract_twamomy_terms_pym6(initializer):
domain = Domain.Domain(initializer)
plot_loc = 'vl'
with mfdset(initializer.fil) as fh, mfdset(initializer.fil2) as fh2:
h = (gv('h_Cv', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0))
vr = (gv('vh',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc,
divisor='h_Cv').read_array(divide_by_dx=True, filled=0))
vrx = (gv('vh',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc,
divisor='h_Cv').xsm().xep().read_array(
extend_kwargs={
'method': 'vorticity'
},
divide_by_dx=True,
filled=0).ddx(3).move_to(plot_loc))
vry = (gv('vh',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc,
divisor='h_Cv').ysm().yep().read_array(
extend_kwargs={
'method': 'symmetric'
},
divide_by_dx=True,
filled=0).ddx(2).move_to(plot_loc))
humx = (gv('uh', domain, 'ul', fh2, fh,
plot_loc=plot_loc).xsm().yep().read_array(
extend_kwargs={
'method': 'vorticity'
}, filled=0).ddx(3, div_by_area=True).move_to(plot_loc))
hum = (gv('uh', domain, 'ul', fh2, fh,
plot_loc=plot_loc).xsm().yep().read_array(
divide_by_dy=True,
filled=0,
extend_kwargs={
'method': 'vorticity'
}).move_to('hl').move_to(plot_loc))
hvmy = (gv('vh', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).ysm().yep().read_array(
filled=0, extend_kwargs={
'method': 'symmetric'
}).ddx(2, div_by_area=True).move_to(plot_loc))
huvxphvvym = (
(gv('twa_huvxpt', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0)) +
(gv('twa_hvvymt', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0)))
hvvym = (gv('hvv_Cv', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).ysm().yep().read_array(
extend_kwargs={
'method': 'symmetric'
}, filled=0).ddx(2,
div_by_area=True).move_to(plot_loc))
huvxm = -(huvxphvvym + hvvym)
vrb = (gv('vh',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc,
divisor='h_Cv').lsm().lep().read_array(
extend_kwargs={
'method': 'mirror'
},
divide_by_dx=True,
filled=0).ddx(1).move_to(plot_loc))
hwb = (gv('wd', domain, 'hi', fh2, fh,
plot_loc=plot_loc).yep().read_array(
filled=0, extend_kwargs={
'method': 'vorticity'
}).o1diff(1).move_to(plot_loc))
hwm = (gv('wd', domain, 'hi', fh2, fh,
plot_loc=plot_loc).yep().read_array(
filled=0, extend_kwargs={
'method': 'vorticity'
}).move_to('hl').move_to(plot_loc)) * domain.db
esq = (gv('esq', domain, 'hl', fh2, fh,
plot_loc=plot_loc).yep().read_array(
extend_kwargs={'method': 'mirror'}))
e = (gv('e', domain, 'hi', fh2, fh,
plot_loc=plot_loc).yep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('hl'))
edlsqm = esq - e * e
edlsqmy = edlsqm.ddx(2)
hpfv = (gv('twa_hpfv', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0))
pfvm = (gv('PFv', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0))
edpfvdmb = -hpfv + h * pfvm - edlsqmy * domain.db * 0.5
hmfum = (gv('twa_hmfu', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0))
hvwbm = (gv('twa_hvwb', domain, plot_loc, fh2, fh,
plot_loc=plot_loc).read_array(filled=0))
hdiffvm = (gv('twa_hdiffv',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc).read_array(filled=0))
hdvdtviscm = (gv('twa_hdvdtvisc',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc).read_array(filled=0))
advx = hum * vrx / h
advy = vr * vry
advb = hwm * vrb / h
cor = hmfum / h
pfvm = pfvm
xdivep1 = -huvxm / h
xdivep2 = advx
xdivep3 = vr * humx / h
xdivep = (xdivep1 + xdivep2 + xdivep3)
ydivep1 = -hvvym / h
ydivep2 = advy
ydivep3 = vr * hvmy / h
ydivep4 = -edlsqmy / 2 * domain.db / h
ydivep = (ydivep1 + ydivep2 + ydivep3 + ydivep4)
bdivep1 = hvwbm / h
bdivep2 = advb
bdivep3 = vr * hwb / h
bdivep4 = -edpfvdmb / h
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
Y1twa = hdiffvm / h
Y2twa = hdvdtviscm / h
budgetlist = [
-advx, -advy, -advb, cor, pfvm, xdivep, ydivep1 + ydivep2 + ydivep3,
ydivep4, bdivep1 + bdivep2 + bdivep3, bdivep4, Y1twa, Y2twa
]
ti = [
'(a)', '(b)', '(c)', '(d)', '(e)', '(f)', '(g)', '(h)', '(i)', '(j)',
'(k)', '(l)'
]
lab = [
r'$-\hat{u}\hat{v}_{\tilde{x}}$', r'$-\hat{v}\hat{v}_{\tilde{y}}$',
r'$-\hat{\varpi}\hat{v}_{\tilde{b}}$', r'$-f\hat{u}$',
r'$-\overline{m_{\tilde{y}}}$',
r"""$-\frac{1}{\overline{h}}(\overline{h}\widehat{u^{\prime \prime}v^{\prime \prime}})_{\tilde{x}}$""",
r"""$-\frac{1}{\overline{h}}(\overline{h}\widehat{v^{\prime \prime}v^{\prime \prime}})_{\tilde{y}}$""",
r"""$-\frac{1}{2\overline{\zeta_{\tilde{b}}}}(\overline{\zeta^{\prime 2}})_{\tilde{y}}$""",
r"""$-\frac{1}{\overline{h}}(\overline{h}\widehat{v^{\prime \prime}\varpi ^{\prime \prime}})_{\tilde{b}}$""",
r"""$-\frac{1}{\overline{\zeta_{\tilde{b}}}}(\overline{\zeta^\prime m_{\tilde{y}}^\prime})_{\tilde{b}}$""",
r'$\widehat{Y^H}$', r'$\widehat{Y^V}$'
]
for i, var in enumerate(budgetlist):
var.name = lab[i]
return budgetlist
def plot_forcing(fil, fil2=None):
y = np.linspace(10, 60, 100)
sflat = 20
nflat = 40
T = np.zeros(y.shape)
Tnorth = 10
Tsouth = 20
T[y >= sflat] = (Tsouth - Tnorth) / (sflat - nflat) * (y[y >= sflat] -
nflat) + Tnorth
T[y < sflat] = Tsouth
T[y > nflat] = Tnorth
rho0 = 1031
drhodt = -0.2
rho = rho0 + drhodt * T
gs = gridspec.GridSpec(2, 3, width_ratios=[1, 4, 1], height_ratios=[1, 4])
fig = plt.figure(figsize=(9, 4))
#ax = fig.add_subplot(121)
ax = plt.subplot(gs[1])
ax.plot(y, T, 'k', lw=1)
# ax.set_ylabel('T ($^{\circ}$C)')
ax.set_xlabel('y ($^{\circ}$N)')
ax.set_ylim(8, 22)
ax.grid(True)
ax.xaxis.tick_top()
ax.xaxis.set_label_position("top")
ax2 = ax.twinx()
ax2.set_ylabel(r"""$\rho (kgm^{-3})$""")
ax2.set_ylim(ax.get_ylim())
yticks = ax.get_yticks()
ax2.set_yticks(yticks[::2])
yticks = rho0 + drhodt * yticks[::2]
ax2.set_yticklabels(['{}'.format(i) for i in yticks])
z = np.linspace(0, 3000)
dabyss = 1500
Tabyss = 5
Tz = (Tabyss - Tsouth) / dabyss * z + Tsouth
Tz[z > dabyss] = Tabyss
#ax = fig.add_subplot(122)
ax = plt.subplot(gs[3])
ax.plot(Tz, z, 'k', lw=1)
ax.set_xlabel('T ($^{\circ}$C)')
ax.set_ylabel('z (m)')
ax.set_xlim(4, 21)
ax.grid(True)
ax.invert_yaxis()
# ax.yaxis.tick_right()
# ax.yaxis.set_label_position("right")
# ax2 = ax.twiny()
# ax2.set_xlabel(r"""$\rho (kgm^{-3})$""")
# ax2.set_xlim(ax.get_xlim())
# xticks = ax.get_xticks()
# ax2.set_xticks(xticks[::2])
# xticks = rho0 + drhodt*xticks[::2]
# ax2.set_xticklabels(['{}'.format(i) for i in xticks])
with mfdset(fil) as fh, mfdset(fil2) as fh2:
geofil = 'ocean_geometry.nc'
vgeofil = 'Vertical_coordinate.nc'
domain = Domain.Domain2(geofil,
vgeofil,
-0.5,
0,
10,
60,
ls=0,
le=None)
e = gv('e', domain, 'hi', fh).read_array()
ax = plt.subplot(gs[4])
ax.plot(domain.lath[e.plot_slice[2, 0]:e.plot_slice[2, 1]],
np.mean(e.values, axis=(0, 3)).T[:, ::2],
'k',
lw=1)
ax.plot(domain.lath[e.plot_slice[2, 0]:e.plot_slice[2, 1]],
np.mean(e.values, axis=(0, 3)).T[:, -1],
'k',
lw=1)
ax.set_yticklabels('')
ax.set_xlabel('y ($^{\circ}$N)')
ax.axvline(x=38.5, color='k')
ax.grid()
swash = gv('islayerdeep', domain, 'ql', fh2,
plot_loc='hl').xsm().ysm().read_array(
extend_kwargs=dict(method='mirror'),
tmean=False).move_to('ul').move_to('hl')
swash0 = fh2.variables['islayerdeep'][-1, 0, 0, 320]
swash = (-swash + swash0) * (100 / swash0)
z = np.linspace(-3000, 0)
swash = swash.toz(z, e=e)
ax.contour(swash.dom.lath[e.plot_slice[2, 0]:e.plot_slice[2, 1]],
z,
np.nanmean(swash.values, axis=(0, 3)),
levels=[1],
colors='r')
domain = Domain.Domain2(geofil,
vgeofil,
-0.5,
0,
38,
39,
ls=0,
le=None)
ax = plt.subplot(gs[5])
e = gv('e', domain, 'hi', fh).read_array()
ax.plot(domain.lonh[e.plot_slice[3, 0]:e.plot_slice[3, 1]],
np.mean(e.values, axis=(0, 2)).T[:, ::2],
'k',
lw=1)
ax.plot(domain.lonh[e.plot_slice[3, 0]:e.plot_slice[3, 1]],
np.mean(e.values, axis=(0, 2)).T[:, -1],
'k',
lw=1)
ax.set_yticklabels('')
ax.set_xlabel('x ($^{\circ}$)')
swash = gv('islayerdeep', domain, 'ql', fh2,
plot_loc='hl').xsm().ysm().read_array(
tmean=False).move_to('ul').move_to('hl')
swash0 = fh2.variables['islayerdeep'][-1, 0, 0, 320]
swash = (-swash + swash0) * (100 / swash0)
z = np.linspace(-3000, 0)
swash = swash.toz(z, e=e)
ax.contour(swash.dom.lonh[swash.plot_slice[3, 0]:swash.plot_slice[3,
1]],
z,
np.nanmean(swash.values, axis=(0, 2)),
levels=[1],
colors='r')
xdegtokm(ax, (38 + 39) / 2)
ax.grid()
#swash.plot('nanmean',(0,2),zcoord=True,e=e,z=np.linspace(-3000,0),cbar=True,ax=ax)
return fig
def extract_ep_terms_pym6(initializer):
domain = Domain.Domain(initializer)
plot_loc = 'vl'
with mfdset(initializer.fil) as fh, mfdset(initializer.fil2) as fh2:
#initializer.z = np.linspace(-500, 0, 20)
e = (gv('e', domain, 'hi', fh2, fh, plot_loc=plot_loc)
.yep().read_array(
extend_kwargs={'method': 'mirror'}).move_to('vi'))
h = (gv('h_Cv', domain, plot_loc, fh2, fh, plot_loc=plot_loc)
.read_array(filled=0))
vr = (gv('vh',
domain,
plot_loc,
fh2,
fh,
plot_loc=plot_loc,
divisor='h_Cv').read_array(
divide_by_dx=True, filled=0))
hum = (gv('uh', domain, 'ul', fh2, fh, plot_loc=plot_loc).xsm().yep()
.read_array(
divide_by_dy=True,
filled=0,
extend_kwargs={'method': 'vorticity'})
.move_to('hl').move_to(plot_loc))
huvm = (gv('huv_Bu', domain, 'ql', fh2, fh, plot_loc=plot_loc).xsm()
.read_array(
filled=0,
extend_kwargs={'method': 'vorticity'}).move_to(plot_loc))
uv_e_flux = (huvm - vr * hum) / h
uv_e_flux = uv_e_flux.mean(axis=initializer.meanax).toz(initializer.z,
e=e)
ex = (gv('e', domain, 'hi', fh2, fh, plot_loc=plot_loc)
.xsm().xep().yep().read_array(extend_kwargs={'method': 'mirror'})
.ddx(3).move_to('hi').move_to('vi').move_to('vl'))
ex = ex.mean(axis=initializer.meanax).toz(initializer.z, e=e)
e1 = (gv('e', domain, 'hi', fh2, fh, plot_loc=plot_loc)
.yep().read_array(extend_kwargs={'method': 'mirror'})
.move_to('hl'))
esq = (gv('esq', domain, 'hl', fh2, fh, plot_loc=plot_loc)
.yep().read_array(extend_kwargs={'method': 'mirror'}))
edlsqm = esq - e1 * e1
edlsqmy = edlsqm.ddx(2)
hpfv = (gv('twa_hpfv', domain, plot_loc, fh2, fh, plot_loc=plot_loc)
.read_array(filled=0))
pfvm = (gv('PFv', domain, plot_loc, fh2, fh, plot_loc=plot_loc)
.read_array(filled=0))
edpfvdmb = -hpfv + h * pfvm - edlsqmy * domain.db * 0.5
edpfvdmb1 = copy.copy(edpfvdmb)
edpfvdmb1 = edpfvdmb1.mean(axis=initializer.meanax).toz(initializer.z,
e=e)
edpfvdmb = edpfvdmb / domain.db
edpfvdm = edpfvdmb.vert_integral()
edpfvdm = edpfvdm.mean(axis=initializer.meanax).toz(initializer.z, e=e)
flux_list = [uv_e_flux, edpfvdm]
lab = [
r"""$\widehat{u^{\prime \prime}v^{\prime \prime}}(\hat{i} + \bar{\zeta_{\tilde{x}}}\hat{k})$""",
r"""$\overline{\zeta^\prime m_{\tilde{y}}^\prime} \hat{k}$""",
r"""$\widehat{u^{\prime \prime}v^{\prime \prime}}(\hat{i} + \bar{\zeta_{\tilde{x}}}\hat{k}) +\overline{\zeta^\prime m_{\tilde{y}}^\prime} \hat{k}$"""
]
dz = np.diff(initializer.z)[0]
dx = np.radians(
np.diff(domain.lonh[ex._plot_slice[3, 0]:ex._plot_slice[3, 1]])[0]
) * 6378000 * np.cos(
np.radians(0.5 * (initializer.slat + initializer.nlat)))
fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(10, 5))
units = 'xy'
angles = 'xy'
q1 = ax[0].quiver(
domain.lonh[ex._plot_slice[3, 0]:ex._plot_slice[3, 1]],
initializer.z[::2],
uv_e_flux.values.squeeze()[::2] / dx * np.sqrt(dx**2 + dz**2),
uv_e_flux.values.squeeze()[::2] * ex.values.squeeze()[::2] / dz *
np.sqrt(dx**2 + dz**2),
units=units)
# angles=angles)
tx = ax[0].text(0.05, 0.1, lab[0], transform=ax[0].transAxes)
tx.set_fontsize(15)
tx.set_bbox(dict(facecolor='white', alpha=1, edgecolor='white'))
ax[0].quiverkey(
q1, 0.75, 1.05, 1e-2, r'$1e-2$', labelpos='E', coordinates='axes')
xdegtokm(ax[0], 0.5 * (initializer.slat + initializer.nlat))
ax[0].set_ylabel('z (m)')
q2 = ax[1].quiver(
domain.lonh[ex._plot_slice[3, 0]:ex._plot_slice[3, 1]],
initializer.z[::2],
np.zeros(edpfvdm.values.squeeze()[::2].shape),
edpfvdm.values.squeeze()[::2] / dx * np.sqrt(dx**2 + dz**2),
units=units)
# angles=angles)
tx = ax[1].text(0.05, 0.1, lab[1], transform=ax[1].transAxes)
tx.set_fontsize(15)
tx.set_bbox(dict(facecolor='white', alpha=1, edgecolor='white'))
ax[1].quiverkey(
q2, 0.75, 1.05, 1e-3, r'$1e-3$', labelpos='E', coordinates='axes')
xdegtokm(ax[1], 0.5 * (initializer.slat + initializer.nlat))
q3 = ax[2].quiver(
domain.lonh[ex._plot_slice[3, 0]:ex._plot_slice[3, 1]],
initializer.z[::2],
np.zeros(edpfvdm.values.squeeze()[::2].shape) +
uv_e_flux.values.squeeze()[::2] / dx * np.sqrt(dx**2 + dz**2),
(edpfvdm.values.squeeze()[::2] + uv_e_flux.values.squeeze()[::2] *
ex.values.squeeze()[::2]) / dz * np.sqrt(dx**2 + dz**2),
units=units)
# angles=angles)
tx = ax[2].text(0.05, 0.1, lab[2], transform=ax[2].transAxes)
tx.set_fontsize(15)
tx.set_bbox(dict(facecolor='white', alpha=1, edgecolor='white'))
ax[2].quiverkey(
q3, 0.75, 1.05, 1e-2, r'$1e-2$', labelpos='E', coordinates='axes')
xdegtokm(ax[2], 0.5 * (initializer.slat + initializer.nlat))
return fig
def extract_buoy_terms_pym6(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
ls,
le,
fil3=None,
z=np.linspace(-3000, 0, 100),
htol=1e-3,
whichterms=None):
domain = Domain.Domain(geofil,
vgeofil,
xstart,
xend,
ystart,
yend,
ls=ls,
le=le)
with mfdset(fil) as fh, mfdset(fil2) as fh2:
h = -gv('e', domain, 'hi', fh2,
units='m').read_array().o1diff(1).values
h[h < htol] = np.nan
sigma = gv('e', domain, 'hi', fh2).read_array().ddx(1).values
ex = gv('e', domain, 'hi',
fh2).xsm().xep().read_array(extend_kwargs={
'method': 'mirror'
}).ddx(3).move_to('ul')
u = gv('uh', domain, 'ul', fh2, fh, plot_loc='hl',
divisor='h_Cu').xsm().read_array(divide_by_dy=True)
uex = u * ex
uex = uex.move_to('hl') #.values
ubx = -uex * (1 / sigma)
ey = gv('e', domain, 'hi',
fh2).ysm().yep().read_array().ddx(2).move_to('vl')
v = gv('vh', domain, 'vl', fh2, fh, plot_loc='hl',
divisor='h_Cv').ysm().read_array(divide_by_dx=True)
vey = v * ey
vey = vey.move_to('hl') #.values
vby = -vey * (1 / sigma)
wd = gv('wd', domain, 'hi', fh2).read_array().move_to('hl') #.values
whash = wd + uex + vey
wbz = whash * (1 / sigma)
hwb = -gv('wd', domain, 'hi', fh2).read_array().o1diff(1) * domain.db
wtwa = hwb * (1 / h)
budgetlist = [-ubx, -vby, -wbz, -wtwa]
name = [
r'$-\hat{u}b^{\#}_x$', r'$-\hat{v}b^{\#}_y$', r'$-w^{\#}b^{\#}_z$',
r'$\hat{\varpi}$'
]
if whichterms:
budgetlist = budgetlist[whichterms]
name = name[whichterms]
for i, var in enumerate(budgetlist):
var.units = r'$ms^{-3}$'
var.name = name[i]
z = np.linspace(-2000, 0)
e = gv('e', domain, 'hi', fh2).read_array()
plot_kwargs = dict(cmap='RdBu_r')
plotter_kwargs = dict(zcoord=True, z=z, e=e, isop_mean=True)
fig = Plotter.budget_plot(budgetlist, (0, 2),
plot_kwargs=plot_kwargs,
plotter_kwargs=plotter_kwargs,
perc=90,
individual_cbars=False)
return fig
def plot_eddy_velocities(initializer, ts=0, te=1, z=None):
if z:
initializer.z = z
domain = Domain.Domain(initializer)
initializer_instantaneous = copy.copy(initializer)
initializer_instantaneous.ts = ts
initializer_instantaneous.te = te
domain1 = Domain.Domain(initializer_instantaneous)
with mfdset(initializer.fil) as fh, mfdset(
initializer.fil2) as fh2, mfdset(initializer.fil3) as fh3:
e_cu = (gv('e', domain, 'hi', fh2, fh,
plot_loc='ui').xep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('ui'))
ur = gv('uh_masked',
domain,
'ul',
fh2,
fh,
plot_loc='ul',
divisor='h_Cu',
name=r'$\hat{u}$').read_array(divide_by_dy=True, filled=0)
u = gv('u', domain1, 'ul', fh3, plot_loc='ul',
name=r'$u$').read_array(filled=0, tmean=False)
udd = u - ur
udd = udd.toz(initializer.z, e=e_cu)
udd.name = r"$u''$"
u = u.toz(initializer.z, e=e_cu)
um = gv('u', domain, 'ul', fh2, fh, plot_loc='ul',
name=r'$\bar{u}$').read_array(filled=0).toz(initializer.z,
e=e_cu)
ud = u - um
ud.name = r"$u'$"
e_cv = (gv('e', domain, 'hi', fh2, fh,
plot_loc='vi').yep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('vi'))
vr = gv('vh_masked',
domain,
'vl',
fh2,
fh,
plot_loc='vl',
divisor='h_Cv',
name=r'$\hat{v}$').read_array(divide_by_dy=True, filled=0)
v = gv('v', domain1, 'vl', fh3, plot_loc='vl',
name=r'$v$').read_array(filled=0, tmean=False)
vdd = v - vr
vdd = vdd.toz(initializer.z, e=e_cv)
vdd.name = r"$v''$"
v = v.toz(initializer.z, e=e_cv)
vm = gv('v', domain, 'vl', fh2, fh, plot_loc='vl',
name=r'$\bar{v}$').read_array(filled=0).toz(initializer.z,
e=e_cv)
vd = v - vm
vd.name = r"$v'$"
budgetlist = [u, ud, udd, v, vd, vdd]
ax, fig = Plotter.budget_plot(budgetlist,
initializer.meanax,
ncols=3,
plot_kwargs=dict(cmap='RdBu_r'),
individual_cbars=False)
# fig, ax = plt.subplots(2, 3, sharex=True, sharey=True)
# u.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[0, 0],
# cmap='RdBu_r',
# cbar=True)
# ud.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[0, 1],
# cmap='RdBu_r',
# cbar=True)
# udd.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[0, 2],
# cmap='RdBu_r',
# cbar=True)
# v.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[1, 0],
# cmap='RdBu_r',
# cbar=True)
# vd.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[1, 1],
# cmap='RdBu_r',
# cbar=True)
# vdd.plot(
# 'nanmean',
# initializer.meanax,
# ax=ax[1, 2],
# cmap='RdBu_r',
# cbar=True)
return fig
def extract_velocities(initializer):
htol = initializer.htol #1e-3
z = initializer.z
meanax = initializer.meanax
domain = Domain.Domain(initializer)
fig, ax = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(8, 8))
ax = ax.ravel()
with mfdset(initializer.fil) as fh, mfdset(initializer.fil2) as fh2:
h = -gv('e', domain, 'hi', fh2,
units='m').read_array().o1diff(1).values
h[h < htol] = np.nan
sigma = gv('e', domain, 'hi', fh2).read_array().ddx(1).values
e = gv('e', domain, 'hi',
fh2).xsm().xep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('ul')
ur = gv('uh',
domain,
'ul',
fh2,
fh,
plot_loc='hl',
divisor='h_Cu',
name=r'$\hat{u}$').xsm().read_array(divide_by_dy=True,
filled=0)
u = gv('u', domain, 'ul', fh2, fh, plot_loc='hl',
name=r'$\bar{u}$').xsm().read_array(filled=0)
ub = ur - u
ub.name = r'$\frac{\overline{u^{\prime}h^{\prime}}}{\bar{h}}$'
ax[2], im = ur.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
plot_kwargs=dict(cmap='RdBu_r'),
clevs=[-0.01, 0.01],
contour=True,
ax=ax[2],
perc=99)
ax[2].plot(domain.lonq[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
vmin, vmax = im.get_clim()
u.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.01, 0.01],
plot_kwargs=dict(cmap='RdBu_r', vmin=vmin, vmax=vmax),
cbar=False,
contour=True,
ax=ax[0],
perc=99)
ax[0].plot(domain.lonq[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
ub.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.01, 0.01],
plot_kwargs=dict(cmap='RdBu_r', vmin=vmin, vmax=vmax),
cbar=False,
contour=True,
ax=ax[1],
perc=99)
ax[1].plot(domain.lonq[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
cbar = fig.colorbar(im, ax=ax[:3].tolist())
cbar.formatter.set_powerlimits((-2, 2))
cbar.update_ticks()
ex = gv('e', domain, 'hi',
fh2).xsm().xep().read_array(extend_kwargs={
'method': 'mirror'
}).ddx(3).move_to('ul')
uex = ur * ex
uex = uex.move_to('hl') #.values
e = gv('e', domain, 'hi',
fh2).ysm().yep().read_array(extend_kwargs={
'method': 'mirror'
}).move_to('vl')
vr = gv('vh',
domain,
'vl',
fh2,
fh,
plot_loc='hl',
divisor='h_Cv',
name=r'$\hat{v}$').ysm().read_array(divide_by_dx=True,
filled=0)
v = gv('v', domain, 'vl', fh2, fh, plot_loc='hl',
name=r'$\bar{v}$').ysm().read_array(filled=0)
vb = vr - v
vb.name = r'$\frac{\overline{v^{\prime}h^{\prime}}}{\bar{h}}$'
ax[5], im = vr.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
plot_kwargs=dict(cmap='RdBu_r'),
clevs=[-0.12, 0.12],
contour=True,
ax=ax[5],
perc=99)
vmin, vmax = im.get_clim()
ax[5].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
v.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.12, 0.12],
plot_kwargs=dict(cmap='RdBu_r', vmin=vmin, vmax=vmax),
contour=True,
ax=ax[3],
perc=99)
ax[3].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
vb.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.12, 0.12],
plot_kwargs=dict(cmap='RdBu_r'),
contour=True,
ax=ax[4],
perc=99)
ax[4].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
cbar = fig.colorbar(im, ax=ax[3:6].tolist())
cbar.formatter.set_powerlimits((-2, 2))
cbar.update_ticks()
ey = gv('e', domain, 'hi',
fh2).ysm().yep().read_array().ddx(2).move_to('vl')
vey = vr * ey
vey = vey.move_to('hl') #.values
e = gv('e', domain, 'hi', fh2).read_array() #.values
wd = gv('wd',
domain,
'hi',
fh2,
name=r'$\hat{\varpi}\bar{\zeta}_{\tilde{b}}$').read_array(
).move_to('hl') #.values
conv = uex + vey
conv.name = r'$\hat{u}\bar{\zeta}_{\tilde{x}}+\hat{v}\bar{\zeta}_{\tilde{y}}$'
whash = wd + conv
whash.name = r'$w^{\#}$'
ax[8], im = whash.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
plot_kwargs=dict(cmap='RdBu_r'),
clevs=[-0.0001, 0.0001],
contour=True,
fmt='%.0e',
ax=ax[8],
perc=99)
vmin, vmax = im.get_clim()
ax[8].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
wd.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.0001, 0.0001],
plot_kwargs=dict(cmap='RdBu_r', vmin=vmin, vmax=vmax),
contour=True,
fmt='%.0e',
ax=ax[7],
perc=99)
ax[7].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
conv.plot('nanmean',
meanax,
zcoord=True,
z=z,
e=e,
clevs=[-0.0001, 0.0001],
plot_kwargs=dict(cmap='RdBu_r', vmin=vmin, vmax=vmax),
contour=True,
fmt='%.0e',
ax=ax[6],
perc=99,
xtokm=True)
ax[6].plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values[:, ::4], axis=2).squeeze().T, 'k')
cbar = fig.colorbar(im, ax=ax[6:9].tolist())
cbar.formatter.set_powerlimits((-2, 2))
cbar.update_ticks()
for i, axc in enumerate(ax):
axc.set_ylabel('z (m)') if i % 3 == 0 else axc.set_ylabel('')
axc.set_xlabel('x from EB (km)') if i > 5 else axc.set_xlabel('')
fig2, ax = plt.subplots(1, 1, figsize=(4, 1))
ax.plot(domain.lonh[e._plot_slice[3, 0]:e._plot_slice[3, 1]],
np.nanmean(e.values, axis=2).squeeze().T, 'k')
xdegtokm(ax, 0.5 * (initializer.slat + initializer.nlat))
ax.set_ylim(-500, 0)
ax.set_ylabel('z (m)')
return fig, fig2