Source code for improver.synthetic_data.generate_metadata
# (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.
"""Module to generate a metadata cube."""
from datetime import datetime, timedelta
from typing import Any, List, Optional, Tuple
import numpy as np
from iris.cube import Cube
from iris.std_names import STD_NAMES
from iris.util import squeeze
from numpy import ndarray
from improver.metadata.constants.attributes import MANDATORY_ATTRIBUTES
from improver.synthetic_data.set_up_test_cubes import (
set_up_percentile_cube,
set_up_probability_cube,
set_up_variable_cube,
)
DEFAULT_GRID_SPACING = {"latlon": 0.02, "equalarea": 2000}
DEFAULT_SPATIAL_GRID = "latlon"
DEFAULT_TIME = datetime(2017, 11, 10, 4, 0)
CUBE_TYPES = ("variable", "percentile", "probability")
[docs]
def _get_units(name: str) -> str:
"""Get output variable units from iris.std_names.STD_NAMES"""
try:
units = STD_NAMES[name]["canonical_units"]
except KeyError:
raise ValueError("Units of {} are not known.".format(name))
return units
[docs]
def _create_time_bounds(time: datetime, time_period: int) -> Tuple[datetime, datetime]:
"""Create time bounds using time - time_period as the lower bound and time as the
upper bound"""
lower_bound = time - timedelta(minutes=time_period)
upper_bound = time
return (lower_bound, upper_bound)
[docs]
def _create_data_array(
ensemble_members: int,
leading_dimension: Optional[List[float]],
npoints: int,
vertical_levels: Optional[List[float]],
) -> ndarray:
"""Create data array of specified shape filled with zeros"""
if leading_dimension is not None:
nleading_dimension = len(leading_dimension)
else:
nleading_dimension = ensemble_members
if vertical_levels is not None:
nvertical_levels = len(vertical_levels)
else:
nvertical_levels = None
data_shape = []
if nleading_dimension > 1:
data_shape.append(nleading_dimension)
if nvertical_levels is not None:
data_shape.append(nvertical_levels)
data_shape.append(npoints)
data_shape.append(npoints)
return np.zeros(data_shape, dtype=np.float32)
[docs]
def generate_metadata(
mandatory_attributes: dict,
name: str = "air_pressure_at_sea_level",
units: Optional[str] = None,
time_period: Optional[int] = None,
ensemble_members: int = 8,
leading_dimension: Optional[List[float]] = None,
cube_type: str = "variable",
spp__relative_to_threshold: str = "greater_than",
npoints: int = 71,
**kwargs: Any,
) -> Cube:
"""Generate a cube with metadata only.
Args:
mandatory_attributes:
Specifies the values of the mandatory attributes, title, institution and
source.
name:
Output variable name, or if creating a probability cube the name of the
underlying variable to which the probability field applies.
units:
Output variable units, or if creating a probability cube the units of the
underlying variable / threshold.
time_period:
The period in minutes between the time bounds. This is used to calculate
the lower time bound. If unset the diagnostic will be instantaneous, i.e.
without time bounds.
ensemble_members:
Number of ensemble members. Default 8, unless percentile or probability set
to True.
leading_dimension:
List of realizations, percentiles or thresholds.
cube_type:
The type of cube to be generated. Permitted values are "variable",
"percentile" or "probability".
spp__relative_to_threshold:
Value of the attribute "spp__relative_to_threshold" which is required for
IMPROVER probability cubes.
npoints:
Number of points along each of the y and x spatial axes.
**kwargs:
Additional keyword arguments to pass to the required cube setup function.
Raises:
ValueError:
If any options are not supported
KeyError:
If mandatory_attributes does not contain all the required keys
Returns:
Output of set_up_variable_cube(), set_up_percentile_cube() or
set_up_probability_cube()
"""
if cube_type not in CUBE_TYPES:
raise ValueError(
(
"Cube type {} not supported. "
'Specify one of "variable", "percentile" or "probability".'
).format(cube_type)
)
if "spatial_grid" in kwargs and kwargs["spatial_grid"] not in (
"latlon",
"equalarea",
):
raise ValueError(
"Spatial grid {} not supported. Specify either latlon or equalarea.".format(
kwargs["spatial_grid"]
)
)
if (
"domain_corner" in kwargs
and kwargs["domain_corner"] is not None
and len(kwargs["domain_corner"]) != 2
):
raise ValueError("Domain corner must be a list or tuple of length 2.")
if units is None:
units = _get_units(name)
# If time_period specified, create time bounds using time as upper bound
if time_period is not None:
if "time" not in kwargs:
kwargs["time"] = DEFAULT_TIME
time_bounds = _create_time_bounds(kwargs["time"], time_period)
kwargs["time_bounds"] = time_bounds
# If grid_spacing not specified, use default for requested spatial grid
for spacing_axis in ["x_grid_spacing", "y_grid_spacing"]:
if spacing_axis not in kwargs or kwargs[spacing_axis] is None:
if "spatial_grid" not in kwargs:
kwargs["spatial_grid"] = DEFAULT_SPATIAL_GRID
kwargs[spacing_axis] = DEFAULT_GRID_SPACING[kwargs["spatial_grid"]]
# Create ndimensional array of zeros
if "vertical_levels" not in kwargs:
kwargs["vertical_levels"] = None
data = _create_data_array(
ensemble_members, leading_dimension, npoints, kwargs["vertical_levels"]
)
missing_mandatory_attributes = MANDATORY_ATTRIBUTES - mandatory_attributes.keys()
if missing_mandatory_attributes:
raise KeyError(
f"No values for these mandatory attributes: {missing_mandatory_attributes}"
)
if "attributes" in kwargs:
kwargs["attributes"] = kwargs["attributes"].copy()
else:
kwargs["attributes"] = {}
kwargs["attributes"].update(mandatory_attributes)
# Set up requested cube
if cube_type == "percentile":
metadata_cube = set_up_percentile_cube(
data, percentiles=leading_dimension, name=name, units=units, **kwargs
)
elif cube_type == "probability":
metadata_cube = set_up_probability_cube(
data,
leading_dimension,
variable_name=name,
threshold_units=units,
spp__relative_to_threshold=spp__relative_to_threshold,
**kwargs,
)
else:
metadata_cube = set_up_variable_cube(
data, name=name, units=units, realizations=leading_dimension, **kwargs
)
metadata_cube = squeeze(metadata_cube)
return metadata_cube