Source code for plotnine.geoms.geom_rect

from __future__ import annotations

import typing

import numpy as np
import pandas as pd

from ..doctools import document
from ..utils import SIZE_FACTOR, to_rgba
from .geom import geom
from .geom_polygon import geom_polygon

if typing.TYPE_CHECKING:
    from typing import Any

    from plotnine.iapi import panel_view
    from plotnine.typing import Axes, Coord


[docs]@document class geom_rect(geom): """ Rectangles {usage} Parameters ---------- {common_parameters} """ DEFAULT_AES = { "color": None, "fill": "#595959", "linetype": "solid", "size": 0.5, "alpha": 1, } REQUIRED_AES = {"xmax", "xmin", "ymax", "ymin"} DEFAULT_PARAMS = { "stat": "identity", "position": "identity", "na_rm": False, } draw_legend = staticmethod(geom_polygon.draw_legend) def draw_panel( self, data: pd.DataFrame, panel_params: panel_view, coord: Coord, ax: Axes, **params: Any, ): """ Plot all groups """ if not coord.is_linear: data = _rectangles_to_polygons(data) for _, gdata in data.groupby("group"): gdata.reset_index(inplace=True, drop=True) geom_polygon.draw_group( gdata, panel_params, coord, ax, **params ) else: self.draw_group(data, panel_params, coord, ax, **params) @staticmethod def draw_group( data: pd.DataFrame, panel_params: panel_view, coord: Coord, ax: Axes, **params: Any, ): from matplotlib.collections import PolyCollection data = coord.transform(data, panel_params, munch=True) data["size"] *= SIZE_FACTOR limits = zip(data["xmin"], data["xmax"], data["ymin"], data["ymax"]) verts = [[(l, b), (l, t), (r, t), (r, b)] for (l, r, b, t) in limits] fill = to_rgba(data["fill"], data["alpha"]) color = data["color"] # prevent unnecessary borders if all(color.isna()): color = "none" col = PolyCollection( verts, facecolors=fill, edgecolors=color, linestyles=data["linetype"], linewidths=data["size"], zorder=params["zorder"], rasterized=params["raster"], ) ax.add_collection(col)
def _rectangles_to_polygons(df: pd.DataFrame) -> pd.DataFrame: """ Convert rect data to polygons Paramters --------- df : dataframe Dataframe with *xmin*, *xmax*, *ymin* and *ymax* columns, plus others for aesthetics ... Returns ------- data : dataframe Dataframe with *x* and *y* columns, plus others for aesthetics ... """ n = len(df) # Helper indexing arrays xmin_idx = np.tile([True, True, False, False], n) xmax_idx = ~xmin_idx ymin_idx = np.tile([True, False, False, True], n) ymax_idx = ~ymin_idx # There are 2 x and 2 y values for each of xmin, xmax, ymin & ymax # The positions are as layed out in the indexing arrays # x and y values x = np.empty(n * 4) y = np.empty(n * 4) x[xmin_idx] = df["xmin"].repeat(2) x[xmax_idx] = df["xmax"].repeat(2) y[ymin_idx] = df["ymin"].repeat(2) y[ymax_idx] = df["ymax"].repeat(2) # Aesthetic columns and others other_cols = df.columns.difference( ["x", "y", "xmin", "xmax", "ymin", "ymax"] ) d = {str(col): np.repeat(df[col].to_numpy(), 4) for col in other_cols} data = pd.DataFrame({"x": x, "y": y, **d}) return data