diff --git a/examples/timeseries/anomaly.py b/examples/timeseries/anomaly.py index 047692301..23d4b81c9 100644 --- a/examples/timeseries/anomaly.py +++ b/examples/timeseries/anomaly.py @@ -9,10 +9,19 @@ time = torch.arange(200, device=device, dtype=torch.float32) x = torch.stack([torch.sin(time / 12), torch.cos(time / 20)], dim=-1) -x[140:149] += 3 # Add anomalies. +x[140:150] += 3 # Add anomalies. model = sdm.models.KumoAnomaly(device=device) +out = model(x) -# Higher reconstruction error scores indicate more likely anomalies: -score = model(x).numerical.squeeze(-1) -print(score[131:158]) +# Higher MAE scores indicate more likely anomalies: +print("Reconstruction Error (MAE):") +print(out["mae"].numerical[130:160, 0]) + +# SCS flags anomalies using adaptive segments of the preprocessed series: +print("Anomaly Labels (SCS):") +print(out["scs"].categorical.decode()[130:160, 0]) + +# MACS flags anomalies using variance-weighted score bounds: +print("Anomaly Labels (MACS):") +print(out["macs"].categorical.decode()[130:160, 0]) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index f15f06a94..e2d669e17 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -14,7 +14,7 @@ import torch.nn.functional as F from torch import Tensor -from sdm import Recipe, Stype, TableTensor +from sdm import CategoricalTensor, Recipe, Stype, TableTensor from sdm._inference import inference_mode from sdm.models import Callback, Model from sdm.models._huggingface import download_checkpoint @@ -24,6 +24,10 @@ 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 +from sdm.models.kumo.timeseries.anomaly.threshold import ( + macs_threshold, + scs_threshold, +) from sdm.processing.execution import MemberQuery, RecipeExecution @@ -99,7 +103,7 @@ def forward( num_steps: int = 20, callbacks: Sequence[Callback] | None = None, generator: torch.Generator | None = None, - ) -> TableTensor: # [..., R, 1] + ) -> TableTensor: # [..., R, 3] r"""The forward pass. Args: @@ -116,9 +120,13 @@ def forward( pre-processing and model execution. Returns: - The anomaly scores with shape ``[..., R, 1]``, denoting the mean - absolute reconstruction error across preprocessed features at each - time step. + A table with shape ``[..., R, 3]`` containing numerical ``mae`` + reconstruction scores and categorical boolean ``scs`` and ``macs`` + anomaly labels at each time step. SCS (Segmented Confidence + Sequences) segments the preprocessed series and thresholds scores + using expanding bounds within each segment. MACS (Multi-Scale + Adaptive Confidence Segments) uses variance-weighted expanding + score bounds. """ callbacks = () if callbacks is None else callbacks requires_grad = self.training @@ -158,9 +166,19 @@ def forward( generator=generator, ) + scs = scs_threshold(out, query.x.numerical).long() + macs = macs_threshold(out).long() + categories = torch.tensor([False, True], device=out.device) out = TableTensor( - columns={Stype.numerical: ("mae",)}, + columns={ + Stype.numerical: ("mae",), + Stype.categorical: ("scs", "macs"), + }, numerical=out.unsqueeze(-1), + categorical=CategoricalTensor( + code=torch.stack([scs, macs], dim=-1), + categories=(categories, categories), + ), ) for callback in callbacks: diff --git a/sdm/models/kumo/timeseries/anomaly/threshold.py b/sdm/models/kumo/timeseries/anomaly/threshold.py new file mode 100644 index 000000000..072e19ed3 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/threshold.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from NVIDIA/Kumo-TS, Kumo-Anomaly/utils/adaptive_threshold.py +# at af74798bbf7d020d3298f545fb9be383560a18ce. + +import torch +from torch import Tensor + +from sdm import TableTensor + + +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( + [torch.zeros_like(scores[..., :1]), scores.cumsum(-1)], dim=-1 + ) + squares = torch.cat( + [torch.zeros_like(scores[..., :1]), scores.square().cumsum(-1)], dim=-1 + ) + counts = ( + torch.arange(1, scores.size(-1) + 1, device=scores.device) - starts + ) + mean = (sums[..., 1:] - sums.gather(-1, starts)) / counts + mean = torch.where(counts == 1, scores, mean) + variance = ( + squares[..., 1:] - squares.gather(-1, 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[..., :1] + sums = torch.cat( + [torch.zeros_like(scores[..., :1]), centered.cumsum(-1)], dim=-1 + ) + squares = torch.cat( + [torch.zeros_like(scores[..., :1]), centered.square().cumsum(-1)], + dim=-1, + ) + end = torch.arange(1, scores.size(-1) + 1, device=scores.device) + start = (end - window).clamp_min(0) + count = end - start + total = sums[..., end] - sums[..., start] + mean = total / count + scores[..., :1] + 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) + + +def _segment_starts(data: Tensor, min_segment_size: int) -> 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 * min_segment_size: + break + first = start + min_segment_size + last = length - min_segment_size + if last <= first: + break + 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 scs_threshold( + scores: Tensor, + target: Tensor | TableTensor, + *, + confidence_level: float = 0.99, + min_segment_size: int = 50, + percentile_filter: float = 95.0, +) -> Tensor: + """Flag anomalies using Segmented Confidence Sequences (SCS). + + Segment preprocessed observations with APCA, compute expanding score + bounds per segment, and retain violations above a whole-series percentile. + Defaults follow the Kumo-Anomaly SDK. Decisions use future observations; + the confidence setting is a heuristic, not a calibrated guarantee. + Segmentation copies feature means to CPU; flags stay on the score device. + + Args: + scores: Finite reconstruction scores with shape ``[..., R]``. + target: Finite preprocessed observations of shape ``[..., R, D]`` + without padded features. ``R`` is the number of time steps and + ``D`` the number of features. + confidence_level: Bound width setting. + min_segment_size: Minimum observations on each side of an APCA split. + percentile_filter: Global score percentile used as a strict cutoff. + + Returns: + Boolean anomaly flags with shape ``[..., R]``. + """ + scores = scores.double() + data = target.numerical if isinstance(target, TableTensor) else target + if data.shape[:-1] != scores.shape: + raise ValueError("Target batch and row dimensions must match scores") + if scores.numel() == 0: + return scores.bool() + data = data.mean(-1).to(device="cpu", dtype=torch.float64) + starts = ( + torch.stack( + [ + _segment_starts(series, min_segment_size) + for series in data.reshape(-1, data.size(-1)) + ] + ) + .reshape(scores.shape) + .to(scores.device) + ) + lower, upper = _bounds(scores, starts, confidence_level) + cutoff = scores.quantile(percentile_filter / 100, dim=-1, keepdim=True) + return ((scores < lower) | (scores > upper)) & (scores > cutoff) + + +@torch.no_grad() +def macs_threshold( + scores: Tensor, + *, + short_window: int = 10, + medium_window: int = 50, + long_window: int = 250, + confidence_level: float = 0.95, + percentile_filter: float = 97.0, +) -> Tensor: + """Flag anomalies using Multi-Scale Adaptive Confidence Segments (MACS). + + Combine expanding score bounds with variance-based weights and a + regime-change check, then retain violations above a whole-series + percentile. Defaults follow the SDK, including its shared expanding bounds + across scales. Decisions use future observations; the confidence setting + is a heuristic, not a calibrated guarantee. + + Args: + scores: Finite reconstruction scores with shape ``[..., R]``, + where ``R`` is the number of time steps. + short_window: Window used for variance-based attention weights. + medium_window: SDK scale setting; bounds remain expanding. + long_window: Window used to identify regime changes. + confidence_level: Bound width setting. + percentile_filter: Global score percentile used as a strict cutoff. + + Returns: + Boolean anomaly flags with shape ``[..., R]``. + """ + scores = scores.double() + if scores.numel() == 0: + return scores.bool() + starts = torch.zeros_like(scores, dtype=torch.long) + lower, upper = _bounds(scores, starts, confidence_level) + window = max(2, min(short_window, scores.size(-1) // 10)) + _, variance = _rolling_stats(scores, window) + maximum = variance.amax(dim=-1, keepdim=True) + 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, long_window) + std = variance.sqrt() + previous_mean = torch.cat( + [torch.zeros_like(scores[..., :1]), mean[..., :-1]], dim=-1 + ) + previous_std = torch.cat( + [torch.zeros_like(scores[..., :1]), std[..., :-1]], dim=-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(percentile_filter / 100, dim=-1, keepdim=True) + return anomalies & (scores > cutoff) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py index 849dd2874..6e1c6c71e 100644 --- a/test/models/kumo/timeseries/anomaly/test_model.py +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -13,8 +13,13 @@ def test_forward(device: torch.device) -> None: x = torch.randn(2, 130, 3, device=device) out = model(x, num_samples=1, num_steps=1) - assert out.size() == (2, 130, 1) + assert out.size() == (2, 130, 3) assert out.columns["numerical"] == ("mae",) + assert out.columns["categorical"] == ("scs", "macs") assert out.device == device assert out.numerical.isfinite().all() assert (out.numerical >= 0).all() + for categories in out.categorical.categories: + torch.testing.assert_close( + categories, torch.tensor([False, True], device=device) + ) 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 000000000..700dc5d30 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_threshold.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm.models.kumo.timeseries.anomaly.threshold import ( + macs_threshold, + scs_threshold, +) +from sdm.testing import withCUDA + + +@withCUDA +def test_threshold(device: torch.device) -> None: + scores = torch.full((2, 400), 0.1, device=device) + scores[0, [0, 50, 200, 350]] = 2 + target = torch.ones(2, 400, 1, device=device) + + expected = torch.zeros_like(scores, dtype=torch.bool) + expected[0, [50, 350]] = True + torch.testing.assert_close(scs_threshold(scores, target), expected) + + expected[0, 200] = True + torch.testing.assert_close(macs_threshold(scores), expected)