Support antialiased lines in *_n reductions

datashader/antialias.py

@@ -32,11 +32,13 @@ def two_stage_agg(antialias_stage_2: UnzippedAntialiasStage2 | None):
         # Not using antialiased lines, doesn't matter what is returned.
         return False, False
 
+    aa_combinations = antialias_stage_2[0]
+
     # A single combination in (SUM_2AGG, FIRST, LAST, MIN) means that a 2-stage
     # aggregation will be used, otherwise use a 1-stage aggregation that is
     # faster.
     use_2_stage_agg = False
-    for comb in antialias_stage_2[0]:
+    for comb in aa_combinations:
         if comb in (AntialiasCombination.SUM_2AGG, AntialiasCombination.MIN,
                     AntialiasCombination.FIRST, AntialiasCombination.LAST):
             use_2_stage_agg = True
@@ -47,7 +49,7 @@ def two_stage_agg(antialias_stage_2: UnzippedAntialiasStage2 | None):
     # complicated correction algorithm. Prefer overwrite=True for speed, but
     # any SUM_1AGG implies overwrite=False.
     overwrite = True
-    for comb in antialias_stage_2[0]:
+    for comb in aa_combinations:
         if comb == AntialiasCombination.SUM_1AGG:
             overwrite = False
             break

datashader/compiler.py

@@ -9,8 +9,12 @@
 
 from .antialias import AntialiasCombination
 from .reductions import SpecialColumn, by, category_codes, summary
-from .utils import (isnull, ngjit, parallel_fill, nanmax_in_place, nanmin_in_place, nansum_in_place,
-    nanfirst_in_place, nanlast_in_place, row_max_in_place, row_min_in_place
+from .utils import (isnull, ngjit, parallel_fill,
+    nanmax_in_place, nanmin_in_place, nansum_in_place, nanfirst_in_place, nanlast_in_place,
+    nanmax_n_in_place_3d, nanmax_n_in_place_4d, nanmin_n_in_place_3d, nanmin_n_in_place_4d,
+    nanfirst_n_in_place_3d, nanfirst_n_in_place_4d, nanlast_n_in_place_3d, nanlast_n_in_place_4d,
+    row_min_in_place, row_min_n_in_place_3d, row_min_n_in_place_4d,
+    row_max_in_place, row_max_n_in_place_3d, row_max_n_in_place_4d,
 )
 
 try:
@@ -142,7 +146,26 @@ def compile_components(agg, schema, glyph, *, antialias=False, cuda=False, parti
 
 def _get_antialias_stage_2_combine_func(combination: AntialiasCombination, zero: float,
                                         n_reduction: bool, categorical: bool):
-    if not n_reduction:
+    if n_reduction:
+        if zero == -1:
+            if combination == AntialiasCombination.MAX:
+                return row_max_n_in_place_4d if categorical else row_max_n_in_place_3d
+            elif combination == AntialiasCombination.MIN:
+                return row_min_n_in_place_4d if categorical else row_min_n_in_place_3d
+            else:
+                raise NotImplementedError
+        else:
+            if combination == AntialiasCombination.MAX:
+                return nanmax_n_in_place_4d if categorical else nanmax_n_in_place_3d
+            elif combination == AntialiasCombination.MIN:
+                return nanmin_n_in_place_4d if categorical else nanmin_n_in_place_3d
+            elif combination == AntialiasCombination.FIRST:
+                return nanfirst_n_in_place_4d if categorical else nanfirst_n_in_place_3d
+            elif combination == AntialiasCombination.LAST:
+                return nanlast_n_in_place_4d if categorical else nanlast_n_in_place_3d
+            else:
+                raise NotImplementedError
+    else:
         # The aggs to combine here are either 3D (ny, nx, ncat) if categorical is True or
         # 2D (ny, nx) if categorical is False. The same combination functions can be for both
         # as all elements are independent.
@@ -151,6 +174,8 @@ def _get_antialias_stage_2_combine_func(combination: AntialiasCombination, zero:
                 return row_max_in_place
             elif combination == AntialiasCombination.MIN:
                 return row_min_in_place
+            else:
+                raise NotImplementedError
         else:
             if combination == AntialiasCombination.MAX:
                 return nanmax_in_place
@@ -163,8 +188,6 @@ def _get_antialias_stage_2_combine_func(combination: AntialiasCombination, zero:
             else:
                 return nansum_in_place
 
-    raise NotImplementedError
-
 
 def make_antialias_stage_2_functions(antialias_stage_2):
     aa_combinations, aa_zeroes, aa_n_reductions, aa_categorical = antialias_stage_2

datashader/core.py

@@ -440,7 +440,8 @@ def line(self, source, x=None, y=None, agg=None, axis=0, geometry=None,
 
             if not isinstance(non_cat_agg, (
                 rd.any, rd.count, rd.max, rd.min, rd.sum, rd.summary, rd._sum_zero,
-                rd.mean, rd._first_or_last, rd._max_or_min_row_index,
+                rd._first_or_last, rd.mean, rd.max_n, rd.min_n, rd._first_n_or_last_n,
+                rd._max_or_min_row_index, rd._max_n_or_min_n_row_index
             )):
                 raise NotImplementedError(
                     f"{type(non_cat_agg)} reduction not implemented for antialiased lines")

datashader/reductions.py

@@ -1463,10 +1463,30 @@ def _append(x, y, agg, field):
             # Could do binary search instead but not expecting n to be large.
             for i in range(n):
                 if isnull(agg[y, x, i]):
+                    # Nothing to shift.
                     agg[y, x, i] = field
                     return i
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        value = field*aa_factor
+        if not isnull(value):
+            # Check final value first for quick abort.
+            n = agg.shape[2]
+            if not isnull(agg[y, x, n-1]):
+                return -1
+
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            for i in range(n):
+                if isnull(agg[y, x, i]):
+                    # Nothing to shift.
+                    agg[y, x, i] = value
+                    return i
+        return -1
+
     def _create_row_index_selector(self):
         return _min_n_row_index(n=self.n)
 
@@ -1485,6 +1505,16 @@ def _append(x, y, agg, field):
             return 0
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        value = field*aa_factor
+        if not isnull(value):
+            # Always inserts at front of agg's third dimension.
+            shift_and_insert(agg[y, x], value, 0)
+            return 0
+        return -1
+
     def _create_row_index_selector(self):
         return _max_n_row_index(n=self.n)
 
@@ -1510,6 +1540,20 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        value = field*aa_factor
+        if not isnull(value):
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            for i in range(n):
+                if isnull(agg[y, x, i]) or value > agg[y, x, i]:
+                    shift_and_insert(agg[y, x], value, i)
+                    return i
+        return -1
+
     # GPU append functions
     @staticmethod
     @nb_cuda.jit(device=True)
@@ -1576,6 +1620,20 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        value = field*aa_factor
+        if not isnull(value):
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            for i in range(n):
+                if isnull(agg[y, x, i]) or value < agg[y, x, i]:
+                    shift_and_insert(agg[y, x], value, i)
+                    return i
+        return -1
+
     # GPU append functions
     @staticmethod
     @nb_cuda.jit(device=True)
@@ -2165,12 +2223,6 @@ def uses_cuda_mutex(self):
     def uses_row_index(self, cuda, partitioned):
         return True
 
-    def _build_append(self, dshape, schema, cuda, antialias, self_intersect):
-        # Doesn't yet support antialiasing
-        if cuda:
-            return self._append_cuda
-        else:
-            return self._append
 
     def _build_combine(self, dshape, antialias, cuda, partitioned):
         if cuda:
@@ -2186,6 +2238,9 @@ class _max_n_row_index(_max_n_or_min_n_row_index):
     user code. It is primarily purpose is to support the use of ``last_n``
     reductions using dask and/or CUDA.
     """
+    def _antialias_stage_2(self, self_intersect, array_module) -> tuple[AntialiasStage2]:
+        return (AntialiasStage2(AntialiasCombination.MAX, -1, n_reduction=True),)
+
     @staticmethod
     @ngjit
     def _append(x, y, agg, field):
@@ -2200,6 +2255,24 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        # field is int64 row index
+        # Ignoring aa_factor
+        if field != -1:
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            for i in range(n):
+                if agg[y, x, i] == -1 or field > agg[y, x, i]:
+                    # Bump previous values along to make room for new value.
+                    for j in range(n-1, i, -1):
+                        agg[y, x, j] = agg[y, x, j-1]
+                    agg[y, x, i] = field
+                    return i
+        return -1
+
     # GPU append functions
     @staticmethod
     @nb_cuda.jit(device=True)
@@ -2245,6 +2318,12 @@ class _min_n_row_index(_max_n_or_min_n_row_index):
     user code. It is primarily purpose is to support the use of ``first_n``
     reductions using dask and/or CUDA.
     """
+    def _antialias_requires_2_stages(self):
+        return True
+
+    def _antialias_stage_2(self, self_intersect, array_module) -> tuple[AntialiasStage2]:
+        return (AntialiasStage2(AntialiasCombination.MIN, -1, n_reduction=True),)
+
     @staticmethod
     @ngjit
     def _append(x, y, agg, field):
@@ -2259,6 +2338,21 @@ def _append(x, y, agg, field):
                     return i
         return -1
 
+    @staticmethod
+    @ngjit
+    def _append_antialias(x, y, agg, field, aa_factor):
+        # field is int64 row index
+        # Ignoring aa_factor
+        if field != -1:
+            # Linear walk along stored values.
+            # Could do binary search instead but not expecting n to be large.
+            n = agg.shape[2]
+            for i in range(n):
+                if agg[y, x, i] == -1 or field < agg[y, x, i]:
+                    shift_and_insert(agg[y, x], field, i)
+                    return i
+        return -1
+
     @staticmethod
     @nb_cuda.jit(device=True)
     def _append_cuda(x, y, agg, field):

datashader/tests/test_pandas.py

@@ -2494,6 +2494,45 @@ def test_line_antialias():
     agg = cvs.line(agg=ds._max_row_index(), **kwargs)
     assert_eq_ndarray(agg.data, line_antialias_sol_max_index_0)
 
+    agg = cvs.line(agg=ds._min_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, line_antialias_sol_min_index_0)
+    sol = np.full((11, 11), -1)
+    sol[(4, 5, 5, 5, 6, 1, 2), (5, 4, 5, 6, 5, 9, 9)] = 2
+    sol[8:10, 9] = 1
+    assert_eq_ndarray(agg[:, :, 1].data, sol)
+
+    agg = cvs.line(agg=ds._max_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, line_antialias_sol_max_index_0)
+    sol = np.full((11, 11), -1)
+    sol[(4, 5, 5, 5, 6, 8, 9), (5, 4, 5, 6, 5, 9, 9)] = 0
+    sol[1:3, 9] = 1
+    assert_eq_ndarray(agg[:, :, 1].data, sol)
+
+    agg = cvs.line(agg=ds.max_n("value", n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, 3*line_antialias_sol_0, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[(1, 5, 9), (9, 5, 9)] = 3.0
+    sol[(2, 4, 5, 5, 6, 8), (9, 5, 4, 6, 5, 9)] = 0.878680
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
+    agg = cvs.line(agg=ds.min_n("value", n=2), **kwargs)
+    sol = 3*line_antialias_sol_0
+    sol[(2, 8), (9, 9)] = 0.878680
+    assert_eq_ndarray(agg[:, :, 0].data, sol, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[(1, 2, 5, 8, 9), (9, 9, 5, 9, 9)] = 3.0
+    sol[(4, 5, 5, 6), (5, 4, 6, 5)] = 0.878680
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
+    agg = cvs.line(agg=ds.first_n("value", n=2), **kwargs)
+    sol = 3*line_antialias_sol_0
+    sol[8, 9] = 0.878680
+    assert_eq_ndarray(agg[:, :, 0].data, sol, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[(1, 5, 8, 9), (9, 5, 9, 9)] = 3.0
+    sol[(2, 4, 5, 5, 6), (9, 5, 4, 6, 5)] = 0.878680
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
     # Second line only, doesn't self-intersect
     kwargs = dict(source=line_antialias_df, x="x1", y="y1", line_width=1)
     agg = cvs.line(agg=ds.any(), **kwargs)
@@ -2533,6 +2572,51 @@ def test_line_antialias():
     agg = cvs.line(agg=ds._max_row_index(), **kwargs)
     assert_eq_ndarray(agg.data, line_antialias_sol_max_index_1)
 
+    agg = cvs.line(agg=ds._min_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, line_antialias_sol_min_index_1)
+    sol = np.full((11, 11), -1)
+    sol[(2, 2, 3), (3, 4, 4)] = 1
+    sol[2:4, 6:8] = 2
+    assert_eq_ndarray(agg[:, :, 1].data, sol)
+
+    agg = cvs.line(agg=ds._max_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, line_antialias_sol_max_index_1)
+    sol = np.full((11, 11), -1)
+    sol[(2, 2, 3), (3, 4, 4)] = 0
+    sol[2:4, 6:8] = 1
+    assert_eq_ndarray(agg[:, :, 1].data, sol)
+
+    agg = cvs.line(agg=ds.max_n("value", n=2), **kwargs)
+    assert_eq_ndarray(agg[:, :, 0].data, 3*line_antialias_sol_1, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[2, 3:8] = (0.911939, 1.833810, nan, 1.437950, 0.667619)
+    sol[(3, 3, 3), (4, 6, 7)] = (0.4, 0.940874, 0.309275)
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
+    agg = cvs.line(agg=ds.min_n("value", n=2), **kwargs)
+    sol = np.full((11, 11), np.nan)
+    sol[2, 1:-1] = [3.0, 2.775630, 0.911939, 1.833810, 2.1025216, 1.437950, 0.667619, 1.429411, 1.205041]
+    sol[3, 1:-1] = [0.008402, 0.232772, 0.457142, 0.4, 0.905881, 0.940874, 0.309275, 1.578991, 1.8]
+    assert_eq_ndarray(agg[:, :, 0].data, sol, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[2, 3:8] = (2.551260, 2.326890, nan, 1.878151, 1.653781)
+    sol[(3, 3, 3), (4, 6, 7)] = (0.681512, 1.130251, 1.354621)
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
+    agg = cvs.line(agg=ds.first_n("value", n=2), **kwargs)
+    sol = 3*line_antialias_sol_1
+    sol[(2, 2, 3, 3), (4, 7, 4, 7)] = (1.833810, 0.667619, 0.4, 0.309275)
+    assert_eq_ndarray(agg[:, :, 0].data, sol, close=True)
+    sol = np.full((11, 11), np.nan)
+    sol[2, 3:8] = (0.911939, 2.326890, nan, 1.437950, 1.653781)
+    sol[(3, 3, 3), (4, 6, 7)] = (0.681512, 0.940874, 1.354621)
+    assert_eq_ndarray(agg[:, :, 1].data, sol, close=True)
+
+    agg = cvs.line(agg=ds.last_n("value", n=2), **kwargs)
+    sol = 3*line_antialias_sol_1
+    sol[(2, 2, 3), (3, 6, 6)] = (0.911939, 1.437950, 0.940874)
+    assert_eq_ndarray(agg[:, :, 0].data, sol, close=True)
+
     # Both lines.
     kwargs = dict(source=line_antialias_df, x=["x0", "x1"], y=["y0", "y1"], line_width=1)
     agg = cvs.line(agg=ds.any(), **kwargs)
@@ -2580,6 +2664,12 @@ def test_line_antialias():
     sol_max_row = rowmax(line_antialias_sol_max_index_0, line_antialias_sol_max_index_1)
     assert_eq_ndarray(agg.data, sol_max_row)
 
+    agg = cvs.line(agg=ds._min_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg.data[:, :, 0], sol_min_row)
+
+    agg = cvs.line(agg=ds._max_n_row_index(n=2), **kwargs)
+    assert_eq_ndarray(agg.data[:, :, 0], sol_max_row)
+
     assert_eq_ndarray(agg.x_range, x_range, close=True)
     assert_eq_ndarray(agg.y_range, y_range, close=True)