Coverage for src/jquantstats/_plots/_specs/_cumulative.py: 100%
44 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-23 04:11 +0000
1"""Spec builders for the cumulative-return and equity-curve charts.
3Each builder takes a dataset and returns a `FigureSpec`. The arithmetic lives
4here once; `jquantstats._plots._render` decides how it is drawn.
5"""
7from __future__ import annotations
9import math
10from typing import TYPE_CHECKING
12import polars as pl
14from .._spec import Axis, Dash, FigureSpec, HoverSpec, LineSeries, Panel, TickFormat
15from .._style import ticker_colors
17if TYPE_CHECKING:
18 from jquantstats._protocol import DataLike
20__all__ = [
21 "compare_spec",
22 "cumulative_returns_spec",
23 "earnings_spec",
24 "log_returns_spec",
25]
28def _split_columns(frame: pl.DataFrame) -> tuple[str, list[str]]:
29 """Separate the date column from the value columns.
31 Args:
32 frame: A frame whose first column is the date axis.
34 Returns:
35 tuple[str, list[str]]: The date column name and every other column.
37 """
38 date_col = frame.columns[0]
39 return date_col, [c for c in frame.columns if c != date_col]
42def _lines(
43 frame: pl.DataFrame,
44 date_col: str,
45 tickers: list[str],
46 colors: dict[str, str],
47 value_format: TickFormat,
48 *,
49 prefix: str = "",
50 suffix: str = "",
51 width: float = 2,
52 dash: Dash | None = None,
53) -> list[LineSeries]:
54 """Build one line series per ticker, all sharing a style.
56 Args:
57 frame: The prepared frame holding the plotted values.
58 date_col: Name of the date column.
59 tickers: Columns to draw, in order.
60 colors: Ticker to hex colour mapping.
61 value_format: How tooltip values are rendered.
62 prefix: Written before each tooltip value.
63 suffix: Written after each tooltip value.
64 width: Stroke width.
65 dash: Stroke pattern.
67 Returns:
68 list[LineSeries]: One series per ticker, in the given order.
70 """
71 return [
72 LineSeries(
73 name=ticker,
74 x=frame[date_col],
75 y=frame[ticker],
76 color=colors[ticker],
77 width=width,
78 dash=dash,
79 hover=HoverSpec(label=ticker, value_format=value_format, prefix=prefix, suffix=suffix),
80 )
81 for ticker in tickers
82 ]
85def cumulative_returns_spec(data: DataLike, title: str, log_scale: bool) -> FigureSpec:
86 """Describe the cumulative compounded-returns chart.
88 Args:
89 data: The dataset to plot.
90 title: Chart title.
91 log_scale: Use a logarithmic y-axis.
93 Returns:
94 FigureSpec: One line per column of ``data.all``.
96 """
97 df = data.all
98 date_col, tickers = _split_columns(df)
99 prices = df.with_columns([(1.0 + pl.col(t)).cum_prod().alias(t) for t in tickers])
101 panel = Panel(
102 # The "x" suffix reads the value as a growth multiple: "1.42x".
103 lines=tuple(_lines(prices, date_col, tickers, ticker_colors(tickers), "float2", suffix="x")),
104 yaxis=Axis(title="Cumulative Return", tick_format="float2", log=log_scale),
105 )
106 return FigureSpec(title=title, panels=(panel,))
109def compare_spec(data: DataLike, title: str, figsize: tuple[int, int] | None) -> FigureSpec:
110 """Describe the asset-versus-benchmark comparison chart.
112 Benchmarks are drawn after the assets, slightly heavier and dashed, so they
113 read as the reference rather than as another asset.
115 Args:
116 data: The dataset to plot. Must carry benchmark columns.
117 title: Chart title.
118 figsize: Optional ``(width, height)`` in pixels.
120 Returns:
121 FigureSpec: Asset lines followed by benchmark lines.
123 Raises:
124 AttributeError: If no benchmark data is available.
126 """
127 benchmark_df = getattr(data, "benchmark", None)
128 if benchmark_df is None:
129 raise AttributeError("compare() requires benchmark data to be set") # noqa: TRY003
131 df = data.all
132 date_col, _ = _split_columns(df)
133 assets = list(data.returns.columns)
134 benchmarks = list(benchmark_df.columns)
136 colors = ticker_colors(assets + benchmarks)
137 prices = df.with_columns([(1.0 + pl.col(col)).cum_prod().alias(col) for col in assets + benchmarks])
139 panel = Panel(
140 lines=(
141 *_lines(prices, date_col, assets, colors, "float2", suffix="x"),
142 *_lines(prices, date_col, benchmarks, colors, "float2", suffix="x", width=2.5, dash="dash"),
143 ),
144 yaxis=Axis(title="Cumulative Return", tick_format="float2"),
145 )
146 return FigureSpec(title=title, panels=(panel,), figsize=figsize)
149def log_returns_spec(data: DataLike, title: str, figsize: tuple[int, int] | None) -> FigureSpec:
150 """Describe the cumulative log-returns chart.
152 Args:
153 data: The dataset to plot.
154 title: Chart title.
155 figsize: Optional ``(width, height)`` in pixels.
157 Returns:
158 FigureSpec: One line per column, on a linear axis of log values.
160 """
161 df = data.all
162 date_col, tickers = _split_columns(df)
163 log_prices = df.with_columns([(1.0 + pl.col(t)).cum_prod().log(math.e).alias(t) for t in tickers])
165 panel = Panel(
166 lines=tuple(_lines(log_prices, date_col, tickers, ticker_colors(tickers), "float4")),
167 yaxis=Axis(title="Log Return"),
168 )
169 return FigureSpec(title=title, panels=(panel,), figsize=figsize)
172def earnings_spec(data: DataLike, start_balance: float, title: str, compounded: bool) -> FigureSpec:
173 """Describe the dollar equity curve.
175 Args:
176 data: The dataset to plot.
177 start_balance: Starting portfolio value in currency units.
178 title: Chart title.
179 compounded: Compound the returns; when False they are summed.
181 Returns:
182 FigureSpec: One line per column, scaled to *start_balance*.
184 """
185 df = data.all
186 date_col, tickers = _split_columns(df)
188 if compounded:
189 equity = df.with_columns([(start_balance * (1.0 + pl.col(t)).cum_prod()).alias(t) for t in tickers])
190 else:
191 equity = df.with_columns([(start_balance * (1.0 + pl.col(t).cum_sum())).alias(t) for t in tickers])
193 panel = Panel(
194 lines=tuple(_lines(equity, date_col, tickers, ticker_colors(tickers), "currency0", prefix="$")),
195 yaxis=Axis(
196 title=f"Portfolio Value (starting ${start_balance:,.0f})",
197 tick_format="currency0",
198 tick_prefix="$",
199 ),
200 )
201 return FigureSpec(title=title, panels=(panel,))