예제 #1
0
 def __init__(self, nlon, nlat, rsphere=1e2, gridtype="regular", legfunc="stored"):
     Spharmt.__init__(self, nlon, nlat, rsphere, gridtype, legfunc)
     if self.gridtype == "gaussian":
         raise Exception, "Can't deal with Gaussian grid at the moment."
     else:
         deltalat = 180.0 / (nlat - 1)
         self.lats = 90 - deltalat * np.arange(nlat)
     deltalon = 360.0 / nlon
     self.lons = deltalon * np.arange(nlon)
     return
예제 #2
0
 def __init__(self,nlon,nlat,rsphere=1e2,gridtype='regular',legfunc='stored'):
     Spharmt.__init__(self,nlon,nlat,rsphere,gridtype,legfunc)
     if self.gridtype=='gaussian':
         raise Exception, "Can't deal with Gaussian grid at the moment."
     else:
         deltalat = 180./(nlat-1)
         self.lats = 90 - deltalat*np.arange(nlat)
     deltalon = 360./nlon
     self.lons = deltalon*np.arange(nlon)
     return
예제 #3
0
def regrid_sphere(nlat,nlon,Nens,X):

    """
    Truncate lat,lon grid to another resolution in spherical harmonic space. Triangular truncation

    Inputs:
    nlat            : number of latitudes
    nlon            : number of longitudes
    Nens            : number of ensemble members
    X               : data array of shape (nlat*nlon,Nens)
    ntrunc          : triangular truncation (e.g., use 42 for T42)

    Outputs :
    lat_new : 2D latitude array on the new grid (nlat_new,nlon_new)
    lon_new : 2D longitude array on the new grid (nlat_new,nlon_new)
    X_new   : truncated data array of shape (nlat_new*nlon_new, Nens)
    """
    # Originator: Greg Hakim
    #             University of Washington
    #             May 2015

    # create the spectral object on the original grid
    specob_lmr = Spharmt(nlon,nlat,gridtype='regular',legfunc='computed')

    # truncate to a lower resolution grid (triangular truncation)
    # ifix = np.remainder(ntrunc,2.0).astype(int)
    # nlat_new = ntrunc + ifix
    # nlon_new = int(nlat_new*1.5)
    ntrunc = 42
    nlat_new = 64
    nlon_new = 128

    # create the spectral object on the new grid
    specob_new = Spharmt(nlon_new,nlat_new,gridtype='regular',legfunc='computed')

    # create new lat,lon grid arrays
    dlat = 90./((nlat_new-1)/2.)
    dlon = 360./nlon_new
    veclat = np.arange(-90.,90.+dlat,dlat)
    veclon = np.arange(0.,360.,dlon)
    blank = np.zeros([nlat_new,nlon_new])
    lat_new = (veclat + blank.T).T
    lon_new = (veclon + blank)

    # transform each ensemble member, one at a time
    X_new = np.zeros([nlat_new*nlon_new,Nens])
    for k in range(Nens):
        X_lalo = np.reshape(X[:,k],(nlat,nlon))
        Xbtrunc = regrid(specob_lmr, specob_new, X_lalo, ntrunc=nlat_new-1, smooth=None)
        vectmp = Xbtrunc.flatten()
        X_new[:,k] = vectmp

    return X_new,lat_new,lon_new
예제 #4
0
def regrid_field(field, lat, lon, lat_new, lon_new):
    nlat_old, nlon_old = np.size(lat), np.size(lon)
    nlat_new, nlon_new = np.size(lat_new), np.size(lon_new)
    spec_old = Spharmt(nlon_old, nlat_old, gridtype='regular', legfunc='computed')
    spec_new = Spharmt(nlon_new, nlat_new, gridtype='regular', legfunc='computed')
    #remove nans
    field[np.isnan(field)] = 0
    field_new = []
    for field_old in field:
        regridded_field =  regrid(spec_old, spec_new, field_old, ntrunc=None, smooth=None)
        field_new.append(regridded_field)

    field_new = np.array(field_new)
    return field_new
예제 #5
0
파일: gridded.py 프로젝트: fzhu2e/LMRt
    def regrid(self, ntrunc, inplace=False):
        old_spec = Spharmt(self.nlon,
                           self.nlat,
                           gridtype='regular',
                           legfunc='computed')
        ifix = ntrunc % 2
        new_nlat = ntrunc + ifix
        new_nlon = int(new_nlat * 1.5)
        new_spec = Spharmt(new_nlon,
                           new_nlat,
                           gridtype='regular',
                           legfunc='computed')
        include_poles = False if new_nlat % 2 == 0 else True
        new_lat_2d, new_lon_2d, _, _ = generate_latlon(
            new_nlat, new_nlon, include_endpts=include_poles)
        new_lat = new_lat_2d[:, 0]
        new_lon = new_lon_2d[0, :]

        # new_value = []
        # for old_value in self.value:
        old_value = np.moveaxis(self.value, 0, -1)
        regridded_value = regrid(old_spec,
                                 new_spec,
                                 old_value,
                                 ntrunc=new_nlat - 1,
                                 smooth=None)
        new_value = np.moveaxis(regridded_value, -1, 0)
        # new_value.append(regridded_value)

        new_value = np.array(new_value)

        if inplace:
            self.value = new_value
            self.lat = new_lat
            self.lon = new_lon
            self.nlat = np.size(new_lat)
            self.nlon = np.size(new_lon)
            self.ntrunc = ntrunc
        else:
            new_field = self.copy()
            new_field.value = new_value
            new_field.lat = new_lat
            new_field.lon = new_lon
            new_field.nlat = np.size(new_lat)
            new_field.nlon = np.size(new_lon)
            new_field.ntrunc = ntrunc
            return new_field
예제 #6
0
    def __init__(self,
                 lat,
                 lon,
                 rsphere=6.3712e6,
                 legfunc='stored',
                 trunc=None):

        # 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
예제 #7
0
 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
예제 #8
0
    def __init__(self, H, W):
        super(PowerSpectrum, self).__init__()
        self.spharm = Spharmt(W, H, legfunc='stored')
        self.spectrums = []

        self.needs_sh = True
예제 #9
0
wavenum = 4.0
omega = 7.848e-6
re = 6.371e6
delat = 2.*math.pi/nlons_reg
lats = (0.5*math.pi-delat*numpy.indices(lats_reg.shape)[0,:,:])
lons = (delat*numpy.indices(lons_reg.shape)[1,:,:])
psi_reg_exact = rhwave(wavenum,omega,re,lats,lons)
delat = 2.*math.pi/nlons_gau
lats = (math.pi/180.)*numpy.transpose(gaulats*numpy.ones((nlons_gau,nlats_gau),'d'))
lons = (delat*numpy.indices(lons_gau.shape)[1,:,:])
psi_gau = rhwave(wavenum,omega,re,lats,lons)

# create Spharmt instances for regular and gaussian grids.

reggrid = Spharmt(nlons_reg,nlats_reg,gridtype='regular')
gaugrid = Spharmt(nlons_gau,nlats_gau,gridtype='gaussian')

# regrid from gaussian to regular grid.

psi_reg = regrid(gaugrid,reggrid,psi_gau)

print('reggrid error (should be less than 1.e-6):')
print(numpy.fabs(psi_reg-psi_reg_exact).max()/numpy.fabs(psi_reg_exact).max())

# spectrally interpolate to geodesic grid.

ntrunc = nlats_reg-1
latpts,lonpts = getgeodesicpts(7) # compute geodesic points
dataspec = reggrid.grdtospec(psi_reg_exact,ntrunc) # compute spectral coeffs
nlat = 0
예제 #10
0
class VectorWind(object):
    """Vector Wind computations (standard `numpy` interface)."""

    def __init__(self, u, v, gridtype='regular', rsphere=6.3712e6):
        """Initialize a VectorWind instance.

        **Arguments:**

        *u*, *v*
            Zonal and meridional wind components respectively. Their
            types should be either `numpy.ndarray` or
            `numpy.ma.MaskedArray`. *u* and *v* must have matching
            shapes and contain no missing values. *u* and *v* may be 2
            or 3-dimensional with shape (nlat, nlon) or
            (nlat, nlon, nt), where nlat and nlon are the number of
            latitudes and longitudes respectively and nt is the number
            of fields. The latitude dimension must be oriented
            north-to-south. The longitude dimension should be
            oriented west-to-east.

        **Optional argument:**

        *gridtype*
            Type of the input grid, either 'regular' for evenly-spaced
            grids, or 'gaussian' for Gaussian grids. Defaults to
            'regular'.

        **See also:**

        `~windspharm.tools.prep_data`,
        `~windspharm.tools.recover_data`,
        `~windspharm.tools.get_recovery`,
        `~windspharm.tools.reverse_latdim`,
        `~windspharm.tools.order_latdim`.

        **Examples:**

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind on the default regular
        (evenly-spaced) grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v)

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind specified on a Gaussian grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v, gridtype='gaussian')

        """
        # For both the input components check if there are missing values by
        # attempting to fill missing values with NaN and detect them. If the
        # inputs are not masked arrays then take copies and check for NaN.
        try:
            self.u = u.filled(fill_value=np.nan)
        except AttributeError:
            self.u = u.copy()
        try:
            self.v = v.filled(fill_value=np.nan)
        except AttributeError:
            self.v = v.copy()
        if np.isnan(self.u).any() or np.isnan(self.v).any():
            raise ValueError('u and v cannot contain missing values')
        # Make sure the shapes of the two components match.
        if u.shape != v.shape:
            raise ValueError('u and v must be the same shape')
        if len(u.shape) not in (2, 3):
            raise ValueError('u and v must be rank 2 or 3 arrays')
        nlat = u.shape[0]
        nlon = u.shape[1]
        try:
            # Create a Spharmt object to do the computations.
            self.gridtype = gridtype.lower()
            self.s = Spharmt(nlon, nlat, gridtype=self.gridtype,
                             rsphere=rsphere)
        except ValueError:
            if self.gridtype not in ('regular', 'gaussian'):
                err = 'invalid grid type: {0:s}'.format(repr(gridtype))
            else:
                err = 'invalid input dimensions'
            raise ValueError(err)
        # Method aliases.
        self.rotationalcomponent = self.nondivergentcomponent
        self.divergentcomponent = self.irrotationalcomponent

    def magnitude(self):
        """Wind speed (magnitude of vector wind).

        **Returns:**

        *speed*
            The wind speed.

        **Example:**

        Magnitude of the vector wind::

            spd = w.magnitude()

        """
        return (self.u ** 2 + self.v ** 2) ** 0.5

    def vrtdiv(self, truncation=None):
        """Relative vorticity and horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vrt*, *div*
            The relative vorticity and divergence respectively.

        **See also:**

        `~VectorWind.vorticity`, `~VectorWind.divergence`.

        **Examples:**

        Compute the relative vorticity and divergence::

            vrt, div = w.vrtdiv()

        Compute the relative vorticity and divergence and apply spectral
        truncation at triangular T13::

            vrtT13, divT13 = w.vrtdiv(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        divgrid = self.s.spectogrd(divspec)
        return vrtgrid, divgrid

    def vorticity(self, truncation=None):
        """Relative vorticity.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vrt*
            The relative vorticity.

        **See also:**

        `~VectorWind.vrtdiv`, `~VectorWind.absolutevorticity`.

        **Examples:**

        Compute the relative vorticity::

            vrt = w.vorticity()

        Compute the relative vorticity and apply spectral truncation at
        triangular T13::

            vrtT13 = w.vorticity(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        return vrtgrid

    def divergence(self, truncation=None):
        """Horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *div*
            The divergence.

        **See also:**

        `~VectorWind.vrtdiv`.

        **Examples:**

        Compute the divergence::

            div = w.divergence()

        Compute the divergence and apply spectral truncation at
        triangular T13::

            divT13 = w.divergence(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        divgrid = self.s.spectogrd(divspec)
        return divgrid

    def planetaryvorticity(self, omega=None):
        """Planetary vorticity (Coriolis parameter).

        **Optional argument:**

        *omega*
            Earth's angular velocity. The default value if not specified
            is 7.292x10**-5 s**-1.

        **Returns:**

        *pvorticity*
            The planetary vorticity.

        **See also:**

        `~VectorWind.absolutevorticity`.

        **Example:**

        Compute planetary vorticity using default values::

            pvrt = w.planetaryvorticity()

        Override the default value for Earth's angular velocity::

            pvrt = w.planetaryvorticity(omega=7.2921150)

        """
        if omega is None:
            # Define the Earth's angular velocity.
            omega = 7.292e-05
        nlat = self.s.nlat
        if self.gridtype == 'gaussian':
            lat, wts = gaussian_lats_wts(nlat)
        else:
            if nlat % 2:
                lat = np.linspace(90, -90, nlat)
            else:
                dlat = 180. / nlat
                lat = np.arange(90 - dlat / 2., -90, -dlat)
        try:
            cp = 2. * omega * np.sin(np.deg2rad(lat))
        except (TypeError, ValueError):
            raise ValueError('invalid value for omega: {!r}'.format(omega))
        indices = [slice(0, None)] + [np.newaxis] * (len(self.u.shape) - 1)
        f = cp[indices] * np.ones(self.u.shape, dtype=np.float32)
        return f

    def absolutevorticity(self, omega=None, truncation=None):
        """Absolute vorticity (sum of relative and planetary vorticity).

        **Optional arguments:**

        *omega*
            Earth's angular velocity. The default value if not specified
            is 7.292x10**-5 s**-1.

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *avorticity*
            The absolute (relative + planetary) vorticity.

        **See also:**

        `~VectorWind.vorticity`, `~VectorWind.planetaryvorticity`.

        **Examples:**

        Compute absolute vorticity::

            avrt = w.absolutevorticity()

        Compute absolute vorticity and apply spectral truncation at
        triangular T13, also override the default value for Earth's
        angular velocity::

            avrt = w.absolutevorticity(omega=7.2921150, truncation=13)

        """
        pvrt = self.planetaryvorticity(omega=omega)
        rvrt = self.vorticity(truncation=truncation)
        return pvrt + rvrt

    def sfvp(self, truncation=None):
        """Streamfunction and velocity potential.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *sf*, *vp*
            The streamfunction and velocity potential respectively.

        **See also:**

        `~VectorWind.streamfunction`, `~VectorWind.velocitypotential`.

        **Examples:**

        Compute streamfunction and velocity potential::

            sf, vp = w.sfvp()

        Compute streamfunction and velocity potential and apply spectral
        truncation at triangular T13::

            sfT13, vpT13 = w.sfvp(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        return psigrid, chigrid

    def streamfunction(self, truncation=None):
        """Streamfunction.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *sf*
            The streamfunction.

        **See also:**

        `~VectorWind.sfvp`.

        **Examples:**

        Compute streamfunction::

            sf = w.streamfunction()

        Compute streamfunction and apply spectral truncation at
        triangular T13::

            sfT13 = w.streamfunction(truncation=13)

        """
        psigrid, chigrid = self.sfvp(truncation=truncation)
        return psigrid

    def velocitypotential(self, truncation=None):
        """Velocity potential.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vp*
            The velocity potential.

        **See also:**

        `~VectorWind.sfvp`.

        **Examples:**

        Compute velocity potential::

            vp = w.velocity potential()

        Compute velocity potential and apply spectral truncation at
        triangular T13::

            vpT13 = w.velocity potential(truncation=13)

        """
        psigrid, chigrid = self.sfvp(truncation=truncation)
        return chigrid

    def helmholtz(self, truncation=None):
        """Irrotational and non-divergent components of the vector wind.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*, *upsi*, *vpsi*
            Zonal and meridional components of irrotational and
            non-divergent wind components respectively.

        **See also:**

        `~VectorWind.irrotationalcomponent`,
        `~VectorWind.nondivergentcomponent`.

        **Examples:**

        Compute the irrotational and non-divergent components of the
        vector wind::

            uchi, vchi, upsi, vpsi = w.helmholtz()

        Compute the irrotational and non-divergent components of the
        vector wind and apply spectral truncation at triangular T13::

            uchiT13, vchiT13, upsiT13, vpsiT13 = w.helmholtz(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        psispec = self.s.grdtospec(psigrid)
        chispec = self.s.grdtospec(chigrid)
        vpsi, upsi = self.s.getgrad(psispec)
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi, -upsi, vpsi

    def irrotationalcomponent(self, truncation=None):
        """Irrotational (divergent) component of the vector wind.

        .. note::

           If both the irrotational and non-divergent components are
           required then `~VectorWind.helmholtz` should be used instead.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*
            The zonal and meridional components of the irrotational wind
            respectively.

        **See also:**

        `~VectorWind.helmholtz`.

        **Examples:**

        Compute the irrotational component of the vector wind::

            uchi, vchi = w.irrotationalcomponent()

        Compute the irrotational component of the vector wind and apply
        spectral truncation at triangular T13::

            uchiT13, vchiT13 = w.irrotationalcomponent(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        chispec = self.s.grdtospec(chigrid)
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi

    def nondivergentcomponent(self, truncation=None):
        """Non-divergent (rotational) component of the vector wind.

        .. note::

           If both the non-divergent and irrotational components are
           required then `~VectorWind.helmholtz` should be used instead.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *upsi*, *vpsi*
            The zonal and meridional components of the non-divergent
            wind respectively.

        **See also:**

        `~VectorWind.helmholtz`.

        **Examples:**

        Compute the non-divergent component of the vector wind::

            upsi, vpsi = w.nondivergentcomponent()

        Compute the non-divergent component of the vector wind and apply
        spectral truncation at triangular T13::

            upsiT13, vpsiT13 = w.nondivergentcomponent(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        psispec = self.s.grdtospec(psigrid)
        vpsi, upsi = self.s.getgrad(psispec)
        return -upsi, vpsi

    def gradient(self, chi, truncation=None):
        """Computes the vector gradient of a scalar field on the sphere.

        **Argument:**

        *chi*
            A scalar field. Its shape must be either (nlat, nlon) or
            (nlat, nlon, nfields) where nlat and nlon are the same
            as those for the vector wind components that initialized the
            `VectorWind` instance.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*
            The zonal and meridional components of the vector gradient
            respectively.

        **Examples:**

        Compute the vector gradient of absolute vorticity::

            avrt = w.absolutevorticity()
            avrt_zonal, avrt_meridional = w.gradient(avrt)

        Compute the vector gradient of absolute vorticity and apply
        spectral truncation at triangular T13::

            avrt = w.absolutevorticity()
            avrt_zonalT13, avrt_meridionalT13 = w.gradient(avrt, truncation=13)

        """
        try:
            chi = chi.filled(fill_value=np.nan)
        except AttributeError:
            pass
        if np.isnan(chi).any():
            raise ValueError('chi cannot contain missing values')
        try:
            chispec = self.s.grdtospec(chi, ntrunc=truncation)
        except ValueError:
            raise ValueError('input field is not compatitble')
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi

    def truncate(self, field, truncation=None):
        """Apply spectral truncation to a scalar field.

        This is useful to represent other fields in a way consistent
        with the output of other `VectorWind` methods.

        **Argument:**

        *field*
            A scalar field. Its shape must be either (nlat, nlon) or
            (nlat, nlon, nfields) where nlat and nlon are the same
            as those for the vector wind components that initialized the
            `VectorWind` instance.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation. If not specified it will default to
            *nlats - 1* where *nlats* is the number of latitudes.

        **Returns:**

        *truncated_field*
            The field with spectral truncation applied.

        **Examples:**

        Truncate a scalar field to the computational resolution of the
        `VectorWind` instance::

            scalar_field_truncated = w.truncate(scalar_field)

        Truncate a scalar field to T21::

            scalar_field_T21 = w.truncate(scalar_field, truncation=21)

        """
        try:
            field = field.filled(fill_value=np.nan)
        except AttributeError:
            pass
        if np.isnan(field).any():
            raise ValueError('field cannot contain missing values')
        try:
            fieldspec = self.s.grdtospec(field, ntrunc=truncation)
        except ValueError:
            raise ValueError('field is not compatible')
        fieldtrunc = self.s.spectogrd(fieldspec)
        return fieldtrunc
예제 #11
0
class VectorWind(object):
    """Vector Wind computations (standard `numpy` interface)."""
    def __init__(self, u, v, gridtype='regular', rsphere=6.3712e6):
        """Initialize a VectorWind instance.

        **Arguments:**

        *u*, *v*
            Zonal and meridional wind components respectively. Their
            types should be either `numpy.ndarray` or
            `numpy.ma.MaskedArray`. *u* and *v* must have matching
            shapes and contain no missing values. *u* and *v* may be 2
            or 3-dimensional with shape (nlat, nlon) or
            (nlat, nlon, nt), where nlat and nlon are the number of
            latitudes and longitudes respectively and nt is the number
            of fields. The latitude dimension must be oriented
            north-to-south. The longitude dimension should be
            oriented west-to-east.

        **Optional arguments:**

        *gridtype*
            Type of the input grid, either 'regular' for evenly-spaced
            grids, or 'gaussian' for Gaussian grids. Defaults to
            'regular'.

        *rsphere*
            The radius in metres of the sphere used in the spherical
            harmonic computations. Default is 6371200 m, the approximate
            mean spherical Earth radius.

        **See also:**

        `~windspharm.tools.prep_data`,
        `~windspharm.tools.recover_data`,
        `~windspharm.tools.get_recovery`,
        `~windspharm.tools.reverse_latdim`,
        `~windspharm.tools.order_latdim`.

        **Examples:**

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind on the default regular
        (evenly-spaced) grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v)

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind specified on a Gaussian grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v, gridtype='gaussian')

        """
        # For both the input components check if there are missing values by
        # attempting to fill missing values with NaN and detect them. If the
        # inputs are not masked arrays then take copies and check for NaN.
        try:
            self.u = u.filled(fill_value=np.nan)
        except AttributeError:
            self.u = u.copy()
        try:
            self.v = v.filled(fill_value=np.nan)
        except AttributeError:
            self.v = v.copy()
        if np.isnan(self.u).any() or np.isnan(self.v).any():
            raise ValueError('u and v cannot contain missing values')
        # Make sure the shapes of the two components match.
        if u.shape != v.shape:
            raise ValueError('u and v must be the same shape')
        if len(u.shape) not in (2, 3):
            raise ValueError('u and v must be rank 2 or 3 arrays')
        nlat = u.shape[0]
        nlon = u.shape[1]
        try:
            # Create a Spharmt object to do the computations.
            self.gridtype = gridtype.lower()
            self.s = Spharmt(nlon,
                             nlat,
                             gridtype=self.gridtype,
                             rsphere=rsphere)
        except ValueError:
            if self.gridtype not in ('regular', 'gaussian'):
                err = 'invalid grid type: {0:s}'.format(repr(gridtype))
            else:
                err = 'invalid input dimensions'
            raise ValueError(err)
        # Method aliases.
        self.rotationalcomponent = self.nondivergentcomponent
        self.divergentcomponent = self.irrotationalcomponent

    def magnitude(self):
        """Wind speed (magnitude of vector wind).

        **Returns:**

        *speed*
            The wind speed.

        **Example:**

        Magnitude of the vector wind::

            spd = w.magnitude()

        """
        return (self.u**2 + self.v**2)**0.5

    def vrtdiv(self, truncation=None):
        """Relative vorticity and horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vrt*, *div*
            The relative vorticity and divergence respectively.

        **See also:**

        `~VectorWind.vorticity`, `~VectorWind.divergence`.

        **Examples:**

        Compute the relative vorticity and divergence::

            vrt, div = w.vrtdiv()

        Compute the relative vorticity and divergence and apply spectral
        truncation at triangular T13::

            vrtT13, divT13 = w.vrtdiv(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        divgrid = self.s.spectogrd(divspec)
        return vrtgrid, divgrid

    def vorticity(self, truncation=None):
        """Relative vorticity.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vrt*
            The relative vorticity.

        **See also:**

        `~VectorWind.vrtdiv`, `~VectorWind.absolutevorticity`.

        **Examples:**

        Compute the relative vorticity::

            vrt = w.vorticity()

        Compute the relative vorticity and apply spectral truncation at
        triangular T13::

            vrtT13 = w.vorticity(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        return vrtgrid

    def divergence(self, truncation=None):
        """Horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *div*
            The divergence.

        **See also:**

        `~VectorWind.vrtdiv`.

        **Examples:**

        Compute the divergence::

            div = w.divergence()

        Compute the divergence and apply spectral truncation at
        triangular T13::

            divT13 = w.divergence(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        divgrid = self.s.spectogrd(divspec)
        return divgrid

    def planetaryvorticity(self, omega=None):
        """Planetary vorticity (Coriolis parameter).

        **Optional argument:**

        *omega*
            Earth's angular velocity. The default value if not specified
            is 7.292x10**-5 s**-1.

        **Returns:**

        *pvorticity*
            The planetary vorticity.

        **See also:**

        `~VectorWind.absolutevorticity`.

        **Example:**

        Compute planetary vorticity using default values::

            pvrt = w.planetaryvorticity()

        Override the default value for Earth's angular velocity::

            pvrt = w.planetaryvorticity(omega=7.2921150)

        """
        if omega is None:
            # Define the Earth's angular velocity.
            omega = 7.292e-05
        nlat = self.s.nlat
        if self.gridtype == 'gaussian':
            lat, wts = gaussian_lats_wts(nlat)
        else:
            if nlat % 2:
                lat = np.linspace(90, -90, nlat)
            else:
                dlat = 180. / nlat
                lat = np.arange(90 - dlat / 2., -90, -dlat)
        try:
            cp = 2. * omega * np.sin(np.deg2rad(lat))
        except (TypeError, ValueError):
            raise ValueError('invalid value for omega: {!r}'.format(omega))
        indices = [slice(0, None)] + [np.newaxis] * (len(self.u.shape) - 1)
        f = cp[indices] * np.ones(self.u.shape, dtype=np.float32)
        return f

    def absolutevorticity(self, omega=None, truncation=None):
        """Absolute vorticity (sum of relative and planetary vorticity).

        **Optional arguments:**

        *omega*
            Earth's angular velocity. The default value if not specified
            is 7.292x10**-5 s**-1.

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *avorticity*
            The absolute (relative + planetary) vorticity.

        **See also:**

        `~VectorWind.vorticity`, `~VectorWind.planetaryvorticity`.

        **Examples:**

        Compute absolute vorticity::

            avrt = w.absolutevorticity()

        Compute absolute vorticity and apply spectral truncation at
        triangular T13, also override the default value for Earth's
        angular velocity::

            avrt = w.absolutevorticity(omega=7.2921150, truncation=13)

        """
        pvrt = self.planetaryvorticity(omega=omega)
        rvrt = self.vorticity(truncation=truncation)
        return pvrt + rvrt

    def sfvp(self, truncation=None):
        """Streamfunction and velocity potential.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *sf*, *vp*
            The streamfunction and velocity potential respectively.

        **See also:**

        `~VectorWind.streamfunction`, `~VectorWind.velocitypotential`.

        **Examples:**

        Compute streamfunction and velocity potential::

            sf, vp = w.sfvp()

        Compute streamfunction and velocity potential and apply spectral
        truncation at triangular T13::

            sfT13, vpT13 = w.sfvp(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        return psigrid, chigrid

    def streamfunction(self, truncation=None):
        """Streamfunction.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *sf*
            The streamfunction.

        **See also:**

        `~VectorWind.sfvp`.

        **Examples:**

        Compute streamfunction::

            sf = w.streamfunction()

        Compute streamfunction and apply spectral truncation at
        triangular T13::

            sfT13 = w.streamfunction(truncation=13)

        """
        psigrid, chigrid = self.sfvp(truncation=truncation)
        return psigrid

    def velocitypotential(self, truncation=None):
        """Velocity potential.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *vp*
            The velocity potential.

        **See also:**

        `~VectorWind.sfvp`.

        **Examples:**

        Compute velocity potential::

            vp = w.velocity potential()

        Compute velocity potential and apply spectral truncation at
        triangular T13::

            vpT13 = w.velocity potential(truncation=13)

        """
        psigrid, chigrid = self.sfvp(truncation=truncation)
        return chigrid

    def helmholtz(self, truncation=None):
        """Irrotational and non-divergent components of the vector wind.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*, *upsi*, *vpsi*
            Zonal and meridional components of irrotational and
            non-divergent wind components respectively.

        **See also:**

        `~VectorWind.irrotationalcomponent`,
        `~VectorWind.nondivergentcomponent`.

        **Examples:**

        Compute the irrotational and non-divergent components of the
        vector wind::

            uchi, vchi, upsi, vpsi = w.helmholtz()

        Compute the irrotational and non-divergent components of the
        vector wind and apply spectral truncation at triangular T13::

            uchiT13, vchiT13, upsiT13, vpsiT13 = w.helmholtz(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        psispec = self.s.grdtospec(psigrid)
        chispec = self.s.grdtospec(chigrid)
        vpsi, upsi = self.s.getgrad(psispec)
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi, -upsi, vpsi

    def irrotationalcomponent(self, truncation=None):
        """Irrotational (divergent) component of the vector wind.

        .. note::

           If both the irrotational and non-divergent components are
           required then `~VectorWind.helmholtz` should be used instead.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*
            The zonal and meridional components of the irrotational wind
            respectively.

        **See also:**

        `~VectorWind.helmholtz`.

        **Examples:**

        Compute the irrotational component of the vector wind::

            uchi, vchi = w.irrotationalcomponent()

        Compute the irrotational component of the vector wind and apply
        spectral truncation at triangular T13::

            uchiT13, vchiT13 = w.irrotationalcomponent(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        chispec = self.s.grdtospec(chigrid)
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi

    def nondivergentcomponent(self, truncation=None):
        """Non-divergent (rotational) component of the vector wind.

        .. note::

           If both the non-divergent and irrotational components are
           required then `~VectorWind.helmholtz` should be used instead.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *upsi*, *vpsi*
            The zonal and meridional components of the non-divergent
            wind respectively.

        **See also:**

        `~VectorWind.helmholtz`.

        **Examples:**

        Compute the non-divergent component of the vector wind::

            upsi, vpsi = w.nondivergentcomponent()

        Compute the non-divergent component of the vector wind and apply
        spectral truncation at triangular T13::

            upsiT13, vpsiT13 = w.nondivergentcomponent(truncation=13)

        """
        psigrid, chigrid = self.s.getpsichi(self.u, self.v, ntrunc=truncation)
        psispec = self.s.grdtospec(psigrid)
        vpsi, upsi = self.s.getgrad(psispec)
        return -upsi, vpsi

    def gradient(self, chi, truncation=None):
        """Computes the vector gradient of a scalar field on the sphere.

        **Argument:**

        *chi*
            A scalar field. Its shape must be either (nlat, nlon) or
            (nlat, nlon, nfields) where nlat and nlon are the same
            as those for the vector wind components that initialized the
            `VectorWind` instance.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Returns:**

        *uchi*, *vchi*
            The zonal and meridional components of the vector gradient
            respectively.

        **Examples:**

        Compute the vector gradient of absolute vorticity::

            avrt = w.absolutevorticity()
            avrt_zonal, avrt_meridional = w.gradient(avrt)

        Compute the vector gradient of absolute vorticity and apply
        spectral truncation at triangular T13::

            avrt = w.absolutevorticity()
            avrt_zonalT13, avrt_meridionalT13 = w.gradient(avrt, truncation=13)

        """
        try:
            chi = chi.filled(fill_value=np.nan)
        except AttributeError:
            pass
        if np.isnan(chi).any():
            raise ValueError('chi cannot contain missing values')
        try:
            chispec = self.s.grdtospec(chi, ntrunc=truncation)
        except ValueError:
            raise ValueError('input field is not compatitble')
        uchi, vchi = self.s.getgrad(chispec)
        return uchi, vchi

    def truncate(self, field, truncation=None):
        """Apply spectral truncation to a scalar field.

        This is useful to represent other fields in a way consistent
        with the output of other `VectorWind` methods.

        **Argument:**

        *field*
            A scalar field. Its shape must be either (nlat, nlon) or
            (nlat, nlon, nfields) where nlat and nlon are the same
            as those for the vector wind components that initialized the
            `VectorWind` instance.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation. If not specified it will default to
            *nlats - 1* where *nlats* is the number of latitudes.

        **Returns:**

        *truncated_field*
            The field with spectral truncation applied.

        **Examples:**

        Truncate a scalar field to the computational resolution of the
        `VectorWind` instance::

            scalar_field_truncated = w.truncate(scalar_field)

        Truncate a scalar field to T21::

            scalar_field_T21 = w.truncate(scalar_field, truncation=21)

        """
        try:
            field = field.filled(fill_value=np.nan)
        except AttributeError:
            pass
        if np.isnan(field).any():
            raise ValueError('field cannot contain missing values')
        try:
            fieldspec = self.s.grdtospec(field, ntrunc=truncation)
        except ValueError:
            raise ValueError('field is not compatible')
        fieldtrunc = self.s.spectogrd(fieldspec)
        return fieldtrunc
예제 #12
0
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, rsphere=rsphere, gridtype=gridtype)
delta = 2. * np.pi / nlons
lats1d = 0.5 * np.pi - delta * np.arange(nlats)
lons1d = np.arange(-np.pi, np.pi, delta)
lons, lats = np.meshgrid(lons1d, lats1d)
f = 2. * omega * np.sin(lats)  # coriolis

# zonal jet.
vg = np.zeros((nlats, nlons), np.float32)
u1 = (umax / en) * np.exp(1. / ((lats1d - phi0) * (lats1d - phi1)))
ug = np.zeros((nlats), np.float32)
ug = np.where(np.logical_and(lats1d < phi1, lats1d > phi0), u1, ug)
ug.shape = (nlats, 1)
ug = ug * np.ones(
    (nlats, nlons), dtype=np.float32)  # broadcast to shape (nlats,nlonss)
# height perturbation.
예제 #13
0
    # ERA
    smatch, ematch = find_date_indices(ERA20C_time,stime,etime)
    ERA20C = ERA20C - np.mean(ERA20C[smatch:ematch,:,:],axis=0)


    # -----------------------------------
    # Regridding the data for comparisons
    # -----------------------------------
    print('\n regridding data to a common T42 grid...\n')

    iplot_loc= False
    #iplot_loc= True

    # create instance of the spherical harmonics object for each grid
    specob_lmr = Spharmt(nlon,nlat,gridtype='regular',legfunc='computed')
    specob_tcr = Spharmt(nlon_TCR,nlat_TCR,gridtype='regular',legfunc='computed')
    specob_era20c = Spharmt(nlon_ERA20C,nlat_ERA20C,gridtype='regular',legfunc='computed')

    # truncate to a lower resolution grid (common:21, 42, 62, 63, 85, 106, 255, 382, 799)
    ntrunc_new = 42 # T42
    ifix = np.remainder(ntrunc_new,2.0).astype(int)
    nlat_new = ntrunc_new + ifix
    nlon_new = int(nlat_new*1.5)
    # lat, lon grid in the truncated space
    dlat = 90./((nlat_new-1)/2.)
    dlon = 360./nlon_new
    veclat = np.arange(-90.,90.+dlat,dlat)
    veclon = np.arange(0.,360.,dlon)
    blank = np.zeros([nlat_new,nlon_new])
    lat2_new = (veclat + blank.T).T  
예제 #14
0
    # grid, time step info
    nlons = 128  # number of longitudes
    ntrunc = 42 # spectral truncation (for alias-free computations)
    nlats = (nlons/2)+1 # for regular grid.
    gridtype = 'regular'
    #nlats = nlons/2 # for gaussian grid.
    #gridtype = 'gaussian'
    dt = 900 # time step in seconds
    itmax = 8*(86400/dt) # integration length in days
    umax = 50. # jet speed
    jetexp = 10 # parameter controlling jet width

    # create spherical harmonic instance.
    rsphere = 6.37122e6 # earth radius
    sp = Spharmt(nlons,nlats,rsphere=rsphere,gridtype=gridtype)

    # create model instance using default parameters.
    model = TwoLayer(sp,dt,ntrunc)

    # vort, div initial conditions
    psipert = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    psipert[:,:,1] = 5.e6*np.sin((model.lons-np.pi))**12*np.sin(2.*model.lats)**12
    psipert = np.where(model.lons[:,:,np.newaxis] > 0., 0, psipert)
    ug = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    vg = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    ug[:,:,1] = umax*np.sin(2.*model.lats)**jetexp
    vrtspec, divspec = sp.getvrtdivspec(ug,vg,model.ntrunc)
    vrtspec = vrtspec + model.lap*sp.grdtospec(psipert,model.ntrunc)
    vrtg = sp.spectogrd(vrtspec)
    lyrthkspec = model.nlbalance(vrtspec)
예제 #15
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
예제 #16
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)
예제 #17
0
from spharm import Spharmt, getspecindx
import matplotlib.pyplot as plt
import numpy as np
# set up orthographic map projection.
map = Basemap(projection='ortho',lat_0=30,lon_0=-60,resolution='l')
# draw coastlines, country boundaries, fill continents.
map.drawcoastlines()
# draw the edge of the map projection region (the projection limb)
map.drawmapboundary()
# draw lat/lon grid lines every 30 degrees.
map.drawmeridians(np.arange(0,360,30))
map.drawparallels(np.arange(-90,90,30))
min = int(raw_input('input degree (m) of legendre function to plot:'))
nin = int(raw_input('input order  (n) of legendre function to plot:'))
nlons = 720; nlats = 361
x = Spharmt(nlons,nlats,legfunc='computed')
ntrunc = nlats-1
indxm, indxn = getspecindx(ntrunc)
nm = -1
i = 0
for m,n in zip(indxm,indxn):
    if m  == min and n == nin:
        nm = i
        exit
    else:
        i = i + 1
if nm < 0:
    raise ValueError('invalid m,n - must fit within triangular truncation at wavenumber '+repr(ntrunc))
coeffs = np.zeros((ntrunc+1)*(ntrunc+2)/2,np.complex)
coeffs[nm] = 1.
spharmonic = x.spectogrd(coeffs)
예제 #18
0
    # grid, time step info
    nlons = 128  # number of longitudes
    ntrunc = 42 # spectral truncation (for alias-free computations)
    nlats = (nlons/2)+1 # for regular grid.
    gridtype = 'regular'
    #nlats = nlons/2 # for gaussian grid.
    #gridtype = 'gaussian'
    dt = 1800 # time step in seconds
    itmax = int(5*(86400/dt)) # integration length in days
    umax = 50. # jet speed
    jetexp = 6 # parameter controlling jet width

    # create spherical harmonic instance.
    rsphere = 6.37122e6 # earth radius
    sp = Spharmt(nlons,nlats,rsphere=rsphere,gridtype=gridtype)

    # create model instance using default parameters.
    model = TwoLevel(sp,dt,ntrunc)

    # vort, div initial conditions
    psipert = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    psipert[:,:,1] = 5.e6*np.sin((model.lons-np.pi))**12*np.sin(2.*model.lats)**12
    psipert = np.where(model.lons[:,:,np.newaxis] > 0., 0, psipert)
    ug = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    vg = np.zeros((sp.nlat,sp.nlon,2),np.float32)
    ug[:,:,1] = umax*np.sin(2.*model.lats)**jetexp
    vrtspec, tmpspec = sp.getvrtdivspec(ug,vg,model.ntrunc)
    vrtspec = vrtspec + model.lap[:,np.newaxis]*sp.grdtospec(psipert,model.ntrunc)
    thetaspec = model.nlbalance(vrtspec)
    divspec = np.zeros(thetaspec.shape, thetaspec.dtype)
예제 #19
0
spec1 = None; spec2 = None; spec3 = None; spec4 = None
for date in dates:

    # get first guess for expt 1 (hybrid gain)
    datapath = '/scratch3/BMC/gsienkf/whitaker/%s' % expt1
    print datapath
    print date
    filename = os.path.join(os.path.join(datapath,date),'sfg_%s_fhr06_ensmean.nc4' % date)
    nc = Dataset(filename)
    if spec1 is None:
        lons = nc['lon'][:]
        lats = nc['lat'][::-1]
        nlats = len(lats); nlons = len(lons)
        lons2, lats2 = np.meshgrid(lons, lats)
        re = 6.3712e6; ntrunc=nlats-1
        spec = Spharmt(nlons,nlats,rsphere=re,gridtype='regular',legfunc='computed')
        indxm, indxn = getspecindx(ntrunc)
        degree = indxn.astype(np.float)
    if var == 'pressfc':
        fg1 = nc[var][0,::-1,...]
    else:
        fg1 = nc[var][0,nlev,::-1,...]
    nc.close()

    # get first guess for expt 2 (hybrid cov/envar)
    datapath = '/scratch3/BMC/gsienkf/whitaker/%s' % expt2
    print datapath
    print date
    filename = os.path.join(os.path.join(datapath,date),'sfg_%s_fhr06_ensmean.nc4' % date)
    nc = Dataset(filename)
    if var == 'pressfc':
예제 #20
0
# set model parameters.
nlons = 128  # number of longitudes
ntrunc = 42  # spectral truncation (for alias-free computations)
nlats = (nlons / 2) + 1  # for regular grid.
gridtype = 'regular'
dt = 900  # time step in seconds
tdiab = 12. * 86400  # thermal relaxation time scale
tdrag = 4. * 86400.  # lower layer drag
efold = 4 * dt  # hyperdiffusion time scale
rsphere = 6.37122e6  # earth radius
jetexp = 2
umax = 40
moistfact = 0.1

# create spherical harmonic instance.
sp = Spharmt(nlons, nlats, rsphere=rsphere, gridtype=gridtype)

# create model instance.
model =\
TwoLevel(sp,dt,ntrunc,efold=efold,tdiab=tdiab,tdrag=tdrag,jetexp=jetexp,umax=umax,moistfact=moistfact)

# initial state is equilbrium jet + random noise.
vg = np.zeros((sp.nlat, sp.nlon, 2), np.float32)
ug = model.uref
vrtspec, divspec = sp.getvrtdivspec(ug, vg, model.ntrunc)
psispec = np.zeros(vrtspec.shape, vrtspec.dtype)
psispec.real += npran.normal(scale=1.e4, size=(psispec.shape))
psispec.imag += npran.normal(scale=1.e4, size=(psispec.shape))
vrtspec = vrtspec + model.lap[:, np.newaxis] * psispec
thetaspec = model.nlbalance(vrtspec)
divspec = np.zeros(thetaspec.shape, thetaspec.dtype)
예제 #21
0
import matplotlib.pyplot as plt
import numpy as np
# set up orthographic map projection.
map = Basemap(projection='ortho', lat_0=30, lon_0=-60, resolution='l')
# draw coastlines, country boundaries, fill continents.
map.drawcoastlines()
# draw the edge of the map projection region (the projection limb)
map.drawmapboundary()
# draw lat/lon grid lines every 30 degrees.
map.drawmeridians(np.arange(0, 360, 30))
map.drawparallels(np.arange(-90, 90, 30))
min = int(raw_input('input degree (m) of legendre function to plot:'))
nin = int(raw_input('input order  (n) of legendre function to plot:'))
nlons = 720
nlats = 361
x = Spharmt(nlons, nlats, legfunc='computed')
ntrunc = nlats - 1
indxm, indxn = getspecindx(ntrunc)
nm = -1
i = 0
for m, n in zip(indxm, indxn):
    if m == min and n == nin:
        nm = i
        exit
    else:
        i = i + 1
if nm < 0:
    raise ValueError(
        'invalid m,n - must fit within triangular truncation at wavenumber ' +
        repr(ntrunc))
coeffs = np.zeros((ntrunc + 1) * (ntrunc + 2) / 2, np.complex)
   f = f - f.mean(axis=0)
   a = a - a.mean(axis=0)
   covfa = (f*a).mean(axis=0)
   varf = (f**2).mean(axis=0)
   vara = (a**2).mean(axis=0)
   return covfa/(np.sqrt(varf)*np.sqrt(vara))

def getmse(f,a):
   return ((f-a)**2).mean(axis=0)

nlonsin = 800; nlatsin = 400
nlons = 360; nlats = 181
ntrunc = 20
latbound = 20.

sin = Spharmt(nlonsin,nlatsin)
sout = Spharmt(nlons,nlats)

delta = 360./nlons
lats = 90.-delta*np.arange(nlats)
coslats = np.cos((np.pi/180.)*lats)
coslats = coslats[:,np.newaxis]*np.ones((nlats,nlons))
latnh = lats.tolist().index(latbound)
latsh = lats.tolist().index(-latbound)
coslatsnh = coslats[0:latnh+1,:]
coslatssh = coslats[latsh:,:]
coslatstr = coslats[latnh:latsh+1,:]
lons = delta*np.arange(nlons)
lons, lats = np.meshgrid(lons,lats[::-1])

filename = varshort
예제 #23
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()
예제 #24
0
    def __init__(self, u, v, gridtype='regular', rsphere=6.3712e6):
        """Initialize a VectorWind instance.

        **Arguments:**

        *u*, *v*
            Zonal and meridional wind components respectively. Their
            types should be either `numpy.ndarray` or
            `numpy.ma.MaskedArray`. *u* and *v* must have matching
            shapes and contain no missing values. *u* and *v* may be 2
            or 3-dimensional with shape (nlat, nlon) or
            (nlat, nlon, nt), where nlat and nlon are the number of
            latitudes and longitudes respectively and nt is the number
            of fields. The latitude dimension must be oriented
            north-to-south. The longitude dimension should be
            oriented west-to-east.

        **Optional arguments:**

        *gridtype*
            Type of the input grid, either 'regular' for evenly-spaced
            grids, or 'gaussian' for Gaussian grids. Defaults to
            'regular'.

        *rsphere*
            The radius in metres of the sphere used in the spherical
            harmonic computations. Default is 6371200 m, the approximate
            mean spherical Earth radius.

        **See also:**

        `~windspharm.tools.prep_data`,
        `~windspharm.tools.recover_data`,
        `~windspharm.tools.get_recovery`,
        `~windspharm.tools.reverse_latdim`,
        `~windspharm.tools.order_latdim`.

        **Examples:**

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind on the default regular
        (evenly-spaced) grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v)

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind specified on a Gaussian grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v, gridtype='gaussian')

        """
        # For both the input components check if there are missing values by
        # attempting to fill missing values with NaN and detect them. If the
        # inputs are not masked arrays then take copies and check for NaN.
        try:
            self.u = u.filled(fill_value=np.nan)
        except AttributeError:
            self.u = u.copy()
        try:
            self.v = v.filled(fill_value=np.nan)
        except AttributeError:
            self.v = v.copy()
        if np.isnan(self.u).any() or np.isnan(self.v).any():
            raise ValueError('u and v cannot contain missing values')
        # Make sure the shapes of the two components match.
        if u.shape != v.shape:
            raise ValueError('u and v must be the same shape')
        if len(u.shape) not in (2, 3):
            raise ValueError('u and v must be rank 2 or 3 arrays')
        nlat = u.shape[0]
        nlon = u.shape[1]
        try:
            # Create a Spharmt object to do the computations.
            self.gridtype = gridtype.lower()
            self.s = Spharmt(nlon,
                             nlat,
                             gridtype=self.gridtype,
                             rsphere=rsphere)
        except ValueError:
            if self.gridtype not in ('regular', 'gaussian'):
                err = 'invalid grid type: {0:s}'.format(repr(gridtype))
            else:
                err = 'invalid input dimensions'
            raise ValueError(err)
        # Method aliases.
        self.rotationalcomponent = self.nondivergentcomponent
        self.divergentcomponent = self.irrotationalcomponent
예제 #25
0
class VectorWind(object):
    """Vector Wind computations (standard :py:mod:`numpy` interface)."""

    def __init__(self, u, v, gridtype='regular'):
        """Initialize a VectorWind instance.

        **Arguments:**

        *u*, *v*
            Zonal and meridional wind components respectively. Their
            types should be either :py:class:`numpy.ndarray` or
            :py:class:`numpy.ma.MaskedArray`. *u* and *v* must have
            matching shapes and contain no missing values. *u* and *v*
            may be 2 or 3-dimensional with shape (nlat, nlon) or
            (nlat, nlon, nt), where nlat and nlon are the number of
            latitudes and longitudes respectively and nt is the number
            of fields. The latitude dimension must be oriented
            north-to-south. The longitude dimension should be
            oriented west-to-east.

        **Optional argument:**

        *gridtype*
            Type of the input grid, either 'regular' for evenly-spaced
            grids, or 'gaussian' for Gaussian grids. Defaults to
            'regular'.

        **Examples:**

        Initialize a VectorWind instance with zonal and meridional
        components of the vector wind on the default regular
        (evenly-spaced) grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v)

        Initialize a VectorWind instance with zonal and meridional
        components of the vector wind specified on a Gaussian grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v, gridtype='gaussian')

        """
        # For both the input components check if there are missing values by
        # attempting to fill missing values and detect them. If the inputs are
        # not masked arrays then this check isn't needed so take a copy.
        try:
            self.u = u.filled()
            if (self.u == u.fill_value).any():
                raise ValueError('u and v cannot contain missing values')
        except AttributeError:
            self.u = u.copy()
        try:
            self.v = v.filled()
            if (self.v == v.fill_value).any():
                raise ValueError('u and v cannot contain missing values')
        except AttributeError:
            self.v = v.copy()
        # Check for NaN values.
        if np.isnan(self.u).any() or np.isnan(self.v).any():
            raise ValueError('u and v cannot contain missing values')
        # Make sure the shapes of the two components match.
        if u.shape != v.shape:
            raise ValueError('u and v must be the same shape')
        if len(u.shape) not in (2, 3):
            raise ValueError('u and v must be rank 2 or 3 arrays')
        try:
            # Get the number of latitudes and longitudes.
            nlat = u.shape[0]
            nlon = u.shape[1]
        except AssertionError:
            raise ValueError('nlon must be >= 4 and nlat must be >= 3')
        try:
            # Create a Spharmt object to do the computations.
            self.gridtype = gridtype.lower()
            self.s = Spharmt(nlon, nlat, gridtype=self.gridtype)
        except ValueError:
            if self.gridtype not in ('regular', 'gaussian'):
                err = 'invalid grid type: {0:s}'.format(repr(gridtype))
            else:
                err = 'invalid input dimensions'
            raise ValueError(err)
        # Method aliases.
        self.rotationalcomponent = self.nondivergentcomponent
        self.divergentcomponent = self.irrotationalcomponent

    def magnitude(self):
        """Wind speed (magnitude of vector wind).

        **Example:**

            spd = w.magnitude()

        """
        return (u ** 2 + v ** 2) ** 0.5

    def vrtdiv(self, truncation=None):
        """Relative vorticity and horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Examples:**

        Compute the relative vorticity and divergence:

            vrt, div = w.vrtdiv()

        Compute the relative vorticity and divergence and apply spectral
        truncation at triangular T13:

            vrtT13, divT13 = w.vrtdiv(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        divgrid = self.s.spectogrd(divspec)
        return vrtgrid, divgrid

    def vorticity(self, truncation=None):
        """Relative vorticity.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Examples:**

        Compute the relative vorticity:

            vrt = w.vorticity()

        Compute the relative vorticity and apply spectral truncation at
        triangular T13:

            vrtT13 = w.vorticity(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        vrtgrid = self.s.spectogrd(vrtspec)
        return vrtgrid

    def divergence(self, truncation=None):
        """Horizontal divergence.

        **Optional argument:**

        *truncation*
            Truncation limit (triangular truncation) for the spherical
            harmonic computation.

        **Examples:**

        Compute the divergence:

            div = w.divergence()

        Compute the divergence and apply spectral truncation at
        triangular T13:

            divT13 = w.divergence(truncation=13)

        """
        vrtspec, divspec = self.s.getvrtdivspec(self.u,
                                                self.v,
                                                ntrunc=truncation)
        divgrid = self.s.spectogrd(divspec)
        return divgrid

    def planetaryvorticity(self, omega=None):
        """Planetary vorticity (Coriolis parameter).

        **Optional argument:**

        *omega*
            Earth's angular velocity. The default value if not specified
            is 7.292x10**-5 s**-1.

        **Example:**

        Compute planetary vorticity using default values:

            pvrt = w.planetaryvorticity()

        Override the default value for Earth's angular velocity:

            pvrt = w.planetaryvorticity(omega=7.2921150)

        """
        if omega is None:
            # Define the Earth's angular velocity.
            omega = 7.292e-05
        nlat = self.s.nlat
        if self.gridtype == 'gaussian':
            lat, wts = gaussian_lats_wts(nlat)
        else:
            if nlat % 2:
                lat = np.linspace(90, -90, nlat)
            else:
                dlat = 180. / nlat
                lat = np.arange(90 - dlat / 2., -90, -dlat)
        try:
            cp = 2. * omega * np.sin(np.deg2rad(lat))
        except TypeError, ValueError:
            raise ValueError('invalid value for omega: {!r}'.format(omega))
        indices = [slice(0, None)] + [np.newaxis] * (len(self.u.shape) - 1)
        f = cp[indices] * np.ones(self.u.shape, dtype=np.float32)
        return f
예제 #26
0
    def __init__(self, u, v, gridtype='regular', rsphere=6.3712e6):
        """Initialize a VectorWind instance.

        **Arguments:**

        *u*, *v*
            Zonal and meridional wind components respectively. Their
            types should be either `numpy.ndarray` or
            `numpy.ma.MaskedArray`. *u* and *v* must have matching
            shapes and contain no missing values. *u* and *v* may be 2
            or 3-dimensional with shape (nlat, nlon) or
            (nlat, nlon, nt), where nlat and nlon are the number of
            latitudes and longitudes respectively and nt is the number
            of fields. The latitude dimension must be oriented
            north-to-south. The longitude dimension should be
            oriented west-to-east.

        **Optional argument:**

        *gridtype*
            Type of the input grid, either 'regular' for evenly-spaced
            grids, or 'gaussian' for Gaussian grids. Defaults to
            'regular'.

        **See also:**

        `~windspharm.tools.prep_data`,
        `~windspharm.tools.recover_data`,
        `~windspharm.tools.get_recovery`,
        `~windspharm.tools.reverse_latdim`,
        `~windspharm.tools.order_latdim`.

        **Examples:**

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind on the default regular
        (evenly-spaced) grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v)

        Initialize a `VectorWind` instance with zonal and meridional
        components of the vector wind specified on a Gaussian grid:

            from windspharm.standard import VectorWind
            w = VectorWind(u, v, gridtype='gaussian')

        """
        # For both the input components check if there are missing values by
        # attempting to fill missing values with NaN and detect them. If the
        # inputs are not masked arrays then take copies and check for NaN.
        try:
            self.u = u.filled(fill_value=np.nan)
        except AttributeError:
            self.u = u.copy()
        try:
            self.v = v.filled(fill_value=np.nan)
        except AttributeError:
            self.v = v.copy()
        if np.isnan(self.u).any() or np.isnan(self.v).any():
            raise ValueError('u and v cannot contain missing values')
        # Make sure the shapes of the two components match.
        if u.shape != v.shape:
            raise ValueError('u and v must be the same shape')
        if len(u.shape) not in (2, 3):
            raise ValueError('u and v must be rank 2 or 3 arrays')
        nlat = u.shape[0]
        nlon = u.shape[1]
        try:
            # Create a Spharmt object to do the computations.
            self.gridtype = gridtype.lower()
            self.s = Spharmt(nlon, nlat, gridtype=self.gridtype,
                             rsphere=rsphere)
        except ValueError:
            if self.gridtype not in ('regular', 'gaussian'):
                err = 'invalid grid type: {0:s}'.format(repr(gridtype))
            else:
                err = 'invalid input dimensions'
            raise ValueError(err)
        # Method aliases.
        self.rotationalcomponent = self.nondivergentcomponent
        self.divergentcomponent = self.irrotationalcomponent
예제 #27
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')
예제 #28
0
                                                             date, date)
 nc = Dataset(fcstfile)
 if ntime is None:
     times = nc['time'][:].tolist()
     levels = nc['plev'][:].tolist()
     ntime = times.index(fhour)
     nlev = levels.index(level)
     lons = nc['longitude'][:]
     lats = nc['latitude'][:]
     nlons = len(lons)
     nlats = len(lats)
     re = 6.3712e6
     ntrunc = nlats - 1
     spec = Spharmt(nlons,
                    nlats,
                    rsphere=re,
                    gridtype='regular',
                    legfunc='computed')
     indxm, indxn = getspecindx(ntrunc)
     degree = indxn.astype(np.float)
 if int(nc['time'][ntime]) != fhour:
     raise ValueError('incorrect forecast time')
 fcst_data1 = nc[varnc][ntime, nlev, ...]
 nc.close()
 if fhour > 9:
     fcstfile = '%s/%s/fv3longcontrol2_historyp_%s_latlon.nc' % (datapath2,
                                                                 date, date)
 else:
     fcstfile = '%s/%s/fv3control2_historyp_%s_latlon.nc' % (datapath2,
                                                             date, date)
 nc = Dataset(fcstfile)
예제 #29
0
# 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,rsphere=rsphere,gridtype=gridtype)
delta = 2.*np.pi/nlons
lats1d = 0.5*np.pi-delta*np.arange(nlats)
lons1d = np.arange(-np.pi,np.pi,delta)
lons,lats = np.meshgrid(lons1d,lats1d)
f = 2.*omega*np.sin(lats) # coriolis

# zonal jet.
vg = np.zeros((nlats,nlons),np.float32)
u1 = (umax/en)*np.exp(1./((lats1d-phi0)*(lats1d-phi1)))
ug = np.zeros((nlats),np.float32)
ug = np.where(np.logical_and(lats1d < phi1, lats1d > phi0), u1, ug)
ug.shape = (nlats,1)
ug = ug*np.ones((nlats,nlons),dtype=np.float32) # broadcast to shape (nlats,nlonss)
# height perturbation.
hbump = hamp*np.cos(lats)*np.exp(-(lons/alpha)**2)*np.exp(-(phi2-lats)**2/beta)