Update the nexus module

message_ix_models/data/water/set.yaml

@@ -100,6 +100,7 @@ nexus:
     - all
     add:
     - Mf # efficient mode
+    - M1 # required for bare_res test
 
   technology:
     remove:
@@ -347,7 +348,7 @@ nexus:
 
   map_tec_addon:
     add:
-    - [bio_hpl,cooling__bio_hpl]
+    - [bio_hpl, cooling__bio_hpl]
     - [bio_istig, cooling__bio_istig]
     - [bio_istig_ccs, cooling__bio_istig_ccs]
     - [bio_ppl,cooling__bio_ppl]
@@ -478,15 +479,16 @@ nexus:
     - share_low_lim_GWat_share
     - share_wat_recycle_total
     - share_wat_recycle_share
+    - water_resource_extraction # required for bare_res test
 
   cat_tec:
     add:
     - [share_low_lim_GWat_total, extract_surfacewater]
     - [share_low_lim_GWat_total, extract_groundwater]
     - [share_low_lim_GWat_share, extract_groundwater]
-    - [share_wat_recycle_total,urban_recycle]
-    - [share_wat_recycle_total,urban_discharge]
-    - [share_wat_recycle_share,urban_recycle]
+    - [share_wat_recycle_total, urban_recycle]
+    - [share_wat_recycle_total, urban_discharge]
+    - [share_wat_recycle_share, urban_recycle]
     - [water_distribution, urban_t_d]
     - [water_distribution, rural_t_d]
     - [water_distribution, urban_unconnected]

message_ix_models/model/water/__init__.py

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

message_ix_models/model/water/build.py

@@ -5,7 +5,7 @@
 import pandas as pd
 from sdmx.model.v21 import Code
 
-from message_ix_models import ScenarioInfo
+from message_ix_models import Context, ScenarioInfo
 from message_ix_models.model import build
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util import package_data_path
@@ -15,7 +15,7 @@
 log = logging.getLogger(__name__)
 
 
-def get_spec(context) -> Mapping[str, ScenarioInfo]:
+def get_spec(context: Context) -> Mapping[str, ScenarioInfo]:
     """Return the specification for nexus implementation
 
     Parameters
@@ -94,7 +94,7 @@ def get_spec(context) -> Mapping[str, ScenarioInfo]:
 
         df_share2 = pd.DataFrame(data=d2)
 
-        df_share = df_share.append(df_share2)
+        df_share = pd.concat([df_share, df_share2])
         df_list = df_share.values.tolist()
 
         results["map_shares_commodity_total"] = df_list
@@ -144,7 +144,7 @@ def get_spec(context) -> Mapping[str, ScenarioInfo]:
 
         df_share2 = pd.DataFrame(data=d2)
 
-        df_share = df_share.append(df_share2)
+        df_share = pd.concat([df_share, df_share2])
         df_list = df_share.values.tolist()
 
         results["map_shares_commodity_total"] = df_list
@@ -171,8 +171,6 @@ def get_spec(context) -> Mapping[str, ScenarioInfo]:
 
     return dict(require=require, remove=remove, add=add)
 
-    return dict(require=require, remove=remove, add=add)
-
 
 @lru_cache()
 def generate_set_elements(set_name, match=None):
@@ -190,7 +188,7 @@ def generate_set_elements(set_name, match=None):
     return results
 
 
-def map_basin(context) -> Mapping[str, ScenarioInfo]:
+def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
     """Return specification for mapping basins to regions
 
     The basins are spatially consolidated from HydroSHEDS basins delineation
@@ -217,10 +215,11 @@ def map_basin(context) -> Mapping[str, ScenarioInfo]:
     # Assigning proper nomenclature
     df["node"] = "B" + df["BCU_name"].astype(str)
     df["mode"] = "M" + df["BCU_name"].astype(str)
-    if context.type_reg == "country":
-        df["region"] = context.map_ISO_c[context.regions]
-    else:
-        df["region"] = f"{context.regions}_" + df["REGION"].astype(str)
+    df["region"] = (
+        context.map_ISO_c[context.regions]
+        if context.type_reg == "country"
+        else f"{context.regions}_" + df["REGION"].astype(str)
+    )
 
     results["node"] = df["node"]
     results["mode"] = df["mode"]
@@ -236,13 +235,13 @@ def map_basin(context) -> Mapping[str, ScenarioInfo]:
     context.all_nodes = df["node"]
 
     for set_name, config in results.items():
-        # Sets  to add
+        # Sets to add
         add.set[set_name].extend(config)
 
     return dict(require=require, remove=remove, add=add)
 
 
-def main(context, scenario, **options):
+def main(context: Context, scenario, **options):
     """Set up MESSAGEix-Nexus on `scenario`.
 
     See also

message_ix_models/model/water/cli.py

@@ -2,6 +2,7 @@
 
 import click
 
+from message_ix_models import Context
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util.click import common_params
 
@@ -13,12 +14,12 @@
 @common_params("regions")
 @click.option("--time", help="Manually defined time")
 @click.pass_obj
-def cli(context, regions, time):
+def cli(context: "Context", regions, time):
     """MESSAGEix-Water and Nexus variant."""
     water_ini(context, regions, time)
 
 
-def water_ini(context, regions, time):
+def water_ini(context: "Context", regions, time):
     """Add components of the MESSAGEix-Nexus module
 
     This function modifies model name & scenario name
@@ -61,7 +62,7 @@
     context.regions = regions
 
     # create a mapping ISO code :
-    # region name, for other scripts
+    # a region name, for other scripts
     # only needed for 1-country models
     nodes = get_codes(f"node/{context.regions}")
     nodes = list(map(str, nodes[nodes.index("World")].child))
@@ -106,7 +107,7 @@
     help="Defines whether the model solves with macro",
 )
 @common_params("regions")
-def nexus_cli(context, regions, rcps, sdgs, rels, macro=False):
+def nexus_cli(context: "Context", regions, rcps, sdgs, rels, macro=False):
     """
     Add basin structure connected to the energy sector and
     water balance linking different water demands to supply.
@@ -115,7 +116,7 @@
     nexus(context, regions, rcps, sdgs, rels, macro)
 
 
-def nexus(context, regions, rcps, sdgs, rels, macro=False):
+def nexus(context: "Context", regions, rcps, sdgs, rels, macro=False):
     """Add basin structure connected to the energy sector and
     water balance linking different water demands to supply.
 
@@ -259,7 +260,7 @@
     help="Default running legacy and water (full) otherwise only water, if specified",
 )
 @common_params("output_model")
-def report_cli(context, output_model, sdgs, water=False):
+def report_cli(context: "Context", output_model, sdgs, water=False):
     """function to run the water report_full from cli to the
     scenario defined by the user with --url
 
@@ -272,13 +273,13 @@
     SDG : Str
         Defines if and what water SDG measures are activated
     """
-    reg = context.regions
+    reg = context.model.regions
     sc = context.get_scenario()
     if water:
-        from message_ix_models.model.water.reporting import report
+        from message_ix_models.model.water.report import report
 
         report(sc, reg, sdgs)
     else:
-        from message_ix_models.model.water.reporting import report_full
+        from message_ix_models.model.water.report import report_full
 
         report_full(sc, reg, sdgs)

message_ix_models/model/water/data/__init__.py

@@ -1,6 +1,7 @@
 """Generate input data."""
 
 import logging
+from typing import TYPE_CHECKING
 
 from message_ix_models import ScenarioInfo
 from message_ix_models.util import add_par_data
@@ -11,6 +12,9 @@
 from .water_for_ppl import cool_tech, non_cooling_tec
 from .water_supply import add_e_flow, add_water_supply
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 log = logging.getLogger(__name__)
 
 DATA_FUNCTIONS = [
@@ -40,7 +44,7 @@
 ]
 
 
-def add_data(scenario, context, dry_run=False):
+def add_data(scenario, context: "Context", dry_run=False):
     """Populate `scenario` with MESSAGEix-Nexus data."""
 
     info = ScenarioInfo(scenario)