scripts/dev/smoke_end_to_end.py

"""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()