feat: add explicit legacy compatibility adapters

docs/compatibility.md

@@ -0,0 +1,22 @@
+# Compatibility boundary
+
+The package's canonical APIs use named `EventMetadata` fields and the version
+1 checkpoint schema. Compatibility is explicit and one-way: old artifacts are
+adapted into the new representation and are never rewritten implicitly.
+
+| Artifact or behavior | Supported | Boundary | Notes |
+| --- | :---: | --- | --- |
+| New checkpoint (`schema_version: 1`) | yes | `CheckpointManager` | Full model/lifecycle resume when state is present |
+| Legacy `model_epoch_N.pt` checkpoint | yes | `gnn4colliders.compat.load_legacy_checkpoint` | Model state and active early-stop fields are adapted |
+| Legacy DDP/compiled prefixes | yes | `normalize_legacy_state_dict_keys` | Supports `module.` and `_orig_mod.` |
+| Legacy ROOT-GNN classifier name | yes | `map_legacy_edge_network_state_dict` | Maps `classify` to `classifier` |
+| Legacy optimizer state | partial | legacy checkpoint adapter | Loaded when present; scheduler state is not available in the historical format |
+| Legacy scheduler state | no | — | Historical checkpoints do not carry a supported scheduler state |
+| Positional tracking rows | yes, at ingestion boundary | `EventMetadata.from_legacy_tracking` | Exactly `tracking[0] = fold`, `tracking[1] = weight`; shorter rows fail |
+| Generic/unknown tracking layouts | no | — | The package does not guess historical column meanings |
+| Modern NPZ output | yes | `inference.write_npz` | Named fields: `sample_id`, `logits`, `scores`, `predictions`, and available metadata |
+| Historical `tracking_info` NPZ output | no | — | No active consumer remains; positional output is intentionally unsupported |
+| Legacy YAML `module`/`class`/`args` | compatibility only | configuration boundary | Accepted only where the semantic factory can safely interpret it; new configs use semantic model names |
+
+The frozen `legacy/` tree remains available to parity tests and historical
+investigation. Production modules do not import executable code from it.

src/gnn4colliders/compat/__init__.py

@@ -0,0 +1,17 @@
+"""Explicit adapters for artifacts from the frozen legacy implementation."""
+
+from .checkpoint import (
+    LEGACY_CHECKPOINT_SCHEMA_VERSION,
+    load_legacy_checkpoint,
+    map_legacy_edge_network_state_dict,
+    normalize_legacy_state_dict_keys,
+)
+from .metadata import event_metadata_from_legacy_tracking
+
+__all__ = [
+    "LEGACY_CHECKPOINT_SCHEMA_VERSION",
+    "event_metadata_from_legacy_tracking",
+    "load_legacy_checkpoint",
+    "map_legacy_edge_network_state_dict",
+    "normalize_legacy_state_dict_keys",
+]

src/gnn4colliders/compat/checkpoint.py

@@ -0,0 +1,72 @@
+"""Checkpoint adapters for the supported historical ROOT-GNN artifacts."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+
+LEGACY_CHECKPOINT_SCHEMA_VERSION = 0
+
+
+def normalize_legacy_state_dict_keys(state: Mapping[str, Any]) -> dict[str, Any]:
+    """Remove the supported DDP and compilation wrappers deterministically."""
+    return {
+        key.removeprefix("module.").removeprefix("_orig_mod."): value
+        for key, value in state.items()
+    }
+
+
+def map_legacy_edge_network_state_dict(
+    state: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Map historical ``classify`` keys to the modern ``classifier`` name."""
+    normalized = normalize_legacy_state_dict_keys(state)
+    return {
+        key.replace(".classify.", ".classifier.").replace(
+            "classify.", "classifier."
+        ): value
+        for key, value in normalized.items()
+    }
+
+
+def load_legacy_checkpoint(
+    source: str | Path | Mapping[str, Any], *, map_location: Any = "cpu"
+) -> dict[str, Any]:
+    """Adapt an active ``model_epoch_N.pt`` payload without rewriting it."""
+    payload = (
+        torch.load(source, map_location=map_location, weights_only=False)
+        if not isinstance(source, Mapping)
+        else dict(source)
+    )
+    if not isinstance(payload, Mapping) or "model_state_dict" not in payload:
+        raise ValueError("not a supported legacy ROOT-GNN checkpoint")
+    early = payload.get("early_stop")
+    if isinstance(early, Mapping):
+        early = {
+            "patience": early.get("patience", 15),
+            "min_delta": early.get("threshold", 1e-8),
+            "mode": early.get("mode", "min"),
+            "best": early.get("current_best", float("inf")),
+            "num_bad_epochs": early.get("count", 0),
+            "should_stop": early.get("should_stop", False),
+        }
+    epoch = int(payload.get("epoch", -1))
+    return {
+        "schema_version": LEGACY_CHECKPOINT_SCHEMA_VERSION,
+        "legacy": True,
+        "epoch": epoch,
+        "global_step": 0,
+        "model_state_dict": normalize_legacy_state_dict_keys(
+            payload["model_state_dict"]
+        ),
+        "optimizer_state_dict": payload.get("optimizer_state_dict"),
+        "scheduler_state_dict": None,
+        "early_stopping_state": early,
+        "trainer_state": {"epoch": epoch, "global_step": 0},
+        "model_config": None,
+        "task_config": None,
+        "metadata": {"legacy_format": True},
+    }

src/gnn4colliders/compat/metadata.py

@@ -0,0 +1,32 @@
+"""Adapters for the historical positional event metadata layout."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from gnn4colliders.data.metadata import EventMetadata
+
+
+def event_metadata_from_legacy_tracking(
+    tracking_row: Sequence[Any],
+    *,
+    sample_id: str,
+    extra: Mapping[str, Any] | None = None,
+) -> EventMetadata:
+    """Convert the only supported legacy layout into named metadata.
+
+    The historical contract is strictly ``tracking[0] = fold`` and
+    ``tracking[1] = weight``.  Additional historical columns are not
+    interpreted or propagated.
+    """
+    if len(tracking_row) < 2:
+        raise ValueError(
+            "legacy tracking must contain at least two values: fold and weight"
+        )
+    return EventMetadata(
+        fold=int(tracking_row[0]),
+        weight=float(tracking_row[1]),
+        sample_id=sample_id,
+        extra=dict(extra or {}),
+    )

src/gnn4colliders/config/factories.py

@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+from gnn4colliders.distributed import DistributedContext
 from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
 from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
 from gnn4colliders.training import (
@@ -101,7 +102,11 @@ def build_task(config: Mapping[str, Any]) -> Any:
 
 
 def build_trainer(
-    config: Mapping[str, Any], *, model: torch.nn.Module, task: Any
+    config: Mapping[str, Any],
+    *,
+    model: torch.nn.Module,
+    task: Any,
+    distributed_context: DistributedContext | None = None,
 ) -> Trainer:
     optimizer_config = _get(config, "optimizer", {}) or {}
     optimizer = build_optimizer(model, **dict(optimizer_config))
@@ -136,4 +141,5 @@ def build_trainer(
         device=str(_get(config, "device", "cpu")),
         early_stopping=early,
         scheduler_step=scheduler_step,
+        distributed_context=distributed_context,
     )

tests/unit/compat/test_adapters.py

@@ -0,0 +1,29 @@
+import pytest
+import torch
+
+from gnn4colliders.compat import (
+    event_metadata_from_legacy_tracking,
+    map_legacy_edge_network_state_dict,
+)
+from gnn4colliders.data import EventMetadata
+
+
+def test_legacy_tracking_is_converted_to_named_metadata():
+    metadata = event_metadata_from_legacy_tracking(
+        [3.0, -2.5, 99.0], sample_id="event:7"
+    )
+    assert metadata == EventMetadata(fold=3, weight=-2.5, sample_id="event:7")
+
+
+def test_legacy_tracking_rejects_ambiguous_short_rows():
+    with pytest.raises(ValueError, match="fold and weight"):
+        EventMetadata.from_legacy_tracking([3.0], sample_id="event:7")
+
+
+def test_legacy_state_dict_mapping_is_canonical_and_deterministic():
+    state = {
+        "module._orig_mod.encoder.weight": torch.ones(2, 2),
+        "module._orig_mod.classify.bias": torch.zeros(1),
+    }
+    mapped = map_legacy_edge_network_state_dict(state)
+    assert set(mapped) == {"encoder.weight", "classifier.bias"}