Coverage for src/tinycta/engine.py: 100%
59 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-09-15 05:23 +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.
88 Example:
89 >>> import math
90 >>> import polars as pl
91 >>> from tinycta.config import Config
92 >>> from tinycta.engine import Engine
93 >>> dates = list(range(1, 11))
94 >>> prices = pl.DataFrame(
95 ... {
96 ... "date": dates,
97 ... "A": [100.0, 101.5, 100.8, 102.3, 103.1, 102.0, 104.5, 105.2, 104.1, 106.0],
98 ... "B": [50.0, 49.2, 50.4, 49.8, 51.1, 50.3, 49.5, 50.8, 51.6, 50.9],
99 ... }
100 ... )
101 >>> mu = pl.DataFrame({"date": dates, "A": [0.1] * 10, "B": [-0.05] * 10})
102 >>> cfg = Config(vola=3, corr=3, clip=4.2, shrink=0.5)
103 >>> cor = Engine(prices=prices, mu=mu, cfg=cfg).cor
105 The first ``cfg.corr + 1`` timestamps are omitted, so the mapping is
106 keyed by the surviving dates rather than by position:
108 >>> sorted(cor)
109 [5, 6, 7, 8, 9, 10]
111 Each value is a correlation matrix with a unit diagonal:
113 >>> mat = cor[5]
114 >>> mat.shape
115 (2, 2)
116 >>> [round(float(v), 6) for v in mat.diagonal()]
117 [1.0, 1.0]
118 >>> bool(-1.0 <= mat[0, 1] <= 1.0)
119 True
121 A zero-variance asset yields ``NaN`` rather than a divide-by-zero. Here
122 ``B`` never moves, so every cell touching it is ``NaN`` while ``A`` keeps
123 its unit diagonal:
125 >>> flat = prices.with_columns(pl.lit(50.0).alias("B"))
126 >>> flat_cor = Engine(prices=flat, mu=mu, cfg=cfg).cor[5]
127 >>> math.isnan(flat_cor[1, 1]), math.isnan(flat_cor[0, 1])
128 (True, True)
129 >>> round(float(flat_cor[0, 0]), 6)
130 1.0
131 """
132 cov = _ewm_covariance(
133 self.ret_adj,
134 assets=self.assets,
135 index_col="date",
136 window=2 * self.cfg.corr + 1,
137 warmup=self.cfg.corr,
138 )
139 result: dict[Hashable, np.ndarray] = {}
140 for k, mat in cov.items():
141 std = np.sqrt(np.abs(np.diag(mat)))
142 outer = np.outer(std, std)
143 # Divide only where the variance product is positive; zero-variance
144 # cells stay NaN. Computing ``mat / outer`` eagerly (before masking)
145 # would divide by zero on those cells and emit a spurious RuntimeWarning.
146 result[k] = np.divide(mat, outer, out=np.full(mat.shape, np.nan), where=outer > 0)
147 return result
149 @property
150 def cash_position(self) -> pl.DataFrame:
151 """Correlation-shrinkage-optimized cash positions for each timestamp.
153 Walks forward through time, and at each timestamp ``t``:
155 1. **Mask** assets with a finite price at ``t`` so the optimisation only
156 sees currently-tradable instruments.
157 2. **Shrink** the EWMA correlation matrix towards the identity by
158 ``cfg.shrink`` (via :func:`~tinycta.signal.shrink2id`) for numerical
159 stability, then restrict it to the masked assets.
160 3. **Solve** the shrunk system for the expected returns ``mu`` and
161 normalise by ``inv_a_norm(mu, matrix)`` so the raw risk position has
162 unit norm under the correlation metric (zeroed when the denominator
163 is non-finite/degenerate or ``mu`` is all-zero).
164 4. **Scale** the risk position by a running EWMA estimate of realised
165 profit variance (decay ``lamb=0.99``), which down-weights positions
166 after volatile P&L, then divide by per-asset EWMA volatility
167 (``self.vola``) to convert the risk position into a cash position.
169 The per-timestamp walk itself is delegated to
170 :func:`tinycta._kernel.forward_walk`, which operates purely on NumPy arrays.
172 Returns:
173 pl.DataFrame: The input ``prices`` frame (including its ``date``
174 column) with each asset column replaced by its per-timestamp
175 cash position. Warmup rows are ``NaN``.
177 Example:
178 >>> import math
179 >>> import polars as pl
180 >>> from tinycta.config import Config
181 >>> from tinycta.engine import Engine
182 >>> dates = list(range(1, 11))
183 >>> prices = pl.DataFrame(
184 ... {
185 ... "date": dates,
186 ... "A": [100.0, 101.5, 100.8, 102.3, 103.1, 102.0, 104.5, 105.2, 104.1, 106.0],
187 ... "B": [50.0, 49.2, 50.4, 49.8, 51.1, 50.3, 49.5, 50.8, 51.6, 50.9],
188 ... }
189 ... )
190 >>> mu = pl.DataFrame({"date": dates, "A": [0.1] * 10, "B": [-0.05] * 10})
191 >>> engine = Engine(prices=prices, mu=mu, cfg=Config(vola=3, corr=3, clip=4.2, shrink=0.5))
192 >>> positions = engine.cash_position
194 The frame keeps its shape and its ``date`` column; only the asset
195 columns are replaced:
197 >>> positions.shape
198 (10, 3)
199 >>> positions.columns
200 ['date', 'A', 'B']
202 The leading rows are warmup and come back ``NaN`` (here ``cfg.corr``
203 observations of :attr:`ret_adj`, which itself starts on the third row):
205 >>> sum(1 for v in positions["A"] if math.isnan(v))
206 4
208 After warmup the position takes the sign of the expected return, so a
209 positive ``mu`` is held long and a negative one short:
211 >>> [v > 0 for v in positions["A"][4:]]
212 [True, True, True, True, True, True]
213 >>> [v < 0 for v in positions["B"][4:]]
214 [True, True, True, True, True, True]
215 """
216 cor = self.cor
217 assets = self.assets
219 prices_num = self.prices.select(assets).to_numpy()
220 returns_num = np.zeros_like(prices_num, dtype=float)
221 returns_num[1:] = prices_num[1:] / prices_num[:-1] - 1.0
223 mu = self.mu.select(assets).to_numpy()
224 risk_pos_np = np.full_like(mu, fill_value=np.nan, dtype=float)
225 cash_pos_np = np.full_like(mu, fill_value=np.nan, dtype=float)
226 vola_np = self.vola.select(assets).to_numpy()
228 # ``cor`` is keyed by the post-warmup dates. Map each key back to its row
229 # in prices/mu/vola so the correlation matrix for date ``t`` is paired with
230 # (and stored at) that same date, rather than at a positional offset of
231 # ``corr`` rows — otherwise the most recent dates never receive a position.
232 row_of = {date: idx for idx, date in enumerate(self.prices["date"].to_list())}
234 _forward_walk(cor, prices_num, returns_num, mu, vola_np, risk_pos_np, cash_pos_np, row_of, self.cfg.shrink)
236 return self.prices.with_columns([(pl.lit(cash_pos_np[:, i]).alias(asset)) for i, asset in enumerate(assets)])