Antialiased lines for categorical aggregates (#1081)

datashader/core.py

@@ -443,10 +443,7 @@ def line(self, source, x=None, y=None, agg=None, axis=0, geometry=None,
             glyph.set_antialias_combination(antialias_combination)
 
             # Switch agg to floating point.
-            if isinstance(agg, rd.count):
-                agg = rd.count_f32(self_intersect=agg.self_intersect)
-            elif isinstance(agg, rd.any):
-                agg = rd.any_f32()
+            agg = rd._reduction_to_floating_point(agg)
 
         return bypixel(source, self, glyph, agg)
 

datashader/glyphs/line.py

@@ -8,6 +8,7 @@
 from datashader.utils import (isnull, isreal, ngjit, nanmax_in_place,
                               nanmin_in_place, nansum_in_place, parallel_fill)
 from numba import cuda
+from numba.extending import overload
 
 
 try:
@@ -692,6 +693,44 @@ def _linearstep(edge0, edge1, x):
     return t
 
 
+def _agg2d_with_scale(aggs_and_cols, i):
+    # Python implementation for use when Numba is disabled.
+    agg2or3d = aggs_and_cols[0]
+    if agg2or3d.ndim == 2:
+        agg = aggs_and_cols[0]  # 2D array
+        # Scale by column value if present.
+        scale = 1.0 if len(aggs_and_cols) == 1 else aggs_and_cols[1][i]
+        return agg, scale
+    elif agg2or3d.ndim == 3:
+        cat_index = aggs_and_cols[1][i]
+        agg = aggs_and_cols[0][:, :, cat_index]  # 2D array
+        return agg, 1.0
+    else:
+        raise TypeError("Not supported")
+
+
+@overload(_agg2d_with_scale)
+def _overload_agg2d_with_scale(aggs_and_cols, i):  # pragma: no cover
+    # Return different implementation based on whether the first array in
+    # aggs_and_cols is 2D or 3D.
+    agg2or3d = aggs_and_cols[0]
+    if agg2or3d.ndim == 2:
+        def impl(aggs_and_cols, i):
+            agg = aggs_and_cols[0]  # 2D array
+            # Scale by column value if present.
+            scale = 1.0 if len(aggs_and_cols) == 1 else aggs_and_cols[1][i]
+            return agg, scale
+        return impl
+    elif agg2or3d.ndim == 3:
+        def impl(aggs_and_cols, i):
+            cat_index = aggs_and_cols[1][i]
+            agg = aggs_and_cols[0][:, :, cat_index]  # 2D array
+            return agg, 1.0
+        return impl
+    else:
+        raise TypeError("Not supported")
+
+
 @ngjit
 def _full_antialias(line_width, antialias_combination, i, x0, x1, y0, y1,
                     segment_start, segment_end, xm, ym, *aggs_and_cols):
@@ -708,18 +747,13 @@ def _full_antialias(line_width, antialias_combination, i, x0, x1, y0, y1,
         x1, y1 = y1, x1
         xm, ym = ym, xm
 
-    agg = aggs_and_cols[0]
+    agg, scale = _agg2d_with_scale(aggs_and_cols, i)
 
     # line_width less than 1 is rendered as 1 but with lower intensity.
-    scale = 1.0
     if line_width < 1.0:
-        scale = line_width
+        scale *= line_width
         line_width = 1.0
 
-    # Scale by column value, if required.
-    if len(aggs_and_cols) > 1:
-        scale *= aggs_and_cols[1][i]
-
     aa = 1.0
     halfwidth = 0.5*(line_width + aa)
 

datashader/reductions.py

@@ -998,6 +998,17 @@ def inputs(self):
         return tuple(unique(concat(v.inputs for v in self.values)))
 
 
+def _reduction_to_floating_point(reduction):
+    # Reductions need to be floating-point when using antialiasing.
+    if isinstance(reduction, count):
+        reduction = count_f32(self_intersect=reduction.self_intersect)
+    elif isinstance(reduction, any):
+        reduction = any_f32()
+    elif isinstance(reduction, by):
+        reduction.reduction = _reduction_to_floating_point(reduction.reduction)
+
+    return reduction
+
 
 __all__ = list(set([_k for _k,_v in locals().items()
                     if isinstance(_v,type) and (issubclass(_v,Reduction) or _v is summary)

datashader/tests/test_pandas.py

@@ -1755,7 +1755,6 @@ def test_area_to_line_autorange_gap():
     assert_eq_xr(agg, out)
 
 
-
 # Using local versions of nan-aware combinations rather than those in
 # utils.py.  These versions are not always applicable, e.g. if summming
 # a positive and negative value to total exactly zero will be wrong here.
@@ -1902,3 +1901,21 @@ def test_line_antialias():
     agg = cvs.line(source=line_antialias_df, x=["x0", "x1"], y=["y0", "y1"], agg=ds.min("value"), line_width=1)
     sol_min = nanmin(line_antialias_sol_0, line_antialias_sol_1)
     assert_eq_ndarray(agg.data, 3*sol_min, close=True)
+
+
+def test_line_antialias_categorical():
+    df = pd.DataFrame(dict(
+        x=np.asarray([0, 1, 1, 0, np.nan, 0, 1/3.0, 2/3.0, 1]),
+        y=np.asarray([0, 1, 0, 1, np.nan, 0.125, 0.15, 0.175, 0.2]),
+        cat=[1, 1, 1, 1, 1, 2, 2, 2, 2],
+    ))
+    df["cat"] = df["cat"].astype("category")
+
+    x_range = y_range = (-0.1875, 1.1875)
+    cvs = ds.Canvas(plot_width=11, plot_height=11, x_range=x_range, y_range=y_range)
+
+    for self_intersect in [False, True]:
+        agg = cvs.line(source=df, x="x", y="y", line_width=1,
+                       agg=ds.by("cat", ds.count(self_intersect=self_intersect)))
+        assert_eq_ndarray(agg.data[:, :, 0], line_antialias_sol_0, close=True)
+        assert_eq_ndarray(agg.data[:, :, 1], line_antialias_sol_1, close=True)