diff --git a/message_ix/util/__init__.py b/message_ix/util/__init__.py index 113fe3abf..c4f6d96cb 100644 --- a/message_ix/util/__init__.py +++ b/message_ix/util/__init__.py @@ -6,7 +6,7 @@ import pandas as pd from ixmp.backend import ItemType from pandas.api.types import is_scalar -from pyam.utils import get_variable_components as gvc +from pyam.str import get_variable_components as gvc from message_ix.core import Scenario from message_ix.models import MACRO, MESSAGE @@ -234,21 +234,17 @@ def expand_dims(scenario: Scenario, name, **data): # Add the expanded data scenario.add_par(name, new_data) -def sankey_mapper(df, year, region, flows_not_needed=[], variables_not_needed = []): - df_filtered = df.filter( - region=region+'*', - year=year) +def sankey_mapper(df, year, region, flows_not_needed=[], variables_not_needed=[]): + df_filtered = df.filter(region=region + "*", year=year) mapping = {} - for var in df_filtered['variable']: - if gvc(var, 0) == 'in': - mapping[var] = (gvc(var, [1,2], join=True), gvc(var, [3,4], join=True)) - if gvc(var, 0) == 'out': - mapping[var] = (gvc(var, [3,4], join=True), gvc(var, [1,2], join=True)) - - - + for var in df_filtered["variable"]: + if gvc(var, 0) == "in": + mapping[var] = (gvc(var, [1, 2], join=True), gvc(var, [3, 4], join=True)) + if gvc(var, 0) == "out": + mapping[var] = (gvc(var, [3, 4], join=True), gvc(var, [1, 2], join=True)) + for k in mapping.keys(): for flow in flows_not_needed: if flow in k: @@ -256,4 +252,4 @@ def sankey_mapper(df, year, region, flows_not_needed=[], variables_not_needed = for var in variables_not_needed: del mapping[var] - return mapping \ No newline at end of file + return mapping diff --git a/tutorial/westeros/westeros_sankey.ipynb b/tutorial/westeros/westeros_sankey.ipynb index 8acc18ddc..cda97290d 100644 --- a/tutorial/westeros/westeros_sankey.ipynb +++ b/tutorial/westeros/westeros_sankey.ipynb @@ -23,12 +23,13 @@ "metadata": {}, "outputs": [], "source": [ - "from message_ix.report import Reporter\n", "import ixmp\n", + "\n", "import message_ix\n", + "from message_ix.report import Reporter\n", "\n", "mp = ixmp.Platform()\n", - "scenario = message_ix.Scenario(mp, model=\"Westeros Electrified\", scenario=\"baseline\")\n" + "scenario = message_ix.Scenario(mp, model=\"Westeros Electrified\", scenario=\"baseline\")" ] }, { @@ -61,9 +62,7 @@ "source": [ "rep = Reporter.from_scenario(scenario)\n", "\n", - "rep.configure(units={\"replace\": {\"-\": \"\"}})\n", - "\n", - "\n" + "rep.configure(units={\"replace\": {\"-\": \"\"}})" ] }, { @@ -79,7 +78,7 @@ "metadata": {}, "outputs": [], "source": [ - "df=rep.get(\"message::sankey\")" + "df = rep.get(\"message::sankey\")" ] }, { @@ -97,7 +96,7 @@ "metadata": {}, "outputs": [], "source": [ - "mapping=message_ix.util.sankey_mapper(df,year=700,region=\"Westeros\")" + "mapping = message_ix.util.sankey_mapper(df, year=700, region=\"Westeros\")" ] }, { @@ -116,8 +115,9 @@ "outputs": [], "source": [ "import plotly.offline as pof\n", + "\n", "fig = df.filter(year=700).plot.sankey(mapping=mapping)\n", - "pof.plot(fig) # opens a new window in your browser" + "pof.plot(fig) # opens a new window in your browser" ] }, {