# (C) Crown Copyright, Met Office. All rights reserved.
#
# This file is part of 'IMPROVER' and is released under the BSD 3-Clause license.
# See LICENSE in the root of the repository for full licensing details.
"""Provides support utilities for checking cubes."""
from typing import List, Optional, Union
import iris
import numpy as np
from iris.cube import Cube, CubeList
from iris.exceptions import CoordinateNotFoundError
[docs]
def check_for_x_and_y_axes(cube: Cube, require_dim_coords: bool = False) -> None:
"""
Check whether the cube has an x and y axis, otherwise raise an error.
Args:
cube:
Cube to be checked for x and y axes.
require_dim_coords:
If true the x and y coordinates must be dimension coordinates.
Raises:
ValueError : Raise an error if non-uniform increments exist between
grid points.
"""
for axis in ["x", "y"]:
if require_dim_coords:
coord = cube.coords(axis=axis, dim_coords=True)
else:
coord = cube.coords(axis=axis)
if coord:
pass
else:
msg = "The cube does not contain the expected {} coordinates.".format(axis)
raise ValueError(msg)
[docs]
def validate_cube_dimensions(
cube: Cube,
required_dimensions: Optional[List[str]] = None,
forbidden_dimensions: Optional[List[str]] = None,
exact_match: bool = True,
) -> None:
"""
Validate cube dimension coordinates.
Notes:
- 't', 'x', 'y', and 'z' are treated as axes to resolve. This prevents instances
where the dimension could be called 'latitude', 'projection_x_coordinate', etc.
from being misidentified as non-dimension coordinates.
- All other entries are treated as explicit dimension coord names.
Args:
cube:
The cube to validate.
required_dimensions:
Dimension names that must be present on the cube.
forbidden_dimensions:
Dimension names that must not be present on the cube.
exact_match:
Validation mode.
- True: the cube must have exactly the required dimensions
(no more, no less).
- False: the cube must have at least the required dimensions,
but can have additional ones.
Raises:
ValueError:
- An invalid mode is specified.
- Required dimensions are missing.
- Forbidden dimensions are present.
"""
required_dimensions = list(required_dimensions or [])
forbidden_dimensions = list(forbidden_dimensions or [])
dim_coord_name_set = {coord.name() for coord in cube.dim_coords}
def _resolve_dimension(dim: str) -> str:
"""
Resolve dimension labels to actual dimension coordinate names. 't', 'x', 'y',
and 'z' are treated as axis labels to resolve, while all other entries are
treated as explicit dimension coordinate names.
Args:
dim:
Dimension label to resolve.
Returns:
Resolved dimension coordinate name.
Raises:
ValueError:
If an axis label is not found on the cube.
"""
label = dim.lower()
if label in "txyz":
try:
return cube.coord(axis=label, dim_coords=True).name()
# If the axis label is not found, we return dim as is, which will be
# caught as a missing required dimension or an unexpected forbidden
# dimension in the main validation logic, rather than raising an error here.
except CoordinateNotFoundError:
pass
return dim
required_set = {_resolve_dimension(dim) for dim in required_dimensions}
forbidden_set = {_resolve_dimension(dim) for dim in forbidden_dimensions}
# Common dimensions check
common_dims = required_set & forbidden_set
if common_dims:
raise ValueError(
f"Dimension(s) cannot be both required and forbidden: "
f"{sorted(common_dims)}. "
)
# Forbidden dimensions check
forbidden_dims_present = sorted(forbidden_set & dim_coord_name_set)
if forbidden_dims_present:
raise ValueError(
f"Forbidden dimension(s) present: {forbidden_dims_present}. "
f"Cube dim coords are: {sorted(dim_coord_name_set)}"
)
# Required dimensions check
missing_required_dims = sorted(required_set - dim_coord_name_set)
if missing_required_dims:
raise ValueError(
f"Missing required dimension(s): {missing_required_dims}. "
f"Found: {sorted(dim_coord_name_set)}"
)
# Exact match check
if exact_match:
extra_dims = sorted(dim_coord_name_set - required_set)
if extra_dims:
raise ValueError(
f"Extra dimension(s) present: {extra_dims}. "
"and exact_match is True. Remove these dimensions from the cube "
"or set exact_match to False if these dimensions are acceptable."
)
[docs]
def check_cube_coordinates(
cube: Cube, new_cube: Cube, exception_coordinates: Optional[List[str]] = None
) -> Cube:
"""Find and promote to dimension coordinates any scalar coordinates in
new_cube that were originally dimension coordinates in the progenitor
cube. If coordinate is in new_cube that is not in the old cube, keep
coordinate in its current position.
Args:
cube:
The input cube that will be checked to identify the preferred
coordinate order for the output cube.
new_cube:
The cube that must be checked and adjusted using the coordinate
order from the original cube.
exception_coordinates:
The names of the coordinates that are permitted to be within the
new_cube but are not available within the original cube.
Returns:
Modified cube with relevant scalar coordinates promoted to
dimension coordinates with the dimension coordinates re-ordered,
as best as can be done based on the original cube.
Raises:
CoordinateNotFoundError : Raised if the final dimension
coordinates of the returned cube do not match the input cube.
CoordinateNotFoundError : If a coordinate is within in the permitted
exceptions but is not in the new_cube.
"""
if exception_coordinates is None:
exception_coordinates = []
# Promote available and relevant scalar coordinates
cube_dim_names = [coord.name() for coord in cube.dim_coords]
for coord in new_cube.aux_coords[::-1]:
if coord.name() in cube_dim_names:
new_cube = iris.util.new_axis(new_cube, coord)
new_cube_dim_names = [coord.name() for coord in new_cube.dim_coords]
# If we have the wrong number of dimensions then raise an error.
if len(cube.dim_coords) + len(exception_coordinates) != len(new_cube.dim_coords):
msg = (
"The number of dimension coordinates within the new cube "
"do not match the number of dimension coordinates within the "
"original cube plus the number of exception coordinates. "
"\n input cube dimensions {}, new cube dimensions {}".format(
cube_dim_names, new_cube_dim_names
)
)
raise CoordinateNotFoundError(msg)
# Ensure dimension order matches
new_cube_dimension_order = {
coord.name(): new_cube.coord_dims(coord.name())[0]
for coord in new_cube.dim_coords
}
correct_order = []
new_cube_only_dims = []
for coord_name in cube_dim_names:
correct_order.append(new_cube_dimension_order[coord_name])
for coord_name in exception_coordinates:
try:
new_coord_dim = new_cube.coord_dims(coord_name)[0]
new_cube_only_dims.append(new_coord_dim)
except CoordinateNotFoundError:
msg = (
"All permitted exception_coordinates must be on the"
" new_cube. In this case, coordinate {0} within the list "
"of permitted exception_coordinates ({1}) is not available"
" on the new_cube."
).format(coord_name, exception_coordinates)
raise CoordinateNotFoundError(msg)
correct_order = np.array(correct_order)
for dim in new_cube_only_dims:
correct_order = np.insert(correct_order, dim, dim)
new_cube.transpose(correct_order)
return new_cube
[docs]
def find_dimension_coordinate_mismatch(
first_cube: Cube, second_cube: Cube, two_way_mismatch: bool = True
) -> List[str]:
"""Determine if there is a mismatch between the dimension coordinates in
two cubes.
Args:
first_cube:
First cube to compare.
second_cube:
Second cube to compare.
two_way_mismatch:
If True, a two way mismatch is calculated e.g.
second_cube - first_cube AND
first_cube - second_cube
If False, a one way mismatch is calculated e.g.
second_cube - first_cube
Returns:
List of the dimension coordinates that are only present in
one out of the two cubes.
"""
first_dim_names = [coord.name() for coord in first_cube.dim_coords]
second_dim_names = [coord.name() for coord in second_cube.dim_coords]
if two_way_mismatch:
mismatch = list(set(second_dim_names) - set(first_dim_names)) + list(
set(first_dim_names) - set(second_dim_names)
)
else:
mismatch = list(set(second_dim_names) - set(first_dim_names))
return mismatch
[docs]
def spatial_coords_match(cubes: Union[List[Cube], CubeList]) -> bool:
"""
Determine if the x and y coords of all the input cubes are the same.
Args:
cubes:
A list of cubes to compare.
Returns:
True if the x and y coords are exactly the same to the
precision of the floating-point values (this should be true for
any cubes derived using cube.regrid()), otherwise False.
"""
ref = cubes[0]
match = True
for cube in cubes[1:]:
match = (
cube.coord(axis="x") == ref.coord(axis="x")
and cube.coord(axis="y") == ref.coord(axis="y")
and match
)
return match
[docs]
def assert_time_coords_valid(inputs: List[Cube], time_bounds: bool):
"""
Raises appropriate ValueError if
- Any input cube has or is missing time bounds (depending on time_bounds)
- Input cube times do not match
- Input cube forecast_reference_times do not match (unless blend_time is present)
Note that blend_time coordinates do not have to match as it is likely that data
from nearby blends will be used together.
Args:
inputs:
List of Cubes where times should match
time_bounds:
When True, time bounds are checked for and compared on the input cubes.
When False, the absence of time bounds is checked for.
Raises:
ValueError: If any of the stated checks fail.
"""
if len(inputs) <= 1:
raise ValueError(f"Need at least 2 cubes to check. Found {len(inputs)}")
cubes_not_matching_time_bounds = [
c.name() for c in inputs if c.coord("time").has_bounds() != time_bounds
]
if cubes_not_matching_time_bounds:
str_bool = "" if time_bounds else "not "
msg = f"{' and '.join(cubes_not_matching_time_bounds)} must {str_bool}have time bounds"
raise ValueError(msg)
if inputs[0].coords("blend_time"):
time_coords_to_check = ["time"]
else:
time_coords_to_check = ["time", "forecast_reference_time"]
for time_coord_name in time_coords_to_check:
time_coords = [c.coord(time_coord_name) for c in inputs]
if not all([tc == time_coords[0] for tc in time_coords[1:]]):
msg = f"{time_coord_name} coordinates do not match. \n " + "\n ".join(
[f"{c.name()}: {c.coord('time')}" for c in inputs]
)
raise ValueError(msg)
[docs]
def assert_spatial_coords_match(cubes: Union[List, CubeList]):
"""
Raises an Exception if `spatial_coords_match` returns False.
Args:
cubes:
A list of cubes to compare.
Raises:
ValueError if spatial coords do not match.
"""
if not spatial_coords_match(cubes):
raise ValueError(
f"Mismatched spatial coords for {', '.join([c.name() for c in cubes])}"
)
