diff --git a/climate.py b/climate.py index d901839..b0a9314 100644 --- a/climate.py +++ b/climate.py @@ -729,6 +729,7 @@ def ComputeWstar(data, slice='all', omega='omega', temp='temp', vcomp='vcomp', p return w_bar + np.squeeze(R*vt_bar) else: return w_bar + R*vt_bar + ############################################################################################## def ComputeEPfluxDiv(lat,pres,u,v,t,w=None,do_ubar=False,wave=-1): """ Compute the EP-flux vectors and divergence terms. @@ -834,6 +835,116 @@ def ComputeEPfluxDiv(lat,pres,u,v,t,w=None,do_ubar=False,wave=-1): # return ep1_cart,ep2_cart,div1,div2 +############################################################################################## +def ComputeWaveActivityFlux(phi_or_u,phiref_or_v,uref,vref,lat='lat',lon='lon',pres='level',tref=None,qg=False): + ''' + Compute Wave Activity Flux as in Takaya & Nakamura GRL 1997 and Takaya & Nakamura JAS 2001. + Results checked against plots at http://www.atmos.rcast.u-tokyo.ac.jp/nishii/programs/ + This is 3D wave flux in a non-zonally symmetric base flow, optionally QG approximated. + Inputs are xarray DataArrays. + Latitude, longitude in degrees, pressure in hPa. + + INPUTS: + phi_or_u : geopotential [m2/s2] (qg = True) OR zonal wind [m/s] (qg = False) + phiref_or_v: reference geopotential [m2/s2] (qg = True) + OR (full) meridional wind [m/s] (qg = False) + uref : reference zonal wind [m/s] + vref : reference meridional wind [m/s] + lat : name of latitude in DataArrays + lon : name of longitude in DataArrays + pres : name of pressure in DataArrays [hPa] + tref : reference temperature [K] for static stability parameter S2. + If None, only compute horizontal wave flux. + Else, used to compute S2. + If xr.Dataset or xr.DataArray, compute S2 assuming tref is temperature [K]. + Else, assume S2 = tref. + Note that S2 should be a function of + pressure only (see Vallis 2017, Eq 5.127) + qg : use QG streamfunction (phi-phiref)/f? Otherwise, use windspharm.streamfunction. + Note that if qg=False, phi_or_u = u and phiref_or_v = v to compute the streamfunction. + OUTPUTS: + Wx,Wy : Activity Vectors along lon [m2/s2], lat [m2/s2] + Wx,Wy,Wz : Activity Vectors along lon [m2/s2], lat [m2/s2], pres [hPa.m/s2] + ''' + import numpy as np + from xarray import DataArray + for var in [phi_or_u,uref,vref,phiref_or_v,tref]: + if not isinstance(var,DataArray): + if var is None or np.isscalar(var): + pass + else: + raise ValueError('all inputs have to be xarray.DataArrays!') + a0 = 6376000.0 + Rd = 287.04 + kappa = 2./7 + p0 = 1.e3 + radlat = np.deg2rad(phi_or_u[lat]) + coslat = np.cos(radlat) + one_over_coslat2 = coslat**(-2) + one_over_a2 = a0**(-2) + mag_u = np.sqrt(uref**2+vref**2) + f = 2*2*np.pi/86400.*np.sin(radlat) + + if qg: + psi = (phi_or_u-phiref_or_v)/f + else: + from windspharm.xarray import VectorWind + psi = VectorWind(phi_or_u-uref,phiref_or_v-vref).streamfunction() + + # psi.differentiate(lon) == np.gradient(psi)/np.gradient(lon) [psi/lon] + dpsi_dlon = psi.differentiate(lon,edge_order=2).reduce(np.nan_to_num) + dpsi_dlat = psi.differentiate(lat,edge_order=2).reduce(np.nan_to_num) + d2psi_dlon2 = dpsi_dlon.differentiate(lon,edge_order=2) + d2psi_dlat2 = dpsi_dlat.differentiate(lat,edge_order=2) + d2psi_dlon_dlat = dpsi_dlon.differentiate(lat,edge_order=2) + + wx = uref*one_over_coslat2*one_over_a2*(dpsi_dlon**2 - psi*d2psi_dlon2) \ + + vref*one_over_a2/coslat*(dpsi_dlon*dpsi_dlat - psi*d2psi_dlon_dlat) + wy = uref*one_over_a2/coslat*(dpsi_dlon*dpsi_dlat - psi*d2psi_dlon_dlat) \ + + vref*one_over_a2*(dpsi_dlat**2 - psi*d2psi_dlat2) + + coeff = coslat/2/mag_u + + rad2deg = 180/np.pi + + # get the vectors in physical units of m2/s2, correcting for radians vs. degrees + wx = coeff*wx*rad2deg*rad2deg + wy = coeff*wy*rad2deg*rad2deg + + wx.name = 'wx' + wx.attrs['units'] = 'm2/s2' + wy.name = 'wy' + wy.attrs['units'] = 'm2/s2' + + if tref is None: + return wx,wy + else: + # psi.differentiate(pres) == np.gradient(psi)/np.gradient(pres) [psi/pres] + dpsi_dpres = psi.differentiate(pres,edge_order=2).reduce(np.nan_to_num) + d2psi_dlon_dpres = dpsi_dlon.differentiate(pres,edge_order=2) + d2psi_dlat_dpres = dpsi_dlat.differentiate(pres,edge_order=2) + # S2 = -\alpha*\partial_p\ln\theta, \alpha = 1/\rho = Rd*T/p + # = R/p*(p/p0)**\kappa d\theta/dp, Vallis (2017, p. 192 (eq. 5.127)) + # this should be a reference profile and a function of pressure only! + pressure = uref[pres] + if isinstance(tref,Dataset) or isinstance(tref,DataArray): + pp0 = (p0/pressure)**kappa + theta = pp0*tref + S2 = -Rd/pressure*(pressure/p0)**kappa*theta.differentiate(pres,edge_order=2) + # S2 = ExtractArray(S2) + # S2 = S2.where(S2>1e-7,1e-7) # avoid division by zero + else: + S2 = tref + + wz = f**2/S2*( uref/a0/coslat*(dpsi_dlon*dpsi_dpres - psi*d2psi_dlon_dpres) + vref/a0*(dpsi_dlat*dpsi_dpres - psi*d2psi_dlat_dpres) ) + + wz = coeff*wz*rad2deg + # wz = ExtractArray(wz) + wz.name = 'wz' + wz.attrs['units'] = 'hPa.m/s2' + + return wx,wy,wz + ############################################################################################## def PlotEPfluxArrows(x,y,ep1,ep2,fig,ax,xlim=None,ylim=None,xscale='linear',yscale='linear',invert_y=True, newax=False, pivot='tail',scale=None): """Correctly scales the Eliassen-Palm flux vectors for plotting on a latitude-pressure or latitude-height axis. @@ -903,7 +1014,7 @@ def Deltas(z,zlim): if yscale == 'linear': dy = y_sign*height/delta_y elif yscale == 'log': - dy = y_sign*height/np.log(10)/y/np.log10(np.max(y)/np.min(y)) + dy = y_sign*height/y/np.log(np.max(y)/np.min(y)) # # plot the arrows onto axis quivArgs = {'angles':'uv','scale_units':'inches','pivot':pivot} @@ -915,7 +1026,7 @@ def Deltas(z,zlim): try: ax.quiver(x,y,Fphi*dx,Fp*dy,**quivArgs) except: - ax.quiver(x,y,Fphi.transpose()*dx,Fp.transpose()*dy,**quivArgs) + ax.quiver(x,y,dx*Fphi.transpose(),dy*Fp.transpose(),**quivArgs) if invert_y: ax.invert_yaxis() if xlim is not None: diff --git a/inout.py b/inout.py index 38b84bd..919851e 100644 --- a/inout.py +++ b/inout.py @@ -9,6 +9,16 @@ from __future__ import print_function ## find compression parameters as in ncpdq def DefCompress(x,varName=None): + """Produce encoding dictionary for to_netcdf(encoding=encodeDict). + + INPUTS: + x : either a xarray.DataArray or xarray.Dataset + varName : only encode that variable. If None, encode all variables + OUTPUTS: + encodeDict : dictionary containing compressing information for each + variable. To be used as x.to_netcdf(encoding=encodeDict) + """ + import numpy as np # check whether x is a Dataset or DataArray try: keys = x.variables.keys() @@ -34,11 +44,13 @@ def DefCompress(x,varName=None): fillVal = -2**(bytes-1) for var in vars: if is_dataset: - dataMin = float(x[var].min()) - dataMax = float(x[var].max()) + filtr = np.isfinite(x[var]) + dataMin = x[var].where(filtr).min().values + dataMax = x[var].where(filtr).max().values else: - dataMin = float(x.min()) - dataMax = float(x.max()) + filtr = np.isfinite(x) + dataMin = x.where(filtr).min().values + dataMax = x.where(filtr).max().values scale_factor=(dataMax - dataMin) / (2**bytes - 2) add_offset = (dataMax + dataMin) / 2 encodeDict[var] = {