test: add parity and end-to-end workflow coverage

examples/README.md

@@ -0,0 +1,18 @@
+# Examples
+
+The supported experiment interface is the root CLI with Hydra configuration;
+examples do not define a second application framework. The repository's
+reference groups cover multiclass pretraining, binary fine-tuning, resume,
+evaluation, and NPZ prediction. Start with the commands in the top-level
+[`README.md`](../README.md), then copy a config group into a local config file
+when a dataset needs more than command-line overrides.
+
+For a complete CPU smoke workflow using generated temporary ROOT data:
+
+```bash
+uv run python scripts/dev/smoke_end_to_end.py
+```
+
+Values such as input ROOT paths, cache locations, and checkpoint paths are
+deliberately local placeholders. Production datasets and generated outputs do
+not belong in this directory.

scripts/dev/shrink_root_sample.py

@@ -0,0 +1,57 @@
+"""Create a small ROOT fixture from a Delphes sample."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import uproot
+
+BRANCHES = (
+    "jet_pt",
+    "jet_eta",
+    "jet_phi",
+    "jet_btag",
+    "ph_pt",
+    "ph_eta",
+    "ph_phi",
+    "ele_pt",
+    "ele_eta",
+    "ele_phi",
+    "ele_charge",
+    "mu_pt",
+    "mu_eta",
+    "mu_phi",
+    "mu_charge",
+    "MET_met",
+    "MET_phi",
+    "weight",
+    "Number",
+)
+
+
+def shrink_sample(source: Path, target: Path, entries: int) -> None:
+    """Copy the active branches and first ``entries`` events to ``target``."""
+    if entries < 1:
+        raise ValueError("entries must be positive")
+    target.parent.mkdir(parents=True, exist_ok=True)
+    with uproot.open(source) as source_file:
+        arrays = source_file["output"].arrays(
+            BRANCHES, entry_start=0, entry_stop=entries, library="ak"
+        )
+        with uproot.recreate(target) as target_file:
+            target_file["output"] = arrays
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("source", type=Path)
+    parser.add_argument("target", type=Path)
+    parser.add_argument("--entries", type=int, default=64)
+    args = parser.parse_args()
+    shrink_sample(args.source, args.target, args.entries)
+    print(f"Wrote {args.entries} events to {args.target}")
+
+
+if __name__ == "__main__":
+    main()

scripts/dev/smoke_end_to_end.py

@@ -0,0 +1,131 @@
+"""Run the documented new-stack workflow on a tiny temporary ROOT sample.
+
+This is an integration smoke test, not a physics example.  It requires the
+validated ``root-gnn`` extra and uses only the public CLI for preparation,
+training, evaluation, and prediction.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import awkward as ak
+import numpy as np
+import uproot
+
+
+def _run(root: Path, *arguments: str) -> None:
+    command = [sys.executable, "-m", "gnn4colliders.cli", *arguments]
+    subprocess.run(command, cwd=root, check=True)
+
+
+def _write_root(path: Path) -> None:
+    events = 5
+    with uproot.recreate(path) as output:
+        output["Events"] = {
+            "jet_pt": ak.Array([[40.0, 25.0]] * events),
+            "jet_eta": ak.Array([[-1.0, 1.0]] * events),
+            "jet_phi": ak.Array([[-2.5, 2.5]] * events),
+            "eventNumber": np.arange(events, dtype=np.int64),
+            "weight": np.ones(events, dtype=np.float32),
+        }
+
+
+def _prepare(root_file: Path, cache: Path) -> None:
+    _run(
+        root_file.parent,
+        "prepare",
+        f"data.files=[{root_file}]",
+        "data.tree_name=Events",
+        f"data.cache.path={cache}",
+        'data.feature_branches=[["jet_pt"],["jet_eta"],["jet_phi"],CALC_E,[1.0],[0.0],NODE_TYPE]',
+        "data.object_types=[vector]",
+        "data.scales=[1,1,1,1,1,1,1]",
+        "data.fold_var=eventNumber",
+        "data.weight_var=weight",
+    )
+
+
+def _train(root: Path, cache: Path, output: Path, *extra: str) -> Path:
+    _run(
+        root,
+        "train",
+        f"data.cache.path={cache}",
+        "trainer.max_epochs=1",
+        "trainer.device=cpu",
+        "data.batch_size=1",
+        "model.hid_size=8",
+        "model.n_layers=1",
+        "model.n_proc_steps=1",
+        f"environment.output_root={output}",
+        *extra,
+    )
+    return output / "checkpoints" / "epoch_0000.pt"
+
+
+def main() -> None:
+    try:
+        import dgl  # noqa: F401
+    except ImportError as error:  # pragma: no cover - environment-dependent
+        raise SystemExit(
+            "install the root-gnn extra before running this smoke test"
+        ) from error
+
+    with tempfile.TemporaryDirectory(prefix="gnn4colliders-smoke-") as directory:
+        root = Path(directory)
+        root_file = root / "events.root"
+        cache = root / "graphs.pt"
+        target_cache = root / "target.pt"
+        _write_root(root_file)
+        _prepare(root_file, cache)
+        pretrained = _train(root, cache, root / "pretrain")
+        _prepare(root_file, target_cache)
+        fine_tuned = _train(
+            root,
+            target_cache,
+            root / "finetune",
+            "model=root_gnn/fine_tuned_edge_network",
+            "task=binary_classification",
+            f"checkpoint.pretrained={pretrained}",
+            "model.freeze_backbone=true",
+        )
+        _run(
+            root,
+            "evaluate",
+            f"data.cache.path={target_cache}",
+            "inference.split=test",
+            f"inference.checkpoint={fine_tuned}",
+            "model=root_gnn/fine_tuned_edge_network",
+            "task=binary_classification",
+            f"checkpoint.pretrained={pretrained}",
+            "model.hid_size=8",
+            "model.n_layers=1",
+            "model.n_proc_steps=1",
+            "trainer.device=cpu",
+        )
+        prediction = root / "predictions.npz"
+        _run(
+            root,
+            "predict",
+            f"data.cache.path={target_cache}",
+            "inference.split=test",
+            f"inference.checkpoint={fine_tuned}",
+            "model=root_gnn/fine_tuned_edge_network",
+            "task=binary_classification",
+            f"checkpoint.pretrained={pretrained}",
+            "model.hid_size=8",
+            "model.n_layers=1",
+            "model.n_proc_steps=1",
+            f"inference.output={prediction}",
+            "trainer.device=cpu",
+        )
+        if not prediction.is_file():
+            raise RuntimeError("smoke workflow did not produce predictions.npz")
+        print(f"smoke workflow succeeded in {root}")
+
+
+if __name__ == "__main__":
+    main()

tests/parity/conftest.py

@@ -0,0 +1,88 @@
+"""Helpers for importing the legacy ROOT-GNN implementation in parity tests."""
+
+import os
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+
+def _dgl_is_importable():
+    try:
+        import dgl  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+@pytest.fixture(scope="session", autouse=True)
+def require_root_gnn_dependencies():
+    if (
+        os.environ.get("GNN4COLLIDERS_REQUIRE_ROOT_GNN") == "1"
+        and not _dgl_is_importable()
+    ):
+        pytest.fail(
+            "ROOT-GNN parity requires an importable DGL installation",
+            pytrace=False,
+        )
+
+
+@pytest.fixture(scope="session")
+def legacy_dataset_module():
+    """Return the active legacy dataset module when DGL is available."""
+    pytest.importorskip("dgl")
+    legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl"
+    legacy_root_string = str(legacy_root)
+    if legacy_root_string not in sys.path:
+        sys.path.insert(0, legacy_root_string)
+    from root_gnn_base import dataset
+
+    return dataset
+
+
+@pytest.fixture(scope="session")
+def legacy_model_module():
+    """Return the legacy active model module for fixed-weight parity tests."""
+    pytest.importorskip("dgl")
+    legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl"
+    legacy_root_string = str(legacy_root)
+    if legacy_root_string not in sys.path:
+        sys.path.insert(0, legacy_root_string)
+    from models import GCN
+
+    return GCN
+
+
+@pytest.fixture(scope="session")
+def legacy_dataset_module_without_dgl():
+    """Import pure legacy preprocessing with a minimal DGL import shim."""
+    legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl"
+    legacy_root_string = str(legacy_root)
+    if legacy_root_string not in sys.path:
+        sys.path.insert(0, legacy_root_string)
+    try:
+        import dgl  # noqa: F401
+    except ImportError:
+        dgl_module = types.ModuleType("dgl")
+        dgl_data_module = types.ModuleType("dgl.data")
+        dgl_data_module.DGLDataset = type("DGLDataset", (), {})
+        dgl_module.data = dgl_data_module
+        sys.modules["dgl"] = dgl_module
+        sys.modules["dgl.data"] = dgl_data_module
+        matplotlib_module = types.ModuleType("matplotlib")
+        matplotlib_pyplot_module = types.ModuleType("matplotlib.pyplot")
+        matplotlib_module.pyplot = matplotlib_pyplot_module
+        sys.modules["matplotlib"] = matplotlib_module
+        sys.modules["matplotlib.pyplot"] = matplotlib_pyplot_module
+    from root_gnn_base import dataset
+
+    return dataset
+
+
+@pytest.fixture(scope="session")
+def legacy_utils_module(legacy_dataset_module_without_dgl):
+    """Return legacy fold-selection helpers after the package path is set up."""
+    from root_gnn_base import utils
+
+    return utils

tests/parity/test_legacy_node_features.py

@@ -0,0 +1,100 @@
+"""Characterization tests for the active legacy node feature builder."""
+
+import numpy as np
+import pytest
+import torch
+
+
+@pytest.fixture
+def event():
+    return {
+        "jet_pt": np.array([100.0, 50.0], dtype=np.float32),
+        "ele_pt": np.array([20.0], dtype=np.float32),
+        "mu_pt": np.array([30.0], dtype=np.float32),
+        "ph_pt": np.array([40.0], dtype=np.float32),
+        "MET_met": np.float32(25.0),
+        "jet_eta": np.array([1.0, -0.5], dtype=np.float32),
+        "ele_eta": np.array([0.25], dtype=np.float32),
+        "mu_eta": np.array([-0.75], dtype=np.float32),
+        "ph_eta": np.array([0.5], dtype=np.float32),
+        "jet_phi": np.array([3.0, -3.0], dtype=np.float32),
+        "ele_phi": np.array([0.2], dtype=np.float32),
+        "mu_phi": np.array([-0.4], dtype=np.float32),
+        "ph_phi": np.array([1.0], dtype=np.float32),
+        "MET_phi": np.float32(-1.2),
+        "jet_btag": np.array([0.8, 0.1], dtype=np.float32),
+        "ele_charge": np.array([-1.0], dtype=np.float32),
+        "mu_charge": np.array([1.0], dtype=np.float32),
+    }
+
+
+@pytest.fixture
+def node_schema():
+    return (
+        [
+            ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"],
+            ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0],
+            ["jet_phi", "ele_phi", "mu_phi", "ph_phi", "MET_phi"],
+            "CALC_E",
+            ["jet_btag", 0, 0, 0, 0],
+            [0, "ele_charge", "mu_charge", 0, 0],
+            "NODE_TYPE",
+        ],
+        ["vector", "vector", "vector", "vector", "single"],
+        [0.1, 1, 1, 0.1, 1, 1, 1],
+    )
+
+
+def test_node_features_have_active_schema_and_order(
+    legacy_dataset_module_without_dgl, event, node_schema
+):
+    names, types, scales = node_schema
+    features, lengths = legacy_dataset_module_without_dgl.node_features_from_tree(
+        event, names, types, torch.tensor(scales, dtype=torch.float32)
+    )
+    expected = np.array(
+        [
+            [10.0, 1.0, 3.0, 15.431, 0.8, 0.0, 0.0],
+            [5.0, -0.5, -3.0, 5.638, 0.1, 0.0, 0.0],
+            [2.0, 0.25, 0.2, 2.063, 0.0, -1.0, 1.0],
+            [3.0, -0.75, -0.4, 3.884, 0.0, 1.0, 2.0],
+            [4.0, 0.5, 1.0, 4.511, 0.0, 0.0, 3.0],
+            [2.5, 0.0, -1.2, 2.500, 0.0, 0.0, 4.0],
+        ],
+        dtype=np.float32,
+    )
+    assert lengths == [2, 1, 1, 1, 1]
+    assert features.shape == (6, 7)
+    assert features.dtype == torch.float32
+    np.testing.assert_allclose(features.numpy(), expected, rtol=0, atol=2e-3)
+
+
+def test_node_feature_builder_preserves_zero_length_vector(
+    legacy_dataset_module_without_dgl, node_schema
+):
+    names, types, scales = node_schema
+    event = {
+        "jet_pt": np.array([], dtype=np.float32),
+        "ele_pt": np.array([20.0], dtype=np.float32),
+        "mu_pt": np.array([], dtype=np.float32),
+        "ph_pt": np.array([], dtype=np.float32),
+        "MET_met": np.float32(25.0),
+        "jet_btag": np.array([], dtype=np.float32),
+        "ele_charge": np.array([-1.0], dtype=np.float32),
+        "mu_charge": np.array([], dtype=np.float32),
+        "jet_eta": np.array([], dtype=np.float32),
+        "ele_eta": np.array([0.25], dtype=np.float32),
+        "mu_eta": np.array([], dtype=np.float32),
+        "ph_eta": np.array([], dtype=np.float32),
+        "jet_phi": np.array([], dtype=np.float32),
+        "ele_phi": np.array([0.2], dtype=np.float32),
+        "mu_phi": np.array([], dtype=np.float32),
+        "ph_phi": np.array([], dtype=np.float32),
+        "MET_phi": np.float32(-1.2),
+    }
+    features, lengths = legacy_dataset_module_without_dgl.node_features_from_tree(
+        event, names, types, torch.tensor(scales)
+    )
+    assert lengths == [0, 1, 0, 0, 1]
+    assert features.shape == (2, 7)
+    assert features[:, 6].tolist() == [1.0, 4.0]

tests/parity/test_model_parity.py

@@ -0,0 +1,110 @@
+"""Fixed-weight parity tests for the active legacy and rewritten models."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import (
+    EdgeNetwork,
+    FineTunedEdgeNetwork,
+    load_legacy_edge_network_state_dict,
+)
+
+pytest.importorskip("dgl")
+
+
+def _graph_and_globals():
+    graph = build_dgl_graph(
+        torch.tensor(
+            [[10.0, 0.2, 0.1], [8.0, -0.3, -0.2], [4.0, 0.1, 2.8]],
+            dtype=torch.float32,
+        )
+    )
+    return graph, torch.tensor([[1.0, 2.0]], dtype=torch.float32)
+
+
+def _new_model(graph, globals_):
+    return EdgeNetwork(graph, globals_, 8, 3, 2, 2, dropout=0.0).eval()
+
+
+def test_edge_network_matches_legacy_forward_and_representation(
+    legacy_model_module,
+):
+    graph, globals_ = _graph_and_globals()
+    new_model = _new_model(graph, globals_)
+    legacy_model = legacy_model_module.Edge_Network(
+        graph, globals_, 8, 3, 2, 2, dropout=0.0
+    ).eval()
+    legacy_state = {
+        key.replace("classifier.", "classify."): value
+        for key, value in new_model.state_dict().items()
+    }
+    legacy_model.load_state_dict(legacy_state)
+
+    with torch.no_grad():
+        new_representation = new_model.forward_features(graph, globals_)
+        legacy_representation = legacy_model.representation(graph.clone(), globals_)[1]
+        new_logits = new_model(graph, globals_)
+        legacy_logits = legacy_model(graph.clone(), globals_)
+    assert torch.equal(new_representation, legacy_representation)
+    assert torch.equal(new_logits, legacy_logits)
+
+
+def test_transfer_forward_matches_legacy_transfer(legacy_model_module, tmp_path):
+    graph, _ = _graph_and_globals()
+    globals_ = torch.empty((0, 0), dtype=torch.float32)
+    new_model = EdgeNetwork(graph, globals_, 8, 3, 2, 2, dropout=0.0).eval()
+    legacy_pretrained = legacy_model_module.Edge_Network(
+        graph, globals_, 8, 3, 2, 2, dropout=0.0
+    ).eval()
+    legacy_state = {
+        key.replace("classifier.", "classify."): value
+        for key, value in new_model.state_dict().items()
+    }
+    legacy_pretrained.load_state_dict(legacy_state)
+    checkpoint_path = tmp_path / "pretrained.pt"
+    torch.save({"model_state_dict": legacy_pretrained.state_dict()}, checkpoint_path)
+
+    config = {
+        "module": "models.GCN",
+        "class": "Edge_Network",
+        "args": {
+            "hid_size": 8,
+            "out_size": 3,
+            "n_layers": 2,
+            "n_proc_steps": 2,
+            "dropout": 0.0,
+        },
+    }
+    legacy_transfer = legacy_model_module.Transferred_Learning_Finetuning(
+        str(checkpoint_path), config, graph, globals_, 8, 1, 2, 2, dropout=0.0
+    ).eval()
+    transfer = FineTunedEdgeNetwork.from_pretrained(new_model, 1)
+    transfer.classifier.load_state_dict(legacy_transfer.classify.state_dict())
+
+    with torch.no_grad():
+        expected = legacy_transfer(graph.clone(), None)
+        actual = transfer(graph, None)
+    assert torch.equal(actual, expected)
+    assert torch.equal(
+        transfer.representation(graph, None), transfer.forward_features(graph, None)
+    )
+
+
+def test_legacy_prefix_loader_handles_checkpoint_prefixes():
+    graph, globals_ = _graph_and_globals()
+    source = _new_model(graph, globals_)
+    target = _new_model(graph, globals_)
+    prefixed = {
+        "module._orig_mod." + key.replace("classifier.", "classify."): value
+        for key, value in source.state_dict().items()
+    }
+    load_legacy_edge_network_state_dict(target, {"model_state_dict": prefixed})
+    for actual, expected in zip(target.parameters(), source.parameters()):
+        assert torch.equal(actual, expected)
+    transfer = FineTunedEdgeNetwork.from_pretrained(
+        source, 1, state_dict={"model_state_dict": prefixed}
+    )
+    assert isinstance(transfer.backbone.classifier, torch.nn.Identity)