def windprops(ua, va, lat, lon, exps=['All-Nat', 'SST-Nat']): from windspharm.standard import VectorWind def diverg(u, v, lat, lon): meshlon, meshlat = np.meshgrid(lon, lat) phi = meshlat * np.pi / 180 cs_phi = np.cos(phi) dphi = lat * np.pi / 180 dtheta = lon * np.pi / 180 a = 6.371e6 v_y = np.gradient(v * cs_phi, dphi, axis=1) / (a * cs_phi) u_x = np.gradient(u, dtheta, axis=2) / (a * cs_phi) div = u_x + v_y return div u = ua[exps[0]] v = va[exps[0]] u_ghg = ua[exps[1]] v_ghg = va[exps[1]] div = diverg(u, v, lat, lon) div_ghg = diverg(u_ghg, v_ghg, lat, lon) u = np.transpose(u, (1, 2, 0)) v = np.transpose(v, (1, 2, 0)) u_ghg = np.transpose(u_ghg, (1, 2, 0)) v_ghg = np.transpose(v_ghg, (1, 2, 0)) w = VectorWind(u, v) w_ghg = VectorWind(u_ghg, v_ghg) vp = w.velocitypotential() vp_ghg = w_ghg.velocitypotential() vp = np.transpose(vp, (2, 0, 1)) vp_ghg = np.transpose(vp_ghg, (2, 0, 1)) delta_vp = vp_ghg - vp delta_div = div_ghg - div return div, delta_div, vp, delta_vp
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
f=Dataset(model, mode='r') uwnd = f.variables['ua'][0,0,:,:] vwnd = f.variables['va'][0,0,:,:] uwnd2 = np.array(uwnd) vwnd2 = np.array(vwnd) lons = f.variables['lon'][:] lats = f.variables['lat'][:] f.close() plots=plt.subplot(2,3, count+1) print count # velocity potential calculation - requires some data formatting using Windspharm's tools uwnd, uwnd_info = prep_data(uwnd, 'yx') vwnd, vwnd_info = prep_data(vwnd, 'yx') lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) w = VectorWind(uwnd, vwnd, gridtype='gaussian') vp = w.velocitypotential() vp = recover_data(vp, uwnd_info) vp = vp*1e-06 # divergence calculation divg = w.divergence() divg = recover_data(divg, uwnd_info) divg = divg*1e6 #divergent wind compotnents calculation uchi,vchi = w.irrotationalcomponent() uchi = np.squeeze(uchi) vchi = np.squeeze(vchi) # print some data dimensions - useful for colourbar range etc print np.shape(lons) print np.shape(lats) print np.shape(uchi) print np.shape(vchi)