Coverage for src/jquantstats/_plots/_data/_periodic.py: 100%
80 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-06 04:52 +0000
1"""Periodic-return bar charts and the monthly-return heatmap."""
3from __future__ import annotations
5from typing import TYPE_CHECKING
7import plotly.graph_objects as go
8import polars as pl
10from ._styling import _apply_base_layout, _bar_colors, _ticker_colors, _yearly_bar_colors
12if TYPE_CHECKING:
13 from jquantstats._protocol import DataLike
16def _period_agg_exprs(tickers: list[str], compounded: bool) -> list[pl.Expr]:
17 """Per-ticker aggregation expressions for a period bucket.
19 Args:
20 tickers: Asset column names to aggregate.
21 compounded: Compound returns within the bucket when True, sum them
22 when False.
24 Returns:
25 One aliased expression per ticker.
26 """
27 if compounded:
28 return [((1.0 + pl.col(t)).product() - 1.0).alias(t) for t in tickers]
29 return [pl.col(t).sum().alias(t) for t in tickers]
32def _monthly_heatmap_matrix(
33 monthly: pl.DataFrame, years: list[int]
34) -> tuple[list[list[float | None]], list[list[str]]]:
35 """Build the year-by-month value and label grids for the monthly heatmap.
37 Args:
38 monthly: Aggregated frame with ``_year``, ``_month`` and ``ret`` columns.
39 years: Sorted unique years, defining the row order of the output grids.
41 Returns:
42 A ``(z, text)`` tuple: ``z`` holds returns scaled to percent (``None``
43 for missing cells) and ``text`` the formatted per-cell labels.
45 """
46 year_idx = {y: i for i, y in enumerate(years)}
47 z: list[list[float | None]] = [[None] * 12 for _ in years]
48 text: list[list[str]] = [[""] * 12 for _ in years]
49 for row in monthly.iter_rows(named=True):
50 yi = year_idx[row["_year"]]
51 mi = row["_month"] - 1
52 val = row["ret"]
53 z[yi][mi] = val * 100 if val is not None else None
54 text[yi][mi] = f"{val:.1%}" if val is not None else ""
55 return z, text
58class _PeriodicPlotsMixin:
59 """Daily/monthly/yearly bar charts and the monthly heatmap for :class:`DataPlots`."""
61 __slots__ = ()
63 _data: DataLike
65 def daily_returns(self, title: str = "Daily Returns") -> go.Figure:
66 """Daily returns as a bar chart.
68 Each bar is coloured green for positive returns and red for negative
69 returns. When multiple assets are present each asset gets its own
70 trace in the palette colour with opacity used for positive/negative
71 differentiation.
73 Args:
74 title: Chart title. Defaults to ``"Daily Returns"``.
76 Returns:
77 go.Figure: Interactive Plotly bar chart.
79 """
80 df = self._data.all
81 date_col = df.columns[0]
82 tickers = [c for c in df.columns if c != date_col]
83 colors = _ticker_colors(tickers)
84 single = len(tickers) == 1
86 fig = go.Figure()
87 for ticker in tickers:
88 values = df[ticker].to_list()
89 bar_colors = _bar_colors(values, colors[ticker], single_asset=single)
91 fig.add_trace(
92 go.Bar(
93 x=df[date_col],
94 y=df[ticker],
95 name=ticker,
96 marker={"color": bar_colors, "line": {"width": 0}},
97 opacity=0.85,
98 hovertemplate=f"{ticker}: %{{y:.2%}}",
99 )
100 )
102 _apply_base_layout(fig, title)
103 fig.update_yaxes(title_text="Return", tickformat=".1%")
104 return fig
106 def yearly_returns(self, title: str = "Yearly Returns", compounded: bool = True) -> go.Figure:
107 """Annual compounded (or summed) returns as a grouped bar chart.
109 Args:
110 title: Chart title. Defaults to ``"Yearly Returns"``.
111 compounded: Compound returns within each year. Defaults to True.
113 Returns:
114 go.Figure: Interactive Plotly grouped bar chart.
116 """
117 df = self._data.all
118 date_col = df.columns[0]
119 tickers = [c for c in df.columns if c != date_col]
120 colors = _ticker_colors(tickers)
122 agg_exprs = _period_agg_exprs(tickers, compounded)
123 yearly = (
124 df.with_columns(pl.col(date_col).dt.year().alias("_year")).group_by("_year").agg(agg_exprs).sort("_year")
125 )
127 fig = go.Figure()
128 for ticker in tickers:
129 bar_colors = _yearly_bar_colors(yearly[ticker].to_list(), colors[ticker])
130 fig.add_trace(
131 go.Bar(
132 x=yearly["_year"],
133 y=yearly[ticker],
134 name=ticker,
135 marker={"color": bar_colors, "line": {"width": 0}},
136 opacity=0.85,
137 hovertemplate=f"{ticker}: %{{y:.2%}}",
138 )
139 )
141 _apply_base_layout(fig, title, with_range_selector=False)
142 fig.update_layout(barmode="group", xaxis_title="Year")
143 fig.update_yaxes(title_text="Annual Return", tickformat=".1%")
144 return fig
146 def monthly_returns(self, title: str = "Monthly Returns", compounded: bool = True) -> go.Figure:
147 """Monthly compounded (or summed) returns as a bar chart.
149 Args:
150 title: Chart title. Defaults to ``"Monthly Returns"``.
151 compounded: Compound returns within each month. Defaults to True.
153 Returns:
154 go.Figure: Interactive Plotly bar chart.
156 """
157 df = self._data.all
158 date_col = df.columns[0]
159 tickers = [c for c in df.columns if c != date_col]
160 colors = _ticker_colors(tickers)
161 single = len(tickers) == 1
163 monthly = df.group_by_dynamic(
164 index_column=date_col, every="1mo", period="1mo", closed="right", label="right"
165 ).agg(_period_agg_exprs(tickers, compounded))
167 fig = go.Figure()
168 for ticker in tickers:
169 values = monthly[ticker].to_list()
170 bar_colors = _bar_colors(values, colors[ticker], single_asset=single)
172 fig.add_trace(
173 go.Bar(
174 x=monthly[date_col],
175 y=monthly[ticker],
176 name=ticker,
177 marker={"color": bar_colors, "line": {"width": 0}},
178 opacity=0.85,
179 hovertemplate=f"{ticker}: %{{y:.2%}}",
180 )
181 )
183 _apply_base_layout(fig, title)
184 fig.update_yaxes(title_text="Monthly Return", tickformat=".1%")
185 return fig
187 def monthly_heatmap(
188 self,
189 title: str = "Monthly Returns Heatmap",
190 compounded: bool = True,
191 asset: str | None = None,
192 ) -> go.Figure:
193 """Monthly returns calendar heatmap (year x month).
195 One heatmap is produced per call for a single asset. Green cells
196 indicate positive months; red cells indicate negative months.
198 Args:
199 title: Chart title. Defaults to ``"Monthly Returns Heatmap"``.
200 compounded: Compound intra-month returns. Defaults to True.
201 asset: Asset column name to display. Defaults to the first
202 non-date column in the dataset.
204 Returns:
205 go.Figure: Interactive Plotly heatmap.
207 """
208 df = self._data.all
209 date_col = df.columns[0]
210 tickers = [c for c in df.columns if c != date_col]
211 col = asset if asset in tickers else tickers[0]
213 month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
215 agg_expr = ((1.0 + pl.col(col)).product() - 1.0).alias("ret") if compounded else pl.col(col).sum().alias("ret")
216 monthly = (
217 df.with_columns(
218 [
219 pl.col(date_col).dt.year().alias("_year"),
220 pl.col(date_col).dt.month().alias("_month"),
221 ]
222 )
223 .group_by(["_year", "_month"])
224 .agg(agg_expr.alias("ret"))
225 .sort(["_year", "_month"])
226 )
228 years = sorted(monthly["_year"].unique().to_list())
229 z, text = _monthly_heatmap_matrix(monthly, years)
231 fig = go.Figure(
232 go.Heatmap(
233 x=month_names,
234 y=[str(y) for y in years],
235 z=z,
236 text=text,
237 texttemplate="%{text}",
238 colorscale=[[0, "#d62728"], [0.5, "#ffffff"], [1, "#2ca02c"]],
239 zmid=0,
240 showscale=True,
241 colorbar={"title": "Return (%)"},
242 hovertemplate="<b>%{y} %{x}</b><br>Return: %{text}<extra></extra>",
243 )
244 )
246 fig.update_layout(
247 title=f"{title} — {col}",
248 height=max(300, 40 * len(years) + 100),
249 plot_bgcolor="white",
250 xaxis={"side": "top"},
251 )
252 return fig