Repository navigation
feat: block preconditioners for IDAKLU - #6009
Draft
MarcBerliner wants to merge 1 commit into
Draft
MarcBerliner wants to merge 1 commit into
MarcBerliner wants to merge 1 commit into
Conversation
- pybamm.preconditioners: block solvers (KLU, ILU, Jacobi, Tridiagonal) keyed by state variable, staged Schur eliminations and additive or multiplicative stages, compiled at discretisation; extensible through a C block-solver plugin ABI. - Jacobian-free path: an iterative linear solver with preconditioner='block' computes only the Jacobian entries the preconditioner reads, from a structural sparsity pattern with bidirectional colouring, folds linear operators and computes input-only values once per solve. - BaseModel.jacobian_sparsity and spy; IDAKLUSolver.spy is deprecated. - Fixes for newton_mode='full' initial conditions, first-step error reporting, discretised model serialisation of the preconditioner, and experiment state mappers with per-cell inputs.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| void example_destroy(void *context) { delete static_cast<ExampleBlockSolver *>(context); } | ||
|
|
||
| pybamm_block_solver example_block_solver = { |
| PYBAMM_BLOCK_SOLVER_CAPSULE); | ||
| }, "Capsule of an example block solver plugin (Jacobi), for tests"); | ||
|
|
||
| py::class_<BlockAnalysis, std::shared_ptr<BlockAnalysis>>(m, "BlockAnalysis") |
| }, "Sorted positions of the iteration matrix's CSC data that the tree reads") | ||
| .def_readonly("size", &BlockAnalysis::size); | ||
|
|
||
| py::class_<BlockPreconditioner>(m, "BlockPreconditioner") |
| std::string("preconditioner tree node is missing \"") + key + "\""); | ||
| } | ||
| auto array = node[key].cast< | ||
| py::array_t<int64_t, py::array::c_style | py::array::forcecast>>(); |
| return static_cast<PreconditionerSolverContent *>(S->content); | ||
| } | ||
|
|
||
| SUNLinearSolver_Type get_type(SUNLinearSolver) { return SUNLINEARSOLVER_DIRECT; } |
|
|
||
| py::class_<BlockAnalysis, std::shared_ptr<BlockAnalysis>>(m, "BlockAnalysis") | ||
| .def(py::init([](const py::dict &tree, | ||
| py::array_t<int64_t, py::array::c_style | py::array::forcecast> indptr, |
| py::class_<BlockAnalysis, std::shared_ptr<BlockAnalysis>>(m, "BlockAnalysis") | ||
| .def(py::init([](const py::dict &tree, | ||
| py::array_t<int64_t, py::array::c_style | py::array::forcecast> indptr, | ||
| py::array_t<int64_t, py::array::c_style | py::array::forcecast> indices) { |
| .def_readonly("size", &BlockAnalysis::size); | ||
|
|
||
| py::class_<BlockPreconditioner>(m, "BlockPreconditioner") | ||
| .def(py::init([](std::shared_ptr<BlockAnalysis> analysis) { |
| .def(py::init([](std::shared_ptr<BlockAnalysis> analysis) { | ||
| return std::make_unique<BlockPreconditioner>(analysis); | ||
| }), py::arg("analysis")) | ||
| .def("setup", [](BlockPreconditioner &self, np_array matrix_data) { |
| return self.setup(matrix_data.data()); | ||
| }, "Set up from the iteration matrix's CSC data; 0 on success", | ||
| py::arg("matrix_data")) | ||
| .def("solve", [](BlockPreconditioner &self, np_array r) { |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds block preconditioners to
IDAKLUSolver, so models with large spatial fields (for example a 3D thermal field coupled to an electrochemical model) solve with an iterative linear solver whose set-up, solve time and memory scale roughly linearly with the number of states.Block preconditioners (
pybamm.preconditioners)KLU,ILU,Jacobi,Tridiagonal. Unlisted variables get an exactKLUblock in the last stage, with that stage's settings.model.default_preconditionerholds the model's default;Simulation(preconditioner=...)andDiscretisation(preconditioner=...)override it, as withvar_pts. It is compiled against the states at discretisation (model.preconditioner); nothing needs the model before then.Stagenests Schur eliminations and combines a stage's blocks additively (block Jacobi) or multiplicatively (block Gauss-Seidel).options={"linear_solver": "SUNLinSol_SPGMR", "preconditioner": "block"}applies it to a Krylov solver;options={"linear_solver": "block"}applies an exact one directly.pybammsolvers.get_include()/pybamm_block_solver.h) and subclassBlockSolver. Block solvers, the compiled tree, and models carrying them all serialise (JSON and pickle).Jacobian-free path
With an iterative linear solver and
preconditioner="block", Krylov products come from the Jacobian action, so only the Jacobian entries the preconditioner reads are computed:pybamm.JacobianSparsitypropagates a structural sparsity pattern up the expression tree (each symbol class defines its rule), avoiding CasADi's sparsity detection, which is quadratic when a Jacobian has both a dense row and a dense column (a field mean read by every equation).pybamm.casadi_jacobiandifferentiates the densest rows in reverse mode and the rest in forward mode. On a 51k-state model with such a mean, forward colouring alone needs 50,562 Jacobian-vector products per evaluation; bidirectional colouring needs 18.pybamm.CasadiConversionfolds linear operators with input-only coefficients into one sparse product and computes input-only values once per solve rather than at every evaluation.The default direct path (
SUNLinSol_KLU) is unchanged: the same Jacobian, evaluation cost and results asmain.Other changes
BaseModel.jacobian_sparsity()andBaseModel.spy()work on any discretised model;IDAKLUSolver.spy/get_jacobian_sparsityare deprecated.newton_mode="full"now converges consistent initial conditions (identity rows for differential states) and no longer segfaults with iterative linear solvers;BBDP, which reads IDA state that is unset there, raises aSolverErrorwith it.IDA_CONV_FAIL) instead ofIDA_BAD_K.save_modelkeep their compiled preconditioner.Benchmarks
SPM, SPMe and DFN coupled to a 3D heat equation on an unstructured hex mesh of a cylindrical cell; 600 s discharge; block Gauss-Seidel with
ILU(fill=1)on the temperature field andKLUelsewhere.For comparison, the default
SUNLinSol_KLUpath takes about 15 s to set up and 10 s to solve at 6k states. A CC-CV cycle (1C discharge, rest, C/2 charge, CV hold to C/50) at about 100k states runs in 27 s (SPM) and 64 s (DFN), peak 3.8 GB.On standard models the default path keeps
main's Jacobian and results; CC-CV with the block preconditioner on the DFN matchesSUNLinSol_KLUto within 2 mV.Type of change
Feature, with one breaking change (the
"preconditioner"option defaults to"block", which only affects iterative linear solvers) and one deprecation; seeCHANGELOG.md.Important checks:
Please confirm the following before marking the PR as ready for review:
nox -s pre-commitnox -s testsnox -s doctests