plot-io 0.1.0__py3-none-any.whl

plot_io/plot.py

@@ -0,0 +1,435 @@
+"""Plots: one area of a figure, with axes, series, title, legend and colorbars."""
+
+from typing import Any
+
+from .axis import Axis
+from .settings import (Setting, Settings, boolean, canonical_kwargs, color, either, literal,
+                       number, one_of, route, table, text)
+from .numbers import number_settings
+from .maps import check_cmap
+from .style import check_palette
+
+__all__ = ["Plot", "Title", "Legend"]
+
+_SPINES = ("left", "bottom", "right", "top")
+_POSITIONS = ("auto", "top right", "top left", "bottom right", "bottom left", "top", "bottom",
+              "outside right", "outside top", "outside bottom", "none")
+
+
+def _spines(v: Any) -> frozenset[str]:
+    if isinstance(v, str):
+        s = v.strip().lower()
+        if s == "all":
+            return frozenset(_SPINES)
+        if s == "none":
+            return frozenset()
+        parts = s.replace(",", " ").split()
+    else:
+        parts = [str(p).lower() for p in v]
+    bad = [p for p in parts if p not in _SPINES]
+    if bad:
+        raise ValueError(f"unknown spine {bad[0]!r} (use left, bottom, right, top, all, none)")
+    return frozenset(parts)
+
+
+def _position(v: Any) -> str:
+    s = str(getattr(v, "value", v)).strip().lower().replace("_", " ").replace("-", " ")
+    s = " ".join(s.split())
+    if s not in _POSITIONS:
+        raise ValueError(f"expected one of {', '.join(map(repr, _POSITIONS))}, got {v!r}")
+    return s
+
+
+def _plot(v: Any) -> Any:
+    if not isinstance(v, Plot):
+        raise ValueError(f"expected a Plot, got {v!r}")
+    return v
+
+
+_INHERITED = dict(
+    font_size=Setting(number(positive=True), inherited=True, doc="Base font size (pt).", length=True),
+    font_family=Setting(str, inherited=True, doc="Font family."),
+    text_color=Setting(color, inherited=True, doc="Text color."),
+    line_width=Setting(number(lo=0), inherited=True, doc="Default line width (pt).", length=True),
+    palette=Setting(check_palette, inherited=True, doc="Colors of series without one."),
+    alpha=Setting(number(0, 1), inherited=True, doc="Default opacity of series."),
+    rasterize=Setting(either(boolean, number(positive=True)), inherited=True,
+                      doc="`True` or a dpi: series drawn as images inside the PDF."),
+    simplify=Setting(either(literal("auto"), boolean, number(positive=True)), inherited=True,
+                     doc="pdf_io's `simplify` for every series: `\"auto\"`, `True`, a dpi or "
+                         "`False`."),
+    color_map=Setting(check_cmap, inherited=True,
+                      doc="Default color map (alias `cmap`): of heatmaps, contours and "
+                          "numbers given as colors."),
+    data=Setting(lambda v: v, inherited=True,
+                 doc="A table (dict, pandas or polars data frame, numpy structured array): "
+                     "strings in the data arguments of series name its columns."),
+    **number_settings(),
+)
+
+
+class Title(Settings):
+    """The title of a plot (``plot.title``)."""
+
+    SETTINGS = table(
+        text=Setting(text, doc="Title text: a string or rich content."),
+        font_size=Setting(number(positive=True),
+                          doc="Font size (pt); default 1.2 × the base font size.", length=True),
+        font_weight=Setting(one_of("normal", "bold"), doc="`normal` or `bold`."),
+        color=Setting(color, doc="Text color."),
+        align=Setting(one_of("center", "left", "right"),
+                      doc="Alignment over the axes."),
+        pad=Setting(number(lo=0), doc="Gap above the axes (pt).", length=True),
+        over=Setting(one_of("axes", "plot"),
+                     doc="Width the title is aligned within: the axes, or the axes and the "
+                         "y axis labels."),
+        font_family=Setting(str, inherited=True, doc="Font family."),
+    )
+
+    def __init__(self, plot: "Plot") -> None:
+        self._plot = plot
+        super().__init__()
+
+    def _parent(self) -> Settings:
+        return self._plot
+
+    def __call__(self, text: Any = None, **settings: Any) -> "Title":
+        if text is not None:
+            settings["text"] = text
+        return self.set(**settings)
+
+    def __repr__(self) -> str:
+        return f"Title({self.get('text')!r})"
+
+
+class Legend(Settings):
+    """The legend of a plot (``plot.legend``): an entry per named series."""
+
+    SETTINGS = table(
+        position=Setting(_position,
+                         doc="`auto`, `top right`, `top left`, `bottom right`, `bottom left`, "
+                             "`top`, `bottom`, `outside right`, `outside top`, `outside bottom` "
+                             "or `none`."),
+        columns=Setting(number(lo=1, integer=True), doc="Number of columns."),
+        frame=Setting(boolean, doc="Draw a frame and background."),
+        title=Setting(text,
+                      doc="Legend title; default the `group` column name."),
+        font_size=Setting(number(positive=True),
+                          doc="Font size (pt).", length=True),
+        color=Setting(color, doc="Text color."),
+        background=Setting(color, doc="Background of a framed legend."),
+        frame_color=Setting(color, doc="Frame color."),
+        pad=Setting(number(lo=0), doc="Gap to the axes (pt).", length=True),
+        frame_width=Setting(number(lo=0), doc="Frame line width (pt).", length=True),
+        frame_radius=Setting(number(lo=0), doc="Corner radius of the frame (pt).",
+                             length=True),
+        frame_alpha=Setting(number(0, 1), doc="Opacity of the frame and background."),
+        glyph_alpha=Setting(number(0, 1),
+                            doc="Fill opacity of `box` and `circle` symbols, outlined in the "
+                                "opaque color; default: the series' own fill."),
+        glyph=Setting(one_of("series", "box", "circle", "line"),
+                      doc="Entry symbol: as the series draws itself, or a filled `box`, a "
+                          "`circle` or a `line` in the series color."),
+        orientation=Setting(one_of("auto", "vertical", "horizontal"),
+                            doc="`horizontal`: entries side by side; `auto`: horizontal at "
+                                "`top`, `bottom`, `outside top` and `outside bottom`."),
+        glyph_width=Setting(number(positive=True),
+                            doc="Width of an entry's symbol, in font sizes."),
+        row_height=Setting(number(positive=True), doc="Height of an entry, in font sizes."),
+        font_family=Setting(str, inherited=True, doc="Font family."),
+    )
+
+    def __init__(self, plot: "Plot") -> None:
+        self._plot = plot
+        super().__init__()
+
+    def _parent(self) -> Settings:
+        return self._plot
+
+    def __call__(self, position: Any = None, **settings: Any) -> "Legend":
+        if position is not None:
+            settings["position"] = position
+        return self.set(**settings)
+
+    def __repr__(self) -> str:
+        return f"Legend({self.get('position')!r})"
+
+
+class Plot(Settings):
+    """One plot: axes, series and decorations.
+
+    ``pl.Plot(**settings)`` creates a one-plot figure (``plot.figure``);
+    ``figure.plot(…)`` adds a plot to an existing figure.
+    """
+
+    SETTINGS = table(
+        spines=Setting(_spines,
+                       doc="Axis lines drawn: any of `left`, `bottom`, `right`, `top`, or "
+                           "`all`, `none`."),
+        aspect=Setting(number(positive=True),
+                       doc="Data aspect ratio: 1 makes one unit as long on both axes."),
+        background=Setting(color, doc="Background color of the axes area."),
+        border=Setting(color, doc="Color of a frame around the axes area (drawn under the "
+                                  "axis lines); default none."),
+        border_width=Setting(number(lo=0), doc="Width of the frame (pt).", length=True),
+        clip=Setting(boolean, doc="Clip series to the axes area."),
+        axes=Setting(boolean, doc="Draw axes; `False` hides both axes, drops padding "
+                                        "and fixes the aspect to 1 (set by pie charts)."),
+        share_x=Setting(_plot, doc="Another plot of the figure whose x axis this plot "
+                                   "shares (groups chain); overrides the figure's `share_x`."),
+        share_y=Setting(_plot, doc="Another plot whose y axis this plot shares."),
+        share=Setting(_plot, doc="Another plot whose x and y axes this plot shares."),
+        **_INHERITED,
+    )
+    #: Keywords accepted besides settings (``title="…"`` sets the title text).
+    EXTRA = ("title",)
+
+    def __init__(self, figure: Any = None, id: Any = None, _twin: Any = None,
+                 **settings: Any) -> None:
+        from .figure import Figure
+
+        own, fig_kw = route(settings, [Plot, Figure], extra={0: Plot.EXTRA})
+        title = own.pop("title", None)
+        if figure is None:
+            figure = Figure(**fig_kw)
+        elif fig_kw:
+            figure.set(**fig_kw)
+        self.figure = figure
+        self.x_axis = Axis(self, "x")
+        self.y_axis = Axis(self, "y")
+        self.title = Title(self)
+        self.legend = Legend(self)
+        self.series: list[Any] = []
+        self.decorations: list[Any] = []
+        self.colorbars: list[Any] = []
+        self.twins: list[Plot] = []
+        #: ``(main plot, "y" or "x")`` for a twin: the axis it has of its own.
+        self.twin_of: tuple[Plot, str] | None = _twin
+        super().__init__(**own)
+        if title is not None:
+            self.title.set(text=title)
+        self.id = figure._attach(self, id) if _twin is None else None
+
+    def _parent(self) -> Settings:
+        return self.twin_of[0] if self.twin_of is not None else self.figure
+
+    def twin_y(self, **settings: Any) -> "Plot":
+        """A plot on the same axes with its own y axis (on the right) and the x axis of
+        this plot: add series to it and set its ``y_axis``."""
+        t = Plot(figure=self.figure, _twin=(self, "y"), **settings)
+        t.x_axis.set(visible=False)
+        if not t.y_axis.is_set("side"):
+            t.y_axis.set(side="right")
+        if not t.is_set("spines"):
+            t.set(spines="left")                 # drawn on the right with the axis
+        self.twins.append(t)
+        return t
+
+    def twin_x(self, **settings: Any) -> "Plot":
+        """A plot on the same axes with its own x axis (on top) and the y axis of this
+        plot."""
+        t = Plot(figure=self.figure, _twin=(self, "x"), **settings)
+        t.y_axis.set(visible=False)
+        if not t.x_axis.is_set("side"):
+            t.x_axis.set(side="top")
+        if not t.is_set("spines"):
+            t.set(spines="bottom")
+        self.twins.append(t)
+        return t
+
+    def set(self, **settings: Any) -> "Plot":
+        settings = canonical_kwargs(settings, self.SETTINGS, "Plot")
+        fig_share = {k: settings.pop(k) for k in ("share_x", "share_y", "share")
+                     if k in settings and settings[k] is not None
+                     and not isinstance(settings[k], Plot)}
+        if fig_share:
+            self.figure.set(**fig_share)
+        title = settings.pop("title", None)
+        if title is not None:
+            self.title.set(text=title)
+        super().set(**settings)
+        return self
+
+    def __repr__(self) -> str:
+        return f"Plot({self.id!r})"
+
+    def own_share(self, role: str) -> Any:
+        """The plot this plot's ``role`` axis is paired with, or ``None``."""
+        return self.get(f"share_{role}") or self.get("share")
+
+    def _add_series(self, s: Any) -> None:
+        self.series.append(s)
+
+    # -- series ----------------------------------------------------------------------
+
+    def scatter(self, x: Any, y: Any, **settings: Any) -> Any:
+        """Markers at points ``(x, y)``: see :class:`~plot_io.Scatter`."""
+        from .series.scatter import Scatter
+        return Scatter(x, y, plot=self, **settings)
+
+    def line(self, x: Any, y: Any = None, **settings: Any) -> Any:
+        """A line through points: see :class:`~plot_io.Line`."""
+        from .series.line import Line
+        return Line(x, y, plot=self, **settings)
+
+    def fill(self, x: Any, y1: Any, y2: Any = 0.0, **settings: Any) -> Any:
+        """The area between two curves: see :class:`~plot_io.Fill`."""
+        from .series.fill import Fill
+        return Fill(x, y1, y2, plot=self, **settings)
+
+    def box(self, values: Any, **settings: Any) -> Any:
+        """Box plots: see :class:`~plot_io.Box`."""
+        from .series.box import Box
+        return Box(values, plot=self, **settings)
+
+    def heatmap(self, values: Any, **settings: Any) -> Any:
+        """A grid of values as colored cells: see :class:`~plot_io.Heatmap`."""
+        from .series.heatmap import Heatmap
+        return Heatmap(values, plot=self, **settings)
+
+    def bar(self, x: Any, height: Any = None, **settings: Any) -> Any:
+        """Bars: see :class:`~plot_io.Bar`."""
+        from .series.bar import Bar
+        return Bar(x, height, plot=self, **settings)
+
+    def hist(self, values: Any, **settings: Any) -> Any:
+        """A histogram: see :class:`~plot_io.Hist`."""
+        from .series.hist import Hist
+        return Hist(values, plot=self, **settings)
+
+    def kde(self, values: Any, y: Any = None, **settings: Any) -> Any:
+        """A kernel density estimate: see :class:`~plot_io.Kde`; with ``y``, 2-D
+        density contours (:class:`~plot_io.Kde2d`)."""
+        if y is not None:
+            from .series.kde2d import Kde2d
+            return Kde2d(values, y, plot=self, **settings)
+        from .series.kde import Kde
+        return Kde(values, plot=self, **settings)
+
+    def ecdf(self, values: Any, **settings: Any) -> Any:
+        """An empirical cumulative distribution: see :class:`~plot_io.Ecdf`."""
+        from .series.ecdf import Ecdf
+        return Ecdf(values, plot=self, **settings)
+
+    def violin(self, values: Any, **settings: Any) -> Any:
+        """Violin plots: see :class:`~plot_io.Violin`."""
+        from .series.violin import Violin
+        return Violin(values, plot=self, **settings)
+
+    def strip(self, values: Any, **settings: Any) -> Any:
+        """Jittered points per category: see :class:`~plot_io.Strip`."""
+        from .series.catpoints import Strip
+        return Strip(values, plot=self, **settings)
+
+    def swarm(self, values: Any, **settings: Any) -> Any:
+        """Non-overlapping points per category: see :class:`~plot_io.Swarm`."""
+        from .series.catpoints import Swarm
+        return Swarm(values, plot=self, **settings)
+
+    def area(self, x: Any, ys: Any, **settings: Any) -> Any:
+        """Stacked areas: see :class:`~plot_io.Area`."""
+        from .series.area import Area
+        return Area(x, ys, plot=self, **settings)
+
+    def errorbar(self, x: Any, y: Any, **settings: Any) -> Any:
+        """Error bars: see :class:`~plot_io.Errorbar`."""
+        from .series.errorbar import Errorbar
+        return Errorbar(x, y, plot=self, **settings)
+
+    def stem(self, x: Any, y: Any = None, **settings: Any) -> Any:
+        """Stems with markers: see :class:`~plot_io.Stem`."""
+        from .series.stem import Stem
+        return Stem(x, y, plot=self, **settings)
+
+    def regression(self, x: Any, y: Any, **settings: Any) -> Any:
+        """A fitted line with a confidence band: see :class:`~plot_io.Regression`."""
+        from .series.regression import Regression
+        return Regression(x, y, plot=self, **settings)
+
+    def hist2d(self, x: Any, y: Any, **settings: Any) -> Any:
+        """A 2-D histogram: see :class:`~plot_io.Hist2d`."""
+        from .series.hist2d import Hist2d
+        return Hist2d(x, y, plot=self, **settings)
+
+    def hexbin(self, x: Any, y: Any, **settings: Any) -> Any:
+        """A hexagonal 2-D histogram: see :class:`~plot_io.Hexbin`."""
+        from .series.hexbin import Hexbin
+        return Hexbin(x, y, plot=self, **settings)
+
+    def contour(self, values: Any, **settings: Any) -> Any:
+        """Contour lines and bands of a grid: see :class:`~plot_io.Contour`."""
+        from .series.contour import Contour
+        return Contour(values, plot=self, **settings)
+
+    def pie(self, values: Any, **settings: Any) -> Any:
+        """A pie or donut chart: see :class:`~plot_io.Pie`."""
+        from .series.pie import Pie
+        return Pie(values, plot=self, **settings)
+
+    # -- decorations -----------------------------------------------------------------
+
+    def hline(self, y: Any, **settings: Any) -> Any:
+        """A horizontal reference line at ``y``."""
+        from .decorations import HLine
+        return HLine(self, y, **settings)
+
+    def vline(self, x: Any, **settings: Any) -> Any:
+        """A vertical reference line at ``x``."""
+        from .decorations import VLine
+        return VLine(self, x, **settings)
+
+    def axline(self, point: Any, slope: Any = None, point2: Any = None,
+               **settings: Any) -> Any:
+        """An infinite line through ``point`` with ``slope``, or through ``point2``."""
+        from .decorations import AxLine
+        return AxLine(self, point, slope, point2, **settings)
+
+    def hspan(self, y0: Any, y1: Any, **settings: Any) -> Any:
+        """A horizontal band from ``y0`` to ``y1``."""
+        from .decorations import HSpan
+        return HSpan(self, y0, y1, **settings)
+
+    def vspan(self, x0: Any, x1: Any, **settings: Any) -> Any:
+        """A vertical band from ``x0`` to ``x1``."""
+        from .decorations import VSpan
+        return VSpan(self, x0, x1, **settings)
+
+    def text(self, x: Any, y: Any, content: Any, **settings: Any) -> Any:
+        """Text at ``(x, y)``."""
+        from .decorations import Text
+        return Text(self, x, y, content, **settings)
+
+    def annotate(self, content: Any, xy: Any, xytext: Any = None, **settings: Any) -> Any:
+        """Text at ``xytext`` with an arrow to ``xy``."""
+        from .decorations import Annotate
+        return Annotate(self, content, xy, xytext, **settings)
+
+    def colorbar(self, source: Any = None, **settings: Any) -> Any:
+        """A colorbar attached to this plot (``source``: a mapped series or a
+        ``ColorMap``; default: the last mapped series)."""
+        from .colorbar import Colorbar
+        cb = Colorbar(source, plot=self, **settings)
+        self.colorbars.append(cb)
+        return cb
+
+    # -- output ----------------------------------------------------------------------
+
+    def draw(self, canvas: Any, x: float = 0.0, y: float = 0.0) -> tuple[float, float]:
+        """Draw the plot alone on a pdf_io canvas with its top-left corner at ``(x, y)``
+        (page units); returns its size in points."""
+        from . import render
+        return render.draw(self.figure, canvas, x, y, only=self)
+
+    def save(self, path: Any, **options: Any) -> None:
+        """Save the plot alone: PDF, PNG, JPEG or JPEG XL by the extension (see
+        :meth:`Figure.save`)."""
+        from . import render
+        render.save(self.figure, path, only=self, **options)
+
+    def to_bytes(self, format: str = "pdf", **options: Any) -> bytes:
+        """The plot alone as ``"pdf"``, ``"png"``, ``"jpeg"`` or ``"jxl"`` bytes."""
+        from . import render
+        return render.to_bytes(self.figure, format, only=self, **options)
+
+    def _repr_png_(self) -> bytes:
+        return self.to_bytes("png")