Pytests for CuPy zonal stats (#658)

* adds unit-tests for cupy-zonal

* fix bug when using zone_id list

* fixes bug related to using zone_ids list

* removes comment-out code

* flake8 compatible

Author: kiliakis

xrspatial/tests/test_zonal.py

@@ -13,6 +13,8 @@
 from xrspatial import trim
 from xrspatial import crop
 from xrspatial.zonal import regions
+from xrspatial.utils import doesnt_have_cuda
+
 
 from xrspatial.tests.general_checks import create_test_raster
 
@@ -169,26 +171,35 @@ def check_results(backend, df_result, expected_results_dict):
         np.testing.assert_allclose(df_result[col], expected_results_dict[col])
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_default_stats(backend, data_zones, data_values_2d, result_default_stats):
+    if backend == 'cupy' and doesnt_have_cuda():
+        pytest.skip("CUDA Device not Available")
     df_result = stats(zones=data_zones, values=data_values_2d)
     check_results(backend, df_result, result_default_stats)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_zone_ids_stats(backend, data_zones, data_values_2d, result_zone_ids_stats):
+    if backend == 'cupy' and doesnt_have_cuda():
+        pytest.skip("CUDA Device not Available")
     zone_ids, expected_result = result_zone_ids_stats
-    df_result = stats(zones=data_zones, values=data_values_2d, zone_ids=zone_ids)
+    df_result = stats(zones=data_zones, values=data_values_2d,
+                      zone_ids=zone_ids)
     check_results(backend, df_result, expected_result)
 
 
-@pytest.mark.parametrize("backend", ['numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'cupy'])
 def test_custom_stats(backend, data_zones, data_values_2d, result_custom_stats):
-    # ---- custom stats (NumPy only) ----
+    # ---- custom stats (NumPy and CuPy only) ----
+    if backend == 'cupy' and doesnt_have_cuda():
+        pytest.skip("CUDA Device not Available")
+
     custom_stats = {
         'double_sum': _double_sum,
         'range': _range,
     }
+
     nodata_values, zone_ids, expected_result = result_custom_stats
     df_result = stats(
         zones=data_zones, values=data_values_2d, stats_funcs=custom_stats,
@@ -219,7 +230,8 @@ def test_percentage_crosstab_2d(backend, data_zones, data_values_2d, result_perc
 @pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
 def test_crosstab_3d(backend, data_zones, data_values_3d, result_crosstab_3d):
     layer, zone_ids, expected_result = result_crosstab_3d
-    df_result = crosstab(zones=data_zones, values=data_values_3d, zone_ids=zone_ids, layer=layer)
+    df_result = crosstab(zones=data_zones, values=data_values_3d,
+                         zone_ids=zone_ids, layer=layer)
     check_results(backend, df_result, expected_result)
 
 

xrspatial/zonal.py

@@ -324,16 +324,30 @@ def _stats_cupy(
     sorted_zones = sorted_zones[filter_values]
 
     # Now I need to find the unique zones, and zone breaks
-    unique_zones, unique_index = cupy.unique(sorted_zones, return_index=True)
+    unique_zones, unique_index, unique_counts = cupy.unique(
+        sorted_zones, return_index=True, return_counts=True)
 
     # Transfer to the host
     unique_index = unique_index.get()
-    if zone_ids is None:
-        unique_zones = unique_zones.get()
-    else:
+    unique_counts = unique_counts.get()
+    unique_zones = unique_zones.get()
+
+    if zone_ids is not None:
+        # We need to extract the index and element count
+        # only for the elements in zone_ids
+        unique_index_lst = []
+        unique_counts_lst = []
+        unique_zones = list(unique_zones)
+        for z in zone_ids:
+            try:
+                idx = unique_zones.index(z)
+                unique_index_lst.append(unique_index[idx])
+                unique_counts_lst.append(unique_counts[idx])
+            except ValueError:
+                continue
         unique_zones = zone_ids
-    # unique_zones = list(map(_to_int, unique_zones))
-    unique_zones = np.asarray(unique_zones)
+        unique_counts = unique_counts_lst
+        unique_index = unique_index_lst
 
     # stats columns
     stats_dict = {'zone': []}
@@ -347,11 +361,9 @@ def _stats_cupy(
             continue
 
         stats_dict['zone'].append(zone_id)
+
         # extract zone_values
-        if i < len(unique_zones) - 1:
-            zone_values = values_by_zone[unique_index[i]:unique_index[i+1]]
-        else:
-            zone_values = values_by_zone[unique_index[i]:]
+        zone_values = values_by_zone[unique_index[i]:unique_index[i]+unique_counts[i]]
 
         # apply stats on the zone data
         for j, stats in enumerate(stats_funcs):
@@ -362,7 +374,7 @@ def _stats_cupy(
 
             assert(len(result.shape) == 0)
 
-            stats_dict[stats].append(cupy.float(result))
+            stats_dict[stats].append(cupy.float_(result))
 
     stats_df = pd.DataFrame(stats_dict)
     stats_df.set_index("zone")