Coverage for src/jquantstats/_plots/_specs/_periodic.py: 100%
59 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-23 04:11 +0000
1"""Spec builders for the periodic bar charts and the monthly calendar.
3The bar charts colour each bar by the sign of its value rather than by series,
4which is why `~jquantstats._plots._spec.BarSeries` carries a colour per bar.
5"""
7from __future__ import annotations
9from typing import TYPE_CHECKING
11import polars as pl
13from .._spec import Axis, BarSeries, FigureSpec, HeatmapGrid, HoverSpec, Panel
14from .._style import bar_colors, ticker_colors, yearly_bar_colors
16if TYPE_CHECKING:
17 from jquantstats._protocol import DataLike
19__all__ = [
20 "daily_returns_spec",
21 "monthly_heatmap_spec",
22 "monthly_returns_spec",
23 "yearly_returns_spec",
24]
26# Bars are drawn slightly translucent so overlapping series stay readable.
27_BAR_OPACITY = 0.85
29_MONTH_NAMES = ("Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec")
31# A calendar row per year, plus room for the title and colour bar.
32_HEATMAP_ROW_PX = 40
33_HEATMAP_CHROME_PX = 100
34_HEATMAP_MIN_PX = 300
37def _split_columns(frame: pl.DataFrame) -> tuple[str, list[str]]:
38 """Separate the date column from the value columns.
40 Args:
41 frame: A frame whose first column is the date axis.
43 Returns:
44 tuple[str, list[str]]: The date column name and every other column.
46 """
47 date_col = frame.columns[0]
48 return date_col, [c for c in frame.columns if c != date_col]
51def period_agg_exprs(tickers: list[str], compounded: bool) -> list[pl.Expr]:
52 """Per-ticker aggregation expressions for a period bucket.
54 Args:
55 tickers: Asset column names to aggregate.
56 compounded: Compound returns within the bucket when True, sum them
57 when False.
59 Returns:
60 list[pl.Expr]: One aliased expression per ticker.
62 """
63 if compounded:
64 return [((1.0 + pl.col(t)).product() - 1.0).alias(t) for t in tickers]
65 return [pl.col(t).sum().alias(t) for t in tickers]
68def _sign_coloured_bars(
69 frame: pl.DataFrame,
70 x_col: str,
71 tickers: list[str],
72 colors: dict[str, str],
73 *,
74 single: bool,
75) -> tuple[BarSeries, ...]:
76 """Build one bar series per ticker, each bar coloured by its sign.
78 Args:
79 frame: The prepared frame holding the plotted values.
80 x_col: Column giving the bar positions.
81 tickers: Columns to draw, in order.
82 colors: Ticker to hex colour mapping.
83 single: Whether the dataset holds exactly one asset, which selects
84 the plain green/red palette over the faded per-asset one.
86 Returns:
87 tuple[BarSeries, ...]: One series per ticker, in the given order.
89 """
90 return tuple(
91 BarSeries(
92 name=ticker,
93 x=frame[x_col],
94 y=frame[ticker],
95 colors=tuple(bar_colors(frame[ticker].to_list(), colors[ticker], single_asset=single)),
96 opacity=_BAR_OPACITY,
97 hover=HoverSpec(label=ticker, value_format="percent2", date_header=False),
98 )
99 for ticker in tickers
100 )
103def daily_returns_spec(data: DataLike, title: str) -> FigureSpec:
104 """Describe the daily-returns bar chart.
106 Args:
107 data: The dataset to plot.
108 title: Chart title.
110 Returns:
111 FigureSpec: One bar series per asset, coloured by sign.
113 """
114 df = data.all
115 date_col, tickers = _split_columns(df)
117 panel = Panel(
118 bars=_sign_coloured_bars(df, date_col, tickers, ticker_colors(tickers), single=len(tickers) == 1),
119 yaxis=Axis(title="Return", tick_format="percent1"),
120 )
121 return FigureSpec(title=title, panels=(panel,))
124def yearly_returns_spec(data: DataLike, title: str, compounded: bool) -> FigureSpec:
125 """Describe the annual-returns grouped bar chart.
127 Args:
128 data: The dataset to plot.
129 title: Chart title.
130 compounded: Compound returns within each year.
132 Returns:
133 FigureSpec: One grouped bar series per asset, over a year axis.
135 """
136 df = data.all
137 date_col, tickers = _split_columns(df)
138 colors = ticker_colors(tickers)
140 yearly = (
141 df.with_columns(pl.col(date_col).dt.year().alias("_year"))
142 .group_by("_year")
143 .agg(period_agg_exprs(tickers, compounded))
144 .sort("_year")
145 )
147 panel = Panel(
148 bars=tuple(
149 BarSeries(
150 name=ticker,
151 x=yearly["_year"],
152 y=yearly[ticker],
153 # A flat zero year counts as positive here, unlike the other
154 # bar charts — see `yearly_bar_colors`.
155 colors=tuple(yearly_bar_colors(yearly[ticker].to_list(), colors[ticker])),
156 opacity=_BAR_OPACITY,
157 hover=HoverSpec(label=ticker, value_format="percent2", date_header=False),
158 )
159 for ticker in tickers
160 ),
161 xaxis=Axis(title="Year"),
162 yaxis=Axis(title="Annual Return", tick_format="percent1"),
163 )
164 # No range selector: the axis is calendar years, not a continuous date line.
165 return FigureSpec(title=title, panels=(panel,), date_range_selector=False, bar_mode="group")
168def monthly_returns_spec(data: DataLike, title: str, compounded: bool) -> FigureSpec:
169 """Describe the monthly-returns bar chart.
171 Args:
172 data: The dataset to plot.
173 title: Chart title.
174 compounded: Compound returns within each month.
176 Returns:
177 FigureSpec: One bar series per asset, coloured by sign.
179 """
180 df = data.all
181 date_col, tickers = _split_columns(df)
183 monthly = df.group_by_dynamic(index_column=date_col, every="1mo", period="1mo", closed="right", label="right").agg(
184 period_agg_exprs(tickers, compounded)
185 )
187 panel = Panel(
188 bars=_sign_coloured_bars(monthly, date_col, tickers, ticker_colors(tickers), single=len(tickers) == 1),
189 yaxis=Axis(title="Monthly Return", tick_format="percent1"),
190 )
191 return FigureSpec(title=title, panels=(panel,))
194def _heatmap_grids(
195 monthly: pl.DataFrame, years: list[int]
196) -> tuple[tuple[tuple[float | None, ...], ...], tuple[tuple[str, ...], ...]]:
197 """Build the year-by-month value and label grids.
199 Args:
200 monthly: Aggregated frame with ``_year``, ``_month`` and ``ret`` columns.
201 years: Sorted unique years, defining the row order of the output grids.
203 Returns:
204 A ``(z, text)`` pair: ``z`` holds returns scaled to percent, with None
205 for months the data does not cover, and ``text`` the formatted labels.
207 """
208 year_idx = {y: i for i, y in enumerate(years)}
209 z: list[list[float | None]] = [[None] * 12 for _ in years]
210 text: list[list[str]] = [[""] * 12 for _ in years]
211 for row in monthly.iter_rows(named=True):
212 yi = year_idx[row["_year"]]
213 mi = row["_month"] - 1
214 val = row["ret"]
215 z[yi][mi] = val * 100 if val is not None else None
216 text[yi][mi] = f"{val:.1%}" if val is not None else ""
217 return tuple(tuple(r) for r in z), tuple(tuple(r) for r in text)
220def monthly_heatmap_spec(data: DataLike, title: str, compounded: bool, asset: str | None) -> FigureSpec:
221 """Describe the monthly-returns calendar heatmap.
223 One asset per chart: the grid is already two-dimensional, so a second
224 asset would need a second grid.
226 Args:
227 data: The dataset to plot.
228 title: Chart title, which gains the asset name.
229 compounded: Compound returns within each month.
230 asset: Asset column to display, or None for the first.
232 Returns:
233 FigureSpec: A single year-by-month grid.
235 """
236 df = data.all
237 date_col, tickers = _split_columns(df)
238 col = asset if asset in tickers else tickers[0]
240 agg = ((1.0 + pl.col(col)).product() - 1.0) if compounded else pl.col(col).sum()
241 monthly = (
242 df.with_columns(
243 pl.col(date_col).dt.year().alias("_year"),
244 pl.col(date_col).dt.month().alias("_month"),
245 )
246 .group_by(["_year", "_month"])
247 .agg(agg.alias("ret"))
248 .sort(["_year", "_month"])
249 )
251 years = sorted(monthly["_year"].unique().to_list())
252 z, text = _heatmap_grids(monthly, years)
254 panel = Panel(
255 heatmap=HeatmapGrid(
256 x_labels=_MONTH_NAMES,
257 y_labels=tuple(str(y) for y in years),
258 z=z,
259 text=text,
260 colorscale="red_white_green",
261 colorbar_title="Return (%)",
262 hover_label="Return",
263 ),
264 # Months read as column headings, so they belong along the top.
265 xaxis=Axis(opposite_side=True),
266 )
267 return FigureSpec(
268 title=f"{title} — {col}",
269 panels=(panel,),
270 height=max(_HEATMAP_MIN_PX, _HEATMAP_ROW_PX * len(years) + _HEATMAP_CHROME_PX),
271 chrome="bare",
272 )