Convolve_2d gpu fixes (#702)

* fix indexing * fix indexing * fix convolve 2d cupy case * correct tests * clean
makepath · May 4, 2022 · 9a5cf1e · 9a5cf1e
1 parent faf501b
commit 9a5cf1e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 14 deletions.
diff --git a/xrspatial/convolution.py b/xrspatial/convolution.py
@@ -351,13 +351,13 @@ def _convolve_2d_cuda(data, kernel, out):
         return
 
     # The out at coordinates (i, j) is equal to
-    # sum_{k, h} kernel[k, h] * data[i - k + delta_rows, j - h + delta_cols]
+    # sum_{k, h} kernel[k, h] * data[i + k - delta_rows, j + h - delta_cols]
     # with k and h going through the whole kernel array:
     s = 0
     for k in range(kernel.shape[0]):
         for h in range(kernel.shape[1]):
-            i_k = i - k + delta_rows
-            j_h = j - h + delta_cols
+            i_k = i + k - delta_rows
+            j_h = j + h - delta_cols
             # (-4-) Check if (i_k, j_h) coordinates are inside the array:
             if (i_k >= 0) and (i_k < data_rows) and \
                     (j_h >= 0) and (j_h < data_cols):
@@ -496,8 +496,8 @@ def convolution_2d(agg, kernel, name='convolution_2d'):
         >>> convolved_agg
         <xarray.DataArray 'convolution_2d' (dim_0: 4, dim_1: 6)>
         array([[ nan,  nan,  nan,  nan,  nan,  nan],
-               [ nan,  56.,  64.,  72.,  80.,  nan],
-               [ nan, 104., 112., 120., 128.,  nan],
+               [ nan,  50.,  58.,  66.,  74.,  nan],
+               [ nan,  98., 106., 114., 122.,  nan],
                [ nan,  nan,  nan,  nan,  nan,  nan]], dtype=float32)
         Dimensions without coordinates: dim_0, dim_1
     """

diff --git a/xrspatial/tests/test_focal.py b/xrspatial/tests/test_focal.py
@@ -135,6 +135,20 @@ def convolution_kernel_annulus_2_2_1():
     return expected_result
 
 
+@pytest.fixture
+def convolution_custom_kernel():
+    kernel = np.array([[1, 0, 0], [1, 1, 0], [1, 0, 0]])
+    expected_result = np.array([
+        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
+        [np.nan, 2., 3., 3., 4., np.nan],
+        [np.nan, 4., np.nan, np.nan, np.nan, np.nan],
+        [np.nan, 4., np.nan, np.nan, np.nan, np.nan],
+        [np.nan, 4., np.nan, np.nan, np.nan, np.nan],
+        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
+    ])
+    return kernel, expected_result
+
+
 def test_kernel_custom_kernel_invalid_type():
     kernel = [1, 0, 0]  # only arrays are accepted, not lists
     with pytest.raises(ValueError):
@@ -159,16 +173,18 @@ def test_kernel(kernel_circle_1_1_1, kernel_annulus_2_2_2_1):
 
 def test_convolution_numpy(
     convolve_2d_data,
+    convolution_custom_kernel,
     kernel_circle_1_1_1,
     convolution_kernel_circle_1_1_1,
     kernel_annulus_2_2_2_1,
     convolution_kernel_annulus_2_2_1
 ):
-    kernel_custom = np.ones((1, 1))
+    kernel_custom, expected_result_custom = convolution_custom_kernel
     result_kernel_custom = convolve_2d(convolve_2d_data, kernel_custom)
     assert isinstance(result_kernel_custom, np.ndarray)
-    # kernel is [[1]], thus the result equals input data
-    np.testing.assert_allclose(result_kernel_custom, convolve_2d_data, equal_nan=True)
+    np.testing.assert_allclose(
+        result_kernel_custom, expected_result_custom, equal_nan=True
+    )
 
     result_kernel_circle = convolve_2d(convolve_2d_data, kernel_circle_1_1_1)
     assert isinstance(result_kernel_circle, np.ndarray)
@@ -185,17 +201,20 @@ def test_convolution_numpy(
 
 def test_convolution_dask_numpy(
     convolve_2d_data,
+    convolution_custom_kernel,
     kernel_circle_1_1_1,
     convolution_kernel_circle_1_1_1,
     kernel_annulus_2_2_2_1,
     convolution_kernel_annulus_2_2_1
 ):
     dask_agg = create_test_raster(convolve_2d_data, backend='dask+numpy')
-    kernel_custom = np.ones((1, 1))
+
+    kernel_custom, expected_result_custom = convolution_custom_kernel
     result_kernel_custom = convolution_2d(dask_agg, kernel_custom)
     assert isinstance(result_kernel_custom.data, da.Array)
-    # kernel is [[1]], thus the result equals input data
-    np.testing.assert_allclose(result_kernel_custom.compute(), convolve_2d_data, equal_nan=True)
+    np.testing.assert_allclose(
+        result_kernel_custom.compute(), expected_result_custom, equal_nan=True
+    )
 
     result_kernel_circle = convolution_2d(dask_agg, kernel_circle_1_1_1)
     assert isinstance(result_kernel_circle.data, da.Array)
@@ -213,6 +232,7 @@ def test_convolution_dask_numpy(
 @cuda_and_cupy_available
 def test_2d_convolution_gpu(
     convolve_2d_data,
+    convolution_custom_kernel,
     kernel_circle_1_1_1,
     convolution_kernel_circle_1_1_1,
     kernel_annulus_2_2_2_1,
@@ -221,11 +241,12 @@ def test_2d_convolution_gpu(
     import cupy
     cupy_data = cupy.asarray(convolve_2d_data)
 
-    kernel_custom = np.ones((1, 1))
+    kernel_custom, expected_result_custom = convolution_custom_kernel
     result_kernel_custom = convolve_2d(cupy_data, kernel_custom)
     assert isinstance(result_kernel_custom, cupy.ndarray)
-    # kernel is [[1]], thus the result equals input data
-    np.testing.assert_allclose(result_kernel_custom.get(), convolve_2d_data, equal_nan=True)
+    np.testing.assert_allclose(
+        result_kernel_custom.get(), expected_result_custom, equal_nan=True
+    )
 
     result_kernel_circle = convolve_2d(cupy_data, kernel_circle_1_1_1)
     assert isinstance(result_kernel_circle, cupy.ndarray)