Use improved .report.plot features in .transport

Author: khaeru

message_ix_models/model/transport/build.py

@@ -15,6 +15,7 @@
 from message_ix_models import Context, ScenarioInfo
 from message_ix_models.model import bare, build
 from message_ix_models.model.structure import get_codelist
+from message_ix_models.model.workflow import STAGE
 from message_ix_models.util import (
     MappingAdapter,
     WildcardAdapter,
@@ -42,7 +43,9 @@
 def add_debug(c: Computer) -> None:
     """Add tasks for debugging the build."""
 
-    from . import plot
+    from message_ix_models.report import plot
+
+    from . import plot as plots
     from .key import gdp_cap, ms, pdt_nyt
 
     context: Context = c.graph["context"]
@@ -85,16 +88,19 @@ def add_debug(c: Computer) -> None:
     ):
         c.add(k_debug[i], "write_report_debug", key, output_dir.joinpath(f"{stem}.csv"))
 
-    c.add("transport build debug", "summarize", *k_debug.generated)
-    plot.prepare_computer(c, kind=plot.Kind.BUILD, target="transport build debug")
+    k_debug = Key("transport build debug")
+    c.add(k_debug["data"], "summarize", *k_debug.generated)
+    plot.prepare_computer(c, plots, k_debug["plot"], stage=STAGE.BUILD, single=True)
 
     # Also generate these debugging outputs when building the scenario
-    c.graph["add transport data"].append("transport build debug")
+    c.graph["add transport data"].extend(k_debug["data"], k_debug["plot"])
 
 
 def debug_multi(context: Context, *paths: Path) -> None:
     """Generate plots comparing data from multiple build debug directories."""
-    from . import plot
+    from message_ix_models.report import plot
+
+    from . import plot as plots
 
     if isinstance(paths[0], Scenario):
         # Workflow was called with --from="…", so paths from the previous step are not
@@ -105,10 +111,10 @@ def debug_multi(context: Context, *paths: Path) -> None:
             )
         )
 
-    c = Computer(config={"transport build debug dir": paths[0].parent})
+    c = Computer(config={"build debug dir": paths[0].parent})
     c.require_compat("message_ix_models.report.operator")
 
-    plot.prepare_computer(c, kind=plot.Kind.BUILD_MULTI, target="debug multi")
+    plot.prepare_computer(c, plots, "debug multi", stage=STAGE.BUILD, single=False)
 
     c.get("debug multi")
 

message_ix_models/model/transport/plot.py

@@ -1,22 +1,29 @@
 """Plots for MESSAGEix-Transport reporting."""
 
 from pathlib import Path
+from typing import cast
 
 import pandas as pd
 import plotnine as p9
 from iam_units import registry
 
+from message_ix_models.model.workflow import STAGE
+from message_ix_models.report.plot import COMMON, LabelFirst, Plot, PlotFromIAMC
 
 from .key import gdp_cap, pdt_nyt
 
+#: Common, static settings
+STATIC = [
+    COMMON["A4 landscape"],
+]
 
 
 class BaseEnergy0(Plot):
     """Transport final energy intensity of GDP."""
 
     basename = "base-fe-intensity-gdp"
     inputs = ["fe intensity:nl-ya:units"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="value", color="nl"),
         p9.geom_line(),
         p9.geom_point(),
@@ -35,7 +42,7 @@ class CapNewLDV(Plot):
 
     basename = "cap-new-t-ldv"
     inputs = ["historical_new_capacity:nl-t-yv:ldv", "CAP_NEW:nl-t-yv:ldv"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="yv", y="value", color="t"),
         p9.geom_vline(xintercept=2020, size=4, color="white"),
         p9.geom_line(),
@@ -96,11 +103,10 @@ class ComparePDT(Plot):
     - One line with points per scenario, coloured by scenario.
     """
 
-    kind = Kind.BUILD_MULTI
-
     basename = "compare-pdt"
-
-    static = Plot.static + [
+    stage = STAGE.BUILD
+    single = False
+    static = STATIC + [
         p9.aes(x="y", y="value", color="scenario"),
         p9.facet_wrap("t", ncol=5),
         p9.geom_line(),
@@ -118,7 +124,10 @@ class ComparePDT(Plot):
     factor = 1e6
 
     def generate(self, *paths: Path):
-        data = read_csvs(self.measure, *paths).eval("value = value / @self.factor")
+        data = cast(
+            pd.DataFrame,
+            read_csvs(self.measure, *paths).eval("value = value / @self.factor"),
+        )
 
         # Add factor to the unit expression
         if self.factor != 1.0:
@@ -140,7 +149,7 @@ class ComparePDTCap0(ComparePDT):
 
 
 #: Common layers for :class:`ComparePDTCap1` and :class:`DemandExoCap1`.
-PDT_CAP_GDP_STATIC = Plot.static + [
+PDT_CAP_GDP_STATIC = STATIC + [
     p9.aes(x="gdp", y="value", color="t"),
     p9.geom_line(),
     p9.geom_point(),
@@ -160,8 +169,9 @@ class ComparePDTCap1(Plot):
     - One line with points per |t| (=transport mode), coloured by mode.
     """
 
-    kind = Kind.BUILD_MULTI
     basename = "compare-pdt-capita-gdp"
+    stage = STAGE.BUILD
+    single = False
     static = PDT_CAP_GDP_STATIC + [
         p9.facet_wrap("scenario", ncol=5),
     ]
@@ -192,7 +202,7 @@ class InvCost0(Plot):
 
     basename = "inv-cost-transport"
     inputs = ["inv_cost:nl-t-yv:transport all"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="yv", y="inv_cost", color="t"),
         p9.geom_line(),
         p9.geom_point(),
@@ -231,7 +241,7 @@ class FixCost(Plot):
 
     basename = "fix-cost"
     inputs = ["fix_cost:nl-t-yv-ya:transport all"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="fix_cost", color="t", group="t * yv"),
         p9.geom_line(),
         p9.geom_point(),
@@ -250,7 +260,7 @@ class VarCost(Plot):
 
     basename = "var-cost"
     inputs = ["var_cost:nl-t-yv-ya:transport all"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="var_cost", color="t", group="t * yv"),
         p9.geom_line(),
         p9.geom_point(),
@@ -269,7 +279,7 @@ class LDV_IO(Plot):
 
     basename = "ldv-efficiency"
     inputs = ["input:nl-t-yv-ya:transport all"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="input", color="t"),
         # TODO remove typing exclusion once plotnine >0.12.4 is released
         p9.facet_wrap(
@@ -291,7 +301,7 @@ class OutShareLDV0(Plot):
 
     basename = "out-share-t-ldv"
     inputs = ["out:nl-t-ya:ldv+units"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="value", fill="t"),
         p9.geom_bar(stat="identity", width=4),
         # # Select a palette with up to 12 colors
@@ -314,7 +324,7 @@ class OutShareLDV1(Plot):
 
     basename = "out-share-t-cg-ldv"
     inputs = ["out:nl-t-ya-c", "cg"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="value", fill="t"),
         p9.facet_wrap(["c"], ncol=5),
         p9.geom_bar(stat="identity", width=4),
@@ -352,7 +362,7 @@ class Demand0(Plot):
 
     basename = "demand"
     inputs = ["demand:n-c-y", "c::transport", "cg"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="y", y="demand", fill="c_group"),
         p9.geom_bar(stat="identity", width=4),
         p9.labs(x="Period", y="", fill="Transport mode"),
@@ -394,7 +404,7 @@ class DemandCap(Plot):
 
     basename = "demand-capita"
     inputs = ["demand:n-c-y:capita", "c::transport", "cg"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="y", y="value", fill="c"),
         p9.geom_bar(stat="identity", width=4),
         p9.labs(x="Period", y="", fill="Transport mode group"),
@@ -412,18 +422,20 @@ def _reduce_units(df: pd.DataFrame, target_units) -> tuple[pd.DataFrame, str]:
     assert 1 == len(df_units)
     tmp = registry.Quantity(1.0, df_units[0]).to(target_units)
     return (
-        df.eval("value = value * @tmp.magnitude").assign(unit=f"{tmp.units:~}"),
+        cast(pd.DataFrame, df.eval("value = value * @tmp.magnitude")).assign(
+            unit=f"{tmp.units:~}"
+        ),
         f"{tmp.units:~}",
     )
 
 
 class DemandExo(Plot):
     """Passenger transport activity."""
 
-    kind = Kind.BUILD
     basename = "demand-exo"
+    stage = STAGE.BUILD
     inputs = [pdt_nyt]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="y", y="value", fill="t"),
         p9.geom_bar(stat="identity", width=4),
         p9.labs(x="Period", y="", fill="Mode (tech group)"),
@@ -442,10 +454,10 @@ def generate(self, data):
 class DemandExoCap0(Plot):
     """Passenger transport activity per person."""
 
-    kind = Kind.BUILD
     basename = "demand-exo-capita"
+    stage = STAGE.BUILD
     inputs = [pdt_nyt + "capita+post"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="y", y="value", fill="t"),
         p9.geom_bar(stat="identity", width=4),
         p9.labs(x="Period", y="", fill="Transport mode"),
@@ -467,8 +479,8 @@ class DemandExoCap1(DemandExoCap0):
     Unlike :class:`DemandExoCap0`, this uses GDP per capita as the abscissa/x-aesthetic.
     """
 
-    kind = Kind.BUILD
     basename = "demand-exo-capita-gdp"
+    stage = STAGE.BUILD
     inputs = [pdt_nyt + "capita+post", gdp_cap]
     static = PDT_CAP_GDP_STATIC
 
@@ -494,7 +506,7 @@ class EnergyCmdty0(Plot):
 
     basename = "energy-c"
     inputs = ["y0", "in:nl-ya-c:transport all"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="value", fill="c"),
         p9.geom_bar(stat="identity", width=5, color="black"),
         p9.labs(x="Period", y="Energy", fill="Commodity"),
@@ -538,20 +550,22 @@ def generate(self, y0: int, data):
             yield ggplot
 
 
-class MultiStock(Plot):
+class MultiStock(PlotFromIAMC):
     """LDV technology stock."""
 
-    kind = Kind.REPORT_MULTI
     basename = "multi-stock"
+    single = False
+    inputs = ["all:n-s-UNIT-v-y"]
+    iamc_variable_pattern = r"Transport\|Stock\|Road\|Passenger\|LDV\|(.*)"
 
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="y", y="value", color="v"),
         p9.facet_wrap("s", ncol=3, nrow=5),
         p9.geom_point(),
         p9.geom_line(),
         p9.scale_color_brewer(type="qualitative", palette="Paired"),
         p9.labs(x="Period", y="", color="Technology"),
-        p9.theme(figure_size=(11.7, 16.6)),
+        COMMON["A3 portrait"],
     ]
 
     def generate(self, data):
@@ -567,12 +581,11 @@ class Scale1Diff(Plot):
     """scale-1 factor in y=2020; changes between 2 scenarios."""
 
     basename = "scale-1-diff"
-
-    kind = Kind.BUILD
+    stage = STAGE.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"
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x=_v, y=_y, yend=_y, group=_s, color=_s, shape=_s),
         p9.facet_wrap("nl", scales="free_x"),
         p9.geom_vline(p9.aes(xintercept=_v), pd.DataFrame([[1.0]], columns=[_v])),
@@ -611,7 +624,7 @@ class Stock0(Plot):
     basename = "stock-ldv"
     # Partial sum over driver_type dimension
     inputs = ["CAP:nl-t-ya:ldv+units"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="CAP", color="t"),
         p9.geom_line(),
         p9.geom_point(),
@@ -634,7 +647,7 @@ class Stock1(Plot):
 
     basename = "stock-non-ldv"
     inputs = ["CAP:nl-t-ya:non-ldv+units"]
-    static = Plot.static + [
+    static = STATIC + [
         p9.aes(x="ya", y="CAP", color="t"),
         p9.geom_line(),
         p9.geom_point(),

message_ix_models/model/transport/report.py

@@ -3,7 +3,6 @@
 import logging
 import re
 from copy import deepcopy
-from types import SimpleNamespace
 from typing import TYPE_CHECKING
 
 import pandas as pd
@@ -12,6 +11,7 @@
 from message_ix import Reporter
 
 from message_ix_models import Context, ScenarioInfo
+from message_ix_models.model.workflow import STAGE
 from message_ix_models.report.util import add_replacements
 
 from . import Config
@@ -418,9 +418,10 @@ def misc(c: "Computer") -> None:
 def multi(context: Context, targets: list[str], *, use_platform: bool = False) -> None:
     """Report outputs from multiple scenarios."""
 
+    from message_ix_models.report import plot
     from message_ix_models.tools.iamc import to_quantity
 
-    from . import plot
+    from . import plot as plots
 
     k = Keys(
         in_="in::pd",
@@ -430,16 +431,14 @@ def multi(context: Context, targets: list[str], *, use_platform: bool = False) -
         plot_data="plot data:n-s-y",
     )
 
-    c = Computer()
+    # Computer, with configuration expected by Plot.add_tasks
+    c = Computer(config=dict(output_dir=context.get_local_path("report")))
     # Order is important: use genno.compat.pyam.operator.quantity_from_iamc over local
     c.require_compat("message_ix_models.report.operator")
     c.require_compat("pyam")
 
     c.add("context", context)
 
-    # Expected by .transport.plot.Plot.add_tasks
-    c.graph["config"]["transport build debug dir"] = context.get_local_path("report")
-
     # Retrieve all data for each of the `targets` as pd.DataFrame
     kw0 = dict(filename="transport.csv", use={"file"})
     for i, info in enumerate(map(ScenarioInfo.from_url, targets)):
@@ -459,16 +458,8 @@ def multi(context: Context, targets: list[str], *, use_platform: bool = False) -
     # Rename dimensions
     c.add(k.all1, "rename_dims", k.all0, name_dict={"SCENARIO": "s", "VARIABLE": "v"})
 
-    # Select a subset of data and reduce dimensionality
-    expr = r"Transport\|Stock\|Road\|Passenger\|LDV\|(.*)"
-    c.add(k.plot_data[0], "quantity_from_iamc", k.all1, variable=expr)
-
-    # Key used by .plot.Plot for formatting titles
-    c.add("scenario", SimpleNamespace(url="Multiple scenarios"))
-
     # Collect all plots
-    c.add("multi", "summarize")
-    plot.prepare_computer(c, kind=plot.Kind.REPORT_MULTI, target="multi")
+    plot.prepare_computer(c, plots, "multi", stage=STAGE.REPORT, single=False)
 
     return c.get("multi")