"""Operations on geographical grid."""
from warnings import warn
import iris
from iris.analysis.cartography import wrap_lons
from iris.util import is_regular
import numpy as np
from .const import get_planet_radius
from .coord_utils import UM_LATLON, ensure_bounds
from .exceptions import AeolusWarning
__all__ = (
"roll_cube_0_360",
"roll_cube_pm180",
"area_weights_cube",
"add_binned_lon_lat",
"coarsen_cube",
)
def _is_longitude_global(lon_points):
"""Return True if array of longitudes covers the whole sphere."""
dx = np.diff(lon_points)[0] # assume regular grid
case_0_360 = ((lon_points[0] - dx) <= 0) and ((lon_points[-1] + dx) >= 360)
case_pm180 = ((lon_points[0] - dx) <= -180) and ((lon_points[-1] + dx) >= 180)
return case_0_360 or case_pm180
[docs]def roll_cube_pm180(cube_in, coord_name=UM_LATLON[1], inplace=False):
"""
Take a cube spanning 0...360 degrees in longitude and roll it to -180...180 degrees.
Works with global model output, and in some cases for regional.
Parameters
----------
cube: iris.cube.Cube
Cube with longitude and latitude coordinates.
coord_name: str, optional
Name of the longitude coordinate.
inplace: bool, optional
Do this in-place or copy the cube.
Returns
-------
iris.cube.Cube
See also
--------
aeolus.grid.roll_cube_0_360
"""
if inplace:
cube = cube_in
else:
cube = cube_in.copy()
xcoord = cube.coord(coord_name)
if (xcoord.points >= 0.0).all():
assert is_regular(xcoord), "Operation is only valid for a regularly spaced coordinate."
if _is_longitude_global(xcoord.points):
# Shift data symmetrically only when dealing with global cubes
cube.data = np.roll(cube.data, len(xcoord.points) // 2, axis=-1)
if xcoord.has_bounds():
bounds = wrap_lons(xcoord.bounds, -180, 360) # + subtract
bounds = bounds[bounds[:, 0].argsort(axis=0)]
else:
bounds = None
cube.replace_coord(
xcoord.copy(points=np.sort(wrap_lons(xcoord.points, -180, 360)), bounds=bounds)
)
else:
# Nothing to do, the cube is already centered on 0 longitude
# unless there is something wrong with longitude
msg = f"Incorrect {coord_name} values: from {xcoord.points.min()} to {xcoord.points.max()}"
assert ((xcoord.points >= -180.0) & (xcoord.points <= 180.0)).all(), msg
if not inplace:
return cube
[docs]def roll_cube_0_360(cube_in, inplace=False):
"""
Take a cube spanning -180...180 degrees in longitude and roll it to 0...360 degrees.
Works with global model output, and in some cases for regional.
Parameters
----------
cube: iris.cube.Cube
Cube with longitude and latitude coordinates.
coord_name: str, optional
Name of the longitude coordinate.
inplace: bool, optional
Do this in-place or copy the cube.
Returns
-------
iris.cube.Cube
See also
--------
aeolus.grid.roll_cube_pm180
"""
if inplace:
cube = cube_in
else:
cube = cube_in.copy()
lon = cube.coord("longitude")
if (lon.points < 0.0).any():
add = 180
cube.data = np.roll(cube.data, len(lon.points) // 2, axis=-1)
if lon.has_bounds():
bounds = lon.bounds + add
else:
bounds = None
cube.replace_coord(lon.copy(points=lon.points + add, bounds=bounds))
if not inplace:
return cube
[docs]def area_weights_cube(cube, r_planet=None, normalize=False):
"""
Create a cube of area weights for an arbitrary planet.
Parameters
----------
cube: iris.cube.Cube
Cube with longitude and latitude coordinates
r_planet: float, optional
Radius of the planet.
normalize: bool, optional
Normalize areas.
Returns
-------
iris.cube.Cube
Cube of area weights with the same metadata as the input cube
"""
cube = cube.copy()
ensure_bounds(cube)
aw = iris.analysis.cartography.area_weights(cube, normalize=normalize)
if normalize:
aw = cube.copy(data=aw)
aw.rename("normalized_grid_cell_area")
aw.units = "1"
else:
if r_planet is None:
r = get_planet_radius(cube)
else:
r = r_planet
aw *= (r / iris.fileformats.pp.EARTH_RADIUS) ** 2
aw = cube.copy(data=aw)
aw.rename("grid_cell_area")
aw.units = "m**2"
return aw
def _cell_bounds(points, bound_position=0.5):
"""
Calculate coordinate cell boundaries.
Taken from SciTools iris package.
Parameters
----------
points: numpy.array
One-dimensional array of uniformy spaced values of shape (M,)
bound_position: bool, optional
The desired position of the bounds relative to the position
of the points.
Returns
-------
bounds: numpy.array
Array of shape (M+1,)
Examples
--------
>>> a = np.arange(-1, 2.5, 0.5)
>>> a
array([-1. , -0.5, 0. , 0.5, 1. , 1.5, 2. ])
>>> cell_bounds(a)
array([-1.25, -0.75, -0.25, 0.25, 0.75, 1.25, 1.75, 2.25])
See Also
--------
aeolus.grid._cell_centres
"""
assert points.ndim == 1, "Only 1D points are allowed"
diffs = np.diff(points)
if not np.allclose(diffs, diffs[0]):
warn("_cell_bounds() is supposed to work only for uniformly spaced points", AeolusWarning)
delta = diffs[0] * bound_position
bounds = np.concatenate([[points[0] - delta], points + delta])
return bounds
def _cell_centres(bounds, bound_position=0.5):
"""
Calculate coordinate cell centres.
Taken from SciTools iris package.
Parameters
----------
bounds: numpy.array
One-dimensional array of cell boundaries of shape (M,)
bound_position: bool, optional
The desired position of the bounds relative to the position
of the points.
Returns
-------
centres: numpy.array
Array of shape (M-1,)
Examples
--------
>>> a = np.arange(-1, 3., 1.)
>>> a
array([-1, 0, 1, 2])
>>> cell_centres(a)
array([-0.5, 0.5, 1.5])
See Also
--------
aeolus.grid._cell_bounds
"""
assert bounds.ndim == 1, "Only 1D points are allowed"
deltas = np.diff(bounds) * bound_position
centres = bounds[:-1] + deltas
return centres
[docs]def add_binned_lon_lat(cube, lon_bins, lat_bins, coord_names=UM_LATLON, inplace=False):
"""
Add binned longitude and latitude as auxiliary coordinates to a cube.
Parameters
----------
cube: iris.cube.Cube
Cube with longitude and latitude coordinates.
lon_bins: array-like
Longitude bins.
lat_bins: array-like
Latitude bins.
coord_names: list, optional
List of latitude and longitude labels.
inplace: bool, optional
Do this in-place or copy the cube.
Returns
-------
iris.cube.Cube
"""
if inplace:
cube_out = cube
else:
cube_out = cube.copy()
for name, target_points in zip(coord_names, (lat_bins, lon_bins)):
binned_points = np.digitize(cube_out.coord(name).points, target_points)
binned_points = np.clip(binned_points, 0, len(target_points) - 1)
new_coord = iris.coords.AuxCoord(binned_points, long_name=f"{name}_binned")
cube_out.add_aux_coord(new_coord, cube_out.coord_dims(name))
return cube_out
[docs]def coarsen_cube(cube, lon_bins, lat_bins, coord_names=UM_LATLON, inplace=False):
"""
Block-average cube in longitude and latitude.
Note: no weighting is applied!
Parameters
----------
cube: iris.cube.Cube
Cube with longitude and latitude coordinates.
lon_bins: array-like
Longitude bins.
lat_bins: array-like
Latitude bins.
coord_names: list, optional
List of latitude and longitude labels.
inplace: bool, optional
Do this in-place or copy the cube.
Returns
-------
iris.cube.Cube
"""
if inplace:
cube_out = cube
else:
cube_out = cube.copy()
add_binned_lon_lat(cube_out, lon_bins, lat_bins, coord_names=coord_names, inplace=True)
# To avoid oversampling on the edges, extract subset within the boundaries of target coords
for coord, target_points in zip(coord_names, (lat_bins, lon_bins)):
cube_out = cube_out.extract(
iris.Constraint(**{coord: lambda p: target_points.min() <= p <= target_points.max()})
)
for coord in coord_names:
cube_out = cube_out.aggregated_by([f"{coord}_binned"], iris.analysis.MEAN)
for coord, target_points in zip(coord_names, (lat_bins, lon_bins)):
dim = cube_out.coord_dims(coord)
units = cube_out.coord(coord).units
cube_out.remove_coord(coord)
aux = cube_out.coord(f"{coord}_binned")
new_points = target_points[aux.points]
# if len(aux.points) < len(target_points):
# new_points =
new_coord = iris.coords.DimCoord.from_coord(aux.copy(points=new_points, bounds=None))
cube_out.remove_coord(f"{coord}_binned")
new_coord.rename(coord)
new_coord.units = units
cube_out.add_dim_coord(new_coord, dim)
return cube_out