Coverage for src/cc_liquid/portfolio/sizing.py: 91%

1from __future__ import annotations 

2 

3from typing import Dict, Iterable, Literal, Mapping, Sequence 

4 

5import polars as pl 

6 

7Scheme = Literal["equal", "rank_power"] 

8 

9 

10def weights_from_ranks( 

11 latest_preds: pl.DataFrame | Sequence[tuple[str, float]] | Mapping[str, float], 

12 *, 

13 id_col: str = "id", 

14 pred_col: str = "pred", 

15 long_assets: Sequence[str], 

16 short_assets: Sequence[str], 

17 target_gross: float, 

18 scheme: Scheme = "rank_power", 

19 power: float = 1.5, 

20) -> Dict[str, float]: 

21 """Convert ranks (higher = stronger long) to signed weights. 

22 

23 The output weights sum in absolute value to ``target_gross``. 

24 

25 ``latest_preds`` accepts either: 

26 - Polars DataFrame with ``id_col``/``pred_col`` 

27 - Iterable mapping asset id to score 

28 - Dict[str, float] 

29 """ 

30 

31 # Normalize predictions into a simple dict for quick lookups. 

32 if isinstance(latest_preds, pl.DataFrame): 

33 preds_dict = dict(zip(latest_preds[id_col], latest_preds[pred_col])) 

34 elif isinstance(latest_preds, Mapping): 

35 preds_dict = dict(latest_preds) 

36 else: 

37 preds_dict = {asset: score for asset, score in latest_preds} 

38 

39 n_long, n_short = len(long_assets), len(short_assets) 

40 total_positions = n_long + n_short 

41 if total_positions == 0 or target_gross <= 0: 

42 return {} 

43 

44 gross_long = target_gross * (n_long / total_positions) 

45 gross_short = target_gross * (n_short / total_positions) 

46 

47 def _side(ids: Iterable[str], gross: float, sign: float) -> Dict[str, float]: 

48 ids_list = [i for i in ids if i in preds_dict] 

49 n = len(ids_list) 

50 if n == 0 or gross <= 0: 

51 return {} 

52 

53 # Fetch scores and rank within this side (best first). 

54 scored = sorted( 

55 ((preds_dict[i], i) for i in ids_list), 

56 key=lambda x: x[0], 

57 reverse=sign > 0, 

58 ) 

59 

60 raw: list[float] 

61 if scheme == "equal": 

62 raw = [1.0] * n 

63 elif scheme == "rank_power": 

64 p = max(1e-6, float(power)) 

65 raw = [((n - idx) / n) ** p for idx in range(n)] 

66 else: 

67 raise ValueError( 

68 f"Invalid weighting scheme: {scheme}. Must be 'equal' or 'rank_power'" 

69 ) 

70 

71 denom = sum(raw) or 1.0 

72 scale = gross / denom 

73 

74 return {asset: sign * raw[idx] * scale for idx, (_, asset) in enumerate(scored)} 

75 

76 weights_long = _side(long_assets, gross_long, +1.0) 

77 weights_short = _side(short_assets, gross_short, -1.0) 

78 

79 return {**weights_long, **weights_short}