Coverage for src/jointview/plot.py: 100%
63 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-17 06:44 +0000
1"""Two price or NAV series drawn as two lines on one pair of axes.
3The two columns are put on a shared y-axis rather than one axis each: two scales on
4one plot invent a relationship that is not in the data. Where the levels are far
5apart, ``rebase`` indexes both to the same starting value instead, which is the
6honest way to compare a series priced at 12 with one priced at 4,000.
8Choosing the columns and cutting them to their common sample happens before any of
9this, in :mod:`jointview.columns`.
10"""
12from __future__ import annotations
14import math
15from typing import Literal, cast
17import altair as alt
18import polars as pl
20from jointview.columns import PERIOD, aligned, date_column
22# Categorical slots 1 and 2 (blue, orange). The pair clears the contrast, chroma and
23# colour-vision separation floors against both the light (#fcfcfb) and the dark
24# (#1a1a19) chart surface, so it survives marimo's theme switch without being
25# redefined. CONTEXT is chrome — the crosshair — not a series.
26SERIES = ("#2a78d6", "#d95926")
27CONTEXT = "#8a8a84"
29BASE = 100.0
30MAX_POINTS = 4_000
32# Narrower than `str`, because altair's `type=` only accepts its four measurement kinds
33# and these are the two an x-axis of periods can be. Inside one function the literal was
34# inferred; passing it between them needs the type spelled out.
35XType = Literal["temporal", "quantitative"]
38def line_frame(
39 frame: pl.DataFrame,
40 a: str,
41 b: str,
42 *,
43 rebase: bool = True,
44 base: float = BASE,
45 max_points: int = MAX_POINTS,
46) -> pl.DataFrame:
47 """What the chart draws: one ``period`` column and one column per named series.
49 Wide rather than long, because the crosshair reads every series at the hovered
50 period out of a single row.
52 >>> import polars as pl
53 >>> frame = pl.DataFrame({"cash": [1.0, 1.01, 1.02], "balanced": [1450.0, 1479.0, 1465.0]})
54 >>> drawn = line_frame(frame, "cash", "balanced")
55 >>> drawn.columns
56 ['period', 'cash', 'balanced']
58 Rebasing is what lets those two share a y-axis at all: both leave the first
59 period at ``base``, whatever they were priced at.
61 >>> round(drawn["cash"][0], 6), round(drawn["balanced"][0], 6)
62 (100.0, 100.0)
64 A column against itself is one line rather than two identical ones:
66 >>> line_frame(frame, "cash", "cash").columns
67 ['period', 'cash']
69 A series starting at zero cannot be indexed — a cumulative P&L curve starts there
70 by construction — so the pair keeps its own levels instead:
72 >>> pnl = pl.DataFrame({"strategy": [0.0, 5.0, 3.0], "benchmark": [0.0, 2.0, 4.0]})
73 >>> line_frame(pnl, "strategy", "benchmark")["strategy"].to_list()
74 [0.0, 5.0, 3.0]
75 """
76 wide, _ = _wide(frame, a, b, rebase=rebase, base=base, max_points=max_points)
77 return wide
80def line_chart(
81 frame: pl.DataFrame,
82 a: str,
83 b: str,
84 *,
85 rebase: bool = True,
86 base: float = BASE,
87 width: int | str = "container",
88 height: int | str = 700,
89 max_points: int = MAX_POINTS,
90) -> alt.LayerChart:
91 """Draw columns ``a`` and ``b`` of ``frame`` as two lines against time.
93 Four layers over one plotting area — the crosshair, the lines, the hover markers
94 and the end labels — handed back as a plain Altair chart, so nothing here needs
95 marimo to draw it:
97 >>> import polars as pl
98 >>> frame = pl.DataFrame({"cash": [1.0, 1.01, 1.02], "balanced": [1450.0, 1479.0, 1465.0]})
99 >>> chart = line_chart(frame, "cash", "balanced")
100 >>> type(chart).__name__
101 'LayerChart'
102 >>> len(chart.to_dict()["layer"])
103 4
105 Width defaults to ``"container"``: the plot takes whatever the column around it
106 gives it, which is the point of a full-width app. Height stays a number, because
107 nothing in the page has a height for a chart to follow — 700 fills a laptop window
108 once the notebook margins are out of the way, without spilling off a short one.
110 Asking to rebase a pair that cannot be indexed draws the raw levels, and the
111 y-axis says ``level`` rather than claiming otherwise — see :func:`_rebasable`.
112 The title is read back out of the compiled spec, because that is the only place it
113 exists; layer 1 is :func:`_lines`, the series themselves:
115 >>> chart.to_dict()["layer"][1]["encoding"]["y"]["title"]
116 'indexed to 100'
117 >>> pnl = pl.DataFrame({"strategy": [0.0, 5.0, 3.0], "benchmark": [0.0, 2.0, 4.0]})
118 >>> line_chart(pnl, "strategy", "benchmark").to_dict()["layer"][1]["encoding"]["y"]["title"]
119 'level'
120 """
121 wide, rebased = _wide(frame, a, b, rebase=rebase, base=base, max_points=max_points)
122 names = [name for name in wide.columns if name != PERIOD]
124 date = date_column(frame)
125 x_type: XType = "temporal" if date else "quantitative"
126 x_title = date or "row"
127 # `rebased`, not `rebase`: the title names what the numbers underneath actually are.
128 y_title = f"indexed to {base:g}" if rebased else "level"
129 x = alt.X(PERIOD, type=x_type, title=x_title)
131 # Made here rather than inside a layer because two of them share it: the crosshair
132 # carries the parameter, the markers only read it.
133 hover = _hover()
134 lines = _lines(wide, names, x, y_title)
135 return (
136 alt.layer(
137 _crosshair(wide, names, x, x_type, x_title, hover),
138 lines,
139 _markers(lines, hover),
140 _end_labels(wide, names, x),
141 )
142 .resolve_scale(color="shared")
143 .configure_axis(grid=True, gridOpacity=0.3, domain=False, labelPadding=4, tickSize=4)
144 .configure_view(stroke=None)
145 .configure_legend(labelFontSize=12)
146 .properties(
147 # One plotting area for all four layers, sized at the top level so a
148 # container width is measured once rather than per layer.
149 width=width,
150 height=height,
151 # Only the right margin earns its keep: it is where the end labels go.
152 padding={"left": 0, "top": 0, "bottom": 0, "right": 76},
153 # "pad", the default, grows the figure past the box it was given and puts
154 # the gutter back; fitting spends the padding out of the size instead.
155 autosize=_autosize(width, height),
156 )
157 )
160def _wide(
161 frame: pl.DataFrame,
162 a: str,
163 b: str,
164 *,
165 rebase: bool,
166 base: float,
167 max_points: int,
168) -> tuple[pl.DataFrame, bool]:
169 """The frame the chart is drawn from, and whether it ended up indexed after all.
171 One decision point for two callers. :func:`line_frame` wants the frame;
172 :func:`line_chart` wants the answer too, because an axis labelled "indexed to 100"
173 over unindexed levels is a wrong label rather than a missing one.
174 """
175 data = aligned(frame, a, b)
176 names = [a] if a == b else [a, b]
177 # `aligned` always writes both, so a column against itself takes only the first.
178 sources = ["a", "b"][: len(names)]
179 rebased = rebase and _rebasable(data, sources)
181 columns = [pl.col(source).alias(name) for source, name in zip(sources, names, strict=True)]
182 if rebased:
183 columns = [column / column.first() * base for column in columns]
185 return _thin(data.select(PERIOD, *columns), max_points), rebased
188def _rebasable(data: pl.DataFrame, sources: list[str]) -> bool:
189 """Whether dividing by the first value leaves a number for every series.
191 A first value of zero is not a broken frame — a cumulative P&L curve starts there
192 by construction — but dividing by it turns the rest of the line into infinities,
193 and Vega-Lite drops those silently. The reader picks two series, gets one, and
194 nothing on the plot says where the other went.
196 Answered for the pair together rather than per series. Indexing one to 100 while
197 the other kept its own units would put two unrelated scales on one axis, which is
198 the relationship this module exists in order not to invent.
200 A pair with no overlap has no first value to divide by, so there is nothing to
201 check and nothing to break: an empty frame comes out empty either way.
202 """
203 firsts = [value for row in data.head(1).select(sources).rows() for value in row]
204 return all(value != 0.0 and math.isfinite(value) for value in firsts)
207def _hover() -> alt.Parameter:
208 """Which period the pointer is nearest, shared by the crosshair and the markers.
210 Nearest-point rather than a hit on the mark itself: nobody can be asked to hover a
211 2px line. It resolves to a period, not to a series, which is what lets one tooltip
212 report every line at that x.
213 """
214 return alt.selection_point(
215 name="hover",
216 fields=[PERIOD],
217 nearest=True,
218 on="pointerover",
219 clear="pointerout",
220 empty=False,
221 )
224def _lines(wide: pl.DataFrame, names: list[str], x: alt.X, y_title: str) -> alt.Chart:
225 """The series themselves — the layer everything else is chrome around.
227 Drawn from the long form, because one line per ``series`` value is what lets a
228 single colour encoding paint both.
229 """
230 long = wide.unpivot(index=PERIOD, variable_name="series", value_name="value")
231 colour = alt.Color(
232 "series",
233 type="nominal",
234 title=None,
235 # Bound to the names, so picking a different pair never repaints a series that
236 # stayed on screen.
237 scale=alt.Scale(domain=names, range=list(SERIES[: len(names)])),
238 legend=alt.Legend(orient="top", offset=4, symbolType="stroke", symbolStrokeWidth=2),
239 )
240 # The `ty: ignore` here and in the layers below is altair's `mark_*` returning an
241 # unresolved TypeVar rather than a chart, so the checker cannot see `.encode` on it.
242 # It is a limitation of the stubs, not of the call — the same chain is what altair's
243 # own documentation shows.
244 return (
245 alt.Chart(long) # ty: ignore[unresolved-attribute]
246 .mark_line(strokeWidth=2, clip=True)
247 .encode(
248 x,
249 # The levels frame the data; a zero baseline on a NAV is wasted panel.
250 alt.Y("value", type="quantitative", title=y_title, scale=alt.Scale(zero=False)),
251 colour,
252 )
253 )
256def _crosshair(
257 wide: pl.DataFrame,
258 names: list[str],
259 x: alt.X,
260 x_type: XType,
261 x_title: str,
262 hover: alt.Parameter,
263) -> alt.Chart:
264 """The vertical rule under the pointer, carrying the tooltip for every series.
266 Drawn from the wide form: the tooltip reads all the series at the hovered period
267 out of a single row, which is the whole reason :func:`line_frame` is wide.
268 """
269 return (
270 alt.Chart(wide) # ty: ignore[unresolved-attribute]
271 .mark_rule(color=CONTEXT, strokeWidth=1)
272 .encode(
273 x,
274 opacity=alt.condition(hover, alt.value(0.6), alt.value(0.0)),
275 tooltip=[
276 alt.Tooltip(PERIOD, type=x_type, title=x_title),
277 *(alt.Tooltip(name, type="quantitative", format=",.2f") for name in names),
278 ],
279 )
280 .add_params(hover)
281 )
284def _markers(lines: alt.Chart, hover: alt.Parameter) -> alt.Chart:
285 """A dot on each line at the hovered period — the lines' own marks, made visible.
287 Built off ``lines`` rather than from scratch so the points inherit its data and
288 colour encoding, and cannot drift from the curve they sit on.
289 """
290 return lines.mark_point(size=64, filled=True).encode( # ty: ignore[unresolved-attribute]
291 opacity=alt.condition(hover, alt.value(1.0), alt.value(0.0))
292 )
295def _autosize(width: int | str, height: int | str) -> alt.AutoSizeParams:
296 """Fit whichever axes were asked to follow their container, and leave the rest.
298 A figure given two numbers keeps Vega-Lite's own ``pad``: it was asked for a
299 drawing of exactly that size, and fitting would quietly shrink it.
300 """
301 follows = ((width == "container", "x"), (height == "container", "y"))
302 axes = "".join(axis for container, axis in follows if container)
303 # Built from the axes rather than spelled out, so the cast is what tells the type
304 # checker that the four strings this can produce are exactly Vega-Lite's four.
305 kind = cast(
306 "Literal['pad', 'fit', 'fit-x', 'fit-y']",
307 {"": "pad", "xy": "fit"}.get(axes, f"fit-{axes}"),
308 )
309 return alt.AutoSizeParams(type=kind, contains="padding")
312def _end_labels(wide: pl.DataFrame, names: list[str], x: alt.X) -> alt.LayerChart:
313 """The series name at the end of its own line, so identity is never colour alone.
315 The labels wear chrome ink rather than the series colour — the line they sit on
316 carries the identity — and the upper one is nudged up, the lower one down, so a
317 pair that ends at the same level does not print on top of itself.
318 """
319 last = wide.tail(1)
320 order = sorted(names, key=lambda name: -float(last[name][0]))
321 dodge = dict(zip(order, (-8, 8), strict=False)) if len(order) > 1 else {order[0]: 0}
323 # alt.layer widens to LayerChart | FacetChart for the general case; none of these
324 # labels is faceted, so the layer it returns is always the former.
325 return cast(
326 "alt.LayerChart",
327 alt.layer(
328 *(
329 alt.Chart(last.select(PERIOD, pl.col(name).alias("value"))) # ty: ignore[unresolved-attribute]
330 .mark_text(align="left", dx=8, dy=dodge[name], fontSize=11, fontWeight=600, color=CONTEXT)
331 .encode(x, alt.Y("value", type="quantitative"), text=alt.value(name))
332 for name in names
333 )
334 ),
335 )
338def _thin(frame: pl.DataFrame, max_points: int) -> pl.DataFrame:
339 """Every k-th row of an over-long curve, with the last one kept.
341 A line is a shape, not a scatter: dropping intermediate points leaves the shape
342 intact, and it is the browser rather than the reader that notices the difference.
343 """
344 if frame.height <= max_points or max_points < 2:
345 return frame
347 # Counting gaps rather than rows: keeping the last point costs one of the budget.
348 stride = -(-(frame.height - 1) // (max_points - 1))
349 index = pl.int_range(pl.len())
350 return frame.filter((index % stride == 0) | (index == frame.height - 1))