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

1"""Pure numerical helpers for the incremental (streaming) optimiser. 

2 

3These functions carry no stream state — they are factored out of 

4:mod:`basanos.math._stream` so the EWMA recurrences and the per-step input 

5validation can be read, reused, and tested in isolation. 

6""" 

7 

8from __future__ import annotations 

9 

10import numpy as np 

11from scipy.signal import lfilter 

12 

13 

14def _ewm_std_from_state( 

15 s_x: np.ndarray, 

16 s_x2: np.ndarray, 

17 s_w: np.ndarray, 

18 s_w2: np.ndarray, 

19 count: np.ndarray, 

20 min_samples: int, 

21) -> np.ndarray: 

22 r"""Compute the unbiased EWMA standard deviation from running accumulators. 

23 

24 Implements the same Bessel-corrected formula used by 

25 ``polars.Expr.ewm_std(adjust=True)``:: 

26 

27 var_biased = s_x2/s_w - (s_x/s_w)^2 

28 correction = s_w^2 / (s_w^2 - s_w2) # Bessel correction 

29 var_unbiased = var_biased * correction 

30 std = sqrt(max(0, var_unbiased)) 

31 

32 where ``s_w2 = sum(wi^2)`` is the sum of squared EWM weights. 

33 

34 Parameters 

35 ---------- 

36 s_x, s_x2, s_w, s_w2: 

37 Running accumulators, each of shape ``(N,)``. 

38 count: 

39 Integer count of finite observations per asset, shape ``(N,)``. 

40 min_samples: 

41 Minimum number of finite observations required before returning a 

42 non-NaN value. 

43 

44 Returns: 

45 ------- 

46 np.ndarray of shape ``(N,)`` with per-asset standard deviations. 

47 NaN is returned for assets where ``count < min_samples``. 

48 """ 

49 n = len(s_x) 

50 result = np.full(n, np.nan, dtype=float) 

51 ok = count >= min_samples 

52 if not ok.any(): 

53 return result 

54 

55 with np.errstate(divide="ignore", invalid="ignore"): 

56 mean = np.where(s_w > 0, s_x / s_w, 0.0) 

57 mean_sq = np.where(s_w > 0, s_x2 / s_w, 0.0) 

58 var_biased = np.maximum(mean_sq - mean**2, 0.0) 

59 denom_corr = s_w**2 - s_w2 

60 # denom_corr > 0 iff count >= 2; equals 0 when count == 1 

61 var_unbiased = np.where(denom_corr > 0, var_biased * s_w**2 / denom_corr, 0.0) 

62 std = np.sqrt(var_unbiased) 

63 

64 return np.where(ok, std, np.nan) 

65 

66 

67def _ewm_vol_accumulators_from_batch( 

68 returns: np.ndarray, 

69 beta: float, 

70 beta_sq: float, 

71) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 

72 r"""Compute final EWMA volatility accumulators from a batch of returns. 

73 

74 Implements the same IIR recurrence as ``BasanosStream.step`` but 

75 vectorised over *T* timesteps using ``scipy.signal.lfilter``. The five 

76 returned arrays are identical to the accumulators that would result from 

77 feeding each row of *returns* through the scalar step-by-step recurrence:: 

78 

79 s_x[t] = beta * s_x[t-1] + (x[t] if finite else 0) 

80 s_x2[t] = beta * s_x2[t-1] + (x[t]^2 if finite else 0) 

81 s_w[t] = beta * s_w[t-1] + (1 if finite else 0) 

82 s_w2[t] = beta^2 * s_w2[t-1] + (1 if finite else 0) 

83 

84 Parameters 

85 ---------- 

86 

87 Returns: 

88 Float array of shape ``(T, N)``. NaN entries are treated as missing 

89 observations — they contribute nothing to the numerator sums and do 

90 not increment the weight accumulators. 

91 beta: 

92 EWM decay factor for ``s_x``, ``s_x2``, and ``s_w`` 

93 (``beta = (com) / (1 + com)`` for ``com = cfg.vola - 1``). 

94 beta_sq: 

95 Squared decay factor used for ``s_w2``. Must equal ``beta ** 2``. 

96 

97 Returns: 

98 ------- 

99 s_x, s_x2, s_w, s_w2 : np.ndarray of shape ``(N,)`` 

100 Final EWMA running accumulators after processing all *T* rows. 

101 count : np.ndarray of shape ``(N,)`` dtype int 

102 Number of finite observations per asset. 

103 

104 Notes: 

105 ----- 

106 This function is the shared implementation used by 

107 ``BasanosStream.from_warmup`` for both the log-return (``vola_*``) 

108 and pct-return (``pct_*``) accumulators. Keeping a single implementation 

109 here guarantees that the batch and incremental paths stay in sync when the 

110 recurrence definition changes. 

111 """ 

112 fin = np.isfinite(returns).astype(np.float64) # (T, N) 

113 x_z = np.where(fin.astype(bool), returns, 0.0) # (T, N) 

114 filt_a = np.array([1.0, -beta]) 

115 filt_a2 = np.array([1.0, -beta_sq]) 

116 

117 s_x: np.ndarray = lfilter([1.0], filt_a, x_z, axis=0)[-1] 

118 s_x2: np.ndarray = lfilter([1.0], filt_a, x_z**2, axis=0)[-1] 

119 s_w: np.ndarray = lfilter([1.0], filt_a, fin, axis=0)[-1] 

120 s_w2: np.ndarray = lfilter([1.0], filt_a2, fin, axis=0)[-1] 

121 count: np.ndarray = fin.sum(axis=0).astype(int) 

122 

123 return s_x, s_x2, s_w, s_w2, count 

124 

125 

126def _resolve_step_vector( 

127 values: np.ndarray | dict[str, float], 

128 assets: list[str], 

129 n_assets: int, 

130 arg_name: str, 

131) -> np.ndarray: 

132 """Resolve one step input to a validated ``(N,)`` float vector. 

133 

134 Args: 

135 values: Raw input provided to ``step`` as either dict or array-like. 

136 assets: Ordered asset names used when ``values`` is a mapping. 

137 n_assets: Expected vector length. 

138 arg_name: Argument label used in shape-mismatch errors. 

139 

140 Returns: 

141 A float64 numpy vector of shape ``(n_assets,)``. 

142 

143 Raises: 

144 ValueError: If the resolved vector does not match ``(n_assets,)``. 

145 """ 

146 if isinstance(values, dict): 

147 vector = np.array([float(values[a]) for a in assets], dtype=float) 

148 else: 

149 vector = np.asarray(values, dtype=float).ravel() 

150 if vector.shape != (n_assets,): 

151 raise ValueError(f"{arg_name} must have shape ({n_assets},); got {vector.shape}") # noqa: TRY003 

152 return vector