def calculate_relvor(self, lev=850): from windspharm.standard import VectorWind # Load the full U and V fields u_full = self['uzonal_{}hPa'.format(lev)].values v_full = self['umeridional_{}hPa'.format(lev)].values relvor = np.zeros(np.shape(u_full)) # Loop through all the valid times and ensemble members for i1 in range(np.shape(u_full)[0]): for i2 in range(np.shape(u_full)[1]): # Create the spherical harmonics vector object u = u_full[i1, i2, ::-1, :] # lats must go from N to S v = v_full[i1, i2, ::-1, :] # lats must go from N to S wnd = VectorWind(u, v, gridtype='regular') # Calculate the relative vorticity relvor[i1, i2, :, :] = wnd.vorticity()[::-1, :] # Assign relative vorticity as a new variable vorvar = xarray.DataArray(relvor, dims=self['uzonal_{}hPa'.format(lev)].dims) assignvar = {'relvor_{}hPa'.format(lev): vorvar} self.update(self.assign(**assignvar))
def calc_quantity(uwnd, vwnd, quantity, lat_axis, lon_axis, axis_order): """Calculate a single wind quantity. Args: uwnd (numpy.ndarray): Zonal wind vwnd (numpy.ndarray): Meridional wind quantity (str): Quantity to be calculated lat_axis (list): Latitude axis values lon_axis (list): Longitude axis values axis_order (str): e.g. tyx Design: windsparm requires the input data to be on a global grid (due to the spherical harmonic representation used), latitude and longitude to be the leading axes and the latitude axis must go from 90 to -90. The cdms2 interface is supposed to adjust for these things but I've found that things come back upsidedown if the lat axis isn't right, so I've just used the standard interface here instead. Reference: ajdawson.github.io/windspharm """ check_global(lat_axis, lon_axis) # Make latitude and longitude the leading coordinates uwnd, uwnd_info = prep_data(numpy.array(uwnd), axis_order) vwnd, vwnd_info = prep_data(numpy.array(vwnd), axis_order) # Make sure latitude dimension is north-to-south lats, uwnd, vwnd = order_latdim(lat_axis, uwnd, vwnd) flip_lat = False if lats[0] == lat_axis[0] else True w = VectorWind(uwnd, vwnd) data_out = {} if quantity == 'rossbywavesource': eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) data_out['rws1'] = (-eta * div) / (1.e-11) data_out['rws2'] = (-(uchi * etax + vchi * etay)) / (1.e-11) data_out['rws'] = data_out['rws1'] + data_out['rws2'] elif quantity == 'magnitude': data_out['spd'] = w.magnitude() elif quantity == 'vorticity': data_out['vrt'] = w.vorticity() elif quantity == 'divergence': div = w.divergence() data_out['div'] = div / (1.e-6) elif quantity == 'absolutevorticity': avrt = w.absolutevorticity() data_out['avrt'] = avrt / (1.e-5) elif quantity == 'absolutevorticitygradient': avrt = w.absolutevorticity() ugrad, vgrad = w.gradient(avrt) avrtgrad = numpy.sqrt(numpy.square(ugrad) + numpy.square(vgrad)) data_out['avrtgrad'] = avrtgrad / (1.e-5) elif quantity == 'planetaryvorticity': data_out['pvrt'] = w.planetaryvorticity() elif quantity == 'irrotationalcomponent': data_out['uchi'], data_out['vchi'] = w.irrotationalcomponent() elif quantity == 'nondivergentcomponent': data_out['upsi'], data_out['vpsi'] = w.nondivergentcomponent() elif quantity == 'streamfunction': sf = w.streamfunction() data_out['sf'] = sf / (1.e+6) elif quantity == 'velocitypotential': vp = w.velocitypotential() data_out['vp'] = vp / (1.e+6) else: sys.exit('Wind quantity not recognised') # Return data to its original shape for key in data_out.keys(): data_out[key] = recover_structure(data_out[key], flip_lat, uwnd_info) return data_out