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.)
def windsfarm(): #Prepare wund for vector w = VectorWind(uwnd, vwnd) #Calculating basic variables eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) planvor = w.planetaryvorticity S = -eta * div - uchi * etax - vchi * etay S = recover_data(S, uwnd_info) return(w, eta, div, uchi, vchi, etax, etay, planvor, S)
def recover_structure(data, flip_lat, orig_info): """Return data to its original structure. Args: data (numpy.ndarray): Data flip_lat (bool): Flag for flipping latitude axis (axis 0) orig_info: Obtained from windspharm.tools.prep_data """ if flip_lat: data = data[::-1, ...] data = recover_data(data, orig_info) return data
# 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) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). S_dec, lons_c = add_cyclic_point(S[11], lons) # Plot Rossby wave source. ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=180)) clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30] fill = ax.contourf(lons_c, lats, S_dec * 1e11, clevs, cmap=plt.cm.RdBu_r, transform=ccrs.PlateCarree(), extend="both") ax.coastlines() ax.gridlines() ax.set_xticks([0, 60, 120, 180, 240, 300, 359.99], crs=ccrs.PlateCarree()) ax.set_yticks([-90, -60, -30, 0, 30, 60, 90], crs=ccrs.PlateCarree()) lon_formatter = LongitudeFormatter(zero_direction_label=True, number_format=".0f") lat_formatter = LatitudeFormatter()
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
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) vp = recover_data(vp, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). sf_dec, lons_c = addcyclic(sf[11], lons) vp_dec, lons_c = addcyclic(vp[11], lons) # Plot streamfunction. m = Basemap(projection='cyl', resolution='c', llcrnrlon=0, llcrnrlat=-90, urcrnrlon=360.01, urcrnrlat=90) x, y = m(*np.meshgrid(lons_c, lats)) clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120] m.contourf(x, y, sf_dec*1e-06, clevs, cmap=plt.cm.RdBu_r, extend='both') m.drawcoastlines()
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)
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) print np.mean(uchi)
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 #convert vort + div from A-grid model to u,v using Spharm vrt, vrt_info = prep_data(vort_n[4:-4, 4:-4], 'xyp') div, div_info = prep_data(div_n[4:-4, 4:-4], 'xyp') lats, vrt, div = order_latdim(phi, vrt, div) vrtspec = s.grdtospec(vrt) divspec = s.grdtospec(div)
### Calculate VelocityWind object wpos = VectorWind(uwnd_pos, vwnd_pos) wneg = VectorWind(uwnd_neg, vwnd_neg) hitwpos = VectorWind(hituwnd_pos, hitvwnd_pos) hitwneg = VectorWind(hituwnd_neg, hitvwnd_neg) ### Calculate Absolute Velocity avrt_pos = wpos.absolutevorticity() avrt_neg = wneg.absolutevorticity() hitavrt_pos = hitwpos.absolutevorticity() hitavrt_neg = hitwneg.absolutevorticity() ### Change dimensions fictapos = recover_data(avrt_pos, uwndinfo_pos) fictaneg = recover_data(avrt_neg, uwndinfo_neg) hitapos = recover_data(hitavrt_pos, hituwndinfo_pos) hitaneg = recover_data(hitavrt_neg, hituwndinfo_neg) ### Calculate dy fictdyapos = np.gradient(fictapos, axis=2) fictdyaneg = np.gradient(fictaneg, axis=2) hitdyapos = np.gradient(hitapos, axis=2) hitdyaneg = np.gradient(hitaneg, axis=2) ### Average over time timeq = np.arange(90, 120) fictdyapos_t = np.nanmean(fictdyapos[timeq, :, :], axis=0)
# 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) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). S_dec, lons_c = addcyclic(S[11], lons) # Plot Rossby wave source. m = Basemap(projection='cyl', resolution='c', llcrnrlon=0, llcrnrlat=-90, urcrnrlon=360.01, urcrnrlat=90) x, y = m(*np.meshgrid(lons_c, lats)) clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30] m.contourf(x, y, S_dec * 1e11, clevs, cmap=plt.cm.RdBu_r, extend='both')
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) vp = recover_data(vp, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). sf_dec, lons_c = add_cyclic_point(sf[11], lons) vp_dec, lons_c = add_cyclic_point(vp[11], lons) # Plot streamfunction. ax1 = plt.axes(projection=ccrs.PlateCarree(central_longitude=180)) clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120] sf_fill = ax1.contourf(lons_c, lats, sf_dec * 1e-06, clevs, transform=ccrs.PlateCarree(),
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