Add xarray/dask bilinear resampling

satpy/resample.py

@@ -38,6 +38,7 @@
     "nearest", "Nearest Neighbor", :class:`~satpy.resample.KDTreeResampler`
     "ewa", "Elliptical Weighted Averaging", :class:`~satpy.resample.EWAResampler`
     "native", "Native", :class:`~satpy.resample.NativeResampler`
+    "bilinear", "Bilinear", :class`~satpy.resample.BilinearResampler`
 
 The resampling algorithm used can be specified with the ``resampler`` keyword
 argument and defaults to ``nearest``:
@@ -84,7 +85,7 @@
 SatPy will do its best to reuse calculations performed to resample datasets,
 but it can only do this for the current processing and will lose this
 information when the process/script ends. Some resampling algorithms, like
-``nearest``, can benefit by caching intermediate data on disk in the directory
+``nearest`` and ``bilinear``, can benefit by caching intermediate data on disk in the directory
 specified by `cache_dir` and using it next time. This is most beneficial with
 geostationary satellite data where the locations of the source data and the
 target pixels don't change over time.
@@ -137,10 +138,10 @@
 import dask
 import dask.array as da
 
-from pyresample.bilinear import get_bil_info, get_sample_from_bil_info
 from pyresample.ewa import fornav, ll2cr
 from pyresample.geometry import SwathDefinition, AreaDefinition
 from pyresample.kd_tree import XArrayResamplerNN
+from pyresample.bilinear.xarr import XArrayResamplerBilinear
 from satpy import CHUNK_SIZE
 from satpy.config import config_search_paths, get_config_path
 
@@ -587,76 +588,84 @@ def compute(self, data, cache_id=None, fill_value=0, weight_count=10000,
 
 
 class BilinearResampler(BaseResampler):
+
     """Resample using bilinear."""
 
-    def precompute(self, mask=None, radius_of_influence=50000,
-                   cache_dir=None, **kwargs):
+    def __init__(self, source_geo_def, target_geo_def):
+        super(BilinearResampler, self).__init__(source_geo_def, target_geo_def)
+        self.resampler = None
+
+    def precompute(self, mask=None, radius_of_influence=50000, epsilon=0,
+                   reduce_data=True, nprocs=1,
+                   cache_dir=False, **kwargs):
         """Create bilinear coefficients and store them for later use.
 
         Note: The `mask` keyword should be provided if geolocation may be valid
         where data points are invalid. This defaults to the `mask` attribute of
         the `data` numpy masked array passed to the `resample` method.
         """
 
-        raise NotImplementedError("Bilinear interpolation has not been "
-                                  "converted to XArray/Dask yet.")
-
         del kwargs
-        source_geo_def = self.source_geo_def
 
-        bil_hash = self.get_hash(source_geo_def=source_geo_def,
-                                 radius_of_influence=radius_of_influence,
-                                 mode="bilinear")
+        if self.resampler is None:
+            kwargs = dict(source_geo_def=self.source_geo_def,
+                          target_geo_def=self.target_geo_def,
+                          radius_of_influence=radius_of_influence,
+                          neighbours=32,
+                          epsilon=epsilon,
+                          reduce_data=reduce_data)
 
-        filename = self._create_cache_filename(cache_dir, bil_hash)
-        self._read_params_from_cache(cache_dir, bil_hash, filename)
+            self.resampler = XArrayResamplerBilinear(**kwargs)
 
-        if self.cache is not None:
-            LOG.debug("Loaded bilinear parameters")
-            return self.cache
-        else:
-            LOG.debug("Computing bilinear parameters")
+            try:
+                self.load_bil_info(cache_dir, **kwargs)
+                LOG.debug("Loaded bilinear parameters")
+            except IOError:
+                LOG.debug("Computing bilinear parameters")
+                self.resampler.get_bil_info()
+                self.save_bil_info(cache_dir, **kwargs)
 
-        bilinear_t, bilinear_s, input_idxs, idx_arr = get_bil_info(source_geo_def, self.target_geo_def,
-                                                                   radius_of_influence, neighbours=32,
-                                                                   masked=False)
-        self.cache = {'bilinear_s': bilinear_s,
-                      'bilinear_t': bilinear_t,
-                      'input_idxs': input_idxs,
-                      'idx_arr': idx_arr}
+    def load_bil_info(self, cache_dir, **kwargs):
 
-        self._update_caches(bil_hash, cache_dir, filename)
+        if cache_dir:
+            filename = self._create_cache_filename(cache_dir,
+                                                   prefix='resample_lut_bil_',
+                                                   **kwargs)
+            cache = np.load(filename)
+            for elt in ['bilinear_s', 'bilinear_t', 'valid_input_index',
+                        'index_array']:
+                if isinstance(cache[elt], tuple):
+                    setattr(self.resampler, elt, cache[elt][0])
+                else:
+                    setattr(self.resampler, elt, cache[elt])
+            cache.close()
+        else:
+            raise IOError
 
-        return self.cache
+    def save_bil_info(self, cache_dir, **kwargs):
+        if cache_dir:
+            filename = self._create_cache_filename(cache_dir,
+                                                   prefix='resample_lut_bil_',
+                                                   **kwargs)
+            LOG.info('Saving kd_tree neighbour info to %s', filename)
+            cache = {'bilinear_s': self.resampler.bilinear_s,
+                     'bilinear_t': self.resampler.bilinear_t,
+                     'valid_input_index': self.resampler.valid_input_index,
+                     'index_array': self.resampler.index_array}
+
+            np.savez(filename, **cache)
 
     def compute(self, data, fill_value=None, **kwargs):
         """Resample the given data using bilinear interpolation"""
         del kwargs
 
+        if fill_value is None:
+            fill_value = data.attrs.get('_FillValue')
         target_shape = self.target_geo_def.shape
-        if data.ndim == 3:
-            output_shape = list(target_shape)
-            output_shape.append(data.shape[-1])
-            res = np.zeros(output_shape, dtype=data.dtype)
-            for i in range(data.shape[-1]):
-                res[:, :, i] = get_sample_from_bil_info(data[:, :, i].ravel(),
-                                                        self.cache[
-                                                            'bilinear_t'],
-                                                        self.cache[
-                                                            'bilinear_s'],
-                                                        self.cache[
-                                                            'input_idxs'],
-                                                        self.cache['idx_arr'],
-                                                        output_shape=target_shape)
 
-        else:
-            res = get_sample_from_bil_info(data.ravel(),
-                                           self.cache['bilinear_t'],
-                                           self.cache['bilinear_s'],
-                                           self.cache['input_idxs'],
-                                           self.cache['idx_arr'],
-                                           output_shape=target_shape)
-        res = np.ma.masked_invalid(res)
+        res = self.resampler.get_sample_from_bil_info(data,
+                                                      fill_value=fill_value,
+                                                      output_shape=target_shape)
 
         return res
 
@@ -796,7 +805,7 @@ def compute(self, data, expand=True, **kwargs):
 RESAMPLERS = {"kd_tree": KDTreeResampler,
               "nearest": KDTreeResampler,
               "ewa": EWAResampler,
-              # "bilinear": BilinearResampler,
+              "bilinear": BilinearResampler,
               "native": NativeResampler,
               }
 

satpy/tests/test_resample.py

@@ -384,6 +384,47 @@ def test_expand_without_dims_4D(self):
         self.assertRaises(ValueError, resampler.resample, ds1)
 
 
+class TestBilinearResampler(unittest.TestCase):
+    """Test the bilinear resampler."""
+
+    @mock.patch('satpy.resample.BilinearResampler.load_bil_info')
+    @mock.patch('satpy.resample.np.savez')
+    @mock.patch('satpy.resample.np.load')
+    @mock.patch('satpy.resample.BilinearResampler._create_cache_filename')
+    @mock.patch('satpy.resample.XArrayResamplerBilinear')
+    def test_bil_resampling(self, resampler, create_filename, load, savez,
+                            load_bil_info):
+        """Test the bilinear resampler."""
+        import numpy as np
+        import dask.array as da
+        from satpy.resample import BilinearResampler
+        from pyresample.geometry import SwathDefinition
+        source_area = mock.MagicMock()
+        source_swath = SwathDefinition(
+            da.arange(5, chunks=5), da.arange(5, chunks=5))
+        target_area = mock.MagicMock()
+
+        # Test that bilinear resampling info calculation is called,
+        # and the info is saved
+        load_bil_info.side_effect = IOError()
+        resampler = BilinearResampler(source_swath, target_area)
+        resampler.precompute(
+            mask=da.arange(5, chunks=5).astype(np.bool))
+        resampler.resampler.get_bil_info.assert_called()
+        resampler.resampler.get_bil_info.assert_called_with()
+        self.assertFalse(len(savez.mock_calls), 1)
+        resampler.resampler.reset_mock()
+        load_bil_info.reset_mock()
+
+        # Test that get_sample_from_bil_info is called properly
+        data = mock.MagicMock()
+        data.name = 'foo'
+        data.data = [1, 2, 3]
+        fill_value = 8
+        resampler.compute(data, fill_value=fill_value)
+        resampler.resampler.get_sample_from_bil_info.assert_called_with(data, fill_value=fill_value, output_shape=target_area.shape)
+
+
 def suite():
     """The test suite for test_scene.
     """