feat: add Hydra configuration and CLI workflows

configs/checkpoint/default.yaml

@@ -0,0 +1,6 @@
+directory: ${environment.output_root}/checkpoints
+resume: null
+pretrained: null
+save_every_epochs: 1
+save_best: false
+save_last: true

configs/config.yaml

@@ -0,0 +1,18 @@
+defaults:
+  - data: local_test
+  - model: root_gnn/edge_network
+  - task: pretraining_multiclass
+  - trainer: default
+  - checkpoint: default
+  - inference: default
+  - environment: local
+  - _self_
+experiment:
+  name: pretraining_multiclass
+logging:
+  level: INFO
+hydra:
+  run:
+    dir: .
+  job:
+    chdir: false

configs/data/delphes.yaml

@@ -0,0 +1,4 @@
+defaults:
+  - /data/local_test
+  - _self_
+tree_name: Events

configs/data/local_test.yaml

@@ -0,0 +1,12 @@
+files: []
+tree_name: Events
+batch_size: 32
+num_workers: 0
+shuffle: true
+seed: 42
+cache:
+  path: null
+splits:
+  train_folds: [0, 1, 2]
+  validation_folds: [3]
+  test_folds: [4]

configs/environment/local.yaml

@@ -0,0 +1,3 @@
+device: cpu
+data_root: null
+output_root: outputs/${experiment.name}

configs/environment/perlmutter.yaml

@@ -0,0 +1,3 @@
+device: cuda
+data_root: null
+output_root: outputs/${experiment.name}

configs/inference/default.yaml

@@ -0,0 +1,4 @@
+checkpoint: null
+output: ${environment.output_root}/predictions/predictions.npz
+format: npz
+split: test

configs/model/root_gnn/edge_network.yaml

@@ -0,0 +1,7 @@
+family: root_gnn
+name: edge_network
+hid_size: 128
+out_size: 12
+n_layers: 2
+n_proc_steps: 4
+dropout: 0.1

configs/model/root_gnn/fine_tuned_edge_network.yaml

@@ -0,0 +1,9 @@
+family: root_gnn
+name: fine_tuned_edge_network
+hid_size: 128
+out_size: 1
+n_layers: 2
+n_proc_steps: 4
+dropout: 0.1
+freeze_backbone: true
+pretrained_checkpoint: null

configs/task/binary_classification.yaml

@@ -0,0 +1,3 @@
+type: binary_classification
+threshold: 0.5
+absolute_weights: false

configs/task/pretraining_multiclass.yaml

@@ -0,0 +1,3 @@
+type: multiclass_classification
+num_classes: 12
+absolute_weights: false

configs/task/tth_cp_finetune.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - /task/binary_classification
+  - _self_

configs/trainer/debug.yaml

@@ -0,0 +1,4 @@
+defaults:
+  - /trainer/default
+  - _self_
+max_epochs: 1

configs/trainer/default.yaml

@@ -0,0 +1,15 @@
+max_epochs: 100
+device: ${environment.device}
+seed: 42
+optimizer:
+  name: adam
+  learning_rate: 0.001
+  weight_decay: 0.0
+scheduler:
+  name: null
+early_stopping:
+  enabled: false
+  monitor: loss
+  mode: min
+  patience: 10
+  min_delta: 1.0e-8

pyproject.toml

@@ -22,6 +22,9 @@ ml = [
     "torch>=2.0",
 ]
 
+[project.scripts]
+gnn4colliders = "gnn4colliders.cli:main"
+
 [dependency-groups]
 dev = [
     "pytest>=8",

src/gnn4colliders/cli/__init__.py

@@ -0,0 +1,71 @@
+"""Single, thin command-line entry point for GNN4Colliders."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+
+from hydra import compose, initialize_config_dir
+from omegaconf import OmegaConf
+
+from gnn4colliders.config.application import predict_or_evaluate, prepare, train
+
+_COMMANDS = {"prepare", "train", "evaluate", "predict"}
+
+
+def _help(command: str | None = None) -> None:
+    if command:
+        print(f"Usage: gnn4colliders {command} [key=value ...]")
+        print("Compose Hydra configuration with semantic overrides.")
+    else:
+        print("Usage: gnn4colliders <prepare|train|evaluate|predict> [key=value ...]")
+        print("Use gnn4colliders <command> --help for command help.")
+
+
+def _config(overrides: list[str]):
+    config_dir = Path(__file__).resolve().parents[3] / "configs"
+    with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
+        return compose(config_name="config", overrides=overrides)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = list(sys.argv[1:] if argv is None else argv)
+    if not args or args[0] in {"-h", "--help"}:
+        _help()
+        return 0
+    command = args.pop(0)
+    if command not in _COMMANDS:
+        print(
+            f"unknown command {command!r}; choose from {', '.join(sorted(_COMMANDS))}",
+            file=sys.stderr,
+        )
+        return 2
+    if "--help" in args or "-h" in args:
+        _help(command)
+        return 0
+    try:
+        config = _config(args)
+        logging.basicConfig(
+            level=getattr(logging, str(config.logging.level).upper(), logging.INFO)
+        )
+        if command == "prepare":
+            print(f"prepared cache: {prepare(config)}")
+        elif command == "train":
+            print(f"saved checkpoint: {train(config)}")
+        elif command == "evaluate":
+            result = predict_or_evaluate(config, evaluate=True)
+            print(OmegaConf.to_yaml(result.metrics))
+        else:
+            result = predict_or_evaluate(config)
+            print(
+                f"wrote predictions: {config.inference.output} "
+                f"({len(result.sample_ids)} events)"
+            )
+        return 0
+    except (ValueError, FileNotFoundError, KeyError, ImportError) as error:
+        print(f"gnn4colliders: {error}", file=sys.stderr)
+        return 2
+
+
+__all__ = ["main"]

src/gnn4colliders/config/__init__.py

@@ -0,0 +1,6 @@
+"""Semantic configuration factories used by the command-line applications."""
+
+from .factories import build_model, build_task, build_trainer
+from .validation import validate_config
+
+__all__ = ["build_model", "build_task", "build_trainer", "validate_config"]

src/gnn4colliders/config/application.py

@@ -0,0 +1,187 @@
+"""Application-level orchestration kept separate from the CLI parser."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from gnn4colliders.data import (
+    GraphDataLoader,
+    GraphDataset,
+    GraphSample,
+    GraphSampleCache,
+    RootEventDataset,
+    SplitDefinition,
+    select_split,
+)
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.inference import Predictor, write_npz
+from gnn4colliders.training import (
+    CheckpointManager,
+    restore_training_state,
+    seed_everything,
+)
+
+from .factories import build_model, build_task, build_trainer
+from .validation import validate_config
+
+logger = logging.getLogger(__name__)
+
+
+def plain(config: DictConfig | dict[str, Any]) -> dict[str, Any]:
+    return OmegaConf.to_container(config, resolve=True)  # type: ignore[return-value]
+
+
+def save_resolved(config: DictConfig, output_root: str | Path) -> Path:
+    target = Path(output_root) / "resolved_config.yaml"
+    target.parent.mkdir(parents=True, exist_ok=True)
+    OmegaConf.save(config, target, resolve=True)
+    return target
+
+
+def load_samples(config: DictConfig) -> GraphDataset:
+    data = config.data
+    cache_path = data.cache.path
+    if not cache_path:
+        raise ValueError("data.cache.path is required for train/evaluate/predict")
+    cache = GraphSampleCache(cache_path)
+    if not cache.exists():
+        raise FileNotFoundError(f"graph sample cache does not exist: {cache_path}")
+    return GraphDataset(cache.load())
+
+
+def loaders(config: DictConfig) -> dict[str, GraphDataLoader]:
+    samples = load_samples(config)
+    split = SplitDefinition(
+        train_folds=frozenset(config.data.splits.train_folds),
+        validation_folds=frozenset(config.data.splits.validation_folds),
+        test_folds=frozenset(config.data.splits.test_folds),
+    )
+    result = {}
+    for name in ("train", "validation", "test"):
+        selected = GraphDataset(select_split(samples.samples, split, name))
+        result[name] = GraphDataLoader(
+            selected,
+            int(config.data.batch_size),
+            shuffle=bool(config.data.shuffle) if name == "train" else False,
+            seed=int(config.data.seed),
+        )
+    return result
+
+
+def prepare(config: DictConfig) -> Path:
+    data = config.data
+    if not data.files:
+        raise ValueError("prepare requires data.files")
+    feature_branches = data.get("feature_branches")
+    object_types = data.get("object_types")
+    scales = data.get("scales")
+    if feature_branches is None or object_types is None or scales is None:
+        raise ValueError(
+            "prepare requires data.feature_branches, object_types, and scales"
+        )
+    source = RootEventDataset(
+        data.files,
+        tree_name=str(data.tree_name),
+        label=data.get("label", 1),
+        feature_branches=feature_branches,
+        tracking_info=data.get("tracking_info", []),
+        global_features=data.get("global_features", []),
+        fold_var=str(data.get("fold_var", "eventNumber")),
+        weight_var=data.get("weight_var"),
+    )
+    samples: list[GraphSample] = []
+    for event in source:
+        features, _ = build_node_features(
+            event.objects, feature_branches, object_types, scales
+        )
+        graph = build_dgl_graph(features)
+        globals_ = torch.as_tensor(event.global_features, dtype=torch.float32)
+        samples.append(
+            GraphSample(
+                graph,
+                torch.as_tensor(event.label),
+                globals_ if globals_.numel() else None,
+                event.event_metadata,
+            )
+        )
+    path = Path(data.cache.path)
+    GraphSampleCache(path).save(samples)
+    return path
+
+
+def _model_and_task(config: DictConfig, loader: GraphDataLoader):
+    first = next(iter(loader), None)
+    if first is None:
+        raise ValueError("selected data split is empty")
+    model_config = plain(config.model)
+    pretrained = config.checkpoint.pretrained
+    if pretrained:
+        model_config["pretrained_checkpoint"] = pretrained
+    model = build_model(
+        model_config, sample_graph=first.graph, sample_global=first.global_features
+    )
+    return model, build_task(plain(config.task))
+
+
+def train(config: DictConfig) -> Path:
+    validate_config(plain(config))
+    seed_everything(int(config.trainer.seed))
+    run_root = Path(config.environment.output_root)
+    save_resolved(config, run_root)
+    split_loaders = loaders(config)
+    model, task = _model_and_task(config, split_loaders["train"])
+    trainer = build_trainer(plain(config.trainer), model=model, task=task)
+    if config.checkpoint.resume:
+        payload = CheckpointManager.load(config.checkpoint.resume, map_location="cpu")
+        restore_training_state(
+            payload,
+            model=model,
+            trainer=trainer,
+            optimizer=trainer.optimizer,
+            scheduler=trainer.scheduler,
+            early_stopping=trainer.early_stopping,
+        )
+    validation_loader = (
+        split_loaders["validation"] if len(split_loaders["validation"]) else None
+    )
+    trainer.fit(
+        split_loaders["train"], validation_loader, epochs=int(config.trainer.max_epochs)
+    )
+    manager = CheckpointManager(config.checkpoint.directory)
+    return manager.save(
+        model=model,
+        trainer_state=trainer.state,
+        optimizer=trainer.optimizer,
+        scheduler=trainer.scheduler,
+        early_stopping=trainer.early_stopping,
+        task_config=plain(config.task),
+        model_config=plain(config.model),
+    )
+
+
+def predict_or_evaluate(config: DictConfig, *, evaluate: bool = False) -> Any:
+    validate_config(plain(config))
+    split_loaders = loaders(config)
+    model, task = _model_and_task(config, split_loaders[config.inference.split])
+    payload = CheckpointManager.load(config.inference.checkpoint, map_location="cpu")
+    from gnn4colliders.training import load_model_weights
+
+    load_model_weights(model, payload)
+    predictor = Predictor(model, task, device=str(config.trainer.device))
+    result = (
+        predictor.evaluate(split_loaders[config.inference.split])
+        if evaluate
+        else predictor.predict(split_loaders[config.inference.split])
+    )
+    if not evaluate:
+        output = Path(config.inference.output)
+        if str(config.inference.format).lower() != "npz":
+            raise ValueError("only npz inference output is currently supported")
+        write_npz(result, output)
+    return result

src/gnn4colliders/config/factories.py

@@ -0,0 +1,139 @@
+"""Small allow-listed factories for configured application components."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
+from gnn4colliders.training import (
+    EarlyStopping,
+    Trainer,
+    build_optimizer,
+    build_scheduler,
+)
+
+
+def _get(config: Mapping[str, Any], name: str, default: Any = None) -> Any:
+    value = config.get(name, default)
+    return value
+
+
+def build_model(
+    config: Mapping[str, Any],
+    *,
+    sample_graph: Any,
+    sample_global: torch.Tensor | None = None,
+) -> torch.nn.Module:
+    """Build one of the supported models from semantic configuration."""
+    family = str(_get(config, "family", "root_gnn")).lower()
+    name = str(_get(config, "name", "")).lower()
+    if not name:
+        name = (
+            "fine_tuned_edge_network"
+            if "fine" in str(_get(config, "class", "")).lower()
+            else "edge_network"
+        )
+    if family != "root_gnn":
+        raise ValueError(f"unsupported model family {family!r}")
+    kwargs = {
+        "hid_size": int(_get(config, "hid_size", _get(config, "hidden_dim", 128))),
+        "out_size": int(_get(config, "out_size", 1)),
+        "n_layers": int(_get(config, "n_layers", 2)),
+        "n_proc_steps": int(
+            _get(config, "n_proc_steps", _get(config, "processing_steps", 4))
+        ),
+        "dropout": float(_get(config, "dropout", 0.0)),
+    }
+    if name in {"edge_network", "edge"}:
+        return EdgeNetwork(sample_graph, sample_global, **kwargs)
+    if name in {"fine_tuned_edge_network", "finetuned_edge_network", "fine_tuned"}:
+        checkpoint = _get(config, "pretrained_checkpoint")
+        if not checkpoint:
+            raise ValueError(
+                "fine-tuned model requires model.pretrained_checkpoint or "
+                "checkpoint.pretrained"
+            )
+        from gnn4colliders.training import CheckpointManager, load_model_weights
+
+        payload = CheckpointManager.load(checkpoint, map_location="cpu")
+        base_config = dict(payload.get("model_config") or {})
+        base_config.update(
+            {key: value for key, value in kwargs.items() if key != "out_size"}
+        )
+        backbone = EdgeNetwork(
+            sample_graph,
+            sample_global,
+            hid_size=int(base_config.get("hid_size", kwargs["hid_size"])),
+            out_size=int(base_config.get("out_size", 1)),
+            n_layers=int(base_config.get("n_layers", kwargs["n_layers"])),
+            n_proc_steps=int(base_config.get("n_proc_steps", kwargs["n_proc_steps"])),
+            dropout=float(base_config.get("dropout", kwargs["dropout"])),
+        )
+        load_model_weights(backbone, payload)
+        return FineTunedEdgeNetwork(
+            backbone,
+            kwargs["out_size"],
+            freeze_backbone=bool(_get(config, "freeze_backbone", False)),
+        )
+    raise ValueError(f"unsupported root_gnn model {name!r}")
+
+
+def build_task(config: Mapping[str, Any]) -> Any:
+    task_type = str(_get(config, "type", "multiclass_classification")).lower()
+    kwargs = {
+        "absolute_weights": bool(
+            _get(
+                config, "absolute_weights", _get(config, "use_absolute_weights", False)
+            )
+        )
+    }
+    if task_type in {"binary", "binary_classification"}:
+        return BinaryClassificationTask(
+            threshold=float(_get(config, "threshold", 0.5)), **kwargs
+        )
+    if task_type in {"multiclass", "multiclass_classification"}:
+        return MulticlassClassificationTask(**kwargs)
+    raise ValueError(f"unsupported task type {task_type!r}")
+
+
+def build_trainer(
+    config: Mapping[str, Any], *, model: torch.nn.Module, task: Any
+) -> Trainer:
+    optimizer_config = _get(config, "optimizer", {}) or {}
+    optimizer = build_optimizer(model, **dict(optimizer_config))
+    scheduler_config = _get(config, "scheduler", {}) or {}
+    scheduler = None
+    scheduler_step = "epoch"
+    if _get(scheduler_config, "name") not in (None, "", "none"):
+        scheduler_step = str(_get(scheduler_config, "step", "epoch"))
+        scheduler_kwargs = {
+            key: value
+            for key, value in dict(scheduler_config).items()
+            if key not in {"name", "step"}
+        }
+        scheduler = build_scheduler(
+            optimizer, name=str(_get(scheduler_config, "name")), **scheduler_kwargs
+        )
+    early_config = _get(config, "early_stopping", {}) or {}
+    early = None
+    if bool(_get(early_config, "enabled", False)):
+        early = EarlyStopping(
+            **{
+                key: value
+                for key, value in dict(early_config).items()
+                if key != "enabled"
+            }
+        )
+    return Trainer(
+        model,
+        task,
+        optimizer,
+        scheduler,
+        device=str(_get(config, "device", "cpu")),
+        early_stopping=early,
+        scheduler_step=scheduler_step,
+    )

src/gnn4colliders/config/validation.py

@@ -0,0 +1,41 @@
+"""Early validation for cross-group configuration relationships."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+
+def validate_config(config: Mapping[str, object]) -> None:
+    data = config.get("data", {}) or {}
+    trainer = config.get("trainer", {}) or {}
+    model = config.get("model", {}) or {}
+    task = config.get("task", {}) or {}
+    if int(data.get("batch_size", 1)) < 1:
+        raise ValueError("data.batch_size must be positive")
+    if int(trainer.get("max_epochs", 1)) < 1:
+        raise ValueError("trainer.max_epochs must be positive")
+    task_type = str(task.get("type", "")).lower()
+    out_size = int(model.get("out_size", 1))
+    if "binary" in task_type and out_size != 1:
+        raise ValueError("binary classification requires model.out_size=1")
+    if "multi" in task_type and out_size != int(task.get("num_classes", out_size)):
+        raise ValueError("multiclass model.out_size must equal task.num_classes")
+    if str(model.get("name", "")).lower() in {"fine_tuned_edge_network", "fine_tuned"}:
+        checkpoint = model.get("pretrained_checkpoint") or (
+            config.get("checkpoint", {}) or {}
+        ).get("pretrained")
+        if not checkpoint:
+            raise ValueError("fine-tuned models require checkpoint.pretrained")
+    checkpoint = config.get("checkpoint", {}) or {}
+    if checkpoint.get("resume") and checkpoint.get("pretrained"):
+        raise ValueError(
+            "checkpoint.resume and checkpoint.pretrained are distinct "
+            "workflows and cannot both be set"
+        )
+    splits = data.get("splits", {}) or {}
+    groups = [
+        set(splits.get(name, []))
+        for name in ("train_folds", "validation_folds", "test_folds")
+    ]
+    if any(groups[i] & groups[j] for i in range(3) for j in range(i + 1, 3)):
+        raise ValueError("train, validation, and test folds must be disjoint")