File size: 4,127 Bytes
ecc81b3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
"""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})"
        )