Source code for genno.compat.plotnine.plot

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, Hashable, Optional, Sequence
from warnings import warn

import plotnine as p9

from import Computer
from genno.core.key import KeyLike
from genno.core.quantity import Quantity

log = logging.getLogger(__name__)

[docs] class Plot(ABC): """Class for plotting using :doc:`plotnine <plotnine:index>`.""" #: File name base for saving the plot. basename = "" #: File extension; determines file format. suffix = ".pdf" #: Path for file output. If it is not set, :meth:`save` will populate it with a #: value constructed from :py:`config["output_dir"]`, :attr:`basename`, and #: :attr:`suffix`. The implementation of :meth:`generate` in a Plot sub-class may #: assign any other value, for instance one constructed at runtime from the #: :attr:`inputs`. path: Optional[Path] = None #: :class:`Keys <.Key>` referring to :class:`Quantities <.Quantity>` or other inputs #: accepted by :meth:`generate`. inputs: Sequence[Hashable] = [] #: Keyword arguments for :meth:``. save_args: Dict[str, Any] = dict(verbose=False) # TODO add static geoms automatically in generate() __static: Sequence = []
[docs] def save(self, config, *args, **kwargs) -> Optional[Path]: """Prepare data, call :meth:`.generate`, and save to file. This method is used as the callable in the task generated by :meth:`.add_tasks`. """ self.path = self.path or ( config["output_dir"] / f"{self.basename}{self.suffix}" ) missing = tuple(filter(lambda arg: isinstance(arg, str), args)) if len(missing): log.error( f"Missing input(s) {missing!r} to plot {self.basename!r}; no output" ) return None # Convert Quantity arguments to pd.DataFrame for use with plotnine _args = map( lambda arg: arg if not isinstance(arg, Quantity) else arg.to_series() .rename( or "value") .reset_index() .assign(unit=f"{arg.units:~}"), args, ) plot_or_plots = self.generate(*_args, **kwargs) if not plot_or_plots: f"{self.__class__.__name__}.generate() returned {plot_or_plots!r}; no " "output" ) return None"Save to {self.path}") try: # Single plot, **self.save_args) except AttributeError: # Iterator containing 0 or more plots p9.save_as_pdf_pages(plot_or_plots, self.path, **self.save_args) return self.path
[docs] @classmethod def make_task(cls, *inputs): """Return a task :class:`tuple` to add to a Computer. .. deprecated:: 1.18.0 Use :func:`add_tasks` instead. Parameters ---------- *inputs : `.Key` or str or hashable, optional If provided, overrides the :attr:`inputs` property of the class. Returns ------- tuple - The first, callable element of the task is :meth:`save`. - The second element is ``"config"``, to access the configuration of the Computer. - The third and following elements are the `inputs`. """ inputs_repr = ",".join(map(repr, inputs)) warn( f"Plot.make_task(…). Use: Computer.add(…, {cls.__name__}" + (", " if inputs_repr else "") + f"{inputs_repr})", DeprecationWarning, ) return tuple([cls().save, "config"] + (list(inputs) if inputs else cls.inputs))
[docs] @classmethod def add_tasks( cls, c: Computer, key: KeyLike, *inputs, strict: bool = False ) -> KeyLike: """Add a task to `c` to generate and save the Plot. Analogous to :meth:`.Operator.add_tasks`. """ _inputs = list(inputs if inputs else cls.inputs) if strict: _inputs = c.check_keys(*_inputs) return c.add_single(key, cls().save, "config", *_inputs)
[docs] @abstractmethod def generate(self, *args, **kwargs): """Generate and return the plot. A subclass of Plot **must** implement this method. Parameters ---------- args : sequence of pandas.DataFrame or other One argument is given corresponding to each of the :attr:`inputs`. Because :doc:`plotnine <plotnine:index>` operates on pandas data structures, :meth:`save` automatically converts any :class:`.Quantity` inputs to :class:`pandas.DataFrame` before they are passed to :meth:`generate`. """