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

1 from __future__ import annotations

2

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

4

5 import polars as pl

6

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

8

9

10 def 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}