diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c4f304934..8ea4f696ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ - 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)) - `IDAKLUSolver` solves `convert_to_format = "rust"` models through a Rust `CompiledModel`, one evaluator per thread, so `num_threads` parallelises input sets. `output_variables` support sensitivities and time integrals, and `NonlinearSolver` compiles its own model. ([#5962](https://github.com/pybamm-team/PyBaMM/pull/5962)) - `IDAKLUSolver` observes rust-model solutions natively: variables and their sensitivities are evaluated on tapes compiled into the solve's Rust graph, with no CasADi function, and `Solution.update([...])` compiles a set of variables into one tape per segment. ([#5963](https://github.com/pybamm-team/PyBaMM/pull/5963)) +- `pybamm-core` gains an in-process DAE solver built on diffsol's BDF, with forward sensitivities under error control, batched one lane per parameter, and a rayon batch layer for input-set sweeps. ([#5964](https://github.com/pybamm-team/PyBaMM/pull/5964)) ## Bug fixes diff --git a/packages/pybamm-rust/Cargo.lock b/packages/pybamm-rust/Cargo.lock index 5177d0a95a..177327fbcf 100644 --- a/packages/pybamm-rust/Cargo.lock +++ b/packages/pybamm-rust/Cargo.lock @@ -23,6 +23,25 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "atomic-wait" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55b94919229f2c42292fd71ffa4b75e83193bffdd77b1e858cd55fd2d0b0ea8" +dependencies = [ + "libc", + "windows-sys 0.42.0", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -70,6 +89,26 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "cast" @@ -77,6 +116,16 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "cc" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -171,6 +220,28 @@ dependencies = [ "itertools", ] +[[package]] +name = "crossbeam" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -190,6 +261,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" @@ -202,12 +282,148 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "defer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "930c7171c8df9fb1782bdf9b918ed9ed2d33d1d22300abb754f9085bc48bf8e8" + +[[package]] +name = "diffsol" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1133b60811dbb117af7e1186b2794a8493425ed3e9f3cd8c6073b27a3bffcb77" +dependencies = [ + "diffsol-la", + "diffsol-nl", + "log", + "num-traits", + "petgraph", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "diffsol-la" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5214d0559ce59380d0ef141601c48b0a5bcbd12db8a647e128a98acb5c270545" +dependencies = [ + "faer", + "faer-traits", + "nalgebra", + "num-traits", + "thiserror 2.0.19", +] + +[[package]] +name = "diffsol-nl" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8798c6a2508b0b284fa7600f0430049ecf59233f4c3f7b8f2eaefb58b3a397e8" +dependencies = [ + "diffsol-la", + "log", + "num-traits", + "thiserror 2.0.19", +] + +[[package]] +name = "dyn-stack" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c4713e43e2886ba72b8271aa66c93d722116acf7a75555cce11dcde84388fe8" +dependencies = [ + "bytemuck", + "dyn-stack-macros", +] + +[[package]] +name = "dyn-stack-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" + [[package]] name = "either" version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "equator" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c35da53b5a021d2484a7cc49b2ac7f2d840f8236a286f84202369bd338d761ea" +dependencies = [ + "equator-macro 0.2.1", +] + +[[package]] +name = "equator" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" +dependencies = [ + "equator-macro 0.4.2", +] + +[[package]] +name = "equator" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02da895aab06bbebefb6b2595f6d637b18c9ff629b4cd840965bb3164e4194b0" +dependencies = [ + "equator-macro 0.6.0", +] + +[[package]] +name = "equator-macro" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bf679796c0322556351f287a51b49e48f7c4986e727b5dd78c972d30e2e16cc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "equator-macro" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "equator-macro" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b14b339eb76d07f052cdbad76ca7c1310e56173a138095d3bf42a23c06ef5d8" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "errno" version = "0.3.14" @@ -215,7 +431,50 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "faer" +version = "0.24.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ab6df3dd147fe8d702a288b95bcd8fcc499ab572fc80da6828f60cd4d524d67" +dependencies = [ + "bytemuck", + "dyn-stack", + "equator 0.6.0", + "faer-traits", + "gemm", + "generativity", + "libm", + "nano-gemm", + "npyz", + "num-complex", + "num-traits", + "private-gemm-x86", + "pulp", + "rand 0.9.5", + "rand_distr", + "rayon", + "reborrow", + "spindle", +] + +[[package]] +name = "faer-traits" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b87d23ed7ab1f26c0cba0e5b9e061a796fbb7dc170fa8bee6970055a1308bb0f" +dependencies = [ + "bytemuck", + "dyn-stack", + "generativity", + "libm", + "num-complex", + "num-traits", + "pulp", + "qd", + "reborrow", ] [[package]] @@ -224,12 +483,30 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + [[package]] name = "fnv" version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "futures-core" version = "0.3.33" @@ -254,6 +531,146 @@ dependencies = [ "slab", ] +[[package]] +name = "gemm" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb" +dependencies = [ + "dyn-stack", + "gemm-c32", + "gemm-c64", + "gemm-common", + "gemm-f16", + "gemm-f32", + "gemm-f64", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "086936dbdcb99e37aad81d320f98f670e53c1e55a98bee70573e83f95beb128c" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20c8aeeeec425959bda4d9827664029ba1501a90a0d1e6228e48bef741db3a3f" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88027625910cc9b1085aaaa1c4bc46bb3a36aad323452b33c25b5e4e7c8e2a3e" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3df7a55202e6cd6739d82ae3399c8e0c7e1402859b30e4cb780e61525d9486e" +dependencies = [ + "dyn-stack", + "gemm-common", + "gemm-f32", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02e0b8c9da1fbec6e3e3ab2ce6bc259ef18eb5f6f0d3e4edf54b75f9fd41a81c" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "056131e8f2a521bfab322f804ccd652520c79700d81209e9d9275bbdecaadc6a" +dependencies = [ + "dyn-stack", + "gemm-common", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "generativity" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2c81fb5260e37854d09d5c87183309fd8c555b75289427884b25660bc87a85e" + +[[package]] +name = "generator" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3b854b0e584ead1a33f18b2fcad7cf7be18b3875c78816b753639aa501513ae" +dependencies = [ + "cc", + "cfg-if", + "libc", + "log", + "rustversion", + "windows-link", + "windows-result", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -277,17 +694,58 @@ dependencies = [ "r-efi 6.0.0", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "half" version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", + "num-traits", "zerocopy", ] +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -300,6 +758,27 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "interpol" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb58032ba748f4010d15912a1855a8a0b1ba9eaad3395b0c171c09b3b356ae50" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "is-terminal" version = "0.4.17" @@ -308,7 +787,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -337,6 +816,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.189" @@ -355,6 +840,34 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "loom" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" +dependencies = [ + "cfg-if", + "generator", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "matrixmultiply" version = "0.3.11" @@ -371,6 +884,107 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "nano-gemm" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e04345dc84b498ff89fe0d38543d1f170da9e43a2c2bcee73a0f9069f72d081" +dependencies = [ + "equator 0.2.2", + "nano-gemm-c32", + "nano-gemm-c64", + "nano-gemm-codegen", + "nano-gemm-core", + "nano-gemm-f32", + "nano-gemm-f64", + "num-complex", +] + +[[package]] +name = "nano-gemm-c32" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0775b1e2520e64deee8fc78b7732e3091fb7585017c0b0f9f4b451757bbbc562" +dependencies = [ + "nano-gemm-codegen", + "nano-gemm-core", + "num-complex", +] + +[[package]] +name = "nano-gemm-c64" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9af49a20d58816e6b5ee65f64142e50edb5eba152678d4bb7377fcbf63f8437a" +dependencies = [ + "nano-gemm-codegen", + "nano-gemm-core", + "num-complex", +] + +[[package]] +name = "nano-gemm-codegen" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cc8d495c791627779477a2cf5df60049f5b165342610eb0d76bee5ff5c5d74c" + +[[package]] +name = "nano-gemm-core" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d998dfa644de87a0f8660e5ea511d7cb5c33b5a2d9847b7af57a2565105089f0" + +[[package]] +name = "nano-gemm-f32" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879d962e79bc8952e4ad21ca4845a21132540ed3f5e01184b2ff7f720e666523" +dependencies = [ + "nano-gemm-codegen", + "nano-gemm-core", +] + +[[package]] +name = "nano-gemm-f64" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9a513473dce7dc00c7e7c318481ca4494034e76997218d8dad51bd9f007a815" +dependencies = [ + "nano-gemm-codegen", + "nano-gemm-core", +] + [[package]] name = "ndarray" version = "0.17.2" @@ -386,6 +1000,36 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "npyz" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f0e759e014e630f90af745101b614f761306ddc541681e546649068e25ec1b9" +dependencies = [ + "byteorder", + "num-bigint", + "py_literal", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -394,6 +1038,7 @@ checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" dependencies = [ "bytemuck", "num-traits", + "rand 0.8.7", ] [[package]] @@ -405,6 +1050,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -412,6 +1068,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", ] [[package]] @@ -448,6 +1115,60 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pest" +version = "2.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df728be843c7070fab6ab7c328c4e9e9d78e23bf749c0669c86ee7ebfa050a2" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e2dd6fc3b26b3462ee188aac870f5a41d398f1cd5e2408d16531bd71c9591fd" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a7a9205cfb6f596a9e8b689c0a15f9ceb7a1aafae7aaf788150ac65b29975b6" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pest_meta" +version = "2.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85abd351c0de1e8384fc791a0737111a350394937e92b956b743dac12429f57c" +dependencies = [ + "pest", +] + +[[package]] +name = "petgraph" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" +dependencies = [ + "fixedbitset", + "hashbrown 0.15.5", + "indexmap", + "serde", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -506,6 +1227,22 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "private-gemm-x86" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0af8c3e5087969c323f667ccb4b789fa0954f5aa650550e38e81cf9108be21b5" +dependencies = [ + "crossbeam", + "defer", + "interpol", + "num_cpus", + "raw-cpuid", + "rayon", + "spindle", + "sysctl", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -525,7 +1262,7 @@ dependencies = [ "bit-vec", "bitflags", "num-traits", - "rand", + "rand 0.9.5", "rand_chacha", "rand_xorshift", "regex-syntax", @@ -557,18 +1294,34 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d8f70e07b9c3962945a74e59ca1c511bba65b6419468acc217c457d93f3c740" +[[package]] +name = "py_literal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1" +dependencies = [ + "num-bigint", + "num-complex", + "num-traits", + "pest", + "pest_derive", +] + [[package]] name = "pybamm-core" version = "0.1.0" dependencies = [ "bincode", "criterion", + "diffsol", + "faer", "libm", "proptest", "pulp", + "rayon", "rustc-hash", "serde", - "thiserror", + "thiserror 2.0.19", ] [[package]] @@ -639,6 +1392,18 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "qd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15f1304a5aecdcfe9ee72fbba90aa37b3aa067a69d14cb7f3d9deada0be7c07c" +dependencies = [ + "bytemuck", + "libm", + "num-traits", + "pulp", +] + [[package]] name = "quick-error" version = "1.2.3" @@ -666,6 +1431,15 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" +dependencies = [ + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.5" @@ -673,7 +1447,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha", - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -683,9 +1457,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", ] +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" + [[package]] name = "rand_core" version = "0.9.5" @@ -695,13 +1475,23 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.5", +] + [[package]] name = "rand_xorshift" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -790,7 +1580,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -811,6 +1601,15 @@ dependencies = [ "wait-timeout", ] +[[package]] +name = "safe_arch" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a52ec151f024d703f9fd65abb7cbe81e7cdb39f18917a3a37e3014470dc7c59" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -820,6 +1619,18 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.229" @@ -863,12 +1674,69 @@ dependencies = [ "zmij", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "slab" version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spindle" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aaca3d8aa5387a6eba861fbf984af5348d9df5d940c25c6366b19556fdf64" +dependencies = [ + "atomic-wait", + "crossbeam", + "equator 0.4.2", + "loom", + "rayon", +] + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.119" @@ -891,6 +1759,20 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "sysctl" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc" +dependencies = [ + "bitflags", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + [[package]] name = "target-lexicon" version = "0.13.5" @@ -907,7 +1789,16 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", ] [[package]] @@ -916,7 +1807,18 @@ version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.19", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -930,6 +1832,15 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tinytemplate" version = "1.2.1" @@ -940,6 +1851,67 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + [[package]] name = "unarray" version = "0.1.4" @@ -952,6 +1924,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "version_check" version = "0.9.5" @@ -1041,13 +2019,23 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi-util" version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1056,6 +2044,30 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -1065,6 +2077,48 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/packages/pybamm-rust/Cargo.toml b/packages/pybamm-rust/Cargo.toml index 30221d8cdb..b84572e81d 100644 --- a/packages/pybamm-rust/Cargo.toml +++ b/packages/pybamm-rust/Cargo.toml @@ -5,6 +5,7 @@ members = [ "pybamm-python", ] +# Floor set by the dep graph (nalgebra 0.35 needs 1.89), not by edition 2024. [workspace.package] rust-version = "1.89" diff --git a/packages/pybamm-rust/README.md b/packages/pybamm-rust/README.md index 82d7624ba5..6a464d3559 100644 --- a/packages/pybamm-rust/README.md +++ b/packages/pybamm-rust/README.md @@ -2,7 +2,8 @@ Expression compiler and evaluator for PyBaMM models, filling the role CasADi plays on PyBaMM's other backends. Python hands over a discretised model as an expression DAG; -this workspace turns it into flat instruction tapes it can evaluate and differentiate. +this workspace turns it into flat instruction tapes it can evaluate, differentiate and +solve. Models opt in with `model.convert_to_format = "rust"`. @@ -25,6 +26,10 @@ module, and explains how the Jacobian is assembled (see ## Who consumes it +Two paths, and they differ in how they reach the crate: + +- **diffsol**, in-process. The `solver` module implements one operator trait per + callback diffsol needs and runs the integration entirely in Rust. - **IDAKLU**, through the C ABI in `ffi.rs`. `pybammsolvers` cannot link against this crate, because PyBaMM depends on `pybammsolvers` and not the reverse. So the FFI entry points are compiled into the Python extension, and the C++ side resolves them at @@ -97,5 +102,6 @@ doc comments are the reference; this README only orients you. - **MSRV and dependency versions** live in the workspace `Cargo.toml`. The Rust floor is declared there rather than implied by the edition, so check it before assuming a newer language feature is available. -- **Feature flags**: `serialize` enables snapshot support that some tests and benches require. - A `required-features` error means the target needs it. +- **Feature flags**: `diffsol` (on by default) pulls in the integrator and its matrix + backend; `serialize` enables snapshot support that some tests and benches require. A + `required-features` error means the target needs one of these. diff --git a/packages/pybamm-rust/pybamm-core/Cargo.toml b/packages/pybamm-rust/pybamm-core/Cargo.toml index 9d58f30d46..51e4182c6b 100644 --- a/packages/pybamm-rust/pybamm-core/Cargo.toml +++ b/packages/pybamm-rust/pybamm-core/Cargo.toml @@ -12,6 +12,9 @@ thiserror = "2" libm = "0.2" serde = { version = "1", features = ["derive"], optional = true } bincode = { version = "1", optional = true } +diffsol = { version = "0.18.0", optional = true, default-features = false, features = ["faer"] } +faer = { version = "0.24", optional = true } +rayon = { version = "1", optional = true } [target.'cfg(target_arch = "x86_64")'.dependencies] # Compiles matrix products for the detected instruction set, so `accumulate` needs no `unsafe`. @@ -22,7 +25,13 @@ criterion = { version = "0.5", features = ["html_reports"] } proptest = "1" [features] +default = ["diffsol"] serialize = ["serde", "bincode"] +diffsol = ["dep:diffsol", "dep:faer", "dep:rayon"] + +[[test]] +name = "test_faer_global_parallelism" +required-features = ["diffsol"] [[bench]] name = "compile" @@ -37,6 +46,16 @@ name = "sparsity" harness = false required-features = ["serialize"] +[[bench]] +name = "solver_setup" +harness = false +required-features = ["serialize", "diffsol"] + +[[bench]] +name = "sens_params" +harness = false +required-features = ["diffsol"] + [[bench]] name = "derivatives" harness = false diff --git a/packages/pybamm-rust/pybamm-core/benches/sens_params.rs b/packages/pybamm-rust/pybamm-core/benches/sens_params.rs new file mode 100644 index 0000000000..e46cbb34e3 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/sens_params.rs @@ -0,0 +1,82 @@ +//! Benchmark: the saving from integrating only the requested sensitivity columns. + +mod helpers; + +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use helpers::identity_mass_matrix; +use pybamm_core::solver::solve::{InputSet, PreparedSolver, SolveRequest}; +use pybamm_core::{Arena, CompiledModel, CompiledModelOptions, Node, PackedWidth, ParamIdx}; + +const N_STATES: usize = 8; +const N_PARAMS: usize = 16; + +/// `dy_i/dt = -(sum_j p_j) * y_i`: every state depends on every parameter. +fn build_dense_decay_model(sens: &[ParamIdx]) -> CompiledModel { + let mut arena = Arena::new(); + let mut sum = arena.alloc(Node::Scalar(0.0)); + for index in 0..N_PARAMS { + let p = arena.alloc(Node::InputParameter { + name: format!("p{index}"), + index, + offset: index, + width: 1, + }); + sum = arena.alloc(Node::Add(sum, p)); + } + let neg = arena.alloc(Node::Scalar(-1.0)); + let rate = arena.alloc(Node::Mul(neg, sum)); + + let rows: Vec<_> = (0..N_STATES) + .map(|i| { + let y_i = arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }); + arena.alloc(Node::Mul(rate, y_i)) + }) + .collect(); + let rhs = arena.alloc(Node::Concat(rows)); + + CompiledModel::new( + &arena, + rhs, + identity_mass_matrix(N_STATES), + N_STATES, + PackedWidth::new(N_PARAMS), + CompiledModelOptions { + sensitivities: sens, + algebraic: None, + }, + ) + .expect("a valid fixture") +} + +fn bench_sens_params(c: &mut Criterion) { + let mut group = c.benchmark_group("sens_params"); + let all: Vec = (0..N_PARAMS).map(ParamIdx::from_usize).collect(); + let y0 = vec![1.0; N_STATES]; + let inputs = vec![1.0 / N_PARAMS as f64; N_PARAMS]; + let atol = vec![1e-8; N_STATES]; + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.05).collect(); + + for k in [1usize, 4, N_PARAMS] { + let prepared = PreparedSolver::new(build_dense_decay_model(&all[..k]), 1e-8, &atol) + .expect("PreparedSolver failed"); + group.bench_function(BenchmarkId::from_parameter(k), |b| { + b.iter(|| { + black_box( + prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("sensitivity solve failed"), + ) + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_sens_params); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/benches/solver_setup.rs b/packages/pybamm-rust/pybamm-core/benches/solver_setup.rs new file mode 100644 index 0000000000..37a0e61cce --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/solver_setup.rs @@ -0,0 +1,146 @@ +//! Benchmark: prepare-once cost, and reuse versus no reuse of that handle. +//! Per-solve diffsol setup runs inside `solve`; read it off `SolverStatistics`. + +mod helpers; + +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use helpers::{FIXTURES, compile_fixture, fixture_state}; +use pybamm_core::DagSnapshot; +use pybamm_core::solver::solve::{InputSet, PreparedSolver, SolveRequest}; + +/// Benchmark A: `PreparedSolver::new()` only, excluding model compilation. +fn bench_solver_setup(c: &mut Criterion) { + let mut group = c.benchmark_group("solver_setup"); + group.sample_size(100); + + for case in FIXTURES { + let snap = DagSnapshot::from_bytes(case.bytes).expect("fixture decodes"); + + group.bench_with_input( + BenchmarkId::from_parameter(case.name), + &case.name, + |b, _| { + b.iter_batched( + || compile_fixture(&snap), + |model| { + let n_states = model.n_states(); + let atol = vec![1e-6; n_states]; + let prepared = PreparedSolver::new(model, 1e-6, &atol).unwrap(); + black_box(prepared); + }, + criterion::BatchSize::SmallInput, + ); + }, + ); + } + + group.finish(); +} + +/// Benchmark B: `solve()` on an already-constructed `PreparedSolver` (reuse path) +fn bench_solve_reuse(c: &mut Criterion) { + let mut group = c.benchmark_group("solve_reuse"); + group.sample_size(50); + + for case in FIXTURES { + let snap = DagSnapshot::from_bytes(case.bytes).expect("fixture decodes"); + let model = compile_fixture(&snap); + let n_states = model.n_states(); + let (y0, inputs) = fixture_state(n_states, model.input_width().get()); + let atol = vec![1e-6; n_states]; + + let prepared = PreparedSolver::new(model, 1e-6, &atol).unwrap(); + + let t_eval: Vec = (0..100).map(|i| f64::from(i) * 36.0).collect(); + + if prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .is_err() + { + eprintln!( + "Skipping {} solve_reuse: dummy y0 doesn't converge", + case.name + ); + continue; + } + + group.bench_with_input( + BenchmarkId::from_parameter(case.name), + &case.name, + |b, _| { + b.iter(|| { + let result = prepared + .solve( + SolveRequest::new(black_box(&t_eval)), + InputSet::new(black_box(&y0), black_box(&inputs)), + ) + .unwrap(); + black_box(&result); + }); + }, + ); + } + + group.finish(); +} + +/// Benchmark C: fresh `PreparedSolver::new()` + `solve()` (no-reuse path). +fn bench_fresh_setup_and_solve(c: &mut Criterion) { + let mut group = c.benchmark_group("fresh_setup_and_solve"); + group.sample_size(50); + + for case in FIXTURES { + let snap = DagSnapshot::from_bytes(case.bytes).expect("fixture decodes"); + let n_states = snap.n_states; + let (y0, inputs) = fixture_state(n_states, snap.packed_input_width.get()); + let t_eval: Vec = (0..100).map(|i| f64::from(i) * 36.0).collect(); + + { + let model = compile_fixture(&snap); + let atol = vec![1e-6; n_states]; + let prepared = PreparedSolver::new(model, 1e-6, &atol).unwrap(); + if prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .is_err() + { + eprintln!( + "Skipping {} fresh_setup_and_solve: dummy y0 doesn't converge", + case.name + ); + continue; + } + } + + group.bench_with_input( + BenchmarkId::from_parameter(case.name), + &case.name, + |b, _| { + b.iter_batched( + || compile_fixture(&snap), + |model| { + let atol = vec![1e-6; n_states]; + let prepared = PreparedSolver::new(model, 1e-6, &atol).unwrap(); + let result = prepared + .solve( + SolveRequest::new(black_box(&t_eval)), + InputSet::new(black_box(&y0), black_box(&inputs)), + ) + .unwrap(); + black_box(&result); + }, + criterion::BatchSize::SmallInput, + ); + }, + ); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_solver_setup, + bench_solve_reuse, + bench_fresh_setup_and_solve, +); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/src/error.rs b/packages/pybamm-rust/pybamm-core/src/error.rs index 63febe3cec..c95ebf8b61 100644 --- a/packages/pybamm-rust/pybamm-core/src/error.rs +++ b/packages/pybamm-rust/pybamm-core/src/error.rs @@ -26,6 +26,30 @@ pub enum CoreError { #[error("InterpolantNd: expected {axes} children (one per axis), got {children}")] InterpolantNdArity { axes: usize, children: usize }, + #[error("t_eval must not be empty")] + EmptyTimePoints, + + /// The evaluation time grid decreased at some point. + #[error("t_eval must increase monotonically, but t_eval[{index}] = {got} follows {previous}")] + UnsortedTimePoints { + index: usize, + got: f64, + previous: f64, + }, + + /// The evaluation time grid held a NaN or an infinity. + #[error("t_eval must be finite, but t_eval[{index}] = {got}")] + NonFiniteTimePoint { index: usize, got: f64 }, + + #[error("y0 has length {got} but the model has {expected} states")] + Y0Length { got: usize, expected: usize }, + + #[error("inputs has length {got} but the model expects {expected} packed parameter value(s)")] + InputsLength { got: usize, expected: usize }, + + #[error("atol has length {got} but the model has {expected} states")] + AtolLength { got: usize, expected: usize }, + /// A sensitivity index names no registered input parameter. #[error( "sensitivity index {index} names no input parameter; the model registers {n_registered}" @@ -154,6 +178,45 @@ pub enum CoreError { )] AlgebraicBlockNotSquare { rows: usize, columns: usize }, + #[error("no sensitivity parameters were requested when the model was compiled")] + NoSensitivityParams, + + #[error("no output variables were registered when the model was compiled")] + NoOutputVariables, + + #[error( + "y0_sens has length {got} but must be empty or {expected} (n_states x n_sens_params, column-major)" + )] + Y0SensLength { got: usize, expected: usize }, + + #[error("sens_atol_factor must be finite and > 0, got {got}")] + SensAtolFactor { got: f64 }, + + /// A solver option held a value diffsol would accept and then misbehave on. + #[error("solver option {name} must be {requirement}, got {got}")] + SolverOption { + name: &'static str, + requirement: &'static str, + got: f64, + }, + + /// A tolerance IDAS would reject: non-finite, negative, or zero in both `rtol` and `atol`. + #[error( + "{name}{} = {got} is not a valid tolerance: rtol and atol must be finite and >= 0, and not both zero", + .index.map_or_else(String::new, |i| format!("[{i}]")) + )] + Tolerance { + name: &'static str, + index: Option, + got: f64, + }, + + /// Both the error-controlled sensitivity solve and its relaxed retry failed. + #[error( + "sensitivity solve failed under error control ({controlled}); the retry with sensitivities excluded from error control also failed ({relaxed})" + )] + SensRetryFailed { controlled: String, relaxed: String }, + #[error("invalid CSR matrix: {0}")] Csr(String), @@ -162,4 +225,8 @@ pub enum CoreError { #[error("invalid interpolant: {0}")] Interpolant(String), + + #[cfg(feature = "diffsol")] + #[error(transparent)] + Diffsol(#[from] diffsol::DiffsolError), } diff --git a/packages/pybamm-rust/pybamm-core/src/ffi.rs b/packages/pybamm-rust/pybamm-core/src/ffi.rs index 22edc73ad6..bf2ca57080 100644 --- a/packages/pybamm-rust/pybamm-core/src/ffi.rs +++ b/packages/pybamm-rust/pybamm-core/src/ffi.rs @@ -103,6 +103,7 @@ where compiled, workspace, } = unsafe { &mut *user_data.cast::() }; + let workspace = workspace.get_mut(); let input_width = compiled.input_width().get(); if input_width != 0 && inputs.is_null() { return ERROR_NULL_POINTER; diff --git a/packages/pybamm-rust/pybamm-core/src/lib.rs b/packages/pybamm-rust/pybamm-core/src/lib.rs index 7a1907b1c3..621af374c5 100644 --- a/packages/pybamm-rust/pybamm-core/src/lib.rs +++ b/packages/pybamm-rust/pybamm-core/src/lib.rs @@ -6,7 +6,7 @@ //! 3. **Differentiate**: [`tangent`] in forward mode, and [`adjoint`] for wide Jacobian rows. //! 4. **Lower**: [`ir`] flattens the DAG into fixed-size instructions over one scratch buffer. //! 5. **Evaluate**: [`eval`] runs a tape, [`CompiledExpr::eval_batch`] at several time points. -//! 6. **Assemble**: [`CompiledModel`] and its colored [`jacobian`], driven by IDAKLU via [`ffi`]. +//! 6. **Assemble**: [`CompiledModel`] and its colored [`jacobian`], driven by `solver` or [`ffi`]. #![deny(unused_must_use)] @@ -36,6 +36,9 @@ pub mod tangent; pub mod tangent_batch; mod zero_propagate; +#[cfg(feature = "diffsol")] +pub mod solver; + pub use arena::{Arena, Dag, DagRef, NodeId, ReadExtents, resolve_state_count}; pub use error::CoreError; pub use eval::{CompiledExpr, SplitCompiledExpr, TangentInputs}; diff --git a/packages/pybamm-rust/pybamm-core/src/model/mass.rs b/packages/pybamm-rust/pybamm-core/src/model/mass.rs index e44033388d..0e21c0d71c 100644 --- a/packages/pybamm-rust/pybamm-core/src/model/mass.rs +++ b/packages/pybamm-rust/pybamm-core/src/model/mass.rs @@ -135,6 +135,36 @@ impl CompiledModel { } } + /// Compute `y = M @ x + beta * y` in place. + /// + /// When `beta == 0`, `y` is overwritten without being read, as in BLAS. + pub fn mass_gemv(&self, x: &[f64], beta: f64, y: &mut [f64]) { + if beta == 0.0 { + self.mass_action(x, y); + return; + } + let n = self.n_states; + match &self.mass_kind { + MassKind::Identity => { + for (y_i, &x_i) in y[..n].iter_mut().zip(&x[..n]) { + *y_i = beta * *y_i + x_i; + } + }, + MassKind::DiagonalSelector => { + let mask = self.state_kinds.algebraic_mask(); + for ((y_i, &x_i), &algebraic) in y[..n].iter_mut().zip(&x[..n]).zip(mask) { + let mx = if algebraic { 0.0 } else { x_i }; + *y_i = beta * *y_i + mx; + } + }, + MassKind::General => { + for (row, y_i) in y[..n].iter_mut().enumerate() { + *y_i = beta * *y_i + csr_row_dot(&self.mass_matrix, row, x); + } + }, + } + } + /// Compute mass matrix action: M @ x (no beta term). pub(crate) fn mass_action(&self, x: &[f64], mv: &mut [f64]) { let n = self.n_states; @@ -361,6 +391,28 @@ mod tests { }) } + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point + fn mass_gemv_is_m_x_plus_beta_y_on_every_mass_kind() { + let x = [3.0, -2.0]; + for mass in masses_of_every_kind() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let model = compile(&arena, y, mass, 2, 0); + + let mut mx = [0.0; 2]; + model.mass_action(&x, &mut mx); + let mut out = [10.0, 20.0]; + model.mass_gemv(&x, 0.5, &mut out); + assert_eq!(out, [0.5 * 10.0 + mx[0], 0.5 * 20.0 + mx[1]]); + + // beta == 0 must not read y, so a NaN there cannot leak through. + let mut out = [f64::NAN; 2]; + model.mass_gemv(&x, 0.0, &mut out); + assert_eq!(out, mx); + } + } + #[test] fn residual_and_assembly_apply_mass_action_on_every_mass_kind() { let (y, yp, cj) = ([0.3, -1.2], [2.0, -0.7], 0.25); diff --git a/packages/pybamm-rust/pybamm-core/src/model/mod.rs b/packages/pybamm-rust/pybamm-core/src/model/mod.rs index 828b0a48d2..25d58b923f 100644 --- a/packages/pybamm-rust/pybamm-core/src/model/mod.rs +++ b/packages/pybamm-rust/pybamm-core/src/model/mod.rs @@ -12,6 +12,7 @@ mod sens; #[cfg(test)] pub(crate) mod test_fixtures; +use std::cell::RefCell; use std::sync::Arc; use crate::arena::{Arena, DagRef, NodeId, ReadExtents, resolve_state_count}; @@ -112,17 +113,19 @@ impl Workspace { } /// A shared [`CompiledModel`] and one caller's [`Workspace`]: the `user_data` behind every -/// [`crate::ffi`] entry point. +/// [`crate::ffi`] entry point, and what the diffsol solver's operators evaluate through. #[derive(Debug)] pub struct ModelEvaluator { pub(crate) compiled: Arc, - pub(crate) workspace: Workspace, + /// A `RefCell` so operators sharing the evaluator borrow it per call; the FFI, which + /// holds the evaluator exclusively, reaches it with `get_mut`. + pub(crate) workspace: RefCell, } impl ModelEvaluator { /// Bind a fresh [`Workspace`] to an already-compiled model, without recompiling. pub fn from_compiled(compiled: Arc) -> Self { - let workspace = compiled.create_workspace(); + let workspace = RefCell::new(compiled.create_workspace()); Self { compiled, workspace, @@ -358,7 +361,7 @@ impl CompiledModel { &self.state_kinds } - /// CSC pattern of the assembled Jacobian `df/dy - cj*M`, the layout KLU expects. + /// CSC pattern of the assembled Jacobian `df/dy - cj*M`, the layout KLU and diffsol expect. #[inline] pub const fn csc_sparsity(&self) -> &CscPattern { &self.csc_sparsity @@ -406,7 +409,7 @@ impl CompiledModel { self.apply_mass_postpass(cj, jac_data); } - /// Assemble `df/dy` alone in CSC, without the mass post-pass. + /// Assemble `df/dy` alone in CSC, without the mass post-pass (diffsol applies mass itself). pub fn assemble_jacobian_csc_no_mass( &self, ws: &mut Workspace, diff --git a/packages/pybamm-rust/pybamm-core/src/model/poison_probe.rs b/packages/pybamm-rust/pybamm-core/src/model/poison_probe.rs index 9f4679769d..2ee7d566ea 100644 --- a/packages/pybamm-rust/pybamm-core/src/model/poison_probe.rs +++ b/packages/pybamm-rust/pybamm-core/src/model/poison_probe.rs @@ -1,6 +1,6 @@ //! No entry point trusts per-caller scratch left by an earlier call, finished or unwound. //! -//! Every entry point the FFI reaches runs on a fresh [`Workspace`], then on one that +//! Every entry point the FFI and diffsol reach runs on a fresh [`Workspace`], then on one that //! every entry point has dirtied and that was then overwritten whole; both give the same bits. use super::test_fixtures::{ diff --git a/packages/pybamm-rust/pybamm-core/src/params.rs b/packages/pybamm-rust/pybamm-core/src/params.rs index ae9aacfd6c..be093cfd1a 100644 --- a/packages/pybamm-rust/pybamm-core/src/params.rs +++ b/packages/pybamm-rust/pybamm-core/src/params.rs @@ -4,6 +4,8 @@ //! diverge once any input is vector-valued, so each is its own type. use std::collections::HashMap; +#[cfg(feature = "diffsol")] +use std::sync::Arc; use crate::arena::Arena; use crate::error::CoreError; @@ -298,6 +300,94 @@ pub const fn resolve_input_width( } } +/// What a solve's parameter vector stands for. +/// +/// Either the packed input array itself, or one entry per differentiated parameter. +#[cfg(feature = "diffsol")] +#[derive(Clone, Debug)] +pub struct ParamBinding { + kind: Binding, +} + +#[cfg(feature = "diffsol")] +#[derive(Clone, Debug)] +enum Binding { + PackedInputs { + width: PackedWidth, + }, + /// Where each sensitivity column's parameter sits in the packed input array. + Sensitivities { + packed_offsets: Arc<[InputOffset]>, + }, +} + +#[cfg(feature = "diffsol")] +impl ParamBinding { + /// The vector is the packed input array itself. + pub const fn packed_inputs(width: PackedWidth) -> Self { + Self { + kind: Binding::PackedInputs { width }, + } + } + + /// One entry per differentiated parameter, in sensitivity-column order. + pub fn sensitivities(params: &ParamRegistry, columns: &[ParamIdx]) -> Self { + let packed_offsets: Arc<[InputOffset]> = columns + .iter() + .map(|¶m| { + params + .slot(param) + .unwrap_or_else(|| { + panic!( + "sensitivity column {} names no input parameter", + param.get() + ) + }) + .offset + }) + .collect(); + Self { + kind: Binding::Sensitivities { packed_offsets }, + } + } + + /// How many entries the parameter vector carries. + pub fn len(&self) -> usize { + match &self.kind { + Binding::PackedInputs { width } => width.get(), + Binding::Sensitivities { packed_offsets } => packed_offsets.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Write the parameter vector into the packed input array, leaving other inputs as given. + pub fn splice_into(&self, inputs: &mut [f64], values: &[f64]) { + match &self.kind { + Binding::PackedInputs { .. } => inputs.copy_from_slice(values), + Binding::Sensitivities { packed_offsets } => { + for (&offset, &value) in packed_offsets.iter().zip(values) { + inputs[offset.index()] = value; + } + }, + } + } + + /// Read the parameter vector out of the packed input array. + pub fn gather_from(&self, inputs: &[f64], values: &mut [f64]) { + match &self.kind { + Binding::PackedInputs { .. } => values.copy_from_slice(inputs), + Binding::Sensitivities { packed_offsets } => { + for (dst, &offset) in values.iter_mut().zip(packed_offsets.iter()) { + *dst = inputs[offset.index()]; + } + }, + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -574,6 +664,34 @@ mod tests { assert_eq!(registry.packed_width(), PackedWidth::new(5)); } + #[test] + #[cfg(feature = "diffsol")] + fn a_sens_binding_splices_at_packed_offsets() { + let registry = registry_of(&wide_then_scalar()); + + let binding = ParamBinding::sensitivities(®istry, &[ParamIdx::new(1)]); + assert_eq!(binding.len(), 1); + + let mut inputs = vec![1.0, 2.0, 3.0, 4.0]; + binding.splice_into(&mut inputs, &[9.0]); + assert_eq!(inputs, vec![1.0, 2.0, 3.0, 9.0]); + + let mut values = vec![0.0]; + binding.gather_from(&inputs, &mut values); + assert_eq!(values, vec![9.0]); + } + + #[test] + #[cfg(feature = "diffsol")] + fn a_packed_binding_is_the_whole_input_array() { + let binding = ParamBinding::packed_inputs(PackedWidth::new(4)); + assert_eq!(binding.len(), 4); + + let mut inputs = vec![0.0; 4]; + binding.splice_into(&mut inputs, &[1.0, 2.0, 3.0, 4.0]); + assert_eq!(inputs, vec![1.0, 2.0, 3.0, 4.0]); + } + #[test] fn an_arena_without_inputs_reports_both_spaces_empty() { let registry = registry_of(&Arena::new()); diff --git a/packages/pybamm-rust/pybamm-core/src/solver/batch.rs b/packages/pybamm-rust/pybamm-core/src/solver/batch.rs new file mode 100644 index 0000000000..cfb3312a8e --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/batch.rs @@ -0,0 +1,178 @@ +//! Solving many input sets concurrently. +//! +//! A fan-out over [`PreparedSolver::solve`], so a batch is bit-identical to the serial loop. +//! Runs on the ambient rayon pool; wrap the call in `ThreadPool::install` to set its width. + +use rayon::prelude::*; + +use super::solve::{InputSet, PreparedSolver, SolveOutcome, SolveRequest}; +use crate::error::CoreError; + +impl PreparedSolver { + /// Solve one trajectory per input set, concurrently, each factorising serially. + /// + /// A set that fails keeps its own `Err` at its own index. + pub fn solve_batch( + &self, + request: SolveRequest<'_>, + sets: &[InputSet<'_>], + ) -> Vec> { + // The sets already occupy the pool; a parallel LU in each would oversubscribe it. + let request = request.with_lu_threads(1); + sets.par_iter() + .map(|set| self.solve(request, *set)) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::solver::test_fixtures::decay_with_output_and_event; + + /// The event at `y = 0.4` exercises root-finding across sets. + fn prepared() -> PreparedSolver { + PreparedSolver::new(decay_with_output_and_event(), 1e-8, &[1e-10]).expect("setup failed") + } + + fn t_eval() -> Vec { + (0..=40).map(|i| f64::from(i) * 0.05).collect() + } + + /// Descending solve cost: the first set decays slowest. + fn rates(n: usize) -> Vec> { + (0..n).map(|i| vec![0.4 + 0.35 * (i as f64)]).collect() + } + + fn sets<'a>(y0: &'a [f64], rates: &'a [Vec]) -> Vec> { + rates + .iter() + .map(|inputs| InputSet::new(y0, inputs)) + .collect() + } + + fn run_in_pool(threads: usize, body: impl FnOnce() -> T + Send) -> T { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .expect("pool build failed") + .install(body) + } + + // Bit-identical is the property under test. + #[allow(clippy::float_cmp)] + fn assert_same_outcome(got: &SolveOutcome, want: &SolveOutcome, context: &str) { + assert_eq!(got.t, want.t, "{context}: times differ"); + assert_eq!(got.y, want.y, "{context}: trajectories differ"); + assert_eq!(got.n_rows, want.n_rows, "{context}: row counts differ"); + assert_eq!(got.yp, want.yp, "{context}: yp differs"); + assert_eq!( + got.sensitivities, want.sensitivities, + "{context}: sensitivities differ" + ); + assert_eq!( + got.termination, want.termination, + "{context}: terminations differ" + ); + assert_eq!(got.y_event, want.y_event, "{context}: root states differ"); + assert_eq!( + got.statistics.solver.number_of_steps, want.statistics.solver.number_of_steps, + "{context}: step counts differ", + ); + } + + #[test] + fn batch_is_bit_identical_to_the_serial_loop() { + let solver = prepared(); + let times = t_eval(); + let rates = rates(8); + let y0 = [1.0]; + let sets = sets(&y0, &rates); + let request = SolveRequest::new(×); + + let serial: Vec<_> = sets + .iter() + .map(|set| solver.solve(request, *set).expect("solve failed")) + .collect(); + + for threads in [1, 2, 8] { + let batched = run_in_pool(threads, || solver.solve_batch(request, &sets)); + assert_eq!(batched.len(), serial.len()); + for (i, (got, want)) in batched.into_iter().zip(&serial).enumerate() { + let got = got.expect("set failed"); + assert_same_outcome(&got, want, &format!("set {i} at {threads} thread(s)")); + } + } + } + + #[test] + fn a_parallel_lu_solve_is_bit_identical_to_the_serial_one() { + let solver = prepared(); + let times = t_eval(); + let rates = rates(1); + let y0 = [1.0]; + let set = sets(&y0, &rates)[0]; + let request = SolveRequest::new(×); + + let serial = solver.solve(request, set).expect("solve failed"); + let parallel = + run_in_pool(4, || solver.solve(request.with_lu_threads(4), set)).expect("solve failed"); + assert_same_outcome(¶llel, &serial, "4 LU threads"); + } + + #[test] + fn a_sensitivity_batch_is_bit_identical_to_the_serial_loop() { + let solver = prepared(); + let times = t_eval(); + let rates = rates(4); + let y0 = [1.0]; + let seed = [0.25]; + let sets: Vec> = rates + .iter() + .map(|inputs| InputSet::new(&y0, inputs).with_sens_seed(&seed)) + .collect(); + let request = SolveRequest::new(×).with_sensitivities(); + + let serial: Vec<_> = sets + .iter() + .map(|set| solver.solve(request, *set).expect("solve failed")) + .collect(); + + let batched = run_in_pool(4, || solver.solve_batch(request, &sets)); + for (i, (got, want)) in batched.into_iter().zip(&serial).enumerate() { + let got = got.expect("set failed"); + assert!(got.sensitivities.is_some(), "set {i} dropped its blocks"); + assert_same_outcome(&got, want, &format!("sensitivity set {i}")); + } + } + + #[test] + fn one_failing_set_leaves_the_others_intact() { + let solver = prepared(); + let times = t_eval(); + let good = [1.0]; + let y0 = [1.0]; + // Set 2 gets the wrong input width: a failure independent of the integrator. + let sets = [ + InputSet::new(&y0, &good), + InputSet::new(&y0, &good), + InputSet::new(&y0, &[]), + InputSet::new(&y0, &good), + InputSet::new(&y0, &good), + ]; + + let batched = run_in_pool(4, || solver.solve_batch(SolveRequest::new(×), &sets)); + + assert_eq!(batched.len(), 5); + for (i, result) in batched.into_iter().enumerate() { + if i == 2 { + assert!( + matches!(result, Err(CoreError::InputsLength { .. })), + "set 2 should carry its own error", + ); + } else { + assert!(result.is_ok(), "set {i} failed alongside set 2"); + } + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/equations.rs b/packages/pybamm-rust/pybamm-core/src/solver/equations.rs new file mode 100644 index 0000000000..844b2c3eb7 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/equations.rs @@ -0,0 +1,139 @@ +//! The equation bundle diffsol solves. +//! +//! Operator views are minted per call, since diffsol's operator traits borrow per call. + +use std::sync::Arc; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + FaerContext, OdeEquations, OdeEquationsImplicitSens, OdeEquationsRef, Op, UnitCallable, +}; + +use super::init::InitOp; +use super::mass::MassOp; +use super::observable::ObservableOp; +use super::rhs::RhsOp; +use super::{FaerSparsity, lane, lane_mut}; +use crate::model::ModelEvaluator; +use crate::params::ParamBinding; + +/// Local implementation of diffsol's `OdeEquations` trait. +/// +/// Views borrow from here rather than owning, so minting one per callback is cheap. +pub struct Equations { + /// The model and this solve's workspace, which the operator views borrow. + pub(crate) evaluator: ModelEvaluator, + pub(crate) params: Vec, + /// The packed inputs, or one differentiated parameter per entry on the sens path. + pub(crate) binding: ParamBinding, + pub(crate) y0: Vec, + /// `dy0/dp`, column-major `n_states x binding.len()`; empty means zero. + pub(crate) y0_sens: Vec, + pub(crate) jac_sparsity: Arc, + pub(crate) mass_sparsity: Arc, + /// M's values in `mass_sparsity` order. + pub(crate) mass_csc_values: Arc<[f64]>, + pub(crate) context: FaerContext, +} + +impl std::fmt::Debug for Equations { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Equations") + .field("n_states", &self.evaluator.compiled.n_states()) + .field("nparams", &self.binding.len()) + .finish_non_exhaustive() + } +} + +impl Op for Equations { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.evaluator.compiled.n_states() + } + fn nout(&self) -> usize { + self.evaluator.compiled.n_states() + } + fn nparams(&self) -> usize { + self.binding.len() + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl<'a> OdeEquationsRef<'a> for Equations { + type Rhs = RhsOp<'a>; + type Mass = MassOp<'a>; + type Root = ObservableOp<'a>; + type Init = InitOp<'a>; + /// Never constructed: the solve loop evaluates output variables itself. + type Out = UnitCallable>; + /// Never constructed: `PyBaMM` events terminate the solve rather than reset it. + type Reset = UnitCallable>; +} + +impl OdeEquations for Equations { + fn rhs(&self) -> RhsOp<'_> { + RhsOp { + compiled: &self.evaluator.compiled, + ws: &self.evaluator.workspace, + inputs: &self.params, + binding: &self.binding, + jac_sparsity: &self.jac_sparsity, + context: self.context, + } + } + + fn mass(&self) -> Option> { + Some(MassOp { + compiled: &self.evaluator.compiled, + sparsity: &self.mass_sparsity, + csc_values: &self.mass_csc_values, + context: self.context, + }) + } + + fn root(&self) -> Option> { + (self.evaluator.compiled.total_event_len() > 0).then(|| ObservableOp { + compiled: &self.evaluator.compiled, + ws: &self.evaluator.workspace, + inputs: &self.params, + binding: &self.binding, + context: self.context, + }) + } + + fn out(&self) -> Option>> { + None + } + + fn init(&self) -> InitOp<'_> { + InitOp { + y0: &self.y0, + y0_sens: &self.y0_sens, + n_sens_params: self.binding.len(), + context: self.context, + } + } + + /// Splice `p` into the packed inputs, leaving those it does not cover unchanged. + fn set_params(&mut self, p: &FaerVec) { + self.binding.splice_into(&mut self.params, lane(p)); + } + + /// Read `p` back out of the packed input vector. + fn get_params(&self, p: &mut FaerVec) { + self.binding.gather_from(&self.params, lane_mut(p)); + } +} + +// Compile-time check that `Equations` meets the `OdeEquationsImplicitSens` bounds. +const _: fn() = || { + const fn assert_implicit_sens() {} + assert_implicit_sens::(); +}; diff --git a/packages/pybamm-rust/pybamm-core/src/solver/init.rs b/packages/pybamm-rust/pybamm-core/src/solver/init.rs new file mode 100644 index 0000000000..84a3103f37 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/init.rs @@ -0,0 +1,172 @@ +//! Initial-condition operator, with a per-solve `dy0/dp` seed (empty means zero). + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ConstantOp, ConstantOpSens, FaerContext, Op}; + +use super::{FaerSparsity, dense_faer_sparsity, lane, lane_mut, values_mut}; + +/// Initial-condition operator returning the stored `y0`. +#[derive(Debug)] +pub struct InitOp<'a> { + pub y0: &'a [f64], + /// `dy0/dp`, column-major `n_states x n_sens_params`; empty means zero. + pub y0_sens: &'a [f64], + pub n_sens_params: usize, + pub context: FaerContext, +} + +impl Op for InitOp<'_> { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + 0 + } + fn nout(&self) -> usize { + self.y0.len() + } + fn nparams(&self) -> usize { + self.n_sens_params + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl ConstantOp for InitOp<'_> { + fn call_inplace(&self, _t: f64, y: &mut FaerVec) { + lane_mut(y).copy_from_slice(self.y0); + } +} + +impl ConstantOpSens for InitOp<'_> { + /// Write `dy0/dp · v` into `y`, or zeros if no seed was given. + fn sens_mul_inplace(&self, _t: f64, v: &FaerVec, y: &mut FaerVec) { + let out = lane_mut(y); + out.fill(0.0); + if self.y0_sens.is_empty() { + return; + } + let n_states = self.y0.len(); + for (col, &v_col) in lane(v).iter().enumerate() { + if v_col == 0.0 { + continue; + } + let offset = col * n_states; + for (out_i, &s_i) in out.iter_mut().zip(&self.y0_sens[offset..offset + n_states]) { + *out_i += v_col * s_i; + } + } + } + + /// Write the full `dy0/dp` matrix into `y`, or zeros if no seed was given. + /// + /// `y` uses the dense column-major pattern of `y0_sens`, so values copy directly. + fn sens_inplace(&self, _t: f64, y: &mut FaerSparseMat) { + let values = values_mut(y); + if self.y0_sens.is_empty() { + values.fill(0.0); + } else { + values.copy_from_slice(self.y0_sens); + } + } + + fn sens_sparsity(&self) -> Option { + Some(dense_faer_sparsity(self.y0.len(), self.n_sens_params)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use diffsol::{Context, Matrix, Vector}; + + fn init_op<'a>(y0: &'a [f64], y0_sens: &'a [f64], n_sens_params: usize) -> InitOp<'a> { + InitOp { + y0, + y0_sens, + n_sens_params, + context: FaerContext::default(), + } + } + + fn sens_column(op: &InitOp<'_>, col: usize) -> Vec { + let ctx = FaerContext::default(); + let mut v = ctx.vector_zeros::>(op.n_sens_params); + v.set_index(col, 1.0); + let mut y = ctx.vector_zeros::>(op.y0.len()); + op.sens_mul_inplace(0.0, &v, &mut y); + lane(&y).to_vec() + } + + #[test] + fn empty_seed_is_all_zero() { + let y0 = [0.0; 3]; + let op = init_op(&y0, &[], 2); + assert_eq!(sens_column(&op, 0), vec![0.0; 3]); + assert_eq!(sens_column(&op, 1), vec![0.0; 3]); + } + + #[test] + fn mixed_direction_is_the_matvec() { + let y0 = [0.0; 3]; + let seed = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let op = init_op(&y0, &seed, 2); + let ctx = FaerContext::default(); + let mut v = ctx.vector_zeros::>(2); + v.set_index(0, 2.0); + v.set_index(1, -1.0); + let mut y = ctx.vector_zeros::>(3); + op.sens_mul_inplace(0.0, &v, &mut y); + assert_eq!(lane(&y), &[2.0 - 4.0, 4.0 - 5.0, 6.0 - 6.0]); + } + + fn dense_sens_matrix(op: &InitOp<'_>) -> FaerSparseMat { + FaerSparseMat::new_from_sparsity( + op.y0.len(), + op.n_sens_params, + op.sens_sparsity(), + FaerContext::default(), + ) + } + + #[test] + fn sens_inplace_copies_the_seed_the_probing_default_reads() { + let y0 = [0.0; 3]; + let seed = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let op = init_op(&y0, &seed, 2); + let mut copied = dense_sens_matrix(&op); + let mut probed = dense_sens_matrix(&op); + op.sens_inplace(0.0, &mut copied); + op._default_sens_inplace(0.0, &mut probed); + assert_eq!(values_mut(&mut copied), &seed); + assert_eq!(values_mut(&mut copied), values_mut(&mut probed)); + } + + #[test] + fn sens_inplace_of_an_empty_seed_overwrites_with_zeros() { + let y0 = [0.0; 3]; + let op = init_op(&y0, &[], 2); + let mut m = dense_sens_matrix(&op); + values_mut(&mut m).fill(9.0); + op.sens_inplace(0.0, &mut m); + assert_eq!(values_mut(&mut m), &[0.0; 6]); + } + + #[test] + fn output_is_overwritten_not_accumulated() { + let y0 = [0.0; 3]; + let seed = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let op = init_op(&y0, &seed, 2); + let ctx = FaerContext::default(); + let mut v = ctx.vector_zeros::>(2); + v.set_index(0, 1.0); + let mut y = ctx.vector_zeros::>(3); + y.set_index(0, 99.0); + op.sens_mul_inplace(0.0, &v, &mut y); + assert_eq!(lane(&y), &[1.0, 2.0, 3.0]); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/mass.rs b/packages/pybamm-rust/pybamm-core/src/solver/mass.rs new file mode 100644 index 0000000000..172c684f10 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/mass.rs @@ -0,0 +1,66 @@ +//! Mass matrix operator; `M` is constant and arrives from Python in CSR. +//! +//! A zero *column* `i` makes state `i` algebraic. diffsol's init partition tests the diagonal +//! instead, which disagrees only on a zero diagonal beside an off-diagonal non-zero. + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, LinearOp, Matrix, Op, Vector}; + +use super::{FaerSparsity, values_mut}; +use crate::model::CompiledModel; + +/// Mass operator `M`, constant in `t` and `p`. +pub struct MassOp<'a> { + pub compiled: &'a CompiledModel, + pub sparsity: &'a FaerSparsity, + /// M's values in the CSC order of `sparsity`. + pub csc_values: &'a [f64], + pub context: FaerContext, +} + +impl std::fmt::Debug for MassOp<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MassOp") + .field("n_states", &self.compiled.n_states()) + .finish_non_exhaustive() + } +} + +impl Op for MassOp<'_> { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.compiled.n_states() + } + fn nout(&self) -> usize { + self.compiled.n_states() + } + fn nparams(&self) -> usize { + 0 + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl LinearOp for MassOp<'_> { + /// Compute `y = M @ x + beta * y` for each lane of `x` and `y`. + fn gemv_inplace(&self, x: &FaerVec, _t: f64, beta: f64, y: &mut FaerVec) { + FaerVec::for_each_batch_mut_host([y], [x], |[y], [x], _| { + self.compiled.mass_gemv(x, beta, y); + }); + } + + /// Copy M's precomputed values into `y`, which uses this operator's `sparsity`. + fn matrix_inplace(&self, _t: f64, y: &mut FaerSparseMat) { + values_mut(y).copy_from_slice(self.csc_values); + } + + fn sparsity(&self) -> Option< as Matrix>::Sparsity> { + Some(self.sparsity.clone()) + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/mod.rs b/packages/pybamm-rust/pybamm-core/src/solver/mod.rs new file mode 100644 index 0000000000..7c0cff3f5a --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/mod.rs @@ -0,0 +1,502 @@ +//! In-process DAE solving through diffsol. +//! +//! diffsol applies `RhsOp::jac_mul_inplace` and `MassOp::gemv_inplace` to `S`, one lane per +//! parameter; every other callback expects one lane, read with `lane` or `lane_mut`. + +pub mod batch; +mod equations; +mod init; +mod mass; +mod observable; +pub mod options; +mod rhs; +pub mod solve; +#[cfg(test)] +mod test_fixtures; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{Context, Vector, VectorCommon}; +use faer::sparse::{SparseRowMatRef, SymbolicSparseColMat, SymbolicSparseRowMatRef}; + +use crate::jacobian::CscPattern; +use crate::node::CsrData; + +pub use self::options::SolverOptions; + +/// Sparsity type of the faer sparse matrix diffsol is instantiated with. +pub(crate) type FaerSparsity = as diffsol::Matrix>::Sparsity; + +/// Borrow the values of a single-lane vector. +pub(crate) fn lane(v: &FaerVec) -> &[f64] { + assert_single_lane(v.context().nbatch()); + v.inner().col_as_slice(0) +} + +/// Mutably borrow the values of a single-lane vector. +pub(crate) fn lane_mut(v: &mut FaerVec) -> &mut [f64] { + assert_single_lane(v.context().nbatch()); + v.inner_mut().col_as_slice_mut(0) +} + +/// Borrow lane `j` of a batched vector. +#[cfg(test)] +pub(crate) fn lane_at(v: &FaerVec, j: usize) -> &[f64] { + v.inner().col_as_slice(j) +} + +/// Iterate over the lanes of a batched vector. +pub(crate) fn lanes(v: &FaerVec) -> impl ExactSizeIterator { + let values = v.inner(); + (0..values.ncols()).map(move |j| values.col_as_slice(j)) +} + +/// Iterate mutably over the lanes of a batched vector. +pub(crate) fn lanes_mut(v: &mut FaerVec) -> impl ExactSizeIterator { + v.inner_mut().col_iter_mut().map(|column| { + column + .try_as_col_major_mut() + .expect("a faer vector lane is one contiguous column") + .as_slice_mut() + }) +} + +/// Mutably borrow the stored values of a single-lane sparse matrix, in CSC order. +pub(crate) fn values_mut(m: &mut FaerSparseMat) -> &mut [f64] { + let lanes = diffsol::Matrix::inner_mut(m); + assert_single_lane(lanes.len()); + lanes[0].val_mut() +} + +fn assert_single_lane(nbatch: usize) { + assert_eq!( + nbatch, 1, + "a single-lane callback was handed {nbatch} lanes" + ); +} + +/// Reinterpret a CSC Jacobian pattern as faer's symbolic sparsity. +pub(crate) fn csc_to_faer_sparsity(csc: &CscPattern) -> FaerSparsity { + let col_ptrs: Vec = csc.colptr().to_vec(); + let row_indices: Vec = csc.rowind().to_vec(); + SymbolicSparseColMat::new_checked(csc.nrows(), csc.ncols(), col_ptrs, None, row_indices) +} + +/// Build a dense CSC sparsity for an `nrows × ncols` matrix. +pub(crate) fn dense_faer_sparsity(nrows: usize, ncols: usize) -> FaerSparsity { + let col_ptrs: Vec = (0..=ncols).map(|c| c * nrows).collect(); + let row_indices: Vec = (0..ncols).flat_map(|_| 0..nrows).collect(); + SymbolicSparseColMat::new_checked(nrows, ncols, col_ptrs, None, row_indices) +} + +/// Transpose a CSR mass matrix into faer's CSC sparsity plus the values in that order. +pub(crate) fn csr_mass_to_faer_csc(mass: &CsrData) -> (FaerSparsity, Vec) { + let shape = mass.shape(); + let symbolic = SymbolicSparseRowMatRef::new_unsorted_checked( + shape.rows, + shape.cols, + mass.indptr(), + None, + mass.indices(), + ); + SparseRowMatRef::new(symbolic, mass.data()) + .to_col_major() + .expect("transposing M fails only to allocate") + .into_parts() +} + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + use std::sync::Arc; + + use super::init::InitOp; + use super::mass::MassOp; + use super::rhs::RhsOp; + use super::solve::{InputSet, PreparedSolver, SolveRequest, Termination}; + use super::test_fixtures::{ + decay, decay_with_output_and_event, differential_then_algebraic_mass, + }; + use super::*; + use crate::arena::Arena; + use crate::model::test_fixtures::{compile, identity_mass_matrix}; + use crate::model::{CompiledModel, Workspace}; + use crate::node::{CsrData, Node, Shape}; + use crate::params::ParamBinding; + use diffsol::vector::faer_serial::FaerVec; + use diffsol::{ + ConstantOp, Context, FaerContext, LinearOp, Matrix, NonLinearOp, NonLinearOpJacobian, Op, + }; + + /// Owns what a borrowed operator view points at, so one op can be exercised alone. + struct OpFixture { + compiled: Arc, + ws: RefCell, + jac_sparsity: FaerSparsity, + mass_sparsity: FaerSparsity, + mass_csc_values: Vec, + inputs: Vec, + binding: ParamBinding, + n_states: usize, + context: FaerContext, + } + + impl OpFixture { + fn new(model: CompiledModel) -> Self { + let compiled = Arc::new(model); + let ws = RefCell::new(compiled.create_workspace()); + let jac_sparsity = csc_to_faer_sparsity(compiled.csc_sparsity()); + let (mass_sparsity, mass_csc_values) = csr_mass_to_faer_csc(compiled.mass_matrix()); + let width = compiled.input_width(); + let binding = ParamBinding::packed_inputs(width); + Self { + n_states: compiled.n_states(), + compiled, + ws, + jac_sparsity, + mass_sparsity, + mass_csc_values, + inputs: vec![0.0; width.get()], + binding, + context: FaerContext::default(), + } + } + + fn rhs(&self) -> RhsOp<'_> { + RhsOp { + compiled: &self.compiled, + ws: &self.ws, + inputs: &self.inputs, + binding: &self.binding, + jac_sparsity: &self.jac_sparsity, + context: self.context, + } + } + + fn mass(&self) -> MassOp<'_> { + MassOp { + compiled: &self.compiled, + sparsity: &self.mass_sparsity, + csc_values: &self.mass_csc_values, + context: self.context, + } + } + } + + fn build_2state_model() -> CompiledModel { + // dy0/dt = -y0, dy1/dt = -2*y1 + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let neg_one = arena.alloc(Node::Scalar(-1.0)); + let neg_two = arena.alloc(Node::Scalar(-2.0)); + let rhs0 = arena.alloc(Node::Mul(neg_one, sv0)); + let rhs1 = arena.alloc(Node::Mul(neg_two, sv1)); + let rhs = arena.alloc(Node::Concat(vec![rhs0, rhs1])); + compile(&arena, rhs, identity_mass_matrix(2), 2, 0) + } + + #[test] + fn csc_to_faer_preserves_pattern() { + let model = build_2state_model(); + let csc = model.csc_sparsity(); + let sparsity = csc_to_faer_sparsity(csc); + assert_eq!(sparsity.col_ptr(), csc.colptr()); + assert_eq!(sparsity.row_idx(), csc.rowind()); + } + + #[test] + // Values are moved, never computed. + #[allow(clippy::float_cmp)] + fn csr_mass_to_faer_moves_each_value_to_its_column() { + // Row 1 lists its columns out of order; each CSC column still runs down the rows. + let mass = CsrData { + indptr: vec![0, 2, 4, 5], + indices: vec![0, 2, 2, 1, 0], + data: vec![1.0, 2.0, 3.0, 4.0, 5.0], + shape: Shape::matrix(3, 3), + }; + let (sparsity, values) = csr_mass_to_faer_csc(&mass); + assert_eq!(sparsity.col_ptr(), [0, 2, 3, 5]); + assert_eq!(sparsity.row_idx(), [0, 2, 1, 0, 1]); + assert_eq!(values, [1.0, 5.0, 4.0, 2.0, 3.0]); + } + + #[test] + fn a_two_state_decay_matches_the_analytic_solution() { + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.05).collect(); + let result = PreparedSolver::new(build_2state_model(), 1e-8, &[1e-10, 1e-10]) + .expect("prepare") + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0, 1.0], &[])) + .expect("solve"); + + assert_eq!(result.termination, Termination::Completed); + assert!(result.y_event.is_none(), "no event state expected"); + assert_eq!(result.t, t_eval); + for (j, &t) in result.t.iter().enumerate() { + let exact = [(-t).exp(), (-2.0 * t).exp()]; + for (row, exact) in exact.into_iter().enumerate() { + let got = result.y[j * result.n_rows + row]; + assert!( + (got - exact).abs() < 1e-6, + "y{row}({t}) = {got}, exact {exact}" + ); + } + } + } + + #[test] + // Repeating an input set must repeat its trajectory bit for bit. + #[allow(clippy::float_cmp)] + fn one_prepared_solver_keeps_its_solves_independent() { + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = PreparedSolver::new(decay(), 1e-8, &[1e-10]).expect("prepare"); + let [first, scaled, again] = [1.0, 5.0, 1.0].map(|y0| { + prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[y0], &[])) + .expect("solve") + }); + + assert_eq!(first.t, again.t); + assert_eq!(first.y, again.y); + for (&t, &y) in scaled.t.iter().zip(&scaled.y) { + let exact = 5.0 * (-t).exp(); + assert!((y - exact).abs() < 5e-6, "y({t}) = {y}, exact {exact}"); + } + } + + #[test] + fn rhs_eval_2state_decay() { + let fixture = OpFixture::new(build_2state_model()); + let rhs = fixture.rhs(); + let ctx = FaerContext::default(); + assert_eq!(rhs.nstates(), 2); + assert_eq!(rhs.nout(), 2); + + let mut x = ctx.vector_zeros::>(2); + lane_mut(&mut x).copy_from_slice(&[3.0, 5.0]); + let mut y = ctx.vector_zeros::>(2); + rhs.call_inplace(&x, 0.0, &mut y); + + assert!((lane(&y)[0] - (-3.0)).abs() < 1e-12); + assert!((lane(&y)[1] - (-10.0)).abs() < 1e-12); + } + + fn build_2state_nonlinear_model() -> CompiledModel { + // dy0/dt = -y0 * y1, dy1/dt = y0 * y0 + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let neg_one = arena.alloc(Node::Scalar(-1.0)); + let product = arena.alloc(Node::Mul(sv0, sv1)); + let rhs0 = arena.alloc(Node::Mul(neg_one, product)); + let rhs1 = arena.alloc(Node::Mul(sv0, sv0)); + let rhs = arena.alloc(Node::Concat(vec![rhs0, rhs1])); + compile(&arena, rhs, identity_mass_matrix(2), 2, 0) + } + + fn batched(lanes: &[&[f64]]) -> FaerVec { + let ctx = FaerContext::with_nbatch(lanes.len()); + let mut v = ctx.vector_zeros::>(lanes[0].len()); + for (j, values) in lanes.iter().enumerate() { + v.inner_mut().col_as_slice_mut(j).copy_from_slice(values); + } + v + } + + #[test] + fn rhs_jac_mul_takes_every_lane_at_one_state() { + // Lane counts cover the scalar path, a padded batch, a full batch and a tail. + let fixture = OpFixture::new(build_2state_nonlinear_model()); + let rhs = fixture.rhs(); + let state = [0.7, -1.3]; + let mut x = FaerContext::default().vector_zeros::>(2); + lane_mut(&mut x).copy_from_slice(&state); + + for n_lanes in [1usize, 2, 3, 8, 9] { + let directions: Vec<[f64; 2]> = (0..n_lanes) + .map(|j| { + let j = j as f64; + [j - 1.5, 0.25 * j * j + 1.0] + }) + .collect(); + let as_slices: Vec<&[f64]> = directions.iter().map(<[f64; 2]>::as_slice).collect(); + let v = batched(&as_slices); + let zeros: &[f64] = &[0.0; 2]; + let mut y = batched(&vec![zeros; n_lanes]); + rhs.jac_mul_inplace(&x, 0.0, &v, &mut y); + + for (j, direction) in directions.iter().enumerate() { + let mut expected = [0.0; 2]; + fixture.compiled.jac_action( + &mut fixture.ws.borrow_mut(), + 0.0, + &state, + &fixture.inputs, + direction, + &mut expected, + ); + assert_eq!(lane_at(&y, j), expected, "lane {j} of {n_lanes}"); + } + } + } + + #[test] + #[should_panic(expected = "a single-lane callback was handed 2 lanes")] + fn a_single_lane_callback_rejects_a_batched_state() { + let fixture = OpFixture::new(build_2state_model()); + let x = batched(&[&[1.0, 1.0], &[2.0, 2.0]]); + let mut y = batched(&[&[0.0; 2], &[0.0; 2]]); + fixture.rhs().call_inplace(&x, 0.0, &mut y); + } + + #[test] + fn rhs_jacobian_inplace_2state() { + let fixture = OpFixture::new(build_2state_model()); + let rhs = fixture.rhs(); + let ctx = FaerContext::default(); + + let mut x = ctx.vector_zeros::>(2); + lane_mut(&mut x).copy_from_slice(&[1.0, 1.0]); + + let jac = rhs.jacobian(&x, 0.0); + + let mut diag = [0.0f64; 2]; + let (indices, values) = jac.triplet_iter(); + for ((row, col), val) in indices.zip(values) { + assert_eq!(row, col, "jacobian must be diagonal, got ({row},{col})"); + diag[row] = val; + } + assert!((diag[0] - (-1.0)).abs() < 1e-12, "J[0,0]={}", diag[0]); + assert!((diag[1] - (-2.0)).abs() < 1e-12, "J[1,1]={}", diag[1]); + } + + #[test] + fn mass_gemv_applies_m_to_every_lane() { + let fixture = OpFixture::new(build_dae_model()); + let mass = fixture.mass(); + let x = batched(&[&[1.0, 2.0], &[3.0, 4.0]]); + let mut y = batched(&[&[10.0, 20.0], &[30.0, 40.0]]); + + mass.gemv_inplace(&x, 0.0, 3.0, &mut y); + assert_eq!(lane_at(&y, 0), [31.0, 60.0]); + assert_eq!(lane_at(&y, 1), [93.0, 120.0]); + } + + #[test] + fn mass_matrix_inplace_matches_probing_default() { + for model in [build_2state_model(), build_dae_model()] { + let fixture = OpFixture::new(model); + let mass = fixture.mass(); + let n = fixture.n_states; + let ctx = FaerContext::default(); + let mut fast = FaerSparseMat::::new_from_sparsity(n, n, mass.sparsity(), ctx); + let mut probed = FaerSparseMat::::new_from_sparsity(n, n, mass.sparsity(), ctx); + mass.matrix_inplace(0.0, &mut fast); + mass._default_matrix_inplace(0.0, &mut probed); + + let (fast_idx, fast_vals) = fast.triplet_iter(); + let (probed_idx, probed_vals) = probed.triplet_iter(); + let fast_entries: Vec<_> = fast_idx.zip(fast_vals).collect(); + let probed_entries: Vec<_> = probed_idx.zip(probed_vals).collect(); + assert_eq!(fast_entries, probed_entries); + } + } + + #[test] + fn init_copies_y0_exactly() { + let ctx = FaerContext::default(); + let init = InitOp { + y0: &[1.5, -2.5, 3.0], + y0_sens: &[], + n_sens_params: 0, + context: ctx, + }; + let mut y = ctx.vector_zeros::>(3); + init.call_inplace(0.0, &mut y); + assert_eq!(lane(&y), &[1.5, -2.5, 3.0]); + } + + #[test] + fn prepared_problem_is_send_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + } + + #[test] + fn an_outputs_request_at_final_time_carries_the_terminal_full_state() { + // Without a trajectory, y_event is the only state a caller can restart from. + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.1).collect(); + let prepared = PreparedSolver::new(decay_with_output_and_event(), 1e-8, &[1e-10]).unwrap(); + let set = InputSet::new(&[1.0], &[1.0]); + let full = prepared.solve(SolveRequest::new(&t_eval), set).unwrap(); + let r = prepared + .solve(SolveRequest::new(&t_eval).with_outputs(), set) + .unwrap(); + + assert_eq!(r.termination, Termination::Completed); + let y_event = r.y_event.expect("terminal state missing"); + assert_eq!(y_event.len(), 1); + let full_last = full.y[(full.t.len() - 1) * full.n_rows]; + assert!( + (y_event[0] - full_last).abs() < 1e-8, + "terminal state {} diverges from full solve {full_last}", + y_event[0] + ); + } + + #[test] + #[allow(clippy::float_cmp)] // the root is reported at the first requested time itself + fn an_event_non_positive_at_the_start_ends_the_solve_there() { + // The event starts below zero and stays there: no sign change for root finding. + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = PreparedSolver::new(decay_with_output_and_event(), 1e-8, &[1e-10]).unwrap(); + for request in [ + SolveRequest::new(&t_eval), + SolveRequest::new(&t_eval).with_outputs(), + ] { + let r = prepared + .solve(request, InputSet::new(&[0.3], &[1.0])) + .unwrap(); + + assert_eq!(r.termination, Termination::Event { t: 0.0, index: 0 }); + assert_eq!(r.t, vec![0.0]); + assert_eq!(r.y_event.expect("y_event missing"), vec![0.3]); + } + } + + fn build_dae_model() -> CompiledModel { + // y0' = -y0 and algebraic 0 = 2*y0 - y1, with M = diag(1, 0). + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let neg1 = arena.alloc(Node::Scalar(-1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let r0 = arena.alloc(Node::Mul(neg1, sv0)); + let two_y0 = arena.alloc(Node::Mul(two, sv0)); + let r1 = arena.alloc(Node::Sub(two_y0, sv1)); + let rhs = arena.alloc(Node::Concat(vec![r0, r1])); + compile(&arena, rhs, differential_then_algebraic_mass(), 2, 0) + } + + #[test] + fn dae_consistent_ic_and_solve() { + let atol = vec![1e-10, 1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = PreparedSolver::new(build_dae_model(), 1e-8, &atol).unwrap(); + // The algebraic y1 starts deliberately inconsistent (should be 2.0). + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0, 0.0], &[])) + .unwrap(); + assert_eq!(r.termination, Termination::Completed); + for (j, &t) in r.t.iter().enumerate() { + let y0 = r.y[j * r.n_rows]; + let y1 = r.y[j * r.n_rows + 1]; + assert!((y0 - (-t).exp()).abs() < 1e-6, "t={t}: y0={y0}"); + assert!( + 2.0f64.mul_add(-(-t).exp(), y1).abs() < 1e-6, + "t={t}: y1={y1} violates y1=2*y0" + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/observable.rs b/packages/pybamm-rust/pybamm-core/src/solver/observable.rs new file mode 100644 index 0000000000..587fe7c85d --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/observable.rs @@ -0,0 +1,95 @@ +//! Observable operator: the model's events as diffsol's `root` function. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op}; + +use super::{lane, lane_mut}; + +use crate::model::{CompiledModel, Workspace}; +use crate::observable::ObservableKind; +use crate::params::ParamBinding; + +/// The concatenated events as `H(t, y; p)`. +pub struct ObservableOp<'a> { + pub compiled: &'a CompiledModel, + pub ws: &'a RefCell, + pub inputs: &'a [f64], + /// What the solve's parameter vector stands for. + pub binding: &'a ParamBinding, + pub context: FaerContext, +} + +/// The only family a solve hands diffsol. +const KIND: ObservableKind = ObservableKind::Events; + +impl std::fmt::Debug for ObservableOp<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ObservableOp") + .field("n_out", &self.nout()) + .finish_non_exhaustive() + } +} + +impl Op for ObservableOp<'_> { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.compiled.n_states() + } + /// The events' concatenated length. + fn nout(&self) -> usize { + self.compiled.observables(KIND).total_len() + } + fn nparams(&self) -> usize { + self.binding.len() + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl NonLinearOp for ObservableOp<'_> { + fn call_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled + .eval_observables(&mut ws, KIND, t, lane(x), self.inputs, lane_mut(y)); + } +} + +impl NonLinearOpJacobian for ObservableOp<'_> { + /// dH/dy · v. + fn jac_mul_inplace(&self, x: &FaerVec, t: f64, v: &FaerVec, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled.observable_jac_action( + &mut ws, + KIND, + t, + lane(x), + self.inputs, + lane(v), + lane_mut(y), + ); + } +} + +impl NonLinearOpSens for ObservableOp<'_> { + /// dH/dp · v. + fn sens_mul_inplace(&self, x: &FaerVec, t: f64, v: &FaerVec, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled.observable_sens_action( + &mut ws, + KIND, + t, + lane(x), + self.inputs, + lane(v), + lane_mut(y), + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/options.rs b/packages/pybamm-rust/pybamm-core/src/solver/options.rs new file mode 100644 index 0000000000..29cbede492 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/options.rs @@ -0,0 +1,230 @@ +//! Integrator tuning, one field per diffsol `OdeSolverOptions` knob. +//! +//! Defaults are read from diffsol, except [`SolverOptions::max_nonlinear_solver_failures`]. + +use diffsol::{OdeEquations, OdeSolverOptions, OdeSolverProblem}; + +use crate::error::CoreError; + +/// Our replacement for diffsol's whole-solve nonlinear-failure budget. +/// +/// The counter is cumulative over a solve, so it bounds solve length, not divergence; +/// `min_timestep` catches divergence. +const MAX_NONLINEAR_SOLVER_FAILURES: usize = 100_000; + +/// Call `$callback!` with every option, each as `$(#[doc])* name: Type,`. +/// +/// The one list that [`SolverOptions`] and each mirror of it, such as the bindings, expand. +#[macro_export] +macro_rules! for_each_solver_option { + ($callback:ident) => { + $callback! { + /// Newton iterations allowed per nonlinear solve. + max_nonlinear_solver_iterations: usize, + /// Consecutive step rejections allowed before the solve aborts. + max_error_test_failures: usize, + /// Nonlinear-solver convergence failures allowed over the whole solve. + max_nonlinear_solver_failures: usize, + /// Newton convergence-test scaling factor. + nonlinear_solver_tolerance: f64, + /// Smallest step the controller may take before the solve aborts. + min_timestep: f64, + /// Upper bound on step-size growth; `None` keeps diffsol's own. + max_timestep_growth: Option, + /// Lower bound of the step-growth dead zone; `None` keeps diffsol's own. + min_timestep_growth: Option, + /// Upper bound of the step-shrink dead zone; `None` keeps diffsol's own. + max_timestep_shrink: Option, + /// Absolute lower bound on step-size reduction; `None` keeps diffsol's own. + min_timestep_shrink: Option, + /// Steps between linear-solver setups. + update_jacobian_after_steps: usize, + /// Steps between full RHS Jacobian re-evaluations. + update_rhs_jacobian_after_steps: usize, + /// Relative step-size change that forces a Jacobian update. + threshold_to_update_jacobian: f64, + /// Relative step-size change that forces an RHS Jacobian update. + threshold_to_update_rhs_jacobian: f64, + /// Proportional gain of the step-size PI controller. + pi_control_proportional: f64, + /// Integral gain of the step-size PI controller. + pi_control_integral: f64, + } + }; +} + +/// Declare [`SolverOptions`] with its copies from and onto diffsol's own options. +macro_rules! define_solver_options { + ($($(#[$doc:meta])* $field:ident: $ty:ty,)*) => { + /// Tuning for the diffsol BDF integrator. + #[derive(Debug, Clone, PartialEq)] + pub struct SolverOptions { + $($(#[$doc])* pub $field: $ty,)* + } + + impl SolverOptions { + /// diffsol's own defaults, including the failure budget we otherwise raise. + #[must_use] + pub fn diffsol_defaults() -> Self { + let defaults = OdeSolverOptions::::default(); + Self { $($field: defaults.$field,)* } + } + + /// Stamp these options onto a freshly built problem, before the solver is created. + #[must_use] + pub(crate) const fn apply>( + &self, + mut problem: OdeSolverProblem, + ) -> OdeSolverProblem { + $(problem.ode_options.$field = self.$field;)* + problem + } + } + }; +} + +for_each_solver_option!(define_solver_options); + +impl Default for SolverOptions { + fn default() -> Self { + Self { + max_nonlinear_solver_failures: MAX_NONLINEAR_SOLVER_FAILURES, + ..Self::diffsol_defaults() + } + } +} + +impl SolverOptions { + /// Reject values diffsol would take at face value and then misbehave on. + pub fn validate(&self) -> Result<(), CoreError> { + // `(name, value, must be positive)`; an absent optional bound is not checked. + let checks = [ + ( + "nonlinear_solver_tolerance", + Some(self.nonlinear_solver_tolerance), + true, + ), + ("min_timestep", Some(self.min_timestep), true), + ( + "threshold_to_update_jacobian", + Some(self.threshold_to_update_jacobian), + true, + ), + ( + "threshold_to_update_rhs_jacobian", + Some(self.threshold_to_update_rhs_jacobian), + true, + ), + ("max_timestep_growth", self.max_timestep_growth, true), + ("min_timestep_growth", self.min_timestep_growth, true), + ("max_timestep_shrink", self.max_timestep_shrink, true), + ("min_timestep_shrink", self.min_timestep_shrink, true), + // A PI gain of zero disables that term, so only finiteness applies. + ( + "pi_control_proportional", + Some(self.pi_control_proportional), + false, + ), + ("pi_control_integral", Some(self.pi_control_integral), false), + ]; + for (name, value, positive) in checks { + if let Some(got) = value + && !(got.is_finite() && (!positive || got > 0.0)) + { + let requirement = if positive { "finite and > 0" } else { "finite" }; + return Err(CoreError::SolverOption { + name, + requirement, + got, + }); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_defaults_are_valid_and_depart_from_diffsol_only_in_the_failure_budget() { + let ours = SolverOptions::default(); + ours.validate().expect("defaults invalid"); + // The raised budget clears the measured worst case. + assert!(ours.max_nonlinear_solver_failures > 10_000); + let theirs = SolverOptions::diffsol_defaults(); + assert_ne!( + ours.max_nonlinear_solver_failures, + theirs.max_nonlinear_solver_failures + ); + assert_eq!( + SolverOptions { + max_nonlinear_solver_failures: theirs.max_nonlinear_solver_failures, + ..ours + }, + theirs, + "a second departure from diffsol's defaults needs its own rationale", + ); + } + + #[test] + fn values_diffsol_would_misbehave_on_are_rejected() { + let tolerance = |got| SolverOptions { + nonlinear_solver_tolerance: got, + ..Default::default() + }; + let positive = ("nonlinear_solver_tolerance", "finite and > 0"); + let cases = [ + (tolerance(0.0), positive), + (tolerance(-1.0), positive), + (tolerance(f64::NAN), positive), + (tolerance(f64::INFINITY), positive), + // A present optional bound is checked as a required one is. + ( + SolverOptions { + max_timestep_growth: Some(-2.0), + ..Default::default() + }, + ("max_timestep_growth", "finite and > 0"), + ), + ( + SolverOptions { + pi_control_integral: f64::NAN, + ..Default::default() + }, + ("pi_control_integral", "finite"), + ), + ]; + let mut message = String::new(); + for (options, (name, requirement)) in cases { + let err = options.validate().expect_err(name); + assert!( + matches!(err, CoreError::SolverOption { name: n, requirement: r, .. } if n == name && r == requirement), + "{err:?}" + ); + message = err.to_string(); + } + assert_eq!( + message, + "solver option pi_control_integral must be finite, got NaN" + ); + } + + #[test] + fn an_absent_bound_and_a_zero_pi_gain_are_accepted() { + // A zero PI gain disables its term rather than breaking the controller. + for options in [ + SolverOptions { + max_timestep_growth: None, + ..Default::default() + }, + SolverOptions { + pi_control_proportional: 0.0, + ..Default::default() + }, + ] { + options.validate().expect("accepted"); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/rhs.rs b/packages/pybamm-rust/pybamm-core/src/solver/rhs.rs new file mode 100644 index 0000000000..faf6ae18a3 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/rhs.rs @@ -0,0 +1,118 @@ +//! Right-hand-side operator: `f(t, y; p)`, its sparse `df/dy`, and `df/dp`. +//! +//! `df/dp` columns are numbered by position in the parameter subset, not registration index. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + Context, FaerContext, Matrix, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op, Vector, +}; + +use super::{FaerSparsity, dense_faer_sparsity, lane, lane_mut, lanes, lanes_mut, values_mut}; +use crate::model::{CompiledModel, Workspace}; +use crate::params::ParamBinding; + +/// RHS operator `f(t, y; p)`. +pub struct RhsOp<'a> { + pub compiled: &'a CompiledModel, + pub ws: &'a RefCell, + pub inputs: &'a [f64], + /// What the solve's parameter vector stands for. + pub binding: &'a ParamBinding, + pub jac_sparsity: &'a FaerSparsity, + pub context: FaerContext, +} + +impl std::fmt::Debug for RhsOp<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RhsOp") + .field("n_states", &self.compiled.n_states()) + .field("nparams", &self.binding.len()) + .finish_non_exhaustive() + } +} + +impl Op for RhsOp<'_> { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.compiled.n_states() + } + fn nout(&self) -> usize { + self.compiled.n_states() + } + fn nparams(&self) -> usize { + self.binding.len() + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl NonLinearOp for RhsOp<'_> { + fn call_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled + .eval_rhs(&mut ws, t, lane(x), self.inputs, lane_mut(y)); + } +} + +impl NonLinearOpJacobian for RhsOp<'_> { + /// Write `df/dy · v` into `y` for each lane of `v`, all at state `x`. + fn jac_mul_inplace(&self, x: &FaerVec, t: f64, v: &FaerVec, y: &mut FaerVec) { + assert_eq!( + v.context().nbatch(), + y.context().nbatch(), + "jac_mul direction and output lanes differ" + ); + let mut ws = self.ws.borrow_mut(); + self.compiled + .jac_action_at(&mut ws, t, lane(x), self.inputs) + .apply_lanes(lanes(v), lanes_mut(y)); + } + + fn jacobian_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerSparseMat) { + let mut ws = self.ws.borrow_mut(); + self.compiled.assemble_jacobian_csc_no_mass( + &mut ws, + t, + lane(x), + self.inputs, + values_mut(y), + ); + } + + fn jacobian_sparsity(&self) -> Option< as Matrix>::Sparsity> { + Some(self.jac_sparsity.clone()) + } +} + +impl NonLinearOpSens for RhsOp<'_> { + fn sens_mul_inplace(&self, x: &FaerVec, t: f64, v: &FaerVec, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled + .sens_action(&mut ws, t, lane(x), self.inputs, lane(v), lane_mut(y)); + } + + /// Batched `df/dp`: one primal pass, then a tangent-only sweep per column. + /// + /// `y` must use the dense pattern of [`Self::sens_sparsity`]. + fn sens_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerSparseMat) { + let mut ws = self.ws.borrow_mut(); + self.compiled + .eval_sens_all(&mut ws, t, lane(x), self.inputs, values_mut(y)); + } + + /// Dense pattern for `df/dp`: every state may depend on every parameter. + fn sens_sparsity(&self) -> Option { + Some(dense_faer_sparsity( + self.compiled.n_states(), + self.binding.len(), + )) + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/solve.rs b/packages/pybamm-rust/pybamm-core/src/solver/solve.rs new file mode 100644 index 0000000000..f62a6da16c --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/solve.rs @@ -0,0 +1,3359 @@ +//! Problem setup and the solve loop. +//! +//! [`PreparedSolver`] holds everything reusable about a model; each solve builds its own +//! workspace, problem and diffsol instance, keeping the handle `Send + Sync`. + +use std::cell::RefCell; +use std::ops::Range; +use std::sync::Arc; +use std::time::Instant; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::ode_solver::OdeSolverStatistics; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + AugmentedOdeSolverMethod, BacktrackingLineSearch, BdfState, FaerContext, FaerSparseLU, + InitialConditionSolverOptions, NewtonNonlinearSolver, NoLineSearch, NonLinearOp, OdeBuilder, + OdeEquations, OdeSolverMethod, OdeSolverProblem, OdeSolverStopReason, Op, SensEquations, + StateRef, StateRefMut, Vector, +}; +use faer::dyn_stack::{MemBuffer, MemStack}; +use faer::sparse::linalg::lu::{ + LuRef, LuSymbolicParams, NumericLu, SymbolicLu, factorize_symbolic_lu, +}; +use faer::sparse::{SparseColMatRef, SymbolicSparseColMat}; +use faer::{Conj, Mat, MatMut, Par, Spec}; + +use super::equations::Equations; +use super::options::SolverOptions; +use super::{FaerSparsity, csc_to_faer_sparsity, csr_mass_to_faer_csc, lane, lane_mut, lanes}; +use crate::arena::{Arena, DagRef, NodeId}; +use crate::error::CoreError; +use crate::eval::CompiledExpr; +use crate::jacobian::CscPattern; +use crate::model::{CompiledModel, ModelEvaluator, Workspace}; +use crate::node::{CsrData, Node}; +use crate::params::ParamBinding; +use crate::simplify::{dce, simplify_pipeline}; +use crate::tangent::derivative_wrt_time; + +/// Sparse LU used by every Newton iteration. +type SparseLu = FaerSparseLU; + +/// Default multiplier on differential-state `atol` for the sensitivity `atol` floor. +/// +/// 1.0 matches IDAS's `IDASensEEtolerances`: `atol / |p_j|`, with `sens_rtol = rtol`. +pub const DEFAULT_SENS_ATOL_FACTOR: f64 = 1.0; + +/// Solver statistics from a BDF solve. +#[derive(Debug, Clone, Default)] +pub struct SolverStatistics { + /// diffsol's step, Newton and linear-solver counters. + pub solver: OdeSolverStatistics, + /// Time spent building this solve's workspace, `Equations` and problem. + pub problem_setup_time_secs: f64, + /// Time spent computing consistent initial conditions. + pub ic_time_secs: f64, + /// Time spent building the BDF solver around the consistent initial state. + pub solver_setup_time_secs: f64, + /// Wall-clock time for the whole solve, excluding FFI marshalling. + pub integration_time_secs: f64, + /// True when the sensitivity solve was retried with sensitivities out of error control. + pub sens_error_control_relaxed: bool, +} + +/// Seconds since `mark`, moving `mark` to now. +fn lap(mark: &mut Instant) -> f64 { + let now = Instant::now(); + let secs = now.duration_since(*mark).as_secs_f64(); + *mark = now; + secs +} + +/// How a solve ended. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Termination { + Completed, + /// The root of event `index`, in registration order, stopped the solve at time `t`. + Event { + t: f64, + index: usize, + }, +} + +/// What one solve returns, whichever payloads were asked for. +/// +/// On an event root, `t` ends at the root time rather than the last requested point. +#[derive(Debug)] +pub struct SolveOutcome { + /// Output times, one per trajectory column. + pub t: Vec, + /// Flat column-major trajectory: `y[i + j * n_rows]` is row `i` at time `j`. + /// Rows are states, or output variables when [`SolveRequest::outputs`] is set. + pub y: Vec, + /// Rows in `y`, and in each sensitivity block. + pub n_rows: usize, + /// Row time derivatives in `y`'s layout: only on a state trajectory with + /// [`PreparedSolver::with_store_yp`] set and at most `MAX_HERMITE_COLUMNS` columns. + pub yp: Option>, + /// One block per sensitivity parameter in `y`'s layout and request order. + /// On an outputs request each block is the full `dg/dp + dg/dy · y_s`. + pub sensitivities: Option>>, + /// The full state at a root, or at the end on an outputs request; never an outputs row. + pub y_event: Option>, + pub termination: Termination, + /// The state after the initial-condition solve, consistent even if `y0` was a guess. + pub y_start: Vec, + /// State sensitivities at `y_start`, one `n_states` block per sensitivity parameter. + pub y_start_sens: Option>>, + /// State sensitivities at `y_event`, laid out as `y_start_sens`. + pub y_event_sens: Option>>, + pub statistics: SolverStatistics, +} + +/// Copy every lane of a batched sensitivity vector, one block per parameter. +fn sens_blocks(s: &FaerVec) -> Vec> { + lanes(s).map(<[f64]>::to_vec).collect() +} + +/// What to integrate over and which payloads to report. +#[derive(Clone, Copy, Debug)] +#[must_use] +pub struct SolveRequest<'a> { + pub t_eval: &'a [f64], + /// Discontinuity times the integrator lands on and restarts from. + /// Entries not also in `t_eval` are integrated through. + pub t_stop: &'a [f64], + /// Report the model's output variables instead of the full state. + pub outputs: bool, + /// Solve the forward-sensitivity system, seeded per set by [`InputSet::y0_sens`]. + pub sensitivities: bool, + /// Workers of the current rayon pool the sparse LU may use; 1 factorises serially. + pub lu_threads: usize, +} + +impl<'a> SolveRequest<'a> { + /// A state-trajectory request over `t_eval`: no stop times, no sensitivities. + pub const fn new(t_eval: &'a [f64]) -> Self { + Self { + t_eval, + t_stop: &[], + outputs: false, + sensitivities: false, + lu_threads: 1, + } + } + + /// Land on, and restart from, each of `t_stop`. + pub const fn with_stop_times(mut self, t_stop: &'a [f64]) -> Self { + self.t_stop = t_stop; + self + } + + /// Report output-variable rows rather than states. + pub const fn with_outputs(mut self) -> Self { + self.outputs = true; + self + } + + /// Solve the forward-sensitivity system alongside the trajectory. + pub const fn with_sensitivities(mut self) -> Self { + self.sensitivities = true; + self + } + + /// Factorise with up to `threads` workers of the current rayon pool. + pub const fn with_lu_threads(mut self, threads: usize) -> Self { + self.lu_threads = threads; + self + } + + fn lu_par(&self) -> Par { + if self.lu_threads > 1 { + Par::rayon(self.lu_threads) + } else { + Par::Seq + } + } +} + +/// One input set of a solve: where the state starts and what the parameters are. +#[derive(Clone, Copy, Debug)] +#[must_use] +pub struct InputSet<'a> { + pub y0: &'a [f64], + /// Flat input-parameter vector. + pub inputs: &'a [f64], + /// `dy0/dp`, column-major `n_states x k` over the sensitivity subset; empty is zero. + pub y0_sens: &'a [f64], +} + +impl<'a> InputSet<'a> { + /// An input set carrying the all-zero `dy0/dp` seed. + pub const fn new(y0: &'a [f64], inputs: &'a [f64]) -> Self { + Self { + y0, + inputs, + y0_sens: &[], + } + } + + /// Seed the sensitivity system with `y0_sens` rather than with zeros. + pub const fn with_sens_seed(mut self, y0_sens: &'a [f64]) -> Self { + self.y0_sens = y0_sens; + self + } +} + +/// Prepare-once/execute-many handle for repeated solves of the same model. +/// +/// Holds no per-solve mutable state, so it is `Send + Sync`. +pub struct PreparedSolver { + compiled: Arc, + rtol: f64, + atol: Vec, + /// Multiplier applied to `atol` on differential states only. + sens_atol_factor: f64, + options: SolverOptions, + store_yp: bool, + /// `dF/dt` at fixed states and inputs, for time-dependent algebraic constraints. + time_derivative: Option, + /// How a pure ODE reads `dy/dt` off `f`; `None` on a DAE. + ode_mass: Option, + /// `g_v`'s layout, built once [`Self::with_store_yp`] needs it on a DAE. + algebraic_block: Option, + jac_sparsity: Arc, + mass_sparsity: Arc, + /// M's values in `mass_sparsity` order. + mass_csc_values: Arc<[f64]>, + /// The configured sensitivity subset. + sens_binding: ParamBinding, +} + +impl std::fmt::Debug for PreparedSolver { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PreparedSolver") + .field("n_states", &self.compiled.n_states()) + .field("input_width", &self.compiled.input_width()) + .field("n_events", &self.compiled.n_events()) + .field("n_outputs", &self.compiled.total_output_len()) + .finish_non_exhaustive() + } +} + +impl PreparedSolver { + /// Build a prepared problem from a compiled model and tolerances. + pub fn new( + model: impl Into>, + rtol: f64, + atol: &[f64], + ) -> Result { + let compiled = model.into(); + let n_states = compiled.n_states(); + if atol.len() != n_states { + return Err(CoreError::AtolLength { + got: atol.len(), + expected: n_states, + }); + } + validate_tolerances(rtol, atol)?; + + let mass = compiled.mass_matrix(); + let (mass_sparsity, mass_csc_values) = csr_mass_to_faer_csc(mass); + let ode_mass = (compiled.state_kinds().n_algebraic() == 0) + .then(|| OdeMass::new(mass, &mass_sparsity, &mass_csc_values)) + .flatten(); + let sens_binding = + ParamBinding::sensitivities(compiled.params(), compiled.sens_param_indices()); + + Ok(Self { + jac_sparsity: Arc::new(csc_to_faer_sparsity(compiled.csc_sparsity())), + compiled, + rtol, + atol: atol.to_vec(), + sens_atol_factor: DEFAULT_SENS_ATOL_FACTOR, + options: SolverOptions::default(), + store_yp: false, + time_derivative: None, + ode_mass, + algebraic_block: None, + mass_sparsity: Arc::new(mass_sparsity), + mass_csc_values: mass_csc_values.into(), + sens_binding, + }) + } + + fn n_states(&self) -> usize { + self.compiled.n_states() + } + + fn n_outputs(&self) -> usize { + self.compiled.total_output_len() + } + + /// The plain path's binding: the packed input array itself. + fn packed_binding(&self) -> ParamBinding { + ParamBinding::packed_inputs(self.compiled.input_width()) + } + + fn build_eqn(&self, set: InputSet<'_>, binding: &ParamBinding, lu_par: Par) -> Equations { + let mut context = FaerContext::default(); + context.par = lu_par; + Equations { + evaluator: ModelEvaluator::from_compiled(Arc::clone(&self.compiled)), + params: set.inputs.to_vec(), + binding: binding.clone(), + y0: set.y0.to_vec(), + y0_sens: set.y0_sens.to_vec(), + jac_sparsity: Arc::clone(&self.jac_sparsity), + mass_sparsity: Arc::clone(&self.mass_sparsity), + mass_csc_values: Arc::clone(&self.mass_csc_values), + context, + } + } + + /// One solve's diffsol problem, owning its workspace, over `binding`'s parameter vector. + /// + /// `sens_error_control` puts the sensitivities under error control, scaled by `|p_j|`. + fn build_problem( + &self, + set: InputSet<'_>, + t0: f64, + binding: &ParamBinding, + sens_error_control: bool, + lu_par: Par, + ) -> Result, CoreError> { + let eqn = self.build_eqn(set, binding, lu_par); + // The ops report the binding's width, so p narrows with it. + let mut p = vec![0.0; binding.len()]; + binding.gather_from(set.inputs, &mut p); + let mut builder = OdeBuilder::>::new() + .rtol(self.rtol) + .atol(self.atol.clone()); + if sens_error_control { + // param_scales plays the role of IDAS's pbar. + builder = builder + .sens_rtol(self.rtol) + .sens_atol(self.sens_atol()) + .param_scales(Self::param_scales(&p)); + } + Ok(self.options.apply(builder.t0(t0).p(p).build_from_eqn(eqn)?)) + } + + /// Set the multiplier on differential-state `atol` for the sensitivity tolerance floor. + pub fn with_sens_atol_factor(mut self, factor: f64) -> Result { + if !factor.is_finite() || factor <= 0.0 { + return Err(CoreError::SensAtolFactor { got: factor }); + } + self.sens_atol_factor = factor; + Ok(self) + } + + /// Set the integrator tuning applied to every problem this template builds. + pub fn with_options(mut self, options: SolverOptions) -> Result { + options.validate()?; + self.options = options; + Ok(self) + } + + /// Store state time derivatives (`yp`) alongside `y` on state-trajectory paths. + /// + /// Doubles trajectory memory, so it is opt-in. + #[must_use] + pub fn with_store_yp(mut self, store_yp: bool) -> Self { + self.store_yp = store_yp; + let has_algebraic = self.compiled.state_kinds().n_algebraic() > 0; + if store_yp && has_algebraic && self.algebraic_block.is_none() { + let block = AlgebraicBlock::new(self.compiled.csc_sparsity(), self.algebraic_states()); + self.algebraic_block = Some(block); + } + self + } + + /// Compile `dF/dt` from `rhs`, the root this model was compiled from in `arena`. + /// + /// Only compiled when there are algebraic states and `rhs` reads `t`. + #[must_use] + pub fn with_time_derivative(mut self, arena: &Arena, rhs: NodeId) -> Self { + let reads_time = || { + arena + .topological_order(rhs) + .iter() + .any(|&id| matches!(arena.get(id), Node::Time)) + }; + if self.compiled.state_kinds().n_algebraic() > 0 && reads_time() { + // Copies only what `rhs` reaches, not the whole graph. + let mut derivative = dce(DagRef::new(arena, rhs)); + derivative.rewrite_in_place(derivative_wrt_time); + self.time_derivative = Some(CompiledExpr::new(simplify_pipeline(derivative).as_ref())); + } + self + } + + /// Whether a solve reporting `n_columns` columns stores `yp` alongside `y`. + const fn store_yp(&self, n_columns: usize) -> bool { + self.store_yp && n_columns <= MAX_HERMITE_COLUMNS + } + + fn algebraic_states(&self) -> &[bool] { + self.compiled.state_kinds().algebraic_mask() + } + + /// Per-state sensitivity absolute tolerance. + /// + /// Algebraic states keep the state `atol`: tightening them fails DAE sensitivity solves. + fn sens_atol(&self) -> Vec { + self.atol + .iter() + .zip(self.algebraic_states()) + .map(|(a, is_algebraic)| { + if *is_algebraic { + *a + } else { + a * self.sens_atol_factor + } + }) + .collect() + } + + /// Per-parameter scales `|p_j|` for the sensitivity absolute tolerances. + /// + /// `sens_inputs` is the length-`k` sensitivity subset, not the packed inputs; zero + /// or non-finite values fall back to 1.0, since diffsol rejects them. + fn param_scales(sens_inputs: &[f64]) -> Vec { + sens_inputs + .iter() + .map(|v| { + if v.is_finite() && *v != 0.0 { + v.abs() + } else { + 1.0 + } + }) + .collect() + } + + /// An empty trajectory over `n_points` columns of states, or of outputs when `with_output`. + fn empty_trajectory( + &self, + n_points: usize, + with_output: bool, + y_start: &[f64], + ) -> DenseTrajectory { + // Output mode never stores yp: its trajectory holds outputs, not states. + let store_yp = self.store_yp(n_points) && !with_output; + let n_rows = if with_output { + self.n_outputs() + } else { + self.n_states() + }; + let mut trajectory = DenseTrajectory::with_capacity(n_rows, n_points, store_yp); + trajectory.y_start = y_start.to_vec(); + trajectory + } + + /// Validate caller-supplied solve arguments against the model dimensions. + fn validate_args(&self, y0: &[f64], inputs: &[f64]) -> Result<(), CoreError> { + if y0.len() != self.n_states() { + return Err(CoreError::Y0Length { + got: y0.len(), + expected: self.n_states(), + }); + } + let width = self.compiled.input_width().get(); + if inputs.len() != width { + return Err(CoreError::InputsLength { + got: inputs.len(), + expected: width, + }); + } + Ok(()) + } + + fn validate_has_outputs(&self) -> Result<(), CoreError> { + if self.n_outputs() == 0 { + return Err(CoreError::NoOutputVariables); + } + Ok(()) + } + + /// Validate the caller-supplied `dy0/dp` seed: empty, or column-major `n_states x k`. + fn validate_y0_sens(&self, y0_sens: &[f64]) -> Result<(), CoreError> { + if self.sens_binding.is_empty() { + return Err(CoreError::NoSensitivityParams); + } + let expected = self.n_states() * self.sens_binding.len(); + if y0_sens.is_empty() || y0_sens.len() == expected { + return Ok(()); + } + Err(CoreError::Y0SensLength { + got: y0_sens.len(), + expected, + }) + } + + /// Integrate the model, returning whichever payloads `request` asked for. + pub fn solve( + &self, + request: SolveRequest<'_>, + set: InputSet<'_>, + ) -> Result { + let times = TimePlan::new(request.t_eval, request.t_stop)?; + self.validate_args(set.y0, set.inputs)?; + if request.sensitivities { + self.validate_y0_sens(set.y0_sens)?; + } + if request.outputs { + self.validate_has_outputs()?; + } + let lu_par = request.lu_par(); + + if request.sensitivities { + self.solve_sensitivities(times, set, request.outputs, lu_par) + } else { + self.solve_states(times, set, request.outputs, lu_par) + } + } + + /// Integrate without sensitivities; reports output rows when `with_output` is set. + fn solve_states( + &self, + times: TimePlan<'_>, + set: InputSet<'_>, + with_output: bool, + lu_par: Par, + ) -> Result { + let started = Instant::now(); + let mut mark = started; + let mut statistics = SolverStatistics::default(); + // Nothing here reads a `dy0/dp` seed. + let set = InputSet::new(set.y0, set.inputs); + let problem = self.build_problem(set, times.t0(), &self.packed_binding(), false, lu_par)?; + statistics.problem_setup_time_secs = lap(&mut mark); + let state = problem.bdf_state::()?; + statistics.ic_time_secs = lap(&mut mark); + let mut solver = problem.bdf_solver::(state)?; + statistics.solver_setup_time_secs = lap(&mut mark); + + let run = Run { + times, + inputs: set.inputs, + with_output, + started, + lu_par, + }; + let ws = &problem.eqn.evaluator.workspace; + let mut recorder = Recorder::new(self, run, ws, false, &solver.state()); + let mut restarter = BreakpointRestarter::new(self, &problem.ic_options); + let termination = drive_segments(&mut solver, times, &mut recorder, |solver, t_next| { + restarter.restart(solver, t_next) + })?; + statistics.solver = solver.get_statistics().clone(); + Ok(recorder.finish(termination, &solver.state(), statistics)) + } + + /// Integrate with sensitivities, retrying once without them in error control. + fn solve_sensitivities( + &self, + times: TimePlan<'_>, + set: InputSet<'_>, + with_output: bool, + lu_par: Par, + ) -> Result { + let started = Instant::now(); + let mut mark = started; + let mut statistics = SolverStatistics::default(); + let mut problem = self.build_problem(set, times.t0(), &self.sens_binding, true, lu_par)?; + statistics.problem_setup_time_secs = lap(&mut mark); + // Only the state tolerances shape the IC, so both attempts start from this state. + let state = problem.bdf_state_sens::()?; + statistics.ic_time_secs = lap(&mut mark); + + let run = Run { + times, + inputs: set.inputs, + with_output, + started, + lu_par, + }; + match self.integrate_sensitivities(&problem, state.clone(), run, statistics.clone()) { + // Only an integration failure depends on the sensitivities' error control. + Err(CoreError::Diffsol(controlled)) => { + // `SensEquations` reads these as the solver is built, so this relaxes it. + problem.sens_rtol = None; + problem.sens_atol = None; + let mut outcome = self + .integrate_sensitivities(&problem, state, run, statistics) + .map_err(|relaxed| CoreError::SensRetryFailed { + controlled: controlled.to_string(), + relaxed: relaxed.to_string(), + })?; + outcome.statistics.sens_error_control_relaxed = true; + Ok(outcome) + }, + other => other, + } + } + + /// One sensitivity integration from `state`, under `problem`'s sensitivity tolerances. + /// + /// `statistics` carries the setup times the attempts share. + fn integrate_sensitivities( + &self, + problem: &OdeSolverProblem, + state: BdfState>, + run: Run<'_>, + mut statistics: SolverStatistics, + ) -> Result { + let mut mark = Instant::now(); + let mut solver = problem.bdf_solver_sens::(state)?; + statistics.solver_setup_time_secs = lap(&mut mark); + let ws = &problem.eqn.evaluator.workspace; + let mut recorder = Recorder::new(self, run, ws, true, &solver.state()); + let mut restarter = BreakpointRestarter::new(self, &problem.ic_options); + let termination = + drive_segments(&mut solver, run.times, &mut recorder, |solver, t_next| { + restarter.restart_sens(solver, t_next) + })?; + statistics.solver = solver.get_statistics().clone(); + Ok(recorder.finish(termination, &solver.state(), statistics)) + } +} + +/// What one solve integrates over. +#[derive(Clone, Copy)] +struct Run<'r> { + times: TimePlan<'r>, + inputs: &'r [f64], + /// Report output rows rather than states. + with_output: bool, + /// When the solve started, for its integration time. + started: Instant, + /// The parallelism every sparse LU of the solve runs at. + lu_par: Par, +} + +/// Reads each output column off the solver and files its rows and sensitivities. +struct Recorder<'p> { + prepared: &'p PreparedSolver, + ws: &'p RefCell, + inputs: &'p [f64], + started: Instant, + lu_par: Par, + /// The time of the last column filed. + last_time: Option, + /// The column being read: the state, and `dy` when the trajectory stores `yp`. + y: FaerVec, + dy: Option>, + /// `(t, dy)` made consistent for the column at `t` that starts a segment. + held_dy: Option<(f64, Vec)>, + hold_scratch: HoldScratch, + rows: Rows, + sens: Option, +} + +/// Where a solve files its rows. +enum Rows { + /// Straight into a state trajectory. + States(DenseTrajectory), + /// Staged for batched evaluation into an outputs trajectory. + Outputs(OutputBatchWindow), +} + +/// The sensitivity half of a solve's columns. +struct SensColumns { + /// `S` as read for the current column, one lane per parameter. + lanes: FaerVec, + /// One block per parameter in the trajectory's layout. + blocks: Vec>, + /// `S` at the start, one `n_states` block per parameter. + start: Vec>, + /// Projects `S` onto the output rows, on an outputs request. + outputs: Option, +} + +impl SensColumns { + /// File the sensitivities of the column at `(t, y)`, read into `self.lanes`. + fn push( + &mut self, + compiled: &CompiledModel, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + ) { + if let Some(outputs) = &mut self.outputs { + outputs.project(compiled, ws, t, y, inputs, &self.lanes); + for (block, column) in self.blocks.iter_mut().zip(outputs.sens_columns()) { + block.extend_from_slice(column); + } + } else { + for (block, lane) in self.blocks.iter_mut().zip(lanes(&self.lanes)) { + block.extend_from_slice(lane); + } + } + } +} + +impl<'p> Recorder<'p> { + /// A recorder for `run`, whose solver starts at `start`. + fn new( + prepared: &'p PreparedSolver, + run: Run<'p>, + ws: &'p RefCell, + sensitivities: bool, + start: &StateRef<'_, FaerVec>, + ) -> Self { + let n_points = run.times.eval.len(); + let trajectory = prepared.empty_trajectory(n_points, run.with_output, lane(start.y)); + let n_columns = prepared.sens_binding.len(); + let sens = sensitivities.then(|| SensColumns { + lanes: start.s.clone(), + // One reservation per block: `vec![v; n]` would clone away all but one. + blocks: (0..n_columns) + .map(|_| Vec::with_capacity(trajectory.n_rows * n_points)) + .collect(), + start: sens_blocks(start.s), + outputs: run + .with_output + .then(|| OutputSensBuffers::new(prepared, n_columns)), + }); + let dy = trajectory.stores_yp().then(|| start.dy.clone()); + let rows = if run.with_output { + Rows::Outputs(OutputBatchWindow::new(prepared.n_states(), trajectory)) + } else { + Rows::States(trajectory) + }; + Self { + prepared, + ws, + inputs: run.inputs, + started: run.started, + lu_par: run.lu_par, + last_time: None, + y: start.y.clone(), + dy, + held_dy: None, + hold_scratch: HoldScratch::default(), + rows, + sens, + } + } + + /// Hold the consistent derivative for a segment starting at `t`, if the solver sits on it. + fn begin(&mut self, state: &StateRef<'_, FaerVec>, t: f64) { + self.held_dy = if self.dy.is_some() && within_round_off(state.t, t) { + let mut ws = self.ws.borrow_mut(); + let held = self.prepared.consistent_dy( + &mut ws, + self.inputs, + state, + t, + self.lu_par, + &mut self.hold_scratch, + ); + held.map(|dy| (t, dy)) + } else { + None + }; + } + + /// Read one column off `solver` and file it. + fn record<'a, S: OdeSolverMethod<'a, Equations>>( + &mut self, + solver: &S, + column: ColumnSource, + ) -> Result<(), CoreError> { + let sens = self.sens.as_mut().map(|sens| &mut sens.lanes); + let t = match column { + ColumnSource::Interpolated(t) => { + solver.interpolate_inplace(t, &mut self.y)?; + if let Some(s) = sens { + solver.interpolate_sens_inplace(t, s)?; + } + if let Some(dy) = self.dy.as_mut() { + solver.interpolate_dy_inplace(t, dy)?; + } + t + }, + // The state sits at, or within round-off of, this column, past interpolation's reach. + ColumnSource::CurrentState(t) => { + let state = solver.state(); + self.y.copy_from(state.y); + if let Some(s) = sens { + s.copy_from(state.s); + } + if let Some(dy) = self.dy.as_mut() { + dy.copy_from(state.dy); + } + t + }, + }; + #[allow(clippy::float_cmp)] // the column is read at exactly the held time + if let (Some(dy), Some((_, held))) = ( + self.dy.as_mut(), + self.held_dy.take_if(|(t_held, _)| *t_held == t), + ) { + lane_mut(dy).copy_from_slice(&held); + } + + let (compiled, y) = (&self.prepared.compiled, lane(&self.y)); + match &mut self.rows { + Rows::States(trajectory) => trajectory.push_column(t, y, self.dy.as_ref().map(lane)), + Rows::Outputs(window) => window.push(t, y, compiled, self.ws, self.inputs), + } + if let Some(sens) = &mut self.sens { + sens.push(compiled, &mut self.ws.borrow_mut(), t, y, self.inputs); + } + self.last_time = Some(t); + Ok(()) + } + + /// Package every filed column as the outcome of a solve that ended as `termination` says. + /// + /// `end` is the solver's final state, wound back to the root on an event. + fn finish( + self, + termination: Termination, + end: &StateRef<'_, FaerVec>, + mut statistics: SolverStatistics, + ) -> SolveOutcome { + let rows_are_outputs = matches!(self.rows, Rows::Outputs(_)); + let trajectory = match self.rows { + Rows::States(trajectory) => trajectory, + Rows::Outputs(window) => window.finish( + &self.prepared.compiled, + &mut self.ws.borrow_mut(), + self.inputs, + ), + }; + statistics.integration_time_secs = self.started.elapsed().as_secs_f64(); + let mut outcome = + trajectory.into_outcome(termination, lane(end.y), rows_are_outputs, statistics); + if let Some(sens) = self.sens { + // Read where y_event was, so the two describe the same point. + outcome.y_event_sens = outcome.y_event.is_some().then(|| sens_blocks(end.s)); + outcome.sensitivities = Some(sens.blocks); + outcome.y_start_sens = Some(sens.start); + } + outcome + } +} + +/// Interpolation points staged before one batched output evaluation. +const OUTPUT_BATCH_LANES: usize = 128; + +/// Staging buffer for the batched output-variable path, filling an outputs trajectory. +struct OutputBatchWindow { + staged_t: Vec, + /// `(n_states, k)` F-contiguous staged states. + staged_y: Vec, + /// `(n_rows, k)` F-contiguous per-window results. + window_out: Vec, + trajectory: DenseTrajectory, +} + +impl OutputBatchWindow { + fn new(n_states: usize, trajectory: DenseTrajectory) -> Self { + Self { + staged_t: Vec::with_capacity(OUTPUT_BATCH_LANES), + staged_y: Vec::with_capacity(n_states * OUTPUT_BATCH_LANES), + window_out: vec![0.0; trajectory.n_rows * OUTPUT_BATCH_LANES], + trajectory, + } + } + + /// Stage the state `y`, all `n_states` of it, at `t`; evaluates the window once it is full. + fn push( + &mut self, + t: f64, + y: &[f64], + compiled: &CompiledModel, + ws: &RefCell, + inputs: &[f64], + ) { + self.staged_t.push(t); + self.staged_y.extend_from_slice(y); + if self.staged_t.len() == OUTPUT_BATCH_LANES { + self.flush(compiled, &mut ws.borrow_mut(), inputs); + } + } + + fn flush(&mut self, compiled: &CompiledModel, ws: &mut Workspace, inputs: &[f64]) { + let k = self.staged_t.len(); + if k == 0 { + return; + } + let window_out = &mut self.window_out[..self.trajectory.n_rows * k]; + compiled.eval_outputs_batch(ws, k, &self.staged_t, &self.staged_y, inputs, window_out); + self.trajectory.push_columns(&self.staged_t, window_out); + self.staged_t.clear(); + self.staged_y.clear(); + } + + /// Flush the last, partial window and hand back the filled trajectory. + fn finish( + mut self, + compiled: &CompiledModel, + ws: &mut Workspace, + inputs: &[f64], + ) -> DenseTrajectory { + self.flush(compiled, ws, inputs); + self.trajectory + } +} + +/// Buffers projecting state sensitivities onto the outputs: `dg/dp + dg/dy · S`. +struct OutputSensBuffers { + n_outputs: usize, + /// `S` gathered from its lanes, `n_states x k` column-major. + state_sens: Vec, + /// `dg/dp + dg/dy · S`, `n_outputs x k` column-major. + output_sens: Vec, +} + +impl OutputSensBuffers { + fn new(prepared: &PreparedSolver, n_columns: usize) -> Self { + let n_outputs = prepared.n_outputs(); + Self { + n_outputs, + state_sens: vec![0.0; prepared.n_states() * n_columns], + output_sens: vec![0.0; n_outputs * n_columns], + } + } + + /// Project `S`, held in `sens_lanes`, onto the outputs at `(t, y)`. + fn project( + &mut self, + compiled: &CompiledModel, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + sens_lanes: &FaerVec, + ) { + for (block, lane) in self + .state_sens + .chunks_exact_mut(y.len()) + .zip(lanes(sens_lanes)) + { + block.copy_from_slice(lane); + } + compiled.output_sens_project(ws, t, y, inputs, &self.state_sens, &mut self.output_sens); + } + + /// The sensitivity columns from the last [`Self::project`], one per parameter. + fn sens_columns(&self) -> std::slice::ChunksExact<'_, f64> { + self.output_sens.chunks_exact(self.n_outputs) + } +} + +/// Reject what IDAS rejects: a non-finite or negative tolerance, or a zero `rtol` and `atol` pair. +/// +/// diffsol takes any value, and a NaN weight stalls the step controller rather than failing. +fn validate_tolerances(rtol: f64, atol: &[f64]) -> Result<(), CoreError> { + let invalid = |value: f64| !value.is_finite() || value < 0.0; + if invalid(rtol) { + return Err(CoreError::Tolerance { + name: "rtol", + index: None, + got: rtol, + }); + } + let bad_atol = atol + .iter() + .position(|&value| invalid(value) || (value == 0.0 && rtol == 0.0)); + bad_atol.map_or(Ok(()), |index| { + Err(CoreError::Tolerance { + name: "atol", + index: Some(index), + got: atol[index], + }) + }) +} + +/// A sparse LU's factors, applied at an explicit [`Par`] rather than faer's global one. +struct LuFactors { + symbolic: SymbolicLu, + numeric: NumericLu, +} + +impl LuFactors { + /// Factor `matrix` at `par`, or `None` if it is singular. + fn new(matrix: SparseColMatRef<'_, usize, f64>, par: Par) -> Option { + let symbolic = + factorize_symbolic_lu(matrix.symbolic(), LuSymbolicParams::default()).ok()?; + let numeric = factor_numeric(&symbolic, matrix, par)?; + Some(Self { symbolic, numeric }) + } + + /// Overwrite `rhs` with `A⁻¹ rhs`. + fn solve_in_place(&self, rhs: MatMut<'_, f64>, par: Par) { + lu_solve_in_place(&self.symbolic, &self.numeric, rhs, par); + } +} + +/// Factor `matrix`, which has the pattern `symbolic` was built from, or `None` if it is singular. +fn factor_numeric( + symbolic: &SymbolicLu, + matrix: SparseColMatRef<'_, usize, f64>, + par: Par, +) -> Option> { + let mut numeric = NumericLu::new(); + let scratch = symbolic.factorize_numeric_lu_scratch::(par, Spec::default()); + let mut scratch = MemBuffer::try_new(scratch).ok()?; + symbolic + .factorize_numeric_lu( + &mut numeric, + matrix, + par, + MemStack::new(&mut scratch), + Spec::default(), + ) + .ok()?; + Some(numeric) +} + +/// Overwrite `rhs` with `A⁻¹ rhs`, for `A` factored into `symbolic` and `numeric`. +fn lu_solve_in_place( + symbolic: &SymbolicLu, + numeric: &NumericLu, + rhs: MatMut<'_, f64>, + par: Par, +) { + let mut scratch = MemBuffer::new(symbolic.solve_in_place_scratch::(rhs.ncols(), par)); + LuRef::new_unchecked(symbolic, numeric).solve_in_place_with_conj( + Conj::No, + rhs, + par, + MemStack::new(&mut scratch), + ); +} + +/// How a pure ODE's `dy/dt` follows from `M dy/dt = f`. +enum OdeMass { + /// `M` has nothing nonzero off its diagonal. + Diagonal(Vec), + /// `M`'s LU, for an `M` with off-diagonal entries. + General(Box), +} + +impl OdeMass { + /// `None` for a non-diagonal `M` that cannot be factored. + /// + /// Factored serially: it runs once, as the solver is prepared. + fn new(mass: &CsrData, sparsity: &FaerSparsity, values: &[f64]) -> Option { + mass_diagonal(mass).map(Self::Diagonal).or_else(|| { + let mass = SparseColMatRef::new(sparsity.as_ref(), values); + LuFactors::new(mass, Par::Seq).map(|lu| Self::General(Box::new(lu))) + }) + } + + fn diagonal(&self) -> Option<&[f64]> { + match self { + Self::Diagonal(diagonal) => Some(diagonal), + Self::General(_) => None, + } + } + + /// Overwrite `f` with `M⁻¹ f`. + fn solve_in_place(&self, f: &mut [f64], par: Par) { + match self { + Self::Diagonal(diagonal) => { + for (f, m) in f.iter_mut().zip(diagonal) { + *f /= m; + } + }, + Self::General(lu) => { + let n = f.len(); + lu.solve_in_place(MatMut::from_column_major_slice_mut(f, n, 1), par); + }, + } + } +} + +/// M's diagonal, or `None` if any stored nonzero sits off it. +fn mass_diagonal(mass: &CsrData) -> Option> { + let indptr = mass.indptr(); + let mut diagonal = vec![0.0; indptr.len().saturating_sub(1)]; + for (row, entry) in diagonal.iter_mut().enumerate() { + for k in indptr[row]..indptr[row + 1] { + if mass.indices()[k] == row { + *entry = mass.data()[k]; + } else if mass.data()[k] != 0.0 { + return None; + } + } + } + Some(diagonal) +} + +/// The Jacobian's algebraic block `g_v`, laid out once for every held derivative. +struct AlgebraicBlock { + /// The algebraic states, ascending: `g_v`'s rows and columns in order. + states: Vec, + /// Where each `g_v` entry sits in the Jacobian's CSC values, in `g_v`'s CSC order. + slots: Vec, + pattern: SymbolicSparseColMat, + /// `None` when `g_v` is structurally singular. + symbolic_lu: Option>, +} + +impl AlgebraicBlock { + fn new(csc: &CscPattern, algebraic: &[bool]) -> Self { + let states: Vec = (0..algebraic.len()).filter(|&i| algebraic[i]).collect(); + let mut local = vec![usize::MAX; algebraic.len()]; + for (k, &i) in states.iter().enumerate() { + local[i] = k; + } + let (mut col_ptr, mut row_idx, mut slots) = (vec![0], Vec::new(), Vec::new()); + for &col in &states { + for slot in csc.colptr()[col]..csc.colptr()[col + 1] { + let row = csc.rowind()[slot]; + if algebraic[row] { + row_idx.push(local[row]); + slots.push(slot); + } + } + col_ptr.push(row_idx.len()); + } + let n = states.len(); + let pattern = SymbolicSparseColMat::new_checked(n, n, col_ptr, None, row_idx); + let symbolic_lu = factorize_symbolic_lu(pattern.as_ref(), LuSymbolicParams::default()).ok(); + Self { + states, + slots, + pattern, + symbolic_lu, + } + } + + /// Overwrite `rhs` with `g_v⁻¹ rhs`, `g_v` read out of the full Jacobian's CSC `values`. + /// + /// `None`, leaving `rhs` untouched, when `g_v` is singular. + fn solve_in_place(&self, values: &[f64], rhs: MatMut<'_, f64>, par: Par) -> Option<()> { + let symbolic = self.symbolic_lu.as_ref()?; + let g_v: Vec = self.slots.iter().map(|&slot| values[slot]).collect(); + let g_v = SparseColMatRef::new(self.pattern.as_ref(), &g_v); + let numeric = factor_numeric(symbolic, g_v, par)?; + lu_solve_in_place(symbolic, &numeric, rhs, par); + Some(()) + } +} + +/// The Newton solve behind a consistent state, line-searched like diffsol's default. +type IcNewton = + NewtonNonlinearSolver, SparseLu, BacktrackingLineSearch>>; + +/// An [`IcNewton`] configured as diffsol's `new_and_consistent` configures its own. +/// +/// Always line-searched: no solve here clears `use_linesearch`. +fn ic_newton(options: &InitialConditionSolverOptions) -> IcNewton { + let mut line_search = BacktrackingLineSearch::default(); + line_search.c = options.armijo_constant; + line_search.max_iter = options.max_linesearch_iterations; + line_search.tau = options.step_reduction_factor; + NewtonNonlinearSolver::new(SparseLu::default(), line_search) +} + +/// Newton solvers and buffers reused by every breakpoint restart in one solve. +/// +/// Only the state solve is line-searched; the sensitivity system is linear, so its +/// undamped step is already exact. +struct BreakpointRestarter<'p> { + root_solver: IcNewton, + sens_solver: NewtonNonlinearSolver, SparseLu, NoLineSearch>, + /// M's diagonal for a pure ODE with diagonal M; `None` leaves `dy` to `set_consistent`. + ode_mass_diagonal: Option<&'p [f64]>, +} + +impl<'p> BreakpointRestarter<'p> { + fn new(prepared: &'p PreparedSolver, ic_options: &InitialConditionSolverOptions) -> Self { + Self { + root_solver: ic_newton(ic_options), + sens_solver: NewtonNonlinearSolver::new(SparseLu::default(), NoLineSearch), + ode_mass_diagonal: prepared.ode_mass.as_ref().and_then(OdeMass::diagonal), + } + } + + /// Rebuild `dy` from `M dy = f(t, y)` when nothing else will. + /// + /// diffsol's `set_consistent` skips an ODE, leaving `dy` at the pre-discontinuity slope. + fn refresh_ode_dy( + &self, + problem: &OdeSolverProblem, + state: &mut StateRefMut<'_, FaerVec>, + ) { + let Some(diagonal) = self.ode_mass_diagonal else { + return; + }; + problem.eqn.rhs().call_inplace(state.y, *state.t, state.dy); + for (dy, m) in lane_mut(state.dy).iter_mut().zip(diagonal) { + *dy /= m; + } + } + + /// Move onto the far side of the corner at `t_next` and make `y` and `dy` consistent. + fn make_consistent<'a, S: OdeSolverMethod<'a, Equations>>( + &mut self, + solver: &mut S, + t_next: f64, + ) -> Result<(), CoreError> { + let problem = solver.problem(); + let mut state = solver.state_mut(); + cross_corner(&mut state, t_next); + state.set_consistent(problem, &mut self.root_solver)?; + self.refresh_ode_dy(problem, &mut state); + Ok(()) + } + + /// Restart the integrator at a breakpoint, as IDAKLU's `HandleBreakpoint` does. + fn restart<'a, S: OdeSolverMethod<'a, Equations>>( + &mut self, + solver: &mut S, + t_next: f64, + ) -> Result<(), CoreError> { + self.make_consistent(solver, t_next)?; + reset_step_size(solver); + Ok(()) + } + + /// [`Self::restart`] for a solve carrying forward sensitivities. + fn restart_sens<'a, S>(&mut self, solver: &mut S, t_next: f64) -> Result<(), CoreError> + where + S: AugmentedOdeSolverMethod<'a, Equations, SensEquations<'a, Equations>>, + { + self.make_consistent(solver, t_next)?; + // Reads the y and dy just made consistent, so it has to follow them. + let problem = solver.problem(); + if let Some((mut state, augmented_eqn)) = solver.state_and_augmented_eqn_mut() { + state.set_consistent_augmented(problem, augmented_eqn, &mut self.sens_solver)?; + } + reset_step_size(solver); + Ok(()) + } +} + +/// Pick a fresh first-order step size, as diffsol does after its own IC solve. +fn reset_step_size<'a, S: OdeSolverMethod<'a, Equations>>(solver: &mut S) { + let problem = solver.problem(); + solver + .state_mut() + .set_step_size(problem.h0, &problem.atol, problem.rtol, &problem.eqn, 1); +} + +/// Buffers reused by every consistent derivative a solve holds. +#[derive(Default)] +struct HoldScratch { + /// The Jacobian's CSC values. + jacobian: Vec, + /// The `dF/dt` tape's scratch. + time_scratch: Vec, +} + +impl PreparedSolver { + /// `dy` at `state` made consistent, or `None` to keep the interpolated slope. + /// + /// `None` when `g_v`, or a pure ODE's non-diagonal `M`, cannot be factored. + fn consistent_dy( + &self, + ws: &mut Workspace, + inputs: &[f64], + state: &StateRef<'_, FaerVec>, + t: f64, + lu_par: Par, + scratch: &mut HoldScratch, + ) -> Option> { + let compiled = &self.compiled; + let algebraic = self.algebraic_states(); + let (y, mut dy) = (lane(state.y), lane(state.dy).to_vec()); + if !algebraic.contains(&true) { + // diffsol seeds an ODE's dy with f itself, never solving M dy = f. + let ode_mass = self.ode_mass.as_ref()?; + compiled.eval_rhs(ws, t, y, inputs, &mut dy); + ode_mass.solve_in_place(&mut dy, lu_par); + return Some(dy); + } + + let block = self.algebraic_block.as_ref()?; + let n = self.n_states(); + let du: Vec = dy + .iter() + .zip(algebraic) + .map(|(&d, &is_algebraic)| if is_algebraic { 0.0 } else { d }) + .collect(); + let mut g_u_du = vec![0.0; n]; + compiled.jac_action(ws, t, y, inputs, &du, &mut g_u_du); + let mut time_derivative = vec![0.0; n]; + if let Some(expr) = &self.time_derivative { + scratch.time_scratch.clear(); + scratch.time_scratch.resize(expr.scratch_len(), 0.0); + let result = expr.eval(&mut scratch.time_scratch, t, y, &[], inputs); + time_derivative[..result.len()].copy_from_slice(result); + } + + scratch.jacobian.clear(); + scratch.jacobian.resize(compiled.csc_sparsity().nnz(), 0.0); + compiled.assemble_jacobian_csc_no_mass(ws, t, y, inputs, &mut scratch.jacobian); + let mut dv = Mat::from_fn(block.states.len(), 1, |k, _| { + let i = block.states[k]; + -(g_u_du[i] + time_derivative[i]) + }); + block.solve_in_place(&scratch.jacobian, dv.as_mut(), lu_par)?; + for (k, &i) in block.states.iter().enumerate() { + dy[i] = dv[(k, 0)]; + } + dy.iter().all(|d| d.is_finite()).then_some(dy) + } +} + +/// The times a solve reports at, together with the stop times inside them. +/// +/// `eval` is non-empty and never decreases: [`drive_segments`] would otherwise report a +/// truncated solve. +#[derive(Clone, Copy, Debug)] +struct TimePlan<'a> { + eval: &'a [f64], + stop: &'a [f64], +} + +impl<'a> TimePlan<'a> { + /// Check `eval`; equal times are discontinuity brackets, so only a decrease or a + /// non-finite time fails. + const fn new(eval: &'a [f64], stop: &'a [f64]) -> Result { + if eval.is_empty() { + return Err(CoreError::EmptyTimePoints); + } + let mut i = 0; + while i < eval.len() { + let got = eval[i]; + if !got.is_finite() { + return Err(CoreError::NonFiniteTimePoint { index: i, got }); + } + if i > 0 && got < eval[i - 1] { + return Err(CoreError::UnsortedTimePoints { + index: i, + got, + previous: eval[i - 1], + }); + } + i += 1; + } + Ok(Self { eval, stop }) + } + + /// The first output time, where integration starts. + const fn t0(&self) -> f64 { + self.eval[0] + } + + fn segments(&self) -> Vec> { + segment_ranges(self.eval, self.stop) + } +} + +/// Where one output column's values come from. +/// +/// Interpolated once the solver has stepped on or past it; read off the state only for an +/// event's root and for columns left within round-off past where the state stopped. +#[derive(Clone, Copy, Debug)] +enum ColumnSource { + Interpolated(f64), + CurrentState(f64), +} + +/// Step one `t_stop` segment at a time, recording every output column. +/// +/// On an event, the solver is left wound back to its root. +fn drive_segments<'a, S: OdeSolverMethod<'a, Equations>>( + solver: &mut S, + times: TimePlan<'_>, + recorder: &mut Recorder<'_>, + mut restart: impl FnMut(&mut S, f64) -> Result<(), CoreError>, +) -> Result { + let t_eval = times.eval; + let mut col = 0; + for (i, range) in times.segments().into_iter().enumerate() { + let last = range.end - 1; + if i > 0 { + // The far side of the ULP-bracketed corner, where the new branch is evaluable. + restart(solver, t_eval[range.start])?; + } + recorder.begin(&solver.state(), t_eval[range.start]); + // A consistent-state solve can jump past a root the root finder would then miss. + if let Some(index) = violated_event(solver) { + let t = solver.state().t; + // A restart that stays on the last column's time ends the solve without a new one. + if let Some(last_time) = recorder.last_time.filter(|&last_time| t <= last_time) { + return Ok(Termination::Event { + t: last_time, + index, + }); + } + recorder.record(solver, ColumnSource::CurrentState(t))?; + return Ok(Termination::Event { t, index }); + } + // `set_stop_time` rejects a stop diffsol already counts as reached. + if !stop_time_reached(&solver.state(), t_eval[last]) { + solver.set_stop_time(t_eval[last])?; + loop { + let reason = solver.step()?; + let drain_until = match reason { + OdeSolverStopReason::InternalTimestep | OdeSolverStopReason::TstopReached => { + solver.state().t + }, + OdeSolverStopReason::RootFound(t_root, _) => t_root, + }; + while col <= last && t_eval[col] <= drain_until { + recorder.record(solver, ColumnSource::Interpolated(t_eval[col]))?; + col += 1; + } + match reason { + OdeSolverStopReason::InternalTimestep => {}, + OdeSolverStopReason::TstopReached => break, + OdeSolverStopReason::RootFound(t_root, index) => { + solver.state_mut_back(t_root)?; + if col <= last { + recorder.record(solver, ColumnSource::CurrentState(t_root))?; + } + return Ok(Termination::Event { t: t_root, index }); + }, + } + } + } + // Whatever is left lies within diffsol's stop round-off, past where the state stopped. + for &t in &t_eval[col..=last] { + recorder.record(solver, ColumnSource::CurrentState(t))?; + } + col = last + 1; + } + Ok(Termination::Completed) +} + +/// diffsol's own test for a stop at `t_stop` being reached, round-off included. +fn stop_time_reached(state: &StateRef<'_, FaerVec>, t_stop: f64) -> bool { + (state.t - t_stop).abs() <= 100.0 * f64::EPSILON * (state.t.abs() + state.h.abs()) +} + +/// The most negative termination event at the solver's current state, if any is non-positive. +fn violated_event<'a, S: OdeSolverMethod<'a, Equations>>(solver: &S) -> Option { + let events = solver.problem().eqn.root()?; + let state = solver.state(); + let mut values = FaerVec::zeros(events.nout(), *state.y.context()); + events.call_inplace(state.y, state.t, &mut values); + lane(&values) + .iter() + .enumerate() + .filter(|(_, value)| **value <= 0.0) + .min_by(|(_, a), (_, b)| a.total_cmp(b)) + .map(|(index, _)| index) +} + +/// Move a restart at a corner onto its far side, `t_next`, where the new branch holds. +fn cross_corner(state: &mut StateRefMut<'_, FaerVec>, t_next: f64) { + if within_round_off(*state.t, t_next) { + *state.t = t_next; + } +} + +/// Whether `from` and `to` are one instant to diffsol's stop round-off, without its step term. +/// +/// `PyBaMM`'s ULP-apart discontinuity brackets collapse to the earlier, pre-discontinuity time. +fn within_round_off(from: f64, to: f64) -> bool { + (to - from).abs() <= 100.0 * f64::EPSILON * from.abs().max(to.abs()) +} + +/// One range of `t_eval` per stretch between consecutive stop times. +/// +/// A stop time not on the output grid is integrated through rather than ending a segment. +fn segment_ranges(t_eval: &[f64], t_stop: &[f64]) -> Vec> { + let last = t_eval.len().saturating_sub(1); + let mut cuts: Vec = Vec::new(); + for stop in t_stop { + let idx = t_eval.partition_point(|t| t < stop); + // Bit equality, not tolerance: the value was copied from this grid. + let on_grid = idx < t_eval.len() && t_eval[idx].to_bits() == stop.to_bits(); + if !(on_grid && idx > 0 && idx < last) { + continue; + } + // Keeps cuts increasing, so an unsorted `t_stop` cannot cause a slice panic. + let previous = cuts.last().copied().unwrap_or(0); + if idx <= previous { + continue; + } + if within_round_off(t_eval[previous], t_eval[idx]) { + continue; + } + cuts.push(idx); + } + + let mut ranges = Vec::with_capacity(cuts.len() + 1); + let mut start = 0; + for cut in cuts { + ranges.push(start..cut + 1); + start = cut + 1; + } + ranges.push(start..t_eval.len()); + ranges +} + +/// Output columns above which a solve gives up storing `yp`. +const MAX_HERMITE_COLUMNS: usize = 4096; + +/// A column-major trajectory accumulated one output column at a time. +struct DenseTrajectory { + t: Vec, + y: Vec, + /// `dy/dt` columns sharing `y`'s layout, when stored. + yp: Option>, + /// Trajectory width: states, or outputs on the output paths. + n_rows: usize, + y_start: Vec, +} + +impl DenseTrajectory { + fn with_capacity(n_rows: usize, n_points: usize, store_yp: bool) -> Self { + Self { + t: Vec::with_capacity(n_points), + y: Vec::with_capacity(n_rows * n_points), + yp: store_yp.then(|| Vec::with_capacity(n_rows * n_points)), + n_rows, + y_start: Vec::new(), + } + } + + /// Package the trajectory as the outcome of a solve that ended as `termination` says. + /// + /// `final_state` is the full state where the trajectory ends. It is reported at a root, + /// and always when the rows are outputs, as then it is the caller's only state. + fn into_outcome( + self, + termination: Termination, + final_state: &[f64], + rows_are_outputs: bool, + statistics: SolverStatistics, + ) -> SolveOutcome { + let at_root = matches!(termination, Termination::Event { .. }); + SolveOutcome { + n_rows: self.n_rows, + t: self.t, + y: self.y, + yp: self.yp, + sensitivities: None, + y_event: (at_root || rows_are_outputs).then(|| final_state.to_vec()), + termination, + y_start: self.y_start, + y_start_sens: None, + y_event_sens: None, + statistics, + } + } + + fn push_column(&mut self, t: f64, y: &[f64], yp: Option<&[f64]>) { + debug_assert_eq!(y.len(), self.n_rows, "column width mismatch"); + debug_assert_eq!( + yp.is_some(), + self.yp.is_some(), + "every column must carry yp exactly when the trajectory stores it" + ); + self.t.push(t); + self.y.extend_from_slice(y); + if let (Some(buffer), Some(column)) = (self.yp.as_mut(), yp) { + debug_assert_eq!(column.len(), self.n_rows, "yp column width mismatch"); + buffer.extend_from_slice(column); + } + } + + /// Append one column per entry of `t`, stored back to back in `y`, storing no `yp`. + fn push_columns(&mut self, t: &[f64], y: &[f64]) { + debug_assert_eq!( + y.len(), + t.len() * self.n_rows, + "column block width mismatch" + ); + debug_assert!(self.yp.is_none(), "a column block carries no yp"); + self.t.extend_from_slice(t); + self.y.extend_from_slice(y); + } + + const fn stores_yp(&self) -> bool { + self.yp.is_some() + } +} + +#[cfg(test)] +mod segment_tests { + use super::segment_ranges; + + #[test] + fn no_stop_times_is_one_segment() { + let t = [0.0, 1.0, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[]), vec![0..4]); + } + + #[test] + fn an_interior_stop_time_ends_a_segment_on_itself() { + let t = [0.0, 1.0, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[2.0]), vec![0..3, 3..4]); + } + + #[test] + fn the_span_endpoints_are_not_interior_stop_times() { + let t = [0.0, 1.0, 2.0]; + assert_eq!(segment_ranges(&t, &[0.0, 2.0]), vec![0..3]); + } + + #[test] + fn a_stop_time_on_the_last_point_leaves_no_trailing_segment() { + let t = [0.0, 1.0, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[2.0, 3.0]), vec![0..3, 3..4]); + } + + #[test] + fn stop_times_absent_from_the_output_grid_are_ignored() { + let t = [0.0, 1.0, 2.0]; + assert_eq!(segment_ranges(&t, &[1.5]), vec![0..3]); + } + + #[test] + fn consecutive_output_points_can_each_be_a_stop_time() { + let t = [0.0, 1.0, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[1.0, 2.0]), vec![0..2, 2..3, 3..4]); + } + + #[test] + fn a_discontinuity_bracket_collapses_to_its_earlier_side() { + // diffsol rejects the second of a ULP-apart pair as a stop already reached. + let below = f64::from_bits(2.0f64.to_bits() - 1); + let t = [0.0, 1.0, below, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[below, 2.0]), vec![0..3, 3..5]); + } + + #[test] + fn unsorted_stop_times_are_skipped_not_reversed_into_a_panic() { + let t = [0.0, 1.0, 2.0, 3.0, 4.0]; + assert_eq!(segment_ranges(&t, &[3.0, 1.0]), vec![0..4, 4..5]); + } + + #[test] + fn a_repeated_stop_time_does_not_cut_twice() { + let t = [0.0, 1.0, 2.0, 3.0]; + assert_eq!(segment_ranges(&t, &[2.0, 2.0]), vec![0..3, 3..4]); + } + + #[test] + fn a_bracket_at_the_end_of_the_span_still_cuts() { + // The final point restarts on the far side of the corner, as a mid-span bracket does. + let below = f64::from_bits(3.0f64.to_bits() - 1); + let t = [0.0, 1.0, below, 3.0]; + assert_eq!(segment_ranges(&t, &[below]), vec![0..3, 3..4]); + } +} +#[cfg(test)] +mod tests { + use diffsol::{Matrix, NonLinearOpSens, OdeEquations, Op, Vector}; + use faer::linalg::solvers::Solve; + use faer::sparse::{SparseColMat, Triplet}; + + use super::*; + use crate::arena::Arena; + use crate::model::test_fixtures::{compile, compile_with_sens, identity_mass_matrix, input}; + use crate::node::{CsrData, Node, Shape}; + use crate::observable::ObservableKind; + use crate::params::ParamIdx; + use crate::solver::test_fixtures::{ + decay, decay_with_output_and_event, differential_then_algebraic_mass, + }; + + /// `dy/dt = -a * y` and `0 = y - z`, with sensitivities wrt `a`. + fn build_small_dae_with_param() -> CompiledModel { + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg_a_y = { + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + arena.alloc(Node::Mul(neg_a, sv0)) + }; + let algebraic_residual = arena.alloc(Node::Sub(sv0, sv1)); + let rhs = arena.alloc(Node::Concat(vec![neg_a_y, algebraic_residual])); + let mass = differential_then_algebraic_mass(); + compile_with_sens(&arena, rhs, mass, 2, 1, &[ParamIdx::new(0)]) + } + + /// A Heaviside node constructor, such as `Node::EqualHeaviside`. + type Switch = fn(NodeId, NodeId) -> Node; + + /// `dv/dt = -a*v` with `0 = z - v*(1 - 2*H)`, `H = switch(0.5, t)`: `z` jumps from `v` to `-v`. + fn build_dae_with_algebraic_jump(switch: Switch, terminate: bool) -> CompiledModel { + let mut arena = Arena::new(); + let v = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dvdt = arena.alloc(Node::Mul(neg_a, v)); + let t = arena.alloc(Node::Time); + let corner = arena.alloc(Node::Scalar(0.5)); + let after = arena.alloc(switch(corner, t)); + let two = arena.alloc(Node::Scalar(2.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let two_after = arena.alloc(Node::Mul(two, after)); + let sign = arena.alloc(Node::Sub(one, two_after)); + let target = arena.alloc(Node::Mul(v, sign)); + let constraint = arena.alloc(Node::Sub(z, target)); + let rhs = arena.alloc(Node::Concat(vec![dvdt, constraint])); + let mass = differential_then_algebraic_mass(); + let mut model = compile_with_sens(&arena, rhs, mass, 2, 1, &[ParamIdx::new(0)]); + model + .set_observables(ObservableKind::Outputs, &mut arena, &[z]) + .expect("a valid output"); + if terminate { + model + .set_observables(ObservableKind::Events, &mut arena, &[z]) + .expect("a valid event"); + } + model + } + + #[test] + fn a_corner_restart_lands_the_algebraic_state_on_the_new_branch() { + // PyBaMM brackets `t > 0.5` with 0.5 and the time one ULP above it. + let above = f64::from_bits(0.5f64.to_bits() + 1); + let t_eval = [0.0, 0.25, 0.5, above, 0.75, 1.0, 1.5, 2.0]; + let a = 0.1; + let request = SolveRequest::new(&t_eval).with_stop_times(&t_eval); + for (store_yp, request) in [ + (false, request), + (true, request), + (false, request.with_sensitivities()), + ] { + let prepared = PreparedSolver::new( + build_dae_with_algebraic_jump(Node::NotEqualHeaviside, false), + 1e-6, + &[1e-6, 1e-6], + ) + .expect("prepare") + .with_store_yp(store_yp); + let result = prepared + .solve(request, InputSet::new(&[1.0, 1.0], &[a])) + .expect("a jump in an algebraic state at a corner must not stall the solve"); + + for (j, &t) in t_eval.iter().enumerate() { + let sign = if t > 0.5 { -1.0 } else { 1.0 }; + let v = (-a * t).exp(); + let got = &result.y[j * result.n_rows..(j + 1) * result.n_rows]; + assert!( + (got[0] - v).abs() < 1e-5, + "v({t}) = {} but expected {v}", + got[0] + ); + assert!((got[1] - sign * v).abs() < 1e-5, "z({t}) = {}", got[1]); + if let Some(sens) = &result.sensitivities { + let dv_da = -t * v; + let got = &sens[0][j * result.n_rows..(j + 1) * result.n_rows]; + assert!((got[0] - dv_da).abs() < 1e-5, "dv/da({t}) = {}", got[0]); + assert!( + (got[1] - sign * dv_da).abs() < 1e-5, + "dz/da({t}) = {}", + got[1] + ); + } + } + } + } + + #[test] + fn a_corner_restart_reports_a_termination_event_crossed_by_the_jump() { + let above = f64::from_bits(0.5f64.to_bits() + 1); + let times = [0.0, 0.25, 0.5, above, 0.75, 1.0]; + let request = SolveRequest::new(×).with_stop_times(×); + let prepared = PreparedSolver::new( + build_dae_with_algebraic_jump(Node::NotEqualHeaviside, true), + 1e-8, + &[1e-8, 1e-8], + ) + .expect("prepare"); + + for request in [ + request, + request.with_outputs(), + request.with_sensitivities(), + request.with_outputs().with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&[1.0, 1.0], &[0.1])) + .expect("solve"); + assert_eq!( + result.termination, + Termination::Event { t: above, index: 0 } + ); + assert_eq!(result.t, times[..4]); + let event_state = result.y_event.expect("event state"); + assert!((event_state[0] - (-0.05_f64).exp()).abs() < 1e-6); + assert!((event_state[1] + (-0.05_f64).exp()).abs() < 1e-6); + } + } + + #[test] + fn a_termination_column_at_a_restart_carries_the_consistent_derivative() { + // On the new branch z = -v, so dz/dt = a v where set_consistent leaves zero. + let above = f64::from_bits(0.5f64.to_bits() + 1); + let times = [0.0, 0.25, 0.5, above, 0.75, 1.0]; + let a = 0.1; + let prepared = PreparedSolver::new( + build_dae_with_algebraic_jump(Node::NotEqualHeaviside, true), + 1e-8, + &[1e-8, 1e-8], + ) + .expect("prepare") + .with_store_yp(true); + let request = SolveRequest::new(×).with_stop_times(×); + for request in [request, request.with_sensitivities()] { + let result = prepared + .solve(request, InputSet::new(&[1.0, 1.0], &[a])) + .expect("solve"); + assert_eq!(result.t, times[..4], "{request:?}"); + let yp = result.yp.expect("yp requested"); + let root = &yp[3 * result.n_rows..]; + let v = (-a * above).exp(); + assert!((root[0] + a * v).abs() < 1e-6, "dv/dt = {}", root[0]); + assert!((root[1] - a * v).abs() < 1e-6, "dz/dt = {}", root[1]); + } + } + + #[test] + // The termination is reported at the time of the column already filed. + #[allow(clippy::float_cmp)] + fn an_event_violated_by_a_restart_in_place_adds_no_column() { + // 0.5 is no corner bracket, so the restart stays at the time the last segment ended. + let prepared = build_decay_param_output_event().with_store_yp(true); + let t_eval = [0.0, 0.25, 0.5, 0.75, 1.0]; + let times = TimePlan::new(&t_eval, &[0.5]).expect("sorted grid"); + let set = InputSet::new(&[1.0], &[1.0]); + let problem = prepared + .build_problem(set, 0.0, &prepared.packed_binding(), false, Par::Seq) + .expect("problem"); + let state = problem.bdf_state::().expect("IC"); + let mut solver = problem.bdf_solver::(state).expect("solver"); + let run = Run { + times, + inputs: set.inputs, + with_output: false, + started: Instant::now(), + lu_par: Par::Seq, + }; + let ws = &problem.eqn.evaluator.workspace; + let mut recorder = Recorder::new(&prepared, run, ws, false, &solver.state()); + // Drops the state past the event at y = 0.4, as a jump in an algebraic state can. + let termination = drive_segments(&mut solver, times, &mut recorder, |solver, _| { + lane_mut(solver.state_mut().y)[0] = 0.3; + Ok(()) + }) + .expect("solve"); + let result = recorder.finish(termination, &solver.state(), SolverStatistics::default()); + + assert_eq!(result.termination, Termination::Event { t: 0.5, index: 0 }); + assert_eq!(result.t, t_eval[..3]); + assert_eq!(result.yp.expect("yp requested").len(), 3); + assert_eq!(result.y_event, Some(vec![0.3])); + } + + #[test] + fn a_bracket_at_the_end_of_the_span_restarts_on_the_new_branch() { + // PyBaMM brackets a switch at t_end as [0.5, above] for `t > 0.5`, [below, 0.5] for `>=`. + let below = f64::from_bits(0.5f64.to_bits() - 1); + let above = f64::from_bits(0.5f64.to_bits() + 1); + let a = 0.1; + let v = (-a * 0.5f64).exp(); + let brackets: [(Switch, [f64; 2]); 2] = [ + (Node::NotEqualHeaviside, [0.5, above]), + (Node::EqualHeaviside, [below, 0.5]), + ]; + for (switch, bracket) in brackets { + let t_eval = [0.0, 0.25, bracket[0], bracket[1]]; + let prepared = PreparedSolver::new( + build_dae_with_algebraic_jump(switch, false), + 1e-8, + &[1e-8, 1e-8], + ) + .expect("prepare"); + let request = SolveRequest::new(&t_eval).with_stop_times(&t_eval); + for request in [request, request.with_sensitivities()] { + let result = prepared + .solve(request, InputSet::new(&[1.0, 1.0], &[a])) + .expect("a switch at the final time must not stall the solve"); + let end = &result.y[3 * result.n_rows..]; + assert!((end[0] - v).abs() < 1e-6, "v = {}, {request:?}", end[0]); + assert!((end[1] + v).abs() < 1e-6, "z = {}, {request:?}", end[1]); + if let Some(blocks) = &result.sensitivities { + // dz/da = -dv/da = t v on the new branch. + let dz_da = blocks[0][3 * result.n_rows + 1]; + assert!((dz_da - 0.5 * v).abs() < 1e-5, "dz/da = {dz_da}"); + } + } + } + } + + /// `0 = atan(z - v - 3*(t > 0.5))`: a jump too far for an undamped Newton step. + fn build_dae_with_a_far_algebraic_jump() -> CompiledModel { + let mut arena = Arena::new(); + let v = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dvdt = arena.alloc(Node::Mul(neg_a, v)); + let t = arena.alloc(Node::Time); + let corner = arena.alloc(Node::Scalar(0.5)); + let after = arena.alloc(Node::NotEqualHeaviside(corner, t)); + let three = arena.alloc(Node::Scalar(3.0)); + let jump = arena.alloc(Node::Mul(three, after)); + let target = arena.alloc(Node::Add(v, jump)); + let offset = arena.alloc(Node::Sub(z, target)); + let constraint = arena.alloc(Node::Arctan(offset)); + let rhs = arena.alloc(Node::Concat(vec![dvdt, constraint])); + compile(&arena, rhs, differential_then_algebraic_mass(), 2, 1) + } + + #[test] + fn a_corner_restart_line_searches_a_far_algebraic_jump() { + let above = f64::from_bits(0.5f64.to_bits() + 1); + let t_eval = [0.0, 0.25, 0.5, above, 1.0]; + let a = 0.1; + let prepared = + PreparedSolver::new(build_dae_with_a_far_algebraic_jump(), 1e-6, &[1e-6, 1e-6]) + .expect("prepare"); + let result = prepared + .solve( + SolveRequest::new(&t_eval).with_stop_times(&t_eval), + InputSet::new(&[1.0, 1.0], &[a]), + ) + .expect("the restart's Newton must reach a far branch, as diffsol's own IC does"); + + let j = t_eval.len() - 1; + let z = result.y[j * result.n_rows + 1]; + let expected = (-a).exp() + 3.0; + assert!( + (z - expected).abs() < 1e-5, + "z(1) = {z} but expected {expected}" + ); + } + + /// `dy/dt = -a*y` with `0 = 2y - z - z^3`, so the IC Newton stops short of consistency. + fn build_nonlinear_dae() -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dydt = arena.alloc(Node::Mul(neg_a, y)); + let two = arena.alloc(Node::Scalar(2.0)); + let two_y = arena.alloc(Node::Mul(two, y)); + let z2 = arena.alloc(Node::Mul(z, z)); + let z3 = arena.alloc(Node::Mul(z2, z)); + let z_plus_z3 = arena.alloc(Node::Add(z, z3)); + let constraint = arena.alloc(Node::Sub(two_y, z_plus_z3)); + let rhs = arena.alloc(Node::Concat(vec![dydt, constraint])); + compile(&arena, rhs, differential_then_algebraic_mass(), 2, 1) + } + + #[test] + fn the_first_column_stores_the_consistent_derivative() { + // From y = 1 the constraint gives z = 1 and dz/dt = 2 dy/dt / (1 + 3z^2) = -0.5. + let prepared = PreparedSolver::new(build_nonlinear_dae(), 1e-6, &[1e-6, 1e-6]) + .expect("prepare") + .with_store_yp(true); + let t_eval: Vec = (0..=4).map(|i| f64::from(i) * 0.5).collect(); + let result = prepared + .solve( + SolveRequest::new(&t_eval), + InputSet::new(&[1.0, 3.0], &[1.0]), + ) + .expect("solve"); + + assert!( + (result.y_start[1] - 1.0).abs() < 1e-5, + "z(0) = {}", + result.y_start[1] + ); + let yp = result.yp.expect("yp requested"); + assert!((yp[0] + 1.0).abs() < 1e-9, "dy/dt(0) = {}", yp[0]); + assert!((yp[1] + 0.5).abs() < 1e-6, "dz/dt(0) = {}", yp[1]); + } + + #[test] + // The same held derivative fills both columns, bit for bit. + #[allow(clippy::float_cmp)] + fn a_single_column_solve_stores_the_consistent_derivative_too() { + let prepared = PreparedSolver::new(build_nonlinear_dae(), 1e-6, &[1e-6, 1e-6]) + .expect("prepare") + .with_store_yp(true); + let set = InputSet::new(&[1.0, 3.0], &[1.0]); + let full = prepared + .solve(SolveRequest::new(&[0.0, 0.5, 1.0]), set) + .expect("solve"); + let single = prepared + .solve(SolveRequest::new(&[0.0]), set) + .expect("a one-point span must solve"); + assert_eq!(single.y, single.y_start); + assert_eq!( + single.yp.expect("yp requested"), + full.yp.expect("yp requested")[..2] + ); + } + + /// States `[u0, v0, u1, v1]` with `v` algebraic, coupled through a non-diagonal `g_v`. + fn build_coupled_dae() -> CompiledModel { + let mut arena = Arena::new(); + let [u0, v0, u1, v1] = [0, 1, 2, 3].map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let half = arena.alloc(Node::Scalar(0.5)); + let t = arena.alloc(Node::Time); + let du0 = arena.alloc(Node::Mul(neg, v1)); + let u1v0 = arena.alloc(Node::Mul(u1, v0)); + let du1 = arena.alloc(Node::Mul(neg, u1v0)); + let hv1 = arena.alloc(Node::Mul(half, v1)); + let g0 = arena.alloc(Node::Add(v0, hv1)); + let g0 = arena.alloc(Node::Sub(g0, u0)); + let sin_t = arena.alloc(Node::Sin(t)); + let g0 = arena.alloc(Node::Sub(g0, sin_t)); + let v0v1 = arena.alloc(Node::Mul(v0, v1)); + let g1 = arena.alloc(Node::Add(v1, v0v1)); + let g1 = arena.alloc(Node::Sub(g1, u1)); + let rhs = arena.alloc(Node::Concat(vec![du0, g0, du1, g1])); + let mass = CsrData { + indptr: vec![0, 1, 1, 2, 2], + indices: vec![0, 2], + data: vec![1.0, 1.0], + shape: Shape::matrix(4, 4), + }; + compile(&arena, rhs, mass, 4, 0) + } + + #[test] + // Bit-identical is the property under test. + #[allow(clippy::float_cmp)] + fn the_cached_algebraic_block_factors_as_a_fresh_triplet_build_does() { + let prepared = PreparedSolver::new(build_coupled_dae(), 1e-8, &[1e-8; 4]) + .expect("prepare") + .with_store_yp(true); + let block = prepared.algebraic_block.as_ref().expect("built for yp"); + let (compiled, csc) = (&prepared.compiled, prepared.compiled.csc_sparsity()); + let algebraic = prepared.algebraic_states(); + let local = |i: usize| { + block + .states + .iter() + .position(|&s| s == i) + .expect("algebraic") + }; + let mut ws = compiled.create_workspace(); + for (t, y) in [(0.3, [1.0, 0.2, 0.5, -0.3]), (2.0, [0.3, -1.1, 2.0, 0.4])] { + let mut values = vec![0.0; csc.nnz()]; + compiled.assemble_jacobian_csc_no_mass(&mut ws, t, &y, &[], &mut values); + // The per-hold build this replaced: triplets, a fresh pattern and a full `sp_lu`. + let mut triplets = Vec::new(); + for &col in &block.states { + let entries = csc.colptr()[col]..csc.colptr()[col + 1]; + for (&row, &value) in csc.rowind()[entries.clone()].iter().zip(&values[entries]) { + if algebraic[row] { + triplets.push(Triplet::new(local(row), local(col), value)); + } + } + } + let n = block.states.len(); + let g_v = SparseColMat::try_new_from_triplets(n, n, &triplets).expect("g_v"); + let rhs = Mat::from_fn(n, 1, |i, _| 1.5 - i as f64); + let mut expected = rhs.clone(); + g_v.sp_lu().expect("lu").solve_in_place(expected.as_mut()); + let mut got = rhs; + block + .solve_in_place(&values, got.as_mut(), Par::Seq) + .expect("lu"); + assert_eq!(got, expected, "t = {t}"); + } + } + + #[test] + fn a_time_dependent_constraint_takes_its_first_slope_from_the_time_derivative() { + // dy/dt = -a*y with 0 = z - t: only dF/dt gives dz/dt = 1 at t0. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dydt = arena.alloc(Node::Mul(neg_a, y)); + let t = arena.alloc(Node::Time); + let constraint = arena.alloc(Node::Sub(z, t)); + let rhs = arena.alloc(Node::Concat(vec![dydt, constraint])); + let mass = differential_then_algebraic_mass(); + let model = compile_with_sens(&arena, rhs, mass, 2, 1, &[ParamIdx::new(0)]); + let prepared = PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]) + .expect("prepare") + .with_store_yp(true) + .with_time_derivative(&arena, rhs); + let t_eval = [0.0, 0.5, 1.0]; + + for request in [ + SolveRequest::new(&t_eval), + SolveRequest::new(&t_eval).with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&[1.0, 0.0], &[1.0])) + .expect("solve"); + let yp = result.yp.expect("yp requested"); + assert!((yp[1] - 1.0).abs() < 1e-12, "dz/dt(0) = {}", yp[1]); + } + } + + #[test] + fn every_state_path_stores_the_derivative_of_every_column() { + // With z = y every slope is -a*y on both rows; a stop at every time carries none over. + let prepared = PreparedSolver::new(build_small_dae_with_param(), 1e-10, &[1e-10, 1e-10]) + .expect("prepare") + .with_store_yp(true); + let (y0, inputs, t_eval) = small_dae_setup(); + let plain = SolveRequest::new(&t_eval); + let stops = plain.with_stop_times(&t_eval); + for request in [ + plain, + stops, + plain.with_sensitivities(), + stops.with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&y0, &inputs)) + .expect("solve"); + let yp = result.yp.expect("yp requested"); + assert_eq!(yp.len(), result.y.len()); + for j in 0..result.t.len() { + let expected = -inputs[0] * result.y[j * result.n_rows]; + for row in 0..result.n_rows { + let got = yp[j * result.n_rows + row]; + assert!( + (got - expected).abs() <= 1e-6 * expected.abs(), + "yp[{row}, {j}] = {got} but -a*y = {expected}, {request:?}", + ); + } + } + } + } + + #[test] + fn y_start_is_the_consistent_state_whatever_the_payload() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, t_eval) = small_dae_setup(); + for request in [ + SolveRequest::new(&t_eval), + SolveRequest::new(&t_eval).with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&y0, &inputs)) + .expect("solve"); + assert!((result.y_start[0] - 1.0).abs() < 1e-12); + assert!( + (result.y_start[1] - 1.0).abs() < 1e-6, + "z(0) = {}", + result.y_start[1] + ); + } + } + + /// Like `build_small_dae_with_param` with a second input: `dy/dt = -a*y + b`. + fn build_small_dae_two_inputs(sens: &[ParamIdx]) -> CompiledModel { + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let b = input(&mut arena, "b", 1, 1, 1); + let ode = { + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let neg_a_y = arena.alloc(Node::Mul(neg_a, sv0)); + arena.alloc(Node::Add(neg_a_y, b)) + }; + let algebraic_residual = arena.alloc(Node::Sub(sv0, sv1)); + let rhs = arena.alloc(Node::Concat(vec![ode, algebraic_residual])); + compile_with_sens(&arena, rhs, differential_then_algebraic_mass(), 2, 2, sens) + } + + fn build_small_dae_prepared_with_sens() -> PreparedSolver { + let model = build_small_dae_with_param(); + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed") + } + + fn blocks(outcome: &SolveOutcome) -> &[Vec] { + outcome + .sensitivities + .as_deref() + .expect("the request asked for sensitivities") + } + + fn small_dae_setup() -> (Vec, Vec, Vec) { + let y0 = vec![1.0, 0.0]; // algebraic IC deliberately inconsistent + let inputs = vec![1.0]; + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + (y0, inputs, t_eval) + } + + /// One state, `dy/dt = -a * y`: no algebraic row, so `set_consistent` returns early. + fn build_small_ode_with_param() -> CompiledModel { + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let rhs = arena.alloc(Node::Mul(neg_a, sv0)); + compile(&arena, rhs, identity_mass_matrix(1), 1, 1) + } + + /// `M dy/dt = -y` on two states, with the given mass matrix. + fn build_ode_with_mass(mass: CsrData) -> CompiledModel { + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let f0 = arena.alloc(Node::Mul(neg, y0)); + let f1 = arena.alloc(Node::Mul(neg, y1)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1])); + compile(&arena, rhs, mass, 2, 1) + } + + #[test] + fn the_first_column_of_a_pure_ode_divides_out_a_diagonal_mass() { + // M = diag(2, 4), so dy/dt = (-y0 / 2, -y1 / 4), not f = -y. + let mass = CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 4.0], + shape: Shape::matrix(2, 2), + }; + let prepared = PreparedSolver::new(build_ode_with_mass(mass), 1e-8, &[1e-8, 1e-8]) + .expect("prepare") + .with_store_yp(true); + let t_eval = [0.0, 0.5, 1.0]; + for request in [ + SolveRequest::new(&t_eval), + SolveRequest::new(&t_eval).with_stop_times(&t_eval), + ] { + let result = prepared + .solve(request, InputSet::new(&[1.0, 2.0], &[0.0])) + .expect("solve"); + let yp = result.yp.expect("yp requested"); + assert_eq!(yp[..2], [-0.5, -0.5], "{request:?}"); + } + } + + #[test] + fn the_first_column_of_a_pure_ode_solves_against_a_non_diagonal_mass() { + // M = [[2, 1], [0, 1]] gives dy/dt = ((y1 - y0) / 2, -y1), far from f = -y. + let mass = CsrData { + indptr: vec![0, 2, 3], + indices: vec![0, 1, 1], + data: vec![2.0, 1.0, 1.0], + shape: Shape::matrix(2, 2), + }; + let prepared = PreparedSolver::new(build_ode_with_mass(mass), 1e-10, &[1e-10, 1e-10]) + .expect("prepare") + .with_store_yp(true); + let result = prepared + .solve( + SolveRequest::new(&[0.0, 0.5, 1.0]), + InputSet::new(&[1.0, 0.5], &[0.0]), + ) + .expect("solve"); + let yp = result.yp.expect("yp requested"); + assert!((yp[0] + 0.25).abs() < 1e-15, "dy0/dt(0) = {}", yp[0]); + assert!((yp[1] + 0.5).abs() < 1e-15, "dy1/dt(0) = {}", yp[1]); + } + + #[test] + fn a_breakpoint_restart_refreshes_dy_on_a_pure_ode() { + let prepared = + PreparedSolver::new(build_small_ode_with_param(), 1e-8, &[1e-8]).expect("prepare"); + let inputs = [2.0_f64]; + let y0 = [1.0_f64]; + + let set = InputSet::new(&y0, &inputs); + let problem = prepared + .build_problem(set, 0.0, &prepared.packed_binding(), false, Par::Seq) + .expect("build"); + let state = problem.bdf_state::().expect("state"); + let mut solver = problem.bdf_solver::(state).expect("solver"); + // Step away from t0 so `dy` holds a BDF slope rather than its initial seed. + solver.solve_dense(&[0.0, 0.5, 1.0]).expect("solve"); + + let stale = solver.state().dy[0]; + BreakpointRestarter::new(&prepared, &problem.ic_options) + .restart(&mut solver, 2.0) + .expect("restart"); + + let state = solver.state(); + let expected = -inputs[0] * state.y[0]; + assert!( + (stale - expected).abs() > 1e-14, + "fixture is vacuous: dy was already exact before the restart", + ); + assert!( + (state.dy[0] - expected).abs() <= 1e-13 * expected.abs(), + "dy = {} after the restart but f(t, y) = {expected}", + state.dy[0], + ); + } + + #[test] + // Every column is the untouched start state. + #[allow(clippy::float_cmp)] + fn a_span_that_never_leaves_t0_reports_every_column_at_the_start() { + let prepared = build_decay_param_output_event().with_store_yp(true); + for t_eval in [&[0.0][..], &[0.0, 0.0]] { + let base = SolveRequest::new(t_eval).with_stop_times(t_eval); + for request in [ + base, + base.with_outputs(), + base.with_sensitivities(), + base.with_outputs().with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&[1.0], &[1.0])) + .expect("a span of zero length must not be an error"); + assert_eq!(result.t, t_eval, "{request:?}"); + assert_eq!(result.termination, Termination::Completed); + let expected = if request.outputs { 2.0 } else { 1.0 }; + assert!(result.y.iter().all(|&y| y == expected), "{request:?}"); + if let Some(blocks) = &result.sensitivities { + assert!(blocks[0].iter().all(|&s| s == 0.0), "{request:?}"); + } + } + } + } + + #[test] + fn output_points_within_round_off_of_a_segment_end_are_all_reported() { + // One step spans each short span, so diffsol can stop an ULP short of tstop. + let prepared = build_decay_param_output_event(); + for i in 0..100 { + let t_end = 1e-6 * (0.2 + 0.0123 * f64::from(i)); + for k in 0..=4 { + let below = f64::from_bits(t_end.to_bits() - k); + let t_eval = [0.0, below, t_end]; + for request in [ + SolveRequest::new(&t_eval), + SolveRequest::new(&t_eval).with_sensitivities(), + ] { + let result = prepared + .solve(request, InputSet::new(&[1.0], &[1e-3])) + .expect("solve"); + assert_eq!(result.t, t_eval, "t_end = {t_end}, k = {k}"); + } + } + } + } + + #[test] + fn solve_empty_t_eval_returns_error_not_panic() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _t) = small_dae_setup(); + let err = prepared + .solve(SolveRequest::new(&[]), InputSet::new(&y0, &inputs)) + .expect_err("empty t_eval must be a returned error, not a panic"); + assert!( + matches!(err, CoreError::EmptyTimePoints), + "expected EmptyTimePoints, got: {err:?}" + ); + } + + #[test] + fn an_unsorted_or_non_finite_t_eval_is_an_error_not_a_wrong_trajectory() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _t) = small_dae_setup(); + let rejected = |t_eval: &[f64]| { + prepared + .solve(SolveRequest::new(t_eval), InputSet::new(&y0, &inputs)) + .expect_err("a bad t_eval must be an error, not a trajectory") + }; + for t_eval in [[0.0, 2.0, 1.0], [0.0, 0.5, 0.2]] { + let err = rejected(&t_eval); + assert!( + matches!(err, CoreError::UnsortedTimePoints { index: 2, .. }), + "{t_eval:?}: {err:?}" + ); + } + // -inf once reported every column at the start state; +inf integrated to y = 0. + for (t_eval, index) in [ + (vec![0.0, f64::NAN, 1.0], 1), + (vec![f64::NEG_INFINITY, 0.0], 0), + (vec![0.0, f64::INFINITY], 1), + (vec![f64::NAN], 0), + ] { + let err = rejected(&t_eval); + assert!( + matches!(err, CoreError::NonFiniteTimePoint { index: i, .. } if i == index), + "{t_eval:?}: {err:?}" + ); + } + } + + #[test] + fn output_solves_reject_a_model_without_output_variables() { + let prepared = build_small_dae_prepared_with_sens(); // no outputs + let (y0, inputs, t_eval) = small_dae_setup(); + let err = prepared + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&y0, &inputs), + ) + .expect_err("an outputs solve on a model with none must be an error"); + assert!( + matches!(err, CoreError::NoOutputVariables), + "expected NoOutputVariables, got: {err:?}" + ); + let err = prepared + .solve( + SolveRequest::new(&t_eval) + .with_outputs() + .with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect_err("an outputs sensitivity solve on a model with none must be an error"); + assert!( + matches!(err, CoreError::NoOutputVariables), + "expected NoOutputVariables, got: {err:?}" + ); + } + + #[test] + fn solve_repeated_t_eval_still_solves() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _t) = small_dae_setup(); + let t_eval = vec![0.0, 0.2, 0.2, 0.4]; + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .expect("a repeated time must still solve"); + assert_eq!(r.t.len(), t_eval.len()); + } + + #[test] + fn solve_y0_length_mismatch_returns_error_not_panic() { + let prepared = build_small_dae_prepared_with_sens(); // n_states = 2 + let (_y0, inputs, t_eval) = small_dae_setup(); + let err = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &inputs)) + .expect_err("y0 length mismatch must be a returned error, not a panic"); + assert!( + matches!( + err, + CoreError::Y0Length { + got: 1, + expected: 2 + } + ), + "expected Y0Length {{ got: 1, expected: 2 }}, got: {err:?}" + ); + } + + #[test] + fn solve_inputs_length_mismatch_is_a_clear_input_error() { + let prepared = build_small_dae_prepared_with_sens(); // needs 1 packed input + let (y0, _inputs, t_eval) = small_dae_setup(); + let err = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &[])) + .expect_err("too-short inputs must be an error"); + assert!( + matches!( + err, + CoreError::InputsLength { + got: 0, + expected: 1 + } + ), + "expected InputsLength {{ got: 0, expected: 1 }}, got: {err:?}" + ); + let msg = err.to_string(); + assert!( + msg.contains("inputs") && msg.contains("parameter"), + "expected a clear inputs-length error, got: {msg}" + ); + } + + #[test] + fn prepared_problem_atol_length_mismatch_is_error() { + let model = build_small_dae_with_param(); // n_states = 2 + let err = PreparedSolver::new(model, 1e-8, &[1e-8]) + .expect_err("atol length mismatch must be a returned error"); + assert!( + matches!( + err, + CoreError::AtolLength { + got: 1, + expected: 2 + } + ), + "expected AtolLength {{ got: 1, expected: 2 }}, got: {err:?}" + ); + } + + #[test] + fn tolerances_idas_would_reject_are_rejected_up_front() { + let good = [1e-8, 1e-8]; + let cases = [ + (f64::NAN, good, "rtol", None), + (f64::INFINITY, good, "rtol", None), + (-1e-8, good, "rtol", None), + (1e-8, [1e-8, f64::NAN], "atol", Some(1)), + (1e-8, [f64::NEG_INFINITY, 1e-8], "atol", Some(0)), + (1e-8, [1e-8, -1e-8], "atol", Some(1)), + (0.0, [1e-8, 0.0], "atol", Some(1)), + ]; + for (rtol, atol, name, index) in cases { + let err = PreparedSolver::new(build_small_dae_with_param(), rtol, &atol) + .expect_err("an invalid tolerance must be rejected"); + assert!( + matches!(err, CoreError::Tolerance { name: n, index: i, .. } if n == name && i == index), + "rtol={rtol}, atol={atol:?}: got {err:?}" + ); + } + let err = PreparedSolver::new(build_small_dae_with_param(), 0.0, &[1e-8, 0.0]) + .expect_err("rtol and atol both zero"); + assert_eq!( + err.to_string(), + "atol[1] = 0 is not a valid tolerance: rtol and atol must be finite and >= 0, and not both zero" + ); + } + + #[test] + fn a_zero_rtol_or_a_zero_atol_alone_is_a_valid_tolerance() { + for (rtol, atol) in [(0.0, [1e-8, 1e-8]), (1e-8, [0.0, 1e-8])] { + PreparedSolver::new(build_small_dae_with_param(), rtol, &atol) + .expect("IDAS accepts a zero on one side"); + } + } + + /// Check sensitivity block 0 against a central difference at the final time. + /// + /// `perturbed(h)` is the `(y0, inputs)` the parameter moved by `h` gives. + fn assert_matches_fd( + prepared: &PreparedSolver, + set: InputSet<'_>, + perturbed: impl Fn(f64) -> (Vec, Vec), + ) -> SolveOutcome { + let (_, _, t_eval) = small_dae_setup(); + let outcome = prepared + .solve(SolveRequest::new(&t_eval).with_sensitivities(), set) + .expect("sensitivity solve failed"); + let h = 1e-6; + let [plus, minus] = [h, -h].map(|step| { + let (y0, inputs) = perturbed(step); + prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .expect("perturbed solve failed") + }); + let last = (t_eval.len() - 1) * outcome.n_rows; + for row in 0..outcome.n_rows { + let fd = (plus.y[last + row] - minus.y[last + row]) / (2.0 * h); + let got = blocks(&outcome)[0][last + row]; + assert!((got - fd).abs() < 1e-4, "row {row}: got {got}, fd {fd}"); + } + outcome + } + + #[test] + fn a_sensitivity_request_matches_finite_difference_on_inputs() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _) = small_dae_setup(); + let outcome = assert_matches_fd(&prepared, InputSet::new(&y0, &inputs), |h| { + (y0.clone(), vec![inputs[0] + h]) + }); + assert!(!outcome.statistics.sens_error_control_relaxed); + } + + #[test] + fn a_subset_sens_solve_integrates_only_the_requested_column() { + // [b] is a non-prefix subset: it catches a truncating seed that [a] would hide. + let (y0, inputs) = ([1.0, 0.0], [1.0, 0.5]); + for param in [0, 1] { + let model = build_small_dae_two_inputs(&[ParamIdx::from_usize(param)]); + let prepared = PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("prepare"); + let outcome = assert_matches_fd(&prepared, InputSet::new(&y0, &inputs), |h| { + let mut moved = inputs.to_vec(); + moved[param] += h; + (y0.to_vec(), moved) + }); + assert_eq!(blocks(&outcome).len(), 1, "2 inputs, 1 requested"); + } + } + + #[test] + fn seeded_y0_sens_matches_finite_difference_through_the_initial_condition() { + // y0 itself is `a`, so d/da picks up the seed as well as the rhs term. + let prepared = build_small_dae_prepared_with_sens(); + let a = 1.0; + let (y0, inputs) = ([a, a], [a]); + let seeded = assert_matches_fd( + &prepared, + InputSet::new(&y0, &inputs).with_sens_seed(&[1.0, 1.0]), + |h| (vec![a + h, a + h], vec![a + h]), + ); + + let (_, _, t_eval) = small_dae_setup(); + let unseeded = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("unseeded sensitivity solve failed"); + let last = (t_eval.len() - 1) * seeded.n_rows; + assert!( + (blocks(&unseeded)[0][last] - blocks(&seeded)[0][last]).abs() > 1e-3, + "the seed has to move the answer, else the test proves nothing", + ); + } + + #[test] + fn a_sensitivity_request_without_sensitivity_params_is_an_error() { + // Reported as itself, not via the relaxed retry, and ahead of the seed and row space. + let prepared = PreparedSolver::new(decay(), 1e-8, &[1e-8]).expect("prepare"); + let request = SolveRequest::new(&[0.0, 0.1]).with_sensitivities(); + let set = InputSet::new(&[1.0], &[]); + for (request, set) in [ + (request, set), + (request, set.with_sens_seed(&[0.5])), + (request.with_outputs(), set), + ] { + let err = prepared + .solve(request, set) + .expect_err("a sensitivity request the model cannot supply must be an error"); + assert!( + matches!(err, CoreError::NoSensitivityParams), + "{request:?}, {set:?}: got {err:?}" + ); + } + } + + /// The shared decay fixture, prepared at tolerances tight enough for analytic checks. + fn build_decay_param_output_event() -> PreparedSolver { + PreparedSolver::new(decay_with_output_and_event(), 1e-10, &[1e-10]) + .expect("PreparedSolver failed") + } + + /// Two states, three outputs, one sensitivity parameter and an event: row spaces differ. + fn build_two_state_three_output_with_sens() -> PreparedSolver { + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let two = arena.alloc(Node::Scalar(2.0)); + let neg_two_a = arena.alloc(Node::Mul(neg_a, two)); + let rhs = { + let d0 = arena.alloc(Node::Mul(neg_a, y0)); + let d1 = arena.alloc(Node::Mul(neg_two_a, y1)); + arena.alloc(Node::Concat(vec![d0, d1])) + }; + let mass = identity_mass_matrix(2); + let mut model = compile_with_sens(&arena, rhs, mass, 2, 1, &[ParamIdx::new(0)]); + let out0 = arena.alloc(Node::Mul(two, y0)); + let out1 = arena.alloc(Node::Add(y0, y1)); + let out2 = arena.alloc(Node::Mul(two, y1)); + model + .set_observables(ObservableKind::Outputs, &mut arena, &[out0, out1, out2]) + .expect("valid outputs"); + let threshold = arena.alloc(Node::Scalar(0.4)); + let event = arena.alloc(Node::Sub(y0, threshold)); + model + .set_observables(ObservableKind::Events, &mut arena, &[event]) + .expect("a valid event"); + + PreparedSolver::new(model, 1e-10, &[1e-10, 1e-10]) + .expect("PreparedSolver failed") + .with_store_yp(true) + } + + #[test] + fn every_payload_combination_honours_one_layout_contract() { + let prepared = build_two_state_three_output_with_sens(); + let (n_states, n_outputs) = (2, 3); + // Runs past the event at y0 = 0.4, so every combination terminates on a root. + let t_eval: Vec = (0..=40).map(|i| f64::from(i) * 0.05).collect(); + let set = InputSet::new(&[1.0, 1.0], &[1.0]); + + for outputs in [false, true] { + for sensitivities in [false, true] { + let mut request = SolveRequest::new(&t_eval); + if outputs { + request = request.with_outputs(); + } + if sensitivities { + request = request.with_sensitivities(); + } + let label = format!("outputs={outputs}, sensitivities={sensitivities}"); + let outcome = prepared.solve(request, set).expect(&label); + + let expected_rows = if outputs { n_outputs } else { n_states }; + assert_eq!(outcome.n_rows, expected_rows, "{label}: row space"); + assert_eq!( + outcome.y.len(), + outcome.n_rows * outcome.t.len(), + "{label}: trajectory size" + ); + + assert_eq!(outcome.yp.is_some(), !outputs, "{label}: yp presence"); + if let Some(yp) = &outcome.yp { + assert_eq!(yp.len(), outcome.y.len(), "{label}: yp layout"); + } + + assert_eq!( + outcome.sensitivities.is_some(), + sensitivities, + "{label}: sensitivity presence" + ); + if let Some(blocks) = &outcome.sensitivities { + assert_eq!(blocks.len(), 1, "{label}: one block per parameter"); + for (i, block) in blocks.iter().enumerate() { + assert_eq!(block.len(), outcome.y.len(), "{label}: block {i} layout"); + } + } + + let Termination::Event { t: t_event, .. } = outcome.termination else { + panic!("{label}: the event should have fired"); + }; + assert!( + (t_event - outcome.t[outcome.t.len() - 1]).abs() < 1e-12, + "{label}: the trajectory should end at the root" + ); + let y_event = outcome.y_event.as_ref().expect("y_event missing"); + assert_eq!(y_event.len(), n_states, "{label}: y_event is a full state"); + if outputs { + // The root column holds the outputs of the very state y_event reports. + let root = &outcome.y[outcome.y.len() - n_outputs..]; + let (y0, y1) = (y_event[0], y_event[1]); + assert_eq!(root, [2.0 * y0, y0 + y1, 2.0 * y1], "{label}"); + } + } + } + } + + #[test] + // Bit-identical is the property under test. + #[allow(clippy::float_cmp)] + fn a_row_space_change_leaves_the_termination_fields_untouched() { + let prepared = build_two_state_three_output_with_sens(); + let t_eval: Vec = (0..=40).map(|i| f64::from(i) * 0.05).collect(); + let set = InputSet::new(&[1.0, 1.0], &[1.0]); + + let states = prepared + .solve(SolveRequest::new(&t_eval), set) + .expect("state solve failed"); + let outputs = prepared + .solve(SolveRequest::new(&t_eval).with_outputs(), set) + .expect("output solve failed"); + + assert_eq!(states.t, outputs.t); + assert_eq!(states.t.len(), outputs.t.len()); + assert_eq!(states.termination, outputs.termination); + assert_eq!(states.y_event, outputs.y_event); + assert_eq!( + states.statistics.solver.number_of_steps, + outputs.statistics.solver.number_of_steps + ); + } + + #[test] + // Bit-identical is the property under test: doubling is exact. + #[allow(clippy::float_cmp)] + fn an_outputs_request_keeps_every_column_across_batch_windows() { + // Two full windows and a partial one, ending before the event at t = ln 2.5. + let prepared = build_decay_param_output_event(); + let n_points = 2 * OUTPUT_BATCH_LANES + 3; + let t_eval: Vec = (0..n_points).map(|i| i as f64 * 1e-3).collect(); + let set = InputSet::new(&[1.0], &[1.0]); + + let states = prepared + .solve(SolveRequest::new(&t_eval), set) + .expect("state solve failed"); + let outputs = prepared + .solve(SolveRequest::new(&t_eval).with_outputs(), set) + .expect("output solve failed"); + + assert_eq!(outputs.termination, Termination::Completed); + assert_eq!(outputs.t, t_eval); + let doubled: Vec = states.y.iter().map(|y| 2.0 * y).collect(); + assert_eq!(outputs.y, doubled); + } + + #[test] + fn an_outputs_sensitivity_request_matches_analytic() { + // Stops before the event at t = ln 2.5, so this is the final-time path. + let prepared = build_decay_param_output_event(); + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.1).collect(); + let r = prepared + .solve( + SolveRequest::new(&t_eval) + .with_outputs() + .with_sensitivities(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("outputs sensitivity solve failed"); + + assert_eq!(r.termination, Termination::Completed); + assert_eq!(r.n_rows, 1); + assert_eq!(blocks(&r).len(), 1); + assert_eq!(blocks(&r)[0].len(), r.t.len()); + for (j, &t) in r.t.iter().enumerate() { + let out = r.y[j]; + let sens = blocks(&r)[0][j]; + let expected_out = 2.0 * (-t).exp(); + let expected_sens = -2.0 * t * (-t).exp(); + assert!( + (out - expected_out).abs() < 1e-6, + "t={t}: out={out}, want {expected_out}" + ); + assert!( + (sens - expected_sens).abs() < 1e-4, + "t={t}: d(out)/da={sens}, want {expected_sens}" + ); + } + let y_event = r.y_event.expect("terminal state missing"); + assert_eq!(y_event.len(), 1); + assert!( + (y_event[0] - (-0.5f64).exp()).abs() < 1e-6, + "terminal state {} diverges from exp(-0.5)", + y_event[0] + ); + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: columns keep their times + fn solve_event_reports_grid_times_then_the_root() { + let prepared = build_decay_param_output_event(); + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.1).collect(); + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[1.0])) + .expect("solve failed"); + + let Termination::Event { t: t_event, .. } = r.termination else { + panic!("expected event termination"); + }; + assert!((t_event - 2.5f64.ln()).abs() < 1e-6, "t_event = {t_event}"); + let drained = t_eval.iter().filter(|&&t| t <= t_event).count(); + assert_eq!( + r.t.len(), + drained + 1, + "one root column past the grid points" + ); + assert_eq!(r.n_rows, 1); + for (j, (&got, &want)) in r.t.iter().zip(&t_eval).take(drained).enumerate() { + assert_eq!(got, want, "column {j} moved off its requested time"); + } + assert_eq!(r.t[drained], t_event, "last column is not the root time"); + let y_root = r.y[drained]; + assert!((y_root - 0.4).abs() < 1e-6, "y at the root={y_root}"); + let y_event = r.y_event.expect("y_event missing"); + assert!((y_event[0] - 0.4).abs() < 1e-6, "y_event={}", y_event[0]); + } + + #[test] + fn the_termination_names_the_event_that_ended_the_solve() { + // dy/dt = -y with events y - 0.1 and y - 0.5: the second is crossed first. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs = arena.alloc(Node::Mul(neg, y)); + let mut model = compile(&arena, rhs, identity_mass_matrix(1), 1, 0); + let events = [0.1, 0.5].map(|threshold| { + let threshold = arena.alloc(Node::Scalar(threshold)); + arena.alloc(Node::Sub(y, threshold)) + }); + model + .set_observables(ObservableKind::Events, &mut arena, &events) + .expect("valid events"); + let prepared = PreparedSolver::new(model, 1e-10, &[1e-10]).expect("prepare"); + let t_eval: Vec = (0..=30).map(|i| f64::from(i) * 0.1).collect(); + + // From 0.05 both are violated at the start; the further below zero is reported. + for (y0, t_root) in [(1.0, 2.0f64.ln()), (0.3, 0.0), (0.05, 0.0)] { + let result = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[y0], &[])) + .expect("solve"); + let Termination::Event { t, index } = result.termination else { + panic!("y0 = {y0}: an event should have fired"); + }; + assert_eq!(index, 1, "y0 = {y0}"); + assert!((t - t_root).abs() < 1e-8, "y0 = {y0}: t = {t}"); + } + } + + #[test] + fn every_solve_path_reports_its_own_integration_time() { + let prepared = build_decay_param_output_event(); + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.1).collect(); + let base = SolveRequest::new(&t_eval); + for request in [ + base, + base.with_outputs(), + base.with_sensitivities(), + base.with_outputs().with_sensitivities(), + ] { + let statistics = prepared + .solve(request, InputSet::new(&[1.0], &[1.0])) + .expect("solve") + .statistics; + assert!( + statistics.integration_time_secs > 0.0, + "{request:?} reported no integration time", + ); + assert!( + statistics.integration_time_secs + >= statistics.ic_time_secs + statistics.solver_setup_time_secs, + "{request:?} reported less time than the phases it contains", + ); + } + } + + fn grid(n: usize) -> Vec { + assert!(n > 1, "a grid needs at least two columns"); + let last = (n - 1) as f64; + (0..n).map(|i| i as f64 / last).collect() + } + + #[test] + fn yp_is_stored_up_to_the_hermite_column_limit_on_every_state_path() { + let prepared = build_small_dae_prepared_with_sens().with_store_yp(true); + let (y0, inputs, _) = small_dae_setup(); + let (at, past) = (grid(MAX_HERMITE_COLUMNS), grid(MAX_HERMITE_COLUMNS + 1)); + for (request, stores_yp) in [ + (SolveRequest::new(&at), true), + (SolveRequest::new(&past), false), + (SolveRequest::new(&past).with_sensitivities(), false), + ] { + let result = prepared + .solve(request, InputSet::new(&y0, &inputs)) + .expect("solve"); + let columns = request.t_eval.len(); + assert_eq!(result.t.len(), columns, "every column is still reported"); + assert_eq!(result.yp.is_some(), stores_yp, "{columns} columns"); + } + } + + #[test] + fn yp_is_absent_unless_requested() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, t_eval) = small_dae_setup(); + let result = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .expect("solve"); + assert!(result.yp.is_none()); + } + + #[test] + fn the_root_column_carries_the_wound_back_derivative() { + // yp must be the root-time slope after state_mut_back, not the overshot step's. + let prepared = build_decay_param_output_event().with_store_yp(true); + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.1).collect(); + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[1.0])) + .expect("solve failed"); + assert!(matches!(r.termination, Termination::Event { .. })); + let yp = r.yp.expect("yp requested"); + let root = r.t.len() - 1; + assert!( + (yp[root] - (-0.4)).abs() < 1e-6, + "yp at the y = 0.4 root is {}, want -0.4", + yp[root], + ); + } + + #[test] + fn an_outputs_sensitivity_request_ends_at_the_event_root() { + // Runs past the event at y = 0.4 (t = ln 2.5), where dy/da = -t y. + let prepared = build_decay_param_output_event(); + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.1).collect(); + let r = prepared + .solve( + SolveRequest::new(&t_eval) + .with_outputs() + .with_sensitivities(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("outputs sensitivity solve failed"); + + let Termination::Event { t: t_event, .. } = r.termination else { + panic!("expected event termination"); + }; + let ln2p5 = 2.5f64.ln(); + assert!( + (t_event - ln2p5).abs() < 1e-6, + "t_event={t_event}, want {ln2p5}" + ); + assert_eq!(blocks(&r)[0].len(), r.t.len()); + let y_event = r.y_event.as_deref().expect("y_event missing"); + assert!( + (y_event[0] - 0.4).abs() < 1e-6, + "y_event={}, expected state 0.4", + y_event[0] + ); + let out_last = r.y[r.t.len() - 1]; + assert!((out_last - 0.8).abs() < 1e-6, "out at root={out_last}"); + let sens_last = blocks(&r)[0][r.t.len() - 1]; + let expected = -2.0 * t_event * 0.4; + assert!( + (sens_last - expected).abs() < 1e-4, + "d(out)/da at root={sens_last}, want {expected}" + ); + + let start = r.y_start_sens.as_deref().expect("y_start_sens missing"); + let end = r.y_event_sens.as_deref().expect("y_event_sens missing"); + assert_eq!((start.len(), start[0].len()), (1, 1)); + assert_eq!((end.len(), end[0].len()), (1, 1)); + assert!(start[0][0].abs() < 1e-12, "dy/da at start={}", start[0][0]); + let expected = -t_event * 0.4; + assert!( + (end[0][0] - expected).abs() < 1e-4, + "dy/da at root={}, want {expected}", + end[0][0] + ); + } + + #[test] + fn end_state_sensitivities_follow_y_event() { + // Ends before the event: only the outputs request reports a (terminal) y_event. + let prepared = build_decay_param_output_event(); + let t_eval = [0.0, 0.25, 0.5]; + for outputs in [false, true] { + let request = SolveRequest::new(&t_eval).with_sensitivities(); + let request = if outputs { + request.with_outputs() + } else { + request + }; + let r = prepared + .solve(request, InputSet::new(&[1.0], &[1.0])) + .expect("sensitivity solve failed"); + assert!(r.y_start_sens.is_some()); + assert_eq!(r.y_event.is_some(), outputs); + assert_eq!(r.y_event_sens.is_some(), outputs); + if let Some(end) = r.y_event_sens.as_deref() { + let expected = -0.5 * (-0.5f64).exp(); + assert!( + (end[0][0] - expected).abs() < 1e-6, + "dy/da at t=0.5 is {}, want {expected}", + end[0][0] + ); + } + } + + let r = prepared + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("outputs solve failed"); + assert!(r.y_start_sens.is_none() && r.y_event_sens.is_none()); + } + + /// The equations over the two-input fixture with the given sensitivity subset. + fn two_input_rhs_at_y0(sens: &[ParamIdx]) -> Equations { + let model = build_small_dae_two_inputs(sens); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + assert_eq!(prepared.compiled.n_sens_params(), sens.len()); + prepared.build_eqn( + InputSet::new(&[1.0, 0.0], &[1.0, 0.5]), + &prepared.sens_binding, + Par::Seq, + ) + } + + fn dense_columns(m: &FaerSparseMat, n: usize, ncols: usize) -> Vec { + let mut got = vec![0.0; n * ncols]; + let (indices, values) = m.triplet_iter(); + for ((i, j), val) in indices.zip(values) { + got[j * n + i] = val; + } + got + } + + #[test] + fn rhs_sens_inplace_assembles_the_requested_subset_in_order() { + // Reordered subset [b, a]: column order follows the subset, not the input vector. + let eqn = two_input_rhs_at_y0(&[ParamIdx::new(1), ParamIdx::new(0)]); + let rhs = eqn.rhs(); + let ctx = *rhs.context(); + let mut x = FaerVec::::zeros(2, ctx); + x.set_index(0, 1.0); + let (n, np) = (rhs.nout(), rhs.nparams()); + assert_eq!(np, 2); + + let mut m = FaerSparseMat::::new_from_sparsity(n, np, rhs.sens_sparsity(), ctx); + rhs.sens_inplace(&x, 0.0, &mut m); + + // f = [-a*y + b, y - z], so df/db = [1, 0] and df/da = [-y, 0] = [-1, 0]. + assert_eq!(dense_columns(&m, n, np), vec![1.0, 0.0, -1.0, 0.0]); + } + + #[test] + fn rhs_sens_inplace_matches_per_column_sens_mul_on_a_subset() { + let eqn = two_input_rhs_at_y0(&[ParamIdx::new(1), ParamIdx::new(0)]); + let rhs = eqn.rhs(); + let ctx = *rhs.context(); + let mut x = FaerVec::::zeros(2, ctx); + x.set_index(0, 1.0); + let (n, np) = (rhs.nout(), rhs.nparams()); + + let mut expected = vec![0.0; n * np]; + let mut v = FaerVec::::zeros(np, ctx); + let mut col = FaerVec::::zeros(n, ctx); + for j in 0..np { + v.set_index(j, 1.0); + rhs.sens_mul_inplace(&x, 0.0, &v, &mut col); + expected[j * n..(j + 1) * n].copy_from_slice(lane(&col)); + v.set_index(j, 0.0); + } + + let mut m = FaerSparseMat::::new_from_sparsity(n, np, rhs.sens_sparsity(), ctx); + rhs.sens_inplace(&x, 0.0, &mut m); + + assert_eq!(dense_columns(&m, n, np), expected); + } + + #[test] + fn sens_inplace_runs_one_primal_pass_regardless_of_k() { + let one = [ParamIdx::new(0)]; + let two = [ParamIdx::new(1), ParamIdx::new(0)]; + for sens in [&one[..], &two[..]] { + let eqn = two_input_rhs_at_y0(sens); + let rhs = eqn.rhs(); + let ctx = *rhs.context(); + let mut x = FaerVec::::zeros(2, ctx); + x.set_index(0, 1.0); + + eqn.evaluator.workspace.borrow_mut().sens_primal_passes = 0; + let (n, np) = (rhs.nout(), rhs.nparams()); + assert_eq!(np, sens.len(), "subset width should equal k"); + let mut m = FaerSparseMat::::new_from_sparsity(n, np, rhs.sens_sparsity(), ctx); + rhs.sens_inplace(&x, 0.0, &mut m); + + assert_eq!( + eqn.evaluator.workspace.borrow().sens_primal_passes, + 1, + "k = {np}" + ); + } + } + + #[test] + fn a_wrongly_sized_y0_sens_is_an_error_not_a_panic() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, t_eval) = small_dae_setup(); + let err = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs).with_sens_seed(&[1.0]), + ) + .expect_err("a 1-entry seed for a 2-state model must be rejected"); + assert!(matches!( + err, + CoreError::Y0SensLength { + got: 1, + expected: 2 + } + )); + } + + #[test] + fn sens_atol_tightens_only_differential_states() { + let model = build_small_dae_two_inputs(&[ParamIdx::new(0)]); + let prepared = PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]) + .expect("PreparedSolver failed") + .with_sens_atol_factor(1e-3) + .expect("factor rejected"); + let sens_atol = prepared.sens_atol(); + assert!( + (sens_atol[0] / 1e-11 - 1.0).abs() < 1e-12, + "differential row was not tightened: {}", + sens_atol[0] + ); + assert!( + (sens_atol[1] / 1e-8 - 1.0).abs() < 1e-12, + "algebraic row must keep the state atol: {}", + sens_atol[1] + ); + } + + #[test] + fn with_sens_atol_factor_rejects_non_positive_and_non_finite() { + for bad in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let model = build_small_dae_two_inputs(&[ParamIdx::new(0)]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + assert!( + prepared.with_sens_atol_factor(bad).is_err(), + "factor {bad} should be rejected" + ); + } + } + + #[test] + fn a_sensitivity_solve_reports_an_initial_condition_failure_as_itself() { + // `0 = z^2 + 1` has no real root, whatever the sensitivity tolerances are. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dydt = arena.alloc(Node::Mul(neg_a, y)); + let one = arena.alloc(Node::Scalar(1.0)); + let z2 = arena.alloc(Node::Mul(z, z)); + let constraint = arena.alloc(Node::Add(z2, one)); + let rhs = arena.alloc(Node::Concat(vec![dydt, constraint])); + let mass = differential_then_algebraic_mass(); + let model = compile_with_sens(&arena, rhs, mass, 2, 1, &[ParamIdx::new(0)]); + let prepared = PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("prepare"); + let err = prepared + .solve( + SolveRequest::new(&[0.0, 1.0]).with_sensitivities(), + InputSet::new(&[1.0, 0.0], &[1.0]), + ) + .expect_err("an inconsistent IC must fail"); + assert!( + matches!(err, CoreError::Diffsol(_)), + "expected the IC failure itself, got: {err:?}" + ); + } + + /// `dy/dt = -y + k*p*(t > 0.5)` with sensitivities wrt `p`. + /// + /// At `p = 0` the state never sees the switch, but `ds/dt` jumps by `k` at `t = 0.5`. + fn build_decay_with_a_switched_sensitivity(k: f64) -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let p = input(&mut arena, "p", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let decay = arena.alloc(Node::Mul(neg, y)); + let t = arena.alloc(Node::Time); + let corner = arena.alloc(Node::Scalar(0.5)); + let after = arena.alloc(Node::NotEqualHeaviside(corner, t)); + let k = arena.alloc(Node::Scalar(k)); + let kp = arena.alloc(Node::Mul(k, p)); + let forcing = arena.alloc(Node::Mul(kp, after)); + let rhs = arena.alloc(Node::Add(decay, forcing)); + let mass = identity_mass_matrix(1); + compile_with_sens(&arena, rhs, mass, 1, 1, &[ParamIdx::new(0)]) + } + + /// A sensitivity solve built without sensitivity error control, as the retry ends up. + fn relaxed_from_scratch( + prepared: &PreparedSolver, + t_eval: &[f64], + set: InputSet<'_>, + ) -> SolveOutcome { + let times = TimePlan::new(t_eval, &[]).expect("sorted grid"); + let binding = &prepared.sens_binding; + let problem = prepared + .build_problem(set, times.t0(), binding, false, Par::Seq) + .expect("problem"); + let state = problem.bdf_state_sens::().expect("IC"); + let run = Run { + times, + inputs: set.inputs, + with_output: false, + started: Instant::now(), + lu_par: Par::Seq, + }; + prepared + .integrate_sensitivities(&problem, state, run, SolverStatistics::default()) + .expect("relaxed solve") + } + + #[test] + // Bit-identical is the property under test. + #[allow(clippy::float_cmp)] + fn a_sensitivity_too_sharp_for_error_control_is_retried_without_it() { + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let set = InputSet::new(&[1.0], &[0.0]); + let request = SolveRequest::new(&t_eval).with_sensitivities(); + let prepare = |k| { + PreparedSolver::new(build_decay_with_a_switched_sensitivity(k), 1e-8, &[1e-8]) + .expect("prepare") + }; + let gentle = prepare(1.0).solve(request, set).expect("solve"); + assert!( + !gentle.statistics.sens_error_control_relaxed, + "a small jump must solve under error control, else the fixture proves nothing" + ); + + // Resolving a jump of 1e6 in ds/dt takes a step below diffsol's minimum. + let k = 1e6; + let prepared = prepare(k); + let outcome = prepared + .solve(request, set) + .expect("the relaxed retry must succeed"); + assert!(outcome.statistics.sens_error_control_relaxed); + for (&t, &y) in t_eval.iter().zip(&outcome.y) { + assert!((y - (-t).exp()).abs() < 1e-6, "y({t}) = {y}"); + } + // Out of error control, the sensitivity is still integrated, not dropped. + let s_end = blocks(&outcome)[0][t_eval.len() - 1]; + let exact = k * (1.0 - (-0.5f64).exp()); + assert!( + (s_end - exact).abs() < 0.1 * exact, + "s(1) = {s_end}, exact {exact}" + ); + + // The retry starts from the failed attempt's IC, yet is a relaxed solve from scratch. + let relaxed = relaxed_from_scratch(&prepared, &t_eval, set); + assert_eq!(outcome.t, relaxed.t); + assert_eq!(outcome.y, relaxed.y); + assert_eq!(outcome.sensitivities, relaxed.sensitivities); + assert_eq!(outcome.y_start_sens, relaxed.y_start_sens); + assert_eq!( + outcome.statistics.solver.number_of_steps, + relaxed.statistics.solver.number_of_steps + ); + } + + #[test] + fn a_retry_that_fails_too_reports_both_failures() { + // dy/dt = a*y^2 from y = 1 blows up at t = 1/a, whatever the sensitivity tolerances. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let a = input(&mut arena, "a", 0, 0, 1); + let y_squared = arena.alloc(Node::Mul(y, y)); + let rhs = arena.alloc(Node::Mul(a, y_squared)); + let mass = identity_mass_matrix(1); + let model = compile_with_sens(&arena, rhs, mass, 1, 1, &[ParamIdx::new(0)]); + let prepared = PreparedSolver::new(model, 1e-8, &[1e-8]).expect("prepare"); + let err = prepared + .solve( + SolveRequest::new(&[0.0, 2.0]).with_sensitivities(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect_err("a blow-up must fail with and without sensitivity error control"); + + let CoreError::SensRetryFailed { + controlled, + relaxed, + } = &err + else { + panic!("expected both attempts' failures, got: {err:?}"); + }; + let message = err.to_string(); + for attempt in [controlled, relaxed] { + assert!( + !attempt.is_empty() && message.contains(attempt.as_str()), + "{message}" + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/test_fixtures.rs b/packages/pybamm-rust/pybamm-core/src/solver/test_fixtures.rs new file mode 100644 index 0000000000..a51d8b4279 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/test_fixtures.rs @@ -0,0 +1,52 @@ +//! Models shared by the solver's test modules. + +use crate::arena::Arena; +use crate::model::CompiledModel; +use crate::model::test_fixtures::{compile, compile_with_sens, identity_mass_matrix, input}; +use crate::node::{CsrData, Node, Shape}; +use crate::observable::ObservableKind; +use crate::params::ParamIdx; + +/// `diag(1, 0)`: a differential state, then an algebraic one. +pub fn differential_then_algebraic_mass() -> CsrData { + CsrData { + indptr: vec![0, 1, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(2, 2), + } +} + +/// `dy/dt = -y` on one state, with no inputs and so no sensitivity parameters. +pub fn decay() -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs = arena.alloc(Node::Mul(neg, y)); + compile(&arena, rhs, identity_mass_matrix(1), 1, 0) +} + +/// `dy/dt = -a*y`, output `2*y`, event `y - 0.4`, sensitivities wrt `a`. +/// +/// From `y = 1` with `a = 1` the event fires at `t = ln 2.5`. +pub fn decay_with_output_and_event() -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let a = input(&mut arena, "a", 0, 0, 1); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let rhs = arena.alloc(Node::Mul(neg_a, y)); + let two = arena.alloc(Node::Scalar(2.0)); + let output = arena.alloc(Node::Mul(two, y)); + let threshold = arena.alloc(Node::Scalar(0.4)); + let event = arena.alloc(Node::Sub(y, threshold)); + let mass = identity_mass_matrix(1); + let mut model = compile_with_sens(&arena, rhs, mass, 1, 1, &[ParamIdx::new(0)]); + model + .set_observables(ObservableKind::Outputs, &mut arena, &[output]) + .expect("a valid output"); + model + .set_observables(ObservableKind::Events, &mut arena, &[event]) + .expect("a valid event"); + model +} diff --git a/packages/pybamm-rust/pybamm-core/src/tangent.rs b/packages/pybamm-rust/pybamm-core/src/tangent.rs index 3ed610d269..5bc2f75a42 100644 --- a/packages/pybamm-rust/pybamm-core/src/tangent.rs +++ b/packages/pybamm-rust/pybamm-core/src/tangent.rs @@ -15,16 +15,28 @@ pub enum DiffTarget { Params, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum DiffMode { + Tangent(DiffTarget), + Time, +} + /// Build the tangent expression for `(d expr/dy) @ dy`. #[must_use] pub fn tangent_wrt_states(arena: &mut Arena, expr: NodeId) -> NodeId { - differentiate_root(arena, expr, DiffTarget::States, None) + differentiate_root(arena, expr, DiffMode::Tangent(DiffTarget::States), None) } /// Build the tangent expression for `(d expr/dp) @ dp`. #[must_use] pub fn tangent_wrt_params(arena: &mut Arena, expr: NodeId) -> NodeId { - differentiate_root(arena, expr, DiffTarget::Params, None) + differentiate_root(arena, expr, DiffMode::Tangent(DiffTarget::Params), None) +} + +/// Build the partial derivative `d expr/dt`, holding states and inputs fixed. +#[must_use] +pub fn derivative_wrt_time(arena: &mut Arena, expr: NodeId) -> NodeId { + differentiate_root(arena, expr, DiffMode::Time, None) } /// [`tangent_wrt_states`] seeding only `StateVector` nodes that overlap `active_indices`. @@ -41,14 +53,19 @@ pub(crate) fn tangent_wrt_subset( active_indices.is_sorted(), "active state indices must ascend" ); - differentiate_root(arena, expr, DiffTarget::States, Some(active_indices)) + differentiate_root( + arena, + expr, + DiffMode::Tangent(DiffTarget::States), + Some(active_indices), + ) } /// Differentiate `expr`, seeding only states in `state_filter` when one is given. fn differentiate_root( arena: &mut Arena, expr: NodeId, - mode: DiffTarget, + mode: DiffMode, state_filter: Option<&[usize]>, ) -> NodeId { let mut pending = arena.topological_order(expr); @@ -83,7 +100,7 @@ fn differentiate_root( /// A rule missing a child's derivative queues the child and is retried after it, so nodes /// are allocated in exactly the order a depth-first recursion would allocate them. struct Differentiator<'a> { - mode: DiffTarget, + mode: DiffMode, state_filter: Option<&'a [usize]>, /// Primal output widths. widths: Vec, @@ -200,8 +217,10 @@ impl Differentiator<'_> { let width = self.widths[id.index()]; let result = match *arena.get(id) { + Node::Time if mode == DiffMode::Time => arena.alloc(Node::Scalar(1.0)), + Node::StateVector { start, end } - if mode == DiffTarget::States + if mode == DiffMode::Tangent(DiffTarget::States) && state_filter.is_none_or(|filter| overlaps(filter, start, end)) => { arena.alloc(Node::TangentStateVector { start, end }) @@ -211,7 +230,7 @@ impl Differentiator<'_> { index, width: param_width, .. - } if mode == DiffTarget::Params => { + } if mode == DiffMode::Tangent(DiffTarget::Params) => { // Replicated to the packed width so an `Index` slice stays inside the buffer slot. let tangent = arena.alloc(Node::TangentParameter { index }); if param_width == 1 { @@ -543,10 +562,10 @@ impl Differentiator<'_> { mod tests { use super::*; use crate::adjoint::AdjointTape; - use crate::arena::DagRef; + use crate::arena::{Dag, DagRef}; use crate::eval::{CompiledExpr, TangentInputs}; use crate::node::InterpolantData; - use crate::simplify::{cse, dce, simplify}; + use crate::simplify::{cse, dce, simplify, simplify_pipeline}; use crate::zero_propagate::zero_propagate; /// Run the tangent pipeline end to end and return the compiled tape. @@ -1133,6 +1152,80 @@ mod tests { assert_eq!(result, [0.0]); } + #[test] + fn time_partial_holds_states_and_parameters_fixed_and_preserves_widths() { + let mut arena = Arena::new(); + let states = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let state = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let parameter = arena.alloc(Node::InputParameter { + name: "scale".into(), + index: 0, + offset: 0, + width: 1, + }); + let time = arena.alloc(Node::Time); + let square = arena.alloc(Node::Mul(time, time)); + let scaled = arena.alloc(Node::Mul(parameter, square)); + let product = arena.alloc(Node::Mul(state, scaled)); + let root = arena.alloc(Node::Concat(vec![states, product])); + let derivative = derivative_wrt_time(&mut arena, root); + let derivative = simplify_pipeline(Dag::new(arena, derivative)); + let compiled = CompiledExpr::new(derivative.as_ref()); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut scratch, 3.0, &[4.0, 5.0, 6.0], &[], &[2.0]); + + assert_eq!(result.len(), 4); + for (&actual, expected) in result.iter().zip([0.0, 0.0, 0.0, 48.0]) { + assert!((actual - expected).abs() < 1e-12); + } + } + + #[test] + fn time_partial_uses_the_active_conditional_branch() { + let mut arena = Arena::new(); + let time = arena.alloc(Node::Time); + let corner = arena.alloc(Node::Scalar(1.0)); + let after = arena.alloc(Node::NotEqualHeaviside(corner, time)); + let selector = arena.alloc(Node::Add(corner, after)); + let square = arena.alloc(Node::Mul(time, time)); + let cube = arena.alloc(Node::Mul(square, time)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![square, cube], + }); + let derivative = derivative_wrt_time(&mut arena, root); + let derivative = simplify_pipeline(Dag::new(arena, derivative)); + let compiled = CompiledExpr::new(derivative.as_ref()); + let mut scratch = vec![0.0; compiled.scratch_len()]; + + for (time, expected) in [(0.5, 1.0), (1.0, 2.0), (2.0, 12.0)] { + let result = compiled.eval(&mut scratch, time, &[], &[], &[]); + assert!((result[0] - expected).abs() < 1e-12); + } + } + + #[test] + fn time_partial_uses_the_interpolant_slope_on_each_side_of_a_knot() { + let mut arena = Arena::new(); + let time = arena.alloc(Node::Time); + let root = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 30.0]) + .expect("valid interpolant"), + ), + child: time, + }); + let derivative = derivative_wrt_time(&mut arena, root); + let derivative = simplify_pipeline(Dag::new(arena, derivative)); + let compiled = CompiledExpr::new(derivative.as_ref()); + let mut scratch = vec![0.0; compiled.scratch_len()]; + + for (time, expected) in [(0.5, 10.0), (1.5, 20.0)] { + let result = compiled.eval(&mut scratch, time, &[], &[], &[]); + assert!((result[0] - expected).abs() < 1e-12); + } + } + #[test] #[allow(clippy::float_cmp)] fn modulo_tangent_keeps_a_scalar_dividends_broadcast_width() { diff --git a/packages/pybamm-rust/pybamm-core/tests/test_faer_global_parallelism.rs b/packages/pybamm-rust/pybamm-core/tests/test_faer_global_parallelism.rs new file mode 100644 index 0000000000..2a73cf5009 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_faer_global_parallelism.rs @@ -0,0 +1,141 @@ +//! Solving with faer's process-global parallelism disabled, as a host embedding faer may. +//! +//! A binary of its own: the setting is process-wide, and once it is disabled any faer call +//! that still reads it panics. + +use std::sync::Arc; + +use faer::Par; +use pybamm_core::solver::solve::{InputSet, PreparedSolver, SolveOutcome, SolveRequest}; +use pybamm_core::{Arena, CompiledModel, CompiledModelOptions, CsrData, Node, PackedWidth, Shape}; + +/// Disables faer's global parallelism until dropped, then restores what it was. +struct GlobalParallelismDisabled(Par); + +impl GlobalParallelismDisabled { + fn new() -> Self { + let previous = faer::get_global_parallelism(); + faer::disable_global_parallelism(); + Self(previous) + } +} + +impl Drop for GlobalParallelismDisabled { + fn drop(&mut self) { + faer::set_global_parallelism(self.0); + } +} + +fn compile(arena: &Arena, rhs: pybamm_core::NodeId, mass: CsrData, inputs: usize) -> CompiledModel { + let n_states = mass.shape().rows; + let options = CompiledModelOptions::default(); + CompiledModel::new( + arena, + rhs, + mass, + n_states, + PackedWidth::new(inputs), + options, + ) + .expect("a valid fixture") +} + +/// `dy/dt = -a*y` with `0 = 2y - z - z^3`: each held derivative factors the algebraic block. +fn nonlinear_dae() -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let z = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let dydt = arena.alloc(Node::Mul(neg_a, y)); + let two = arena.alloc(Node::Scalar(2.0)); + let two_y = arena.alloc(Node::Mul(two, y)); + let z2 = arena.alloc(Node::Mul(z, z)); + let z3 = arena.alloc(Node::Mul(z2, z)); + let z_plus_z3 = arena.alloc(Node::Add(z, z3)); + let constraint = arena.alloc(Node::Sub(two_y, z_plus_z3)); + let rhs = arena.alloc(Node::Concat(vec![dydt, constraint])); + let mass = CsrData::try_new(vec![0, 1, 1], vec![0], vec![1.0], Shape::matrix(2, 2)) + .expect("diag(1, 0)"); + compile(&arena, rhs, mass, 1) +} + +/// `M dy/dt = -a*y` with `M = [[2, 1], [0, 1]]`: each held derivative solves against `M`. +fn ode_with_a_non_diagonal_mass() -> CompiledModel { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_a = arena.alloc(Node::Mul(neg, a)); + let rhs = arena.alloc(Node::Mul(neg_a, y)); + let mass = CsrData::try_new( + vec![0, 2, 3], + vec![0, 1, 1], + vec![2.0, 1.0, 1.0], + Shape::matrix(2, 2), + ) + .expect("an upper-triangular M"); + compile(&arena, rhs, mass, 1) +} + +// Bit-identical is the property under test. +#[allow(clippy::float_cmp)] +fn assert_same_trajectory(got: &SolveOutcome, want: &SolveOutcome, context: &str) { + assert_eq!(got.t, want.t, "{context}: times differ"); + assert_eq!(got.y, want.y, "{context}: trajectories differ"); + assert_eq!(got.yp, want.yp, "{context}: yp differs"); +} + +#[test] +fn every_lu_honours_the_requested_parallelism_not_the_global_one() { + let t_eval: Vec = (0..=8).map(|i| f64::from(i) * 0.25).collect(); + // A stop at every time holds a consistent derivative at each one. + let request = SolveRequest::new(&t_eval).with_stop_times(&t_eval); + let rates: Vec<[f64; 1]> = (1..=4).map(|i| [0.5 * f64::from(i)]).collect(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .expect("pool build failed"); + + for (label, model, y0) in [ + ("DAE", nonlinear_dae(), [1.0, 1.0]), + ("non-diagonal M", ode_with_a_non_diagonal_mass(), [1.0, 0.5]), + ] { + let model = Arc::new(model); + let prepare = || { + PreparedSolver::new(Arc::clone(&model), 1e-8, &[1e-8, 1e-8]) + .expect("prepare") + .with_store_yp(true) + }; + let sets: Vec> = rates.iter().map(|a| InputSet::new(&y0, a)).collect(); + let serial: Vec = { + let prepared = prepare(); + sets.iter() + .map(|set| prepared.solve(request, *set).expect("serial solve")) + .collect() + }; + + let _disabled = GlobalParallelismDisabled::new(); + let prepared = prepare(); + let batch = pool.install(|| prepared.solve_batch(request, &sets)); + for (i, (got, want)) in batch.into_iter().zip(&serial).enumerate() { + let got = got.expect("batched solve"); + assert_same_trajectory(&got, want, &format!("{label}, set {i} of the batch")); + } + let threaded = pool + .install(|| prepared.solve(request.with_lu_threads(4), sets[0])) + .expect("threaded solve"); + assert_same_trajectory(&threaded, &serial[0], &format!("{label}, 4 LU threads")); + } +}