Source code for plotnine.qplot

from __future__ import annotations

import typing
from contextlib import suppress
from warnings import warn

import numpy as np
import pandas as pd

from .exceptions import PlotnineError, PlotnineWarning
from .facets import facet_grid, facet_null, facet_wrap
from .facets.facet_grid import parse_grid_facets
from .facets.facet_wrap import parse_wrap_facets
from .ggplot import ggplot
from .labels import labs
from .mapping.aes import ALL_AESTHETICS, SCALED_AESTHETICS, aes
from .scales import lims, scale_x_log10, scale_y_log10
from .themes import theme
from .utils import Registry, array_kind

if typing.TYPE_CHECKING:
    from typing import Any, Iterable, Literal

    from plotnine.typing import DataLike, Ggplot, TupleFloat2

__all__ = ("qplot",)


[docs]def qplot( x: str | Iterable[Any] | range | None = None, y: str | Iterable[Any] | range | None = None, data: DataLike | None = None, facets: str = "", margins: bool | list[str] = False, geom: str | list[str] | tuple[str] = "auto", xlim: TupleFloat2 | None = None, ylim: TupleFloat2 | None = None, log: Literal["x", "y", "xy"] | None = None, main: str | None = None, xlab: str | None = None, ylab: str | None = None, asp: float | None = None, **kwargs: Any, ) -> Ggplot: """ Quick plot Parameters ---------- x : str | array_like x aesthetic y : str | array_like y aesthetic data : dataframe Data frame to use (optional). If not specified, will create one, extracting arrays from the current environment. geom : str | list *geom(s)* to do the drawing. If ``auto``, defaults to 'point' if ``x`` and ``y`` are specified or 'histogram' if only ``x`` is specified. facets : str Facets margins : bool | list[str] variable names to compute margins for. True will compute all possible margins. Depends on the facetting. xlim : tuple x-axis limits ylim : tuple y-axis limits log : str in ``{'x', 'y', 'xy'}`` Which (if any) variables to log transform. main : str Plot title xlab : str x-axis label ylab : str y-axis label asp : str | float The y/x aspect ratio. **kwargs : dict Arguments passed on to the geom. Returns ------- p : ggplot ggplot object """ from patsy.eval import EvalEnvironment # Extract all recognizable aesthetic mappings from the parameters # String values e.g "I('red')", "I(4)" are not treated as mappings environment = EvalEnvironment.capture(1) aesthetics = {} if x is None else {"x": x} if y is not None: aesthetics["y"] = y def is_mapping(value: Any) -> bool: """ Return True if value is not enclosed in I() function """ with suppress(AttributeError): return not (value.startswith("I(") and value.endswith(")")) return True def I(value: Any) -> Any: return value I_env = EvalEnvironment([{"I": I}]) for ae in kwargs.keys() & ALL_AESTHETICS: value = kwargs[ae] if is_mapping(value): aesthetics[ae] = value else: kwargs[ae] = I_env.eval(value) # List of geoms if isinstance(geom, str): geom = [geom] elif isinstance(geom, tuple): geom = list(geom) if data is None: data = pd.DataFrame() # Work out plot data, and modify aesthetics, if necessary def replace_auto(lst: list[str], str2: str) -> list[str]: """ Replace all occurences of 'auto' in with str2 """ for i, value in enumerate(lst): if value == "auto": lst[i] = str2 return lst if "auto" in geom: if "sample" in aesthetics: replace_auto(geom, "qq") elif y is None: # If x is discrete we choose geom_bar & # geom_histogram otherwise. But we need to # evaluate the mapping to find out the dtype env = environment.with_outer_namespace({"factor": pd.Categorical}) if isinstance(aesthetics["x"], str): try: x = env.eval(aesthetics["x"], inner_namespace=data) except Exception: msg = "Could not evaluate aesthetic 'x={}'" raise PlotnineError(msg.format(aesthetics["x"])) elif not hasattr(aesthetics["x"], "dtype"): x = np.asarray(aesthetics["x"]) if array_kind.discrete(x): replace_auto(geom, "bar") else: replace_auto(geom, "histogram") else: if x is None: if isinstance(aesthetics["y"], typing.Sized): aesthetics["x"] = range(len(aesthetics["y"])) xlab = "range(len(y))" ylab = "y" else: # We could solve the issue in layer.compute_asthetics # but it is not worth the extra complexity raise PlotnineError("Cannot infer how long x should be.") replace_auto(geom, "point") p: Ggplot = ggplot(data, aes(**aesthetics), environment=environment) def get_facet_type(facets: str) -> Literal["grid", "wrap", "null"]: with suppress(PlotnineError): parse_grid_facets(facets) return "grid" with suppress(PlotnineError): parse_wrap_facets(facets) return "wrap" warn( "Could not determine the type of faceting, " "therefore no faceting.", PlotnineWarning, ) return "null" if facets: facet_type = get_facet_type(facets) if facet_type == "grid": p += facet_grid(facets, margins=margins) elif facet_type == "wrap": p += facet_wrap(facets) else: p += facet_null() # Add geoms for g in geom: geom_name = f"geom_{g}" geom_klass = Registry[geom_name] stat_name = "stat_{}".format(geom_klass.DEFAULT_PARAMS["stat"]) stat_klass = Registry[stat_name] # find params recognized = kwargs.keys() & ( geom_klass.DEFAULT_PARAMS.keys() | geom_klass.aesthetics() | stat_klass.DEFAULT_PARAMS.keys() | stat_klass.aesthetics() ) recognized = recognized - aesthetics.keys() params = {ae: kwargs[ae] for ae in recognized} p += geom_klass(**params) # pd.Series objects have name attributes. In a dataframe, the # series have the name of the column. labels = {} for ae in SCALED_AESTHETICS & kwargs.keys(): with suppress(AttributeError): labels[ae] = kwargs[ae].name with suppress(AttributeError): labels["x"] = xlab if xlab is not None else x.name # type: ignore with suppress(AttributeError): labels["y"] = ylab if ylab is not None else y.name # type: ignore if main is not None: labels["title"] = main if log: if "x" in log: p += scale_x_log10() if "y" in log: p += scale_y_log10() if labels: p += labs(**labels) if asp: p += theme(aspect_ratio=asp) if xlim: p += lims(x=xlim) if ylim: p += lims(y=ylim) return p