Update cuspatial

holoviz · May 6, 2024 · ea7522a · ea7522a
1 parent a13df21
commit ea7522a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 11 deletions.
diff --git a/holoviews/element/selection.py b/holoviews/element/selection.py
@@ -80,6 +80,38 @@ def spatial_select_gridded(xvals, yvals, geometry):
         sel_mask = spatial_select_columnar(xvals.flatten(), yvals.flatten(), geometry)
         return sel_mask.reshape(xvals.shape)
 
+def _cuspatial_old(xvals, yvals, geometry):
+    import cudf
+    import cuspatial
+
+    result = cuspatial.point_in_polygon(
+        xvals,
+        yvals,
+        cudf.Series([0], index=["selection"]),
+        [0],
+        geometry[:, 0],
+        geometry[:, 1],
+    )
+    return result.values
+
+
+def _cuspatial_new(xvals, yvals, geometry):
+    import cudf
+    import cuspatial
+    import geopandas
+    from shapely.geometry import Polygon
+
+    df = cudf.DataFrame({'x':xvals, 'y':yvals})
+    points = cuspatial.GeoSeries.from_points_xy(
+       df.interleave_columns().astype('float')
+    )
+    polygons = cuspatial.GeoSeries(
+       geopandas.GeoSeries(Polygon(geometry)), index=["selection"]
+    )
+    result = cuspatial.point_in_polygon(points,polygons)
+    return result.values.ravel()
+
+
 def spatial_select_columnar(xvals, yvals, geometry, geom_method=None):
     if 'cudf' in sys.modules:
         import cudf
@@ -88,16 +120,10 @@ def spatial_select_columnar(xvals, yvals, geometry, geom_method=None):
             xvals = xvals.values.astype('float')
             yvals = yvals.values.astype('float')
             try:
-                import cuspatial
-                result = cuspatial.point_in_polygon(  # TODO: now only takes two arguments
-                    xvals,
-                    yvals,
-                    cudf.Series([0], index=["selection"]),
-                    [0],
-                    geometry[:, 0],
-                    geometry[:, 1],
-                )
-                return result.values
+                try:
+                    return _cuspatial_old(xvals, yvals, geometry)
+                except TypeError:
+                    return _cuspatial_new(xvals, yvals, geometry)
             except ImportError:
                 xvals = cp.asnumpy(xvals)
                 yvals = cp.asnumpy(yvals)

diff --git a/holoviews/tests/element/test_selection.py b/holoviews/tests/element/test_selection.py
@@ -692,7 +692,6 @@ def test_cudf(self, geometry, pt_mask, pandas_df, _method, unimport):
             mask = spatial_select_columnar(df.x, df.y, geometry, _method)
             assert np.array_equal(cp.asnumpy(mask), pt_mask)
 
-        @pytest.mark.xfail(reason='cuspatial.point_in_polygon API has changed')
         @pytest.mark.gpu
         def test_cuspatial(self, geometry, pt_mask, pandas_df, _method):
             import cudf