[BradyPlanden]

Description

Adds pybamm.DiffsolSolver, which integrates a rust model entirely in Rust through the solver from #5964, with no CasADi or SUNDIALS on the path. It releases the GIL while it integrates. set_up switches the model to convert_to_format = "rust", lowers it with RustModelLowering (#5962), and reads its solutions natively (#5963).

• Options:
  • rtol, and atol as a scalar or per state.
  • calc_ic=True: diffsol finds consistent initial conditions. With False, they are root-solved first with root_method, which defaults to NonlinearSolver for DAEs.
  • hermite_interpolation=True and sens_atol_factor=1.0.
  • options: a dict of diffsol's 15 OdeSolverOptions, plus num_threads.
  • Unknown keys, an option of the wrong type (a count that is not a non-negative integer, a real that is not a number, None where only the growth and shrink bounds take it), and an invalid num_threads or sens_atol_factor raise SolverError naming the option. NumPy scalars are accepted.
  • on_failure takes only "error", the default; "warn" and "ignore" raise SolverError, since diffsol returns no partial solution.
• Times: t_eval is the discontinuity map, as for IDAKLUSolver. Output times are t_eval ∪ t_interp plus the discontinuity times, and each discontinuity segment gets only the t_interp points inside it.
  • Without t_interp, they also include 100 uniform points over the span, chosen from the times passed in, before discontinuities are merged. A bare [t0, tf] therefore gives 100 points plus the discontinuity times, and a coarse t_eval such as [0, 1800, 3600] still has a grid fine enough for Hermite reads. Pass t_interp=t_eval for output at exactly t_eval, as for IDAKLUSolver.
  • An explicit empty t_interp keeps exactly t_eval when that has three or more points, which drive-cycle experiment steps rely on.
• Threads:
  • num_threads > 1 solves several input sets concurrently with solve_batch, or multithreads a single set's sparse LU. With num_threads = 1, the sets run one after another.
  • Rayon pools are shared process-wide by thread count. The cache keeps the 4 most recently used, and builds a pool outside its lock, so a fork during a build cannot deadlock the child.
• Hermite reads:
  • Off-grid reads interpolate with cubic Hermite from the stored yp, as for IDAKLUSolver, and hermite_interpolation=False drops yp.
  • yp is never stored for output_variables, or above 4096 output points.
• Sensitivities:
  • Forward sensitivities of states, output variables, events and time integrals are computed under error control, with sens_atol_factor scaling the sensitivity floor.
  • When a solve completes only with sensitivities excluded from error control, a warning says so and names a too-small sens_atol_factor as the likely cause. Solution.solver_statistics.sens_error_control_relaxed records it.
  • An output_variables solve keeps the states and state sensitivities at both ends, so first_state, last_state and a continued step have them. Like IDAKLU's, it lists its sensitivity names with zero-row state sensitivities, so step()'s restart checks find them.
• Bindings:
  • pybamm.rust.PreparedSolver, SolveOutcome, SolverStatistics and default_solver_options, with stubs. SolveOutcome.event_index names the event that ended a solve, and solutions set closest_event_idx from it.
  • In pybamm-core, check_batch_widths with its BatchWidths and BatchSensWidth errors, and Termination::as_flag, event_time and event_index, which these bindings use.
  • pybamm.SolverStatistics gains diffsol's counters, including problem_setup_time_secs, and combines sens_error_control_relaxed with OR when solutions are added. Unpickling one saved by 26.10 fills the new fields with their defaults, so it still adds to a new one.
  • A pickled solver drops its native handles and rebuilds them on the next solve.
• Shared with IDAKLUSolver through BaseSolver: the atol check, _internal_initialisation, and the event-violation guard.
• Other: OutputAssembly.stack_parameter_blocks stacks sensitivity blocks, and hatch_build.py builds pybamm.rust with the diffsol feature.
• Docs: the API page api/solvers/diffsol_solver.rst.

Errors and limits:

• An invalid tolerance or option, or a request the compiled model cannot serve, such as sensitivities it was compiled without, raises SolverError with the binding's message. Sensitivity names that are not inputs raise SolverError from the lowering.
• An integration failure raises SolverError naming the input set.
• Sensitivities with respect to an input of width above 1, and TensorField output variables, raise SolverError, as on IDAKLU's Rust path.
• nproc is ignored.

Review on #5732

• Name the pool file rayon_pool.rs: the thread pools are in pybamm-python/src/rayon_pool.rs.

Type of change

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

Key checklist

• No style issues: pre-commit run
• Rust tests pass (582 with --all-features, 1 ignored)
• Unit tests pass (N), as do the integration tests (N)
• Tests added: 171 Python test functions, and 4 Rust tests: the pool cache's eviction, a racing build and the unheld lock, and the batch-width check.
  • Solver:
    • output grids, including the default grid with discontinuities and a coarse t_eval;
    • Hermite reads against the analytic derivative, and the 4096-point limit;
    • first and last states with sensitivities in output_variables solves;
    • failures, initial conditions and event violations;
    • pickling;
    • option defaults equal to Rust's, and SolverError for unknown keys, wrong types, out-of-range values, invalid tolerances, unknown sensitivity parameters and on_failure;
    • a statistics object pickled by 26.10, and closest_event_idx for state and output solves;
    • num_threads: serial equals batch (Hypothesis), result order, shared and evicted pools, a forked child, failures named per set;
    • per-state atol and time discontinuities.
  • Sensitivity parity with a CasADi IDAKLUSolver at rtol 1e-5 and atol 1e-8:
    • states, output variables, subsets, time integrals and events;
    • finite differences;
    • sens_atol_factor, and the relaxed retry at a floor no step can meet;
    • dy0/dp seeding;
    • the DFN under error control at the default floor;
    • stepped experiments.
  • Output variables: layouts, the batch window, experiment stitching and seeded restarts.
  • Batched inputs.
  • GIL release: a worker thread keeps running during a solve.
  • Integration:
    • SPM, DFN, cubic OCP, ECM, a two-step experiment and a pulsed DFN against IDAKLU (rtol 1e-4 to 1e-5);
    • a stepped-current SPM against IDAKLU, with and without Hermite, on a bare and a coarse t_eval: within 2e-6 V on diffsol's own grid, and 1.2 mV (Hermite) or 19 mV (linear) at worst in between;
    • a pulse train;
    • native initial conditions against a root solve, including a 3D model.

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
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 ← this PR
11	#5966	"rust" as the default convert_to_format

[codecov]

Codecov Report

❌ Patch coverage is 92.70073% with 20 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (stack/09-rust-core-diffsol@8a7b5ac). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...ckages/pybamm/src/pybamm/solvers/diffsol_solver.py	91.26%	20 Missing ⚠️

Additional details and impacted files

@@                      Coverage Diff                      @@
##             stack/09-rust-core-diffsol    #5965   +/-   ##
=============================================================
  Coverage                              ?   98.15%           
=============================================================
  Files                                 ?      354           
  Lines                                 ?    37278           
  Branches                              ?        0           
=============================================================
  Hits                                  ?    36592           
  Misses                                ?      686           
  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.

[github-actions]

Bencher Report

Project	PyBaMM
Branch	stack/10-diffsol-solver
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