Add .transport.plot.Kind enumeration

- Add .plot.prepare_computer(…, kind=…, target=…) keyword arguments.
- Filter plots to add using `kind`.

Author: khaeru

message_ix_models/model/transport/key.py

@@ -114,6 +114,7 @@
 #: Keys for :mod:`.transport.report`.
 report = SimpleNamespace(
     all="transport all",
+    plot="transport plot",
     sdmx=Key("transport::sdmx"),
 )
 

message_ix_models/model/transport/plot.py

@@ -2,6 +2,7 @@
 
 import logging
 from datetime import datetime
+from enum import Enum, auto
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -11,6 +12,8 @@
 from genno import Computer
 from iam_units import registry
 
+from message_ix_models.util.genno import append
+
 from .key import gdp_cap, pdt_nyt
 from .key import report as k_report
 
@@ -20,7 +23,6 @@
     # NB The following is itself within an if TYPE_CHECKING block; mypy can't find it
     from plotnine.typing import PlotAddable  # type: ignore [attr-defined]
 
-    from .config import Config
 
 log = logging.getLogger(__name__)
 
@@ -30,6 +32,27 @@
 logging.getLogger("matplotlib.font_manager").setLevel(logging.INFO + 1)
 
 
+class Kind(Enum):
+    """Kinds of transport :class:`Plots <.Plot>`."""
+
+    #: Plots for :func:`.transport.build.get_computer`, that is, during the process of
+    #: building a single scenario. These are used within a :class:`~.genno.Computer`,
+    #: without a built or solved Scenario.
+    BUILD = auto()
+
+    #: Plots for pre-solve or build-process analysis of multiple scenarios. These are
+    #: added by :func:`.transport.build.debug_multi`.
+    BUILD_MULTI = auto()
+
+    #: Plots for post-solve reporting of single MESSAGEix-Transport scenarios. These are
+    #: added to a :class:`.Reporter` via :func:`.transport.report.callback`.
+    REPORT = auto()
+
+    #: Plots of post-solve reporting of multiple MESSAGEix-Transport scenarios, or
+    #: of data reported from these. These are added by :class:`.transport.report.multi`.
+    REPORT_MULTI = auto()
+
+
 class LabelFirst:
     """Labeller that labels the first item using a format string.
 
@@ -54,6 +77,9 @@ class Plot(genno.compat.plotnine.Plot):
     This class extends :class:`genno.compat.plotnine.Plot` with extra features.
     """
 
+    #: Kind of plot.
+    kind: Kind = Kind.REPORT
+
     #: 'Static' geoms: list of plotnine objects that are not dynamic
     static: list["PlotAddable"] = [
         p9.theme(figure_size=(11.7, 8.3)),
@@ -66,9 +92,6 @@ class Plot(genno.compat.plotnine.Plot):
     #: Units expression for plot title.
     unit: str | None = None
 
-    #: :obj:`False` for plots not intended to be run on a solved scenario.
-    runs_on_solved_scenario: bool = True
-
     def ggtitle(self, extra: str | None = None):
         """Return :class:`plotnine.ggtitle` including the current date & time."""
         title_parts = [
@@ -119,7 +142,7 @@ def add_tasks(
         # Output path for this parameter
         k_path = f"plot {cls.basename} path"
         filename = f"{cls.basename}{cls.suffix}"
-        if cls.runs_on_solved_scenario:
+        if cls.kind is Kind.REPORT:
             # Make a path including the Scenario URL
             c.add(k_path, "make_output_path", "config", name=filename)
         else:
@@ -212,7 +235,7 @@ class ComparePDT(Plot):
 
     This plot is used in :func:`.transport.build.debug_multi`, not in ordinary
     reporting. Rather than receiving data from computed quantities already in the graph,
-    it reads them from files named :file:`pdt.csv` (per :attr:`.kind`) in the
+    it reads them from files named :file:`pdt.csv` (per :attr:`.measure`) in the
     directories generated by the workflow steps like "SSP1 debug build" (
     :func:`.transport.build.debug`).
 
@@ -221,7 +244,8 @@ class ComparePDT(Plot):
     - One line with points per scenario, coloured by scenario.
     """
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD_MULTI
+
     basename = "compare-pdt"
 
     static = Plot.static + [
@@ -234,15 +258,15 @@ class ComparePDT(Plot):
     ]
 
     #: Base name for source data files, for instance :file:`pdt.csv`.
-    kind = "pdt"
+    measure = "pdt"
 
     #: Units of input files
     unit = "km/a"
     #: Unit adjustment factor.
     factor = 1e6
 
     def generate(self, *paths: Path):
-        data = read_csvs(self.kind, *paths).eval("value = value / @self.factor")
+        data = read_csvs(self.measure, *paths).eval("value = value / @self.factor")
 
         # Add factor to the unit expression
         if self.factor != 1.0:
@@ -259,7 +283,7 @@ class ComparePDTCap0(ComparePDT):
     """
 
     basename = "compare-pdt-cap"
-    kind = "pdt-cap"
+    measure = "pdt-cap"
     factor = 1e3
 
 
@@ -284,7 +308,7 @@ class ComparePDTCap1(Plot):
     - One line with points per |t| (=transport mode), coloured by mode.
     """
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD_MULTI
     basename = "compare-pdt-capita-gdp"
     static = PDT_CAP_GDP_STATIC + [
         p9.facet_wrap("scenario", ncol=5),
@@ -526,7 +550,7 @@ class DemandCap(Plot):
 
     def generate(self, data, commodities, cg):
         # Convert and select data
-        data = data.query(f"c in {repr(list(map(str, commodities)))}").pipe(c_group, cg)
+        data = data.query(f"c in {list(map(str, commodities))!r}").pipe(c_group, cg)
         for _, ggplot in self.groupby_plot(data, "n"):
             yield ggplot
 
@@ -544,7 +568,7 @@ def _reduce_units(df: pd.DataFrame, target_units) -> tuple[pd.DataFrame, str]:
 class DemandExo(Plot):
     """Passenger transport activity."""
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD
     basename = "demand-exo"
     inputs = [pdt_nyt]
     static = Plot.static + [
@@ -566,7 +590,7 @@ def generate(self, data):
 class DemandExoCap0(Plot):
     """Passenger transport activity per person."""
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD
     basename = "demand-exo-capita"
     inputs = [pdt_nyt + "capita+post"]
     static = Plot.static + [
@@ -591,7 +615,7 @@ class DemandExoCap1(DemandExoCap0):
     Unlike :class:`DemandExoCap0`, this uses GDP per capita as the abscissa/x-aesthetic.
     """
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD
     basename = "demand-exo-capita-gdp"
     inputs = [pdt_nyt + "capita+post", gdp_cap]
     static = PDT_CAP_GDP_STATIC
@@ -663,10 +687,10 @@ def generate(self, y0: int, data):
 
 
 class MultiStock(Plot):
-    """Stock, LDV technologies, SSP × policy."""
+    """LDV technology stock."""
 
+    kind = Kind.REPORT_MULTI
     basename = "multi-stock"
-    runs_on_solved_scenario = False
 
     static = Plot.static + [
         p9.aes(x="y", y="value", color="v"),
@@ -692,7 +716,7 @@ class Scale1Diff(Plot):
 
     basename = "scale-1-diff"
 
-    runs_on_solved_scenario = False
+    kind = Kind.BUILD
     inputs = ["scale-1:nl-t-c-l-h:a", "scale-1:nl-t-c-l-h:b"]
 
     _s, _v, _y = "scenario", "value", "t + ' ' + c"
@@ -776,40 +800,36 @@ def generate(self, data):
             yield ggplot + p9.expand_limits(y=[0, y_max])
 
 
-def prepare_computer(c: Computer):
+def prepare_computer(
+    c: Computer, *, kind: Kind = Kind.REPORT, target: "KeyLike" = k_report.all
+) -> None:
     """Add :data:`.PLOTS` to `c`.
 
     Adds:
 
     - 1 key like **"plot inv-cost"** corresponding to the :attr:`~.Plot.basename` of
-      each :class:`.Plot` subclass defined in this module.
+      each :class:`.Plot` subclass defined in this module with a matching `kind`.
     - The key **"transport plots"** that triggers writing all the plots to file.
     """
+    # Force matplotlib to use a non-interactive backend for plotting
     import matplotlib
 
-    # Force matplotlib to use a non-interactive backend for plotting
     matplotlib.use("pdf")
 
-    keys = []
-
-    config: "Config" = c.graph["config"]["transport"]
-
     # Iterate over the Plot subclasses defined in the current module
+    keys = []
     for plot in filter(
-        lambda cls: isinstance(cls, type) and issubclass(cls, Plot) and cls is not Plot,
+        lambda cls: isinstance(cls, type)
+        and issubclass(cls, Plot)
+        and cls is not Plot
+        and cls.kind is kind,
         globals().values(),
     ):
-        if (not plot.runs_on_solved_scenario and config.with_solution) or (
-            False  # Use True here or uncomment below to skip some or all plots
-            # "stock" not in plot.basename
-        ):
-            log.info(f"Skip {plot}")
-            continue
         keys.append(f"plot {plot.basename}")
         c.add(keys[-1], plot)
 
-    key = "transport plots"
-    log.info(f"Add {repr(key)} collecting {len(keys)} plots")
-    c.add(key, keys)
+    log.info(f"Add {k_report.plot!r} collecting {len(keys)} plots")
+    c.add(k_report.plot, "summarize", keys)
 
-    c.graph[k_report.all].append(key)
+    # Extend the task at `target` with k_report.plot
+    append(c, target, k_report.plot)