Coverage for src/cvx/linalg/covariance/ewm_cov.py: 100%

1"""Exponentially weighted covariance matrix computation. 

2 

3This module requires the optional ``polars`` dependency. Install it with 

4``pip install cvx-linalg[ewm]``. 

5""" 

6 

7from __future__ import annotations 

8 

9from collections.abc import Callable, Hashable 

10 

11import numpy as np 

12 

13from ..core.exceptions import NegativeWarmupError as NegativeWarmupError 

14from ..core.exceptions import NonIntegerWarmupError 

15from ..core.types import Matrix 

16 

17try: 

18 import polars as pl 

19except ImportError as exc: # pragma: no cover 

20 _msg = ( 

21 "polars is required for cvx.linalg.covariance.ewm_cov; " 

22 "install it with `pip install cvx-linalg[ewm]` or `pip install polars`." 

23 ) 

24 raise ImportError(_msg) from exc 

25 

26 

27def _validate_warmup(warmup: int) -> None: 

28 """Check *warmup* is a non-negative integer (booleans rejected).""" 

29 if isinstance(warmup, bool) or not isinstance(warmup, int): 

30 raise NonIntegerWarmupError(warmup) 

31 if warmup < 0: 

32 raise NegativeWarmupError(warmup) 

33 

34 

35def _pairwise_cov_exprs(assets: list[str], ewm: Callable[[pl.Expr], pl.Expr]) -> list[pl.Expr]: 

36 """Build the upper-triangular EWM covariance expression for each asset pair. 

37 

38 Uses ``Cov(X, Y) = EWM(X*Y) - EWM(X)*EWM(Y)`` over the pair's common 

39 non-null observations, masking each mean to where the *other* asset is present. 

40 """ 

41 return [ 

42 ( 

43 ewm(pl.col(a) * pl.col(b)) 

44 - ewm(pl.when(pl.col(b).is_null()).then(None).otherwise(pl.col(a))) 

45 * ewm(pl.when(pl.col(a).is_null()).then(None).otherwise(pl.col(b))) 

46 ).alias(f"{a}_{b}") 

47 for i, a in enumerate(assets) 

48 for b in assets[i:] 

49 ] 

50 

51 

52def ewm_covariance( 

53 data: pl.DataFrame, 

54 assets: list[str], 

55 index_col: str, 

56 window: int = 30, 

57 is_halflife: bool = False, 

58 warmup: int = 0, 

59) -> dict[Hashable, Matrix]: 

60 """Compute the exponentially weighted covariance matrix of returns. 

61 

62 EWM covariance uses the identity 

63 ``Cov(X, Y) = EWM(X*Y) - EWM(X)*EWM(Y)`` applied to the 

64 *common non-null observations* of each pair, which is equivalent 

65 to ``pandas.DataFrame.ewm(span).cov(bias=True)``. 

66 

67 Each date is included in the result as long as at least one 

68 matrix entry is non-NaN. Cells involving a late-starting asset 

69 are ``NaN`` until that asset has enough observations; the date is 

70 never dropped on account of a single asset being unavailable. 

71 Dates where every cell is NaN (before the warmup period is met 

72 for any asset) are omitted. 

73 

74 Args: 

75 data: Polars DataFrame containing the index column and asset columns. 

76 assets: Ordered list of asset column names. 

77 index_col: Name of the index (e.g. date) column in *data*. 

78 window: Span (default) or half-life (when *is_halflife* is 

79 ``True``) of the exponential decay. Defaults to ``30``. 

80 is_halflife: When ``True`` *window* is interpreted as the 

81 half-life; otherwise it is the EWMA span. Defaults to 

82 ``False``. 

83 warmup: Minimum number of common observations required before 

84 a pair's cell is non-NaN. Defaults to ``0`` (cells are 

85 non-NaN from the first shared observation). 

86 

87 Returns: 

88 Dictionary keyed by index value (date or integer) mapping to 

89 a square symmetric ``numpy.ndarray`` of shape ``(n, n)`` 

90 where ``n`` is the number of assets. Row/column order 

91 matches *assets*. Unavailable cells are ``NaN``. 

92 

93 Raises: 

94 NonIntegerWarmupError: If *warmup* is not an integer (booleans included). 

95 NegativeWarmupError: If *warmup* is negative. 

96 

97 """ 

98 _validate_warmup(warmup) 

99 

100 n = len(assets) 

101 min_samples = 1 if warmup == 0 else warmup 

102 

103 def _ewm(expr: pl.Expr) -> pl.Expr: 

104 """Apply EWM mean with the configured span or half-life.""" 

105 if is_halflife: 

106 return expr.ewm_mean(half_life=window, min_samples=min_samples) 

107 return expr.ewm_mean(span=window, min_samples=min_samples) 

108 

109 pair_df = data.with_columns(_pairwise_cov_exprs(assets, _ewm)).drop(assets) 

110 all_keys = pair_df[index_col].to_list() 

111 pair_arr = pair_df.drop(index_col).to_numpy() 

112 

113 ii, jj = np.triu_indices(n) 

114 cube = np.full((len(all_keys), n, n), np.nan) 

115 cube[:, ii, jj] = pair_arr 

116 cube[:, jj, ii] = pair_arr 

117 

118 has_data = ~np.all(np.isnan(cube), axis=(1, 2)) 

119 return {k: cube[t] for t, k in enumerate(all_keys) if has_data[t]}