Skip to content

Add Reservoir DA Functionality for NHF#109

Open
sclaw wants to merge 37 commits into
developmentfrom
feature/reservoir-da
Open

Add Reservoir DA Functionality for NHF#109
sclaw wants to merge 37 commits into
developmentfrom
feature/reservoir-da

Conversation

@sclaw

@sclaw sclaw commented Jun 26, 2026

Copy link
Copy Markdown

This PR adds the option to use data assimilation on reservoirs with an NHF network.

Additions

  • Uncommented the preprocess_data_assimilation() method in nhf_preprocess.py. Modified hyfeaturesnetwork approach to improve readability. Much of the hyfeaturesnetwork functionality was moved to nhf_builds, so the table was more or less a simple load and store operation.

Removals

  • None

Changes

  • Transition from lake_id to nhf_lake_id as primary key on lakes table. lake_id has some non int keys that will likely not be updated in nhf-builds. nhf_lake_id is probably what should have been used from the start.
  • The transition to nhf_lake_id necessitated a change to 64 bit IDs being passed into the mc_reach.pyx file and associated level pool classes.
  • All relevant tests have been updated to support the new reservoir_da table
  • Added AbstractNetwork property setters when encountered
  • Corrected CLI for make_forcing.py . Not scoped within this story, but was a useful fix
  • Update regression test for NHF 1.2.1

Testing

  1. Added test test/nhf/four_lakes/test_reservoir_da.py that checks whether reservoir data will be assimilated from USGS, USACE, RFC, and USBR sources. The test adds large qlats to every reach and then artificially sets USGS, USACE, and USBR gages at a low value. At first, I attempted to do 0 qlat and force reservoir outflow high, but the kernel guarded against negative storage and would not allow this. RFC data is set high because the RFC method can force an outflow directly instead of using the hybrid DA approach.
  2. Verified run_conus_lakes.py is passing
  3. Verified no regression with regression_check.sh
regression check: baseline='regress-base'  candidate='regress-cand'
gates: slowdown <= 1.05x | rel drift <= 1e-03 | mem growth <= 1.05x

  check                    baseline             candidate    speedup  verdict
  -------------------  ------------  --------------------  ---------  -------------
  Tier A wall               42.55 s               42.85 s      0.99x  ok (slower)
  Tier B kernel          2488.87 ms            2507.26 ms      0.99x  ok (slower)
  Accuracy (flow rel)     reference       0.0e+00 (NaN 0)          -  ok (identical)

info:
  - velocity/depth rel drift (informational, gated only on NaN): velocity 0.0e+00, depth 0.0e+00
OK: no regression beyond the gates

With the modifications to compiled kernel code, I had to manually modify this test as follows.
bench_kernel.py

# The waterbody-index args widened int32<->int64 across the two trees under
# test. They are small positional indices (not feature IDs), so the cast is
# value-preserving and the cross-tree regression stays valid.

WIDENED_INT_ARGS = (
    "reservoir_usgs_wbody_idx",
    "reservoir_usace_wbody_idx",
    "reservoir_usbr_wbody_idx",
    "reservoir_rfc_wbody_idx",
)  # <- set to exactly the args you changed in mc_reach.pyx

def coerce_widened(calls, dtype):
    for c in calls:
        for name in WIDENED_INT_ARGS:
            v = c["kwargs"].get(name)
            if isinstance(v, np.ndarray):
                c["kwargs"][name] = v.astype(f"int{dtype}", copy=False)
    return calls
def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--runs", type=int, default=15, help="timed iterations")
    ap.add_argument("--warmup", type=int, default=3, help="discarded warmups")
    ap.add_argument("--label", default=None, help="label for the JSON record")
    ap.add_argument("--json", action="store_true", help="write results/<label>.json")
    ap.add_argument("--int-width", type=int, choices=(32, 64), default=64,
                help="dtype for widened wbody-index args (64 = long kernel)")
    args = ap.parse_args()

    if not CALLS_PKL.exists():
        sys.exit("ERROR: no kernel_calls.pkl. Run: "
                 "python benchmark/harvest_kernel_inputs.py")
    payload = pickle.load(open(CALLS_PKL, "rb"))
    calls = payload["calls"]

regression_check.sh

# Baseline captures its own output as the accuracy reference (--save-golden into a
# private dir); the candidate is then compared against it.
BASE_CMD="$A --label regress-base-A --golden-dir benchmark/.regress_golden --save-golden --json && $B --label regress-base-B --json --int-width 32"
CAND_CMD="$A --label regress-cand-A --golden-dir benchmark/.regress_golden --json && $B --label regress-cand-B --json --int-width 64"

Notes

  • On line 614 of nhf_preprocess, I save source nhf_lake_id saved to og_nhf_lake_id in the waterbodies dataframe. I also considered creating a lookup dictionary for this information instead. Any suggestions on which would be better or an alternate approach?
  • Left gage DA code commented out. Addressed by a different story.

@sclaw
sclaw requested a review from cheginit June 26, 2026 16:18

@cheginit cheginit left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work on this. Getting reservoir DA wired into the NHF path is a meaningful step, and the overall shape is right. A few things I appreciated:

  • Moving to nhf_lake_id as the primary key is the correct call, and the reservoir_da table-driven approach — deriving reservoir_type from da_type and demoting to level pool when a DA flag is off — is clean and easy to follow.
  • The _clean_waterbodies refactor is a real improvement: counting and warning on each dropped category instead of failing the whole run on the first bad row makes hydrofabric debugging much nicer.
  • Diagnosing that the large synthetic lake ids overflow int32 and widening the kernel indices to int64 was the right fix. I built and ran the four_lakes case, and once it's routing, all four DA sources land correctly (USGS/USACE/USBR persistence at their gage values and the RFC reservoir at its forecast), so the core logic is sound.
  • Good instinct covering all four reservoir DA types in a dedicated test.

I did find a few things to address before merge. A couple of them are build/run blockers that only surface on a clean build (modern gcc/clang) and on actually exercising the reservoir-DA path, plus the Great Lakes interaction with the id rename. Details in the inline comments below.

*
*/
void init_levelpool_reach(_Reach* reach, int lake_number,
void init_levelpool_reach(_Reach* reach, long lake_number,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Building this on a modern toolchain (gcc >= 14 / clang >= 16) fails to compile:

levelpool_structs.c: in function 'init_levelpool_reach':
error: passing argument 12 of 'init_lp' from incompatible pointer type [-Wincompatible-pointer-types]

The int -> long widening of lake_number reaches the C structs and the Cython layer, but not the Fortran kernels: bind_lp.f90, bind_hybrid.f90, and bind_rfc.f90 still declare integer, intent(in) :: lake_number (4-byte c_int), while init_levelpool_reach / init_hybrid_reach / init_rfc_reach now pass &lake_number as a long*. No .f90 is touched in this PR.

On older compilers this is just a warning, which is why local builds pass but then Fortran reads 4 bytes of the 8-byte long, so a large nhf_lake_id truncates inside the kernel. On gcc >= 14 / clang >= 16 it's a hard error and the package won't build.

Fix: declare lake_number as integer(c_int64_t) (ISO_C_BINDING) in the three binds and propagate through the module init methods they call (plus any derived-type field that stores it).

Same issue at hybrid_structs.c and rfc_structs.c

@sclaw sclaw Jun 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice find! I believe I have resolved this.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One thing to finish so the id widening is fully platform-agnostic: lake_number and the reservoir index arrays are currently a mix of long and int64_t, which only holds together on Linux.

long isn't a portable 64-bit type. It's 64-bit on Linux and macOS (LP64) but 32-bit on Windows (LLP64). And even among 64-bit platforms, int64_t resolves to a different underlying type — it's long on Linux but long long on macOS — so there long and int64_t are distinct C types. The Fortran binds and the init_* externs in the *_structs.c files use int64_t (C_INT64_T), while the C struct fields, the init_*_reach signatures, and the Cython side still use long. That mixed state compiles on Linux (where int64_t == long) but not elsewhere: on macOS, passing &reach->reach.lp.lake_number (a long*) into an int64_t* parameter is an incompatible pointer type; on Windows, long is too narrow to hold the id at all (truncation).

Using int64_t uniformly for every 64-bit id avoids all of it. Changes needed:

  • levelpool_structs.h, hybrid_structs.h, rfc_structs.h: struct field long lake_number;int64_t lake_number;, and add #include <stdint.h>.
  • levelpool_structs.c, hybrid_structs.c, rfc_structs.c: the init_*_reach(...) definition's long lake_numberint64_t (the int64_t *lake_number externs stay as-is).
  • reach.pxd and the three reservoir .pxd: long lake_numberint64_t lake_number in the _MC_* struct decls, plus from libc.stdint cimport int64_t at the top.
  • levelpool.pyx, hybrid.pyx, rfc.pyx: long lake_numberint64_t in both the init_*_reach extern block and the __init__ signature, plus the same cimport.
  • mc_reach.pyx: the four reservoir_*_wbody_idx memoryviews (const long[:]) and the usgs_idx / usace_idx / usbr_idx / rfc_idx arrays (np.ndarray[long]) → int64_t, plus the cimport. compute.py already builds these as int64, so this makes the buffer binding valid on every platform and a long[:] view can't bind an int64 array where long is 32-bit.

The same longint64_t treatment applies to any other field carrying a 64-bit id (e.g. the reach id and upstream_ids in the reservoir __init__s) if you want it fully airtight on Windows as well.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe I have now addressed this in 61edb1c

Comment thread src/troute-network/troute/nhf_preprocess.py Outdated
Comment thread test/nhf/make_forcing.py Outdated
Comment thread src/troute-network/troute/nhf_preprocess.py
Comment thread test/nhf/four_lakes/test_reservoir_da.py Outdated
Comment thread src/troute-routing/troute/routing/compute.py
Comment thread src/troute-network/troute/nhf_preprocess.py
Comment thread src/troute-network/troute/nhf_preprocess.py Outdated
@sclaw

sclaw commented Jul 1, 2026

Copy link
Copy Markdown
Author

@cheginit , I have updated this. Please review at your convenience.

Tests passing

  • great lakes
  • four lakes
  • conus lakes
  • regression check

@sclaw

sclaw commented Jul 8, 2026

Copy link
Copy Markdown
Author

The following tests are passing

pytest -m integration

python benchmark/run_conus_lakes.py

benchmark/regression_check.sh

@sclaw

sclaw commented Jul 8, 2026

Copy link
Copy Markdown
Author

@cheginit , apologies for the scope creep on this PR. I have added the following features to this PR

Additions

  • Added lakeout feature for NHF. Now, when a user specifies lakeout_output in output_parameters, CF-1.6 compliant netcdf files will be written to that directory containing reservoir inflows, outflows, and water surface elevations. The files are indexed by nhf_lake_id and model timestamps. I have also incorporated these outputs into the ciss_creek, four_lakes, and conus_reservoir_da tests. Adding lakeout to these tests 1) tests the lakeout feature and 2) allows me to do some better assertions on those integration tests.

Removals

  • I have removed the hot brook and lake creek tests. These were initially added to test level pool routing, but the ciss_creek test checks all the same functionallity that they would check.
  • In the refactor or NHF tests, I removed lots of git-tracked files in favor of a setup where all test data is generated via test/nhf/prep_tests.py

Changes

  • I performed a complete refactor of NHF tests following your test feedback. Under the new approach, you must explicitly generate test data via python -m test.nhf.prep_tests. If that data is not present, the tests are skipped. The tests are also all marked as integration tests now, which given their longer runtimes and larger data requirements, allows them to be run separately from the rest of the pytest suite.
    I noticed that many of the tests needed shared utilities to generate their data and perform validation, so I created a shared utils submodule to house these functions.
    Individual integration tests now follow a consistent format of defining setup() and test_*() methods. I made a deliberate effort to reduce boilerplate while retaining enough flexibility for each test to test what it needed.
    Finally, the readme has been updated to reflect these changes and is generally much more accurate and useful IMO.

Testing

As stated above, the following tests are passing for me

pytest -m integration

python benchmark/run_conus_lakes.py

benchmark/regression_check.sh

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants