Add .tools.sankey and tutorial

RELEASE_NOTES.rst

@@ -12,6 +12,8 @@ All changes
 
 - :mod:`message_ix` is tested and compatible with `Python 3.13 <https://www.python.org/downloads/release/python-3130/>`__ (:pull:`881`).
 - Support for Python 3.8 is dropped (:pull:`881`), as it has reached end-of-life.
+- Add :meth:`.Reporter.add_sankey` and :mod:`.tools.sankey` to create Sankey diagrams from solved scenarios (:pull:`770`).
+  The :file:`westeros_sankey.ipynb` :ref:`tutorial <tutorial-westeros>` shows how to use this feature.
 - Add option to :func:`.util.copy_model` from a non-default location of model files (:pull:`877`).
 
 .. _v3.9.0:

doc/api.rst

@@ -300,6 +300,8 @@ Utility methods
 .. automodule:: message_ix.util
    :members: expand_dims, copy_model, make_df
 
+.. automodule:: message_ix.util.sankey
+   :members: map_for_sankey
 
 Testing utilities
 -----------------

doc/conf.py

@@ -185,6 +185,7 @@ def local_inv(name: str, *parts: str) -> Optional[str]:
     "message_doc": ("https://docs.messageix.org/projects/global/en/latest/", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
     "pint": ("https://pint.readthedocs.io/en/stable/", None),
+    "plotly": ("https://plotly.com/python-api-reference", None),
     "plotnine": ("https://plotnine.org", None),
     "pyam": ("https://pyam-iamc.readthedocs.io/en/stable/", None),
     "python": ("https://docs.python.org/3/", None),

doc/reporting.rst

@@ -215,6 +215,7 @@ These automatic contents are prepared using:
    .. autosummary::
       add
       add_queue
+      add_sankey
       add_single
       apply
       check_keys

doc/tools/sankey.rst

@@ -0,0 +1,12 @@
+.. currentmodule:: message_ix.tools.sankey
+
+:mod:`.sankey`: generate Sankey diagrams
+****************************************
+
+See :meth:`.Reporter.add_sankey` and the :file:`westeros_sankey.ipynb` :ref:`tutorial <tutorial-westeros>`.
+
+API reference
+=============
+
+.. automodule:: message_ix.tools.sankey
+   :members:

message_ix/report/__init__.py

@@ -229,6 +229,62 @@ def from_scenario(cls, scenario, **kwargs) -> "Reporter":
 
         return rep
 
+    def add_sankey(
+        self,
+        year: int,
+        node: str,
+        exclude: list[str] = [],
+    ) -> str:
+        """Add the tasks required to produce a Sankey diagram.
+
+        See :func:`.map_for_sankey` for the meaning of the `node`, and `exclude`
+        parameters.
+
+        Parameters
+        ----------
+        year : int
+            The period (year) to be plotted.
+
+        Returns
+        -------
+        str
+            A key like :py:`"sankey figure a1b2c"`, where the last part is a unique hash
+            of the arguments `year`, `node`, and `exclude`. Calling
+            :meth:`.Reporter.get` with this key triggers generation of a
+            :class:`plotly.Figure <plotly.graph_objects.Figure>` with the Sankey
+            diagram.
+
+        See also
+        --------
+        map_for_sankey
+        pyam.figures.sankey
+        """
+        from warnings import filterwarnings
+
+        from genno import KeySeq
+        from genno.caching import hash_args
+        from pyam import IamDataFrame
+        from pyam.figures import sankey
+
+        from message_ix.tools.sankey import map_for_sankey
+
+        # Silence a warning raised by pyam-iamc 3.0.0 with pandas 2.2.3
+        filterwarnings("ignore", "Downcasting behavior", FutureWarning, "pyam.figures")
+
+        # Sequence of similar Keys for individual operations; use a unique hash of the
+        # arguments to avoid conflicts between multiple calls
+        unique = hash_args(year, node, exclude)[:6]
+        k = KeySeq(f"message sankey {unique}")
+
+        # Concatenate 'out' and 'in' data
+        self.add(k[0], "concat", "out::pyam", "in::pyam", strict=True)
+        # `df` argument to pyam.figures.sankey()
+        self.add(k[1], partial(IamDataFrame.filter, year=year), k[0])
+        # `mapping` argument to pyam.figures.sankey()
+        self.add(k[2], map_for_sankey, k[1], node=node, exclude=exclude)
+        # Generate the plotly.Figure object; return the key
+        return str(self.add(f"sankey figure {unique}", sankey, k[1], k[2]))
+
     def add_tasks(self, fail_action: Union[int, str] = "raise") -> None:
         """Add the pre-defined MESSAGEix reporting tasks to the Reporter.
 

message_ix/tests/test_report.py

@@ -20,6 +20,20 @@
 from message_ix.testing import SCENARIO, make_dantzig, make_westeros
 
 
+class TestReporter:
+    def test_add_sankey(self, test_mp, request) -> None:
+        scen = make_westeros(test_mp, solve=True, quiet=True, request=request)
+        rep = Reporter.from_scenario(scen, units={"replace": {"-": ""}})
+
+        # Method runs
+        key = rep.add_sankey(year=700, node="Westeros")
+
+        # Returns an existing key of the expected form
+        assert key.startswith("sankey figure ")
+
+        assert rep.check_keys(key)
+
+
 def test_reporter_no_solution(caplog, message_test_mp):
     scen = Scenario(message_test_mp, **SCENARIO["dantzig"])
 

message_ix/tests/test_tutorials.py

@@ -86,6 +86,7 @@ def _t(group: Union[str, None], basename: str, *, check=None, marks=None):
     _t("w0", f"{W}_addon_technologies"),
     _t("w0", f"{W}_historical_new_capacity"),
     _t("w0", f"{W}_multinode_energy_trade"),
+    _t("w0", f"{W}_sankey"),
     # NB this is the same value as in test_reporter()
     _t(None, f"{W}_report", check=[("len-rep-graph", 13724)]),
     _t("at0", "austria", check=[("solve-objective-value", 206321.90625)]),

message_ix/tests/tools/test_sankey.py

@@ -0,0 +1,50 @@
+from typing import TYPE_CHECKING, cast
+
+from ixmp.testing import assert_logs
+
+from message_ix.report import Reporter
+from message_ix.testing import make_westeros
+from message_ix.tools.sankey import map_for_sankey
+
+if TYPE_CHECKING:
+    import pyam
+
+
+def test_map_for_sankey(caplog, test_mp, request) -> None:
+    from genno.operator import concat
+
+    scen = make_westeros(test_mp, solve=True, request=request)
+    rep = Reporter.from_scenario(scen, units={"replace": {"-": ""}})
+    df = cast(
+        "pyam.IamDataFrame", concat(rep.get("in::pyam"), rep.get("out::pyam"))
+    ).filter(year=700)
+
+    # Set expectations
+    expected_all = {
+        "in|final|electricity|bulb|standard": ("final|electricity", "bulb|standard"),
+        "in|secondary|electricity|grid|standard": (
+            "secondary|electricity",
+            "grid|standard",
+        ),
+        "out|final|electricity|grid|standard": ("grid|standard", "final|electricity"),
+        "out|secondary|electricity|coal_ppl|standard": (
+            "coal_ppl|standard",
+            "secondary|electricity",
+        ),
+        "out|secondary|electricity|wind_ppl|standard": (
+            "wind_ppl|standard",
+            "secondary|electricity",
+        ),
+        "out|useful|light|bulb|standard": ("bulb|standard", "useful|light"),
+    }
+
+    # Load all variables
+    assert expected_all == map_for_sankey(df, node="Westeros")
+
+    x = "final|electricity"
+    assert {k: v for (k, v) in expected_all.items() if x not in v} == map_for_sankey(
+        df, node="Westeros", exclude=[x]
+    )
+
+    with assert_logs(caplog, "No mapping entries generated"):
+        map_for_sankey(df, node="not_a_node")

message_ix/tools/sankey.py

@@ -0,0 +1,70 @@
+import logging
+from typing import TYPE_CHECKING
+
+try:
+    from pyam.str import get_variable_components
+except ImportError:  # Python < 3.10 → pyam-iamc < 3
+    from pyam.utils import get_variable_components
+
+
+if TYPE_CHECKING:
+    import pyam
+
+log = logging.getLogger(__name__)
+
+
+def exclude_flow(flow: tuple[str, str], exclude: list[str]) -> bool:
+    """Return :any:`True` if either the source or target of `flow` is in `exclude`."""
+    return flow[0] in exclude or flow[1] in exclude
+
+
+def get_source_and_target(variable: str) -> tuple[str, str]:
+    """Get source and target for the `variable` flow."""
+    start_idx, end_idx = get_start_and_end_index(variable)
+    return (
+        get_variable_components(variable, start_idx, join=True),
+        get_variable_components(variable, end_idx, join=True),
+    )
+
+
+def get_start_and_end_index(variable: str) -> tuple[list[int], list[int]]:
+    """Get indices of source and target in variable name."""
+    return (
+        ([1, 2], [3, 4])
+        if get_variable_components(variable, 0) == "in"
+        else ([3, 4], [1, 2])
+    )
+
+
+def map_for_sankey(
+    iam_df: "pyam.IamDataFrame", node: str, exclude: list[str] = []
+) -> dict[str, tuple[str, str]]:
+    """Maps input to output flows to enable Sankey diagram.
+
+    Parameters
+    ----------
+    iam_df : :class:`pyam.IamDataframe`
+        Data to plot as Sankey diagram.
+    node : str
+        The node (MESSAGEix) or region (pyam) to plot.
+    exclude : list[str], optional
+        Flows to omit from the diagram. By default, nothing is excluded.
+
+    Returns
+    -------
+    dict
+        mapping from variable names to 2-tuples of their (inputs, output) flows.
+    """
+    result = {
+        var: get_source_and_target(var)
+        for var in iam_df.filter(region=node + "*").variable
+        if not exclude_flow(get_source_and_target(var), exclude)
+    }
+
+    if not result:
+        log.warning(
+            f"No mapping entries generated for {node=}, {exclude=} and data:\n"
+            + repr(iam_df)
+        )
+
+    return result

pyproject.toml

@@ -55,8 +55,9 @@ docs = [
   "sphinx_rtd_theme",
   "sphinxcontrib-bibtex",
 ]
-tutorial = ["jupyter", "matplotlib", "message_ix[report]"]
+tutorial = ["jupyter", "matplotlib", "message_ix[report,sankey]"]
 report = ["ixmp[report]"]
+sankey = ["plotly"]
 tests = [
   "asyncssh",
   "message_ix[docs,tutorial]",
@@ -96,7 +97,7 @@ local_partial_types = true
 [[tool.mypy.overrides]]
 # Packages/modules for which no type hints are available.
 module = [
-  "pyam",
+  "pyam.*",
   "scipy.*",
   # Indirectly via ixmp; this should be a subset of the list in ixmp's pyproject.toml
   "jpype",

tutorial/README.rst

@@ -83,6 +83,8 @@ From the command line
 
        $ jupyter notebook
 
+.. _tutorial-westeros:
+
 Westeros Electrified
 ====================
 
@@ -164,6 +166,10 @@ framework, such as used in global research applications of |MESSAGEix|.
       module to ‘report’ results, e.g. do post-processing, plotting, and other
       calculations (:tut:`westeros/westeros_report.ipynb`).
 
+      #. After familiarizing yourself with ‘reporting’, learn how to quickly assess
+         variable flows by plotting Sankey diagrams
+         (:tut:`westeros/westeros_sankey.ipynb`).
+
    #. Build the baseline scenario using data stored in Excel files to
       populate sets and parameters: