Add Kumo-Anomaly diffusion schedule

sdm/models/kumo/timeseries/anomaly/schedule.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+
+class DiffusionSchedule(torch.nn.Module):
+    """Discrete noise coefficients for Kumo-Anomaly diffusion.
+
+    The one-dimensional buffers ``beta``, ``alpha_hat = 1 - beta``, and
+    ``alpha = cumprod(alpha_hat)`` have length ``num_steps``. Only ``alpha``
+    is clipped to ``[1e-6, 1 - 1e-6]``, matching Kumo-TS. Tables are computed
+    in float64 on CPU before conversion to the requested device and dtype.
+    They are derived from configuration and excluded from the state dict.
+
+    Args:
+        num_steps: Number of diffusion steps, at least one.
+        beta_start: Initial noise variance, between zero and one.
+        beta_end: Final noise variance, between ``beta_start`` and one.
+        schedule: ``"quad"`` interpolates square roots of the noise variances
+            before squaring; ``"linear"`` interpolates the variances directly.
+        device: The device.
+        dtype: Coefficient dtype. Defaults to float64, as in the upstream
+            NumPy schedule. Use float32 or float64 to retain the alpha bounds.
+    """
+
+    beta: Tensor
+    alpha_hat: Tensor
+    alpha: Tensor
+
+    def __init__(
+        self,
+        num_steps: int,
+        beta_start: float,
+        beta_end: float,
+        schedule: Literal["quad", "linear"] = "quad",
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float64,
+    ) -> None:
+        super().__init__()
+        device = torch.get_default_device() if device is None else device
+        self.num_steps = num_steps
+        if schedule == "quad":
+            beta_start = beta_start**0.5
+            beta_end = beta_end**0.5
+        else:
+            assert schedule == "linear"
+        beta = torch.linspace(
+            start=beta_start,
+            end=beta_end,
+            steps=num_steps,
+            device="cpu",
+            dtype=torch.float64,
+        )
+        if schedule == "quad":
+            beta = beta.square()
+        alpha_hat = 1 - beta
+        alpha = alpha_hat.cumprod(dim=0).clamp(min=1e-6, max=1 - 1e-6)
+        for name, value in (
+            ("beta", beta),
+            ("alpha_hat", alpha_hat),
+            ("alpha", alpha),
+        ):
+            self.register_buffer(
+                name=name,
+                tensor=value.to(device=device, dtype=dtype),
+                persistent=False,
+            )
+
+    def forward(self, step: Tensor) -> Tensor:
+        """Look up clipped cumulative signal coefficients.
+
+        Args:
+            step: Integer diffusion indices in ``[0, num_steps)`` with any
+                shape, on the module's device.
+
+        Returns:
+            Cumulative ``alpha`` values with the same shape as ``step`` and
+            the coefficient dtype. Cast to the data dtype and broadcast over
+            feature/time dimensions when constructing noisy observations.
+        """
+        return self.alpha[step]

test/models/kumo/timeseries/anomaly/test_schedule.py

@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+import numpy as np
+import pytest
+import torch
+
+from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule
+from sdm.testing import withCUDA
+
+
+@withCUDA
+@pytest.mark.parametrize(
+    ("schedule", "num_steps", "beta_start", "beta_end", "dtype"),
+    [
+        ("quad", 1000, 1e-4, 0.5, torch.float64),
+        ("linear", 1000, 1e-4, 0.02, torch.float64),
+        ("quad", 1000, 1e-4, 0.5, torch.float32),
+        ("quad", 1, 1e-4, 0.5, torch.float64),
+        ("linear", 100, 1e-10, 0.9, torch.float64),
+    ],
+)
+def test_diffusion_schedule(
+    device: torch.device,
+    schedule: Literal["quad", "linear"],
+    num_steps: int,
+    beta_start: float,
+    beta_end: float,
+    dtype: torch.dtype,
+) -> None:
+    module = DiffusionSchedule(
+        num_steps=num_steps,
+        beta_start=beta_start,
+        beta_end=beta_end,
+        schedule=schedule,
+        device=device,
+        dtype=dtype,
+    )
+    if schedule == "quad":
+        beta = np.linspace(beta_start**0.5, beta_end**0.5, num_steps) ** 2
+    else:
+        beta = np.linspace(beta_start, beta_end, num_steps)
+    alpha_hat = 1 - beta
+    alpha = np.clip(np.cumprod(alpha_hat), 1e-6, 1 - 1e-6)
+    for actual, values in (
+        (module.beta, beta),
+        (module.alpha_hat, alpha_hat),
+        (module.alpha, alpha),
+    ):
+        expected = torch.tensor(values, device=device, dtype=dtype)
+        assert actual.shape == (num_steps,)
+        assert actual.dtype == dtype
+        assert actual.device == device
+        assert not actual.requires_grad
+        torch.testing.assert_close(
+            actual=actual,
+            expected=expected,
+            rtol=1e-12 if dtype == torch.float64 else 1e-6,
+            atol=1e-15,
+        )
+
+    # Unsorted and repeated indices retain their layout, including endpoints.
+    step = torch.tensor(
+        [[num_steps - 1, 0], [num_steps // 2, 0]], device=device
+    )
+    expected = torch.tensor(alpha, device=device, dtype=dtype)[step]
+    torch.testing.assert_close(module(step), expected)
+    if beta_start < 1e-6:
+        assert module.alpha[0] == 1 - 1e-6
+        assert module.alpha[-1] == 1e-6
+        assert module.alpha_hat[0] > 1 - 1e-6
+
+
+@withCUDA
+def test_diffusion_schedule_conversion(device: torch.device) -> None:
+    module = DiffusionSchedule(num_steps=10, beta_start=1e-4, beta_end=0.5)
+    assert module.beta.dtype == torch.float64
+    original = [module.beta, module.alpha_hat, module.alpha]
+    module = module.to(device=device, dtype=torch.float32)
+    for actual, expected in zip(
+        (module.beta, module.alpha_hat, module.alpha), original, strict=True
+    ):
+        torch.testing.assert_close(actual, expected.to(device).float())
+    assert not module.state_dict()
+
+
+def test_diffusion_schedule_default_device() -> None:
+    with torch.device("meta"):
+        module = DiffusionSchedule(num_steps=3, beta_start=0.1, beta_end=0.3)
+        cpu_module = DiffusionSchedule(
+            num_steps=3, beta_start=0.1, beta_end=0.3, device="cpu"
+        )
+    assert all(buffer.device.type == "meta" for buffer in module.buffers())
+    assert all(buffer.device.type == "cpu" for buffer in cpu_module.buffers())