def plotzetaanimation(geofil,
fil,
savfil=None,
fig=None,
wlon=-25,
elon=0,
slat=10,
nlat=60,
zs=0,
ze=1,
fps=5,
bitrate=1000,
**plotkwargs):
fhgeo = dset(geofil)
fh = mfdset(fil)
ani = animatezeta(fhgeo,
fh,
fig=fig,
wlon=wlon,
elon=elon,
slat=slat,
nlat=nlat,
zs=zs,
ze=ze,
**plotkwargs)
if savfil:
mywriter = animation.FFMpegWriter(fps=fps, bitrate=bitrate)
ani.save(savfil + '.mp4', writer=mywriter)
else:
plt.show()
fh.close()
fhgeo.close()
def __init__(self):
fname="/home/jesse/Desktop/Repos/stations/data/GC/GC_area.nc"
from netCDF4 import Dataset as dset
ncd=dset(fname,mode='r')
self.lats=ncd['latitude'][...] # [91]
self.lons=ncd['longitude'][...] # [144]
self.latedges=ncd['latedges'][...]
self.lonedges=ncd['lonedges'][...]
self.area=ncd['area'][...] # square metres [144, 91]
def getwhash(fhgeo, vgeofil, fh, fh2, sl, htol=1e-3):
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
zs, ze = sl[1].start, sl[1].stop
ys, ye = sl[2].start, sl[2].stop
xs, xe = sl[3].start, sl[3].stop
slmx = np.s_[:, zs:ze, ys:ye, xs - 1:xe]
slmy = np.s_[:, zs:ze, ys - 1:ye, xs:xe]
slmpy = np.s_[:, zs:ze, ys - 1:ye + 1, xs:xe]
emforxdiff = getvaravg(fh2, 'e', slmx)
elmforxdiff = 0.5 * (emforxdiff[:, 0:-1, :, :] + emforxdiff[:, 1:, :, :])
elmforxdiff = np.concatenate((elmforxdiff, elmforxdiff[:, :, :, -1:]),
axis=3)
dxcuforxdiff = fhgeo.variables['dxCu'][slmx[2:]]
ex = np.diff(elmforxdiff, axis=3) / dxcuforxdiff
uh = getvaravg(fh2, 'uh', slmx)
h_cu = getvaravg(fh, 'h_Cu', slmx)
h_cu = np.ma.masked_array(h_cu, mask=(h_cu < htol))
dycu = fhgeo.variables['dyCu'][slmx[2:]]
uzx = (uh / h_cu / dycu).filled(0) * ex
uzx = 0.5 * (uzx[:, :, :, 1:] + uzx[:, :, :, :-1])
emforydiff = getvaravg(fh2, 'e', slmpy)
elmforydiff = 0.5 * (emforydiff[:, 0:-1, :, :] + emforydiff[:, 1:, :, :])
dycv = fhgeo.variables['dyCv'][slmy[2:]]
ey = np.diff(elmforydiff, axis=2) / dycv
vh = getvaravg(fh2, 'vh', slmy)
h_cv = getvaravg(fh, 'h_Cv', slmy)
h_cv = np.ma.masked_array(h_cv, mask=(h_cv < htol))
dxcv = fhgeo.variables['dxCv'][slmy[2:]]
vtwa = vh / dxcv / h_cv
vzy = vtwa * ey
vzy = 0.5 * (vzy[:, :, 1:, :] + vzy[:, :, :-1, :])
wd = getvaravg(fh2, 'wd', sl)
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 1:, :, :] + hw[:, :-1, :, :])
wzb = hw / dbl[:, np.newaxis, np.newaxis]
whash = uzx + vzy + wzb
return whash.filled(np.nan)
def getgeom(filename,
wlon=-25,
elon=0,
slat=10,
nlat=60,
xadditional=False,
yadditional=False):
"""
Usage:
| D, (ah,aq), (dxcu,dycu,dxcv,dycv,dxbu,dybu,dxt,dyt), f = getgeom(filename)
| This fuction returns the depth of the domain D,
| cell areas at T and Bu points ah and aq, resp,
| grid spacing at Cu, Cv, Bu and T points,
| and f at Bu points.
"""
fhgeo = dset(filename, mode='r')
lath = fhgeo.variables['lath'][:]
lonh = fhgeo.variables['lonh'][:]
latq = fhgeo.variables['latq'][:]
lonq = fhgeo.variables['lonq'][:]
xsh = (lonh >= wlon).nonzero()[0][0]
xeh = (lonh <= elon).nonzero()[0][-1]
ysh = (lath >= slat).nonzero()[0][0]
yeh = (lath <= nlat).nonzero()[0][-1]
xsq = (lonq >= wlon).nonzero()[0][0]
xeq = (lonq <= elon).nonzero()[0][-1]
ysq = (latq >= slat).nonzero()[0][0]
yeq = (latq <= nlat).nonzero()[0][-1]
if xadditional:
xsh -= 1
xsq -= 1
if yadditional:
ysh -= 1
ysq -= 1
D = fhgeo.variables['D'][ysh:yeh + 1, xsh:xeh + 1]
ah = fhgeo.variables['Ah'][ysh:yeh + 1, xsh:xeh + 1]
aq = fhgeo.variables['Aq'][ysq:yeq + 1, xsq:xeq + 1]
dxcu = fhgeo.variables['dxCu'][ysh:yeh + 1, xsq:xeq + 1]
dycu = fhgeo.variables['dyCu'][ysh:yeh + 1, xsq:xeq + 1]
dxcv = fhgeo.variables['dxCv'][ysq:yeq + 1, xsh:xeh + 1]
dycv = fhgeo.variables['dyCv'][ysq:yeq + 1, xsh:xeh + 1]
dxbu = fhgeo.variables['dxBu'][ysq:yeq + 1, xsq:xeq + 1]
dybu = fhgeo.variables['dyBu'][ysq:yeq + 1, xsq:xeq + 1]
dxt = fhgeo.variables['dxT'][ysh:yeh + 1, xsh:xeh + 1]
dyt = fhgeo.variables['dyT'][ysh:yeh + 1, xsh:xeh + 1]
f = fhgeo.variables['f'][ysq:yeq + 1, xsq:xeq + 1]
fhgeo.close()
return D, (ah, aq), (dxcu, dycu, dxcv, dycv, dxbu, dybu, dxt, dyt), f
def getgeom(filename):
# This fuction returns the depth of the domain D, cell areas at
# T and Bu points ah and aq, and grid spacing at Cu, Cv, Bu and T points
fhgeo = dset(filename, mode='r')
D = fhgeo.variables['D'][:]
ah = fhgeo.variables['Ah'][:]
aq = fhgeo.variables['Aq'][:]
dxcu = fhgeo.variables['dxCu'][:]
dycu = fhgeo.variables['dyCu'][:]
dxcv = fhgeo.variables['dxCv'][:]
dycv = fhgeo.variables['dyCv'][:]
dxbu = fhgeo.variables['dxBu'][:]
dybu = fhgeo.variables['dyBu'][:]
dxt = fhgeo.variables['dxT'][:]
dyt = fhgeo.variables['dyT'][:]
fhgeo.close()
return D, (ah,aq), (dxcu,dycu,dxcv,dycv,dxbu,dybu,dxt,dyt)
def extract_northward_transport(geofil,vgeofil,fil,fil2,xstart,xend,ystart,yend,zs,ze,meanax):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i,)
fhgeo = dset(geofil)
fh = mfdset(fil)
zi = rdp1.getdims(fh)[2][0]
dbl = -np.diff(zi)*9.8/1031
(xs,xe),(ys,ye),dimv = rdp1.getlatlonindx(fh,wlon=xstart,elon=xend,
slat=ystart, nlat=yend,zs=zs,ze=ze,yhyq='yq')
dxcv = rdp1.getgeombyindx(fhgeo,xs,xe,ys,ye)[2][2]
f = rdp1.getgeombyindx(fhgeo,xs,xe,ys,ye)[-1]
nt_const = dimv[0].size
fhgeo.close()
vh = fh.variables['vh'][0:,zs:ze,ys:ye,xs:xe]/dxcv
fh2 = mfdset(fil2)
h = fh2.variables['h_Cv'][0:,zs:ze,ys:ye,xs:xe]
fh2.close()
fh.close()
vh = vh.filled(0)
vhplus = np.where(vh>0,vh,0)
vhminus = np.where(vh<0,vh,0)
R = 6378
#lr = np.sum(np.sqrt(-h*dbl[:,np.newaxis,np.newaxis])/np.pi/f/1e3,axis=1,keepdims=True)/np.radians(R)
lr = np.sum(np.sqrt(-h*dbl[:,np.newaxis,np.newaxis])/np.pi/f/1e3,axis=1,keepdims=True)/np.radians(R*np.cos(np.radians(dimv[2][:,np.newaxis])))
lr = np.mean(lr,axis=0).squeeze()
terms = np.ma.concatenate(( vhplus[:,:,:,:,np.newaxis],
vhminus[:,:,:,:,np.newaxis]),axis=4)
termsm = np.ma.apply_over_axes(np.mean, terms, meanax)
termsm = termsm.squeeze()
X = dimv[keepax[1]]
Y = dimv[keepax[0]]
return X,Y, termsm, lr
def extract_twapv(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
savfil=None,
cmaxscalefactor=None,
plotatz=False,
Z=None):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
(xs, xe), (ys, ye), dimq = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq',
yhyq='yq')
dxbu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][4]
dybu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][5]
dycu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][1]
dxcv = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][2]
f = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[-1]
nt_const = dimq[0].size
fhgeo.close()
em = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
uh = fh.variables['uh'][0:, zs:ze, ys:ye + 1, xs:xe]
h_cu = fh2.variables['h_Cu'][0:, zs:ze, ys:ye + 1, xs:xe]
utwa = uh / h_cu / dycu
h_cu = np.where(h_cu > 1e-3, h_cu, np.nan)
vh = fh.variables['vh'][0:, zs:ze, ys:ye, xs:xe]
h_cv = fh2.variables['h_Cv'][0:, zs:ze, ys:ye, xs:xe]
vtwa = vh / h_cv / dxcv
vtwa = np.concatenate((vtwa, -vtwa[:, :, :, -1:]), axis=3)
h_cv = np.concatenate((h_cv, -h_cv[:, :, :, -1:]), axis=3)
h_cv = np.where(h_cv > 1e-3, h_cv, np.nan)
fh2.close()
fh.close()
hq = 0.25 * (h_cu[:, :, :-1, :] + h_cv[:, :, :, :-1] + h_cu[:, :, 1:, :] +
h_cv[:, :, :, 1:])
pv = f - np.diff(utwa, axis=2) / dybu + np.diff(vtwa, axis=3) / dxbu
pv = pv / (hq / dbl[:, np.newaxis, np.newaxis])
X = dimq[keepax[1]]
Y = dimq[keepax[0]]
if 1 in keepax:
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimq[1])[0]
if plotatz:
pv = getvaratz(pv, Z, em)
pv = np.ma.apply_over_axes(np.nanmean, pv, meanax)
pv = pv.squeeze()
cmax = np.nanmax(np.absolute(pv)) * cmaxscalefactor
im = m6plot((X, Y, pv),
xlabel=r'x ($^{\circ}$ E)',
ylabel=r'y ($^{\circ}$ N)',
vmin=6e-10,
vmax=cmax,
aspect='equal',
bvnorm=True)
if savfil:
plt.savefig(savfil + '.eps',
dpi=300,
facecolor='w',
edgecolor='w',
format='eps',
transparent=False,
bbox_inches='tight')
else:
plt.show()
def extractT(geofil,
fil,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
ts=0,
te=None,
z=None,
drhodt=-0.2,
rho0=1031.0,
savfil=None,
plotit=True,
loop=True):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fh = mfdset(fil)
(xs, xe), (ys, ye), dimh = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze)
fhgeo = dset(geofil)
D = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0]
fhgeo.close()
nt = dimh[0].size
t0 = time.time()
zl = rdp1.getdims(fh)[2][1]
if loop:
print('Reading data in loop...')
e = fh.variables['e'][0:1, zs:ze, ys:ye, xs:xe] / nt
for i in range(nt):
e += fh.variables['e'][i:i + 1, zs:ze, ys:ye, xs:xe] / nt
sys.stdout.write('\r' + str(int((i + 1) / nt * 100)) + '% done...')
sys.stdout.flush()
print('Time taken for data reading: {}s'.format(time.time() - t0))
else:
e = fh.variables['e'][ts:te, zs:ze, ys:ye, xs:xe]
X = dimh[keepax[1]]
Y = dimh[keepax[0]]
if 1 in keepax:
Y = z
if z == None:
z = np.linspace(-np.nanmax(D), -1, num=50)
Y = z
T = getTatz(zl, z, e)
T = (T - rho0) / drhodt
T = np.ma.apply_over_axes(np.nanmean, T, meanax)
P = T.squeeze()
data = (X, Y, P)
if plotit:
Pmax = np.nanmax(P)
Pmin = np.nanmin(P)
im = m6plot(data,
vmax=Pmax,
vmin=Pmin,
title=r'T at 40N ($^{\circ}$C)',
xlabel=r'x ($^{\circ}$)',
ylabel='z (m)',
bvnorm=True,
blevs=15)
if savfil:
plt.savefig(savfil + '.eps',
dpi=300,
facecolor='w',
edgecolor='w',
format='eps',
transparent=False,
bbox_inches='tight')
else:
plt.show()
else:
return data
def extract_twamomy_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
fil3=None,
alreadysaved=False,
xyasindices=False,
calledfrompv=False,
z=np.linspace(-3000, 0, 100),
htol=1e-3):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimv = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xh',
yhyq='yq',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimv = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
yhyq='yq')
sl = np.s_[:, :, ys:ye, xs:xe]
slmx = np.s_[:, :, ys:ye, xs - 1:xe]
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
Dforgetutwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetutwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][6:8]
dxcv, dycv = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][2:4]
dxcvforxdiff, dycvforxdiff = rdp1.getgeombyindx(
fhgeo, xs - 1, xe, ys, ye)[2][2:4]
dxcvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[2][3:4]
dxbu, dybu = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][4:6]
nt_const = dimv[0].size
t0 = time.time()
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
if fil3:
fh3 = mfdset(fil3)
slmxtn = np.s_[-1:, :, ys:ye, xs - 1:xe]
# islayerdeep0 = fh3.variables['islayerdeep'][:,0,0,0].sum()
# islayerdeep = (fh3.variables['islayerdeep'][slmx].filled(np.nan)).sum(axis=0,
# keepdims=True)
islayerdeep0 = fh3.variables['islayerdeep'][-1:, 0, 0, 0]
islayerdeep = (fh3.variables['islayerdeep'][slmxtn].filled(np.nan))
islayerdeep[:, :, :, -1:] = islayerdeep[:, :, :, -2:-1]
swash = (islayerdeep0 - islayerdeep) / islayerdeep0 * 100
swash = 0.5 * (swash[:, :, :, :-1] + swash[:, :, :, 1:])
fh3.close()
else:
swash = None
em = (fh2.variables['e'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
vh = (fh.variables['vh_masked'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
h_cv = (fh.variables['h_Cv'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
h_cv[h_cv < htol] = np.nan
h_vm = h_cv
vtwa = vh / h_cv / dxcv
vhforxdiff = (fh.variables['vh_masked'][0:, zs:ze, ys:ye, xs - 1:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cvforxdiff = (fh.variables['h_Cv'][0:, zs:ze, ys:ye, xs - 1:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cvforxdiff[h_cvforxdiff < htol] = np.nan
vtwaforxdiff = vhforxdiff / h_cvforxdiff #/dxcvforxdiff
vtwaforxdiff = np.concatenate(
(vtwaforxdiff, vtwaforxdiff[:, :, :, -1:]), axis=3)
vhforydiff = (
fh.variables['vh_masked'][0:, zs:ze, ys - 1:ye + 1, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cvforydiff = (fh.variables['h_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cvforydiff[h_cvforydiff < htol] = np.nan
vtwaforydiff = vhforydiff / h_cvforydiff #/dxcvforydiff
vtwax = np.diff(vtwaforxdiff, axis=3) / dxbu / dybu
vtwax = 0.5 * (vtwax[:, :, :, 0:-1] + vtwax[:, :, :, 1:])
vtway = np.diff(vtwaforydiff, axis=2) / dyt / dxt
vtway = 0.5 * (vtway[:, :, 0:-1, :] + vtway[:, :, 1:, :])
hvmy = np.diff(vhforydiff, axis=2) / dyt / dxt
hvmy = 0.5 * (hvmy[:, :, :-1, :] + hvmy[:, :, 1:, :])
uh = (fh.variables['uh_masked'][0:, zs:ze, ys:ye + 1, xs - 1:xe] *
dt).filled(0).sum(axis=0, keepdims=True) / np.sum(dt)
hum = 0.25 * (uh[:, :, :-1, :-1] + uh[:, :, :-1, 1:] +
uh[:, :, 1:, :-1] + uh[:, :, 1:, 1:]) / dycv
humx = np.diff(np.nan_to_num(uh), axis=3) / dxt / dyt
humx = 0.5 * (humx[:, :, :-1, :] + humx[:, :, 1:, :])
huvxphvvym = (
fh.variables['twa_huvxpt'][0:, zs:ze, ys:ye, xs:xe] * dt +
fh.variables['twa_hvvymt'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hvv = (fh.variables['hvv_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
hvvym = np.diff(hvv, axis=2) / dxt / dyt
hvvym = 0.5 * (hvvym[:, :, :-1, :] + hvvym[:, :, 1:, :])
# hvv = (fh.variables['hvv_T'][0:,zs:ze,ys:ye+1,xs:xe]*dt).sum(axis=0,keepdims=True)*dxt/np.sum(dt)
# hvvym = np.diff(hvv,axis=2)/dycv/dxcv
huvxm = -(huvxphvvym + hvvym)
# hvv = (fh.variables['hvv_Cv'][0:,zs:ze,ys-1:ye+1,xs:xe]*dt).sum(axis=0,keepdims=True)/np.sum(dt)
# hvvym = np.diff(hvv,axis=2)/dxt/dyt
# hvvym = 0.5*(hvvym[:,:,:-1,:] + hvvym[:,:,1:,:])
# huv = (fh.variables['huv_Bu'][0:,zs:ze,ys:ye,xs-1:xe].filled(0)*dt).sum(axis=0,keepdims=True)/np.sum(dt)
# huvxm = np.diff(huv,axis=3)/dxcv
vtwaforvdiff = np.concatenate((vtwa[:, [0], :, :], vtwa), axis=1)
vtwab = np.diff(vtwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
vtwab = np.concatenate(
(vtwab,
np.zeros([vtwab.shape[0], 1, vtwab.shape[2], vtwab.shape[3]])),
axis=1)
vtwab = 0.5 * (vtwab[:, 0:-1, :, :] + vtwab[:, 1:, :, :])
hwb = (fh2.variables['wd'][0:, zs:ze, ys:ye + 1, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hwb = np.diff(hwb, axis=1)
hwb_v = 0.5 * (hwb[:, :, :-1, :] + hwb[:, :, 1:, :])
hwb = (fh2.variables['wd'][0:, zs:ze, ys:ye + 1, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hwm = hwb_v * dbl[:, np.newaxis, np.newaxis]
hwm = 0.5 * (hwb[:, :-1] + hwb[:, 1:]) * dbl[:, np.newaxis, np.newaxis]
hwm_v = 0.5 * (hwm[:, :, :-1, :] + hwm[:, :, 1:, :])
#hwb_v = fh.variables['hwb_Cv'][0:,zs:ze,ys:ye,xs:xe]
#hwm_v = fh.variables['hw_Cv'][0:,zs:ze,ys:ye,xs:xe]
esq = (fh.variables['esq'][0:, zs:ze, ys:ye + 1, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
emforydiff = (fh2.variables['e'][0:, zs:ze, ys:ye + 1, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
elmforydiff = 0.5 * (emforydiff[:, 0:-1, :, :] +
emforydiff[:, 1:, :, :])
edlsqm = (esq - elmforydiff**2)
edlsqmy = np.diff(edlsqm, axis=2) / dycv
hpfv = (fh.variables['twa_hpfv'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
pfvm = (fh2.variables['PFv'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
edpfvdmb = -hpfv + h_cv * pfvm - 0.5 * edlsqmy * dbl[:, np.newaxis,
np.newaxis]
hmfum = (fh.variables['twa_hmfu'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hvwbm = (fh.variables['twa_hvwb'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hdiffvm = (fh.variables['twa_hdiffv'][0:, zs:ze, ys:ye, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
hdvdtviscm = (fh.variables['twa_hdvdtvisc'][0:, zs:ze, ys:ye, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
fh2.close()
fh.close()
advx = hum * vtwax / h_vm
advy = vtwa * vtway
advb = hwm_v * vtwab / h_vm
cor = hmfum / h_vm
pfvm = pfvm
xdivep1 = -huvxm / h_vm
xdivep2 = advx
xdivep3 = vtwa * humx / h_vm
xdivep = (xdivep1 + xdivep2 + xdivep3)
ydivep1 = -hvvym / h_vm
ydivep2 = advy
ydivep3 = vtwa * hvmy / h_vm
ydivep4 = -0.5 * edlsqmy * dbl[:, np.newaxis, np.newaxis] / h_vm
ydivep = (ydivep1 + ydivep2 + ydivep3 + ydivep4)
bdivep1 = hvwbm / h_vm
bdivep2 = advb
bdivep3 = vtwa * hwb_v / h_vm
bdivep4 = -edpfvdmb / h_vm
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
Y1twa = hdiffvm / h_vm
Y2twa = hdvdtviscm / h_vm
terms = np.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfvm[:, :, :, :, np.newaxis], xdivep[:, :, :, :, np.newaxis],
ydivep[:, :, :, :, np.newaxis], bdivep[:, :, :, :, np.newaxis],
Y1twa[:, :, :, :, np.newaxis], Y2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.concatenate(
(xdivep1[:, :, :, :, np.newaxis], xdivep3[:, :, :, :, np.newaxis],
ydivep1[:, :, :, :, np.newaxis], ydivep3[:, :, :, :, np.newaxis],
ydivep4[:, :, :, :, np.newaxis], bdivep1[:, :, :, :, np.newaxis],
bdivep3[:, :, :, :, np.newaxis], bdivep4[:, :, :, :, np.newaxis]),
axis=4)
termsm = np.nanmean(terms, meanax, keepdims=True)
termsepm = np.nanmean(termsep, meanax, keepdims=True)
X = dimv[keepax[1]]
Y = dimv[keepax[0]]
if 1 in keepax and not calledfrompv:
em = np.nanmean(em, axis=meanax, keepdims=True)
elm = np.nanmean(elm, axis=meanax, keepdims=True)
#z = np.linspace(-3000,0,100)
Y = z
P = getvaratzc5(termsm.astype(np.float32), z.astype(np.float32),
em.astype(np.float32))
Pep = getvaratzc5(termsepm.astype(np.float32),
z.astype(np.float32), em.astype(np.float32))
if fil3:
swash = np.nanmean(swash, meanax, keepdims=True)
swash = getvaratzc(swash.astype(np.float32),
z.astype(np.float32),
em.astype(np.float32)).squeeze()
else:
P = termsm.squeeze()
Pep = termsepm.squeeze()
if not calledfrompv:
np.savez('twamomy_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomy_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep, swash, em.squeeze())
def extract_buoy_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
fil3=None,
z=np.linspace(-3000, 0, 100),
htol=1e-3):
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
dz = np.diff(z)[0]
(xs, xe), (ys, ye), dimh = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
ts=0,
te=None)
sl = np.s_[:, zs:ze, ys:ye, xs:xe]
sl2d = sl[2:]
slmx = np.s_[:, zs:ze, ys:ye, xs - 1:xe]
slmx2d = slmx[2:]
slmpy = np.s_[:, zs:ze, ys - 1:ye + 1, xs:xe]
slmpy2d = slmpy[2:]
slmy = np.s_[:, zs:ze, ys - 1:ye, xs:xe]
slmy2d = slmy[2:]
dxcu = fhgeo.variables['dxCu'][slmx2d]
dycv = fhgeo.variables['dyCv'][slmy2d]
e = getvaravg(fh2, 'e', sl)
h = -np.diff(e, axis=1)
h = np.ma.masked_array(h, mask=(h < htol)).filled(np.nan)
zi = fh.variables['zi'][:]
bi = 9.8 * (1 - zi / 1031)
dbl = np.diff(zi) * 9.8 / 1031
bz = -(np.diff(bi)[:, np.newaxis, np.newaxis] / h)
bz = getvaratzc(bz.astype(np.float32), z.astype(np.float32),
e.astype(np.float32))
whash = getwhash(fhgeo, vgeofil, fh, fh2, sl, htol=htol)
whash = getvaratzc(whash.astype(np.float32), z.astype(np.float32),
e.astype(np.float32))
wbz = whash * bz
efxd = getvaravg(fh2, 'e', slmx)
bfxd = getTatzc2(bi.astype(np.float32), z.astype(np.float32),
efxd.astype(np.float32))
bfxd = np.concatenate((bfxd, bfxd[:, :, :, -1:]), axis=3)
bx = np.diff(bfxd, axis=3) / dxcu
u = getutwa(fhgeo, fh, fh2, slmx, htol=htol)
e_cu = np.concatenate((efxd, efxd[:, :, :, -1:]), axis=3)
e_cu = 0.5 * (e_cu[:, :, :, :-1] + e_cu[:, :, :, 1:])
u = getvaratzc(u.astype(np.float32), z.astype(np.float32),
e_cu.astype(np.float32))
ubx = u * bx
ubx = 0.5 * (ubx[:, :, :, :-1] + ubx[:, :, :, 1:])
efyd = getvaravg(fh2, 'e', slmpy)
bfyd = getTatzc2(bi.astype(np.float32), z.astype(np.float32),
efyd.astype(np.float32))
by = np.diff(bfyd, axis=2) / dycv
e_cv = 0.5 * (efyd[:, :, :-1, :] + efyd[:, :, 1:, :])
v = getvtwa(fhgeo, fh, fh2, slmy, htol=htol)
v = getvaratzc(v.astype(np.float32), z.astype(np.float32),
e_cv.astype(np.float32))
vby = v * by
vby = 0.5 * (vby[:, :, :-1, :] + vby[:, :, 1:, :])
hwb = getvaravg(fh2, 'wd', sl)
hwb = -np.diff(hwb, axis=1)
hwm = hwb * dbl[:, np.newaxis, np.newaxis]
wtwa = hwm / h
wtwa = getvaratzc(wtwa.astype(np.float32), z.astype(np.float32),
e.astype(np.float32))
terms = np.ma.concatenate(
(ubx[:, :, :, :, np.newaxis], vby[:, :, :, :, np.newaxis],
wbz[:, :, :, :, np.newaxis], -wtwa[:, :, :, :, np.newaxis]),
axis=4)
return dimh, terms
def get_pveddy_coeffs(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs=0,
ze=None,
z=np.linspace(-3000, 0, 100),
percx=0,
percy=0):
fhgeo = dset(geofil)
fhvgeo = dset(vgeofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
db = fhvgeo.variables['g'][:]
fhvgeo.close()
dbi = np.append(db, 0)
zi = fh.variables['zi'][:]
zl = fh.variables['zl'][:]
dbl = np.diff(zi) * 9.8 / 1031
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
sl, dimq = rdp1.getslice(fh,
xstart,
xend,
ystart,
yend,
xhxq='xq',
yhyq='yq')
xs = sl[3].start
xe = sl[3].stop
ys = sl[2].start
ye = sl[2].stop
zs = sl[1].start
ze = sl[1].stop
slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
(sl[3].start - 1):sl[3].stop]
slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
sl2d = sl[2:]
slmx2d = slmx[2:]
slpx2d = slpx[2:]
dxbu = fhgeo.variables['dxBu'][sl2d]
dxcv = fhgeo.variables['dxCv'][sl2d]
e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
axis=0, keepdims=True)
e = np.concatenate((e, e[:, :, :, -1:]), axis=3)
eatq = 0.25 * (e[:, :, :-1, :-1] + e[:, :, 1:, 1:] + e[:, :, 1:, :-1] +
e[:, :, :-1, 1:])
(_, _, _, q, _, _) = pv2.extract_twapv_terms(geofil,
vgeofil,
fil,
fil2,
xs - 1,
xe,
ys,
ye,
zs,
ze,
meanax=(0, ),
fil3=None,
calledfromgeteddycoeffs=True)
(_, _, _, _, qbud,
_) = pv2.extract_twapv_terms(geofil,
vgeofil,
fil,
fil2,
xs,
xe,
ys,
ye,
zs,
ze,
meanax=(0, ),
fil3=None,
calledfromgeteddycoeffs=True)
fd = qbud[:, :, :, :, 12]
q[np.isnan(q)] = 0
qx = np.diff(q, axis=3) / dxcv
qx = np.concatenate((qx, qx[:, :, :, -1:]), axis=3)
qxx = np.diff(qx, axis=3) / dxbu
fdaz = getvaratzc(fd.astype(np.float32), z.astype(np.float32),
eatq.astype(np.float32))
qxxaz = getvaratzc(qxx.astype(np.float32), z.astype(np.float32),
eatq.astype(np.float32))
qxxazm = np.apply_over_axes(np.nanmean, qxxaz, (0, 2))
fdazm = np.apply_over_axes(np.nanmean, fdaz, (0, 2))
fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
cmax = np.fabs(np.nanpercentile(qxxazm, (1, 99))).max()
im = ax[0].pcolormesh(dimq[3],
z,
qxxazm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
cb = fig.colorbar(im, ax=ax[0])
cb.formatter.set_powerlimits((0, 0))
cb.update_ticks()
cmax = np.fabs(np.nanpercentile(fdazm, (1, 99))).max()
im = ax[1].pcolormesh(dimq[3],
z,
fdazm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
cb = fig.colorbar(im, ax=ax[1])
cb.formatter.set_powerlimits((0, 0))
cb.update_ticks()
for axc in ax:
xdegtokm(axc, 0.5 * (ystart + yend))
axc.set_xlabel('x (km)')
axc.grid()
ax[0].set_ylabel('z (m)')
fig2, ax = plt.subplots(1, 1, sharey=True)
ax.linfit(qxx, fd, percx=percx, percy=percy)
ax.set_xlabel('qxx')
ax.set_ylabel('grad fd')
ax.grid()
return fig, fig2
def pvbudget(firstrun,geofil,fil,savfil=None):
#def pvbudget(firstrun,savfil=None):
# geofil = ("/home/sbire/MOM6-examples/ocean_only/buoy_forced_gyre_highfluxc"
# "_2km_fromut_kvkd1em4_kh1e3_khth1e3_lowbz_smd_veryhires_3yr/ocean_geometry.nc")
# fil = ("/home/sbire/MOM6-examples/ocean_only/buoy_forced_gyre_highfluxc"
# "_2km_fromut_kvkd1em4_kh1e3_khth1e3_lowbz_smd_veryhires_3yr/output__0031_07.nc")
D, (ah,aq), (dxcu,dycu,dxcv,dycv,dxbu,dybu,dxt,dyt) = rdp.getgeom(geofil)
(xh,yh), (xq,yq), (zi,zl), time = rdp.getdims(fil)
nx = len(xh)
ny = len(yh)
nz = len(zl)
nzp1 = len(zi)
nt = len(time)
if firstrun == 1:
hm_cu = np.zeros((nz,ny,nx))
hm_cv = np.zeros((nz,ny,nx))
hm = np.zeros((nz,ny,nx))
em = np.zeros((nzp1,ny,nx))
elm = np.zeros((nz,ny,nx))
um = np.zeros((nz,ny,nx))
vm = np.zeros((nz,ny,nx))
msigpiuxm = np.zeros((nz,ny,nx))
Yxm = np.zeros((nz,ny,nx))
wvbxm = np.zeros((nz,ny,nx))
sigpivym = np.zeros((nz,ny,nx))
Xym = np.zeros((nz,ny,nx))
wubym = np.zeros((nz,ny,nx))
#nt = 2
fh = dset(fil, mode='r')
for i in range(nt):
e = fh.variables['e'][i,:,:,:]
el = (e[0:-1,:,:] + e[1:,:,:])/2;
u = fh.variables['u'][i,:,:,:]
v = fh.variables['v'][i,:,:,:]
frhatv = fh.variables['frhatv'][i,:,:,:]
elm += el/nt
um += u/nt
vm += v/nt
cav = fh.variables['CAv'][i,:,:,:]
gkev = fh.variables['gKEv'][i,:,:,:]
msigpiu = cav - gkev
msigpiu = np.ma.filled(msigpiu.astype(float), 0)
msigpiux = np.diff(np.concatenate((msigpiu,msigpiu[:,:,[-1]]),axis=2),axis = 2)/dxBu
Y = fh.variables['diffv'][i,:,:,:] + fh.variables['dv_dt_visc'][i,:,:,:]
Y = np.ma.filled(Y.astype(float), 0)
Yx = np.diff(np.concatenate((Y,Y[:,:,[-1]]),axis=2),axis = 2)/dxBu
wvb = fh.variables['dvdt_dia'][i,:,:,:]
wvb = np.ma.filled(wvb.astype(float), 0)
wvbx = np.diff(np.concatenate((wvb,wvb[:,:,[-1]]),axis=2),axis = 2)/dxBu
cau = fh.variables['CAu'][i,:,:,:]
gkeu = fh.variables['gKEu'][i,:,:,:]
sigpiv = cau - gkeu
sigpiv = np.ma.filled(sigpiv.astype(float), 0)
sigpivy = np.diff(np.concatenate((sigpiv,sigpiv[:,[-1],:]),axis=1),axis = 1)/dyBu
X = fh.variables['diffu'][i,:,:,:] + fh.variables['du_dt_visc'][i,:,:,:]
X = np.ma.filled(X.astype(float), 0)
Xy = np.diff(np.concatenate((X,X[:,[-1],:]),axis=1),axis = 1)/dyBu
wub = fh.variables['dudt_dia'][i,:,:,:]
wub = np.ma.filled(wub.astype(float), 0)
wuby = np.diff(np.concatenate((wub,wub[:,[-1],:]),axis=1),axis = 1)/dyBu
msigpiuxm += msigpiux/nt
Yxm += Yx/nt
wvbxm += wvbx/nt
sigpivym += sigpivy/nt
Xym += Xy/nt
wubym += wuby/nt
print((i+1)/nt*100)
fh.close()
terms = np.concatenate((msigpiuxm[:,:,:,np.newaxis],
Yxm[:,:,:,np.newaxis],
wvbxm[:,:,:,np.newaxis],
-sigpivym[:,:,:,np.newaxis],
-Xym[:,:,:,np.newaxis],
-wubym[:,:,:,np.newaxis]),axis=3)
np.savez('PVterms', terms=terms, elm=elm, em=em, vm=vm, um=um)
else:
npzfile = np.load('PVterms.npz')
terms = npzfile['terms']
elm = npzfile['elm']
res = np.sum(terms,axis=3)
xlen = 10
ystart = 30
yend = 40
figa = ['(a)','(b)','(c)','(d)','(e)','(f)']
xs = nx-xlen
xe = nx
ys = [i for i in range(ny) if yq[i] >= ystart and yq[i] <= yend][1]
ye = [i for i in range(ny) if yq[i] >= ystart and yq[i] <= yend][-1]
zs = 0
ze = nz
s = 6400*np.cos(0.5*(ystart+yend)*np.pi/180)*xq*np.pi/180
X, Y = np.meshgrid(s[xs:xe],zl[zs:ze])
Y_dummy = np.mean(elm[zs:ze,ys:ye,xs:xe], axis=1)
#termsy *= 1e4
#res *= 1e4
vmn = -np.amax(np.absolute(np.mean(terms[zs:ze,ys:ye,xs:xe,:], axis=1)))
vmx = np.amax(np.absolute(np.mean(terms[zs:ze,ys:ye,xs:xe,:], axis=1)))
#Vctr = np.linspace(np.floor(vmn),np.ceil(vmx),num=12,endpoint=True)
#Vctr = np.linspace(-5e-4,5e-4,num=12,endpoint=True)
Vctr = np.linspace(vmn,vmx,num=12,endpoint=True)
Vcbar = (Vctr[1:] + Vctr[:-1])/2
#Vcbar = np.arange(Vctr[1],Vctr[-1],1,'int32')
print(Vctr,Vcbar)
plt.figure()
for i in np.arange(6):
ax = plt.subplot(3,2,i+1)
P = np.mean(terms[zs:ze,ys:ye,xs:xe,i], axis=1)
im = ax.contourf(X, Y_dummy, P, Vctr, cmap=plt.cm.RdBu_r)
cbar = plt.colorbar(im, ticks=Vcbar)
ax.text(-61,-300,figa[i])
cbar.formatter.set_powerlimits((-1, 1))
cbar.update_ticks()
if (i+1) % 2 == 1:
plt.ylabel('$z (m)$')
else:
plt.yticks([0,-500,-1000,-1500,-2000],[])
if np.in1d(i+1,[5,6]):
plt.xlabel('$x (km)$')
plt.xticks([-50,-30,-10],[-50,-30,-10])
else:
plt.xticks([-50,-30,-10],[-50,-30,-10])
if savfil:
plt.savefig(savfil+'.eps', dpi=300, facecolor='w', edgecolor='w',
format='eps', transparent=False, bbox_inches='tight')
else:
plt.show()
#plt.savefig('twamomy.eps', dpi=300, facecolor='w', edgecolor='w', format='eps', transparent=False, bbox_inches='tight')
plt.figure()
im = plt.contourf(X, Y_dummy, np.mean(res[zs:ze,ys:ye,xs:xe], axis=1), Vctr, cmap=plt.cm.RdBu_r)
cbar = plt.colorbar(im, ticks=Vcbar)
cbar.formatter.set_powerlimits((-1, 1))
cbar.update_ticks()
plt.show()
def extract_twamomy_terms(geofil,
vgeofil,
fil,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
alreadysaved=False,
xyasindices=False,
calledfrompv=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
zi = rdp1.getdims(fh)[2][0]
dbl = -np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimv = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xh',
yhyq='yq',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimv = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
yhyq='yq')
D = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0]
(ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[1]
Dforgetvtwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetvtwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethuforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys,
ye + 1)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][6:8]
dycv = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][3]
dxbu = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][4]
nt_const = dimv[0].size
t0 = time.time()
print('Reading data using loop...')
v = fh.variables['v'][0:1, zs:ze, ys:ye, xs:xe]
nt = np.ones(v.shape) * nt_const
frhatv = fh.variables['frhatv'][0:1, zs:ze, ys:ye, xs:xe]
h_v = frhatv * D[np.newaxis, np.newaxis, :, :]
h_v = np.ma.masked_array(h_v, mask=(h_v <= 1e-3).astype(int))
nt[h_v <= 1e-3] -= 1
hvm = (h_v * v).filled(0)
hvvm = (h_v * v * v).filled(0)
h_vm = h_v.filled(0)
hvmforxdiff, h_vmforxdiff = getvtwaforxdiff(fh, fhgeo,
Dforgetvtwaforxdiff, 0,
xs - 1, xe, ys, ye, zs, ze)
hvmforydiff, h_vmforydiff, hvvmforydiff = getvtwaforydiff(
fh, fhgeo, Dforgetvtwaforydiff, 0, xs, xe, ys - 1, ye + 1, zs, ze)
cav = fh.variables['CAv'][0:1, zs:ze, ys:ye, xs:xe]
gkev = fh.variables['gKEv'][0:1, zs:ze, ys:ye, xs:xe]
rvxu = fh.variables['rvxu'][0:1, zs:ze, ys:ye, xs:xe]
hmfum = (h_v * (cav - gkev - rvxu)).filled(0)
pfvm = fh.variables['PFv'][0:1, zs:ze, ys:ye, xs:xe]
# pfvm = np.ma.masked_array(pfvm,mask=(h_v<=1e-3).astype(int))
# pfvm = pfvm.filled(0)
hdvdtviscm = (
h_v *
fh.variables['dv_dt_visc'][0:1, zs:ze, ys:ye, xs:xe]).filled(0)
hdiffvm = (h_v *
fh.variables['diffv'][0:1, zs:ze, ys:ye, xs:xe]).filled(0)
dvdtdia = fh.variables['dvdt_dia'][0:1, zs:ze, ys:ye, xs:xe]
wd = fh.variables['wd'][0:1, zs:ze, ys:ye + 1, xs:xe]
wd = np.diff(wd, axis=1)
wdm = wd
hvwbm = (v * (wd[:, :, 0:-1, :] + wd[:, :, 1:, :]) / 2 -
h_v * dvdtdia).filled(0)
uh = fh.variables['uh'][0:1, zs:ze, ys:ye + 1, xs - 1:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh, axis=3) / ah
huvxpTm = (v * (uhx[:, :, 0:-1, :] + uhx[:, :, 1:, :]) / 2 -
h_v * rvxu).filled(0)
vh = fh.variables['vh'][0:1, zs:ze, ys - 1:ye + 1, xs:xe]
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh, axis=2) / ah
hvvymTm = (v * (vhy[:, :, 0:-1, :] + vhy[:, :, 1:, :]) / 2 -
h_v * gkev).filled(0)
u = fh.variables['u'][0:1, zs:ze, ys:ye + 1, xs - 1:xe]
u = 0.25 * (u[:, :, 0:-1, 0:-1] + u[:, :, 1:, 0:-1] +
u[:, :, 0:-1, 1:] + u[:, :, 1:, 1:])
hum = (h_v * u).filled(0)
humforxdiff = gethuforxdiff(fh, fhgeo, Dforgethuforxdiff, 0, xs - 1,
xe, ys, ye + 1, zs, ze)
wd = fh.variables['wd'][i:i + 1, zs:ze, ys:ye + 1, xs:xe]
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 0:-1, :, :] + hw[:, 1:, :, :])
hwm_v = 0.5 * (hw[:, :, 0:-1, :] + hw[:, :, 1:, :])
emforydiff = fh.variables['e'][0:1, zs:ze, ys:ye + 1, xs:xe] / nt_const
if 1 in keepax:
em = fh.variables['e'][0:1, zs:ze, ys:ye, xs:xe] / nt_const
for i in range(1, nt_const):
v = fh.variables['v'][i:i + 1, zs:ze, ys:ye, xs:xe]
frhatv = fh.variables['frhatv'][i:i + 1, zs:ze, ys:ye, xs:xe]
h_v = frhatv * D[np.newaxis, np.newaxis, :, :]
h_v = np.ma.masked_array(h_v, mask=(h_v <= 1e-3).astype(int))
nt[h_v <= 1e-3] -= 1
hvm += (h_v * v).filled(0)
h_vm += h_v.filled(0)
hvforxdiff, h_vforxdiff = getvtwaforxdiff(fh, fhgeo,
Dforgetvtwaforxdiff, i,
xs - 1, xe, ys, ye, zs,
ze)
hvforydiff, h_vforydiff, hvvforydiff = getvtwaforydiff(
fh, fhgeo, Dforgetvtwaforydiff, i, xs, xe, ys - 1, ye + 1, zs,
ze)
hvmforxdiff += hvforxdiff
h_vmforxdiff += h_vforxdiff
hvvmforydiff += hvvforydiff
hvmforydiff += hvforydiff
h_vmforydiff += h_vforydiff
cav = fh.variables['CAv'][i:i + 1, zs:ze, ys:ye, xs:xe]
gkev = fh.variables['gKEv'][i:i + 1, zs:ze, ys:ye, xs:xe]
rvxu = fh.variables['rvxu'][i:i + 1, zs:ze, ys:ye, xs:xe]
hmfu = h_v * (cav - gkev - rvxu)
hmfum += hmfu.filled(0)
pfv = fh.variables['PFv'][i:i + 1, zs:ze, ys:ye, xs:xe]
# pfv = np.ma.masked_array(pfv,mask=(h_v<=1e-3).astype(int))
# pfvm += pfv.filled(0)
pfvm += pfv
hdvdtvisc = h_v * fh.variables['dv_dt_visc'][i:i + 1, zs:ze, ys:ye,
xs:xe]
hdvdtviscm += hdvdtvisc.filled(0)
hdiffv = h_v * fh.variables['diffv'][i:i + 1, zs:ze, ys:ye, xs:xe]
hdiffvm += hdiffv.filled(0)
dvdtdia = fh.variables['dvdt_dia'][i:i + 1, zs:ze, ys:ye, xs:xe]
wd = fh.variables['wd'][i:i + 1, zs:ze, ys:ye + 1, xs:xe]
wd = np.diff(wd, axis=1)
wdm += wd
hvwb = (v * (wd[:, :, 0:-1, :] + wd[:, :, 1:, :]) / 2 -
h_v * dvdtdia)
hvwb = np.ma.masked_array(hvwb, mask=(h_v <= 1e-3).astype(int))
hvwbm += hvwb.filled(0)
uh = fh.variables['uh'][0:1, zs:ze, ys:ye + 1, xs - 1:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh, axis=3) / ah
huvxpT = (v * (uhx[:, :, 0:-1, :] + uhx[:, :, 1:, :]) / 2 -
h_v * rvxu).filled(0)
huvxpTm += huvxpT
vh = fh.variables['vh'][0:1, zs:ze, ys - 1:ye + 1, xs:xe]
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh, axis=2) / ah
hvvymT = (v * (vhy[:, :, 0:-1, :] + vhy[:, :, 1:, :]) / 2 -
h_v * gkev).filled(0)
hvvymTm = hvvymT
u = fh.variables['u'][0:1, zs:ze, ys:ye + 1, xs - 1:xe]
u = 0.25 * (u[:, :, 0:-1, 0:-1] + u[:, :, 1:, 0:-1] +
u[:, :, 0:-1, 1:] + u[:, :, 1:, 1:])
hum += (h_v * u).filled(0)
humforxdiff += gethuforxdiff(fh, fhgeo, Dforgethuforxdiff, i,
xs - 1, xe, ys, ye + 1, zs, ze)
wd = fh.variables['wd'][i:i + 1, zs:ze, ys:ye + 1, xs:xe]
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 0:-1, :, :] + hw[:, 1:, :, :])
hwm_v += 0.5 * (hw[:, :, 0:-1, :] + hw[:, :, 1:, :])
emforydiff += fh.variables['e'][i:i + 1, zs:ze, ys:ye + 1,
xs:xe] / nt_const
if 1 in keepax:
em += fh.variables['e'][i:i + 1, zs:ze, ys:ye,
xs:xe] / nt_const
sys.stdout.write('\r' + str(int((i + 1) / nt_const * 100)) +
'% done...')
sys.stdout.flush()
fhgeo.close()
print('Time taken for data reading: {}s'.format(time.time() - t0))
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
elmforydiff = 0.5 * (emforydiff[:, 0:-1, :, :] +
emforydiff[:, 1:, :, :])
vtwa = hvm / h_vm
vtwaforxdiff = hvmforxdiff / h_vmforxdiff
vtwaforydiff = hvmforydiff / h_vmforydiff
vtwaforxdiff = np.concatenate(
(vtwaforxdiff, -vtwaforxdiff[:, :, :, -1:]), axis=3)
vtwax = np.diff(vtwaforxdiff, axis=3) / dxbu
vtwax = 0.5 * (vtwax[:, :, :, 0:-1] + vtwax[:, :, :, 1:])
vtway = np.diff(vtwaforydiff, axis=2) / dyt
vtway = 0.5 * (vtway[:, :, 0:-1, :] + vtway[:, :, 1:, :])
humx = np.diff(humforxdiff, axis=3) / dxt
humx = 0.5 * (humx[:, :, 0:-1, :] + humx[:, :, 1:, :])
hvmy = np.diff(hvmforydiff, axis=2) / dyt
hvmy = 0.5 * (hvmy[:, :, 0:-1, :] + hvmy[:, :, 1:, :])
huvxphvvym = huvxpTm + hvvymTm
hvvym = np.diff(hvvmforydiff, axis=2) / dyt
hvvym = 0.5 * (hvvym[:, :, 0:-1, :] + hvvym[:, :, 1:, :])
huvxm = huvxphvvym - hvvym
vtwaforvdiff = np.concatenate((vtwa[:, [0], :, :], vtwa), axis=1)
vtwab = np.diff(vtwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
vtwab = np.concatenate(
(vtwab,
np.zeros([vtwab.shape[0], 1, vtwab.shape[2], vtwab.shape[3]])),
axis=1)
vtwab = 0.5 * (vtwab[:, 0:-1, :, :] + vtwab[:, 1:, :, :])
hwb_v = 0.5 * (wdm[:, :, 0:-1, :] + wdm[:, :, 1:, :])
print('Calculating form drag using loop...')
t0 = time.time()
e = fh.variables['e'][0:1, zs:ze, ys:ye + 1, xs:xe]
el = 0.5 * (e[:, 0:-1, :, :] + e[:, 1:, :, :])
ed = e - emforydiff
edl = el - elmforydiff
edlsqm = (edl**2)
pfv = fh.variables['PFv'][0:1, zs:ze, ys:ye, xs:xe]
pfvd = pfv - pfvm / nt_const
geta = 9.8 * ed[:, :1, :, :] / 1031
getay = np.diff(geta, axis=2) / dycv
pfvd = np.concatenate(
(pfvd, np.zeros([pfvd.shape[0], 1, pfvd.shape[2], pfvd.shape[3]])),
axis=1)
pfvd = 0.5 * (pfvd[:, 0:-1, :, :] + pfvd[:, 1:, :, :])
pfvd = np.concatenate((-getay, pfvd), axis=1)
ed = 0.5 * (ed[:, :, 0:-1, :] + ed[:, :, 1:, :])
edpfvdm = ed * pfvd
for i in range(1, nt_const):
e = fh.variables['e'][i:i + 1, zs:ze, ys:ye + 1, xs:xe]
el = 0.5 * (e[:, 0:-1, :, :] + e[:, 1:, :, :])
ed = e - emforydiff
edl = el - elmforydiff
edlsqm += (edl**2)
pfv = fh.variables['PFv'][i:i + 1, zs:ze, ys:ye, xs:xe]
pfvd = pfv - pfvm / nt_const
geta = 9.8 * ed[:, :1, :, :] / 1031
getay = np.diff(geta, axis=2) / dycv
pfvd = np.concatenate(
(pfvd,
np.zeros([pfvd.shape[0], 1, pfvd.shape[2], pfvd.shape[3]])),
axis=1)
pfvd = 0.5 * (pfvd[:, 0:-1, :, :] + pfvd[:, 1:, :, :])
pfvd = np.concatenate((-getay, pfvd), axis=1)
ed = 0.5 * (ed[:, :, 0:-1, :] + ed[:, :, 1:, :])
edpfvdm += ed * pfvd
sys.stdout.write('\r' + str(int((i + 1) / nt_const * 100)) +
'% done...')
sys.stdout.flush()
print('Time taken for data reading: {}s'.format(time.time() - t0))
fh.close()
edlsqmy = np.diff(edlsqm, axis=2) / dycv
advx = hum * vtwax / h_vm
advy = vtwa * vtway
advb = hwm_v * vtwab / h_vm
cor = hmfum / h_vm
pfvm = pfvm / nt_const
xdivep1 = huvxm / h_vm
xdivep2 = -advx
xdivep3 = -vtwa * humx / h_vm
xdivep = (xdivep1 + xdivep2 + xdivep3)
ydivep1 = hvvym / h_vm
ydivep2 = -advy
ydivep3 = -vtwa * hvmy / h_vm
ydivep4 = 0.5 * edlsqmy * dbl[:, np.newaxis, np.newaxis] / h_vm
ydivep = (ydivep1 + ydivep2 + ydivep3 + ydivep4)
bdivep1 = hvwbm / h_vm
bdivep2 = -advb
bdivep3 = -vtwa * hwb_v / h_vm
bdivep4 = np.diff(edpfvdm,
axis=1) / db[:, np.newaxis,
np.newaxis] * dbl[:, np.newaxis,
np.newaxis] / h_vm
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
Y1twa = hdiffvm / h_vm
Y2twa = hdvdtviscm / h_vm
terms = np.ma.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfvm[:, :, :, :, np.newaxis], -xdivep[:, :, :, :, np.newaxis],
-ydivep[:, :, :, :, np.newaxis], -bdivep[:, :, :, :, np.newaxis],
Y1twa[:, :, :, :, np.newaxis], Y2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.ma.concatenate(
(xdivep1[:, :, :, :, np.newaxis], xdivep3[:, :, :, :, np.newaxis],
ydivep1[:, :, :, :, np.newaxis], ydivep3[:, :, :, :, np.newaxis],
ydivep4[:, :, :, :, np.newaxis], bdivep1[:, :, :, :, np.newaxis],
bdivep3[:, :, :, :, np.newaxis], bdivep4[:, :, :, :, np.newaxis]),
axis=4)
terms[np.isinf(terms)] = np.nan
termsep[np.isinf(termsep)] = np.nan
termsm = np.ma.apply_over_axes(np.nanmean, terms, meanax)
termsepm = np.ma.apply_over_axes(np.nanmean, termsep, meanax)
X = dimv[keepax[1]]
Y = dimv[keepax[0]]
if 1 in keepax:
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimv[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
Pep = termsepm.squeeze()
Pep = np.ma.filled(Pep.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
if not calledfrompv:
np.savez('twamomy_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomy_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep)
def get_deddy_coeffs_fromflx(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs=0,
ze=None,
zlim=(-1500, 0),
percx=0,
percy=0,
nsqdep=False,
fil3=None,
swashperc=1):
fhgeo = dset(geofil)
fhvgeo = dset(vgeofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
db = fhvgeo.variables['g'][:]
fhvgeo.close()
dbi = np.append(db, 0)
zi = fh.variables['zi'][:]
zl = fh.variables['zl'][:]
dbl = np.diff(zi) * 9.8 / 1031
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
z = np.linspace(zlim[0], zlim[1], 100)
sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
sl2d = sl[2:]
dycv = fhgeo.variables['dyCv'][sl2d]
slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
f = fhgeo.variables['f'][sl2d]
e = (fh2.variables['e'][slpy] * dt).sum(axis=0,
keepdims=True) / dt.sum(axis=0)
eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
v, h = pv.getvtwa(fhgeo, fh, fh2, sl)
v = v[:, :, :, :-1]
h = h[:, :, :, :-1]
sig = h / dbl[:, np.newaxis, np.newaxis]
hi = 0.5 * (h[:, :-1] + h[:, 1:])
sigi = hi / dbi[1:-1, np.newaxis, np.newaxis]
vb = -np.diff(v, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
#vb = np.concatenate((vb[:,:1,:,:],vb,vb[:,-1:,:,:]),axis=1)
vi = np.concatenate((v[:, :1, :, :], v, -v[:, -1:, :, :]), axis=1)
vi = 0.5 * (vi[:, :-1, :, :] + vi[:, 1:, :, :])
vib = -np.diff(vi, axis=1) / dbl[:, np.newaxis, np.newaxis]
if nsqdep:
fsqvb = f**2 * vb * sigi
else:
fsqvb = f**2 * vb
fsqvbaz = getvaratzc(fsqvb.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
esq = (fh.variables['esq'][slpy] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
elmforydiff = 0.5 * (e[:, 0:-1, :, :] + e[:, 1:, :, :])
edlsqm = (esq - elmforydiff**2)
edlsqmy = np.diff(edlsqm, axis=2) / dycv
hpfv = (fh.variables['twa_hpfv'][sl] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
pfvm = (fh2.variables['PFv'][sl] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
edpfvdmb = -(-hpfv + h * pfvm -
0.5 * edlsqmy * dbl[:, np.newaxis, np.newaxis])
fh.close()
fh2.close()
fhgeo.close()
zdpfvd = sint.cumtrapz(edpfvdmb, dx=-dbl[0], axis=1)
# zdpfvd = np.concatenate((np.zeros(zdpfvd[:,:1,:,:].shape),
# zdpfvd,
# np.zeros(zdpfvd[:,:1,:,:].shape)),axis=1)
# zdpfvd = 0.5*(zdpfvd[:,:-1,:,:]+zdpfvd[:,1:,:,:])
zdpfvdaz = getvaratzc(zdpfvd.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
fsqvbazm = np.apply_over_axes(np.nanmean, fsqvbaz, (0, 2))
zdpfvdazm = np.apply_over_axes(np.nanmean, zdpfvdaz, (0, 2))
fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
cmax = np.fabs(np.percentile(zdpfvdazm, (1, 99))).max()
im = ax[0].pcolormesh(dimv[3],
z,
zdpfvdazm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[0])
cmax = np.fabs(np.percentile(fsqvbazm, (1, 99))).max()
im = ax[1].pcolormesh(dimv[3],
z,
fsqvbazm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[1])
if fil3:
fh3 = mfdset(fil3)
slmxtn = np.s_[-1:, sl[1], sl[2], (sl[3].start - 1):sl[3].stop]
islayerdeep0 = fh3.variables['islayerdeep'][-1:, 0, 0, 0]
islayerdeep = fh3.variables['islayerdeep'][slmxtn]
if np.ma.is_masked(islayerdeep):
islayerdeep = islayerdeep.filled(np.nan)
islayerdeep[:, :, :, -1:] = islayerdeep[:, :, :, -2:-1]
swash = (islayerdeep0 - islayerdeep) / islayerdeep0 * 100
swash = 0.5 * (swash[:, :, :, :-1] + swash[:, :, :, 1:])
fh3.close()
swash = getvaratzc(swash.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
swash = np.apply_over_axes(np.nanmean, swash, (0, 2))
em = np.apply_over_axes(np.nanmean, e, (0, 2))
xx, zz = np.meshgrid(dimv[3], zi)
for axc in ax:
axc.contour(dimv[3],
z,
swash.squeeze(),
np.array([swashperc]),
colors='k')
axc.contour(xx, em.squeeze(), zz, ls='k--')
axc.set_ylim(zlim[0], zlim[1])
fig2, ax = plt.subplots(1, 1, sharey=True, figsize=(12, 3))
ax.linfit(fsqvbaz, zdpfvdaz, percx=percx, percy=percy)
ax.set_xlabel(r'$f^2 v_b$')
ax.set_ylabel(r'$\zeta^{\prime} m_y^{\prime}$')
return fig, fig2
def extract_twapv_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
fil3=None,
calledfromgeteddycoeffs=False,
alreadysaved=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = -np.diff(zi) * 9.8 / 1031
if not calledfromgeteddycoeffs:
(xs, xe), (ys, ye), dimq = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq',
yhyq='yq')
else:
(xs, xe), (ys, ye), dimq = rdp1.getdimsbyindx(fh,
xstart,
xend,
ystart,
yend,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xq',
yhyq='yq')
sl = np.s_[:, zs:ze, ys:ye, xs:xe]
nt_const = dimq[0].size
pvhash, hq = getpv(fhgeo, fh, fh2, xs, xe, ys, ye)
slpy = np.s_[:, zs:ze, ys:ye + 1, xs:xe]
dxcu = fhgeo.variables['dxCu'][slpy[2:]]
dycv = fhgeo.variables['dyCv'][sl[2:]]
dycv = np.concatenate((dycv, dycv[:, -1:]), axis=1)
dxbu = fhgeo.variables['dxBu'][sl[2:]]
dybu = fhgeo.variables['dyBu'][sl[2:]]
aq = fhgeo.variables['Aq'][sl[2:]]
f = fhgeo.variables['f'][sl[2:]]
if fil3:
fh3 = mfdset(fil3)
sltn = np.s_[-1:, zs:ze, ys:ye, xs:xe]
islayerdeep0 = fh3.variables['islayerdeep'][-1:, 0, 0, 0]
islayerdeep = (fh3.variables['islayerdeep'][sltn].filled(np.nan))
swash = (islayerdeep0 - islayerdeep) / islayerdeep0 * 100
fh3.close()
else:
swash = None
xmom = extract_twamomx_terms(geofil,
vgeofil,
fil,
fil2,
xs,
xe,
ys,
ye + 1,
zs,
ze, (0, ),
alreadysaved=False,
xyasindices=True,
calledfrompv=True)[2]
ymom = extract_twamomy_terms(geofil,
vgeofil,
fil,
fil2,
xs,
xe,
ys,
ye,
zs,
ze, (0, ),
alreadysaved=False,
xyasindices=True,
calledfrompv=True)[2]
xmom = xmom[np.newaxis, :, :, :, :]
ymom = ymom[np.newaxis, :, :, :, :]
ymom = np.concatenate((ymom, ymom[:, :, :, -1:]), axis=3)
bxppvflx = np.sum(xmom[:, :, :, :, [0, 1, 3, 4]], axis=4)
pvhash1, _ = getpv(fhgeo, fh, fh2, xs, xe, ys - 1, ye + 1)
sl1 = np.s_[:, zs:ze, ys - 1:ye + 1, xs:xe]
vtwa1, h_cv1 = getvtwa(fhgeo, fh, fh2, sl1)
vtwa1 = 0.5 * (vtwa1[:, :, :, :-1] + vtwa1[:, :, :, 1:])
pvhashvtwa = pvhash1 * vtwa1
sl1 = np.s_[:, zs:ze, ys:ye + 1, xs:xe]
h_cu1 = fh.variables['h_Cu'][sl1].filled(np.nan).mean(axis=0,
keepdims=True)
h_cu1 = np.ma.masked_array(h_cu1, mask=(h_cu1 < 1e-3))
pvflxx = h_cu1 * (pvhashvtwa[:, :, :-1, :] +
pvhashvtwa[:, :, 1:, :]) / 2
byppvflx = np.sum(ymom[:, :, :, :, [0, 1, 3, 4]], axis=4)
pvhash1, _ = getpv(fhgeo, fh, fh2, xs - 1, xe, ys, ye)
sl1 = np.s_[:, zs:ze, ys:ye + 1, xs - 1:xe]
utwa1, h_cu1 = getutwa(fhgeo, fh, fh2, sl1)
utwa1 = 0.5 * (utwa1[:, :, :-1, :] + utwa1[:, :, 1:, :])
pvhashutwa = pvhash1 * utwa1
pvhashutwa[:, :, :, -1] = 0.0
sl1 = np.s_[:, zs:ze, ys:ye, xs:xe]
h_cv1 = fh.variables['h_Cv'][sl1].mean(axis=0, keepdims=True)
h_cv1 = np.ma.masked_array(h_cv1, mask=(h_cv1 < 1e-3))
pvflxy = h_cv1 * (pvhashutwa[:, :, :, :-1] +
pvhashutwa[:, :, :, 1:]) / 2
pvflxy = np.concatenate((pvflxy, pvflxy[:, :, :, -1:]), axis=3)
bx = bxppvflx - pvflxx
by = byppvflx + pvflxy
xmom1 = xmom[:, :, :, :, [2, 5, 6, 7, 8, 9]]
xmom1 = np.concatenate(
(+pvflxx[:, :, :, :, np.newaxis], xmom1, bx[:, :, :, :,
np.newaxis]),
axis=4)
ymom1 = ymom[:, :, :, :, [2, 5, 6, 7, 8, 9]]
ymom1 = np.concatenate(
(-pvflxy[:, :, :, :, np.newaxis], ymom1, by[:, :, :, :,
np.newaxis]),
axis=4)
#pv = (-np.diff(xmom*dxcu[:,:,np.newaxis],axis=2) + np.diff(ymom*dycv[:,:,np.newaxis],axis=3))/aq[:,:,np.newaxis]
pv = -np.diff(xmom, axis=2) / dybu[:, :, np.newaxis] + np.diff(
ymom, axis=3) / dxbu[:, :, np.newaxis]
pv1x = -np.diff(xmom1 * dxcu[:, :, np.newaxis],
axis=2) / aq[:, :, np.newaxis]
pv1y = np.diff(ymom1 * dycv[:, :, np.newaxis], axis=3) / aq[:, :,
np.newaxis]
slyp = np.s_[:, :, ys:ye + 1, xs:xe]
ah1 = fhgeo.variables['Ah'][slyp[2:]]
slxmyp = np.s_[:, :, ys:ye + 1, xs - 1:xe]
uh = fh2.variables['uh'][slxmyp].filled(0).mean(axis=0, keepdims=True)
uhx = np.diff(uh, axis=3) / ah1
uhx = np.concatenate((uhx, uhx[:, :, :, -1:]), axis=3)
uhx = 0.25 * (uhx[:, :, :-1, :-1] + uhx[:, :, :-1, 1:] +
uhx[:, :, 1:, :-1] + uhx[:, :, 1:, 1:])
pv1y[:, :, :, :, 0] += pvhash * uhx
slymp = np.s_[:, :, ys - 1:ye + 1, xs:xe]
vh = fh2.variables['vh'][slymp].mean(axis=0, keepdims=True)
vhy = np.diff(vh, axis=2) / ah1
vhy = np.concatenate((vhy, vhy[:, :, :, -1:]), axis=3)
vhy = 0.25 * (vhy[:, :, :-1, :-1] + vhy[:, :, :-1, 1:] +
vhy[:, :, 1:, :-1] + vhy[:, :, 1:, 1:])
pv1x[:, :, :, :, 0] += pvhash * vhy
wd = fh2.variables['wd'][slyp].mean(axis=0, keepdims=True)
wdb = np.diff(wd, axis=1)
wdb = np.concatenate((wdb, wdb[:, :, :, -1:]), axis=3)
wdb = 0.25 * (wdb[:, :, :-1, :-1] + wdb[:, :, :-1, 1:] +
wdb[:, :, 1:, :-1] + wdb[:, :, 1:, 1:])
pv3 = pvhash * wdb
pv3 = pv3[:, :, :, :, np.newaxis]
#hq[hq<1] = np.nan
pvnew = np.concatenate(
(pv1y[:, :, :, :, :1], pv1x[:, :, :, :, :1], pv3,
pv1x[:, :, :, :, 1:-1], pv1y[:, :, :, :, 1:-1],
pv1x[:, :, :, :, -1:] + pv1y[:, :, :, :, -1:]),
axis=4) / hq[:, :, :, :, np.newaxis]
pv = np.ma.filled(pv.astype(np.float64), np.nan)
pvnew = np.ma.filled(pvnew.astype(np.float64), np.nan)
pvhash = np.ma.filled(pvhash.astype(np.float64), np.nan)
pv = np.nanmean(pv, meanax, keepdims=True)
pvnew = np.nanmean(pvnew, meanax, keepdims=True)
pvhash = np.nanmean(pvhash, meanax, keepdims=True)
X = dimq[keepax[1]]
Y = dimq[keepax[0]]
if 1 in keepax and not calledfromgeteddycoeffs:
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
em = (fh2.variables['e'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
em = np.nanmean(em, meanax, keepdims=True)
z = np.linspace(-3000, 0, 100)
Y = z
P = getvaratzc5(pv.astype(np.float32), z.astype(np.float32),
em.astype(np.float32)).squeeze()
Pnew = getvaratzc5(pvnew.astype(np.float32), z.astype(np.float32),
em.astype(np.float32)).squeeze()
pvhash = getvaratzc(pvhash.astype(np.float32),
z.astype(np.float32),
em.astype(np.float32)).squeeze()
if fil3:
swash = np.nanmean(swash, meanax, keepdims=True)
swash = getvaratzc(swash.astype(np.float32),
z.astype(np.float32),
em.astype(np.float32)).squeeze()
else:
P = pv
Pnew = pvnew
fhgeo.close()
fh.close()
fh2.close()
return (X, Y, P, pvhash, Pnew, swash)
def extract_twamomx_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
fil3=None,
alreadysaved=False,
xyasindices=False,
calledfrompv=False,
htol=1e-3):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimu = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xq',
yhyq='yh',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimu = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq')
sl = np.s_[:, :, ys:ye, xs:xe]
slmy = np.s_[:, :, ys - 1:ye, xs:xe]
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
Dforgetutwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetutwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][6:8]
dxcu, dycu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][0:2]
dycuforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][1:2]
dycuforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[2][1:2]
dxbu, dybu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][4:6]
aq1 = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[1][1]
ah1 = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[1][0]
dxcu1 = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye + 1)[2][0]
nt_const = dimu[0].size
t0 = time.time()
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
if fil3:
fh3 = mfdset(fil3)
slmytn = np.s_[-1:, :, ys - 1:ye, xs:xe]
# islayerdeep0 = fh3.variables['islayerdeep'][:,0,0,0].sum()
# islayerdeep = (fh3.variables['islayerdeep'][slmy].filled(np.nan)).sum(axis=0,
# keepdims=True)
islayerdeep0 = fh3.variables['islayerdeep'][-1:, 0, 0, 0]
islayerdeep = (fh3.variables['islayerdeep'][slmytn].filled(np.nan))
swash = (islayerdeep0 - islayerdeep) / islayerdeep0 * 100
swash = 0.5 * (swash[:, :, :-1, :] + swash[:, :, 1:, :])
fh3.close()
else:
swash = None
em = (fh2.variables['e'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
uh = (fh.variables['uh_masked'][0:, zs:ze, ys:ye, xs:xe].filled(np.nan)
* dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cu = (fh.variables['h_Cu'][0:, zs:ze, ys:ye, xs:xe].filled(0) *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
h_cu[h_cu < htol] = np.nan
h_um = h_cu
utwa = uh / h_cu / dycu
uhforxdiff = (fh.variables['uh_masked'][0:, zs:ze, ys:ye, xs - 1:xe] *
dt).filled(np.nan).sum(axis=0,
keepdims=True) / np.sum(dt)
h_cuforxdiff = (fh.variables['h_Cu'][0:, zs:ze, ys:ye, xs - 1:xe] *
dt).filled(0).sum(axis=0, keepdims=True) / np.sum(dt)
h_cuforxdiff[h_cuforxdiff < htol] = np.nan
utwaforxdiff = uhforxdiff / h_cuforxdiff #/dycuforxdiff
uhforydiff = (
fh.variables['uh_masked'][0:, zs:ze, ys - 1:ye + 1, xs:xe] *
dt).filled(np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
h_cuforydiff = (fh.variables['h_Cu'][0:, zs:ze, ys - 1:ye + 1, xs:xe] *
dt).filled(0).sum(axis=0, keepdims=True) / np.sum(dt)
h_cuforydiff[h_cuforydiff < htol] = np.nan
utwaforydiff = uhforydiff / h_cuforydiff #/dycuforydiff
utwax = np.diff(np.nan_to_num(utwaforxdiff), axis=3) / dxt / dyt
utwax = np.concatenate((utwax, -utwax[:, :, :, [-1]]), axis=3)
utwax = 0.5 * (utwax[:, :, :, 0:-1] + utwax[:, :, :, 1:])
utway = np.diff(utwaforydiff, axis=2) / dxbu / dybu
utway = 0.5 * (utway[:, :, 0:-1, :] + utway[:, :, 1:, :])
humx = np.diff(np.nan_to_num(uhforxdiff), axis=3) / dxt / dyt
humx = np.concatenate((humx, -humx[:, :, :, [-1]]), axis=3)
humx = 0.5 * (humx[:, :, :, 0:-1] + humx[:, :, :, 1:])
hvm = (fh.variables['vh_masked'][0:, zs:ze, ys - 1:ye, xs:xe] *
dt).sum(axis=0, keepdims=True) / np.sum(dt)
hvm = np.concatenate((hvm, -hvm[:, :, :, -1:]), axis=3)
hvm = 0.25 * (hvm[:, :, :-1, :-1] + hvm[:, :, :-1, 1:] +
hvm[:, :, 1:, :-1] + hvm[:, :, 1:, 1:]) / dxcu
hv = (fh.variables['vh_masked'][0:, zs:ze, ys - 1:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hvmy = np.diff(hv, axis=2) / dxt / dyt
hvmy = np.concatenate((hvmy, -hvmy[:, :, :, -1:]), axis=3)
hvmy = 0.5 * (hvmy[:, :, :, :-1] + hvmy[:, :, :, 1:])
huuxphuvym = (
fh.variables['twa_huuxpt'][0:, zs:ze, ys:ye, xs:xe] * dt +
fh.variables['twa_huvymt'][0:, zs:ze, ys:ye, xs:xe] * dt).filled(
np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
#u = (fh.variables['u_masked'][0:,zs:ze,ys:ye,xs-1:xe]*dt).filled(np.nan).sum(axis=0,keepdims=True)/np.sum(dt)
huu = (fh.variables['huu_Cu'][0:, zs:ze, ys:ye, xs - 1:xe] *
dt).filled(np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
huuxm = np.diff(np.nan_to_num(huu), axis=3) / dxt / dyt
huuxm = np.concatenate((huuxm, -huuxm[:, :, :, -1:]), axis=3)
huuxm = 0.5 * (huuxm[:, :, :, :-1] + huuxm[:, :, :, 1:])
# huu = (fh.variables['huu_T'][0:,zs:ze,ys:ye,xs:xe]*dt).sum(axis=0,keepdims=True)/np.sum(dt)*dyt
# huu = np.concatenate((huu,-huu[:,:,:,-1:]),axis=3)
# huuxm = np.diff(huu,axis=3)/dxcu/dycu
huvym = huuxphuvym + huuxm
utwaforvdiff = np.concatenate((utwa[:, [0], :, :], utwa), axis=1)
utwab = np.diff(utwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
utwab = np.concatenate((utwab, np.zeros(utwab[:, :1, :, :].shape)),
axis=1)
utwab = 0.5 * (utwab[:, 0:-1, :, :] + utwab[:, 1:, :, :])
hwb = (fh2.variables['wd'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hwb = np.diff(hwb, axis=1)
hwb = np.concatenate((hwb, -hwb[:, :, :, -1:]), axis=3)
hwb_u = 0.5 * (hwb[:, :, :, :-1] + hwb[:, :, :, 1:])
hwb = (fh2.variables['wd'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hwm = 0.5 * (hwb[:, :-1] + hwb[:, 1:]) * dbl[:, np.newaxis, np.newaxis]
hwm = np.concatenate((hwm, -hwm[:, :, :, -1:]), axis=3)
hwm_u = 0.5 * (hwm[:, :, :, :-1] + hwm[:, :, :, 1:])
esq = (fh.variables['esq'][0:, zs:ze, ys:ye, xs:xe] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
edlsqm = (esq - elm**2)
edlsqm = np.concatenate((edlsqm, edlsqm[:, :, :, -1:]), axis=3)
edlsqmx = np.diff(edlsqm, axis=3) / dxcu
hpfu = (fh.variables['twa_hpfu'][0:, zs:ze, ys:ye, xs:xe] * dt).filled(
np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
pfum = (fh2.variables['PFu'][0:, zs:ze, ys:ye, xs:xe] * dt).filled(
np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
edpfudmb = -hpfu + h_cu * pfum - 0.5 * edlsqmx * dbl[:, np.newaxis,
np.newaxis]
hfvm = (fh.variables['twa_hfv'][0:, zs:ze, ys:ye, xs:xe] * dt).filled(
np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
huwbm = (fh.variables['twa_huwb'][0:, zs:ze, ys:ye, xs:xe] *
dt).filled(np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
hdiffum = (fh.variables['twa_hdiffu'][0:, zs:ze, ys:ye, xs:xe] *
dt).filled(np.nan).sum(axis=0, keepdims=True) / np.sum(dt)
hdudtviscm = (fh.variables['twa_hdudtvisc'][0:, zs:ze, ys:ye, xs:xe] *
dt).filled(np.nan).sum(axis=0,
keepdims=True) / np.sum(dt)
fh2.close()
fh.close()
advx = utwa * utwax
advy = hvm * utway / h_um
advb = hwm_u * utwab / h_um
cor = hfvm / h_um
pfum = pfum
xdivep1 = -huuxm / h_um
xdivep2 = advx
xdivep3 = utwa * humx / h_um
xdivep4 = -0.5 * edlsqmx * dbl[:, np.newaxis, np.newaxis] / h_um
xdivep = (xdivep1 + xdivep2 + xdivep3 + xdivep4)
ydivep1 = huvym / h_um
ydivep2 = advy
ydivep3 = utwa * hvmy / h_um
ydivep = (ydivep1 + ydivep2 + ydivep3)
bdivep1 = huwbm / h_um
bdivep2 = advb
bdivep3 = utwa * hwb_u / h_um
bdivep4 = -edpfudmb / h_um
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
X1twa = hdiffum / h_um
X2twa = hdudtviscm / h_um
terms = np.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfum[:, :, :, :, np.newaxis], xdivep[:, :, :, :, np.newaxis],
ydivep[:, :, :, :, np.newaxis], bdivep[:, :, :, :, np.newaxis],
X1twa[:, :, :, :, np.newaxis], X2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.concatenate(
(xdivep1[:, :, :, :, np.newaxis], xdivep3[:, :, :, :, np.newaxis],
xdivep4[:, :, :, :, np.newaxis], ydivep1[:, :, :, :, np.newaxis],
ydivep3[:, :, :, :, np.newaxis], bdivep1[:, :, :, :, np.newaxis],
bdivep3[:, :, :, :, np.newaxis], bdivep4[:, :, :, :, np.newaxis]),
axis=4)
termsm = np.nanmean(terms, axis=meanax, keepdims=True)
termsepm = np.nanmean(termsep, axis=meanax, keepdims=True)
X = dimu[keepax[1]]
Y = dimu[keepax[0]]
if 1 in keepax and not calledfrompv:
em = np.nanmean(em, axis=meanax, keepdims=True)
elm = np.nanmean(elm, axis=meanax, keepdims=True)
z = np.linspace(-3000, 0, 100)
Y = z
P = getvaratzc5(termsm.astype(np.float32), z.astype(np.float32),
em.astype(np.float32))
Pep = getvaratzc5(termsepm.astype(np.float32),
z.astype(np.float32), em.astype(np.float32))
if fil3:
swash = np.nanmean(swash, meanax, keepdims=True)
swash = getvaratzc(swash.astype(np.float32),
z.astype(np.float32),
em.astype(np.float32)).squeeze()
else:
P = termsm.squeeze()
Pep = termsepm.squeeze()
if not calledfrompv:
np.savez('twamomx_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomx_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep, swash, em.squeeze())
def extract_cb_terms(geofil,fil,xstart,xend,ystart,yend,zs,ze,meanax):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i,)
fh = mfdset(fil)
(xs,xe),(ys,ye),dimh = rdp1.getlatlonindx(fh,wlon=xstart,elon=xend,
slat=ystart, nlat=yend,zs=zs,ze=ze)
fhgeo = dset(geofil)
D, (ah,aq) = rdp1.getgeombyindx(fhgeo,xs,xe,ys,ye)[0:2]
fhgeo.close()
nt = dimh[0].size
t0 = time.time()
uh = fh.variables['uh'][0:1,zs:ze,ys:ye,xs-1:xe]
vh = fh.variables['vh'][0:1,zs:ze,ys-1:ye,xs:xe]
em = fh.variables['e'][0:1,zs:ze,ys:ye,xs:xe]/nt
#dhdtm = fh.variables['dhdt'][0:1,zs:ze,ys:ye,xs:xe]/nt
wd = fh.variables['wd'][0:1,zs:ze,ys:ye,xs:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh,axis = 3)/ah
uhxm = uhx/nt
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh,axis = 2)/ah
vhym = vhy/nt
wdm = np.diff(wd,axis=1)/nt
for i in range(1,nt):
uh = fh.variables['uh'][i:i+1,zs:ze,ys:ye,xs-1:xe]
vh = fh.variables['vh'][i:i+1,zs:ze,ys-1:ye,xs:xe]
em += fh.variables['e'][i:i+1,zs:ze,ys:ye,xs:xe]/nt
#dhdtm += fh.variables['dhdt'][i:i+1,zs:ze,ys:ye,xs:xe]/nt
wd = fh.variables['wd'][i:i+1,zs:ze,ys:ye,xs:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh,axis = 3)/ah
uhxm += uhx/nt
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh,axis = 2)/ah
vhym += vhy/nt
wdm += np.diff(wd,axis=1)/nt
sys.stdout.write('\r'+str(int((i+1)/nt*100))+'% done...')
sys.stdout.flush()
fh.close()
print('Total reading time: {}s'.format(time.time()-t0))
terms = np.ma.concatenate(( uhxm[:,:,:,:,np.newaxis],
vhym[:,:,:,:,np.newaxis],
wdm[:,:,:,:,np.newaxis]),axis=4)
termsm = np.ma.apply_over_axes(np.mean, terms, meanax)
elm = 0.5*(em[:,0:-1,:,:]+em[:,1:,:,:])
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
X = dimh[keepax[1]]
Y = dimh[keepax[0]]
if 1 in keepax:
Y = elm.squeeze()
X = np.meshgrid(X,dimh[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
np.savez('cb_terms', X=X,Y=Y,P=P)
return (X,Y,P)
def extract_momx_terms(geofil,
fil,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
alreadysaved=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhgeo = dset(geofil)
fh = mfdset(fil)
(xs, xe), (ys, ye), dimu = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq')
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
fhgeo.close()
nt = dimu[0].size
t0 = time.time()
print('Reading data using loop...')
#dudtm = fh.variables['dudt'][0:1,zs:ze,ys:ye,xs:xe]/nt
caum = fh.variables['CAu'][0:1, zs:ze, ys:ye, xs:xe] / nt
pfum = fh.variables['PFu'][0:1, zs:ze, ys:ye, xs:xe] / nt
dudtviscm = fh.variables['du_dt_visc'][0:1, zs:ze, ys:ye, xs:xe] / nt
diffum = fh.variables['diffu'][0:1, zs:ze, ys:ye, xs:xe] / nt
dudtdiam = fh.variables['dudt_dia'][0:1, zs:ze, ys:ye, xs:xe] / nt
if 1 in keepax:
em = fh.variables['e'][0:1, zs:ze, ys:ye, xs:xe] / nt
print(caum.shape, em.shape)
for i in range(1, nt):
#dudtm += fh.variables['dudt'][i:i+1,zs:ze,ys:ye,xs:xe]/nt
caum += fh.variables['CAu'][i:i + 1, zs:ze, ys:ye, xs:xe] / nt
pfum += fh.variables['PFu'][i:i + 1, zs:ze, ys:ye, xs:xe] / nt
dudtviscm += fh.variables['du_dt_visc'][i:i + 1, zs:ze, ys:ye,
xs:xe] / nt
diffum += fh.variables['diffu'][i:i + 1, zs:ze, ys:ye, xs:xe] / nt
dudtdiam += fh.variables['dudt_dia'][i:i + 1, zs:ze, ys:ye,
xs:xe] / nt
if 1 in keepax:
em += fh.variables['e'][i:i + 1, zs:ze, ys:ye, xs:xe] / nt
sys.stdout.write('\r' + str(int((i + 1) / nt * 100)) + '% done...')
sys.stdout.flush()
fh.close()
print('Time taken for data reading: {}s'.format(time.time() - t0))
terms = np.ma.concatenate(
(caum[:, :, :, :, np.newaxis], dudtdiam[:, :, :, :, np.newaxis],
pfum[:, :, :, :, np.newaxis], diffum[:, :, :, :, np.newaxis],
dudtviscm[:, :, :, :, np.newaxis]),
axis=4)
terms = terms.filled(0)
termsm = np.ma.apply_over_axes(np.mean, terms, meanax)
X = dimu[keepax[1]]
Y = dimu[keepax[0]]
if 1 in keepax:
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimu[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
np.savez('momx_terms', X=X, Y=Y, P=P)
else:
npzfile = np.load('momx_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
return (X, Y, P)
def extract_twamomx_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
alreadysaved=False,
xyasindices=False,
calledfrompv=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimu = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xq',
yhyq='yh',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimu = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq')
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
Dforgetutwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetutwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][6:8]
dxcu, dycu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][0:2]
dycuforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][1:2]
dycuforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[2][1:2]
dybu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][5]
dyt1 = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][7]
nt_const = dimu[0].size
t0 = time.time()
em = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
uh = fh2.variables['uh'][0:, zs:ze, ys:ye, xs:xe]
h_cu = fh.variables['h_Cu'][0:, zs:ze, ys:ye, xs:xe]
h_um = h_cu
utwa = uh / h_cu / dycu
uhforxdiff = fh2.variables['uh'][0:, zs:ze, ys:ye, xs - 1:xe]
h_cuforxdiff = fh.variables['h_Cu'][0:, zs:ze, ys:ye, xs - 1:xe]
utwaforxdiff = uhforxdiff / h_cuforxdiff / dycuforxdiff
uhforydiff = fh2.variables['uh'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
h_cuforydiff = fh.variables['h_Cu'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
utwaforydiff = uhforydiff / h_cuforydiff / dycuforydiff
utwax = np.diff(utwaforxdiff.filled(0), axis=3) / dxt
utwax = np.concatenate((utwax, -utwax[:, :, :, [-1]]), axis=3)
utwax = 0.5 * (utwax[:, :, :, 0:-1] + utwax[:, :, :, 1:])
utway = np.diff(utwaforydiff, axis=2) / dybu
utway = 0.5 * (utway[:, :, 0:-1, :] + utway[:, :, 1:, :])
humx = np.diff(uhforxdiff.filled(0) / dycuforxdiff, axis=3) / dxt
humx = np.concatenate((humx, -humx[:, :, :, [-1]]), axis=3)
humx = 0.5 * (humx[:, :, :, 0:-1] + humx[:, :, :, 1:])
hvm = fh.variables['hv_Cu'][0:, zs:ze, ys:ye, xs:xe]
hv_cu = fh.variables['hv_Cu'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
hvmy = np.diff(hv_cu, axis=2) / dyt1
hvmy = 0.5 * (hvmy[:, :, :-1, :] + hvmy[:, :, 1:, :])
huuxphuvym = fh.variables['twa_huuxpt'][
0:, zs:ze, ys:ye, xs:xe] + fh.variables['twa_huvymt'][0:, zs:ze,
ys:ye, xs:xe]
huu = fh.variables['huu_Cu'][0:, zs:ze, ys:ye, xs - 1:xe].filled(0)
huuxm = np.diff(huu, axis=3) / dxt
huuxm = np.concatenate((huuxm, -huuxm[:, :, :, -1:]), axis=3)
huuxm = 0.5 * (huuxm[:, :, :, :-1] + huuxm[:, :, :, 1:])
huvym = huuxphuvym - huuxm
utwaforvdiff = np.concatenate((utwa[:, [0], :, :], utwa), axis=1)
utwab = np.diff(utwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
utwab = np.concatenate(
(utwab,
np.zeros([utwab.shape[0], 1, utwab.shape[2], utwab.shape[3]])),
axis=1)
utwab = 0.5 * (utwab[:, 0:-1, :, :] + utwab[:, 1:, :, :])
hwb_u = fh.variables['hwb_Cu'][0:, zs:ze, ys:ye, xs:xe]
hwm_u = fh.variables['hw_Cu'][0:, zs:ze, ys:ye, xs:xe]
esq = fh.variables['esq_Cu'][0:, zs:ze, ys:ye, xs - 1:xe].filled(0)
ecu = fh.variables['e_Cu'][0:, zs:ze, ys:ye, xs - 1:xe].filled(0)
edlsqm = (esq - ecu**2)
edlsqmx = np.diff(edlsqm, axis=3) / dxt
edlsqmx = np.concatenate((edlsqmx, -edlsqmx[:, :, :, [-1]]), axis=3)
edlsqmx = 0.5 * (edlsqmx[:, :, :, :-1] + edlsqmx[:, :, :, 1:])
epfu = fh.variables['epfu'][0:, zs:ze, ys:ye, xs:xe]
ecu = fh.variables['e_Cu'][0:, zs:ze, ys:ye, xs:xe]
pfum = fh.variables['pfu_masked'][0:, zs:ze, ys:ye, xs:xe]
edpfudm = epfu - pfum * ecu
edpfudmb = np.diff(edpfudm, axis=1)
edpfudmb = np.concatenate(
(edpfudmb[:, :1, :, :], edpfudmb, edpfudmb[:, -1:, :, :]), axis=1)
edpfudmb = 0.5 * (edpfudmb[:, :-1, :, :] + edpfudmb[:, 1:, :, :])
hfvm = fh.variables['twa_hfv'][0:1, zs:ze, ys:ye, xs:xe]
huwbm = fh.variables['twa_huwb'][0:1, zs:ze, ys:ye, xs:xe]
hdiffum = fh.variables['twa_hdiffu'][0:1, zs:ze, ys:ye, xs:xe]
hdudtviscm = fh.variables['twa_hdudtvisc'][0:1, zs:ze, ys:ye, xs:xe]
fh2.close()
fh.close()
advx = utwa * utwax
advy = hvm * utway / h_um
advb = hwm_u * utwab / h_um
cor = hfvm / h_um
pfum = pfum
xdivep1 = -huuxm / h_um
xdivep2 = -advx
xdivep3 = -utwa * humx / h_um
xdivep4 = 0.5 * edlsqmx * dbl[:, np.newaxis, np.newaxis] / h_um
xdivep = (xdivep1 + xdivep2 + xdivep3 + xdivep4)
ydivep1 = -huvym / h_um
ydivep2 = -advy
ydivep3 = -utwa * hvmy / h_um
ydivep = (ydivep1 + ydivep2 + ydivep3)
bdivep1 = -huwbm / h_um
bdivep2 = -advb
bdivep3 = -utwa * hwb_u / h_um
bdivep4 = edpfudmb / h_um
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
X1twa = hdiffum / h_um
X2twa = hdudtviscm / h_um
terms = np.ma.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfum[:, :, :, :, np.newaxis], -xdivep[:, :, :, :, np.newaxis],
-ydivep[:, :, :, :, np.newaxis], -bdivep[:, :, :, :, np.newaxis],
X1twa[:, :, :, :, np.newaxis], X2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.ma.concatenate((
-xdivep1[:, :, :, :, np.newaxis], -xdivep3[:, :, :, :, np.newaxis],
-xdivep4[:, :, :, :, np.newaxis], -ydivep1[:, :, :, :, np.newaxis],
-ydivep3[:, :, :, :, np.newaxis], -bdivep1[:, :, :, :, np.newaxis],
-bdivep3[:, :, :, :, np.newaxis],
-bdivep4[:, :, :, :, np.newaxis]),
axis=4)
terms[np.isinf(terms)] = np.nan
termsep[np.isinf(termsep)] = np.nan
termsm = np.ma.apply_over_axes(np.nanmean, terms, meanax)
termsepm = np.ma.apply_over_axes(np.nanmean, termsep, meanax)
X = dimu[keepax[1]]
Y = dimu[keepax[0]]
if 1 in keepax:
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimu[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
Pep = termsepm.squeeze()
Pep = np.ma.filled(Pep.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
if not calledfrompv:
np.savez('twamomx_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomx_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep)
def get_heddy_coeffs_fromflx(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs=0,
ze=None,
z=np.linspace(-3000, 0, 100),
perc=5,
htol=1e-3):
fhgeo = dset(geofil)
fhvgeo = dset(vgeofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
db = fhvgeo.variables['g'][:]
fhvgeo.close()
dbi = np.append(db, 0)
zi = fh.variables['zi'][:]
zl = fh.variables['zl'][:]
dbl = np.diff(zi) * 9.8 / 1031
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
(sl[3].start - 1):sl[3].stop]
slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
slmpy = np.s_[:, :, (sl[2].start - 1):(sl[2].stop + 1), sl[3]]
sl2d = sl[2:]
slmx2d = slmx[2:]
slpx2d = slpx[2:]
slpy2d = slpy[2:]
dxbu = fhgeo.variables['dxBu'][slmx2d]
dybu = fhgeo.variables['dyBu'][slmx2d]
dxcv = fhgeo.variables['dxCv'][sl2d]
dycv = fhgeo.variables['dyCv'][sl2d]
dxt = fhgeo.variables['dxT'][slpy2d]
dyt = fhgeo.variables['dyT'][slpy2d]
v, _ = pv.getvtwa(fhgeo, fh, fh2, slmx)
vx = np.diff(v, axis=3) / dxbu
vx = vx[:, :, :, 1:]
e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
axis=0, keepdims=True)
eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
e = fh2.variables['e'][slmxpy]
eatvmx = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
vh = (fh.variables['vh_masked'][sl] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
h_cv = (fh.variables['h_Cv'][sl] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
h_cv[h_cv < htol] = np.nan
sig = h_cv / dbl[:, np.newaxis, np.newaxis]
h_vm = h_cv
vtwa = vh / h_cv / dxcv
vhforxdiff = (fh.variables['vh_masked'][slmx] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
h_cvforxdiff = (fh.variables['h_Cv'][slmx] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
h_cvforxdiff[h_cvforxdiff < htol] = np.nan
vtwaforxdiff = vhforxdiff / h_cvforxdiff
vtwaforxdiff = np.concatenate((vtwaforxdiff, vtwaforxdiff[:, :, :, -1:]),
axis=3)
vtwax = np.diff(vtwaforxdiff, axis=3) / dxbu / dybu
vtwax = 0.5 * (vtwax[:, :, :, :-1] + vtwax[:, :, :, 1:])
uh = (fh.variables['uh_masked'][slmxpy] * dt).filled(0).sum(
axis=0, keepdims=True) / np.sum(dt)
hum = 0.25 * (uh[:, :, :-1, :-1] + uh[:, :, :-1, 1:] + uh[:, :, 1:, :-1] +
uh[:, :, 1:, 1:]) / dycv
huvxphvvym = (fh.variables['twa_huvxpt'][sl] * dt +
fh.variables['twa_hvvymt'][sl] * dt).sum(
axis=0, keepdims=True) / np.sum(dt)
hvv = (fh.variables['hvv_Cv'][slmpy] * dt).sum(axis=0,
keepdims=True) / np.sum(dt)
hvvym = np.diff(hvv, axis=2) / dxt / dyt
hvvym = 0.5 * (hvvym[:, :, :-1, :] + hvvym[:, :, 1:, :])
huvxm = -(huvxphvvym + hvvym)
advx = hum * vtwax / h_vm
humx = np.diff(np.nan_to_num(uh), axis=3) / dxt / dyt
humx = 0.5 * (humx[:, :, :-1, :] + humx[:, :, 1:, :])
xdivep1 = -huvxm / h_vm
xdivep2 = advx
xdivep3 = vtwa * humx / h_vm
xdivep = (xdivep1 + xdivep2 + xdivep3)
xdivep *= sig
uv = sint.cumtrapz(
xdivep[:, :, :, ::-1], dx=-dxbu[:, 1:-1], axis=3,
initial=0)[:, :, :, ::-1] / sig
uvm = np.apply_over_axes(np.nanmean, uv, (0, 2))
vxm = np.apply_over_axes(np.nanmean, vx, (0, 2))
uvm = getvaratzc(uvm.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
vxm = getvaratzc(vxm.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
cmax = np.fabs(np.percentile(uvm, (1, 99))).max()
im = ax[0].pcolormesh(dimv[3],
z,
uvm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
cmax = np.fabs(np.percentile(vxm, (1, 99))).max()
im = ax[1].pcolormesh(dimv[3],
z,
vxm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
for axc in ax:
xdegtokm(axc, 0.5 * (ystart + yend))
axc.set_xlabel('x (km)')
axc.grid()
ax[0].set_ylabel('z (m)')
fig2, ax = plt.subplots(1, 1, sharey=True)
ax.linfit(vx, uv)
ax.set_xlabel('v_x (m/s)')
ax.set_ylabel('RS (m2/s2)')
ax.grid()
return fig, fig2
def getuv(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
xyasindices=False):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
slh, dimh = rdp1.getslice(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze)
ys, ye = slh[2].start, slh[2].stop
xs, xe = slh[3].start, slh[3].stop
slhmx = np.s_[:, zs:ze, ys:ye, xs - 1:xe]
slhmy = np.s_[:, zs:ze, ys - 1:ye, xs:xe]
slhmpy = np.s_[:, zs:ze, ys - 1:ye + 1, xs:xe]
slu, dimu = rdp1.getslice(fh,
xstart,
xend,
ystart,
yend,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xq',
yhyq='yh',
zlzi='zl')
slv, dimv = rdp1.getslice(fh,
xstart,
xend,
ystart,
yend,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xh',
yhyq='yq',
zlzi='zl')
ys, ye = slv[2].start, slv[2].stop
xs, xe = slv[3].start, slv[3].stop
slvpy = np.s_[:, zs:ze, ys:ye + 1, xs:xe]
uh = getvaravg(fh2, 'uh', slu)
u = getvaravg(fh2, 'u', slu).filled(np.nan)
h_cu = getvaravg(fh, 'h_Cu', slu)
h_cu = np.ma.masked_array(h_cu, mask=(h_cu < 1e-3))
dycu = fhgeo.variables['dyCu'][slu[2:]]
utwa = uh / h_cu / dycu
vh = getvaravg(fh2, 'vh', slv)
v = getvaravg(fh2, 'v', slv)
h_cv = getvaravg(fh, 'h_Cv', slv)
h_cv = np.ma.masked_array(h_cv, mask=(h_cv < 1e-3))
dxcv = fhgeo.variables['dxCv'][slv[2:]]
vtwa = vh / dxcv / h_cv
emforxdiff = getvaravg(fh2, 'e', slhmx)
elmforxdiff = 0.5 * (emforxdiff[:, 0:-1, :, :] + emforxdiff[:, 1:, :, :])
elmforxdiff = np.concatenate((elmforxdiff, elmforxdiff[:, :, :, -1:]),
axis=3)
dxcuforxdiff = fhgeo.variables['dxCu'][slhmx[2:]]
ex = np.diff(elmforxdiff, axis=3) / dxcuforxdiff
uh = getvaravg(fh2, 'uh', slhmx)
h_cu = getvaravg(fh, 'h_Cu', slhmx)
h_cu = np.ma.masked_array(h_cu, mask=(h_cu < 1e-3))
dycu = fhgeo.variables['dyCu'][slhmx[2:]]
uzx = (uh / h_cu / dycu).filled(0) * ex
uzx = 0.5 * (uzx[:, :, :, 1:] + uzx[:, :, :, :-1])
emforydiff = getvaravg(fh2, 'e', slhmpy)
elmforydiff = 0.5 * (emforydiff[:, 0:-1, :, :] + emforydiff[:, 1:, :, :])
dycv = fhgeo.variables['dyCv'][slhmy[2:]]
ey = np.diff(elmforydiff, axis=2) / dycv
vh = getvaravg(fh2, 'vh', slhmy)
h_cv = getvaravg(fh, 'h_Cv', slhmy)
h_cv = np.ma.masked_array(h_cv, mask=(h_cv < 1e-3))
dxcv = fhgeo.variables['dxCv'][slhmy[2:]]
vtwa = vh / dxcv / h_cv
vzy = vtwa * ey
vtwa = 0.5 * (vtwa[:, :, 1:, :] + vtwa[:, :, :-1, :])
vzy = 0.5 * (vzy[:, :, 1:, :] + vzy[:, :, :-1, :])
wd = getvaravg(fh2, 'wd', slh)
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 1:, :, :] + hw[:, :-1, :, :])
wzb = hw / dbl[:, np.newaxis, np.newaxis]
whash = uzx + vzy + wzb
terms = [utwa, vtwa, whash, u, v]
slices = [slu, slvpy, slh, slu, slvpy]
X = [
dimu[keepax[1]], dimv[keepax[1]], dimh[keepax[1]], dimu[keepax[1]],
dimv[keepax[1]]
]
Y = [
dimu[keepax[0]], dimv[keepax[0]], dimh[keepax[0]], dimu[keepax[1]],
dimv[keepax[1]]
]
termsm = []
for item in terms:
try:
item = item.filled(np.nan)
except AttributeError:
item = item
termsm.append(np.ma.apply_over_axes(np.nanmean, item, meanax))
if 1 in keepax:
Y = []
for i in range(len(terms)):
z = np.linspace(-3000, 0, 100)
em = fh2.variables['e'][slices[i]]
if i == 0 or i == 3:
em = np.concatenate((em, em[:, :, :, -1:]), axis=3)
em = 0.5 * (em[:, :, :, :-1] + em[:, :, :, 1:])
elif i == 1 or i == 4:
em = 0.5 * (em[:, :, :-1, :] + em[:, :, 1:, :])
em = np.ma.apply_over_axes(np.mean, em, meanax)
termsm[i] = getvaratzc(termsm[i].astype(np.float32),
z.astype(np.float32), em.astype(np.float32))
Y.append(z)
fh2.close()
fh.close()
fhgeo.close()
P = []
for i, item in enumerate(termsm):
P.append(item)
X[i] = np.ma.filled(X[i].astype(float), np.nan)
Y[i] = np.ma.filled(Y[i].astype(float), np.nan)
return X, Y, P
def get_deddy_coeffs(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs=0,
ze=None,
z=np.linspace(-1500, 0, 100),
perc=5):
fhgeo = dset(geofil)
fhvgeo = dset(vgeofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
db = fhvgeo.variables['g'][:]
fhvgeo.close()
dbi = np.append(db, 0)
zi = fh.variables['zi'][:]
zl = fh.variables['zl'][:]
dbl = np.diff(zi) * 9.8 / 1031
sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
sl2d = sl[2:]
slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
f = fhgeo.variables['f'][sl2d]
e = fh2.variables['e'][slpy]
eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
v, h = pv.getvtwa(fhgeo, fh, fh2, sl)
v = v[:, :, :, :-1]
h = h[:, :, :, :-1]
sig = h / dbl[:, np.newaxis, np.newaxis]
sigaz = getvaratzc(sig.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
vb = -np.diff(v, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
vb = np.concatenate((vb[:, :1, :, :], vb, vb[:, -1:, :, :]), axis=1)
vbb = -np.diff(vb, axis=1) / dbl[:, np.newaxis, np.newaxis]
vbb = getvaratzc(vbb.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
vi = np.concatenate((v[:, :1, :, :], v, -v[:, -1:, :, :]), axis=1)
vi = 0.5 * (vi[:, :-1, :, :] + vi[:, 1:, :, :])
vib = -np.diff(vi, axis=1) / dbl[:, np.newaxis, np.newaxis]
svib = sig * vib
svibb = -np.diff(svib, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
svibb = np.concatenate((svibb[:, :1, :, :], svibb, svibb[:, -1:, :, :]),
axis=1)
svibb = 0.5 * (svibb[:, :-1, :, :] + svibb[:, 1:, :, :])
svibb = getvaratzc(svibb.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
sv = sig * v
svb = -np.diff(sv, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
svb = np.concatenate((svb[:, :1, :, :], svb, svb[:, -1:, :, :]), axis=1)
svbb = -np.diff(svb, axis=1) / dbl[:, np.newaxis, np.newaxis]
svbb = getvaratzc(svbb.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
x, y, P, _, _ = py.extract_twamomy_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze, (0, ),
z=z)
fdrag = P[:, :, :, :, 7]
vbbm = np.apply_over_axes(np.nanmean, vbb, (0, 2))
svibbm = np.apply_over_axes(np.nanmean, svibb, (0, 2))
svbbm = np.apply_over_axes(np.nanmean, svbb, (0, 2))
fdragm = np.apply_over_axes(np.nanmean, fdrag, (0, 2))
fig, ax = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(12, 3))
cmax = np.fabs(np.percentile(vbbm, (1, 99))).max()
im = ax[0].pcolormesh(dimv[3],
z,
vbbm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[0])
cmax = np.fabs(np.percentile(svibbm, (1, 99))).max()
im = ax[1].pcolormesh(dimv[3],
z,
svibbm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[1])
cmax = np.fabs(np.percentile(svbbm, (1, 99))).max()
im = ax[2].pcolormesh(dimv[3],
z,
svbbm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[2])
cmax = np.fabs(np.percentile(fdragm, (1, 99))).max()
im = ax[3].pcolormesh(dimv[3],
z,
fdragm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[3])
fig2, ax = plt.subplots(1, 3, sharey=True, figsize=(12, 3))
ax[0].linfit(vbb, sigaz * fdrag / f**2)
ax[0].set_xlabel(r'$v_{bb}$')
ax[1].linfit(svibb, sigaz * fdrag / f**2)
ax[1].set_xlabel(r'$(\sigma v_b)_b$')
ax[2].linfit(svbb, sigaz * fdrag / f**2)
ax[2].set_xlabel(r'$(\sigma v)_{bb}$')
ax[0].set_ylabel('Form Drag')
for axc in ax:
axc.grid()
return fig, fig2
def extract_uv(geofil,
fil,
fil2,
fil3,
xstart,
xend,
ystart,
yend,
zs,
ze,
savfil=None,
utwamaxscalefactor=1,
vtwamaxscalefactor=1):
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
fhgeo = dset(geofil)
(xs, xe), (ys, ye), dimutwa = rdp1.getlatlonindx(
fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq',
yhyq='yh')
em = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
elm = np.ma.apply_over_axes(np.mean, elm, (0, 2))
dycu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][1]
uh = fh.variables['uh'][0:, zs:ze, ys:ye, xs:xe]
h_cu = fh2.variables['h_Cu'][0:, zs:ze, ys:ye, xs:xe]
utwa = uh / h_cu / dycu
fig = plt.figure()
ax = fig.add_subplot(221)
X = dimutwa[3]
X = np.meshgrid(X, dimutwa[1])[0]
Y = elm.squeeze()
utwa = np.ma.apply_over_axes(np.nanmean, utwa, (0, 2))
utwa = utwa.squeeze()
cmax = np.nanmax(np.absolute(utwa)) * utwamaxscalefactor
im = m6plot(
(X, Y, utwa),
ax=ax,
ylabel='z (m)',
txt=r'$\hat{u}$',
vmin=-cmax,
vmax=cmax,
cmap='RdBu_r',
ylim=(-3000, 0))
ax.set_xticklabels([])
(xs, xe), (ys, ye), dimvtwa = rdp1.getlatlonindx(
fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xh',
yhyq='yq')
em = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
elm = np.ma.apply_over_axes(np.mean, elm, (0, 2))
dxcv = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][2]
vh = fh.variables['vh'][0:, zs:ze, ys:ye, xs:xe]
h_cv = fh2.variables['h_Cv'][0:, zs:ze, ys:ye, xs:xe]
vtwa = vh / h_cv / dxcv
ax = fig.add_subplot(222)
X = dimvtwa[3]
X = np.meshgrid(X, dimvtwa[1])[0]
Y = elm.squeeze()
vtwa = np.ma.apply_over_axes(np.nanmean, vtwa, (0, 2))
vtwa = vtwa.squeeze()
cmax = np.nanmax(np.absolute(vtwa)) * vtwamaxscalefactor
im = m6plot(
(X, Y, vtwa),
ax=ax,
txt=r'$\hat{v}$',
vmin=-cmax,
vmax=cmax,
cmap='RdBu_r',
ylim=(-3000, 0))
ax.set_xticklabels([])
ax.set_yticklabels([])
fh2.close()
fh.close()
fhgeo.close()
fh3 = mfdset(fil3)
(xs, xe), (ys, ye), dimu = rdp1.getlatlonindx(
fh3,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq',
yhyq='yh',
zlzi='zlremap')
u = fh3.variables['u'][0:, zs:ze, ys:ye, xs:xe]
ax = fig.add_subplot(223)
X = dimu[3]
Y = dimu[1]
u = np.ma.apply_over_axes(np.nanmean, u, (0, 2))
u = u.squeeze()
cmax = np.nanmax(np.absolute(u))
im = m6plot(
(X, -Y, u),
ax=ax,
txt='u',
ylabel='z (m)',
vmin=-cmax,
vmax=cmax,
cmap='RdBu_r')
xdegtokm(ax, 0.5 * (ystart + yend))
(xs, xe), (ys, ye), dimv = rdp1.getlatlonindx(
fh3,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xh',
yhyq='yq',
zlzi='zlremap')
v = fh3.variables['v'][0:, zs:ze, ys:ye, xs:xe]
ax = fig.add_subplot(224)
X = dimv[3]
Y = dimv[1]
v = np.ma.apply_over_axes(np.nanmean, v, (0, 2))
v = v.squeeze()
cmax = np.nanmax(np.absolute(v))
im = m6plot(
(X, -Y, v), ax=ax, txt='v', vmin=-cmax, vmax=cmax, cmap='RdBu_r')
xdegtokm(ax, 0.5 * (ystart + yend))
ax.set_yticklabels([])
fh3.close()
if savfil:
plt.savefig(
savfil + '.eps',
dpi=300,
facecolor='w',
edgecolor='w',
format='eps',
transparent=False,
bbox_inches='tight')
else:
plt.show()
def get_heddy_coeffs(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs=0,
ze=None,
z=np.linspace(-3000, 0, 100),
percx=0,
percy=0):
fhgeo = dset(geofil)
fhvgeo = dset(vgeofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
db = fhvgeo.variables['g'][:]
fhvgeo.close()
dbi = np.append(db, 0)
zi = fh.variables['zi'][:]
zl = fh.variables['zl'][:]
dbl = np.diff(zi) * 9.8 / 1031
dt = fh.variables['average_DT'][:]
dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
(sl[3].start - 1):sl[3].stop]
slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
sl2d = sl[2:]
slmx2d = slmx[2:]
slpx2d = slpx[2:]
dxbu = fhgeo.variables['dxBu'][slmx2d]
dxcv = fhgeo.variables['dxCv'][sl2d]
v, _ = pv.getvtwa(fhgeo, fh, fh2, slmx)
vx = np.diff(v, axis=3) / dxbu
vxx = np.diff(vx, axis=3) / dxcv
e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
axis=0, keepdims=True)
eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
vxx = getvaratzc(vxx.astype(np.float32), z.astype(np.float32),
eatv.astype(np.float32))
e = fh2.variables['e'][slmxpy]
eatvmx = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
vaz = getvaratzc(v.astype(np.float32), z.astype(np.float32),
eatvmx.astype(np.float32))
vazx = np.diff(vaz, axis=3) / dxbu
vazxx = np.diff(vazx, axis=3) / dxcv
x, y, P, _, _ = py.extract_twamomy_terms(geofil, vgeofil, fil, fil2,
xstart, xend, ystart, yend, zs,
ze, (0, ))
uv = P[:, :, :, :, 5]
vvalid = vaz[:, :, :, 1:-1]
vazm = np.apply_over_axes(np.nanmean, vvalid, (0, 2))
vxxm = np.apply_over_axes(np.nanmean, vxx, (0, 2))
vazxxm = np.apply_over_axes(np.nanmean, vazxx, (0, 2))
uvm = np.apply_over_axes(np.nanmean, uv, (0, 2))
fig, ax = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(12, 3))
cmax = np.fabs(np.percentile(vxxm, (1, 99))).max()
im = ax[0].pcolormesh(dimv[3],
z,
vxxm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[0])
cmax = np.fabs(np.percentile(vazxxm, (1, 99))).max()
im = ax[1].pcolormesh(dimv[3],
z,
vazxxm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[1])
cmax = np.fabs(np.percentile(vvalid, (0.5, 99.5))).max()
im = ax[2].pcolormesh(dimv[3],
z,
vazm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
fig.colorbar(im, ax=ax[2])
cmax = np.fabs(np.percentile(uvm, (0.5, 99.5))).max()
im = ax[3].pcolormesh(dimv[3],
z,
uvm.squeeze(),
vmax=cmax,
vmin=-cmax,
cmap='RdBu_r')
cb = fig.colorbar(im, ax=ax[3])
cb.formatter.set_powerlimits((0, 0))
cb.update_ticks()
for axc in ax:
xdegtokm(axc, 0.5 * (ystart + yend))
axc.set_xlabel('x (km)')
axc.grid()
ax[0].set_ylabel('z (m)')
fig2, ax = plt.subplots(1, 2, sharey=True)
ax[0].linfit(vvalid, uv, percx=percx, percy=percy)
ax[0].set_xlabel('v (m/s)')
ax[0].set_ylabel('RS div (m/s2)')
ax[0].grid()
ax[1].linfit(vxx, uv, percx=percx, percy=percy)
ax[1].set_xlabel('vxx (1/ms)')
ax[1].grid()
return fig, fig2
def extract_twamomy_terms(geofil,
vgeofil,
fil,
fil2,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
alreadysaved=False,
xyasindices=False,
calledfrompv=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
fh2 = mfdset(fil2)
zi = rdp1.getdims(fh)[2][0]
dbl = np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimv = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xh',
yhyq='yq',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimv = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
yhyq='yq')
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
Dforgetutwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetutwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][6:8]
dxcv, dycv = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][2:4]
dxcvforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][2:3]
dxcvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[2][3:4]
dxbu = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[2][4]
nt_const = dimv[0].size
t0 = time.time()
em = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
elm = 0.5 * (em[:, 0:-1, :, :] + em[:, 1:, :, :])
vh = fh2.variables['vh'][0:, zs:ze, ys:ye, xs:xe]
h_cv = fh.variables['h_Cv'][0:, zs:ze, ys:ye, xs:xe]
h_vm = h_cv
vtwa = vh / h_cv / dxcv
vhforxdiff = fh2.variables['vh'][0:, zs:ze, ys:ye, xs - 1:xe]
h_cvforxdiff = fh.variables['h_Cv'][0:, zs:ze, ys:ye, xs - 1:xe]
vtwaforxdiff = vhforxdiff / h_cvforxdiff / dxcvforxdiff
vtwaforxdiff = np.concatenate(
(vtwaforxdiff, vtwaforxdiff[:, :, :, -1:]), axis=3)
vhforydiff = fh2.variables['vh'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
h_cvforydiff = fh.variables['h_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
vtwaforydiff = vhforydiff / h_cvforydiff / dxcvforydiff
vtwax = np.diff(vtwaforxdiff, axis=3) / dxbu
vtwax = 0.5 * (vtwax[:, :, :, 0:-1] + vtwax[:, :, :, 1:])
vtway = np.diff(vtwaforydiff, axis=2) / dyt
vtway = 0.5 * (vtway[:, :, 0:-1, :] + vtway[:, :, 1:, :])
hvmy = np.diff(vhforydiff / dxcvforydiff, axis=2) / dyt
hvmy = 0.5 * (hvmy[:, :, :-1, :] + hvmy[:, :, 1:, :])
hum = fh.variables['hu_Cv'][0:, zs:ze, ys:ye, xs:xe]
humforxdiff = fh.variables['hu_Cv'][0:, zs:ze, ys:ye, xs - 1:xe]
humforxdiff = np.concatenate((humforxdiff, -humforxdiff[:, :, :, -1:]),
axis=3)
humx = np.diff(humforxdiff, axis=3) / dxbu
humx = 0.5 * (humx[:, :, :, :-1] + humx[:, :, :, 1:])
huvxphvvym = fh.variables['twa_huvxpt'][
0:, zs:ze, ys:ye, xs:xe] + fh.variables['twa_hvvymt'][0:, zs:ze,
ys:ye, xs:xe]
hvv = fh.variables['hvv_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
hvvym = np.diff(hvv, axis=2) / dyt
hvvym = 0.5 * (hvvym[:, :, :-1, :] + hvvym[:, :, 1:, :])
huvxm = huvxphvvym - hvvym
vtwaforvdiff = np.concatenate((vtwa[:, [0], :, :], vtwa), axis=1)
vtwab = np.diff(vtwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
vtwab = np.concatenate(
(vtwab,
np.zeros([vtwab.shape[0], 1, vtwab.shape[2], vtwab.shape[3]])),
axis=1)
vtwab = 0.5 * (vtwab[:, 0:-1, :, :] + vtwab[:, 1:, :, :])
hwb_v = fh.variables['hwb_Cv'][0:, zs:ze, ys:ye, xs:xe]
hwm_v = fh.variables['hw_Cv'][0:, zs:ze, ys:ye, xs:xe]
esq = fh.variables['esq_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
ecv = fh.variables['e_Cv'][0:, zs:ze, ys - 1:ye + 1, xs:xe]
edlsqm = (esq - ecv**2)
edlsqmy = np.diff(edlsqm, axis=2) / dyt
edlsqmy = 0.5 * (edlsqmy[:, :, :-1, :] + edlsqmy[:, :, 1:, :])
epfv = fh.variables['epfv'][0:, zs:ze, ys:ye, xs:xe]
ecv = fh.variables['e_Cv'][0:, zs:ze, ys:ye, xs:xe]
pfvm = fh.variables['pfv_masked'][0:, zs:ze, ys:ye, xs:xe]
edpfvdm = epfv - pfvm * ecv
edpfvdmb = np.diff(edpfvdm, axis=1)
edpfvdmb = np.concatenate(
(edpfvdmb[:, :1, :, :], edpfvdmb, edpfvdmb[:, -1:, :, :]), axis=1)
edpfvdmb = 0.5 * (edpfvdmb[:, :-1, :, :] + edpfvdmb[:, 1:, :, :])
hmfum = fh.variables['twa_hmfu'][0:1, zs:ze, ys:ye, xs:xe]
hvwbm = fh.variables['twa_hvwb'][0:1, zs:ze, ys:ye, xs:xe]
hdiffvm = fh.variables['twa_hdiffv'][0:1, zs:ze, ys:ye, xs:xe]
hdvdtviscm = fh.variables['twa_hdvdtvisc'][0:1, zs:ze, ys:ye, xs:xe]
fh2.close()
fh.close()
advx = hum * vtwax / h_vm
advy = vtwa * vtway
advb = hwm_v * vtwab / h_vm
cor = hmfum / h_vm
pfvm = pfvm
xdivep1 = -huvxm / h_vm
xdivep2 = -advx
xdivep3 = -vtwa * humx / h_vm
xdivep = (xdivep1 + xdivep2 + xdivep3)
ydivep1 = -hvvym / h_vm
ydivep2 = -advy
ydivep3 = -vtwa * hvmy / h_vm
ydivep4 = 0.5 * edlsqmy * dbl[:, np.newaxis, np.newaxis] / h_vm
ydivep = (ydivep1 + ydivep2 + ydivep3 + ydivep4)
bdivep1 = -hvwbm / h_vm
bdivep2 = -advb
bdivep3 = -vtwa * hwb_v / h_vm
bdivep4 = edpfvdmb / h_vm
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
Y1twa = hdiffvm / h_vm
Y2twa = hdvdtviscm / h_vm
terms = np.ma.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfvm[:, :, :, :, np.newaxis], -xdivep[:, :, :, :, np.newaxis],
-ydivep[:, :, :, :, np.newaxis], -bdivep[:, :, :, :, np.newaxis],
Y1twa[:, :, :, :, np.newaxis], Y2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.ma.concatenate((
-xdivep1[:, :, :, :, np.newaxis], -xdivep3[:, :, :, :, np.newaxis],
-ydivep1[:, :, :, :, np.newaxis], -ydivep3[:, :, :, :, np.newaxis],
-ydivep4[:, :, :, :, np.newaxis], -bdivep1[:, :, :, :, np.newaxis],
-bdivep3[:, :, :, :, np.newaxis],
-bdivep4[:, :, :, :, np.newaxis]),
axis=4)
terms[np.isinf(terms)] = np.nan
termsep[np.isinf(termsep)] = np.nan
termsm = np.ma.apply_over_axes(np.nanmean, terms, meanax)
termsepm = np.ma.apply_over_axes(np.nanmean, termsep, meanax)
X = dimv[keepax[1]]
Y = dimv[keepax[0]]
if 1 in keepax:
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimv[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
Pep = termsepm.squeeze()
Pep = np.ma.filled(Pep.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
if not calledfrompv:
np.savez('twamomy_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomy_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep)
def extract_twamomx_terms(geofil,
vgeofil,
fil,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
alreadysaved=False,
xyasindices=False,
calledfrompv=False):
if not alreadysaved:
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fhvgeo = dset(vgeofil)
db = -fhvgeo.variables['g'][:]
dbi = np.append(db, 0)
fhvgeo.close()
fhgeo = dset(geofil)
fh = mfdset(fil)
zi = rdp1.getdims(fh)[2][0]
dbl = -np.diff(zi) * 9.8 / 1031
if xyasindices:
(xs, xe), (ys, ye) = (xstart, xend), (ystart, yend)
_, _, dimu = rdp1.getdimsbyindx(fh,
xs,
xe,
ys,
ye,
zs=zs,
ze=ze,
ts=0,
te=None,
xhxq='xq',
yhyq='yh',
zlzi='zl')
else:
(xs, xe), (ys, ye), dimu = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
xhxq='xq')
D, (ah, aq) = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0:2]
Dforgetutwaforxdiff = rdp1.getgeombyindx(fhgeo, xs - 1, xe, ys, ye)[0]
Dforgetutwaforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1,
ye + 1)[0]
Dforgethvforydiff = rdp1.getgeombyindx(fhgeo, xs, xe, ys - 1, ye)[0]
dxt, dyt = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][6:8]
dxcu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[2][0]
dybu = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye + 1)[2][5]
nt_const = dimu[0].size
t0 = time.time()
print('Reading data using loop...')
h = fh.variables['h'][0:1, zs:ze, ys:ye, xs:xe]
u = fh.variables['u'][0:1, zs:ze, ys:ye, xs:xe]
nt = np.ones(u.shape) * nt_const
frhatu = fh.variables['frhatu'][0:1, zs:ze, ys:ye, xs:xe]
h_u = frhatu * D[np.newaxis, np.newaxis, :, :]
h_u = np.ma.masked_array(h_u, mask=(h <= 1e-3).astype(int))
nt[h <= 1e-3] -= 1
hum = (h_u * u).filled(0)
huum = (h_u * u * u).filled(0)
h_um = h_u.filled(0)
humforxdiff, h_umforxdiff, huumforxdiff = getutwaforxdiff(
fh, fhgeo, Dforgetutwaforxdiff, 0, xs - 1, xe, ys, ye, zs, ze)
humforydiff, h_umforydiff = getutwaforydiff(fh, fhgeo,
Dforgetutwaforydiff, 0, xs,
xe, ys - 1, ye + 1, zs, ze)
cau = fh.variables['CAu'][0:1, zs:ze, ys:ye, xs:xe]
gkeu = fh.variables['gKEu'][0:1, zs:ze, ys:ye, xs:xe]
rvxv = fh.variables['rvxv'][0:1, zs:ze, ys:ye, xs:xe]
hfvm = h_u * (cau - gkeu - rvxv).filled(0)
pfu = fh.variables['PFu'][0:1, zs:ze, ys:ye, xs:xe]
pfu = np.ma.masked_array(pfu, mask=(h <= 1e-3).astype(int))
pfum = pfu.filled(0)
hdudtviscm = (
h_u *
fh.variables['du_dt_visc'][0:1, zs:ze, ys:ye, xs:xe]).filled(0)
hdiffum = (h_u *
fh.variables['diffu'][0:1, zs:ze, ys:ye, xs:xe]).filled(0)
dudtdia = fh.variables['dudt_dia'][0:1, zs:ze, ys:ye, xs:xe]
wd = fh.variables['wd'][0:1, zs:ze, ys:ye, xs:xe]
wd = np.diff(wd, axis=1)
wd = np.concatenate((wd, wd[:, :, :, -1:]), axis=3)
wdm = 0.5 * (wd[:, :, :, :-1] + wd[:, :, :, 1:])
wdm = np.ma.masked_array(wdm, mask=(h <= 1e-3).astype(int))
wdm = wdm.filled(0)
huwbm = (u * (wd[:, :, :, 0:-1] + wd[:, :, :, 1:]) / 2 -
h_u * dudtdia).filled(0)
uh = fh.variables['uh'][0:1, zs:ze, ys:ye, xs - 1:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh, axis=3) / ah
uhx = np.concatenate((uhx, uhx[:, :, :, -1:]), axis=3)
huuxpTm = (u * (uhx[:, :, :, 0:-1] + uhx[:, :, :, 1:]) / 2 -
h_u * gkeu).filled(0)
vh = fh.variables['vh'][0:1, zs:ze, ys - 1:ye, xs:xe]
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh, axis=2) / ah
vhy = np.concatenate((vhy, vhy[:, :, :, -1:]), axis=3)
huvymTm = (u * (vhy[:, :, :, 0:-1] + vhy[:, :, :, 1:]) / 2 -
h_u * rvxv).filled(0)
v = fh.variables['v'][0:1, zs:ze, ys - 1:ye, xs:xe]
v = np.concatenate((v, -v[:, :, :, [-1]]), axis=3)
v = 0.25 * (v[:, :, 0:-1, 0:-1] + v[:, :, 1:, 0:-1] +
v[:, :, 0:-1, 1:] + v[:, :, 1:, 1:])
hvm = (h_u * v).filled(0)
hvmforydiff = gethvforydiff(fh, fhgeo, Dforgethvforydiff, 0, xs, xe,
ys - 1, ye, zs, ze)
wd = fh.variables['wd'][0:1, zs:ze, ys:ye, xs:xe]
wd = np.concatenate((wd, -wd[:, :, :, -1:]), axis=3)
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 0:-1, :, :] + hw[:, 1:, :, :])
hwm_u = 0.5 * (hw[:, :, :, 0:-1] + hw[:, :, :, 1:])
hwm_u = np.ma.masked_array(hwm_u, mask=(h <= 1e-3).astype(int))
hwm_u = hwm_u.filled(0)
e = fh.variables['e'][0:1, zs:ze, ys:ye, xs:xe]
em = e / nt_const
el = 0.5 * (e[:, :-1, :, :] + e[:, 1:, :, :])
el = np.ma.masked_array(el, mask=(h <= 1e-3).astype(int))
elm = el.filled(0)
esq = el**2
esq = np.ma.masked_array(esq, mask=(h <= 1e-3).astype(int))
esqm = esq.filled(0)
el = np.concatenate((el, el[:, :, :, -1:]), axis=3)
el = 0.5 * (el[:, :, :, :-1] + el[:, :, :, 1:])
elmatu = el.filled(0)
epfum = (pfu * el).filled(0)
for i in range(1, nt_const):
h = fh.variables['h'][i:i + 1, zs:ze, ys:ye, xs:xe]
u = fh.variables['u'][i:i + 1, zs:ze, ys:ye, xs:xe]
frhatu = fh.variables['frhatu'][i:i + 1, zs:ze, ys:ye, xs:xe]
h_u = frhatu * D[np.newaxis, np.newaxis, :, :]
h_u = np.ma.masked_array(h_u, mask=(h <= 1e-3).astype(int))
nt[h <= 1e-3] -= 1
hum += (h_u * u).filled(0)
h_um += h_u.filled(0)
huforxdiff, h_uforxdiff, huuforxdiff = getutwaforxdiff(
fh, fhgeo, Dforgetutwaforxdiff, i, xs - 1, xe, ys, ye, zs, ze)
huforydiff, h_uforydiff = getutwaforydiff(fh, fhgeo,
Dforgetutwaforydiff, i,
xs, xe, ys - 1, ye + 1,
zs, ze)
humforxdiff += huforxdiff
h_umforxdiff += h_uforxdiff
huumforxdiff += huuforxdiff
humforydiff += huforydiff
h_umforydiff += h_uforydiff
cau = fh.variables['CAu'][i:i + 1, zs:ze, ys:ye, xs:xe]
gkeu = fh.variables['gKEu'][i:i + 1, zs:ze, ys:ye, xs:xe]
rvxv = fh.variables['rvxv'][i:i + 1, zs:ze, ys:ye, xs:xe]
hfv = h_u * (cau - gkeu - rvxv)
hfvm += hfv.filled(0)
pfu = fh.variables['PFu'][i:i + 1, zs:ze, ys:ye, xs:xe]
pfu = np.ma.masked_array(pfu, mask=(h <= 1e-3).astype(int))
pfum += pfu.filled(0)
hdudtvisc = h_u * fh.variables['du_dt_visc'][i:i + 1, zs:ze, ys:ye,
xs:xe]
hdudtviscm += hdudtvisc.filled(0)
hdiffu = h_u * fh.variables['diffu'][i:i + 1, zs:ze, ys:ye, xs:xe]
hdiffum += hdiffu.filled(0)
dudtdia = fh.variables['dudt_dia'][i:i + 1, zs:ze, ys:ye, xs:xe]
wd = fh.variables['wd'][i:i + 1, zs:ze, ys:ye, xs:xe]
wd = np.diff(wd, axis=1)
wd = np.concatenate((wd, wd[:, :, :, -1:]), axis=3)
huwb = (u * (wd[:, :, :, 0:-1] + wd[:, :, :, 1:]) / 2 -
h_u * dudtdia)
huwb = np.ma.masked_array(huwb, mask=(h_u <= 1e-3).astype(int))
huwbm += huwb.filled(0)
wd = 0.5 * (wd[:, :, :, :-1] + wd[:, :, :, 1:])
wd = np.ma.masked_array(wd, mask=(h <= 1e-3).astype(int))
wd = wd.filled(0)
wdm += wd
uh = fh.variables['uh'][i:i + 1, zs:ze, ys:ye, xs - 1:xe]
uh = np.ma.filled(uh.astype(float), 0)
uhx = np.diff(uh, axis=3) / ah
uhx = np.concatenate((uhx, uhx[:, :, :, -1:]), axis=3)
huuxpT = (u * (uhx[:, :, :, 0:-1] + uhx[:, :, :, 1:]) / 2 -
h_u * gkeu)
huuxpT = np.ma.masked_array(huuxpT, mask=(h_u <= 1e-3).astype(int))
huuxpTm += huuxpT.filled(0)
vh = fh.variables['vh'][i:i + 1, zs:ze, ys - 1:ye, xs:xe]
vh = np.ma.filled(vh.astype(float), 0)
vhy = np.diff(vh, axis=2) / ah
vhy = np.concatenate((vhy, vhy[:, :, :, -1:]), axis=3)
huvymT = (u * (vhy[:, :, :, 0:-1] + vhy[:, :, :, 1:]) / 2 -
h_u * rvxv)
huvymT = np.ma.masked_array(huvymT, mask=(h_u <= 1e-3).astype(int))
huvymTm += huvymT.filled(0)
v = fh.variables['v'][i:i + 1, zs:ze, ys - 1:ye, xs:xe]
v = np.concatenate((v, -v[:, :, :, [-1]]), axis=3)
v = 0.25 * (v[:, :, 0:-1, 0:-1] + v[:, :, 1:, 0:-1] +
v[:, :, 0:-1, 1:] + v[:, :, 1:, 1:])
hv = h_u * v
hvm += hv.filled(0)
hvmforydiff += gethvforydiff(fh, fhgeo, Dforgethvforydiff, i, xs,
xe, ys - 1, ye, zs, ze)
wd = fh.variables['wd'][i:i + 1, zs:ze, ys:ye, xs:xe]
wd = np.concatenate((wd, -wd[:, :, :, -1:]), axis=3)
hw = wd * dbi[:, np.newaxis, np.newaxis]
hw = 0.5 * (hw[:, 0:-1, :, :] + hw[:, 1:, :, :])
hw_u = 0.5 * (hw[:, :, :, 0:-1] + hw[:, :, :, 1:])
hw_u = np.ma.masked_array(hw_u, mask=(h <= 1e-3).astype(int))
hw_u = hw_u.filled(0)
hwm_u += hw_u
e = fh.variables['e'][i:i + 1, zs:ze, ys:ye, xs:xe]
em += e / nt_const
el = 0.5 * (e[:, :-1, :, :] + e[:, 1:, :, :])
el = np.ma.masked_array(el, mask=(h <= 1e-3).astype(int))
elm += el.filled(0)
esq = el**2
esqm += esq.filled(0)
el = np.concatenate((el, el[:, :, :, -1:]), axis=3)
el = 0.5 * (el[:, :, :, :-1] + el[:, :, :, 1:])
elmatu += el.filled(0)
epfum += (pfu * el).filled(0)
sys.stdout.write('\r' + str(int((i + 1) / nt_const * 100)) +
'% done...')
sys.stdout.flush()
fhgeo.close()
fh.close()
print('Time taken for data reading: {}s'.format(time.time() - t0))
utwa = hum / h_um
utwaforxdiff = humforxdiff / h_umforxdiff
utwaforydiff = humforydiff / h_umforydiff
utwaforxdiff[np.isnan(utwaforxdiff)] = 0
utwax = np.diff(utwaforxdiff, axis=3) / dxt
utwax = np.concatenate((utwax, -utwax[:, :, :, [-1]]), axis=3)
utwax = 0.5 * (utwax[:, :, :, 0:-1] + utwax[:, :, :, 1:])
utway = np.diff(utwaforydiff, axis=2) / dybu
utway = 0.5 * (utway[:, :, 0:-1, :] + utway[:, :, 1:, :])
humx = np.diff(humforxdiff, axis=3) / dxt
humx = np.concatenate((humx, -humx[:, :, :, [-1]]), axis=3)
humx = 0.5 * (humx[:, :, :, 0:-1] + humx[:, :, :, 1:])
hvmy = np.diff(hvmforydiff, axis=2) / dyt
hvmy = np.concatenate((hvmy, -hvmy[:, :, :, [-1]]), axis=3)
hvmy = 0.5 * (hvmy[:, :, :, 0:-1] + hvmy[:, :, :, 1:])
huuxphuvym = huuxpTm + huvymTm
huuxm = np.diff(huumforxdiff, axis=3) / dxt
huuxm = np.concatenate((huuxm, -huuxm[:, :, :, [-1]]), axis=3)
huuxm = 0.5 * (huuxm[:, :, :, 0:-1] + huuxm[:, :, :, 1:])
huvym = huuxphuvym - huuxm
utwaforvdiff = np.concatenate((utwa[:, [0], :, :], utwa), axis=1)
utwab = np.diff(utwaforvdiff, axis=1) / db[:, np.newaxis, np.newaxis]
utwab = np.concatenate(
(utwab,
np.zeros([utwab.shape[0], 1, utwab.shape[2], utwab.shape[3]])),
axis=1)
utwab = 0.5 * (utwab[:, 0:-1, :, :] + utwab[:, 1:, :, :])
hwb_u = wdm
edlsqm = esqm / nt - (elm / nt)**2
edlsqm = np.concatenate((edlsqm, edlsqm[:, :, :, [-1]]), axis=3)
edlsqmx = np.diff(edlsqm, axis=3) / dxcu * nt
edpfudm = (epfum / nt - elmatu / nt * pfum / nt) * nt
edpfudmb = np.diff(edpfudm, axis=1)
# edpfudmb = np.diff(edpfudm,axis=1)/db[1:,np.newaxis,np.newaxis]
edpfudmb = np.concatenate(
(edpfudmb[:, :1, :, :], edpfudmb, edpfudmb[:, -1:, :, :]), axis=1)
edpfudmb = 0.5 * (edpfudmb[:, :-1, :, :] + edpfudmb[:, 1:, :, :])
advx = utwa * utwax
advy = hvm * utway / h_um
advb = hwm_u * utwab / h_um
cor = hfvm / h_um
pfum = pfum / nt
xdivep1 = huuxm / h_um
xdivep2 = -advx
xdivep3 = -utwa * humx / h_um
xdivep4 = 0.5 * edlsqmx * dbl[:, np.newaxis, np.newaxis] / h_um
xdivep = (xdivep1 + xdivep2 + xdivep3 + xdivep4)
ydivep1 = huvym / h_um
ydivep2 = -advy
ydivep3 = -utwa * hvmy / h_um
ydivep = (ydivep1 + ydivep2 + ydivep3)
bdivep1 = huwbm / h_um
bdivep2 = -advb
bdivep3 = -utwa * hwb_u / h_um
bdivep4 = edpfudmb / h_um
bdivep = (bdivep1 + bdivep2 + bdivep3 + bdivep4)
X1twa = hdiffum / h_um
X2twa = hdudtviscm / h_um
terms = np.ma.concatenate(
(-advx[:, :, :, :, np.newaxis], -advy[:, :, :, :, np.newaxis],
-advb[:, :, :, :, np.newaxis], cor[:, :, :, :, np.newaxis],
pfum[:, :, :, :, np.newaxis], -xdivep[:, :, :, :, np.newaxis],
-ydivep[:, :, :, :, np.newaxis], -bdivep[:, :, :, :, np.newaxis],
X1twa[:, :, :, :, np.newaxis], X2twa[:, :, :, :, np.newaxis]),
axis=4)
termsep = np.ma.concatenate(
(xdivep1[:, :, :, :, np.newaxis], xdivep3[:, :, :, :, np.newaxis],
xdivep4[:, :, :, :, np.newaxis], ydivep1[:, :, :, :, np.newaxis],
ydivep3[:, :, :, :, np.newaxis], bdivep1[:, :, :, :, np.newaxis],
bdivep3[:, :, :, :, np.newaxis], bdivep4[:, :, :, :, np.newaxis]),
axis=4)
terms[np.isinf(terms)] = np.nan
termsep[np.isinf(termsep)] = np.nan
termsm = np.ma.apply_over_axes(np.nanmean, terms, meanax)
termsepm = np.ma.apply_over_axes(np.nanmean, termsep, meanax)
elm = 0.5 * (em[:, :-1, :, :] + em[:, 1:, :, :])
X = dimu[keepax[1]]
Y = dimu[keepax[0]]
if 1 in keepax:
em = np.ma.apply_over_axes(np.mean, em, meanax)
elm = np.ma.apply_over_axes(np.mean, elm, meanax)
Y = elm.squeeze()
X = np.meshgrid(X, dimu[1])[0]
P = termsm.squeeze()
P = np.ma.filled(P.astype(float), np.nan)
Pep = termsepm.squeeze()
Pep = np.ma.filled(Pep.astype(float), np.nan)
X = np.ma.filled(X.astype(float), np.nan)
Y = np.ma.filled(Y.astype(float), np.nan)
if not calledfrompv:
np.savez('twamomx_complete_terms', X=X, Y=Y, P=P, Pep=Pep)
else:
npzfile = np.load('twamomx_complete_terms.npz')
X = npzfile['X']
Y = npzfile['Y']
P = npzfile['P']
Pep = npzfile['Pep']
return (X, Y, P, Pep)
def plot_eb_transport(geofil,
vgeofil,
fil,
xstart,
xend,
ystart,
yend,
zs,
ze,
meanax,
savfil=None):
keepax = ()
for i in range(4):
if i not in meanax:
keepax += (i, )
fh = mfdset(fil)
(xs, xe), (ys, ye), dimv = rdp1.getlatlonindx(fh,
wlon=xstart,
elon=xend,
slat=ystart,
nlat=yend,
zs=zs,
ze=ze,
yhyq='yq')
fhgeo = dset(geofil)
D = rdp1.getgeombyindx(fhgeo, xs, xe, ys, ye)[0]
fhgeo.close()
nt_const = dimv[0].size
vh = fh.variables['vh'][0:, zs:ze, ys:ye, xs:xe]
e = fh.variables['e'][0:, zs:ze, ys:ye, xs:xe]
fh.close()
vh = vh.filled(0)
X = dimv[keepax[1]]
Y = dimv[keepax[0]]
if 1 in keepax:
z = np.linspace(-np.nanmax([2000]), -1, num=50)
Y = z
P = getvaratzc(vh, z, e)
P = np.ma.apply_over_axes(np.mean, P, meanax)
P = P.squeeze()
#P = np.ma.apply_over_axes(np.mean, vh, meanax).squeeze()
im = m6plot((X, Y, P),
titl='Transport near EB',
ylab='z (m)',
xlab='y (Deg)')
if savfil:
plt.savefig(savfil + '.eps',
dpi=300,
facecolor='w',
edgecolor='w',
format='eps',
transparent=False,
bbox_inches='tight')
else:
im = m6plot((X, Y, P),
titl='Transport near EB',
ylab='z (m)',
xlab='y (Deg)')
plt.show()
