Source code for plotnine.guides.guides

from __future__ import annotations

import typing
from copy import deepcopy
from warnings import warn

import numpy as np
import pandas as pd

from ..exceptions import PlotnineError, PlotnineWarning
from ..utils import Registry, is_string
from .guide import guide as guide_class

if typing.TYPE_CHECKING:
    from typing import Literal

    from plotnine.typing import TupleFloat2

# Terminology
# -----------
# - A guide is either a legend or colorbar.
#
# - A guide definition (gdef) is an instantiated guide as it
#   is used in the process of creating the legend
#
# - The guides class holds all guides that will appear in the
#   plot
#
# - A guide box is a fully drawn out guide.
#   It is of subclass matplotlib.offsetbox.Offsetbox


[docs]class guides(dict): """ Guides for each scale Used to assign a particular guide to an aesthetic(s). Parameters ---------- kwargs : dict aesthetic - guide pairings. e.g ``color=guide_colorbar()`` """ # Determined from the theme when the guides are # getting built position: Literal["left", "right", "top", "bottom"] | TupleFloat2 box_direction: Literal["horizontal", "vertical", "auto"] box_align: Literal["left", "right", "top", "bottom", "center", "auto"] box_margin: int spacing: float def __init__(self, **kwargs): aes_names = { "alpha", "color", "fill", "linetype", "shape", "size", "stroke", } if "colour" in kwargs: kwargs["color"] = kwargs.pop("colour") dict.__init__( self, ((ae, kwargs[ae]) for ae in kwargs if ae in aes_names) ) def __radd__(self, gg): """ Add guides to the plot Parameters ---------- gg : ggplot ggplot object being created Returns ------- out : gglot ggplot object with guides. If *inplace* is **False** this is a copy of the original ggplot object. """ new_guides = {} for k in self: new_guides[k] = deepcopy(self[k]) gg.guides.update(new_guides) return gg def build(self, plot): """ Build the guides Parameters ---------- plot : ggplot ggplot object being drawn Returns ------- box : matplotlib.offsetbox.Offsetbox | None A box that contains all the guides for the plot. If there are no guides, **None** is returned. """ _property = plot.theme.themeables.property self.box_direction = _property("legend_box") self.position = _property("legend_position") self.box_align = _property("legend_box_just") self.box_margin = _property("legend_box_margin") self.spacing = _property("legend_spacing") if self.position == "none": return # No Legend # Direction if self.box_direction == "auto": if self.position in ("right", "left"): self.box_direction = "vertical" else: self.box_direction = "horizontal" # Justification of legend boxes if self.box_align == "auto": if self.position in ("right", "left"): self.box_align = "left" else: self.box_align = "right" gdefs = self.train(plot) if not gdefs: return gdefs = self.merge(gdefs) gdefs = self.create_geoms(gdefs, plot) if not gdefs: return gboxes = self.draw(gdefs, plot.theme) bigbox = self.assemble(gboxes, gdefs, plot.theme) return bigbox def train(self, plot): """ Compute all the required guides Parameters ---------- plot : ggplot ggplot object Returns ------- gdefs : list Guides for the plots """ gdefs = [] for scale in plot.scales: for output in scale.aesthetics: # The guide for aesthetic 'xxx' is stored # in plot.guides['xxx']. The priority for # the guides depends on how they are created # 1. ... + guides(xxx=guide_blah()) # 2. ... + scale_xxx(guide=guide_blah()) # 3. default(either guide_legend or guide_colorbar # depending on the scale type) # output = scale.aesthetics[0] guide = self.get(output, scale.guide) if guide is None or guide is False: continue # check the validity of guide. # if guide is character, then find the guide object guide = self.validate(guide) # check the consistency of the guide and scale. if ( "any" not in guide.available_aes and scale.aesthetics[0] not in guide.available_aes ): raise PlotnineError( f"{guide.__class__.__name__} cannot be used for " f"{scale.aesthetics}" ) # title if not hasattr(guide, "title"): if scale.name: guide.title = scale.name else: guide.title = getattr(plot.labels, output) if guide.title is None: warn( f"Cannot generate legend for the {output!r} " "aesthetic. Make sure you have mapped a " "variable to it", PlotnineWarning, ) # each guide object trains scale within the object, # so Guides (i.e., the container of guides) # need not to know about them guide = guide.train(scale, output) if guide is not None: gdefs.append(guide) return gdefs def validate(self, guide): """ Validate guide object """ if is_string(guide): guide = Registry[f"guide_{guide}"]() if not isinstance(guide, guide_class): raise PlotnineError(f"Unknown guide: {guide}") return guide def merge(self, gdefs): """ Merge overlapped guides For example:: from plotnine import * gg = ggplot(mtcars, aes(y='wt', x='mpg', colour='factor(cyl)')) gg = gg + stat_smooth(aes(fill='factor(cyl)'), method='lm') gg = gg + geom_point() gg This would create two guides with the same hash """ # group guide definitions by hash, and # reduce each group to a single guide # using the guide.merge method definitions = pd.DataFrame( {"gdef": gdefs, "hash": [g.hash for g in gdefs]} ) grouped = definitions.groupby("hash", sort=False) gdefs = [] for name, group in grouped: # merge gdef = group["gdef"].iloc[0] for g in group["gdef"].iloc[1:]: gdef = gdef.merge(g) gdefs.append(gdef) return gdefs def create_geoms(self, gdefs, plot): """ Add geoms to the guide definitions """ new_gdefs = [] for gdef in gdefs: gdef = gdef.create_geoms(plot) if gdef: new_gdefs.append(gdef) return new_gdefs def draw(self, gdefs, theme): """ Draw out each guide definition Parameters ---------- gdefs : list of guide_legend|guide_colorbar guide definitions theme : theme Plot theme Returns ------- out : list of matplotlib.offsetbox.Offsetbox A drawing of each legend """ for g in gdefs: g._set_defaults(theme) return [g.draw() for g in gdefs] def assemble(self, gboxes, gdefs, theme): """ Put together all the guide boxes Parameters ---------- gboxes : list List of :class:`~matplotlib.offsetbox.Offsetbox`, where each item is a legend for a single aesthetic. gdefs : list of guide_legend|guide_colorbar guide definitions theme : theme Plot theme Returns ------- box : OffsetBox A box than can be placed onto a plot """ from matplotlib.offsetbox import HPacker, VPacker # place the guides according to the guide.order # 0 do not sort # 1-99 sort for gdef in gdefs: if gdef.order == 0: gdef.order = 100 elif not 0 <= gdef.order <= 99: raise PlotnineError( "'order' for a guide should be " "between 0 and 99" ) orders = [gdef.order for gdef in gdefs] idx = np.argsort(orders) gboxes = [gboxes[i] for i in idx] # direction when more than legend if self.box_direction == "vertical": packer = VPacker elif self.box_direction == "horizontal": packer = HPacker else: raise PlotnineError( "'legend_box' should be either " "'vertical' or 'horizontal'" ) box = packer( children=gboxes, align=self.box_align, pad=self.box_margin, sep=self.spacing, ) return box