From fa7e4135bfac637be18d1710f2c801cfd81a2aaa Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Tue, 22 Sep 2026 22:00:54 +0100 Subject: [PATCH] feat(pybammsolvers): an IDAKLU backend over the Rust C ABI create_rust_solver_group, RustFunctions adapters, the Rust Newton ctor, native output batching/sensitivities and the rust-owner lifetime. --- CHANGELOG.md | 4 + packages/pybamm/pyproject.toml | 2 +- .../src/pybamm/solvers/idaklu_solver.py | 3 +- .../unit/test_solvers/test_idaklu_solver.py | 27 ++ packages/pybammsolvers/CMakeLists.txt | 3 + packages/pybammsolvers/pyproject.toml | 2 +- .../src/pybammsolvers/idaklu.cpp | 138 ++++++++++- .../Expressions/Base/ExpressionSet.hpp | 5 + .../Expressions/Rust/RustFunctions.hpp | 5 + .../idaklu_source/IDAKLUSolverGroup.cpp | 6 + .../idaklu_source/IDAKLUSolverGroup.hpp | 14 +- .../idaklu_source/IDAKLUSolverOpenMP.hpp | 50 ++++ .../idaklu_source/IDAKLUSolverOpenMP.inl | 234 +++++++++++++++++- .../pybammsolvers/idaklu_source/SolverLog.hpp | 5 + .../idaklu_source/StandaloneNewtonSolver.cpp | 49 ++++ .../idaklu_source/StandaloneNewtonSolver.hpp | 17 ++ .../idaklu_source/sundials_functions.inl | 53 +++- .../src/pybammsolvers/version.py | 2 +- .../tests/test_consistent_initialization.py | 57 +++++ packages/pybammsolvers/vcpkg.json | 2 +- uv.lock | 2 +- 21 files changed, 648 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62e30f742c..33c24fc4c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Breaking changes +- `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)) ## Deprecated @@ -19,6 +20,7 @@ - `pybamm.rust` exposes `ExprGraph`, `CompiledFunction` (one expression, or a named set via `compile_outputs`, with lane-batched trajectory and Hermite evaluation) and `CompiledJacobian`, with hand-written type stubs that `mypy.stubtest` pins against the built extension. ([#5958](https://github.com/pybamm-team/PyBaMM/pull/5958)) - `Symbol.to_rust` lowers a discretised expression tree into a `pybamm.rust.ExprGraph`, covering the operators, matrix products, concatenations, `Conditional`, the specific functions, `RegPower` and 1D-3D interpolants. Spatial and tensor operators, such as `Norm` and `Component`, raise `TypeError`. ([#5959](https://github.com/pybamm-team/PyBaMM/pull/5959)) - `model.convert_to_format = "rust"` evaluates a model's rhs, algebraic, Jacobian and event functions through compiled Rust evaluators, for `ScipySolver`, `AlgebraicSolver` and EIS. ([#5960](https://github.com/pybamm-team/PyBaMM/pull/5960)) +- pybammsolvers gains an IDAKLU backend that evaluates a Rust `CompiledModel` through its C ABI, resolved at runtime so IDAKLU still loads without the Rust extension, with a Rust-backed Newton solver and output sensitivities projected natively in-solver. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961)) ## Bug fixes @@ -41,6 +43,8 @@ - Fixed copying a function without simplification, as `ParameterValues` does for a parameter set loaded with `Serialise.load_parameters`, returning a generic `pybamm.Function`: `pybamm.exp` and the other specific functions lost their class, and an `Arcsinh2` copy could not be evaluated. ([#5954](https://github.com/pybamm-team/PyBaMM/pull/5954)) - `RegPower` now evaluates the same operations in `evaluate` and CasADi, taking a literal exponent of 0.5 or 2 as square roots, so they agree bitwise; CasADi used a `hypot` form. With `delta=0`, `evaluate` and the Jacobian use `base ** exponent`, as CasADi did. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955)) - CasADi linear and `pchip` interpolants on a uniform grid now find a point's segment by binary search. The `"exact"` lookup could pick the neighbouring segment at or one ulp beside a knot, giving the Jacobian that segment's slope. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955)) +- `IDAKLUSolver` no longer reads one value past the inputs of a model without input parameters. Inputs with a different number of values than at set-up now raise a `SolverError` instead of being read at the set-up length. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961)) +- `IDAKLUSolver` now finds a DAE's consistent initial sensitivities with one linear solve; at tight tolerances `IDACalcIC` could leave them zero at the first time point. A failed `IDACalcIC` is logged, and raises a `SolverError` when computing sensitivities. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961)) ## Optimizations diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index 9429fbf0ad..9909f3af03 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -28,7 +28,7 @@ classifiers = [ "Topic :: Scientific/Engineering", ] dependencies = [ - "pybammsolvers>=0.11.0,<0.12.0", + "pybammsolvers>=0.12.0,<0.13.0", "black", "numpy>=2.0.0", "scipy>=1.11.4,!=1.18.0", # 1.18.0 breaks pickling/deepcopy of interpolators, fixed in 1.18.1 (scipy#25489) diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index d85258757c..c6cec0d254 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -677,7 +677,8 @@ def to_idaklu(fn): self._setup = { "number_of_states": len(y0), - "inputs": len(stacked_inputs), + # Not len: an empty stack is 1 x 0, and each solved row must be this wide + "inputs": stacked_inputs.size, "jac_bandwidth_upper": jac_bw_upper, "jac_bandwidth_lower": jac_bw_lower, "jac_times_cjmass_colptrs": jac_times_cjmass_colptrs, 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 c83d32f265..1c96e8ec33 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -380,6 +380,33 @@ def test_model_events(self): atol=1e-5, ) + def test_an_input_free_model_takes_zero_width_inputs_rows(self): + # An empty input stack is 1 x 0, so each solver must copy no values from a row + model = pybamm.BaseModel() + model.convert_to_format = "casadi" + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1} + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver() + t_eval = np.array([0.0, 1.0]) + solution = solver.solve(model, t_eval) + np.testing.assert_allclose( + solution["u"](t_eval), np.exp(-t_eval), rtol=1e-3, atol=1e-6 + ) + + group = solver._setup["solver"] + with pytest.raises( + ValueError, match=r"wrong number of cols\. Expected 0 but got 1" + ): + group.solve( + t_eval, + np.array([]), + np.ones((1, 1)), + np.zeros((1, 1)), + np.ones((1, 1)), + ) + def test_input_params(self): # test a mix of scalar and vector input params model = pybamm.BaseModel() diff --git a/packages/pybammsolvers/CMakeLists.txt b/packages/pybammsolvers/CMakeLists.txt index af58ccc45b..a638af2866 100644 --- a/packages/pybammsolvers/CMakeLists.txt +++ b/packages/pybammsolvers/CMakeLists.txt @@ -124,6 +124,9 @@ pybind11_add_module(idaklu src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionTypes.hpp + # IDAKLU expressions - Rust backend (header-only) + src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp + src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h # IDAKLU expressions - concrete implementations ${IDAKLU_EXPR_CASADI_SOURCE_FILES} ) diff --git a/packages/pybammsolvers/pyproject.toml b/packages/pybammsolvers/pyproject.toml index b30efb676b..9743fd586d 100644 --- a/packages/pybammsolvers/pyproject.toml +++ b/packages/pybammsolvers/pyproject.toml @@ -18,7 +18,7 @@ requires-python = ">=3.10,<3.16" license-files = ["LICENSE"] readme = "README.md" # version.py stays the source of truth; a pre-commit hook mirrors it here. -version = "0.11.1" +version = "0.12.0" dependencies = [ "casadi==3.8.1", "numpy>=2.0.0", diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp index 6f9d03d054..93567319f6 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp @@ -15,6 +15,7 @@ #include "idaklu_source/IdakluJax.hpp" #include "idaklu_source/common.hpp" #include "idaklu_source/Expressions/Casadi/CasadiFunctions.hpp" +#include "idaklu_source/Expressions/Rust/RustFunctions.hpp" #include "idaklu_source/sundials_error_handler.hpp" #include "idaklu_source/reduce.hpp" #include "idaklu_source/StandaloneNewtonSolver.hpp" @@ -119,8 +120,110 @@ IDAKLUSolverGroup *create_casadi_solver_group( ); } + // No owner: CasADi's expression sets hold their own functions. return new IDAKLUSolverGroup( - std::move(solvers), number_of_states, number_of_parameters); + std::move(solvers), number_of_states, number_of_parameters, inputs_length, py::none()); +} + +// An int leads: without pybind11 headers, CodeQL parses a leading `py::object` as a variable initialiser +IDAKLUSolverGroup *create_rust_solver_group( + int number_of_states, + py::object rust_evaluators, + int number_of_inputs, + int number_of_sens_params, + int number_of_events, + int number_of_algebraic_states, + const std::vector &output_lens, + const np_array_int &jac_colptrs, + const np_array_int &jac_rowvals, + int jac_nnz, + int jac_bandwidth_lower, + int jac_bandwidth_upper, + const np_array_int &alg_jac_rowvals, + const np_array_int &alg_jac_colvals, + int alg_jac_nnz, + const np_array &rhs_alg_id, + const np_array &atol_np, + double rel_tol, + py::dict py_opts +) { + auto setup_opts = SetupOptions(py_opts); + auto solver_opts = SolverOptions(py_opts); + + // Resolve before any parallel region: a throwing static initialiser is retried per thread. + rust_ffi(); + + const int num_solvers = setup_opts.num_solvers; + const int pool_size = static_cast(py::len(rust_evaluators)); + if (pool_size != num_solvers) { + throw std::invalid_argument( + "rust_evaluators holds " + std::to_string(pool_size) + + " evaluator(s) but num_solvers is " + std::to_string(num_solvers) + + "; each solver needs its own evaluator."); + } + // The sub-block IC hands Rust a residual buffer sized from the algebraic ids. + if (number_of_algebraic_states != 0) { + auto ids = rhs_alg_id.unchecked<1>(); + int n_alg_ids = 0; + for (py::ssize_t i = 0; i < ids.shape(0); ++i) { + n_alg_ids += is_algebraic(ids(i)) ? 1 : 0; + } + if (n_alg_ids != number_of_algebraic_states) { + throw std::invalid_argument( + "n_alg is " + std::to_string(number_of_algebraic_states) + + " but rhs_alg_id marks " + std::to_string(n_alg_ids) + " states algebraic."); + } + } + + auto colptrs_ptr = jac_colptrs.unchecked<1>(); + auto rowvals_ptr = jac_rowvals.unchecked<1>(); + std::vector colptrs(colptrs_ptr.data(0), + colptrs_ptr.data(0) + colptrs_ptr.size()); + std::vector rowvals(rowvals_ptr.data(0), + rowvals_ptr.data(0) + rowvals_ptr.size()); + auto alg_rowvals_ptr = alg_jac_rowvals.unchecked<1>(); + auto alg_colvals_ptr = alg_jac_colvals.unchecked<1>(); + std::vector alg_rowvals( + alg_rowvals_ptr.data(0), + alg_rowvals_ptr.data(0) + alg_rowvals_ptr.size() + ); + std::vector alg_colvals( + alg_colvals_ptr.data(0), + alg_colvals_ptr.data(0) + alg_colvals_ptr.size() + ); + + std::vector> solvers; + solvers.reserve(num_solvers); + for (int i = 0; i < num_solvers; i++) { + // One evaluator per solver: they share the compiled tape but not their scratch. + void* model_ptr = pybamm_rust_detail::model_from_address( + rust_evaluators.attr("as_ptr")(i).cast()); + + auto functions = std::make_unique( + model_ptr, number_of_states, number_of_inputs, number_of_sens_params, + number_of_algebraic_states, number_of_events, jac_nnz, colptrs, rowvals, + alg_jac_nnz, alg_rowvals, alg_colvals, output_lens, setup_opts + ); + + solvers.emplace_back( + std::unique_ptr( + create_idaklu_solver( + std::move(functions), + number_of_sens_params, + jac_colptrs, jac_rowvals, jac_nnz, + jac_bandwidth_lower, jac_bandwidth_upper, + number_of_events, + rhs_alg_id, atol_np, rel_tol, + number_of_inputs, + solver_opts, setup_opts + ) + ) + ); + } + + return new IDAKLUSolverGroup( + std::move(solvers), number_of_states, number_of_sens_params, number_of_inputs, + std::move(rust_evaluators)); } PYBIND11_MAKE_OPAQUE(std::vector); @@ -176,6 +279,29 @@ PYBIND11_MODULE(idaklu, m) py::arg("alg_jac"), py::return_value_policy::take_ownership); + m.def("create_rust_solver_group", &create_rust_solver_group, + "Create a Rust-backed IDAKLU solver group", + py::arg("number_of_states"), + py::arg("rust_evaluators"), + py::arg("number_of_inputs"), + py::arg("number_of_sens_params"), + py::arg("number_of_events"), + py::arg("n_alg"), + py::arg("output_lens"), + py::arg("jac_colptrs"), + py::arg("jac_rowvals"), + py::arg("jac_nnz"), + py::arg("jac_bandwidth_lower"), + py::arg("jac_bandwidth_upper"), + py::arg("alg_jac_rowvals"), + py::arg("alg_jac_colvals"), + py::arg("alg_jac_nnz"), + py::arg("rhs_alg_id"), + py::arg("atol"), + py::arg("rtol"), + py::arg("options"), + py::return_value_policy::take_ownership); + m.def("observe", &observe, "Observe variables", py::arg("ts"), @@ -266,6 +392,16 @@ PYBIND11_MODULE(idaklu, m) py::arg("max_iter"), py::arg("max_backtracks"), py::arg("eps_newt"), py::arg("use_sparse"), py::arg("klu_pivot_tolerance") = 0.1) + .def(py::init, std::vector, + std::vector, sunrealtype, sunrealtype, + int, int, sunrealtype, bool, sunrealtype>(), + py::arg("rust_model"), py::arg("n_rhs"), py::arg("n_alg"), + py::arg("jac_rows"), py::arg("jac_cols"), + py::arg("atol"), py::arg("rtol"), py::arg("step_tol"), + py::arg("max_iter"), py::arg("max_backtracks"), + py::arg("eps_newt"), py::arg("use_sparse"), + py::arg("klu_pivot_tolerance") = 0.1) .def("solve", &StandaloneNewtonSolver::solve, py::arg("t"), py::arg("y0"), py::arg("inputs"), py::arg("logger") = py::none(), diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp index b1b47013e0..c3ce757b73 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp @@ -10,6 +10,11 @@ class ExpressionSet { public: + // Compile-time markers, overridden by RustFunctions. + static constexpr bool kNativeOutputSensitivities = false; + static constexpr bool kNativeBatchedOutputs = false; + static constexpr bool kNativeBatchedJacActions = false; + /** * @brief Constructor */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp index e1f6b79eaf..ca1936a3c7 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp @@ -15,6 +15,11 @@ namespace pybamm_rust_detail { +/* Rust models reach C++ as the integer address Python holds, so the cast is unavoidable. */ +inline void* model_from_address(std::uintptr_t address) { + return reinterpret_cast(address); // NOLINT(performance-no-int-to-ptr) +} + /* Throw unless `given`, a count C++ sizes buffers from, equals the Rust model's `rust`. */ inline void check_dim(const char* name, int64_t given, int64_t rust) { if (given != rust) { diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index d31d174504..f9bcd2603a 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -100,6 +100,12 @@ std::vector IDAKLUSolverGroup::solve( "inputs has wrong number of rows. Expected " + std::to_string(number_of_groups) + " but got " + std::to_string(inputs.shape()[0])); + // Each solver copies `number_of_inputs` values from its row, so a narrower row is over-read + if (inputs.shape()[1] != number_of_inputs) + throw std::domain_error( + "inputs has wrong number of cols. Expected " + std::to_string(number_of_inputs) + + " but got " + std::to_string(inputs.shape()[1])); + // pbar is optional: an empty array leaves IDAS at its unit default. const bool has_pbar = pbar.size() > 0; if (has_pbar) { diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp index e795fafb3d..9815199f4a 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp @@ -13,11 +13,18 @@ class IDAKLUSolverGroup /** * @brief Default constructor + * + * `number_of_inputs` is the width each solver copies from an inputs row. `rust_owner` keeps + * the Python object owning the expression sets' memory alive with the group (none for + * CasADi); pybind11 destroys the group with the GIL held. */ - IDAKLUSolverGroup(std::vector> solvers, int number_of_states, int number_of_parameters): + IDAKLUSolverGroup(std::vector> solvers, int number_of_states, int number_of_parameters, + int number_of_inputs, py::object rust_owner): + m_rust_owner(std::move(rust_owner)), m_solvers(std::move(solvers)), number_of_states(number_of_states), - number_of_parameters(number_of_parameters) + number_of_parameters(number_of_parameters), + number_of_inputs(number_of_inputs) {} // no copy constructor (unique_ptr cannot be copied) @@ -54,9 +61,12 @@ class IDAKLUSolverGroup */ void set_loggers(const py::object &logger); + // Declared first so it is destroyed last, after the solvers that borrow its memory. + py::object m_rust_owner; std::vector> m_solvers; int number_of_states; int number_of_parameters; + int number_of_inputs; }; #endif // PYBAMM_IDAKLU_SOLVER_GROUP_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp index d80ac45890..7ca18c9434 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp @@ -95,6 +95,10 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver vector res_dvar_dp; // Reused scratch for sparse->dense output sensitivity scatter (outputs-only mode) vector dvar_dp_dense; + // Staged (t, y) points for batched output evaluation (Rust ExprSet only). + vector stage_out_t_; + vector stage_out_y_; + size_t stage_out_first_ = 0; // i_save_ index of the first staged point bool const sensitivity; bool const save_outputs_only; bool save_hermite; @@ -172,6 +176,21 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver }; std::unique_ptr alg_state_; + // Scratch for SensitivityInitialConditions, allocated only when it can run + struct SensitivityICResources { + SUNMatrix mass_jac = nullptr; // J at cj = 1 + N_Vector* res = nullptr; + N_Vector delta = nullptr; + int number_of_parameters = 0; + + ~SensitivityICResources() { + if (res) N_VDestroyVectorArray(res, number_of_parameters); + if (delta) N_VDestroy(delta); + if (mass_jac) SUNMatDestroy(mass_jac); + } + }; + std::unique_ptr sens_ic_; + int len_rhs_; int len_alg_; std::vector event_values_; @@ -331,6 +350,27 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ bool NonlinearSolverInitialConditions(const sunrealtype& t_val); + /** + * @brief Run IDACalcIC, returning true on success. + * + * A failure is logged, and raised when with_sensitivities: IDACalcIC then + * leaves the sensitivities at their unconverged guess. + */ + bool CalcIC( + const sunrealtype& t_val, + const sunrealtype& t_next_perturbed, + const int& icopt, + bool with_sensitivities); + + /** + * @brief Make the algebraic sensitivities and differential sensitivity + * derivatives consistent at t_val; returns true on success. + * + * Solves the linear YA_YDP_INIT sensitivity system once per parameter with + * IDA's matrix and direct linear solver. + */ + bool SensitivityInitialConditions(const sunrealtype& t_val); + /** * @brief Set a consistent initialization for ODEs */ @@ -480,11 +520,21 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ void SetStepOutput(sunrealtype &tval); + /** @brief Evaluate and store any staged batched-output points */ + void FlushOutputBatch(); + /** * @brief Save the output function sensitivities at the requested time */ void SetStepOutputSensitivities(sunrealtype &tval); + /** @brief Save output sensitivities via native projection, in the CasADi path's yS layout */ + void NativeSetStepOutputSensitivities(sunrealtype &tval); + + // Native-projection scratch: [n_params * n_states] and [n_params * length_of_return_vector]. + std::vector yS_flat_; + std::vector out_sens_flat_; + /** * @brief Save Hermite interpolation derivatives at the requested time */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl index b001058198..fe0839fa37 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl @@ -293,6 +293,16 @@ void IDAKLUSolverOpenMP::Initialize() { if (len_alg_ > 0 && solver_opts.calc_ic) { BuildAlgebraicSolver(id_val); } + + // The sensitivity IC solve needs an assembled matrix and a direct solver + if (sensitivity && len_alg_ > 0 && solver_opts.calc_ic && J != nullptr && + SUNLinSolGetType(LS) == SUNLINEARSOLVER_DIRECT) { + sens_ic_ = std::make_unique(); + sens_ic_->number_of_parameters = number_of_parameters; + sens_ic_->mass_jac = SUNMatClone(J); + sens_ic_->res = N_VCloneVectorArray(number_of_parameters, yy); + sens_ic_->delta = N_VClone(yy); + } } template @@ -300,6 +310,7 @@ IDAKLUSolverOpenMP::~IDAKLUSolverOpenMP() { DEBUG("IDAKLUSolverOpenMP::~IDAKLUSolverOpenMP"); alg_state_.reset(); + sens_ic_.reset(); if (sensitivity) { IDASensFree(ida_mem); @@ -466,6 +477,9 @@ void IDAKLUSolverOpenMP::InitializeSolveStorage( length_of_return_vector = ReturnVectorLength(); i_save_ = 0; + stage_out_t_.clear(); + stage_out_y_.clear(); + stage_out_first_ = 0; // Allocate output arrays. Pre-allocate 64 elements for initial storage. int est = std::max(n_evals + n_interps, 64); @@ -683,6 +697,11 @@ SolutionData IDAKLUSolverOpenMP::BuildSolutionData(int retval) { CaptureStats(); } + // Evaluate any tail of staged batched-output points before y is finalized. + if (save_outputs_only) { + FlushOutputBatch(); + } + // Finalize output arrays if (use_knot_reduction_) { knot_reducer->Finalize(); @@ -886,15 +905,142 @@ void IDAKLUSolverOpenMP::ConsistentInitializationDAE( // If the model is a standard-form DAE, we can use the Newton solver // to solve for the consistent initialization + bool newton_success = false; if (alg_state_ && alg_state_->solver && icopt == IDA_YA_YDP_INIT) { - const bool newton_success = NonlinearSolverInitialConditions(t_val); - // Sensitivity requires IDACalcIC to compute its initial conditions - if (newton_success && !sensitivity) { - return; + newton_success = NonlinearSolverInitialConditions(t_val); + } + + // IDACalcIC weights sensitivity corrections by their guess, which is zero for + // algebraic states, so they are solved for directly once the states are set + const bool solve_sensitivities = sens_ic_ && icopt == IDA_YA_YDP_INIT; + const bool calc_ic_sensitivities = sensitivity && !solve_sensitivities; + if (!newton_success || calc_ic_sensitivities) { + if (solve_sensitivities) { + IDASensToggleOff(ida_mem); + } + const bool calc_ic_success = + CalcIC(t_val, t_next_perturbed, icopt, calc_ic_sensitivities); + if (calc_ic_success && solve_sensitivities) { + CheckErrors(IDAGetConsistentIC(ida_mem, yy, yyp), "IDAGetConsistentIC"); } } - IDACalcIC(ida_mem, icopt, t_next_perturbed); + if (solve_sensitivities) { + const bool solved = SensitivityInitialConditions(t_val); + // Also turns the sensitivities back on + ReinitializeIntegrator(t_val); + if (!solved) { + CalcIC(t_val, t_next_perturbed, icopt, true); + } + } +} + +template +bool IDAKLUSolverOpenMP::CalcIC( + const sunrealtype& t_val, + const sunrealtype& t_next_perturbed, + const int& icopt, + const bool with_sensitivities) { + const int flag = IDACalcIC(ida_mem, icopt, t_next_perturbed); + if (flag >= 0) { + return true; + } + log_.log_consistent_init_failed(t_val, sundials_error_message(flag)); + // The first step corrects the states, but the sensitivities stored at + // t_val would keep IDACalcIC's unconverged guess + if (with_sensitivities) { + CheckErrors(flag, "IDACalcIC"); + } + return false; +} + +template +bool IDAKLUSolverOpenMP::SensitivityInitialConditions( + const sunrealtype& t_val) { + DEBUG("IDAKLUSolver::SensitivityInitialConditions"); + + auto& sic = *sens_ic_; + // A direct solver may factor in place, so both are zeroed first, as IDA does + SUNMatZero(J); + SUNMatZero(sic.mass_jac); + // The unknowns are the algebraic yS and differential ypS, so the matrix is + // df/dy (J at cj = 0) in algebraic columns and -M (J1 - J0) elsewhere + jacobian_eval(t_val, SUN_RCONST(0.0), yy, yyp, sic.delta, J, + functions.get(), sic.delta, sic.delta, sic.delta); + jacobian_eval(t_val, SUN_RCONST(1.0), yy, yyp, sic.delta, + sic.mass_jac, functions.get(), sic.delta, sic.delta, sic.delta); + + const sunrealtype* id_val = N_VGetArrayPointer(id); + const auto replace_column = [](sunrealtype* col0, const sunrealtype* col1, + sunindextype begin, sunindextype end) { + for (sunindextype k = begin; k < end; k++) { + col0[k] = col1[k] - col0[k]; + } + }; + switch (SUNMatGetID(J)) { + case SUNMATRIX_SPARSE: { + if (SUNSparseMatrix_SparseType(J) != CSC_MAT) { + return false; + } + const sunindextype* colptrs = SUNSparseMatrix_IndexPointers(J); + for (int j = 0; j < number_of_states; j++) { + if (is_differential(id_val[j])) { + replace_column(SUNSparseMatrix_Data(J), + SUNSparseMatrix_Data(sic.mass_jac), colptrs[j], colptrs[j + 1]); + } + } + break; + } + case SUNMATRIX_DENSE: + for (int j = 0; j < number_of_states; j++) { + if (is_differential(id_val[j])) { + replace_column(SUNDenseMatrix_Column(J, j), + SUNDenseMatrix_Column(sic.mass_jac, j), 0, number_of_states); + } + } + break; + case SUNMATRIX_BAND: { + // Band columns point at their diagonal, so rows index from -upper + const sunindextype upper = SUNBandMatrix_UpperBandwidth(J); + const sunindextype lower = SUNBandMatrix_LowerBandwidth(J); + for (int j = 0; j < number_of_states; j++) { + if (is_differential(id_val[j])) { + replace_column(SUNBandMatrix_Column(J, j), + SUNBandMatrix_Column(sic.mass_jac, j), + -std::min(upper, j), + std::min(lower, number_of_states - 1 - j) + 1); + } + } + break; + } + default: + return false; + } + + // Overwrites IDA's factorization, which its next lsetup after a reinit redoes + if (SUNLinSolSetup(LS, J) != 0) { + return false; + } + + // The system is linear, so one Newton step from the current guess solves it + sensitivities_eval(number_of_parameters, t_val, yy, yyp, nullptr, + yyS, yypS, sic.res, functions.get(), nullptr, nullptr, nullptr); + for (int p = 0; p < number_of_parameters; p++) { + if (SUNLinSolSolve(LS, J, sic.delta, sic.res[p], SUN_RCONST(0.0)) != 0) { + return false; + } + const sunrealtype* delta = N_VGetArrayPointer(sic.delta); + sunrealtype* yS_p = N_VGetArrayPointer(yyS[p]); + sunrealtype* ypS_p = N_VGetArrayPointer(yypS[p]); + for (int i = 0; i < number_of_states; i++) { + if (is_differential(id_val[i])) { + ypS_p[i] -= delta[i]; + } else { + yS_p[i] -= delta[i]; + } + } + } + return true; } template @@ -1004,6 +1150,9 @@ void IDAKLUSolverOpenMP::SetStepFullSensitivities( } } +// Points per batched output evaluation, bounding the staged state copy. +static constexpr size_t kOutputBatchPoints = 128; + template void IDAKLUSolverOpenMP::SetStepOutput( sunrealtype &tval @@ -1011,17 +1160,52 @@ void IDAKLUSolverOpenMP::SetStepOutput( DEBUG("IDAKLUSolver::SetStepOutput"); // FLAT STORAGE: Write output variables to y[i_save_ * stride + j] - sunrealtype* y_dest = &y[i_save_ * length_of_return_vector]; - size_t j = 0; - for (auto& var_fcn : functions->var_fcns) { - (*var_fcn)({&tval, y_val_, functions->inputs.data()}, {&res[0]}); - for (size_t jj = 0; jj < var_fcn->nnz_out(); jj++) { - y_dest[j++] = res[jj]; + if constexpr (ExprSet::kNativeBatchedOutputs) { + // Outputs read only y, so they are staged even when the sensitivities below run per point. + if (stage_out_t_.empty()) { + stage_out_first_ = i_save_; + stage_out_t_.reserve(kOutputBatchPoints); + stage_out_y_.reserve(kOutputBatchPoints * number_of_states); + } + stage_out_t_.push_back(tval); + stage_out_y_.insert(stage_out_y_.end(), y_val_, y_val_ + number_of_states); + if (stage_out_t_.size() == kOutputBatchPoints) { + FlushOutputBatch(); + } + } else { + sunrealtype* y_dest = &y[i_save_ * length_of_return_vector]; + size_t j = 0; + for (auto& var_fcn : functions->var_fcns) { + (*var_fcn)({&tval, y_val_, functions->inputs.data()}, {&res[0]}); + for (size_t jj = 0; jj < var_fcn->nnz_out(); jj++) { + y_dest[j++] = res[jj]; + } } } if (sensitivity) { - SetStepOutputSensitivities(tval); + if constexpr (ExprSet::kNativeOutputSensitivities) { + NativeSetStepOutputSensitivities(tval); + } else { + SetStepOutputSensitivities(tval); + } + } +} + +template +void IDAKLUSolverOpenMP::FlushOutputBatch() { + if constexpr (ExprSet::kNativeBatchedOutputs) { + const size_t k = stage_out_t_.size(); + if (k == 0) { + return; + } + DEBUG("IDAKLUSolver::FlushOutputBatch"); + // The staged points' y rows are contiguous from stage_out_first_. + functions->eval_outputs_batch( + stage_out_t_.data(), stage_out_y_.data(), static_cast(k), + &y[stage_out_first_ * length_of_return_vector]); + stage_out_t_.clear(); + stage_out_y_.clear(); } } @@ -1092,6 +1276,32 @@ void IDAKLUSolverOpenMP::SetStepOutputSensitivities( } } +template +void IDAKLUSolverOpenMP::NativeSetStepOutputSensitivities( + sunrealtype &tval +) { + DEBUG("IDAKLUSolver::NativeSetStepOutputSensitivities"); + const size_t n_params = static_cast(number_of_parameters); + const size_t stride = length_of_return_vector; + const size_t yS_base = i_save_ * n_params * stride; + + yS_flat_.resize(n_params * number_of_states); + for (size_t p = 0; p < n_params; ++p) { + std::copy(yS_val_[p], yS_val_[p] + number_of_states, + yS_flat_.begin() + p * number_of_states); + } + out_sens_flat_.resize(n_params * stride); + functions->project_output_sensitivities( + tval, y_val_, functions->inputs.data(), + yS_flat_.data(), out_sens_flat_.data()); + // The projection emits (param, output); yS is (output, param), as the CasADi path writes it. + for (size_t p = 0; p < n_params; ++p) { + for (size_t o = 0; o < stride; ++o) { + yS[yS_base + o * n_params + p] = out_sens_flat_[p * stride + o]; + } + } +} + template void IDAKLUSolverOpenMP::SetStepHermite( sunrealtype &tval diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolverLog.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolverLog.hpp index 5245e3b92d..d5a6c80715 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolverLog.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolverLog.hpp @@ -60,6 +60,11 @@ class SolverLog { append(format("Consistent initialization at t = %.17e", t_val)); } + void log_consistent_init_failed(double t_val, const char* reason) { + if (!enabled()) return; + append(format("IDACalcIC FAILED at t = %.17e (%s)", t_val, reason)); + } + void log_breakpoint(double t_val) { if (!enabled()) return; append(format("Breakpoint at t = %.17e, reinitializing", t_val)); diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp index 48a62535bf..1a330f4242 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp @@ -1,6 +1,10 @@ #include "StandaloneNewtonSolver.hpp" +#include "Expressions/Rust/RustFunctions.hpp" + #include +#include +#include // ────────────────────── StandaloneAlgebraicSystem ────────────────────── @@ -208,6 +212,51 @@ StandaloneNewtonSolver::StandaloneNewtonSolver( solver_.set_log(&log_); } +namespace { + +// `solver_` reads one tolerance per algebraic variable as it initialises. +int checked_atol_width(const std::vector& atol, int n_alg) { + if (atol.size() != static_cast(n_alg)) { + throw std::invalid_argument( + "pybammsolvers: atol has length " + std::to_string(atol.size()) + + " but n_alg is " + std::to_string(n_alg) + "."); + } + return n_alg; +} + +} // namespace + +StandaloneNewtonSolver::StandaloneNewtonSolver( + std::uintptr_t rust_model_ptr, + int n_rhs, + int n_alg, + const std::vector& jac_rows, + const std::vector& jac_cols, + const std::vector& atol, + sunrealtype rtol, + sunrealtype step_tol, + int max_iter, + int max_backtracks, + sunrealtype epsNewt, + bool use_sparse, + sunrealtype klu_pivot_tolerance) + : system_( + std::make_unique( + pybamm_rust_detail::model_from_address(rust_model_ptr), n_rhs, n_alg), + std::make_unique( + pybamm_rust_detail::model_from_address(rust_model_ptr), n_rhs, n_alg, + static_cast(jac_rows.size()), jac_rows, jac_cols), + checked_atol_width(atol, n_alg), + use_sparse, + klu_pivot_tolerance), + solver_(system_, system_.n_vars(), atol.data(), rtol, step_tol, + max_iter, max_backtracks, epsNewt), + n_vars_(system_.n_vars()), + y_work_(n_vars_) +{ + solver_.set_log(&log_); +} + std::pair StandaloneNewtonSolver::solve( sunrealtype t, const np_array& y0_np, diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp index 682377d44d..2fb889c51c 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp @@ -4,6 +4,7 @@ #include "NonlinearSolver.hpp" #include "Expressions/Casadi/CasadiFunctions.hpp" #include "common.hpp" +#include #include #include @@ -98,6 +99,22 @@ class StandaloneNewtonSolver { bool use_sparse, sunrealtype klu_pivot_tolerance); + /** @brief Rust-backed constructor from an opaque Rust model handle. */ + StandaloneNewtonSolver( + std::uintptr_t rust_model_ptr, + int n_rhs, + int n_alg, + const std::vector& jac_rows, + const std::vector& jac_cols, + const std::vector& atol, + sunrealtype rtol, + sunrealtype step_tol, + int max_iter, + int max_backtracks, + sunrealtype epsNewt, + bool use_sparse, + sunrealtype klu_pivot_tolerance); + /** * @brief Solve F(t, y, inputs) = 0 starting from y0. * @return (success, y_solution) with zero-copy numpy output. diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl index a96afa7bcd..27047c49ba 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl @@ -157,6 +157,8 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, // create pointer to jac data, column pointers, and row values sunrealtype *jac_data; + // Dense SUNMatrix with a sparse-emitting (Rust) expression: eval to tmp, then scatter below. + bool dense_needs_scatter = false; if (p_python_functions->setup_opts.using_sparse_matrix) { jac_data = SUNSparseMatrix_Data(JJ); @@ -166,7 +168,11 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, } else { - jac_data = SUNDenseMatrix_Data(JJ); + const sunindextype n_dense = SUNDenseMatrix_Rows(JJ) * SUNDenseMatrix_Columns(JJ); + dense_needs_scatter = p_python_functions->jac_times_cjmass->nnz_out() < n_dense; + jac_data = dense_needs_scatter + ? p_python_functions->get_tmp_sparse_jacobian_data() + : SUNDenseMatrix_Data(JJ); } DEBUG_VECTORn(yy, 100); @@ -254,6 +260,21 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, } else throw std::runtime_error("Unknown matrix format detected (Expected CSC or CSR)"); } + else if (dense_needs_scatter) + { + auto jac_colptrs = p_python_functions->jac_times_cjmass_colptrs.data(); + auto jac_rowvals = p_python_functions->jac_times_cjmass_rowvals.data(); + sunrealtype *dense = SUNDenseMatrix_Data(JJ); + const sunindextype n_rows = SUNDenseMatrix_Rows(JJ); + const sunindextype n_cols = SUNDenseMatrix_Columns(JJ); + std::memset(dense, 0, static_cast(n_rows) * n_cols * sizeof(sunrealtype)); + for (sunindextype col_ij = 0; col_ij < n_cols; col_ij++) { + for (auto data_i = jac_colptrs[col_ij]; data_i < jac_colptrs[col_ij + 1]; data_i++) { + const auto row_ij = jac_rowvals[data_i]; + dense[static_cast(col_ij) * n_rows + row_ij] = jac_data[data_i]; + } + } + } return (0); } @@ -319,19 +340,29 @@ int sensitivities_eval(int Ns, sunrealtype t, N_Vector yy, N_Vector yp, // resvalsS now has (∂F/∂p i ) (*p_python_functions->sens)(); + const int ns = p_python_functions->number_of_states; + if constexpr (T::kNativeBatchedJacActions) { + // Every (∂F/∂y)s i (t) at once, sharing one primal pass; bitwise the per-i products. + const sunrealtype *jac_yS = + p_python_functions->jac_action_sensitivities(t, NV_DATA(yy), yS); + for (int i = 0; i < np; i++) { + axpy(ns, 1., jac_yS + static_cast(i) * ns, NV_DATA(resvalS[i])); + } + } + for (int i = 0; i < np; i++) { - // put (∂F/∂y)s i (t) in tmp sunrealtype *tmp = p_python_functions->get_tmp_state_vector(); - p_python_functions->jac_action->m_arg[0] = &t; - p_python_functions->jac_action->m_arg[1] = NV_DATA(yy); - p_python_functions->jac_action->m_arg[2] = p_python_functions->inputs.data(); - p_python_functions->jac_action->m_arg[3] = NV_DATA(yS[i]); - p_python_functions->jac_action->m_res[0] = tmp; - (*p_python_functions->jac_action)(); - - const int ns = p_python_functions->number_of_states; - axpy(ns, 1., tmp, NV_DATA(resvalS[i])); + if constexpr (!T::kNativeBatchedJacActions) { + // put (∂F/∂y)s i (t) in tmp + p_python_functions->jac_action->m_arg[0] = &t; + p_python_functions->jac_action->m_arg[1] = NV_DATA(yy); + p_python_functions->jac_action->m_arg[2] = p_python_functions->inputs.data(); + p_python_functions->jac_action->m_arg[3] = NV_DATA(yS[i]); + p_python_functions->jac_action->m_res[0] = tmp; + (*p_python_functions->jac_action)(); + axpy(ns, 1., tmp, NV_DATA(resvalS[i])); + } // put -(∂F/∂ ẏ) ṡ i (t) in tmp2 p_python_functions->mass_action->m_arg[0] = NV_DATA(ypS[i]); diff --git a/packages/pybammsolvers/src/pybammsolvers/version.py b/packages/pybammsolvers/src/pybammsolvers/version.py index adad20be2e..5c726d306e 100644 --- a/packages/pybammsolvers/src/pybammsolvers/version.py +++ b/packages/pybammsolvers/src/pybammsolvers/version.py @@ -1,3 +1,3 @@ # Single source of truth for the pybammsolvers version. Bump this at release; # a pre-commit hook mirrors it into pyproject.toml (see .pre-commit-config.yaml). -__version__ = "0.11.1" +__version__ = "0.12.0" diff --git a/packages/pybammsolvers/tests/test_consistent_initialization.py b/packages/pybammsolvers/tests/test_consistent_initialization.py index 691e59e5f3..f10acc356d 100644 --- a/packages/pybammsolvers/tests/test_consistent_initialization.py +++ b/packages/pybammsolvers/tests/test_consistent_initialization.py @@ -2,6 +2,7 @@ from __future__ import annotations +import numpy as np import pytest pytest.importorskip("pybamm", reason="PyBaMM not installed") @@ -19,3 +20,59 @@ def test_internal_ic_keeps_converged_state_at_tiny_atol(): sol = sim.solve([0, 3600]) assert sol["Voltage [V]"].entries.size > 1 + + +class TestSensitivityInitialConditions: + """Initial sensitivities of algebraic states, at tolerances near roundoff. + + IDACalcIC's sensitivity correction used to start from the zero algebraic + guess, failing at some inputs and leaving those sensitivities zero at t0. + """ + + INPUT = "Current function [A]" + + @staticmethod + def _model(): + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values.update({"Current function [A]": "[input]"}) + simulation = pybamm.Simulation( + pybamm.lithium_ion.DFN({"contact resistance": "true"}), + parameter_values=parameter_values, + var_pts={"x_n": 6, "x_s": 5, "x_p": 6, "r_n": 8, "r_p": 8}, + ) + simulation.build() + model = simulation.built_model + model.convert_to_format = "casadi" + return model + + @pytest.mark.parametrize("newton_mode", ["auto", "full"]) + def test_algebraic_sensitivities_at_t0_match_finite_differences(self, newton_mode): + model = self._model() + solver = pybamm.IDAKLUSolver( + rtol=1e-11, atol=1e-13, options={"newton_mode": newton_mode} + ) + n_states = model.len_rhs_and_alg + algebraic = slice(model.len_rhs, n_states) + for current in (4.5, 4.875, 5.0, 5.125, 5.375, 5.5): + solution = solver.solve( + model, + [0, 1], + inputs={self.INPUT: current}, + calculate_sensitivities=True, + ) + # Stored time-major, so the first rows are the states at t0 + sensitivity = solution.sensitivities[self.INPUT][:n_states, 0] + + step = 1e-3 * current + plus, minus = ( + solver.solve(model, [0, 1], inputs={self.INPUT: current + sign * step}) + for sign in (1, -1) + ) + expected = (plus.y[algebraic, 0] - minus.y[algebraic, 0]) / (2 * step) + np.testing.assert_allclose( + sensitivity[algebraic], + expected, + rtol=0, + atol=1e-5 * np.abs(expected).max(), + err_msg=f"current {current}", + ) diff --git a/packages/pybammsolvers/vcpkg.json b/packages/pybammsolvers/vcpkg.json index f35c6e3bdf..b06184163f 100644 --- a/packages/pybammsolvers/vcpkg.json +++ b/packages/pybammsolvers/vcpkg.json @@ -1,6 +1,6 @@ { "name": "pybammsolvers", - "version-string": "0.11.1", + "version-string": "0.12.0", "dependencies": [ "casadi", { diff --git a/uv.lock b/uv.lock index 5102be82f7..67dd5d88e1 100644 --- a/uv.lock +++ b/uv.lock @@ -4366,7 +4366,7 @@ provides-extras = ["zoo-all", "zoo-multilayer-3d-thermal"] [[package]] name = "pybammsolvers" -version = "0.11.1" +version = "0.12.0" source = { editable = "packages/pybammsolvers" } dependencies = [ { name = "casadi" },