feat: add GraphTransforms and graph_transform support to GraphDataset

src/arborist/data/augmentation.py

@@ -12,6 +12,40 @@
 import random
 
 
+class GraphTransforms:
+    """
+    Applies a sequence of transforms to all node coordinates of a subgraph.
+
+    Unlike CurveTransforms (which acts on one curve at a time), this class
+    operates on the full (N_nodes, 3) node_xyz array of a rooted subgraph so
+    that every curve in the graph receives the same rotation and mirror flip,
+    preserving inter-curve spatial relationships.
+    """
+
+    def __init__(self):
+        self.transforms = [
+            RandomRotation3D(),
+            RandomMirror3D(),
+            RandomJitter3D(),
+        ]
+
+    def __call__(self, xyz):
+        """
+        Parameters
+        ----------
+        xyz : numpy.ndarray
+            Shape (N, 3) — coordinates of all nodes in the subgraph.
+
+        Returns
+        -------
+        numpy.ndarray
+            Augmented coordinates, same shape.
+        """
+        for transform in self.transforms:
+            xyz = transform(xyz)
+        return xyz
+
+
 class CurveTransforms:
     """
     Class that applies a sequence of transforms to a 3D space curve.

src/arborist/data/datasets.py

@@ -98,7 +98,7 @@ class GraphDataset(Dataset):
     connectivity between them.
     """
 
-    def __init__(self, graph, root_nodes, max_depth=None, transform=None):
+    def __init__(self, graph, root_nodes, max_depth=None, transform=None, graph_transform=None):
         """
         Instantiates a GraphDataset object.
 
@@ -113,6 +113,12 @@ def __init__(self, graph, root_nodes, max_depth=None, transform=None):
         transform : callable, optional
             Applied to each raw xyz array before differencing (e.g.
             CurveTransforms for augmentation). Default is None.
+        graph_transform : callable, optional
+            Applied to the full (N_nodes, 3) node_xyz array of the rooted
+            subgraph before path decomposition (e.g. GraphTransforms). Because
+            it acts on all node coordinates at once, every curve in the
+            subgraph receives the same rotation and mirror flip, preserving
+            inter-curve spatial relationships. Default is None.
         """
         # Call parent class
         super().__init__()
@@ -122,16 +128,20 @@ def __init__(self, graph, root_nodes, max_depth=None, transform=None):
         self.root_nodes = root_nodes
         self.max_depth = max_depth
         self.transform = transform
+        self.graph_transform = graph_transform
         self.config = {
             "max_depth": max_depth,
             "transform": type(transform).__name__ if transform else None,
+            "graph_transform": type(graph_transform).__name__ if graph_transform else None,
         }
 
 
     def __getitem__(self, i):
         # Extract tree sample components
         root = self.root_nodes[i]
         subgraph = self.graph.rooted_subgraph(root, self.max_depth)
+        if self.graph_transform:
+            subgraph.node_xyz = self.graph_transform(subgraph.node_xyz)
         _, paths, topo_edge_index = topological_decomposition(subgraph)
 
         # Create list of curves

src/arborist/models/graph_transformer.py

@@ -104,8 +104,7 @@ def _build_attn_mask(self, n, edge_index, device):
         True at (i, j) blocks attention from node i to node j.
         Every node attends to itself and all immediate neighbors.
         """
-        mask = torch.ones(n, n, dtype=torch.bool, device=device)
-        mask.fill_diagonal_(False)
+        mask = ~torch.eye(n, dtype=torch.bool, device=device)
         if edge_index.shape[1] > 0:
             mask[edge_index[0], edge_index[1]] = False
         return mask