com4FlowPy: optional Numba compute engine (bit-identical, ~30-45x faster)#1319
com4FlowPy: optional Numba compute engine (bit-identical, ~30-45x faster)#1319jmasseysykes wants to merge 2 commits into
Conversation
Adds an @njit compute engine for the per-cell BFS, selected via a new `engine` config key (python | numba; default python). The BFS is expressed as a single @njit function over flat arrays; tiling, multiprocessing, forest handling, variable parameters, I/O and merging are unchanged. infra/back-calculation, previewMode and relId outputs fall back to the Python engine, as does a missing numba install (declared as an optional dependency). Faithful to the fluxDistOldVersion=False path: float persistence, g=9.81, all forest modules + forestInteraction, variable alpha/max_z/exponent, and the default flux distribution. The per-BFS queue is sized modestly and grows/retries on overflow, so a path is never silently truncated. Bit-identical to the Python engine on single-tile domains (Connaught 21 m; BFW 10 m, also matching an independent BFW run's stored zDelta/fpTravelAngleMax); ~30-45x faster at 8 cores with memory on par (PSS). On fine-resolution flat runout terrain a few cells differ — inherent numerical routing sensitivity (the Python result itself shifts ~175 cells under a 1 micron DEM perturbation), characterized in benchmark/REPORT.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Runs synthetic cases (no forest, forestFriction, forestDetrainment, variableAlpha+Umax) through both engines and asserts bit-identical output. Skips automatically if numba is not installed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
@jmasseysykes : Thanks for your PR! 3 things from my side (i.e. purely maintainer-wise, not com4FlowPy specific):
|
|
Thanks @fso42 — really glad this is useful, and happy to answer. Quick honest caveat up front: I've done all my own testing on Ubuntu, and I'm newer to the packaging/cross-platform side, so I've flagged below where I'm relaying how the integration is designed vs. what I've personally verified. Happy for you to do the final squash/rebase — I'll keep the branch current in the meantime. And good to know re: Standardtest. On numba:
Per your AI policy, this is a Level 5 contribution (AI-assisted via Claude Code, with every line human-reviewed and validated against the Python engine and BFW's stored outputs). |
Summary
Adds an optional Numba (
@njit) compute engine for com4FlowPy's per-cell BFS,selected via a new
engineconfig key (python|numba; defaultpython).The BFS is expressed as a single
@njitfunction over flat arrays; tiling,multiprocessing, forest handling, variable parameters, I/O and merging are
unchanged. infra/back-calculation, previewMode and relId outputs fall back to the
Python engine, as does a missing
numbainstall — which is declared as anoptional dependency (
pip install avaframe[numba]), so nothing changes forexisting installs.
Results
against an independent BFW com4FlowPy run's stored
zDelta/fpTravelAngleMax.95 min → 193 s), with peak memory on par with the Python engine (measured via PSS).
avaframe/tests/test_com4FlowPyNumba.py) asserts engine equivalenceacross forest modules and variable alpha/Umax; it skips automatically if
numbais not installed.
Python engine — this is inherent numerical sensitivity of the routing (the
Python engine's own result shifts ~175 cells under a physically meaningless
±1 µm DEM perturbation), not an engine defect; both are valid IEEE-754 solutions
of the same model.
Full methodology, benchmarks and characterization:
avaframe/com4FlowPy/benchmark/REPORT.md.Notes / questions for maintainers
snake_case; happy to switchto
lowerCamelCaseper the style guide if you prefer.com4FlowPy/benchmark/; happy torelocate (e.g. under
tests/ordocs/).memory benefit, larger error at high resolution) — details in the report.
with every line human-reviewed and validated against the Python engine and BFW's
stored outputs.
PR Checklist
pytestlocally without fails — com4FlowPy suite 12/12 (full suite via CI)com4FlowPyCfg.ini; happy to add.rstdocs🤖 Generated with Claude Code