Merge pull request #405 from davidhassell/dask-Units

dask: `Data.Units`

Author: davidhassell

cf/data/dask_utils.py

@@ -14,8 +14,7 @@
 from ..cfdatetime import dt2rt, rt2dt
 from ..functions import atol as cf_atol
 from ..functions import rtol as cf_rtol
-
-# from dask.utils import deepmap  # Apply function inside nested lists
+from ..units import Units
 
 
 def _da_ma_allclose(x, y, masked_equal=True, rtol=None, atol=None):
@@ -594,3 +593,45 @@ def cf_dt2rt(a, units):
 
     """
     return dt2rt(a, units_out=units, units_in=None)
+
+
+def cf_units(a, from_units, to_units):
+    """Convert array values to have different equivalent units.
+
+    .. versionadded:: TODODASK
+
+    .. seealso:: `cf.Data.Units`
+
+    :Parameters:
+
+        a: `numpy.ndarray`
+            The array.
+
+        from_units: `Units`
+            The existing units of the array.
+
+        to_units: `Units`
+            The units that the array should be converted to. Must be
+            equivalent to *from_units*.
+
+    :Returns:
+
+        `numpy.ndarray`
+            An array containing values in the new units. In order to
+            represent the new units, the returned data type may be
+            different from that of the input array. For instance, if
+            *a* has an integer data type, *from_units* are kilometres,
+            and *to_units* are ``'miles'`` then the returned array
+            will have a float data type.
+
+    **Examples**
+
+    >>> import numpy as np
+    >>> a = np.array([1, 2])
+    >>> print(cf.data.dask_utils.cf_units(a, cf.Units('km'), cf.Units('m')))
+    [1000. 2000.]
+
+    """
+    return Units.conform(
+        a, from_units=from_units, to_units=to_units, inplace=False
+    )

cf/data/data.py

@@ -50,6 +50,7 @@
     cf_percentile,
     cf_rt2dt,
     cf_soften_mask,
+    cf_units,
     cf_where,
 )
 from .mixin import DataClassDeprecationsMixin
@@ -4699,11 +4700,8 @@ def Units(self, value):
                     "Consider using the override_units method instead."
                 )
 
-            if not old_units:
-                self.override_units(value, inplace=True)
-                return
-
-            if self.Units.equals(value):
+            if not old_units or self.Units.equals(value):
+                self._Units = value
                 return
 
         dtype = self.dtype
@@ -4713,13 +4711,11 @@ def Units(self, value):
             else:
                 dtype = _dtype_float
 
-        def cf_Units(x):
-            return Units.conform(
-                x=x, from_units=old_units, to_units=value, inplace=False
-            )
-
         dx = self.to_dask_array()
-        dx = dx.map_blocks(cf_Units, dtype=dtype)
+        dx = dx.map_blocks(
+            partial(cf_units, from_units=old_units, to_units=value),
+            dtype=dtype,
+        )
         self._set_dask(dx, reset_mask_hardness=False)
 
         self._Units = value

cf/test/test_Data.py

@@ -3961,6 +3961,22 @@ def test_Data_fits_in_memory(self):
         d = cf.Data.empty((size,), dtype=float)
         self.assertFalse(d.fits_in_memory())
 
+    def test_Data_Units(self):
+        d = cf.Data(100, "m")
+        self.assertEqual(d.Units, cf.Units("m"))
+
+        d.Units = cf.Units("km")
+        self.assertEqual(d.Units, cf.Units("km"))
+        self.assertEqual(d.array, 0.1)
+
+        # Assign non-equivalent units
+        with self.assertRaises(ValueError):
+            d.Units = cf.Units("watt")
+
+        # Delete units
+        with self.assertRaises(ValueError):
+            del d.Units
+
     def test_Data_get_data(self):
         d = cf.Data(9)
         self.assertIs(d, d.get_data())