def readin(): ncu = (NetCDFFile('/home/scottyiu/Desktop/work/data/model/50yrs/seasonal/3D/uwind/uwind_sea_' + myjob + '_yseasmean.nc')) uwnd = ncu.variables['u'][:,myheight,:,:] #7 for model, 14 for era: 200hPa lons = ncu.variables['longitude'][:] lats = ncu.variables['latitude'][:] ncu.close() ncv = (NetCDFFile('/home/scottyiu/Desktop/work/data/model/50yrs/seasonal/3D/vwind/vwind_sea_' + myjob + '_yseasmean.nc')) vwnd = ncv.variables['v'][:,myheight,:,:] #7 for model, 14 for era: 200hPa ncv.close() uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) return (ncu,uwnd,lons,lats,ncv,vwnd,uwnd_info,vwnd_info)
def test_prep_recover_data(self): """prepared and recovered data matches original?""" u = np.random.rand(12, 17, 73, 144) up, uinfo = prep_data(u, 'tzyx') ur = recover_data(up, uinfo) err = error(u, ur) assert_almost_equal(err, 0.)
def test_get_recovery(self): # recovery helper should produce the same result as the manual method u = np.random.rand(12, 17, 73, 144) up, uinfo = prep_data(u, 'tzyx') ur1 = recover_data(up, uinfo) recover = get_recovery(uinfo) ur2, = recover(up) assert_array_equal(ur1, ur2)
def test_get_recovery(self): """ get_recovery(info)(pdata) matches recover_data(pdata, info)? """ u = np.random.rand(12, 17, 73, 144) up, uinfo = prep_data(u, 'tzyx') ur1 = recover_data(up, uinfo) recover = get_recovery(uinfo) ur2 = recover(up) err = error(ur1, ur2) assert_almost_equal(err, 0.)
# files. ncu = Dataset(example_data_path('uwnd_mean.nc'), 'r') uwnd = ncu.variables['uwnd'][:] lons = ncu.variables['longitude'][:] lats = ncu.variables['latitude'][:] ncu.close() ncv = Dataset(example_data_path('vwnd_mean.nc'), 'r') vwnd = ncv.variables['vwnd'][:] ncv.close() # The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. Also use the bundled # tools to re-shape the outputs to the 4D shape of the wind components as they # were read off files. sf, vp = w.sfvp() sf = recover_data(sf, uwnd_info)
# files. ncu = Dataset(example_data_path("uwnd_mean.nc"), "r") uwnd = ncu.variables["uwnd"][:] lons = ncu.variables["longitude"][:] lats = ncu.variables["latitude"][:] ncu.close() ncv = Dataset(example_data_path("vwnd_mean.nc"), "r") vwnd = ncv.variables["vwnd"][:] ncv.close() # The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, "tyx") vwnd, vwnd_info = prep_data(vwnd, "tyx") # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # Create a VectorWind instance to handle the computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta)
def mass_stream_func(uu, vv, lat, lvl, Global=True): ''' Name: MASS_STREAM_FUNC Purpose: An IDL procedure to compute the zonal-mean meridional stream function Inputs: uu : 3-D array of zonal winds [nlvl, nlat, nlon] vv : 3-D array of meridional winds [nlvl, nlat, nlon] lat : A 1- or 2-D array of latitude values (degrees) [nlat] lvl : A 1- or 2-D array of pressure levels (Pa) [nlvl] Outputs: psi : The meridional mass stream function (kg s**-1) p : Pressures for the stream function (Pa) Keywords: Global : Sets if computing for global or for a zonal subset using method of Zhang and Wang (2013) Author and history: Kyle R. Wodzicki ''' if (not Global) and (not VectorWind): raise Exception('Failed to import windspharm.standard.VectorWind!!!') elif (not Global): uu, uu_info = prep_data(uu, 'pyx') # Prepare data for VectorWind class vv, vv_info = prep_data(vv, 'pyx') # Prepare data for VectorWind class lat, uu, vv = order_latdim(lat, uu, vv) # Fix latitude order VW = VectorWind(uu, vv) # Initialize vector wind class uchi, vchi = VW.irrotationalcomponent() # Get irrotational components vv = recover_data( vchi, vv_info) # Convert v-irrot component back to input order vv = np.nanmean(vv, axis=2) # Average over longitude (last dimension) dims = vv.shape # Get dimensions of VV if (lat.ndim == 1): # If the latitude is only 1-D if (dims[0] == lat.size): vv = vv.transpose() dims = vv.shape lat = np.repeat(lat.reshape(1, dims[1]), dims[0], axis=0) # Reshape to match vv array lat = lat[:-1, :] if (lvl.ndim == 1): # If the level is only 1-D lvl = np.repeat(lvl.reshape(dims[0], 1), dims[1], axis=1) # Reshape to match vv array revFlat = False if (lvl[0, 0] > lvl[-1, 0]): # If pressure levels are decending revFlat = True lvl = lvl[::-1, :] # Reverse to ascending vv = vv[::-1, :] # Reverse vv too dp = lvl[1:, :] - lvl[:-1, :] # Compute change in pressure between levels dv = (vv[1:, :] + vv[:-1, :]) / 2.0 # Compute mean wind for level p = ((lvl[1:, 0] + lvl[:-1, 0]) / 2.0).flatten() # Reform mean pressure for level for output psi = 2.0 * np.pi * R_e * np.cos( np.radians(lat)) / g # Compute scalar for psi equation psi *= np.cumsum(dv * dp, axis=0) # Multiply scalar by the integeral of vv * dp return psi, p # Return stream function and pressure
coslat=np.cos(dtr*lats) #coslat[0]=0 # a very small number is used instead #coslat[-1]=0 # ----"""---- # velocity in the Mercator projection um=u/coslat[:,None] vm=v/coslat[:,None] # um checked!!! for j in range(0,np.size(um,axis=0)) : print lats[j], ym[j],u[j,i],um[j,i] # Create a VectorWind instance to handle the computations. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') #w = VectorWind(uwnd, vwnd) # Compute absolute vorticity w = VectorWind(u, v) q = w.absolutevorticity() qbar = q #qbar = np.average(q[:,:,nt[nt>0]],axis=2) print "qbar(4,0)=",qbar[4,0] #qbar checked!!! print "------------------------------"
def test_prep_recover_data(self): # applying preparation and recovery should yield an identical data set u = np.random.rand(12, 17, 73, 144) up, uinfo = prep_data(u, 'tzyx') ur = recover_data(up, uinfo) assert_array_equal(u, ur)
#ncv.close() # find 300 hPa ilev = np.where(levs==300)[0] print("Data uploaded for %d"%iyr) #print uwnd.shape #print uwnd[1,1,1] # The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd1, uwnd_info = prep_data(np.squeeze(uwnd[:,ilev,:,:]), 'tyx') vwnd1, vwnd_info = prep_data(np.squeeze(vwnd[:,ilev,:,:]), 'tyx') # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd1, vwnd1 = order_latdim(lats, uwnd1, vwnd1) # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd1, vwnd1) # Compute the streamfunction and velocity potential. Also use the bundled # tools to re-shape the outputs to the 4D shape of the wind components as they # were read off files. sf, vp = w.sfvp() sf = recover_data(sf, uwnd_info)
A = A1 + A2 + A3 A = scipy.sparse.csc_matrix(A) F0 = F[i, :] U0 = scipy.sparse.linalg.spsolve(A, F0) U[i, :] = U0 test = np.fft.ifft(U) from windspharm.standard import VectorWind from windspharm.tools import prep_data, recover_data, order_latdim uwnd, uwnd_info = prep_data(u_n[4:-4, 4:-4], 'xyp') vwnd, vwnd_info = prep_data(v_n[4:-4, 4:-4], 'xyp') lats, uwnd, vwnd = order_latdim(phi, uwnd, vwnd) w = VectorWind(uwnd, vwnd) sf, vp = w.sfvp() sf = recover_data(sf, uwnd_info) vp = recover_data(vp, vwnd_info) uchi, vchi, upsi, vpsi = w.helmholtz() uchi = recover_data(uchi, uwnd_info) vchi = recover_data(vchi, uwnd_info) upsi = recover_data(upsi, uwnd_info) vpsi = recover_data(vpsi, uwnd_info) s = w.s
## Calculate ensemble mean over QBO phases varum_pos = np.nanmean(varu_pos, axis=0) varum_neg = np.nanmean(varu_neg, axis=0) varvm_pos = np.nanmean(varv_pos, axis=0) varvm_neg = np.nanmean(varv_neg, axis=0) hitum_pos = np.nanmean(hitu_pos, axis=0) hitum_neg = np.nanmean(hitu_neg, axis=0) hitvm_pos = np.nanmean(hitv_pos, axis=0) hitvm_neg = np.nanmean(hitv_neg, axis=0) ### Prepare data for calculating uwnd_pos, uwndinfo_pos = prep_data(varum_pos, 'tyx') vwnd_pos, vwndinfo_pos = prep_data(varvm_pos, 'tyx') uwnd_neg, uwndinfo_neg = prep_data(varum_neg, 'tyx') vwnd_neg, vwndinfo_neg = prep_data(varvm_neg, 'tyx') latn, uwnd_pos, vwnd_pos = order_latdim(lat, uwnd_pos, vwnd_pos) latn, uwnd_neg, vwnd_neg = order_latdim(lat, uwnd_neg, vwnd_neg) hituwnd_pos, hituwndinfo_pos = prep_data(hitum_pos, 'tyx') hitvwnd_pos, hitvwndinfo_pos = prep_data(hitvm_pos, 'tyx') hituwnd_neg, hituwndinfo_neg = prep_data(hitum_neg, 'tyx') hitvwnd_neg, hitvwndinfo_neg = prep_data(hitvm_neg, 'tyx') latn, hituwnd_pos, hitvwnd_pos = order_latdim(lat, hituwnd_pos, hitvwnd_pos)
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
def compute_rws(ds_u, ds_v, lat_coord='lat', lon_coord='lon', time_coord='time'): """ Computation of absolute vorticity, divergence, and Rossby wave source. Outputs xarray datasets of each. Args: ds_u (xarray data array): Zonal (u) wind (m/s). ds_v (xarray data array): Meridional (v) wind (m/s). lat_coord (str): Latitude coordinate. Defaults to ``lat``. lon_coord (str): Longitude coordinate. Defaults to ``lon``. time_coord (str): Time coordinate. Defaults to ``time``. Returns: Xarray datasets for absolute vorticity, divergence, and Rossby wave source. """ from windspharm.standard import VectorWind from windspharm.tools import prep_data, recover_data, order_latdim # grab lat and lon coords lats = ds_u.coords[lat_coord].values lons = ds_u.coords[lon_coord].values time = ds_u.coords[time_coord].values _, wnd_info = prep_data(ds_u.values, 'tyx') # reorder dims into lat, lon, time uwnd = ds_u.transpose(lat_coord, lon_coord, time_coord).values vwnd = ds_v.transpose(lat_coord, lon_coord, time_coord).values # reorder lats to north-south direction lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # initialize wind vector instance w = VectorWind(uwnd, vwnd) # Absolute vorticity (sum of relative and planetary vorticity). eta = w.absolutevorticity() # Horizontal divergence. div = w.divergence() # Irrotational (divergent) component of the vector wind. uchi, vchi = w.irrotationalcomponent() # Computes the vector gradient of a scalar field on the sphere. etax, etay = w.gradient(eta) # Compute rossby wave source S = -eta * div - (uchi * etax + vchi * etay) # recover data shape S = recover_data(S, wnd_info) div = recover_data(div, wnd_info) eta = recover_data(eta, wnd_info) # assemble xarray datasets data_rws = xr.Dataset({ 'rws': (['time', 'lat', 'lon'], S), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={'long_name': 'Rossby wave source'}) data_div = xr.Dataset( { 'div': (['time', 'lat', 'lon'], div), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={'long_name': 'Horizontal divergence (300-mb)'}) data_eta = xr.Dataset( { 'eta': (['time', 'lat', 'lon'], eta), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={ 'long_name': 'Absolute vorticity (sum of relative and planetary vorticity)' }) return data_eta, data_div, data_rws