Coverage for src/basanos/math/_stream.py: 100%
132 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-08-04 07:53 +0000
1"""Incremental (streaming) API for BasanosEngine.
3This private module defines the `BasanosStream` façade: an incremental
4optimiser with a `from_warmup` classmethod and a `step` method. After warming
5up on a historical batch, each `step` call advances the internal state by
6exactly one row in O(N^2) time without revisiting the warmup history.
8The supporting pieces live in sibling modules so this file stays focused on the
9streaming loop itself:
11* :mod:`basanos.math._stream_state` — the mutable `_StreamState` carrier and the
12 frozen `StepResult` output.
13* :mod:`basanos.math._stream_math` — the pure EWMA recurrences and input
14 validation helpers.
15* :mod:`basanos.math._stream_solve` — the per-step position solvers for the
16 EWM and sliding-window covariance modes.
17* :mod:`basanos.math._stream_io` — save/load of the full stream state.
19EWM correlation state model
20----------------------------
21In EWM mode the correlation at each step is recomputed by calling
22``ewm_covariance`` from ``cvx.linalg`` over the full growing history of
23vol-adjusted returns stored in ``corr_ret_buf``. This keeps the incremental
24and batch paths numerically identical at the cost of O(T·N²) time per step
25(acceptable for small N or short warmup histories).
27The volatility accumulators (``vola_*``, ``pct_*``) use a simpler scalar
28recurrence and store the running sums directly as ``(N,)`` arrays.
30Memory
31------
32Total incremental state is O(T·N) for the growing history buffer plus
338x(N,) + O(1) scalars. For the SlidingWindowConfig the buffer is a fixed
34(W, N) array independent of T.
35"""
37from __future__ import annotations
39import dataclasses
40import os
41from typing import Any, cast
43import numpy as np
44import polars as pl
46from ..exceptions import MissingDateColumnError
47from ._config import BasanosConfig, EwmaShrinkConfig, SlidingWindowConfig
48from ._engine_solve import SolveStatus, _SolveMixin
49from ._stream_io import load_stream_archive, save_stream_archive
50from ._stream_math import _ewm_std_from_state, _ewm_vol_accumulators_from_batch, _resolve_step_vector
51from ._stream_solve import solve_ewma_position, solve_sliding_window_position
52from ._stream_state import StepResult as StepResult
53from ._stream_state import _StreamState
54from .optimizer import BasanosEngine
56# Number of leading rows for which ``vol_adj`` cannot produce a value: row 0
57# has no log return, and row 1 has a single observation, for which the
58# bias-corrected EWMA std is undefined. The correlation buffers seeded from
59# ``ret_adj`` therefore carry this many NaN prefix rows.
60_RET_ADJ_LEAD_IN = 2
63class BasanosStream:
64 """Incremental (streaming) optimiser backed by a single `_StreamState`.
66 After warming up on a historical batch via `from_warmup`, each call
67 to `step` advances the internal state by exactly one row in
68 O(N^2) time — without revisiting the full warmup history.
70 Attributes:
71 assets: Ordered list of asset column names (read-only).
73 Examples:
74 >>> import numpy as np
75 >>> import polars as pl
76 >>> from datetime import date, timedelta
77 >>> from basanos.math import BasanosConfig, BasanosStream
78 >>> rng = np.random.default_rng(0)
79 >>> warmup_len = 60
80 >>> dates = pl.date_range(
81 ... start=date(2024, 1, 1),
82 ... end=date(2024, 1, 1) + timedelta(days=warmup_len),
83 ... interval="1d",
84 ... eager=True,
85 ... )
86 >>> prices = pl.DataFrame({
87 ... "date": dates,
88 ... "A": np.cumprod(1 + rng.normal(0.001, 0.02, warmup_len + 1)) * 100.0,
89 ... "B": np.cumprod(1 + rng.normal(0.001, 0.02, warmup_len + 1)) * 150.0,
90 ... })
91 >>> mu = pl.DataFrame({
92 ... "date": dates,
93 ... "A": rng.normal(0, 0.5, warmup_len + 1),
94 ... "B": rng.normal(0, 0.5, warmup_len + 1),
95 ... })
96 >>> cfg = BasanosConfig(vola=5, corr=10, clip=3.0, shrink=0.5, aum=1e6)
97 >>> stream = BasanosStream.from_warmup(prices.head(warmup_len), mu.head(warmup_len), cfg)
98 >>> result = stream.step(
99 ... prices.select(["A", "B"]).to_numpy()[warmup_len],
100 ... mu.select(["A", "B"]).to_numpy()[warmup_len],
101 ... prices["date"][warmup_len],
102 ... )
103 >>> isinstance(result, StepResult)
104 True
105 >>> result.cash_position.shape
106 (2,)
107 """
109 _cfg: BasanosConfig
110 _assets: list[str]
111 _state: _StreamState
113 def __init__(self, cfg: BasanosConfig, assets: list[str], state: _StreamState) -> None:
114 """Initialise from an explicit config, asset list, and state container."""
115 object.__setattr__(self, "_cfg", cfg)
116 object.__setattr__(self, "_assets", assets)
117 object.__setattr__(self, "_state", state)
119 def __setattr__(self, name: str, value: object) -> None:
120 """Prevent accidental attribute mutation — BasanosStream is immutable."""
121 raise dataclasses.FrozenInstanceError(f"{type(self).__name__}.{name}")
123 @property
124 def assets(self) -> list[str]:
125 """Ordered list of asset column names."""
126 return self._assets
128 # ------------------------------------------------------------------
129 # from_warmup
130 # ------------------------------------------------------------------
132 @classmethod
133 def from_warmup(
134 cls,
135 prices: pl.DataFrame,
136 mu: pl.DataFrame,
137 cfg: BasanosConfig,
138 ) -> BasanosStream:
139 """Build a `BasanosStream` from a historical warmup batch.
141 Runs `BasanosEngine` on the full warmup batch
142 exactly once and extracts the minimal IIR-filter state required for
143 subsequent `step` calls. After this call, each `step`
144 advances the optimiser in O(N^2) time without touching the warmup
145 data again.
147 Parameters
148 ----------
149 prices:
150 Historical price DataFrame. Must contain a ``'date'`` column and
151 at least one numeric asset column with strictly positive,
152 non-monotonic values.
153 mu:
154 Expected-return signal DataFrame aligned row-by-row with
155 ``prices``.
156 cfg:
157 Engine configuration. Both `EwmaShrinkConfig`
158 and `SlidingWindowConfig` are supported.
160 Returns:
161 -------
162 BasanosStream
163 A stream instance whose `step` method is ready to accept the
164 row immediately following the last warmup row.
166 Notes:
167 ------
168 **Short-warmup behaviour with** ``SlidingWindowConfig``: when
169 ``len(prices) < cfg.covariance_config.window``, the internal rolling
170 buffer (``sw_ret_buf``) is NaN-padded for the missing prefix rows.
171 `step` returns ``StepResult(status="warmup")`` for each of the
172 first ``window - len(prices)`` calls, exactly matching the EWM warmup
173 semantics. By the time `step` returns the first non-warmup
174 result the buffer contains only real data — no NaN-padded rows remain.
176 Raises:
177 ------
178 MissingDateColumnError
179 If ``'date'`` is absent from ``prices``.
180 """
181 # 1. Validate -------------------------------------------------------
182 if "date" not in prices.columns:
183 raise MissingDateColumnError("prices")
185 # 2. Build the engine on the full warmup batch ----------------------
186 engine = BasanosEngine(prices=prices, mu=mu, cfg=cfg)
187 assets = engine.assets
188 n_assets = len(assets)
189 n_rows = prices.height
190 prices_np = prices.select(assets).to_numpy() # (n_rows, n_assets)
192 # 3. Extract mode-specific state from WarmupState --------------------
193 ws = engine.warmup_state()
194 if isinstance(cfg.covariance_config, EwmaShrinkConfig):
195 # EWM: seed the growing history buffer from engine.ret_adj so that
196 # each subsequent step() can call ewm_covariance over the full history.
197 ret_adj_np = engine.ret_adj.select(assets).to_numpy()
198 corr_ret_buf: np.ndarray | None = ret_adj_np
199 sw_ret_buf: np.ndarray | None = None
200 else:
201 # SW: carry the last W vol-adjusted returns as a rolling buffer.
202 sw_config = cfg.covariance_config
203 win_w = sw_config.window
204 ret_adj_np = engine.ret_adj.select(assets).to_numpy() # (n_rows, N)
205 if n_rows >= win_w:
206 sw_ret_buf = ret_adj_np[-win_w:].copy()
207 else:
208 sw_ret_buf = np.full((win_w, n_assets), np.nan)
209 sw_ret_buf[-n_rows:] = ret_adj_np
210 corr_ret_buf = None
212 # 4. Derive EWMA volatility accumulators (vectorised) ---------------
213 # Both log-return (for vol_adj) and pct-return (for vola) use the
214 # same beta = (vola-1)/vola. NaN observations (leading NaN at row 0
215 # from diff/pct_change) are skipped — the filter input is 0 for NaN
216 # rows and the weight accumulator (s_w) only increments for finite
217 # observations, matching Polars' effective behaviour for a
218 # leading-NaN series.
219 #
220 # Delegate to the shared helper _ewm_vol_accumulators_from_batch so
221 # that the batch and incremental recurrences share a single definition.
222 beta_vola: float = (cfg.vola - 1) / cfg.vola
223 beta_vola_sq: float = beta_vola**2
225 log_ret = np.full((n_rows, n_assets), np.nan, dtype=float)
226 pct_ret = np.full((n_rows, n_assets), np.nan, dtype=float)
227 if n_rows > 1:
228 with np.errstate(divide="ignore", invalid="ignore"):
229 log_ret[1:] = np.log(prices_np[1:] / prices_np[:-1])
230 pct_ret[1:] = prices_np[1:] / prices_np[:-1] - 1.0
232 vola_s_x, vola_s_x2, vola_s_w, vola_s_w2, vola_count = _ewm_vol_accumulators_from_batch(
233 log_ret, beta_vola, beta_vola_sq
234 )
235 pct_s_x, pct_s_x2, pct_s_w, pct_s_w2, pct_count = _ewm_vol_accumulators_from_batch(
236 pct_ret, beta_vola, beta_vola_sq
237 )
239 # 5. Extract prev_cash_pos from WarmupState --------------------------
240 prev_cash_pos: np.ndarray = ws.prev_cash_pos
241 prev_price: np.ndarray = prices_np[-1].copy()
243 # 6. Construct _StreamState and return ------------------------------
244 state = _StreamState(
245 corr_ret_buf=corr_ret_buf,
246 vola_s_x=vola_s_x,
247 vola_s_x2=vola_s_x2,
248 vola_s_w=vola_s_w,
249 vola_s_w2=vola_s_w2,
250 vola_count=vola_count,
251 pct_s_x=pct_s_x,
252 pct_s_x2=pct_s_x2,
253 pct_s_w=pct_s_w,
254 pct_s_w2=pct_s_w2,
255 pct_count=pct_count,
256 prev_price=prev_price,
257 prev_cash_pos=prev_cash_pos,
258 step_count=n_rows,
259 sw_ret_buf=sw_ret_buf,
260 )
261 return cls(cfg=cfg, assets=assets, state=state)
263 # ------------------------------------------------------------------
264 # step
265 # ------------------------------------------------------------------
267 @staticmethod
268 def _warmup_threshold(cfg: BasanosConfig) -> int:
269 """Return the step count at which warmup ends for the configured mode.
271 The sliding-window threshold adds ``_RET_ADJ_LEAD_IN`` so that
272 ``sw_ret_buf`` holds no NaN rows once warmup ends — ``ret_adj`` only
273 starts at row 2, and ``_fit_sliding_factor_model`` would otherwise
274 zero-fill the remaining NaN row into a fabricated observation.
276 EwmaShrink needs no such adjustment: ``ewm_covariance`` counts non-null
277 rows itself, so an under-warmed buffer yields a NaN correlation matrix
278 and a ``degenerate`` status — matching what `BasanosEngine` reports
279 for the same row.
280 """
281 if isinstance(cfg.covariance_config, SlidingWindowConfig):
282 return cfg.covariance_config.window + _RET_ADJ_LEAD_IN
283 return cfg.corr
285 @staticmethod
286 def _warmup_result(n_assets: int, date: Any) -> StepResult:
287 """Build a standard warmup ``StepResult`` payload."""
288 return StepResult(
289 date=date,
290 cash_position=np.full(n_assets, np.nan),
291 status=SolveStatus.WARMUP,
292 vola=np.full(n_assets, np.nan),
293 )
295 @staticmethod
296 def _apply_step_turnover(
297 cfg: BasanosConfig,
298 status: SolveStatus,
299 new_cash_pos: np.ndarray,
300 mask: np.ndarray,
301 prev_cash_pos: np.ndarray,
302 ) -> None:
303 """Cap the active-asset position change in place when a turnover limit is set."""
304 if cfg.max_turnover is not None and status == SolveStatus.VALID:
305 new_cash_pos[mask] = _SolveMixin._apply_turnover_constraint(
306 new_cash_pos[mask],
307 prev_cash_pos[mask],
308 cfg.max_turnover,
309 )
311 def step(
312 self,
313 new_prices: np.ndarray | dict[str, float],
314 new_mu: np.ndarray | dict[str, float],
315 date: Any = None,
316 ) -> StepResult:
317 """Advance the stream by one row and return the new optimised position.
319 Parameters
320 ----------
321 new_prices:
322 Per-asset prices for the new timestep. Either a numpy array of
323 shape ``(N,)`` (assets ordered as in `assets`) or a dict
324 mapping asset names to price values.
325 new_mu:
326 Per-asset expected-return signals, same format as ``new_prices``.
327 date:
328 Timestamp for this step (stored in `date`
329 verbatim; not used in any computation).
331 Returns:
332 -------
333 StepResult
334 Frozen dataclass with ``cash_position``, ``vola``, ``status``, and
335 ``date`` for this timestep.
336 """
337 cfg = self._cfg
338 assets = self._assets
339 state = self._state
340 n_assets = len(assets)
342 # ── Check if still in the warmup period ──────────────────────────────
343 # step_count is initialised to n_rows in from_warmup.
344 #
345 # EwmaShrinkConfig: in_warmup is True for the first (cfg.corr - n_rows)
346 # calls when the warmup batch was shorter than cfg.corr (not enough rows
347 # to populate the EWM correlation matrix).
348 #
349 # SlidingWindowConfig: in_warmup is True for the first
350 # (window + _RET_ADJ_LEAD_IN - n_rows) calls when the warmup batch was
351 # shorter than that. During this period sw_ret_buf still contains NaN
352 # rows — both the NaN padding and the ret_adj lead-in; each step shifts
353 # one NaN out and appends a real row, so the buffer is fully populated
354 # with real data exactly when in_warmup becomes False. This matters
355 # because _fit_sliding_factor_model zero-fills non-finite entries, so a
356 # residual NaN row would silently enter the covariance estimate as a
357 # fabricated all-zero observation.
358 #
359 # In both modes all accumulators are still updated during warmup so that
360 # the state is ready the moment the warmup period ends.
361 _warmup_thresh = self._warmup_threshold(cfg)
362 in_warmup: bool = state.step_count < _warmup_thresh
364 # ── Resolve inputs to (N,) float64 arrays ──────────────────────────
365 new_p = _resolve_step_vector(new_prices, assets, n_assets, "new_prices")
366 new_m = _resolve_step_vector(new_mu, assets, n_assets, "new_mu")
368 prev_p = state.prev_price
369 beta_vola: float = (cfg.vola - 1) / cfg.vola
370 beta_vola_sq: float = beta_vola**2
372 # ── Compute new log-returns and pct-returns ─────────────────────────
373 with np.errstate(divide="ignore", invalid="ignore"):
374 ratio = np.where(
375 np.isfinite(new_p) & np.isfinite(prev_p) & (prev_p > 0),
376 new_p / prev_p,
377 np.nan,
378 )
379 log_ret = np.log(ratio)
380 pct_ret = ratio - 1.0
382 # ── Update log-return EWMA accumulators ────────────────────────────
383 fin_log = np.isfinite(log_ret)
384 vola_s_x = beta_vola * state.vola_s_x + np.where(fin_log, log_ret, 0.0)
385 vola_s_x2 = beta_vola * state.vola_s_x2 + np.where(fin_log, log_ret**2, 0.0)
386 vola_s_w = beta_vola * state.vola_s_w + fin_log.astype(float)
387 vola_s_w2 = beta_vola_sq * state.vola_s_w2 + fin_log.astype(float)
388 vola_count = state.vola_count + fin_log.astype(int)
390 # ── Update pct-return EWMA accumulators ────────────────────────────
391 fin_pct = np.isfinite(pct_ret)
392 pct_s_x = beta_vola * state.pct_s_x + np.where(fin_pct, pct_ret, 0.0)
393 pct_s_x2 = beta_vola * state.pct_s_x2 + np.where(fin_pct, pct_ret**2, 0.0)
394 pct_s_w = beta_vola * state.pct_s_w + fin_pct.astype(float)
395 pct_s_w2 = beta_vola_sq * state.pct_s_w2 + fin_pct.astype(float)
396 pct_count = state.pct_count + fin_pct.astype(int)
398 # ── Compute vol-adjusted return (for the correlation IIR input) ─────
399 log_vol = _ewm_std_from_state(vola_s_x, vola_s_x2, vola_s_w, vola_s_w2, vola_count, min_samples=2)
400 # min_samples=2 mirrors vol_adj: a single observation has no defined
401 # bias-corrected std, so log_vol is NaN and vol_adj_val stays NaN.
402 with np.errstate(divide="ignore", invalid="ignore"):
403 vol_adj_val = np.where(
404 fin_log,
405 np.clip(log_ret / log_vol, -cfg.clip, cfg.clip),
406 np.nan,
407 )
409 # ── Mode-specific correlation state update ───────────────────────────
410 if isinstance(cfg.covariance_config, SlidingWindowConfig):
411 # SW: shift the rolling window buffer in-place and append this row.
412 buf = cast(np.ndarray, state.sw_ret_buf) # (W, N), already owned by state
413 buf[:-1] = buf[1:]
414 buf[-1] = vol_adj_val
415 corr_ret_buf = state.corr_ret_buf # None for SW; pass through
416 else:
417 # EWM: append new vol-adjusted return to the growing history buffer.
418 new_row = vol_adj_val[np.newaxis] # (1, N)
419 corr_ret_buf = np.vstack([cast(np.ndarray, state.corr_ret_buf), new_row])
421 # ── Early return during EWM warmup period ───────────────────────────
422 # All accumulators are already updated above; skip the O(N²) matrix
423 # reconstruction and O(N³) Cholesky solve which are wasteful during
424 # warmup — the computed positions would be discarded anyway.
425 if in_warmup:
426 state.persist(
427 corr_ret_buf=corr_ret_buf,
428 vola_s_x=vola_s_x,
429 vola_s_x2=vola_s_x2,
430 vola_s_w=vola_s_w,
431 vola_s_w2=vola_s_w2,
432 vola_count=vola_count,
433 pct_s_x=pct_s_x,
434 pct_s_x2=pct_s_x2,
435 pct_s_w=pct_s_w,
436 pct_s_w2=pct_s_w2,
437 pct_count=pct_count,
438 new_price=new_p,
439 )
440 return self._warmup_result(n_assets, date)
442 # ── Compute EWMA volatility (pct-return std) — shared ───────────────
443 vola_vec = _ewm_std_from_state(pct_s_x, pct_s_x2, pct_s_w, pct_s_w2, pct_count, min_samples=cfg.vola)
445 # ── Solve for position ───────────────────────────────────────────────
446 mask = np.isfinite(new_p)
447 if isinstance(cfg.covariance_config, SlidingWindowConfig):
448 new_cash_pos, status = solve_sliding_window_position(
449 cfg=cfg,
450 state=state,
451 mask=mask,
452 new_m=new_m,
453 vola_vec=vola_vec,
454 n_assets=n_assets,
455 date=date,
456 )
457 else:
458 new_cash_pos, status = solve_ewma_position(
459 cfg=cfg,
460 state=state,
461 corr_ret_buf=cast(np.ndarray, corr_ret_buf),
462 mask=mask,
463 new_m=new_m,
464 vola_vec=vola_vec,
465 assets=list(self._assets),
466 n_assets=n_assets,
467 date=date,
468 )
470 # ── Apply turnover constraint ─────────────────────────────────────────
471 self._apply_step_turnover(cfg, status, new_cash_pos, mask, state.prev_cash_pos)
473 # ── Persist updated state ───────────────────────────────────────────
474 state.persist(
475 corr_ret_buf=corr_ret_buf,
476 vola_s_x=vola_s_x,
477 vola_s_x2=vola_s_x2,
478 vola_s_w=vola_s_w,
479 vola_s_w2=vola_s_w2,
480 vola_count=vola_count,
481 pct_s_x=pct_s_x,
482 pct_s_x2=pct_s_x2,
483 pct_s_w=pct_s_w,
484 pct_s_w2=pct_s_w2,
485 pct_count=pct_count,
486 new_price=new_p,
487 new_cash_pos=new_cash_pos,
488 )
490 return StepResult(
491 date=date,
492 cash_position=new_cash_pos,
493 status=status,
494 vola=vola_vec,
495 )
497 # ------------------------------------------------------------------
498 # persistence
499 # ------------------------------------------------------------------
501 def save(self, path: str | os.PathLike[str]) -> None:
502 """Serialise the stream to a ``.npz`` archive at *path*.
504 Delegates to `basanos.math._stream_io.save_stream_archive`. A stream
505 restored via `load` produces bit-for-bit identical `step` output.
507 Args:
508 path: Destination file path. ``np.savez`` appends ``.npz``
509 automatically when the suffix is absent.
510 """
511 save_stream_archive(self._cfg, self._assets, self._state, path)
513 @classmethod
514 def load(cls, path: str | os.PathLike[str]) -> BasanosStream:
515 """Restore a stream previously saved with `save`.
517 Delegates to `basanos.math._stream_io.load_stream_archive`.
519 Args:
520 path: Path to a ``.npz`` archive written by `save`.
522 Returns:
523 A `BasanosStream` whose `step` output is bit-for-bit identical to
524 the original stream at the time `save` was called.
525 """
526 cfg, assets, state = load_stream_archive(path)
527 return cls(cfg=cfg, assets=assets, state=state)