"""Helper functions for `matplotlib`."""
from __future__ import annotations
import os
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.collections import LineCollection, PathCollection
from matplotlib.contour import QuadContourSet
[docs]def add_watermark(
ax: Axes,
x: float = 0.03,
y: float = 0.03,
fontsize: int | None = None,
**kwargs,
) -> None:
text = "LHCb\n" R"$1.7\mathrm{~fb}^{-1}$"
ax.text(x, y, text, size=fontsize, transform=ax.transAxes, **kwargs)
[docs]def get_contour_line(contour_set: QuadContourSet) -> LineCollection:
(line_collection, *_), _ = contour_set.legend_elements()
return line_collection
[docs]def use_mpl_latex_fonts(reset_mpl: bool = True) -> None:
# cspell:ignore dejavusans fontset mathtext usetex
if not _is_latex_allowed():
return
if reset_mpl:
_wake_up_matplotlib()
plt.rc("font", family="serif", serif="Helvetica")
plt.rc("mathtext", fontset="dejavusans")
plt.rc("text", usetex=True)
def _is_latex_allowed() -> bool:
if "BINDER_LAUNCH_HOST" in os.environ:
return False
return True
def _wake_up_matplotlib() -> None:
"""Somehow `plt.rc` does not work if a figure hasn't been created before..."""
plt.figure()
plt.close()
[docs]def stylize_contour(
contour_set: QuadContourSet,
*,
edgecolor=None,
label: str | None = None,
linestyle: str | None = None,
linewidth: float | None = None,
) -> None:
contour_line: PathCollection = contour_set.collections[0]
if edgecolor is not None:
contour_line.set_edgecolor(edgecolor)
if label is not None:
contour_line.set_label(label)
if linestyle is not None:
contour_line.set_linestyle(linestyle)
if linewidth is not None:
contour_line.set_linewidth(linewidth)
