apex/ttm/forecast.py

"""Frozen TTM-r2 zero-shot forecast wrapper (Phase 2 Day 4 task 2.8).

Three responsibilities (in dependency order):

  1. `aggregate_to_1hz`: collapse raw N Hz telemetry to the macroscopic 1 Hz
     mini-sector backbone (D-011 path A; the wave-30 horizon contract is
     30 steps at 1 Hz per shapes.HORIZON). Per-channel aggregation rules
     respect driver-input semantics: brake pressure and throttle pct
     preserve peaks (a 20 ms brake spike must not be averaged away), the
     gear channel preserves the last value of each second, everything
     else uses the mean.

  2. `shape_ttm_input`: align an aggregated telemetry array to the
     TTM-r2 context-window contract. Pads short telemetry by repeating
     the first row (edge-pad, matching the G1 smoke convention at
     `logs/day-03-g1-ttm-smoke.md`) and tail-truncates long telemetry so
     the most recent context drives the prediction.

  3. `TtmForecaster`: the actual frozen-model holder. Loaded once;
     `.forecast()` returns a `(B, HORIZON, CHANNEL_COUNT)` tensor. Heavy
     dependency (torch + tsfm_public + 600MB HF download) so it lives in
     a class that's only instantiated when a forecast is actually
     needed. Unit tests cover surfaces 1 + 2; the integration test
     (task 2.10) covers surface 3 end-to-end.

The engine-agnostic boundary lives downstream in `apex.physics.validator`
and `apex.shared.contracts.violations`; this module produces tensors,
not violation logs.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Final

import numpy as np

from apex.shared.contracts import CHANNEL_COUNT, CHANNELS, HORIZON, channel_index


# ---- Aggregation rules: which channels peak, which last-value, rest mean ---

_DEFAULT_PEAK_CHANNELS: Final[tuple[str, ...]] = (
    "brake_pa",         # driver-input peak preserves brief spikes
    "throttle_pct",     # driver-input peak preserves shifts
    "lat_g",            # acceleration peaks matter for friction-ellipse audit
    "long_g",
)

_DEFAULT_LAST_CHANNELS: Final[tuple[str, ...]] = (
    "gear",             # discrete; averaging is nonsense
    "coa_overlap_flag", # discrete {0, 1}
)


@dataclass(frozen=True)
class AggregationConfig:
    """Per-channel aggregation rule overrides.

    Default rules apply when this dataclass is left at its defaults. Callers
    that need a different aggregation policy (e.g. the polyphase 50 Hz path
    B will pass an instance with empty peak_channels because it preserves
    the original sample rate) construct a custom instance.
    """

    peak_channels: tuple[str, ...] = field(
        default_factory=lambda: _DEFAULT_PEAK_CHANNELS
    )
    last_value_channels: tuple[str, ...] = field(
        default_factory=lambda: _DEFAULT_LAST_CHANNELS
    )


def aggregate_to_1hz(
    telemetry: np.ndarray,
    *,
    source_hz: int,
    config: AggregationConfig | None = None,
) -> np.ndarray:
    """Collapse `telemetry` from `source_hz` to 1 Hz mini-sector rows.

    Args:
      telemetry: (T, CHANNEL_COUNT) raw array in CHANNELS column order.
      source_hz: positive integer source sample rate. `source_hz=1` is a
        no-op pass-through.
      config: aggregation rule overrides; defaults applied when None.

    Returns: (floor(T / source_hz), CHANNEL_COUNT) float64 array. Partial
      trailing windows are dropped; the wave-30 D-011 path A is anchored
      on full-second mini-sectors, so a 2.4s capture yields 2 rows.
    """
    if source_hz <= 0:
        raise ValueError(f"source_hz must be positive; got {source_hz}.")
    if telemetry.ndim != 2 or telemetry.shape[1] != CHANNEL_COUNT:
        raise ValueError(
            f"aggregate_to_1hz expects (T, {CHANNEL_COUNT}) channels; "
            f"got {telemetry.shape}."
        )

    cfg = config or AggregationConfig()
    full_seconds = telemetry.shape[0] // source_hz
    if full_seconds == 0:
        return np.zeros((0, CHANNEL_COUNT), dtype=np.float64)

    # Reshape into (seconds, source_hz, channels) for vectorized aggregation.
    trimmed = telemetry[: full_seconds * source_hz].astype(np.float64, copy=False)
    windowed = trimmed.reshape(full_seconds, source_hz, CHANNEL_COUNT)

    peak_idx = {channel_index(c) for c in cfg.peak_channels if c in CHANNELS}
    last_idx = {channel_index(c) for c in cfg.last_value_channels if c in CHANNELS}

    out = np.empty((full_seconds, CHANNEL_COUNT), dtype=np.float64)
    for ch in range(CHANNEL_COUNT):
        if ch in peak_idx:
            out[:, ch] = windowed[:, :, ch].max(axis=1)
        elif ch in last_idx:
            out[:, ch] = windowed[:, -1, ch]
        else:
            out[:, ch] = windowed[:, :, ch].mean(axis=1)
    return out


def shape_ttm_input(
    telemetry: np.ndarray,
    *,
    context_length: int,
    dtype: np.dtype = np.float32,
) -> np.ndarray:
    """Align `telemetry` to TTM-r2's (1, context_length, CHANNEL_COUNT) input.

    Pads short telemetry by repeating the first row (edge-pad, matching
    the G1 smoke at `logs/day-03-g1-ttm-smoke.md`). Truncates long
    telemetry from the head so the tail (most recent samples) drives
    the prediction.
    """
    if telemetry.ndim != 2 or telemetry.shape[1] != CHANNEL_COUNT:
        raise ValueError(
            f"shape_ttm_input expects (T, {CHANNEL_COUNT}) channels; "
            f"got {telemetry.shape}."
        )

    T = telemetry.shape[0]
    if T < context_length:
        pad = np.repeat(telemetry[:1], context_length - T, axis=0)
        aligned = np.concatenate([pad, telemetry], axis=0)
    else:
        aligned = telemetry[-context_length:]

    return aligned.astype(dtype, copy=False)[None, :, :]


# ---- TtmForecaster: heavy class, loaded lazily -------------------------

class TtmForecaster:
    """Frozen Granite TimeSeries TTM-r2 zero-shot forecaster.

    Loads `ibm-granite/granite-timeseries-ttm-r2` once per instance.
    `.forecast(telemetry, source_hz=N)` aggregates -> shapes -> forwards
    and returns a `(1, HORIZON, CHANNEL_COUNT)` numpy array matching
    `shapes.TENSOR_SHAPE` (with batch=1).

    This class is NOT imported at module load; callers must construct it
    explicitly. The unit-test suite covers `aggregate_to_1hz` +
    `shape_ttm_input` without instantiating this class; the integration
    test (task 2.10) instantiates it and runs a real forward pass.
    """

    def __init__(self, model_id: str = "ibm-granite/granite-timeseries-ttm-r2"):
        import torch
        from tsfm_public import TinyTimeMixerForPrediction

        self._torch = torch
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._model = TinyTimeMixerForPrediction.from_pretrained(
            model_id,
            num_input_channels=CHANNEL_COUNT,
            prediction_filter_length=HORIZON,
        ).to(self._device).eval()
        self._context_length: int = int(self._model.config.context_length)

    @property
    def context_length(self) -> int:
        return self._context_length

    def forecast(
        self,
        telemetry: np.ndarray,
        *,
        source_hz: int,
        config: AggregationConfig | None = None,
    ) -> np.ndarray:
        """End-to-end zero-shot forecast: aggregate -> shape -> forward."""
        aggregated = aggregate_to_1hz(telemetry, source_hz=source_hz, config=config)
        shaped = shape_ttm_input(aggregated, context_length=self._context_length)
        x = self._torch.from_numpy(shaped).to(self._device)
        with self._torch.no_grad():
            out = self._model(past_values=x)
        return out.prediction_outputs.detach().cpu().numpy()


__all__ = [
    "AggregationConfig",
    "TtmForecaster",
    "aggregate_to_1hz",
    "shape_ttm_input",
]