From 748a889559a0350eca79a216b74704ed96646f8f Mon Sep 17 00:00:00 2001
From: Fridolin Glatter <glatter@iiasa.ac.at>
Date: Wed, 10 Jan 2024 13:11:07 +0100
Subject: [PATCH] Apply formatting and fix import from pyam

* In recent pyam versions, get_variable_components moved to pyam.str
---
 message_ix/util/__init__.py             | 24 ++++++++++--------------
 tutorial/westeros/westeros_sankey.ipynb | 16 ++++++++--------
 2 files changed, 18 insertions(+), 22 deletions(-)

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"
    ]
   },
   {