Add unit support to cf-xarray (#197)

Co-authored-by: keewis <[email protected]>
Co-authored-by: dcherian <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>
Co-authored-by: Keewis <[email protected]>

Author: 3 people

.github/workflows/ci.yaml

@@ -48,6 +48,29 @@ jobs:
           name: codecov-umbrella
           fail_ci_if_error: false
 
+  no-optional-deps:
+    name: no-optional-deps
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          channels: conda-forge
+          mamba-version: "*"
+          activate-environment: cf_xarray_test
+          auto-update-conda: false
+          python-version: ${{ matrix.python-version }}
+      - name: Set up conda environment
+        shell: bash -l {0}
+        run: |
+          mamba env update -f ci/environment-no-optional-deps.yml
+          python -m pip install -e .
+          conda list
+      - name: Run Tests
+        shell: bash -l {0}
+        run: |
+          pytest -n 2
+
   upstream-dev:
     name: upstream-dev
     runs-on: ubuntu-latest

cf_xarray/tests/test_units.py

@@ -0,0 +1,60 @@
+r"""Tests the operation of cf_xarray's ported unit support code.
+
+Reused with modification from MetPy under the terms of the BSD 3-Clause License.
+Copyright (c) 2017 MetPy Developers.
+"""
+
+import pytest
+
+pytest.importorskip("pint")
+
+from ..units import units
+
+
+def test_added_degrees_units():
+    """Test that our added degrees units are present in the registry."""
+    # Test equivalence of abbreviations/aliases to our defined names
+    assert str(units("degrees_N").units) == "degrees_north"
+    assert str(units("degreesN").units) == "degrees_north"
+    assert str(units("degree_north").units) == "degrees_north"
+    assert str(units("degree_N").units) == "degrees_north"
+    assert str(units("degreeN").units) == "degrees_north"
+    assert str(units("degrees_E").units) == "degrees_east"
+    assert str(units("degreesE").units) == "degrees_east"
+    assert str(units("degree_east").units) == "degrees_east"
+    assert str(units("degree_E").units) == "degrees_east"
+    assert str(units("degreeE").units) == "degrees_east"
+
+    # Test equivalence of our defined units to base units
+    assert units("degrees_north") == units("degrees")
+    assert units("degrees_north").to_base_units().units == units.radian
+    assert units("degrees_east") == units("degrees")
+    assert units("degrees_east").to_base_units().units == units.radian
+
+
+def test_gpm_unit():
+    """Test that the gpm unit does alias to meters."""
+    x = 1 * units("gpm")
+    assert str(x.units) == "meter"
+
+
+def test_psu_unit():
+    """Test that the psu unit are present in the registry."""
+    x = 1 * units("psu")
+    assert str(x.units) == "practical_salinity_unit"
+
+
+def test_percent_units():
+    """Test that percent sign units are properly parsed and interpreted."""
+    assert str(units("%").units) == "percent"
+
+
+@pytest.mark.xfail(reason="not supported by pint, yet: hgrecco/pint#1295")
+def test_udunits_power_syntax():
+    """Test that UDUNITS style powers are properly parsed and interpreted."""
+    assert units("m2 s-2").units == units.m ** 2 / units.s ** 2
+
+
+def test_udunits_power_syntax_parse_units():
+    """Test that UDUNITS style powers are properly parsed and interpreted."""
+    assert units.parse_units("m2 s-2") == units.m ** 2 / units.s ** 2

cf_xarray/units.py

@@ -0,0 +1,58 @@
+r"""Module to provide unit support via pint approximating UDUNITS/CF.
+
+Reused with modification from MetPy under the terms of the BSD 3-Clause License.
+Copyright (c) 2015,2017,2019 MetPy Developers.
+"""
+import functools
+import re
+import warnings
+
+import pint
+from pint import (  # noqa: F401
+    DimensionalityError,
+    UndefinedUnitError,
+    UnitStrippedWarning,
+)
+
+# Create registry, with preprocessors for UDUNITS-style powers (m2 s-2) and percent signs
+units = pint.UnitRegistry(
+    autoconvert_offset_to_baseunit=True,
+    preprocessors=[
+        functools.partial(
+            re.compile(
+                r"(?<=[A-Za-z])(?![A-Za-z])(?<![0-9\-][eE])(?<![0-9\-])(?=[0-9\-])"
+            ).sub,
+            "**",
+        ),
+        lambda string: string.replace("%", "percent"),
+    ],
+    force_ndarray_like=True,
+)
+
+units.define(
+    pint.unit.UnitDefinition("percent", "%", (), pint.converters.ScaleConverter(0.01))
+)
+
+# Define commonly encoutered units (both CF and non-CF) not defined by pint
+units.define(
+    "degrees_north = degree = degrees_N = degreesN = degree_north = degree_N = degreeN"
+)
+units.define(
+    "degrees_east = degree = degrees_E = degreesE = degree_east = degree_E = degreeE"
+)
+units.define("@alias meter = gpm")
+units.define("practical_salinity_unit = [] = psu")
+
+# Enable pint's built-in matplotlib support
+try:
+    units.setup_matplotlib()
+except ImportError:
+    warnings.warn(
+        "Import(s) unavailable to set up matplotlib support...skipping this portion "
+        "of the setup."
+    )
+
+# Set as application registry
+pint.set_application_registry(units)
+
+del pint

ci/environment-no-optional-deps.yml

@@ -0,0 +1,13 @@
+name: cf_xarray_test
+channels:
+  - conda-forge
+dependencies:
+  - pytest-cov
+  - pytest
+  - pytest-xdist
+  - dask
+  - matplotlib-base
+  - netcdf4
+  - pandas
+  - pooch
+  - xarray

ci/environment.yml

@@ -9,5 +9,6 @@ dependencies:
   - matplotlib-base
   - netcdf4
   - pandas
+  - pint
   - pooch
   - xarray

ci/upstream-dev-env.yml

@@ -12,3 +12,4 @@ dependencies:
   - pooch
   - pip:
     - git+https://github.com/pydata/xarray
+    - git+https://github.com/hgrecco/pint

doc/whats-new.rst

@@ -3,6 +3,11 @@
 What's New
 ----------
 
+v0.5.3 (unreleased)
+===================
+- Begin adding support for units with a unit registry for pint arrays. :pr:`197`.
+  By `Jon Thielen`_ and `Justus Magin`_.
+
 v0.5.2 (May 11, 2021)
 =====================
 
@@ -102,6 +107,8 @@ v0.1.3
 - Support expanding key to multiple dimension names.
 
 .. _`Mattia Almansi`: https://github.com/malmans2
+.. _`Justus Magin`: https://github.com/keewis
+.. _`Jon Thielen`: https://github.com/jthielen
 .. _`Anderson Banihirwe`: https://github.com/andersy005
 .. _`Pascal Bourgault`: https://github.com/aulemahal
 .. _`Deepak Cherian`: https://github.com/dcherian

setup.cfg

@@ -41,7 +41,7 @@ skip_gitignore = true
 force_to_top = true
 default_section = THIRDPARTY
 known_first_party = cf_xarray
-known_third_party = dask,matplotlib,numpy,pandas,pkg_resources,pytest,setuptools,sphinx_autosummary_accessors,xarray
+known_third_party = dask,matplotlib,numpy,pandas,pint,pkg_resources,pytest,setuptools,sphinx_autosummary_accessors,xarray
 
 # Most of the numerical computing stack doesn't have type annotations yet.
 [mypy-affine.*]