Make Scene bilinear reduction halo-aware and single-stage

Author: Zaczero

satpy/scene.py

@@ -40,6 +40,8 @@
 
 LOG = logging.getLogger(__name__)
 
+BILINEAR_REDUCTION_HALO = 1
+
 
 def _get_area_resolution(area):
     """Attempt to retrieve resolution from AreaDefinition."""
@@ -60,6 +62,13 @@ def _aggregate_data_array(data_array, func, **coarsen_kwargs):
     return out
 
 
+def _expand_slice(slc: slice, size: int, padding: int):
+    """Expand a slice by a symmetric padding within array bounds."""
+    start = 0 if slc.start is None else max(0, slc.start - padding)
+    stop = size if slc.stop is None else min(size, slc.stop + padding)
+    return slice(start, stop, slc.step)
+
+
 class DelayedGeneration(KeyError):
     """Mark that a dataset can't be generated without further modification."""
 
@@ -873,6 +882,7 @@ def _resampled_scene(self, new_scn, destination_area, reduce_data=True,
         new_datasets = {}
         datasets = list(new_scn._datasets.values())
         destination_area = self._get_finalized_destination_area(destination_area, new_scn)
+        resampler_kwargs = self._get_resampler_kwargs(resample_kwargs)
 
         resamplers = {}
         reductions = {}
@@ -896,8 +906,8 @@ def _resampled_scene(self, new_scn, destination_area, reduce_data=True,
             source_area = dataset.attrs["area"]
             dataset, source_area = self._reduce_data(dataset, source_area, destination_area,
                                                      reduce_data, reductions, resample_kwargs)
-            self._prepare_resampler(source_area, destination_area, resamplers, resample_kwargs)
-            kwargs = resample_kwargs.copy()
+            self._prepare_resampler(source_area, destination_area, resamplers, resampler_kwargs)
+            kwargs = resampler_kwargs.copy()
             kwargs["resampler"] = resamplers[source_area]
             res = resample_dataset(dataset, destination_area, **kwargs)
             new_datasets[ds_id] = res
@@ -927,23 +937,48 @@ def _prepare_resampler(self, source_area, destination_area, resamplers, resample
             resamplers[source_area] = resampler
             self._resamplers[key] = resampler
 
+    @staticmethod
+    def _uses_bilinear_resampler(resampler: Any):
+        if resampler == "bilinear":
+            return True
+        try:
+            from satpy.resample.kdtree import BilinearResampler
+        except ImportError:
+            return False
+        if isinstance(resampler, BilinearResampler):
+            return True
+        if isinstance(resampler, type):
+            return issubclass(resampler, BilinearResampler)
+        return False
+
+    def _get_resampler_kwargs(self, resample_kwargs: dict[str, Any]):
+        resampler_kwargs = resample_kwargs.copy()
+        if self._uses_bilinear_resampler(resampler_kwargs.get("resampler")):
+            resampler_kwargs["reduce_data"] = False
+        return resampler_kwargs
+
+    def _get_area_slice_kwargs(self, resample_kwargs: dict[str, Any]):
+        if resample_kwargs.get("resampler") == "gradient_search":
+            return {"shape_divisible_by": resample_kwargs.get("shape_divisible_by", 2)}
+        return {}
+
+    def _pad_reduce_data_slices(self, source_area, slice_x, slice_y, resample_kwargs):
+        if self._uses_bilinear_resampler(resample_kwargs.get("resampler")):
+            slice_x = _expand_slice(slice_x, source_area.width, BILINEAR_REDUCTION_HALO)
+            slice_y = _expand_slice(slice_y, source_area.height, BILINEAR_REDUCTION_HALO)
+        return slice_x, slice_y
+
     def _reduce_data(self, dataset, source_area, destination_area, reduce_data, reductions, resample_kwargs):
         try:
             if reduce_data:
                 key = source_area
                 try:
                     (slice_x, slice_y), source_area = reductions[key]
                 except KeyError:
-                    if resample_kwargs.get("resampler") == "gradient_search":
-                        factor = resample_kwargs.get("shape_divisible_by", 2)
-                    else:
-                        factor = None
-                    try:
-                        slice_x, slice_y = source_area.get_area_slices(
-                            destination_area, shape_divisible_by=factor)
-                    except TypeError:
-                        slice_x, slice_y = source_area.get_area_slices(
-                            destination_area)
+                    slice_kwargs = self._get_area_slice_kwargs(resample_kwargs)
+                    slice_x, slice_y = source_area.get_area_slices(destination_area, **slice_kwargs)
+                    slice_x, slice_y = self._pad_reduce_data_slices(
+                        source_area, slice_x, slice_y, resample_kwargs)
                     source_area = source_area[slice_y, slice_x]
                     reductions[key] = (slice_x, slice_y), source_area
                 dataset = self._slice_data(source_area, (slice_x, slice_y), dataset)

satpy/tests/scene_tests/test_resampling.py

@@ -194,6 +194,47 @@ def _fake_resample_dataset_force_20x20(self, dataset, dest_area, **kwargs):
             attrs=attrs,
         )
 
+    @staticmethod
+    def _geos_cross_projection_areas():
+        from pyresample.geometry import AreaDefinition
+
+        source_area = AreaDefinition(
+            "src",
+            "src",
+            "src",
+            {
+                "proj": "geos",
+                "a": 6378169.0,
+                "b": 6356583.8,
+                "h": 35785831.0,
+                "lon_0": 0.0,
+                "units": "m",
+            },
+            3712,
+            1392,
+            (5568748.0, 5568748.0, -5568748.0, 1392187.0),
+        )
+        target_area = AreaDefinition(
+            "dst",
+            "dst",
+            "dst",
+            {"proj": "latlong", "datum": "WGS84"},
+            768,
+            768,
+            (7.456086060029671, 43.98125198600542, 33.461915849571966, 59.24271064133026),
+        )
+        return source_area, target_area
+
+    @staticmethod
+    def _geos_cross_projection_data(source_area):
+        y = da.arange(source_area.height, chunks=256)[:, None]
+        x = da.arange(source_area.width, chunks=256)[None, :]
+        return xr.DataArray(
+            (y * 10000 + x).astype(np.float64),
+            dims=("y", "x"),
+            attrs={"area": source_area, "name": "edge"},
+        )
+
     @mock.patch("satpy.resample.base.resample_dataset")
     @pytest.mark.parametrize("datasets", [
         None,
@@ -335,6 +376,61 @@ def test_resample_reduce_data_toggle(self, rs):
             assert get_area_slices.call_count == 2
             assert get_area_slices_big.call_count == 2
 
+    @mock.patch("satpy.resample.base.resample_dataset")
+    def test_resample_bilinear_reduce_data_uses_halo_and_disables_internal_reduce_data(self, rs):
+        """Test bilinear Scene reduction pads the crop and avoids internal reduction."""
+        from pyresample.geometry import AreaDefinition
+
+        rs.side_effect = self._fake_resample_dataset_force_20x20
+        proj_str = ("+proj=lcc +datum=WGS84 +ellps=WGS84 "
+                    "+lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs")
+        target_area = AreaDefinition("test", "test", "test", proj_str, 4, 4, (-1000., -1500., 1000., 1500.))
+        area_def = AreaDefinition("test", "test", "test", proj_str, 5, 5, (-1000., -1500., 1000., 1500.))
+        area_def.get_area_slices = mock.MagicMock(return_value=(slice(1, 3, None), slice(1, 3, None)))
+
+        scene = Scene(filenames=["fake1_1.txt"], reader="fake1")
+        scene.load(["comp19"])
+        scene["comp19"].attrs["area"] = area_def
+
+        scene.resample(target_area, resampler="bilinear", reduce_data=True)
+
+        dataset = rs.call_args.args[0]
+        assert dataset.sizes["y"] == 4
+        assert dataset.sizes["x"] == 4
+        assert dataset.attrs["area"].width == 4
+        assert dataset.attrs["area"].height == 4
+        assert rs.call_args.kwargs["reduce_data"] is False
+
+    def test_resample_bilinear_reduce_data_halo_matches_full_source_on_edge_case(self):
+        """Test bilinear Scene reduction keeps edge pixels with a 1-pixel halo."""
+        from pyresample.bilinear import XArrayBilinearResampler
+
+        source_area, target_area = self._geos_cross_projection_areas()
+        scene = Scene()
+        scene["edge"] = self._geos_cross_projection_data(source_area)
+
+        baseline = scene.resample(target_area, resampler="bilinear", reduce_data=False)["edge"].compute()
+        reduced = scene.resample(target_area, resampler="bilinear", reduce_data=True)["edge"].compute()
+
+        slice_x, slice_y = source_area.get_area_slices(target_area)
+        tight_data = scene["edge"].isel(x=slice_x, y=slice_y)
+        tight_area = source_area[slice_y, slice_x]
+        tight = XArrayBilinearResampler(
+            tight_area,
+            target_area,
+            50000,
+            reduce_data=False,
+        ).resample(tight_data).compute()
+
+        baseline_values = baseline.values
+        reduced_values = reduced.values
+        tight_values = tight.values
+        finite_both = np.isfinite(reduced_values) & np.isfinite(baseline_values)
+        max_abs_diff = np.max(np.abs(reduced_values[finite_both] - baseline_values[finite_both]))
+        np.testing.assert_array_equal(np.isfinite(reduced_values), np.isfinite(baseline_values))
+        assert max_abs_diff < 0.02
+        assert np.count_nonzero(np.isfinite(baseline_values) & ~np.isfinite(tight_values)) == 4
+
     def test_resample_ancillary(self):
         """Test that the Scene reducing data does not affect final output."""
         from pyresample.geometry import AreaDefinition