Fix dtype for swath -> area resampling with gradient search

pyresample/gradient/__init__.py

@@ -414,14 +414,15 @@ def parallel_gradient_search(data, src_x, src_y, dst_x, dst_y,
         else:
             is_pad = False
             res = dask.delayed(_gradient_resample_data)(
-                arr.astype(np.float64),
+                arr,
                 src_x[i], src_y[i],
                 src_gradient_xl[i], src_gradient_xp[i],
                 src_gradient_yl[i], src_gradient_yp[i],
                 dst_x[i], dst_y[i],
                 method=method)
             res = da.from_delayed(res, (num_bands, ) + dst_x[i].shape,
-                                  dtype=np.float64)
+                                  meta=np.array((), dtype=arr.dtype),
+                                  dtype=arr.dtype)
         if dst_mosaic_locations[i] in chunks:
             if not is_pad:
                 chunks[dst_mosaic_locations[i]].append(res)

pyresample/gradient/_gradient_search.pyx

@@ -23,18 +23,22 @@
 
 import numpy as np
 
-cimport numpy as np
-
-DTYPE = np.double
-ctypedef np.double_t DTYPE_t
 cimport cython
+cimport numpy as np
 from libc.math cimport fabs, isinf
 
+ctypedef fused data_type:
+    np.float64_t
+    np.float32_t
+
+ctypedef np.float64_t float_index
+float_index_dtype = np.float64
+
 np.import_array()
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline void nn(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, double dp, int lmax, int pmax, DTYPE_t[:] res) noexcept nogil:
+cdef inline void nn(const data_type[:, :, :] data, int l0, int p0, float_index dl, float_index dp, int lmax, int pmax, data_type[:] res) noexcept nogil:
     cdef int nnl, nnp
     cdef size_t z_size = res.shape[0]
     cdef size_t i
@@ -54,9 +58,9 @@ cdef inline void nn(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, doub
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline void bil(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, double dp, int lmax, int pmax, DTYPE_t[:] res) noexcept nogil:
+cdef inline void bil(const data_type[:, :, :] data, int l0, int p0, float_index dl, float_index dp, int lmax, int pmax, data_type[:] res) noexcept nogil:
     cdef int l_a, l_b, p_a, p_b
-    cdef double w_l, w_p
+    cdef float_index w_l, w_p
     cdef size_t z_size = res.shape[0]
     cdef size_t i
     if dl < 0:
@@ -84,26 +88,28 @@ cdef inline void bil(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, dou
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline void indices_xy(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, double dp, int lmax, int pmax, DTYPE_t[:] res) noexcept nogil:
+cdef inline void indices_xy(const data_type[:, :, :] data, int l0, int p0, float_index dl, float_index dp, int lmax, int pmax, data_type[:] res) noexcept nogil:
     cdef int nnl, nnp
     cdef size_t z_size = res.shape[0]
     cdef size_t i
     res[1] = dl + l0
     res[0] = dp + p0
 
-ctypedef void (*FN)(const DTYPE_t[:, :, :] data, int l0, int p0, double dl, double dp, int lmax, int pmax, DTYPE_t[:] res) noexcept nogil
+
+ctypedef void (*FN)(const data_type[:, :, :] data, int l0, int p0, float_index dl, float_index dp, int lmax, int pmax, data_type[:] res) noexcept nogil
+
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cpdef one_step_gradient_search(const DTYPE_t [:, :, :] data,
-                               DTYPE_t [:, :] src_x,
-                               DTYPE_t [:, :] src_y,
-                               DTYPE_t [:, :] xl,
-                               DTYPE_t [:, :] xp,
-                               DTYPE_t [:, :] yl,
-                               DTYPE_t [:, :] yp,
-                               DTYPE_t [:, :] dst_x,
-                               DTYPE_t [:, :] dst_y,
+cpdef one_step_gradient_search(const data_type[:, :, :] data,
+                               float_index [:, :] src_x,
+                               float_index [:, :] src_y,
+                               float_index [:, :] xl,
+                               float_index [:, :] xp,
+                               float_index [:, :] yl,
+                               float_index [:, :] yp,
+                               float_index [:, :] dst_x,
+                               float_index [:, :] dst_y,
                                str method='bilinear'):
     """Gradient search, simple case variant."""
     cdef FN fun
@@ -118,10 +124,14 @@ cpdef one_step_gradient_search(const DTYPE_t [:, :, :] data,
     cdef size_t y_size = dst_y.shape[0]
     cdef size_t x_size = dst_x.shape[1]
 
+    if data_type is double:
+        dtype = np.float64
+    else:
+        dtype = np.float32
 
     # output image array --> needs to be (lines, pixels) --> y,x
-    image = np.full([z_size, y_size, x_size], np.nan, dtype=DTYPE)
-    cdef DTYPE_t [:, :, :] image_view = image
+    image = np.full([z_size, y_size, x_size], np.nan, dtype=dtype)
+    cdef data_type[:, :, :] image_view = image
     with nogil:
         one_step_gradient_search_no_gil(data,
                                         src_x, src_y,
@@ -132,22 +142,23 @@ cpdef one_step_gradient_search(const DTYPE_t [:, :, :] data,
     # return the output image
     return image
 
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 @cython.cdivision(True)
-cdef void one_step_gradient_search_no_gil(const DTYPE_t[:, :, :] data,
-                                          const DTYPE_t[:, :] src_x,
-                                          const DTYPE_t[:, :] src_y,
-                                          const DTYPE_t[:, :] xl,
-                                          const DTYPE_t[:, :] xp,
-                                          const DTYPE_t[:, :] yl,
-                                          const DTYPE_t[:, :] yp,
-                                          const DTYPE_t[:, :] dst_x,
-                                          const DTYPE_t[:, :] dst_y,
+cdef void one_step_gradient_search_no_gil(const data_type[:, :, :] data,
+                                          const float_index[:, :] src_x,
+                                          const float_index[:, :] src_y,
+                                          const float_index[:, :] xl,
+                                          const float_index[:, :] xp,
+                                          const float_index[:, :] yl,
+                                          const float_index[:, :] yp,
+                                          const float_index[:, :] dst_x,
+                                          const float_index[:, :] dst_y,
                                           const size_t x_size,
                                           const size_t y_size,
                                           FN fun,
-                                          DTYPE_t[:, :, :] result_array) noexcept nogil:
+                                          data_type[:, :, :] result_array) noexcept nogil:
 
     # pixel max ---> data is expected in [lines, pixels]
     cdef int pmax = src_x.shape[1] - 1
@@ -161,7 +172,7 @@ cdef void one_step_gradient_search_no_gil(const DTYPE_t[:, :, :] data,
     # intermediate variables:
     cdef int l_a, l_b, p_a, p_b
     cdef size_t i, j, elt
-    cdef double dx, dy, d, dl, dp
+    cdef float_index dx, dy, d, dl, dp
     cdef int col_step = -1
     # number of iterations
     cdef int cnt = 0
@@ -218,16 +229,17 @@ cdef void one_step_gradient_search_no_gil(const DTYPE_t[:, :, :] data,
                     p0 = int(p0 + dp)
             j += col_step
 
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cpdef one_step_gradient_indices(DTYPE_t [:, :] src_x,
-                                DTYPE_t [:, :] src_y,
-                                DTYPE_t [:, :] xl,
-                                DTYPE_t [:, :] xp,
-                                DTYPE_t [:, :] yl,
-                                DTYPE_t [:, :] yp,
-                                DTYPE_t [:, :] dst_x,
-                                DTYPE_t [:, :] dst_y):
+cpdef one_step_gradient_indices(float_index [:, :] src_x,
+                                float_index [:, :] src_y,
+                                float_index [:, :] xl,
+                                float_index [:, :] xp,
+                                float_index [:, :] yl,
+                                float_index [:, :] yp,
+                                float_index [:, :] dst_x,
+                                float_index [:, :] dst_y):
     """Gradient search, simple case variant, returning float indices.
 
     This is appropriate for monotonous gradients only, i.e. not modis or viirs in satellite projection.
@@ -238,18 +250,19 @@ cpdef one_step_gradient_indices(DTYPE_t [:, :] src_x,
     cdef size_t y_size = dst_y.shape[0]
     cdef size_t x_size = dst_x.shape[1]
 
+
     # output indices arrays --> needs to be (lines, pixels) --> y,x
-    indices = np.full([2, y_size, x_size], np.nan, dtype=DTYPE)
-    cdef DTYPE_t [:, :, :] indices_view = indices
+    indices = np.full([2, y_size, x_size], np.nan, dtype=float_index_dtype)
+    cdef float_index [:, :, :] indices_view_result = indices
 
     # fake_data is not going to be used anyway as we just fill in the indices
-    cdef DTYPE_t [:, :, :] fake_data = np.full([1, 1, 1], np.nan, dtype=DTYPE)
+    cdef float_index [:, :, :] fake_data = np.full([1, 1, 1], np.nan, dtype=float_index_dtype)
 
     with nogil:
-        one_step_gradient_search_no_gil(fake_data,
-                                        src_x, src_y,
-                                        xl, xp, yl, yp,
-                                        dst_x, dst_y,
-                                        x_size, y_size,
-                                        indices_xy, indices_view)
+        one_step_gradient_search_no_gil[float_index](fake_data,
+                                                     src_x, src_y,
+                                                     xl, xp, yl, yp,
+                                                     dst_x, dst_y,
+                                                     x_size, y_size,
+                                                     indices_xy, indices_view_result)
     return indices

pyresample/test/test_gradient.py

@@ -249,28 +249,42 @@ def test_resample_area_to_area_3d_single_channel(self):
         assert res.shape == (1, ) + self.dst_area.shape
         assert np.allclose(res[0, :, :], 1.0)
 
-    def test_resample_swath_to_area_2d(self):
+    @pytest.mark.parametrize("input_dtype", (np.float32, np.float64))
+    def test_resample_swath_to_area_2d(self, input_dtype):
         """Resample swath to area, 2d."""
-        data = xr.DataArray(da.ones(self.src_swath.shape, dtype=np.float64),
+        data = xr.DataArray(da.ones(self.src_swath.shape, dtype=input_dtype),
                             dims=['y', 'x'])
         with np.errstate(invalid="ignore"):  # 'inf' space pixels cause runtime warnings
-            res = self.swath_resampler.compute(
-                data, method='bil').compute(scheduler='single-threaded')
-        assert res.shape == self.dst_area.shape
-        assert not np.all(np.isnan(res))
-
-    def test_resample_swath_to_area_3d(self):
+            res_xr = self.swath_resampler.compute(data, method='bil')
+            res_np = res_xr.compute(scheduler='single-threaded')
+
+        assert res_xr.dtype == data.dtype
+        assert res_np.dtype == data.dtype
+        assert res_xr.shape == self.dst_area.shape
+        assert res_np.shape == self.dst_area.shape
+        assert type(res_xr) is type(data)
+        assert type(res_xr.data) is type(data.data)
+        assert not np.all(np.isnan(res_np))
+
+    @pytest.mark.parametrize("input_dtype", (np.float32, np.float64))
+    def test_resample_swath_to_area_3d(self, input_dtype):
         """Resample area to area, 3d."""
         data = xr.DataArray(da.ones((3, ) + self.src_swath.shape,
-                                    dtype=np.float64) *
+                                    dtype=input_dtype) *
                             np.array([1, 2, 3])[:, np.newaxis, np.newaxis],
                             dims=['bands', 'y', 'x'])
         with np.errstate(invalid="ignore"):  # 'inf' space pixels cause runtime warnings
-            res = self.swath_resampler.compute(
-                data, method='bil').compute(scheduler='single-threaded')
-        assert res.shape == (3, ) + self.dst_area.shape
-        for i in range(res.shape[0]):
-            arr = np.ravel(res[i, :, :])
+            res_xr = self.swath_resampler.compute(data, method='bil')
+            res_np = res_xr.compute(scheduler='single-threaded')
+
+        assert res_xr.dtype == data.dtype
+        assert res_np.dtype == data.dtype
+        assert res_xr.shape == (3, ) + self.dst_area.shape
+        assert res_np.shape == (3, ) + self.dst_area.shape
+        assert type(res_xr) is type(data)
+        assert type(res_xr.data) is type(data.data)
+        for i in range(res_np.shape[0]):
+            arr = np.ravel(res_np[i, :, :])
             assert np.allclose(arr[np.isfinite(arr)], float(i + 1))
 
 
@@ -496,7 +510,7 @@ def test_resample_area_to_area_nn(self):
 
 
 class TestRBGradientSearchResamplerArea2Swath:
-    """Test RBGradientSearchResampler for the Swath to Area case."""
+    """Test RBGradientSearchResampler for the Area to Swath case."""
 
     def setup_method(self):
         """Set up the test case."""