feat: one trainer, a CLI, and a coverage matrix that cannot go stale

.gitignore

@@ -46,6 +46,9 @@ MANIFEST
 # migration table, the test fixtures and the migrator itself.
 !unit_tests/fixtures/**/*.json
 !migration/*.json
+# Generated by `smauglab template --write` and checked by CI; it is source,
+# not per-experiment output.
+!smauglab/configs/all_augmentations.json
 *.yaml
 *.nii
 *.nii.gz

README.md

@@ -168,3 +168,47 @@ If you use SmaugLab, please make sure to cite the following paper:
   year={2026}
 }
 ```
+
+## Available augmentations
+
+Which augmentations exist, and which backends implement each one. An empty cell
+means no implementation on that backend yet. Regenerate with `smauglab matrix --write`.
+
+<!-- BEGIN AUG MATRIX (generated by `smauglab matrix --write`; do not edit) -->
+| Augmentation | Group | GPU | CPU | MONAI |
+| --- | --- | --- | --- | --- |
+| flip | GEO | `RandomFlipTransformGPU` | — | — |
+| affine | GEO | `RandomAffineGPU` | — | — |
+| crop | GEO | `RandomCropTransformGPU` | — | — |
+| spatial | GEO | — | `SpatialTransform` | — |
+| gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` | — |
+| gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` | — |
+| brightness | GE | `RandomBrightnessGPU` | `MultiplicativeBrightnessTransform` | — |
+| contrast | GE | `RandomContrastGPU` | `ContrastTransform` | — |
+| gamma | GE | `RandomGammaGPU` | `GammaTransform` | — |
+| inv_gamma | GE | `RandomInvGammaGPU` | `InvertedGammaTransform` | — |
+| clamp | GE | `RandomClampGPU` | — | — |
+| low_res | GE | `RandomLowResTransformGPU` | `SimulateLowResolutionTransform` | — |
+| acq | GE | `RandomAcqTransformGPU` | — | — |
+| zscore | GE | `ZscoreNormalizationGPU` | `ZscoreNormalization` | — |
+| mirror | GEO | — | `MirrorTransform` | — |
+| scharr | TA | `RandomScharrGPU` | `ScharrConvTransform` | — |
+| laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` | — |
+| unsharp_mask | TA | `RandomUnsharpMaskGPU` | — | — |
+| rand_conv | TA | `RandomRandConvGPU` | — | — |
+| bias_field | TA | `RandomBiasFieldGPU` | — | — |
+| inverse | TA | `RandomInverseGPU` | — | — |
+| histogram_equal | TA | `RandomHistogramEqualizationGPU` | `HistogramEqualTransform` | — |
+| redistribute_seg | TA | `RandomRedistributeSegGPU` | `RedistributeTransform` | — |
+| palette | TA | `RandomPaletteGPU` | — | — |
+| domain_transfer | TA | `RandomDomainTransferGPU` | — | — |
+| synthseg | TA | `RandomSynthSegGPU` | — | — |
+| artifact | TA | — | `ArtifactTransform` | — |
+| spatial_custom | GEO | — | `SpatialCustomTransform` | — |
+| shape | GE | — | `ShapeTransform` | — |
+| func_log1p | TA | `RandomLog1pGPU` | `Log1pTransform` | — |
+| func_sqrt | TA | `RandomSqrtGPU` | `SqrtTransform` | — |
+| func_sin | TA | `RandomSinGPU` | `SinTransform` | — |
+| func_exp | TA | `RandomExpGPU` | `ExpTransform` | — |
+| func_sigmoid | TA | `RandomSigmoidGPU` | `SigmoidTransform` | — |
+<!-- END AUG MATRIX -->

migration/migrate.py

@@ -389,6 +389,26 @@ def migrate(payload: dict, source: str = "<config>") -> tuple[dict, list[str]]:
             target["SpatialTransform"] = _migrate_block("SpatialTransform", spatial, Backend.CPU, source, legacy_key="SpatialTransform")
             notes.append("moved nnUNetSpatialTransform -> CPU.SpatialTransform")
 
+        if backend is Backend.CPU and "SpatialTransform" not in migrated and "SpatialTransform" not in out.get("CPU", {}):
+            # nnU-Net's own SpatialTransform used to be appended by the trainer with
+            # these values hardcoded, so no config named it. The trainer builds its CPU
+            # pipeline from the config now, so it has to be in the file or it silently
+            # stops running. Values are verbatim from get_training_transforms;
+            # patch_size and rotation stay out because nnU-Net supplies them at runtime
+            # (they are context_params on the registry entry).
+            migrated["SpatialTransform"] = {
+                "patch_center_dist_from_border": 0,
+                "random_crop": False,
+                "p_elastic_deform": 0,
+                "p_rotation": 0,
+                "p_scaling": 0,
+                "scaling": [0.7, 1.4],
+                "p_synchronize_scaling_across_axes": 1,
+                "bg_style_seg_sampling": False,
+                "mode_seg": "nearest",
+            }
+            notes.append("added CPU.SpatialTransform, which the trainer used to hardcode")
+
         if backend is Backend.CPU:
             axes = section.get("mirror_axes")
             if axes:
@@ -398,6 +418,11 @@ def migrate(payload: dict, source: str = "<config>") -> tuple[dict, list[str]]:
         choose = section.get("RandomChooseXTransforms")
         if isinstance(choose, dict):
             out.setdefault("pipeline", {})["random_choose"] = choose
+            # Which pipeline ran used to be the *trainer class*: a config carrying this
+            # block was only ever used with the list trainer. That choice is
+            # pipeline.mode now, so it has to be written down.
+            out["pipeline"].setdefault("mode", "random_order")
+            notes.append("moved RandomChooseXTransforms -> pipeline.random_choose, mode=random_order")
             notes.append("moved RandomChooseXTransforms -> pipeline.random_choose")
 
         if migrated:

pyproject.toml

@@ -94,6 +94,7 @@ all = ["monai", "tqdm", "wandb"]
 dev = ["build", "coverage", "mypy", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"]
 
 [tool.poetry.scripts]
+smauglab = "smauglab.cli:main"
 smauglab_add_nnunettrainer = "smauglab.add_trainer:main"
 
 [build-system]

smauglab/cli.py

@@ -0,0 +1,279 @@
+"""The `smauglab` command line: what augmentations exist, and are my configs valid?
+
+Before the registry, both questions could only be answered by reading four
+hand-written dispatch ladders side by side. These subcommands read the registry, so
+they cannot go out of date:
+
+    smauglab list --backend gpu       what a GPU config can name, in pipeline order
+    smauglab matrix                   which backends implement each augmentation
+    smauglab show RandomScharrGPU     one augmentation's parameters and defaults
+    smauglab validate config.json     strict check, every problem at once
+    smauglab template --backend gpu   a config naming everything, at defaults
+    smauglab hash config.json         content-addressed config identity
+
+Bringing a pre-registry config forward is a one-time job and is not a subcommand: the
+migrator lives in `migration/` in the repository, not in the wheel.
+
+`hash` is the only one that does not need the registry; the rest import torch and
+kornia to populate it, which takes a few seconds on first use.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+REPO = Path(__file__).resolve().parent.parent
+README = REPO / "README.md"
+TEMPLATE_PATH = Path(__file__).resolve().parent / "configs" / "all_augmentations.json"
+BEGIN = "<!-- BEGIN AUG MATRIX (generated by `smauglab matrix --write`; do not edit) -->"
+END = "<!-- END AUG MATRIX -->"
+
+
+def _backend(name: str):
+    from smauglab.registry import Backend
+
+    return Backend[name.upper()]
+
+
+# --- generated artefacts ----------------------------------------------------------
+
+
+def matrix_block() -> str:
+    from smauglab import registry
+
+    return f"{BEGIN}\n{registry.render_matrix('md')}\n{END}"
+
+
+def readme_with_matrix(text: str) -> str:
+    """Replace the marked block, or append the section if it is not there yet."""
+    if BEGIN in text and END in text:
+        head, rest = text.split(BEGIN, 1)
+        _, tail = rest.split(END, 1)
+        return f"{head}{matrix_block()}{tail}"
+    section = (
+        "\n## Available augmentations\n\n"
+        "Which augmentations exist, and which backends implement each one. An empty cell\n"
+        "means no implementation on that backend yet. Regenerate with `smauglab matrix --write`.\n\n" + matrix_block() + "\n"
+    )
+    return text.rstrip("\n") + "\n" + section
+
+
+def template_json() -> str:
+    from smauglab import registry
+    from smauglab.registry import Backend
+
+    payload = {b.value: registry.render_template(b) for b in (Backend.GPU, Backend.CPU)}
+    return json.dumps(payload, indent=4) + "\n"
+
+
+def _sync(targets: dict[Path, str], *, write: bool, check: bool) -> int:
+    stale = [path for path, content in targets.items() if not path.is_file() or path.read_text() != content]
+    if check:
+        for path in stale:
+            print(f"out of date: {path.relative_to(REPO)}")
+        if stale:
+            print("Run `smauglab matrix --write` / `smauglab template --write` and commit the result.")
+        return 1 if stale else 0
+    if write:
+        for path in stale:
+            path.write_text(targets[path])
+            print(f"wrote {path.relative_to(REPO)}")
+        if not stale:
+            print("already up to date")
+    return 0
+
+
+# --- subcommands ------------------------------------------------------------------
+
+
+def cmd_list(args) -> int:
+    from smauglab import registry
+    from smauglab.registry import AugType
+
+    backend = _backend(args.backend) if args.backend else None
+    group = AugType[args.group.upper()] if args.group else None
+    entries = registry.entries(backend=backend, group=group)
+
+    if args.json:
+        print(
+            json.dumps(
+                [
+                    {
+                        "name": e.name,
+                        "backend": e.backend.value,
+                        "aug_id": e.aug_id.value,
+                        "group": e.group.value,
+                        "position": registry.pipeline_position(e),
+                        "summary": e.summary,
+                    }
+                    for e in entries
+                ],
+                indent=2,
+            )
+        )
+        return 0
+
+    if not entries:
+        print("no augmentations match")
+        return 0
+    width = max(len(e.name) for e in entries)
+    current = None
+    for entry in entries:
+        if entry.backend is not current:
+            current = entry.backend
+            print(f"\n{current.value} ({len([e for e in entries if e.backend is current])}), in pipeline order:")
+        print(f"  {registry.pipeline_position(entry):>4}  {entry.name:<{width}}  {entry.group.value:<4} {entry.summary}")
+    return 0
+
+
+def cmd_matrix(args) -> int:
+    from smauglab import registry
+
+    if args.write or args.check:
+        return _sync(
+            {README: readme_with_matrix(README.read_text()), TEMPLATE_PATH: template_json()},
+            write=args.write,
+            check=args.check,
+        )
+    if args.format == "json":
+        table = registry.matrix()
+        print(
+            json.dumps(
+                {aug_id.value: {b.value: (e.name if e else None) for b, e in row.items()} for aug_id, row in table.items()},
+                indent=2,
+            )
+        )
+        return 0
+    print(registry.render_matrix(args.format))
+    return 0
+
+
+def cmd_show(args) -> int:
+    from smauglab import registry
+
+    try:
+        entry = registry.get(args.name)
+    except registry.UnknownAugmentationError as exc:
+        print(str(exc), file=sys.stderr)
+        return 1
+
+    print(
+        f"{entry.name}  ({entry.backend.value}, group {entry.group.value}, position {registry.pipeline_position(entry)}, aug_id {entry.aug_id.value})"
+    )
+    if entry.summary:
+        print(f"  {entry.summary}")
+    if entry.forwards_to is not None:
+        print(f"  forwards extra parameters to {entry.forwards_to.__name__}")
+    if entry.external_asset:
+        print(f"  needs an external asset; set ${entry.external_asset}")
+    print(f"\n  module: {entry.cls.__module__}")
+
+    params = registry.accepted_params(entry)
+    required = registry.required_params(entry)
+    print(f"\n  parameters ({len(params)}):")
+    for name in sorted(params):
+        default = params[name].default
+        shown = "REQUIRED" if name in required else repr(default)
+        print(f"    {name:<34} {shown}")
+    if entry.context_params:
+        print(f"\n  supplied by the trainer, not the config: {', '.join(entry.context_params)}")
+    return 0
+
+
+def cmd_validate(args) -> int:
+    from smauglab.config import validate_file
+
+    failed = 0
+    for path in args.configs:
+        problems = validate_file(path)
+        if problems:
+            failed += 1
+            print(f"{path}: {len(problems)} problem(s)")
+            for problem in problems:
+                print(f"  - {problem}")
+        elif not args.quiet:
+            print(f"{path}: ok")
+    return 1 if failed else 0
+
+
+def cmd_template(args) -> int:
+    from smauglab import registry
+
+    if args.write or args.check:
+        return _sync({TEMPLATE_PATH: template_json()}, write=args.write, check=args.check)
+    section = registry.render_template(_backend(args.backend))
+    text = json.dumps({args.backend.upper(): section}, indent=4) + "\n"
+    if args.output:
+        Path(args.output).write_text(text)
+        print(f"wrote {args.output}")
+    else:
+        print(text, end="")
+    return 0
+
+
+def cmd_hash(args) -> int:
+    from smauglab.config import config_hash
+
+    for path in args.configs:
+        payload = json.loads(Path(path).read_text())
+        print(f"{config_hash(payload)[:8]}  {path}")
+    return 0
+
+
+# --- wiring -----------------------------------------------------------------------
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="smauglab", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p = sub.add_parser("list", help="registered augmentations, in pipeline order")
+    p.add_argument("--backend", choices=["gpu", "cpu", "monai"])
+    p.add_argument("--group", choices=["geo", "ge", "ta"])
+    p.add_argument("--json", action="store_true")
+    p.set_defaults(func=cmd_list)
+
+    p = sub.add_parser("matrix", help="which backends implement each augmentation")
+    p.add_argument("--format", choices=["md", "table", "json"], default="table")
+    p.add_argument("--write", action="store_true", help="update README.md and the template config")
+    p.add_argument("--check", action="store_true", help="exit 1 if either is out of date")
+    p.set_defaults(func=cmd_matrix)
+
+    p = sub.add_parser("show", help="one augmentation's parameters and defaults")
+    p.add_argument("name")
+    p.set_defaults(func=cmd_show)
+
+    p = sub.add_parser("validate", help="strict config check, reporting every problem")
+    p.add_argument("configs", nargs="+")
+    p.add_argument("-q", "--quiet", action="store_true", help="only report failures")
+    p.set_defaults(func=cmd_validate)
+
+    p = sub.add_parser("template", help="a config naming every augmentation at its defaults")
+    p.add_argument("--backend", choices=["gpu", "cpu"], default="gpu")
+    p.add_argument("-o", "--output")
+    p.add_argument("--write", action="store_true", help="update the shipped template config")
+    p.add_argument("--check", action="store_true", help="exit 1 if it is out of date")
+    p.set_defaults(func=cmd_template)
+
+    p = sub.add_parser("hash", help="content-addressed config identity")
+    p.add_argument("configs", nargs="+")
+    p.set_defaults(func=cmd_hash)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    if getattr(args, "output", None) and len(getattr(args, "configs", [])) > 1:
+        print("-o takes a single input config", file=sys.stderr)
+        return 2
+    result: Any = args.func(args)
+    return int(result or 0)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())