Update usage of genno.KeySeq

Use features of Key directly.

Author: khaeru

message_ix_models/model/transport/demand.py

@@ -115,13 +115,13 @@ def dummy(
     ((price[1], "select", price[0]), dict(indexers=dict(c="transport"), drop=True)),
     (price[2], "price_units", price[1]),
     # Smooth prices to avoid zig-zag in share projections
-    (price.base, "smooth", price[2]),
+    (price, "smooth", price[2]),
     # Interpolate speed data
     (("speed:scenario-n-t-y:0", "interpolate", exo.speed, "y::coords"), EXTRAPOLATE),
     # Select speed data
     ("speed:n-t-y", "select", "speed:scenario-n-t-y:0", "indexers:scenario"),
     # Cost of transport (n, t, y)
-    (cost, "cost", price.base, gdp_cap, "whour:", "speed:n-t-y", "votm:n-y", y),
+    (cost, "cost", price, gdp_cap, "whour:", "speed:n-t-y", "votm:n-y", y),
     # Share weights (n, t, y)
     (
         sw,

message_ix_models/model/transport/key.py

@@ -2,7 +2,7 @@
 
 from types import SimpleNamespace
 
-from genno import Key, KeySeq
+from genno import Key
 
 from message_ix_models.report.key import GDP, PRICE_COMMODITY
 from message_ix_models.util.genno import Keys
@@ -109,7 +109,7 @@
 pdt_nyt = _pdt * "t"
 
 #: Prices.
-price = KeySeq(PRICE_COMMODITY / ("h", "l") + "transport")
+price = Key(PRICE_COMMODITY / ("h", "l") + "transport")
 
 #: Keys for :mod:`.transport.report`.
 report = SimpleNamespace(

message_ix_models/model/transport/operator.py

@@ -12,7 +12,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from genno import Computer, KeySeq, Operator, quote
+from genno import Computer, Key, Operator, quote
 from genno.operator import apply_units, as_quantity, rename_dims
 from genno.testing import assert_qty_allclose, assert_units
 from scipy import integrate
@@ -355,9 +355,9 @@ def distance_nonldv(context: "Context") -> "AnyQuantity":
     source_kw = dict(measure="Vehicle use", aggregate=True)
     keys = exo_data.prepare_computer(context, c, "IEA EEI", source_kw)
 
-    ks = KeySeq(keys[0])
+    ks = Key(keys[0])
 
-    c.add(ks[0], "select", ks.base, indexers={"SECTOR": "transport"}, drop=True)
+    c.add(ks[0], "select", ks, indexers={"SECTOR": "transport"}, drop=True)
     c.add(ks[1], "rename", ks[0], quote({"Mode/vehicle type": "t"}))
     # Replace IEA EEI technology codes with MESSAGEix-Transport ones
     c.add(ks[2], "relabel", ks[1], labels=dict(t=EEI_TECH_MAP))

message_ix_models/model/transport/passenger.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 import pandas as pd
-from genno import Computer, Key, KeySeq, Quantity, quote
+from genno import Computer, Key, Quantity, quote
 from message_ix import make_df
 from sdmx.model.v21 import Code
 
@@ -96,8 +96,8 @@ def prepare_computer(c: Computer):
 
     #### NB lines below duplicated from .transport.base
     e_iea = Key("energy:n-y-product-flow:iea")
-    e_fnp = KeySeq(e_iea.drop("y"))
-    e = KeySeq("energy:commodity-flow-node_loc:iea")
+    e_fnp = Key(e_iea.drop("y"))
+    e = Key("energy:commodity-flow-node_loc:iea")
 
     # Transform IEA EWEB data for comparison
 
@@ -219,7 +219,7 @@ def _(nodes, technologies, y0, config: dict) -> Quantity:
             units="GWa",
         )
 
-    k = KeySeq("bound_activity_lo:n-t-y:transport minimum")
+    k = Key("bound_activity_lo:n-t-y:transport minimum")
     c.add(next(k), _, "n::ex world", "t::transport", "y0", "config")
 
     # Produce MESSAGE parameter bound_activity_lo:nl-t-ya-m-h

message_ix_models/model/transport/report.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 import pandas as pd
-from genno import Computer, Key, Keys, KeySeq, MissingKeyError
+from genno import Computer, Key, Keys, MissingKeyError
 from genno.core.key import single_key
 from message_ix import Reporter
 
@@ -133,16 +133,13 @@ def add_iamc_store_write(c: Computer, base_key) -> "Key":
 
     .. todo:: Move upstream, to :mod:`message_ix_models`.
     """
-    k = KeySeq(base_key)
+    k = Key(base_key)
 
     file_keys = []
     for suffix in ("csv", "xlsx"):
         # Create the path
         path = c.add(
-            k[f"{suffix} path"],
-            "make_output_path",
-            "config",
-            name=f"{k.base.name}.{suffix}",
+            k[f"{suffix} path"], "make_output_path", "config", name=f"{k.name}.{suffix}"
         )
         # Write `key` to the path
         file_keys.append(c.add(k[suffix], "write_report", base_key, path))
@@ -163,12 +160,12 @@ def aggregate(c: "Computer") -> None:
 
     config: Config = c.graph["config"]["transport"]
 
-    for key in map(lambda s: KeySeq(c.infer_keys(s)), "emi in out".split()):
+    for key in map(lambda s: Key(c.infer_keys(s)), "emi in out".split()):
         try:
             # Reference the function to avoid the genno magic which would treat as sum()
             # NB aggregation on the nl dimension *could* come first, but this can use a
             #    lot of memory when applied to e.g. out:*: for a full global model.
-            c.add(key[0], func, key.base, "t::transport agg", keep=False)
+            c.add(key[0], func, key, "t::transport agg", keep=False)
             c.add(key[1], func, key[0], "nl::world agg", keep=False)
             c.add(key["transport"], "select", key[1], "t::transport modes 1", sums=True)
         except MissingKeyError:

message_ix_models/tests/model/transport/test_base.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from genno import Computer, KeySeq
+from genno import Computer, Key
 from genno.operator import random_qty, relabel
 
 from message_ix_models.model.structure import get_codes
@@ -23,7 +23,7 @@ def test_smooth(recwarn) -> None:
     c.add("y0", y0)
 
     # Expected base key for input
-    k = KeySeq("ue:nl-ya-c-l-h-t")
+    k = Key("ue:nl-ya-c-l-h-t")
 
     # Prepare input data
     periods = list(

message_ix_models/tests/model/transport/test_build.py

@@ -191,7 +191,7 @@ def run(self, obj):
     "cg share:n-y-cg": (HasUnits(""),),
     "GDP:n-y:PPP+capita": (HasUnits("kUSD / passenger / year"),),
     "votm:n-y": (HasUnits(""),),
-    key.price.base: (HasUnits("USD / km"),),
+    key.price: (HasUnits("USD / km"),),
     "cost:n-y-c-t": (HasUnits("USD / km"),),
     key.pdt_nyt[0]: (HasUnits("passenger km / year"),),
     key.pdt_nyt[1]: (HasUnits("passenger km / year"),),
@@ -396,7 +396,7 @@ def test_debug(
     options: dict,
     N_node: int,
     *,
-    verbosity: Literal[0, 1, 2, 3] = 0,
+    verbosity: Literal[0, 1, 2, 3] = 1,
 ):
     """Check and debug particular steps in the transport build process.
 
@@ -440,6 +440,7 @@ def test_debug(
 
     # DEBUG Show and compute a different key
     # k = key.pdt_cny
+    k = "transport::policy+ixmp"
 
     # Show what will be computed
     # verbosity = True  # DEBUG Force printing the description
@@ -452,7 +453,7 @@ def test_debug(
     tmp = c.get(k)
 
     # DEBUG Handle a subset of the result for inspection
-    # print(tmp)
+    print(tmp)
 
     assert result, "1 or more checks failed"
     del tmp

message_ix_models/tools/costs/gdp.py

@@ -3,7 +3,6 @@
 
 import numpy as np
 import pandas as pd
-from genno import KeySeq
 
 from message_ix_models import Context
 
@@ -82,13 +81,13 @@ def broadcast_qty(s) -> Quantity:
     # Concatenate single-scenario data
     k_pop = Key("pop", dims)
     c.add(k_pop, "concat", *keys["pop"])
-    k_gdp = KeySeq("gdp", dims)
-    c.add(k_gdp.base, "concat", *keys["gdp"])
+    k_gdp = Key("gdp", dims)
+    c.add(k_gdp, "concat", *keys["gdp"])
 
     # Further calculations
 
     # GDP per capita
-    c.add(k_gdp["cap"], "div", k_gdp.base, k_pop)
+    c.add(k_gdp["cap"], "div", k_gdp, k_pop)
 
     # Ratio to reference region value
     c.add(
@@ -114,7 +113,7 @@ def merge(*dfs: pd.DataFrame) -> pd.DataFrame:
         )
 
     k_result = "data::pandas"
-    c.add(k_result, merge, k_pop, k_gdp.base, k_gdp["cap"], k_gdp["indexed"])
+    c.add(k_result, merge, k_pop, k_gdp, k_gdp["cap"], k_gdp["indexed"])
 
     # log.debug(c.describe(k_result))  # DEBUG Show what would be done
     result = c.get(k_result)