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

1"""Spec builders for the drawdown charts. 

2 

3Two views of the same thing: `drawdown_spec` plots the decline from the running 

4peak directly, while `drawdowns_periods_spec` leaves the equity curve alone and 

5shades the worst episodes on top of it. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING, Any 

11 

12import polars as pl 

13 

14from .._spec import Axis, Band, FigureSpec, HoverSpec, LineSeries, Panel, RefLine 

15from .._style import PALETTE, hex_to_rgba, ticker_colors 

16 

17if TYPE_CHECKING: 

18 from jquantstats._protocol import DataLike 

19 

20__all__ = ["compute_drawdown_periods", "drawdown_spec", "drawdowns_periods_spec"] 

21 

22# The underwater curve's fill, and the shading over a drawdown episode. Both are 

23# translucent so the line underneath stays readable. 

24_FILL_ALPHA = 0.3 

25_BAND_ALPHA = 0.2 

26 

27# The equity curve in the periods chart is a single fixed colour rather than a 

28# palette entry: there is only ever one asset on it, so nothing to distinguish. 

29_EQUITY_COLOR = "#1f77b4" 

30 

31 

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

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

34 

35 Args: 

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

37 

38 Returns: 

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

40 

41 """ 

42 date_col = frame.columns[0] 

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

44 

45 

46def compute_drawdown_periods(prices: list[float], n: int) -> list[dict[str, Any]]: 

47 """Identify the *n* worst drawdown periods in a cumulative price series. 

48 

49 A period runs from the point the series falls below its running peak until 

50 it recovers to it. Periods are ranked by depth, worst first. 

51 

52 Args: 

53 prices: Cumulative price (NAV) values. 

54 n: Maximum number of periods to return. 

55 

56 Returns: 

57 list[dict[str, Any]]: Dicts with ``start_idx``, ``end_idx``, 

58 ``valley_idx`` and ``max_drawdown`` (a fraction <= 0), worst first. 

59 

60 """ 

61 length = len(prices) 

62 hwm: list[float] = [0.0] * length 

63 hwm[0] = prices[0] 

64 for i in range(1, length): 

65 hwm[i] = max(hwm[i - 1], prices[i]) 

66 

67 in_dd = [prices[i] < hwm[i] for i in range(length)] 

68 periods: list[dict[str, Any]] = [] 

69 i = 0 

70 while i < length: 

71 if not in_dd[i]: 

72 i += 1 

73 continue 

74 start = i 

75 while i < length and in_dd[i]: 

76 i += 1 

77 end = i - 1 

78 valley = start + min(range(end - start + 1), key=lambda k: prices[start + k]) 

79 max_dd = (prices[valley] - hwm[valley]) / hwm[valley] 

80 periods.append({"start_idx": start, "end_idx": end, "valley_idx": valley, "max_drawdown": max_dd}) 

81 

82 periods.sort(key=lambda p: p["max_drawdown"]) 

83 return periods[:n] 

84 

85 

86def drawdown_spec(data: DataLike, title: str) -> FigureSpec: 

87 """Describe the underwater equity curve. 

88 

89 Args: 

90 data: The dataset to plot. 

91 title: Chart title. 

92 

93 Returns: 

94 FigureSpec: One filled series per column, with a break-even line. 

95 

96 """ 

97 df = data.all 

98 date_col, tickers = _split_columns(df) 

99 colors = ticker_colors(tickers) 

100 prices = df.with_columns([(1.0 + pl.col(t)).cum_prod().alias(t) for t in tickers]) 

101 

102 lines = [] 

103 for ticker in tickers: 

104 price_s = prices[ticker] 

105 hwm = price_s.cum_max() 

106 lines.append( 

107 LineSeries( 

108 name=ticker, 

109 x=prices[date_col], 

110 y=((price_s - hwm) / hwm).to_list(), 

111 color=colors[ticker], 

112 width=1.5, 

113 fill=True, 

114 fill_color=hex_to_rgba(colors[ticker], _FILL_ALPHA), 

115 hover=HoverSpec(label=ticker, value_format="percent2", date_header=False), 

116 ) 

117 ) 

118 

119 panel = Panel( 

120 lines=tuple(lines), 

121 ref_lines=(RefLine(value=0),), 

122 yaxis=Axis(title="Drawdown", tick_format="percent0"), 

123 ) 

124 return FigureSpec(title=title, panels=(panel,)) 

125 

126 

127def drawdowns_periods_spec(data: DataLike, n: int, title: str, asset: str | None) -> FigureSpec: 

128 """Describe the equity curve with its worst drawdown episodes shaded. 

129 

130 One asset per chart: overlapping shaded spans from several assets would be 

131 unreadable. 

132 

133 Args: 

134 data: The dataset to plot. 

135 n: How many of the worst episodes to shade. 

136 title: Chart title, which gains the asset name. 

137 asset: Asset column to display, or None for the first. 

138 

139 Returns: 

140 FigureSpec: The equity curve, with one band per episode. 

141 

142 """ 

143 df = data.all 

144 date_col, tickers = _split_columns(df) 

145 col = asset if asset in tickers else tickers[0] 

146 

147 price_list = (1.0 + df[col].cast(pl.Float64)).cum_prod().to_list() 

148 dates = df[date_col].to_list() 

149 

150 bands = [] 

151 for i, period in enumerate(compute_drawdown_periods(price_list, n)): 

152 bands.append( 

153 Band( 

154 x0=dates[period["start_idx"]], 

155 # Extend to the next point so the span covers the final day of 

156 # the episode rather than stopping at its left edge. 

157 x1=dates[min(period["end_idx"] + 1, len(dates) - 1)], 

158 color=hex_to_rgba(PALETTE[i % len(PALETTE)], alpha=_BAND_ALPHA), 

159 label=f"#{i + 1} {period['max_drawdown']:.1%}", 

160 ) 

161 ) 

162 

163 panel = Panel( 

164 lines=( 

165 LineSeries( 

166 name=col, 

167 x=dates, 

168 y=price_list, 

169 color=_EQUITY_COLOR, 

170 hover=HoverSpec(label=col, value_format="float2", suffix="x"), 

171 ), 

172 ), 

173 bands=tuple(bands), 

174 yaxis=Axis(title="Cumulative Return", tick_format="float2"), 

175 ) 

176 return FigureSpec(title=f"{title} — {col}", panels=(panel,))