From acaa4124430d67010daf769a772ca4da3aef3595 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <prudiger@anaconda.com>
Date: Mon, 24 Jul 2023 19:24:38 +0200
Subject: [PATCH] Implement collapse operation for NdMapping of Overlays
 (#5825)

---
 holoviews/core/ndmapping.py            | 16 +++++++---
 holoviews/tests/core/test_ndmapping.py | 41 +++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/holoviews/core/ndmapping.py b/holoviews/core/ndmapping.py
index 94d11fa8f6..9358f371a9 100644
--- a/holoviews/core/ndmapping.py
+++ b/holoviews/core/ndmapping.py
@@ -835,6 +835,8 @@ def collapse(self, dimensions=None, function=None, spreadfn=None, **kwargs):
             elements
         """
         from .data import concat
+        from .overlay import CompositeOverlay
+
         if not dimensions:
             dimensions = self.kdims
         if not isinstance(dimensions, list): dimensions = [dimensions]
@@ -860,11 +862,17 @@ def collapse(self, dimensions=None, function=None, spreadfn=None, **kwargs):
                 if function:
                     agg = group_data.aggregate(group.last.kdims, function, spreadfn, **kwargs)
                     group_data = group.type(agg)
+            elif issubclass(group.type, CompositeOverlay) and hasattr(self, '_split_overlays'):
+                keys, maps = self._split_overlays()
+                group_data = group.type(OrderedDict([
+                    (key, ndmap.collapse(function=function, spreadfn=spreadfn, **kwargs))
+                    for key, ndmap in zip(keys, maps)
+                ]))
             else:
-                group_data = [el.data for el in group]
-                args = (group_data, function, group.last.kdims)
-                data = group.type.collapse_data(*args, **kwargs)
-                group_data = group.last.clone(data)
+                raise ValueError(
+                    "Could not determine correct collapse operation "
+                    "for items of type: {group.type!r}."
+                )
             collapsed[key] = group_data
         return collapsed if self.ndims-len(dimensions) else collapsed.last
 
diff --git a/holoviews/tests/core/test_ndmapping.py b/holoviews/tests/core/test_ndmapping.py
index 04a1be890a..2a51d3e449 100644
--- a/holoviews/tests/core/test_ndmapping.py
+++ b/holoviews/tests/core/test_ndmapping.py
@@ -1,12 +1,15 @@
 from collections import OrderedDict
 
+import numpy as np
+
 from holoviews.core import Dimension
 from holoviews.core.ndmapping import (
     MultiDimensionalMapping, NdMapping, UniformNdMapping
 )
+from holoviews.core.overlay import Overlay
+from holoviews.element import Curve
 from holoviews.element.comparison import ComparisonTestCase
 from holoviews import HoloMap, Dataset
-import numpy as np
 
 class DimensionTest(ComparisonTestCase):
 
@@ -269,3 +272,39 @@ def test_holomap_hist_two_dims(self):
         hists = hmap.hist(dimension=['x', 'y'])
         self.assertEqual(hists['right'].last.kdims, ['y'])
         self.assertEqual(hists['top'].last.kdims, ['x'])
+
+    def test_holomap_collapse_overlay_no_function(self):
+        hmap = HoloMap({
+            (1,0): Curve(np.arange(8)) * Curve(-np.arange(8)),
+            (2,0): Curve(np.arange(8)**2) * Curve(-np.arange(8)**3)
+        }, kdims=["A","B"])
+        self.assertEqual(hmap.collapse(), Overlay([
+            (('Curve', 'I'), Dataset({
+                'A': np.concatenate([np.ones(8), np.ones(8)*2]),
+                'B': np.zeros(16),
+                'x': np.tile(np.arange(8), 2),
+                'y': np.concatenate([np.arange(8), np.arange(8)**2])
+            }, kdims=['A', 'B', 'x'], vdims=['y'])),
+            (('Curve', 'II'), Dataset({
+                'A': np.concatenate([np.ones(8), np.ones(8)*2]),
+                'B': np.zeros(16),
+                'x': np.tile(np.arange(8), 2),
+                'y': np.concatenate([-np.arange(8), -np.arange(8)**3])
+            }, kdims=['A', 'B', 'x'], vdims=['y']))
+        ]))
+
+    def test_holomap_collapse_overlay_max(self):
+        hmap = HoloMap({
+            (1,0): Curve(np.arange(8)) * Curve(-np.arange(8)),
+            (2,0): Curve(np.arange(8)**2) * Curve(-np.arange(8)**3)
+        }, kdims=["A","B"])
+        self.assertEqual(hmap.collapse(function=np.max), Overlay([
+            (('Curve', 'I'), Curve({
+                'x': np.arange(8),
+                'y': np.arange(8)**2
+            }, kdims=['x'], vdims=['y'])),
+            (('Curve', 'II'), Curve({
+                'x': np.arange(8),
+                'y': -np.arange(8)
+            }, kdims=['x'], vdims=['y']))
+        ]))