Add reduced basin functionality

[Wegatriespython]

Reduced basin filtering for local nexus development

Adds basin filtering to enable running the nexus module locally for development.
Reduced basin mode is not intended for production, basin reduction is not
physically realistic. However for development of scenarios, model capabilities or policy design, the reduced basin scenario needs to give a good proxy of the whole model, hence the granular control over its implementation.

Addresses #414. Builds on #405 (now merged).

Design

Basin selection is two-step and compositional:

1. Automatic selection per region: either first_k (head N from CSV order)
   or stress (N basins sampled across the demand/supply ratio spectrum).
2. filter_list augmentation: if specific basins are named, they are added
   to the automatic set (union, not replacement). E.g. automatic gives
   {A1, B1, C1}, filter_list [B1, B2, B3] gives {A1, B1, B2, B3, C1}.

All regions are guaranteed coverage by the automatic step.

New CLI options (on mix-models water-ix nexus)

--reduced-basin / --no-reduced-basin
                                Toggle for A/B testing between full and
                                reduced basin sets.
--num-basins INTEGER            Basins per region (default 3).
--basin-selection [first_k|stress]
                                Automatic selection method (default first_k).
                                stress ranks basins by demand/supply ratio.
--filter-list TEXT              Additional basins on top of automatic
                                selection (repeatable).

Python API

# Toggle for A/B comment one line to switch
ctx.reduced_basin = True

ctx.basin_selection = "stress"   # or "first_k"
ctx.num_basins = 3               # per region
# Additive on automatic set:
ctx.filter_list = ["10|SAS", "115|SAS"]

Files changed

File	What changed
utils.py	filter_basins_by_region() (two-step selection), compute_basin_demand_ratio(), _select_by_stress(), _diversity_select()
cli.py	New click options on nexus_cli
build.py	Applies filter_basins_by_region() in map_basin(), stores context.valid_basins
demands.py, infrastructure.py, irrigation.py, water_supply.py	Filter data to context.valid_basins
conftest.py	setup_valid_basins() helper for tests
test_utils.py	Tests for stress selection, demand ratio, filter_list additive behavior

How to review

• Core logic: filter_basins_by_region() in utils.py - verify the two-step
  selection and the additive union of filter_list.
• Stress selection: compute_basin_demand_ratio() and _select_by_stress() -
  verify demand/supply ranking and diversity sampling.
• Data modules: check that context.valid_basins filtering is applied
  consistently across demands, infrastructure, irrigation, water_supply.
• CLI: verify options propagate correctly to context in cli.py.

PR checklist

• Continuous integration checks all pass
• Add or expand tests; coverage checks pass
• Update doc/whatsnew

[codecov]

Codecov Report

❌ Patch coverage is 91.42857% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.6%. Comparing base (9c5cc24) to head (df6c5dc).
⚠️ Report is 242 commits behind head on main.

Files with missing lines	Patch %	Lines
message_ix_models/model/water/cli.py	45.4%	6 Missing ⚠️
message_ix_models/model/water/data/water_supply.py	68.4%	6 Missing ⚠️
message_ix_models/model/water/utils.py	93.6%	5 Missing ⚠️
message_ix_models/model/water/data/demands.py	93.7%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##            main    #432     +/-   ##
=======================================
+ Coverage   69.5%   69.6%   +0.1%     
=======================================
  Files        300     300             
  Lines      24598   24771    +173     
=======================================
+ Hits       17104   17260    +156     
- Misses      7494    7511     +17     

Files with missing lines	Coverage Δ
message_ix_models/model/water/build.py	74.4% <100.0%> (+0.2%)	⬆️
...ssage_ix_models/model/water/data/infrastructure.py	100.0% <100.0%> (ø)
message_ix_models/model/water/data/irrigation.py	100.0% <100.0%> (ø)
...essage_ix_models/model/water/data/water_for_ppl.py	91.5% <ø> (ø)
message_ix_models/tests/model/water/conftest.py	98.7% <100.0%> (+0.3%)	⬆️
...e_ix_models/tests/model/water/data/test_demands.py	100.0% <100.0%> (ø)
...dels/tests/model/water/data/test_infrastructure.py	100.0% <100.0%> (ø)
...x_models/tests/model/water/data/test_irrigation.py	100.0% <100.0%> (ø)
...odels/tests/model/water/data/test_water_for_ppl.py	97.7% <100.0%> (-0.3%)	⬇️
...models/tests/model/water/data/test_water_supply.py	100.0% <100.0%> (ø)
... and 5 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[adrivinca]

I did not have time to check all the "how to review" points yet.
but a question about the cli options.

would it be possible to simplify a bit, given that this is secondary functionality. and at the same time make sure contradictory option are not possible.

1. is it possible to get rid of --reduced-basin / --no-reduced-basin and have the feature active if one adds one of the other commands?
2. are some of these command contradictory (like filter list and num-basin?), or some that do require others (like basin-selection needs filter list?)? make it clear or think if they could be merged

                                  multiple times)
  --num-basins INTEGER            Number of basins per region to keep when
                                  reduced-basin is enabled
  --basin-selection [first_k|stress]```

[Wegatriespython]

On your two points:

1. --reduced-basin I use this primarily in scripts for A/B testing between full basin set vs reduced. Commenting out one boolean between runs is the toggle, since the basin filter-list is tedious to repeat, and the extra layer prevents accidents.
2. The options were not meant to be contradictory, I have adjusted so that the selection works in two steps:

• Automatic selection runs first (--basin-selection picks the method, --num-basins picks how many per region) and produces a base set like {A1, B1, C1}.
• --filter-list adds basins on top of that base set - e.g. passing [B1, B2, B3] gives the union {A1, B1, B2, B3, C1}.

Updated the PR description to clarify the option relationships and show only the new CLI options.

[adrivinca]

I think this is a useful feature, both for quickly testing new feature on the water module without having extremely long runs, and for application that focus on specific regions (e.g. Eurpoean projects), where we can neglect the energy-water dynamics in other continents and reduce run time.

To me the proposed changes look good. I would ask to please update the doc/water files, in particular index.rst, to update the cli command options and, if possible check and update obsolete text and reference to utility functions.

[Wegatriespython]

@adrivinca @khaeru All outstanding comments are addressed, hopefully once the test suite runs it should be ready!

[adrivinca]

To me it looks good, also thanks for the documentation udates.
Just nee to clarify why tests fail, maybe rebase needed?

[khaeru]

Just nee to clarify why tests fail, maybe rebase needed?

@adrivinca @Wegatriespython —this failure is new and not related to this PR. See #469 (comment) and #482, just opened, which should resolve it.

After I merge that one commit to main, this branch can be rebased and the failure should disappear.

[khaeru]

#482, just opened, which should resolve it.

Now merged. I'll update this branch with rebase.

[khaeru]

I have not looked in detail at the code, but appreciate the expanded tests covering this new functionality and I see that comments from @adrivinca have been addressed.

Two comments that do not block merging:

1. I found the text/examples (under "Design" and "Python API") in the PR description to be helpful. In the future it would be great to simply put such examples in the documentation directly. That way, people finding and reading the docs do not need to hunt back for older PRs to get a complete description of how the code is supposed to work.
2. The "Python API" example seems to indicate there is a growing collection of Context settings that are recognized by .model.water. It may make sense to add a class like .model.water.config.Config, analogous to .model.transport.config.Config or .report.config.Config, that collects all settings that are internal to this module. Then one instance of this class can be stored/expected at Context.water.

This will be good to merge once CI checks all pass.

[khaeru]

This will be good to merge once CI checks all pass.

The one job macos-latest-py3.14-upstream-main has timed out a few times. This is a 'flakiness' in the test suite (#391) not related to the changes in this PR, so I will merge despite it.

[khaeru]

Based on the feedback I have implemented the following changes.

I think I am missing some context here: this PR/branch was merged in March, and I don't see new commits on it. Does the comment from Monday refer to something else?

[Wegatriespython]

Oops had this PR open on the side, had meant to post the reply on #479. Sorry for the confusion.