Skip to content

feat(pybammsolvers)!: an IDAKLU backend over the Rust C ABI - #5961

Open
BradyPlanden wants to merge 1 commit into
stack/05-rust-generic-evaluatorsfrom
stack/06-pybammsolvers-rust-backend
Open

BradyPlanden wants to merge 1 commit into
stack/05-rust-generic-evaluatorsfrom
stack/06-pybammsolvers-rust-backend

Conversation

@BradyPlanden

Copy link
Copy Markdown
Member

Description

Lets IDAKLU evaluate a model through pybamm-core's C ABI instead of CasADi. This PR compiles the adapters that #5957 added (RustFunctions.hpp and pybamm_rust_ffi.h) into idaklu, with a Rust-backed solver group and Newton solver. Python reaches them in #5962, which adds the evaluator pool they take. It also fixes two IDAKLU bugs that main has on the CasADi path: the input width of an input-free model, and inconsistent initial sensitivities.

For pybammsolvers maintainers: this bumps pybammsolvers to 0.12.0 and pins pybammsolvers>=0.12.0,<0.13.0. pybammsolvers-v0.12.0 has to be released before a pybamm release that contains this PR. The pin matters in both directions: a released pybamm sends an input width of 1 for an input-free model, which 0.12.0's width check rejects. pybamm's release job installs its wheel on Python 3.15 with dependencies from PyPI, so 0.12.0's cp315 wheels must be published too; the solver wheel job builds them.

  • Solver group:
    • idaklu.create_rust_solver_group takes an evaluator pool and builds one RustFunctions per solver from rust_evaluators.as_ptr(i), so the solvers share the compiled tape but not their scratch.
    • The pool's length must equal num_solvers, else ValueError, as is a non-zero n_alg that differs from the number of algebraic entries in rhs_alg_id. The group keeps the pool alive, and declares it before the solvers so that it outlives them.
    • The ABI table is resolved before any solver is built, since a static initialiser that throws would otherwise be retried on every thread.
  • Runtime resolution:
    • idaklu has no link-time dependency on Rust. The 21 symbols are resolved at runtime with dlsym(RTLD_DEFAULT), or GetProcAddress over the loaded modules on Windows.
    • Resolution starts with pybamm_rust_abi_version(), which must equal PYBAMM_RUST_ABI_VERSION (1).
    • IDAKLU still loads without the Rust extension.
  • Capability flags: ExpressionSet gains three, all false for CasADi and true for RustFunctions. With them:
    • sensitivities make one jac_action_batch call for all parameters;
    • outputs-only solves evaluate up to 128 time points per output_eval_batch call;
    • output sensitivities are projected in Rust with output_sens_project.
  • Dense Jacobian: with a dense linear solver and a sparse Rust Jacobian, jacobian_eval scatters the CSC entries into a zeroed dense matrix. CasADi densifies its own Jacobian, so it never takes this branch.
  • Input width:
    • IDAKLUSolverGroup takes the input width and raises ValueError when an inputs row has a different width.
    • IDAKLUSolver sends stacked_inputs.size, so an input-free model sends 0. The previous value of 1 read one entry past the end; this is under ## Bug fixes too.
  • Newton solver: StandaloneNewtonSolver gains a constructor over a Rust model pointer, which takes no ownership. It rejects an atol whose length is not n_alg, and its expressions check the other sizes against the model (feat(rust): sparsity, coloured Jacobians, compiled models and the C ABI #5957), raising ValueError.
  • Consistent initial sensitivities, on every backend:
    • With sensitivities, IDAKLU always called IDACalcIC, which starts the algebraic sensitivities from Python's zero guess. Its correction test then weights them by atol alone, below roundoff at tight tolerances, so convergence came down to the last bits: on main, CasADi fails it at 4 of 9 nearby currents on a DFN with contact resistance. The return value was ignored, so the algebraic sensitivities stayed 0 at t0 (or, with newton_mode="full", the first step failed and surfaced as IDA_BAD_K).
    • The states are now made consistent by the Newton iteration, falling back to IDACalcIC with sensitivities off. The sensitivities then come from one linear solve per parameter with IDA's own matrix and direct linear solver (SensitivityInitialConditions), at the cost of one extra copy of the Jacobian and one vector per parameter.
    • IDACalcIC still computes the sensitivities with a matrix-free linear solver, with init_all_y_ic, or if the linear solve fails; a failure then raises SolverError (IDACalcIC: ...). Any other failed IDACalcIC is logged at debug level, since the first step corrects the states, as on main. This is under ## Bug fixes.

Nothing in Python exposes an evaluator pointer until #5962, so no test reaches the Rust paths in this PR. #5962's tests exercise each of them:

  • parallel input sets and one solver per thread;
  • the pool hand-out;
  • the input-width check, on both backends;
  • the dense-Jacobian scatter;
  • batched outputs and output sensitivities against CasADi;
  • the Rust Newton solver.

The existing IDAKLU suites cover the CasADi side, and the ABI drift test from #5957 keeps the headers in step with ffi.rs. #5962 also checks the Rust backend's t0 sensitivities and size checks.

Type of change

New feature, under ## Features in CHANGELOG.md, and two bug fixes, under ## Bug fixes. The pybammsolvers>=0.12.0 requirement is under ## Breaking changes, with an upgrade note: pybamm uses no Rust backend until #5962, so what breaks here is the input-width contract, under which pybamm sends an input-free model's inputs at width 0.

Key checklist

  • No style issues: pre-commit run
  • pybammsolvers builds on every wheel platform
  • Unit tests pass (N), as do the integration tests (N)
  • Tests added: 2 Python test functions, on CasADi; feat(rust): solve rust models with IDAKLUSolver #5962 adds the tests listed above.
    • an input-free model solves with zero-width inputs rows, and its solver group rejects a width-1 row;
    • algebraic sensitivities at t0 against central finite differences, for newton_mode "auto" and "full", in pybammsolvers' tests.
Rust execution engine stack: proposal #5667, showcase #5732

Each PR is one commit on the one above it.

# PR Adds
0 #5954 Function.create_copy keeps the function's class
0b #5955 CasADi's RegPower and uniform-grid interpolant lookups match evaluate
1 #5956 pybamm-core: expression graph, typed IR, interpreters, forward-mode AD
2 #5957 pybamm-core: CompiledModel, coloured Jacobians, the C ABI
3 #5958 pybamm.rust bindings, built into the pybamm wheel
4 #5959 Symbol.to_rust
5 #5960 convert_to_format = "rust" for ScipySolver, AlgebraicSolver and EIS
6 #5961 pybammsolvers: IDAKLU over the C ABI ← this PR
7 #5962 IDAKLUSolver on rust models
8 #5963 variables and sensitivities read in Rust
9 #5964 pybamm-core: diffsol DAE solver
10 #5965 pybamm.DiffsolSolver
11 #5966 "rust" as the default convert_to_format

@codecov

codecov Bot commented Oct 22, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (stack/05-rust-generic-evaluators@b408f4d). Learn more about missing BASE report.

Additional details and impacted files
@@                         Coverage Diff                         @@
##             stack/05-rust-generic-evaluators    #5961   +/-   ##
===================================================================
  Coverage                                    ?   98.18%           
===================================================================
  Files                                       ?      353           
  Lines                                       ?    36497           
  Branches                                    ?        0           
===================================================================
  Hits                                        ?    35833           
  Misses                                      ?      664           
  Partials                                    ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Comment thread packages/pybammsolvers/src/pybammsolvers/idaklu.cpp Fixed
py::arg("alg_jac"),
py::return_value_policy::take_ownership);

m.def("create_rust_solver_group", &create_rust_solver_group,
@BradyPlanden
BradyPlanden force-pushed the stack/06-pybammsolvers-rust-backend branch 2 times, most recently from b315704 to 6ffa49d Compare October 28, 2026 15:28
@BradyPlanden
BradyPlanden force-pushed the stack/06-pybammsolvers-rust-backend branch from 6ffa49d to 22f5ae5 Compare October 29, 2026 08:01
@BradyPlanden
BradyPlanden marked this pull request as ready for review October 29, 2026 08:55
@BradyPlanden
BradyPlanden requested a review from a team as a code owner October 29, 2026 08:55
@github-actions

github-actions Bot commented Oct 29, 2026 •

Copy link
Copy Markdown
Contributor

🐰 Bencher Report

ProjectPyBaMM
Branchstack/06-pybammsolvers-rust-backend
Testbedbare-metal

⚠️ WARNING: Truncated view!

The full continuous benchmarking report exceeds the maximum length allowed on this platform.

🐰 View full continuous benchmarking report in Bencher

create_rust_solver_group, RustFunctions adapters, the Rust Newton ctor,
native output batching/sensitivities and the rust-owner lifetime.
@BradyPlanden
BradyPlanden force-pushed the stack/06-pybammsolvers-rust-backend branch from 22f5ae5 to fa7e413 Compare October 29, 2026 17:26

This branch has not been deployed

No deployments
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