Address review comments

- Add 'sankey' optional dependencies set.
- Reporter.add_sankey()
  - Sort methods in alpha order.
  - Include all steps for figure generation.
  - Expand docstring.
- .tools.sankey —rename from .util.sankey
  - Sort methods in order.
  - Simplify type hints.
  - Remove year= parameter from map_for_sankey()
  - Add warning if map_for_sankey() gives an empty result.
- Reorganize tutorial to align with simplified interface.
- Simplify tests
- Update docs
  - Add doc/tools/sankey.rst
  - Add plotly to intersphinx config.
  - Remove trailing whitespace in tutorial/README.rst
  - Link docs, tutorial in release notes.

RELEASE_NOTES.rst

@@ -12,7 +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 functionality to create Sankey diagrams from :class:`.Reporter` together with a new tutorial showcase (:pull:`770`).
+- 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/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/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.
 
@@ -243,24 +299,3 @@ def add_tasks(self, fail_action: Union[int, str] = "raise") -> None:
 
         # Use a queue pattern via Reporter.add_queue()
         self.add_queue(get_tasks(), fail=fail_action)
-
-    def add_sankey(self, fail_action: Union[int, str] = "raise") -> None:
-        """Add the calculations required to produce Sankey plots.
-
-        Parameters
-        ----------
-        fail_action : "raise" or int
-            :mod:`logging` level or level name, passed to the `fail` argument of
-            :meth:`.Reporter.add_queue`.
-        """
-        # NOTE This includes just one task for the base version, but could later be
-        # expanded.
-        self.add_queue(
-            [
-                (
-                    ("message::sankey", "concat", "out::pyam", "in::pyam"),
-                    dict(strict=True),
-                )
-            ],
-            fail=fail_action,
-        )

message_ix/tests/test_report.py

@@ -292,17 +292,13 @@ def add_tm(df, name="Activity"):
 
 
 def test_reporter_add_sankey(test_mp, request):
-    scen = make_westeros(
-        test_mp, emissions=True, solve=True, quiet=True, request=request
-    )
-
-    # Reporter.from_scenario can handle Westeros example model
-    rep = Reporter.from_scenario(scen)
+    scen = make_westeros(test_mp, solve=True, quiet=True, request=request)
+    rep = Reporter.from_scenario(scen, units={"replace": {"-": ""}})
 
-    # Westeros-specific configuration: '-' is a reserved character in pint
-    configure(units={"replace": {"-": ""}})
+    # Method runs
+    key = rep.add_sankey()
 
-    # Add Sankey calculation(s)
-    rep.add_sankey()
+    # Returns an existing key of the expected form
+    assert key.startswith("sankey figure ")
 
-    assert rep.check_keys("message::sankey")
+    assert rep.check_keys(key)

message_ix/tests/test_util.py

@@ -4,9 +4,7 @@
 import pytest
 
 from message_ix import Scenario, make_df
-from message_ix.report import Reporter
 from message_ix.testing import make_dantzig, make_westeros
-from message_ix.util.sankey import map_for_sankey
 
 
 def test_make_df():
@@ -61,46 +59,3 @@ def test_testing_make_scenario(test_mp, request):
     # Westeros model can be created
     scen = make_westeros(test_mp, solve=True, request=request)
     assert isinstance(scen, Scenario)
-
-
-def test_map_for_sankey(test_mp, request):
-    # NB: we actually only need a pyam.IamDataFrame that has the same form as the result
-    # of these setup steps, so maybe this can be simplified
-    scen = make_westeros(test_mp, solve=True, request=request)
-    rep = Reporter.from_scenario(scen)
-    rep.configure(units={"replace": {"-": ""}})
-    rep.add_sankey()
-    df = rep.get("message::sankey")
-
-    # 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"),
-    }
-    expected_without_final_electricity = {
-        key: value
-        for (key, value) in expected_all.items()
-        if "final|electricity" not in value
-    }
-
-    # Load all variables
-    mapping_all = map_for_sankey(df, year=700, region="Westeros")
-    assert mapping_all == expected_all
-
-    mapping_without_final_electricity = map_for_sankey(
-        df, year=700, region="Westeros", exclude=["final|electricity"]
-    )
-    assert mapping_without_final_electricity == expected_without_final_electricity

message_ix/tests/tools/test_sankey.py

@@ -0,0 +1,38 @@
+from message_ix.report import Reporter
+from message_ix.testing import make_westeros
+from message_ix.tools.sankey import map_for_sankey
+
+
+def test_map_for_sankey(test_mp, request):
+    from genno.operator import concat
+
+    scen = make_westeros(test_mp, solve=True, request=request)
+    rep = Reporter.from_scenario(scen, units={"replace": {"-": ""}})
+    df = 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]
+    )

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]", "plotly"]
+tutorial = ["jupyter", "matplotlib", "message_ix[report,sankey]"]
 report = ["ixmp[report]"]
+sankey = ["plotly"]
 tests = [
   "asyncssh",
   "message_ix[docs,tutorial]",