feat: noise relative to the sample's own spread with RandomGaussianNoiseGPU(relative=True)

smauglab/configs/all_augmentations.json

@@ -318,6 +318,7 @@
             "out_seg": 0.0,
             "p": 1.0,
             "p_batch": 1.0,
+            "relative": false,
             "same_on_batch": false,
             "std": 1.0
         },

smauglab/transforms/gpu/contrast.py

@@ -718,7 +718,15 @@ class RandomGaussianNoiseGPU(ImageOnlyTransform):
 
     Args:
         mean (float): Mean of the Gaussian noise. Default is 0.0.
-        std (float): Standard deviation of the Gaussian noise. Default is 0.1.
+        std (float): Upper bound of the noise standard deviation, which is drawn from
+            `U(0, std)` per sample (or once, with `same_on_batch`). Default is 1.0.
+        relative (bool): Read `std` as a fraction of each sample's own standard deviation
+            in that channel rather than in absolute intensity units. The absolute form
+            makes the same config a different augmentation on every intensity scale: a
+            z-scored MR patch (std ~0.9), a CT patch under nnU-Net's CT normalisation
+            (std ~0.17) and a `[0, 1]` SynthSeg synthesis (std ~0.18) see noise at the
+            same absolute std, so the last two are mostly noise. `mean` stays absolute.
+            Default False.
         same_on_batch (bool): Apply the same transformation across the batch. Default is False.
         p (float): Probability of applying the transform. Default is 1.0.
         keepdim (bool): Whether to keep the number of dimensions. Default is False.
@@ -731,6 +739,7 @@ def __init__(
         self,
         mean: float = 0.0,
         std: float = 1.0,
+        relative: bool = False,
         apply_to_channel: Sequence[int] = (0,),  # Apply to first channel by default
         same_on_batch: bool = False,
         in_seg: float = 0.0,
@@ -744,6 +753,7 @@ def __init__(
         self.apply_to_channel = apply_to_channel
         self.mean = mean
         self.std = std
+        self.relative = relative
         self.in_seg = in_seg
         self.out_seg = out_seg
         self.mix_in_out = mix_in_out
@@ -757,13 +767,16 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[
             if self.same_on_batch:
                 std = torch.rand(1, device=input.device, dtype=input.dtype) * self.std
                 noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype)
-                noise = noise * std + self.mean
             else:
                 std = torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * self.std
                 noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype)
-                # Broadcast the per-sample std instead of writing one row at a time:
-                # same draws, same arithmetic, one kernel instead of N.
-                noise = noise * std.view(-1, *([1] * (noise.dim() - 1))) + self.mean
+            if self.relative:
+                # Scaled after both draws, so the random stream is the absolute form's.
+                std = std * input[:, c].reshape(input.shape[0], -1).std(dim=1)
+            # Broadcast the per-sample std instead of writing one row at a time:
+            # same draws, same arithmetic, one kernel instead of N. A one-element std
+            # (`same_on_batch`, absolute) broadcasts the same way it did unreshaped.
+            noise = noise * std.view(-1, *([1] * (noise.dim() - 1))) + self.mean
 
             orig = input[:, c]
             x = orig + noise

unit_tests/test_gaussian_noise_relative.py

@@ -0,0 +1,71 @@
+"""`RandomGaussianNoiseGPU(relative=True)` scales the noise with the image.
+
+The std is drawn from `U(0, std)` in absolute intensity units, so the same config is a
+different augmentation on every intensity scale: a z-scored MR patch (std ~0.9), a CT
+patch under nnU-Net's CT normalisation (std ~0.17) and a `[0, 1]` SynthSeg synthesis
+(std ~0.18) all get noise of the same absolute std. `relative=True` reads `std` as a
+fraction of each sample's own std, which makes the transform scale-equivariant.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from smauglab.registry import Backend, accepted_params, get
+from smauglab.transforms.gpu.contrast import RandomGaussianNoiseGPU
+from unit_tests.helpers import SmaugLabTestCase
+
+SHAPE = (4, 1, 12, 12, 12)
+
+
+def noised(x: torch.Tensor, relative: bool, same_on_batch: bool, seed: int = 0) -> torch.Tensor:
+    transform = RandomGaussianNoiseGPU(std=0.3, relative=relative, same_on_batch=same_on_batch, p=1.0)
+    torch.manual_seed(seed)
+    return transform(x.clone())
+
+
+def image() -> torch.Tensor:
+    torch.manual_seed(123)
+    # Four samples on four different intensity scales.
+    return torch.randn(*SHAPE) * torch.tensor([0.17, 0.9, 2.5, 30.0]).view(-1, 1, 1, 1, 1)
+
+
+class TestRelativeGaussianNoise(SmaugLabTestCase):
+    def test_the_absolute_form_is_unchanged(self):
+        """Same draws in the same order, same arithmetic as before the option existed."""
+        x = image()
+        for same_on_batch in (False, True):
+            with self.subTest(same_on_batch=same_on_batch):
+                torch.manual_seed(0)
+                n = 1 if same_on_batch else SHAPE[0]
+                std = torch.rand(n) * 0.3
+                noise = torch.randn_like(x[:, 0])
+                scale = std if same_on_batch else std.view(-1, 1, 1, 1)
+                expected = x[:, 0] + (noise * scale + 0.0)
+                out = noised(x, relative=False, same_on_batch=same_on_batch)
+                self.assertTrue(torch.equal(out[:, 0], expected))
+
+    def test_relative_noise_is_scale_equivariant(self):
+        x = image()
+        for same_on_batch in (False, True):
+            with self.subTest(same_on_batch=same_on_batch):
+                out = noised(x, relative=True, same_on_batch=same_on_batch)
+                out_scaled = noised(7.0 * x, relative=True, same_on_batch=same_on_batch)
+                torch.testing.assert_close(out_scaled, 7.0 * out, rtol=1e-5, atol=1e-5)
+
+    def test_relative_is_the_absolute_draw_times_the_sample_std(self):
+        """No extra random draw: the noise field is the absolute one, rescaled per sample.
+
+        In float64, because the noise is read back as `out - x`, which in float32 loses
+        about ulp(x) -- 2e-6 at the std-30 sample, then multiplied by 30.
+        """
+        x = image().double()
+        sample_std = x.reshape(SHAPE[0], -1).std(dim=1).view(-1, 1, 1, 1, 1)
+        for same_on_batch in (False, True):
+            with self.subTest(same_on_batch=same_on_batch):
+                absolute = noised(x, relative=False, same_on_batch=same_on_batch) - x
+                relative = noised(x, relative=True, same_on_batch=same_on_batch) - x
+                torch.testing.assert_close(relative, absolute * sample_std, rtol=1e-4, atol=1e-5)
+
+    def test_the_config_surface_accepts_it(self):
+        self.assertIn("relative", accepted_params(get("RandomGaussianNoiseGPU", Backend.GPU)))