Add .util.sankey and tutorial

RELEASE_NOTES.rst

@@ -12,6 +12,7 @@ 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 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/reporting.rst

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

message_ix/report/__init__.py

@@ -243,3 +243,24 @@ 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

@@ -272,3 +272,20 @@ def add_tm(df, name="Activity"):
     # Results have the expected units
     assert all(df5["unit"] == "centiUSD / case")
     assert_series_equal(df4["value"], df5["value"] / 100.0)
+
+
+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)
+
+    # Westeros-specific configuration: '-' is a reserved character in pint
+    configure(units={"replace": {"-": ""}})
+
+    # Add Sankey calculation(s)
+    rep.add_sankey()
+
+    assert rep.check_keys("message::sankey")

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/test_util.py

@@ -4,7 +4,9 @@
 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():
@@ -59,3 +61,46 @@ 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/util/sankey.py

@@ -0,0 +1,75 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from pyam import IamDataFrame
+
+try:
+    from pyam.str import get_variable_components
+except ImportError:  # Python < 3.10, pandas < 2.0
+    from pyam.utils import get_variable_components
+
+try:
+    from typing import LiteralString
+except ImportError:  # Python < 3.11
+    from typing_extensions import LiteralString
+
+
+def map_for_sankey(
+    iam_df: IamDataFrame,
+    year: int,
+    region: str,
+    exclude: List[Optional[str]] = [],
+) -> Dict[str, Tuple[Union[List, Any, LiteralString], Union[List, Any, LiteralString]]]:
+    """Maps input to output flows to enable Sankey plots.
+
+    Parameters
+    ----------
+    iam_df: :class:`pyam.IamDataframe`
+        The IAMC-format DataFrame holding the data to plot as Sankey diagrams.
+    year: int
+        The year to display in the Sankey diagram.
+    region: str
+        The region to display in the Sankey diagram.
+    exclude: list[str], optional
+        If provided, exclude these keys from the Sankey diagram. Defaults to an empty
+        list, i.e. showing all flows.
+
+    Returns
+    -------
+    mapping: dict
+        A mapping from variable names to their inputs and outputs.
+    """
+    return {
+        var: get_source_and_target(var)
+        for var in iam_df.filter(region=region + "*", year=year).variable
+        if not exclude_flow(get_source_and_target(var), exclude)
+    }
+
+
+def get_source_and_target(
+    variable: str,
+) -> Tuple[Union[List, Any, LiteralString], Union[List, Any, LiteralString]]:
+    """Get source and target for the `variable` flow."""
+    start_idx, end_idx = set_start_and_end_index(variable)
+    return (
+        get_variable_components(variable, start_idx, join=True),
+        get_variable_components(variable, end_idx, join=True),
+    )
+
+
+def set_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 exclude_flow(
+    flow: Tuple[Union[List, Any, LiteralString], Union[List, Any, LiteralString]],
+    exclude: List[Optional[str]],
+) -> bool:
+    """Exclude sources or targets of variable flow if requested."""
+    if flow[0] in exclude or flow[1] in exclude:
+        return True
+    return False

pyproject.toml

@@ -55,7 +55,7 @@ docs = [
   "sphinx_rtd_theme",
   "sphinxcontrib-bibtex",
 ]
-tutorial = ["jupyter", "matplotlib", "message_ix[report]"]
+tutorial = ["jupyter", "matplotlib", "message_ix[report]", "plotly"]
 report = ["ixmp[report]"]
 tests = [
   "asyncssh",
@@ -96,7 +96,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

@@ -164,6 +164,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: