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):
# 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')
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, 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
