Apply formatting and fix import from pyam

* In recent pyam versions, get_variable_components moved to pyam.str
iiasa · Jan 10, 2024 · 748a889 · 748a889
1 parent 9656c81
commit 748a889
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 22 deletions.
diff --git 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,26 +234,22 @@ 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:
                 variables_not_needed.append(k)
     for var in variables_not_needed:
         del mapping[var]
 
-    return mapping
+    return mapping
diff --git 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"
    ]
   },
   {