diff --git a/sdm/models/kumo/timeseries/anomaly/schedule.py b/sdm/models/kumo/timeseries/anomaly/schedule.py new file mode 100644 index 0000000000..c7e1d983ac --- /dev/null +++ b/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] diff --git a/test/models/kumo/timeseries/anomaly/test_schedule.py b/test/models/kumo/timeseries/anomaly/test_schedule.py new file mode 100644 index 0000000000..21e67ee870 --- /dev/null +++ b/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())