Quantify/Dequantify working & tested on DataArrays

Author: TomNicholas

pintxarray/accessors.py

@@ -1,6 +1,6 @@
 # TODO is it possible to import pint-xarray from within xarray if pint is present?
 from xarray import (register_dataarray_accessor, register_dataset_accessor,
-                    DataArray, Dataset)
+                    DataArray, Dataset, Variable)
 from xarray.core.npcompat import IS_NEP18_ACTIVE
 
 import numpy as np
@@ -69,6 +69,13 @@ def _array_attach_units(data, unit, convert_from=None):
     return quantity
 
 
+def _dequantify_variable(var):
+    new_var = Variable(dims=var.dims, data=var.data.magnitude,
+                  attrs=var.attrs)
+    new_var.attrs['units'] = str(var.data.units)
+    return new_var
+
+
 @register_dataarray_accessor("pint")
 class PintDataArrayAccessor:
     """
@@ -115,17 +122,18 @@ def quantify(self, units=None, unit_registry=None, registry_kwargs=None):
         * wavelength  (wavelength) np.array 1e-4, 2e-4, 4e-4, 6e-4, 1e-3, 2e-3
         """
 
+        # TODO should also quantify coordinates (once explicit indexes ready)
+
         if isinstance(self.da.data, Quantity):
-            raise ValueError
+            raise ValueError(f"Cannot attach unit {units} to quantity: data "
+                             f"already has units {self.da.data.units}")
 
         if unit_registry is None:
             if registry_kwargs is None:
                 registry_kwargs = {}
+            registry_kwargs.update(force_ndarray=True)
+            # TODO should this registry object then be stored somewhere global?
             unit_registry = pint.UnitRegistry(**registry_kwargs)
-        else:
-            if registry_kwargs is not None:
-                raise ValueError("Cannot supply registry kwargs without "
-                                 "supplying a registry")
 
         if units is None:
             # TODO option to read and decode units according to CF conventions (see MetPy)?
@@ -134,6 +142,7 @@ def quantify(self, units=None, unit_registry=None, registry_kwargs=None):
         elif isinstance(units, Unit):
             # TODO do we have to check what happens if someone passes a Unit instance
             # without creating a unit registry?
+            # TODO and what happens if they pass in a Unit from a different registry
             pass
         else:
             units = unit_registry.Unit(units)
@@ -142,13 +151,29 @@ def quantify(self, units=None, unit_registry=None, registry_kwargs=None):
 
         # TODO should we (temporarily) remove the attrs here so that they don't become inconsistent?
         return DataArray(dims=self.da.dims, data=quantity,
-                         coords=self.da.coords, attrs=self.da.attrs )
+                         coords=self.da.coords, attrs=self.da.attrs)
+
+    def dequantify(self):
+        """
+        Removes units from the DataArray and it's coordinates.
+
+        Will replace `.attrs['units']` on each variable with a string
+        representation of the `pint.Unit` instance.
 
-    def dequantify(self, encode_cf=True):
-        da = DataArray(dim=self.da.dims, data=self.da.pint.magnitude,
-                       coords=self.da.coords, attrs=self.da.attrs,
-                       encoding=self.da.encoding)
-        da.attrs['units'] = self.da.pint.units
+        Returns
+        -------
+        dequantified - DataArray whose array data is unitless, and of the type
+        that was previously wrapped by `pint.Quantity`.
+        """
+
+        if not isinstance(self.da.data, Quantity):
+            raise ValueError("Cannot remove units from data that does not have"
+                             " units")
+
+        # TODO also dequantify coords (once explicit indexes ready)
+        da = DataArray(dims=self.da.dims, data=self.da.pint.magnitude,
+                       coords=self.da.coords, attrs=self.da.attrs)
+        da.attrs['units'] = str(self.da.data.units)
         return da
 
     @property
@@ -173,6 +198,15 @@ def units(self, units):
     def dimensionality(self):
         return self.da.data.dimensionality
 
+    @property
+    def registry(self):
+        # TODO is this a bad idea? (see GH issue #1071 in pint)
+        return self.data._REGISTRY
+
+    @registry.setter
+    def registry(self, _):
+        raise AttributeError("Don't try to change the registry once created")
+
     def to(self, units):
         quantity = self.da.data.to(units)
         return DataArray(dim=self.da.dims, data=quantity,

pintxarray/tests/test_accessors.py

@@ -8,10 +8,10 @@
 
 from pint import UnitRegistry
 from pint.unit import Unit
-from pint.errors import DimensionalityError
+from pint.errors import UndefinedUnitError, DimensionalityError
 
 from pintxarray.accessors import PintDataArrayAccessor, PintDatasetAccessor
-from .utils import extract_units
+from .utils import extract_units, raises_regex
 
 
 # make sure scalars are converted to 0d arrays so quantities can
@@ -37,44 +37,70 @@ def example_quantity_da():
 
 
 class TestQuantifyDataArray:
-    @pytest.mark.skip(reason="Not yet implemented")
-    def test_attach_units_given_registry(self):
-        ...
+    def test_attach_units_given_registry(self, example_unitless_da):
+        orig = example_unitless_da
+        ureg = UnitRegistry(force_ndarray=True)
+        result = orig.pint.quantify('m', unit_registry=ureg)
+        assert_array_equal(result.data.magnitude, orig.data)
+        assert result.data.units == ureg.Unit('m')
 
     def test_attach_units(self, example_unitless_da):
         orig = example_unitless_da
         result = orig.pint.quantify('m')
         assert_array_equal(result.data.magnitude, orig.data)
+        # TODO better comparisons for when you can't access the unit_registry?
         assert str(result.data.units) == 'meter'
 
-    def test_attach_units_from_attrs(self):
-        ...
+    def test_attach_units_from_attrs(self, example_unitless_da):
+        orig = example_unitless_da
+        result = orig.pint.quantify()
+        assert_array_equal(result.data.magnitude, orig.data)
+        assert str(result.data.units) == 'meter'
 
-    def test_attach_unit_class(self):
-        ...
+    def test_attach_unit_class(self, example_unitless_da):
+        orig = example_unitless_da
+        ureg = UnitRegistry(force_ndarray=True)
+        result = orig.pint.quantify(ureg.Unit('m'), unit_registry=ureg)
+        assert_array_equal(result.data.magnitude, orig.data)
+        assert result.data.units == ureg.Unit('m')
 
     def test_error_when_already_units(self, example_quantity_da):
         da = example_quantity_da
-        with pytest.raises(ValueError):
+        with raises_regex(ValueError, "already has units"):
             da.pint.quantify()
 
-    def test_error_on_nonsense_units(self):
-        ...
+    def test_error_on_nonsense_units(self, example_unitless_da):
+        da = example_unitless_da
+        with pytest.raises(UndefinedUnitError):
+            da.pint.quantify(units='aecjhbav')
 
-    def test_registry_kwargs(self):
-        ...
+    def test_registry_kwargs(self, example_unitless_da):
+        orig = example_unitless_da
+        result = orig.pint.quantify(registry_kwargs=
+                                    {'auto_reduce_dimensions': True})
+        assert(result.data._REGISTRY.auto_reduce_dimensions) == True
 
 
-@pytest.mark.skip(reason="Not yet implemented")
 class TestDequantifyDataArray:
-    def test_strip_units(self):
-        ...
+    def test_strip_units(self, example_quantity_da):
+        result = example_quantity_da.pint.dequantify()
+        assert isinstance(result.data, np.ndarray)
+        assert isinstance(result.coords['x'].data, np.ndarray)
 
-    def test_error_if_no_units(self):
-        ...
+    def test_error_if_no_units(self, example_unitless_da):
+        with raises_regex(ValueError, "does not have units"):
+            example_unitless_da.pint.dequantify()
 
-    def test_roundtrip_attrs(self):
-        ...
+    def test_attrs_reinstated(self, example_quantity_da):
+        da = example_quantity_da
+        result = da.pint.dequantify()
+        assert result.attrs['units'] == 'meter'
+
+    def test_roundtrip_data(self, example_unitless_da):
+        orig = example_unitless_da
+        quantified = orig.pint.quantify()
+        result = quantified.pint.dequantify()
+        assert_equal(result, orig)
 
 
 @pytest.mark.skip(reason="Not yet implemented")

pintxarray/tests/utils.py

@@ -1,8 +1,24 @@
+from contextlib import contextmanager
+import re
+
+import pytest
+
 import xarray as xr
 
 from pint.quantity import Quantity
 
 
+@contextmanager
+def raises_regex(error, pattern):
+    __tracebackhide__ = True
+    with pytest.raises(error) as excinfo:
+        yield
+    message = str(excinfo.value)
+    if not re.search(pattern, message):
+        raise AssertionError(
+            f"exception {excinfo.value!r} did not match pattern {pattern!r}"
+        )
+
 def array_extract_units(obj):
     if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
         obj = obj.data