[BradyPlanden]

Description

Builds the model layer of pybamm-core on the compiler from #5956. A CompiledModel holds a DAE M y' = f(t, y; p) with its Jacobian, sensitivities, output variables and events, and a C ABI lets pybammsolvers' IDAKLU evaluate it. Nothing in Python changes.

• Sparsity: detect_sparsity_per_output reads each output's Jacobian pattern from the graph without evaluating it. The pattern is conservative: an entry it leaves out is zero for every y.
• Jacobian assembly:
  • Columns are coloured with DSATUR.
  • Entries that classify_constant_entries proves independent of (t, y, y', p) are written from a table instead of swept. On the pinned fixtures, SPM has 116 such entries and needs 0 colours, the SPMe 116 and 3, and the DFN 2316 and 9.
  • LanePlan sweeps the colours 8 or 4 lanes at a time, from a cost model fitted on an Apple M4, within a 32 MiB scratch budget.
  • Rows too dense to colour well are filled one reverse-mode pass per row (AdjointTape). That split is kept only when it lowers the colour count.
• CompiledModel, Workspace and ModelEvaluator:
  • A CompiledModel is built once, from CompiledModelOptions { sensitivities, algebraic }. Once its output variables and events are set, it is shared read-only.
  • Each caller holds its own Workspace. ModelEvaluator pairs an Arc<CompiledModel> with one Workspace, and it is what the C ABI receives as user_data.
  • The mass matrix is classified as the identity, a diagonal selector or general.
  • The rhs is compiled without simplify, as a fold that shifts a value by one ulp stalls IDA's Newton iteration.
• Output variables and events: each set is fused into one tape, with output sensitivities projected as dg/dp + dg/dy·S.
• Snapshots: DagSnapshot, behind the serialize feature, is the bincode form of a graph. The benches and the golden and bit-exactness tests read it.
• C ABI (ffi.rs):
  • 21 extern "C" entry points. The count queries return a count or a negative status; every other entry point returns SUCCESS (0) or a negative status: ERROR_NULL_POINTER, ERROR_PANIC, ERROR_INVALID_PARAM, ERROR_INVALID_OUTPUT, ERROR_NO_SENS, ERROR_NO_OUTPUTS, ERROR_NO_ALG or ERROR_NO_EVENTS.
  • Each checks user_data and every required buffer for null, and catches a panic, rebuilding the workspace after one.
  • Buffer sizes, and exclusive use of the evaluator for the call, are the caller's contract. The algebraic block's sizes have their own count queries, pybamm_rust_n_algebraic and pybamm_rust_alg_jac_nnz.
  • pybamm_rust_abi_version() returns RUST_ABI_VERSION, which is 1.
• pybammsolvers consumer headers:
  • pybamm_rust_ffi.h resolves the symbols at runtime: dlsym(RTLD_DEFAULT), or GetProcAddress over the loaded modules on Windows. It checks the version first, and throws on any failed status.
  • RustFunctions.hpp adapts a model to IDAKLU's Expression interface. Its algebraic and Newton adapters check n_alg, n_rhs + n_alg and the Jacobian pattern and nnz against those queries, and throw std::invalid_argument before Rust writes through a buffer.
  • Both land here, uncompiled, so that the drift test can read both sides of the ABI; feat(pybammsolvers)!: an IDAKLU backend over the Rust C ABI #5961 compiles them.
• ABI drift test:
  • test_ffi_abi_contract parses ffi.rs and both headers, and checks that:
    • the versions match;
    • each bound symbol's signature and name match its export;
    • every export is bound and called;
    • every call checks its status;
    • the status codes agree.
  • EXPECTED_ABI_HASH pins the exported signatures, so a change to them fails until the hash is updated. The version numbers are only checked against each other.
  • The rust_tests CI job also runs when only packages/pybammsolvers changes, since the test reads its headers.
• libm: test_primitives_match_libm (feat(rust): expression graph, typed IR, evaluator and forward-mode AD #5956) gains reverse mode: the adjoint partials take libm's values too.

Known limits:

• Sensitivities:
  • Only width-1 inputs are supported; a vector-width input gives SensParamVector.
  • df/dp is assembled from a dense pattern.
• Algebraic block: it must be square, strictly ascending and made only of algebraic states.
• Workspaces: a workspace from another model panics, which the C ABI returns as ERROR_PANIC.
• Lanes: the lane widths are 8 and 4, with costs tuned on one machine.
• Linux loading: the extension must be loaded with RTLD_GLOBAL for the symbols to resolve, which pybamm.rust does in feat(rust)!: ship pybamm.rust in the pybamm wheel #5958.

Review on #5732

• Make the relevant parts of the core thread safe: the shared CompiledModel holds no scratch. Each caller's Workspace is separate, and ModelEvaluator pairs the two. The bindings' pools in feat(rust)!: ship pybamm.rust in the pybamm wheel #5958 and feat(rust): solve rust models with IDAKLUSolver #5962 follow the same split.

Type of change

New feature, under ## Features in CHANGELOG.md.

Key checklist

• No style issues: pre-commit run, including cargo fmt and cargo clippy
• Rust tests pass (478 with --all-features, 1 ignored: the golden-fixture regenerator)
• MSRV check passes: cargo +1.89 check --locked --all-targets --all-features
• Tests added: 274 Rust tests.
  • Unit tests:
    • colouring and Jacobian assembly recover every entry;
    • sparsity and constant-entry classification;
    • row extraction;
    • mass matrices and the algebraic block;
    • sensitivities;
    • reverse mode;
    • every C ABI entry point with a null user_data and each null buffer, a panic, null inputs, negative indices and overflowing counts, and the count queries, the algebraic block's included;
    • a dirtied workspace is never trusted.
  • Integration tests:
    • forward and reverse mode against finite differences;
    • sparsity and constant-entry soundness proptests over graphs with inputs, powers, min/max, signed zeros, sparse products, conditionals on literal, input or state-dependent selectors, linear interpolants, reductions, step functions and time, probed at varying t, y and inputs;
    • the reverse pass skips inactive branch blocks;
    • dense-row splits match the unsplit build and central finite differences;
    • SPM, SPMe and DFN sparsity goldens;
    • residuals bitwise reproducible;
    • reverse-mode partials take libm's values;
    • the five ABI contract checks.
  • Benches: four criterion benches cover compilation, evaluation, sparsity detection and Jacobian assembly.

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 ← this PR
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
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 Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.23%. Comparing base (73c0262) to head (5f19006).

Additional details and impacted files

@@                     Coverage Diff                      @@
##           stack/01-rust-core-graph    #5957      +/-   ##
============================================================
- Coverage                     98.24%   98.23%   -0.01%     
============================================================
  Files                           350      350              
  Lines                         36104    36104              
============================================================
- Hits                          35469    35468       -1     
- Misses                          635      636       +1     

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

[github-actions]

Bencher Report

Project	PyBaMM
Branch	stack/02-rust-core-model-ffi
Testbed	bare-metal

⚠️ WARNING: Truncated view!

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

🐰 View full continuous benchmarking report in Bencher