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()