From 638aa963082914660c8f6ee47ec0e7be4ccab052 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Tue, 22 Sep 2026 22:00:56 +0100 Subject: [PATCH] feat!: default convert_to_format to rust BaseModel defaults to "rust"; CasADi-specific tests pin their format; user guide and multithreading docs describe the rust backend. --- CHANGELOG.md | 1 + docs/source/developer/rust_idaklu_build.md | 6 + .../performance/07-multithreading.ipynb | 16 ++ docs/source/user_guide/index.md | 1 + docs/source/user_guide/rust_backend.rst | 168 ++++++++++++++++++ packages/pybamm-rust/README.md | 4 +- .../expression_tree/operations/serialise.py | 4 +- .../pybamm/src/pybamm/models/base_model.py | 4 +- .../pybamm/src/pybamm/solvers/base_solver.py | 11 +- .../test_models/standard_output_tests.py | 3 +- .../test_asymptotics_convergence.py | 2 + .../test_compare_basic_models.py | 2 + .../test_lead_acid/test_compare_outputs.py | 2 + .../test_solvers/test_casadi_mode.py | 2 + .../integration/test_solvers/test_idaklu.py | 1 + .../integration/test_voltage_as_state.py | 1 + .../pybamm/tests/unit/test_base_simulation.py | 3 +- .../test_simulation_with_experiment.py | 2 + .../unit/test_plotting/test_quick_plot.py | 2 + .../test_serialisation/test_serialisation.py | 6 +- packages/pybamm/tests/unit/test_simulation.py | 1 + .../unit/test_solvers/test_base_solver.py | 1 + .../test_casadi_algebraic_solver.py | 6 + .../unit/test_solvers/test_casadi_solver.py | 6 + .../test_solvers/test_composite_solver.py | 1 + .../unit/test_solvers/test_idaklu_solver.py | 4 + .../test_solvers/test_nonlinear_solver.py | 1 + .../unit/test_solvers/test_observation.py | 3 + .../unit/test_solvers/test_process_rust.py | 20 ++- .../tests/unit/test_solvers/test_solution.py | 27 ++- .../test_solvers/test_store_first_last.py | 2 + 31 files changed, 292 insertions(+), 21 deletions(-) create mode 100644 docs/source/user_guide/rust_backend.rst diff --git a/CHANGELOG.md b/CHANGELOG.md index ccfe051af8..12a1159249 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - `pybamm` now requires `pybammsolvers>=0.12.0`, whose IDAKLU backend can evaluate a Rust `CompiledModel`. Upgrade `pybammsolvers` together with `pybamm`. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961)) - `pybamm` now ships platform wheels with its Rust core compiled in as `pybamm.rust._core`: one `cp310-abi3` wheel per platform for Python 3.10-3.15. Building from source or the sdist requires `cargo` 1.89+; `pip install --only-binary pybamm pybamm` forces a wheel. ([#5958](https://github.com/pybamm-team/PyBaMM/pull/5958)) +- Models now default to `convert_to_format = "rust"` instead of `"casadi"`, so every model is lowered to the Rust compute core unless it opts out. Set `model.convert_to_format = "casadi"` to restore the previous backend. ([#5966](https://github.com/pybamm-team/PyBaMM/pull/5966)) ## Deprecated diff --git a/docs/source/developer/rust_idaklu_build.md b/docs/source/developer/rust_idaklu_build.md index dbaf312f0e..37dc0f054b 100644 --- a/docs/source/developer/rust_idaklu_build.md +++ b/docs/source/developer/rust_idaklu_build.md @@ -163,3 +163,9 @@ uv sync --extra all --group dev --reinstall-package pybamm An ABI mismatch between components — rebuild both extensions from the same tree (see "Rebuilding after changes"). + +## Supported features + +Feature coverage (events, DAEs, forward sensitivities, output variables, +experiments) and the expression types the backend can lower are documented in +the [Rust backend user guide](../user_guide/rust_backend.rst). diff --git a/docs/source/examples/notebooks/performance/07-multithreading.ipynb b/docs/source/examples/notebooks/performance/07-multithreading.ipynb index 5d548372ac..11489fde69 100644 --- a/docs/source/examples/notebooks/performance/07-multithreading.ipynb +++ b/docs/source/examples/notebooks/performance/07-multithreading.ipynb @@ -183,6 +183,22 @@ "source": [ "So in this case the speed-up for using multiple threads to solve 1000 SPM simulations is much less than for the DFN simulations, and above 4 threads no speed-up is observed at all." ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## The same option on the other backends\n", + "\n", + "Given several input sets, `num_threads` means the same thing wherever it appears: *this many input sets solve at once*.\n", + "\n", + "- `IDAKLUSolver` schedules them with OpenMP, one solver per thread, for both `convert_to_format=\"casadi\"` and `\"rust\"`. Work is handed out dynamically, so a sweep whose sets terminate at different times (a current sweep hitting a voltage cut-off, say) does not leave threads idle waiting on the slowest block.\n", + "- `DiffsolSolver` takes the same option and schedules with rayon: `pybamm.DiffsolSolver(options={\"num_threads\": 8})`. Pools are shared process-wide by thread count, so ten solvers asking for 8 threads share one pool of 8 rather than spawning eighty. With a single input set it gives the threads to that set's sparse LU instead, which pays off only for large models (around 50k states and up).\n", + "\n", + "Either way, leave `num_threads` at its default of 1 if something outside PyBaMM is already parallelising over your sweep — a `ThreadPoolExecutor`, joblib, or a parameter-estimation library — or the two thread counts multiply.\n", + "\n", + "The timings above are from one machine and one PyBaMM version; treat them as the shape of the curve rather than numbers to reproduce." + ] } ], "metadata": { diff --git a/docs/source/user_guide/index.md b/docs/source/user_guide/index.md index 7b9558df37..108c585479 100644 --- a/docs/source/user_guide/index.md +++ b/docs/source/user_guide/index.md @@ -23,6 +23,7 @@ maxdepth: 2 fundamentals/index fundamentals/battery_models fundamentals/public_api +rust_backend ``` ```{toctree} diff --git a/docs/source/user_guide/rust_backend.rst b/docs/source/user_guide/rust_backend.rst new file mode 100644 index 0000000000..578a3d3380 --- /dev/null +++ b/docs/source/user_guide/rust_backend.rst @@ -0,0 +1,168 @@ +.. _rust-backend: + +Rust Backend +============ + +PyBaMM ships a Rust compute core (``pybamm.rust``) that compiles a discretised +model's expression trees: right-hand sides, algebraic residuals, Jacobians, +events, and output variables into a form evaluated entirely outside Python. +Two solvers use it: + +- :class:`pybamm.IDAKLUSolver` runs the usual SUNDIALS IDA integrator with + every callback evaluated by the Rust core instead of CasADi. It supports + forward sensitivities, ``output_variables``, events, and experiments, with + the differences listed in `Differences from the CasADi backend`_. +- :class:`pybamm.DiffsolSolver` integrates with the pure-Rust BDF + implementation from the `diffsol `_ + crate, with no CasADi or SUNDIALS involvement. It supports forward + sensitivities under solver error control, ``output_variables``-only solves, + events, and ``t_interp``. + +Selecting the backend +--------------------- + +Which backend a solver uses follows the model's ``convert_to_format`` +attribute (documented on :class:`pybamm.BaseModel`), which defaults to +``"rust"``: + +.. code-block:: python + + import pybamm + + model = pybamm.lithium_ion.SPM() + + sim = pybamm.Simulation(model) # the default IDAKLUSolver, on the Rust core + sol = sim.solve([0, 3600]) + +Set ``convert_to_format = "casadi"`` before solving to use the CasADi backend +instead: + +.. code-block:: python + + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + +:class:`pybamm.DiffsolSolver` always uses the Rust backend and converts the +model itself: + +.. code-block:: python + + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve([0, 3600]) + +Differences from the CasADi backend +----------------------------------- + +On the Rust backend, :class:`pybamm.IDAKLUSolver`: + +- raises :class:`pybamm.SolverError` for sensitivities with respect to an input + parameter that holds several values (``pybamm.InputParameter`` with an + ``expected_size`` above 1); +- runs one solver per thread, so it solves with ``num_solvers`` equal to + ``num_threads``, with a warning when they differ. + +Set ``convert_to_format = "casadi"`` for either. + +Output times with ``DiffsolSolver`` +----------------------------------- + +Without ``output_variables``, both solvers store the full state at every +output time; they differ in which times those are. ``IDAKLUSolver`` uses its +internal integrator steps as the output grid when ``t_interp`` is omitted, +whereas diffsol evaluates its error-controlled dense output at a grid of its +own: ``t_eval`` and the times of any discontinuities, merged with 100 points +uniform over the span. Pass ``t_interp`` to choose the output times exactly; +they are then ``t_interp`` merged with ``t_eval`` and any discontinuity times. + +Off-grid reads — ``sol["Voltage [V]"](t)`` at a ``t`` between output points — +interpolate with cubic Hermite, as they do for ``IDAKLUSolver``: the solver +stores the state time derivatives alongside the solution by default. Pass +``hermite_interpolation=False`` to drop them and halve trajectory memory, at +the cost of off-grid reads falling back to linear interpolation. IDAKLU's +Hermite knots are its internal steps, while diffsol's are its output times, so +on a coarse ``t_interp`` IDAKLU's off-grid reads remain the more accurate of +the two. + +Supported expression types +-------------------------- + +The Rust backend supports the following expression types: + +**Leaf nodes:** + +- ``Scalar`` — constant scalar values +- ``Array`` / ``Vector`` / ``Matrix`` — dense arrays and matrices +- ``SparseMatrix`` — CSR sparse matrices +- ``StateVector`` / ``StateVectorDot`` — state variable slices +- ``InputParameter`` — named input parameters +- ``Time`` — simulation time + +**Binary operators:** + +- Arithmetic: ``+``, ``-``, ``*``, ``/``, ``**`` +- Matrix: ``@`` (matrix multiplication) +- Comparison: ``minimum``, ``maximum`` +- Other: ``modulo``, ``hypot``, ``EqualHeaviside``, ``NotEqualHeaviside`` + +**Unary operators:** + +- ``-`` (negation), ``abs`` +- Math functions: ``sqrt``, ``exp``, ``log``, ``sin``, ``cos``, ``tanh``, + ``sinh``, ``cosh``, ``arcsinh``, ``arctan``, ``erf``, ``sign``, ``floor``, + ``ceiling`` +- Reductions: ``max``, ``min`` (over arrays), differentiated to the + argmax/argmin subgradient + +**Structural:** + +- ``Index`` — array slicing +- ``Concatenation`` — combining arrays +- ``Conditional`` — branch selection +- ``VectorField`` — stacked components (read scalar components back with + ``pybamm.Component``) + +**Interpolation:** + +- 1D ``linear``, ``cubic``, and ``pchip`` interpolation +- 2D and 3D ``linear`` and ``cubic`` regular-grid interpolation + +A symbol the backend cannot convert raises an error naming the symbol at +solver set-up, so an unsupported model fails loudly rather than +mis-evaluating. + +Adding new expression types +--------------------------- + +A symbol reaches the Rust core through its ``_to_rust(graph, rust_symbols)`` +method, which allocates nodes on a ``pybamm.rust.ExprGraph``. Subclasses of +:class:`pybamm.Function` override ``_rust_evaluate(graph, *children)`` instead, +which receives the children already converted. + +1. **If existing graph methods can express it**, only the Python side + changes: compose them in ``_rust_evaluate`` or ``_to_rust``, as + ``Arcsinh2`` in ``functions.py`` does. + +2. **Otherwise, add a primitive** under ``packages/pybamm-rust``: + + - add a ``Node`` variant in ``pybamm-core/src/node.rs``, the vocabulary of + the graph Python builds; + - lower it to the instruction tape in ``pybamm-core/src/ir/mod.rs``, adding + a ``UnaryOp`` or ``BinaryOp`` variant for an element-wise operation; + - define that op's value in ``pybamm-core/src/eval.rs`` + (``with_unary_op!`` or ``with_binary_op!``), which every interpreter of + the tape shares; + - give it a derivative in ``tangent.rs`` (forward mode) and ``adjoint.rs`` + (reverse mode), and an arm in each pass whose ``match`` the compiler then + reports as non-exhaustive (``simplify.rs``, ``sparsity.rs``, + ``zero_propagate.rs``, ``row_extract.rs``, ...); + - expose it as an ``ExprGraph`` method in ``pybamm-python/src/expr.rs``, + and declare that method in the hand-written stubs, + ``packages/pybamm/src/pybamm/rust/_core.pyi``, in the same change + (``tests/unit/test_rust_stubs.py`` checks them against the extension). + +3. **Test** the lowering in + ``packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_rust.py`` + and its values against CasADi in + ``packages/pybamm/tests/integration/test_rust_parity.py``. A new primitive + also joins the generators in ``pybamm-core/tests/common/cases.rs``, so the + Rust property tests exercise its values and derivatives. diff --git a/packages/pybamm-rust/README.md b/packages/pybamm-rust/README.md index 6a464d3559..78fcb35b96 100644 --- a/packages/pybamm-rust/README.md +++ b/packages/pybamm-rust/README.md @@ -5,7 +5,9 @@ PyBaMM's other backends. Python hands over a discretised model as an expression this workspace turns it into flat instruction tapes it can evaluate, differentiate and solve. -Models opt in with `model.convert_to_format = "rust"`. +Models use it by default; set `model.convert_to_format = "casadi"` to opt out. The +[Rust backend user guide](../../docs/source/user_guide/rust_backend.rst) covers choosing +the backend and how it differs from CasADi. ## Layout diff --git a/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py b/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py index faf4558c50..84604caaa6 100644 --- a/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py +++ b/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py @@ -483,7 +483,9 @@ def serialise_custom_model(model: pybamm.BaseModel, compress: bool = False) -> d "base_class": base_cls_str, "base_class_mro": base_class_mro, "options": getattr(model, "options", {}), - "convert_to_format": getattr(model, "convert_to_format", "casadi"), + "convert_to_format": getattr( + model, "convert_to_format", pybamm.BaseModel._DEFAULT_CONVERT_TO_FORMAT + ), "rhs": [ ( convert_symbol_to_json(variable), diff --git a/packages/pybamm/src/pybamm/models/base_model.py b/packages/pybamm/src/pybamm/models/base_model.py index 58f544b2a8..99c86be2c8 100644 --- a/packages/pybamm/src/pybamm/models/base_model.py +++ b/packages/pybamm/src/pybamm/models/base_model.py @@ -95,7 +95,7 @@ class BaseModel: different submodels in the full concatenated solution vector. """ - _DEFAULT_CONVERT_TO_FORMAT = "casadi" + _DEFAULT_CONVERT_TO_FORMAT = "rust" _VALID_CONVERT_TO_FORMATS = ("python", "casadi", "jax", "rust") # Set by discretisation, deserialisation and BaseSolver.set_up. Evaluators are CasADi @@ -280,7 +280,7 @@ def convert_to_format(self) -> str | None: - "jax": convert to JAX expression tree. - "rust": convert to a compiled Rust expression graph. - Default is "casadi". ``None`` is deprecated: it evaluates as "python" and + Default is "rust". ``None`` is deprecated: it evaluates as "python" and raises a ``DeprecationWarning`` when set. """ return self._convert_to_format diff --git a/packages/pybamm/src/pybamm/solvers/base_solver.py b/packages/pybamm/src/pybamm/solvers/base_solver.py index ce57cee85a..4558b768c6 100644 --- a/packages/pybamm/src/pybamm/solvers/base_solver.py +++ b/packages/pybamm/src/pybamm/solvers/base_solver.py @@ -601,10 +601,15 @@ def _check_and_prepare_model_inplace(self, model): if ( isinstance(self, pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver) - ) and model.convert_to_format != "casadi": - pybamm.logger.warning( - f"Converting {model.name} to CasADi for solving with CasADi solver" + and model.convert_to_format != "casadi" + ): + # Converting a model left at the default format is routine + log = ( + pybamm.logger.info + if model.convert_to_format == model._DEFAULT_CONVERT_TO_FORMAT + else pybamm.logger.warning ) + log(f"Converting {model.name} to CasADi for solving with CasADi solver") model.convert_to_format = "casadi" if isinstance(self.root_method, pybamm.CasadiAlgebraicSolver): if model.convert_to_format == "rust": diff --git a/packages/pybamm/tests/integration/test_models/standard_output_tests.py b/packages/pybamm/tests/integration/test_models/standard_output_tests.py index 5eb4680b7f..0498fe6a36 100644 --- a/packages/pybamm/tests/integration/test_models/standard_output_tests.py +++ b/packages/pybamm/tests/integration/test_models/standard_output_tests.py @@ -186,8 +186,9 @@ def test_overpotentials(self): elif self.operating_condition == "off": np.testing.assert_allclose(self.eta_r_av(self.t), 0, rtol=1e-7, atol=1e-6) np.testing.assert_allclose(self.eta_e_av(self.t), 0, rtol=1e-12, atol=1e-11) + # Analytically zero; round-off of ~4 V potentials varies by backend. np.testing.assert_allclose( - self.delta_phi_s_av(self.t), 0, atol=2e-14, rtol=1e-16 + self.delta_phi_s_av(self.t), 0, atol=5e-14, rtol=1e-16 ) def test_ocps(self): diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py index 1d90f93792..4639bebb6e 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py @@ -15,6 +15,8 @@ def test_leading_order_convergence(self): # Create models leading_order_model = pybamm.lead_acid.LOQS() full_model = pybamm.lead_acid.Full() + for model in (leading_order_model, full_model): + model.convert_to_format = "casadi" # Same parameters, same geometry parameter_values = full_model.default_parameter_values diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py index 21858a3d4f..3cea18cd03 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py @@ -10,6 +10,8 @@ class TestCompareBasicModels: def test_compare_full(self): basic_full = pybamm.lead_acid.BasicFull() full = pybamm.lead_acid.Full() + for model in (basic_full, full): + model.convert_to_format = "casadi" parameter_values = pybamm.ParameterValues("Sulzer2019") parameter_values["Current function [A]"] = 10 diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py index 82c43b8c60..9d6655ed14 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py @@ -19,6 +19,7 @@ def test_compare_averages_asymptotics(self): param = models[0].default_parameter_values param.update({"Current function [A]": 1}) for model in models: + model.convert_to_format = "casadi" param.process_model(model) # set mesh @@ -61,6 +62,7 @@ def test_compare_outputs_surface_form(self): param = models[0].default_parameter_values param.update({"Current function [A]": 1}) for model in models: + model.convert_to_format = "casadi" param.process_model(model) # set mesh diff --git a/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py b/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py index 5e24c85860..3e2031179c 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py +++ b/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py @@ -25,6 +25,7 @@ def test_casadi_solver_mode(self): for solver in solvers: # define model model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" # solve simulation sim = pybamm.Simulation( @@ -57,6 +58,7 @@ def test_casadi_fast_matches_safe_default_ode(self): solutions = [] for mode in ["safe", "fast"]: model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(mode=mode, atol=1e-6, rtol=1e-6) sim = pybamm.Simulation(model, solver=solver) solutions.append(sim.solve(t_eval)) diff --git a/packages/pybamm/tests/integration/test_solvers/test_idaklu.py b/packages/pybamm/tests/integration/test_solvers/test_idaklu.py index 12e4ee1b7e..4ddf0ebaec 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_idaklu.py +++ b/packages/pybamm/tests/integration/test_solvers/test_idaklu.py @@ -323,6 +323,7 @@ class TestIDAKLUSolverAOTCompilation: @staticmethod def _build_model(): model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "casadi" geometry = model.default_geometry param = model.default_parameter_values param.process_model(model) diff --git a/packages/pybamm/tests/integration/test_voltage_as_state.py b/packages/pybamm/tests/integration/test_voltage_as_state.py index f188ce48cc..289e568c37 100644 --- a/packages/pybamm/tests/integration/test_voltage_as_state.py +++ b/packages/pybamm/tests/integration/test_voltage_as_state.py @@ -64,6 +64,7 @@ def test_default_voltage_is_expression(self, model_cls): @pytest.mark.parametrize("model_cls", REDUCED_MODELS) def test_opt_in_solvable_by_casadi_safe(self, model_cls): model = model_cls(options=OPT_IN) + model.convert_to_format = "casadi" sim = pybamm.Simulation(model, solver=pybamm.CasadiSolver(mode="safe")) sol = sim.solve([0, 3600]) v = sol["Voltage [V]"].entries diff --git a/packages/pybamm/tests/unit/test_base_simulation.py b/packages/pybamm/tests/unit/test_base_simulation.py index 5995b739bd..0f194d9738 100644 --- a/packages/pybamm/tests/unit/test_base_simulation.py +++ b/packages/pybamm/tests/unit/test_base_simulation.py @@ -176,6 +176,7 @@ def test_simulation_pickled_with_public_convert_to_format(self, monkeypatch): # in each model's instance dict under the public name. sim = pybamm.Simulation(pybamm.lithium_ion.SPM()) sim.build() + convert_to_format = sim.built_model.convert_to_format def public_name_state(self): state = self.__dict__.copy() @@ -189,7 +190,7 @@ def public_name_state(self): monkeypatch.undo() loaded = pickle.loads(data) - assert loaded.built_model.convert_to_format == "casadi" + assert loaded.built_model.convert_to_format == convert_to_format solution = loaded.solve([0, 100]) np.testing.assert_allclose( solution["Voltage [V]"](100.0), diff --git a/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py b/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py index 6c45123bb5..efcf2d2883 100644 --- a/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py +++ b/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py @@ -1463,6 +1463,8 @@ def test_run_experiment_cccv_solvers(self): solutions = {} for name, solver in solvers.items(): model = pybamm.lithium_ion.SPM() + if isinstance(solver, pybamm.CasadiSolver): + model.convert_to_format = "casadi" sim = pybamm.Simulation(model, experiment=experiment_2step, solver=solver) solution = sim.solve() assert solution.t[-1] == pytest.approx(3600 * len(experiment_2step.steps)) diff --git a/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py b/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py index 45956fc020..91666e7505 100644 --- a/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py +++ b/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py @@ -12,6 +12,8 @@ class TestQuickPlot: @pytest.mark.parametrize("solver", _solver_args) def test_simple_ode_model(self, solver): model = pybamm.lithium_ion.BaseModel(name="Simple ODE Model") + if isinstance(solver, pybamm.CasadiSolver): + model.convert_to_format = "casadi" whole_cell = ["negative electrode", "separator", "positive electrode"] # Create variables: domain is explicitly empty since these variables are only diff --git a/packages/pybamm/tests/unit/test_serialisation/test_serialisation.py b/packages/pybamm/tests/unit/test_serialisation/test_serialisation.py index 2fe4c64966..5126c3dedc 100644 --- a/packages/pybamm/tests/unit/test_serialisation/test_serialisation.py +++ b/packages/pybamm/tests/unit/test_serialisation/test_serialisation.py @@ -259,7 +259,8 @@ def test_discretised_model_round_trips_convert_to_format(self, convert_to_format # files written before the key existed load with the model's default del data["convert_to_format"] - assert Serialise().load_model(data).convert_to_format == "casadi" + default = pybamm.BaseModel._DEFAULT_CONVERT_TO_FORMAT + assert Serialise().load_model(data).convert_to_format == default @pytest.mark.parametrize("convert_to_format", ["python", "jax", "rust"]) def test_custom_model_round_trips_convert_to_format(self, convert_to_format): @@ -275,7 +276,8 @@ def test_custom_model_round_trips_convert_to_format(self, convert_to_format): assert loaded.convert_to_format == convert_to_format del data["model"]["convert_to_format"] - assert Serialise.load_custom_model(data).convert_to_format == "casadi" + default = pybamm.BaseModel._DEFAULT_CONVERT_TO_FORMAT + assert Serialise.load_custom_model(data).convert_to_format == default def test_save_experiment_model_error(self): model = pybamm.lithium_ion.SPM() diff --git a/packages/pybamm/tests/unit/test_simulation.py b/packages/pybamm/tests/unit/test_simulation.py index cd37f9173c..09bdf72312 100644 --- a/packages/pybamm/tests/unit/test_simulation.py +++ b/packages/pybamm/tests/unit/test_simulation.py @@ -893,6 +893,7 @@ def test_drive_cycle_t_eval_warnings_for_missing_points_and_resolution(self): param["Current function [A]"] = pybamm.Interpolant( drive_cycle[:, 0], drive_cycle[:, 1], pybamm.t ) + model.convert_to_format = "casadi" sim = pybamm.Simulation( model, parameter_values=param, solver=pybamm.CasadiSolver() ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_base_solver.py b/packages/pybamm/tests/unit/test_solvers/test_base_solver.py index 56b9f73a93..f8c355eeac 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_base_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_base_solver.py @@ -1003,6 +1003,7 @@ def test_input_dependent_discontinuity_times_match_casadi(self, discontinuity): } else: model = _modulo_ramp(time) + model.convert_to_format = "casadi" inputs = {"a": 2.0, "k": 0.5, "t_switch": 100.0, "b": 0.0} solver = self._solver() solver.set_up(model, [inputs]) diff --git a/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py b/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py index beadb8894a..8b86884349 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py @@ -6,6 +6,12 @@ from tests import get_discretisation_for_testing +@pytest.fixture(autouse=True) +def _casadi_default(monkeypatch): + # This file tests the CasADi solvers, so build CasADi models instead of converting. + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi") + + class TestCasadiAlgebraicSolver: def test_algebraic_solver_init(self): solver = pybamm.CasadiAlgebraicSolver(step_tol=1e-6, tol=1e-4) diff --git a/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py b/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py index 08acb90a29..12e86c395b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py @@ -7,6 +7,12 @@ from tests import get_discretisation_for_testing, get_mesh_for_testing +@pytest.fixture(autouse=True) +def _casadi_default(monkeypatch): + # This file tests the CasADi solvers, so build CasADi models instead of converting. + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi") + + class TestCasadiSolver: def test_no_sensitivities_error(self): model = pybamm.lithium_ion.SPM() diff --git a/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py b/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py index 76538f760f..e9969e7cc7 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py @@ -59,6 +59,7 @@ def solve(self, *args, **kwargs): failing_solver = FailingSolver() working_solver = pybamm.CasadiAlgebraicSolver() + model.convert_to_format = "casadi" composite = pybamm.CompositeSolver([failing_solver, working_solver]) solution = composite.solve(model) diff --git a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py index f957ec64b0..8d0a66daae 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -176,6 +176,8 @@ def test_multiple_inputs(self): rate = pybamm.InputParameter("rate") model.rhs = {var: -rate * var} model.initial_conditions = {var: 2} + # Rust runs one solver per thread, so only casadi has num_solvers < num_threads + model.convert_to_format = "casadi" disc = pybamm.Discretisation() disc.process_model(model) @@ -1518,6 +1520,8 @@ def test_with_sparse_output_variables_and_sensitivities(self): output_variables=["Negative particle flux [mol.m-2.s-1]"], ) model = pybamm.lithium_ion.DFN() + # Rust outputs are dense, so only the casadi path has this guard + model.convert_to_format = "casadi" params = model.default_parameter_values params.update({"Current function [A]": "[input]"}) sim = pybamm.Simulation(model, solver=solver, parameter_values=params) diff --git a/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py index 5afbb50abd..aa5954faef 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py @@ -227,6 +227,7 @@ def test_matches_casadi_algebraic_solver(self): model2 = _ElectrodeSOH(param=li_param) parameter_values.process_model(model2) + model2.convert_to_format = "casadi" sol_casadi = pybamm.CasadiAlgebraicSolver(tol=1e-10).solve( model2, [0], diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation.py b/packages/pybamm/tests/unit/test_solvers/test_observation.py index ebcc7ca26d..5f884f8777 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation.py @@ -56,6 +56,8 @@ def _decay_model(factor): model.initial_conditions = {u: 1} model.variables = {"u": u, "w": factor * u} pybamm.Discretisation().process_model(model) + # CasADi whatever the default, since most of these tests read CasADi leaves + model.convert_to_format = "casadi" return model @@ -128,6 +130,7 @@ def test_a_time_integral_is_converted_once_per_model(self): "Q squared": pybamm.ExplicitTimeIntegral(u, pybamm.Scalar(0)) ** 2, } pybamm.Discretisation().process_model(model) + model.convert_to_format = "casadi" solver = pybamm.IDAKLUSolver() convert = pybamm.Solution._convert_to_casadi diff --git a/packages/pybamm/tests/unit/test_solvers/test_process_rust.py b/packages/pybamm/tests/unit/test_solvers/test_process_rust.py index 18fbe45eb8..37e68dc184 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_process_rust.py +++ b/packages/pybamm/tests/unit/test_solvers/test_process_rust.py @@ -500,12 +500,24 @@ def _run(self, *a, **kw): # pragma: no cover class TestRustSolverGuards: - def test_casadi_solver_converts_rust_model_with_a_warning(self, caplog): - model = _toy_dae("rust") - with caplog.at_level(logging.WARNING, logger=pybamm.logger.name): + @pytest.mark.parametrize( + ("convert_to_format", "level"), + [("rust", logging.INFO), ("python", logging.WARNING)], + ids=["default", "chosen"], + ) + def test_casadi_solver_warns_only_when_converting_a_chosen_format( + self, caplog, convert_to_format, level + ): + model = _toy_dae(convert_to_format) + with caplog.at_level(logging.INFO, logger=pybamm.logger.name): pybamm.CasadiSolver()._check_and_prepare_model_inplace(model) assert model.convert_to_format == "casadi" - assert "Converting" in caplog.text + levels = [ + record.levelno + for record in caplog.records + if record.getMessage().startswith("Converting") + ] + assert levels == [level] def test_casadi_root_method_is_kept_while_rust_models_use_newton(self): class _DaeSolver(pybamm.BaseSolver): diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 1632612f9e..0b62169681 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -1184,8 +1184,11 @@ def test_pickle_first_states_across_processes(self, tmp_path): assert save_result.stdout == load_result.stdout def test_load_solution_pickled_before_variable_observers(self): + # Those pickles hold CasADi leaves, so the model must read through CasADi + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" sim = pybamm.Simulation( - pybamm.lithium_ion.SPM(), + model, experiment=pybamm.Experiment( ["Discharge at 1C for 1 minute", "Rest for 1 minute"] ), @@ -1219,9 +1222,9 @@ def test_load_solution_pickled_before_variable_observers(self): def test_load_processed_variable_pickled_before_observers(self): parameter_values = pybamm.ParameterValues("Chen2020") parameter_values.update({"Current function [A]": "[input]"}) - sim = pybamm.Simulation( - pybamm.lithium_ion.SPM(), parameter_values=parameter_values - ) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + sim = pybamm.Simulation(model, parameter_values=parameter_values) solution = sim.solve( [0, 600], inputs={"Current function [A]": 1.0}, @@ -1351,7 +1354,13 @@ def test_discrete_data_sum_errors(self): @pytest.mark.parametrize( "solver_class,use_post_sum,use_output_var", _solver_classes ) - def test_discrete_data_sum(self, solver_class, use_post_sum, use_output_var): + def test_discrete_data_sum( + self, solver_class, use_post_sum, use_output_var, monkeypatch + ): + if solver_class is pybamm.CasadiSolver: + monkeypatch.setattr( + pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi" + ) model = pybamm.BaseModel(name="test_model") c = pybamm.Variable("c") model.rhs = {c: -2 * c} @@ -1470,7 +1479,13 @@ def test_discrete_data_sum(self, solver_class, use_post_sum, use_output_var): @pytest.mark.parametrize( "solver_class,use_post_sum,use_output_var", _solver_classes ) - def test_explicit_time_integral(self, solver_class, use_post_sum, use_output_var): + def test_explicit_time_integral( + self, solver_class, use_post_sum, use_output_var, monkeypatch + ): + if solver_class is pybamm.CasadiSolver: + monkeypatch.setattr( + pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi" + ) times = np.linspace(0, 1, 10) c = pybamm.Variable("c") if solver_class == pybamm.IDAKLUSolver: diff --git a/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py b/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py index 7b0e308e0f..d2baa4255b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py +++ b/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py @@ -114,6 +114,7 @@ def test_composes_with_output_variables(self): def test_non_idaklu_warns_and_no_ops(self): model = _build_simple_dae_model() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(store_first_last=True) t_eval = np.linspace(0, 1, 11) @@ -125,6 +126,7 @@ def test_non_idaklu_warns_and_no_ops(self): def test_non_idaklu_warning_fires_once_per_instance(self): model = _build_simple_dae_model() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(store_first_last=True) t_eval = np.linspace(0, 1, 5)