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
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) model.vrt = sp.spectogrd(vrtspec) model.theta = sp.spectogrd(thetaspec) model.heat = np.zeros(model.theta.shape, model.theta.dtype) # animate vorticity from solution as it is running # first, create initial frame fig = plt.figure(figsize=(8, 10)) ax1 = fig.add_subplot(2, 1, 1) ax2 = fig.add_subplot(2, 1, 2)
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
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
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) # initial vorticity, divergence in spectral space vrtspec, divspec = x.getvrtdivspec(ug, vg, ntrunc) # create spectral indexing arrays, laplacian operator and its inverse. indxm, indxn = getspecindx(ntrunc) lap = -(indxn * (indxn + 1.0) / rsphere**2).astype(np.float32) ilap = np.zeros(lap.shape, np.float32) ilap[1:] = 1. / lap[1:] hyperdiff_fact = np.exp((-dt / efold) * (lap / 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, ntrunc) tmpspec2 = x.grdtospec(0.5 * (ug**2 + vg**2), ntrunc)
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()
# 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) lyrthkg = sp.spectogrd(lyrthkspec) print(lyrthkg[:,:,0].min(), lyrthkg[:,:,0].max()) print(lyrthkg[:,:,1].min(), lyrthkg[:,:,1].max()) if lyrthkg.min() < 0: raise ValueError('negative layer thickness! adjust jet parameters') # time loop. time1 = time.clock() for ncycle in range(itmax+1): t = ncycle*model.dt vrtspec, divspec, lyrthkspec = model.rk4step(vrtspec, divspec, lyrthkspec) pvg = (0.5*model.zmid/model.omega)*(model.vrt + model.f)/model.lyrthk
class TransformsEngine(object): """A spectral transforms engine based on pyspharm.""" def __init__(self, nlon, nlat, truncation, radius=6371200.): """ Initialize the spectral transforms engine. Arguments: * nlon: int Number of longitudes in the transform grid. * nlat: int Number of latitudes in the transform grid. * truncation: int The spectral truncation (triangular). This is the maximum number of spherical harmonic modes retained in the discrete truncation. More modes means higher resolution. """ self.sh = Spharmt(nlon, nlat, gridtype='regular', rsphere=radius) self.radius = radius self.nlon = nlon self.nlat = nlat self.truncation = truncation def vrtdiv_spec_from_uv_grid(self, u, v): """ Compute spectral vorticity and divergence from grid u and v. """ try: vrt, div = self.sh.getvrtdivspec(u, v, ntrunc=self.truncation) except ValueError: msg = ('u and v must be 2d or 3d arrays with shape ({y}, {x}) ' 'or ({y}, {x}, :)'.format(y=self.nlat, x=self.nlon)) raise ValueError(msg) return vrt, div def uv_grid_from_vrtdiv_spec(self, vrt, div): """ Compute grid u and v from spectral vorticity and divergence. """ try: u, v = self.sh.getuv(vrt, div) except ValueError: nspec = (self.truncation + 1) * (self.truncation + 2) // 2 msg = ('vrt and div must be 1d or 2d arrays with shape ' '(n) or (n, :) where n <= {}'.format(nspec)) raise ValueError(msg) return u, v def spec_to_grid(self, scalar_spec): """ Transform a scalar field from spectral to grid space. """ try: scalar_grid = self.sh.spectogrd(scalar_spec) except ValueError: nspec = (self.truncation + 1) * (self.truncation + 2) // 2 msg = ('scalar_spec must be a 1d or 2d array with shape ' '(n) or (n, :) where n <= {}'.format(nspec)) raise ValueError(msg) return scalar_grid def grid_to_spec(self, scalar_grid): """ Transform a scalar field from grid to spectral space. """ try: scalar_spec = self.sh.grdtospec(scalar_grid, ntrunc=self.truncation) except ValueError: msg = ('scalar_grid must be a 2d or 3d array with shape ' '({y}, {x}) or ({y}, {x}, :)'.format(y=self.nlat, x=self.nlon)) raise ValueError(msg) return scalar_spec def grad_of_spec(self, scalar_spec): """ Return zonal and meridional gradients of a spectral field. """ try: dsdx, dsdy = self.sh.getgrad(scalar_spec) except ValueError: nspec = (self.truncation + 1) * (self.truncation + 2) // 2 msg = ('scalar_spec must be a 1d or 2d array with shape ' '(n) or (n, :) where n <= {}'.format(nspec)) raise ValueError(msg) return dsdx, dsdy @property def wavenumbers(self): """ Wavenumbers corresponding to the spectral fields. """ return getspecindx(self.truncation) @property def grid_latlon(self): """ Return the latitude and longitude coordinate vectors of the model grid. """ lats, _ = gaussian_lats_wts(self.nlat) lons = np.arange(0., 360., 360. / self.nlon) return lats, lons
class spectral: def __init__(self, lat, lon, rsphere=6.3712e6, legfunc='stored'): # Length of lat/lon arrays self.nlat = len(lat) self.nlon = len(lon) if self.nlat % 2: gridtype = 'gaussian' else: gridtype = 'regular' self.s = Spharmt(self.nlon, self.nlat, gridtype=gridtype, rsphere=rsphere, legfunc=legfunc) # Reverse latitude array if necessary # self.ReverseLat = False # if lat[0] < lat[-1]: # lat = self._reverse_lat(lat) # self.ReverseLat = True # lat/lon in degrees self.glat = lat self.glon = lon # lat/lon in radians self.rlat = np.deg2rad(lat) self.rlon = np.deg2rad(lon) self.rlons, self.rlats = np.meshgrid(self.rlon, self.rlat) # Constants # Earth's angular velocity self.omega = 7.292e-05 # unit: s-1 # Gravitational acceleration self.g = 9.8 # unit: m2/s # Misc self.dtype = np.float32 def _reverse_lat(self, var): if len(np.shape(var)) == 1: return (var[::-1]) if len(np.shape(var)) == 2: return (var[::-1, :]) if len(np.shape(var)) == 3: return (var[:, ::-1, :]) def planetaryvorticity(self, omega=None): pass def uv2vrt(self, u, v, trunc=None): """ Calculate relative vorticity from horizonal wind field Input: u and v (grid) Output: relative vorticity (grid) """ vrts, _ = self.s.getvrtdivspec(u, v, ntrunc=trunc) vrtg = self.s.spectogrd(vrts) return (vrtg) def uv2div(self, u, v, trunc=None): pass def uv2sfvp(self, u, v, trunc=None): """ Calculate geostrophic streamfuncion and velocity potential from u and v winds """ psig, chig = self.s.getpsichi(u, v, ntrunc=trunc) return (psig, chig) def uv2vrtdiv(self, u, v, trunc=None): """ Calculate relative vorticity and divergence from u and v winds """ vrts, divs = self.s.getvrtdivspec(u, v, ntrunc=trunc) vrtg = self.s.spectogrd(vrts) divg = self.s.spectogrd(divs) return (vrtg, divg) def vrtdiv2uv(self, vrt, div, realm='grid', trunc=None): """ # Get u,v from vrt, div fields # Input either in grid space # or in spectral space """ if realm in ['g', 'grid']: vrts = self.s.grdtospec(vrt, trunc) divs = self.s.grdtospec(div, trunc) elif realm in ['s', 'spec', 'spectral']: vrts = vrt divs = div ug, vg = self.s.getuv(vrts, divs) return (ug, vg) def gradient(self, var, trunc=None): """ Calculate horizontal gradients """ # if self.ReverseLat is True: # var = self._reverse_lat(var) try: var = var.filled(fill_value=np.nan) except AttributeError: pass if np.isnan(var).any(): raise ValueError('var cannot contain missing values') try: varspec = self.s.grdtospec(var, ntrunc=trunc) except ValueError: raise ValueError('input field is not compatitble') dxvarg, dyvarg = self.s.getgrad(varspec) return (dxvarg, dyvarg)
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) # initial vorticity, divergence in spectral space vrtspec, divspec = x.getvrtdivspec(ug,vg,ntrunc) # create spectral indexing arrays, laplacian operator and its inverse. indxm, indxn = getspecindx(ntrunc) lap = -(indxn*(indxn+1.0)/rsphere**2).astype(np.float32) ilap = np.zeros(lap.shape, np.float32) ilap[1:] = 1./lap[1:] hyperdiff_fact = np.exp((-dt/efold)*(lap/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,ntrunc) tmpspec2 = x.grdtospec(0.5*(ug**2+vg**2),ntrunc) phispec = ilap*tmpspec1 - tmpspec2