Units

setup.py

@@ -37,6 +37,7 @@ def read(filename):
         "rich-click",
         "versioneer",
         "xarray",
+        "chemicals"
     ],
     development_requires=[
         "black",

src/pymorize/pipeline.py

@@ -197,4 +197,5 @@ def __init__(self):
         super().__init__(
             # FIXME: Fill in with appropriate steps
             get_callable_by_name("pymorize.generic.load_data"),
+            get_callable_by_name("pymorize.units.handle_unit_conversion"),
         )

src/pymorize/units.py

@@ -0,0 +1,75 @@
+"""
+This module deals with the auto-unit conversion in the cmorization process.
+In case the units in model files differ from CMIP Tables, this module attempts to
+convert them automatically.
+
+In case of missing units in either model files or CMIP Tables, this module can
+not convert from a dimentionless base to something with dimension. Dealing with
+such thing have to done with `action` section in the Rules module on a per
+variable basis.
+"""
+
+import re
+from typing import Pattern, Union
+
+import cf_xarray.units
+import pint_xarray
+import xarray as xr
+from chemicals import periodic_table
+from loguru import logger
+
+ureg = pint_xarray.unit_registry
+
+
+def handle_chemicals(
+    s: Union[str, None] = None, pattern: Pattern = re.compile(r"mol(?P<symbol>\w+)")
+):
+    """Registers known chemical elements definitions to global ureg (unit registry)"""
+    if s is None:
+        return
+    match = pattern.search(s)
+    if match:
+        d = match.groupdict()
+        try:
+            element = getattr(periodic_table, d["symbol"])
+        except AttributeError:
+            raise ValueError(
+                f"Unknown chemical element {d.groupdict()['symbol']} in {d.group()}"
+            )
+        else:
+            try:
+                ureg(s)
+            except pint_xarray.pint.errors.UndefinedUnitError:
+                logger.debug(f"Chemical element {element.name} detected in units {s}.")
+                logger.debug(
+                    f"Registering definition: {match.group()} = {element.MW} * g"
+                )
+                ureg.define(f"{match.group()} = {element.MW} * g")
+
+
+def handle_unit_conversion(
+    da: xr.DataArray, unit: str, source_unit: Union[str, None] = None
+) -> xr.DataArray:
+    """Performs the unit-aware data conversion.
+
+    If `source_unit` is provided, it is used instead of the unit from DataArray.
+
+    Parameters:
+    -----------
+    da: xr.DataArray
+    unit: unit to convert data to
+    source_unit: Override the unit on xr.DataArray if needed.
+    """
+    from_unit = da.attrs.get("units")
+    if source_unit is not None:
+        logger.debug(
+            f"using user defined unit ({source_unit}) instead of ({from_unit}) from DataArray "
+        )
+        from_unit = source_unit
+    handle_chemicals(from_unit)
+    handle_chemicals(unit)
+    new_da = da.pint.quantify(from_unit)
+    new_da = new_da.pint.to(unit).pint.dequantify()
+    logger.debug(f"setting units on DataArray: {unit}")
+    new_da.attrs["units"] = unit
+    return new_da

tests/test_units.py

@@ -0,0 +1,142 @@
+import numpy as np
+import pint
+import pytest
+import xarray as xr
+from chemicals import periodic_table
+
+from pymorize.units import handle_chemicals, handle_unit_conversion, ureg
+
+#  input samples that are found in CMIP6 tables and in fesom1 (recom)
+allunits = [
+    "%",
+    "0.001",
+    "1",
+    "1.e6 J m-1 s-1",
+    "1e-06",
+    "1e-3 kg m-2",
+    "1e3 km3",
+    "J m-2",
+    "K",
+    "K Pa s-1",
+    "K s-1",
+    "K2",
+    "Pa2 s-2",
+    "W m^-2",
+    "W/m2",
+    "W/m^2",
+    "day",
+    "degC",
+    "degC kg m-2",
+    "degC2",
+    "degree",
+    "degrees_east",
+    "degrees_north",
+    "kg kg-1",
+    "kg m-2 s-1",
+    "kg m-3",
+    "kg s-1",
+    "km-2 s-1",
+    "m-1 sr-1",
+    "m-2",
+    "m^-3",
+    "m^2",
+    "mol/kg",
+    "mol/m2",
+    "mol m-2",
+    "mol m^-2",
+    "(mol/kg) / atm",
+    "mmol/m2/d",
+    "uatm",
+    "year",
+    "yr",
+]
+
+
+@pytest.mark.parametrize("test_input", allunits)
+def test_can_read_units(test_input):
+    ureg(test_input)
+
+
+units_with_chemical_element = [
+    "mmolC/(m2*d)",
+    "mmolC/d",
+    "mmolC/m2/d",
+    "mmolN/(m2*d)",
+    "mmolN/d",
+    "umolFe/m2/s",
+]
+
+
+@pytest.mark.parametrize("test_input", units_with_chemical_element)
+def test_handle_chemicals(test_input):
+    """Ensures the unit registry can add new units when parsed by ``handle_chemicals``."""
+    handle_chemicals(test_input)
+    ureg(test_input)
+
+
+def test_can_handle_simple_chemical_elements():
+    from_unit = "molC"
+    to_unit = "g"
+    da = xr.DataArray(10, attrs={"units": from_unit})
+    new_da = handle_unit_conversion(da, to_unit)
+    assert new_da.data == np.array(periodic_table.Carbon.MW * 10)
+    assert new_da.attrs["units"] == to_unit
+
+
+def test_can_handle_chemical_elements():
+    from_unit = "mmolC/m2/d"
+    to_unit = "kg m-2 s-1"
+    da = xr.DataArray(10, attrs={"units": from_unit})
+    new_da = handle_unit_conversion(da, to_unit)
+    assert np.allclose(new_da.data, np.array(1.39012731e-09))
+    assert new_da.attrs["units"] == to_unit
+
+
+def test_user_defined_units_takes_precedence_over_units_in_dataarray():
+    from_unit = "molC"
+    to_unit = "g"
+    da = xr.DataArray(10, attrs={"units": "kg"})
+    # here, "molC" will be used instead of "kg"
+    new_da = handle_unit_conversion(da, to_unit, from_unit)
+    assert new_da.data == np.array(periodic_table.Carbon.MW * 10)
+    assert new_da.attrs["units"] == to_unit
+
+
+def test_without_defining_uraninum_to_weight_conversion_raises_error():
+    """Checks that only elements we added are defined"""
+    with pytest.raises(pint.errors.UndefinedUnitError):
+        ureg("mmolU/m**2/d")
+
+
+def test_recognizes_previous_defined_chemical_elements():
+    assert "mmolC/m^2/d" in ureg
+
+
+def test_works_when_both_units_are_None():
+    to_unit = None
+    da = xr.DataArray(10, attrs={"units": None})
+    new_da = handle_unit_conversion(da, to_unit)
+    assert new_da.attrs["units"] == to_unit
+
+
+def test_works_when_both_units_are_empty_string():
+    to_unit = ""
+    da = xr.DataArray(10, attrs={"units": ""})
+    new_da = handle_unit_conversion(da, to_unit)
+    assert new_da.attrs["units"] == to_unit
+
+
+@pytest.mark.parametrize("from_unit", ["m/s", None, ""])
+def test_when_target_units_is_None_overrides_existing_units(from_unit):
+    to_unit = None
+    da = xr.DataArray(10, attrs={"units": from_unit})
+    new_da = handle_unit_conversion(da, to_unit)
+    assert new_da.attrs["units"] == to_unit
+
+
+@pytest.mark.parametrize("from_unit", ["m/s", None])
+def test_when_tartget_unit_is_empty_string_raises_error(from_unit):
+    to_unit = ""
+    da = xr.DataArray(10, attrs={"units": from_unit})
+    with pytest.raises(ValueError):
+        handle_unit_conversion(da, to_unit)