Coverage for src/basanos/math/_stream_state.py: 100%
45 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-08-04 07:53 +0000
1"""State containers for the incremental (streaming) optimiser.
3Defines the mutable per-instance state carrier (`_StreamState`), the frozen
4per-step output (`StepResult`), and the on-disk format version that the
5persistence layer in :mod:`basanos.math._stream_io` writes and validates.
7Keeping these types in their own module lets the state layout be read and
8tested in isolation, and lets the persistence and solver layers depend on the
9state representation without importing the `BasanosStream` façade.
10"""
12from __future__ import annotations
14import dataclasses
16import numpy as np
18from ._engine_solve import SolveStatus
20#: Increment this when the ``save`` archive layout changes in
21#: a backward-incompatible way. ``load`` asserts the stored
22#: value matches before deserialising anything, so callers get a clear error
23#: instead of a silent ``KeyError`` or wrong state.
24_SAVE_FORMAT_VERSION: int = 3
27@dataclasses.dataclass
28class _StreamState:
29 """Mutable state carrier for one `BasanosStream` instance.
31 All arrays are updated in-place (or replaced) by ``BasanosStream.step()``.
32 The class is intentionally *not* frozen so that the step method can modify
33 fields directly without creating a new object on every tick.
35 EWM correlation state
36 ~~~~~~~~~~~~~~~~~~~~~
37 ``corr_ret_buf`` holds the growing history of vol-adjusted returns used by
38 ``ewm_covariance`` to recompute the correlation matrix on each step. It
39 is ``None`` for ``SlidingWindowConfig`` (which uses ``sw_ret_buf`` instead).
41 EWM accumulator state (volatility)
42 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
43 ``vola_*`` and ``pct_*`` accumulate the running weighted sums needed to
44 compute exponentially-weighted standard deviations:
46 * ``s_x`` — EWM sum of x (numerator of the mean)
47 * ``s_x2`` — EWM sum of x² (numerator of the second moment)
48 * ``s_w`` — EWM sum of weights (denominator)
49 * ``s_w2`` — EWM sum of squared weights (for bias correction)
51 ``beta_vola = (cfg.vola - 1) / cfg.vola`` (from ``com = cfg.vola - 1``)
53 Attributes:
54 corr_ret_buf: Growing history of vol-adjusted returns used by
55 ``ewm_covariance``; shape ``(T, N)`` for EwmaShrinkConfig.
56 ``None`` for SlidingWindowConfig.
57 vola_s_x: EWM sum of volatility-adjusted log-returns; shape ``(N,)``.
58 vola_s_x2: EWM sum of squared vol-adj log-returns; shape ``(N,)``.
59 vola_s_w: EWM weight sum for vol accumulators; shape ``(N,)``.
60 vola_s_w2: EWM squared-weight sum for vol accumulators; shape ``(N,)``.
61 vola_count: Cumulative finite observation count for vol; shape ``(N,)``
62 dtype int.
63 pct_s_x: EWM sum of pct-returns; shape ``(N,)``.
64 pct_s_x2: EWM sum of squared pct-returns; shape ``(N,)``.
65 pct_s_w: EWM weight sum for pct accumulators; shape ``(N,)``.
66 pct_s_w2: EWM squared-weight sum for pct accumulators; shape ``(N,)``.
67 pct_count: Cumulative finite observation count for pct-return vol;
68 shape ``(N,)`` dtype int.
69 prev_price: Last price row seen, used to compute returns on the next
70 step; shape ``(N,)``.
71 prev_cash_pos: Last cash position, used to apply the turnover constraint
72 on the next step; shape ``(N,)``.
73 step_count: Number of steps processed so far (0 before first step).
74 """
76 # ── EWM correlation history buffer — (T, N) for EwmaShrinkConfig ─────────
77 corr_ret_buf: np.ndarray | None # (T, N) growing history of vol-adj returns; None for SlidingWindowConfig
79 # ── EWMA accumulators for vol_adj (log-return std; com=vola-1, min_samples=1) ──
80 vola_s_x: np.ndarray # (N,)
81 vola_s_x2: np.ndarray # (N,)
82 vola_s_w: np.ndarray # (N,)
83 vola_s_w2: np.ndarray # (N,)
84 vola_count: np.ndarray # (N,) int
86 # ── EWMA accumulators for vola (pct-return std; com=vola-1, min_samples=vola) ──
87 pct_s_x: np.ndarray # (N,)
88 pct_s_x2: np.ndarray # (N,)
89 pct_s_w: np.ndarray # (N,)
90 pct_s_w2: np.ndarray # (N,)
91 pct_count: np.ndarray # (N,) int
93 # ── Scalars ───────────────────────────────────────────────────────────────
94 prev_price: np.ndarray # (N,) last price row (to compute returns at next step)
95 prev_cash_pos: np.ndarray # (N,) last cash position (for turnover constraint at next step)
96 step_count: int
98 # ── SlidingWindowConfig state — None for EwmaShrinkConfig ────────────────
99 # shape (W, N): last W vol-adjusted returns (oldest row first); None when
100 # using EwmaShrinkConfig. corr_ret_buf above is unused (None) in this mode;
101 # sw_ret_buf carries all the correlation state instead.
102 sw_ret_buf: np.ndarray | None = None # (W, N) rolling buffer, or None
104 def persist(
105 self,
106 *,
107 corr_ret_buf: np.ndarray | None,
108 vola_s_x: np.ndarray,
109 vola_s_x2: np.ndarray,
110 vola_s_w: np.ndarray,
111 vola_s_w2: np.ndarray,
112 vola_count: np.ndarray,
113 pct_s_x: np.ndarray,
114 pct_s_x2: np.ndarray,
115 pct_s_w: np.ndarray,
116 pct_s_w2: np.ndarray,
117 pct_count: np.ndarray,
118 new_price: np.ndarray,
119 new_cash_pos: np.ndarray | None = None,
120 ) -> None:
121 """Persist accumulators, last-seen vectors, and increment step count."""
122 self.corr_ret_buf = corr_ret_buf
123 self.vola_s_x = vola_s_x
124 self.vola_s_x2 = vola_s_x2
125 self.vola_s_w = vola_s_w
126 self.vola_s_w2 = vola_s_w2
127 self.vola_count = vola_count
128 self.pct_s_x = pct_s_x
129 self.pct_s_x2 = pct_s_x2
130 self.pct_s_w = pct_s_w
131 self.pct_s_w2 = pct_s_w2
132 self.pct_count = pct_count
133 self.prev_price = new_price.copy()
134 if new_cash_pos is not None:
135 self.prev_cash_pos = new_cash_pos.copy()
136 self.step_count += 1
139#: Keys that ``save`` writes to the ``.npz`` archive for
140#: `_StreamState` fields. Derived automatically from
141#: `dataclasses.fields` so that adding a new field to ``_StreamState``
142#: is sufficient — no manual update here is required.
143#:
144#: The three non-state keys (``format_version``, ``cfg_json``, ``assets``) are
145#: added explicitly because they are not fields of ``_StreamState`` itself.
146_REQUIRED_KEYS: frozenset[str] = frozenset(
147 {f.name for f in dataclasses.fields(_StreamState)} | {"format_version", "cfg_json", "assets"}
148)
151@dataclasses.dataclass(frozen=True)
152class StepResult:
153 """Frozen dataclass representing the output of a single ``BasanosStream`` step.
155 Each call to ``BasanosStream.step()`` returns one ``StepResult`` capturing
156 the optimised cash positions, the per-asset volatility estimate, the step
157 date, and a status label that describes the solver outcome for that
158 timestep.
160 Attributes:
161 date: The timestamp or date label for this step. The type mirrors
162 whatever is stored in the ``'date'`` column of the input prices
163 DataFrame (typically a Python `date`,
164 `datetime`, or a Polars temporal scalar).
165 cash_position: Optimised cash-position vector, shape ``(N,)``.
166 Entries are ``NaN`` for assets that are still in the EWMA warmup
167 period or that are otherwise inactive at this step.
168 status: Solver outcome label for this timestep
169 (`SolveStatus`). Since `SolveStatus`
170 is a ``StrEnum``, values compare equal to their string equivalents
171 (e.g. ``result.status == "valid"`` is ``True``):
173 * ``'warmup'`` — fewer rows have been seen than the EWMA warmup
174 requires; all positions are ``NaN``.
175 * ``'zero_signal'`` — the expected-return signal vector ``mu`` is
176 identically zero; positions are set to zero rather than solved.
177 * ``'degenerate'`` — the covariance matrix is ill-conditioned or
178 numerically singular; positions cannot be computed reliably and
179 are returned as ``NaN``.
180 * ``'valid'`` — normal operation; ``cash_position`` holds the
181 optimised allocations.
182 vola: Per-asset EWMA percentage-return volatility, shape ``(N,)``.
183 Values are ``NaN`` during the warmup period before the EWMA has
184 accumulated sufficient history.
186 Examples:
187 >>> import numpy as np
188 >>> result = StepResult(
189 ... date="2024-01-02",
190 ... cash_position=np.array([1000.0, -500.0]),
191 ... status="valid",
192 ... vola=np.array([0.012, 0.018]),
193 ... )
194 >>> result.status
195 'valid'
196 >>> result.cash_position.shape
197 (2,)
198 """
200 date: object
201 cash_position: np.ndarray
202 status: SolveStatus
203 vola: np.ndarray