Coverage for src/basanos/math/_engine_core.py: 100%

1"""Core data-access mixin for `BasanosEngine`. 

2 

3Provides the volatility-adjusted returns, EWMA volatility, and per-timestamp 

4correlation properties as a reusable mixin so that ``optimizer.py`` stays 

5focused on the position-solving facade. 

6 

7Classes in this module are **private implementation details**. The public API 

8is `BasanosEngine`, which inherits from `_CoreDataMixin`. 

9""" 

10 

11from __future__ import annotations 

12 

13import datetime 

14from typing import TYPE_CHECKING 

15 

16import numpy as np 

17import polars as pl 

18from cvx.linalg import cov_to_corr 

19from cvx.linalg.covariance.ewm_cov import ewm_covariance 

20 

21from ._signal import vol_adj 

22 

23if TYPE_CHECKING: 

24 from ._engine_protocol import _EngineProtocol 

25 

26 

27class _CoreDataMixin: 

28 """Mixin providing the core data-access properties of `BasanosEngine`. 

29 

30 The consuming class must satisfy `_EngineProtocol`, i.e. it must expose 

31 ``prices`` (Polars DataFrame with a ``'date'`` column) and ``cfg`` 

32 (a `BasanosConfig`). 

33 """ 

34 

35 @property 

36 def assets(self: _EngineProtocol) -> list[str]: 

37 """List asset column names (numeric columns excluding 'date').""" 

38 return [c for c in self.prices.columns if c != "date" and self.prices[c].dtype.is_numeric()] 

39 

40 @property 

41 def ret_adj(self: _EngineProtocol) -> pl.DataFrame: 

42 """Return per-asset volatility-adjusted log returns clipped by cfg.clip. 

43 

44 Uses an EWMA volatility estimate with lookback ``cfg.vola`` to 

45 standardize log returns for each numeric asset column. 

46 """ 

47 return self.prices.with_columns( 

48 [vol_adj(pl.col(asset), vola=self.cfg.vola, clip=self.cfg.clip) for asset in self.assets] 

49 ) 

50 

51 @property 

52 def vola(self: _EngineProtocol) -> pl.DataFrame: 

53 """Per-asset EWMA volatility of percentage returns. 

54 

55 Computes percent changes for each numeric asset column and applies an 

56 exponentially weighted standard deviation using the lookback specified 

57 by ``cfg.vola``. The result is a DataFrame aligned with ``self.prices`` 

58 whose numeric columns hold per-asset volatility estimates. 

59 """ 

60 return self.prices.with_columns( 

61 pl.col(asset) 

62 .pct_change() 

63 .ewm_std(com=self.cfg.vola - 1, adjust=True, min_samples=self.cfg.vola) 

64 .alias(asset) 

65 for asset in self.assets 

66 ) 

67 

68 @property 

69 def cor(self: _EngineProtocol) -> dict[datetime.date, np.ndarray]: 

70 """Compute per-timestamp EWM correlation matrices. 

71 

72 Builds volatility-adjusted returns for all assets, computes an 

73 exponentially weighted correlation using a pure NumPy implementation 

74 (with window ``cfg.corr``), and returns a mapping from each timestamp 

75 to the corresponding correlation matrix as a NumPy array. 

76 

77 Returns: 

78 dict: Mapping ``date -> np.ndarray`` of shape (n_assets, n_assets). 

79 

80 Performance: 

81 Delegates to ``ewm_covariance`` from ``cvx.linalg``. 

82 For large *N* or *T*, prefer ``cor_tensor`` to keep a single 

83 contiguous array rather than building a Python dict. 

84 """ 

85 assets = list(self.assets) 

86 n = len(assets) 

87 span = 2 * self.cfg.corr + 1 

88 cov_dict = ewm_covariance( 

89 self.ret_adj, 

90 assets=assets, 

91 index_col="date", 

92 window=span, 

93 warmup=self.cfg.corr, 

94 ) 

95 nan_mat = np.full((n, n), np.nan) 

96 return { 

97 date: cov_to_corr(cov_dict[date], self.cfg.min_corr_denom) if date in cov_dict else nan_mat.copy() 

98 for date in self.prices["date"].to_list() 

99 } 

100 

101 @property 

102 def cor_tensor(self: _EngineProtocol) -> np.ndarray: 

103 """Return all correlation matrices stacked as a 3-D tensor. 

104 

105 Converts the per-timestamp correlation dict (see `cor`) into a 

106 single contiguous NumPy array so that the full history can be saved to 

107 a flat ``.npy`` file with `save` and reloaded with 

108 `load`. 

109 

110 Returns: 

111 np.ndarray: Array of shape ``(T, N, N)`` where *T* is the number of 

112 timestamps and *N* the number of assets. ``tensor[t]`` is the 

113 correlation matrix for the *t*-th date (same ordering as 

114 ``self.prices["date"]``). 

115 

116 Examples: 

117 >>> import tempfile, pathlib 

118 >>> import numpy as np 

119 >>> import polars as pl 

120 >>> from basanos.math.optimizer import BasanosConfig, BasanosEngine 

121 >>> dates = pl.Series("date", list(range(100))) 

122 >>> rng0 = np.random.default_rng(0).lognormal(size=100) 

123 >>> rng1 = np.random.default_rng(1).lognormal(size=100) 

124 >>> prices = pl.DataFrame({"date": dates, "A": rng0, "B": rng1}) 

125 >>> rng2 = np.random.default_rng(2).normal(size=100) 

126 >>> rng3 = np.random.default_rng(3).normal(size=100) 

127 >>> mu = pl.DataFrame({"date": dates, "A": rng2, "B": rng3}) 

128 >>> cfg = BasanosConfig(vola=10, corr=20, clip=3.0, shrink=0.5, aum=1e6) 

129 >>> engine = BasanosEngine(prices=prices, mu=mu, cfg=cfg) 

130 >>> tensor = engine.cor_tensor 

131 >>> with tempfile.TemporaryDirectory() as td: 

132 ... path = pathlib.Path(td) / "cor.npy" 

133 ... np.save(path, tensor) 

134 ... loaded = np.load(path) 

135 >>> np.testing.assert_array_equal(tensor, loaded) 

136 """ 

137 return np.stack(list(self.cor.values()), axis=0)