Update use of integer keys in .transport.demand

- Replace usage of genno.KeySeq with Key indexing.

Author: khaeru

message_ix_models/model/transport/demand.py

@@ -6,13 +6,13 @@
 import genno
 import numpy as np
 import pandas as pd
-from dask.core import literal
-from genno import Computer, KeySeq
+from genno import Key, literal
 from message_ix import make_df
 
 from message_ix_models.report.key import GDP
 from message_ix_models.util import broadcast
 
+from . import factor
 from .key import (
     cg,
     cost,
@@ -36,6 +36,7 @@
 )
 
 if TYPE_CHECKING:
+    from genno import Computer
     from genno.types import AnyQuantity
 
     from .config import Config
@@ -136,15 +137,15 @@ def dummy(
     # Mode shares
     ((ms, "logit", cost, sw, "lambda:", y), dict(dim="t")),
     # Total PDT (n, t, y), with modes for the 't' dimension
-    (pdt_nyt + "0", "mul", pdt_ny, ms),
+    (pdt_nyt[0], "mul", pdt_ny, ms),
     # Scenario-specific adjustment factors
     ("pdt factor:n-y-t", "factor_pdt", n, y, t_modes, "config"),
     # Only the LDV values
     (
         ("ldv pdt factor:n-y", "select", "pdt factor:n-y-t"),
         dict(indexers=dict(t="LDV"), drop=True),
     ),
-    (pdt_nyt, "mul", pdt_nyt + "0", "pdt factor:n-y-t"),
+    (pdt_nyt, "mul", pdt_nyt[0], "pdt factor:n-y-t"),
     # Per capita (for validation)
     (pdt_nyt + "capita+post", "div", pdt_nyt, pop),
     # LDV PDT only (n, y)
@@ -163,18 +164,18 @@ def dummy(
     # LDV PDT shared out by consumer group (cg, n, y)
     (ldv_nycg, "mul", ldv_ny + "total", cg),
     # Select only non-LDV PDT
-    ((pdt_nyt + "1", "select", pdt_nyt), dict(indexers=dict(t=["LDV"]), inverse=True)),
+    ((pdt_nyt[1], "select", pdt_nyt), dict(indexers=dict(t=["LDV"]), inverse=True)),
     # Relabel PDT
     (
-        (pdt_cny + "0", "relabel2", pdt_nyt + "1"),
+        (pdt_cny[0], "relabel2", pdt_nyt[1]),
         dict(new_dims={"c": "transport pax {t.lower()}"}),
     ),
-    (pdt_cny, "convert_units", pdt_cny + "0", "Gp km / a"),
+    (pdt_cny, "convert_units", pdt_cny[0], "Gp km / a"),
     # Convert to ixmp format
     (("demand::P+ixmp", "as_message_df", pdt_cny), _DEMAND_KW),
     # Relabel ldv pdt:n-y-cg
-    ((ldv_cny + "0", "relabel2", ldv_nycg), dict(new_dims={"c": "transport pax {cg}"})),
-    (ldv_cny, "convert_units", ldv_cny + "0", "Gp km / a"),
+    ((ldv_cny[0], "relabel2", ldv_nycg), dict(new_dims={"c": "transport pax {cg}"})),
+    (ldv_cny, "convert_units", ldv_cny[0], "Gp km / a"),
     (("demand::LDV+ixmp", "as_message_df", ldv_cny), _DEMAND_KW),
     # Dummy demands, if these are configured
     ("demand::dummy+ixmp", dummy, "c::transport", "nodes::ex world", y, "config"),
@@ -189,7 +190,7 @@ def dummy(
 ]
 
 
-def pdt_per_capita(c: Computer) -> None:
+def pdt_per_capita(c: "Computer") -> None:
     """Set up calculation of :data:`~.key.pdt_cap`.
 
     Per Schäfer et al. (2009) Figure 2.5: linear interpolation between log GDP PPP per
@@ -202,12 +203,12 @@ def pdt_per_capita(c: Computer) -> None:
     between projected, log GDP in each future period and the log GDP in the reference
     year.
     """
-    gdp = KeySeq(GDP)
-    pdt = KeySeq("_pdt:n-y")
+    gdp = Key(GDP)
+    pdt = Key("_pdt:n-y")
 
     # GDP expressed in PPP. In the SSP(2024) input files, this conversion is already
     # applied, so no need to multiply by a mer_to_ppp factor here → simple alias.
-    c.add(gdp["PPP"], gdp.base)
+    c.add(gdp["PPP"], gdp)
 
     # GDP PPP per capita
     c.add(gdp["capita"], "div", gdp["PPP"], pop)
@@ -229,7 +230,7 @@ def _delta(qty: "AnyQuantity", y0: int) -> "AnyQuantity":
     # Same transformation for both quantities
     for x, reference_values in ((gdp, gdp["capita"]), (pdt, pdt["ref"])):
         # Retrieve value from configuration
-        k = KeySeq(f"{x.name}::fixed")
+        k = Key(x.name, (), "fixed")
         c.add(k[0], "quantity_from_config", "config", name=f"fixed_{x.name.strip('_')}")
         # Broadcast on `n` dimension
         c.add(k[1] * "n", "mul", k[0], "n:n:ex world")
@@ -286,15 +287,13 @@ def _delta(qty: "AnyQuantity", y0: int) -> "AnyQuantity":
     c.add(gdp_cap + "adj", gdp[6])
 
 
-def prepare_computer(c: Computer) -> None:
+def prepare_computer(c: "Computer") -> None:
     """Prepare `c` to calculate and add transport demand data.
 
     See also
     --------
     TASKS
     """
-    from . import factor
-
     config: "Config" = c.graph["context"].transport
 
     # Compute total PDT per capita
@@ -313,6 +312,6 @@ def prepare_computer(c: Computer) -> None:
         c.add(pdt_cap * "t", "select", exo.pdt_cap_proj, indexers=dict(scenario="LED"))
 
         # Multiply by population for the total
-        c.add(pdt_nyt + "0", "mul", pdt_cap * "t", pop)
+        c.add(pdt_nyt[0], "mul", pdt_cap * "t", pop)
 
     c.add("transport_data", __name__, key="transport demand::ixmp")