Coverage for src/jquantstats/_plots/_specs/_distribution.py: 100%

1"""Spec builders for the return-distribution charts. 

2 

3`histogram_spec` overlays the whole return distribution of each series on one 

4pair of axes. `distribution_spec` instead asks how the distribution *widens* as 

5the holding period lengthens, which needs a panel per asset. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11 

12import polars as pl 

13 

14from .._spec import Axis, BoxSeries, FigureSpec, HistogramSeries, HoverSpec, Panel 

15from .._style import ticker_colors 

16 

17if TYPE_CHECKING: 

18 from jquantstats._protocol import DataLike 

19 

20__all__ = ["distribution_spec", "histogram_spec"] 

21 

22# Overlaid distributions have to be seen through one another. 

23_OVERLAY_OPACITY = 0.6 

24 

25# Holding periods, shortest first, so the widening reads left to right. 

26# None means "no aggregation": the raw per-observation returns. 

27_PERIODS: tuple[tuple[str, str | None], ...] = ( 

28 ("Daily", None), 

29 ("Weekly", "1w"), 

30 ("Monthly", "1mo"), 

31 ("Quarterly", "3mo"), 

32 ("Yearly", "1y"), 

33) 

34 

35_PANEL_HEIGHT_PX = 500 

36 

37 

38def _split_columns(frame: pl.DataFrame) -> tuple[str, list[str]]: 

39 """Separate the date column from the value columns. 

40 

41 Args: 

42 frame: A frame whose first column is the date axis. 

43 

44 Returns: 

45 tuple[str, list[str]]: The date column name and every other column. 

46 

47 """ 

48 date_col = frame.columns[0] 

49 return date_col, [c for c in frame.columns if c != date_col] 

50 

51 

52def histogram_spec(data: DataLike, title: str, bins: int) -> FigureSpec: 

53 """Describe the overlaid return-distribution histogram. 

54 

55 Args: 

56 data: The dataset to plot. 

57 title: Chart title. 

58 bins: Number of histogram bins. 

59 

60 Returns: 

61 FigureSpec: One translucent histogram per column, on shared axes. 

62 

63 """ 

64 df = data.all 

65 _, tickers = _split_columns(df) 

66 colors = ticker_colors(tickers) 

67 

68 panel = Panel( 

69 histograms=tuple( 

70 HistogramSeries( 

71 name=ticker, 

72 values=df[ticker].drop_nulls().to_list(), 

73 color=colors[ticker], 

74 bins=bins, 

75 opacity=_OVERLAY_OPACITY, 

76 hover=HoverSpec( 

77 label=ticker, 

78 value_format="percent2", 

79 date_header=False, 

80 axis="x", 

81 hide_extra=True, 

82 ), 

83 ) 

84 for ticker in tickers 

85 ), 

86 xaxis=Axis(title="Return", tick_format="percent1"), 

87 yaxis=Axis(title="Count"), 

88 ) 

89 # No range selector: the x-axis is return magnitude, not time. 

90 return FigureSpec(title=title, panels=(panel,), date_range_selector=False, bar_mode="overlay") 

91 

92 

93def _period_values(df: pl.DataFrame, date_col: str, ticker: str, trunc: str | None, compounded: bool) -> list[float]: 

94 """Aggregate one ticker's returns into buckets of a given length. 

95 

96 Args: 

97 df: The combined index/returns frame. 

98 date_col: Name of the date column. 

99 ticker: Column to aggregate. 

100 trunc: Polars duration to bucket by, or None for no aggregation. 

101 compounded: Compound returns within each bucket. 

102 

103 Returns: 

104 list[float]: One value per bucket, nulls dropped. 

105 

106 """ 

107 if trunc is None: 

108 return df[ticker].drop_nulls().to_list() 

109 

110 agg = ((1.0 + pl.col(ticker)).product() - 1.0) if compounded else pl.col(ticker).sum() 

111 bucketed = ( 

112 df.with_columns(pl.col(date_col).dt.truncate(trunc).alias("_period")).group_by("_period").agg(agg.alias("ret")) 

113 ) 

114 return bucketed["ret"].drop_nulls().to_list() 

115 

116 

117def distribution_spec(data: DataLike, title: str, compounded: bool) -> FigureSpec: 

118 """Describe the by-holding-period distribution chart. 

119 

120 One panel per asset, each holding a box per period, so the widening of the 

121 distribution with holding length can be compared across assets. 

122 

123 Args: 

124 data: The dataset to plot. 

125 title: Chart title. 

126 compounded: Compound returns within each period. 

127 

128 Returns: 

129 FigureSpec: One side-by-side panel per asset, sharing a vertical scale. 

130 

131 """ 

132 df = data.all 

133 date_col, tickers = _split_columns(df) 

134 colors = ticker_colors(tickers) 

135 

136 panels = tuple( 

137 Panel( 

138 boxes=tuple( 

139 BoxSeries( 

140 name=period_name, 

141 values=_period_values(df, date_col, ticker, trunc, compounded), 

142 color=colors[ticker], 

143 # The same periods repeat in every panel, so only the 

144 # first names them; the groups tie them together. 

145 show_legend=(index == 0), 

146 legend_group=period_name, 

147 hover=HoverSpec( 

148 label=period_name, 

149 value_format="percent2", 

150 date_header=False, 

151 hide_extra=True, 

152 ), 

153 ) 

154 for period_name, trunc in _PERIODS 

155 ), 

156 title=ticker, 

157 yaxis=Axis(tick_format="percent1"), 

158 ) 

159 for index, ticker in enumerate(tickers) 

160 ) 

161 

162 return FigureSpec( 

163 title=title, 

164 panels=panels, 

165 height=_PANEL_HEIGHT_PX, 

166 chrome="panels", 

167 arrangement="side_by_side", 

168 shared_y=True, 

169 )