diff --git a/message_ix/report/__init__.py b/message_ix/report/__init__.py index bfc4be567..3b72aa006 100644 --- a/message_ix/report/__init__.py +++ b/message_ix/report/__init__.py @@ -141,6 +141,7 @@ "message::costs", "message::emissions", ), + ("message::sankey", "concat", "out::pyam", "in::pyam"), ) diff --git a/message_ix/tests/test_util.py b/message_ix/tests/test_util.py index 5a64d704f..48225b1da 100644 --- a/message_ix/tests/test_util.py +++ b/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 sankey_mapper def test_make_df(): @@ -59,3 +61,55 @@ def test_testing_make_scenario(test_mp): # Westeros model can be created scen = make_westeros(test_mp, solve=True) assert isinstance(scen, Scenario) + + +def test_sankey_mapper(test_mp): + # NB: we actually only need a pd.DataFrame 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) + rep = Reporter.from_scenario(scen) + rep.configure(units={"replace": {"-": ""}}) + 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 = { + "in|secondary|electricity|grid|standard": ( + "secondary|electricity", + "grid|standard", + ), + "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 + mapping_all = sankey_mapper(df, year=700, region="Westeros") + assert mapping_all == expected_all + + mapping_without_final_electricity = sankey_mapper( + df, year=700, region="Westeros", exclude=["final|electricity"] + ) + assert mapping_without_final_electricity == expected_without_final_electricity diff --git a/message_ix/util/sankey.py b/message_ix/util/sankey.py new file mode 100644 index 000000000..c41c8e4df --- /dev/null +++ b/message_ix/util/sankey.py @@ -0,0 +1,28 @@ +from typing import Any, Optional + +import pandas as pd + +try: + from pyam.str import get_variable_components as gvc +except ImportError: # Python < 3.10, pandas < 2.0 + from pyam.utils import get_variable_components as gvc + + +def sankey_mapper( + df: pd.DataFrame, + year: int, + region: str, + exclude: list[Optional[str]] = [], +) -> dict[str, Any]: + mapping = {} + + for var in df.filter(region=region + "*", year=year).variable: + is_input = gvc(var, 0) == "in" + (start_idx, end_idx) = ([1, 2], [3, 4]) if is_input else ([3, 4], [1, 2]) + source = gvc(var, start_idx, join=True) + target = gvc(var, end_idx, join=True) + if source in exclude or target in exclude: + continue + mapping[var] = (source, target) + + return mapping diff --git a/pyproject.toml b/pyproject.toml index 51037542c..712ae7f48 100644 --- a/pyproject.toml +++ b/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", diff --git a/tutorial/westeros/westeros_sankey.ipynb b/tutorial/westeros/westeros_sankey.ipynb new file mode 100644 index 000000000..957058142 --- /dev/null +++ b/tutorial/westeros/westeros_sankey.ipynb @@ -0,0 +1,199 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Westeros Tutorial - Introducing Sankey diagrams\n", + "\n", + "Sankey diagrams are a useful technique to visualize energy flow accounts.\n", + "\n", + "This tutorial introduces the sankey feature provided by the ``pyam`` packages.\n", + "\n", + "\n", + "**Pre-requisites**\n", + "- You have the *MESSAGEix* framework installed and working\n", + " In particular, you should have installed ``message_ix``, ``pyam``, and ``plotly``\n", + "- Complete tutorial Part 1 (``westeros_baseline.ipynb``) and Introducing Reporting (``westeros_report.ipynb``)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import ixmp\n", + "\n", + "from message_ix import Scenario\n", + "\n", + "mp = ixmp.Platform()\n", + "scenario = Scenario(mp, model=\"Westeros Electrified\", scenario=\"baseline\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Remove any existing solution\n", + "try:\n", + " scenario.remove_solution()\n", + "except ValueError:\n", + " pass\n", + "\n", + "scenario.solve()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the reporter object. (Since ``\"-\"`` is not a unit, we replace it by ``\"\"``.)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from message_ix.report import Reporter\n", + "\n", + "rep = Reporter.from_scenario(scenario)\n", + "\n", + "rep.configure(units={\"replace\": {\"-\": \"\"}})" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the `message::sankey` reporter option to generate a pyam.dataframe including the reqiured input (`in::pyam`) and output flows (`out::pyam`) in iamc format.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "df = rep.get(\"message::sankey\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The utility function `sankey_mapper(df, year, region, exclude=[])` can be used to create the required mapping for the `figures.sankey()` function of the `pyam` package.\n", + "\n", + "In some models it might be necessary to exclude variables and flow to get meaningful sankey diagrams. But let´s try with all!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from message_ix.util.sankey import sankey_mapper\n", + "\n", + "mapping = sankey_mapper(df, year=700, region=\"Westeros\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pyam function `plot.sankey(mapping)`returns a plotly sankey figure object that can be further modified.\n", + "\n", + "To plot it as an interactive diagram in your web browser, you can do the following." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pyam.figures import sankey\n", + "\n", + "fig = sankey(df=df.filter(year=700), mapping=mapping)\n", + "fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mapping_without_final_electricity = sankey_mapper(\n", + " df, year=700, region=\"Westeros\", exclude=[\"final|electricity\"]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = sankey(df=df.filter(year=700), mapping=mapping_without_final_electricity)\n", + "fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mapping_without_wind_ppl_standard = sankey_mapper(\n", + " df, year=700, region=\"Westeros\", exclude=[\"wind_ppl|standard\"]\n", + ")\n", + "fig = sankey(df=df.filter(year=700), mapping=mapping_without_wind_ppl_standard)\n", + "fig.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Do not forget to close the database ;-) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mp.close_db()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "messageix", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +}