Skip to content

com4FlowPy: optional Numba compute engine (bit-identical, ~30-45x faster)#1319

Open
jmasseysykes wants to merge 2 commits into
OpenNHM:masterfrom
jmasseysykes:feat/com4flowpy-numba-engine
Open

com4FlowPy: optional Numba compute engine (bit-identical, ~30-45x faster)#1319
jmasseysykes wants to merge 2 commits into
OpenNHM:masterfrom
jmasseysykes:feat/com4flowpy-numba-engine

Conversation

@jmasseysykes

@jmasseysykes jmasseysykes commented Jul 14, 2026

Copy link
Copy Markdown

Summary

Adds an optional Numba (@njit) compute engine for com4FlowPy's 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 — which is declared as an
optional dependency (pip install avaframe[numba]), so nothing changes for
existing installs.

Results

  • Bit-identical to the Python engine on single-tile domains — including
    against an independent BFW com4FlowPy run's stored zDelta / fpTravelAngleMax.
  • ~30–45× faster at 8 cores (BFW 10 m: 302 s → 6.7 s; 5 m Connaught extreme:
    95 min → 193 s), with peak memory on par with the Python engine (measured via PSS).
  • New pytest (avaframe/tests/test_com4FlowPyNumba.py) asserts engine equivalence
    across forest modules and variable alpha/Umax; it skips automatically if numba
    is not installed.
  • On fine-resolution flat runout terrain a small set of cells differ from the
    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

  • New identifiers follow the surrounding com4FlowPy snake_case; happy to switch
    to lowerCamelCase per the style guide if you prefer.
  • The benchmark harness + report live under com4FlowPy/benchmark/; happy to
    relocate (e.g. under tests/ or docs/).
  • A single-precision (float32) variant was evaluated and rejected (no speed or
    memory benefit, larger error at high resolution) — details in the report.
  • AI assistance: Level 5 per the project policy — AI-assisted via Claude Code,
    with every line human-reviewed and validated against the Python engine and BFW's
    stored outputs.

PR Checklist

  • I ran pytest locally without fails — com4FlowPy suite 12/12 (full suite via CI)
  • I added/updated tests where needed
  • I updated documentation where needed — config documented in com4FlowPyCfg.ini; happy to add .rst docs
  • Commits squashed/rebased onto latest master — 2 commits; will rebase/squash to your preference before merge
  • Standardtest run without difference — com1DFA/cython paths are untouched by this change

🤖 Generated with Claude Code

jmasseysykes and others added 2 commits July 14, 2026 14:01
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
jmasseysykes marked this pull request as ready for review July 14, 2026 23:45
@PaulaSp3
PaulaSp3 requested review from PaulaSp3 and ahuber-bfw July 15, 2026 06:46
@PaulaSp3 PaulaSp3 added the flowPyDev Ideas for future development label Jul 15, 2026
@fso42

fso42 commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

@jmasseysykes : Thanks for your PR!

3 things from my side (i.e. purely maintainer-wise, not com4FlowPy specific):

  • We can do the squash/rebase if you don't want to deal with it.

  • In general: you don't have to worry about the Standardtest stuff at all because we handle it. :-) (I should probably make that clear in the automatic PR list)

  • I would be interested in your experiences with the numba package: How much hassle is it to install? Is it working reliably on different platforms? Have you got experiences using it in a python package prepared for pypi installation?

@jmasseysykes

Copy link
Copy Markdown
Author

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:

  • Install (verified): frictionless on Ubuntu. pip install avaframe[numba] pulls a prebuilt wheel — llvmlite bundles its own LLVM, so no system compiler and no extra system deps. Andi saw the same in a separate pixi env.
  • Cross-platform (by design, not personally tested): numba ships manylinux/macOS/Windows wheels with LLVM bundled, so it's normally toolchain-free — but I've only run it on Ubuntu x86-64 myself. The integration is deliberately low-risk regardless: numba is an optional extra, and if its import or JIT ever fails the code falls back to the pure-Python engine (with a warning), so a missing or broken numba never blocks a run.
  • PyPI packaging: declared under [project.optional-dependencies], so the default pip install avaframe is unchanged — only avaframe[numba] adds it. I haven't shipped numba in a published package before, so no first-hand wheel-resolution experience there; I'd defer to yours.

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).

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

Labels

flowPyDev Ideas for future development

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants