Upload ts-anomaly-benchmark/models/moment_detector.py

ts-anomaly-benchmark/models/moment_detector.py

@@ -0,0 +1,214 @@
+"""
+models/moment_detector.py — MOMENT foundation model wrapper
+=============================================================
+
+MOMENT uses a reconstruction-based approach:
+1. Feed your time series into the pre-trained model.
+2. The model reconstructs what it thinks the series "should" look like.
+3. The reconstruction error at each timestep is the anomaly score.
+
+High reconstruction error → the model couldn't predict that pattern
+from its learned baseline → likely anomalous.
+
+Reference: https://github.com/moment-timeseries-foundation-model/moment
+"""
+
+import numpy as np
+import torch
+from models.base import AnomalyDetector
+
+
+class MOMENTDetector(AnomalyDetector):
+    """
+    Wraps the MOMENT foundation model for anomaly detection
+    via reconstruction error scoring.
+    """
+
+    def __init__(self, config: dict):
+        super().__init__(name="MOMENT", config=config)
+        self.model = None
+        self._model_loaded = False
+
+    def _load_model(self):
+        """
+        Lazy-load the MOMENT model on first use.
+
+        Why lazy? The model is large (~1GB+). We only want to load it
+        if it's actually going to be used in this benchmark run.
+        """
+        if self._model_loaded:
+            return
+
+        try:
+            from momentfm import MOMENTPipeline
+
+            pretrained = self.config.get("pretrained", "AutonLab/MOMENT-1-large")
+            task = self.config.get("task", "reconstruction")
+
+            print(f"  Loading MOMENT from '{pretrained}'...")
+
+            # ── Step 1: Create the pipeline ──
+            # MOMENTPipeline.from_pretrained() downloads the model weights
+            # from Hugging Face and sets up the architecture.
+            # model_kwargs tells it which task head to use.
+            self.model = MOMENTPipeline.from_pretrained(
+                pretrained,
+                model_kwargs={"task_name": task},
+            )
+
+            # ── Step 2: Initialise the model ──
+            # .init() prepares internal buffers and moves to the
+            # appropriate device (GPU if available, else CPU).
+            self.model.init()
+
+            self._model_loaded = True
+            print(f"  MOMENT loaded successfully.")
+
+        except ImportError:
+            raise ImportError(
+                "momentfm is not installed. Install it with:\n"
+                "  pip install momentfm\n"
+                "Then re-run the benchmark."
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MOMENT model: {e}")
+
+    def detect(self, series: np.ndarray) -> np.ndarray:
+        """
+        Run MOMENT anomaly detection on a 1-D time series.
+
+        Parameters
+        ----------
+        series : np.ndarray
+            1-D array of float values (your time-series data).
+
+        Returns
+        -------
+        scores : np.ndarray
+            Normalised anomaly scores in [0, 1], same length as input.
+
+        How it works (step by step)
+        ---------------------------
+        1. We reshape your 1-D array into the tensor shape MOMENT expects:
+           (batch_size=1, n_channels=1, sequence_length).
+        2. We pass it through the model, which returns a reconstruction.
+        3. We compute the absolute difference between the original
+           and the reconstruction at every timestep.
+        4. We normalise those errors to [0, 1].
+        """
+        self._load_model()
+
+        # ── Step 1: Prepare the input tensor ──
+        # MOMENT expects shape: (batch, channels, sequence_length)
+        # For univariate data: batch=1, channels=1
+        input_tensor = torch.tensor(
+            series, dtype=torch.float32
+        ).unsqueeze(0).unsqueeze(0)
+        # Shape is now: (1, 1, len(series))
+
+        # ── Handle sequence length requirements ──
+        # MOMENT has a fixed input length (typically 512).
+        # If your series is longer, we process it in sliding windows
+        # and stitch the results together.
+        model_seq_len = 512  # MOMENT's expected input length
+
+        if len(series) <= model_seq_len:
+            # Short enough — process in one go
+            scores = self._detect_single(series, model_seq_len)
+        else:
+            # Longer series — use sliding window with overlap
+            scores = self._detect_sliding(series, model_seq_len)
+
+        return self.normalise_scores(scores)
+
+    def _detect_single(self, series: np.ndarray, model_seq_len: int) -> np.ndarray:
+        """
+        Process a series that fits within MOMENT's context window.
+
+        If the series is shorter than model_seq_len, we pad it
+        with the mean value, run detection, then trim the result.
+        """
+        original_len = len(series)
+
+        # Pad if shorter than expected length
+        if len(series) < model_seq_len:
+            pad_val = np.mean(series)
+            padded = np.pad(
+                series,
+                (0, model_seq_len - len(series)),
+                mode="constant",
+                constant_values=pad_val,
+            )
+        else:
+            padded = series[:model_seq_len]
+
+        # Create input tensor: (batch=1, channels=1, seq_len)
+        input_tensor = torch.tensor(
+            padded, dtype=torch.float32
+        ).unsqueeze(0).unsqueeze(0)
+
+        # Run the model
+        with torch.no_grad():
+            output = self.model(x_enc=input_tensor)
+
+        # Extract reconstruction and compute error
+        reconstruction = output.reconstruction.squeeze().cpu().numpy()
+        errors = np.abs(padded - reconstruction)
+
+        # Trim back to original length
+        return errors[:original_len]
+
+    def _detect_sliding(self, series: np.ndarray, model_seq_len: int) -> np.ndarray:
+        """
+        Process a long series using overlapping sliding windows.
+
+        Strategy:
+        - Slide a window of model_seq_len across the series
+        - Overlap by 50% to ensure smooth transitions
+        - Average the error scores where windows overlap
+        """
+        step = model_seq_len // 2  # 50% overlap
+        n = len(series)
+
+        # Accumulators for averaging overlapping regions
+        score_sum = np.zeros(n)
+        score_count = np.zeros(n)
+
+        # Slide across the series
+        start = 0
+        while start < n:
+            end = min(start + model_seq_len, n)
+            window = series[start:end]
+
+            # Pad if the last window is shorter
+            if len(window) < model_seq_len:
+                pad_val = np.mean(window)
+                window = np.pad(
+                    window,
+                    (0, model_seq_len - len(window)),
+                    mode="constant",
+                    constant_values=pad_val,
+                )
+
+            # Create input tensor
+            input_tensor = torch.tensor(
+                window, dtype=torch.float32
+            ).unsqueeze(0).unsqueeze(0)
+
+            # Run the model
+            with torch.no_grad():
+                output = self.model(x_enc=input_tensor)
+
+            reconstruction = output.reconstruction.squeeze().cpu().numpy()
+            errors = np.abs(window - reconstruction)
+
+            # Only accumulate scores for the actual (non-padded) region
+            actual_len = min(model_seq_len, n - start)
+            score_sum[start : start + actual_len] += errors[:actual_len]
+            score_count[start : start + actual_len] += 1
+
+            start += step
+
+        # Average overlapping scores
+        score_count[score_count == 0] = 1  # Avoid division by zero
+        return score_sum / score_count