Merge pull request #87 from pp-mo/dask_masked_fix

Dask masked with iris

Author: pelson

iris_grib/__init__.py

@@ -36,7 +36,7 @@
 import numpy.ma as ma
 
 import iris
-from iris._lazy_data import as_lazy_data, convert_nans_array
+from iris._lazy_data import as_lazy_data
 import iris.coord_systems as coord_systems
 from iris.exceptions import TranslationError, NotYetImplementedError
 
@@ -145,7 +145,6 @@ class GribWrapper(object):
     def __init__(self, grib_message, grib_fh=None):
         """Store the grib message and compute our extra keys."""
         self.grib_message = grib_message
-        self.realised_dtype = np.array([0.]).dtype
 
         if self.edition != 1:
             emsg = 'GRIB edition {} is not supported by {!r}.'
@@ -184,14 +183,11 @@ def __init__(self, grib_message, grib_fh=None):
             # The byte offset requires to be reset back to the first byte
             # of this message. The file pointer offset is always at the end
             # of the current message due to the grib-api reading the message.
-            proxy = GribDataProxy(shape, self.realised_dtype, grib_fh.name,
+            proxy = GribDataProxy(shape, np.array([0.]).dtype, grib_fh.name,
                                   offset - message_length)
             self._data = as_lazy_data(proxy)
         else:
-            values_array = _message_values(grib_message, shape)
-            # mask where the values are nan
-            self.data = convert_nans_array(values_array,
-                                           nans_replacement=ma.masked)
+            self.data = _message_values(grib_message, shape)
 
     def _confirm_in_scope(self):
         """Ensure we have a grib flavour that we choose to support."""
@@ -698,6 +694,10 @@ def _message_values(grib_message, shape):
     data = gribapi.grib_get_double_array(grib_message, 'values')
     data = data.reshape(shape)
 
+    # Handle missing values in a sensible way.
+    mask = np.isnan(data)
+    if mask.any():
+        data = ma.array(data, mask=mask, fill_value=np.nan)
     return data
 
 

iris_grib/_save_rules.py

@@ -1157,17 +1157,20 @@ def product_definition_section(cube, grib):
 def data_section(cube, grib):
     # Masked data?
     if ma.isMaskedArray(cube.data):
-        fill_value = cube.fill_value
-        if fill_value is None or np.isnan(cube.fill_value):
+        if not np.isnan(cube.data.fill_value):
+            # Use the data's fill value.
+            fill_value = float(cube.data.fill_value)
+        else:
             # We can't use the cube's fill value if it's NaN,
             # the GRIB API doesn't like it.
             # Calculate an MDI outside the data range.
             min, max = cube.data.min(), cube.data.max()
             fill_value = min - (max - min) * 0.1
+        # Prepare the unmasked data array, using fill_value as the MDI.
+        data = cube.data.filled(fill_value)
     else:
         fill_value = None
-
-    data = cube.data
+        data = cube.data
 
     # units scaling
     grib2_info = gptx.cf_phenom_to_grib2_info(cube.standard_name,

iris_grib/message.py

@@ -30,7 +30,7 @@
 import numpy as np
 import numpy.ma as ma
 
-from iris._lazy_data import array_masked_to_nans, as_lazy_data
+from iris._lazy_data import as_lazy_data
 from iris.exceptions import TranslationError
 
 
@@ -120,10 +120,6 @@ def bmdi(self):
         # Default for fill value is None.
         return None
 
-    @property
-    def realised_dtype(self):
-        return np.dtype('f8')
-
     def core_data(self):
         return self.data
 
@@ -164,8 +160,7 @@ def data(self):
                 shape = (grid_section['numberOfDataPoints'],)
             else:
                 shape = (grid_section['Nj'], grid_section['Ni'])
-            proxy = _DataProxy(shape, self.realised_dtype, np.nan,
-                               self._recreate_raw)
+            proxy = _DataProxy(shape, np.dtype('f8'), self._recreate_raw)
             data = as_lazy_data(proxy)
         else:
             fmt = 'Grid definition template {} is not supported'
@@ -196,9 +191,7 @@ class _DataProxy(object):
 
     __slots__ = ('shape', 'dtype', 'recreate_raw')
 
-    def __init__(self, shape, dtype, fill_value, recreate_raw):
-        # TODO: I (@pelson) have no idea why fill_value remains an argument as
-        # it isn't used. It was copied verbatim from iris' dask branch.
+    def __init__(self, shape, dtype, recreate_raw):
         self.shape = shape
         self.dtype = dtype
         self.recreate_raw = recreate_raw
@@ -260,10 +253,10 @@ def __getitem__(self, keys):
                 # Only the non-masked values are included in codedValues.
                 _data = np.empty(shape=bitmap.shape)
                 _data[bitmap.astype(bool)] = data
-                # Use nan where input = 1, the opposite of the behaviour
-                # specified by the GRIB spec.
-                _data[np.logical_not(bitmap.astype(bool))] = np.nan
-                data = _data
+                # `ma.masked_array` masks where input = 1, the opposite of
+                # the behaviour specified by the GRIB spec.
+                data = ma.masked_array(_data, mask=np.logical_not(bitmap),
+                                       fill_value=np.nan)
             else:
                 msg = 'Shapes of data and bitmap do not match.'
                 raise TranslationError(msg)

iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube core-dtype="float64" dtype="float64" standard_name="geopotential" units="m2 s-2">
+  <cube dtype="float64" standard_name="geopotential" units="m2 s-2">
     <attributes>
       <attribute name="centre" value="European Centre for Medium Range Weather Forecasts"/>
     </attributes>

iris_grib/tests/unit/load_convert/test_reference_time_coord.py

@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2014 - 2016, Met Office
+# (C) British Crown Copyright 2014 - 2017, Met Office
 #
 # This file is part of iris-grib.
 #
@@ -60,12 +60,12 @@ def _check(self, section, standard_name=None):
         coord = reference_time_coord(section)
         self.assertEqual(coord, expected)
 
-    def test_start_of_forecast(self):
+    def test_start_of_forecast__0(self):
         section = deepcopy(self.section)
         section['significanceOfReferenceTime'] = 0
         self._check(section, 'forecast_reference_time')
 
-    def test_start_of_forecast(self):
+    def test_start_of_forecast__1(self):
         section = deepcopy(self.section)
         section['significanceOfReferenceTime'] = 1
         self._check(section, 'forecast_reference_time')

iris_grib/tests/unit/message/test_GribMessage.py

@@ -92,38 +92,16 @@ def test_no_bitmap(self):
         self.assertEqual(result.shape, self.shape)
         self.assertArrayEqual(result, expected)
 
-    def test_bitmap_present_int_data(self):
+    def test_bitmap_present(self):
         # Test the behaviour where bitmap and codedValues shapes
-        # are not equal, and codedValues is integer data.
-        data_type = np.int64
+        # are not equal.
         input_values = np.arange(5)
-        output_values = np.array([-1, -1, 0, 1, -1, -1, -1, 2, -1, 3, -1, 4],
-                                 dtype=data_type)
+        output_values = np.array([-1, -1, 0, 1, -1, -1, -1, 2, -1, 3, -1, 4])
         message = _make_test_message({3: self._section_3,
                                       6: {'bitMapIndicator': 0,
                                           'bitmap': self.bitmap},
                                       7: {'codedValues': input_values}})
-        result = as_concrete_data(message.data, nans_replacement=ma.masked,
-                                  result_dtype=data_type)
-        expected = ma.masked_array(output_values,
-                                   np.logical_not(self.bitmap))
-        expected = expected.reshape(self.shape)
-        self.assertMaskedArrayEqual(result, expected)
-        self.assertEqual(result.dtype, data_type)
-
-    def test_bitmap_present_float_data(self):
-        # Test the behaviour where bitmap and codedValues shapes
-        # are not equal, and codedValues is float data.
-        data_type = np.float32
-        input_values = np.arange(5, dtype=np.float32) + 5
-        output_values = np.array([-1, -1, 5, 6, -1, -1, -1, 7, -1, 8, -1, 9],
-                                 dtype=data_type)
-        message = _make_test_message({3: self._section_3,
-                                      6: {'bitMapIndicator': 0,
-                                          'bitmap': self.bitmap},
-                                      7: {'codedValues': input_values}})
-        result = as_concrete_data(message.data, nans_replacement=ma.masked,
-                                  result_dtype=data_type)
+        result = as_concrete_data(message.data)
         expected = ma.masked_array(output_values,
                                    np.logical_not(self.bitmap))
         expected = expected.reshape(self.shape)

iris_grib/tests/unit/message/test__DataProxy.py

@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2014 - 2016, Met Office
+# (C) British Crown Copyright 2014 - 2017, Met Office
 #
 # This file is part of iris-grib.
 #
@@ -37,20 +37,20 @@
 class Test__bitmap(tests.IrisGribTest):
     def test_no_bitmap(self):
         section_6 = {'bitMapIndicator': 255, 'bitmap': None}
-        data_proxy = _DataProxy(0, 0, 0, 0)
+        data_proxy = _DataProxy(0, 0, 0)
         result = data_proxy._bitmap(section_6)
         self.assertIsNone(result)
 
     def test_bitmap_present(self):
         bitmap = randint(2, size=(12))
         section_6 = {'bitMapIndicator': 0, 'bitmap': bitmap}
-        data_proxy = _DataProxy(0, 0, 0, 0)
+        data_proxy = _DataProxy(0, 0, 0)
         result = data_proxy._bitmap(section_6)
         self.assertArrayEqual(bitmap, result)
 
     def test_bitmap__invalid_indicator(self):
         section_6 = {'bitMapIndicator': 100, 'bitmap': None}
-        data_proxy = _DataProxy(0, 0, 0, 0)
+        data_proxy = _DataProxy(0, 0, 0)
         with self.assertRaisesRegexp(TranslationError, 'unsupported bitmap'):
             data_proxy._bitmap(section_6)
 

iris_grib/tests/unit/save_rules/test_data_section.py

@@ -94,8 +94,8 @@ def test_simple(self):
 
     def test_masked_with_finite_fill_value(self):
         cube = iris.cube.Cube(np.ma.MaskedArray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0],
-                                                mask=[0, 0, 0, 1, 1, 1]),
-                              fill_value=2000)
+                                                mask=[0, 0, 0, 1, 1, 1],
+                                                fill_value=2000))
         grib_message = mock.sentinel.GRIB_MESSAGE
         with mock.patch(GRIB_API) as grib_api:
             data_section(cube, grib_message)
@@ -107,8 +107,8 @@ def test_masked_with_finite_fill_value(self):
 
     def test_masked_with_nan_fill_value(self):
         cube = iris.cube.Cube(np.ma.MaskedArray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0],
-                                                mask=[0, 0, 0, 1, 1, 1]),
-                              fill_value=np.nan)
+                                                mask=[0, 0, 0, 1, 1, 1],
+                                                fill_value=np.nan))
         grib_message = mock.sentinel.GRIB_MESSAGE
         with mock.patch(GRIB_API) as grib_api:
             data_section(cube, grib_message)
@@ -135,8 +135,8 @@ def test_scaled_with_finite_fill_value(self):
         # When re-scaling masked data with a finite fill value, ensure
         # the fill value and any filled values are also re-scaled.
         cube = iris.cube.Cube(np.ma.MaskedArray([1.0, 2.0, 3.0, 1.0, 2.0, 3.0],
-                                                mask=[0, 0, 0, 1, 1, 1]),
-                              fill_value=2000,
+                                                mask=[0, 0, 0, 1, 1, 1],
+                                                fill_value=2000),
                               standard_name='geopotential_height', units='km')
         grib_message = mock.sentinel.GRIB_MESSAGE
         with mock.patch(GRIB_API) as grib_api:
@@ -152,8 +152,8 @@ def test_scaled_with_nan_fill_value(self):
         # a fill value is chosen which allows for the scaling, and any
         # filled values match the chosen fill value.
         cube = iris.cube.Cube(np.ma.MaskedArray([-1.0, 2.0, -1.0, 2.0],
-                                                mask=[0, 0, 1, 1]),
-                              fill_value=np.nan,
+                                                mask=[0, 0, 1, 1],
+                                                fill_value=np.nan),
                               standard_name='geopotential_height', units='km')
         grib_message = mock.sentinel.GRIB_MESSAGE
         with mock.patch(GRIB_API) as grib_api:

iris_grib/tests/unit/test_load_cubes.py

@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2014 - 2016, Met Office
+# (C) British Crown Copyright 2014 - 2017, Met Office
 #
 # This file is part of iris-grib.
 #
@@ -41,7 +41,7 @@ def test(self):
             rules_load.return_value = expected_result
             result = load_cubes(files, callback)
             kwargs = {}
-            loader = Loader(generator, kwargs, converter, None)
+            loader = Loader(generator, kwargs, converter)
             rules_load.assert_called_once_with(files, callback, loader)
             self.assertIs(result, expected_result)