feat: Make quadmesh support bandwise 2D (#1472)

Author: hoxbro

datashader/core.py

@@ -780,7 +780,23 @@ def polygons(self, source, geometry, agg=None):
         return bypixel(source, self, glyph, agg)
 
     def quadmesh(self, source, x=None, y=None, agg=None):
-        """Samples a recti- or curvi-linear quadmesh by canvas size and bounds.
+        r"""Samples a raster, rectilinear or curvilinear quadmesh by canvas size and bounds.
+
+        +---------------------+---------------------+---------------------+
+        |   RASTER            |   RECTILINEAR       |   CURVILINEAR       |
+        +---------------------+---------------------+---------------------+
+        | Regular spacing     | Variable spacing    | Variable 2D spacing |
+        | o---o---o---o---o   | o-o---o----o-o      |    o---o---o---o    |
+        | |   |   |   |   |   | | |   |    | |      |   /   /   /   /     |
+        | o---o---o---o---o   | o-o---o----o-o      |  o----o--o---o      |
+        | |   |   |   |   |   | | |   |    | |      | /   /    /  /       |
+        | o---o---o---o---o   | o-o---o----o-o      | o--o----o---o       |
+        | |   |   |   |   |   | | |   |    | |      |  \   \    \  \      |
+        | o---o---o---o---o   | o-o---o----o-o      |   o---o---o---o     |
+        | `dx = dy = constant`  | `dx` & `dy` vary in 1D  | `dx` & `dy` vary in 2D  |
+        | `x[i] = i * dx`       | `x[i], y[j]`          | `x[i,j], y[i,j]`      |
+        | `y[j] = j * dy`       |                     |                     |
+        +---------------------+---------------------+---------------------+
 
         Parameters
         ----------
@@ -795,10 +811,12 @@ def quadmesh(self, source, x=None, y=None, agg=None):
         Returns
         -------
         data : xarray.DataArray
+
+        Note
+        ----
+        Table from EarthMover
         """
         from .glyphs import QuadMeshRaster, QuadMeshRectilinear, QuadMeshCurvilinear
-
-        # Determine reduction operation
         from .reductions import mean as mean_rnd
 
         if isinstance(source, Dataset):
@@ -819,7 +837,10 @@ def quadmesh(self, source, x=None, y=None, agg=None):
             agg = mean_rnd(name)
 
         if x is None and y is None:
-            y, x = source[name].dims
+            if len(source[name].dims) == 2:
+                y, x = source[name].dims
+            else:
+                raise ValueError("x and y must be specified if dims is not 2D.")
         elif not x or not y:
             raise ValueError("Either specify both x and y coordinates"
                              "or allow them to be inferred.")

datashader/data_libraries/dask_xarray.py

@@ -4,15 +4,29 @@
 from datashader.glyphs.quadmesh import (
     QuadMeshRaster, QuadMeshRectilinear, QuadMeshCurvilinear, build_scale_translate
 )
+from .xarray import _extract_third_dim
 from datashader.utils import apply
 import dask
 import numpy as np
 import xarray as xr
-from dask.base import tokenize, compute
+from dask.base import tokenize, compute, flatten
 from dask.array.overlap import overlap
 dask_glyph_dispatch = Dispatcher()
 
 
+def _prepare_3d_coords_and_dims(third_dim, xr_ds, axis, glyph):
+    dims_list = [glyph.y_label, glyph.x_label]
+
+    if third_dim:
+        coords_dict = axis.copy()
+        coords_dict[third_dim] = xr_ds.coords[third_dim]
+        dims_list = [third_dim, glyph.y_label, glyph.x_label]
+    else:
+        coords_dict = axis
+
+    return coords_dict, dims_list
+
+
 def dask_xarray_pipeline(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
     dsk, name = dask_glyph_dispatch(
         glyph, xr_ds, schema, canvas, summary, antialias=antialias, cuda=cuda)
@@ -54,9 +68,19 @@ def shape_bounds_st_and_axis(xr_ds, canvas, glyph):
 
     return shape, bounds, st, axis
 
+def _data_info_3d(xr_ds, canvas, glyph):
+    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
+
+    third_dim = _extract_third_dim(glyph, xr_ds)
+    if third_dim:
+        height, width = shape
+        shape = (len(xr_ds.coords[third_dim]), height, width)
+
+    return shape, bounds, st, axis, third_dim
+
 
 def dask_rectilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
-    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
+    shape, bounds, st, axis, third_dim = _data_info_3d(xr_ds, canvas, glyph)
 
     # Compile functions
     create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
@@ -129,17 +153,19 @@ def chunk(np_arr, *inds):
         return aggs
 
     name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
-    keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
+    keys = tuple(flatten(xr_ds.__dask_keys__()[0]))
     keys2 = [(name, i) for i in range(len(keys))]
-    dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
+    dsk = dict((k2, (chunk, k, *k[1:])) for (k2, k) in zip(keys2, keys))
+
+    coords_dict, dims_list = _prepare_3d_coords_and_dims(third_dim, xr_ds, axis, glyph)
     dsk[name] = (apply, finalize, [(combine, keys2)],
-                 dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
+                 dict(cuda=cuda, coords=coords_dict, dims=dims_list,
                       attrs=dict(x_range=x_range, y_range=y_range)))
     return dsk, name
 
 
 def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
-    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
+    shape, bounds, st, axis, third_dim = _data_info_3d(xr_ds, canvas, glyph)
 
     # Compile functions
     create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
@@ -177,7 +203,7 @@ def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=
     ybinsize = abs(float(xr_ds[y_name][1] - xr_ds[y_name][0]))
 
     # Compute scale/translate
-    out_h, out_w = shape
+    out_h, out_w = shape[-2:]
     src_h, src_w = [xr_ds[glyph.name].shape[i] for i in [ydim_ind, xdim_ind]]
     out_x0, out_x1, out_y0, out_y1 = bounds
     scale_y, translate_y = build_scale_translate(
@@ -224,19 +250,20 @@ def chunk(np_arr, *inds):
         return aggs
 
     name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
-    keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
+    keys = tuple(flatten(xr_ds.__dask_keys__()[0]))
     keys2 = [(name, i) for i in range(len(keys))]
-    dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
+    dsk = dict((k2, (chunk, k, *k[1:])) for (k2, k) in zip(keys2, keys))
+
+    coords_dict, dims_list = _prepare_3d_coords_and_dims(third_dim, xr_ds, axis, glyph)
     dsk[name] = (apply, finalize, [(combine, keys2)],
-                 dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
+                 dict(cuda=cuda, coords=coords_dict, dims=dims_list,
                       attrs=dict(x_range=x_range, y_range=y_range)))
     return dsk, name
 
 
 def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
-    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
+    shape, bounds, st, axis, third_dim = _data_info_3d(xr_ds, canvas, glyph)
 
-    # Compile functions
     create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
         compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
     x_mapper = canvas.x_axis.mapper
@@ -259,20 +286,29 @@ def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False,
 
     var_name = list(xr_ds.data_vars.keys())[0]
 
-    # Validate coordinates
+    # Validate coordinates - for 3D, compare only spatial dimensions (exclude third_dim)
+    if third_dim:
+        expected_dims = tuple(d for d in xr_ds[glyph.name].dims if d != third_dim)
+        expected_chunks = tuple(
+            c for d, c in zip(xr_ds[glyph.name].dims, xr_ds[glyph.name].chunks) if d != third_dim
+        )
+    else:
+        expected_dims = xr_ds[glyph.name].dims
+        expected_chunks = xr_ds[glyph.name].chunks
+
     err_msg = (
         "DataArray {name} is backed by a Dask array, \n"
         "but coordinate {coord} is not backed by a Dask array with identical \n"
         "dimension order and chunks"
     )
     if (not isinstance(x_centers, dask.array.Array) or
-            xr_ds[glyph.name].dims != xr_ds[glyph.x].dims or
-            xr_ds[glyph.name].chunks != xr_ds[glyph.x].chunks):
+            xr_ds[glyph.x].dims != expected_dims or
+            xr_ds[glyph.x].chunks != expected_chunks):
         raise ValueError(err_msg.format(name=glyph.name, coord=glyph.x))
 
     if (not isinstance(y_centers, dask.array.Array) or
-            xr_ds[glyph.name].dims != xr_ds[glyph.y].dims or
-            xr_ds[glyph.name].chunks != xr_ds[glyph.y].chunks):
+            xr_ds[glyph.y].dims != expected_dims or
+            xr_ds[glyph.y].chunks != expected_chunks):
         raise ValueError(err_msg.format(name=glyph.name, coord=glyph.y))
 
     # Make sure coordinates are floats so that overlap with nan will behave properly
@@ -324,9 +360,9 @@ def chunk(np_zs, np_x_centers, np_y_centers):
 
     result_name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
 
-    z_keys = [k for row in zs.__dask_keys__() for k in row]
-    x_overlap_keys = [k for row in x_overlapped_centers.__dask_keys__() for k in row]
-    y_overlap_keys = [k for row in y_overlapped_centers.__dask_keys__() for k in row]
+    z_keys = tuple(flatten(zs.__dask_keys__()))
+    x_overlap_keys = tuple(flatten(x_overlapped_centers.__dask_keys__()))
+    y_overlap_keys = tuple(flatten(y_overlapped_centers.__dask_keys__()))
 
     result_keys = [(result_name, i) for i in range(len(z_keys))]
 
@@ -337,9 +373,10 @@ def chunk(np_zs, np_x_centers, np_y_centers):
         )
     )
 
+    coords_dict, dims_list = _prepare_3d_coords_and_dims(third_dim, xr_ds, axis, glyph)
     dsk[result_name] = (
         apply, finalize, [(combine, result_keys)],
-        dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
+        dict(cuda=cuda, coords=coords_dict, dims=dims_list,
              attrs=dict(x_range=x_range, y_range=y_range))
     )
 

datashader/data_libraries/xarray.py

@@ -5,6 +5,7 @@
 from datashader.core import bypixel
 import xarray as xr
 from datashader.utils import Dispatcher
+from datashader.compiler import compile_components
 
 
 try:
@@ -33,6 +34,61 @@ def xarray_pipeline(xr_ds, schema, canvas, glyph, summary, *, antialias=False):
             glyph, xr_ds, schema, canvas, summary, antialias=antialias, cuda=cuda)
 
 
+def _extract_third_dim(glyph, source):
+    x_dims = set(source.coords[glyph.x].dims) if glyph.x in source.coords else {glyph.x}
+    y_dims = set(source.coords[glyph.y].dims) if glyph.y in source.coords else {glyph.y}
+    dims = set(source.dims) - (x_dims | y_dims)
+    match len(dims):
+        case 0:
+            return None
+        case 1:
+            return next(iter(dims))
+        case _:
+            raise ValueError("Only one additional dimension supported for QuadMesh glyphs.")
+
+
+@glyph_dispatch.register(_QuadMeshLike)
+def quadmesh_default(glyph, source, schema, canvas, summary, *, antialias=False, cuda=False):
+    third_dim = _extract_third_dim(glyph, source)
+    if not third_dim:
+        return default(glyph, source, schema, canvas, summary, antialias=antialias, cuda=cuda)
+
+    create, info, append, _, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
+        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda,
+                           partitioned=False)
+    x_mapper = canvas.x_axis.mapper
+    y_mapper = canvas.y_axis.mapper
+    extend = glyph._build_extend(
+        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
+
+    x_range = canvas.x_range or glyph.compute_x_bounds(source)
+    y_range = canvas.y_range or glyph.compute_y_bounds(source)
+    canvas.validate_ranges(x_range, y_range)
+
+    width = canvas.plot_width
+    height = canvas.plot_height
+
+    x_st = canvas.x_axis.compute_scale_and_translate(x_range, width)
+    y_st = canvas.y_axis.compute_scale_and_translate(y_range, height)
+
+    x_axis = canvas.x_axis.compute_index(x_st, width)
+    y_axis = canvas.y_axis.compute_index(y_st, height)
+
+    bases = create((len(source.coords[third_dim]), height, width))
+
+    extend(bases, source, x_st + y_st, x_range + y_range)
+
+    return finalize(
+        bases,
+        cuda=cuda,
+        coords=dict([
+            (third_dim, source.coords[third_dim]),
+            (glyph.x_label, x_axis),
+            (glyph.y_label, y_axis),
+        ]),
+        dims=[third_dim, glyph.y_label, glyph.x_label],
+        attrs=dict(x_range=x_range, y_range=y_range)
+    )
+
 # Default to default pandas implementation
-glyph_dispatch.register(_QuadMeshLike)(default)
 glyph_dispatch.register(LinesXarrayCommonX)(default)