feat: graph learning dataset, model

src/arborist/data/datasets.py

@@ -8,6 +8,7 @@
 
 """
 
+from collections import defaultdict
 from copy import deepcopy
 from torch.utils.data import Dataset, DataLoader, Sampler
 
@@ -16,6 +17,8 @@
 import pandas as pd
 import torch
 
+from arborist.utils.graph_utils import topological_decomposition
+
 
 # --- Dataset Classes ---
 class CurveDataset(Dataset):
@@ -83,84 +86,207 @@ def __repr__(self):
         )
 
 
-class CurveDatasetCollection(Dataset):
+class GraphDataset(Dataset):
+    """
+    Dataset over rooted subgraphs of a single SkeletonGraph.
+
+    Each item corresponds to one root node. __getitem__ extracts the
+    subgraph within "max_depth" microns, decomposes it into irreducible paths
+    via topological decomposition, computes first-order finite differences for
+    each path, and returns a TreeSample with those curves and the line-graph
+    connectivity between them.
+    """
 
-    def __init__(self, datasets, is_val=False, n_val_examples=1000, seed=42):
+    def __init__(self, graph, root_nodes, max_depth=None, transform=None):
         """
+        Instantiates a GraphDataset object.
+
         Parameters
         ----------
-        datasets : List[CurveDataset]
-            List of PathsDataset instances, one per brain.
-        is_val : bool, optional
-            If True, precomputes a fixed set of examples at construction time.
-            Default is False.
-        n_val_examples : int, optional
-            Number of fixed validation examples to precompute. Default is 1000.
-        seed : int, optional
-            Random seed for reproducible val set. Default is 42.
+        graph : SkeletonGraph
+            The full skeleton graph to sample from.
+        root_nodes : List[int]
+            One root node per dataset item.
+        max_depth : float
+            Depth in microns for rooted subgraph extraction.
+        transform : callable, optional
+            Applied to each raw xyz array before differencing (e.g.
+            CurveTransforms for augmentation). Default is None.
         """
+        # Call parent class
+        super().__init__()
+
         # Instance attributes
+        self.graph = graph
+        self.root_nodes = root_nodes
+        self.max_depth = max_depth
+        self.transform = transform
+
+    def __getitem__(self, i):
+        # Extract tree sample components
+        root = self.root_nodes[i]
+        subgraph = self.graph.rooted_subgraph(root, self.max_depth)
+        _, paths, topo_edge_index = topological_decomposition(subgraph)
+
+        # Create list of curves
+        curves = []
+        for path in paths:
+            xyz = subgraph.node_xyz[path].copy()
+            if self.transform:
+                xyz = self.transform(xyz)
+            xyz -= xyz[0]
+            xyz[1:] -= xyz[:-1].copy()
+            curves.append(xyz)
+
+        # Create TreeSample
+        edge_index = _build_line_graph_edge_index(topo_edge_index)
+        return TreeSample(curves=curves, edge_index=edge_index)
+
+    def __len__(self):
+        return len(self.root_nodes)
+
+
+class TreeSample:
+    """
+    A rooted subgraph ready to pass through CurveEncoder then GraphTransformer.
+
+    Attributes
+    ----------
+    curves : List[numpy.ndarray]
+        One array per irreducible path, each of shape (N_i, 3). Values are
+        first-order finite differences with a leading zero row, matching the
+        convention expected by CurveEncoder.
+    edge_index : numpy.ndarray
+        Shape (2, E), dtype int64. Line-graph adjacency: two curves share an
+        edge when they meet at a topological node (branch point or leaf),
+        so message passing over this graph communicates between neighboring
+        branches.
+    """
+
+    def __init__(self, curves, edge_index):
+        self.curves = curves
+        self.edge_index = edge_index
+
+    def __repr__(self):
+        return (
+            f"TreeSample("
+            f"n_curves={len(self.curves)}, "
+            f"n_edges={self.edge_index.shape[1]})"
+        )
+
+
+def _build_line_graph_edge_index(topo_edge_index):
+    """
+    Converts topological-graph edge pairs into line-graph edge pairs.
+
+    In the topological graph each node is a branching/leaf point and each
+    edge is an irreducible path (a curve). In the line graph each curve
+    becomes a node and two curve-nodes are connected when they share a
+    topological endpoint.
+
+    Parameters
+    ----------
+    topo_edge_index : List[Tuple[int, int]]
+        Edges of the topological graph as (src_topo_idx, dst_topo_idx) pairs,
+        parallel to the list of curves.
+
+    Returns
+    -------
+    numpy.ndarray
+        Shape (2, E), int64.
+    """
+    topo_to_curves = defaultdict(list)
+    for curve_idx, (u, v) in enumerate(topo_edge_index):
+        topo_to_curves[u].append(curve_idx)
+        topo_to_curves[v].append(curve_idx)
+
+    src, dst = [], []
+    for neighbors in topo_to_curves.values():
+        for i in neighbors:
+            for j in neighbors:
+                if i != j:
+                    src.append(i)
+                    dst.append(j)
+
+    if not src:
+        return np.zeros((2, 0), dtype=np.int64)
+    return np.array([src, dst], dtype=np.int64)
+
+
+class DatasetCollection(Dataset):
+    """
+    A flat, indexable view over multiple datasets (one per brain/specimen).
+
+    Parameters
+    ----------
+    datasets : List[Dataset]
+        Constituent datasets to combine.
+    weight_fn : callable, optional
+        Maps a dataset to a 1-D array of per-item sampling weights. Used by
+        samplers for non-uniform drawing (e.g. length-weighted curve sampling).
+        Defaults to uniform weights when None. Default is None.
+    is_val : bool, optional
+        If True, precomputes a fixed set of examples at construction time.
+        Default is False.
+    n_val_examples : int, optional
+        Number of validation examples to precompute. Default is 1000.
+    seed : int, optional
+        Random seed for reproducible val set. Default is 42.
+    """
+
+    def __init__(
+        self,
+        datasets,
+        weight_fn=None,
+        is_val=False,
+        n_val_examples=1000,
+        seed=42,
+    ):
         self.datasets = datasets
         self.is_val = is_val
-        self.set_examples_df()
-
-        # Check whether to set validation examples
+        self._build_index(weight_fn)
         if is_val:
-            self.val_examples = self.set_val_examples(n_val_examples, seed)
+            self.val_examples = self._precompute_val(n_val_examples, seed)
 
-    def set_examples_df(self):
+    def _build_index(self, weight_fn):
         rows = []
         for ds_idx, dataset in enumerate(self.datasets):
-            ds_idxs = np.full(len(dataset), ds_idx)
-            p_idxs = np.arange(len(dataset))
-            ds_lengths = dataset.curve_lengths()
-            rows.append(
-                pd.DataFrame(
-                    {
-                        "ds_idx": ds_idxs,
-                        "path_idx": p_idxs,
-                        "length": ds_lengths,
-                    }
-                )
-            )
-        self.examples_df = pd.concat(rows, ignore_index=True)
-
-    def set_val_examples(self, n, seed):
-        """
-        Samples n examples with fixed seed, strips transforms, and caches
-        the resulting examples.
-        """
+            n = len(dataset)
+            weights = weight_fn(dataset) if weight_fn is not None else np.ones(n)
+            rows.append(pd.DataFrame({
+                "ds_idx": np.full(n, ds_idx, dtype=int),
+                "item_idx": np.arange(n),
+                "weight": weights,
+            }))
+        self.index = pd.concat(rows, ignore_index=True)
+
+    def _precompute_val(self, n, seed):
         rng = np.random.default_rng(seed)
-        indices = rng.choice(len(self.examples_df), size=n, replace=False)
+        idxs = rng.choice(len(self.index), size=n, replace=False)
         examples = []
-        for i in indices:
-            ds_idx = self.examples_df["ds_idx"][i]
-            path_idx = self.examples_df["path_idx"][i]
-            dataset = self.datasets[ds_idx]
-            examples.append(dataset[path_idx])
+        for i in idxs:
+            ds_idx = self.index["ds_idx"][i]
+            item_idx = self.index["item_idx"][i]
+            examples.append(self.datasets[ds_idx][item_idx])
         return examples
 
-    # --- Data Fetching ---
     def __getitem__(self, i):
-        # Case 1: validation example
         if self.is_val:
             return self.val_examples[i]
-
-        # Case 2: train example
-        ds_idx = self.examples_df["ds_idx"][i]
-        path_idx = self.examples_df["path_idx"][i]
-        return self.datasets[ds_idx][path_idx]
+        ds_idx = self.index["ds_idx"][i]
+        item_idx = self.index["item_idx"][i]
+        return self.datasets[ds_idx][item_idx]
 
     def __len__(self):
         if self.is_val:
             return len(self.val_examples)
-        return len(self.examples_df)
+        return len(self.index)
 
     def __repr__(self):
         return (
-            f"CurveDatasetCollection("
-            f"num_brains={len(self.datasets)}, "
-            f"num_curves={len(self.examples_df)}) "
+            f"DatasetCollection("
+            f"num_datasets={len(self.datasets)}, "
+            f"num_items={len(self.index)})"
         )
 
 
@@ -178,8 +304,8 @@ def __init__(self, dataset, examples_per_epoch):
         self.examples_per_epoch = examples_per_epoch
 
     def __iter__(self):
-        idxs = self.dataset.examples_df.sample(
-            self.examples_per_epoch, replace=True, weights="length"
+        idxs = self.dataset.index.sample(
+            self.examples_per_epoch, replace=True, weights="weight"
         ).index
         return iter(np.array(idxs))