Пример #1
0
# 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)
Пример #2
0
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')
Пример #3
0
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()
Пример #4
0
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
Пример #5
0
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)
Пример #6
0
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]