Source code for genno.compat.plotnine.plot

import logging
from abc import ABC, abstractmethod
from typing import Hashable, Sequence

import plotnine as p9

from genno.core.quantity import Quantity

log = logging.getLogger(__name__)


[docs]class Plot(ABC): """Class for plotting using :mod:`plotnine`.""" #: Filename base for saving the plot. basename = "" #: File extension; determines file format. suffix = ".pdf" #: Keys for quantities needed by :meth:`generate`. inputs: Sequence[Hashable] = [] #: Keyword arguments for :meth:`plotnine.ggplot.save`. save_args = dict(verbose=False) # TODO add static geoms automatically in generate() __static: Sequence = []
[docs] def save(self, config, *args, **kwargs): """Prepare data, call :meth:`.generate`, and save to file. This method is used as the callable in the task generated by :meth:`.make_task`. """ path = 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 # 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(arg.name or "value") .reset_index() .assign(unit=f"{arg.units:~}"), args, ) plot_or_plots = self.generate(*args, **kwargs) if not plot_or_plots: log.info( f"{self.__class__.__name__}.generate() returned {plot_or_plots!r}; no " "output" ) return log.info(f"Save to {path}") try: # Single plot plot_or_plots.save(path, **self.save_args) except AttributeError: # Iterator containing 0 or more plots p9.save_as_pdf_pages(plot_or_plots, path, **self.save_args) return path
[docs] @classmethod def make_task(cls, *inputs): """Return a task :class:`tuple` to add to a Computer. Parameters ---------- inputs : sequence of :class:`.Key`, :class:`str`, or other 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`. """ return tuple([cls().save, "config"] + (list(inputs) if inputs else cls.inputs))
[docs] @abstractmethod def generate(self, *args, **kwargs): """Generate and return the plot. Must be implemented by subclasses. Parameters ---------- args : sequence of :class:`pandas.DataFrame` Because :mod:`plotnine` operates on pandas data structures, :meth:`save` automatically converts :obj:`Quantity` before being provided to :meth:`generate`. """