Multiscene blend with weights

satpy/multiscene.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (c) 2016-2019 Satpy developers
+# Copyright (c) 2016-2019, 2022 Satpy developers
 #
 # This file is part of satpy.
 #
@@ -45,6 +45,16 @@
 log = logging.getLogger(__name__)
 
 
+def weighted(datasets, weights=None):
+    """Blend datasets using weights."""
+    indices = da.argmax(da.dstack(weights), axis=-1)
+    dims = datasets[0].dims
+    attrs = datasets[0].attrs
+    weighted_array = xr.DataArray(da.choose(indices, datasets),
+                                  dims=dims, attrs=attrs)
+    return weighted_array
+
+
 def stack(datasets):
     """Overlay series of datasets on top of each other."""
     base = datasets[0].copy()
@@ -339,7 +349,7 @@ def resample(self, destination=None, **kwargs):
         """Resample the multiscene."""
         return self._generate_scene_func(self._scenes, 'resample', True, destination=destination, **kwargs)
 
-    def blend(self, blend_function=stack):
+    def blend(self, blend_function=stack, **kwargs):
         """Blend the datasets into one scene.
 
         Reduce the :class:`MultiScene` to a single :class:`~satpy.scene.Scene`.  Datasets
@@ -364,7 +374,7 @@ def blend(self, blend_function=stack):
         common_datasets = self.shared_dataset_ids
         for ds_id in common_datasets:
             datasets = [scn[ds_id] for scn in self.scenes if ds_id in scn]
-            new_scn[ds_id] = blend_function(datasets)
+            new_scn[ds_id] = blend_function(datasets, **kwargs)
 
         return new_scn
 
@@ -508,7 +518,7 @@ def _get_single_frame(self, ds, enh_args, fill_value):
         enh_args = enh_args.copy()  # don't change caller's dict!
         if "decorate" in enh_args:
             enh_args["decorate"] = self._format_decoration(
-                    ds, enh_args["decorate"])
+                ds, enh_args["decorate"])
         img = get_enhanced_image(ds, **enh_args)
         data, mode = img.finalize(fill_value=fill_value)
         if data.ndim == 3:
@@ -632,7 +642,7 @@ def _get_writers_and_frames(
             info_datasets = [scn.get(dataset_id) for scn in info_scenes]
             this_fn, shape, this_fill = self._get_animation_info(info_datasets, filename, fill_value=fill_value)
             data_to_write = self._get_animation_frames(
-                    all_datasets, shape, this_fill, ignore_missing, enh_args)
+                all_datasets, shape, this_fill, ignore_missing, enh_args)
 
             writer = imageio.get_writer(this_fn, **imio_args)
             frames[dataset_id] = data_to_write
@@ -703,8 +713,8 @@ def save_animation(self, filename, datasets=None, fps=10, fill_value=None,
             raise ImportError("Missing required 'imageio' library")
 
         (writers, frames) = self._get_writers_and_frames(
-                filename, datasets, fill_value, ignore_missing,
-                enh_args, imio_args={"fps": fps, **kwargs})
+            filename, datasets, fill_value, ignore_missing,
+            enh_args, imio_args={"fps": fps, **kwargs})
 
         client = self._get_client(client=client)
         # get an ordered list of frames

satpy/tests/test_multiscene.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (c) 2018 Satpy developers
+# Copyright (c) 2018, 2022 Satpy developers
 #
 # This file is part of satpy.
 #
@@ -24,8 +24,10 @@
 from datetime import datetime
 from unittest import mock
 
+import dask.array as da
 import pytest
 import xarray as xr
+from pyresample.geometry import AreaDefinition
 
 from satpy import DataQuery
 from satpy.dataset.dataid import DataID, ModifierTuple, WavelengthRange
@@ -166,7 +168,7 @@ def test_from_files(self):
             "OR_ABI-L1b-RadC-M3C01_G16_s20171171517203_e20171171519577_c20171171520019.nc",
             "OR_ABI-L1b-RadC-M3C01_G16_s20171171522203_e20171171524576_c20171171525020.nc",
             "OR_ABI-L1b-RadC-M3C01_G16_s20171171527203_e20171171529576_c20171171530017.nc",
-            ]
+        ]
         input_files_glm = [
             "OR_GLM-L2-GLMC-M3_G16_s20171171500000_e20171171501000_c20380190314080.nc",
             "OR_GLM-L2-GLMC-M3_G16_s20171171501000_e20171171502000_c20380190314080.nc",
@@ -179,9 +181,9 @@ def test_from_files(self):
         ]
         with mock.patch('satpy.multiscene.Scene') as scn_mock:
             mscn = MultiScene.from_files(
-                    input_files_abi,
-                    reader='abi_l1b',
-                    scene_kwargs={"reader_kwargs": {}})
+                input_files_abi,
+                reader='abi_l1b',
+                scene_kwargs={"reader_kwargs": {}})
             assert len(mscn.scenes) == 6
             calls = [mock.call(
                 filenames={'abi_l1b': [in_file_abi]},
@@ -192,11 +194,11 @@ def test_from_files(self):
             scn_mock.reset_mock()
             with pytest.warns(DeprecationWarning):
                 mscn = MultiScene.from_files(
-                        input_files_abi + input_files_glm,
-                        reader=('abi_l1b', "glm_l2"),
-                        group_keys=["start_time"],
-                        ensure_all_readers=True,
-                        time_threshold=30)
+                    input_files_abi + input_files_glm,
+                    reader=('abi_l1b', "glm_l2"),
+                    group_keys=["start_time"],
+                    ensure_all_readers=True,
+                    time_threshold=30)
             assert len(mscn.scenes) == 2
             calls = [mock.call(
                 filenames={'abi_l1b': [in_file_abi], 'glm_l2': [in_file_glm]})
@@ -206,11 +208,11 @@ def test_from_files(self):
             scn_mock.assert_has_calls(calls)
             scn_mock.reset_mock()
             mscn = MultiScene.from_files(
-                    input_files_abi + input_files_glm,
-                    reader=('abi_l1b', "glm_l2"),
-                    group_keys=["start_time"],
-                    ensure_all_readers=False,
-                    time_threshold=30)
+                input_files_abi + input_files_glm,
+                reader=('abi_l1b', "glm_l2"),
+                group_keys=["start_time"],
+                ensure_all_readers=False,
+                time_threshold=30)
             assert len(mscn.scenes) == 12
 
 
@@ -531,7 +533,7 @@ def test_crop(self):
             x_size // 2,
             y_size // 2,
             area_extent,
-            )
+        )
         scene1["1"] = DataArray(np.zeros((y_size, x_size)))
         scene1["2"] = DataArray(np.zeros((y_size, x_size)), dims=('y', 'x'))
         scene1["3"] = DataArray(np.zeros((y_size, x_size)), dims=('y', 'x'),
@@ -562,19 +564,29 @@ def setUp(self):
         import dask.array as da
         import xarray as xr
         from pyresample.geometry import AreaDefinition
+        shape = (8, 12)
         area = AreaDefinition('test', 'test', 'test',
                               {'proj': 'geos', 'lon_0': -95.5, 'h': 35786023.0},
-                              2, 2, [-200, -200, 200, 200])
-        ds1 = xr.DataArray(da.zeros((2, 2), chunks=-1), dims=('y', 'x'),
+                              shape[1], shape[0], [-200, -200, 200, 200])
+        ds1 = xr.DataArray(da.ones(shape, chunks=-1), dims=('y', 'x'),
                            attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
         self.ds1 = ds1
-        ds2 = xr.DataArray(da.zeros((2, 2), chunks=-1), dims=('y', 'x'),
+        wgt1 = xr.DataArray(da.ones(shape, chunks=-1), dims=('y', 'x'),
+                            attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+        self.ds1_wgt = wgt1
+        ds2 = xr.DataArray(da.ones(shape, chunks=-1) * 2, dims=('y', 'x'),
                            attrs={'start_time': datetime(2018, 1, 1, 1, 0, 0), 'area': area})
         self.ds2 = ds2
-        ds3 = xr.DataArray(da.zeros((2, 2), chunks=-1), dims=('y', 'time'),
+        wgt2 = xr.DataArray(da.zeros(shape, chunks=-1), dims=('y', 'x'),
+                            attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+        self.line = 2
+        wgt2[self.line, :] = 2
+        self.ds2_wgt = wgt2
+
+        ds3 = xr.DataArray(da.zeros(shape, chunks=-1), dims=('y', 'time'),
                            attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
         self.ds3 = ds3
-        ds4 = xr.DataArray(da.zeros((2, 2), chunks=-1), dims=('y', 'time'),
+        ds4 = xr.DataArray(da.zeros(shape, chunks=-1), dims=('y', 'time'),
                            attrs={'start_time': datetime(2018, 1, 1, 1, 0, 0), 'area': area})
         self.ds4 = ds4
 
@@ -597,6 +609,78 @@ def test_timeseries(self):
         self.assertTupleEqual((self.ds3.shape[0], self.ds3.shape[1]+self.ds4.shape[1]), res2.shape)
 
 
+@pytest.fixture
+def datasets_and_weights():
+    """X-Array datasets with area definition plus weights for input to tests."""
+    shape = (8, 12)
+    area = AreaDefinition('test', 'test', 'test',
+                          {'proj': 'geos', 'lon_0': -95.5, 'h': 35786023.0},
+                          shape[1], shape[0], [-200, -200, 200, 200])
+    ds1 = xr.DataArray(da.ones(shape, chunks=-1), dims=('y', 'x'),
+                       attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+    ds2 = xr.DataArray(da.ones(shape, chunks=-1) * 2, dims=('y', 'x'),
+                       attrs={'start_time': datetime(2018, 1, 1, 1, 0, 0), 'area': area})
+    ds3 = xr.DataArray(da.ones(shape, chunks=-1) * 3, dims=('y', 'x'),
+                       attrs={'start_time': datetime(2018, 1, 1, 1, 0, 0), 'area': area})
+
+    wgt1 = xr.DataArray(da.ones(shape, chunks=-1), dims=('y', 'x'),
+                        attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+    wgt2 = xr.DataArray(da.zeros(shape, chunks=-1), dims=('y', 'x'),
+                        attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+    wgt3 = xr.DataArray(da.zeros(shape, chunks=-1), dims=('y', 'x'),
+                        attrs={'start_time': datetime(2018, 1, 1, 0, 0, 0), 'area': area})
+
+    datastruct = {'shape': shape,
+                  'area': area,
+                  'datasets': [ds1, ds2, ds3],
+                  'weights': [wgt1, wgt2, wgt3]}
+    return datastruct
+
+
+@pytest.mark.parametrize(('line', 'column',),
+                         [(2, 3), (4, 5)]
+                         )
+def test_blend_function_weighted(datasets_and_weights, line, column):
+    """Test the 'weighted' function."""
+    from satpy.multiscene import weighted
+
+    input_data = datasets_and_weights
+
+    input_data['weights'][1][line, :] = 2
+    input_data['weights'][2][:, column] = 2
+
+    blend_result = weighted(input_data['datasets'], input_data['weights'])
+
+    ds1 = input_data['datasets'][0]
+    ds2 = input_data['datasets'][1]
+    ds3 = input_data['datasets'][2]
+    expected = ds1.copy()
+    expected[:, column] = ds3[:, column]
+    expected[line, :] = ds2[line, :]
+
+    xr.testing.assert_equal(blend_result.compute(), expected.compute())
+    assert expected.attrs == blend_result.attrs
+
+
+def test_blend_function_stack(datasets_and_weights):
+    """Test the 'stack' function."""
+    from satpy.multiscene import stack
+
+    input_data = datasets_and_weights
+
+    ds1 = input_data['datasets'][0]
+    ds2 = input_data['datasets'][1]
+
+    res = stack([ds1, ds2])
+    expected = ds2.copy()
+
+    xr.testing.assert_equal(res.compute(), expected.compute())
+    # assert expected.attrs == res.attrs
+    # FIXME! Looks like the attributes are taken from the first dataset. Should
+    # be like that?  So in this case the datetime is different from "expected"
+    # (= in this case the last dataset in the stack, the one on top)
+
+
 @mock.patch('satpy.multiscene.get_enhanced_image')
 def test_save_mp4(smg, tmp_path):
     """Save a series of fake scenes to an mp4 video."""
@@ -656,11 +740,11 @@ def test_save_mp4(smg, tmp_path):
     with mock.patch('satpy.multiscene.imageio.get_writer') as get_writer:
         get_writer.return_value = writer_mock
         mscn.save_animation(
-                fn, client=False,
-                enh_args={"decorate": {
-                    "decorate": [{
-                        "text": {
-                            "txt":
+            fn, client=False,
+            enh_args={"decorate": {
+                "decorate": [{
+                    "text": {
+                        "txt":
                             "Test {start_time:%Y-%m-%d %H:%M} - "
                             "{end_time:%Y-%m-%d %H:%M}"}}]}})
     assert writer_mock.append_data.call_count == 2 + 2