Solves the non-linear barotropically unstable shallow water test case in Python.
14 class Spharmt(object):
16 wrapper class for commonly used spectral transform operations in
17 atmospheric models. Provides an interface to shtns compatible
18 with pyspharm (pyspharm.googlecode.com).
20 def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
22 nlons: number of longitudes
23 nlats: number of latitudes"""
24 self._shtns = shtns.sht(ntrunc, ntrunc, 1, \
25 shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE)
26 if gridtype ==
'gaussian':
28 self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
29 elif gridtype ==
'regular':
30 self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
31 self.lats = np.arcsin(self._shtns.cos_theta)
32 self.lons = (2.*np.pi/nlons)*np.arange(nlons)
36 self.nlm = self._shtns.nlm
37 self.degree = self._shtns.l
38 self.lap = -self.degree*(self.degree+1.0).astype(np.complex)
39 self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
40 self.invlap[1:] = 1./self.lap[1:]
41 self.rsphere = rsphere
42 self.lap = self.lap/rsphere**2
43 self.invlap = self.invlap*rsphere**2
44 def grdtospec(self,data):
45 """compute spectral coefficients from gridded data"""
46 return self._shtns.analys(data)
47 def spectogrd(self,dataspec):
48 """compute gridded data from spectral coefficients"""
49 return self._shtns.synth(dataspec)
50 def getuv(self,vrtspec,divspec):
51 """compute wind vector from spectral coeffs of vorticity and divergence"""
52 return self._shtns.synth((self.invlap/self.rsphere)*vrtspec, (self.invlap/self.rsphere)*divspec)
53 def getvrtdivspec(self,u,v):
54 """compute spectral coeffs of vorticity and divergence from wind vector"""
55 vrtspec, divspec = self._shtns.analys(u, v)
56 return self.lap*self.rsphere*vrtspec, self.lap*rsphere*divspec
57 def getgrad(self,divspec):
58 """compute gradient vector from spectral coeffs"""
59 vrtspec = np.zeros(divspec.shape, dtype=np.complex)
60 u,v = self._shtns.synth(vrtspec,divspec)
61 return u/rsphere, v/rsphere
63 if __name__ ==
"__main__":
64 import matplotlib.pyplot
as plt
80 itmax = 6*int(86400/dt)
88 phi0 = np.pi/7.; phi1 = 0.5*np.pi - phi0; phi2 = 0.25*np.pi
89 en = np.exp(-4.0/(phi1-phi0)**2)
90 alpha = 1./3.; beta = 1./15.
96 x = Spharmt(nlons,nlats,ntrunc,rsphere,gridtype=
'gaussian')
97 lons,lats = np.meshgrid(x.lons, x.lats)
98 f = 2.*omega*np.sin(lats)
101 vg = np.zeros((nlats,nlons),np.float)
102 u1 = (umax/en)*np.exp(1./((x.lats-phi0)*(x.lats-phi1)))
103 ug = np.zeros((nlats),np.float)
104 ug = np.where(np.logical_and(x.lats < phi1, x.lats > phi0), u1, ug)
106 ug = ug*np.ones((nlats,nlons),dtype=np.float)
108 hbump = hamp*np.cos(lats)*np.exp(-(lons/alpha)**2)*np.exp(-(phi2-lats)**2/beta)
111 vrtspec, divspec = x.getvrtdivspec(ug,vg)
112 vrtg = x.spectogrd(vrtspec)
113 divg = x.spectogrd(divspec)
116 hyperdiff_fact = np.exp((-dt/efold)*(x.lap/x.lap[-1])**(ndiss/2))
120 vrtg = x.spectogrd(vrtspec)
121 tmpg1 = ug*(vrtg+f); tmpg2 = vg*(vrtg+f)
122 tmpspec1, tmpspec2 = x.getvrtdivspec(tmpg1,tmpg2)
123 tmpspec2 = x.grdtospec(0.5*(ug**2+vg**2))
124 phispec = x.invlap*tmpspec1 - tmpspec2
125 phig = grav*(hbar + hbump) + x.spectogrd(phispec)
126 phispec = x.grdtospec(phig)
129 ddivdtspec = np.zeros(vrtspec.shape+(3,), np.complex)
130 dvrtdtspec = np.zeros(vrtspec.shape+(3,), np.complex)
131 dphidtspec = np.zeros(vrtspec.shape+(3,), np.complex)
132 nnew = 0; nnow = 1; nold = 2
136 for ncycle
in range(itmax+1):
139 vrtg = x.spectogrd(vrtspec)
140 ug,vg = x.getuv(vrtspec,divspec)
141 phig = x.spectogrd(phispec)
142 print(
't=%6.2f hours: min/max %6.2f, %6.2f' % (t/3600.,vg.min(), vg.max()))
144 tmpg1 = ug*(vrtg+f); tmpg2 = vg*(vrtg+f)
145 ddivdtspec[:,nnew], dvrtdtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2)
146 dvrtdtspec[:,nnew] *= -1
147 tmpg = x.spectogrd(ddivdtspec[:,nnew])
148 tmpg1 = ug*phig; tmpg2 = vg*phig
149 tmpspec, dphidtspec[:,nnew] = x.getvrtdivspec(tmpg1,tmpg2)
150 dphidtspec[:,nnew] *= -1
151 tmpspec = x.grdtospec(phig+0.5*(ug**2+vg**2))
152 ddivdtspec[:,nnew] += -x.lap*tmpspec
156 dvrtdtspec[:,nnow] = dvrtdtspec[:,nnew]
157 dvrtdtspec[:,nold] = dvrtdtspec[:,nnew]
158 ddivdtspec[:,nnow] = ddivdtspec[:,nnew]
159 ddivdtspec[:,nold] = ddivdtspec[:,nnew]
160 dphidtspec[:,nnow] = dphidtspec[:,nnew]
161 dphidtspec[:,nold] = dphidtspec[:,nnew]
163 dvrtdtspec[:,nold] = dvrtdtspec[:,nnew]
164 ddivdtspec[:,nold] = ddivdtspec[:,nnew]
165 dphidtspec[:,nold] = dphidtspec[:,nnew]
167 (23./12.)*dvrtdtspec[:,nnew] - (16./12.)*dvrtdtspec[:,nnow]+ \
168 (5./12.)*dvrtdtspec[:,nold] )
170 (23./12.)*ddivdtspec[:,nnew] - (16./12.)*ddivdtspec[:,nnow]+ \
171 (5./12.)*ddivdtspec[:,nold] )
173 (23./12.)*dphidtspec[:,nnew] - (16./12.)*dphidtspec[:,nnow]+ \
174 (5./12.)*dphidtspec[:,nold] )
176 vrtspec *= hyperdiff_fact
177 divspec *= hyperdiff_fact
179 nsav1 = nnew; nsav2 = nnow
180 nnew = nold; nnow = nsav1; nold = nsav2
183 print(
'CPU time = ',time2-time1)
186 fig = plt.figure(figsize=(12,4))
188 pvg = (0.5*hbar*grav/omega)*(vrtg+f)/phig
189 print(
'max/min PV',pvg.min(), pvg.max())
190 lons1d = (180./np.pi)*x.lons-180.; lats1d = (180./np.pi)*x.lats
191 levs = np.arange(-0.2,1.801,0.1)
192 cs=plt.contourf(lons1d,lats1d,pvg,levs,cmap=plt.cm.spectral,extend=
'both')
193 cb=plt.colorbar(cs,orientation=
'horizontal')
194 cb.set_label(
'potential vorticity')
196 plt.xlabel(
'degrees longitude')
197 plt.ylabel(
'degrees latitude')
198 plt.xticks(np.arange(-180,181,60))
199 plt.yticks(np.arange(-5,81,20))
202 plt.ylim(0,lats1d[0])
203 plt.title(
'PV (T%s with hyperdiffusion, hour %6.2f)' % (ntrunc,t/3600.))