Coverage for src/jquantstats/_plots/_render/_mpl.py: 100%
163 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-23 04:11 +0000
1"""Render a `FigureSpec` with matplotlib.
3Figures are built by constructing `matplotlib.figure.Figure` directly rather
4than through ``pyplot``. pyplot keeps every figure it creates alive in a global
5registry, so a script looping over hundreds of portfolios accumulates figures
6until matplotlib warns and memory grows without bound — the resource cost behind
7issue #628. A directly constructed Figure is owned by its caller, garbage
8collected normally, and still supports ``fig.savefig(...)``. Importing this
9module also never selects a GUI backend or mutates ``rcParams``.
11This backend reproduces the same data and the same static design as the Plotly
12one. It does not emulate interactivity: hover tooltips and the date
13range-selector buttons have no matplotlib equivalent and are silently absent.
14Everything determining *what* is plotted — series, colours, axis scales and tick
15formats — is reproduced.
16"""
18from __future__ import annotations
20from typing import TYPE_CHECKING, Any
22import numpy as np
23import polars as pl
24from matplotlib import colormaps
25from matplotlib.axes import Axes
26from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm, to_rgba
27from matplotlib.figure import Figure
28from matplotlib.ticker import EngFormatter, MultipleLocator, StrMethodFormatter
30from .._spec import (
31 Axis,
32 Band,
33 BarSeries,
34 BoxSeries,
35 Chrome,
36 ColorScale,
37 FigureSpec,
38 HeatmapGrid,
39 HistogramSeries,
40 LineSeries,
41 Panel,
42 RefLine,
43 TickFormat,
44 Values,
45)
47if TYPE_CHECKING:
48 from matplotlib.ticker import Formatter
50__all__ = ["render_mpl"]
52# Pixels per inch used to reinterpret the pixel sizes the specs carry, so
53# `figsize=(920, 420)` frames the same chart on either backend.
54_DPI = 100
56# Sizes for a spec that names none. Plotly lets a figure size itself to its
57# container; matplotlib needs numbers up front, so charts that never fix a
58# dimension get these.
59_DEFAULT_WIDTH_PX = 1000
60_DEFAULT_HEIGHT_PX = 600
62# Semantic tick format -> a str.format field spec. Python's format mini-language
63# covers percentages too, so `StrMethodFormatter` serves every case and no
64# separate `PercentFormatter` is needed.
65_FORMATS = {
66 "float2": "{x:.2f}",
67 "float4": "{x:.4f}",
68 "percent0": "{x:.0%}",
69 "percent1": "{x:.1%}",
70 "percent2": "{x:.2%}",
71 "currency0": "{x:,.0f}",
72}
74# Named colour ramp -> how matplotlib names it. The custom one is built from
75# its stops; the other two are matplotlib built-ins that Plotly borrowed the
76# names from, so they need no translation beyond case.
77_COLORSCALES: dict[ColorScale, tuple[str, ...]] = {
78 "red_white_green": ("#d62728", "#ffffff", "#2ca02c"),
79}
80_BUILTIN_COLORMAPS: dict[ColorScale, str] = {"rdylgn": "RdYlGn", "rdbu_r": "RdBu_r"}
82# Semantic dash style -> matplotlib's linestyle vocabulary.
83_LINESTYLES = {"solid": "-", "dash": "--", None: "-"}
85# The Plotly charts draw on white with a light grey grid; mirror that here so a
86# figure is recognisably the same chart whichever backend drew it.
87_GRID_COLOR = "lightgrey"
88_GRID_WIDTH = 0.5
91def _tick_formatter(tick_format: TickFormat, prefix: str) -> Formatter:
92 """Build a tick formatter for one axis.
94 Args:
95 tick_format: How to render the number.
96 prefix: Written before each tick value, e.g. a currency sign.
98 Returns:
99 Formatter: A formatter applying the requested format and prefix.
101 """
102 if tick_format == "si2":
103 # Plotly's ".2s" abbreviates with an SI suffix — 1500000 reads as
104 # "1.5M". `EngFormatter` is matplotlib's equivalent; `sep=""` keeps the
105 # suffix tight against the number as d3 writes it.
106 return EngFormatter(places=1, sep="")
107 return StrMethodFormatter(f"{prefix}{_FORMATS[tick_format]}")
110def _as_array(values: Values) -> Any:
111 """Convert plotted values to a numpy array.
113 Specs carry either a polars Series or a plain list, whichever the chart
114 already used (see `~jquantstats._plots._spec.Values`). Going via numpy
115 rather than a list also turns a polars null into NaN, which matplotlib
116 draws as a gap; a list of ``None`` would instead force a slow object-dtype
117 array it cannot interpolate across.
119 Args:
120 values: A polars Series or a Python list.
122 Returns:
123 The values as a numpy array — floats where they convert, so ``None``
124 becomes NaN, and otherwise left as-is for x-axes holding dates.
126 """
127 if isinstance(values, pl.Series):
128 return values.to_numpy()
129 try:
130 return np.asarray(values, dtype=float)
131 except TypeError:
132 # A date axis: leave the objects alone, matplotlib understands them.
133 return np.asarray(values)
136def _draw_line(ax: Axes, line: LineSeries) -> None:
137 """Draw one series onto *ax*.
139 The series' `~jquantstats._spec.HoverSpec` is ignored: tooltips are
140 interactive, and this backend is static.
142 Args:
143 ax: The axes to draw on.
144 line: The series to draw.
146 """
147 y = _as_array(line.y)
148 # A series with no x is drawn against its own index, which is what
149 # Plotly does with an unset `x` too.
150 x = np.arange(len(y)) if line.x is None else _as_array(line.x)
151 # A line that names no colour is left to matplotlib's own colour cycle,
152 # matching how Plotly treats an unset `line.color`.
153 styling: dict[str, Any] = {} if line.color is None else {"color": mpl_color(line.color)}
154 if line.markers:
155 styling["marker"] = "o"
156 if line.marker_size is not None:
157 # matplotlib sizes a marker by diameter in points, Plotly by area-ish
158 # "size"; the square root keeps the two visually comparable.
159 styling["markersize"] = line.marker_size**0.5 * 2
160 ax.plot(
161 x,
162 y,
163 linewidth=line.width,
164 linestyle=_LINESTYLES[line.dash],
165 label=line.name,
166 **styling,
167 )
168 if line.fill:
169 # `where` keeps the fill out of the gaps a NaN leaves in the line,
170 # which otherwise get shaded as though the value were zero. An unnamed
171 # fill colour is left to matplotlib, which matches it to the line.
172 shading: dict[str, Any] = {} if line.fill_color is None else {"color": mpl_color(line.fill_color)}
173 ax.fill_between(x, y, 0, where=~np.isnan(y), linewidth=0, **shading)
176def mpl_color(color: str) -> str | tuple[float, float, float, float]:
177 """Translate a spec colour into something matplotlib accepts.
179 Specs spell translucent colours the CSS way — ``rgba(99, 110, 250, 0.4)``
180 — because that is what Plotly emits and the fidelity snapshots pin.
181 matplotlib does not parse that form, so it is converted here; plain hex
182 passes straight through.
184 Args:
185 color: A hex string or a CSS ``rgba()`` string.
187 Returns:
188 The colour as matplotlib understands it.
190 Examples:
191 >>> mpl_color("#636EFA")
192 '#636EFA'
193 >>> mpl_color("rgba(99, 110, 250, 0.4)")
194 (0.38823529411764707, 0.43137254901960786, 0.9803921568627451, 0.4)
196 """
197 if not color.startswith("rgba("):
198 return color
199 r, g, b, a = (part.strip() for part in color[len("rgba(") : -1].split(","))
200 return (int(r) / 255, int(g) / 255, int(b) / 255, float(a))
203def _faded(color: str, opacity: float) -> tuple[float, float, float, float]:
204 """Scale a colour's alpha by *opacity*.
206 Matches Plotly, which multiplies a trace's opacity by the alpha already in
207 its colour, rather than matplotlib's ``alpha=`` argument, which replaces it.
209 Args:
210 color: A hex or CSS ``rgba()`` colour.
211 opacity: Factor to scale the alpha channel by.
213 Returns:
214 The colour as RGBA floats, alpha scaled.
216 Examples:
217 >>> _faded("rgba(0, 0, 0, 0.4)", 0.5)
218 (0.0, 0.0, 0.0, 0.2)
220 """
221 r, g, b, a = to_rgba(mpl_color(color))
222 return (r, g, b, a * opacity)
225def _draw_bars(ax: Axes, bar: BarSeries) -> None:
226 """Draw one bar series onto *ax*.
228 Args:
229 ax: The axes to draw on.
230 bar: The series to draw.
232 """
233 styling: dict[str, Any] = {}
234 if bar.colors is not None:
235 # Opacity is folded into each colour rather than passed as `alpha=`.
236 # matplotlib's alpha argument *replaces* a colour's own alpha channel,
237 # whereas Plotly multiplies its trace opacity by it — so passing it
238 # separately would render the faded negative bars at the wrong strength.
239 styling["color"] = [_faded(c, bar.opacity if bar.opacity is not None else 1.0) for c in bar.colors]
240 elif bar.opacity is not None:
241 # No colours named, so there is no alpha to fold into and nothing to
242 # get wrong: matplotlib's own argument is the right tool.
243 styling["alpha"] = bar.opacity
245 ax.bar(_as_array(bar.x), _as_array(bar.y), linewidth=0, label=bar.name, **styling)
248def _draw_ref_line(ax: Axes, ref: RefLine) -> None:
249 """Draw a fixed-value marker line onto *ax*.
251 Args:
252 ax: The axes to draw on.
253 ref: The line to draw.
255 """
256 draw = ax.axhline if ref.orientation == "h" else ax.axvline
257 draw(ref.value, color=mpl_color(ref.color), linewidth=ref.width, linestyle=_LINESTYLES[ref.dash])
258 if ref.label is not None:
259 # Anchored to the line in data coordinates and to the top of the axes,
260 # matching where Plotly places its annotation.
261 anchor = (0.0, ref.value) if ref.orientation == "h" else (ref.value, 1.0)
262 coords = ("axes fraction", "data") if ref.orientation == "h" else ("data", "axes fraction")
263 ax.annotate(
264 ref.label,
265 xy=anchor,
266 xycoords=coords,
267 xytext=(-2, -2),
268 textcoords="offset points",
269 ha="right",
270 va="top",
271 fontsize=ref.label_size,
272 )
275def _draw_band(ax: Axes, band: Band) -> None:
276 """Draw a shaded vertical span onto *ax*.
278 Args:
279 ax: The axes to draw on.
280 band: The span to draw.
282 """
283 ax.axvspan(band.x0, band.x1, color=mpl_color(band.color), linewidth=0)
284 if band.label is not None:
285 # Placed at the top of the span in axes coordinates, matching where
286 # Plotly puts its annotation.
287 ax.annotate(
288 band.label,
289 xy=(band.x0, 1.0),
290 xycoords=("data", "axes fraction"),
291 xytext=(2, -2),
292 textcoords="offset points",
293 ha="left",
294 va="top",
295 fontsize=band.label_size,
296 )
299def _draw_histogram(ax: Axes, hist: HistogramSeries) -> None:
300 """Draw one histogram series onto *ax*.
302 Args:
303 ax: The axes to draw on.
304 hist: The series to bin and draw.
306 """
307 # `bins=None` is matplotlib's own "use the default count", so an unset bin
308 # count passes straight through rather than needing a separate call.
309 ax.hist(
310 _as_array(hist.values),
311 bins=hist.bins,
312 color=mpl_color(hist.color),
313 alpha=hist.opacity,
314 label=hist.name,
315 )
318def _draw_boxes(ax: Axes, boxes: tuple[BoxSeries, ...]) -> None:
319 """Draw a panel's box-and-whisker series onto *ax*.
321 Drawn together rather than one at a time: matplotlib positions boxes by
322 index within a single call, and wants the category labels alongside.
324 Args:
325 ax: The axes to draw on.
326 boxes: The series to summarise, left to right.
328 """
329 if not boxes:
330 return
331 artists = ax.boxplot(
332 [_as_array(box.values) for box in boxes],
333 tick_labels=[box.name for box in boxes],
334 showfliers=True,
335 patch_artist=True,
336 )
337 for patch, box in zip(artists["boxes"], boxes, strict=True):
338 patch.set_facecolor(mpl_color(box.color))
341def _colormap(scale: ColorScale) -> Any:
342 """Resolve a named colour ramp to a matplotlib colormap.
344 Args:
345 scale: The ramp's semantic name.
347 Returns:
348 The colormap, built from explicit stops or looked up by name.
350 """
351 builtin = _BUILTIN_COLORMAPS.get(scale)
352 if builtin is not None:
353 return colormaps[builtin]
354 return LinearSegmentedColormap.from_list(scale, _COLORSCALES[scale])
357def _heatmap_norm(masked: Any, grid: HeatmapGrid) -> TwoSlopeNorm | None:
358 """Build the diverging normalisation a matrix with a pinned centre needs.
360 Args:
361 masked: The matrix's values, with uncovered cells masked out.
362 grid: The matrix being drawn, for its pinned bounds.
364 Returns:
365 The normalisation, or None when the grid sets no centre or has no
366 unmasked values to scale.
368 """
369 if grid.zmid is None or not masked.count():
370 return None
371 # Pinned bounds win over the data's own range: a correlation runs -1 to
372 # 1 whatever this particular matrix happens to span.
373 low = float(masked.min()) if grid.zmin is None else grid.zmin
374 high = float(masked.max()) if grid.zmax is None else grid.zmax
375 # TwoSlopeNorm needs the centre strictly inside the range; widen a
376 # one-sided or degenerate span so an all-positive year still renders.
377 low = min(low, grid.zmid - 1e-9)
378 high = max(high, grid.zmid + 1e-9)
379 return TwoSlopeNorm(vmin=low, vcenter=grid.zmid, vmax=high)
382def _draw_cell_labels(ax: Axes, grid: HeatmapGrid) -> None:
383 """Write each cell's label into the middle of its square.
385 Args:
386 ax: The axes to draw on.
387 grid: The matrix being drawn, for its per-cell text.
389 """
390 for row, labels in enumerate(grid.text):
391 for col, label in enumerate(labels):
392 if label:
393 ax.text(col, row, label, ha="center", va="center", fontsize=8)
396def _draw_heatmap(fig: Figure, ax: Axes, grid: HeatmapGrid) -> None:
397 """Draw a value matrix onto *ax*, with per-cell labels and a colour bar.
399 Args:
400 fig: The figure owning *ax*, needed to attach the colour bar.
401 ax: The axes to draw on.
402 grid: The matrix to draw.
404 """
405 # None marks a month the data does not cover. NaN carries that through
406 # numpy, and a masked array keeps those cells unpainted rather than
407 # colouring them as if they were zero.
408 values = np.array([[float("nan") if v is None else v for v in row] for row in grid.z], dtype=float)
409 masked = np.ma.masked_invalid(values)
411 image = ax.imshow(masked, cmap=_colormap(grid.colorscale), norm=_heatmap_norm(masked, grid), aspect="auto")
413 ax.set_xticks(range(len(grid.x_labels)), labels=list(grid.x_labels))
414 ax.set_yticks(range(len(grid.y_labels)), labels=list(grid.y_labels))
415 _draw_cell_labels(ax, grid)
417 fig.colorbar(image, ax=ax, label=grid.colorbar_title)
420def _apply_axis(ax: Axes, axis: Axis, *, vertical: bool) -> None:
421 """Apply one axis configuration to *ax*.
423 Only properties the spec actually set are touched, mirroring the Plotly
424 renderer so the two backends agree on what "unset" means.
426 Args:
427 ax: The axes to configure.
428 axis: The requested configuration.
429 vertical: Configure the y-axis rather than the x-axis.
431 """
432 target = ax.yaxis if vertical else ax.xaxis
433 if axis.title is not None:
434 target.set_label_text(axis.title)
435 if axis.tick_format is not None:
436 target.set_major_formatter(_tick_formatter(axis.tick_format, axis.tick_prefix))
437 if axis.log:
438 set_scale = ax.set_yscale if vertical else ax.set_xscale
439 set_scale("log")
440 if axis.dtick is not None:
441 target.set_major_locator(MultipleLocator(axis.dtick))
442 if axis.opposite_side:
443 # Each axis has its own vocabulary for "the far edge", which is why
444 # the spec says `opposite_side` rather than naming a compass point.
445 if vertical:
446 ax.yaxis.set_ticks_position("right")
447 ax.yaxis.set_label_position("right")
448 else:
449 ax.xaxis.set_ticks_position("top")
450 ax.xaxis.set_label_position("top")
453def render_mpl(spec: FigureSpec) -> Figure:
454 """Render *spec* as a static matplotlib figure.
456 Args:
457 spec: The chart to draw.
459 Returns:
460 Figure: The rendered figure. It is not registered with pyplot, so the
461 caller owns it and nothing accumulates between calls.
463 """
464 width, height = spec.figsize if spec.figsize is not None else (_DEFAULT_WIDTH_PX, spec.height)
465 fig = Figure(figsize=(width / _DPI, (height or _DEFAULT_HEIGHT_PX) / _DPI), dpi=_DPI)
466 axes = _make_axes(fig, spec)
468 for ax, panel in zip(axes, spec.panels, strict=True):
469 _draw_panel(fig, ax, panel, spec)
470 fig.suptitle(spec.title)
471 return fig
474def _make_axes(fig: Figure, spec: FigureSpec) -> Any:
475 """Create one Axes per panel, laid out as the spec asks.
477 Args:
478 fig: The figure to add axes to.
479 spec: The chart being rendered.
481 Returns:
482 The axes, flattened to one per panel in spec order.
484 """
485 if spec.arrangement == "stacked":
486 # Stacked panels share the time axis and split the height in the
487 # proportions the panels ask for, so the headline view gets the room.
488 return fig.subplots(
489 nrows=len(spec.panels),
490 sharex=spec.shared_x,
491 height_ratios=[panel.height_ratio for panel in spec.panels],
492 squeeze=False,
493 )[:, 0]
494 return fig.subplots(ncols=len(spec.panels), sharey=spec.shared_y, squeeze=False)[0]
497def _draw_marks(fig: Figure, ax: Axes, panel: Panel) -> None:
498 """Draw every mark the panel carries, in back-to-front order.
500 Reference lines and bands come last so they sit above the series they
501 annotate.
503 Args:
504 fig: The figure owning *ax*, needed to attach a colour bar.
505 ax: The axes to draw on.
506 panel: The panel whose marks to draw.
508 """
509 for line in panel.lines:
510 _draw_line(ax, line)
511 for bar in panel.bars:
512 _draw_bars(ax, bar)
513 for hist in panel.histograms:
514 _draw_histogram(ax, hist)
515 _draw_boxes(ax, panel.boxes)
516 if panel.heatmap is not None:
517 _draw_heatmap(fig, ax, panel.heatmap)
518 for ref in panel.ref_lines:
519 _draw_ref_line(ax, ref)
520 for band in panel.bands:
521 _draw_band(ax, band)
524def _apply_grid(ax: Axes, chrome: Chrome) -> None:
525 """Apply the grid the chart's chrome calls for.
527 Stated either way rather than leaning on ``rcParams["axes.grid"]``: that
528 default is global and third-party libraries flip it on import (quantstats
529 does), which would otherwise put a grid behind a matrix chart depending on
530 what else the caller happened to import.
532 Args:
533 ax: The axes to configure.
534 chrome: The chart's figure-wide chrome setting.
536 """
537 if chrome == "bare":
538 ax.grid(visible=False)
539 elif chrome == "panels":
540 # Horizontal rules help compare box heights across panels; vertical
541 # ones would only clutter what is a categorical axis.
542 ax.grid(visible=True, axis="y", color=_GRID_COLOR, linewidth=_GRID_WIDTH)
543 else:
544 ax.grid(visible=True, color=_GRID_COLOR, linewidth=_GRID_WIDTH)
547def _draw_panel(fig: Figure, ax: Axes, panel: Panel, spec: FigureSpec) -> None:
548 """Draw one panel's marks and configure its axes.
550 Args:
551 fig: The figure owning *ax*, needed to attach a colour bar.
552 ax: The axes to draw on.
553 panel: The panel to draw.
554 spec: The chart being rendered, for its figure-wide chrome.
556 """
557 _draw_marks(fig, ax, panel)
559 ax.set_facecolor("white")
560 _apply_grid(ax, spec.chrome)
562 if panel.title is not None:
563 ax.set_title(panel.title)
564 _apply_axis(ax, panel.xaxis, vertical=False)
565 _apply_axis(ax, panel.yaxis, vertical=True)
567 # Colour carries the value on a matrix chart, so its series names would
568 # label nothing; only series charts get a legend.
569 series_count = len(panel.lines) + len(panel.bars) + len(panel.histograms)
570 if series_count and spec.chrome != "bare":
571 ax.legend(loc="upper left", frameon=False, ncols=series_count)