Source code for aeolus.coord_utils

# -*- coding: utf-8 -*-
"""Functionality related to coordinates of cubes."""
import iris
from iris.util import broadcast_to_shape, guess_coord_axis

from .exceptions import ArgumentError, NotFoundError


__all__ = (
    "UM_TIME",
    "UM_HGT",
    "UM_LEV",
    "UM_LATLON",
    "UM_Z_COORDS",
    "UM_TIME_COORDS",
    "get_cube_datetimes",
    "nearest_coord_value",
    "coord_to_cube",
    "ensure_bounds",
    "not_equal_coord_axes",
    "regrid_3d",
    "get_dim_coord",
)

UM_TIME = "time"
UM_LEV = "model_level_number"
UM_HGT = "level_height"
UM_LATLON = ["latitude", "longitude"]
UM_Z_COORDS = ["sigma", UM_LEV]
UM_TIME_COORDS = ["forecast_reference_time", "forecast_period", UM_TIME]


def get_cube_datetimes(cube):
    """Get a list of `iris.cube.Cube`'s time points as `datetime.datetime`s."""
    return cube.coord("time").units.num2date(cube.coord("time").points)


[docs]def nearest_coord_value(cube, coord_name, val): """ Get the nearest value of a coordinate. Parameters ---------- cube: iris.cube.Cube Cube with the coordinate coord_name: str or iris.coords.Coord Coordinate where to look the nearest point up val: int or float The value to find Returns ------- int or float element of the coordinate array closest to the given `val` """ coord = cube.coord(coord_name) i = coord.nearest_neighbour_index(val) return coord.points[i]
[docs]def coord_to_cube(cube, coord): """ Convert coordinate points to a cube of the same dimension as the given cube. Parameters ---------- cube: iris.cube.Cube Cube containing the coordinate to be broadcast. coord: str or iris.coords.Coord Coordinate to be broadcast Returns ------- iris.cube.Cube Cube of broadcast coordinate """ if isinstance(coord, str): _coord = cube.coord(coord) else: _coord = coord dim_map = cube.coord_dims(_coord.name()) _data = _coord.points if len(dim_map) > 0: _data = broadcast_to_shape(_data, cube.shape, dim_map) dc = [(c.copy(), cube.coord_dims(c)) for c in cube.dim_coords] ac = [(c.copy(), cube.coord_dims(c)) for c in cube.aux_coords] new_cube = iris.cube.Cube( data=_data, units=_coord.units, long_name=_coord.name(), dim_coords_and_dims=dc, aux_coords_and_dims=ac, ) else: new_cube = iris.cube.Cube(data=_data, standard_name=_coord.name(), units=_coord.units) return new_cube
[docs]def ensure_bounds(cube, coords=UM_LATLON): """Auto-generate bounds for cube coordinates.""" for coord_name in coords: c = cube.coord(coord_name) if not c.has_bounds(): if len(c.points) > 1: c.guess_bounds()
[docs]def not_equal_coord_axes(cube1, cube2): """Given 2 cubes, return axes of unequal dimensional coordinates.""" coord_comp = iris.analysis.coord_comparison(cube1, cube2) neq_dim_coords = set(coord_comp["not_equal"]).intersection(set(coord_comp["dimensioned"])) dims = [] for coord_pair in neq_dim_coords: for coord in coord_pair: dims.append(iris.util.guess_coord_axis(coord)) return set(filter(None, dims))
[docs]def regrid_3d(cube, target, vert_coord=None): """ Regrid a cube in the horizontal and in the vertical on to coordinates of the target cube. Adapted from https://github.com/LSaffin/iris-extensions Parameters ---------- cube: iris.cube.Cube The cube to be regridded. target: iris.cube.Cube The cube to regrid to. vert_coord: str or iris.coords.Coord, optional The coordinate for the vertical interpolation. If not given, the target's z-axis `iris.coord.DimCoord` is used. Returns ------- iris.cube.Cube """ neq_axes = not_equal_coord_axes(cube, target) if neq_axes.intersection(["X", "Y"]): cube = cube.regrid(target, iris.analysis.Linear()) # Interpolate in the vertical if needed if "Z" in neq_axes: if vert_coord is None: z = get_dim_coord(target, "z") else: z = target.coord(vert_coord) cube = cube.interpolate([(z.name(), z.points)], iris.analysis.Linear()) # Match coordinate information # XXX is this needed? # newcube = target.copy(data=cube.data) # newcube.rename(cube.name()) # newcube.units = cube.units # Put back correct time information # for coord in newcube.aux_coords: # if iris.util.guess_coord_axis(coord) == "T": # newcube.remove_coord(coord) # for coord in cube.aux_coords: # if iris.util.guess_coord_axis(coord) == "T": # newcube.add_aux_coord(coord) return cube
[docs]def get_dim_coord(cube, axis): """ Return a coordinate from a cube based on the axis it represents. Uses :py:func:`iris.util.guess_coord_axis` to heuristically match a dimensional coordinate with the requested axis. Adapted from https://github.com/LSaffin/iris-extensions Parameters ---------- cube: iris.cube.Cube Cube with the desired coordinate. axis: str The co-ordinate axis to take from the cube. Must be one of X, Y, Z, T. Returns ------- iris.coords.DimCoord The dimensional coordinate matching the requested axis on the given cube. Raises ------ ArgumentError: If axis is not one of {X, Y, Z, T}. NotFoundError: If the cube does not contain a coord with the requested axis. """ _allowed = ["X", "Y", "Z", "T"] axis = axis.upper() # If the axis supplied is not correct raise an error if axis not in _allowed: raise ArgumentError(f"Axis must be one of {_allowed}, {axis} is given.") # Loop over dimensional coords in the cube for coord in cube.dim_coords: # Return the coordinate if it matches the axis if axis == guess_coord_axis(coord): return coord # If no coordinate matches raise an error raise NotFoundError(f"Cube has no coordinate for axis {axis}")
# def z_interp_cube(cube, z=None, z_coord_name=UM_HGT, replace_z_coord=True): # """ # Interpolate cube to z points in vertical and use height as dim coord. # # Parameters # ---------- # cube: iris.cube.Cube # Input cube # z: numpy.array, optional # Array of z-points as target for interpolation. # z_coord_name: str, optional # Vertical coordinate for interpolation. # replace_z_coord: bool, optional # Replace model levels with level height. # # Returns # ------- # iris.cube.Cube # Interpolated cube with updated coordinate data # """ # cube_out = cube.copy() # # Remove model level numbers and other redundant coords # if replace_z_coord: # for coord in UM_Z_COORDS: # try: # cube_out.remove_coord(coord) # except iris.exceptions.CoordinateNotFoundError: # pass # # and use level height as dim coord # try: # iris.util.promote_aux_coord_to_dim_coord(cube_out, z_coord_name) # except ValueError: # cube_out.coord(z_coord_name).bounds = None # cube_out.coord(z_coord_name).guess_bounds() # iris.util.promote_aux_coord_to_dim_coord(cube_out, z_coord_name) # # if z is not None: # height_target = [(z_coord_name, z)] # cube_out = cube_out.interpolate(height_target, iris.analysis.Linear()) # # # Code below is for rescaling z-coordinate # # lh = cube_out.coord(z_coord_name) # # dim = cube_out.coord_dims(z_coord_name)[0] # # cube_out.remove_coord(z_coord_name) # # lh = lh.copy(lh.points * zscale, bounds=None) # # cube_out.add_dim_coord(lh, dim) # return cube_out