Update .report.sim for iiasa/ixmp#500, iiasa/message_ix#761

- Import RENAME_DIMS through a function that provides backwards compatibility. - Use MESSAGE.items and MACRO.items; remove redefinitions. - Simplify simulate_qty() to accept a list of dimensions directly.
iiasa · Nov 22, 2023 · 17be421 · 17be421
1 parent de30f64
commit 17be421
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 152 deletions.
diff --git a/message_ix_models/report/sim.py b/message_ix_models/report/sim.py
@@ -4,17 +4,21 @@
 from collections.abc import Mapping
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
 
 import pandas as pd
 from dask.core import quote
 from genno import Key, KeyExistsError, Quantity, configure
 from message_ix import Reporter
-from message_ix.models import MESSAGE_ITEMS, item
+from message_ix.models import MACRO, MESSAGE
 from pandas.api.types import is_scalar
 
 from message_ix_models import ScenarioInfo
 from message_ix_models.util._logging import mark_time, silence_log
+from message_ix_models.util.ixmp import rename_dims
+
+if TYPE_CHECKING:
+    from message_ix.models import Item
 
 __all__ = [
     "SIMULATE_ITEMS",
@@ -25,138 +29,10 @@
 
 log = logging.getLogger(__name__)
 
-# Copied and expanded from message_ix.models.MESSAGE_ITEMS, where these entries are
-# commented because of JDBCBackend limitations.
-# TODO read from that location once possible
-MESSAGE_VARS = {
-    # Activity
-    "ACT": item("var", "nl t yv ya m h"),
-    # "ACTIVITY_BOUND_ALL_MODES_LO": item("var", ""),
-    # "ACTIVITY_BOUND_ALL_MODES_UP": item("var", ""),
-    # "ACTIVITY_BOUND_LO": item("var", ""),
-    # "ACTIVITY_BOUND_UP": item("var", ""),
-    # "ACTIVITY_BY_RATING": item("var", ""),
-    # "ACTIVITY_CONSTRAINT_LO": item("var", ""),
-    # "ACTIVITY_CONSTRAINT_UP": item("var", ""),
-    # "ACTIVITY_RATING_TOTAL": item("var", ""),
-    # "ACTIVITY_SOFT_CONSTRAINT_LO": item("var", ""),
-    # "ACTIVITY_SOFT_CONSTRAINT_UP": item("var", ""),
-    # "ACT_LO": item("var", ""),
-    # "ACT_RATING": item("var", ""),
-    # "ACT_UP": item("var", ""),
-    # "ADDON_ACTIVITY_LO": item("var", ""),
-    # "ADDON_ACTIVITY_UP": item("var", ""),
-    # Maintained capacity
-    "CAP": item("var", "nl t yv ya"),
-    # "CAPACITY_CONSTRAINT": item("var", ""),
-    # "CAPACITY_MAINTENANCE": item("var", ""),
-    # "CAPACITY_MAINTENANCE_HIST": item("var", ""),
-    # "CAPACITY_MAINTENANCE_NEW": item("var", ""),
-    # "CAP_FIRM": item("var", ""),
-    # New capacity
-    "CAP_NEW": item("var", "nl t yv"),
-    # "CAP_NEW_LO": item("var", ""),
-    # "CAP_NEW_UP": item("var", ""),
-    # "COMMODITY_BALANCE_GT": item("var", ""),
-    # "COMMODITY_BALANCE_LT": item("var", ""),
-    # "COMMODITY_USE": item("var", ""),
-    # "COMMODITY_USE_LEVEL": item("var", ""),
-    # "COST_ACCOUNTING_NODAL": item("var", ""),
-    # "COST_NODAL": item("var", ""),
-    # "COST_NODAL_NET": item("var", ""),
-    # "DEMAND": item("var", ""),
-    # "DYNAMIC_LAND_SCEN_CONSTRAINT_LO": item("var", ""),
-    # "DYNAMIC_LAND_SCEN_CONSTRAINT_UP": item("var", ""),
-    # "DYNAMIC_LAND_TYPE_CONSTRAINT_LO": item("var", ""),
-    # "DYNAMIC_LAND_TYPE_CONSTRAINT_UP": item("var", ""),
-    # Emissions
-    "EMISS": item("var", "n e type_tec y"),
-    # "EMISSION_CONSTRAINT": item("var", ""),
-    # "EMISSION_EQUIVALENCE": item("var", ""),
-    # Extraction
-    "EXT": item("var", "n c g y"),
-    # "EXTRACTION_BOUND_UP": item("var", ""),
-    # "EXTRACTION_EQUIVALENCE": item("var", ""),
-    # "FIRM_CAPACITY_PROVISION": item("var", ""),
-    # "GDP": item("var", ""),
-    # Land scenario share
-    "LAND": item("var", "n land_scenario y"),
-    # "LAND_CONSTRAINT": item("var", ""),
-    # "MIN_UTILIZATION_CONSTRAINT": item("var", ""),
-    # "NEW_CAPACITY_BOUND_LO": item("var", ""),
-    # "NEW_CAPACITY_BOUND_UP": item("var", ""),
-    # "NEW_CAPACITY_CONSTRAINT_LO": item("var", ""),
-    # "NEW_CAPACITY_CONSTRAINT_UP": item("var", ""),
-    # "NEW_CAPACITY_SOFT_CONSTRAINT_LO": item("var", ""),
-    # "NEW_CAPACITY_SOFT_CONSTRAINT_UP": item("var", ""),
-    # Objective (scalar)
-    "OBJ": dict(ix_type="var", idx_sets=[]),
-    # "OBJECTIVE": item("var", ""),
-    # "OPERATION_CONSTRAINT": item("var", ""),
-    # Price of emissions
-    "PRICE_COMMODITY": item("var", "n c l y h"),
-    # Price of emissions
-    "PRICE_EMISSION": item("var", "n e t y"),
-    # Relation (lhs)
-    "REL": item("var", "relation nr yr"),
-    # "RELATION_CONSTRAINT_LO": item("var", ""),
-    # "RELATION_CONSTRAINT_UP": item("var", ""),
-    # "RELATION_EQUIVALENCE": item("var", ""),
-    # "REN": item("var", ""),
-    # "RENEWABLES_CAPACITY_REQUIREMENT": item("var", ""),
-    # "RENEWABLES_EQUIVALENCE": item("var", ""),
-    # "RENEWABLES_POTENTIAL_CONSTRAINT": item("var", ""),
-    # "RESOURCE_CONSTRAINT": item("var", ""),
-    # "RESOURCE_HORIZON": item("var", ""),
-    # "SHARE_CONSTRAINT_COMMODITY_LO": item("var", ""),
-    # "SHARE_CONSTRAINT_COMMODITY_UP": item("var", ""),
-    # "SHARE_CONSTRAINT_MODE_LO": item("var", ""),
-    # "SHARE_CONSTRAINT_MODE_UP": item("var", ""),
-    # "SLACK_ACT_BOUND_LO": item("var", ""),
-    # "SLACK_ACT_BOUND_UP": item("var", ""),
-    # "SLACK_ACT_DYNAMIC_LO": item("var", ""),
-    # "SLACK_ACT_DYNAMIC_UP": item("var", ""),
-    # "SLACK_CAP_NEW_BOUND_LO": item("var", ""),
-    # "SLACK_CAP_NEW_BOUND_UP": item("var", ""),
-    # "SLACK_CAP_NEW_DYNAMIC_LO": item("var", ""),
-    # "SLACK_CAP_NEW_DYNAMIC_UP": item("var", ""),
-    # "SLACK_CAP_TOTAL_BOUND_LO": item("var", ""),
-    # "SLACK_CAP_TOTAL_BOUND_UP": item("var", ""),
-    # "SLACK_COMMODITY_EQUIVALENCE_LO": item("var", ""),
-    # "SLACK_COMMODITY_EQUIVALENCE_UP": item("var", ""),
-    # "SLACK_LAND_SCEN_LO": item("var", ""),
-    # "SLACK_LAND_SCEN_UP": item("var", ""),
-    # "SLACK_LAND_TYPE_LO": item("var", ""),
-    # "SLACK_LAND_TYPE_UP": item("var", ""),
-    # "SLACK_RELATION_BOUND_LO": item("var", ""),
-    # "SLACK_RELATION_BOUND_UP": item("var", ""),
-    # Stock
-    "STOCK": item("var", "n c l y"),
-    # "STOCK_CHG": item("var", ""),
-    # "STOCKS_BALANCE": item("var", ""),
-    # "STORAGE": item("var", ""),
-    # "STORAGE_BALANCE": item("var", ""),
-    # "STORAGE_BALANCE_INIT": item("var", ""),
-    # "STORAGE_CHANGE": item("var", ""),
-    # "STORAGE_CHARGE": item("var", ""),
-    # "STORAGE_INPUT": item("var", ""),
-    # "SYSTEM_FLEXIBILITY_CONSTRAINT": item("var", ""),
-    # "SYSTEM_RELIABILITY_CONSTRAINT": item("var", ""),
-    # "TOTAL_CAPACITY_BOUND_LO": item("var", ""),
-    # "TOTAL_CAPACITY_BOUND_UP": item("var", ""),
-}
-
 # Items to included in a simulated solution: MESSAGE sets and parameters; some variables
-SIMULATE_ITEMS = deepcopy(MESSAGE_ITEMS)
-# Other MESSAGE variables
-SIMULATE_ITEMS.update(MESSAGE_VARS)
+SIMULATE_ITEMS = deepcopy(MESSAGE.items)
 # MACRO variables
-SIMULATE_ITEMS.update(
-    {
-        "GDP": item("var", "n y"),
-        "MERtoPPP": item("var", "n y"),
-    }
-)
+SIMULATE_ITEMS.update({k: MACRO.items[k] for k in ("GDP", "MERtoPPP")})
 
 configure(
     rename_dims=dict(
@@ -166,25 +42,23 @@
 )
 
 
-def dims_of(info: dict) -> Dict[str, str]:
+def dims_of(info: "Item") -> Dict[str, str]:
     """Return a mapping from the full index names to short dimension IDs of `info`."""
-    from ixmp.report import RENAME_DIMS
-
-    return {
-        d: RENAME_DIMS.get(d, d)
-        for d in (info.get("idx_names") or info.get("idx_sets") or [])
-    }
+    return {d: rename_dims().get(d, d) for d in (info.dims or info.coords or [])}
 
 
 def simulate_qty(
-    name: str, info: dict, item_data: Union[dict, pd.DataFrame]
+    name: str, dims: List[str], item_data: Union[dict, pd.DataFrame]
 ) -> Quantity:
-    """Return simulated data for item `name`."""
-    from ixmp.report import RENAME_DIMS
-
-    # Dimensions of the resulting quantity
-    dims = list(dims_of(info).values())
+    """Return simulated data for item `name`.
 
+    Parameters
+    ----------
+    dims :
+        Dimensions of the resulting quantity.
+    item_data :
+        Optional data for the quantity.
+    """
     if isinstance(item_data, dict):
         # NB this is code lightly modified from make_df
 
@@ -211,7 +85,7 @@ def simulate_qty(
         df = pd.DataFrame(**args)
     else:
         # Provided complete data frame
-        df = item_data.rename(columns=RENAME_DIMS)
+        df = item_data.rename(columns=rename_dims())
 
     # Data must be entirely empty, or complete
     assert not df.isna().any().any() or df.isna().all().all(), data
@@ -279,7 +153,7 @@ def add_simulated_solution(
     """
     from importlib.metadata import version
 
-    from ixmp.report import RENAME_DIMS
+    from ixmp.backend import ItemType
 
     if version("message_ix") < "3.6":
         raise NotImplementedError(
@@ -296,7 +170,8 @@ def add_simulated_solution(
     # Add simulated data
     data = data or dict()
     for name, item_info in SIMULATE_ITEMS.items():
-        key = Key(name, list(dims_of(item_info).values()))
+        dims = list(dims_of(item_info).values())
+        key = Key(name, dims)
 
         # Add a task to load data from a file in `path`, if it exists
         try:
@@ -310,10 +185,10 @@ def add_simulated_solution(
             rep.add(key, data_from_file, p, name=name, dims=key.dims, sums=True)
             continue
 
-        if item_info["ix_type"] == "set":
+        if item_info.type == ItemType.SET:
             # Add the set elements from `info`
-            rep.add(RENAME_DIMS.get(name, name), quote(info.set[name]))
-        elif item_info["ix_type"] in ("par", "var"):
+            rep.add(rename_dims().get(name, name), quote(info.set[name]))
+        elif item_info.type in (ItemType.PAR, ItemType.VAR):
             # Retrieve an existing key for `name`
             try:
                 full_key = rep.full_key(name)
@@ -332,7 +207,7 @@ def add_simulated_solution(
                 key,
                 simulate_qty,
                 name=name,
-                info=item_info,
+                dims=dims,
                 item_data=item_data,
                 sums=True,
             )

diff --git a/message_ix_models/util/ixmp.py b/message_ix_models/util/ixmp.py
@@ -0,0 +1,20 @@
+from typing import Dict
+
+
+def rename_dims() -> Dict[str, str]:
+    """Access :data:`.ixmp.report.common.RENAME_DIMS`.
+
+    This provides backwards-compatibility with ixmp versions 3.7.0 and earlier. It can
+    be removed when message-ix-models no longer supports versions of ixmp older than
+    3.8.0.
+    """
+    try:
+        # ixmp 3.8.0 and later
+        import ixmp.report.common
+    except ImportError:
+        # ixmp <= 3.7.0
+        import ixmp.reporting.util  # type: ignore [import-not-found]
+
+        return ixmp.reporting.util.RENAME_DIMS
+    else:
+        return ixmp.report.common.RENAME_DIMS