Add pass through of xr compute, persist and chunk to Scene

satpy/scene.py

@@ -1143,6 +1143,47 @@ def save_datasets(self, writer=None, filename=None, datasets=None, compute=True,
                                           **kwargs)
         return writer.save_datasets(dataarrays, compute=compute, **save_kwargs)
 
+    def compute(self, **kwargs):
+        """Call `compute` on all Scene data arrays.
+
+        See :meth:`xarray.DataArray.compute` for more details.
+        Note that this will convert the contents of the DataArray to numpy arrays which
+        may not work with all parts of Satpy which may expect dask arrays.
+        """
+        from dask import compute
+        new_scn = self.copy()
+        datasets = compute(*(new_scn._datasets.values()), **kwargs)
+
+        for i, k in enumerate(new_scn._datasets.keys()):
+            new_scn[k] = datasets[i]
+
+        return new_scn
+
+    def persist(self, **kwargs):
+        """Call `persist` on all Scene data arrays.
+
+        See :meth:`xarray.DataArray.persist` for more details.
+        """
+        from dask import persist
+        new_scn = self.copy()
+        datasets = persist(*(new_scn._datasets.values()), **kwargs)
+
+        for i, k in enumerate(new_scn._datasets.keys()):
+            new_scn[k] = datasets[i]
+
+        return new_scn
+
+    def chunk(self, **kwargs):
+        """Call `chunk` on all Scene  data arrays.
+
+        See :meth:`xarray.DataArray.chunk` for more details.
+        """
+        new_scn = self.copy()
+        for k in new_scn._datasets.keys():
+            new_scn[k] = new_scn[k].chunk(**kwargs)
+
+        return new_scn
+
     @staticmethod
     def _get_writer_by_ext(extension):
         """Find the writer matching the ``extension``.

satpy/tests/test_scene.py

@@ -1339,6 +1339,31 @@ def test_available_comps_no_deps(self):
         available_comp_ids = scene.available_composite_ids()
         assert make_cid(name='static_image') in available_comp_ids
 
+    def test_compute_pass_through(self):
+        """Test pass through of xarray compute."""
+        import numpy as np
+        scene = Scene(filenames=['fake1_1.txt'], reader='fake1')
+        scene.load(['ds1'])
+        scene = scene.compute()
+        assert isinstance(scene['ds1'].data, np.ndarray)
+
+    def test_persist_pass_through(self):
+        """Test pass through of xarray persist."""
+        from dask.array.utils import assert_eq
+        scene = Scene(filenames=['fake1_1.txt'], reader='fake1')
+        scene.load(['ds1'])
+        scenep = scene.persist()
+        assert_eq(scene['ds1'].data, scenep['ds1'].data)
+        assert set(scenep['ds1'].data.dask).issubset(scene['ds1'].data.dask)
+        assert len(scenep["ds1"].data.dask) == scenep['ds1'].data.npartitions
+
+    def test_chunk_pass_through(self):
+        """Test pass through of xarray chunk."""
+        scene = Scene(filenames=['fake1_1.txt'], reader='fake1')
+        scene.load(['ds1'])
+        scene = scene.chunk(chunks=2)
+        assert scene['ds1'].data.chunksize == (2, 2)
+
 
 class TestSceneResampling:
     """Test resampling a Scene to another Scene object."""