Use da.map_blocks : Tidy and rename.

lib/iris/_lazy_data.py

@@ -160,61 +160,26 @@ def multidim_lazy_stack(stack):
     return result
 
 
-class _ArraylikeWithPointwise(object):
+def wrap_lazy_elementwise(lazy_array, elementwise_op):
     """
-    An Array-like wrapper to a dask array which applies a pointwise array
-    operation.
+    Apply a (numpy-style) elementwise array operation to a lazy array.
 
-    The purpose of this is to apply a pointwise operation that works on numpy
-    arrays but cannot use dask arrays.
-
-    The advantage over dask.delayed operation is that an array-like interface
-    allows us to fetch and operate on only a required subarray of the input.
-
-    The downside is that its reference to the 'dask_array' argument is hidden
-    behind a compute() call within the __getitem__ method :  Because this
-    reference is not visible as part of any larger dask graph, it cannot be
-    efficiently combined with any other accesses to 'dask_array'.
-    This could lead to unnecessary re-computes of 'dask_array'.
-
-    """
-    def __init__(self, dask_array, pointwise_function):
-        """
-        Create an array-like from a dask array
-
-        Args:
-
-        * dask_array:
-            The dask array object to operate on.
-        * pointwise_function:
-            The pointwise (numpy) array operation to apply.
-
-        """
-        self._dask_array = dask_array
-        self._op = pointwise_function
-        self.shape = dask_array.shape
-        self.dtype = dask_array.dtype
-
-    def __getitem__(self, keys):
-        inner_result = self._dask_array[keys].compute()
-        return self._op(inner_result)
-
-
-def lazy_wrap_pointwise(lazy_array, pointwise_op):
-    """
-    Apply a (numpy-style) pointwise array operation to a lazy array.
-
-    See _ArraylikeWithPointwise for pros and cons of the method used.
+    Elementwise means that it performs a independent calculation at each point
+    of the input, producing a result array of the same shape.
 
     Args:
 
     * lazy_array:
         The lazy array object to operate on.
-    * pointwise_op:
-        The pointwise (numpy) array operation to apply.
+    * elementwise_op:
+        The elementwise (numpy) array operation to apply.
 
     """
-    pointwise_arraylike = _ArraylikeWithPointwise(lazy_array, pointwise_op)
-    new_lazy_array = da.from_array(pointwise_arraylike,
-                                   chunks=lazy_array.chunks)
-    return new_lazy_array
+    # This is just an Iris wrapper for the Dask operation.
+    # For now, we support only a single argument array, and assume that the
+    # output dtype is the same as the input.  This scope is intentionally
+    # minimal : we can extend it later as needed.
+
+    # Note: pass dtype, to prevent Dask making a test call to work it out.
+    return da.map_blocks(elementwise_op,
+                         lazy_array, dtype=lazy_array.dtype)

lib/iris/coords.py

@@ -39,7 +39,7 @@
 from iris._data_manager import DataManager
 from iris._deprecation import warn_deprecated
 from iris._lazy_data import (as_concrete_data, is_lazy_data,
-                             multidim_lazy_stack, lazy_wrap_pointwise)
+                             multidim_lazy_stack, wrap_lazy_elementwise)
 import iris.aux_factory
 import iris.exceptions
 import iris.time
@@ -917,15 +917,15 @@ def pointwise_convert(values):
                 return old_unit.convert(values, new_unit)
 
         if self.has_lazy_points():
-            new_points = lazy_wrap_pointwise(self.lazy_points(),
-                                             pointwise_convert)
+            new_points = wrap_lazy_elementwise(self.lazy_points(),
+                                               pointwise_convert)
         else:
             new_points = self.units.convert(self.points, unit)
         self.points = new_points
         if self.has_bounds():
             if self.has_lazy_bounds():
-                new_bounds = lazy_wrap_pointwise(self.lazy_bounds(),
-                                                 pointwise_convert)
+                new_bounds = wrap_lazy_elementwise(self.lazy_bounds(),
+                                                   pointwise_convert)
             else:
                 new_bounds = self.units.convert(self.bounds, unit)
             self.bounds = new_bounds

lib/iris/cube.py

@@ -41,7 +41,7 @@
 import iris._concatenate
 import iris._constraints
 from iris._data_manager import DataManager
-from iris._lazy_data import lazy_wrap_pointwise
+from iris._lazy_data import wrap_lazy_elementwise
 
 import iris._merge
 import iris.analysis
@@ -881,7 +881,8 @@ def convert_units(self, unit):
             def pointwise_convert(values):
                 return old_unit.convert(values, new_unit)
 
-            new_data = lazy_wrap_pointwise(self.lazy_data(), pointwise_convert)
+            new_data = wrap_lazy_elementwise(self.lazy_data(),
+                                             pointwise_convert)
         else:
             new_data = self.units.convert(self.data, unit)
         self.data = new_data

lib/iris/tests/unit/coords/test_AuxCoord.py

@@ -606,7 +606,7 @@ def test_set_bounds_with_lazy_points(self):
 
 
 class Test_convert_units(tests.IrisTest):
-    def test_convert_preserves_lazy(self):
+    def test_preserves_lazy(self):
         test_bounds = np.array([[[11.0, 12.0], [12.0, 13.0], [13.0, 14.0]],
                                 [[21.0, 22.0], [22.0, 23.0], [23.0, 24.0]]])
         test_points = np.array([[11.1, 12.2, 13.3],

lib/iris/tests/unit/cube/test_Cube.py

@@ -1712,7 +1712,7 @@ def test_convert_unknown_units(self):
         with self.assertRaisesRegexp(UnitConversionError, emsg):
             cube.convert_units('mm day-1')
 
-    def test_convert_preserves_lazy(self):
+    def test_preserves_lazy(self):
         real_data = np.arange(12.).reshape((3, 4))
         lazy_data = as_lazy_data(real_data)
         cube = iris.cube.Cube(lazy_data, units='m')