Coverage for src/tinycta/engine.py: 100%
59 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-30 04:16 +0000
1"""Engine for correlation-aware risk position optimization.
3This module is the Polars-facing orchestration layer: :class:`Engine` validates and holds
4the aligned ``prices``/``mu`` frames, derives the volatility-adjusted returns and per-timestamp
5EWMA correlation matrices, and hands the resulting NumPy arrays to the pure-numeric kernel in
6:mod:`tinycta._kernel` for the forward walk.
7"""
9from __future__ import annotations
11import dataclasses
12from collections.abc import Hashable
14import numpy as np
15import polars as pl
17from ._kernel import forward_walk as _forward_walk
18from .config import Config
19from .ewm_cov import ewm_covariance as _ewm_covariance
20from .util import vol_adj as _vol_adj
23@dataclasses.dataclass(frozen=True)
24class Engine:
25 """Correlation-aware risk position optimizer (Basanos engine)."""
27 prices: pl.DataFrame
28 mu: pl.DataFrame
29 cfg: Config
31 def __post_init__(self) -> None:
32 """Validate that prices and mu are aligned and both contain a date column."""
33 if "date" not in self.prices.columns:
34 msg = "prices must contain a 'date' column"
35 raise ValueError(msg)
36 if "date" not in self.mu.columns:
37 msg = "mu must contain a 'date' column"
38 raise ValueError(msg)
39 if self.prices.shape != self.mu.shape:
40 msg = f"prices and mu must share the same shape, got {self.prices.shape} and {self.mu.shape}"
41 raise ValueError(msg)
42 if set(self.prices.columns) != set(self.mu.columns):
43 msg = "prices and mu must share identical columns"
44 raise ValueError(msg)
46 @property
47 def assets(self) -> list[str]:
48 """List numeric asset column names, excluding the date column."""
49 return [c for c in self.prices.columns if c != "date" and self.prices[c].dtype.is_numeric()]
51 @property
52 def ret_adj(self) -> pl.DataFrame:
53 """Per-asset EWMA-volatility-adjusted log returns clipped by cfg.clip."""
54 return self.prices.with_columns(
55 [_vol_adj(pl.col(asset), vola=self.cfg.vola, clip=self.cfg.clip) for asset in self.assets]
56 )
58 @property
59 def vola(self) -> pl.DataFrame:
60 """Per-asset EWMA volatility of percentage returns."""
61 return self.prices.with_columns(
62 pl.col(asset)
63 .pct_change()
64 .ewm_std(com=self.cfg.vola - 1, adjust=True, min_samples=self.cfg.vola)
65 .alias(asset)
66 for asset in self.assets
67 )
69 @property
70 def cor(self) -> dict[Hashable, np.ndarray]:
71 """Per-timestamp EWMA correlation matrices, keyed by index value.
73 Each key is a value of the ``date`` column (a ``datetime.date`` in normal
74 use, but any hashable index value such as an integer is supported, hence
75 the ``Hashable`` key type). Each value is the EWMA covariance matrix at
76 that timestamp normalised to a correlation matrix (unit diagonal).
78 Contract:
79 - **Warmup:** the first ``cfg.corr + 1`` timestamps are omitted — a
80 key exists only once at least one matrix cell is finite (see
81 :func:`~tinycta.ewm_cov.ewm_covariance`). That takes ``cfg.corr``
82 observations of :attr:`ret_adj`, which itself starts on the third
83 row because ``vol_adj`` needs two log returns to standardise one.
84 - **NaN cells:** a cell is ``NaN`` while either asset is still in its
85 own warmup, and a zero-variance asset (``outer == 0``) yields ``NaN``
86 correlations rather than a divide-by-zero.
87 """
88 cov = _ewm_covariance(
89 self.ret_adj,
90 assets=self.assets,
91 index_col="date",
92 window=2 * self.cfg.corr + 1,
93 warmup=self.cfg.corr,
94 )
95 result: dict[Hashable, np.ndarray] = {}
96 for k, mat in cov.items():
97 std = np.sqrt(np.abs(np.diag(mat)))
98 outer = np.outer(std, std)
99 # Divide only where the variance product is positive; zero-variance
100 # cells stay NaN. Computing ``mat / outer`` eagerly (before masking)
101 # would divide by zero on those cells and emit a spurious RuntimeWarning.
102 result[k] = np.divide(mat, outer, out=np.full(mat.shape, np.nan), where=outer > 0)
103 return result
105 @property
106 def cash_position(self) -> pl.DataFrame:
107 """Correlation-shrinkage-optimized cash positions for each timestamp.
109 Walks forward through time, and at each timestamp ``t``:
111 1. **Mask** assets with a finite price at ``t`` so the optimisation only
112 sees currently-tradable instruments.
113 2. **Shrink** the EWMA correlation matrix towards the identity by
114 ``cfg.shrink`` (via :func:`~tinycta.signal.shrink2id`) for numerical
115 stability, then restrict it to the masked assets.
116 3. **Solve** the shrunk system for the expected returns ``mu`` and
117 normalise by ``inv_a_norm(mu, matrix)`` so the raw risk position has
118 unit norm under the correlation metric (zeroed when the denominator
119 is non-finite/degenerate or ``mu`` is all-zero).
120 4. **Scale** the risk position by a running EWMA estimate of realised
121 profit variance (decay ``lamb=0.99``), which down-weights positions
122 after volatile P&L, then divide by per-asset EWMA volatility
123 (``self.vola``) to convert the risk position into a cash position.
125 The per-timestamp walk itself is delegated to
126 :func:`tinycta._kernel.forward_walk`, which operates purely on NumPy arrays.
128 Returns:
129 pl.DataFrame: The input ``prices`` frame (including its ``date``
130 column) with each asset column replaced by its per-timestamp
131 cash position. Warmup rows are ``NaN``.
133 Example:
134 >>> import polars as pl
135 >>> from tinycta.config import Config
136 >>> from tinycta.engine import Engine
137 >>> prices = pl.DataFrame({"date": [1, 2, 3], "A": [100.0, 101.0, 102.0]})
138 >>> mu = pl.DataFrame({"date": [1, 2, 3], "A": [0.0, 0.1, 0.2]})
139 >>> engine = Engine(prices=prices, mu=mu, cfg=Config(vola=2, corr=2, clip=4.2, shrink=0.5))
140 >>> positions = engine.cash_position
141 """
142 cor = self.cor
143 assets = self.assets
145 prices_num = self.prices.select(assets).to_numpy()
146 returns_num = np.zeros_like(prices_num, dtype=float)
147 returns_num[1:] = prices_num[1:] / prices_num[:-1] - 1.0
149 mu = self.mu.select(assets).to_numpy()
150 risk_pos_np = np.full_like(mu, fill_value=np.nan, dtype=float)
151 cash_pos_np = np.full_like(mu, fill_value=np.nan, dtype=float)
152 vola_np = self.vola.select(assets).to_numpy()
154 # ``cor`` is keyed by the post-warmup dates. Map each key back to its row
155 # in prices/mu/vola so the correlation matrix for date ``t`` is paired with
156 # (and stored at) that same date, rather than at a positional offset of
157 # ``corr`` rows — otherwise the most recent dates never receive a position.
158 row_of = {date: idx for idx, date in enumerate(self.prices["date"].to_list())}
160 _forward_walk(cor, prices_num, returns_num, mu, vola_np, risk_pos_np, cash_pos_np, row_of, self.cfg.shrink)
162 return self.prices.with_columns([(pl.lit(cash_pos_np[:, i]).alias(asset)) for i, asset in enumerate(assets)])