Merge pull request #309 from iiasa/issue/304

Adapt to modified "Emissions|*|Aviation|*" variable names in `.ssp.transport`

Author: khaeru

doc/project/ssp.rst

@@ -80,7 +80,22 @@ Transport
 
    2. To process an existing :class:`pandas.DataFrame` from other code, call :func:`.process_df`, passing the input dataframe and the `method` parameter.
 
+   As of 2025-03-07 / :pull:`309`, the set of required "variable" codes handled includes::
+
+       Emissions|.*|Energy|Bunkers
+       Emissions|.*|Energy|Bunkers|International Aviation
+       Emissions|.*|Energy|Demand|Transportation
+       Emissions|.*|Energy|Demand|Transportation|Road Rail and Domestic Shipping
+
+   The previous set is no longer supported.
+
    As of 2025-01-25:
 
+   - The set of "variable" codes modified includes::
+
+       Emissions|.*|Energy|Demand|Transportation|Aviation
+       Emissions|.*|Energy|Demand|Transportation|Aviation|International
+       Emissions|.*|Energy|Demand|Transportation|Road Rail and Domestic Shipping
+
    - Method 'B' (that is, :func:`.prepare_method_B`; see its documentation) is the preferred method.
    - The code is tested on :file:`.xlsx` files in the (internal) directories under `SharePoint > ECE > Documents > SharedSocioeconomicPathways2023 > Scenario_Vetting <https://iiasahub.sharepoint.com/sites/eceprog/Shared%20Documents/Forms/AllItems.aspx?csf=1&web=1&e=APKv0Z&CID=23fa0a51%2Dc303%2D4381%2D8c6d%2D143305cbc5a1&FolderCTID=0x012000AA9481BF7BE9264E85B14105F7F082FF&id=%2Fsites%2Feceprog%2FShared%20Documents%2FSharedSocioEconomicPathways2023%2FScenario%5FVetting&viewid=956acd8a%2De1e7%2D4ae9%2Dab1b%2D0506911bae11>`_, for example :file:`v2.1_Internal_version_Dec13_2024/Reporting_output/SSP_SSP2_v2.1_baseline.xlsx`.

doc/whatsnew.rst

@@ -26,8 +26,9 @@ SSP :ref:`ssp-2024`/ScenarioMIP
 
 Improve :mod:`.ssp.transport`:
 
-- Add :func:`.prepare_method_B` and make this the default1 (:pull:`259`).
+- Add :func:`.prepare_method_B` and make this the default (:pull:`259`).
 - Add :func:`~.ssp.transport.process_df` (:pull:`303`).
+- Adapt to revised ‘variable’ codes (:pull:`309`, :issue:`304`).
 
 Transport
 ---------

message_ix_models/data/test/report/SSP_LED_v2.3.1_baseline.csv

@@ -3,11 +3,11 @@
 import logging
 import re
 from collections.abc import Hashable
-from typing import TYPE_CHECKING
+from functools import cache
+from typing import TYPE_CHECKING, Literal, Optional
 
 import genno
 import pandas as pd
-import xarray as xr
 from genno import Key
 from genno.core.key import single_key
 
@@ -28,9 +28,10 @@
 #: Dimensions of several quantities.
 DIMS = "e n t y UNIT".split()
 
-#: Expression for IAMC ‘variable’ names used in :func:`main`.
-EXPR_EMI = r"^Emissions\|(?P<e>[^\|]+)\|Energy\|Demand\|Transportation(?:\|(?P<t>.*))?$"
-EXPR_FE = r"^Final Energy\|Transportation\|(?P<c>Liquids\|Oil)$"
+EXPR_EMI = re.compile(
+    r"^Emissions\|(?P<e>[^\|]+)\|Energy\|Demand\|(?P<t>(Bunkers|Transportation).*)$"
+)
+EXPR_FE = re.compile(r"^Final Energy\|Transportation\|(?P<c>Liquids\|Oil)$")
 
 #: :class:`.IEA_EWEB` flow codes used in the current file.
 FLOWS = ["AVBUNK", "DOMESAIR", "TOTTRANS"]
@@ -63,27 +64,53 @@ def aviation_share(ref: "TQuantity") -> "TQuantity":
     )
 
 
-def broadcast_t(include_international: bool) -> "AnyQuantity":
+def broadcast_t(version: Literal[1, 2], include_international: bool) -> "AnyQuantity":
     """Quantity to re-add the |t| dimension.
 
     Parameters
     ----------
-    include_international
-        If :any:`True`, include "Aviation|International" with magnitude 1.0. Otherwise,
-        omit
+    version :
+        Version of ‘variable’ names supported by the current module.
+    include_international :
+        If :any:`True`, include "Transportation|Aviation|International" with magnitude
+        1.0. Otherwise, omit.
 
     Return
     ------
     genno.Quantity
-        with dimension "t" and the values:
+        with dimension "t".
 
-        - +1.0 for t="Aviation", a label with missing data.
-        - -1.0 for t="Road Rail and Domestic Shipping", a label with existing data from
-          which the aviation total should be subtracted.
+        If :py:`version=1`, the values include:
+
+        - +1.0 for t="Transportation|Aviation", a label with missing data.
+        - -1.0 for t="Transportation|Road Rail and Domestic Shipping", a label with
+          existing data from which the aviation total must be subtracted.
+
+        If :py:`version=2`, the values include:
+
+        - +1.0 for t="Bunkers" and t="Bunkers|International Aviation", labels with zeros
+          in the input data file.
+        - -1.0 for t="Transportation" and t="Transportation|Road Rail and Domestic
+          Shipping", labels with existing data from which the aviation total must be
+          subtracted.
     """
-    value = [1, -1, 1]
-    t = ["Aviation", "Road Rail and Domestic Shipping", "Aviation|International"]
-    idx = slice(None) if include_international else slice(-1)
+    if version == 1:
+        value = [1, -1, 1]
+        t = [
+            "Transportation|Aviation",
+            "Transportation|Road Rail and Domestic Shipping",
+            "Transportation|Aviation|International",
+        ]
+        idx = slice(None) if include_international else slice(-1)
+    elif version == 2:
+        value = [1, 1, -1, -1]
+        t = [
+            "Bunkers",
+            "Bunkers|International Aviation",
+            "Transportation",
+            "Transportation|Road Rail and Domestic Shipping",
+        ]
+        idx = slice(None)
 
     return genno.Quantity(value[idx], coords={"t": t[idx]})
 
@@ -119,59 +146,6 @@ def e_UNIT(cl_emission: "sdmx.model.common.Codelist") -> "AnyQuantity":
     )
 
 
-def extract_dims(
-    qty: "TQuantity", dim_expr: dict, *, drop: bool = True, fillna: str = "_T"
-) -> "TQuantity":
-    """Extract dimensions from IAMC-like ‘variable’ names using regular expressions."""
-    dims = list(qty.dims)
-
-    dfs = [qty.to_series().rename("value").reset_index()]
-    for dim, expr in dim_expr.items():
-        pattern = re.compile(expr)
-        dfs.append(dfs[0][dim].str.extract(pattern).fillna(fillna))
-        dims.extend(pattern.groupindex)
-        if drop:
-            dims.remove(dim)
-
-    return genno.Quantity(pd.concat(dfs, axis=1).set_index(dims)["value"])
-
-
-def extract_dims1(qty: "TQuantity", dim: dict) -> "TQuantity":  # pragma: no cover
-    """Extract dimensions from IAMC-like ‘variable’ names expressions.
-
-    .. note:: This incomplete, non-working version of :func:`extract_dims` uses
-       :mod:`xarray` semantics.
-    """
-    from collections import defaultdict
-
-    result = qty
-    for d0, expr in dim.items():
-        d0_new = f"{d0}_new"
-        pattern = re.compile(expr)
-
-        indexers: dict[Hashable, list[Hashable]] = {g: [] for g in pattern.groupindex}
-        indexers[d0_new] = []
-
-        coords = qty.coords[d0].data.astype(str)
-        for coord in coords:
-            if match := pattern.match(coord):
-                groupdict = match.groupdict()
-                coord_new = coord[match.span()[1] :]
-            else:
-                groupdict = defaultdict(None)
-                coord_new = coord
-
-            for g in pattern.groupindex:
-                indexers[g].append(groupdict[g])
-            indexers[d0_new].append(coord_new)
-
-        for d1, labels in indexers.items():
-            i2 = {d0: xr.DataArray(coords, coords={d1: labels})}
-            result = result.sel(i2)
-
-    return result
-
-
 def finalize(
     q_all: "TQuantity", q_update: "TQuantity", model_name: str, scenario_name: str
 ) -> pd.DataFrame:
@@ -210,9 +184,7 @@ def _expand(qty):
         .to_frame()
         .reset_index()
         .assign(
-            Variable=lambda df: (
-                "Emissions|" + df["e"] + "|Energy|Demand|Transportation|" + df["t"]
-            ).str.replace("|_T", ""),
+            Variable=lambda df: "Emissions|" + df["e"] + "|Energy|Demand|" + df["t"]
         )
         .drop(["e", "t"], axis=1)
         .set_index(s_all.index.names)[0]
@@ -239,17 +211,19 @@ def prepare_computer(c: "Computer", k_input: Key, method: str) -> "KeyLike":
     str
         "target". Calling :py:`c.get("target")` triggers the calculation.
     """
+    c.require_compat("message_ix_models.report.operator")
+
     # Common structure and utility quantities used by prepare_method_[AB]
-    c.add(f"broadcast:t:{L}", broadcast_t, include_international=method == "A")
+    c.add(
+        f"broadcast:t:{L}", broadcast_t, version=2, include_international=method == "A"
+    )
 
-    k_emi_in, e_t = Key(L, DIMS, "input"), tuple("et")
+    k_emi_in = Key(L, DIMS, "input")
 
     # Select and transform data matching EXPR_EMI
-    # Filter on "VARIABLE"
-    c.add(k_emi_in[0] / e_t, select_re, k_input, indexers={"VARIABLE": EXPR_EMI})
-    # Extract the "e" and "t" dimensions from "VARIABLE"
-    c.add(k_emi_in[1], extract_dims, k_emi_in[0] / e_t, dim_expr={"VARIABLE": EXPR_EMI})
-    c.add(k_emi_in[2], "assign_units", k_emi_in[1], units="Mt/year")
+    # Filter on "VARIABLE", expand the (e, t) dimensions from "VARIABLE"
+    c.add(k_emi_in[0], "select_expand", k_input, dim_cb={"VARIABLE": v_to_emi_coords})
+    c.add(k_emi_in[1], "assign_units", k_emi_in[0], units="Mt/year")
 
     # Call a function to prepare the remaining calculations
     prepare_func = {"A": prepare_method_A, "B": prepare_method_B}[method]
@@ -318,11 +292,8 @@ def prepare_method_B(
        is, final energy use by aviation.
     6. Load emissions intensity of aviation final energy use from the file
        :ref:`transport-input-emi-intensity`.
-    7. Multiply (4) × (5) × (6) to compute the estimate of
-       ``Emissions|*|Energy|Demand|Transportation|Aviation``.
-    8. Estimate an additive adjustment to
-       ``Emissions|*|Energy|Demand|Transportation|Road Rail and Domestic Shipping`` as
-       the negative of (7).
+    7. Multiply (4) × (5) × (6) to compute the estimate of aviation emissions.
+    8. Estimate adjustments according to :func:`broadcast_t`.
     9. Adjust `k_emi_in` by adding (7) and (8).
     """
     from message_ix_models.model.transport import build
@@ -379,14 +350,11 @@ def prepare_method_B(
 
     ### Prepare data from the input data file: total transport consumption of light oil
 
-    # Filter on "VARIABLE"
-    c.add(k_fe_in[0] / "c", select_re, k_input, indexers={"VARIABLE": EXPR_FE})
-
-    # Extract the "e" dimensions from "VARIABLE"
-    c.add(k_fe_in[1], extract_dims, k_fe_in[0] / "c", dim_expr={"VARIABLE": EXPR_FE})
+    # Filter on "VARIABLE", extract (e) dimension
+    c.add(k_fe_in[0], "select_expand", k_input, dim_cb={"VARIABLE": v_to_fe_coords})
 
     # Convert "UNIT" dim labels to Quantity.units
-    c.add(k_fe_in[2] / "UNIT", "unique_units_from_dim", k_fe_in[1], dim="UNIT")
+    c.add(k_fe_in[1] / "UNIT", "unique_units_from_dim", k_fe_in[0], dim="UNIT")
 
     # Relabel:
     # - c[ommodity]: 'Liquids|Oil' (IAMC 'variable' component) → 'lightoil'
@@ -395,7 +363,7 @@ def prepare_method_B(
         c={"Liquids|Oil": "lightoil"},
         n={n.id.partition("_")[2]: n.id for n in get_codelist("node/R12")},
     )
-    c.add(k_fe_in[3] / "UNIT", "relabel", k_fe_in[2] / "UNIT", labels=labels)
+    c.add(k_fe_in[2] / "UNIT", "relabel", k_fe_in[1] / "UNIT", labels=labels)
 
     ### Compute estimate of emissions
     # Product of aviation share and FE of total transport → FE of aviation
@@ -500,11 +468,19 @@ def process_file(
     c.get("target").to_csv(path_out, index=False)
 
 
-def select_re(qty: "AnyQuantity", indexers: dict) -> "AnyQuantity":
-    """Select from `qty` using regular expressions for each dimension."""
-    new_indexers = dict()
-    for dim, expr in indexers.items():
-        new_indexers[dim] = list(
-            map(str, filter(re.compile(expr).match, qty.coords[dim].data.astype(str)))
-        )
-    return qty.sel(new_indexers)
+@cache
+def v_to_fe_coords(value: Hashable) -> Optional[dict[str, str]]:
+    """Match ‘variable’ names used in :func:`main`."""
+    if match := EXPR_FE.fullmatch(str(value)):
+        return match.groupdict()
+    else:
+        return None
+
+
+@cache
+def v_to_emi_coords(value: Hashable) -> Optional[dict[str, str]]:
+    """Match ‘variable’ names used in :func:`main`."""
+    if match := EXPR_EMI.fullmatch(str(value)):
+        return match.groupdict()
+    else:
+        return None