Add basin filtering

Use basin filtering everywhere

More filtering

More filtering

Fix accidental removal

Add basin reduction options

Clean up filtering code

Add min 1 basin region pair req

Add reg_to_basin return tech

consolidate return logic

Fix failing tests

Tests failure due to missing valid_basins list. Adding list to context
to fix issue.

Author: Wegatriespython

message_ix_models/model/water/build.py

@@ -11,7 +11,7 @@
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util import broadcast, package_data_path
 
-from .utils import read_config
+from .utils import filter_basins_by_region, read_config
 
 log = logging.getLogger(__name__)
 
@@ -557,6 +557,10 @@ def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
     PATH = package_data_path("water", "delineation", FILE)
 
     df = pd.read_csv(PATH)
+
+    # Apply basin filter to reduce number of basins per region
+    df = filter_basins_by_region(df, context)
+
     # Assigning proper nomenclature
     df["node"] = "B" + df["BCU_name"].astype(str)
     df["mode"] = "M" + df["BCU_name"].astype(str)
@@ -578,6 +582,8 @@ def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
     results["map_node"] = nodes
 
     context.all_nodes = df["node"]
+    # Store the filtered basin names for use in other functions
+    context.valid_basins = set(df["BCU_name"].astype(str))
 
     for set_name, config in results.items():
         # Sets to add

message_ix_models/model/water/cli.py

@@ -119,13 +119,44 @@ def water_ini(context: "Context", regions, time):
     is_flag=True,
     help="Defines whether the model solves with macro",
 )
+@click.option(
+    "--reduced-basin/--no-reduced-basin",
+    default=False,
+    help="Enable reduced basin filtering",
+)
+@click.option(
+    "--filter-list",
+    multiple=True,
+    help="Specific basins to include (can be used multiple times)",
+)
+@click.option(
+    "--num-basins",
+    type=int,
+    help="Number of basins per region to keep when reduced-basin is enabled",
+)
 @common_params("regions")
 @scenario_param("--ssp")
-def nexus_cli(context: "Context", regions, rcps, sdgs, rels, macro=False):
+def nexus_cli(
+    context: "Context",
+    regions,
+    rcps,
+    sdgs,
+    rels,
+    macro=False,
+    reduced_basin=False,
+    filter_list=None,
+    num_basins=None,
+):
     """
     Add basin structure connected to the energy sector and
     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
 
     nexus(context, regions, rcps, sdgs, rels, macro)
 

message_ix_models/model/water/data/demands.py

@@ -218,6 +218,9 @@ def add_sectoral_demands(context: "Context") -> dict[str, pd.DataFrame]:
     df_dmds.sort_values(["year", "node", "variable", "value"], inplace=True)
 
     df_dmds["time"] = "year"
+    
+    # Filter to only include basins that exist after basin filtering
+    df_dmds = df_dmds[df_dmds["node"].isin(context.valid_basins)]
 
     # Write final interpolated values as csv
     # df2_f.to_csv('final_interpolated_values.csv')
@@ -240,6 +243,9 @@ def add_sectoral_demands(context: "Context") -> dict[str, pd.DataFrame]:
         )
         df_m = df_m[["year", "pid", "variable", "value", "month"]]
         df_m.columns = pd.Index(["year", "node", "variable", "value", "time"])
+        
+        # Filter monthly data to only include valid basins
+        df_m = df_m[df_m["node"].isin(context.valid_basins)]
 
         # remove yearly parts from df_dms
         df_dmds = df_dmds[
@@ -769,13 +775,11 @@ def read_water_availability(context: "Context") -> Sequence[pd.DataFrame]:
         "water", "delineation", f"basins_by_region_simpl_{context.regions}.csv"
     )
     df_x = pd.read_csv(PATH)
+    
+    # Filter to only include valid basins
+    df_x = df_x[df_x["BCU_name"].isin(context.valid_basins)]
 
     if "year" in context.time:
-        # path for reading basin delineation file
-        PATH = package_data_path(
-            "water", "delineation", f"basins_by_region_simpl_{context.regions}.csv"
-        )
-        df_x = pd.read_csv(PATH)
         # Adding freshwater supply constraints
         # Reading data, the data is spatially and temprally aggregated from GHMs
         path1 = package_data_path(
@@ -786,6 +790,14 @@ def read_water_availability(context: "Context") -> Sequence[pd.DataFrame]:
         # Read rcp 2.6 data
         df_sw = pd.read_csv(path1)
         df_sw.drop(["Unnamed: 0"], axis=1, inplace=True)
+        
+        # Filter columns to only include valid basins
+        # The columns are years, so we need to filter rows based on the original basin order
+        # First, get the indices of valid basins from the original full list
+        full_basin_df = pd.read_csv(PATH)  # Read full basin list again
+        valid_indices = full_basin_df[full_basin_df["BCU_name"].isin(context.valid_basins)].index
+        df_sw = df_sw.iloc[valid_indices]  # Keep only rows for valid basins
+        df_sw.reset_index(drop=True, inplace=True)
 
         df_sw.index = df_x["BCU_name"].index
         df_sw = df_sw.stack().reset_index()
@@ -811,6 +823,11 @@ def read_water_availability(context: "Context") -> Sequence[pd.DataFrame]:
         # Read groundwater data
         df_gw = pd.read_csv(path1)
         df_gw.drop(["Unnamed: 0"], axis=1, inplace=True)
+        
+        # Filter to only include valid basins (same as df_sw)
+        df_gw = df_gw.iloc[valid_indices]  # Use same valid_indices from above
+        df_gw.reset_index(drop=True, inplace=True)
+        
         df_gw.index = df_x["BCU_name"].index
         df_gw = df_gw.stack().reset_index()
         df_gw.columns = pd.Index(["Region", "years", "value"])
@@ -834,6 +851,12 @@ def read_water_availability(context: "Context") -> Sequence[pd.DataFrame]:
         )
         df_sw = pd.read_csv(path1)
         df_sw.drop(["Unnamed: 0"], axis=1, inplace=True)
+        
+        # Filter to only include valid basins
+        full_basin_df = pd.read_csv(PATH)  # Read full basin list again
+        valid_indices = full_basin_df[full_basin_df["BCU_name"].isin(context.valid_basins)].index
+        df_sw = df_sw.iloc[valid_indices]
+        df_sw.reset_index(drop=True, inplace=True)
 
         df_sw.index = df_x["BCU_name"].index
         df_sw = df_sw.stack().reset_index()
@@ -857,6 +880,10 @@ def read_water_availability(context: "Context") -> Sequence[pd.DataFrame]:
         )
         df_gw = pd.read_csv(path1)
         df_gw.drop(["Unnamed: 0"], axis=1, inplace=True)
+        
+        # Filter to only include valid basins (same as df_sw)
+        df_gw = df_gw.iloc[valid_indices]  # Use same valid_indices from above
+        df_gw.reset_index(drop=True, inplace=True)
 
         df_gw.index = df_x["BCU_name"].index
         df_gw = df_gw.stack().reset_index()

message_ix_models/model/water/data/infrastructure.py

@@ -224,6 +224,10 @@ def add_infrastructure_techs(context: "Context") -> dict[str, pd.DataFrame]:
     PATH = package_data_path("water", "delineation", FILE2)
 
     df_node = pd.read_csv(PATH)
+    
+    # Filter to only valid basins (already filtered in map_basin)
+    df_node = df_node[df_node["BCU_name"].isin(context.valid_basins)]
+    
     # Assigning proper nomenclature
     df_node["node"] = "B" + df_node["BCU_name"].astype(str)
     df_node["mode"] = "M" + df_node["BCU_name"].astype(str)
@@ -825,6 +829,10 @@ def add_desalination(context: "Context") -> dict[str, pd.DataFrame]:
     PATH = package_data_path("water", "delineation", FILE2)
 
     df_node = pd.read_csv(PATH)
+    
+    # Filter to only valid basins (already filtered in map_basin)
+    df_node = df_node[df_node["BCU_name"].isin(context.valid_basins)]
+    
     # Assigning proper nomenclature
     df_node["node"] = "B" + df_node["BCU_name"].astype(str)
     df_node["mode"] = "M" + df_node["BCU_name"].astype(str)
@@ -833,6 +841,11 @@ def add_desalination(context: "Context") -> dict[str, pd.DataFrame]:
         if context.type_reg == "country"
         else f"{context.regions}_" + df_node["REGION"].astype(str)
     )
+    
+    # Filter historical and projected data to only include basins that exist after filtering
+    df_hist = df_hist[df_hist["BCU_name"].isin(context.valid_basins)]
+    df_proj = df_proj[df_proj["BCU_name"].isin(context.valid_basins)]
+    
     # output dataframe linking to desal tech types
     out_df = (
         make_df(

message_ix_models/model/water/data/irrigation.py

@@ -32,6 +32,10 @@ def add_irr_structure(context: "Context") -> dict[str, pd.DataFrame]:
     FILE2 = f"basins_by_region_simpl_{context.regions}.csv"
     PATH = package_data_path("water", "delineation", FILE2)
     df_node = pd.read_csv(PATH)
+    
+    # Filter to only include valid basins
+    df_node = df_node[df_node["BCU_name"].isin(context.valid_basins)]
+    
     # Assigning proper nomenclature
     df_node["node"] = "B" + df_node["BCU_name"].astype(str)
     df_node["mode"] = "M" + df_node["BCU_name"].astype(str)

message_ix_models/model/water/data/water_supply.py

@@ -11,6 +11,7 @@
     KM3_TO_MCM,
     USD_KM3_TO_USD_MCM,
     GWa_KM3_TO_GWa_MCM,
+    filter_basins_by_region,
     get_vintage_and_active_years,
 )
 from message_ix_models.util import (
@@ -42,6 +43,9 @@ def map_basin_region_wat(context: "Context") -> pd.DataFrame:
             "water", "delineation", f"basins_by_region_simpl_{context.regions}.csv"
         )
         df_x = pd.read_csv(PATH)
+        
+        # Filter to only include valid basins
+        df_x = df_x[df_x["BCU_name"].isin(context.valid_basins)]
         # Adding freshwater supply constraints
         # Reading data, the data is spatially and temprally aggregated from GHMs
         path1 = package_data_path(
@@ -94,6 +98,9 @@ def map_basin_region_wat(context: "Context") -> pd.DataFrame:
             "water", "delineation", f"basins_by_region_simpl_{context.regions}.csv"
         )
         df_x = pd.read_csv(PATH)
+        
+        # Filter to only include valid basins
+        df_x = df_x[df_x["BCU_name"].isin(context.valid_basins)]
 
         # Reading data, the data is spatially and temporally aggregated from GHMs
         df_sw["BCU_name"] = df_x["BCU_name"]
@@ -165,6 +172,10 @@ def add_water_supply(context: "Context") -> dict[str, pd.DataFrame]:
     PATH = package_data_path("water", "delineation", FILE)
 
     df_node = pd.read_csv(PATH)
+    
+    # Apply basin filter to reduce number of basins per region
+    df_node = filter_basins_by_region(df_node, context)
+    
     # Assigning proper nomenclature
     df_node["node"] = "B" + df_node["BCU_name"].astype(str)
     df_node["mode"] = "M" + df_node["BCU_name"].astype(str)
@@ -191,6 +202,10 @@ def add_water_supply(context: "Context") -> dict[str, pd.DataFrame]:
     FILE2 = f"historical_new_cap_gw_sw_km3_year_{context.regions}.csv"
     PATH2 = package_data_path("water", "availability", FILE2)
     df_hist = pd.read_csv(PATH2)
+    
+    # Filter to only include valid basins
+    df_hist = df_hist[df_hist["BCU_name"].isin(context.valid_basins)]
+    
     df_hist["BCU_name"] = "B" + df_hist["BCU_name"].astype(str)
 
     if context.nexus_set == "cooling":