Add files via upload

mace-bench/3rdparty/SevenNet/sevenn/train/collate.py

+from typing import Any, List, Optional, Sequence
+
+from ase.atoms import Atoms
+from torch_geometric.loader.dataloader import Collater
+
+from sevenn.atom_graph_data import AtomGraphData
+
+from .dataload import atoms_to_graph
+
+
+class AtomsToGraphCollater(Collater):
+
+    def __init__(
+        self,
+        dataset: Sequence[Atoms],
+        cutoff: float,
+        transfer_info: bool = False,
+        follow_batch: Optional[List[str]] = None,
+        exclude_keys: Optional[List[str]] = None,
+        y_from_calc: bool = True,
+    ):
+        # quite original collator's type mismatch with []
+        super().__init__([], follow_batch, exclude_keys)
+        self.dataset = dataset
+        self.cutoff = cutoff
+        self.transfer_info = transfer_info
+        self.y_from_calc = y_from_calc
+
+    def __call__(self, batch: List[Any]) -> Any:
+        # build list of graph
+        graph_list = []
+        for stct in batch:
+            graph = atoms_to_graph(
+                stct,
+                self.cutoff,
+                transfer_info=self.transfer_info,
+                y_from_calc=self.y_from_calc,
+            )
+            graph = AtomGraphData.from_numpy_dict(graph)
+            graph_list.append(graph)
+        return super().__call__(graph_list)