Hm/fix augmentation correctness

.github/ISSUE_TEMPLATE/bug_report.md

@@ -13,7 +13,7 @@ A clear and concise description of what goes wrong.
 Steps to reproduce the behaviour, ideally with the exact command:
 
 ```bash
-# e.g. SMAUGLAB_PARAMS_GPU_JSON=/abs/path/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU
+# e.g. SMAUGLAB_PARAMS_JSON=/abs/path/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU
 ```
 
 **Config JSON**

.github/workflows/lint.yml

@@ -68,3 +68,11 @@ jobs:
 
       - name: Run mypy
         run: mypy smauglab/
+
+      # The README augmentation matrix and smauglab/configs/all_augmentations.json are
+      # generated from the registry. Checked here rather than in a pre-commit hook: the
+      # generator has to import torch and kornia to populate the registry, which is far
+      # too slow to run on every commit. Contributors run the script when they add an
+      # augmentation; this is the backstop when they forget.
+      - name: Check generated registry artifacts are current
+        run: smauglab matrix --check

.gitignore

@@ -40,11 +40,29 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+# Repo-wide, so per-experiment config JSONs and personal scratch directories stay
+# out (see 7ff2088). The cost is that a genuinely new *package* file needs an
+# explicit un-ignore below -- `git add -f` alone is not enough, because anything
+# still ignored is invisible to `git status` and gets silently left behind.
 *.json
 *.yaml
 *.nii
 *.nii.gz
 
+# Generated from the registry and checked by CI (`smauglab matrix --check`), and
+# shipped as package data by pyproject's `include`.
+!smauglab/configs/all_augmentations.json
+
+# Test fixtures. unit_tests/test_cli.py points LEGACY at one of these.
+!unit_tests/fixtures/**
+
+# The pre-registry config migrator, which config.py's MIGRATE_HINT and the README
+# both point users at. Listed by extension rather than `migration/**`, which would
+# re-include migration/__pycache__ -- an un-ignore beats the earlier __pycache__/ rule.
+!migration/*.py
+!migration/*.json
+!migration/*.md
+
 # PyInstaller
 #  Usually these files are written by a python script from a template
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.

CONTRIBUTING.md

@@ -32,8 +32,18 @@ pytest -m "not slow"         # skip the wheel-building packaging tests
 pre-commit run --all-files   # everything CI's lint job runs
 ruff check .                 # lint only
 ruff format .                # format in place
+mypy smauglab/               # CI's typecheck job
+smauglab matrix --check      # are the generated README matrix and template current?
 ```
 
+`smauglab matrix --check` is a CI step rather than a pre-commit hook: populating
+the registry imports torch and kornia, which is far too slow to pay on every
+commit. Run `smauglab matrix --write` when you add an augmentation.
+
+Some tests need the domain-transfer LUT bank, which is built offline and not
+shipped. They skip without it; export `SMAUGLAB_DOMAIN_BANK=/path/to/bank.npz`
+to run them.
+
 ## The test suite
 
 `unit_tests/` runs entirely on CPU with 24×24×24 volumes and needs no image
@@ -43,7 +53,13 @@ data on disk, so it is fast enough to gate every pull request.
 | --- | --- |
 | `helpers.py` | `SmaugLabTestCase` base class (RNG seeding, test volumes) and config lookup |
 | `test_imports.py` | Every module under `smauglab/` imports cleanly |
+| `test_registry.py` | Registry mechanics, against synthetic classes |
+| `test_registered_augmentations.py` | The real registry, and that the generated matrix and template are current |
+| `test_variant_leaves.py` | Each variant leaf fixes its variant and hides it from the config surface |
+| `test_builder.py` | Pipeline order, bucketing, runtime context, and strict validation |
+| `test_migration.py` | Migrated configs reproduce the pre-registry builder exactly |
 | `test_configs.py` | Every shipped config parses, builds a pipeline, runs a forward pass, and is reproducible under a fixed seed |
+| `test_trainers.py` | The nnU-Net trainer builds what each config asks for, and puts deep supervision in the right place |
 | `test_transforms_gpu.py` | Each GPU transform in isolation |
 | `test_packaging.py` | Builds the real wheel and checks its contents |
 
@@ -60,16 +76,90 @@ does not hide the rest and the failure names the offending item — look for
 `SmaugLabTestCase` to get seeded RNGs and the shared `tiny_volume()` /
 `tiny_seg()` helpers.
 
-Transforms in `test_transforms_gpu.py` are discovered by introspection, so a
-new transform class is covered as soon as it lands — as long as it can be built
-with default arguments. If yours needs configuration, cover it by adding a
-config JSON under `smauglab/configs/`, which `test_configs.py` picks up
-automatically.
+Transforms in `test_transforms_gpu.py` come from the registry, so a new
+augmentation is covered the moment it is registered — no list to update, and no
+denylist of helper classes. Transforms with a required constructor argument are
+covered through `test_configs.py` instead, which picks up any config JSON added
+under `smauglab/configs/`.
+
+`test_migration.py` replays `fixtures/legacy_effective_kwargs.json`: a record of
+every constructor call the pre-registry builder made, captured while it still
+existed. It is what proves the config rename changed no behaviour, so treat it as
+append-only — regenerating it from current code would make it prove nothing.
 
 Note that these are smoke and contract tests: they check that transforms run,
 preserve shape, stay finite, and do not corrupt the segmentation labels. They
 do not verify that an augmentation is *visually* or *statistically* correct.
 
+## The nnU-Net trainer
+
+There is one trainer class, `nnUNetTrainerDAExtGPU`. It reads a config and builds
+whatever the config names: the `CPU` section runs in the dataloader worker, the `GPU`
+section runs on the batch in `train_step`. Three classes used to encode that split in
+their names; the config already carried it.
+
+The name is load-bearing and must not change: nnU-Net writes the trainer class name
+into every checkpoint and resolves the class from it at inference, so renaming it
+would make every previously trained model unloadable.
+
+Two behaviours are decided from the config rather than hardcoded:
+
+* **Deep supervision** is downsampled in `train_step` when there are GPU
+  augmentations (the mask is deformed there) and in the dataloader when there are
+  not. Getting this backwards trains against targets that no longer match the image.
+* **Pipeline arrangement** comes from `pipeline.mode` -- `sequential`,
+  `random_order` or `random_order_ta`.
+
+## Adding an augmentation
+
+Config keys are class names, exactly, and parameters are constructor arguments,
+exactly. Both are checked against the registry, so an augmentation is reachable
+from a config only once it is registered.
+
+1. **Write the class.** No `**kwargs` — parameter validation reads the signature,
+   so anything hidden behind it is invisible to a config and to `smauglab show`.
+   Declare `p` and `p_batch` explicitly on GPU transforms. Give every parameter a
+   default that is the value you actually want, not a `None` sentinel you resolve
+   in the body: the signature is what the generated template advertises.
+
+2. **Register it.**
+
+   ```python
+   @register(
+       aug_id=AugId.SCHARR,       # add a member if the concept is new
+       backend=Backend.GPU,
+       group=AugType.TA,          # GEO / GE / TA, used for random-order bucketing
+       order=90,                  # pipeline position; 10-spaced, unique per backend
+   )
+   class RandomScharrGPU(_RandomConvBaseGPU): ...
+   ```
+
+   Third-party classes cannot be decorated; register those in
+   `smauglab/transforms/cpu/external.py` instead.
+
+   Less common fields: `forwards_to` when the constructor genuinely passes kwargs
+   on to another class, `context_params` for values the trainer supplies at
+   runtime, `param_adapters` when a value must be wrapped before use, and
+   `external_asset` when it needs a file the wheel does not ship.
+
+3. **Regenerate and commit the artefacts.**
+
+   ```bash
+   smauglab matrix --write
+   ```
+
+4. **Check it.**
+
+   ```bash
+   smauglab show YourTransform
+   pytest unit_tests/test_registry.py unit_tests/test_registered_augmentations.py
+   ```
+
+If a variant differs only by one fixed argument — a kernel, a function, an
+inversion flag — give it its own thin subclass rather than exposing the argument.
+One class per config key is what keeps a config from expressing the same
+augmentation two different ways.
+
 ## Style
 
 Ruff handles both linting and formatting; the configuration lives in

README.md

@@ -81,10 +81,14 @@ Then, when you run nnUNet training as usual, specifying the SmaugLab trainer, fo
 nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans
 ```
 
-You can also specify your data augmentation parameters by providing a JSON file using the environment variable `SMAUGLAB_PARAMS_GPU_JSON`:
+There is one trainer. Which augmentations run is decided by the config, not by the
+trainer: a config's `CPU` section runs in the dataloader worker and its `GPU` section
+runs on the batch, so the same trainer covers CPU-only, GPU-only and mixed setups.
+
+Point it at a config with `SMAUGLAB_PARAMS_JSON`:
 > **Note:** By default [smauglab/configs/transform_params_gpu.json](https://github.com/neuropoly/SmaugLab/blob/main/smauglab/configs/transform_params_gpu.json) is used if no file is specified.
 ```bash
-SMAUGLAB_PARAMS_GPU_JSON=/path/to/your/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans
+SMAUGLAB_PARAMS_JSON=/path/to/your/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans
 ```
 
 > ⚠️ **Warning** : To avoid any paths issues, please specify an absolute path to your JSON file.
@@ -168,3 +172,101 @@ If you use SmaugLab, please make sure to cite the following paper:
   year={2026}
 }
 ```
+
+## Working with augmentations
+
+The `smauglab` command answers what exists and whether a config is valid, reading
+the registry directly so it cannot go out of date:
+
+```bash
+smauglab list --backend gpu            # what a GPU config can name, in pipeline order
+smauglab list --group TA               # just the transfer augmentations
+smauglab show RandomScharrGPU          # one augmentation's parameters and defaults
+smauglab validate my_config.json       # strict check, every problem reported at once
+smauglab template --backend gpu        # a config naming everything, at defaults
+smauglab hash my_config.json           # content-addressed config identity
+```
+
+Bringing a pre-registry config forward is a one-time job, so it is not a subcommand:
+the migrator lives in `migration/` in this repository rather than in the wheel.
+
+Config keys are class names, exactly, and parameters are constructor arguments,
+exactly. Anything else is an error rather than a silent no-op:
+
+```
+my_config.json: 2 problem(s)
+  - GPU.unknown GPU augmentation 'ScharrTransform'.
+      Did you mean: RandomScharrGPU?
+  - GPU.RandomGaussianNoiseGPU: unknown parameter 'probability'.
+      'probability' -> p
+```
+
+## Augmentation behaviour changed in the registry release
+
+Several augmentations were silently producing something other than what they claimed.
+Fixing them changes what the pipeline emits, so **models trained before this release
+saw different augmentations and are not bit-reproducible against it**. Configs are
+unaffected -- no key, parameter or default changed. The hash discontinuity is already
+recorded in `migration/hash_migration.json`.
+
+What changed, and why each was wrong:
+
+| Augmentation | Was | Now |
+|---|---|---|
+| `RandomFlipTransformGPU` | Ignored the flip flags its generator sampled, so it flipped every configured axis, identically, on every call and for every batch element. Seeded runs gave byte-identical output. | Reads `params["flip"]`: each batch element flips an independently sampled subset. |
+| `RandomGaussianBlurGPU`, `RandomUnsharpMaskGPU` | The Gaussian kernel was sampled at `arange(k)` rather than a centred range, so its peak sat at index 0 and the "blur" also translated the image about a voxel — relative to a mask that was not translated. | Centred kernel; blurring an impulse leaves its centre of mass in place. |
+| Every GPU contrast transform with `in_seg` / `out_seg` | Reduced the mask's class axis with `argmax(...) > 0`. For an ordinary single-channel mask that is always false, so `in_seg` applied the transform *nowhere* and `out_seg` applied it *everywhere*; for a one-hot mask it dropped the first foreground class. | `amax(...) > 0`, matching `collapse_onehot_to_index`. Both knobs now do what they say. |
+| `RandomAcqTransformGPU` (and any `one_dim` generator) | Drew its "random" axis in `make_samplers`, which kornia calls once and caches — the same axis was degraded for the entire training run. `CropGenerator3D` also drew *separate* axes for the crop and for its position. | Drawn per call and per batch element, with one axis shared by crop and position. |
+| `ScharrConvTransform` (CPU, 2-D only) | The x-kernel's middle row was `[-10, 0, -10]`, so it summed to −20 and was not a gradient operator. The 3-D CPU and both GPU kernels were correct. | Middle row is `[-10, 0, 10]`. |
+| `RandomLog1pGPU`, `RandomSqrtGPU`, `RandomSinGPU`, `RandomExpGPU`, `RandomSigmoidGPU` | Normalised with a batch-wide `x.min()` / `x.max()`, so a volume's augmentation depended on which other volumes shared its batch. | Per-sample min/max, as every other transform in the file. |
+| `RandomHistogramEqualizationGPU` | Wrote through a view of the input, so its non-finite guard skipped over values already in the batch. | Works on a clone; the guard is effective. |
+| `RandomChooseXTransformsGPU` | Wrote into the caller's batch in place. Transforms with a kornia parameter generator raised `params must contain 'scale'` inside a bucket. | Clones; samples parameters for generator-based transforms. |
+| Transforms drawing blur sigmas / kernel sizes | Used Python's `random`, which `torch.manual_seed` does not reach and which diverges across DDP ranks. | `smauglab.transforms.rng.shared_choice`, driven by torch's RNG. |
+
+The regression tests are in `unit_tests/test_region_mode.py` and
+`unit_tests/test_transform_randomness.py`; each one fails against the previous
+implementation.
+
+## 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/README.md

@@ -0,0 +1,29 @@
+# docs
+
+## `hash_migration.json`
+
+`segtransferaug/run_trainings.py` names each experiment directory
+
+    Dataset{id}-{trainer}-aug-{transform_hash}-c-{config_hash}
+
+where `config_hash` is a sha256 over the canonical JSON of the config and
+`transform_hash` is a sha256 over the text of `smauglab/transforms/gpu/transforms.py`.
+
+The registry migration moved both: every config was rewritten to class-name keys, and
+the GPU builder that file contained was replaced by a registry-driven one. Runs made
+before the migration keep their old directory names, so anything that looks a run up
+by hash needs this table. The canonicalisation itself is unchanged and is asserted
+byte-for-byte against the original implementation, so a config that did not change
+would still hash the same.
+
+Two old hashes map to configs that are now byte-identical:
+
+* `0b639cf1` covered `transform_params_gpu_default01-23.json` and
+  `transform_params_one-sequence-to-segment-them-all.json`, which had the same content
+  all along under different names.
+* `transform_params_gpu_default01-23_ICGT_plus.json` had its own hash only because of
+  the dead `ImageContrastGPUTransform` block; with that gone it is a duplicate of
+  `default01-23`.
+
+Known consumer, updated to accept both: `segtransferaug/refinement/config.py`, whose
+`DOMAIN_TRANSFER_HASHES` selects the fusion source pool.