Merge pull request #3093 from TAlonglong/issue3092

Fix 3D masks with size 1 dimension in MaskingCompositor

Author: djhoese

satpy/composites/__init__.py

@@ -1908,6 +1908,9 @@ def __call__(self, projectables, *args, **kwargs):
         data_in = projectables[0]
         mask_in = projectables[1]
 
+        # remove "bands" dimension for single band masks (ex. "L")
+        mask_in = mask_in.squeeze(drop=True)
+
         alpha_attrs = data_in.attrs.copy()
         data = self._select_data_bands(data_in)
 

satpy/tests/test_composites.py

@@ -1636,6 +1636,16 @@ def conditions_v2(self):
                  "value": 2,
                  "transparency": 50}]
 
+    @pytest.fixture
+    def conditions_v3(self):
+        """Masking conditions with other numerical values."""
+        return [{"method": "equal",
+                 "value": 0,
+                 "transparency": 100},
+                {"method": "equal",
+                 "value": 1,
+                 "transparency": 0}]
+
     @pytest.fixture
     def test_data(self):
         """Test data to use with masking compositors."""
@@ -1654,6 +1664,30 @@ def test_ct_data(self):
         ct_data.attrs["flag_values"] = flag_values
         return ct_data
 
+    @pytest.fixture
+    def value_3d_data(self):
+        """Test 3D data array."""
+        value_3d_data = da.array([[[1, 0, 0],
+                                   [0, 1, 0],
+                                   [0, 0, 1]]])
+        value_3d_data = xr.DataArray(value_3d_data, dims=["bands", "y", "x"])
+        return value_3d_data
+
+    @pytest.fixture
+    def value_3d_data_bands(self):
+        """Test 3D data array."""
+        value_3d_data = da.array([[[1, 0, 0],
+                                   [0, 1, 0],
+                                   [0, 0, 1]],
+                                  [[1, 0, 0],
+                                   [0, 1, 0],
+                                   [0, 0, 1]],
+                                  [[1, 0, 0],
+                                   [0, 1, 0],
+                                   [0, 0, 1]]])
+        value_3d_data = xr.DataArray(value_3d_data, dims=["bands", "y", "x"])
+        return value_3d_data
+
     @pytest.fixture
     def test_ct_data_v3(self, test_ct_data):
         """Set ct data to NaN where it originally is 1."""
@@ -1736,6 +1770,44 @@ def test_call_numerical_transparency_data(
             np.testing.assert_allclose(res.sel(bands=m), reference_data)
         np.testing.assert_allclose(res.sel(bands="A"), reference_alpha)
 
+    @pytest.mark.parametrize("mode", ["LA", "RGBA"])
+    def test_call_numerical_transparency_data_with_3d_mask_data(
+            self, test_data, value_3d_data, conditions_v3, mode):
+        """Test call the compositor with numerical transparency data.
+
+        Use parameterisation to test different image modes.
+        """
+        from satpy.composites import MaskingCompositor
+
+        reference_data_v3 = test_data.where(value_3d_data[0] > 0)
+        reference_alpha_v3 = xr.DataArray([[1., 0., 0.],
+                                           [0., 1., 0.],
+                                           [0., 0., 1.]])
+
+        # Test with numerical transparency data using 3d test mask data which can be squeezed
+        comp = MaskingCompositor("name", conditions=conditions_v3,
+                                 mode=mode)
+        res = comp([test_data, value_3d_data])
+        assert res.mode == mode
+        for m in mode.rstrip("A"):
+            np.testing.assert_allclose(res.sel(bands=m), reference_data_v3)
+        np.testing.assert_allclose(res.sel(bands="A"), reference_alpha_v3)
+
+    @pytest.mark.parametrize("mode", ["LA", "RGBA"])
+    def test_call_numerical_transparency_data_with_3d_mask_data_exception(
+            self, test_data, value_3d_data_bands, conditions_v3, mode):
+        """Test call the compositor with numerical transparency data, too many dimensions to squeeze.
+
+        Use parameterisation to test different image modes.
+        """
+        from satpy.composites import MaskingCompositor
+
+        # Test with numerical transparency data using 3d test mask data which can not be squeezed
+        comp = MaskingCompositor("name", conditions=conditions_v3,
+                                 mode=mode)
+        with pytest.raises(ValueError, match=".*Received 3 dimension\\(s\\) but expected 2.*"):
+            comp([test_data, value_3d_data_bands])
+
     def test_call_named_fields(self, conditions_v2, test_data, test_ct_data,
                                reference_data, reference_alpha):
         """Test with named fields."""