Add Kumo-Anomaly DPM-Solver noise schedule

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

@@ -0,0 +1,103 @@
+# Portions derived from DPM-Solver:
+# Copyright (c) 2022 Cheng Lu
+# Licensed under the MIT License; see third_party/dpm-solver/LICENSE.
+
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+
+class NoiseScheduleVP(torch.nn.Module):
+    """Continuous-time VP coefficients for Kumo-Anomaly's DPM-Solver++.
+
+    Interpolate ``log(alpha_t) = 0.5 * cumsum(log(1 - beta))`` between
+    discrete steps. As in Kumo-TS, discard the tail where half-logSNR is
+    below ``clipped_lambda`` (default -5.1) and assign the retained steps
+    times ``(i + 1) / total_N``.
+    ``total_N`` is the retained step count and ``T = 1`` is the final time.
+    Unlike :class:`DiffusionSchedule`, cumulative alphas are not clamped.
+
+    Methods accept scalar or batched tensors on the module's device and
+    return flattened, one-dimensional tensors. Outside the retained time
+    interval ``[1 / total_N, T]``, interpolation uses the boundary segment
+    for linear extrapolation. The derived tables are nonpersistent buffers.
+
+    Args:
+        betas: One-dimensional diffusion variances strictly between zero
+            and one, with at least two steps remaining after truncation.
+            Coefficients retain the input device and dtype (float32 or
+            float64). Pass ``DiffusionSchedule.beta.float()`` to match
+            Kumo-TS inference.
+        clipped_lambda: Minimum retained half-logSNR. Defaults to -5.1,
+            matching Kumo-TS's numerical-stability cutoff.
+    """
+
+    _log_alpha: Tensor
+    _time: Tensor
+
+    def __init__(self, betas: Tensor, *, clipped_lambda: float = -5.1) -> None:
+        super().__init__()
+        self.schedule = "discrete"
+        self.T = 1.0
+        log_alpha = 0.5 * (1 - betas.detach()).log().cumsum(dim=0)
+        log_sigma = 0.5 * (1 - (2 * log_alpha).exp()).log()
+        # A single host synchronization at construction determines table size.
+        num_clipped = int(
+            torch.searchsorted((log_alpha - log_sigma).flip(0), clipped_lambda)
+        )
+        self.total_N = log_alpha.numel() - num_clipped
+        if self.total_N < 2:
+            raise ValueError(
+                "At least two diffusion steps must remain after truncation"
+            )
+        self.register_buffer(
+            name="_log_alpha",
+            tensor=log_alpha[: self.total_N],
+            persistent=False,
+        )
+        # Match upstream's float32 CPU grid before dtype/device conversion.
+        time = torch.linspace(
+            start=0,
+            end=self.T,
+            steps=self.total_N + 1,
+            device="cpu",
+            dtype=torch.float32,
+        )[1:].to(betas)
+        self.register_buffer("_time", time, persistent=False)
+
+    def marginal_log_mean_coeff(self, t: Tensor) -> Tensor:
+        """Return interpolated log signal coefficients at times ``t``."""
+        return _interpolate(t.reshape(-1), self._time, self._log_alpha)
+
+    def marginal_alpha(self, t: Tensor) -> Tensor:
+        """Return signal coefficients ``alpha_t`` at continuous times ``t``."""
+        return self.marginal_log_mean_coeff(t).exp()
+
+    def marginal_std(self, t: Tensor) -> Tensor:
+        """Return noise standard deviations ``sqrt(1 - alpha_t**2)``."""
+        return (1 - (2 * self.marginal_log_mean_coeff(t)).exp()).sqrt()
+
+    def marginal_lambda(self, t: Tensor) -> Tensor:
+        """Return half-logSNR ``log(alpha_t) - log(sigma_t)``."""
+        log_alpha = self.marginal_log_mean_coeff(t)
+        return log_alpha - 0.5 * (1 - (2 * log_alpha).exp()).log()
+
+    def inverse_lambda(self, lamb: Tensor) -> Tensor:
+        """Return continuous times for the supplied half-logSNR values."""
+        log_alpha = -0.5 * torch.logaddexp(torch.zeros_like(lamb), -2 * lamb)
+        return _interpolate(
+            x=log_alpha.reshape(-1),
+            xp=self._log_alpha.flip(0),
+            yp=self._time.flip(0),
+        )
+
+
+def _interpolate(x: Tensor, xp: Tensor, yp: Tensor) -> Tensor:
+    """Interpolate sorted 1D knots, extrapolating at either boundary."""
+    right = torch.searchsorted(xp, x.contiguous()).clamp(1, xp.numel() - 1)
+    left = right - 1
+    return yp[left] + (x - xp[left]) * (yp[right] - yp[left]) / (
+        xp[right] - xp[left]
+    )

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

@@ -0,0 +1,136 @@
+# 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.noise_schedule import NoiseScheduleVP
+from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule
+from sdm.testing import withCUDA
+
+
+@withCUDA
+@pytest.mark.parametrize(
+    ("schedule", "dtype", "clipped_lambda"),
+    [
+        ("quad", torch.float32, -5.1),
+        ("quad", torch.float64, -5.1),
+        ("linear", torch.float32, -5.1),
+        ("linear", torch.float32, -2.0),
+    ],
+)
+def test_noise_schedule(
+    device: torch.device,
+    schedule: Literal["quad", "linear"],
+    dtype: torch.dtype,
+    clipped_lambda: float,
+) -> None:
+    """Match VP coefficients, tail truncation, and interpolation."""
+    discrete = DiffusionSchedule(
+        num_steps=1000,
+        beta_start=1e-4,
+        beta_end=0.5 if schedule == "quad" else 0.02,
+        schedule=schedule,
+        dtype=dtype,
+        device=device,
+    )
+    module = NoiseScheduleVP(discrete.beta, clipped_lambda=clipped_lambda)
+    log_alpha = 0.5 * (1 - discrete.beta).log().cumsum(0)
+    half_log_snr = log_alpha - 0.5 * (1 - (2 * log_alpha).exp()).log()
+    retained = int((half_log_snr >= clipped_lambda).sum())
+    assert module.total_N == retained
+    assert (retained < 1000) == (schedule == "quad" or clipped_lambda == -2.0)
+
+    grid = torch.linspace(0, 1, retained + 1, dtype=torch.float32)[1:].to(
+        discrete.beta
+    )
+    torch.testing.assert_close(
+        actual=module.marginal_log_mean_coeff(grid),
+        expected=log_alpha[:retained],
+    )
+    # Unsorted queries, repeated values, and both interval endpoints.
+    t = torch.tensor(
+        data=[1.0, 0.435, 1 / retained, 0.731, 0.435],
+        device=device,
+        dtype=dtype,
+    )
+    expected = torch.tensor(
+        data=np.interp(
+            x=t.cpu().numpy(),
+            xp=grid.cpu().numpy(),
+            fp=log_alpha[:retained].cpu().numpy(),
+        ),
+        device=device,
+        dtype=dtype,
+    )
+    alpha = expected.exp()
+    sigma = (1 - (2 * expected).exp()).sqrt()
+    for actual, target in (
+        (module.marginal_log_mean_coeff(t), expected),
+        (module.marginal_alpha(t), alpha),
+        (module.marginal_std(t), sigma),
+        (module.marginal_lambda(t), expected - sigma.log()),
+        (module.inverse_lambda(module.marginal_lambda(t)), t),
+    ):
+        assert actual.device == device
+        assert actual.dtype == dtype
+        torch.testing.assert_close(actual, target)
+
+
+@withCUDA
+def test_noise_schedule_extrapolation(device: torch.device) -> None:
+    """Preserve boundary extrapolation and query-time gradients."""
+    betas = torch.tensor(
+        data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64
+    )
+    module = NoiseScheduleVP(betas)
+    logs = 0.5 * (1 - betas).log().cumsum(0)
+    t = torch.tensor(
+        data=[[0.2, 0.25], [0.375, 1.1]],
+        device=device,
+        dtype=betas.dtype,
+        requires_grad=True,
+    )
+    expected = torch.stack(
+        tensors=[
+            logs[0] - 0.2 * (logs[1] - logs[0]),
+            logs[0],
+            (logs[0] + logs[1]) / 2,
+            logs[-1] + 0.4 * (logs[-1] - logs[-2]),
+        ],
+    )
+    torch.testing.assert_close(module.marginal_log_mean_coeff(t), expected)
+    torch.testing.assert_close(
+        actual=module.inverse_lambda(module.marginal_lambda(t)),
+        expected=t.flatten(),
+    )
+    (gradient,) = torch.autograd.grad(
+        module.marginal_log_mean_coeff(t).sum(), t
+    )
+    assert torch.isfinite(gradient).all()
+    assert (gradient != 0).any()
+
+
+@withCUDA
+def test_noise_schedule_conversion(device: torch.device) -> None:
+    """Move coefficients while keeping outputs and checkpoint contents."""
+    betas = torch.tensor(
+        data=[0.1, 0.2, 0.3, 0.4], dtype=torch.float64, requires_grad=True
+    )
+    module = NoiseScheduleVP(betas)
+    t = torch.tensor([0.3, 0.6, 0.9], dtype=betas.dtype)
+    expected = module.marginal_lambda(t).detach()
+    module = module.to(device=device, dtype=torch.float32)
+    actual = module.marginal_lambda(t.to(device).float())
+    torch.testing.assert_close(actual, expected.to(device).float())
+    assert not actual.requires_grad
+    assert not module.state_dict()
+
+
+def test_noise_schedule_too_short() -> None:
+    """Reject schedules without two usable interpolation knots."""
+    with pytest.raises(ValueError, match="At least two diffusion steps"):
+        NoiseScheduleVP(torch.tensor([0.99999, 0.99999]))