From d000981d24c59bd6107be9266e7802fa6e0cac09 Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Fri, 16 Oct 2026 09:19:13 -0700 Subject: [PATCH 1/3] Add Kumo-Anomaly DPM-Solver++ update Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/sampling.py | 62 +++++++ .../kumo/timeseries/anomaly/test_sampling.py | 151 ++++++++++++++++++ 2 files changed, 213 insertions(+) create mode 100644 sdm/models/kumo/timeseries/anomaly/sampling.py create mode 100644 test/models/kumo/timeseries/anomaly/test_sampling.py diff --git a/sdm/models/kumo/timeseries/anomaly/sampling.py b/sdm/models/kumo/timeseries/anomaly/sampling.py new file mode 100644 index 0000000000..bc2a6c258e --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/sampling.py @@ -0,0 +1,62 @@ +# 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 + +from torch import Tensor + + +def dpm_solver_update( + x: Tensor, + data_start: Tensor, + *, + alpha_end: Tensor, + sigma_start: Tensor, + sigma_end: Tensor, + lambda_delta: Tensor, + data_mid: Tensor | None = None, + r1: Tensor | float = 0.5, +) -> Tensor: + """Apply a first- or second-order single-step DPM-Solver++ update. + + This implements Kumo-TS's ``algorithm_type="dpmsolver++"`` and + ``solver_type="dpmsolver"`` formulas. Predictions are clean data, not + noise: convert predicted noise ``eps`` using ``(x - sigma * eps) / alpha`` + at the corresponding time before calling this function. + + Without ``data_mid``, take a first-order step. For second order, first + use this function to advance from the start to the intermediate time, + then evaluate the model there and pass its clean-data prediction as + ``data_mid`` together with the original ``x`` and full-interval + coefficients. Model calls, time selection, and random draws are owned + by the caller; this function is deterministic and differentiable. + + Args: + x: Noisy sample at the start time, for example shape ``[B, K, L]``. + data_start: Clean-data prediction at the start, with ``x``'s shape. + alpha_end: Signal amplitude at the end time, not cumulative alpha + squared as stored by ``DiffusionSchedule.alpha``. + sigma_start: Positive noise standard deviation at the start time. + sigma_end: Noise standard deviation at the end time. + lambda_delta: End minus start half-logSNR, where + ``lambda = log(alpha) - log(sigma)``. + data_mid: Clean-data prediction at the intermediate sample/time, + with ``x``'s shape. Omit for first order. + r1: Intermediate fraction of the half-logSNR interval, strictly + between zero and one. Ignored for first order. Uniform time + spacing generally does not imply ``r1 = 0.5``; the caller must + compute ``(lambda_mid - lambda_start) / lambda_delta``. + + Returns: + Updated sample with ``x``'s shape. All tensors must be on the same + device with compatible dtypes; standard PyTorch promotion applies. + Coefficients and tensor ``r1`` must broadcast to ``x``: use scalars + for shared times or ``[B, 1, 1]`` for per-window coefficients. + """ + factor = alpha_end * (-lambda_delta).expm1() + out = (sigma_end / sigma_start) * x - factor * data_start + if data_mid is not None: + out = out - (0.5 / r1) * factor * (data_mid - data_start) + return out diff --git a/test/models/kumo/timeseries/anomaly/test_sampling.py b/test/models/kumo/timeseries/anomaly/test_sampling.py new file mode 100644 index 0000000000..7582643353 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_sampling.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from functools import partial + +import pytest +import torch +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.sampling import dpm_solver_update +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("second_order", [False, True]) +@pytest.mark.parametrize("batched", [False, True]) +def test_dpm_solver_update( + device: torch.device, + dtype: torch.dtype, + second_order: bool, + batched: bool, +) -> None: + """Match the reference formulas with shared or per-window coefficients.""" + x = torch.randn(2, 3, 5, device=device, dtype=dtype) + data_start = torch.randn_like(x) + data_mid = torch.randn_like(x) + # Each row is alpha_end, sigma_start, sigma_end, lambda_delta, r1. + values = x.new_tensor( + data=[[0.8, 0.9, 0.6, 0.6, 0.25], [0.6, 0.95, 0.8, 0.2, 0.75]] + ) + if batched: + alpha, sigma_start, sigma_end, delta, r1 = values.T[:, :, None, None] + else: + alpha, sigma_start, sigma_end, delta, r1 = values[0] + before = [value.clone() for value in (x, data_start, data_mid, values)] + actual = dpm_solver_update( + x=x, + data_start=data_start, + alpha_end=alpha, + sigma_start=sigma_start, + sigma_end=sigma_end, + lambda_delta=delta, + data_mid=data_mid if second_order else None, + r1=r1, + ) + expected = torch.empty_like(x) + for i in range(x.size(0)): + a, s, t, h, r = values[i if batched else 0].tolist() + factor = a * math.expm1(-h) + expected[i] = (t / s) * x[i] - factor * data_start[i] + if second_order: + expected[i] -= (0.5 / r) * factor * (data_mid[i] - data_start[i]) + torch.testing.assert_close(actual, expected) + assert actual.shape == x.shape + assert actual.dtype == dtype + assert actual.device == device + for value, original in zip( + (x, data_start, data_mid, values), before, strict=True + ): + torch.testing.assert_close(value, original) + + +@withCUDA +def test_dpm_solver_update_limits(device: torch.device) -> None: + """Check constant predictions, zero intervals, and midpoint limits.""" + x = torch.randn(2, 3, 5, device=device) + predicted = torch.randn_like(x) + update = partial( + dpm_solver_update, + x=x, + data_start=predicted, + alpha_end=x.new_tensor(0.8), + sigma_start=x.new_tensor(0.6), + sigma_end=x.new_tensor(0.6), + ) + delta = x.new_tensor(0.3) + torch.testing.assert_close( + actual=update(lambda_delta=delta, data_mid=predicted, r1=0.3), + expected=update(lambda_delta=delta), + ) + torch.testing.assert_close( + actual=update(lambda_delta=x.new_zeros(()), data_mid=predicted + 1), + expected=x, + rtol=0, + atol=0, + ) + # A half-logSNR midpoint cancels the start prediction in the final step. + mid = predicted + 1 + torch.testing.assert_close( + actual=update(lambda_delta=delta, data_mid=mid), + expected=x - 0.8 * (-delta).expm1() * mid, + ) + + +@withCUDA +@pytest.mark.parametrize("second_order", [False, True]) +def test_dpm_solver_update_small_interval( + device: torch.device, + second_order: bool, +) -> None: + """Retain updates too small for an exp-minus-one computation in float32.""" + x = torch.zeros(2, 1, 3, device=device) + out = dpm_solver_update( + x=x, + data_start=torch.ones_like(x), + alpha_end=x.new_tensor(1.0), + sigma_start=x.new_tensor(1.0), + sigma_end=x.new_tensor(1.0), + lambda_delta=x.new_tensor(1e-10), + data_mid=torch.full_like(x, 3) if second_order else None, + r1=0.25, + ) + expected = torch.full_like(x, 5e-10 if second_order else 1e-10) + torch.testing.assert_close( + actual=out, expected=expected, rtol=1e-6, atol=0 + ) + + +@withCUDA +@pytest.mark.parametrize("second_order", [False, True]) +def test_dpm_solver_update_gradients( + device: torch.device, + second_order: bool, +) -> None: + """Differentiate samples and clean-data predictions for either order.""" + inputs = tuple( + torch.randn( + size=(2, 1, 2), + device=device, + dtype=torch.float64, + requires_grad=True, + ) + for _ in range(3) + ) + + def evaluate(x: Tensor, start: Tensor, mid: Tensor) -> Tensor: + """Evaluate with fixed coefficients for gradient checking.""" + return dpm_solver_update( + x=x, + data_start=start, + alpha_end=x.new_tensor(0.8), + sigma_start=x.new_tensor(0.9), + sigma_end=x.new_tensor(0.6), + lambda_delta=x.new_tensor(0.6), + data_mid=mid if second_order else None, + r1=0.3, + ) + + assert torch.autograd.gradcheck(evaluate, inputs) From 917b27278b3e5419d083243d0164bf56e77d2a56 Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Fri, 16 Oct 2026 10:31:48 -0700 Subject: [PATCH 2/3] Add Kumo-Anomaly pretrained inference Signed-off-by: Aditi Gautam --- docs/source/api/models.rst | 13 +- examples/kumo_anomaly.py | 31 ++ sdm/models/__init__.py | 8 +- sdm/models/kumo/__init__.py | 3 +- sdm/models/kumo/timeseries/__init__.py | 3 + .../kumo/timeseries/anomaly/__init__.py | 5 + sdm/models/kumo/timeseries/anomaly/ckpt.py | 47 +++ .../kumo/timeseries/anomaly/denoiser.py | 148 ++++++++ sdm/models/kumo/timeseries/anomaly/model.py | 333 ++++++++++++++++++ .../kumo/timeseries/anomaly/noise_schedule.py | 100 ++++++ .../kumo/timeseries/anomaly/preprocessing.py | 91 +++++ sdm/models/kumo/timeseries/anomaly/recipe.py | 21 ++ .../kumo/timeseries/anomaly/sampling.py | 125 +++++++ .../kumo/timeseries/anomaly/test_denoiser.py | 122 +++++++ .../kumo/timeseries/anomaly/test_model.py | 233 ++++++++++++ .../timeseries/anomaly/test_noise_schedule.py | 153 ++++++++ .../timeseries/anomaly/test_preprocessing.py | 63 ++++ .../kumo/timeseries/anomaly/test_solver.py | 160 +++++++++ 18 files changed, 1656 insertions(+), 3 deletions(-) create mode 100644 examples/kumo_anomaly.py create mode 100644 sdm/models/kumo/timeseries/anomaly/ckpt.py create mode 100644 sdm/models/kumo/timeseries/anomaly/denoiser.py create mode 100644 sdm/models/kumo/timeseries/anomaly/model.py create mode 100644 sdm/models/kumo/timeseries/anomaly/noise_schedule.py create mode 100644 sdm/models/kumo/timeseries/anomaly/preprocessing.py create mode 100644 sdm/models/kumo/timeseries/anomaly/recipe.py create mode 100644 test/models/kumo/timeseries/anomaly/test_denoiser.py create mode 100644 test/models/kumo/timeseries/anomaly/test_model.py create mode 100644 test/models/kumo/timeseries/anomaly/test_noise_schedule.py create mode 100644 test/models/kumo/timeseries/anomaly/test_preprocessing.py create mode 100644 test/models/kumo/timeseries/anomaly/test_solver.py diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 48306b949d..9b5688055b 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -58,6 +58,17 @@ Model API :template: model_class :nosignatures: -{% for name in api_names("sdm.models") if name != "ECOC" %} +{% for name in api_names("sdm.models") if name not in ("ECOC", "KumoAnomaly") %} ~sdm.models.{{ name }} {% endfor %} + +Anomaly Detection +----------------- + +Unlike supervised ICL models, Kumo-Anomaly reconstructs an ordered series +without labels. Its scores are reconstruction errors, not probabilities. + +.. autoclass:: sdm.models.KumoAnomaly + :members: + +.. autoclass:: sdm.models.kumo.timeseries.anomaly.AnomalyOutput diff --git a/examples/kumo_anomaly.py b/examples/kumo_anomaly.py new file mode 100644 index 0000000000..0d74955207 --- /dev/null +++ b/examples/kumo_anomaly.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Score an ordered series using reference-only feature preprocessing.""" + +import torch + +import sdm +from sdm.models import KumoAnomaly + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +time = torch.arange(600, device=device, dtype=torch.float32) +values = torch.stack([torch.sin(time / 12), torch.cos(time / 20)], dim=-1) +values[450:460] += 3 +series = sdm.TableTensor.from_tensor(values) +model = KumoAnomaly(device=device) + +# Reference rows fit PCA/min-max only; model weights are already pretrained. +# Reconstruction uses future rows within each query window, so this is an +# offline detector, not a causal streaming evaluation. +result = model( + x=series[300:], + reference=series[:300], + num_samples=3, + num_steps=20, + generator=torch.Generator(device=device).manual_seed(42), +) +print(result.scores) +print(result.reconstruction) +# Scores and reconstructions use scaled model-space units. Thresholds need +# separate validation; the corrected DPM timestep mapping changes SDK scores. diff --git a/sdm/models/__init__.py b/sdm/models/__init__.py index 57a27fa914..f5cf736ad9 100644 --- a/sdm/models/__init__.py +++ b/sdm/models/__init__.py @@ -14,7 +14,12 @@ from sdm.models.ecoc import ECOC from sdm.models.tabiclv2 import TabICLv2 from sdm.models.tabfm import TabFM -from sdm.models.kumo import KumoTabular, KumoRelational, KumoForecast +from sdm.models.kumo import ( + KumoTabular, + KumoRelational, + KumoForecast, + KumoAnomaly, +) from sdm.models.timesfm3 import TimesFM3 @@ -26,5 +31,6 @@ "KumoTabular", "KumoRelational", "KumoForecast", + "KumoAnomaly", "TimesFM3", ] diff --git a/sdm/models/kumo/__init__.py b/sdm/models/kumo/__init__.py index aa3df72107..17a53b7f7e 100644 --- a/sdm/models/kumo/__init__.py +++ b/sdm/models/kumo/__init__.py @@ -11,11 +11,12 @@ from sdm.models.kumo.tabular import KumoTabular from sdm.models.kumo.relational import KumoRelational -from sdm.models.kumo.timeseries import KumoForecast +from sdm.models.kumo.timeseries import KumoForecast, KumoAnomaly __all__ = [ "KumoTabular", "KumoRelational", "KumoForecast", + "KumoAnomaly", ] diff --git a/sdm/models/kumo/timeseries/__init__.py b/sdm/models/kumo/timeseries/__init__.py index 8d3fa795b5..6ca084ba71 100644 --- a/sdm/models/kumo/timeseries/__init__.py +++ b/sdm/models/kumo/timeseries/__init__.py @@ -5,11 +5,14 @@ __lazy_modules__ = { "sdm.models.kumo.timeseries.forecast", + "sdm.models.kumo.timeseries.anomaly", } from sdm.models.kumo.timeseries.forecast import KumoForecast +from sdm.models.kumo.timeseries.anomaly import KumoAnomaly __all__ = [ "KumoForecast", + "KumoAnomaly", ] diff --git a/sdm/models/kumo/timeseries/anomaly/__init__.py b/sdm/models/kumo/timeseries/anomaly/__init__.py index dd56fab6f0..7928e4bf29 100644 --- a/sdm/models/kumo/timeseries/anomaly/__init__.py +++ b/sdm/models/kumo/timeseries/anomaly/__init__.py @@ -2,3 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """Kumo time-series anomaly detection model components.""" + +from sdm.models.kumo.timeseries.anomaly.model import KumoAnomaly, AnomalyOutput + + +__all__ = ["KumoAnomaly", "AnomalyOutput"] diff --git a/sdm/models/kumo/timeseries/anomaly/ckpt.py b/sdm/models/kumo/timeseries/anomaly/ckpt.py new file mode 100644 index 0000000000..0600293763 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/ckpt.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from torch import Tensor + + +def model_kwargs(config: dict[str, Any]) -> dict[str, Any]: + """Read architecture from the checkpoint's embedded training config. + + The separately published YAML describes a different feature width. + Only the released conditional model is supported; an unconditional + checkpoint must not silently load into a different conditioning path. + """ + model, diffusion = config["model"], config["diffusion"] + if model["is_unconditional"]: + raise ValueError("KumoAnomaly requires a conditional checkpoint") + return { + "num_features": model["target_dim"], + "time_channels": model["timeemb"], + "feature_channels": model["featureemb"], + "channels": diffusion["channels"], + "num_layers": diffusion["layers"], + "num_heads": diffusion["nheads"], + "embedding_channels": diffusion["diffusion_embedding_dim"], + "num_steps": diffusion["num_steps"], + "beta_start": diffusion["beta_start"], + "beta_end": diffusion["beta_end"], + "schedule": diffusion["schedule"], + } + + +def remap_ckpt(state: dict[str, Tensor]) -> dict[str, Tensor]: + """Map upstream parameter names, preserving all keys for strict loading.""" + out = {} + for key, value in state.items(): + if key.startswith("diffmodel."): + key = "denoiser." + key.removeprefix("diffmodel.") + key = key.replace("time_layer.layers.0.", "time_layer.") + key = key.replace("feature_layer.layers.0.", "feature_layer.") + elif key.startswith( + ("embed_layer.", "embed_norm.", "time_embed_norm.") + ): + key = "side_embedding." + key + out[key] = value + return out diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py new file mode 100644 index 0000000000..eaf22276b7 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Conv1d, Embedding, LayerNorm, ModuleList + +from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock +from sdm.models.kumo.timeseries.anomaly.embedding import DiffusionEmbedding + + +class DiffusionDenoiser(torch.nn.Module): + """Predict noise for one Kumo-Anomaly diffusion step. + + Callers supply normalized noisy/observed inputs and precomputed side + information. This component does not sample or compute anomaly scores. + + Args: + channels: Width of the residual blocks. + side_channels: Width of the precomputed conditioning information. + num_steps: Number of discrete diffusion steps. + num_layers: Number of residual blocks, at least one. + num_heads: Attention heads for each axis. + embedding_channels: Even diffusion and strategy embedding width, + at least four. + input_channels: Input width (two for conditional, one unconditional). + hidden_channels: Transformer feedforward width. + dropout: Transformer attention and feedforward dropout probability. + device: The device. + dtype: The parameter and output dtype. + """ + + def __init__( + self, + channels: int, + side_channels: int, + num_steps: int, + num_layers: int, + num_heads: int, + embedding_channels: int = 128, + input_channels: int = 2, + hidden_channels: int = 64, + dropout: float = 0.1, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.diffusion_embedding = DiffusionEmbedding( + num_steps=num_steps, + channels=embedding_channels, + **factory_kwargs, + ) + self.strategy_embedding = Embedding( + num_embeddings=2, + embedding_dim=embedding_channels, + **factory_kwargs, + ) + self.input_projection = Conv1d( + in_channels=input_channels, + out_channels=channels, + kernel_size=1, + **factory_kwargs, + ) + self.output_projection1 = Conv1d( + in_channels=channels, + out_channels=channels, + kernel_size=1, + **factory_kwargs, + ) + self.output_projection2 = Conv1d( + in_channels=channels, + out_channels=1, + kernel_size=1, + **factory_kwargs, + ) + for projection in (self.input_projection, self.output_projection1): + torch.nn.init.kaiming_normal_( + tensor=projection.weight, mode="fan_out", nonlinearity="relu" + ) + with torch.no_grad(): + projection.weight.mul_(0.5) + torch.nn.init.xavier_uniform_( + tensor=self.output_projection2.weight, gain=0.01 + ) + self.residual_layers = ModuleList( + ResidualBlock( + channels=channels, + side_channels=side_channels, + embedding_channels=embedding_channels, + num_heads=num_heads, + hidden_channels=hidden_channels, + dropout=dropout, + **factory_kwargs, + ) + for _ in range(num_layers) + ) + self.norm_after_skip = LayerNorm(channels, **factory_kwargs) + self.norm_after_proj1 = LayerNorm(channels, **factory_kwargs) + + def forward( + self, + x: Tensor, + side_info: Tensor, + step: Tensor, + strategy: Tensor, + ) -> Tensor: + """Predict noise for normalized time-series windows. + + Args: + x: Input with shape ``[B, I, K, L]``, where ``B`` is batch size, + ``I`` input channels, ``K`` features, and ``L`` time steps. + side_info: Conditioning with shape ``[B, S, K, L]``, where + ``S`` is ``side_channels``. + step: Integer diffusion indices in ``[0, num_steps)`` with + shape ``[B]`` or ``[1]``. + strategy: Integer masking-strategy indices in ``{0, 1}`` with + shape ``[B]`` or ``[1]``. Single-entry indices broadcast + across the batch independently for either embedding. + + Returns: + Predicted noise with shape ``[B, K, L]`` in input feature/time + order. No denormalization or mask restoration is applied. + """ + batch, _, features, length = x.shape + x = self.input_projection(x.flatten(2)).relu() + x = x.reshape(batch, -1, features, length) + diffusion_embedding = self.diffusion_embedding(step) + strategy_embedding = self.strategy_embedding(strategy) + skips = [] + for layer in self.residual_layers: + x, skip = layer( + x=x, + side_info=side_info, + diffusion_embedding=diffusion_embedding, + strategy_embedding=strategy_embedding, + ) + skips.append(skip) + x = torch.stack(skips).sum(dim=0) / math.sqrt(len(skips)) + x = self.norm_after_skip(x.flatten(2).transpose(1, 2)) # [B, K * L, C] + x = self.output_projection1(x.transpose(1, 2)) + x = self.norm_after_proj1(x.transpose(1, 2)).transpose(1, 2) + return self.output_projection2(x.relu()).reshape( + batch, features, length + ) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py new file mode 100644 index 0000000000..2c5b054702 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -0,0 +1,333 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence +from functools import partial +from pathlib import Path +from typing import Any, Literal, NamedTuple, cast + +import torch +from torch import Tensor + +from sdm import Recipe, Stype, TableTensor +from sdm.models._huggingface import download_checkpoint +from sdm.models.kumo.timeseries.anomaly.ckpt import model_kwargs, remap_ckpt +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding +from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule + + +class AnomalyOutput(NamedTuple): + """Per-row anomaly scores and model-space reconstruction. + + Attributes: + scores: Mean absolute reconstruction error, shape ``[time]``. + l2_scores: Euclidean reconstruction error, shape ``[time]``. + reconstruction: Reconstructed series in min-max-scaled feature or + PCA space, shape ``[time, min(input_features, model_features)]``. + target: Preprocessed input in the same space as reconstruction. + + Padding features are excluded from every field. Higher scores indicate + larger reconstruction errors, not calibrated anomaly probabilities. + """ + + scores: Tensor + l2_scores: Tensor + reconstruction: TableTensor + target: TableTensor + + +class KumoAnomaly(torch.nn.Module): + """Pretrained Kumo-Anomaly diffusion reconstruction and anomaly scoring. + + Unlike supervised ICL models, this model needs no labels or prediction + task. Pass one ordered numerical series to :meth:`forward`. Alternating + temporal blocks are reconstructed using two complementary conditioning + masks and DPM-Solver++; the sample median supplies the reconstruction. + + The port corrects the original SDK's double timestep conversion. It + can change predictions; SDK score or threshold equivalence is not + claimed. Validate downstream thresholds again. The released conditional + checkpoint is supported; training, DDPM/DDIM sampling, and the SDK's + dataset-specific preprocessing and reporting are not included. + + Args: + pretrained: Load the published checkpoint when true, initialize + random weights when false, or load an upstream-format local + checkpoint when given a path. Loading uses ``weights_only=True``. + window_length: Non-overlapping reconstruction window length. + Defaults to the SDK inference length, not training metadata. + mask_blocks: Even number of alternating blocks per window. Must + divide ``window_length``. The SDK inference default is four. + device: Device for parameters and inference. + """ + + def __init__( + self, + pretrained: bool | str | Path = True, + window_length: int = 100, + mask_blocks: int = 4, + device: torch.device | str | None = None, + ) -> None: + super().__init__() + if ( + window_length < 1 + or mask_blocks < 2 + or mask_blocks % 2 + or window_length % mask_blocks + ): + raise ValueError( + "'mask_blocks' must be even and divide a positive " + "'window_length'" + ) + self.window_length = window_length + self.mask_blocks = mask_blocks + self.scale_factor = 1.0 + if pretrained is not False: + path = ( + download_checkpoint( + repo_id="nvidia/Kumo-Anomaly", + filename="final_model.pth", + revision="226c5003b0582adc1cac433123f47324bfae39f2", + ) + if pretrained is True + else pretrained + ) + checkpoint = torch.load( + path, map_location="cpu", weights_only=True + ) + config = checkpoint["config"] + self.model = _KumoAnomaly(**model_kwargs(config)) + self.model.load_state_dict(remap_ckpt(checkpoint["model"])) + self.scale_factor = config["dataset"].get("scale_factor", 1.0) + else: + self.model = _KumoAnomaly() + self.to(device=device) + self.eval() + + @classmethod + def default_recipe(cls) -> Recipe: + """Return the published checkpoint's feature preprocessing recipe.""" + return default_recipe() + + @torch.no_grad() + def forward( + self, + x: Tensor | TableTensor, + *, + reference: Tensor | TableTensor | None = None, + num_samples: int = 30, + num_steps: int = 20, + batch_size: int = 8, + generator: torch.Generator | None = None, + ) -> AnomalyOutput: + """Reconstruct and score every row of an ordered series. + + Args: + x: Finite numerical observations, shape ``[time, features]``. + Select measured features only, excluding timestamps, IDs, + and labels. No sorting or resampling is performed. + reference: Optional reference rows with the same ordered columns + for fitting PCA and min-max ranges. Without reference, + preprocessing fits on ``x`` (transductive SDK behavior). + Even with reference, reconstruction uses the whole query + window, including future rows; this is not a causal detector. + num_samples: Reconstruction samples per masking strategy. + Even counts use the lower median, matching Kumo-TS. + num_steps: Denoiser evaluations per sample and strategy. + batch_size: Number of windows reconstructed together. Changing + batching changes the assignment of random draws to windows. + generator: Random generator on the model device. Supplied + generators leave the global random state unchanged. + + Returns: + Scores and reconstructions in preprocessed model space, not raw + input units. Short/final partial windows repeat their last row + and are trimmed afterward; unlike the original SDK, no trailing + observations are dropped. Padding can affect that window's + normalization and reconstruction. No anomaly threshold is fitted. + """ + if self.training: + raise RuntimeError("Call 'eval()' before anomaly inference") + if num_samples < 1: + raise ValueError("'num_samples' must be at least one") + x = self._as_table(x) + reference = x if reference is None else self._as_table(reference) + if reference.columns != x.columns: + raise ValueError( + "Reference and query must have the same ordered columns" + ) + recipe = default_recipe( + num_features=self.model.num_features, + scale_factor=self.scale_factor, + ) + recipe.features.fit(reference) + processed = recipe.features.transform(x) + values = processed.numerical + length, features = values.shape + padding = (-length) % self.window_length + if padding: + values = torch.cat([values, values[-1:].expand(padding, -1)]) + windows = values.reshape(-1, self.window_length, features) + sampler = self.model.make_sampler(num_steps).to(values.device) + reconstruction = torch.cat( + [ + self.model( + windows=batch, + sampler=sampler, + num_samples=num_samples, + mask_blocks=self.mask_blocks, + generator=generator, + ) + for batch in windows.split(batch_size) + ], + dim=0, + ).reshape(-1, features)[:length] + count = min(x.size(-1), features) + target = processed.numerical[:, :count] + reconstruction = reconstruction[:, :count] + # The SDK bounds residuals in model space before MAE/L2 reduction. + error = reconstruction.clamp(-1e8, 1e8).nan_to_num( + nan=0.0 + ) - target.clamp(-1e8, 1e8).nan_to_num(nan=0.0) + columns: dict[str, Sequence[str]] = { + Stype.numerical: x.columns[Stype.numerical] + if x.size(-1) <= features + else [f"pca_{i}" for i in range(features)] + } + return AnomalyOutput( + scores=error.abs().mean(dim=-1).clamp(0, 1e8), + l2_scores=error.norm(dim=-1).clamp(0, 1e8), + reconstruction=TableTensor( + columns=columns, numerical=reconstruction + ), + target=TableTensor(columns=columns, numerical=target), + ) + + def _as_table(self, x: Tensor | TableTensor) -> TableTensor: + """Normalize inputs and reject silently invalid features.""" + if not isinstance(x, TableTensor): + x = TableTensor.from_tensor(x) + if x.ndim != 2 or x.size(-2) == 0 or x.size(-1) == 0: + raise ValueError("Expected a nonempty [time, features] series") + if x.numerical.size(-1) != x.size(-1): + raise ValueError( + "Select numerical measurement columns before inference" + ) + weight = self.model.side_embedding.embed_layer.weight + x = cast(TableTensor, x.to(device=weight.device, dtype=weight.dtype)) + if not x.numerical.isfinite().all(): + raise ValueError("Impute non-finite observations before inference") + return x + + +class _KumoAnomaly(torch.nn.Module): + """Compose the checkpoint-compatible conditional diffusion modules.""" + + def __init__( + self, + num_features: int = 40, + time_channels: int = 256, + feature_channels: int = 64, + channels: int = 128, + num_layers: int = 6, + num_heads: int = 8, + embedding_channels: int = 256, + num_steps: int = 500, + beta_start: float = 0.0001, + beta_end: float = 0.01, + schedule: Literal["linear", "quad"] = "linear", + ) -> None: + super().__init__() + self.num_features = num_features + self.schedule_kwargs: dict[str, Any] = { + "num_steps": num_steps, + "beta_start": beta_start, + "beta_end": beta_end, + "schedule": schedule, + } + self.side_embedding = SideEmbedding( + num_features=num_features, + time_channels=time_channels, + feature_channels=feature_channels, + ) + self.denoiser = DiffusionDenoiser( + channels=channels, + side_channels=time_channels + feature_channels, + num_steps=num_steps, + num_layers=num_layers, + num_heads=num_heads, + embedding_channels=embedding_channels, + ) + + def make_sampler(self, num_steps: int) -> DPMSolver: + """Build inference coefficients on CPU once per scored series.""" + schedule = DiffusionSchedule(**self.schedule_kwargs, device="cpu") + return DPMSolver(betas=schedule.beta.float(), num_steps=num_steps) + + def forward( + self, + windows: Tensor, + sampler: DPMSolver, + num_samples: int = 30, + mask_blocks: int = 4, + generator: torch.Generator | None = None, + ) -> Tensor: + """Reconstruct windows of shape ``[batch, time, features]``.""" + observed = windows.transpose(1, 2) # [B, K, L] + batch, _, length = observed.shape + mean = observed.mean(dim=(1, 2), keepdim=True) + std = observed.std(dim=(1, 2), keepdim=True) + 1e-5 + normalized = ((observed - mean) / std).clamp(-10, 10) + time = torch.arange(length, device=windows.device).expand(batch, -1) + blocks = time // (length // mask_blocks) + predictions = [] + for strategy in range(2): + mask = ((blocks % 2) == strategy).to(windows.dtype) + mask = mask.unsqueeze(1).expand_as(observed) + side = self.side_embedding(time, mask) + strategy_index = torch.tensor([strategy], device=windows.device) + + predict_noise = partial( + self._predict_noise, + conditioned=mask * normalized, + target_mask=1 - mask, + side=side, + strategy=strategy_index, + ) + + samples = [] + for _ in range(num_samples): + # Preserve the source's randn_like memory layout so an + # identical generator assigns noise to the same features. + noise = torch.empty_like(normalized).normal_( + generator=generator + ) + sample = sampler(noise, predict_noise) * std + mean + samples.append((1 - mask) * sample) + predictions.append(torch.stack(samples, dim=1)) + return ( + (predictions[0] + predictions[1]) + .median(dim=1) + .values.transpose(1, 2) + ) + + def _predict_noise( + self, + x: Tensor, + step: Tensor, + *, + conditioned: Tensor, + target_mask: Tensor, + side: Tensor, + strategy: Tensor, + ) -> Tensor: + """Condition the denoiser without converting solver indices.""" + return self.denoiser( + x=torch.stack([conditioned, target_mask * x], dim=1), + side_info=side, + step=step, + strategy=strategy, + ) diff --git a/sdm/models/kumo/timeseries/anomaly/noise_schedule.py b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py new file mode 100644 index 0000000000..0a96b4baac --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py @@ -0,0 +1,100 @@ +# 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 -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. + """ + + _log_alpha: Tensor + _time: Tensor + + def __init__(self, betas: Tensor) -> 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. + clipped = int( + torch.searchsorted((log_alpha - log_sigma).flip(0), -5.1) + ) + self.total_N = log_alpha.numel() - 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] + ) diff --git a/sdm/models/kumo/timeseries/anomaly/preprocessing.py b/sdm/models/kumo/timeseries/anomaly/preprocessing.py new file mode 100644 index 0000000000..ad46103867 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/preprocessing.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn.functional as F +from torch import Tensor + +from sdm import Stype, TableTensor +from sdm.processing import Processor + + +class AnomalyPreprocess(Processor): + """Match Kumo-Anomaly's feature width and min-max preprocessing. + + Numerical inputs wider than the checkpoint are centered and projected + with full-SVD PCA, using the largest-loading-positive sign convention. + Narrower inputs are right-padded with zeros after scaling. Constant + columns use a range of one. Query values are not clipped to the fitted + range. Non-numerical columns are not model inputs. + + Fit on reference data for held-out evaluation. Fitting on the entire + scored series is the original SDK's explicitly transductive workflow. + Inputs must be finite. Rows and leading batch dimensions are preserved; + statistics are fitted independently over the row dimension. + + Args: + num_features: Feature width required by the checkpoint. + scale_factor: Multiplier applied after min-max normalization. + """ + + handles_stypes = (Stype.numerical,) + requires_fit = True + + mean: Tensor + components: Tensor + minimum: Tensor + scale: Tensor + + def __init__( + self, num_features: int = 40, scale_factor: float = 1.0 + ) -> None: + super().__init__() + self.num_features = num_features + self.scale_factor = scale_factor + for name in ("mean", "components", "minimum", "scale"): + self.register_buffer(name, torch.empty(0)) + + def _fit( + self, + table: TableTensor, + *, + generator: torch.Generator | None = None, + ) -> None: + """Learn the projection and ranges using reference rows only.""" + x = table.numerical + self._set_buffer("mean", x.new_zeros((*x.shape[:-2], 1, x.size(-1)))) + self._set_buffer("components", x.new_empty(0)) + if x.size(-1) > self.num_features: + if x.size(-2) < self.num_features: + raise ValueError( + "PCA requires at least 'num_features' reference rows" + ) + self._set_buffer("mean", x.mean(dim=-2, keepdim=True)) + _, _, vh = torch.linalg.svd(x - self.mean, full_matrices=False) + vh = vh[..., : self.num_features, :] + signs = vh.gather(-1, vh.abs().argmax(dim=-1, keepdim=True)).sign() + self._set_buffer("components", (vh * signs).transpose(-1, -2)) + x = (x - self.mean) @ self.components + minimum = x.amin(dim=-2, keepdim=True) + scale = x.amax(dim=-2, keepdim=True) - minimum + self._set_buffer("minimum", minimum) + self._set_buffer("scale", torch.where(scale == 0, 1, scale)) + + def _transform(self, table: TableTensor) -> TableTensor: + """Apply fitted preprocessing and append unscored padding columns.""" + x = table.numerical + if x.size(-1) != self.mean.size(-1): + raise ValueError("Reference and query feature counts must match") + if self.components.numel(): + x = (x - self.mean) @ self.components + x = (x - self.minimum) / self.scale * self.scale_factor + x = F.pad(x, (0, self.num_features - x.size(-1))) + # Model-space columns are not raw features when PCA is used. + return TableTensor( + columns={ + Stype.numerical: [ + f"feature_{i}" for i in range(self.num_features) + ] + }, + numerical=x, + ) diff --git a/sdm/models/kumo/timeseries/anomaly/recipe.py b/sdm/models/kumo/timeseries/anomaly/recipe.py new file mode 100644 index 0000000000..0c63066f56 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/recipe.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from sdm import Recipe +from sdm.models.kumo.timeseries.anomaly.preprocessing import AnomalyPreprocess + + +def default_recipe( + num_features: int = 40, scale_factor: float = 1.0 +) -> Recipe: + """Build the feature-matching and min-max recipe for Kumo-Anomaly. + + Args: + num_features: Checkpoint feature width. + scale_factor: Multiplier applied after min-max normalization. + """ + return Recipe( + features=AnomalyPreprocess( + num_features=num_features, scale_factor=scale_factor + ) + ) diff --git a/sdm/models/kumo/timeseries/anomaly/sampling.py b/sdm/models/kumo/timeseries/anomaly/sampling.py index bc2a6c258e..18591fcdb7 100644 --- a/sdm/models/kumo/timeseries/anomaly/sampling.py +++ b/sdm/models/kumo/timeseries/anomaly/sampling.py @@ -5,8 +5,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable + +import torch from torch import Tensor +from sdm.models.kumo.timeseries.anomaly.noise_schedule import NoiseScheduleVP + def dpm_solver_update( x: Tensor, @@ -60,3 +65,123 @@ def dpm_solver_update( if data_mid is not None: out = out - (0.5 / r1) * factor * (data_mid - data_start) return out + + +class DPMSolver(torch.nn.Module): + """Second-order, single-step DPM-Solver++ on a uniform time grid. + + The denoiser receives integer training indices, not the 1000-scaled + coordinates used by the original DPM-Solver model wrapper. Continuous + times map to the nearest trained noise level ``round(t * total_N - 1)``. + This corrects Kumo-TS's double timestep conversion. No final denoising + or clipping of clean-data predictions is applied. + + Args: + betas: Discrete training variances on CPU, in training order. + num_steps: Number of denoiser evaluations, at least one. Odd counts + end with a first-order update. + """ + + alpha: Tensor + sigma: Tensor + delta: Tensor + fraction: Tensor + steps: Tensor + + def __init__(self, betas: Tensor, num_steps: int = 20) -> None: + super().__init__() + if num_steps < 1: + raise ValueError("'num_steps' must be at least one") + schedule = NoiseScheduleVP(betas) + self.orders = (2,) * (num_steps // 2) + (1,) * (num_steps % 2) + time = torch.linspace( + start=schedule.T, + end=1 / schedule.total_N, + steps=num_steps + 1, + device="cpu", + dtype=torch.float32, + ) + intervals = [] + deltas = [] + fractions = [] + offset = 0 + for order in self.orders: + start, end = time[offset], time[offset + order] + # Match the reference solver's inner float32 grid, including + # its round trip through half-logSNR for the intermediate time. + inner = torch.linspace( + start=float(start), + end=float(end), + steps=3, + device="cpu", + dtype=torch.float32, + ) + lamb = schedule.marginal_lambda(inner) + delta = lamb[2] - lamb[0] + fraction = (lamb[1] - lamb[0]) / delta + mid = schedule.inverse_lambda(lamb[0] + fraction * delta)[0] + intervals.append(torch.stack([start, mid, end])) + deltas.append(delta) + fractions.append(fraction) + offset += order + times = torch.stack(intervals) + for name, value in ( + ("alpha", schedule.marginal_alpha(times).reshape(-1, 3)), + ("sigma", schedule.marginal_std(times).reshape(-1, 3)), + ("delta", torch.stack(deltas)), + ("fraction", torch.stack(fractions)), + ( + "steps", + (times * schedule.total_N - 1) + .round() + .long() + .clamp(0, schedule.total_N - 1), + ), + ): + self.register_buffer(name=name, tensor=value, persistent=False) + + def forward( + self, + x: Tensor, + predict_noise: Callable[[Tensor, Tensor], Tensor], + ) -> Tensor: + """Integrate initial noise using a noise-predicting denoiser. + + Args: + x: Initial noise, for example ``[batch, features, time]``. + predict_noise: Callable accepting ``x`` and a one-element + integer timestep tensor, returning noise with ``x``'s shape. + It owns conditioning and must not rescale the timestep. + + Returns: + Reconstructed data with ``x``'s shape. Move this module to + ``x.device`` before calling; integration itself draws no noise. + """ + for i, order in enumerate(self.orders): + alpha, sigma = self.alpha[i], self.sigma[i] + delta, fraction = self.delta[i], self.fraction[i] + eps = predict_noise(x, self.steps[i, :1]) + data = (x - sigma[0] * eps) / alpha[0] + data_mid = None + if order == 2: + mid = dpm_solver_update( + x=x, + data_start=data, + alpha_end=alpha[1], + sigma_start=sigma[0], + sigma_end=sigma[1], + lambda_delta=fraction * delta, + ) + eps_mid = predict_noise(mid, self.steps[i, 1:2]) + data_mid = (mid - sigma[1] * eps_mid) / alpha[1] + x = dpm_solver_update( + x=x, + data_start=data, + alpha_end=alpha[2], + sigma_start=sigma[0], + sigma_end=sigma[2], + lambda_delta=delta, + data_mid=data_mid, + r1=fraction, + ) + return x diff --git a/test/models/kumo/timeseries/anomaly/test_denoiser.py b/test/models/kumo/timeseries/anomaly/test_denoiser.py new file mode 100644 index 0000000000..401bb05411 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_denoiser.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import pytest +import torch +import torch.nn.functional as F +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.testing import withCUDA + + +def reference_denoiser( + model: DiffusionDenoiser, + x: Tensor, + side: Tensor, + step: Tensor, + strategy: Tensor, +) -> Tensor: + """Compose the tested blocks with pointwise linear projections.""" + x = x.permute(0, 2, 3, 1) + projection = model.input_projection + x = F.linear(x, projection.weight.squeeze(-1), projection.bias).relu() + x = x.permute(0, 3, 1, 2) + diffusion = model.diffusion_embedding(step) + strategy_embedding = model.strategy_embedding(strategy) + skips = [] + for block in model.residual_layers: + x, skip = block(x, side, diffusion, strategy_embedding) + skips.append(skip.permute(0, 2, 3, 1)) + x = sum(skips) / math.sqrt(len(skips)) + x = model.norm_after_skip(x) + projection = model.output_projection1 + x = F.linear(x, projection.weight.squeeze(-1), projection.bias) + x = model.norm_after_proj1(x).relu() + projection = model.output_projection2 + x = F.linear(x, projection.weight.squeeze(-1), projection.bias) + return x.squeeze(-1) + + +@withCUDA +@pytest.mark.parametrize( + ( + "input_channels", + "num_layers", + "features", + "length", + "step_batch", + "strategy_batch", + ), + [ + (2, 1, 3, 5, 1, 2), + (1, 3, 3, 5, 2, 1), + (2, 2, 1, 1, 2, 2), + ], +) +def test_denoiser( + device: torch.device, + input_channels: int, + features: int, + length: int, + num_layers: int, + step_batch: int, + strategy_batch: int, +) -> None: + """Match the channels-last reference across input shapes and strategies.""" + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=num_layers, + num_heads=2, + embedding_channels=12, + input_channels=input_channels, + device=device, + ).eval() + x = torch.randn(2, input_channels, features, length, device=device) + side = torch.randn(2, 6, features, length, device=device) + step = torch.tensor([0, 9], device=device)[:step_batch] + strategy = torch.tensor([0, 1], device=device)[:strategy_batch] + with torch.inference_mode(): + actual = model(x, side, step, strategy) + expected = reference_denoiser(model, x, side, step, strategy) + assert actual.shape == (2, features, length) + assert actual.dtype == x.dtype + assert actual.device == device + assert actual.isfinite().all() + torch.testing.assert_close(actual, expected) + + +@withCUDA +def test_denoiser_grad(device: torch.device) -> None: + """Propagate finite gradients through inputs and all learned parameters.""" + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + device=device, + dtype=torch.float64, + ).eval() + x = torch.randn( + 2, 2, 3, 5, device=device, dtype=torch.float64, requires_grad=True + ) + side = torch.randn( + 2, 6, 3, 5, device=device, dtype=torch.float64, requires_grad=True + ) + step = torch.tensor([0, 9], device=device) + strategy = torch.tensor([0, 1], device=device) + out = model(x, side, step, strategy) + assert out.dtype == torch.float64 + out.square().mean().backward() + for value in (x, side, *model.parameters()): + assert value.grad is not None + assert value.grad.isfinite().all() + for value in (x, side): + assert value.grad is not None + assert value.grad.abs().sum() > 0 diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py new file mode 100644 index 0000000000..2a352602bb --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable +from pathlib import Path + +import pytest +import torch +from torch import Tensor + +from sdm import Stype, TableTensor +from sdm.models import KumoAnomaly +from sdm.models.kumo.timeseries.anomaly.ckpt import model_kwargs, remap_ckpt +from sdm.models.kumo.timeseries.anomaly.model import _KumoAnomaly +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.testing import withCUDA + + +@pytest.fixture +def checkpoint(tmp_path: Path) -> Path: + """Save a tiny upstream-format checkpoint without network access.""" + config = { + "model": { + "target_dim": 3, + "timeemb": 8, + "featureemb": 4, + "is_unconditional": 0, + }, + "diffusion": { + "channels": 8, + "layers": 1, + "nheads": 2, + "diffusion_embedding_dim": 8, + "num_steps": 50, + "beta_start": 0.0001, + "beta_end": 0.01, + "schedule": "linear", + }, + "dataset": {"scale_factor": 1.0, "split": 10}, + } + model = _KumoAnomaly(**model_kwargs(config)) + state = {} + for key, value in model.state_dict().items(): + key = key.replace("denoiser.", "diffmodel.") + key = key.replace("time_layer.", "time_layer.layers.0.") + key = key.replace("feature_layer.", "feature_layer.layers.0.") + key = key.removeprefix("side_embedding.") + state[key] = value + path = tmp_path / "checkpoint.pth" + torch.save({"model": state, "config": config}, path) + return path + + +@withCUDA +@pytest.mark.parametrize("length", [1, 8, 11, 16]) +@pytest.mark.parametrize("features", [1, 3, 5]) +def test_anomaly_output( + checkpoint: Path, device: torch.device, length: int, features: int +) -> None: + """Score real rows/features, including partial windows and PCA input.""" + model = KumoAnomaly(pretrained=checkpoint, window_length=8, device=device) + x = torch.randn(length, features, device=device) + reference = torch.randn(12, features, device=device) + out = model( + x, reference=reference, num_samples=2, num_steps=3, batch_size=1 + ) + count = min(features, 3) + assert out.scores.shape == out.l2_scores.shape == (length,) + assert out.reconstruction.shape == out.target.shape == (length, count) + assert out.scores.device == device + assert out.scores.isfinite().all() + assert not out.scores.requires_grad + error = out.reconstruction.numerical - out.target.numerical + torch.testing.assert_close(out.scores, error.abs().mean(-1)) + torch.testing.assert_close(out.l2_scores, error.norm(dim=-1)) + + +@withCUDA +def test_generator_and_state_dict( + checkpoint: Path, device: torch.device +) -> None: + """A supplied generator reproduces output without consuming global RNG.""" + model = KumoAnomaly(pretrained=checkpoint, window_length=8, device=device) + other = KumoAnomaly(pretrained=checkpoint, window_length=8, device=device) + other.load_state_dict(model.state_dict()) + x = torch.arange(16, device=device, dtype=torch.float32).reshape(8, 2) + generator = torch.Generator(device=device).manual_seed(42) + state = generator.get_state() + cpu_state = torch.random.get_rng_state() + device_state = ( + torch.cuda.get_rng_state(device) if device.type == "cuda" else None + ) + first = model(x, num_samples=2, num_steps=3, generator=generator) + torch.testing.assert_close(torch.random.get_rng_state(), cpu_state) + if device_state is not None: + torch.testing.assert_close( + torch.cuda.get_rng_state(device), device_state + ) + generator.set_state(state) + second = other(x, num_samples=2, num_steps=3, generator=generator) + torch.testing.assert_close(first.scores, second.scores, rtol=0, atol=0) + third = model(x, num_samples=2, num_steps=3, generator=generator) + assert not torch.equal(first.scores, third.scores) + + +def test_checkpoint_strict_loading(checkpoint: Path, tmp_path: Path) -> None: + """Reject missing or unknown weights instead of partially initializing.""" + state = torch.load(checkpoint, weights_only=True) + state["model"]["unknown.weight"] = torch.ones(1) + path = tmp_path / "unexpected.pth" + torch.save(state, path) + with pytest.raises(RuntimeError, match="Unexpected key"): + KumoAnomaly(pretrained=path) + del state["model"]["unknown.weight"] + del state["model"]["embed_layer.weight"] + torch.save(state, path) + with pytest.raises(RuntimeError, match="Missing key"): + KumoAnomaly(pretrained=path) + + +def test_published_loader( + checkpoint: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Read architecture from checkpoint metadata and pin the download.""" + calls = [] + + def download(**kwargs: object) -> Path: + """Record the checkpoint request without accessing the network.""" + calls.append(kwargs) + return checkpoint + + monkeypatch.setattr( + "sdm.models.kumo.timeseries.anomaly.model.download_checkpoint", + download, + ) + model = KumoAnomaly() + assert calls == [ + { + "repo_id": "nvidia/Kumo-Anomaly", + "filename": "final_model.pth", + "revision": "226c5003b0582adc1cac433123f47324bfae39f2", + } + ] + assert model.model.num_features == 3 + assert model.mask_blocks == 4 # Inference default, not training split=10. + saved = torch.load(checkpoint, weights_only=True) + expected = remap_ckpt(saved["model"]) + for key, value in model.model.state_dict().items(): + torch.testing.assert_close(value, expected[key], rtol=0, atol=0) + saved["config"]["model"]["is_unconditional"] = 1 + with pytest.raises(ValueError, match="conditional checkpoint"): + model_kwargs(saved["config"]) + + +def test_reference_schema_and_values(checkpoint: Path) -> None: + """Reference preprocessing preserves names and out-of-range values.""" + model = KumoAnomaly(pretrained=checkpoint, window_length=8) + columns: dict[str, list[str]] = {Stype.numerical: ["a", "b"]} + reference = TableTensor( + columns=columns, numerical=torch.tensor([[0.0, 2.0], [2.0, 4.0]]) + ) + x = TableTensor(columns=columns, numerical=torch.tensor([[4.0, 8.0]])) + result = model(x, reference=reference, num_samples=1, num_steps=1) + torch.testing.assert_close( + result.target.numerical, torch.tensor([[2.0, 3.0]]) + ) + assert result.target.columns == x.columns + reordered = TableTensor( + columns={Stype.numerical: ["b", "a"]}, numerical=reference.numerical + ) + with pytest.raises(ValueError, match="ordered columns"): + model(x, reference=reordered, num_samples=1, num_steps=1) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_nonfinite_input(checkpoint: Path, value: float) -> None: + """Reject missing observations instead of returning misleading scores.""" + model = KumoAnomaly(pretrained=checkpoint) + with pytest.raises(ValueError, match="non-finite"): + model(torch.tensor([[value]]), num_samples=1, num_steps=1) + + +@pytest.mark.parametrize( + ("length", "blocks"), [(7, 4), (12, 3), (0, 4), (8, 0)] +) +def test_mask_partition(length: int, blocks: int) -> None: + """Reject masks that would leave rows unconditioned by both strategies.""" + with pytest.raises(ValueError, match="mask_blocks"): + KumoAnomaly(pretrained=False, window_length=length, mask_blocks=blocks) + + +def test_complementary_reconstruction(monkeypatch: pytest.MonkeyPatch) -> None: + """Only each strategy's held-out half contributes to reconstruction.""" + model = _KumoAnomaly( + num_features=2, + time_channels=8, + feature_channels=4, + channels=8, + num_layers=1, + num_heads=2, + embedding_channels=8, + ).eval() + + class StrategyPrediction(torch.nn.Module): + """Return the strategy index as a directly observable prediction.""" + + def forward( + self, x: Tensor, side_info: Tensor, step: Tensor, strategy: Tensor + ) -> Tensor: + """Fill the prediction with its masking strategy index.""" + return torch.ones_like(x[:, 0]) * strategy + + class DirectSampler(DPMSolver): + """Evaluate once without integrating to expose mask composition.""" + + calls: int = 0 + + def forward( + self, x: Tensor, predict_noise: Callable[[Tensor, Tensor], Tensor] + ) -> Tensor: + """Return the provided denoiser's output.""" + offset = 2 * (self.calls % 2) + self.calls += 1 + return predict_noise(x, torch.zeros(1, dtype=torch.long)) + offset + + monkeypatch.setattr(model, "denoiser", StrategyPrediction()) + x = torch.arange(16, dtype=torch.float32).reshape(1, 8, 2) + sampler = DirectSampler(torch.linspace(0.001, 0.01, 10), num_steps=1) + actual = model(x, sampler=sampler, num_samples=2) + expected = torch.ones_like(x) * x.mean() + # Two samples per strategy differ by 2: use the lower median, not mean. + expected[:, [0, 1, 4, 5]] += x.std() + 1e-5 + torch.testing.assert_close(actual, expected) diff --git a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py new file mode 100644 index 0000000000..c17dada8a4 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -0,0 +1,153 @@ +# 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", ["quad", "linear"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_noise_schedule( + device: torch.device, + schedule: Literal["quad", "linear"], + dtype: torch.dtype, +) -> 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) + 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 >= -5.1).sum()) + assert module.total_N == retained + assert module.T == 1.0 + assert module.schedule == "discrete" + assert (retained < 1000) == (schedule == "quad") + + 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: + """Extrapolate boundary segments and flatten scalar/batched inputs.""" + 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 + ) + 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(), + ) + torch.testing.assert_close( + actual=module.marginal_alpha(t[0, 0]), expected=expected[:1].exp() + ) + empty = torch.empty(0, device=device, dtype=betas.dtype) + assert module.marginal_alpha(empty).shape == (0,) + assert module.inverse_lambda(empty).shape == (0,) + + +@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() + + +@withCUDA +def test_noise_schedule_gradients(device: torch.device) -> None: + """Differentiate time and half-logSNR away from interpolation knots.""" + betas = torch.tensor( + data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64 + ) + module = NoiseScheduleVP(betas) + t = torch.tensor( + data=[0.3, 0.6, 0.9], + device=device, + dtype=betas.dtype, + requires_grad=True, + ) + for method in ( + module.marginal_log_mean_coeff, + module.marginal_alpha, + module.marginal_std, + module.marginal_lambda, + ): + assert torch.autograd.gradcheck(method, (t,)) + lamb = module.marginal_lambda(t).detach().requires_grad_() + assert torch.autograd.gradcheck(module.inverse_lambda, (lamb,)) + + +@pytest.mark.parametrize("betas", [[], [0.1], [0.99999, 0.99999]]) +def test_noise_schedule_too_short(betas: list[float]) -> None: + """Reject schedules without two usable interpolation knots.""" + with pytest.raises(ValueError, match="At least two diffusion steps"): + NoiseScheduleVP(torch.tensor(betas)) diff --git a/test/models/kumo/timeseries/anomaly/test_preprocessing.py b/test/models/kumo/timeseries/anomaly/test_preprocessing.py new file mode 100644 index 0000000000..bc903a0eff --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_preprocessing.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import TableTensor +from sdm.models.kumo.timeseries.anomaly.preprocessing import AnomalyPreprocess +from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.testing import withCUDA + + +@withCUDA +def test_scale_and_pad(device: torch.device) -> None: + """Use reference ranges, retain constants, and append zero padding.""" + reference = TableTensor.from_tensor( + torch.tensor([[1.0, 3.0], [5.0, 3.0]], device=device) + ) + query = TableTensor.from_tensor( + torch.tensor([[9.0, 4.0], [-3.0, 3.0]], device=device) + ) + recipe = default_recipe(num_features=3, scale_factor=2) + recipe.features.fit(reference) + out = recipe.features.transform(query) + expected = torch.tensor([[4.0, 2.0, 0.0], [-2.0, 0.0, 0.0]], device=device) + torch.testing.assert_close(out.numerical, expected) + + +@withCUDA +def test_pca_and_state_dict(device: torch.device) -> None: + """PCA and fitted ranges survive serialization and dtype moves.""" + # Orthogonal columns with distinct variances make component signs unique. + x = torch.tensor( + [ + [3.0, 2.0, 1.0], + [3.0, -2.0, -1.0], + [-3.0, 2.0, -1.0], + [-3.0, -2.0, 1.0], + ], + device=device, + ) + processor = AnomalyPreprocess(num_features=2) + table = TableTensor.from_tensor(x) + actual = processor.fit_transform(table) + expected = ( + x[:, :2] + torch.tensor([3.0, 2.0], device=device) + ) / torch.tensor([6.0, 4.0], device=device) + torch.testing.assert_close(actual.numerical, expected) + restored = AnomalyPreprocess(num_features=2).to(device) + restored.load_state_dict(processor.state_dict()) + assert restored.is_fitted + torch.testing.assert_close(restored(table).numerical, actual.numerical) + restored = restored.double() + torch.testing.assert_close( + restored(table.double()).numerical, expected.double() + ) + + +def test_pca_reference_rows() -> None: + """Do not silently reduce checkpoint width when PCA has too few rows.""" + table = TableTensor.from_tensor(torch.ones(2, 5)) + with pytest.raises(ValueError, match="reference rows"): + AnomalyPreprocess(num_features=3).fit(table) diff --git a/test/models/kumo/timeseries/anomaly/test_solver.py b/test/models/kumo/timeseries/anomaly/test_solver.py new file mode 100644 index 0000000000..3535d9c531 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_solver.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.noise_schedule import NoiseScheduleVP +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize( + ("num_steps", "expected"), + [ + ( + 1, + [ + -3.19437909, + -1.92723203, + -0.73826426, + 0.41330332, + 1.60227108, + 2.86941814, + ], + ), + ( + 2, + [ + -3.04774046, + -1.72219038, + -0.56313372, + 0.39971635, + 1.51524186, + 2.82625222, + ], + ), + ( + 3, + [ + -3.08244777, + -1.77646685, + -0.61630511, + 0.40499809, + 1.53892899, + 2.83434820, + ], + ), + ( + 20, + [ + -3.03812027, + -1.72386146, + -0.57848012, + 0.40655187, + 1.52300930, + 2.82519412, + ], + ), + ], +) +def test_nonlinear_reference( + device: torch.device, num_steps: int, expected: list[float] +) -> None: + """Match the pinned solver with the corrected model adapter.""" + # Generated with Kumo-TS af74798bbf7d020d3298f545fb9be383560a18ce, + # utils/dpm_solver_pytorch.py, DPM_Solver.sample(order=2, + # method="singlestep", skip_type="time_uniform"), using noise() below + # and round(t * total_N - 1) as the continuous-time model adapter. + betas = DiffusionSchedule(500, 0.0001, 0.01, "linear").beta.float() + solver = DPMSolver(betas=betas, num_steps=num_steps).to(device) + x = torch.linspace(-1, 1, 6, device=device).reshape(1, 2, 3) + + def noise(sample: Tensor, step: Tensor) -> Tensor: + """Make both the evolving state and the trained timestep observable.""" + return 0.2 * sample.tanh() + 0.1 * (step / 1000).sin().reshape( + -1, 1, 1 + ) + + torch.testing.assert_close( + actual=solver(x, noise).flatten(), + expected=torch.tensor(expected, device=device), + ) + + +@withCUDA +@pytest.mark.parametrize("num_steps", [1, 2, 3, 10, 20, 21]) +def test_constant_clean_prediction( + device: torch.device, num_steps: int +) -> None: + """Both solver orders integrate a constant clean-data prediction.""" + betas = torch.linspace(0.0001, 0.01, 500) + solver = DPMSolver(betas=betas, num_steps=num_steps).to(device) + x = torch.randn(2, 3, 4, device=device) + clean = torch.full_like(x, 0.25) + calls = [] + coefficients = [] + for i, order in enumerate(solver.orders): + coefficients.extend( + [(solver.alpha[i, j], solver.sigma[i, j]) for j in range(order)] + ) + + def noise(sample: Tensor, step: Tensor) -> Tensor: + """Represent a known clean prediction as noise at each solver time.""" + alpha, sigma = coefficients[len(calls)] + calls.append(step) + return (sample - alpha * clean) / sigma + + actual = solver(x, noise) + schedule = NoiseScheduleVP(betas).to(device) + end = torch.tensor([1 / schedule.total_N], device=device) + alpha_end, sigma_end = ( + schedule.marginal_alpha(end), + schedule.marginal_std(end), + ) + sigma_start = schedule.marginal_std(torch.ones_like(end)) + alpha_start = schedule.marginal_alpha(torch.ones_like(end)) + expected = alpha_end * clean + sigma_end / sigma_start * ( + x - alpha_start * clean + ) + torch.testing.assert_close(actual, expected) + assert len(calls) == num_steps + assert solver.state_dict() == {} + + +def test_timestep_mapping() -> None: + """The default path visits decreasing noise levels instead of all 499.""" + solver = DPMSolver(torch.linspace(0.0001, 0.01, 500)) + steps = [] + + def noise(x: Tensor, step: Tensor) -> Tensor: + """Expose the actual integer indices passed to the denoiser.""" + steps.append(step.clone()) + return torch.zeros_like(x) + + solver(torch.zeros(1, 2, 4), noise) + actual = torch.cat(steps) + expected = (torch.linspace(1, 1 / 500, 21)[:-1] * 500 - 1).round().long() + torch.testing.assert_close(actual, expected) + assert actual[0] == 499 + assert actual[-1] == 25 + assert (actual.diff() < 0).all() + + +def test_truncated_schedule_indices() -> None: + """Clipped schedules keep the training indices of retained knots.""" + betas = torch.linspace(0.001, 0.2, 1000) + schedule = NoiseScheduleVP(betas) + assert schedule.total_N < len(betas) + solver = DPMSolver(betas, num_steps=3) + assert solver.steps[0, 0] == schedule.total_N - 1 + assert solver.steps[-1, -1] == 0 + + +def test_solver_nonempty() -> None: + """A zero-evaluation solver must not return unmodified noise as data.""" + with pytest.raises(ValueError, match="num_steps"): + DPMSolver(torch.linspace(0.0001, 0.01, 500), num_steps=0) From ff766f186c95620b6f9fb80600ee45acd14b239b Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Fri, 16 Oct 2026 16:00:26 -0700 Subject: [PATCH 3/3] Add Kumo-Anomaly SCS and MACS thresholding Signed-off-by: Aditi Gautam --- docs/source/api/models.rst | 19 ++ examples/kumo_anomaly.py | 15 +- .../kumo/timeseries/anomaly/__init__.py | 6 +- sdm/models/kumo/timeseries/anomaly/model.py | 7 +- .../kumo/timeseries/anomaly/threshold.py | 243 ++++++++++++++++++ .../kumo/timeseries/anomaly/test_threshold.py | 207 +++++++++++++++ 6 files changed, 491 insertions(+), 6 deletions(-) create mode 100644 sdm/models/kumo/timeseries/anomaly/threshold.py create mode 100644 test/models/kumo/timeseries/anomaly/test_threshold.py diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 9b5688055b..f6483d2b2d 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -67,8 +67,27 @@ Anomaly Detection Unlike supervised ICL models, Kumo-Anomaly reconstructs an ordered series without labels. Its scores are reconstruction errors, not probabilities. +Apply ``SCSThreshold`` to the returned ``scores`` and ``target``, or +``MACSThreshold`` to ``scores``, to obtain boolean anomaly flags. Thresholding +is separate from reconstruction so the same scores can be evaluated with +either decision rule without rerunning the model. + +Both rules use the scored series itself, including a whole-series percentile +filter. Reference-only feature preprocessing does **not** make these thresholds +reference-fitted or causal. Defaults reproduce the released SDK's APCA-based +SCS and MACS paths; optional K-means segmentation is not included. MACS retains +the upstream expanding bounds shared by all three scales, rather than changing +them to rolling bounds. These heuristic settings do not provide calibrated +false-positive guarantees. The corrected DPM timestep mapping can change +scores and flags relative to the unmodified SDK. .. autoclass:: sdm.models.KumoAnomaly :members: .. autoclass:: sdm.models.kumo.timeseries.anomaly.AnomalyOutput + +.. autoclass:: sdm.models.kumo.timeseries.anomaly.SCSThreshold + :members: + +.. autoclass:: sdm.models.kumo.timeseries.anomaly.MACSThreshold + :members: diff --git a/examples/kumo_anomaly.py b/examples/kumo_anomaly.py index 0d74955207..ba08c60bfe 100644 --- a/examples/kumo_anomaly.py +++ b/examples/kumo_anomaly.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Score an ordered series using reference-only feature preprocessing.""" +"""Score and flag anomalies with reference-only feature preprocessing.""" import torch import sdm from sdm.models import KumoAnomaly +from sdm.models.kumo.timeseries.anomaly import MACSThreshold, SCSThreshold device = torch.device("cuda" if torch.cuda.is_available() else "cpu") time = torch.arange(600, device=device, dtype=torch.float32) @@ -26,6 +27,12 @@ generator=torch.Generator(device=device).manual_seed(42), ) print(result.scores) -print(result.reconstruction) -# Scores and reconstructions use scaled model-space units. Thresholds need -# separate validation; the corrected DPM timestep mapping changes SDK scores. +# These SDK rules use the scored query itself, including its global score +# percentile. Reference-only feature scaling does not make them causal or +# reference-calibrated. The corrected solver can change scores and decisions. +anomalies_scs = SCSThreshold()(result.scores, result.target) +anomalies_macs = MACSThreshold()(result.scores) +print("SCS anomaly rows:", anomalies_scs.nonzero().flatten() + 300) +print("MACS anomaly rows:", anomalies_macs.nonzero().flatten() + 300) +# The booleans are detector decisions, not ground truth or probabilities. +# Validate detection quality and threshold settings on your own labeled data. diff --git a/sdm/models/kumo/timeseries/anomaly/__init__.py b/sdm/models/kumo/timeseries/anomaly/__init__.py index 7928e4bf29..ae3a1b4dd6 100644 --- a/sdm/models/kumo/timeseries/anomaly/__init__.py +++ b/sdm/models/kumo/timeseries/anomaly/__init__.py @@ -4,6 +4,10 @@ """Kumo time-series anomaly detection model components.""" from sdm.models.kumo.timeseries.anomaly.model import KumoAnomaly, AnomalyOutput +from sdm.models.kumo.timeseries.anomaly.threshold import ( + SCSThreshold, + MACSThreshold, +) -__all__ = ["KumoAnomaly", "AnomalyOutput"] +__all__ = ["KumoAnomaly", "AnomalyOutput", "SCSThreshold", "MACSThreshold"] diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 2c5b054702..3ad5796b5c 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -31,6 +31,10 @@ class AnomalyOutput(NamedTuple): Padding features are excluded from every field. Higher scores indicate larger reconstruction errors, not calibrated anomaly probabilities. + Apply :class:`~sdm.models.kumo.timeseries.anomaly.SCSThreshold` to + ``scores`` and ``target``, or + :class:`~sdm.models.kumo.timeseries.anomaly.MACSThreshold` to ``scores``, + to obtain boolean anomaly flags with the SDK's offline decision rules. """ scores: Tensor @@ -147,7 +151,8 @@ def forward( input units. Short/final partial windows repeat their last row and are trimmed afterward; unlike the original SDK, no trailing observations are dropped. Padding can affect that window's - normalization and reconstruction. No anomaly threshold is fitted. + normalization and reconstruction. Thresholding is a separate step: + apply ``SCSThreshold`` or ``MACSThreshold`` to the returned scores. """ if self.training: raise RuntimeError("Call 'eval()' before anomaly inference") diff --git a/sdm/models/kumo/timeseries/anomaly/threshold.py b/sdm/models/kumo/timeseries/anomaly/threshold.py new file mode 100644 index 0000000000..37c173b549 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/threshold.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Offline thresholding ported from Kumo-TS's adaptive_threshold.py. + +Source: NVIDIA/Kumo-TS at af74798bbf7d020d3298f545fb9be383560a18ce. +These are the released SDK's decision rules, not calibrated probabilities +or guarantees about false-positive rates. +""" + +import torch +from torch import Tensor + +from sdm import TableTensor + + +def _scores(scores: Tensor) -> Tensor: + """Normalize finite per-row scores at the thresholding API boundary.""" + if scores.ndim != 1 or not scores.isfinite().all(): + raise ValueError("Expected finite scores with shape [time]") + return scores.double() + + +def _bounds( + scores: Tensor, starts: Tensor, confidence_level: float +) -> tuple[Tensor, Tensor]: + """Compute the SDK's expanding bounds, restarting at segment starts.""" + sums = torch.cat([scores.new_zeros(1), scores.cumsum(0)]) + squares = torch.cat([scores.new_zeros(1), scores.square().cumsum(0)]) + counts = torch.arange(1, scores.numel() + 1, device=scores.device) - starts + mean = (sums[1:] - sums[starts]) / counts + mean = torch.where(counts == 1, scores, mean) + variance = (squares[1:] - squares[starts]) / counts - mean.square() + std = variance.clamp_min(0).sqrt() + width = torch.where(std > 0, 1.5 * std, 0.1) + if confidence_level > 0.95: + width = width * 1.2 + elif confidence_level < 0.90: + width = width * 0.8 + # The first score in each segment always has unscaled +/- 0.1 bounds. + width = torch.where(counts == 1, 0.1, width) + return mean - width, mean + width + + +def _rolling_stats(scores: Tensor, window: int) -> tuple[Tensor, Tensor]: + """Return trailing means and sample variances with min_periods=1.""" + # Center before prefix sums to reduce cancellation for offset scores. + centered = scores - scores[0] + sums = torch.cat([scores.new_zeros(1), centered.cumsum(0)]) + squares = torch.cat([scores.new_zeros(1), centered.square().cumsum(0)]) + end = torch.arange(1, scores.numel() + 1, device=scores.device) + start = (end - window).clamp_min(0) + count = end - start + total = sums[end] - sums[start] + mean = total / count + scores[0] + variance = squares[end] - squares[start] - total.square() / count + variance = variance.clamp_min(0) / (count - 1).clamp_min(1) + return mean, torch.where(count > 1, variance, 0) + + +class SCSThreshold(torch.nn.Module): + """Flag anomalies using the SDK's APCA-based SCS decision rule. + + Segment the model-space target using Adaptive Piecewise Constant + Approximation (APCA), compute expanding score bounds per segment, then + retain only scores strictly above the whole-series percentile cutoff. + Defaults match ``SCSThresholdStrategy`` in the released SDK. + + Thresholds use the scored series itself, including future observations. + Passing reference rows to the model does not make this thresholding + reference-fitted or causal. ``confidence_level`` selects the SDK's + heuristic width multiplier, not a calibrated confidence guarantee. + + APCA makes sequential, data-dependent split decisions on CPU. At the + API boundary, the feature-mean series is copied to CPU once; score + bounds and returned flags are computed on the scores' device. + + Args: + confidence_level: Width setting; the SDK default is 0.99. + min_segment_size: Minimum observations on each side of an APCA split. + percentile_filter: Global score percentile used as a strict cutoff. + """ + + def __init__( + self, + confidence_level: float = 0.99, + min_segment_size: int = 50, + percentile_filter: float = 95.0, + ) -> None: + super().__init__() + if min_segment_size < 1: + raise ValueError("'min_segment_size' must be positive") + self.confidence_level = confidence_level + self.min_segment_size = min_segment_size + self.percentile_filter = percentile_filter + + def _segment_starts(self, data: Tensor) -> Tensor: + """Return each row's segment start using the published APCA rule.""" + length = data.numel() + rows = torch.arange(length, device=data.device) + variation = data.std(correction=0) / (data.mean() + 1e-8) + # Preserve the signed coefficient of variation used by the SDK. + if variation < 0.1: + size = max(200, length // 15) + return rows // size * size + + starts = data.new_zeros(length, dtype=torch.long) + sums = torch.cat([data.new_zeros(1), data.cumsum(0)]) + squares = torch.cat([data.new_zeros(1), data.square().cumsum(0)]) + start = 0 + while start < length: + starts[start:] = start + if length - start < 2 * self.min_segment_size: + break + first = start + self.min_segment_size + last = length - self.min_segment_size + 1 + step = max(1, (last - first) // 200) + splits = torch.arange(first, last, step, device=data.device) + left_sum = sums[splits] - sums[start] + right_sum = sums[-1] - sums[splits] + error = ( + squares[-1] + - squares[start] + - left_sum.square() / (splits - start) + - right_sum.square() / (length - splits) + ) + best = error.argmin() + no_split = ( + squares[-1] + - squares[start] + - (sums[-1] - sums[start]).square() / (length - start) + ) + improvement = 0.7 if variation > 0.3 else 0.5 + if error[best] >= no_split * improvement: + break + start = int(splits[best]) + return starts + + @torch.no_grad() + def forward(self, scores: Tensor, target: Tensor | TableTensor) -> Tensor: + """Return boolean anomaly flags with shape ``[time]``. + + Pass ``result.scores`` and ``result.target`` from :class:`KumoAnomaly`. + The target must contain only real, preprocessed features, not padding, + timestamps or labels. Its shape must be ``[time, features]``. + """ + scores = _scores(scores) + data = target.numerical if isinstance(target, TableTensor) else target + if ( + data.ndim != 2 + or data.size(0) != scores.numel() + or not data.size(1) + ): + raise ValueError( + "Target must have shape [time, features] matching scores" + ) + if not data.isfinite().all(): + raise ValueError("Expected finite target values") + if not scores.numel(): + return scores.bool() + data = data.mean(-1).to(device="cpu", dtype=torch.float64) + starts = self._segment_starts(data).to(scores.device) + lower, upper = _bounds(scores, starts, self.confidence_level) + cutoff = scores.quantile(self.percentile_filter / 100) + return ((scores < lower) | (scores > upper)) & (scores > cutoff) + + +class MACSThreshold(torch.nn.Module): + """Flag anomalies using the released SDK's MACS decision rule. + + Combine expanding score bounds with variance-based weights and a + regime-change check, then retain only scores strictly above the global + percentile cutoff. Defaults match ``MACSThresholdStrategy`` in Kumo-TS. + This is offline, transductive thresholding, not reference-only calibration. + + The released implementation uses the same expanding bounds at all three + scales: its bound calculator ignores the window argument. This port + preserves that behavior, including weighting and regime decisions, rather + than introducing a different rolling-bound algorithm. The confidence + setting is a heuristic width multiplier, not a false-positive guarantee. + + Args: + short_window: Window used for variance-based attention weights. + medium_window: Retained SDK scale setting; bounds remain expanding. + long_window: Window used to identify regime changes. + confidence_level: Width setting; the SDK default is 0.95. + percentile_filter: Global score percentile used as a strict cutoff. + """ + + def __init__( + self, + short_window: int = 10, + medium_window: int = 50, + long_window: int = 250, + confidence_level: float = 0.95, + percentile_filter: float = 97.0, + ) -> None: + super().__init__() + if long_window < 1: + raise ValueError("'long_window' must be positive") + self.short_window = short_window + self.medium_window = medium_window + self.long_window = long_window + self.confidence_level = confidence_level + self.percentile_filter = percentile_filter + + @torch.no_grad() + def forward(self, scores: Tensor) -> Tensor: + """Return boolean anomaly flags from finite MAE scores ``[time]``.""" + scores = _scores(scores) + if not scores.numel(): + return scores.bool() + starts = torch.zeros_like(scores, dtype=torch.long) + lower, upper = _bounds(scores, starts, self.confidence_level) + window = max(2, min(self.short_window, scores.numel() // 10)) + _, variance = _rolling_stats(scores, window) + maximum = variance.max() + normalized = variance / torch.where(maximum > 0, maximum, 1) + group = torch.where(normalized > 0.7, 2, (normalized > 0.3).long()) + weights = scores.new_tensor( + [[0.1, 0.3, 0.6], [0.2, 0.6, 0.2], [0.6, 0.3, 0.1]] + )[group] + short_weight, medium_weight, long_weight = weights.unbind(-1) + # Keep the SDK's weighted arithmetic even though the bounds coincide. + combined_lower = ( + short_weight * lower + medium_weight * lower + long_weight * lower + ) + combined_upper = ( + short_weight * upper + medium_weight * upper + long_weight * upper + ) + mean, variance = _rolling_stats(scores, self.long_window) + std = variance.sqrt() + previous_mean = torch.cat([scores.new_zeros(1), mean[:-1]]) + previous_std = torch.cat([scores.new_zeros(1), std[:-1]]) + denominator = previous_std + 1e-8 + regime = ((mean - previous_mean).abs() / denominator > 2) | ( + (std - previous_std).abs() / denominator > 1.5 + ) + violations = (scores < lower) | (scores > upper) + attention = (scores < combined_lower) | (scores > combined_upper) + anomalies = attention & (~regime | violations) + cutoff = scores.quantile(self.percentile_filter / 100) + return anomalies & (scores > cutoff) diff --git a/test/models/kumo/timeseries/anomaly/test_threshold.py b/test/models/kumo/timeseries/anomaly/test_threshold.py new file mode 100644 index 0000000000..5fa9b33426 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_threshold.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch import Tensor + +from sdm import TableTensor +from sdm.models.kumo.timeseries.anomaly import MACSThreshold, SCSThreshold +from sdm.testing import withCUDA + + +def _series( + kind: str, device: torch.device, dtype: torch.dtype +) -> tuple[Tensor, Tensor]: + """Make fixed scores and targets used for upstream decision fixtures.""" + time = torch.arange(600, device=device, dtype=dtype) + target = torch.stack( + [1 + 0.1 * (time / 20).sin(), 1 + 0.1 * (time / 17).cos()], dim=-1 + ) + scores = 0.1 + 0.02 * (time / 7).sin().square() + if kind == "regimes": + target[200:400] += 9 + target[400:] += 2 + scores[200:400] += 0.4 + scores[400:] += 0.1 + elif kind == "oscillating": + target -= 1 + scores[[0, 50, 151, 199, 200, 201, 299, 398, 400, 450, 599]] += 2 + return scores, target + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("kind", ["flat", "regimes", "oscillating"]) +@pytest.mark.parametrize("as_table", [False, True]) +def test_threshold_upstream_decisions( + device: torch.device, dtype: torch.dtype, kind: str, as_table: bool +) -> None: + """Match flags from pinned Kumo-TS, including segment-boundary behavior.""" + # Expected indices generated by the unmodified threshold algorithms in + # NVIDIA/Kumo-TS af74798, using the SDK wrappers' SCS/MACS defaults. + scores, target = _series(kind, device, dtype) + all_spikes = [50, 151, 199, 200, 201, 299, 398, 400, 450, 599] + scs_spikes = ( + all_spikes + if kind == "oscillating" + else [50, 151, 199, 299, 398, 450, 599] + ) + scs = SCSThreshold()( + scores, TableTensor.from_tensor(target) if as_table else target + ) + macs = MACSThreshold()(scores) + for actual, indices in ((scs, scs_spikes), (macs, all_spikes)): + expected = torch.zeros(600, device=device, dtype=torch.bool) + expected[indices] = True + assert actual.device == device + assert actual.dtype == torch.bool + assert not actual.requires_grad + torch.testing.assert_close(actual, expected) + + +@withCUDA +@pytest.mark.parametrize("length", [0, 1, 2, 37, 200, 401]) +@pytest.mark.parametrize("value", [0.0, 1.0, 1e7]) +def test_constant_scores( + device: torch.device, length: int, value: float +) -> None: + """Strict global percentile filtering never flags constant scores.""" + scores = torch.full((length,), value, device=device) + target = torch.ones(length, 2, device=device) + for flags in (SCSThreshold()(scores, target), MACSThreshold()(scores)): + assert flags.shape == (length,) + assert flags.device == device + assert not flags.any() + + +@withCUDA +@pytest.mark.parametrize( + ("confidence", "first"), [(0.85, 2), (0.9, 7), (0.95, 7), (0.99, 20)] +) +def test_confidence_width( + device: torch.device, confidence: float, first: int +) -> None: + """Retain the SDK's discrete width multipliers, not a quantile model.""" + scores = torch.arange(20, device=device, dtype=torch.float64) + target = torch.ones(20, 2, device=device) + expected = scores >= first + scs = SCSThreshold(confidence_level=confidence, percentile_filter=0) + macs = MACSThreshold(confidence_level=confidence, percentile_filter=0) + torch.testing.assert_close(scs(scores, target), expected) + torch.testing.assert_close(macs(scores), expected) + + +@withCUDA +@pytest.mark.parametrize( + ("confidence", "percentile", "minimum", "indices"), + [ + (0.99, 95, 50, [50, 149, 199, 200, 201, 299, 374, 398, 400, 450, 599]), + (0.85, 95, 25, [50, 199, 200, 201, 299, 398, 400, 450, 599]), + (0.99, 99, 50, [199, 200, 201, 299]), + ], +) +def test_apca_regime_changes( + device: torch.device, + confidence: float, + percentile: float, + minimum: int, + indices: list[int], +) -> None: + """Match upstream splits away from the flat-data 200-row boundaries.""" + time = torch.arange(600, device=device, dtype=torch.float64) + target = torch.stack( + [1 + 0.1 * (time / 20).sin(), 1 + 0.1 * (time / 17).cos()], dim=-1 + ) + target[150:375] += 9 + target[375:] += 2 + scores = 0.1 + 0.02 * (time / 7).sin().square() + scores[150:375] += 0.4 + scores[375:] += 0.1 + scores[ + [ + 0, + 50, + 149, + 150, + 151, + 199, + 200, + 201, + 299, + 374, + 375, + 376, + 398, + 400, + 450, + 599, + ] + ] += 2 + expected = torch.zeros(600, device=device, dtype=torch.bool) + expected[indices] = True + threshold = SCSThreshold( + confidence_level=confidence, + percentile_filter=percentile, + min_segment_size=minimum, + ) + torch.testing.assert_close(threshold(scores, target), expected) + + +@withCUDA +def test_default_percentile_filters(device: torch.device) -> None: + """SCS and MACS retain their distinct 95th and 97th percentile filters.""" + scores = torch.arange(100, device=device, dtype=torch.float64) + scs = SCSThreshold(confidence_level=0.95) + torch.testing.assert_close( + scs(scores, torch.ones(100, 1, device=device)), scores >= 95 + ) + torch.testing.assert_close(MACSThreshold()(scores), scores >= 97) + + +@withCUDA +def test_strict_percentile_ties(device: torch.device) -> None: + """Scores equal to the percentile cutoff are not flagged as anomalies.""" + scores = torch.ones(100, device=device) + scores[20:40] = 10 + target = torch.ones(100, 2, device=device) + assert not SCSThreshold()(scores, target).any() + assert not MACSThreshold()(scores).any() + assert SCSThreshold(percentile_filter=70)(scores, target).any() + assert MACSThreshold(percentile_filter=70)(scores).any() + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_nonfinite_scores(value: float) -> None: + """Reject invalid scores rather than silently marking everything normal.""" + scores = torch.tensor([1.0, value]) + with pytest.raises(ValueError, match="finite scores"): + SCSThreshold()(scores, torch.ones(2, 1)) + with pytest.raises(ValueError, match="finite scores"): + MACSThreshold()(scores) + with pytest.raises(ValueError, match="finite target"): + SCSThreshold()(torch.ones(2), scores[:, None]) + + +@pytest.mark.parametrize("shape", [(4,), (3, 2), (4, 0)]) +def test_target_alignment(shape: tuple[int, ...]) -> None: + """Reject targets whose rows or features cannot align with the scores.""" + with pytest.raises(ValueError, match="matching scores"): + SCSThreshold()(torch.ones(4), torch.ones(shape)) + + +def test_score_shape() -> None: + """Do not silently threshold across batch or feature dimensions.""" + scores = torch.ones(4, 1) + with pytest.raises(ValueError, match="shape"): + SCSThreshold()(scores, torch.ones(4, 2)) + with pytest.raises(ValueError, match="shape"): + MACSThreshold()(scores) + + +def test_segment_size() -> None: + """Require an APCA split to advance by at least one observation.""" + with pytest.raises(ValueError, match="min_segment_size"): + SCSThreshold(min_segment_size=0) + with pytest.raises(ValueError, match="long_window"): + MACSThreshold(long_window=0)