Skip to content

Commit

Permalink
Use our own unit conversion function in our fixes (#2560)
Browse files Browse the repository at this point in the history
  • Loading branch information
schlunma authored Nov 5, 2024
1 parent 80fd9a0 commit 7e16883
Show file tree
Hide file tree
Showing 6 changed files with 217 additions and 89 deletions.
4 changes: 2 additions & 2 deletions esmvalcore/cmor/_fixes/fix.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
)
from esmvalcore.cmor.fixes import get_time_bounds
from esmvalcore.cmor.table import get_var_info
from esmvalcore.iris_helpers import has_unstructured_grid
from esmvalcore.iris_helpers import has_unstructured_grid, safe_convert_units

if TYPE_CHECKING:
from esmvalcore.cmor.table import CoordinateInfo, VariableInfo
Expand Down Expand Up @@ -455,7 +455,7 @@ def _fix_units(self, cube: Cube) -> Cube:
if str(cube.units) != units:
old_units = cube.units
try:
cube.convert_units(units)
safe_convert_units(cube, units)
except (ValueError, UnitConversionError):
self._warning_msg(
cube,
Expand Down
8 changes: 2 additions & 6 deletions esmvalcore/cmor/_fixes/native6/era5.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import iris
import numpy as np

from esmvalcore.iris_helpers import date2num
from esmvalcore.iris_helpers import date2num, safe_convert_units

from ...table import CMOR_TABLES
from ..fix import Fix
Expand Down Expand Up @@ -414,10 +414,6 @@ def _fix_monthly_time_coord(cube):
coord.points = 0.5 * (start + end)
coord.bounds = np.column_stack([start, end])

def _fix_units(self, cube):
"""Fix units."""
cube.convert_units(self.vardef.units)

def fix_metadata(self, cubes):
"""Fix metadata."""
fixed_cubes = iris.cube.CubeList()
Expand All @@ -428,7 +424,7 @@ def fix_metadata(self, cubes):
cube.long_name = self.vardef.long_name

cube = self._fix_coordinates(cube)
self._fix_units(cube)
cube = safe_convert_units(cube, self.vardef.units)

cube.data = cube.core_data().astype("float32")
year = datetime.datetime.now().year
Expand Down
4 changes: 3 additions & 1 deletion esmvalcore/cmor/_fixes/native_datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@

from iris import NameConstraint

from esmvalcore.iris_helpers import safe_convert_units

from ..fix import Fix
from .shared import (
add_scalar_height_coord,
Expand Down Expand Up @@ -77,7 +79,7 @@ def fix_var_metadata(self, cube):
f"Failed to fix invalid units '{invalid_units}' for "
f"variable '{self.vardef.short_name}'"
) from exc
cube.convert_units(self.vardef.units)
safe_convert_units(cube, self.vardef.units)

# Fix attributes
if self.vardef.positive != "":
Expand Down
123 changes: 121 additions & 2 deletions esmvalcore/iris_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import iris.cube
import iris.util
import numpy as np
from cf_units import Unit
from iris.coords import Coord
from iris.cube import Cube
from iris.exceptions import CoordinateMultiDimError, CoordinateNotFoundError
Expand Down Expand Up @@ -92,8 +93,8 @@ def date2num(date, unit, dtype=np.float64):
This is a custom version of :meth:`cf_units.Unit.date2num` that
guarantees the correct dtype for the return value.
Arguments
---------
Parameters
----------
date : :class:`datetime.datetime` or :class:`cftime.datetime`
unit : :class:`cf_units.Unit`
dtype : a numpy dtype
Expand Down Expand Up @@ -358,3 +359,121 @@ def has_unstructured_grid(cube: Cube) -> bool:
if cube.coord_dims(lat) != cube.coord_dims(lon):
return False
return True


# List containing special cases for unit conversion. Each list item is another
# list. Each of these sublists defines one special conversion. Each element in
# the sublists is a tuple (standard_name, units). Note: All units for a single
# special case need to be "physically identical", e.g., 1 kg m-2 s-1 "equals" 1
# mm s-1 for precipitation
_SPECIAL_UNIT_CONVERSIONS = [
[
("precipitation_flux", "kg m-2 s-1"),
("lwe_precipitation_rate", "mm s-1"),
],
[
("equivalent_thickness_at_stp_of_atmosphere_ozone_content", "m"),
("equivalent_thickness_at_stp_of_atmosphere_ozone_content", "1e5 DU"),
],
]


def _try_special_unit_conversions(cube: Cube, units: str | Unit) -> bool:
"""Try special unit conversion (in-place).
Parameters
----------
cube:
Input cube (modified in place).
units:
New units
Returns
-------
bool
``True`` if special unit conversion was successful, ``False`` if not.
"""
for special_case in _SPECIAL_UNIT_CONVERSIONS:
for std_name, special_units in special_case:
# Special unit conversion only works if all of the following
# criteria are met:
# - the cube's standard_name is one of the supported
# standard_names
# - the cube's units are convertible to the ones defined for
# that given standard_name
# - the desired target units are convertible to the units of
# one of the other standard_names in that special case

# Step 1: find suitable source name and units
if cube.standard_name == std_name and cube.units.is_convertible(
special_units
):
for target_std_name, target_units in special_case:
if target_units == special_units:
continue

# Step 2: find suitable target name and units
if Unit(units).is_convertible(target_units):
cube.standard_name = target_std_name

# In order to avoid two calls to cube.convert_units,
# determine the conversion factor between the cube's
# units and the source units first and simply add this
# factor to the target units (remember that the source
# units and the target units should be "physically
# identical").
factor = cube.units.convert(1.0, special_units)
cube.units = f"{factor} {target_units}"
cube.convert_units(units)
return True

# If no special case has been detected, return False
return False


def safe_convert_units(cube: Cube, units: str | Unit) -> Cube:
"""Safe unit conversion (change of `standard_name` not allowed; in-place).
This is a safe version of :func:`esmvalcore.preprocessor.convert_units`
that will raise an error if the input cube's
:attr:`~iris.cube.Cube.standard_name` has been changed.
Parameters
----------
cube:
Input cube (modified in place).
units:
New units.
Returns
-------
iris.cube.Cube
Converted cube. Just returned for convenience; input cube is modified
in place.
Raises
------
iris.exceptions.UnitConversionError
Old units are unknown.
ValueError
Old units are not convertible to new units or unit conversion required
change of `standard_name`.
"""
old_units = cube.units
old_standard_name = cube.standard_name

try:
cube.convert_units(units)
except ValueError:
if not _try_special_unit_conversions(cube, units):
raise

if cube.standard_name != old_standard_name:
raise ValueError(
f"Cannot safely convert units from '{old_units}' to '{units}'; "
f"standard_name changed from '{old_standard_name}' to "
f"'{cube.standard_name}'"
)
return cube
108 changes: 31 additions & 77 deletions esmvalcore/preprocessor/_units.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,77 +11,23 @@
import iris
import numpy as np
from cf_units import Unit
from iris.coords import AuxCoord, DimCoord
from iris.cube import Cube

from esmvalcore.iris_helpers import _try_special_unit_conversions

logger = logging.getLogger(__name__)


# List containing special cases for convert_units. Each list item is another
# list. Each of these sublists defines one special conversion. Each element in
# the sublists is a tuple (standard_name, units). Note: All units for a single
# special case need to be "physically identical", e.g., 1 kg m-2 s-1 "equals" 1
# mm s-1 for precipitation
SPECIAL_CASES = [
[
("precipitation_flux", "kg m-2 s-1"),
("lwe_precipitation_rate", "mm s-1"),
],
[
("equivalent_thickness_at_stp_of_atmosphere_ozone_content", "m"),
("equivalent_thickness_at_stp_of_atmosphere_ozone_content", "1e5 DU"),
],
]


def _try_special_conversions(cube, units):
"""Try special conversion."""
for special_case in SPECIAL_CASES:
for std_name, special_units in special_case:
# Special unit conversion only works if all of the following
# criteria are met:
# - the cube's standard_name is one of the supported
# standard_names
# - the cube's units are convertible to the ones defined for
# that given standard_name
# - the desired target units are convertible to the units of
# one of the other standard_names in that special case

# Step 1: find suitable source name and units
if cube.standard_name == std_name and cube.units.is_convertible(
special_units
):
for target_std_name, target_units in special_case:
if target_units == special_units:
continue

# Step 2: find suitable target name and units
if Unit(units).is_convertible(target_units):
cube.standard_name = target_std_name

# In order to avoid two calls to cube.convert_units,
# determine the conversion factor between the cube's
# units and the source units first and simply add this
# factor to the target units (remember that the source
# units and the target units should be "physically
# identical").
factor = cube.units.convert(1.0, special_units)
cube.units = f"{factor} {target_units}"
cube.convert_units(units)
return True

# If no special case has been detected, return False
return False


def convert_units(cube, units):
"""Convert the units of a cube to new ones.
This converts units of a cube.
def convert_units(cube: Cube, units: str | Unit) -> Cube:
"""Convert the units of a cube to new ones (in-place).
Note
----
Allows special unit conversions which transforms one quantity to another
(physically related) quantity. These quantities are identified via their
``standard_name`` and their ``units`` (units convertible to the ones
(physically related) quantity, which may also change the input cube's
:attr:`~iris.cube.Cube.standard_name`. These quantities are identified via
their ``standard_name`` and their ``units`` (units convertible to the ones
defined are also supported). For example, this enables conversions between
precipitation fluxes measured in ``kg m-2 s-1`` and precipitation rates
measured in ``mm day-1`` (and vice versa).
Expand All @@ -103,32 +49,40 @@ def convert_units(cube, units):
Note that for precipitation variables, a water density of ``1000 kg m-3``
is assumed.
Arguments
---------
cube: iris.cube.Cube
Input cube.
units: str
New units in udunits form.
Parameters
----------
cube:
Input cube (modified in place).
units:
New units.
Returns
-------
iris.cube.Cube
converted cube.
Converted cube. Just returned for convenience; input cube is modified
in place.
Raises
------
iris.exceptions.UnitConversionError
Old units are unknown.
ValueError
Old units are not convertible to new units.
"""
try:
cube.convert_units(units)
except ValueError:
if not _try_special_conversions(cube, units):
if not _try_special_unit_conversions(cube, units):
raise

return cube


def accumulate_coordinate(
cube: iris.cube.Cube,
coordinate: str | iris.coords.DimCoord | iris.coords.AuxCoord,
) -> iris.cube.Cube:
cube: Cube,
coordinate: str | DimCoord | AuxCoord,
) -> Cube:
"""Weight data using the bounds from a given coordinate.
The resulting cube will then have units given by
Expand All @@ -137,15 +91,15 @@ def accumulate_coordinate(
Parameters
----------
cube:
Data cube for the flux
Data cube for the flux.
coordinate:
Name of the coordinate that will be used as weights.
Returns
-------
iris.cube.Cube
Cube with the aggregated data
Cube with the aggregated data.
Raises
------
Expand All @@ -169,7 +123,7 @@ def accumulate_coordinate(
)

array_module = da if coord.has_lazy_bounds() else np
factor = iris.coords.AuxCoord(
factor = AuxCoord(
array_module.diff(coord.core_bounds())[..., -1],
var_name=coord.var_name,
long_name=coord.long_name,
Expand Down
Loading

0 comments on commit 7e16883

Please sign in to comment.