[BradyPlanden]

Description

Adds pybamm-python, a PyO3 crate that binds pybamm-core, and builds it into the pybamm wheel as pybamm.rust._core. pybamm therefore becomes a platform-specific wheel, and installing it from source needs a Rust toolchain. import pybamm does not load the extension; it loads on the first import of pybamm.rust.

Bindings

• ExprGraph builds an expression through pybamm-core's checked constructors and ParamRegistry, and rejects an Expr from another graph. A node the IR cannot compile raises NotImplementedError naming convert_to_format='casadi'. A graph pickles through its arena, and __setstate__ restores only into an empty graph.
• CompiledFunction is one compiled expression, or a named set from compile_outputs on one tape, so shared subexpressions run once. It evaluates (__call__, eval, eval_into), takes JVPs, builds and caches its Jacobian with respect to y or p, and evaluates whole trajectories (eval_trajectory, jvp_trajectory, eval_trajectory_hermite) with the GIL released. The value and Hermite trajectories run in lane tiles of about 1 MiB, clamped to 8–64 lanes; jvp_trajectory is not tiled.
• eval_into borrows out read-write and its inputs read-only through numpy's borrow checker, so a read-only out, or one that shares memory with y, y_dot, p or a dict value, raises ValueError rather than a PanicException.
• CompiledJacobian assembles into a scipy.sparse.csc_matrix whose int32 indptr and indices are cached and read-only, so in-place canonicalisation raises instead of corrupting the shared pattern.
• FunctionMetadata checks and packs each call's arguments, taking p as a dict or a flat array. ScratchPool lends each call its own scratch, because &self methods on these frozen classes may run concurrently once the GIL is released; it keeps at most 8 items.
• Every CoreError is raised as ValueError.
• _core.pyi holds hand-written stubs, which test_rust_stubs.py checks against the built extension with mypy.stubtest. AGENTS.md gains the rule that a change to the bindings updates the stubs in the same commit.
• On Linux, pybamm.rust imports _core with RTLD_GLOBAL, so that pybammsolvers can resolve the C ABI with dlsym (feat(pybammsolvers)!: an IDAKLU backend over the Rust C ABI #5961).

Build and release

• hatch_build.py runs maturin build --release --locked --features pyo3/abi3-py310, plus --compatibility linux --auditwheel skip on Linux. It then copies the extension into src/pybamm/rust/ and tags the wheel for the platform, except for editable installs. Without cargo, it raises an error explaining how to install Rust or use --only-binary.
• maturin>=1.14.1,<2.0 joins the build requirements. The sdist includes the Rust sources, Cargo.lock and benches. [tool.uv] cache-keys cover the Rust sources, so uv sync rebuilds after a change.
• publish_pypi.yml builds one cp310-abi3 wheel per platform with cibuildwheel 4.1.1: manylinux x86_64 and aarch64, macOS arm64 and x86_64, and Windows AMD64.
  • PyPy, musllinux and free-threaded builds are skipped, as pybammsolvers already does, and fall back to the sdist.
  • Each wheel is imported on Python 3.10, which builds it, and the Linux x86_64 wheel also on Python 3.15, the newest version it supports. twine check --strict runs on the result.
  • The sdist job fails above 50 MB, or if the sdist contains target/. It then builds a wheel from the sdist, which PyPy, musllinux and free-threaded installs depend on.
  • Checkouts set persist-credentials: false, and the release jobs do not cache Cargo, so a poisoned cache cannot reach a wheel.
  • The wheel jobs run on pybamm-v* releases.
• CI uses the runner's Rust, adding only rust-cache. Read the Docs and both Dockerfiles install Rust (1.89 in the images), and the benchmark image installs a built wheel. The nox and from-source jobs cache Rust builds. benchmarks_pr.yml also runs on changes to the Rust sources and hatch_build.py.
• developer/rust_idaklu_build.md describes the build. The installation guide covers the platform wheels and installing from source, and api/rust.rst documents pybamm.rust.

The abi3 wheels cover Python 3.10-3.15, the range requires-python allows. The release toolchain is the latest stable Rust.

Review on #5732

• One CompiledFunction rather than a separate group class: a named set is a CompiledFunction with an output layout (also).
• Node construction and the input map live in core (feat(rust): expression graph, typed IR, evaluator and forward-mode AD #5956); expr.rs keeps the cross-graph ownership check and the checked_root, check_read_extents and check_supported guards.
• unzip: compile_outputs collects names and ids in one pass.
• The module comment is stated once in lib.rs, and CompiledFunction has a docstring.
• File names: compiled_function.rs, scratch_pool.rs, and function_metadata.rs, holding FunctionMetadata with a docstring.
• OnceLock rather than Option: the class is #[pyclass(frozen)] and fills the tangent tapes through &self, which an Option field cannot do.
• The Mint trait: there is no trait. ScratchPool<S> holds a Mutex<Vec<S>>, and acquire takes a closure that makes an item when the pool is empty.
• The ScratchPool docstring says why the pool exists.
• data and scratch_pool on CompiledJacobian: data is the Jacobian tape shared by every call, and the pool holds each call's scratch.

Type of change

Breaking change: pybamm becomes a platform-specific wheel and a source install needs Rust (cargo 1.89+). This is under ## Breaking changes, with pip install --only-binary pybamm pybamm as the migration note. pybamm.rust is under ## Features.

Key checklist

• No style issues: pre-commit run
• Unit tests pass (N), as do the integration tests (N)
• Wheel and sdist jobs pass in a test release
• Python 3.15: on 3.15.0 (macOS arm64) at the top of the stack, a cp310-abi3 wheel built with Python 3.10 imports and solves a DFN, and the unit tests (5594), pybammsolvers' tests (70) and the Rust and diffsol integration tests (147) pass
• Tests added: 146 Python test functions and 4 Rust tests.
  • CompiledFunction:
    • call metadata, and p as a dict or an array, with validation errors;
    • eval_into with a read-only or aliased out;
    • JVPs;
    • Jacobians against finite differences, with caching and the read-only CSC arrays;
    • trajectories, jvp_trajectory and Hermite evaluation;
    • named outputs and their shared subexpressions;
    • instruction counts;
    • concurrent calls from threads.
  • Graph building: boundary validation of ranges, CSR and interpolant data, conditionals, buffer lengths, and an input re-registered at another width.
  • Lowering contract: the AD contract, a non-constant left matmul operand, vector-width inputs, and rebuild guards.
  • Build and loading:
    • hatch_build.py: wheel tags, extension extraction, missing cargo, the maturin command, and editable builds;
    • RTLD_GLOBAL is set and restored;
    • mypy.stubtest.
  • Rust: scratch-pool reuse and cap, output-layout overflow, and __setstate__ only into an empty graph.

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 ← this PR
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.24%. Comparing base (377f61d) to head (2a667be).

Additional details and impacted files

@@                      Coverage Diff                      @@
##           stack/02-rust-core-model-ffi    #5958   +/-   ##
=============================================================
  Coverage                         98.24%   98.24%           
=============================================================
  Files                               350      351    +1     
  Lines                             36112    36128   +16     
=============================================================
+ Hits                              35477    35493   +16     
  Misses                              635      635           

☔ 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/03-pybamm-rust-packaging
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