Created
October 6, 2012 15:58
-
-
Save jswhit/3845307 to your computer and use it in GitHub Desktop.
shallow water model python test for shtns
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import shtns | |
| import numpy as np | |
| class Spharmt(object): | |
| """ | |
| wrapper class for commonly used spectral transform operations in | |
| atmospheric models. Provides an interface to shtns compatible | |
| with pyspharm (pyspharm.googlecode.com). | |
| """ | |
| def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'): | |
| """initialize | |
| nlons: number of longitudes | |
| nlats: number of latitudes""" | |
| self._shtns = shtns.sht(ntrunc, ntrunc, 1, \ | |
| shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE) | |
| if gridtype == 'gaussian': | |
| #self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS,1.e-10) | |
| self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-10) | |
| elif gridtype == 'regular': | |
| self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-10) | |
| self.lats = np.arcsin(self._shtns.cos_theta) | |
| self.lons = (2.*np.pi/nlons)*np.arange(nlons) | |
| self.nlons = nlons | |
| self.nlats = nlats | |
| self.ntrunc = ntrunc | |
| self.nlm = self._shtns.nlm | |
| self.degree = self._shtns.l | |
| self.lap = -self.degree*(self.degree+1.0).astype(np.complex) | |
| self.invlap = np.zeros(self.lap.shape, self.lap.dtype) | |
| self.invlap[1:] = 1./self.lap[1:] | |
| self.rsphere = rsphere | |
| self.lap = self.lap/rsphere**2 | |
| self.invlap = self.invlap*rsphere**2 | |
| def grdtospec(self,data): | |
| """compute spectral coefficients from gridded data""" | |
| data = np.ascontiguousarray(data, dtype=np.float) | |
| dataspec = np.empty(self.nlm, dtype=np.complex) | |
| self._shtns.spat_to_SH(data, dataspec) | |
| return dataspec | |
| def spectogrd(self,dataspec): | |
| """compute gridded data from spectral coefficients""" | |
| dataspec = np.ascontiguousarray(dataspec, dtype=np.complex) | |
| data = np.empty((self.nlats,self.nlons), dtype=np.float) | |
| self._shtns.SH_to_spat(dataspec, data) | |
| return data | |
| def getuv(self,vrtspec,divspec): | |
| """compute wind vector from spectral coeffs of vorticity and divergence""" | |
| vrtspec = np.ascontiguousarray(vrtspec, dtype=np.complex) | |
| divspec = np.ascontiguousarray(divspec, dtype=np.complex) | |
| u = np.empty((self.nlats,self.nlons), dtype=np.float) | |
| v = np.empty((self.nlats,self.nlons), dtype=np.float) | |
| self._shtns.SHsphtor_to_spat((self.invlap/self.rsphere)*vrtspec,\ | |
| (self.invlap/self.rsphere)*divspec, u, v) | |
| return u,v | |
| def getvrtdivspec(self,u,v): | |
| """compute spectral coeffs of vorticity and divergence from wind vector""" | |
| u = np.ascontiguousarray(u, dtype=np.float) | |
| v = np.ascontiguousarray(v, dtype=np.float) | |
| vrtspec = np.empty(self.nlm, dtype=np.complex) | |
| divspec = np.empty(self.nlm, dtype=np.complex) | |
| self._shtns.spat_to_SHsphtor(u, v, vrtspec, divspec) | |
| return self.lap*self.rsphere*vrtspec, self.lap*rsphere*divspec | |
| def getgrad(self,divspec): | |
| """compute gradient vector from spectral coeffs""" | |
| divspec = np.ascontiguousarray(divspec, dtype=np.complex) | |
| vrtspec = np.zeros(divspec.shape, dtype=np.complex) | |
| u = np.empty((self.nlats,self.nlons), dtype=np.float) | |
| v = np.empty((self.nlats,self.nlons), dtype=np.float) | |
| self._shtns.SHsphtor_to_spat(vrtspec,divspec) | |
| return u/rsphere,v/rsphere | |
| if __name__ == "__main__": | |
| import matplotlib.pyplot as plt | |
| import time | |
| # 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 = nlons/3 # spectral truncation (for alias-free computations) | |
| nlats = nlons/2 # for gaussian grid. | |
| dt = 150 # time step in seconds | |
| itmax = 6*(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/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.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) | |
| 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.clock() | |
| 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,cmap=plt.cm.spectral,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.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Running the script should pop up a window with this image
http://i.imgur.com/ZlxR1.png