Collect water module settings into Config dataclass

doc/water/index.rst

@@ -96,6 +96,9 @@ Build and run
 .. automodule:: message_ix_models.model.water.build
    :members:
 
+.. automodule:: message_ix_models.model.water.config
+   :members:
+
 Data preparation
 ----------------
 

doc/whatsnew.rst

@@ -6,6 +6,7 @@ Next release
 
 - Add IAMC code list :class:`~.iamc.structure.CL_SCENARIO_DIAGNOSTIC` (:pull:`501`).
 - New module :ref:`tools-newclimate` (:pull:`499`).
+- Add :class:`.model.water.Config` to collect water module settings (:pull:`509`).
 - Add :doc:`/api/model-bmt` (:pull:`433`).
 
   - Add

message_ix_models/model/water/__init__.py

@@ -1,4 +1,5 @@
+from .config import Config
 from .data import demands, water_supply
 from .utils import read_config
 
-__all__ = ["demands", "read_config", "water_supply"]
+__all__ = ["Config", "demands", "read_config", "water_supply"]

message_ix_models/model/water/build.py

@@ -10,6 +10,7 @@
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util import broadcast, package_data_path
 
+from .config import Config
 from .utils import filter_basins_by_region, read_config
 
 log = logging.getLogger(__name__)
@@ -184,9 +185,10 @@ def cat_tec_cooling_calib(
             - 'tec': Name of the cooling technology.
         - regions_df: A list of unique region nodes from the scenario.
     """
+    cfg = Config.from_context(context)
     FILE1 = (
         "cooltech_cost_and_shares_"
-        + (f"ssp_msg_{context.regions}" if context.type_reg == "global" else "country")
+        + (f"ssp_msg_{context.regions}" if cfg.type_reg == "global" else "country")
         + ".csv"
     )
     path1 = package_data_path("water", "ppl_cooling_tech", FILE1)
@@ -288,7 +290,9 @@ def get_spec(context: Context) -> Mapping[str, ScenarioInfo]:
         # Elements to add
         add.set[set_name].extend(config.get("add", []))
 
-    if context.nexus_set == "nexus":
+    cfg = Config.from_context(context)
+
+    if cfg.nexus_set == "nexus":
         # Merge technology.yaml with set.yaml
         context["water set"]["nexus"]["technology"]["add"] = context[
             "water technology"
@@ -311,7 +315,7 @@ def get_spec(context: Context) -> Mapping[str, ScenarioInfo]:
 
         # Share commodity for groundwater
         results = {}
-        df_node = context.all_nodes
+        df_node = cfg.all_nodes
         n = len(df_node.values)
 
         d = {
@@ -550,6 +554,7 @@ def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
 
     # define an empty dictionary
     results = {}
+    cfg = Config.from_context(context)
     # read csv file for basin names and region mapping
     # reading basin_delineation
     FILE = f"basins_by_region_simpl_{context.regions}.csv"
@@ -564,8 +569,8 @@ def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
     df["node"] = "B" + df["BCU_name"].astype(str)
     df["mode"] = "M" + df["BCU_name"].astype(str)
     df["region"] = (
-        context.map_ISO_c[context.regions]
-        if context.type_reg == "country"
+        cfg.map_ISO_c[context.regions]
+        if cfg.type_reg == "country"
         else f"{context.regions}_" + df["REGION"].astype(str)
     )
 
@@ -580,9 +585,9 @@ def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
 
     results["map_node"] = nodes
 
-    context.all_nodes = df["node"]
+    cfg.all_nodes = df["node"]
     # Store the filtered basin names for use in other functions
-    context.valid_basins = set(df["BCU_name"].astype(str))
+    cfg.valid_basins = set(df["BCU_name"].astype(str))
 
     for set_name, config in results.items():
         # Sets to add
@@ -604,7 +609,7 @@ def main(context: Context, scenario, **options):
 
     log.info("Set up MESSAGEix-Nexus")
 
-    if context.nexus_set == "nexus":
+    if Config.from_context(context).nexus_set == "nexus":
         # Add water balance
         spec = map_basin(context)
 

message_ix_models/model/water/cli.py

@@ -8,6 +8,8 @@
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util.click import common_params, scenario_param
 
+from .config import Config
+
 if TYPE_CHECKING:
     from message_ix import Scenario
 log = logging.getLogger(__name__)
@@ -39,6 +41,8 @@ def water_ini(context: "Context", regions, time):
 
     from .utils import read_config
 
+    config = Config.from_context(context)
+
     # Ensure water model configuration is loaded
     read_config(context)
     if not context.scenario_info:
@@ -61,34 +65,34 @@ def water_ini(context: "Context", regions, time):
         regions = "R12"
     # add an attribute to distinguish country models
     if regions in ["R11", "R12", "R14", "R32", "RCP"]:
-        context.type_reg = "global"
+        config.type_reg = "global"
     else:
-        context.type_reg = "country"
+        config.type_reg = "country"
     context.regions = regions
 
     # create a mapping ISO code :
     # a region name, for other scripts
     # only needed for 1-country models
     n_codes = get_codes(f"node/{context.regions}")
     nodes = list(map(str, n_codes[n_codes.index(Code(id="World"))].child))
-    if context.type_reg == "country":
+    if config.type_reg == "country":
         map_ISO_c = {context.regions: nodes[0]}
-        context.map_ISO_c = map_ISO_c
-        log.info(f"mapping {context.map_ISO_c[context.regions]}")
+        config.map_ISO_c = map_ISO_c
+        log.info(f"mapping {config.map_ISO_c[context.regions]}")
 
     # deinfe the timestep
     if not time:
         sc_ref = context.get_scenario()
         time = sc_ref.set("time")
         sub_time = list(time[time != "year"])
         if len(sub_time) == 0:
-            context.time = ["year"]
+            config.time = ["year"]
 
         else:
-            context.time = sub_time
+            config.time = sub_time
     else:
-        context.time = [time]
-    log.info(f"Using the following time-step for the water module: {context.time}")
+        config.time = [time]
+    log.info(f"Using the following time-step for the water module: {config.time}")
 
     # setting the time information in context
 
@@ -163,12 +167,11 @@ def nexus_cli(
     water balance linking different water demands to supply.
     """
     # Set basin filtering configuration on context
-    context.reduced_basin = reduced_basin
-    if filter_list:
-        context.filter_list = list(filter_list)
-    if num_basins is not None:
-        context.num_basins = num_basins
-    context.basin_selection = basin_selection
+    config = Config.from_context(context)
+    config.reduced_basin = reduced_basin
+    config.filter_list = list(filter_list or [])
+    config.num_basins = num_basins
+    config.basin_selection = basin_selection
 
     nexus(context, regions, rcps, sdgs, rels, macro)
 
@@ -193,17 +196,18 @@ def nexus(context: "Context", regions, rcps, sdgs, rels, macro=False):
         Specifies the reliability of hydrological data ['low','mid','high']
     """
     # add input information to the class context
-    context.nexus_set = "nexus"
+    config = Config.from_context(context)
+    config.nexus_set = "nexus"
     if not context.regions:
         context.regions = regions
 
-    context.RCP = rcps
-    context.SDG = sdgs
-    context.REL = rels
+    config.RCP = rcps
+    config.SDG = sdgs
+    config.REL = rels
 
     log.info(
-        f"SSP assumption is {context.ssp}. SDG is {context.SDG}. "
-        f"RCP is {context.RCP}. REL is {context.REL}."
+        f"SSP assumption is {context.ssp}. SDG is {config.SDG}. "
+        f"RCP is {config.RCP}. REL is {config.REL}."
     )
 
     from .build import main as build
@@ -298,12 +302,13 @@ def cooling(
         Specifies the climate scenario used ['no_climate','6p0','2p6']
 
     """
-    context.nexus_set = "cooling"
-    context.RCP = rcps
-    context.REL = rels
+    config = Config.from_context(context)
+    config.nexus_set = "cooling"
+    config.RCP = rcps
+    config.REL = rels
 
     log.info(
-        f"SSP assumption is {context.ssp}. RCP is {context.RCP}. REL is {context.REL}."
+        f"SSP assumption is {context.ssp}. RCP is {config.RCP}. REL is {config.REL}."
     )
 
     from .build import main as build

message_ix_models/model/water/config.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal
+
+from message_ix_models.util.config import ConfigHelper
+
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
+
+@dataclass
+class Config(ConfigHelper):
+    """Settings for :mod:`message_ix_models.model.water`."""
+
+    #: Water build variant.
+    nexus_set: Literal["cooling", "nexus"] = "nexus"
+
+    #: Climate scenario used for water availability and cooling data.
+    RCP: str = "no_climate"
+
+    #: Water SDG policy setting.
+    SDG: str = "baseline"
+
+    #: Hydrological-data reliability setting.
+    REL: str = "low"
+
+    #: Time slices handled by the water module.
+    time: list[str] = field(default_factory=lambda: ["year"])
+
+    #: Whether :attr:`regions` names a global node codelist or a single country.
+    type_reg: Literal["country", "global"] = "global"
+
+    #: Single-country model mapping from region code to ISO code.
+    map_ISO_c: dict[str, str] = field(default_factory=dict)
+
+    #: Enable basin filtering.
+    reduced_basin: bool = False
+
+    #: Basins to add to the automatic reduced-basin selection.
+    filter_list: list[str] = field(default_factory=list)
+
+    #: Number of basins per region to keep when :attr:`reduced_basin` is enabled.
+    num_basins: int | None = None
+
+    #: Automatic reduced-basin selection method.
+    basin_selection: Literal["first_k", "stress"] = "first_k"
+
+    #: Basin names retained after optional filtering.
+    valid_basins: set[str] = field(default_factory=set)
+
+    #: Water-module basin node labels, populated during structural mapping.
+    all_nodes: Any = None
+
+    @classmethod
+    def from_context(cls, context: "Context") -> "Config":
+        """Return the shared water configuration for `context`.
+
+        Create `context.water` if missing, or convert a mapping to `Config`.
+        Repeated calls return the same `Config` instance.
+        """
+        if "water" not in context:
+            context["water"] = cls()
+        elif isinstance(context["water"], dict):
+            context["water"] = cls.from_dict(context["water"])
+
+        return context["water"]

message_ix_models/model/water/data/__init__.py

@@ -9,6 +9,7 @@
 import pandas as pd
 
 from message_ix_models import ScenarioInfo
+from message_ix_models.model.water.config import Config
 from message_ix_models.util import add_par_data
 
 from .demands import add_irrigation_demand, add_sectoral_demands, add_water_availability
@@ -57,12 +58,13 @@ def add_data(scenario, context: "Context", dry_run=False):
 
     info = ScenarioInfo(scenario)
     context["water build info"] = info
+    cfg = Config.from_context(context)
 
     data_funcs: list[DataFunc] = (
         [add_water_supply, cool_tech, non_cooling_tec]
-        if context.nexus_set == "cooling"
+        if cfg.nexus_set == "cooling"
         else DATA_FUNCTIONS
-        if context.type_reg == "global"
+        if cfg.type_reg == "global"
         else DATA_FUNCTIONS_COUNTRY
     )