Replace dask 'compute()' usage with a common realisation call. (#2)

lib/iris/_lazy_data.py

@@ -73,6 +73,31 @@ def as_lazy_data(data, chunks=_MAX_CHUNK_SIZE):
     return data
 
 
+def as_concrete_data(data, nans_replacement=None, result_dtype=None):
+    """
+    Return the actual content of a lazy array, as a numpy array.
+
+    If the data is a NumPy ndarray or masked array, return it unchanged.
+
+    If the data is lazy, return the realised result.
+
+    Where lazy data contains NaNs these are translated by filling or conversion
+    to masked data, using the :func:`~iris._lazy_data.convert_nans_array`
+    function.
+    See there for usage of the 'nans_replacement' and 'result_dtype' keys.
+
+    """
+    if is_lazy_data(data):
+        # Realise dask array.
+        data = data.compute()
+        # Convert any missing data as requested.
+        data = convert_nans_array(data,
+                                  nans_replacement=nans_replacement,
+                                  result_dtype=result_dtype)
+
+    return data
+
+
 def array_masked_to_nans(array):
     """
     Convert a masked array to a NumPy `ndarray` filled with NaN values. Input

lib/iris/_merge.py

@@ -33,7 +33,8 @@
 import numpy as np
 import numpy.ma as ma
 
-from iris._lazy_data import as_lazy_data, is_lazy_data, multidim_lazy_stack
+from iris._lazy_data import (as_lazy_data, as_concrete_data, is_lazy_data,
+                             multidim_lazy_stack)
 import iris.cube
 import iris.coords
 import iris.exceptions
@@ -1217,10 +1218,11 @@ def merge(self, unique=True):
             if all_have_data:
                 # All inputs were concrete, so turn the result back into a
                 # normal array.
-                merged_data = merged_data.compute()
-                # Unmask the array only if it is filled.
+                merged_data = as_concrete_data(merged_data,
+                                               nans_replacement=ma.masked)
+                # Unmask the array if it has no masked points.
                 if (ma.isMaskedArray(merged_data) and
-                        ma.count_masked(merged_data) == 0):
+                        not ma.is_masked(merged_data)):
                     merged_data = merged_data.data
             merged_cube = self._get_cube(merged_data)
             merged_cubes.append(merged_cube)

lib/iris/coords.py

@@ -32,11 +32,11 @@
 import warnings
 import zlib
 
-from iris._lazy_data import is_lazy_data
 import dask.array as da
 import netcdftime
 import numpy as np
 
+from iris._lazy_data import as_concrete_data, is_lazy_data
 import iris.aux_factory
 import iris.exceptions
 import iris.time
@@ -1611,7 +1611,11 @@ def _sanitise_array(self, src, ndmin):
     def points(self):
         """Property containing the points values as a numpy array"""
         if is_lazy_data(self._points):
-            self._points = self._points.compute()
+            self._points = as_concrete_data(self._points,
+                                            nans_replacement=np.ma.masked)
+            # NOTE: we probably don't have full support for masked aux-coords.
+            # We certainly *don't* handle a _FillValue attribute (and possibly
+            # the loader will throw one away ?)
         return self._points.view()
 
     @points.setter
@@ -1649,7 +1653,11 @@ def bounds(self):
         if self._bounds is not None:
             bounds = self._bounds
             if is_lazy_data(bounds):
-                bounds = bounds.compute()
+                bounds = as_concrete_data(bounds,
+                                          nans_replacement=np.ma.masked)
+                # NOTE: we probably don't fully support for masked aux-coords.
+                # We certainly *don't* handle a _FillValue attribute (and
+                # possibly the loader will throw one away ?)
                 self._bounds = bounds
             bounds = bounds.view()
         else:
@@ -1740,9 +1748,11 @@ def measure(self):
     @property
     def data(self):
         """Property containing the data values as a numpy array"""
-        data = self._data
         if is_lazy_data(self._data):
-            self._data = self._data.compute()
+            self._data = as_concrete_data(self._data,
+                                          nans_replacement=np.ma.masked)
+            # NOTE: like AuxCoords, we probably don't fully support masks, and
+            # we certainly don't handle any _FillValue attribute.
         return self._data.view()
 
     @data.setter

lib/iris/cube.py

@@ -42,8 +42,8 @@
 import iris._concatenate
 import iris._constraints
 from iris._deprecation import warn_deprecated
-from iris._lazy_data import (array_masked_to_nans, as_lazy_data,
-                             convert_nans_array, is_lazy_data)
+from iris._lazy_data import as_concrete_data, as_lazy_data, is_lazy_data
+
 import iris._merge
 import iris.analysis
 from iris.analysis.cartography import wrap_lons
@@ -1733,13 +1733,11 @@ def data(self):
         """
         if self.has_lazy_data():
             try:
-                data = self._dask_array.compute()
-                # Now convert the data payload from a NaN array to a
-                # masked array, and if appropriate cast to the specified
-                # cube result dtype.
-                result = convert_nans_array(data,
-                                            nans_replacement=ma.masked,
-                                            result_dtype=self.dtype)
+                # Realise the data, convert from a NaN array to a masked array,
+                # and if appropriate cast to the specified cube result dtype.
+                result = as_concrete_data(self._dask_array,
+                                          nans_replacement=ma.masked,
+                                          result_dtype=self.dtype)
                 self._numpy_array = result
                 self.dtype = None
 

lib/iris/fileformats/pp.py

@@ -39,11 +39,13 @@
 import netcdftime
 
 from iris._deprecation import warn_deprecated
+from iris._lazy_data import (array_masked_to_nans, as_concrete_data,
+                             as_lazy_data, is_lazy_data)
 import iris.config
 import iris.fileformats.rules
 import iris.fileformats.pp_rules
 import iris.coord_systems
-from iris._lazy_data import array_masked_to_nans, as_lazy_data, is_lazy_data
+
 
 try:
     import mo_pack
@@ -1280,11 +1282,11 @@ def data(self):
 
         """
         # Cache the real data on first use
-        if is_lazy_data(self._data):
-            self._data = self._data.compute()
         if self._data.dtype.kind == 'i' and self.bmdi == -1e30:
             self.bmdi = -9999
-        self._data[np.isnan(self._data)] = self.bmdi
+        if is_lazy_data(self._data):
+            self._data = as_concrete_data(self._data,
+                                          nans_replacement=self.bmdi)
         return self._data
 
     @data.setter

lib/iris/tests/unit/analysis/test_MEAN.py

@@ -28,6 +28,7 @@
 import numpy.ma as ma
 
 from iris.analysis import MEAN
+from iris._lazy_data import as_concrete_data
 
 
 class Test_lazy_aggregate(tests.IrisTest):
@@ -42,24 +43,21 @@ def setUp(self):
 
     def test_mdtol_default(self):
         agg = MEAN.lazy_aggregate(self.array, axis=self.axis)
-        result = agg.compute()
-        masked_result = ma.masked_array(result, mask=np.isnan(result))
+        masked_result = as_concrete_data(agg, nans_replacement=ma.masked)
         self.assertMaskedArrayAlmostEqual(masked_result,
                                           self.expected_masked)
 
     def test_mdtol_below(self):
         agg = MEAN.lazy_aggregate(self.array, axis=self.axis, mdtol=0.3)
-        result = agg.compute()
-        masked_result = ma.masked_array(result, mask=np.isnan(result))
+        masked_result = as_concrete_data(agg, nans_replacement=ma.masked)
         expected_masked = self.expected_masked
         expected_masked.mask = [False, True, True, True]
         self.assertMaskedArrayAlmostEqual(masked_result,
                                           expected_masked)
 
     def test_mdtol_above(self):
         agg = MEAN.lazy_aggregate(self.array, axis=self.axis, mdtol=0.4)
-        result = agg.compute()
-        masked_result = ma.masked_array(result, mask=np.isnan(result))
+        masked_result = as_concrete_data(agg, nans_replacement=ma.masked)
         self.assertMaskedArrayAlmostEqual(masked_result,
                                           self.expected_masked)
 
@@ -68,8 +66,9 @@ def test_multi_axis(self):
         collapse_axes = (0, 2)
         lazy_data = as_lazy_data(data)
         agg = MEAN.lazy_aggregate(lazy_data, axis=collapse_axes)
+        result = as_concrete_data(agg, nans_replacement=ma.masked)
         expected = np.mean(data, axis=collapse_axes)
-        self.assertArrayAllClose(agg.compute(), expected)
+        self.assertArrayAllClose(result, expected)
 
 
 class Test_name(tests.IrisTest):

lib/iris/tests/unit/analysis/test_STD_DEV.py

@@ -23,9 +23,9 @@
 # importing anything else.
 import iris.tests as tests
 
-from iris._lazy_data import as_lazy_data
 import numpy as np
 
+from iris._lazy_data import as_concrete_data, as_lazy_data
 from iris.analysis import STD_DEV
 
 
@@ -37,21 +37,22 @@ def test_mdtol(self):
                           [1., 2., na, na]])
         array = as_lazy_data(array)
         var = STD_DEV.lazy_aggregate(array, axis=1, mdtol=0.3)
-        result = var.compute()
-        masked_result = np.ma.masked_array(result, mask=np.isnan(result))
+        masked_result = as_concrete_data(var, nans_replacement=np.ma.masked)
         masked_expected = np.ma.masked_array([0.57735, 1., 0.707107],
                                              mask=[0, 0, 1])
         self.assertMaskedArrayAlmostEqual(masked_result, masked_expected)
 
     def test_ddof_one(self):
         array = as_lazy_data(np.arange(8))
         var = STD_DEV.lazy_aggregate(array, axis=0, ddof=1)
-        self.assertArrayAlmostEqual(var.compute(), np.array(2.449489))
+        result = as_concrete_data(var)
+        self.assertArrayAlmostEqual(result, np.array(2.449489))
 
     def test_ddof_zero(self):
         array = as_lazy_data(np.arange(8))
         var = STD_DEV.lazy_aggregate(array, axis=0, ddof=0)
-        self.assertArrayAlmostEqual(var.compute(), np.array(2.291287))
+        result = as_concrete_data(var)
+        self.assertArrayAlmostEqual(result, np.array(2.291287))
 
 
 class Test_name(tests.IrisTest):

lib/iris/tests/unit/analysis/test_VARIANCE.py

@@ -23,13 +23,14 @@
 # importing anything else.
 import iris.tests as tests
 
-from iris._lazy_data import as_lazy_data
 import numpy as np
 import numpy.ma as ma
 
+from iris._lazy_data import as_lazy_data, as_concrete_data
 from iris.analysis import VARIANCE
 import iris.cube
 from iris.coords import DimCoord
+
 from iris.tests import mock
 
 
@@ -71,12 +72,14 @@ class Test_lazy_aggregate(tests.IrisTest):
     def test_ddof_one(self):
         array = as_lazy_data(np.arange(8))
         var = VARIANCE.lazy_aggregate(array, axis=0, ddof=1)
-        self.assertArrayAlmostEqual(var.compute(), np.array(6.0))
+        result = as_concrete_data(var)
+        self.assertArrayAlmostEqual(result, np.array(6.0))
 
     def test_ddof_zero(self):
         array = as_lazy_data(np.arange(8))
         var = VARIANCE.lazy_aggregate(array, axis=0, ddof=0)
-        self.assertArrayAlmostEqual(var.compute(), np.array(5.25))
+        result = as_concrete_data(var)
+        self.assertArrayAlmostEqual(result, np.array(5.25))
 
 
 class Test_name(tests.IrisTest):

lib/iris/tests/unit/lazy_data/test_as_concrete_data.py

@@ -0,0 +1,90 @@
+# (C) British Crown Copyright 2017, Met Office
+#
+# This file is part of Iris.
+#
+# Iris is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Iris is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
+"""Test function :func:`iris._lazy data.as_concrete_data`."""
+
+from __future__ import (absolute_import, division, print_function)
+from six.moves import (filter, input, map, range, zip)  # noqa
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests
+
+import numpy as np
+import numpy.ma as ma
+import dask.array as da
+
+from iris._lazy_data import as_concrete_data, as_lazy_data, is_lazy_data
+from iris.tests import mock
+
+
+class Test_as_concrete_data(tests.IrisTest):
+    def test_concrete_input_data(self):
+        data = np.arange(24).reshape((4, 6))
+        result = as_concrete_data(data)
+        self.assertIs(data, result)
+        self.assertFalse(is_lazy_data(result))
+
+    def test_concrete_masked_input_data(self):
+        data = ma.masked_array([10, 12, 8, 2], mask=[True, True, False, True])
+        result = as_concrete_data(data)
+        self.assertIs(data, result)
+        self.assertFalse(is_lazy_data(result))
+
+    def test_lazy_data(self):
+        # Minimal testing as as_concrete_data is a wrapper to
+        # convert_nans_array
+        data = np.arange(24).reshape((2, 12))
+        lazy_array = as_lazy_data(data)
+        sentinel = mock.sentinel.data
+        with mock.patch('iris._lazy_data.convert_nans_array') as conv_nans:
+            conv_nans.return_value = sentinel
+            result = as_concrete_data(lazy_array)
+        self.assertEqual(sentinel, result)
+
+        # Check call to convert_nans_array
+        conv_nans.assert_called_once()
+        args, kwargs = conv_nans.call_args
+        arg, = args
+        self.assertFalse(is_lazy_data(arg))
+        self.assertArrayEqual(arg, data)
+        self.assertEqual(kwargs, {'result_dtype': None,
+                                  'nans_replacement': None})
+
+    def test_lazy_data_pass_thru_kwargs(self):
+        # Minimal testing as as_concrete_data is a wrapper to
+        # convert_nans_array
+        data = np.arange(24).reshape((2, 12))
+        lazy_array = as_lazy_data(data)
+        sentinel = mock.sentinel.data
+        with mock.patch('iris._lazy_data.convert_nans_array') as conv_nans:
+            conv_nans.return_value = sentinel
+            result = as_concrete_data(lazy_array, nans_replacement=7,
+                                      result_dtype=np.int16)
+        self.assertEqual(sentinel, result)
+
+        # Check call to convert_nans_array
+        conv_nans.assert_called_once()
+        args, kwargs = conv_nans.call_args
+        arg, = args
+        self.assertFalse(is_lazy_data(arg))
+        self.assertArrayEqual(arg, data)
+        self.assertEqual(kwargs, {'nans_replacement': 7,
+                                  'result_dtype': np.int16, })
+
+
+if __name__ == '__main__':
+    tests.main()