Repository navigation
feat(pybammsolvers)!: an IDAKLU backend over the Rust C ABI - #5961
Open
BradyPlanden wants to merge 1 commit into
Open
BradyPlanden wants to merge 1 commit into
BradyPlanden wants to merge 1 commit into
Conversation
This was referenced Oct 22, 2026
This was referenced Oct 22, 2026
Codecov Report✅ All modified and coverable lines are covered by tests. 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. 🚀 New features to boost your workflow:
|
2 of 3 tasks
BradyPlanden
force-pushed
the
stack/06-pybammsolvers-rust-backend
branch
from
October 28, 2026 09:55
b1cc526 to
2ee1b70
Compare
| py::arg("alg_jac"), | ||
| py::return_value_policy::take_ownership); | ||
|
|
||
| m.def("create_rust_solver_group", &create_rust_solver_group, |
BradyPlanden
force-pushed
the
stack/06-pybammsolvers-rust-backend
branch
2 times, most recently
from
October 28, 2026 15:28
b315704 to
6ffa49d
Compare
BradyPlanden
force-pushed
the
stack/06-pybammsolvers-rust-backend
branch
from
October 29, 2026 08:01
6ffa49d to
22f5ae5
Compare
BradyPlanden
marked this pull request as ready for review
October 29, 2026 08:55
Contributor
|
| Project | PyBaMM |
| Branch | stack/06-pybammsolvers-rust-backend |
| Testbed | bare-metal |
🐰 View full continuous benchmarking report in Bencher
⚠️ WARNING: Truncated view!The full continuous benchmarking report exceeds the maximum length allowed on this platform.
create_rust_solver_group, RustFunctions adapters, the Rust Newton ctor, native output batching/sensitivities and the rust-owner lifetime.
BradyPlanden
force-pushed
the
stack/06-pybammsolvers-rust-backend
branch
from
October 29, 2026 17:26
22f5ae5 to
fa7e413
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Lets IDAKLU evaluate a model through
pybamm-core's C ABI instead of CasADi. This PR compiles the adapters that #5957 added (RustFunctions.hppandpybamm_rust_ffi.h) intoidaklu, 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 thatmainhas on the CasADi path: the input width of an input-free model, and inconsistent initial sensitivities.For pybammsolvers maintainers: this bumps
pybammsolversto 0.12.0 and pinspybammsolvers>=0.12.0,<0.13.0.pybammsolvers-v0.12.0has 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.idaklu.create_rust_solver_grouptakes an evaluator pool and builds oneRustFunctionsper solver fromrust_evaluators.as_ptr(i), so the solvers share the compiled tape but not their scratch.num_solvers, elseValueError, as is a non-zeron_algthat differs from the number of algebraic entries inrhs_alg_id. The group keeps the pool alive, and declares it before the solvers so that it outlives them.idakluhas no link-time dependency on Rust. The 21 symbols are resolved at runtime withdlsym(RTLD_DEFAULT), orGetProcAddressover the loaded modules on Windows.pybamm_rust_abi_version(), which must equalPYBAMM_RUST_ABI_VERSION(1).ExpressionSetgains three, all false for CasADi and true forRustFunctions. With them:jac_action_batchcall for all parameters;output_eval_batchcall;output_sens_project.jacobian_evalscatters the CSC entries into a zeroed dense matrix. CasADi densifies its own Jacobian, so it never takes this branch.IDAKLUSolverGrouptakes the input width and raisesValueErrorwhen an inputs row has a different width.IDAKLUSolversendsstacked_inputs.size, so an input-free model sends 0. The previous value of 1 read one entry past the end; this is under## Bug fixestoo.StandaloneNewtonSolvergains a constructor over a Rust model pointer, which takes no ownership. It rejects anatolwhose length is notn_alg, and its expressions check the other sizes against the model (feat(rust): sparsity, coloured Jacobians, compiled models and the C ABI #5957), raisingValueError.IDACalcIC, which starts the algebraic sensitivities from Python's zero guess. Its correction test then weights them byatolalone, below roundoff at tight tolerances, so convergence came down to the last bits: onmain, 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 att0(or, withnewton_mode="full", the first step failed and surfaced asIDA_BAD_K).IDACalcICwith 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.IDACalcICstill computes the sensitivities with a matrix-free linear solver, withinit_all_y_ic, or if the linear solve fails; a failure then raisesSolverError(IDACalcIC: ...). Any other failedIDACalcICis logged at debug level, since the first step corrects the states, as onmain. 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:
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'st0sensitivities and size checks.Type of change
New feature, under
## FeaturesinCHANGELOG.md, and two bug fixes, under## Bug fixes. Thepybammsolvers>=0.12.0requirement 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
pre-commit runt0against central finite differences, fornewton_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.
Function.create_copykeeps the function's classRegPowerand uniform-grid interpolant lookups matchevaluatepybamm-core: expression graph, typed IR, interpreters, forward-mode ADpybamm-core:CompiledModel, coloured Jacobians, the C ABIpybamm.rustbindings, built into thepybammwheelSymbol.to_rustconvert_to_format = "rust"forScipySolver,AlgebraicSolverand EISpybammsolvers: IDAKLU over the C ABI ← this PRIDAKLUSolveron rust modelspybamm-core: diffsol DAE solverpybamm.DiffsolSolver"rust"as the defaultconvert_to_format