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
import genno
from genno.compat.pandas import disable_copy_on_write
from genno.core.computer import Computer
from genno.core.key import KeyLike
log = logging.getLogger(__name__)
[docs]
class Plot(ABC):
"""Class for plotting using :mod:`plotnine`."""
#: 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 :any:`plotnine.ggplot.save`.
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`.
.. versionadded:: 1.24.1
This method uses :func:`.disable_copy_on_write` to work around
`has2k1/mizani#38 <https://github.com/has2k1/mizani/issues/38>`_. This may
cause issues if other computations (for instance, of the inputs to the Plot)
rely on Pandas' copy-on-write behaviour being enabled.
"""
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, genno.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 None
log.info(f"Save to {self.path}")
with disable_copy_on_write(f"{__name__}.Plot.save()"):
try:
# Single plot
plot_or_plots.save(self.path, **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 :mod:`plotnine` operates on pandas data structures, :meth:`save`
automatically converts any :class:`.Quantity` inputs to
:class:`pandas.DataFrame` before they are passed to :meth:`generate`.
"""