Add PyG GNNs Support

docs/data.md

@@ -163,3 +163,62 @@ test_src,  test_tgt  = dm.load("synthetic_graphs", train=False)
 |-----|--------|--------|-------------|
 | `"mnist_noisy_mnist"` | `torchvision.MNIST` | `NoisyMNIST` | `noise_std` (default `0.3`) |
 | `"synthetic_graphs"` | `SyntheticGraphDataset` (noise=0.1) | `SyntheticGraphDataset` (noise=0.5, flip=0.05) | `n_graphs`, `n_nodes`, `feat_dim`, `feature_noise_src`, `feature_noise_tgt`, `edge_flip_prob` |
+| `"pyg_domains"` | User-supplied PyG `Data` | User-supplied PyG `Data` | `source`, `target`, `task_level`, `train_ratio`, `val_ratio`, `split_seed`, `split_mode` |
+
+---
+
+## PyG domains (`pyg_domains`)
+
+Requires `torch-geometric`. Loads external PyG graphs for domain adaptation with automatic **stratified** train/val/test splits when masks are not already on the `Data` objects.
+
+### Node-level (one graph per domain)
+
+Use when labels live on **nodes** (transductive node classification). Each domain is a single `torch_geometric.data.Data` object. Training runs message passing on the full graph; loss and MMD use **train** nodes only; evaluation uses **test** nodes.
+
+```python
+from shiftkit.data import DataManager
+from shiftkit.models import GNN
+from shiftkit.methods import MMDTrainer
+
+dm = DataManager(batch_size=1, num_workers=0)
+train_src, train_tgt = dm.load(
+    "pyg_domains",
+    train=True,
+    task_level="node",
+    source=source_graph,
+    target=target_graph,
+    train_ratio=0.6,
+    val_ratio=0.2,
+    split_seed=42,
+    split_mode="stratified",
+)
+test_src, test_tgt = dm.load("pyg_domains", train=False, ...)
+
+model = GNN(source_graph, "SAGE", hidden_channels=64, num_layers=2,
+            num_classes=10, pool="none")
+trainer = MMDTrainer(model, train_src, train_tgt, mmd_weight=1.0)
+```
+
+Pair with `shiftkit.models.GNN(..., pool="none")` so `encode()` returns per-node embeddings.
+
+### Graph-level (many graphs per domain)
+
+Pass a **list** of `Data` objects per domain. Splits are by graph index (stratified on graph labels when discrete). Use default `pool="mean"` on `GNN`.
+
+```python
+train_src, train_tgt = dm.load(
+    "pyg_domains",
+    train=True,
+    task_level="graph",
+    source=list_of_src_graphs,
+    target=list_of_tgt_graphs,
+    train_ratio=0.6,
+    val_ratio=0.2,
+)
+```
+
+### Existing masks
+
+If `data.train_mask` is already set, automatic splitting is skipped. Loaders use `train_mask` for `train=True` and `test_mask` for `train=False`.
+
+See `examples/pyg_node_mmd.py` for a full node-level example.

examples/pyg_node_mmd.py

@@ -0,0 +1,92 @@
+"""
+Example: node-level domain adaptation on two PyG graphs (one per domain).
+
+Uses DataManager.load("pyg_domains") with stratified node masks and
+shiftkit.models.GNN with pool="none".
+
+Run from repo root:
+    python examples/pyg_node_mmd.py
+"""
+
+import sys
+import os
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import torch
+from torch_geometric.data import Data
+
+from shiftkit.data import DataManager
+from shiftkit.models import GNN
+from shiftkit.methods import MMDTrainer, SourceOnlyTrainer
+
+
+def make_domain_graph(n_nodes: int, feat_dim: int, n_classes: int, seed: int, shift: float = 0.0) -> Data:
+    torch.manual_seed(seed)
+    x = torch.randn(n_nodes, feat_dim) + shift
+    row = torch.arange(n_nodes - 1)
+    edge_index = torch.stack([row, row + 1], dim=0)
+    edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+    y = torch.randint(0, n_classes, (n_nodes,))
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+if __name__ == "__main__":
+    N_NODES = 300
+    FEAT = 8
+    NUM_CLASSES = 3
+    EPOCHS = 30
+
+    source_graph = make_domain_graph(N_NODES, FEAT, NUM_CLASSES, seed=0, shift=0.0)
+    target_graph = make_domain_graph(N_NODES, FEAT, NUM_CLASSES, seed=1, shift=1.5)
+
+    dm = DataManager(batch_size=1, num_workers=0)
+    train_src, train_tgt = dm.load(
+        "pyg_domains",
+        train=True,
+        task_level="node",
+        source=source_graph,
+        target=target_graph,
+        train_ratio=0.6,
+        val_ratio=0.2,
+        split_seed=42,
+        split_mode="stratified",
+    )
+    test_src, test_tgt = dm.load(
+        "pyg_domains",
+        train=False,
+        task_level="node",
+        source=source_graph,
+        target=target_graph,
+        train_ratio=0.6,
+        val_ratio=0.2,
+        split_seed=42,
+        split_mode="stratified",
+    )
+
+    model_so = GNN(
+        source_graph, "SAGE", hidden_channels=32, num_layers=2,
+        num_classes=NUM_CLASSES, pool="none",
+    )
+    model_mmd = GNN(
+        source_graph, "SAGE", hidden_channels=32, num_layers=2,
+        num_classes=NUM_CLASSES, pool="none",
+    )
+
+    print("Training Source-Only...")
+    so = SourceOnlyTrainer(model_so, train_src, train_tgt, lr=1e-3, device="cpu")
+    so.fit(epochs=EPOCHS)
+
+    print("Training MMD...")
+    mmd = MMDTrainer(model_mmd, train_src, train_tgt, mmd_weight=0.5, lr=1e-3, device="cpu")
+    mmd.fit(epochs=EPOCHS)
+
+    for name, trainer in [("Source-Only", so), ("MMD", mmd)]:
+        r_src = trainer.evaluate(test_src, domain="source-test")
+        r_tgt = trainer.evaluate(test_tgt, domain="target-test")
+        print(
+            f"{name:12s}  src acc={r_src['accuracy']*100:.1f}%  "
+            f"tgt acc={r_tgt['accuracy']*100:.1f}%"
+        )
+
+    print("Done.")

requirements.txt

@@ -5,6 +5,9 @@ matplotlib>=3.7.0
 scikit-learn>=1.2.0
 tqdm>=4.65.0
 
+# Optional — required only for shiftkit.models.GNN (PyTorch Geometric)
+# torch-geometric>=2.4.0
+
 # Optional — required only for UMAP projection in diagnostics
 # umap-learn>=0.5.0
 

shiftkit/__init__.py

@@ -12,6 +12,10 @@
 from .data.datasets      import DataManager
 from .models.networks    import MLP, CNN, MLPRegressor
 from .models.gnn         import SimpleGCN
+try:
+    from .models.gnn_pyg import GNN
+except ImportError:
+    GNN = None  # torch-geometric not installed
 from .methods.base       import BaseTrainer, TrainerRegistry
 from .methods.mmd        import MMDLoss, MMDTrainer, SourceOnlyTrainer
 from .methods.lmmd       import LMMDLoss, LMMDTrainer
@@ -22,8 +26,8 @@
 from .methods.kliep      import KLIEPWeightEstimator, KLIEPTrainer
 from .data.datasets      import SineWaveDataset, CaliforniaHousingDataset
 from .diagnostics.plots  import (
-    plot_latent_space, plot_training_history, compare_latent_spaces,
-    plot_confusion_matrix, plot_roc_curve,
+    plot_latent_space, plot_latent_space_domains, plot_training_history,
+    compare_latent_spaces, plot_confusion_matrix, plot_roc_curve,
 )
 
 __version__ = "0.1.0"
@@ -39,6 +43,9 @@
     "SIDDATrainer",
     "SourceOnlyRegressionTrainer", "MMDRegressionTrainer",
     "KLIEPWeightEstimator", "KLIEPTrainer",
-    "plot_latent_space", "plot_training_history", "compare_latent_spaces",
+    "plot_latent_space", "plot_latent_space_domains", "plot_training_history",
+    "compare_latent_spaces",
     "plot_confusion_matrix", "plot_roc_curve",
 ]
+if GNN is not None:
+    __all__.append("GNN")

shiftkit/data/__init__.py

@@ -1,3 +1,12 @@
 from .datasets import DataManager, SineWaveDataset, CaliforniaHousingDataset
 
+try:
+    from .pyg_utils import NodeGraphBatch, ensure_masks, build_pyg_domain_loaders
+except ImportError:
+    NodeGraphBatch = None
+    ensure_masks = None
+    build_pyg_domain_loaders = None
+
 __all__ = ["DataManager", "SineWaveDataset", "CaliforniaHousingDataset"]
+if NodeGraphBatch is not None:
+    __all__ += ["NodeGraphBatch", "ensure_masks", "build_pyg_domain_loaders"]

shiftkit/data/datasets.py

@@ -428,6 +428,49 @@ def _california_housing(root, batch_size, train, num_workers, **kw):
 
     _REGISTRY["california_housing"] = _california_housing
 
+    def _pyg_domains(root, batch_size, train, num_workers, **kw):
+        """
+        PyG source/target domain pair (graph-level or node-level).
+
+        Required kwargs
+        ---------------
+        source, target : PyG ``Data`` (node-level) or list/dataset of ``Data`` (graph-level)
+        task_level     : ``"node"`` or ``"graph"`` (default ``"node"``)
+
+        Optional kwargs
+        -----------------
+        train_ratio, val_ratio, split_seed, split_mode (``"stratified"`` | ``"random"``)
+        """
+        from .pyg_utils import build_pyg_domain_loaders
+
+        source = kw.get("source")
+        target = kw.get("target")
+        if source is None or target is None:
+            raise ValueError(
+                "pyg_domains requires 'source' and 'target' PyG Data object(s). "
+                "Example: dm.load('pyg_domains', source=src_data, target=tgt_data, ...)"
+            )
+        task_level = kw.get("task_level", "node")
+        train_ratio = float(kw.get("train_ratio", 0.6))
+        val_ratio = float(kw.get("val_ratio", 0.2))
+        split_seed = int(kw.get("split_seed", 42))
+        split_mode = kw.get("split_mode", "stratified")
+
+        return build_pyg_domain_loaders(
+            task_level=task_level,
+            source=source,
+            target=target,
+            train=train,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            train_ratio=train_ratio,
+            val_ratio=val_ratio,
+            split_seed=split_seed,
+            split_mode=split_mode,
+        )
+
+    _REGISTRY["pyg_domains"] = _pyg_domains
+
 
 _register_defaults()