# initialize spectral tendency arrays
ddivdtspec = np.zeros(vrtspec.shape + (3, ), np.complex64)
dvrtdtspec = np.zeros(vrtspec.shape + (3, ), np.complex64)
dphidtspec = np.zeros(vrtspec.shape + (3, ), np.complex64)
nnew = 0
nnow = 1
nold = 2
# time loop.
time1 = time.clock()
for ncycle in range(itmax + 1):
t = ncycle * dt
# get vort,u,v,phi on grid
vrtg = x.spectogrd(vrtspec)
ug, vg = x.getuv(vrtspec, divspec)
phig = x.spectogrd(phispec)
print('t=%6.2f hours: min/max %6.2f, %6.2f' %
(t / 3600., vg.min(), vg.max()))
# compute tendencies.
tmpg1 = ug * (vrtg + f)
tmpg2 = vg * (vrtg + f)
ddivdtspec[:,
nnew], dvrtdtspec[:,
nnew] = x.getvrtdivspec(tmpg1, tmpg2, ntrunc)
dvrtdtspec[:, nnew] *= -1
tmpg1 = ug * phig
tmpg2 = vg * phig
tmpspec, dphidtspec[:, nnew] = x.getvrtdivspec(tmpg1, tmpg2, ntrunc)
dphidtspec[:, nnew] *= -1
tmpspec = x.grdtospec(phig + 0.5 * (ug**2 + vg**2), ntrunc)
class PowerSpectrum(EvaluationMethod):
def __init__(self, H, W):
super(PowerSpectrum, self).__init__()
self.spharm = Spharmt(W, H, legfunc='stored')
self.spectrums = []
self.needs_sh = True
def _calc_energy_spectrum(self, div_vort):
sh = self.spharm.grdtospec(div_vort)
u, v = self.spharm.getuv(sh[..., 1], sh[..., 0])
u_sh = pysh.expand.SHExpandDH(u)
u_spectrum = pysh.spectralanalysis.spectrum(u_sh) # <-> u^2
v_sh = pysh.expand.SHExpandDH(v)
v_spectrum = pysh.spectralanalysis.spectrum(v_sh) # <-> v^2
return 0.5 * (u_spectrum + v_spectrum) # <-> 0.5 * (u^2 + v^2)
def set_shape(self, shape):
self.shape = shape
self.spectrums = [[] for _ in range(self.shape[-1])]
self.energy_spectrum = []
def evaluate_SR_sh(self, i, LR, SR, SR_sh):
for c in range(SR_sh.shape[-1]):
self.spectrums[c] += [
pysh.spectralanalysis.spectrum(sh) for sh in SR_sh[..., c]
]
self.energy_spectrum += [self._calc_energy_spectrum(img) for img in SR]
def finalize(self):
spectrums = np.asarray(self.spectrums) # C x N x L
mean_spectrum = np.mean(spectrums, axis=1)
for c, spectrum in enumerate(mean_spectrum):
np.savetxt(self.dir / f'spectrum_channel_{c}.csv', spectrum)
energy_spectrum = np.mean(np.asarray(self.energy_spectrum), axis=0)
np.savetxt(self.dir / 'energy_spectrum.csv', energy_spectrum)
self.summarize({'SR': self.dir}, self.dir)
def summarize(self, paths, outdir):
p = paths[next(iter(paths))]
C = len(glob(str(p / 'spectrum_channel_*.csv')))
for c in range(C):
plt.figure(figsize=(10, 5))
plt.xscale('log')
plt.yscale('log')
plt.xlabel('wavenumber')
plt.ylabel('energy')
for name, path in paths.items():
spectrum = np.loadtxt(path / f'spectrum_channel_{c}.csv')
x = np.arange(1, 1 + spectrum.shape[0])
plt.plot(x, spectrum, label=name)
plt.plot(x, x**(-5 / 3) * spectrum[0], '--')
plt.legend()
plt.savefig(outdir / f'spectrum_channel_{c}.png')
plt.figure(figsize=(10, 5))
plt.xscale('log')
plt.yscale('log')
plt.xlabel('wavenumber')
plt.ylabel('energy')
for name, path in paths.items():
spectrum = np.loadtxt(path / 'energy_spectrum.csv')
x = np.arange(1, 1 + spectrum.shape[0])
plt.plot(x, spectrum, label=name)
plt.plot(x, x**(-5 / 3) * spectrum[0], '--')
plt.legend()
plt.savefig(outdir / 'energy_spectrum.png')
def main():
# non-linear barotropically unstable shallow water test case
# of Galewsky et al (2004, Tellus, 56A, 429-440).
# "An initial-value problem for testing numerical models of the global
# shallow-water equations" DOI: 10.1111/j.1600-0870.2004.00071.x
# http://www-vortex.mcs.st-and.ac.uk/~rks/reprints/galewsky_etal_tellus_2004.pdf
# requires matplotlib for plotting.
# grid, time step info
nlons = 256 # number of longitudes
ntrunc = int(nlons/3) # spectral truncation (for alias-free computations)
nlats = int(nlons/2) # for gaussian grid.
dt = 150 # time step in seconds
itmax = 6*int(86400/dt) # integration length in days
# parameters for test
rsphere = 6.37122e6 # earth radius
omega = 7.292e-5 # rotation rate
grav = 9.80616 # gravity
hbar = 10.e3 # resting depth
umax = 80. # jet speed
phi0 = np.pi/7.; phi1 = 0.5*np.pi - phi0; phi2 = 0.25*np.pi
en = np.exp(-4.0/(phi1-phi0)**2)
alpha = 1./3.; beta = 1./15.
hamp = 120. # amplitude of height perturbation to zonal jet
efold = 3.*3600. # efolding timescale at ntrunc for hyperdiffusion
ndiss = 8 # order for hyperdiffusion
# setup up spherical harmonic instance, set lats/lons of grid
x = Spharmt(nlons,nlats,ntrunc,rsphere,gridtype='gaussian')
lons,lats = np.meshgrid(x.lons,x.lats)
f = 2.*omega*np.sin(lats) # coriolis
# zonal jet.
vg = np.zeros((nlats,nlons),np.float)
u1 = (umax/en)*np.exp(1./((x.lats-phi0)*(x.lats-phi1)))
ug = np.zeros((nlats),np.float)
ug = np.where(np.logical_and(x.lats < phi1, x.lats > phi0), u1, ug)
ug.shape = (nlats,1)
ug = ug*np.ones((nlats,nlons),dtype=np.float) # broadcast to shape (nlats,nlonss)
# height perturbation.
hbump = hamp*np.cos(lats)*np.exp(-((lons-np.pi)/alpha)**2)*np.exp(-(phi2-lats)**2/beta)
# initial vorticity, divergence in spectral space
vrtspec, divspec = x.getvrtdivspec(ug,vg)
vrtg = x.spectogrd(vrtspec)
divg = x.spectogrd(divspec)
# create hyperdiffusion factor
hyperdiff_fact = np.exp((-dt/efold)*(x.lap/x.lap[-1])**(ndiss/2))
# solve nonlinear balance eqn to get initial zonal geopotential,
# add localized bump (not balanced).
vrtg = x.spectogrd(vrtspec)
tmpg1 = ug*(vrtg+f); tmpg2 = vg*(vrtg+f)
tmpspec1, tmpspec2 = x.getvrtdivspec(tmpg1,tmpg2)
tmpspec2 = x.grdtospec(0.5*(ug**2+vg**2))
phispec = x.invlap*tmpspec1 - tmpspec2
phig = grav*(hbar + hbump) + x.spectogrd(phispec)
phispec = x.grdtospec(phig)
# initialize spectral tendency arrays
ddivdtspec = np.zeros(vrtspec.shape+(3,), np.complex)
dvrtdtspec = np.zeros(vrtspec.shape+(3,), np.complex)
dphidtspec = np.zeros(vrtspec.shape+(3,), np.complex)
nnew = 0; nnow = 1; nold = 2
# time loop.
time1 = time.time()
for ncycle in range(itmax+1):
t = ncycle*dt
# get vort,u,v,phi on grid
vrtg = x.spectogrd(vrtspec)
ug,vg = x.getuv(vrtspec,divspec)
phig = x.spectogrd(phispec)
print('t=%6.2f hours: min/max %6.2f, %6.2f' % (t/3600.,vg.min(), vg.max()))
# compute tendencies.
tmpg1 = ug*(vrtg+f); tmpg2 = vg*(vrtg+f)
ddivdtspec[:,nnew], dvrtdtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2)
dvrtdtspec[:,nnew] *= -1
tmpg = x.spectogrd(ddivdtspec[:,nnew])
tmpg1 = ug*phig; tmpg2 = vg*phig
tmpspec, dphidtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2)
dphidtspec[:,nnew] *= -1
tmpspec = x.grdtospec(phig+0.5*(ug**2+vg**2))
ddivdtspec[:,nnew] += -x.lap*tmpspec
# update vort,div,phiv with third-order adams-bashforth.
# forward euler, then 2nd-order adams-bashforth time steps to start.
if ncycle == 0:
dvrtdtspec[:,nnow] = dvrtdtspec[:,nnew]
dvrtdtspec[:,nold] = dvrtdtspec[:,nnew]
ddivdtspec[:,nnow] = ddivdtspec[:,nnew]
ddivdtspec[:,nold] = ddivdtspec[:,nnew]
dphidtspec[:,nnow] = dphidtspec[:,nnew]
dphidtspec[:,nold] = dphidtspec[:,nnew]
elif ncycle == 1:
dvrtdtspec[:,nold] = dvrtdtspec[:,nnew]
ddivdtspec[:,nold] = ddivdtspec[:,nnew]
dphidtspec[:,nold] = dphidtspec[:,nnew]
vrtspec += dt*( \
(23./12.)*dvrtdtspec[:,nnew] - (16./12.)*dvrtdtspec[:,nnow]+ \
(5./12.)*dvrtdtspec[:,nold] )
divspec += dt*( \
(23./12.)*ddivdtspec[:,nnew] - (16./12.)*ddivdtspec[:,nnow]+ \
(5./12.)*ddivdtspec[:,nold] )
phispec += dt*( \
(23./12.)*dphidtspec[:,nnew] - (16./12.)*dphidtspec[:,nnow]+ \
(5./12.)*dphidtspec[:,nold] )
# implicit hyperdiffusion for vort and div.
vrtspec *= hyperdiff_fact
divspec *= hyperdiff_fact
# switch indices, do next time step.
nsav1 = nnew; nsav2 = nnow
nnew = nold; nnow = nsav1; nold = nsav2
time2 = time.time()
print('CPU time = ',time2-time1)
# make a contour plot of potential vorticity in the Northern Hem.
fig = plt.figure(figsize=(12,4))
# dimensionless PV
pvg = (0.5*hbar*grav/omega)*(vrtg+f)/phig
print('max/min PV',pvg.min(), pvg.max())
lons1d = (180./np.pi)*x.lons-180.; lats1d = (180./np.pi)*x.lats
levs = np.arange(-0.2,1.801,0.1)
cs=plt.contourf(lons1d,lats1d,pvg,levs,extend='both')
cb=plt.colorbar(cs,orientation='horizontal') # add colorbar
cb.set_label('potential vorticity')
plt.grid()
plt.xlabel('degrees longitude')
plt.ylabel('degrees latitude')
plt.xticks(np.arange(-180,181,60))
plt.yticks(np.arange(-5,81,20))
plt.axis('equal')
plt.axis('tight')
plt.ylim(0,lats1d[0])
plt.title('PV (T%s with hyperdiffusion, hour %6.2f)' % (ntrunc,t/3600.))
plt.savefig("output_swe.pdf")
plt.show()
class TransformsEngine(object):
"""A spectral transforms engine based on pyspharm."""
def __init__(self, nlon, nlat, truncation, radius=6371200.):
"""
Initialize the spectral transforms engine.
Arguments:
* nlon: int
Number of longitudes in the transform grid.
* nlat: int
Number of latitudes in the transform grid.
* truncation: int
The spectral truncation (triangular). This is the maximum
number of spherical harmonic modes retained in the discrete
truncation. More modes means higher resolution.
"""
self.sh = Spharmt(nlon, nlat, gridtype='regular', rsphere=radius)
self.radius = radius
self.nlon = nlon
self.nlat = nlat
self.truncation = truncation
def vrtdiv_spec_from_uv_grid(self, u, v):
"""
Compute spectral vorticity and divergence from grid u and v.
"""
try:
vrt, div = self.sh.getvrtdivspec(u, v, ntrunc=self.truncation)
except ValueError:
msg = ('u and v must be 2d or 3d arrays with shape ({y}, {x}) '
'or ({y}, {x}, :)'.format(y=self.nlat, x=self.nlon))
raise ValueError(msg)
return vrt, div
def uv_grid_from_vrtdiv_spec(self, vrt, div):
"""
Compute grid u and v from spectral vorticity and divergence.
"""
try:
u, v = self.sh.getuv(vrt, div)
except ValueError:
nspec = (self.truncation + 1) * (self.truncation + 2) // 2
msg = ('vrt and div must be 1d or 2d arrays with shape '
'(n) or (n, :) where n <= {}'.format(nspec))
raise ValueError(msg)
return u, v
def spec_to_grid(self, scalar_spec):
"""
Transform a scalar field from spectral to grid space.
"""
try:
scalar_grid = self.sh.spectogrd(scalar_spec)
except ValueError:
nspec = (self.truncation + 1) * (self.truncation + 2) // 2
msg = ('scalar_spec must be a 1d or 2d array with shape '
'(n) or (n, :) where n <= {}'.format(nspec))
raise ValueError(msg)
return scalar_grid
def grid_to_spec(self, scalar_grid):
"""
Transform a scalar field from grid to spectral space.
"""
try:
scalar_spec = self.sh.grdtospec(scalar_grid,
ntrunc=self.truncation)
except ValueError:
msg = ('scalar_grid must be a 2d or 3d array with shape '
'({y}, {x}) or ({y}, {x}, :)'.format(y=self.nlat,
x=self.nlon))
raise ValueError(msg)
return scalar_spec
def grad_of_spec(self, scalar_spec):
"""
Return zonal and meridional gradients of a spectral field.
"""
try:
dsdx, dsdy = self.sh.getgrad(scalar_spec)
except ValueError:
nspec = (self.truncation + 1) * (self.truncation + 2) // 2
msg = ('scalar_spec must be a 1d or 2d array with shape '
'(n) or (n, :) where n <= {}'.format(nspec))
raise ValueError(msg)
return dsdx, dsdy
@property
def wavenumbers(self):
"""
Wavenumbers corresponding to the spectral fields.
"""
return getspecindx(self.truncation)
@property
def grid_latlon(self):
"""
Return the latitude and longitude coordinate vectors of the
model grid.
"""
lats, _ = gaussian_lats_wts(self.nlat)
lons = np.arange(0., 360., 360. / self.nlon)
return lats, lons
class spectral:
def __init__(self, lat, lon, rsphere=6.3712e6, legfunc='stored'):
# Length of lat/lon arrays
self.nlat = len(lat)
self.nlon = len(lon)
if self.nlat % 2:
gridtype = 'gaussian'
else:
gridtype = 'regular'
self.s = Spharmt(self.nlon,
self.nlat,
gridtype=gridtype,
rsphere=rsphere,
legfunc=legfunc)
# Reverse latitude array if necessary
# self.ReverseLat = False
# if lat[0] < lat[-1]:
# lat = self._reverse_lat(lat)
# self.ReverseLat = True
# lat/lon in degrees
self.glat = lat
self.glon = lon
# lat/lon in radians
self.rlat = np.deg2rad(lat)
self.rlon = np.deg2rad(lon)
self.rlons, self.rlats = np.meshgrid(self.rlon, self.rlat)
# Constants
# Earth's angular velocity
self.omega = 7.292e-05 # unit: s-1
# Gravitational acceleration
self.g = 9.8 # unit: m2/s
# Misc
self.dtype = np.float32
def _reverse_lat(self, var):
if len(np.shape(var)) == 1:
return (var[::-1])
if len(np.shape(var)) == 2:
return (var[::-1, :])
if len(np.shape(var)) == 3:
return (var[:, ::-1, :])
def planetaryvorticity(self, omega=None):
pass
def uv2vrt(self, u, v, trunc=None):
"""
Calculate relative vorticity from horizonal wind field
Input: u and v (grid)
Output: relative vorticity (grid)
"""
vrts, _ = self.s.getvrtdivspec(u, v, ntrunc=trunc)
vrtg = self.s.spectogrd(vrts)
return (vrtg)
def uv2div(self, u, v, trunc=None):
pass
def uv2sfvp(self, u, v, trunc=None):
"""
Calculate geostrophic streamfuncion and
velocity potential from u and v winds
"""
psig, chig = self.s.getpsichi(u, v, ntrunc=trunc)
return (psig, chig)
def uv2vrtdiv(self, u, v, trunc=None):
"""
Calculate relative vorticity and divergence
from u and v winds
"""
vrts, divs = self.s.getvrtdivspec(u, v, ntrunc=trunc)
vrtg = self.s.spectogrd(vrts)
divg = self.s.spectogrd(divs)
return (vrtg, divg)
def vrtdiv2uv(self, vrt, div, realm='grid', trunc=None):
"""
# Get u,v from vrt, div fields
# Input either in grid space
# or in spectral space
"""
if realm in ['g', 'grid']:
vrts = self.s.grdtospec(vrt, trunc)
divs = self.s.grdtospec(div, trunc)
elif realm in ['s', 'spec', 'spectral']:
vrts = vrt
divs = div
ug, vg = self.s.getuv(vrts, divs)
return (ug, vg)
def gradient(self, var, trunc=None):
"""
Calculate horizontal gradients
"""
# if self.ReverseLat is True:
# var = self._reverse_lat(var)
try:
var = var.filled(fill_value=np.nan)
except AttributeError:
pass
if np.isnan(var).any():
raise ValueError('var cannot contain missing values')
try:
varspec = self.s.grdtospec(var, ntrunc=trunc)
except ValueError:
raise ValueError('input field is not compatitble')
dxvarg, dyvarg = self.s.getgrad(varspec)
return (dxvarg, dyvarg)
phispec = x.grdtospec(phig,ntrunc)
# initialize spectral tendency arrays
ddivdtspec = np.zeros(vrtspec.shape+(3,), np.complex64)
dvrtdtspec = np.zeros(vrtspec.shape+(3,), np.complex64)
dphidtspec = np.zeros(vrtspec.shape+(3,), np.complex64)
nnew = 0; nnow = 1; nold = 2
# time loop.
time1 = time.clock()
for ncycle in range(itmax+1):
t = ncycle*dt
# get vort,u,v,phi on grid
vrtg = x.spectogrd(vrtspec)
ug,vg = x.getuv(vrtspec,divspec)
phig = x.spectogrd(phispec)
print('t=%6.2f hours: min/max %6.2f, %6.2f' % (t/3600.,vg.min(), vg.max()))
# compute tendencies.
tmpg1 = ug*(vrtg+f); tmpg2 = vg*(vrtg+f)
ddivdtspec[:,nnew], dvrtdtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2,ntrunc)
dvrtdtspec[:,nnew] *= -1
tmpg1 = ug*phig; tmpg2 = vg*phig
tmpspec, dphidtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2,ntrunc)
dphidtspec[:,nnew] *= -1
tmpspec = x.grdtospec(phig+0.5*(ug**2+vg**2),ntrunc)
ddivdtspec[:,nnew] += -lap*tmpspec
# update vort,div,phiv with third-order adams-bashforth.
# forward euler, then 2nd-order adams-bashforth time steps to start.
if ncycle == 0:
dvrtdtspec[:,nnow] = dvrtdtspec[:,nnew]