Coverage for src/basanos/math/_engine_solve.py: 100%
140 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-08-04 07:53 +0000
1"""Solve/position mixin for BasanosEngine.
3This private module contains `_SolveMixin`, which provides the
4``_iter_matrices`` and ``_iter_solve`` generator methods. Separating them
5from `optimizer` keeps the engine facade lean and makes
6the per-timestamp solve logic independently readable and testable.
7"""
9from __future__ import annotations
11import datetime
12import logging
13from collections.abc import Generator
14from typing import TYPE_CHECKING, cast
16import numpy as np
17from cvx.linalg import SingularMatrixError, inv_a_norm, solve
19from ._config import EwmaShrinkConfig, SlidingWindowConfig
20from ._engine_solve_base import MatrixBundle as MatrixBundle
21from ._engine_solve_base import MatrixYield as MatrixYield
22from ._engine_solve_base import SolveStatus as SolveStatus
23from ._engine_solve_base import SolveYield as SolveYield
24from ._engine_solve_base import WarmupState as WarmupState
25from ._engine_solve_base import _SolvePrimitivesMixin
26from ._factor_model import FactorModel
27from ._signal import shrink2id
29if TYPE_CHECKING:
30 from ._engine_protocol import _EngineProtocol
32_logger = logging.getLogger(__name__)
35class _SolveMixin(_SolvePrimitivesMixin):
36 """Mixin that provides ``_iter_matrices`` and ``_iter_solve`` generators.
38 Inherits the stateless helpers from `_SolvePrimitivesMixin` and adds the
39 per-timestamp solve orchestration. Consumers must also inherit from (or
40 satisfy the interface of) `_EngineProtocol` so that
41 ``self.assets``, ``self.prices``, ``self.mu``, ``self.cfg``, ``self.cor``,
42 and ``self.ret_adj`` are all available.
43 """
45 @staticmethod
46 def _compute_position(
47 i: int,
48 t: datetime.date,
49 mask: np.ndarray,
50 expected_mu: np.ndarray,
51 bundle: MatrixBundle,
52 denom_tol: float,
53 ) -> SolveYield:
54 """Shared solve step used by both covariance branches.
56 Computes the normalisation denominator via `inv_a_norm`
57 and solves the linear system via `solve`, then
58 delegates to `_denom_guard_yield`. Handles
59 :exc:`~basanos.exceptions.SingularMatrixError` from both calls.
61 Accepting a `MatrixBundle` instead of a raw array means future
62 covariance modes can attach auxiliary state to the bundle without
63 changing this method's signature.
65 Args:
66 i: Row index.
67 t: Timestamp.
68 mask: Boolean asset mask of shape ``(n_assets,)``.
69 expected_mu: Masked expected-return vector of shape ``(n_active,)``.
70 bundle: Covariance bundle whose ``matrix`` field is an
71 ``(n_active, n_active)`` covariance matrix for the active assets.
72 denom_tol: Tolerance threshold for the normalisation denominator.
74 Returns:
75 SolveYield: A degenerate or valid ``(i, t, mask, pos, status)`` tuple.
76 """
77 matrix = bundle.matrix
78 try:
79 denom = inv_a_norm(expected_mu, matrix)
80 except SingularMatrixError:
81 denom = float("nan")
82 try:
83 pos = solve(matrix, expected_mu)
84 except SingularMatrixError:
85 return i, t, mask, np.zeros_like(expected_mu), SolveStatus.DEGENERATE
86 return _SolveMixin._denom_guard_yield(i, t, mask, expected_mu, pos, denom, denom_tol)
88 def _replay_positions(
89 self: _EngineProtocol,
90 risk_pos_np: np.ndarray,
91 cash_pos_np: np.ndarray,
92 vola_np: np.ndarray,
93 ) -> None:
94 """Replay positions across all rows, filling position arrays.
96 Iterates `_iter_solve`, writes risk and cash positions into the
97 provided pre-allocated arrays. Both arrays are mutated **in-place**.
99 When `max_turnover` is set, the L1 norm of the
100 position change ``sum(|x_t - x_{t-1}|)`` is capped at that value by
101 proportionally scaling the delta toward the previous position before
102 writing to ``cash_pos_np``.
104 Args:
105 risk_pos_np: Pre-allocated ``(T, N)`` array for risk positions.
106 cash_pos_np: Pre-allocated ``(T, N)`` array for cash positions.
107 vola_np: ``(T, N)`` EWMA volatility array.
108 """
109 max_to: float | None = self.cfg.max_turnover
110 for i, _t, mask, pos, _status in self._iter_solve():
111 if pos is not None:
112 new_cash = _SolveMixin._scale_to_cash(pos, vola_np[i, mask])
113 if max_to is not None and i > 0:
114 new_cash = _SolveMixin._apply_turnover_constraint(new_cash, cash_pos_np[i - 1, mask], max_to)
115 risk_pos_np[i, mask] = new_cash * vola_np[i, mask]
116 cash_pos_np[i, mask] = new_cash
118 def _iter_matrices(self: _EngineProtocol) -> Generator[MatrixYield, None, None]:
119 r"""Yield ``(i, t, mask, bundle)`` for every timestamp.
121 ``bundle`` is a `MatrixBundle` wrapping the effective
122 $(n_{\text{sub}},\ n_{\text{sub}})$ correlation matrix for the
123 active assets (those with finite prices at timestamp *t*). Yields
124 ``None`` when no valid matrix is available (e.g., before the warm-up
125 period has elapsed or when no assets have finite prices).
127 The behaviour depends on `covariance_config`:
129 * `EwmaShrinkConfig`: Applies `shrink2id` to
130 the EWMA correlation matrix (same computation as
131 `cash_position`).
132 * `SlidingWindowConfig`: Builds a
133 `FactorModel` from the last
134 ``cfg.covariance_config.window`` rows of vol-adjusted returns and returns its
135 `covariance`.
137 Yields:
138 tuple: ``(i, t, mask, bundle)`` where
140 * ``i`` (*int*): Row index into ``self.prices``.
141 * ``t``: Timestamp value from ``self.prices["date"]``.
142 * ``mask`` (*np.ndarray[bool]*): Shape ``(n_assets,)``; ``True``
143 for assets with finite prices at row *i*.
144 * ``bundle`` (`MatrixBundle` | ``None``): Covariance bundle
145 of shape ``(mask.sum(), mask.sum())``, or ``None``.
146 """
147 prices_num = self.prices.select(self.assets).to_numpy()
148 dates = self.prices["date"].to_list()
150 if isinstance(self.cfg.covariance_config, EwmaShrinkConfig):
151 yield from _SolveMixin._iter_matrices_ewma(self, prices_num, dates)
152 else:
153 yield from _SolveMixin._iter_matrices_sliding(self, prices_num, dates)
155 def _iter_matrices_ewma(
156 self: _EngineProtocol,
157 prices_num: np.ndarray,
158 dates: list[datetime.date],
159 ) -> Generator[MatrixYield, None, None]:
160 """Yield per-timestamp `MatrixYield` for the `EwmaShrinkConfig` path.
162 Applies `shrink2id` to the EWMA correlation matrix and restricts it to
163 the active-asset sub-matrix; yields ``None`` when no asset has a finite
164 price at that row.
165 """
166 cor = self.cor
167 for i, t in enumerate(dates):
168 mask = _SolveMixin._compute_mask(prices_num[i])
169 if not mask.any():
170 yield i, t, mask, None
171 continue
172 corr_n = cor[t]
173 matrix = shrink2id(corr_n, lamb=self.cfg.shrink)[np.ix_(mask, mask)]
174 yield i, t, mask, MatrixBundle(matrix=matrix)
176 def _iter_matrices_sliding(
177 self: _EngineProtocol,
178 prices_num: np.ndarray,
179 dates: list[datetime.date],
180 ) -> Generator[MatrixYield, None, None]:
181 """Yield per-timestamp `MatrixYield` for the `SlidingWindowConfig` path.
183 Fits a `FactorModel` from the last ``window`` rows of vol-adjusted
184 returns; yields ``None`` during warm-up, when no asset has a finite
185 price, or when the factor-model fit fails.
186 """
187 sw_config = cast(SlidingWindowConfig, self.cfg.covariance_config)
188 win_w: int = sw_config.window
189 win_k: int = sw_config.n_factors
190 ret_adj_np = self.ret_adj.select(self.assets).to_numpy()
191 for i, t in enumerate(dates):
192 mask = _SolveMixin._compute_mask(prices_num[i])
193 if not mask.any() or i + 1 < win_w:
194 yield i, t, mask, None
195 continue
196 window_ret = ret_adj_np[i + 1 - win_w : i + 1][:, mask]
197 window_ret = np.where(np.isfinite(window_ret), window_ret, 0.0)
198 n_sub = int(mask.sum())
199 k_eff = min(win_k, win_w, n_sub)
200 try:
201 fm = FactorModel.from_returns(window_ret, k=k_eff)
202 yield i, t, mask, MatrixBundle(matrix=fm.covariance)
203 except (np.linalg.LinAlgError, ValueError) as exc:
204 _logger.warning("Factor model fit failed at t=%s: %s", t, exc)
205 yield i, t, mask, None
207 @staticmethod
208 def _batched_solve_group(
209 group: list[tuple[int, datetime.date, np.ndarray, np.ndarray, np.ndarray]],
210 denom_tol: float,
211 ) -> dict[int, SolveYield]:
212 """Solve a batch of linear systems sharing the same active-asset mask.
214 Stacks the ``len(group)`` systems into a ``(G, n, n)`` coefficient tensor
215 and a ``(G, n)`` right-hand-side matrix, then dispatches a single
216 ``numpy.linalg.solve`` call (which maps to a single batched LAPACK
217 routine). Denominators are computed directly from the batch result as
218 ``sqrt(mu_i · pos_i)`` — algebraically identical to the per-row
219 `inv_a_norm` call.
221 Falls back to row-by-row `_compute_position` when
222 ``numpy.linalg.solve`` raises ``LinAlgError`` (any matrix in the batch
223 is singular).
225 Args:
226 group: List of ``(i, t, mask, expected_mu, matrix)`` tuples; all
227 entries share the same boolean mask and therefore the same
228 ``n_active x n_active`` matrix shape.
229 denom_tol: Passed through to `_denom_guard_yield`.
231 Returns:
232 dict: Mapping from row index ``i`` to its `SolveYield`.
233 """
234 results: dict[int, SolveYield] = {}
235 a_stack = np.stack([row[4] for row in group]) # (G, n, n)
236 mu_stack = np.stack([row[3] for row in group]) # (G, n)
238 try:
239 # numpy.linalg.solve requires the RHS to be (..., M, K) when a is (..., M, M).
240 # Reshape mu_stack from (G, n) → (G, n, 1) so core dims match, then squeeze.
241 pos_stack = np.linalg.solve(a_stack, mu_stack[..., np.newaxis])[..., 0] # (G, n)
242 except np.linalg.LinAlgError:
243 # At least one matrix is singular — fall back to sequential per-row solve.
244 return _SolveMixin._sequential_solve_group(group, denom_tol)
246 # Denominators: sqrt(mu_i^T A_i^{-1} mu_i) = sqrt(mu_i · pos_i).
247 dots = (mu_stack * pos_stack).sum(axis=1) # (G,)
248 denoms = np.where(dots > 0.0, np.sqrt(dots), np.nan)
250 for (i, t, mask, expected_mu, _matrix), pos, denom in zip(group, pos_stack, denoms, strict=True):
251 results[i] = _SolveMixin._denom_guard_yield(i, t, mask, expected_mu, pos, float(denom), denom_tol)
253 return results
255 @staticmethod
256 def _sequential_solve_group(
257 group: list[tuple[int, datetime.date, np.ndarray, np.ndarray, np.ndarray]],
258 denom_tol: float,
259 ) -> dict[int, SolveYield]:
260 """Row-by-row fallback used when a batched solve hits a singular matrix.
262 Solves each system in *group* independently via `_compute_position` so a
263 single ill-conditioned matrix does not abort the whole batch.
264 """
265 return {
266 i: _SolveMixin._compute_position(i, t, mask, expected_mu, MatrixBundle(matrix=matrix), denom_tol)
267 for i, t, mask, expected_mu, matrix in group
268 }
270 @staticmethod
271 def _iter_solve_ewma_batched(
272 mu_np: np.ndarray,
273 matrix_yields: list[MatrixYield],
274 denom_tol: float,
275 ) -> Generator[SolveYield, None, None]:
276 r"""Vectorised EwmaShrink solve: batch ``numpy.linalg.solve`` across timestamps.
278 Groups rows by their boolean asset mask so all systems within a group
279 share the same ``(n_active, n_active)`` shape, then stacks them into a
280 ``(G, n, n)`` tensor and calls ``numpy.linalg.solve`` once per unique
281 mask pattern. Results are collected in a dict and yielded in original
282 row order.
284 Denominators are derived from the batch solution as
285 $\sqrt{\mu_i \cdot \mathbf{pos}_i} = \sqrt{\mu_i^\top \Sigma_i^{-1} \mu_i}$,
286 matching the scalar `inv_a_norm` result up
287 to float64 rounding.
289 Any group whose batch solve raises ``LinAlgError`` (singular matrix in
290 the batch) falls back to sequential `_compute_position` for that
291 group only.
293 Args:
294 mu_np: Signal matrix, shape ``(T, n_assets)``.
295 matrix_yields: Pre-collected list from `_iter_matrices`
296 (the EwmaShrinkConfig branch).
297 denom_tol: Denominator guard tolerance.
299 Yields:
300 `SolveYield` tuples in original row order.
301 """
302 # First pass: categorise each row as early-exit or a solve candidate.
303 all_results, solve_groups = _SolveMixin._partition_ewma_rows(mu_np, matrix_yields)
305 # Second pass: batch-solve each mask group.
306 for group in solve_groups.values():
307 all_results.update(_SolveMixin._batched_solve_group(group, denom_tol))
309 # Yield in original row order.
310 for i in range(len(matrix_yields)):
311 if i in all_results:
312 yield all_results[i]
314 @staticmethod
315 def _partition_ewma_rows(
316 mu_np: np.ndarray,
317 matrix_yields: list[MatrixYield],
318 ) -> tuple[
319 dict[int, SolveYield],
320 dict[bytes, list[tuple[int, datetime.date, np.ndarray, np.ndarray, np.ndarray]]],
321 ]:
322 """Split rows into resolved early-exits and mask-grouped solve candidates.
324 Returns ``(early_results, solve_groups)`` where ``early_results`` maps a
325 row index to its final `SolveYield` (no-data / warmup or an early
326 mask/signal exit) and ``solve_groups`` maps each ``mask.tobytes()`` key
327 to the rows sharing that active-asset pattern.
328 """
329 early_results: dict[int, SolveYield] = {}
330 # mask.tobytes() → list of (i, t, mask, expected_mu, matrix)
331 solve_groups: dict[bytes, list[tuple[int, datetime.date, np.ndarray, np.ndarray, np.ndarray]]] = {}
333 for i, t, mask, bundle in matrix_yields:
334 if bundle is None:
335 early_results[i] = (i, t, mask, np.zeros(int(mask.sum())), SolveStatus.DEGENERATE)
336 continue
337 expected_mu, early = _SolveMixin._row_early_check(i, t, mask, mu_np[i])
338 if early is not None:
339 early_results[i] = early
340 continue
341 solve_groups.setdefault(mask.tobytes(), []).append((i, t, mask, expected_mu, bundle.matrix))
343 return early_results, solve_groups
345 def _iter_solve(self: _EngineProtocol) -> Generator[SolveYield, None, None]:
346 r"""Yield ``(i, t, mask, pos_or_none, status)`` for every timestamp.
348 Iterates `_iter_matrices` for the per-row covariance sub-matrix,
349 then applies `_row_early_check` (mask/signal guard) and
350 `_compute_position` (linear solve and denominator guard). The two
351 covariance modes differ only in how ``matrix`` is built, which
352 `_iter_matrices` already encapsulates.
354 * ``matrix is None`` → `WARMUP` (sliding-window before
355 sufficient history) or `DEGENERATE` otherwise.
356 * Signal all-zero → `ZERO_SIGNAL`.
357 * Singular or degenerate solve → `DEGENERATE`.
358 * Success → `VALID`.
360 For the `EwmaShrinkConfig` path the solve step is
361 vectorised: rows are grouped by their active-asset mask pattern and each
362 group is solved via a single batched ``numpy.linalg.solve`` call (see
363 `_iter_solve_ewma_batched`). The `SlidingWindowConfig`
364 path retains a sequential per-row solve because the factor-model matrices
365 are constructed lazily and may vary in numerical character across rows.
367 .. note::
369 **Dual-path maintenance obligation**: this method dispatches to two
370 fundamentally different implementations. Any change to solve
371 semantics — a new edge case, a new `SolveStatus` value, or a
372 change to denominator logic — **must be applied to both branches**:
374 * `_iter_solve_ewma_batched` / `_batched_solve_group`
375 (EwmaShrink vectorised path)
376 * The sequential ``_compute_position`` loop below
377 (SlidingWindow path)
379 The cross-path numerical consistency test
380 ``test_ewma_batch_and_sequential_paths_agree`` in
381 ``tests/test_math/test_numerical_regression.py`` will fail
382 whenever the two paths drift apart, surfacing the divergence
383 before it reaches production.
385 Yields:
386 SolveYield: ``(i, t, mask, pos_or_none, status)`` — see
387 `SolveYield` for detailed field descriptions.
388 """
389 mu_np = self.mu.select(self.assets).to_numpy()
390 cov_config = self.cfg.covariance_config
392 if not isinstance(cov_config, SlidingWindowConfig):
393 # EwmaShrinkConfig path: vectorised batch solve grouped by mask pattern.
394 yield from _SolveMixin._iter_solve_ewma_batched(mu_np, list(self._iter_matrices()), self.cfg.denom_tol)
395 return
397 # SlidingWindowConfig path: sequential per-row solve (lazy factor models).
398 yield from _SolveMixin._iter_solve_sliding(self, mu_np, cov_config.window)
400 def _iter_solve_sliding(
401 self: _EngineProtocol,
402 mu_np: np.ndarray,
403 win_w: int,
404 ) -> Generator[SolveYield, None, None]:
405 """Sequential per-row solve for the `SlidingWindowConfig` path.
407 Factor-model matrices are constructed lazily and may vary in numerical
408 character across rows, so each row is solved individually via
409 `_compute_position` rather than batched.
410 """
411 for i, t, mask, bundle in self._iter_matrices():
412 if bundle is None:
413 yield _SolveMixin._sliding_warmup_or_degenerate(i, t, mask, win_w)
414 continue
415 expected_mu, early = _SolveMixin._row_early_check(i, t, mask, mu_np[i])
416 if early is not None:
417 yield early
418 continue
419 yield _SolveMixin._compute_position(i, t, mask, expected_mu, bundle, self.cfg.denom_tol)
421 def warmup_state(self: _EngineProtocol) -> WarmupState:
422 """Return the final `WarmupState` after replaying the full batch.
424 Encapsulates the position replay loop that was previously duplicated
425 inside `from_warmup`. By centralising the loop
426 here, `from_warmup` no longer needs to call the
427 private `_iter_solve` generator directly.
429 Returns:
430 WarmupState: A frozen dataclass with:
432 * ``prev_cash_pos`` - cash-position vector at the last row,
433 shape ``(n_assets,)``.
435 Examples:
436 >>> import numpy as np
437 >>> import polars as pl
438 >>> from basanos.math import BasanosConfig, BasanosEngine
439 >>> rng = np.random.default_rng(0)
440 >>> dates = list(range(30))
441 >>> prices = pl.DataFrame({
442 ... "date": dates,
443 ... "A": np.cumprod(1 + rng.normal(0.001, 0.02, 30)) * 100.0,
444 ... "B": np.cumprod(1 + rng.normal(0.001, 0.02, 30)) * 150.0,
445 ... })
446 >>> mu = pl.DataFrame({
447 ... "date": dates,
448 ... "A": rng.normal(0, 0.5, 30),
449 ... "B": rng.normal(0, 0.5, 30),
450 ... })
451 >>> cfg = BasanosConfig(vola=5, corr=10, clip=3.0, shrink=0.5, aum=1e6)
452 >>> engine = BasanosEngine(prices=prices, mu=mu, cfg=cfg)
453 >>> ws = engine.warmup_state()
454 >>> ws.prev_cash_pos.shape
455 (2,)
456 """
457 assets = self.assets
458 n_rows = self.prices.height
459 vola_np = self.vola.select(assets).to_numpy()
461 risk_pos_np = np.full((n_rows, len(assets)), np.nan, dtype=float)
462 cash_pos_np = np.full((n_rows, len(assets)), np.nan, dtype=float)
464 _SolveMixin._replay_positions(self, risk_pos_np, cash_pos_np, vola_np)
465 prev_cash_pos = cash_pos_np[-1].copy()
466 return WarmupState(prev_cash_pos=prev_cash_pos)