Refactor stage 2 aggregation for antialiased lines (#1145)

* Antialias combinations for all line types

* Remove scalar _AntiAliasedLine.antialias_combination

* 2-stage agg for min and sum reductions

* Overwrite flag, correct other glyph APIs and tests

* Replace None with False for CUDA support

* Linting

* Update datashader/compiler.py

Co-authored-by: James A. Bednar <jbednar@users.noreply.github.com>

Co-authored-by: James A. Bednar <jbednar@users.noreply.github.com>

datashader/compiler.py

@@ -15,7 +15,9 @@
 
 @memoize
 def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
-    """Given a ``Aggregation`` object and a schema, return 5 sub-functions.
+    """Given an ``Aggregation`` object and a schema, return 5 sub-functions
+    and information on how to perform the second stage aggregation if
+    antialiasing is requested,
 
     Parameters
     ----------
@@ -46,6 +48,11 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
     ``finalize(aggs)``
         Given a tuple of base numpy arrays, returns the finalized ``DataArray``
         or ``Dataset``.
+
+    ``antialias_stage_2``
+        If using antialiased lines this is a tuple of the ``AntialiasCombination``
+        values corresponding to the aggs. If not using antialiased lines then
+        this is False.
     """
     reds = list(traverse_aggregation(agg))
 
@@ -65,7 +72,18 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False):
     combine = make_combine(bases, dshapes, temps, antialias)
     finalize = make_finalize(bases, agg, schema, cuda)
 
-    return create, info, append, combine, finalize
+    if antialias:
+        antialias_stage_2 = make_antialias_stage_2(reds, bases)
+        if cuda:
+            import cupy
+            array_module = cupy
+        else:
+            array_module = np
+        antialias_stage_2 = antialias_stage_2(array_module)
+    else:
+        antialias_stage_2 = False
+
+    return create, info, append, combine, finalize, antialias_stage_2
 
 
 def traverse_aggregation(agg):
@@ -195,3 +213,17 @@ def finalize(bases, cuda=False, **kwargs):
         return finalize
     else:
         return agg._build_finalize(schema)
+
+
+def make_antialias_stage_2(reds, bases):
+    # Only called if antialias is True.
+    self_intersect = True
+    for red in reds:
+        if red._antialias_requires_2_stages():
+            self_intersect = False
+            break
+
+    def antialias_stage_2(array_module):
+        return tuple(zip(*concat(b._antialias_stage_2(self_intersect, array_module) for b in bases)))
+
+    return antialias_stage_2

datashader/core.py

@@ -339,8 +339,7 @@ def line(self, source, x=None, y=None, agg=None, axis=0, geometry=None,
         """
         from .glyphs import (LineAxis0, LinesAxis1, LinesAxis1XConstant,
                              LinesAxis1YConstant, LineAxis0Multi,
-                             LinesAxis1Ragged, LineAxis1Geometry,
-                             AntialiasCombination)
+                             LinesAxis1Ragged, LineAxis1Geometry)
 
         validate_xy_or_geometry('Line', x, y, geometry)
 
@@ -433,28 +432,16 @@ def line(self, source, x=None, y=None, agg=None, axis=0, geometry=None,
         glyph.antialiased = (line_width > 0)
 
         if glyph.antialiased:
-            # Eventually this should be replaced with attributes and/or
-            # member functions of Reduction classes.
+            # This is required to identify and report use of reductions that do
+            # not yet support antialiasing.
             non_cat_agg = agg
             if isinstance(non_cat_agg, rd.by):
                 non_cat_agg = non_cat_agg.reduction
 
-            antialias_combination = AntialiasCombination.NONE
-            if isinstance(non_cat_agg, (rd.any, rd.max)):
-                antialias_combination = AntialiasCombination.MAX
-            elif isinstance(non_cat_agg, rd.min):
-                antialias_combination = AntialiasCombination.MIN
-            elif isinstance(non_cat_agg, (rd.count, rd.sum)):
-                if non_cat_agg.self_intersect:
-                    antialias_combination = AntialiasCombination.SUM_1AGG
-                else:
-                    antialias_combination = AntialiasCombination.SUM_2AGG
-            else:
+            if not isinstance(non_cat_agg, (rd.any, rd.count, rd.max, rd.min, rd.sum)):
                 raise NotImplementedError(
                     f"{type(non_cat_agg)} reduction not implemented for antialiased lines")
 
-            glyph.set_antialias_combination(antialias_combination)
-
         return bypixel(source, self, glyph, agg, antialias=glyph.antialiased)
 
     def area(self, source, x, y, agg=None, axis=0, y_stack=None):

datashader/data_libraries/dask.py

@@ -70,11 +70,11 @@ def default(glyph, df, schema, canvas, summary, *, antialias=False, cuda=False):
     shape, bounds, st, axis = shape_bounds_st_and_axis(df, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize = \
+    create, info, append, combine, finalize, antialias_stage_2 = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     # Here be dragons
     # Get the dataframe graph
@@ -185,11 +185,11 @@ def line(glyph, df, schema, canvas, summary, *, antialias=False, cuda=False):
     shape, bounds, st, axis = shape_bounds_st_and_axis(df, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize = \
+    create, info, append, combine, finalize, antialias_stage_2 = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     def chunk(df, df2=None):
         plot_start = True

datashader/data_libraries/dask_xarray.py

@@ -59,11 +59,11 @@ def dask_rectilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False,
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize = \
+    create, info, append, combine, finalize, antialias_stage_2 = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     # Build chunk indices for coordinates
     chunk_inds = {}
@@ -138,11 +138,11 @@ def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize = \
+    create, info, append, combine, finalize, antialias_stage_2 = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     # Build chunk indices for coordinates
     chunk_inds = {}
@@ -229,11 +229,11 @@ def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False,
     shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
 
     # Compile functions
-    create, info, append, combine, finalize = \
+    create, info, append, combine, finalize, antialias_stage_2 = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     x_coord_name = glyph.x
     y_coord_name = glyph.y

datashader/data_libraries/pandas.py

@@ -24,11 +24,11 @@ def pandas_pipeline(df, schema, canvas, glyph, summary, *, antialias=False):
 @glyph_dispatch.register(_GeometryLike)
 @glyph_dispatch.register(_AreaToLineLike)
 def default(glyph, source, schema, canvas, summary, *, antialias=False, cuda=False):
-    create, info, append, _, finalize = compile_components(
+    create, info, append, _, finalize, antialias_stage_2 = compile_components(
         summary, schema, glyph, antialias=antialias, cuda=cuda)
     x_mapper = canvas.x_axis.mapper
     y_mapper = canvas.y_axis.mapper
-    extend = glyph._build_extend(x_mapper, y_mapper, info, append)
+    extend = glyph._build_extend(x_mapper, y_mapper, info, append, antialias_stage_2)
 
     x_range = canvas.x_range or glyph.compute_x_bounds(source)
     y_range = canvas.y_range or glyph.compute_y_bounds(source)

datashader/enums.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+from enum import Enum
+
+
+# Enum used to specify how the second stage aggregation is performed
+# for 2-stage antialiased lines.
+class AntialiasCombination(Enum):
+    SUM_1AGG = 1
+    SUM_2AGG = 2
+    MIN = 3
+    MAX = 4

datashader/glyphs/area.py

@@ -98,7 +98,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -182,7 +182,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -274,7 +274,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -365,7 +365,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -472,7 +472,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -580,7 +580,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -649,7 +649,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -728,7 +728,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -797,7 +797,7 @@ def compute_bounds_dask(self, ddf):
                 self.compute_y_bounds())
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -865,7 +865,7 @@ def compute_bounds_dask(self, ddf):
                 self.compute_y_bounds())
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -950,7 +950,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(
@@ -1030,7 +1030,7 @@ def compute_bounds_dask(self, ddf):
                 self.maybe_expand_bounds(y_extents))
 
     @memoize
-    def _build_extend(self, x_mapper, y_mapper, info, append):
+    def _build_extend(self, x_mapper, y_mapper, info, append, antialias_stage_2):
         expand_aggs_and_cols = self.expand_aggs_and_cols(append)
         map_onto_pixel = _build_map_onto_pixel_for_line(x_mapper, y_mapper)
         draw_trapezoid_y = _build_draw_trapezoid_y(