Fix composites failing when inputs are different chunk sizes

satpy/composites/__init__.py

@@ -19,6 +19,7 @@
 
 import logging
 import os
+import string
 import warnings
 
 import dask.array as da
@@ -159,7 +160,28 @@ def apply_modifier_info(self, origin, destination):
     def match_data_arrays(self, data_arrays):
         """Match data arrays so that they can be used together in a composite."""
         self.check_geolocation(data_arrays)
-        return self.drop_coordinates(data_arrays)
+        new_arrays = self.drop_coordinates(data_arrays)
+        return self.unify_chunks(new_arrays)
+
+    def unify_chunks(self, data_arrays):
+        """Unify chunk sizes across multiple DataArray objects.
+
+        Useful to avoid errors when calling dask's "map_blocks" or other
+        similar functions where inputs are expected to have the same number
+        of chunks.
+
+        """
+        all_dims = set.union(*(set(x.dims) for x in data_arrays))
+        dim_map = dict(zip(all_dims, string.ascii_lowercase))
+        arr_pairs = []
+        for data_arr in data_arrays:
+            dask_arr = data_arr.data
+            dims = data_arr.dims
+            arr_pairs.extend([dask_arr, "".join(dim_map[dim_name] for dim_name in dims)])
+        _, new_dask_arrs = da.core.unify_chunks(*arr_pairs)
+        for data_arr, dask_arr in zip(data_arrays, new_dask_arrs):
+            data_arr.data = dask_arr  # inplace
+        return data_arrays
 
     def drop_coordinates(self, data_arrays):
         """Drop neglible non-dimensional coordinates."""

satpy/etc/composites/modis.yaml

@@ -4,6 +4,7 @@ modifiers:
     modifier: !!python/name:satpy.modifiers.atmosphere.ReflectanceCorrector
     url: "https://www.ssec.wisc.edu/~davidh/polar2grid/modis_crefl/tbase.hdf"
     known_hash: "sha256:ed5183cddce905361c1cac8ae6e3a447212875ea421a05747751efe76f8a068e"
+    dem_sds: "Elevation"
     prerequisites:
     - name: satellite_azimuth_angle
     - name: satellite_zenith_angle

satpy/modifiers/_crefl.py

@@ -137,7 +137,16 @@ def _read_var_from_hdf4_file_pyhdf(local_filename, var_name):
         f = SD(local_filename, SDC.READ)
         var = f.select(var_name)
         data = var[:]
-        return np.ma.MaskedArray(data, data == var.getfillvalue())
+        fill = ReflectanceCorrector._read_fill_value_from_hdf4(var, data.dtype)
+        return np.ma.MaskedArray(data, data == fill)
+
+    @staticmethod
+    def _read_fill_value_from_hdf4(var, dtype):
+        from pyhdf.error import HDF4Error
+        try:
+            return var.getfillvalue()
+        except HDF4Error:
+            return np.iinfo(dtype).min
 
     def _get_data_and_angles(self, datasets, optional_datasets):
         angles = self._extract_angle_data_arrays(datasets, optional_datasets)

satpy/tests/modifier_tests/test_crefl.py

@@ -15,13 +15,60 @@
 """Tests for the CREFL ReflectanceCorrector modifier."""
 import unittest
 from unittest import mock
+from contextlib import contextmanager
+from datetime import datetime
 
 import numpy as np
 import pytest
 from dask import array as da
+import xarray as xr
 from pyresample.geometry import AreaDefinition
 
 
+@contextmanager
+def mock_cmgdem(tmpdir, url):
+    """Create fake file representing CMGDEM.hdf."""
+    yield from _mock_and_create_dem_file(tmpdir, url, "averaged elevation", fill_value=-9999)
+
+
+@contextmanager
+def mock_tbase(tmpdir, url):
+    """Create fake file representing tbase.hdf."""
+    yield from _mock_and_create_dem_file(tmpdir, url, "Elevation")
+
+
+def _mock_and_create_dem_file(tmpdir, url, var_name, fill_value=None):
+    if not url:
+        yield None
+        return
+
+    rmock_obj, dem_fn = _mock_dem_retrieve(tmpdir, url)
+    _create_fake_dem_file(dem_fn, var_name, fill_value)
+
+    try:
+        yield rmock_obj
+    finally:
+        rmock_obj.stop()
+
+
+def _mock_dem_retrieve(tmpdir, url):
+    rmock_obj = mock.patch('satpy.modifiers._crefl.retrieve')
+    rmock = rmock_obj.start()
+    dem_fn = str(tmpdir.join(url))
+    rmock.return_value = dem_fn
+    return rmock_obj, dem_fn
+
+
+def _create_fake_dem_file(dem_fn, var_name, fill_value):
+    from pyhdf.SD import SD, SDC
+    h = SD(dem_fn, SDC.WRITE | SDC.CREATE)
+    dem_var = h.create(var_name, SDC.INT16, (10, 10))
+    dem_var[:] = np.zeros((10, 10), dtype=np.int16)
+    if fill_value is not None:
+        dem_var.setfillvalue(fill_value)
+    h.end()
+
+
 class TestViirsReflectanceCorrectorAnglesTest(unittest.TestCase):
     """Tests for the VIIRS/MODIS Corrected Reflectance modifier handling angles."""
 
@@ -43,8 +90,6 @@ def tearDown(self):
     @mock.patch('satpy.modifiers._crefl.get_satpos')
     def test_get_angles(self, get_satpos):
         """Test sun and satellite angle calculation."""
-        import numpy as np
-        import dask.array as da
         from satpy.modifiers._crefl import ReflectanceCorrector
 
         # Patch methods
@@ -76,6 +121,20 @@ def test_get_angles(self, get_satpos):
         self.assertEqual(args[6], 0)
 
 
+def _make_viirs_xarray(data, area, name, standard_name, wavelength=None, units='degrees', calibration=None):
+    return xr.DataArray(data, dims=('y', 'x'),
+                        attrs={
+                            'start_orbit': 1708, 'end_orbit': 1708, 'wavelength': wavelength,
+                            'modifiers': None, 'calibration': calibration,
+                            'resolution': 371, 'name': name,
+                            'standard_name': standard_name, 'platform_name': 'Suomi-NPP',
+                            'polarization': None, 'sensor': 'viirs', 'units': units,
+                            'start_time': datetime(2012, 2, 25, 18, 1, 24, 570942),
+                            'end_time': datetime(2012, 2, 25, 18, 11, 21, 175760), 'area': area,
+                            'ancillary_variables': []
+                        })
+
+
 class TestReflectanceCorrectorModifier:
     """Test the CREFL modifier."""
 
@@ -91,17 +150,14 @@ def data_area_ref_corrector():
             cols, rows,
             (-5434894.954752679, -5434894.964451744, 5434894.964451744, 5434894.954752679))
 
-        dnb = np.zeros((rows, cols)) + 25
-        dnb[3, :] += 25
-        dnb[4:, :] += 50
-        dnb = da.from_array(dnb, chunks=100)
-        return area, dnb
+        data = np.zeros((rows, cols)) + 25
+        data[3, :] += 25
+        data[4:, :] += 50
+        data = da.from_array(data, chunks=100)
+        return area, data
 
     def test_reflectance_corrector_abi(self):
         """Test ReflectanceCorrector modifier with ABI data."""
-        import xarray as xr
-        import dask.array as da
-        import numpy as np
         from satpy.modifiers._crefl import ReflectanceCorrector
         from satpy.tests.utils import make_dsq
         ref_cor = ReflectanceCorrector(optional_prerequisites=[
@@ -172,13 +228,15 @@ def test_reflectance_corrector_abi(self):
                                             51.909142813383916, 58.8234273736508, 68.84706145641482, 69.91085190887961,
                                             71.10179768327806, 71.33161009169649])
 
-    @pytest.mark.parametrize('url', [None, 'CMGDEM.hdf'])
-    def test_reflectance_corrector_viirs(self, tmpdir, url):
+    @pytest.mark.parametrize(
+        'url,dem_mock_cm,dem_sds',
+        [
+            (None, mock_cmgdem, "average elevation"),
+            ("CMGDEM.hdf", mock_cmgdem, "averaged elevation"),
+            ("tbase.hdf", mock_tbase, "Elevation"),
+        ])
+    def test_reflectance_corrector_viirs(self, tmpdir, url, dem_mock_cm, dem_sds):
         """Test ReflectanceCorrector modifier with VIIRS data."""
-        import xarray as xr
-        import dask.array as da
-        import numpy as np
-        import datetime
         from satpy.modifiers._crefl import ReflectanceCorrector
         from satpy.tests.utils import make_dsq
 
@@ -196,7 +254,9 @@ def test_reflectance_corrector_viirs(self, tmpdir, url):
             calibration='reflectance',
             modifiers=('sunz_corrected_iband', 'rayleigh_corrected_crefl_iband'),
             sensor='viirs',
-            url=url)
+            url=url,
+            dem_sds=dem_sds,
+        )
 
         assert ref_cor.attrs['modifiers'] == ('sunz_corrected_iband', 'rayleigh_corrected_crefl_iband')
         assert ref_cor.attrs['calibration'] == 'reflectance'
@@ -211,32 +271,17 @@ def test_reflectance_corrector_viirs(self, tmpdir, url):
             make_dsq(name='solar_azimuth_angle'),
             make_dsq(name='solar_zenith_angle')]
 
-        area, dnb = self.data_area_ref_corrector()
-
-        def make_xarray(name, standard_name, wavelength=None, units='degrees', calibration=None):
-            return xr.DataArray(dnb, dims=('y', 'x'),
-                                attrs={
-                                    'start_orbit': 1708, 'end_orbit': 1708, 'wavelength': wavelength, 'level': None,
-                                    'modifiers': None, 'calibration': calibration,
-                                    'resolution': 371, 'name': name,
-                                    'standard_name': standard_name, 'platform_name': 'Suomi-NPP',
-                                    'polarization': None, 'sensor': 'viirs', 'units': units,
-                                    'start_time': datetime.datetime(2012, 2, 25, 18, 1, 24, 570942),
-                                    'end_time': datetime.datetime(2012, 2, 25, 18, 11, 21, 175760), 'area': area,
-                                    'ancillary_variables': []
-                                })
-
-        c01 = make_xarray('I01', 'toa_bidirectional_reflectance',
-                          wavelength=(0.6, 0.64, 0.68), units='%',
-                          calibration='reflectance')
-        c02 = make_xarray('satellite_azimuth_angle', 'sensor_azimuth_angle')
-        c03 = make_xarray('satellite_zenith_angle', 'sensor_zenith_angle')
-        c04 = make_xarray('solar_azimuth_angle', 'solar_azimuth_angle')
-        c05 = make_xarray('solar_zenith_angle', 'solar_zenith_angle')
+        area, data = self.data_area_ref_corrector()
+        c01 = _make_viirs_xarray(data, area, 'I01', 'toa_bidirectional_reflectance',
+                                 wavelength=(0.6, 0.64, 0.68), units='%',
+                                 calibration='reflectance')
+        c02 = _make_viirs_xarray(data, area, 'satellite_azimuth_angle', 'sensor_azimuth_angle')
+        c03 = _make_viirs_xarray(data, area, 'satellite_zenith_angle', 'sensor_zenith_angle')
+        c04 = _make_viirs_xarray(data, area, 'solar_azimuth_angle', 'solar_azimuth_angle')
+        c05 = _make_viirs_xarray(data, area, 'solar_zenith_angle', 'solar_zenith_angle')
 
-        rmock_obj = self._start_dem_mock(tmpdir, url)
-        res = ref_cor([c01], [c02, c03, c04, c05])
-        self._stop_dem_mock(rmock_obj)
+        with dem_mock_cm(tmpdir, url):
+            res = ref_cor([c01], [c02, c03, c04, c05])
 
         assert isinstance(res, xr.DataArray)
         assert isinstance(res.data, da.Array)
@@ -249,8 +294,8 @@ def make_xarray(name, standard_name, wavelength=None, units='degrees', calibrati
         assert res.attrs['platform_name'] == 'Suomi-NPP'
         assert res.attrs['sensor'] == 'viirs'
         assert res.attrs['units'] == '%'
-        assert res.attrs['start_time'] == datetime.datetime(2012, 2, 25, 18, 1, 24, 570942)
-        assert res.attrs['end_time'] == datetime.datetime(2012, 2, 25, 18, 11, 21, 175760)
+        assert res.attrs['start_time'] == datetime(2012, 2, 25, 18, 1, 24, 570942)
+        assert res.attrs['end_time'] == datetime(2012, 2, 25, 18, 11, 21, 175760)
         assert res.attrs['area'] == area
         assert res.attrs['ancillary_variables'] == []
         data = res.values
@@ -261,10 +306,6 @@ def make_xarray(name, standard_name, wavelength=None, units='degrees', calibrati
 
     def test_reflectance_corrector_modis(self):
         """Test ReflectanceCorrector modifier with MODIS data."""
-        import xarray as xr
-        import dask.array as da
-        import numpy as np
-        import datetime
         from satpy.modifiers._crefl import ReflectanceCorrector
         from satpy.tests.utils import make_dsq
         sataa_did = make_dsq(name='satellite_azimuth_angle')
@@ -290,17 +331,16 @@ def test_reflectance_corrector_modis(self):
 
         area, dnb = self.data_area_ref_corrector()
 
-        def make_xarray(name, calibration, wavelength=None, modifiers=None, resolution=1000,
-                        file_type='hdf_eos_geo'):
+        def make_xarray(name, calibration, wavelength=None, modifiers=None, resolution=1000):
             return xr.DataArray(dnb,
                                 dims=('y', 'x'),
                                 attrs={
                                     'wavelength': wavelength, 'level': None, 'modifiers': modifiers,
-                                    'calibration': calibration, 'resolution': resolution, 'file_type': file_type,
+                                    'calibration': calibration, 'resolution': resolution,
                                     'name': name, 'coordinates': ['longitude', 'latitude'],
                                     'platform_name': 'EOS-Aqua', 'polarization': None, 'sensor': 'modis',
-                                    'units': '%', 'start_time': datetime.datetime(2012, 8, 13, 18, 46, 1, 439838),
-                                    'end_time': datetime.datetime(2012, 8, 13, 18, 57, 47, 746296), 'area': area,
+                                    'units': '%', 'start_time': datetime(2012, 8, 13, 18, 46, 1, 439838),
+                                    'end_time': datetime(2012, 8, 13, 18, 57, 47, 746296), 'area': area,
                                     'ancillary_variables': []
                                 })
 
@@ -323,8 +363,8 @@ def make_xarray(name, calibration, wavelength=None, modifiers=None, resolution=1
         assert res.attrs['platform_name'] == 'EOS-Aqua'
         assert res.attrs['sensor'] == 'modis'
         assert res.attrs['units'] == '%'
-        assert res.attrs['start_time'] == datetime.datetime(2012, 8, 13, 18, 46, 1, 439838)
-        assert res.attrs['end_time'] == datetime.datetime(2012, 8, 13, 18, 57, 47, 746296)
+        assert res.attrs['start_time'] == datetime(2012, 8, 13, 18, 46, 1, 439838)
+        assert res.attrs['end_time'] == datetime(2012, 8, 13, 18, 57, 47, 746296)
         assert res.attrs['area'] == area
         assert res.attrs['ancillary_variables'] == []
         data = res.values
@@ -342,21 +382,53 @@ def test_reflectance_corrector_bad_prereqs(self):
         pytest.raises(ValueError, ref_cor, [1, 2, 3, 4], [])
         pytest.raises(ValueError, ref_cor, [], [1, 2, 3, 4])
 
-    def _start_dem_mock(self, tmpdir, url):
-        if not url:
-            return
-        rmock_obj = mock.patch('satpy.modifiers._crefl.retrieve')
-        rmock = rmock_obj.start()
-        dem_fn = str(tmpdir.join(url))
-        rmock.return_value = dem_fn
-        from pyhdf.SD import SD, SDC
-
-        h = SD(dem_fn, SDC.WRITE | SDC.CREATE)
-        dem_var = h.create("averaged elevation", SDC.FLOAT32, (10, 10))
-        dem_var.setfillvalue(-999.0)
-        dem_var[:] = np.zeros((10, 10), dtype=np.float32)
-        return rmock_obj
-
-    def _stop_dem_mock(self, rmock_obj):
-        if rmock_obj:
-            rmock_obj.stop()
+    @pytest.mark.parametrize(
+        'url,dem_mock_cm,dem_sds',
+        [
+            (None, mock_cmgdem, "average elevation"),
+            ("CMGDEM.hdf", mock_cmgdem, "averaged elevation"),
+            ("tbase.hdf", mock_tbase, "Elevation"),
+        ])
+    def test_reflectance_corrector_different_chunks(self, tmpdir, url, dem_mock_cm, dem_sds):
+        """Test that the modifier works with different chunk sizes for inputs.
+
+        The modifier uses dask's "map_blocks". If the input chunks aren't the
+        same an error is raised.
+
+        """
+        from satpy.modifiers._crefl import ReflectanceCorrector
+        from satpy.tests.utils import make_dsq
+
+        ref_cor = ReflectanceCorrector(
+            optional_prerequisites=[
+                make_dsq(name='satellite_azimuth_angle'),
+                make_dsq(name='satellite_zenith_angle'),
+                make_dsq(name='solar_azimuth_angle'),
+                make_dsq(name='solar_zenith_angle')
+            ],
+            name='I01',
+            prerequisites=[],
+            wavelength=(0.6, 0.64, 0.68),
+            resolution=371,
+            calibration='reflectance',
+            modifiers=('sunz_corrected_iband', 'rayleigh_corrected_crefl_iband'),
+            sensor='viirs',
+            url=url,
+            dem_sds=dem_sds,
+        )
+
+        area, data = self.data_area_ref_corrector()
+        c01 = _make_viirs_xarray(data, area, 'I01', 'toa_bidirectional_reflectance',
+                                 wavelength=(0.6, 0.64, 0.68), units='%',
+                                 calibration='reflectance')
+        c02 = _make_viirs_xarray(data, area, 'satellite_azimuth_angle', 'sensor_azimuth_angle')
+        c02.data = c02.data.rechunk((1, -1))
+        c03 = _make_viirs_xarray(data, area, 'satellite_zenith_angle', 'sensor_zenith_angle')
+        c04 = _make_viirs_xarray(data, area, 'solar_azimuth_angle', 'solar_azimuth_angle')
+        c05 = _make_viirs_xarray(data, area, 'solar_zenith_angle', 'solar_zenith_angle')
+
+        with dem_mock_cm(tmpdir, url):
+            res = ref_cor([c01], [c02, c03, c04, c05])
+
+        # make sure it can actually compute
+        res.compute()