def test_cube_dim(self):
cube1_reverse0 = reverse(self.cube1, 0)
cube1_reverse1 = reverse(self.cube1, 1)
cube1_reverse_both = reverse(self.cube1, (0, 1))
self.assertArrayEqual(self.cube1.data[::-1], cube1_reverse0.data)
self.assertArrayEqual(
self.cube2.coord("a").points,
cube1_reverse0.coord("a").points)
self.assertArrayEqual(
self.cube1.coord("b").points,
cube1_reverse0.coord("b").points)
self.assertArrayEqual(self.cube1.data[:, ::-1], cube1_reverse1.data)
self.assertArrayEqual(
self.cube1.coord("a").points,
cube1_reverse1.coord("a").points)
self.assertArrayEqual(
self.cube2.coord("b").points,
cube1_reverse1.coord("b").points)
self.assertArrayEqual(self.cube1.data[::-1, ::-1],
cube1_reverse_both.data)
self.assertArrayEqual(
self.cube2.coord("a").points,
cube1_reverse_both.coord("a").points)
self.assertArrayEqual(
self.cube2.coord("b").points,
cube1_reverse_both.coord("b").points)
def __init__(self, u, v, rsphere=6.3712e6):
"""Initialize a VectorWind instance.
**Arguments:**
*u*, *v*
Zonal and meridional components of the vector wind
respectively. Both components should be `~iris.cube.Cube`
instances. The components must have the same dimension
coordinates and contain no missing values.
**Example:**
Initialize a `VectorWind` instance with zonal and meridional
components of the vector wind::
from windspharm.iris import VectorWind
w = VectorWind(u, v)
"""
# Make sure inputs are Iris cubes.
if type(u) is not Cube or type(v) is not Cube:
raise TypeError('u and v must be iris cubes')
# Get the coordinates of each component and make sure they are the
# same.
ucoords = u.dim_coords
vcoords = v.dim_coords
if ucoords != vcoords:
raise ValueError('u and v must have the same dimensions')
# Extract the latitude and longitude dimension coordinates.
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
lon, lon_dim = _dim_coord_and_dim(v, 'longitude')
# Reverse the latitude dimension if necessary.
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
u = reverse(u, lat_dim)
v = reverse(v, lat_dim)
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
# Determine the grid type of the input.
gridtype = inspect_gridtype(lat.points)
# Determine the ordering list (input to transpose) which will put the
# latitude and longitude dimensions at the front of the cube's
# dimensions, and the ordering list which will reverse this process.
apiorder, self._reorder = get_apiorder(u.ndim, lat_dim, lon_dim)
# Re-order the inputs (in-place, so we take a copy first) so latiutude
# and longitude are at the front.
u = u.copy()
v = v.copy()
u.transpose(apiorder)
v.transpose(apiorder)
# Records the current shape and dimension coordinates of the inputs.
self._ishape = u.shape
self._coords = u.dim_coords
# Reshape the inputs so they are compatible with pyspharm.
u = u.data.reshape(u.shape[:2] + (np.prod(u.shape[2:]),))
v = v.data.reshape(v.shape[:2] + (np.prod(v.shape[2:]),))
# Create a base VectorWind instance to do the computations.
self._api = standard.VectorWind(u, v, gridtype=gridtype,
rsphere=rsphere)
def __init__(self, u, v):
"""Initialize a VectorWind instance.
**Arguments:**
*u*, *v*
Zonal and meridional components of the vector wind
respectively. Both components should be `~iris.cube.Cube`
instances. The components must have the same dimension
coordinates and contain no missing values.
**Example:**
Initialize a `VectorWind` instance with zonal and meridional
components of the vector wind::
from windspharm.iris import VectorWind
w = VectorWind(u, v)
"""
# Make sure inputs are Iris cubes.
if type(u) is not Cube or type(v) is not Cube:
raise TypeError('u and v must be iris cubes')
# Get the coordinates of each component and make sure they are the
# same.
ucoords = u.dim_coords
vcoords = v.dim_coords
if ucoords != vcoords:
raise ValueError('u and v must have the same dimensions')
# Extract the latitude and longitude dimension coordinates.
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
lon, lon_dim = _dim_coord_and_dim(v, 'longitude')
# Reverse the latitude dimension if necessary.
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
u = reverse(u, lat_dim)
v = reverse(v, lat_dim)
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
# Determine the grid type of the input.
gridtype = self._gridtype(lat.points)
# Determine the ordering list (input to transpose) which will put the
# latitude and longitude dimensions at the front of the cube's
# dimensions, and the ordering list which will reverse this process.
apiorder, self._reorder = self._get_apiorder_reorder(
u, lat_dim, lon_dim)
# Re-order the inputs (in-place, so we take a copy first) so latiutude
# and longitude are at the front.
u = u.copy()
v = v.copy()
u.transpose(apiorder)
v.transpose(apiorder)
# Records the current shape and dimension coordinates of the inputs.
self._ishape = u.shape
self._coords = u.dim_coords
# Reshape the inputs so they are compatible with pyspharm.
u = u.data.reshape(u.shape[:2] + (np.prod(u.shape[2:]), ))
v = v.data.reshape(v.shape[:2] + (np.prod(v.shape[2:]), ))
# Create a base VectorWind instance to do the computations.
self._api = standard.VectorWind(u, v, gridtype=gridtype)
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. It must be a `~iris.cube.Cube`
with the same latitude and longitude dimensions as 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)
"""
if type(field) is not Cube:
raise TypeError('scalar field must be an iris cube')
lat, lat_dim = _dim_coord_and_dim(field, 'latitude')
lon, lon_dim = _dim_coord_and_dim(field, 'longitude')
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
field = reverse(field, lat_dim)
lat, lat_dim = _dim_coord_and_dim(field, 'latitude')
apiorder, reorder = get_apiorder(field.ndim, lat_dim, lon_dim)
field = field.copy()
field.transpose(apiorder)
ishape = field.shape
coords = field.dim_coords
fielddata = field.data.reshape(
field.shape[:2] + (np.prod(field.shape[2:]),))
fieldtrunc = self._api.truncate(fielddata, truncation=truncation)
field.data = fieldtrunc.reshape(ishape)
field.transpose(reorder)
return field
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. It must be a `~iris.cube.Cube`
with the same latitude and longitude dimensions as 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)
"""
if type(field) is not Cube:
raise TypeError('scalar field must be an iris cube')
lat, lat_dim = _dim_coord_and_dim(field, 'latitude')
lon, lon_dim = _dim_coord_and_dim(field, 'longitude')
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
field = reverse(field, lat_dim)
lat, lat_dim = _dim_coord_and_dim(field, 'latitude')
apiorder, reorder = get_apiorder(field.ndim, lat_dim, lon_dim)
field = field.copy()
field.transpose(apiorder)
ishape = field.shape
coords = field.dim_coords
fielddata = field.data.reshape(field.shape[:2] +
(np.prod(field.shape[2:]), ))
fieldtrunc = self._api.truncate(fielddata, truncation=truncation)
field.data = fieldtrunc.reshape(ishape)
field.transpose(reorder)
return field
def meridional_mass_streamfunction(cubelist, z_coord=UM_HGT):
v = cubelist.extract_strict("y_wind")
const = v.attributes["planet_conf"]
if v.coord(z_coord).units.is_convertible("m"):
rho = cubelist.extract_strict("air_density")
rho.coord(z_coord).bounds = None
v.coord(z_coord).bounds = None
integrand = zonal_mean(v * rho)
integrand = reverse(integrand, z_coord)
res = -1 * vertical_cumsum(integrand, coord=z_coord)
res = reverse(res, z_coord)
elif v.coord(z_coord).units.is_convertible("Pa"):
# pres = cubelist.extract_strict("air_pressure")
# mmstreamf_const /= const.gravity.asc
raise NotImplementedError()
coslat = apply_ufunc(np.cos, apply_ufunc(np.deg2rad, coord_to_cube(res, UM_LATLON[0])))
coslat.units = "1"
mmstreamf_const = 2 * np.pi * coslat * const.radius.asc
res = res * mmstreamf_const
res.rename("meridional_mass_streamfunction")
res.convert_units("kg s^-1")
return res
def gradient(self, chi, truncation=None):
"""Computes the vector gradient of a scalar field on the sphere.
**Argument:**
*chi*
A scalar field. It must be a `~iris.cube.Cube`
with the same latitude and longitude dimensions as 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)
"""
if type(chi) is not Cube:
raise TypeError('scalar field must be an iris cube')
name = chi.name()
lat, lat_dim = _dim_coord_and_dim(chi, 'latitude')
lon, lon_dim = _dim_coord_and_dim(chi, 'longitude')
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
chi = reverse(chi, lat_dim)
lat, lat_dim = _dim_coord_and_dim(chi, 'latitude')
apiorder, reorder = get_apiorder(chi.ndim, lat_dim, lon_dim)
chi = chi.copy()
chi.transpose(apiorder)
ishape = chi.shape
coords = chi.dim_coords
chi = chi.data.reshape(chi.shape[:2] + (np.prod(chi.shape[2:]),))
uchi, vchi = self._api.gradient(chi, truncation=truncation)
uchi = uchi.reshape(ishape)
vchi = vchi.reshape(ishape)
uchi = Cube(
uchi,
dim_coords_and_dims=list(zip(coords, range(uchi.ndim))))
vchi = Cube(
vchi,
dim_coords_and_dims=list(zip(coords, range(vchi.ndim))))
uchi.transpose(reorder)
vchi.transpose(reorder)
uchi.long_name = 'zonal_gradient_of_{!s}'.format(name)
vchi.long_name = 'meridional_gradient_of_{!s}'.format(name)
return uchi, vchi
def gradient(self, chi, truncation=None):
"""Computes the vector gradient of a scalar field on the sphere.
**Argument:**
*chi*
A scalar field. It must be a `~iris.cube.Cube`
with the same latitude and longitude dimensions as 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)
"""
if type(chi) is not Cube:
raise TypeError('scalar field must be an iris cube')
name = chi.name()
lat, lat_dim = _dim_coord_and_dim(chi, 'latitude')
lon, lon_dim = _dim_coord_and_dim(chi, 'longitude')
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
chi = reverse(chi, lat_dim)
lat, lat_dim = _dim_coord_and_dim(chi, 'latitude')
apiorder, reorder = get_apiorder(chi.ndim, lat_dim, lon_dim)
chi = chi.copy()
chi.transpose(apiorder)
ishape = chi.shape
coords = chi.dim_coords
chi = chi.data.reshape(chi.shape[:2] + (np.prod(chi.shape[2:]), ))
uchi, vchi = self._api.gradient(chi, truncation=truncation)
uchi = uchi.reshape(ishape)
vchi = vchi.reshape(ishape)
uchi = Cube(uchi,
dim_coords_and_dims=list(zip(coords, range(uchi.ndim))))
vchi = Cube(vchi,
dim_coords_and_dims=list(zip(coords, range(vchi.ndim))))
uchi.transpose(reorder)
vchi.transpose(reorder)
uchi.long_name = 'zonal_gradient_of_{!s}'.format(name)
vchi.long_name = 'meridional_gradient_of_{!s}'.format(name)
return uchi, vchi
def test_single_array(self):
a = np.arange(36).reshape(3, 4, 3)
self.assertArrayEqual(a[::-1], reverse(a, 0))
self.assertArrayEqual(a[::-1, ::-1], reverse(a, [0, 1]))
self.assertArrayEqual(a[:, ::-1, ::-1], reverse(a, [1, 2]))
self.assertArrayEqual(a[..., ::-1], reverse(a, 2))
msg = "Reverse was expecting a single axis or a 1d array *"
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [])
msg = "An axis value out of range for the number of dimensions *"
with self.assertRaisesRegex(ValueError, msg):
reverse(a, -1)
with self.assertRaisesRegex(ValueError, msg):
reverse(a, 10)
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [-1])
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [0, -1])
with self.assertRaisesRegex(TypeError,
"To reverse an array, provide an int *"):
reverse(a, "latitude")
def test_cube_coord(self):
cube1_reverse0 = reverse(self.cube1, self.a1)
cube1_reverse1 = reverse(self.cube1, "b")
cube1_reverse_both = reverse(self.cube1, (self.a1, self.b1))
cube1_reverse_spanning = reverse(self.cube1, "spanning")
self.assertArrayEqual(self.cube1.data[::-1], cube1_reverse0.data)
self.assertArrayEqual(
self.cube2.coord("a").points,
cube1_reverse0.coord("a").points)
self.assertArrayEqual(
self.cube1.coord("b").points,
cube1_reverse0.coord("b").points)
self.assertArrayEqual(self.cube1.data[:, ::-1], cube1_reverse1.data)
self.assertArrayEqual(
self.cube1.coord("a").points,
cube1_reverse1.coord("a").points)
self.assertArrayEqual(
self.cube2.coord("b").points,
cube1_reverse1.coord("b").points)
self.assertArrayEqual(self.cube1.data[::-1, ::-1],
cube1_reverse_both.data)
self.assertArrayEqual(
self.cube2.coord("a").points,
cube1_reverse_both.coord("a").points)
self.assertArrayEqual(
self.cube2.coord("b").points,
cube1_reverse_both.coord("b").points)
self.assertArrayEqual(self.cube1.data[::-1, ::-1],
cube1_reverse_spanning.data)
self.assertArrayEqual(
self.cube2.coord("a").points,
cube1_reverse_spanning.coord("a").points,
)
self.assertArrayEqual(
self.cube2.coord("b").points,
cube1_reverse_spanning.coord("b").points,
)
self.assertArrayEqual(
self.span.points[::-1, ::-1],
cube1_reverse_spanning.coord("spanning").points,
)
msg = (
"Expected to find exactly 1 'latitude' coordinate, but found none."
)
with self.assertRaisesRegex(iris.exceptions.CoordinateNotFoundError,
msg):
reverse(self.cube1, "latitude")
msg = "Reverse was expecting a single axis or a 1d array *"
with self.assertRaisesRegex(ValueError, msg):
reverse(self.cube1, [])
msg = ("coords_or_dims must be int, str, coordinate or sequence of "
"these. Got cube.")
with self.assertRaisesRegex(TypeError, msg):
reverse(self.cube1, self.cube1)
msg = ("coords_or_dims must be int, str, coordinate or sequence of "
"these.")
with self.assertRaisesRegex(TypeError, msg):
reverse(self.cube1, 3.0)
def __init__(self, u, v, rsphere=6.3712e6, legfunc='stored'):
"""Initialize a VectorWind instance.
**Arguments:**
*u*, *v*
Zonal and meridional components of the vector wind
respectively. Both components should be `~iris.cube.Cube`
instances. The components must have the same dimension
coordinates and contain no missing values.
**Optional argument:**
*rsphere*
The radius in metres of the sphere used in the spherical
harmonic computations. Default is 6371200 m, the approximate
mean spherical Earth radius.
*legfunc*
'stored' (default) or 'computed'. If 'stored', associated legendre
functions are precomputed and stored when the class instance is
created. This uses O(nlat**3) memory, but speeds up the spectral
transforms. If 'computed', associated legendre functions are
computed on the fly when transforms are requested. This uses
O(nlat**2) memory, but slows down the spectral transforms a bit.
**Example:**
Initialize a `VectorWind` instance with zonal and meridional
components of the vector wind::
from windspharm.iris import VectorWind
w = VectorWind(u, v)
"""
# Make sure inputs are Iris cubes.
if type(u) is not Cube or type(v) is not Cube:
raise TypeError('u and v must be iris cubes')
# Get the coordinates of each component and make sure they are the
# same.
ucoords = u.dim_coords
vcoords = v.dim_coords
if ucoords != vcoords:
raise ValueError('u and v must have the same dimensions')
# Extract the latitude and longitude dimension coordinates.
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
lon, lon_dim = _dim_coord_and_dim(v, 'longitude')
# Reverse the latitude dimension if necessary.
if (lat.points[0] < lat.points[1]):
# need to reverse latitude dimension
u = reverse(u, lat_dim)
v = reverse(v, lat_dim)
lat, lat_dim = _dim_coord_and_dim(u, 'latitude')
# Determine the grid type of the input.
gridtype = inspect_gridtype(lat.points)
# Determine the ordering list (input to transpose) which will put the
# latitude and longitude dimensions at the front of the cube's
# dimensions, and the ordering list which will reverse this process.
apiorder, self._reorder = get_apiorder(u.ndim, lat_dim, lon_dim)
# Re-order the inputs (in-place, so we take a copy first) so latiutude
# and longitude are at the front.
u = u.copy()
v = v.copy()
u.transpose(apiorder)
v.transpose(apiorder)
# Records the current shape and dimension coordinates of the inputs.
self._ishape = u.shape
self._coords = u.dim_coords
# Reshape the inputs so they are compatible with pyspharm.
u = to3d(u.data)
v = to3d(v.data)
# Create a base VectorWind instance to do the computations.
self._api = standard.VectorWind(u,
v,
gridtype=gridtype,
rsphere=rsphere,
legfunc=legfunc)
def test_simple_array(self):
a = np.arange(12).reshape(3, 4)
self.assertArrayEqual(a[::-1], reverse(a, 0))
self.assertArrayEqual(a[::-1, ::-1], reverse(a, [0, 1]))
self.assertArrayEqual(a[:, ::-1], reverse(a, 1))
self.assertArrayEqual(a[:, ::-1], reverse(a, [1]))
msg = 'Reverse was expecting a single axis or a 1d array *'
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [])
msg = 'An axis value out of range for the number of dimensions *'
with self.assertRaisesRegex(ValueError, msg):
reverse(a, -1)
with self.assertRaisesRegex(ValueError, msg):
reverse(a, 10)
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [-1])
with self.assertRaisesRegex(ValueError, msg):
reverse(a, [0, -1])
msg = 'To reverse an array, provide an int *'
with self.assertRaisesRegex(TypeError, msg):
reverse(a, 'latitude')
def test_cube_coord(self):
cube1_reverse0 = reverse(self.cube1, self.a1)
cube1_reverse1 = reverse(self.cube1, 'b')
cube1_reverse_both = reverse(self.cube1, (self.a1, self.b1))
cube1_reverse_spanning = reverse(self.cube1, 'spanning')
self.assertArrayEqual(self.cube1.data[::-1], cube1_reverse0.data)
self.assertArrayEqual(
self.cube2.coord('a').points,
cube1_reverse0.coord('a').points)
self.assertArrayEqual(
self.cube1.coord('b').points,
cube1_reverse0.coord('b').points)
self.assertArrayEqual(self.cube1.data[:, ::-1], cube1_reverse1.data)
self.assertArrayEqual(
self.cube1.coord('a').points,
cube1_reverse1.coord('a').points)
self.assertArrayEqual(
self.cube2.coord('b').points,
cube1_reverse1.coord('b').points)
self.assertArrayEqual(self.cube1.data[::-1, ::-1],
cube1_reverse_both.data)
self.assertArrayEqual(
self.cube2.coord('a').points,
cube1_reverse_both.coord('a').points)
self.assertArrayEqual(
self.cube2.coord('b').points,
cube1_reverse_both.coord('b').points)
self.assertArrayEqual(self.cube1.data[::-1, ::-1],
cube1_reverse_spanning.data)
self.assertArrayEqual(
self.cube2.coord('a').points,
cube1_reverse_spanning.coord('a').points)
self.assertArrayEqual(
self.cube2.coord('b').points,
cube1_reverse_spanning.coord('b').points)
self.assertArrayEqual(self.span.points[::-1, ::-1],
cube1_reverse_spanning.coord('spanning').points)
msg = 'Expected to find exactly 1 latitude coordinate, but found none.'
with self.assertRaisesRegex(iris.exceptions.CoordinateNotFoundError,
msg):
reverse(self.cube1, 'latitude')
msg = 'Reverse was expecting a single axis or a 1d array *'
with self.assertRaisesRegex(ValueError, msg):
reverse(self.cube1, [])
msg = ('coords_or_dims must be int, str, coordinate or sequence of '
'these. Got cube.')
with self.assertRaisesRegex(TypeError, msg):
reverse(self.cube1, self.cube1)
msg = ('coords_or_dims must be int, str, coordinate or sequence of '
'these.')
with self.assertRaisesRegex(TypeError, msg):
reverse(self.cube1, 3.)