Source code for plotnine_extra.facets.facet_manual

"""
Manual panel layout using a design string or matrix.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import numpy as np
import pandas as pd
from plotnine._utils import join_keys, match
from plotnine.facets.facet import (
    add_missing_facets,
    combine_vars,
    eval_facet_vars,
    facet,
)
from plotnine.facets.strips import Strips, strip

from ..guides import apply_axis_guides
from .scale_facet import apply_scale_facets

if TYPE_CHECKING:
    from typing import Optional, Sequence, Union

    from matplotlib.axes import Axes
    from plotnine._mpl.gridspec import p9GridSpec
    from plotnine.iapi import layout_details
    from plotnine.scales.scale import scale


def _parse_design(design: str) -> np.ndarray:
    """
    Parse a design string into a 2-D array of panel labels.

    Parameters
    ----------
    design : str
        Multi-line string where each character is a panel
        label. Use ``#`` for empty cells.

    Returns
    -------
    numpy.ndarray
        2-D character array of panel identifiers.

    Examples
    --------
    >>> _parse_design("AB\\nCD")
    array([['A', 'B'],
           ['C', 'D']], dtype='<U1')
    """
    rows = [list(line) for line in design.strip().splitlines()]
    return np.array(rows)


[docs] class facet_manual(facet): """ Manually specify panel layout via a design matrix. Allows arbitrary placement of facet panels using a text design string or a 2-D array. Use ``#`` to denote empty cells in the layout. Parameters ---------- facets : str or list of str, optional Variables to facet by. Each unique combination is assigned a panel label from the design. design : str or array-like, optional Layout specification. A multi-line string like ``"AB\\n#C"`` or a 2-D list / numpy array. scales : str Whether scales are fixed ("fixed"), free ("free", "free_x", "free_y"). shrink : bool Whether to shrink scales to fit output of statistics. labeller : str Labelling function for strip text. as_table : bool If True, facets are laid out like a table. drop : bool Whether to drop unused factor levels. respect : bool If True, panel sizes respect the aspect ratio. """ def __init__( self, facets: Optional[str | Sequence[str]] = None, *, design: Optional[Union[str, list, np.ndarray]] = None, scales: Literal["fixed", "free", "free_x", "free_y"] = "fixed", shrink: bool = True, labeller: Literal[ "label_value", "label_both", "label_context" ] = "label_value", as_table: bool = True, drop: bool = True, respect: bool = False, widths: Optional[Sequence[float]] = None, heights: Optional[Sequence[float]] = None, axes: Literal["all", "x", "y", "margins"] = "margins", remove_labels: Literal["none", "x", "y", "all"] = "none", trim_blank: bool = False, ): super().__init__( scales=scales, shrink=shrink, labeller=labeller, as_table=as_table, drop=drop, ) self.facets = facets self.design = design self.respect = respect self.widths = list(widths) if widths is not None else None self.heights = list(heights) if heights is not None else None self._panel_widths = self.widths self._panel_heights = self.heights self.axes = axes self.remove_labels = remove_labels self.trim_blank = trim_blank self._design_matrix: Optional[np.ndarray] = None if design is not None: if isinstance(design, str): self._design_matrix = _parse_design(design) else: self._design_matrix = np.asarray(design) # Normalize facets to a list if facets is None: self.vars: list[str] = [] elif isinstance(facets, str): self.vars = [facets] else: self.vars = list(facets)
[docs] def compute_layout( self, data: list[pd.DataFrame], ) -> pd.DataFrame: if self._design_matrix is None: msg = "facet_manual requires a 'design' argument" raise ValueError(msg) design = self._design_matrix design, kept_rows, kept_cols = self._trim_design(design) self._set_panel_sizes(kept_rows, kept_cols) # Extract unique panel labels from design, preserving order, # excluding '#' / NA (empty cells) labels: list[str] = [] for row in design: for cell in row: if not _is_empty_design_cell(cell): label = str(cell) if label not in labels: labels.append(label) # Get facet variable combinations from data if self.vars: base = combine_vars( data, self.environment, self.vars, drop=self.drop ) else: base = pd.DataFrame({"_dummy_": [1]}) if len(labels) < len(base): msg = ( "facet_manual design does not provide enough panels " f"for {len(base)} facet combinations" ) raise ValueError(msg) if len(labels) > len(base): msg = ( "facet_manual design provides more panels than facet " f"combinations ({len(labels)} labels for {len(base)} panels)" ) raise ValueError(msg) n_panels = len(base) rows_list = [] for i in range(n_panels): label = labels[i] positions = np.argwhere(design.astype(str) == label) if positions.size == 0: msg = f"facet_manual design label {label!r} not found" raise ValueError(msg) row_min, col_min = positions.min(axis=0) row_max, col_max = positions.max(axis=0) _check_rectangular_span( label, positions, int(row_max - row_min + 1), int(col_max - col_min + 1), ) r_pos = int(row_min + 1) c_pos = int(col_min + 1) row_data: dict = { "PANEL": i + 1, "ROW": r_pos, "COL": c_pos, "ROWSPAN": int(row_max - row_min + 1), "COLSPAN": int(col_max - col_min + 1), "SCALE_X": i + 1 if self.free["x"] else 1, "SCALE_Y": i + 1 if self.free["y"] else 1, "AXIS_X": True, "AXIS_Y": True, } # Add facet variable values if self.vars and i < len(base): for var in self.vars: row_data[var] = base.iloc[i][var] rows_list.append(row_data) layout = pd.DataFrame(rows_list) layout["PANEL"] = pd.Categorical( layout["PANEL"], categories=range(1, n_panels + 1), ) # Remove helper column if no facet vars if "_dummy_" in layout.columns: layout = layout.drop(columns=["_dummy_"]) self.nrow = int(design.shape[0]) self.ncol = int(design.shape[1]) self._apply_axis_flags(layout) return layout
def _trim_design( self, design: np.ndarray, ) -> tuple[np.ndarray, list[int], list[int]]: if not self.trim_blank: return ( design, list(range(design.shape[0])), list(range(design.shape[1])), ) row_mask = [ any(not _is_empty_design_cell(cell) for cell in design[i, :]) for i in range(design.shape[0]) ] col_mask = [ any(not _is_empty_design_cell(cell) for cell in design[:, j]) for j in range(design.shape[1]) ] kept_rows = [i for i, keep in enumerate(row_mask) if keep] kept_cols = [i for i, keep in enumerate(col_mask) if keep] if not kept_rows or not kept_cols: msg = "facet_manual design contains no panels" raise ValueError(msg) trimmed = design[np.ix_(kept_rows, kept_cols)] return trimmed, kept_rows, kept_cols def _set_panel_sizes(self, rows: list[int], cols: list[int]) -> None: if self.widths is not None: _check_size_length( self.widths, len(self._design_matrix[0]), "widths", ) self._panel_widths = [self.widths[i] for i in cols] else: self._panel_widths = None if self.heights is not None: _check_size_length( self.heights, len(self._design_matrix), "heights", ) self._panel_heights = [self.heights[i] for i in rows] else: self._panel_heights = None def _apply_axis_flags(self, layout: pd.DataFrame) -> None: if self.axes == "all": layout["AXIS_X"] = True layout["AXIS_Y"] = True elif self.axes == "x": layout["AXIS_X"] = True layout["AXIS_Y"] = False elif self.axes == "y": layout["AXIS_X"] = False layout["AXIS_Y"] = True if self.remove_labels == "all": layout["AXIS_X"] = False layout["AXIS_Y"] = False elif self.remove_labels == "x": layout["AXIS_X"] = False elif self.remove_labels == "y": layout["AXIS_Y"] = False
[docs] def map(self, data: pd.DataFrame, layout: pd.DataFrame) -> pd.DataFrame: if not len(data): data["PANEL"] = pd.Categorical( [], categories=layout["PANEL"].cat.categories, ordered=True, ) return data if not self.vars: # No facet variables: all data goes to panel 1 data["PANEL"] = pd.Categorical( [1] * len(data), categories=layout["PANEL"].cat.categories, ordered=True, ) return data facet_vals = eval_facet_vars(data, self.vars, self.environment) data, facet_vals = add_missing_facets( data, layout, self.vars, facet_vals ) # Assign each data point to a panel if len(facet_vals) and len(facet_vals.columns): keys = join_keys(facet_vals, layout, self.vars) data["PANEL"] = match(keys["x"], keys["y"], start=1) else: data["PANEL"] = 1 data["PANEL"] = pd.Categorical( data["PANEL"], categories=layout["PANEL"].cat.categories, ordered=True, ) data.reset_index(drop=True, inplace=True) return data
[docs] def make_strips(self, layout_info: layout_details, ax: Axes) -> Strips: if not self.vars: return Strips([]) s = strip(self.vars, layout_info, self, ax, "top") return Strips([s])
def _get_panels_gridspec(self) -> "p9GridSpec": from plotnine._mpl.gridspec import p9GridSpec ratios = {} if self._panel_widths is not None: ratios["width_ratios"] = self._panel_widths if self._panel_heights is not None: ratios["height_ratios"] = self._panel_heights return p9GridSpec( self.nrow, self.ncol, self.figure, nest_into=self.plot._gridspec[0], **ratios, ) def _make_axes(self) -> list["Axes"]: self._panels_gridspec = self._get_panels_gridspec() axs = [] for _, row in self.layout.layout.sort_values("PANEL").iterrows(): r0 = int(row["ROW"]) - 1 c0 = int(row["COL"]) - 1 r1 = r0 + int(row.get("ROWSPAN", 1)) c1 = c0 + int(row.get("COLSPAN", 1)) spec = self._panels_gridspec[r0:r1, c0:c1] axs.append(self.figure.add_subplot(spec)) return axs
[docs] def init_scales( self, layout: pd.DataFrame, x_scale: Optional[scale] = None, y_scale: Optional[scale] = None, ) -> object: import types from plotnine.scales.scales import Scales scales = types.SimpleNamespace() if x_scale is not None: n = layout["SCALE_X"].max() scales.x = Scales([x_scale.clone() for _i in range(n)]) if y_scale is not None: n = layout["SCALE_Y"].max() scales.y = Scales([y_scale.clone() for _i in range(n)]) plot = getattr(self, "plot", None) fps = getattr(plot, "_facetted_pos_scales", None) if fps is not None: fps.apply(scales) scale_facets = getattr(self, "_scale_facets", None) if scale_facets is None: scale_facets = getattr(plot, "_scale_facets", []) apply_scale_facets(scales, layout, scale_facets) return scales
[docs] def set_limits_breaks_and_labels(self, panel_params, ax): super().set_limits_breaks_and_labels(panel_params, ax) apply_axis_guides(self, panel_params, ax)
def _is_empty_design_cell(cell: object) -> bool: if pd.isna(cell): return True text = str(cell) return text == "#" or text.upper() == "NA" def _check_size_length( values: Sequence[float], expected: int, name: str, ) -> None: if len(values) != expected: msg = ( f"facet_manual {name} must have length {expected} to match " "the untrimmed design" ) raise ValueError(msg) def _check_rectangular_span( label: str, positions: np.ndarray, rowspan: int, colspan: int, ) -> None: if len(positions) != rowspan * colspan: msg = ( "facet_manual repeated design label " f"{label!r} must form a rectangle" ) raise ValueError(msg)