Merge branch 'lama-to-dask' into dask-contains

cf/data/dask_utils.py

@@ -9,7 +9,9 @@
 
 import dask.array as da
 import numpy as np
+from dask.core import flatten
 
+from ..cfdatetime import dt2rt, rt2dt
 from ..functions import atol as cf_atol
 from ..functions import rtol as cf_rtol
 
@@ -82,7 +84,13 @@ def allclose(a_blocks, b_blocks, rtol=rtol, atol=atol):
         if not isinstance(b_blocks, list):
             b_blocks = (b_blocks,)
 
-        for a, b in zip(a_blocks, b_blocks):
+        # Note: If a_blocks or b_blocks has more than one chunk in
+        #       more than one dimension they will comprise a nested
+        #       sequence of sequences, that needs to be flattened so
+        #       that we can safely iterate through the actual numpy
+        #       array elements.
+
+        for a, b in zip(flatten(a_blocks), flatten(b_blocks)):
             result &= np.ma.allclose(
                 a,
                 b,
@@ -134,7 +142,7 @@ def cf_contains(a, value):
 
 
 try:
-    from scipy.ndimage.filters import convolve1d
+    from scipy.ndimage import convolve1d
 except ImportError:
     pass
 
@@ -465,3 +473,116 @@ def cf_where(array, condition, x, y, hardmask):
         array.harden_mask()
 
     return array
+
+
+def _getattr(x, attr):
+    return getattr(x, attr, False)
+
+
+_array_getattr = np.vectorize(_getattr, excluded="attr")
+
+
+def cf_YMDhms(a, attr):
+    """Return a date-time component from an array of date-time objects.
+
+    Only applicable for data with reference time units. The returned
+    array will have the same mask hardness as the original array.
+
+    .. versionadded:: TODODASK
+
+    .. seealso:: `~cf.Data.year`, ~cf.Data.month`, `~cf.Data.day`,
+                 `~cf.Data.hour`, `~cf.Data.minute`, `~cf.Data.second`
+
+    :Parameters:
+
+        a: `numpy.ndarray`
+            The array from which to extract date-time component.
+
+        attr: `str`
+            The name of the date-time component, one of ``'year'``,
+            ``'month'``, ``'day'``, ``'hour'``, ``'minute'``,
+            ``'second'``.
+
+    :Returns:
+
+        `numpy.ndarray`
+            The date-time component.
+
+    **Examples**
+
+    >>> import numpy as np
+    >>> a = np.array([
+    ...  cftime.DatetimeGregorian(2000, 1, 1, 0, 0, 0, 0, has_year_zero=False)
+    ...  cftime.DatetimeGregorian(2000, 1, 2, 0, 0, 0, 0, has_year_zero=False)
+    ... ])
+    >>> cf_YMDmhs(a, 'day')
+    array([1, 2])
+
+    """
+    return _array_getattr(a, attr=attr)
+
+
+def cf_rt2dt(a, units):
+    """Convert an array of reference times to date-time objects.
+
+    .. versionadded:: TODODASK
+
+    .. seealso:: `cf._dt2rt`, `cf.Data._asdatetime`
+
+    :Parameters:
+
+        a: `numpy.ndarray`
+            An array of numeric reference times.
+
+        units: `Units`
+            The units for the reference times
+
+
+    :Returns:
+
+        `numpy.ndarray`
+            A array containing date-time objects.
+
+    **Examples**
+
+    >>> import numpy as np
+    >>> print(cf_rt2dt(np.array([0, 1]), cf.Units('days since 2000-01-01')))
+    [cftime.DatetimeGregorian(2000, 1, 1, 0, 0, 0, 0, has_year_zero=False)
+     cftime.DatetimeGregorian(2000, 1, 2, 0, 0, 0, 0, has_year_zero=False)]
+
+    """
+    return rt2dt(a, units_in=units)
+
+
+def cf_dt2rt(a, units):
+    """Convert an array of date-time objects to reference times.
+
+    .. versionadded:: TODODASK
+
+    .. seealso:: `cf._rt2dt`, `cf.Data._asreftime`
+
+    :Parameters:
+
+        a: `numpy.ndarray`
+            An array of date-time objects.
+
+        units: `Units`
+            The units for the reference times
+
+    :Returns:
+
+        `numpy.ndarray`
+            An array containing numeric reference times
+
+    **Examples**
+
+    >>> import numpy as np
+    >>> a = np.array([
+    ...  cftime.DatetimeGregorian(2000, 1, 1, 0, 0, 0, 0, has_year_zero=False)
+    ...  cftime.DatetimeGregorian(2000, 1, 2, 0, 0, 0, 0, has_year_zero=False)
+    ... ])
+    >>> print(cf_dt2rt(a, cf.Units('days since 1999-01-01')))
+    [365 366]
+
+    """
+    return dt2rt(a, units_out=units, units_in=None)