def spht(field, ntrunc=None, gridtype='gaussian'):
"""Transform a field on lat-lon grid to spherical harmonics.
Currently only works for:
- fields defined on the lat-lon points, not latb-lonb.
- triangular truncation.
Returns an xarray.DataArray
"""
nlat, nlon = len(field.lat), len(field.lon)
grid = spharm.Spharmt(nlon, nlat, gridtype=gridtype)
if ntrunc is None:
ntrunc = nlat - 1
other_dims = [d for d in field.dims if d not in ('lat', 'lon')]
# need the field in N-S, E-W form, (lat, lon, other coords)
vfield = (field.sel(lat=sorted(cbbt.lat, reverse=True),
lon=sorted(cbbt.lon)).transpose(
'lat', 'lon', *other_dims))
# calculate the spectral coefficients, given the truncation level
coeffs = grid.grdtospec(vfield.values, ntrunc)
m, n = spharm.getspecindx(ntrunc)
# put the coefficients into a grid, half of which will be empty
# due to the triangular trunctation
cc = np.zeros((ntrunc + 1, ntrunc + 1) + tuple(coeffs.shape[1:]),
dtype=np.complex128)
cc[m, n] = coeffs
coords = [('m', np.arange(0, ntrunc + 1)), ('n', np.arange(0, ntrunc + 1))]
for d in other_dims:
coords.append((d, field.coords[d]))
return xr.DataArray(data=cc, coords=coords, name=field.name)
def sph_filter(field, l_cut, gridtype='gaussian'):
"""Remove all waves above a specific spherical wavenumber l_cut"""
nlat, nlon = len(field.lat), len(field.lon)
grid = spharm.Spharmt(nlon, nlat, gridtype=gridtype)
other_dims = [d for d in field.dims if d not in ('lat', 'lon')]
# need the field in N-S, E-W form, (lat, lon, other coords)
vfield = (field.sel(lat=sorted(cbbt.lat, reverse=True),
lon=sorted(cbbt.lon)).transpose(
'lat', 'lon', *other_dims))
ntrunc = nlat - 1
m, n = spharm.getspecindx(ntrunc)
l = m + n
# transform to spherical modes, eliminate high wavenumbers and transform back
coeffs = grid.grdtospec(vfield.values, ntrunc)
mask = l > l_cut
coeffs[mask] = 0
values = grid.spectogrd(coeffs)
# copy the original field and comply to the same coordinate ordering
nfield = vfield.copy()
nfield.values = values
return nfield.transpose(*field.dims)
def __init__(self,
rsphere=EARTH_RADIUS,
omega=EARTH_OMEGA,
resolution=2.5,
ntrunc=None,
legfunc="stored"):
"""Grid constructor.
- The radius and angular velocity of the planet are given by `rsphere` in
m and `omega` in 1/s, respectively (the default values correspond to
those of the Earth).
- The grid resolution is uniform and specified with the `resolution`
parameter in degrees.
- Uses spherical harmonics for some operations utilizing the
`spharm`/`pyspharm` package. By default (`ntrunc=None`,
`legfunc="stored"`), the spherical harmonics are truncated after (number
of latitudes - 1) functions and precomputed. Consult the documentation of
`spharm.Spharmt` for more information on these parameters.
"""
# Planet parameters (radius, angular velocity)
self.rsphere = rsphere
self.omega = omega
# Setup longitude and latitude grid. Make sure longitude 0° is not
# repeated and that both latitudes 90° and -90° exist. Latitudes start
# at the North pole, which is the convention used by pyspharm.
self.nlon = int(360. / resolution)
self.nlat = int(180. / resolution) + 1
if self.nlat % 2 == 0:
raise ValueError(
"Number of latitudes must be odd but is {}".format(self.nlat))
self.lons = np.linspace(0., 360., self.nlon, endpoint=False)
self.lats = np.linspace(90., -90., self.nlat, endpoint=True)
self.lon, self.lat = np.meshgrid(self.lons, self.lats)
# Grid spacing in degrees
self.dlon = resolution
self.dlat = -resolution
# Spherical coordinate grid (for use with trigonometric functions)
self.lams = np.deg2rad(self.lons)
self.phis = np.deg2rad(self.lats)
self.lam, self.phi = np.meshgrid(self.lams, self.phis)
# Grid spacing in radians
self.dlam = np.deg2rad(self.dlon)
self.dphi = np.deg2rad(self.dlat)
# Precompute Coriolis field
self.fcor = self.coriolis(self.lat)
# Spherical harmonic transform object
self._spharm = spharm.Spharmt(self.nlon,
self.nlat,
rsphere=rsphere,
gridtype="regular",
legfunc=legfunc)
self._ntrunc = (self.nlat - 1) if ntrunc is None else ntrunc
_, self.specindxn = spharm.getspecindx(self._ntrunc)
# Eigenvalues of the horizontal Laplacian for each spherical harmonic.
# Force use of complex64 datatype (= 2 * float32) because spharm will
# cast to float32 components anyway but the multiplication with the
# python scalars results in float64 values.
self.laplacian_eigenvalues = (self.specindxn * (1. + self.specindxn) /
rsphere / rsphere).astype(
np.complex64, casting="same_kind")
def test_case():
"""
Runs an example case: extratropical zonal jets with superimposed sinusoidal NH vorticity
perturbations and a gaussian vorticity tendency forcing.
"""
from time import time
start = time()
# 1) LET'S CREATE SOME INITIAL CONDITIONS
lon_list = list(np.arange(0, 360.1, 2.5))
lat_list = list(np.arange(-87.5, 88, 2.5))[::-1]
lamb, theta = np.meshgrid([x * np.pi / 180. for x in lon_list],
[x * np.pi / 180. for x in lat_list])
# Mean state: zonal extratropical jets
ubar = 25 * np.cos(theta) - 30 * np.cos(theta)**3 + 300 * np.sin(
theta)**2 * np.cos(theta)**6
vbar = np.zeros(np.shape(ubar))
# Initial perturbation: sinusoidal vorticity perturbations
A = 1.5 * 8e-5 # vorticity perturbation amplitude
m = 4 # zonal wavenumber
theta0 = np.deg2rad(45) # center lat = 45 N
thetaW = np.deg2rad(15)
vort_pert = 0.5 * A * np.cos(theta) * np.exp(-(
(theta - theta0) / thetaW)**2) * np.cos(m * lamb)
# Get U' and V' from this vorticity perturbation
s = spharm.Spharmt(len(lon_list),
len(lat_list),
gridtype='regular',
legfunc='computed',
rsphere=NL.Re)
uprime, vprime = s.getuv(s.grdtospec(vort_pert),
np.zeros(np.shape(s.grdtospec(vort_pert))))
# Full initial conditions dictionary:
ics = {
'u_bar': ubar,
'v_bar': vbar,
'u_prime': uprime,
'v_prime': vprime,
'lons': lon_list,
'lats': lat_list,
'start_time': datetime(2017, 1, 1, 0)
}
# 2) LET'S ALSO FEED IN A GAUSSIAN NH RWS FORCING
amplitude = 10e-10 # s^-2
forcing = np.zeros(np.shape(ubar))
x, y = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10))
d = np.sqrt(x * x + y * y)
sigma, mu = 0.5, 0.0
g = np.exp(-((d - mu)**2 / (2.0 * sigma**2))) # GAUSSIAN CURVE
lat_i = np.where(np.array(lat_list) == 35.)[0][0]
lon_i = np.where(np.array(lon_list) == 160.)[0][0]
forcing[lat_i:lat_i + 10, lon_i:lon_i + 10] = g * amplitude
# 3) INTEGRATE!
model = Model(ics, forcing=forcing)
model.integrate()
print('TOTAL INTEGRATION TIME: {:.02f} minutes'.format(
(time() - start) / 60.))
def super_rotation(linearized=False, idate=None, ntrunc=42):
"""
Creates ICs for an idealized case: "zonal flow corresponding to
a super-rotation of the atmosphere with a maximum value of ~15.4 m/s
on the equator." (Sardeshmukh and Hoskins 1988)
Args
----
linearized : bool
True if this is a linearized model.
idate : datetime
Forecast initialization date
ntrunc : int
Triangular trunction (e.g., 42 for T42).
Returns
-------
ics : dict
Model initial state (a dictionary of DataArrays)
"""
# Set the init. date if not provided
if idate is None:
idate = datetime.utcnow().replace(hour=0,
minute=0,
second=0,
microsecond=0)
# Create a gaussian lat/lon grid
if ntrunc not in TRUNC_NLATS.keys():
raise ValueError(
'Truncation T{} is not in the dictionary TRUNC_NLATS'.format(
ntrunc))
nlat = TRUNC_NLATS[ntrunc]
lons, lats = gaussian_latlon_grid(nlat)
theta = lats * np.pi / 180.
lamb = lons * np.pi / 180.
# Mean state: zonal extratropical jets
ubar = Re * Omega * np.cos(theta) / 30.
vbar = np.zeros(ubar.shape)
# Get the mean state vorticity from ubar and vbar
s = spharm.Spharmt(lamb.shape[1],
lamb.shape[0],
gridtype='gaussian',
rsphere=Re,
legfunc='computed')
vortb_spec, _ = s.getvrtdivspec(ubar, vbar)
vort_bar = s.spectogrd(vortb_spec)
vort_prime = np.zeros(vort_bar.shape)
# Generate the state
return _generate_state(idate,
lats,
lons,
vort_bar,
vort_prime,
linearized=linearized)
def __init__(self, ics, forcing=None):
"""
Initializes the model.
Requires:
ics -----> Dictionary of linearized fields and space/time dimensions
keys: u_bar, v_bar, u_prime, v_prime, lats, lons, start_time
forcing -> a 2D array (same shape as model fields) containing a
vorticity tendency [s^-2] to be imposed at each integration time step
"""
# 1) STORE SPACE/TIME VARIABLES (DIMENSIONS)
# Get the latitudes and longitudes (as lists)
self.lats = ics['lats']
self.lons = ics['lons']
self.start_time = ics['start_time'] # datetime
self.curtime = self.start_time
# 2) GENERATE/STORE NONDIVERGENT INITIAL STATE
# Set up the spherical harmonic transform object
self.s = spharm.Spharmt(self.nlons(),
self.nlats(),
rsphere=NL.Re,
gridtype='regular',
legfunc='computed')
# Truncation for the spherical transformation
if NL.M is None:
self.ntrunc = self.nlats()
else:
self.ntrunc = NL.M
# Use the object to get the initial conditions
# First convert to vorticity using spharm object
vortb_spec, div_spec = self.s.getvrtdivspec(ics['u_bar'], ics['v_bar'])
vortp_spec, div_spec = self.s.getvrtdivspec(ics['u_prime'],
ics['v_prime'])
div_spec = np.zeros(
vortb_spec.shape) # Only want NON-DIVERGENT part of wind
# Re-convert this to u-v winds to get the non-divergent component
# of the wind field
self.ub, self.vb = self.s.getuv(vortb_spec, div_spec) # MEAN WINDS
self.up, self.vp = self.s.getuv(vortp_spec,
div_spec) # PERTURBATION WINDS
# Use these winds to get the streamfunction (psi) and
# velocity potential (chi)
self.psib, chi = self.s.getpsichi(self.ub,
self.vb) # MEAN STREAMFUNCTION
self.psip, chi = self.s.getpsichi(
self.up, self.vp) # PERTURBATION STREAMFUNCTION
# Convert the spectral vorticity to grid
self.vort_bar = self.s.spectogrd(vortb_spec) # MEAN RELATIVE VORTICITY
self.vortp = self.s.spectogrd(
vortp_spec) # PERTURBATION RELATIVE VORTICITY
# 3) STORE A COUPLE MORE VARIABLES
# Map projections for plotting
self.bmaps = create_basemaps(self.lons, self.lats)
# Get the vorticity tendency forcing (if any) for integration
self.forcing = forcing
def make_spharmt(self):
try:
self.spharmt = self.vectorwind.s
except AttributeError:
self.spharmt = spharm.Spharmt(self.n_lon,
self.n_lat,
rsphere=self._rsphere,
gridtype=self._gridtype,
legfunc=self._legfunc)
def regrid(lon_in,lat_in,data,nlon_o,nlat_o):
''' Given data(nlat,nlon), regrid the data into (nlat_o,nlon_o)
using the spharm.regrid routine
If the input data is not on a complete sphere, copy the input
onto a complete sphere before regridding
Input:
lon_in - the longitude the input data is on
lat_in - the latitude the input data is on
data - numpy.ndarray
nlon_o - integer
nlat_o - integer
Return:
numpy.ndarray of shape (nlat_o,nlon_o)
'''
if not _SPHARM_INSTALLED:
raise _import_error
assert data.ndim <= 3
# determine if the domain is a complete sphere
dlon = numpy.diff(lon_in[0:2])
AllLon = dlon*(len(lon_in)+1) > 360.
dlat = numpy.diff(lat_in[0:2])
AllLat = dlat*(len(lat_in)+1) > 180.
sliceobj = [slice(None),]*data.ndim
if AllLon and AllLat: # The data covers the whole sphere
inputdata = data
nlat_i = len(lat_in)
nlon_i = len(lon_in)
lat = lat_in
lon = lon_in
else:
inputdata,lon,lat = regional2global(data,lon_in,lat_in)
nlat_i = len(lat)
nlon_i = len(lon)
gridin = spharm.Spharmt(nlon_i,nlat_i)
gridout = spharm.Spharmt(nlon_o,nlat_o)
return spharm.regrid(gridin,gridout,inputdata)
def rws_from_tropical_divergence(state, center=(0., 145.), amp=6e-6, width=12):
# Get desired state variables
lats = state['latitude'].values
lons = state['longitude'].values
vort_bar = state['base_atmosphere_relative_vorticity'].values
s = spharm.Spharmt(lats.shape[1], lons.shape[0], gridtype='regular', rsphere=Re)
vortb_spec = s.grdtospec(vort_bar)
ubar, vbar = s.getuv(vortb_spec, np.zeros(vortb_spec.shape))
divergence = gaussian_blob_2d(lats, lons, center, width, amp)
# Calculate the Rossby Wave Source
# Term 1
zetabar_spec, _ = s.getvrtdivspec(ubar, vbar)
zetabar = s.spectogrd(zetabar_spec) + 2 * Omega * np.sin(np.deg2rad(lats))
term1 = -zetabar * divergence
# Term 2
uchi, vchi = s.getuv(np.zeros(zetabar_spec.shape), s.grdtospec(divergence))
dzeta_dx, dzeta_dy = s.getgrad(s.grdtospec(zetabar))
term2 = - uchi * dzeta_dx - vchi * dzeta_dy
rws = term1 + term2
return rws, lats, lons
def compute_del4vort(vort, ntrunc):
# Compute del^4(vorticity) with spherical harmonics
# Approximating del^4 as: d4_dx4 + d4_dy4 + 2 * (d2_dx2 * d2_dy2)
s = spharm.Spharmt(vort.shape[1],
vort.shape[0],
rsphere=6378100.,
gridtype='gaussian',
legfunc='computed')
vspec = s.grdtospec(vort, ntrunc=ntrunc)
# First order
dvort_dx, dvort_dy = s.getgrad(vspec)
# Second order
d2vort_dx2, _ = s.getgrad(s.grdtospec(dvort_dx, ntrunc=ntrunc))
_, d2vort_dy2 = s.getgrad(s.grdtospec(dvort_dy, ntrunc=ntrunc))
# Fourth order
d4vort_dx4, _ = s.getgrad(
s.grdtospec(s.getgrad(s.grdtospec(d2vort_dx2, ntrunc=ntrunc))[0],
ntrunc=ntrunc))
_, d4vort_dy4 = s.getgrad(
s.grdtospec(s.getgrad(s.grdtospec(d2vort_dy2, ntrunc=ntrunc))[1],
ntrunc=ntrunc))
# Put it all together to approximate del^4
del4vort = d4vort_dx4 + d4vort_dy4 + (2 * d2vort_dx2 * d2vort_dy2)
return del4vort
def __init__(self, earth, ice, grid=None, topo=None):
"""
Compute glacial isostatic adjustment on a globe.
Paramaters
----------
earth : <giapy.code.earth_tools.earthSpherical.SphericalEarth>
ice : <giapy.code.icehistory.IceHistory / PersistentIceHistory>
grid : <giapy.code.map_tools.GridObject>
topo : numpy.ndarray
Methods
-------
performConvolution
"""
self.earth = earth
self.ice = ice
self.nlat, self.nlon = ice.shape
# The grid used is a cylindrical projection (equispaced lat/lon grid
# unless otherwise specified)
if grid is not None:
self.grid = grid
else:
self.grid = GridObject(mapparam={'projection': 'cyl'},
shape=ice.shape)
self.topo = topo
# Precompute and store harmonic transform coefficients, for
# computational efficiency, but at a memory cost.
self.harmTrans = spharm.Spharmt(self.nlon, self.nlat, legfunc='stored')
grbs = pygrib.open('../sampledata/spherical_pressure_level.grib1')
g = grbs[1]
fld = g.values
# ECMWF normalizes the spherical harmonic coeffs differently than NCEP.
# (m=0,n=0 is global mean, instead of sqrt(2)/2 times global mean)
fld = 2. * fld / np.sqrt(2.)
fldr = fld[0::2]
fldi = fld[1::2]
fld = np.zeros(fldr.shape, 'F')
fld.real = fldr
fld.imag = fldi
nlons = 360
nlats = 181
s = spharm.Spharmt(nlons, nlats)
data = s.spectogrd(fld)
lons = (360. / nlons) * np.arange(nlons)
lats = 90. - (180. / (nlats - 1)) * np.arange(nlats)
lons, lats = np.meshgrid(lons, lats)
# stack grids side-by-side (in longitiudinal direction), so
# any range of longitudes (between -360 and 360) may be plotted on a world map.
lons = np.concatenate((lons - 360, lons), 1)
lats = np.concatenate((lats, lats), 1)
data = np.concatenate((data, data), 1)
# setup miller cylindrical map projection.
m = Basemap(llcrnrlon=-180.,llcrnrlat=-90,urcrnrlon=180.,urcrnrlat=90.,\
resolution='l',area_thresh=10000.,projection='mill')
x, y = m(lons, lats)
CS = m.contourf(x, y, data, 15, cmap=plt.cm.jet)
ax = plt.gca()
import cartopy.crs as ccrs
import cartopy.util
import spharm
import numpy as np
import cartopy
import cartopy.crs as ccrs
import cartopy.util
import spharm
import numpy as np
field = np.random.normal(0, 1, (64, 128))
x = spharm.Spharmt(128,
64,
rsphere=6.4e6,
gridtype='gaussian',
legfunc='computed')
spec = x.grdtospec(field)
field_2 = x.spectogrd(spec)
#plt.figure(figsize=(10,10))
fig, ax = plt.subplots(2, 1, figsize=(10, 10))
ax[0].imshow(field)
ax[1].imshow(field_2)
#plt.colorbar()
plt.figure(figsize=(10, 10))
plt.imshow(field_2)
def create_sfn_velpot_vortdiv(reanal):
ds_U = netCDF4.Dataset('reanalysis_clean/{0}.monthly.U.nc'.format(reanal))
ds_VOR = netCDF4.Dataset(
'reanalysis_clean/{0}.monthly.VORTICITY.nc'.format(reanal))
ds_DIV = netCDF4.Dataset(
'reanalysis_clean/{0}.monthly.DIVERGENCE.nc'.format(reanal))
U = ds_U.variables['U']
VOR = ds_VOR.variables['VORTICITY']
DIV = ds_DIV.variables['DIVERGENCE']
ntime, nlev, nlat, nlon = U.shape
assert (U.shape == VOR.shape)
assert (U.shape == DIV.shape)
out_fh_1 = create_output_file(reanal,
'monthly',
'HORIZ_SFN',
'm2 s-1',
"Atmospheric Horizontal Streamfunction",
ds_U.variables['level'][:],
ds_U.variables['latitude'][:],
ds_U.variables['longitude'][:],
noclobber=True,
compress=True)
out_fh_2 = create_output_file(reanal,
'monthly',
'HORIZ_VEL_POT',
'm2 s-1',
"Atmospheric Horizontal Streamfunction",
ds_U.variables['level'][:],
ds_U.variables['latitude'][:],
ds_U.variables['longitude'][:],
noclobber=True,
compress=True)
horiz_sfn = out_fh_1.variables['HORIZ_SFN']
vel_pot = out_fh_2.variables['HORIZ_VEL_POT']
time_1 = out_fh_1.variables['time']
time_2 = out_fh_2.variables['time']
lats = ds_U.variables['latitude'][:]
lons = ds_U.variables['longitude'][:]
# spharm requires grid to go from North to South
# and from 0 to 360 positive
# Some data can be from -180 to +180 but since a sphere is invariant under
# rotations, that should be fine, as long as the direction is eastward.
if lats[-1] > lats[0]:
lats_increasing = True
assert (np.allclose(np.linspace(-90, 90, nlat), lats, 1.E-5, 1.E-5))
else:
lats_increasing = False
assert (np.allclose(np.linspace(90, -90, nlat), lats, 1.E-5, 1.E-5))
if lons[-1] > lons[0]:
lons_increasing = True
else:
lons_increasing = False
# check if the grid is regular and increasing
assert (lons[1] - lons[0] > 0.)
assert (np.allclose(np.diff(lons), lons[1] - lons[0], 1.0E-4))
psi = np.zeros(U[0].shape, dtype=np.float32)
chi = np.zeros(U[0].shape, dtype=np.float32)
# this could be written in a much faster way, but it's fast enough for now
s = spharm.Spharmt(nlon=nlon, nlat=nlat)
for itime in range(ntime):
print("{0}/{1}".format(itime + 1, ntime))
time_1[itime] = ds_U.variables['time'][itime]
time_2[itime] = ds_U.variables['time'][itime]
vor_current = VOR[itime]
div_current = DIV[itime]
if type(vor_current) is np.ma.core.MaskedArray:
vor_current = vor_current.filled(0.)
if type(div_current) is np.ma.core.MaskedArray:
div_current = div_current.filled(0.)
if lats_increasing:
vor_current = vor_current[:, ::-1, :]
div_current = div_current[:, ::-1, :]
if not lons_increasing:
vor_current = vor_current[:, :, ::-1]
div_current = div_current[:, :, ::-1]
for ilev in range(nlev):
print("\t{0}".format(ilev + 1))
vor_spec = s.grdtospec(vor_current[ilev])
div_spec = s.grdtospec(div_current[ilev])
psi_spec = _spherepack.invlap(vor_spec, s.rsphere)
chi_spec = _spherepack.invlap(div_spec, s.rsphere)
psi_current = np.squeeze(s.spectogrd(psi_spec))
chi_current = np.squeeze(s.spectogrd(chi_spec))
print("\t\t{0:12.4g}\t{1:12.4g}".format(np.mean(psi_current),
np.mean(chi_current)))
psi[ilev] = psi_current
chi[ilev] = chi_current
assert (not (np.allclose(psi[ilev], chi[ilev], 1.0E-6, 1.0E-6)))
if lats_increasing:
psi = psi[:, ::-1, :]
chi = chi[:, ::-1, :]
if not lons_increasing:
psi = psi[:, :, ::-1]
chi = chi[:, :, ::-1]
horiz_sfn[itime] = psi
vel_pot[itime] = chi
if itime % 10 == 0:
out_fh_1.sync()
out_fh_2.sync()
out_fh_1.close()
out_fh_2.close()
def test_case():
"""
Runs an example case: extratropical zonal jets with superimposed sinusoidal NH vorticity
perturbations and a gaussian vorticity tendency forcing.
"""
from time import time
start = time()
# 1) LET'S CREATE SOME INITIAL CONDITIONS
lons = np.arange(0, 360.1, 2.5)
lats = np.arange(-87.5, 88, 2.5)[::-1]
lamb, theta = np.meshgrid(lons * np.pi / 180., lats * np.pi / 180.)
# Mean state: zonal extratropical jets
ubar = NL.mag * np.cos(theta) - 30 * np.cos(theta)**3 + 300 * np.sin(
theta)**2 * np.cos(theta)**6
vbar = np.zeros(np.shape(ubar))
# Initial perturbation: sinusoidal vorticity perturbations
theta0 = np.deg2rad(NL.pert_center_lat) # center lat = 45 N
thetaW = np.deg2rad(NL.pert_width) #15
vort_pert = 0.5 * NL.A * np.cos(theta) * np.exp(-(
(theta - theta0) / thetaW)**2) * np.cos(NL.m * lamb)
# Get U' and V' from this vorticity perturbation
s = spharm.Spharmt(len(lons),
len(lats),
gridtype='regular',
legfunc='computed',
rsphere=NL.Re)
uprime, vprime = s.getuv(s.grdtospec(vort_pert),
np.zeros(np.shape(s.grdtospec(vort_pert))))
# Full initial conditions dictionary:
ics = {
'u_bar': ubar,
'v_bar': vbar,
'u_prime': uprime,
'v_prime': vprime,
'lons': lons,
'lats': lats,
'start_time': datetime(2017, 1, 1, 0)
}
# 2) LET'S ALSO FEED IN A GAUSSIAN NH RWS FORCING (CAN BE USED TO CREATE SYNTEHTIC HURRICANES)
amplitude = NL.forcing_amp # s^-2
forcing = np.zeros(np.shape(ubar))
x, y = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10))
d = np.sqrt(x * x + y * y)
sigma, mu = 0.5, 0.0
g = np.exp(-((d - mu)**2 / (2.0 * sigma**2))) # GAUSSIAN CURVE
source_lat = NL.forcing_lat
source_lon = NL.forcing_lon # The location of the forcing
lat_i = np.where(lats == source_lat)[0][0]
lon_i = np.where(lons == source_lon)[0][0]
if NL.use_forcing == True:
forcing[lat_i:lat_i + 10, lon_i:lon_i + 10] = g * amplitude
else:
forcing[:, :] = 0
# 3) INTEGRATE!
model = Model(ics, forcing=forcing)
model.integrate()
print('TOTAL INTEGRATION TIME: {:.02f} minutes'.format(
(time() - start) / 60.))
plt.plot((NL.dt * np.arange(len(model.tot_ke))) / 3600.,
(1 - model.tot_ke / model.expected_ke) * 100)
#plt.plot(np.arange(len(model.tot_ke)), model.tot_ke, 'o-')
plt.title('Model Kinetic Energy Error vs. Time Step')
plt.xlabel("Model Time [hr]")
plt.ylabel(u'Model Kinetic Energy Error [%]')
print("Expected Kinetic Energy [m^2/s^2]:", model.expected_ke)
plt.savefig(NL.figdir + '/model_ke_ts.png', bbox_inches='tight')
def from_u_and_v_winds(lats,
lons,
ubar,
vbar,
uprime=None,
vprime=None,
interp=True,
linearized=False,
ntrunc=None,
idate=None):
"""
Creates ICs from numpy arrays describing the intial wind field.
Args
----
lats : numpy array
1D array of global latitudes (in degrees)
lons : numpy array
1D array of global longitudes (in degrees)
ubar : numpy array
2D array (nlat, nlon) of mean state zonal winds
vbar : numpy array
2D array (nlat, nlon) of mean state meridional winds
uprime : numpy array
2D array (nlat, nlon) of perturbation zonal winds
vprime : numpy array
2D array (nlat, nlon) of perturbation meridional winds
interp : bool
If True, fields will be interpolated to a gaussian grid.
linearized : bool
True if this is a linearized model.
ntrunc : int
Triangular trunction (e.g., 42 for T42).
idate : datetime
Foreast initialization date
Returns
-------
ics : dict
Model initial state (a dictionary of DataArrays)
"""
# Set the init. date if not provided
if idate is None:
idate = datetime.utcnow().replace(hour=0,
minute=0,
second=0,
microsecond=0)
# If only ubar and vbar are provided, then forecast will be nonlinear
if uprime is None or vprime is None:
uprime = np.zeros(ubar.shape)
vprime = np.zeros(vbar.shape)
linearized = False
# Interpolate to a gaussian grid, if necessary
gridlons, gridlats = np.meshgrid(lons, lats)
if interp:
if ntrunc not in TRUNC_NLATS.keys():
raise ValueError(
'Truncation T{} is not in the dictionary TRUNC_NLATS'.format(
ntrunc))
nlat = TRUNC_NLATS[ntrunc]
lats, lons, ubar = interp_to_gaussian(gridlats,
gridlons,
ubar,
nlat=nlat,
return_latlon=True)
vbar = interp_to_gaussian(gridlats, gridlons, vbar, nlat=nlat)
if linearized:
uprime = interp_to_gaussian(gridlats, gridlons, uprime, nlat=nlat)
vprime = interp_to_gaussian(gridlats, gridlons, vprime, nlat=nlat)
else:
uprime = np.zeros(ubar.shape)
vprime = np.zeros(vbar.shape)
else:
lons, lats = gridlons, gridlats
# Get the mean state & perturbation vorticity from the winds
s = spharm.Spharmt(lats.shape[1],
lats.shape[0],
gridtype='gaussian',
rsphere=get_constant('planetary_radius', 'm'),
legfunc='computed')
vortb_spec, _ = s.getvrtdivspec(ubar, vbar, ntrunc=ntrunc)
vort_bar = s.spectogrd(vortb_spec)
vortp_spec, _ = s.getvrtdivspec(uprime, vprime, ntrunc=ntrunc)
vort_prime = s.spectogrd(vortp_spec)
# Generate the state
return _generate_state(idate,
lats,
lons,
vort_bar,
vort_prime,
linearized=linearized)
def integrate(init_cond, bmaps, ntimes=480):
""" Function that integrates the barotropic model using
spherical harmonics
Input:
init_cond : dictionary of intial conditions
containing u and v initial conditions, the latitude
and longitudes describing the grids, and a starting time
bmaps : a dictionary of global and regional basemaps and
projected coordinates
ntimes : Number of timesteps to integrate
"""
# Get the initial u and v wind fields
u = init_cond['u_in']
v = init_cond['v_in']
ntrunc = len(init_cond['lats'])
start_time = init_cond['start_time']
# Create a radian grid
lat_list_r = [x * np.pi / 180. for x in init_cond['lats']]
lon_list_r = [x * np.pi / 180. for x in init_cond['lons']]
# Meshgrid
lons, lats = np.meshgrid(init_cond['lons'], init_cond['lats'])
lamb, theta = np.meshgrid(lon_list_r, lat_list_r)
dlamb = np.gradient(lamb)[1]
dtheta = np.gradient(theta)[0]
# Here is the Coriolis parameter
f = 2 * 7.292E-5 * np.sin(theta)
# Set up the spherical harmonic transform object
s = spharm.Spharmt(len(init_cond['lons']),
len(init_cond['lats']),
rsphere=Re,
gridtype='regular',
legfunc='computed')
# Use the object to plot the initial conditions
# First convert to vorticity using spharm object
vort_spec, div_spec = s.getvrtdivspec(u, v)
div_spec = np.zeros(
vort_spec.shape) # Only want non-divergent part of wind
# Re-convert this to u-v winds to get the non-divergent component
# of the wind field
u, v = s.getuv(vort_spec, div_spec)
# Use these winds to get the streamfunction (psi) and
# velocity potential (chi)
psi, chi = s.getpsichi(u, v)
# Convert the spectral vorticity to grid
vort_now = s.spectogrd(vort_spec)
# Plot Initial Conditions
curtime = start_time
plot_figures(0, curtime, u, v, vort_now, psi, bmaps)
# Now loop through the timesteps
for n in xrange(ntimes):
# Compute spectral vorticity from u and v wind
vort_spec, div_spec = s.getvrtdivspec(u, v)
# Now get the actual vorticity
vort_now = s.spectogrd(vort_spec)
div = np.zeros(
vort_now.shape) # Divergence is zero in barotropic vorticity
# Here we actually compute vorticity tendency
# Compute tendency with beta as only forcing
vort_tend_rough = -2. * 7.292E-5/(Re**2) * d_dlamb(psi,dlamb) -\
Jacobian(psi,vort_now,theta,dtheta,dlamb)
# Apply hyperdiffusion if requested for smoothing
if use_hyperdiffusion:
vort_tend = add_hyperdiffusion(s, vort_now, vort_tend_rough,
ntrunc)
else:
vort_tend = vort_tend_rough
if n == 0:
# First step just do forward difference
# Vorticity at next time is just vort + vort_tend * dt
vort_next = vort_now + vort_tend[:, :, 0] * dt
else:
# Otherwise do leapfrog
vort_next = vort_prev + vort_tend[:, :, 0] * 2 * dt
# Invert this new vort to get the new psi (or rather, uv winds)
# First go back to spectral space
vort_spec = s.grdtospec(vort_next)
div_spec = s.grdtospec(div)
# Now use the spharm methods to get new u and v grid
u, v = s.getuv(vort_spec, div_spec)
psi, chi = s.getpsichi(u, v)
#raw_input()
# Change vort_now to vort_prev
# and if not first step add Robert filter
# (from Held barotropic model)
# to dampen out crazy modes
r = 0.2
if n == 0:
vort_prev = vort_now
else:
vort_prev = (1. - 2. * r) * vort_now + r * (vort_next + vort_prev)
cur_fhour = (n + 1) * dt / 3600.
curtime = start_time + timedelta(hours=cur_fhour)
# Output every six hours
if cur_fhour % output_freq == 0 and plot_output:
# Go from psi to geopotential
#phi = divide(psi * f,9.81)
print("Plotting hour", cur_fhour)
plot_figures(cur_fhour, curtime, u, v, vort_next, psi, bmaps)
def sinusoidal_perts_on_zonal_jet(linearized=False,
idate=None,
ntrunc=42,
amp=8e-5,
m=4,
theta0=45.,
theta_w=15.):
"""
Creates ICs for an idealized case: extratropical zonal jets
with superimposed sinusoidal NH vorticity perturbations.
Taken from the GFDL model documentation (default case)
Args
----
linearized : bool
True if this is a linearized model.
idate : datetime
Forecast initialization date
ntrunc : int
Triangular trunction (e.g., 42 for T42).
amp : float
Vorticity perturbation amplitude [s^-1].
m : int
Vorticity perturbation zonal wavenumber.
theta0 : float
Center latitude [degrees] for the vorticity perturbations.
theta_w : float
Halfwidth [degrees lat/lon] of the vorticity perturbations.
Returns
-------
ics : dict
Model initial state (a dictionary of DataArrays)
"""
# Set the init. date if not provided
if idate is None:
idate = datetime.utcnow().replace(hour=0,
minute=0,
second=0,
microsecond=0)
# Create a gaussian lat/lon grid
if ntrunc not in TRUNC_NLATS.keys():
raise ValueError(
'Truncation T{} is not in the dictionary TRUNC_NLATS'.format(
ntrunc))
nlat = TRUNC_NLATS[ntrunc]
lons, lats = gaussian_latlon_grid(nlat)
theta = np.deg2rad(lats)
lamb = np.deg2rad(lons)
# Mean state: zonal extratropical jets
ubar = 25 * np.cos(theta) - 30 * np.cos(theta)**3 + \
300 * np.sin(theta)**2 * np.cos(theta)**6
vbar = np.zeros(np.shape(ubar))
# Get the mean state vorticity from ubar and vbar
s = spharm.Spharmt(lamb.shape[1],
lamb.shape[0],
gridtype='gaussian',
rsphere=Re,
legfunc='computed')
vortb_spec, _ = s.getvrtdivspec(ubar, vbar, ntrunc=42)
vort_bar = s.spectogrd(vortb_spec)
# Initial perturbation: sinusoidal vorticity perturbations
theta0 = np.deg2rad(theta0) # center lat --> radians
theta_w = np.deg2rad(theta_w) # halfwidth ---> radians
vort_prime = 0.5 * amp * np.cos(theta) * np.exp(-((theta-theta0)/theta_w)**2) * \
np.cos(m*lamb)
vort_prime = s.spectogrd(s.grdtospec(vort_prime, ntrunc=42))
# Generate the state
return _generate_state(idate,
lats,
lons,
vort_bar,
vort_prime,
linearized=linearized)
def __init__(self, nx, ny):
self.spharmt = spharm.Spharmt(nx, ny)