Refactored Water Data Pipeline

[Wegatriespython]

Refactored Water Data Pipeline

Short Desc
Improve readability & testability via DSL(Domain Specific Language)

Long Desc

• Isolate “what” (rules) from “how” (engine) with a DSL
• Replace nested procedural code with declarative Rule objects
• Centralize execution in a single engine for targeted unit tests
• Remove magic numbers, hard-coded constants, and duplicated logic
• Introduce CitationWrapper utility for citations

Example Before
Procedural code with repeated make_df/broadcast steps and chained concat on the same variable names, making it hard to follow flow and test individual pieces:

# surface water
output_df = (
  make_df("output", technology="extract_surfacewater", …)
    .pipe(broadcast, yv_ya_sw, time=…)
    .pipe(same_time)
)
# groundwater (duplicated structure)
output_df = pd.concat([output_df, make_df("output", technology="extract_groundwater", …)
    .pipe(broadcast, yv_ya_gw, time=…)
    .pipe(same_time)])
# saline water (slightly different args, more duplication)
…

Example After

def _process_extraction_output_rules(...):
    """Processes extraction output rules declaratively via DSL."""
    extraction_outputs = []
    base = {"df_node": df_node, "runtime_vals": {"year_wat": year_wat}, "node_loc": node_region}
    for r in EXTRACTION_OUTPUT_RULES.get_rule():
        r["type"] = "output"
        bcast = yv_ya_gw if r["technology"].startswith("extract_gw") else yv_ya_sw
        df_rule = (
            run_standard(r, {**base, **({"sub_time": sub_time_series} if "saline" not in r["technology"] else {})}, broadcast_year=bcast)
        )
        extraction_outputs.append(df_rule)
    return safe_concat(extraction_outputs)

---

How to review

• Feedback on the rules and clarify any questions/ambiguity
• Verify DSL engine unit-tests cover all rule types and hidden control flow
• Confirm readability: each rule object maps clearly to its output
• Ensure no remaining large procedural functions or magic constants

---

PR checklist

• CI checks pass
• GDX file comparison
• Tests added/expanded; coverage ✅
• Documentation updated (whatsnew, docstrings)
• Delete Parity Tests after review

[codecov]

Codecov Report

Attention: Patch coverage is 82.52205% with 535 lines in your changes missing coverage. Please review.

Project coverage is 78.0%. Comparing base (557d739) to head (4143a18).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...essage_ix_models/model/water/data/water_for_ppl.py	20.6%	185 Missing ⚠️
message_ix_models/model/water/data/water_supply.py	50.5%	97 Missing ⚠️
...ssage_ix_models/model/water/data/demands_legacy.py	74.7%	78 Missing ⚠️
message_ix_models/model/water/data/demands.py	83.4%	39 Missing ⚠️
..._ix_models/model/water/data/water_supply_legacy.py	78.5%	39 Missing ⚠️
message_ix_models/util/citation_wrapper.py	53.2%	29 Missing ⚠️
message_ix_models/model/water/dsl_engine.py	86.8%	19 Missing ⚠️
...ix_models/model/water/data/water_for_ppl_legacy.py	94.7%	15 Missing ⚠️
...age_ix_models/tests/model/water/test_dsl_engine.py	86.9%	12 Missing ⚠️
...ests/model/water/data/test_water_for_ppl_parity.py	95.4%	5 Missing ⚠️
... and 7 more

Additional details and impacted files

@@           Coverage Diff           @@
##            main    #336     +/-   ##
=======================================
- Coverage   78.5%   78.0%   -0.5%     
=======================================
  Files        216     236     +20     
  Lines      17409   19951   +2542     
=======================================
+ Hits       13679   15578   +1899     
- Misses      3730    4373    +643     

Files with missing lines	Coverage Δ
message_ix_models/model/water/data/demand_rules.py	100.0% <100.0%> (ø)
...ssage_ix_models/model/water/data/infrastructure.py	100.0% <100.0%> (ø)
...x_models/model/water/data/infrastructure_legacy.py	100.0% <100.0%> (ø)
...ix_models/model/water/data/infrastructure_rules.py	100.0% <100.0%> (ø)
message_ix_models/model/water/data/irrigation.py	100.0% <100.0%> (ø)
...ge_ix_models/model/water/data/irrigation_legacy.py	100.0% <100.0%> (ø)
...age_ix_models/model/water/data/irrigation_rules.py	100.0% <100.0%> (ø)
..._ix_models/model/water/data/water_for_ppl_rules.py	100.0% <100.0%> (ø)
...e_ix_models/model/water/data/water_supply_rules.py	100.0% <100.0%> (ø)
...e_ix_models/tests/model/water/data/test_demands.py	100.0% <100.0%> (ø)
... and 21 more

... and 17 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[adrivinca]

Could you please add some description for each series/rule on what are they for, what should be included and what no?
Does each _rule.py script need to have a xx_CONST series with a different name?

[Wegatriespython]

Will add the descriptions for the rules.

The constants are just to hold the magic numbers which were used in the rules. Some of them are common, but I have left them for now.
WD : Water Demand
WS : Water Supply

[adrivinca]

maybe you could rename this function, as the word "run" has in our context a different meaning. Something like "build" "fill" "parametrize", might be more understandable. not sure

[adrivinca]

just commenting, did not want to make a full review

[adrivinca]

Thanks @Wegatriespython to polish the PR, actually make a new one.
As far as I understood, tests pass apart from python < 3.10
the GDX give identical results

Vignesh still has to:

• edit the documentation pages
• remove _legacy script if we agree with it
  but apart from this it's ready for other reviews