| """Windowing over the time axis. Pure index arithmetic — no tensors, no I/O. |
| |
| Kept free of data so it can be reasoned about and tested on its own: an |
| off-by-one here leaks future timesteps into the input window, which is the |
| quietest possible way to produce excellent and meaningless results. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Iterator, Sequence |
| from typing import NamedTuple |
|
|
| __all__ = ["LatticeWindow", "Window"] |
|
|
|
|
| class Window(NamedTuple): |
| """Half-open index ranges: inputs ``[x0, x1)``, targets ``[y0, y1)``.""" |
|
|
| x0: int |
| x1: int |
| y0: int |
| y1: int |
|
|
|
|
| class LatticeWindow: |
| """The set of windows tiling a time axis. |
| |
| Args: |
| n_time: length of the time axis. |
| input_len: timesteps fed to the model. |
| horizon: timesteps to predict. ``0`` means no target range. |
| stride: step between consecutive window starts. |
| |
| Targets begin exactly where inputs end, so no window ever sees its own |
| target. |
| """ |
|
|
| __slots__ = ("windows", "input_len", "horizon", "stride", "n_time") |
|
|
| def __init__( |
| self, |
| n_time: int, |
| input_len: int, |
| horizon: int = 0, |
| stride: int = 1, |
| *, |
| _windows: Sequence[Window] | None = None, |
| ) -> None: |
| if input_len < 1: |
| raise ValueError(f"input_len must be >= 1; got {input_len}") |
| if horizon < 0: |
| raise ValueError(f"horizon must be >= 0; got {horizon}") |
| if stride < 1: |
| raise ValueError(f"stride must be >= 1; got {stride}") |
| self.n_time, self.input_len, self.horizon, self.stride = ( |
| n_time, |
| input_len, |
| horizon, |
| stride, |
| ) |
| if _windows is not None: |
| self.windows = tuple(_windows) |
| return |
| span = input_len + horizon |
| if span > n_time: |
| raise ValueError( |
| f"input_len + horizon = {span} exceeds the {n_time} available timesteps" |
| ) |
| self.windows = tuple( |
| Window(i, i + input_len, i + input_len, i + span) |
| for i in range(0, n_time - span + 1, stride) |
| ) |
|
|
| def _derive(self, windows: Sequence[Window]) -> LatticeWindow: |
| return LatticeWindow( |
| self.n_time, self.input_len, self.horizon, self.stride, _windows=windows |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.windows) |
|
|
| def __iter__(self) -> Iterator[Window]: |
| return iter(self.windows) |
|
|
| def __getitem__(self, i): |
| if isinstance(i, slice): |
| return self._derive(self.windows[i]) |
| return self.windows[i] |
|
|
| def split(self, at: int) -> tuple[LatticeWindow, LatticeWindow]: |
| """Split into windows ending at or before ``at``, and those starting at |
| or after it. |
| |
| Windows straddling the boundary are dropped by both sides rather than |
| assigned to one. That gap is deliberate: keeping a straddling window |
| would put timesteps from after the cut inside a training input. |
| """ |
| before = [w for w in self.windows if w.y1 <= at] |
| after = [w for w in self.windows if w.x0 >= at] |
| return self._derive(before), self._derive(after) |
|
|
| def split_at_time(self, times: Sequence, value) -> tuple[LatticeWindow, LatticeWindow]: |
| """:meth:`split` by an actual timestamp rather than an index. |
| |
| ``times`` must be sorted — a scan for the first timestamp past the cut |
| is meaningless otherwise, and unsorted input previously produced a |
| silently nonsensical split rather than an error. |
| """ |
| for prev, cur in zip(times, times[1:], strict=False): |
| if cur < prev: |
| raise ValueError(f"times must be sorted; got {cur!r} after {prev!r}") |
| for i, t in enumerate(times): |
| if t >= value: |
| return self.split(i) |
| return self.split(len(times)) |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"LatticeWindow({len(self)} windows, input_len={self.input_len}, " |
| f"horizon={self.horizon}, stride={self.stride})" |
| ) |
|
|