Repository: graphnet-team/graphnet
Issue: DOMCoarsening throws an error

Hi guys, when I try to include the DOMCoarsening in the GNN model and start training I get this error:

```
  File "run_jobs.py", line 333, in <module>
    train_and_predict_on_validation_set(target,selection, database, pulsemap, batch_size, num_workers, n_epochs, device, run_name,archive, k = k, train = True)
  File "run_jobs.py", line 170, in train_and_predict_on_validation_set
    trainer.fit(model, training_dataloader, validation_dataloader)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in fit
    self._call_and_handle_interrupt(
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 723, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 811, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1236, in _run
    results = self._run_stage()
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1323, in _run_stage
    return self._run_train()
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1345, in _run_train
    self._run_sanity_check()
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1413, in _run_sanity_check
    val_loop.run()
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 154, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 128, in advance
    output = self._evaluation_step(**kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 226, in _evaluation_step
    output = self.trainer._call_strategy_hook("validation_step", *kwargs.values())
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 344, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/model.py", line 114, in validation_step
    loss = self.shared_step(val_batch, batch_idx)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/model.py", line 97, in shared_step
    preds = self(batch)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/model.py", line 91, in forward
    data = self._detector(data)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/detector/detector.py", line 82, in forward
    data = self._forward(data)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/detector/icecube.py", line 30, in _forward
    self._validate_features(data)
  File "/home/iwsatlas1/haminh/IceCube/graphnet/src/graphnet/models/detector/detector.py", line 101, in _validate_features
    data_features = data[0].features
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/torch_geometric/data/batch.py", line 146, in __getitem__
    return self.get_example(idx)
  File "/home/iwsatlas1/haminh/.conda/envs/graphnet_stable/lib/python3.8/site-packages/torch_geometric/data/batch.py", line 91, in get_example
    raise RuntimeError(
RuntimeError: Cannot reconstruct 'Data' object from 'Batch' because 'Batch' was not created via 'Batch.from_data_list()'
```

But everything runs through if I remove the coarsening argument.

Changed files:
src/graphnet/models/coarsening.py, src/graphnet/models/gnn/dynedge.py, tests/ml/test_coarsening.py

diff --git a/src/graphnet/models/coarsening.py b/src/graphnet/models/coarsening.py
index e9669c4c..64fed99a 100644
--- a/src/graphnet/models/coarsening.py
+++ b/src/graphnet/models/coarsening.py
@@ -1,11 +1,14 @@
 """Classes for coarsening operations (i.e., clustering, or local pooling."""
 
 from abc import ABC, abstractmethod
+from multiprocessing import pool
+from typing import List, Union
 
 import torch
 from torch import LongTensor, Tensor
-from torch_geometric.data import Data
+from torch_geometric.data import Data, Batch
 
+# from torch_geometric.utils import unbatch_edge_index
 from graphnet.components.pool import (
     group_by,
     avg_pool,
@@ -18,9 +21,35 @@ from graphnet.components.pool import (
     sum_pool_x,
     std_pool_x,
 )
+from graphnet.utilities.logging import LoggerMixin
 
+# Utility method(s)
+from torch_geometric.utils import degree
 
-class Coarsening(ABC):
+# NOTE: From [https://github.com/pyg-team/pytorch_geometric/pull/4903]
+# TODO: Remove once bumping to torch_geometric>=2.1.0
+#       See [https://github.com/pyg-team/pytorch_geometric/blob/master/CHANGELOG.md]
+
+
+def unbatch_edge_index(edge_index: Tensor, batch: Tensor) -> List[Tensor]:
+    r"""Splits the :obj:`edge_index` according to a :obj:`batch` vector.
+    Args:
+        edge_index (Tensor): The edge_index tensor. Must be ordered.
+        batch (LongTensor): The batch vector
+            :math:`\\mathbf{b} \\in {\\{ 0, \\ldots, B-1\\}}^N`, which assigns each
+            node to a specific example. Must be ordered.
+    :rtype: :class:`List[Tensor]`
+    """
+    deg = degree(batch, dtype=torch.int64)
+    ptr = torch.cat([deg.new_zeros(1), deg.cumsum(dim=0)[:-1]], dim=0)
+
+    edge_batch = batch[edge_index[0]]
+    edge_index = edge_index - ptr[edge_batch]
+    sizes = degree(edge_batch, dtype=torch.int64).cpu().tolist()
+    return edge_index.split(sizes, dim=1)
+
+
+class Coarsening(ABC, LoggerMixin):
     """Base class for coarsening operations."""
 
     # Class variables
@@ -54,36 +83,73 @@ class Coarsening(ABC):
         By default the nominal `pooling_method` is used for features as well.
         This method can be overwritten for bespoke coarsening operations.
         """
-        return None
 
     def _transfer_attributes(
         self, cluster: LongTensor, original_data: Data, pooled_data: Data
     ) -> Data:
         """Transfer attributes on `original_data` to `pooled_data`."""
+        # Check(s)
         if not self._do_transfer_attributes:
             return pooled_data
 
+        if (
+            isinstance(original_data, Batch)
+            or isinstance(pooled_data, Batch)
+            or (original_data.batch is not None)
+        ):
+            raise ValueError(
+                "Please coarsen events individually, not in batches."
+            )
+
         attributes = list(original_data._store.keys())
         for attr in attributes:
             if attr not in pooled_data._store:
                 values = getattr(original_data, attr)
 
-                if (
-                    isinstance(values, Tensor)
-                    and values.size() == original_data.batch.size()
-                ):  # Node-level tensor
+                attr_is_node_level_tensor = False
+                if isinstance(values, Tensor):
+                    attr_is_node_level_tensor = (
+                        values.dim() > 1 or values.size(dim=0) > 1
+                    )
+
+                if attr_is_node_level_tensor:
+                    batch = torch.zeros_like(values, dtype=torch.int32)
                     values = self._attribute_reduce_method(
-                        cluster, values, original_data.batch
+                        cluster,
+                        values,
+                        batch=batch,
                     )[0]
 
                 setattr(pooled_data, attr, values)
 
         return pooled_data
 
-    def __call__(self, data: Data) -> Data:
+    def __call__(self, data: Union[Data, Batch]) -> Union[Data, Batch]:
         """Coarsening operation."""
+
+        # Coarsen individual events if a batch is passed.
+        if isinstance(data, Batch):
+            data_list = data.to_data_list()
+
+            # Ensure that edge_index is kept when unbatching
+            if data.edge_index is not None:
+                edge_index_list = unbatch_edge_index(
+                    data.edge_index, data.batch
+                )
+                for i in range(len(data_list)):
+                    data_list[i].edge_index = (
+                        edge_index_list[i].contiguous().to(data.x.device)
+                    )
+
+            for i in range(len(data_list)):
+                data_list[i] = self.__call__(data_list[i])
+
+            batch = Batch.from_data_list(data_list)
+            batch.batch = batch.batch.to(data.batch.device)
+            return batch
+
         # Get tensor of cluster indices for each node.
-        cluster = self._perform_clustering(data)
+        cluster: LongTensor = self._perform_clustering(data)
 
         # Check whether a graph has already been built. Otherwise, set a dummy
         # connectivity, as this is required by pooling functions.
@@ -92,7 +158,13 @@ class Coarsening(ABC):
             data.edge_index = torch.tensor([[]], dtype=torch.int64)
 
         # Pool `data` object, including `x`, `batch`. and `edge_index`.
-        pooled_data = self._reduce_method(cluster, data)
+        pooled_data: Batch = self._reduce_method(cluster, data)
+        if isinstance(pooled_data, Batch):
+            pooled_data = Data(
+                x=pooled_data.x,
+                edge_index=pooled_data.edge_index,
+                edge_attr=pooled_data.edge_attr,
+            )
 
         # Optionally overwrite feature tensor
         x = self._additional_features(cluster, data)
diff --git a/src/graphnet/models/gnn/dynedge.py b/src/graphnet/models/gnn/dynedge.py
index 7e7acab7..5c77efbb 100644
--- a/src/graphnet/models/gnn/dynedge.py
+++ b/src/graphnet/models/gnn/dynedge.py
@@ -12,6 +12,7 @@ from torch_geometric.data import Data
 from torch_geometric.nn import EdgeConv
 from torch_scatter import scatter_max, scatter_mean, scatter_min, scatter_sum
 from graphnet.components.layers import DynEdgeConv
+from graphnet.models.coarsening import DOMCoarsening
 
 from graphnet.models.gnn.gnn import GNN
 from graphnet.models.utils import calculate_xyzt_homophily
@@ -160,6 +161,29 @@ class DynEdge(GNN):
         return x
 
 
+class DOMCoarsenedDynEdge(DynEdge):
+    def __init__(self, nb_inputs, layer_size_scale=4):
+        """DynEdge model.
+
+        Args:
+            nb_inputs (int): Number of input features.
+            layer_size_scale (int, optional): Integer that scales the size of
+                hidden layers. Defaults to 4.
+        """
+
+        # Base class constructor
+        super().__init__(nb_inputs, layer_size_scale)
+
+        # Graph operations
+        self.coarsening = DOMCoarsening()
+
+    def forward(self, data: Data) -> Tensor:
+        """Model forward pass."""
+        data = self.coarsening(data)
+        x = super().forward(data)
+        return x
+
+
 class DynEdge_V2(GNN):
     def __init__(self, nb_inputs, layer_size_scale=4):
         """DynEdge model.
diff --git a/tests/ml/test_coarsening.py b/tests/ml/test_coarsening.py
index 63a298e7..9690341c 100644
--- a/tests/ml/test_coarsening.py
+++ b/tests/ml/test_coarsening.py
@@ -8,81 +8,102 @@ from graphnet.models.coarsening import Coarsening
 
 
 # Utility method(s)
-def _get_test_data() -> Data:
+def _get_test_data() -> Batch:
     """Produce toy data for unit tests."""
 
-    x = torch.tensor(
-        [
-            [1, 0, 0.1],  # 0
-            [2, 1, 0.1],  # 1
-            [1, 1, 0.2],  # 2
-            [2, 1, 0.1],  # 1
-            [1, 1, 0.2],  # 2
-            [1, 0, 0.1],  # 3
-            [2, 1, 0.3],  # 4
-            [1, 1, 0.1],  # 5
-            [2, 1, 0.3],  # 4
-            [1, 1, 0.1],  # 6
-        ],
-        dtype=torch.float32,
+    data1 = Data(
+        x=torch.tensor(
+            [
+                [1, 0, 0.1],  # 0
+                [2, 1, 0.1],  # 1
+                [1, 1, 0.2],  # 2
+                [2, 1, 0.1],  # 1
+                [1, 1, 0.2],  # 2
+            ],
+            dtype=torch.float32,
+        ),
+        attr1=torch.tensor(
+            [
+                10,
+                11,
+                12,
+                13,
+                14,
+            ],
+            dtype=torch.int32,
+        ),
+        attr2=torch.tensor(
+            [
+                1.0,
+                2.0,
+                2.0,
+                3.0,
+                4.0,
+            ],
+            dtype=torch.float32,
+        ),
     )
 
-    batch = torch.tensor(
-        [
-            0,
-            0,
-            0,
-            0,
-            0,
-            1,
-            1,
-            1,
-            1,
-            2,
-        ],
-        dtype=torch.int64,
+    data2 = Data(
+        x=torch.tensor(
+            [
+                [1, 0, 0.1],  # 3
+                [2, 1, 0.3],  # 4
+                [1, 1, 0.1],  # 5
+                [2, 1, 0.3],  # 4
+            ],
+            dtype=torch.float32,
+        ),
+        attr1=torch.tensor(
+            [
+                15,
+                16,
+                17,
+                18,
+            ],
+            dtype=torch.int32,
+        ),
+        attr2=torch.tensor(
+            [
+                1.0,
+                4.0,
+                3.0,
+                1.0,
+            ],
+            dtype=torch.float32,
+        ),
     )
 
-    attr1 = torch.tensor(
-        [
-            10,
-            11,
-            12,
-            13,
-            14,
-            15,
-            16,
-            17,
-            18,
-            19,
-        ],
-        dtype=torch.int32,
+    data3 = Data(
+        x=torch.tensor(
+            [
+                [1, 1, 0.1],  # 6
+            ],
+            dtype=torch.float32,
+        ),
+        attr1=torch.tensor(
+            [
+                19,
+            ],
+            dtype=torch.int32,
+        ),
+        attr2=torch.tensor(
+            [
+                6.0,
+            ],
+            dtype=torch.float32,
+        ),
     )
 
-    attr2 = torch.tensor(
-        [
-            1.0,
-            2.0,
-            2.0,
-            3.0,
-            4.0,
-            1.0,
-            4.0,
-            3.0,
-            1.0,
-            6.0,
-        ],
-        dtype=torch.float32,
-    )
+    data_list = [data1, data2, data3]
+    for ix in range(len(data_list)):
+        data_list[ix]["x0"] = data_list[ix].x[:, 0]
+        data_list[ix]["x1"] = data_list[ix].x[:, 1]
+        data_list[ix]["x2"] = data_list[ix].x[:, 2]
 
-    data = Batch(batch=batch, x=x)
-    data["x0"] = x[:, 0]
-    data["x1"] = x[:, 1]
-    data["x2"] = x[:, 2]
-    data["attr1"] = attr1
-    data["attr2"] = attr2
+    batch = Batch.from_data_list([data1, data2, data3])
 
-    return data
+    return batch
 
 
 class SimpleCoarsening(Coarsening):