plotnine 0.15.2__py3-none-any.whl → 0.16.0a1__py3-none-any.whl

plotnine/composition/_plot_layout.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from itertools import cycle
+from typing import TYPE_CHECKING, Sequence
+
+from ..composition._types import ComposeAddable
+
+if TYPE_CHECKING:
+    from ._compose import Compose
+
+
+@dataclass(kw_only=True)
+class plot_layout(ComposeAddable):
+    """
+    Customise the layout of plots in a composition
+    """
+
+    nrow: int | None = None
+    """
+    Number of rows
+    """
+
+    ncol: int | None = None
+    """
+    Number of columns
+    """
+
+    byrow: bool | None = None
+    """
+    How to place plots into the grid.
+    If None or True, they are placed row by row, left to right.
+    If False, they are placed column by column, top to bottom.
+    """
+
+    widths: Sequence[float] | None = None
+    """
+    Relative widths of each column
+    """
+
+    heights: Sequence[float] | None = None
+    """
+    Relative heights of each column
+    """
+
+    _cmp: Compose = field(init=False, repr=False)
+    """
+    Composition that this layout is attached to
+    """
+
+    def __radd__(self, cmp: Compose) -> Compose:
+        """
+        Add plot layout to composition
+        """
+        cmp.layout = self
+        return cmp
+
+    def _setup(self, cmp: Compose):
+        """
+        Setup default parameters as they are expected by the layout manager
+
+        - Ensure nrow and ncol have values
+        - Ensure the widths & heights are set and normalised to mean=1
+        """
+        from . import Beside, Stack
+
+        # setup nrow & ncol
+        if isinstance(cmp, Beside):
+            if self.ncol is None:
+                self.ncol = len(cmp)
+            elif self.ncol < len(cmp):
+                raise ValueError(
+                    "Composition has more items than the layout columns."
+                )
+            if self.nrow is None:
+                self.nrow = 1
+        elif isinstance(cmp, Stack):
+            if self.nrow is None:
+                self.nrow = len(cmp)
+            elif self.nrow < len(cmp):
+                raise ValueError(
+                    "Composition has more items than the layout rows."
+                )
+
+            if self.ncol is None:
+                self.ncol = 1
+        else:
+            from plotnine.facets.facet_wrap import wrap_dims
+
+            self.nrow, self.ncol = wrap_dims(len(cmp), self.nrow, self.ncol)
+
+        nrow, ncol = self.nrow, self.ncol
+
+        # byrow
+        if self.byrow is None:
+            self.byrow = True
+
+        # setup widths & heights
+        ws, hs = self.widths, self.heights
+        if ws is None:
+            ws = (1 / ncol,) * ncol
+        elif len(ws) != ncol:
+            ws = repeat(ws, ncol)
+
+        if hs is None:
+            hs = (1 / nrow,) * nrow
+        elif len(hs) != nrow:
+            hs = repeat(hs, nrow)
+
+        self.widths = normalise(ws)
+        self.heights = normalise(hs)
+
+    def update(self, other: plot_layout):
+        """
+        Update this layout with the contents of other
+        """
+        if other.widths:
+            self.widths = other.widths
+        if other.heights:
+            self.heights = other.heights
+        if other.ncol:
+            self.ncol = other.ncol
+        if other.nrow:
+            self.nrow = other.nrow
+        if other.byrow is not None:
+            self.byrow = other.byrow
+
+
+def repeat(seq: Sequence[float], n: int) -> list[float]:
+    """
+    Ensure returned sequence has n values, repeat as necessary
+    """
+    return [val for _, val in zip(range(n), cycle(seq))]
+
+
+def normalise(seq: Sequence[float]) -> list[float]:
+    """
+    Normalise seq so that the mean is 1
+    """
+    mean = sum(seq) / len(seq)
+    if mean == 0:
+        raise ValueError("Cannot rescale: mean is zero")
+    return [x / mean for x in seq]

plotnine/composition/_plot_spacer.py

@@ -39,7 +39,7 @@ class plot_spacer(ggplot):
         if fill:
             self.theme += theme(plot_background=element_rect(fill=fill))
 
-    def __add__(self, rhs) -> plot_spacer:
+    def __add__(self, rhs) -> plot_spacer:  # pyright: ignore[reportIncompatibleMethodOverride]
         """
         Add to spacer
 

plotnine/composition/_stack.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from ._compose import Compose
@@ -9,7 +8,6 @@ if TYPE_CHECKING:
     from plotnine.ggplot import ggplot
 
 
-@dataclass(repr=False)
 class Stack(Compose):
     """
     Place plots or compositions on top of each other
@@ -26,25 +24,18 @@ class Stack(Compose):
     See Also
     --------
     plotnine.composition.Beside : To arrange plots side by side
+    plotnine.composition.Wrap : To arrange plots in a grid
     plotnine.composition.plot_spacer : To add a blank space between plots
     plotnine.composition.Compose : For more on composing plots
     """
 
-    @property
-    def nrow(self) -> int:
-        return len(self)
-
-    @property
-    def ncol(self) -> int:
-        return 1
-
     def __truediv__(self, rhs: ggplot | Compose) -> Compose:
         """
         Add rhs as a row
         """
         # This is an adjacent div i.e. (DIV | rhs) so we collapse the
         # operands into a single operation
-        return Stack([*self, rhs])
+        return Stack([*self, rhs]) + self.layout + self.annotation
 
     def __or__(self, rhs: ggplot | Compose) -> Compose:
         """
@@ -53,3 +44,14 @@ class Stack(Compose):
         from ._beside import Beside
 
         return Beside([self, rhs])
+
+    def __add__(self, rhs):
+        """
+        Add rhs into the stacking composition
+        """
+        from plotnine import ggplot
+
+        if not isinstance(rhs, (ggplot, Compose)):
+            return super().__add__(rhs)
+
+        return self / rhs

plotnine/composition/_types.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._compose import Compose
+
+
+class ComposeAddable:
+    """
+    Object that can be added to a ggplot object
+    """
+
+    def __radd__(self, other: Compose) -> Compose:
+        """
+        Add to compose object
+
+        Parameters
+        ----------
+        other :
+            Compose object
+
+        Returns
+        -------
+        :
+            Compose object
+        """
+        return other

plotnine/composition/_wrap.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from ..ggplot import ggplot
+from ._compose import Compose
+
+
+class Wrap(Compose):
+    """
+    Wrap plots or compositions into a grid
+
+    **Usage**
+
+        plot + plot
+        plot + composition
+        composition + plot
+        composition + composition
+
+    Typically, you will use this class through the `+` operator.
+
+    Parameters
+    ----------
+    items:
+        The objects to be arranged (composed)
+    nrow:
+        Number of rows in the composition
+    ncol:
+        Number of cols in the composition
+
+    See Also
+    --------
+    plotnine.composition.Beside : To arrange plots side by side
+    plotnine.composition.Stack : To arrange plots vertically
+    plotnine.composition.plot_spacer : To add a blank space between plots
+    plotnine.composition.Compose : For more on composing plots
+    """
+
+    def __add__(self, rhs):
+        """
+        Add rhs into the wrapping composition
+        """
+        if not isinstance(rhs, (ggplot, Compose)):
+            return super().__add__(rhs)
+
+        return Wrap([*self, rhs]) + self.layout + self.annotation
+
+    def __or__(self, rhs: ggplot | Compose) -> Compose:
+        """
+        Add rhs as a column
+        """
+        from ._beside import Beside
+
+        return Beside([self, rhs])
+
+    def __truediv__(self, rhs: ggplot | Compose) -> Compose:
+        """
+        Add rhs as a row
+        """
+        from ._stack import Stack
+
+        return Stack([self, rhs])

plotnine/facets/facet.py

@@ -93,7 +93,6 @@ class facet:
 
     # Axes
     axs: list[Axes]
-    _panels_gridspec: p9GridSpec
 
     # ggplot object that the facet belongs to
     plot: ggplot
@@ -144,15 +143,15 @@ class facet:
         self.figure = plot.figure
 
         if hasattr(plot, "axs"):
-            self.axs = plot.axs
+            gs, self.axs = plot._sub_gridspec, plot.axs
         else:
-            self.axs = self._make_axes()
+            gs, self.axs = self._make_axes()
 
         self.coordinates = plot.coordinates
         self.theme = plot.theme
         self.layout.axs = self.axs
         self.strips = Strips.from_facet(self)
-        return self.axs
+        return gs, self.axs
 
     def setup_data(self, data: list[pd.DataFrame]) -> list[pd.DataFrame]:
         """
@@ -378,7 +377,7 @@ class facet:
 
         return result
 
-    def _get_panels_gridspec(self) -> p9GridSpec:
+    def _make_gridspec(self):
         """
         Create gridspec for the panels
         """
@@ -388,21 +387,19 @@ class facet:
             self.nrow, self.ncol, self.figure, nest_into=self.plot._gridspec[0]
         )
 
-    def _make_axes(self) -> list[Axes]:
+    def _make_axes(self) -> tuple[p9GridSpec, list[Axes]]:
         """
         Create and return subplot axes
         """
+
         num_panels = len(self.layout.layout)
         axsarr = np.empty((self.nrow, self.ncol), dtype=object)
-
-        self._panels_gridspec = self._get_panels_gridspec()
+        gs = self._make_gridspec()
 
         # Create axes
         it = itertools.product(range(self.nrow), range(self.ncol))
         for i, (row, col) in enumerate(it):
-            axsarr[row, col] = self.figure.add_subplot(
-                self._panels_gridspec[i]
-            )
+            axsarr[row, col] = self.figure.add_subplot(gs[i])
 
         # Rearrange axes
         # They are ordered to match the positions in the layout table
@@ -423,7 +420,7 @@ class facet:
         for ax in axs[num_panels:]:
             self.figure.delaxes(ax)
         axs = axs[:num_panels]
-        return list(axs)
+        return gs, list(axs)
 
     def _aspect_ratio(self) -> Optional[float]:
         """

plotnine/facets/facet_grid.py

@@ -107,7 +107,7 @@ class facet_grid(facet):
         self.space = space
         self.margins = margins
 
-    def _get_panels_gridspec(self):
+    def _make_gridspec(self):
         """
         Create gridspec for the panels
         """
@@ -189,7 +189,7 @@ class facet_grid(facet):
 
         n = len(base)
         panel = ninteraction(base, drop=True)
-        panel = pd.Categorical(panel, categories=range(1, n + 1))
+        panel = pd.Categorical(panel, categories=range(1, n + 1))  # pyright: ignore[reportArgumentType]
 
         if self.rows:
             rows = ninteraction(base[self.rows], drop=True)

plotnine/facets/facet_wrap.py

@@ -115,7 +115,7 @@ class facet_wrap(facet):
 
         layout = pd.DataFrame(
             {
-                "PANEL": pd.Categorical(range(1, n + 1)),
+                "PANEL": pd.Categorical(range(1, n + 1)),  # pyright: ignore[reportArgumentType]
                 "ROW": row.astype(int),
                 "COL": col.astype(int),
             }

plotnine/geoms/geom.py

@@ -166,10 +166,10 @@ class geom(ABC, metaclass=Register):
         shallow = {"data", "_kwargs", "environment"}
         for key, item in old.items():
             if key in shallow:
-                new[key] = item
+                new[key] = item  # pyright: ignore[reportIndexIssue]
                 memo[id(new[key])] = new[key]
             else:
-                new[key] = deepcopy(item, memo)
+                new[key] = deepcopy(item, memo)  # pyright: ignore[reportIndexIssue]
 
         return result
 

plotnine/geoms/geom_map.py

@@ -126,18 +126,18 @@ class geom_map(geom):
         params = self.params
         data.loc[data["color"].isna(), "color"] = "none"
         data.loc[data["fill"].isna(), "fill"] = "none"
-        data["fill"] = to_rgba(data["fill"], data["alpha"])
 
         geom_type = data.geometry.iloc[0].geom_type
         if geom_type in ("Polygon", "MultiPolygon"):
             from matplotlib.collections import PatchCollection
 
             linewidth = data["size"] * SIZE_FACTOR
+            fill = to_rgba(data["fill"], data["alpha"])
             patches = [PolygonPatch(g) for g in data["geometry"]]
             coll = PatchCollection(
                 patches,
                 edgecolor=data["color"],
-                facecolor=data["fill"],
+                facecolor=fill,
                 linestyle=data["linetype"],
                 linewidth=linewidth,
                 zorder=params["zorder"],
@@ -152,7 +152,6 @@ class geom_map(geom):
             data["y"] = arr[:, 1]
             for _, gdata in data.groupby("group"):
                 gdata.reset_index(inplace=True, drop=True)
-                gdata.is_copy = None
                 geom_point.draw_group(gdata, panel_params, coord, ax, params)
         elif geom_type == "MultiPoint":
             # Where n is the length of the dataframe (no. of multipoints),
@@ -173,7 +172,7 @@ class geom_map(geom):
             from matplotlib.collections import LineCollection
 
             linewidth = data["size"] * SIZE_FACTOR
-            data["color"] = to_rgba(data["color"], data["alpha"])
+            color = to_rgba(data["color"], data["alpha"])
             segments = []
             for g in data["geometry"]:
                 if g.geom_type == "LineString":
@@ -183,7 +182,7 @@ class geom_map(geom):
 
             coll = LineCollection(
                 segments,
-                edgecolor=data["color"],
+                edgecolor=color,
                 linewidth=linewidth,
                 linestyle=data["linetype"],
                 zorder=params["zorder"],

plotnine/geoms/geom_path.py

@@ -96,8 +96,9 @@ class geom_path(geom):
         n2 = len(data)
 
         if n2 != n1 and not self.params["na_rm"]:
-            msg = "geom_path: Removed {} rows containing missing values."
-            warn(msg.format(n1 - n2), PlotnineWarning)
+            geom = self.__class__.__name__
+            msg = f"{geom}: Removed {n1 - n2} rows containing missing values."
+            warn(msg, PlotnineWarning)
 
         return data
 
@@ -109,8 +110,9 @@ class geom_path(geom):
         ax: Axes,
     ):
         if not any(data["group"].duplicated()):
+            geom = self.__class__.__name__
             warn(
-                "geom_path: Each group consist of only one "
+                f"{geom}: Each group consist of only one "
                 "observation. Do you need to adjust the "
                 "group aesthetic?",
                 PlotnineWarning,
@@ -300,7 +302,7 @@ class arrow:
         last = self.ends in ("last", "both")
 
         data = data.sort_values("group", kind="mergesort")
-        data["color"] = to_rgba(data["color"], data["alpha"])
+        data["color"] = to_rgba(data["color"], data["alpha"])  # pyright: ignore[reportCallIssue,reportArgumentType]
 
         if self.type == "open":
             data["facecolor"] = "none"
@@ -423,9 +425,8 @@ class arrow:
         # We need the axes dimensions so that we can
         # compute scaling factors
         width, height = _axes_get_size_inches(ax)
-        ranges = coord.range(panel_params)
-        width_ = np.ptp(ranges.x)
-        height_ = np.ptp(ranges.y)
+        width_ = np.ptp(panel_params.x.range)
+        height_ = np.ptp(panel_params.y.range)
 
         # scaling factors to prevent skewed arrowheads
         lx = self.length * width_ / width

plotnine/geoms/geom_rug.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import typing
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
@@ -10,7 +10,7 @@ from ..doctools import document
 from .geom import geom
 from .geom_path import geom_path
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from typing import Any
 
     import pandas as pd
@@ -89,30 +89,26 @@ class geom_rug(geom):
         xheight = (xmax - xmin) * params["length"]
         yheight = (ymax - ymin) * params["length"]
 
-        # Please the type checker
-        x: FloatArray
-        y: FloatArray
-
         if has_x:
+            x = cast("FloatArray", np.repeat(data["x"].to_numpy(), 2))
+
             if "b" in sides:
-                x = np.repeat(data["x"].to_numpy(), 2)
                 y = np.tile([ymin, ymin + yheight], n)
                 rugs.extend(make_line_segments(x, y, ispath=False))
 
             if "t" in sides:
-                x = np.repeat(data["x"].to_numpy(), 2)
                 y = np.tile([ymax - yheight, ymax], n)
                 rugs.extend(make_line_segments(x, y, ispath=False))
 
         if has_y:
+            y = cast("FloatArray", np.repeat(data["y"].to_numpy(), 2))
+
             if "l" in sides:
                 x = np.tile([xmin, xmin + xheight], n)
-                y = np.repeat(data["y"].to_numpy(), 2)
                 rugs.extend(make_line_segments(x, y, ispath=False))
 
             if "r" in sides:
                 x = np.tile([xmax - xheight, xmax], n)
-                y = np.repeat(data["y"].to_numpy(), 2)
                 rugs.extend(make_line_segments(x, y, ispath=False))
 
         color = to_rgba(data["color"], data["alpha"])

plotnine/geoms/geom_text.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-import typing
 from contextlib import suppress
+from typing import TYPE_CHECKING, cast
 from warnings import warn
 
 import numpy as np
@@ -12,7 +12,7 @@ from ..exceptions import PlotnineError, PlotnineWarning
 from ..positions import position_nudge
 from .geom import geom
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from typing import Any, Sequence
 
     import pandas as pd
@@ -244,7 +244,7 @@ class geom_text(geom):
             axis=1,
             inplace=True,
         )
-        plot_data["color"] = color
+        plot_data["color"] = color  # pyright: ignore[reportCallIssue,reportArgumentType]
         plot_data["zorder"] = zorder
         plot_data["rasterized"] = params["raster"]
         plot_data["clip_on"] = True
@@ -255,7 +255,7 @@ class geom_text(geom):
             fill = to_rgba(data.pop("fill"), data["alpha"])
             if isinstance(fill, tuple):
                 fill = [list(fill)] * len(data["x"])
-            plot_data["facecolor"] = fill
+            plot_data["facecolor"] = fill  # pyright: ignore[reportCallIssue,reportArgumentType]
 
             tokens = [params["boxstyle"], f"pad={params['label_padding']}"]
             if params["boxstyle"] in {"round", "round4"}:
@@ -272,7 +272,7 @@ class geom_text(geom):
 
         # For labels add a bbox
         for i in range(len(data)):
-            kw: dict[str, Any] = plot_data.iloc[i].to_dict()
+            kw = cast("dict[str, Any]", plot_data.iloc[i].to_dict())
             if draw_label:
                 kw["bbox"] = bbox
                 kw["bbox"]["edgecolor"] = params["boxcolor"] or kw["color"]