Spaces:
Running on Zero
Running on Zero
| """The ci-proto preprocessing and quantile forecasting path. | |
| Includes postprocessing adapted by Continuous Intelligence from TiRex-2: | |
| https://github.com/NX-AI/tirex-2/blob/bb5427f4a97379a17756e67bc846c4f4b569eb48/src/tirex2/model/component/postprocessor.py | |
| Copyright (c) NXAI GmbH. | |
| Upstream portions are licensed under Apache-2.0; see NOTICE. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| from collections.abc import Callable, Iterable, Mapping, Sequence | |
| from dataclasses import dataclass | |
| from numbers import Integral | |
| import torch | |
| from torch import Tensor, nn | |
| from torch.nn.functional import pad | |
| from ._model import Model, forecast_layout | |
| class Forecaster: | |
| """ci-proto with preprocessing and calibration.""" | |
| def __init__(self, model: Model, tta: dict[str, float]) -> None: | |
| self._model = model | |
| self._tta = tta | |
| self._long_model: Callable[[], Forecaster] | None = None | |
| self._info: dict | None = None | |
| def device(self) -> torch.device: | |
| """Device used for the neural network.""" | |
| return next(self._model.parameters()).device | |
| def quantiles(self) -> Tensor: | |
| """CPU tensor of quantile levels, shape ``[quantiles]``. | |
| ci-proto returns nine quantiles between 0.1 and 0.9 in increments of 0.1. | |
| """ | |
| return self._model.quantiles.detach().cpu().clone() | |
| def info(self) -> dict | None: | |
| """Model description from the bundle's ``config.json``: name, parameter count, quantile | |
| levels and licence. ``None`` when loaded without one.""" | |
| return None if self._info is None else json.loads(json.dumps(self._info)) | |
| def backend(self) -> str: | |
| """Selected inference backend: ``mamba3`` (fused CUDA) or ``torch`` (portable).""" | |
| return getattr(self._model, "backend", "torch") | |
| def forecast( | |
| self, | |
| history: Tensor, | |
| horizon: int = 320, | |
| *, | |
| past_covariates: Tensor | None = None, | |
| future_covariates: Tensor | None = None, | |
| categorical_past: tuple[int, ...] = (), | |
| categorical_future: tuple[int, ...] = (), | |
| timestamps: Tensor | None = None, | |
| freq: str | None = None, | |
| series_ids: Sequence[str] | None = None, | |
| static_features: Mapping[str, Sequence[object]] | None = None, | |
| ) -> Tensor: | |
| """Forecast joint variables, or independent series identified by ``series_ids``. | |
| ``history``: float tensor ``[variables, time]``; NaN denotes missing data. | |
| ``past_covariates``: ``[past_variables, time]``. | |
| ``future_covariates``: ``[known_variables, time + horizon]``; extra future | |
| steps are ignored. Only supply values known at forecast time. Categorical | |
| row indices refer to numeric category codes in the corresponding tensor. | |
| ``timestamps``: int64 epoch nanoseconds over history plus horizon, shared | |
| ``[time + horizon]`` or per-series ``[variables, time + horizon]``. Supply | |
| ``freq`` to generate ci-proto's calendar features. Without timestamps, | |
| calendar features are omitted. Leading padding uses int64's minimum value. | |
| ``series_ids`` enables cohort policies for independent univariate rows. | |
| Supply the complete cohort for a window. ``static_features`` maps column | |
| names to values in input row order. Cohort covariates can also be shaped | |
| ``[variables, covariates, time]``, including the horizon for known values. | |
| Returns CPU float32 quantiles ``[variables, quantiles, horizon]``, in | |
| original units, with levels given by ``self.quantiles``. ci-proto returns | |
| nine quantiles between 0.1 and 0.9 in increments of 0.1. Horizon is any | |
| positive integer and defaults to 320. | |
| """ | |
| model, tta = self._model, self._tta | |
| if type(horizon) is not int or horizon < 1: | |
| raise ValueError("horizon must be a positive integer") | |
| if ( | |
| self._long_model is not None | |
| and isinstance(history, Tensor) | |
| and history.ndim == 2 | |
| and history.shape[-1] > model.config.context_len | |
| ): | |
| long = self._long_model() | |
| model, tta = long._model, long._tta | |
| device = next(model.parameters()).device | |
| history = _rows(history, "history", device) | |
| count, length = history.shape | |
| stamps = _timestamps(timestamps, freq, count, length, horizon) | |
| if stamps is not None and tta["calendar_covariates"] and freq is None: | |
| raise ValueError("freq is required to generate calendar covariates") | |
| if series_ids is not None: | |
| past = _cohort_covariates( | |
| past_covariates, "past_covariates", count, length, 0, device, categorical_past | |
| ) | |
| future = _cohort_covariates( | |
| future_covariates, | |
| "future_covariates", | |
| count, | |
| length, | |
| horizon, | |
| device, | |
| categorical_future, | |
| ) | |
| cases = _cohort_cases( | |
| history, | |
| horizon, | |
| tta, | |
| past, | |
| future, | |
| categorical_past, | |
| categorical_future, | |
| stamps, | |
| freq, | |
| series_ids, | |
| static_features, | |
| ) | |
| return _run_cases(model, cases, count, horizon, tta) | |
| if static_features is not None: | |
| raise ValueError("static_features requires series_ids") | |
| if stamps is not None: | |
| if not torch.equal(stamps, stamps[:1].expand_as(stamps)): | |
| raise ValueError("different timestamp axes require series_ids") | |
| stamps = stamps[0] | |
| start = int((stamps != torch.iinfo(torch.int64).min).long().argmax()) | |
| history, stamps = history[:, start:], stamps[start:] | |
| if past_covariates is not None: | |
| past_covariates = past_covariates[:, start:] | |
| if future_covariates is not None: | |
| future_covariates = future_covariates[:, start:] | |
| context, future = _context( | |
| history, | |
| horizon, | |
| past_covariates, | |
| future_covariates, | |
| categorical_past, | |
| categorical_future, | |
| bool(tta["target_encode_categoricals"]), | |
| ) | |
| context, future = _with_calendar( | |
| context, future, count, stamps, freq, bool(tta["calendar_covariates"]) | |
| ) | |
| return _predict(model, context[None], future[None], count, horizon, tta)[0] | |
| def _rows( | |
| value: Tensor, | |
| name: str, | |
| device: torch.device, | |
| *, | |
| dtype: torch.dtype | None = torch.float32, | |
| length: int | None = None, | |
| ) -> Tensor: | |
| if not isinstance(value, Tensor) or value.ndim != 2 or min(value.shape) < 1: | |
| raise ValueError(f"{name} must be a nonempty [variables, time] tensor") | |
| value = value[:, :length] | |
| if value.is_complex() or bool(torch.isinf(value).any()): | |
| raise ValueError(f"{name} must contain real finite values or nan") | |
| result = value.to(device=device, dtype=dtype) | |
| if bool(torch.isinf(result).any()): | |
| raise ValueError(f"{name} exceeds the float32 range") | |
| return result | |
| def _target_encode(context: Tensor, codes: Tensor) -> Tensor: | |
| """Replace categories with smoothed target means fitted only on history.""" | |
| if codes.dtype == torch.bool: | |
| codes = codes.to(torch.int64) | |
| rows, past = context.shape | |
| observed = torch.isfinite(context) | |
| target = torch.where(observed, context.to(torch.float64), 0.0) | |
| count = observed.sum(dim=-1) | |
| mean = torch.where(count > 0, target.sum(dim=-1) / count.clamp(min=1), 0.0) | |
| missing = torch.isnan(codes) | |
| categories = torch.unique(codes[~missing]) | |
| found = torch.searchsorted(categories, codes.nan_to_num().contiguous()) | |
| slots = torch.where(missing, categories.numel(), found) | |
| width = categories.numel() + 1 | |
| offsets = torch.arange(rows, device=codes.device)[:, None] * width | |
| index = (offsets + slots[:, :past]).flatten() | |
| sums = torch.bincount(index, weights=target.flatten(), minlength=rows * width) | |
| counts = torch.bincount( | |
| index, weights=observed.to(torch.float64).flatten(), minlength=rows * width | |
| ) | |
| lookup = (mean[:, None] + sums.view(rows, width)) / (1.0 + counts.view(rows, width)) | |
| return lookup.gather(1, slots).to(torch.float32) | |
| def _encode_covariates( | |
| rows: Tensor, indices: tuple[int, ...], history: Tensor, name: str, encode: bool | |
| ) -> Tensor: | |
| if len(set(indices)) != len(indices) or any( | |
| not isinstance(index, Integral) or isinstance(index, bool) or not 0 <= index < rows.shape[0] | |
| for index in indices | |
| ): | |
| raise ValueError(f"{name} must contain unique valid covariate row indices") | |
| encoded = rows.to(torch.float32, copy=True) | |
| if indices and encode: | |
| index = [int(value) for value in indices] | |
| context = torch.nanmean(history.to(torch.float64), dim=0, keepdim=True) | |
| encoded[index] = _target_encode(context.expand(len(index), -1), rows[index]) | |
| if bool(torch.isinf(encoded).any()): | |
| raise ValueError("numeric covariates exceed the float32 range") | |
| return encoded | |
| def _difference_mask( | |
| rows: Tensor, config: Mapping[str, float], lengths: Tensor | None = None | |
| ) -> Tensor: | |
| """Detect strong linear trends and monotone running totals on observed history.""" | |
| total = rows.shape[-1] | |
| window = ( | |
| max(1, round(config["trend_window"] * total)) | |
| if lengths is None | |
| else (lengths.double() * config["trend_window"]).round().clamp_min(1) | |
| ) | |
| values = rows.to(torch.float64) | |
| index = torch.arange(total, dtype=torch.float64, device=rows.device).expand_as(values) | |
| start = total - window if lengths is None else total - window[:, None] | |
| observed = torch.isfinite(values) & (index >= start) | |
| v = torch.where(observed, values, torch.zeros_like(values)) | |
| t = torch.where(observed, index, torch.zeros_like(index)) | |
| n = observed.sum(dim=-1).to(torch.float64) | |
| n_safe = n.clamp(min=2.0) | |
| var_t = ((t * t).sum(dim=-1) - t.sum(dim=-1) ** 2 / n_safe).clamp(min=1e-12) | |
| var_v = (v * v).sum(dim=-1) - v.sum(dim=-1) ** 2 / n_safe | |
| cov_tv = (t * v).sum(dim=-1) - t.sum(dim=-1) * v.sum(dim=-1) / n_safe | |
| deviation = torch.sqrt((var_v / (n_safe - 1.0)).clamp(min=0.0)) | |
| score = (cov_tv / var_t / deviation.clamp(min=1e-12)).abs() * window | |
| r_squared = cov_tv**2 / (var_t * var_v.clamp(min=1e-12)) | |
| trending = ( | |
| (score >= config["trend_threshold"]) | |
| & (r_squared >= config["trend_min_r2"]) | |
| & (n >= 4) | |
| & (deviation > 1e-12) | |
| ) | |
| if not config["tta_running_totals"]: | |
| return trending | |
| observed = torch.isfinite(rows) | |
| index = torch.arange(total, device=rows.device).expand_as(rows) | |
| latest = torch.where(observed, index, torch.zeros_like(index)).cummax(dim=-1).values | |
| filled = rows.gather(-1, latest) | |
| steps = filled[..., 1:] - filled[..., :-1] | |
| finite = torch.isfinite(steps) | |
| rising = (~finite | (steps >= 0)).all(dim=-1) | |
| falling = (~finite | (steps <= 0)).all(dim=-1) | |
| counted = finite & observed[..., 1:] | |
| moving = (counted & (steps != 0)).sum(dim=-1) | |
| share = moving / counted.sum(dim=-1).clamp(min=1) | |
| running = (observed.sum(dim=-1) >= 4) & (moving > 0) & (share >= 0.5) | |
| return trending | (running & (rising | falling)) | |
| def _last_observed(history: Tensor) -> Tensor: | |
| observed = torch.isfinite(history) | |
| last_index = observed.long().cumsum(dim=-1).argmax(dim=-1) | |
| return torch.nan_to_num(history.gather(-1, last_index[:, None]).squeeze(-1), nan=0.0) | |
| def _calibrate( | |
| pred: Tensor, | |
| history: Tensor, | |
| differenced: Tensor, | |
| config: Mapping[str, float], | |
| *, | |
| anchors: Tensor | None = None, | |
| has_trend: bool | None = None, | |
| ) -> Tensor: | |
| median = pred[:, pred.shape[1] // 2 : pred.shape[1] // 2 + 1] | |
| steps = torch.arange(1, pred.shape[-1] + 1, dtype=pred.dtype, device=pred.device) | |
| factor = config["raw_band_scale"] ** 0.5 * steps ** (config["raw_band_exponent"] - 0.5) | |
| calibrated = median + factor * (pred - median) | |
| if bool(differenced.any()) if has_trend is None else has_trend: | |
| selected = pred[differenced] | |
| middle = selected[:, selected.shape[1] // 2 : selected.shape[1] // 2 + 1] | |
| offset = selected - middle | |
| band = offset.sum(dim=-1, keepdim=True).sign() * torch.sqrt( | |
| torch.cumsum(offset * offset, dim=-1) | |
| ) | |
| band = band * steps ** (config["diff_band_exponent"] - 0.5) | |
| if anchors is None: | |
| anchors = _last_observed(history) | |
| calibrated[differenced] = ( | |
| torch.cumsum(middle, dim=-1) | |
| + config["diff_band_scale"] ** 0.5 * band | |
| + anchors[differenced, None, None] | |
| ) | |
| return calibrated.sort(dim=1).values | |
| def _forecast_once( | |
| model: nn.Module, | |
| context: Tensor, | |
| future: Tensor, | |
| n_targets: int, | |
| horizon: int, | |
| tta: Mapping[str, float], | |
| lengths: Tensor, | |
| *, | |
| sign_flip: bool = False, | |
| ) -> Tensor: | |
| """Forecast equal-width groups, optionally including their independent negative views.""" | |
| cfg = model.config | |
| patch = cfg.output_patch_size | |
| window_len = max(cfg.future_len, -(-horizon // patch) * patch) | |
| samples, channels, width = context.shape | |
| history = context[:, :n_targets].flatten(0, 1) | |
| row_lengths = lengths.repeat_interleave(n_targets) | |
| differenced = ( | |
| _difference_mask(history, tta, row_lengths) | |
| if tta["postprocess"] and tta["tta_diff"] and horizon <= tta["differencing_max_horizon"] | |
| else torch.zeros(len(history), dtype=torch.bool, device=history.device) | |
| ) | |
| has_trend = bool(differenced.any()) | |
| anchors = _last_observed(history) if has_trend else None | |
| if has_trend: | |
| diff = torch.cat([torch.zeros_like(history[:, :1]), history[:, 1:] - history[:, :-1]], -1) | |
| start = width - row_lengths | |
| valid = torch.arange(width, device=history.device) >= start[:, None] | |
| diff = torch.where(valid, diff, torch.nan).scatter(1, start[:, None], 0.0) | |
| targets = torch.where(differenced[:, None], diff, history).reshape( | |
| samples, n_targets, width | |
| ) | |
| context = torch.cat([targets, context[:, n_targets:]], dim=1) | |
| n_known = future.shape[1] | |
| window = context.new_full((samples, channels, window_len), torch.nan) | |
| window[:, n_targets : n_targets + n_known, :horizon] = future | |
| observed = torch.isfinite(context).flatten(0, 1).any(0) | |
| span = width - int(observed.int().argmax()) if bool(observed.any()) else 1 | |
| context_len = max(patch, 1 << (span - 1).bit_length()) | |
| if cfg.context_len is not None: | |
| context_len = min(context_len, cfg.context_len) | |
| x = torch.cat([context[..., -context_len:], window], dim=-1) | |
| x = pad(x, (context_len + window_len - x.shape[-1], 0), value=torch.nan) | |
| if sign_flip: | |
| x = torch.cat([x, -x]) | |
| differenced = differenced.repeat(2) | |
| if anchors is not None: | |
| anchors = torch.cat([anchors, -anchors]) | |
| layout = forecast_layout(len(x), channels, n_targets, n_known, x.device) | |
| batch = { | |
| "x": x.flatten(0, 1), | |
| "group_vector": layout.groups, | |
| "target_mask": layout.targets, | |
| "future_covariate_mask": layout.future, | |
| "past_covariate_mask": layout.past, | |
| } | |
| with torch.no_grad(): | |
| pred = model(batch, future_len=window_len, horizon=horizon, layout=layout) | |
| pred = pred.reshape(len(x), channels, -1, horizon)[:, :n_targets].flatten(0, 1) | |
| if tta["postprocess"]: | |
| pred = _calibrate(pred, history, differenced, tta, anchors=anchors, has_trend=has_trend) | |
| return pred.reshape(len(x), n_targets, -1, horizon) | |
| def _context( | |
| history: Tensor, | |
| horizon: int, | |
| past_covariates: Tensor | None, | |
| future_covariates: Tensor | None, | |
| categorical_past: tuple[int, ...], | |
| categorical_future: tuple[int, ...], | |
| encode_categoricals: bool, | |
| ) -> tuple[Tensor, Tensor]: | |
| device = history.device | |
| n_targets, length = history.shape | |
| parts = [history] | |
| future = history.new_empty((0, horizon)) | |
| if future_covariates is not None: | |
| rows = _rows( | |
| future_covariates, "future_covariates", device, dtype=None, length=length + horizon | |
| ) | |
| if rows.shape[1] < length + horizon: | |
| raise ValueError("future_covariates must cover history plus horizon") | |
| rows = _encode_covariates( | |
| rows, categorical_future, history, "categorical_future", encode_categoricals | |
| ) | |
| parts.append(rows[:, :length]) | |
| future = rows[:, length:] | |
| elif categorical_future: | |
| raise ValueError("categorical_future requires future_covariates") | |
| if past_covariates is not None: | |
| rows = _rows(past_covariates, "past_covariates", device, dtype=None) | |
| if rows.shape[1] != length: | |
| raise ValueError("past_covariates must have the same history length") | |
| parts.append( | |
| _encode_covariates( | |
| rows, categorical_past, history, "categorical_past", encode_categoricals | |
| ) | |
| ) | |
| elif categorical_past: | |
| raise ValueError("categorical_past requires past_covariates") | |
| return torch.cat(parts, dim=0), future | |
| def _predict( | |
| model: nn.Module, | |
| context: Tensor, | |
| future: Tensor, | |
| n_targets: int, | |
| horizon: int, | |
| tta: Mapping[str, float], | |
| lengths: Tensor | None = None, | |
| ) -> Tensor: | |
| """Forecast ``[samples, channels, time]`` groups with shared sign preprocessing.""" | |
| device = context.device | |
| samples = len(context) | |
| history = context[:, :n_targets].flatten(0, 1) | |
| if lengths is None: | |
| lengths = torch.full((samples,), context.shape[-1], dtype=torch.long, device=device) | |
| sign_flip = bool(tta["postprocess"] and tta["tta_sign_flip"]) | |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=bool(tta["bfloat16"])): | |
| pred = _forecast_once( | |
| model, context, future, n_targets, horizon, tta, lengths, sign_flip=sign_flip | |
| ) | |
| if sign_flip: | |
| positive, negative = pred.chunk(2) | |
| # Loading validates a strictly sorted, symmetric grid centred on the median. | |
| pred = (positive - negative.flip(2)) / 2 | |
| pred = pred.float().flatten(0, 1) | |
| observed = torch.isfinite(history) | |
| if tta["clip_nonnegative"]: | |
| nonnegative = observed.any(dim=-1) & ~(observed & (history < 0)).any(dim=-1) | |
| pred = torch.where(nonnegative[:, None, None], pred.clamp_min(0), pred) | |
| if tta["integer_round"]: | |
| integers = observed.any(dim=-1) & ~(observed & (history != history.round())).any(dim=-1) | |
| pred = torch.where(integers[:, None, None], pred.round(), pred) | |
| expected = (samples * n_targets, model.get_buffer("quantiles").numel(), horizon) | |
| if pred.shape != expected or not bool(torch.isfinite(pred).all()): | |
| raise ValueError("model returned invalid forecast values or shape") | |
| return pred.reshape(samples, n_targets, -1, horizon).detach().cpu() | |
| def _timestamps( | |
| timestamps: Tensor | None, freq: str | None, count: int, length: int, horizon: int | |
| ) -> Tensor | None: | |
| if timestamps is None: | |
| if freq is not None: | |
| raise ValueError("freq requires timestamps") | |
| return None | |
| if not isinstance(timestamps, Tensor) or timestamps.dtype != torch.int64: | |
| raise ValueError("timestamps must be an int64 tensor of epoch nanoseconds") | |
| stamps = timestamps.detach().cpu() | |
| if stamps.ndim == 1: | |
| stamps = stamps[None].expand(count, -1) | |
| if stamps.ndim != 2 or stamps.shape[0] != count or stamps.shape[1] < length + horizon: | |
| raise ValueError("timestamps must cover history plus horizon for every series") | |
| stamps = stamps[:, : length + horizon] | |
| valid = stamps != torch.iinfo(torch.int64).min | |
| if not valid[:, length - 1 :].all(): | |
| raise ValueError("timestamps must include the forecast origin and every future step") | |
| start = valid.long().argmax(dim=1) | |
| if not torch.equal(valid, torch.arange(stamps.shape[1])[None] >= start[:, None]): | |
| raise ValueError("missing timestamps are only allowed in leading history padding") | |
| if ((stamps[:, 1:] <= stamps[:, :-1]) & valid[:, :-1]).any(): | |
| raise ValueError("timestamps must be strictly increasing") | |
| if freq is not None and not isinstance(freq, str): | |
| raise ValueError("freq must be a frequency string") | |
| return stamps | |
| def _with_calendar( | |
| context: Tensor, | |
| future: Tensor, | |
| n_targets: int, | |
| stamps: Tensor | None, | |
| freq: str | None, | |
| enabled: bool, | |
| *, | |
| cache: dict[bytes, Tensor] | None = None, | |
| ) -> tuple[Tensor, Tensor]: | |
| if not enabled or stamps is None: | |
| return context, future | |
| if freq is None: | |
| raise ValueError("freq is required to generate calendar covariates") | |
| from ._calendar import calendar_channels, frequency_step_ns | |
| key = stamps.numpy().tobytes() | |
| calendar = None if cache is None else cache.get(key) | |
| if calendar is None: | |
| calendar = torch.from_numpy(calendar_channels(stamps.numpy(), frequency_step_ns(freq))) | |
| calendar = calendar.to(context) | |
| if cache is not None: | |
| cache[key] = calendar | |
| if not calendar.shape[0]: | |
| return context, future | |
| length = context.shape[-1] | |
| at = n_targets + future.shape[0] | |
| context = torch.cat([context[:at], calendar[:, :length], context[at:]], dim=0) | |
| return context, torch.cat([future, calendar[:, length:]], dim=0) | |
| class _Case: | |
| """One independent forecast window and the caller's output row indices.""" | |
| indices: list[int] | |
| context: Tensor | |
| future: Tensor | |
| n_targets: int | |
| def _cohort_covariates( | |
| value: Tensor | None, | |
| name: str, | |
| count: int, | |
| length: int, | |
| horizon: int, | |
| device: torch.device, | |
| categorical: tuple[int, ...], | |
| ) -> Tensor | None: | |
| if value is None: | |
| if categorical: | |
| raise ValueError(f"categorical indices require {name}") | |
| return None | |
| if not isinstance(value, Tensor) or value.ndim not in (2, 3): | |
| raise ValueError(f"{name} must be a two- or three-dimensional tensor") | |
| rows = value[None].expand(count, -1, -1) if value.ndim == 2 else value | |
| if rows.shape[0] != count or rows.shape[1] < 1: | |
| raise ValueError(f"{name} must have covariates for every series") | |
| width = length + horizon | |
| if (horizon and rows.shape[-1] < width) or (not horizon and rows.shape[-1] != width): | |
| raise ValueError(f"{name} must cover the required time steps") | |
| rows = rows[..., :width] | |
| flat = _rows(rows.reshape(-1, width), name, device, dtype=None) | |
| # Validate indices without fitting an encoding across unrelated series. | |
| if len(set(categorical)) != len(categorical) or any( | |
| not isinstance(index, Integral) or isinstance(index, bool) or not 0 <= index < rows.shape[1] | |
| for index in categorical | |
| ): | |
| raise ValueError(f"categorical_{name} must contain unique valid covariate row indices") | |
| return flat.reshape(count, rows.shape[1], width) | |
| def _cohort_cases( | |
| history: Tensor, | |
| horizon: int, | |
| tta: Mapping[str, float], | |
| past: Tensor | None, | |
| future: Tensor | None, | |
| categorical_past: tuple[int, ...], | |
| categorical_future: tuple[int, ...], | |
| stamps: Tensor | None, | |
| freq: str | None, | |
| series_ids: Sequence[str], | |
| static_features: Mapping[str, Sequence[object]] | None, | |
| ) -> Iterable[_Case]: | |
| from ._groups import chunk_groups, joint_neighbour_rows, static_group_rows | |
| count, length = history.shape | |
| if ( | |
| isinstance(series_ids, str) | |
| or len(series_ids) != count | |
| or any(not isinstance(name, str) or not name for name in series_ids) | |
| or len(set(series_ids)) != count | |
| ): | |
| raise ValueError("series_ids must contain one unique string per history row") | |
| if static_features is not None and not isinstance(static_features, Mapping): | |
| raise ValueError("static_features must map column names to per-series values") | |
| statics = dict(static_features or {}) | |
| for name, values in statics.items(): | |
| if ( | |
| not isinstance(name, str) | |
| or isinstance(values, (str, bytes, Mapping)) | |
| or not hasattr(values, "__len__") | |
| or len(values) != count | |
| ): | |
| raise ValueError("static_features must have one value per series in each column") | |
| statics[name] = ( | |
| values.detach().cpu().tolist() if isinstance(values, Tensor) else tuple(values) | |
| ) | |
| if stamps is None: | |
| stamps = torch.arange(length + horizon)[None].expand(count, -1) | |
| calendar_stamps = None | |
| starts = [0] * count | |
| else: | |
| calendar_stamps = stamps | |
| starts = (stamps != torch.iinfo(torch.int64).min).long().argmax(dim=1).tolist() | |
| order = sorted(range(count), key=lambda i: series_ids[i]) | |
| stamps_np = stamps.numpy() | |
| bare = not statics and past is None and future is None | |
| chunk = int(tta["chunk_joint"]) if bare else 0 | |
| static_rows = None | |
| neighbours = None | |
| if (statics and tta["static_groups"]) or ( | |
| not statics and not chunk and tta["joint_neighbours"] | |
| ): | |
| histories = history.detach().double().cpu().numpy() | |
| histories_in_order = [histories[i, starts[i] :] for i in order] | |
| times = [stamps_np[i, starts[i] : length] for i in order] | |
| if statics and tta["static_groups"]: | |
| columns = [[statics[name][i] for i in order] for name in sorted(statics)] | |
| static_rows = static_group_rows( | |
| histories_in_order, times, columns, int(tta["static_groups"]) | |
| ) | |
| elif not statics and not chunk and tta["joint_neighbours"]: | |
| neighbours = joint_neighbour_rows(histories_in_order, times, int(tta["joint_neighbours"])) | |
| groups = ( | |
| chunk_groups([int(stamps_np[i, length - 1]) for i in order], chunk) | |
| if chunk | |
| else [[i] for i in range(count)] | |
| ) | |
| calendar_cache: dict[bytes, Tensor] = {} | |
| for members in groups: | |
| indices = [order[position] for position in members] | |
| longest = min(indices, key=lambda i: starts[i]) | |
| start = starts[longest] | |
| targets = history[indices, start:] | |
| own_past = None if past is None else past[longest, :, start:] | |
| own_future = None if future is None else future[longest, :, start:] | |
| if static_rows is not None: | |
| extra = torch.from_numpy(static_rows[members[0]]).to(history) | |
| own_past = extra if own_past is None else torch.cat([own_past, extra]) | |
| if neighbours is not None: | |
| targets = torch.from_numpy(neighbours[members[0]]).to(history) | |
| context, revealed = _context( | |
| targets, | |
| horizon, | |
| own_past, | |
| own_future, | |
| categorical_past, | |
| categorical_future, | |
| bool(tta["target_encode_categoricals"]), | |
| ) | |
| own_stamps = None if calendar_stamps is None else calendar_stamps[longest, start:] | |
| context, revealed = _with_calendar( | |
| context, | |
| revealed, | |
| len(targets), | |
| own_stamps, | |
| freq, | |
| bool(tta["calendar_covariates"]), | |
| cache=calendar_cache, | |
| ) | |
| yield _Case(indices, context, revealed, len(targets)) | |
| def _packed_predict( | |
| model: nn.Module, cases: Sequence[_Case], horizon: int, tta: Mapping[str, float] | |
| ) -> Tensor: | |
| """Stack equal-width groups without changing group membership or covariate ownership.""" | |
| width = max(case.context.shape[-1] for case in cases) | |
| context = torch.stack( | |
| [pad(case.context, (width - case.context.shape[-1], 0), value=torch.nan) for case in cases] | |
| ) | |
| lengths = torch.tensor([case.context.shape[-1] for case in cases], device=context.device) | |
| return _predict( | |
| model, | |
| context, | |
| torch.stack([case.future for case in cases]), | |
| cases[0].n_targets, | |
| horizon, | |
| tta, | |
| lengths, | |
| ).flatten(0, 1) | |
| def _run_cases( | |
| model: nn.Module, | |
| cases: Iterable[_Case], | |
| count: int, | |
| horizon: int, | |
| tta: Mapping[str, float], | |
| ) -> Tensor: | |
| out = torch.empty(count, model.get_buffer("quantiles").numel(), horizon) | |
| pending: dict[tuple[int, int, int], list[_Case]] = {} | |
| rows = 0 | |
| views = 2 if tta["postprocess"] and tta["tta_sign_flip"] else 1 | |
| def flush() -> None: | |
| for members in pending.values(): | |
| predictions = _packed_predict(model, members, horizon, tta) | |
| cursor = 0 | |
| for case in members: | |
| out[case.indices] = predictions[cursor : cursor + len(case.indices)] | |
| cursor += case.n_targets | |
| pending.clear() | |
| for case in cases: | |
| if pending and (rows + len(case.context)) * views > 2048: | |
| flush() | |
| rows = 0 | |
| key = (case.n_targets, len(case.future), len(case.context)) | |
| pending.setdefault(key, []).append(case) | |
| rows += len(case.context) | |
| if pending: | |
| flush() | |
| return out | |