Add support for GeoPandas lines and points (#1293)

* Add support for GeoPandas lines and points

* Linting

* Simplify

datashader/core.py

@@ -206,15 +206,19 @@ def points(self, source, x=None, y=None, agg=None, geometry=None):
                 x_range = self.x_range if self.x_range is not None else (None, None)
                 y_range = self.y_range if self.y_range is not None else (None, None)
                 source = source.cx_partitions[slice(*x_range), slice(*y_range)]
+                glyph = MultiPointGeometry(geometry)
             elif spatialpandas and isinstance(source, spatialpandas.GeoDataFrame):
-                pass
+                glyph = MultiPointGeometry(geometry)
+            elif (geopandas_source := self._source_from_geopandas(source)) is not None:
+                source = geopandas_source
+                from datashader.glyphs.points import MultiPointGeoPandas
+                glyph = MultiPointGeoPandas(geometry)
             else:
                 raise ValueError(
-                    "source must be an instance of spatialpandas.GeoDataFrame or \n"
-                    "spatialpandas.dask.DaskGeoDataFrame.\n"
-                    "  Received value of type {typ}".format(typ=type(source)))
-
-            glyph = MultiPointGeometry(geometry)
+                    "source must be an instance of spatialpandas.GeoDataFrame, "
+                    "spatialpandas.dask.DaskGeoDataFrame, geopandas.GeoDataFrame, or "
+                    "dask_geopandas.GeoDataFrame. Received objects of type {typ}".format(
+                        typ=type(source)))
 
         return bypixel(source, self, glyph, agg)
 
@@ -365,15 +369,20 @@ def line(self, source, x=None, y=None, agg=None, axis=0, geometry=None,
                 x_range = self.x_range if self.x_range is not None else (None, None)
                 y_range = self.y_range if self.y_range is not None else (None, None)
                 source = source.cx_partitions[slice(*x_range), slice(*y_range)]
+                glyph = LineAxis1Geometry(geometry)
             elif spatialpandas and isinstance(source, spatialpandas.GeoDataFrame):
-                pass
+                glyph = LineAxis1Geometry(geometry)
+            elif (geopandas_source := self._source_from_geopandas(source)) is not None:
+                source = geopandas_source
+                from datashader.glyphs.line import LineAxis1GeoPandas
+                glyph = LineAxis1GeoPandas(geometry)
             else:
                 raise ValueError(
-                    "source must be an instance of spatialpandas.GeoDataFrame or \n"
-                    "spatialpandas.dask.DaskGeoDataFrame.\n"
-                    "  Received value of type {typ}".format(typ=type(source)))
+                    "source must be an instance of spatialpandas.GeoDataFrame, "
+                    "spatialpandas.dask.DaskGeoDataFrame, geopandas.GeoDataFrame, or "
+                    "dask_geopandas.GeoDataFrame. Received objects of type {typ}".format(
+                        typ=type(source)))
 
-            glyph = LineAxis1Geometry(geometry)
         else:
             # Broadcast column specifications to handle cases where
             # x is a list and y is a string or vice versa

datashader/glyphs/__init__.py

@@ -8,6 +8,7 @@
     LinesAxis1YConstant,
     LinesAxis1Ragged,
     LineAxis1Geometry,
+    LineAxis1GeoPandas,
 )
 from .area import (  # noqa (API import)
     AreaToZeroAxis0,

datashader/glyphs/line.py

@@ -568,6 +568,42 @@ def extend(aggs, df, vt, bounds, plot_start=True):
         return extend
 
 
+class LineAxis1GeoPandas(_GeometryLike, _AntiAliasedLine):
+    # geopandas must be available for a GeoPandasLine to be created.
+    @property
+    def geom_dtypes(self):
+        from geopandas.array import GeometryDtype
+        return (GeometryDtype,)
+
+    @memoize
+    def _internal_build_extend(
+        self, x_mapper, y_mapper, info, append, line_width, antialias_stage_2,
+        antialias_stage_2_funcs,
+    ):
+        expand_aggs_and_cols = self.expand_aggs_and_cols(append)
+        draw_segment, antialias_stage_2_funcs = _line_internal_build_extend(
+            x_mapper, y_mapper, append, line_width, antialias_stage_2, antialias_stage_2_funcs,
+            expand_aggs_and_cols,
+        )
+        perform_extend_cpu = _build_extend_line_axis1_geopandas(
+            draw_segment, expand_aggs_and_cols, antialias_stage_2_funcs,
+        )
+        geometry_name = self.geometry
+
+        def extend(aggs, df, vt, bounds, plot_start=True):
+            sx, tx, sy, ty = vt
+            xmin, xmax, ymin, ymax = bounds
+            aggs_and_cols = aggs + info(df, aggs[0].shape[:2])
+            geom_array = df[geometry_name].array
+
+            perform_extend_cpu(
+                sx, tx, sy, ty, xmin, xmax, ymin, ymax,
+                geom_array, antialias_stage_2, *aggs_and_cols
+            )
+
+        return extend
+
+
 def _build_map_onto_pixel_for_line(x_mapper, y_mapper, want_antialias=False):
     @ngjit
     def map_onto_pixel_snap(sx, tx, sy, ty, xmin, xmax, ymin, ymax, x, y):
@@ -1764,3 +1800,171 @@ def extend_cpu_numba_antialias_2agg(
         return extend_cpu_antialias_2agg
     else:
         return extend_cpu
+
+
+def _build_extend_line_axis1_geopandas(draw_segment, expand_aggs_and_cols, antialias_stage_2_funcs):
+    if antialias_stage_2_funcs is not None:
+        aa_stage_2_accumulate, aa_stage_2_clear, aa_stage_2_copy_back = antialias_stage_2_funcs
+    use_2_stage_agg = antialias_stage_2_funcs is not None
+
+    # Lazy import shapely. Cannot get here if geopandas and shapely are not available.
+    import shapely
+
+    def _process_geometry(geometry):
+        ragged = shapely.to_ragged_array(geometry)
+        geometry_type = ragged[0]
+
+        if geometry_type not in (
+            shapely.GeometryType.LINESTRING, shapely.GeometryType.MULTILINESTRING,
+            shapely.GeometryType.MULTIPOLYGON, shapely.GeometryType.POLYGON,
+        ):
+            raise ValueError(
+                "Canvas.line supports GeoPandas geometry types of LINESTRING, MULTILINESTRING, "
+                f"MULTIPOLYGON and POLYGON, not {repr(geometry_type)}")
+
+        coords = ragged[1].ravel()
+
+        # Use type to decide whether geometry represents closed line loops or open list strips.
+        # Skip the last point for closed geometries so as not to double count the first/last point.
+        if geometry_type == shapely.GeometryType.LINESTRING:
+            offsets = ragged[2][0]
+            outer_offsets = np.arange(len(offsets))
+            closed_rings = False
+        elif geometry_type == shapely.GeometryType.MULTILINESTRING:
+            offsets, outer_offsets = ragged[2]
+            closed_rings = False
+        elif geometry_type == shapely.GeometryType.MULTIPOLYGON:
+            offsets, temp_offsets, outer_offsets = ragged[2]
+            outer_offsets = temp_offsets[outer_offsets]
+            closed_rings = True
+        else:  # geometry_type == shapely.GeometryType.POLYGON:
+            offsets, outer_offsets = ragged[2]
+            closed_rings = True
+
+        return coords, offsets, outer_offsets, closed_rings
+
+    def extend_cpu(
+        sx, tx, sy, ty, xmin, xmax, ymin, ymax, geometry, antialias_stage_2, *aggs_and_cols
+    ):
+        coords, offsets, outer_offsets, closed_rings = _process_geometry(geometry)
+        extend_cpu_numba(
+            sx, tx, sy, ty, xmin, xmax, ymin, ymax, coords, offsets, outer_offsets, closed_rings,
+            antialias_stage_2, *aggs_and_cols)
+
+    @ngjit
+    @expand_aggs_and_cols
+    def extend_cpu_numba(
+            sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, offsets, outer_offsets,
+            closed_rings, antialias_stage_2, *aggs_and_cols
+    ):
+        antialias = antialias_stage_2 is not None
+        buffer = np.empty(8) if antialias else None
+        n_multilines = len(outer_offsets) - 1
+        for i in range(n_multilines):
+            start0 = outer_offsets[i]
+            stop0 = outer_offsets[i + 1]
+
+            for j in range(start0, stop0):
+                start1 = offsets[j]
+                stop1 = offsets[j + 1]
+
+                for k in range(2*start1, 2*stop1 - 2, 2):
+                    x0 = values[k]
+                    y0 = values[k + 1]
+                    x1 = values[k + 2]
+                    y1 = values[k + 3]
+                    if not (np.isfinite(x0) and np.isfinite(y0) and
+                            np.isfinite(x1) and np.isfinite(y1)):
+                        continue
+
+                    segment_start = (
+                            (k == start1 and not closed_rings) or
+                            (k > start1 and
+                             (not np.isfinite(values[k - 2]) or not np.isfinite(values[k - 1])))
+                    )
+
+                    segment_end = (
+                            (not closed_rings and k == stop1-4) or
+                            (k < stop1-4 and
+                             (not np.isfinite(values[k + 4]) or not np.isfinite(values[k + 5])))
+                    )
+
+                    if segment_start or use_2_stage_agg:
+                        xm = 0.0
+                        ym = 0.0
+                    elif k == start1 and closed_rings:
+                        xm = values[stop1-4]
+                        ym = values[stop1-3]
+                    else:
+                        xm = values[k-2]
+                        ym = values[k-1]
+
+                    draw_segment(i, sx, tx, sy, ty, xmin, xmax, ymin, ymax,
+                                 segment_start, segment_end, x0, x1, y0, y1,
+                                 xm, ym, buffer, *aggs_and_cols)
+
+    def extend_cpu_antialias_2agg(
+        sx, tx, sy, ty, xmin, xmax, ymin, ymax, geometry, antialias_stage_2, *aggs_and_cols
+    ):
+        coords, offsets, outer_offsets, closed_rings = _process_geometry(geometry)
+        n_aggs = len(antialias_stage_2[0])
+        aggs_and_accums = tuple((agg, agg.copy()) for agg in aggs_and_cols[:n_aggs])
+
+        extend_cpu_numba_antialias_2agg(
+            sx, tx, sy, ty, xmin, xmax, ymin, ymax, coords, offsets, outer_offsets, closed_rings,
+            antialias_stage_2, aggs_and_accums, *aggs_and_cols)
+
+    @ngjit
+    @expand_aggs_and_cols
+    def extend_cpu_numba_antialias_2agg(
+        sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, offsets, outer_offsets,
+        closed_rings, antialias_stage_2, aggs_and_accums, *aggs_and_cols
+    ):
+        antialias = antialias_stage_2 is not None
+        buffer = np.empty(8) if antialias else None
+        n_multilines = len(outer_offsets) - 1
+        first_pass = True
+        for i in range(n_multilines):
+            start0 = outer_offsets[i]
+            stop0 = outer_offsets[i + 1]
+
+            for j in range(start0, stop0):
+                start1 = offsets[j]
+                stop1 = offsets[j + 1]
+
+                for k in range(2*start1, 2*stop1 - 2, 2):
+                    x0 = values[k]
+                    y0 = values[k + 1]
+                    x1 = values[k + 2]
+                    y1 = values[k + 3]
+                    if not (np.isfinite(x0) and np.isfinite(y0) and
+                            np.isfinite(x1) and np.isfinite(y1)):
+                        continue
+
+                    segment_start = (
+                            (k == start1 and not closed_rings) or
+                            (k > start1 and
+                             (not np.isfinite(values[k - 2]) or not np.isfinite(values[k - 1])))
+                    )
+
+                    segment_end = (
+                            (not closed_rings and k == stop1-4) or
+                            (k < stop1-4 and
+                             (not np.isfinite(values[k + 4]) or not np.isfinite(values[k + 5])))
+                    )
+
+                    draw_segment(i, sx, tx, sy, ty, xmin, xmax, ymin, ymax,
+                                 segment_start, segment_end, x0, x1, y0, y1,
+                                 #xm, ym, buffer, *aggs_and_cols)
+                                 0.0, 0.0, buffer, *aggs_and_cols)
+
+            aa_stage_2_accumulate(aggs_and_accums, first_pass)
+            first_pass = False
+            aa_stage_2_clear(aggs_and_accums)
+
+        aa_stage_2_copy_back(aggs_and_accums)
+
+    if use_2_stage_agg:
+        return extend_cpu_antialias_2agg
+    else:
+        return extend_cpu

datashader/glyphs/points.py

@@ -241,6 +241,86 @@ def extend(aggs, df, vt, bounds):
         return extend
 
 
+class MultiPointGeoPandas(_GeometryLike):
+    # geopandas must be available if a GeoPandasPointGeometry object is created.
+    @property
+    def geom_dtypes(self):
+        from geopandas.array import GeometryDtype
+        return (GeometryDtype,)
+
+    @memoize
+    def _build_extend(
+        self, x_mapper, y_mapper, info, append, _antialias_stage_2, _antialias_stage_2_funcs,
+    ):
+        # Lazy import shapely. Cannot get here if geopandas and shapely are not available.
+        import shapely
+
+        geometry_name = self.geometry
+
+        @ngjit
+        @self.expand_aggs_and_cols(append)
+        def _perform_extend_points(
+            i, j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, *aggs_and_cols
+        ):
+            x = values[j]
+            y = values[j + 1]
+            # points outside bounds are dropped; remainder
+            # are mapped onto pixels
+            if (xmin <= x <= xmax) and (ymin <= y <= ymax):
+                xx = int(x_mapper(x) * sx + tx)
+                yy = int(y_mapper(y) * sy + ty)
+                xi, yi = (xx - 1 if x == xmax else xx,
+                          yy - 1 if y == ymax else yy)
+
+                append(i, xi, yi, *aggs_and_cols)
+
+        def extend(aggs, df, vt, bounds):
+            aggs_and_cols = aggs + info(df, aggs[0].shape[:2])
+            sx, tx, sy, ty = vt
+            xmin, xmax, ymin, ymax = bounds
+            geometry = df[geometry_name].array
+
+            ragged = shapely.to_ragged_array(geometry)
+            geometry_type = ragged[0]
+
+            if geometry_type not in (shapely.GeometryType.MULTIPOINT, shapely.GeometryType.POINT):
+                raise ValueError(
+                    "Canvas.points supports GeoPandas geometry types of POINT and MULTIPOINT, "
+                    f"not {repr(geometry_type)}")
+
+            coords = ragged[1].ravel()  # No offsets required if POINT not MULTIPOINT
+            if geometry_type == shapely.GeometryType.POINT:
+                extend_point_cpu(sx, tx, sy, ty, xmin, xmax, ymin, ymax, coords, *aggs_and_cols)
+            else:
+                offsets = ragged[2][0]
+                extend_multipoint_cpu(
+                    sx, tx, sy, ty, xmin, xmax, ymin, ymax, coords, offsets, *aggs_and_cols)
+
+        @ngjit
+        @self.expand_aggs_and_cols(append)
+        def extend_multipoint_cpu(
+            sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, offsets, *aggs_and_cols,
+        ):
+            for i in range(len(offsets) - 1):
+                start = offsets[i]
+                stop = offsets[i+1]
+                for j in range(start, stop):
+                    _perform_extend_points(
+                        i, 2*j, sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, *aggs_and_cols,
+                    )
+
+        @ngjit
+        @self.expand_aggs_and_cols(append)
+        def extend_point_cpu(sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, *aggs_and_cols):
+            n = len(values) // 2
+            for i in range(n):
+                _perform_extend_points(
+                    i, 2*i, sx, tx, sy, ty, xmin, xmax, ymin, ymax, values, *aggs_and_cols,
+                )
+
+        return extend
+
+
 class MultiPointGeometry(_GeometryLike):
     # spatialpandas must be available if a MultiPointGeometry object is created.