Merge pull request #1130 from ioam/grid_interface_multi_groupby

Fix GridInterface multi-dimensional groupby

holoviews/core/data/xarray.py

@@ -114,7 +114,7 @@ def groupby(cls, dataset, dimensions, container_type, group_type, **kwargs):
                     dataset.data.groupby(index_dims[0].name)]
         else:
             unique_iters = [cls.values(dataset, d, False) for d in group_by]
-            indexes = zip(*[vals.flat for vals in util.cartesian_product(unique_iters)])
+            indexes = zip(*util.cartesian_product(unique_iters))
             data = [(k, group_type(dataset.data.sel(**dict(zip(group_by, k))),
                                    **group_kwargs))
                     for k in indexes]

holoviews/core/util.py

@@ -1057,13 +1057,17 @@ def groupby_python(self_or_cls, ndmapping, dimensions, container_type,
         return container_type(groups, kdims=dimensions)
 
 
-def cartesian_product(arrays):
+def cartesian_product(arrays, flat=True, copy=False):
     """
-    Computes the cartesian product of a list of 1D arrays
-    returning arrays matching the shape defined by all
-    supplied dimensions.
+    Efficient cartesian product of a list of 1D arrays returning the
+    expanded array views for each dimensions. By default arrays are
+    flattened, which may be controlled with the flat flag. The array
+    views can be turned into regular arrays with the copy flag.
     """
-    return np.broadcast_arrays(*np.ix_(*arrays))
+    arrays = np.broadcast_arrays(*np.ix_(*arrays))
+    if flat:
+        return tuple(arr.flatten() if copy else arr.flat for arr in arrays)
+    return tuple(arr.copy() if copy else arr for arr in arrays)
 
 
 def arglexsort(arrays):
@@ -1117,7 +1121,7 @@ def expand_grid_coords(dataset, dim):
     arrays = [dataset.interface.coords(dataset, d.name, True)
               for d in dataset.kdims]
     idx = dataset.get_dimension_index(dim)
-    return cartesian_product(arrays)[idx]
+    return cartesian_product(arrays, flat=False)[idx]
 
 
 def dt64_to_dt(dt64):

holoviews/element/util.py

@@ -125,8 +125,8 @@ def _aggregate_dataset(self, obj, xcoords, ycoords):
         shape = (len(ycoords), len(xcoords))
         nsamples = np.product(shape)
 
-        ys, xs = cartesian_product([ycoords, xcoords])
-        data = {xdim: xs.flatten(), ydim: ys.flatten()}
+        ys, xs = cartesian_product([ycoords, xcoords], copy=True)
+        data = {xdim: xs, ydim: ys}
         for vdim in vdims:
             values = np.empty(nsamples)
             values[:] = np.NaN

holoviews/plotting/bokeh/raster.py

@@ -189,10 +189,9 @@ def get_data(self, element, ranges=None, empty=False):
             yvals = element.dimension_values(1, False)
             widths = np.diff(element.data[0])
             heights = np.diff(element.data[1])
-            xs, ys = cartesian_product([xvals, yvals])
-            ws, hs = cartesian_product([widths, heights])
-            data = {x: xs.flatten(), y: ys.flatten(), z: zvals,
-                    'widths': ws.flatten(), 'heights': hs.flatten()}
+            xs, ys = cartesian_product([xvals, yvals], copy=True)
+            ws, hs = cartesian_product([widths, heights], copy=True)
+            data = {x: xs, y: ys, z: zvals, 'widths': ws, 'heights': hs}
 
         return (data, {'x': x, 'y': y,
                        'fill_color': {'field': z, 'transform': cmapper},

tests/testdataset.py

@@ -3,6 +3,8 @@
 """
 
 from unittest import SkipTest
+from itertools import product
+
 import numpy as np
 from holoviews import Dataset, NdElement, HoloMap, Dimension
 from holoviews.element.comparison import ComparisonTestCase
@@ -833,6 +835,15 @@ def test_dataset_groupby_dynamic_alias(self):
                         kdims=[('y', 'Y')], vdims=[('z', 'Z')])
         self.assertEqual(grouped[0], first)
 
+    def test_dataset_groupby_multiple_dims(self):
+        dataset = Dataset((range(8), range(8), range(8), range(8),
+                           np.random.rand(8, 8, 8, 8)),
+                          kdims=['a', 'b', 'c', 'd'], vdims=['Value'])
+        grouped = dataset.groupby(['c', 'd'])
+        keys = list(product(range(8), range(8)))
+        self.assertEqual(list(grouped.keys()), keys)
+        for c, d in keys:
+            self.assertEqual(grouped[c, d], dataset.select(c=c, d=d).reindex(['a', 'b']))
 
 
 class IrisDatasetTest(GridDatasetTest):