diff --git a/.github/workflows/_nox.yml b/.github/workflows/_nox.yml index 04f8dac8d0..ddc509f989 100644 --- a/.github/workflows/_nox.yml +++ b/.github/workflows/_nox.yml @@ -119,6 +119,11 @@ jobs: with: python-version: ${{ matrix.leg.python }} + # Every session builds the Rust crate from source; cache the cargo build. + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: version: "latest" diff --git a/.github/workflows/publish_pypi.yml b/.github/workflows/publish_pypi.yml index d82d61eced..f1ebe1fa7f 100644 --- a/.github/workflows/publish_pypi.yml +++ b/.github/workflows/publish_pypi.yml @@ -7,10 +7,79 @@ on: permissions: {} jobs: - build: + build_wheels: # Monorepo: this workflow publishes ONLY pybamm. Releases are discriminated # by tag namespace now that pybamm and pybammsolvers share a repository # (the old `github.repository` guard no longer distinguishes them). + if: startsWith(github.event.release.tag_name, 'pybamm-v') + name: Wheels (${{ matrix.os }} ${{ matrix.arch }}) + runs-on: ${{ matrix.os }} + permissions: + contents: read + strategy: + fail-fast: false + matrix: + include: + - { os: ubuntu-latest, arch: x86_64, artifact: wheels_manylinux, crossversion: true } + - { os: ubuntu-24.04-arm, arch: aarch64, artifact: wheels_manylinux_aarch64 } + - { os: macos-latest, arch: arm64, artifact: wheels_macos_arm64, deployment_target: "11.0" } + - { os: macos-15-intel, arch: x86_64, artifact: wheels_macos_x86_64, deployment_target: "10.13" } + - { os: windows-2025, arch: AMD64, artifact: wheels_windows } + + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + # Full history + tags so hatch-vcs can resolve the pybamm-v* version. + fetch-depth: 0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: 3.13 + + # Deliberately no cargo cache here: a cache entry poisoned from a PR could + # be restored into a wheel that gets published. Release builds start clean. + - name: Build wheels + # Pinned: cibuildwheel's abi3audit default is load-bearing here. + run: pipx run cibuildwheel==4.1.1 packages/pybamm --output-dir wheelhouse + env: + CIBW_BUILD_VERBOSITY: 1 + # One abi3 wheel per platform serves CPython 3.10-3.14. + CIBW_BUILD: "cp310-*" + CIBW_SKIP: "pp* *musllinux* *t-*" + CIBW_ARCHS: ${{ matrix.arch }} + MACOSX_DEPLOYMENT_TARGET: ${{ matrix.deployment_target }} + # manylinux containers have no rustup; the crate graph needs cargo >= 1.89. + CIBW_BEFORE_ALL_LINUX: > + curl -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal + CIBW_ENVIRONMENT_LINUX: 'PATH=$HOME/.cargo/bin:$PATH' + CIBW_TEST_COMMAND: > + python -c "import pybamm; from pybamm.rust import ExprGraph; print(pybamm.__version__, ExprGraph.__module__)" + + - name: Check wheel metadata + run: pipx run twine check --strict wheelhouse/*.whl + + - name: Verify the abi3 wheel imports on the newest Python + # CIBW_TEST_COMMAND only exercises the build interpreter (3.10). The point + # of abi3 is that the same wheel serves 3.14, so prove it once. + if: ${{ matrix.crossversion }} + run: | + pipx run --spec 'uv' uv venv --python 3.14 /tmp/abi3-check + VIRTUAL_ENV=/tmp/abi3-check pipx run --spec 'uv' uv pip install wheelhouse/*.whl + /tmp/abi3-check/bin/python -c " + import sys + from pybamm.rust import ExprGraph + assert sys.version_info[:2] == (3, 14), sys.version_info + print(sys.version_info[:2], ExprGraph.__module__) + " + + - name: Upload wheels + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: ${{ matrix.artifact }} + path: wheelhouse/*.whl + if-no-files-found: error + + build_sdist: if: startsWith(github.event.release.tag_name, 'pybamm-v') runs-on: ubuntu-latest permissions: @@ -19,26 +88,36 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: - persist-credentials: true - # Full history + tags so hatch-vcs can resolve the pybamm-v* version. + persist-credentials: false fetch-depth: 0 - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: - python-version: 3.14 + python-version: 3.13 + + - name: Build sdist + run: pipx run build --sdist packages/pybamm --outdir deploy + + - name: Guard the sdist size + # A force-include entry pointing one level too high would vendor the + # multi-gigabyte Cargo target/ tree. Fail rather than publish it. + run: | + size=$(stat -c%s deploy/*.tar.gz) + echo "sdist is $size bytes" + test "$size" -lt 52428800 || { echo "sdist over 50 MB - check sdist force-include"; exit 1; } + ! tar -tzf deploy/*.tar.gz | grep -q "pybamm-rust/target/" - - name: Build wheel - # Build the pybamm workspace package (its pyproject now lives under packages/). - run: pipx run build packages/pybamm --outdir deploy + - name: Check sdist metadata + run: pipx run twine check --strict deploy/*.tar.gz - - name: Upload package + - name: Upload sdist uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: - name: distributions - path: deploy/ + name: pybamm_sdist + path: deploy/*.tar.gz if-no-files-found: error publish: - needs: build + needs: [build_wheels, build_sdist] runs-on: ubuntu-latest environment: pypi permissions: diff --git a/.github/workflows/test_on_push.yml b/.github/workflows/test_on_push.yml index 7dcc6fb8d2..db82ffef45 100644 --- a/.github/workflows/test_on_push.yml +++ b/.github/workflows/test_on_push.yml @@ -55,6 +55,7 @@ jobs: - '.github/workflows/_build_solver_wheels.yml' pybamm: - 'packages/pybamm/**' + - 'packages/pybamm-rust/**' - 'pyproject.toml' - 'uv.lock' - 'noxfile.py' @@ -89,6 +90,50 @@ jobs: echo "macos_runners=$macos" } >> "$GITHUB_OUTPUT" + rust_msrv: + needs: changes + if: ${{ github.event_name == 'push' || needs.changes.outputs.pybamm == 'true' }} + runs-on: ubuntu-latest + permissions: + contents: read + name: Rust MSRV (1.89) + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install the MSRV toolchain + run: rustup toolchain install 1.89 --profile minimal --no-self-update + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + key: msrv-1.89 + - name: Check against the declared MSRV + run: > + cargo +1.89 check --manifest-path packages/pybamm-rust/Cargo.toml + --locked --all-targets + + rust_tests: + needs: changes + if: ${{ github.event_name == 'push' || needs.changes.outputs.pybamm == 'true' }} + runs-on: ubuntu-latest + permissions: + contents: read + name: Rust tests + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + key: tests + # Includes the ABI drift tests, which read the pybammsolvers consumer + # header, so a cross-language contract break fails here. + - name: Run the Rust test suite + run: > + cargo test --manifest-path packages/pybamm-rust/Cargo.toml + --locked --workspace --all-features + style: runs-on: ubuntu-latest permissions: @@ -102,6 +147,9 @@ jobs: with: python-version: 3.12 + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: version: "latest" @@ -202,7 +250,7 @@ jobs: # Guards both the bootstrap and the clean-clone dev experience. from_source_smoke: needs: changes - if: ${{ github.event_name == 'push' || needs.changes.outputs.solver == 'true' }} + if: ${{ github.event_name == 'push' || needs.changes.outputs.solver == 'true' || needs.changes.outputs.pybamm == 'true' }} runs-on: ubuntu-latest permissions: contents: read @@ -225,6 +273,9 @@ jobs: with: python-version: 3.13 + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: version: "latest" @@ -289,6 +340,8 @@ jobs: needs: - changes - style + - rust_msrv + - rust_tests - build_solver - run_unit_tests - run_coverage diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ee37733c13..692ae33f0a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -66,3 +66,16 @@ repos: additional_dependencies: [pyyaml] files: ^\.github/(workflows/test_on_push\.yml|scripts/check_ci_gate\.py)$ pass_filenames: false + - id: cargo-fmt + name: cargo fmt + entry: cargo fmt --manifest-path packages/pybamm-rust/Cargo.toml --all + language: system + types: [rust] + pass_filenames: false + - id: cargo-clippy + name: cargo clippy + entry: cargo clippy --manifest-path packages/pybamm-rust/Cargo.toml --all-targets --fix --allow-dirty --allow-staged + language: system + types: [rust] + pass_filenames: false + args: ["--", "-D", "warnings"] diff --git a/AGENTS.md b/AGENTS.md index a169799237..82c6afdf7c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -149,6 +149,11 @@ gets wrong: `import matplotlib` breaks `import pybamm` for minimal installs. - Public, user-facing objects are re-exported through `packages/pybamm/src/pybamm/__init__.py` (users write `pybamm.X`) and get a `docs/source/api/*.rst` entry. +- **The Rust bindings ship hand-written stubs.** Any change to the Python-visible API in + `packages/pybamm-rust/pybamm-python/src` must update `packages/pybamm/src/pybamm/rust/_core.pyi` + in the same commit. `mypy.stubtest` (run by `tests/unit/test_rust_stubs.py`) machine-checks + names, arities and defaults against the built extension; the types are review-enforced, so keep + them as precise as the bindings' own coercions (e.g. the `p` dict-or-array union). - Every feature or fix adds a `CHANGELOG.md` bullet under `# [Unreleased]` (Keep a Changelog format), ending with the PR link, e.g. `([#1234](https://github.com/pybamm-team/PyBaMM/pull/1234))`. diff --git a/benchmarks/README.md b/benchmarks/README.md index 319d13402b..16dd2e44a7 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -94,3 +94,274 @@ asv dev ``` `asv dev` implies options `--quick`, `--show-stderr`, and `--dry-run` (to avoid updating the `results` directory). + +### Rust observability suite + +For branch-local Rust vs CasADi observability, use the lightweight harness instead of +adding another standalone script: + +```shell +uv run python benchmarks/run_rust_observability.py --lane artifact +uv run python benchmarks/run_rust_observability.py --lane solver +``` + +Useful options: + +```shell +uv run python benchmarks/run_rust_observability.py --lane artifact --artifact-scenarios toy_expr --repeats 1 --warmup 0 +uv run python benchmarks/run_rust_observability.py --lane solver --models SPM DFN --json /tmp/rust-observability.json +uv run python benchmarks/run_rust_observability.py --lane solver --models SPM --output-points 1000 +uv run python benchmarks/run_rust_observability.py --lane solver --protocols drive_cycle pulse_train --aot none +uv run python benchmarks/run_rust_observability.py --lane inference --models DFN --inference-seed 7 +uv run python benchmarks/run_rust_observability.py --lane solver --reference-tolerance 0 +``` + +### The converged reference + +Every row in the solver, sensitivity and inference lanes — **the CasADi baseline +included** — is scored against one converged `casadi_idaklu` solve of the same scenario, +run at `--reference-tolerance` (1e-10 by default) instead of the scenario's 1e-6. The `Δ` +columns are that error. `Base Δ` is the raw difference from `casadi_idaklu` at the +scenario tolerance, reported beside it and gated on nothing: one cell holding whichever +of that row's comparisons came closest to its own tolerance, so on a gradient run it can +be a gradient difference rather than a value one. The per-block breakdown is in the JSON. + +Both numbers are needed because they answer different questions. Two backends at the same +tolerance differ by their *mutual* error, which is not either one's accuracy; when they +are the same integrator taking the same steps the shared error cancels outright, so +`rust_idaklu` reads ~1e-13 against `casadi_idaklu` while both sit a full tolerance unit +from the answer. That makes `Base Δ` the most sensitive port-regression detector in the +suite and a useless accuracy check, and the reference the reverse. Measured on SPM +`cc_discharge`: `casadi_idaklu` and `rust_idaklu` are each 1.0e-06 from the reference and +2.3e-13 from each other; `rust_diffsol` is 5.1e-06 from the reference. Before the +reference existed, that last number was read as diffsol being wrong, when most of it is +the baseline's own error. + +The gate carries two decades of headroom over the scenario tolerance. A tolerance bounds +the local error of one step and the global error accumulates it over thousands, so a +correct solve lands several tolerance units out: measured across the matrix, +`casadi_idaklu` itself reaches ~26 units on states, `rust_diffsol` ~104. A gate at one +tolerance unit would flag the reference integrator on its own scenarios. + +Gradients keep the looser `10 * sqrt(tol)` allowance for the reason in **Inference lane** +below — forward sensitivities are not error-controlled to the state tolerance. The +reference is correspondingly less converged in gradients than in values: SPMe `cc_charge` +improves only 3.5e-03 → 1.0e-04 over four decades of tolerance (the `sqrt(tol)` rate), +and DFN does not improve at all, holding 6.9e-03 at `t = 0` from 1e-6 to 1e-10 while its +value at the same point converges to 2e-10 (see R5 in the issue tracker). A gradient +reference is therefore worth 3-4 decades on SPM and SPMe and about half the gate on DFN, +where a regression below ~1e-02 would not be visible. Values are clear at every model. + +A converged solve is not always reachable. DFN under a ramping current fails IDA's error +test at `t = 0` below 1e-9, so the reference tolerance is loosened a decade at a time +until it converges, never past two decades clear of the scenario. DFN `drive_cycle` +therefore reports against a 1e-09 reference and DFN `pulse_train`, which converges at no +usable tolerance, falls back to the baseline comparison with a logged warning. The +tolerance each row was actually judged against is on the row (`reference_tolerance` in +JSON, listed in the caption above each table). + +`--reference-tolerance 0` drops the reference entirely and restores the older behaviour: +every candidate compared against `casadi_idaklu` at the scenario tolerance, with the +baseline row itself ungated (`baseline` rather than `pass`/`warn`). + +The reference costs one extra untimed solve per scenario in the solver and sensitivity +lanes, and one per draw in the inference lane (~26/160/818 ms for SPM/SPMe/DFN with +gradients). It is solved outside every timed region, so no timing column moves. + +Note the reference shares an integrator family with the baseline, so it cannot by itself +catch a systematic IDA error. What would catch one is `rust_diffsol` — an independent +integrator — failing to converge on the same answer as the tolerance tightens. + +### Protocols + +A scenario is a model crossed with an operating protocol. `--protocols` selects the +protocol axis and defaults to `cc_discharge` alone. + +Every protocol layers its current law on one shared base parameter set, `Chen2020` +(`registry.BASE_PARAMETER_SET`), so a row's timings depend only on the model and the +protocol. Chen2020 rather than each model's own defaults because it gives both particle +diffusivities as scalars, which the inference lane needs to swap for inputs like for +like. Note this differs from the pre-protocol suite, which used `model.default_parameter_values` +— solver-lane numbers recorded before that change are not comparable. + +| Protocol | What it runs | What it exercises beyond a plain discharge | +| --- | --- | --- | +| `cc_discharge` | 1C discharge for 3600 s | the baseline | +| `cc_charge` | 2C charge from 0% SOC over 1800 s | event termination (~730-1150 s), so the trajectory comparison sees a non-`final time` stop | +| `drive_cycle` | triangle-wave current at 50% SOC | an `Interpolant` node inside the rhs graph, in one continuous solve | +| `pulse_train` | ramped pulse/rest train at 50% SOC | the same graph shape with sharp transitions: step-size rejection and recovery | +| `experiment` | discharge / rest / charge / rest | the `solver.step` path: per-step restarts and per-step initial conditions | + +`initial_soc` is applied once at build, never per solve, because passing it to +`Simulation.solve` re-runs `set_initial_state` (and an ElectrodeSOH solve) on every call, +which would swamp the warm timings. The `experiment` protocol takes its output grid from +its step period rather than `--output-points`, so its `Pts` column reads `-` in the solver +and sensitivity lanes; it also builds lazily on first solve, so its `Build` column reads 0 +by design. + +Only `experiment` restarts the integrator: it solves step by step through +`solver.step`, with fresh initial conditions per step. `pulse_train` is one continuous +solve — a linear `Interpolant` raises no discontinuity events (those come only from +`Heaviside`/`Modulo` nodes in `t`), so its breakpoints are `t_eval` stops that the +stepper adapts through rather than restarts at. + +`SENSITIVITY_INPUTS` includes `Current function [A]`, which `drive_cycle` and +`pulse_train` replace with an `Interpolant` and `experiment` supersedes with its own +control. Under those protocols the sensitivity lane differentiates the remaining +parameters only. Which parameters a row actually differentiated is recorded on the row +itself (the `Params` column, and `sensitivity_parameters` in JSON) rather than inferred +from the run's configuration. + +### Inference lane + +`--lane inference` models a parameter-estimation loop rather than a single solve. Four +parameters — both particle diffusivities and both active-material volume fractions — +stay symbolic as `InputParameter`s, and +**every timed repeat solves with a different input vector**, drawn log-uniformly from +`--inference-seed` within a per-parameter half-width (`registry.INFERENCE_SPREADS`). The +same vectors are shared across backends, so repeat *i* is compared like for like. + +The width is per parameter because one figure cannot suit all of them: 20% is routine for +a diffusivity, but the same 20% on an active-material fraction swings porosity between +0.165 and 0.39 about a nominal 0.25, which stalls DFN with `IDA_CONV_FAIL` and drives +SPMe's 2C charge into its voltage ceiling within a couple of output points. Volume +fractions therefore get 5%, holding negative porosity to 0.213-0.285. A parameter added +without a width raises rather than inheriting someone else's. + +The inputs are layered *on top of* the selected protocol, so `--lane inference +--protocols drive_cycle` fits against a drive cycle rather than silently reverting to a +constant current. Because active material and pore volume sum to 1 in the base set, each +electrode's porosity is bound to `1 - eps` (`registry.INFERENCE_COMPLEMENTS`) so sampled +geometries stay feasible. + +A fitted parameter must also leave the initial state valid. The base set gives the +initial concentration as an absolute value, so a fitted *maximum* concentration moves the +initial stoichiometry rather than the capacity: from a nominal 0.90, a -20% draw puts it +past 1.0 and the cell starts above its voltage cutoff. Maximum concentrations are +therefore excluded — they are not usually fitted quantities either. + +Two things differ from the solver lane on purpose. Observation goes through the +interpolating call interface (`solution[var](t)`) at midpoints between solver nodes, +which is the door a fitting loop actually uses, whereas the solver lane deliberately +reads `.data` on the solver's own stored grid. Between them both doors are covered. And +the table splits one-time `Build`/`Setup`/`ColdObs` from the per-evaluation `Eval p50` +with its `p10-p90` spread — under changing inputs that spread is real, not timer noise. + +`ColdObs` is the first, forced materialisation of the observed variable, taken before the +warmup loop. Lazy variable compilation costs several milliseconds and would otherwise +either vanish into a discarded warmup repeat or land inside the first timed one; charging +it explicitly keeps the per-evaluation samples steady-state at any `--warmup`. + +Both full-state and `output_variables` rows run. They are not duplicates: restricting the +solver to the observed variable changes state storage, the observation path and the +per-evaluation cost substantially (often ~2x). Not every backend supports every +combination — diffsol cannot currently stitch an `output_variables` solution across +experiment steps — and those degrade to an in-row `unsupported`. + +Parity is worst-case across every repeat, not just the last, and every backend is read on +the same observation grid. Repeats are ranked by *tolerance-normalised* error, not by +absolute difference: the admissible difference scales with the reference value, so a +repeat that breached tolerance low on the discharge curve must not be masked by one +further up whose larger difference is still permitted. The reference is solved once per +draw and, like every backend, is built from the first input vector: the lane holds `y0` +fixed there, and a reference resolved from a later draw would start the cell at a +different state of charge and move a `cc_charge` event by ~80 s. + +Agreement in value is not sufficient. Each repeat also records its termination reason and +final time, and `Cover` reports the fraction of the shared grid the candidate actually +reached, so a backend that stops early cannot pass on a matching prefix. Where the two +terminations disagree, observation points inside that measured window are excluded from +the value comparison — both trajectories are racing to a cutoff at different moments +there, so those points report the endpoint gap rather than trajectory agreement. The +window is measured, not a fixed point count, so it is zero when the terminations coincide +and does not change meaning with `--output-points`. + +`--inference-sensitivities` adds forward sensitivities for the gradient-based case, and +the gradient is materialised inside the timed observation, so the chain rule is costed +rather than only the sensitivity integration. Gradients are compared as `p . dV/dp`: the +fitted parameters span eighteen orders of magnitude (a diffusivity at 1e-14 beside a +concentration at 1e4), so raw `d/dp` columns are not comparable under one tolerance, while +the scaled form is dimensionless and is what a log-space fitting loop consumes. They are +judged at the square root of the state tolerance, because a solver error-controls the +state to `(atol, rtol)` but not the sensitivities integrated alongside it; that still +catches the order-of-magnitude miss a broken chain rule produces. + +Against the converged reference this lane inverts a conclusion the cross-backend gate used +to reach. On SPMe `cc_charge` with gradients, `casadi_idaklu` and `rust_idaklu` are both +1.8e-01 from the reference while `rust_diffsol` is 4.2e-03 — diffsol's gradients are the +*more* accurate by some forty times, and what the old gate scored as diffsol's error was +mostly the baseline's own. diffsol puts sensitivities under error control with +`sens_atol_factor` tightening the differential-state floor, where IDAS runs them at the +state tolerance. + +The initial state is resolved once, from the first input vector, and held fixed while the +fitted parameters vary. That measures compiled-model reuse under a fixed `y0`, which is +the intended workload, and it is why the fitted set is restricted to parameters that do +not define `y0`. A fit that did vary one would have to re-derive `y0` — and pay an +ElectrodeSOH solve — every evaluation, which is a different workload. + +The lane never profiles AOT in an isolated cache; the `AOT` column instead reports what +the compiler actually did (`miss`, `disk`, `memory`) so a cheap warm `Setup` is never +misread as a fresh compile. Its `Pts` column counts observation timestamps, not the +`--output-points` request. + +### AOT rows + +`--aot` controls which lanes run the CasADi ahead-of-time rows: `solver` (the default), +`all`, or `none`. The sensitivity and inference lanes are off by default because their AOT +rows roughly double the run time; turn them on with `--aot all` when the AOT comparison on +those paths is what you are after. + +The artifact lane times only each backend's native kernel call (format conversion +is excluded) and uses auto-calibrated inner batching, so `--repeats` counts timed +batches rather than individual calls. All solver backends run at the per-scenario +tolerance for an iso-accuracy comparison. + +The suite is intentionally small. It reports parity and timing breakdowns for the +prep-artifact API and representative solver paths. The solver lane always attempts +CasADi IDAKLU, Rust IDAKLU, and Rust diffsol rows, plus separate `output_only` rows +for the observed variable path. Unsupported diffsol configurations are reported in-row +instead of aborting the whole suite. This harness is not a CI gate or a replacement +for ASV. + +Console and JSON both carry `reference_tolerance` per row and in the run metadata, so a +saved run records what its Δ columns were measured against. + +Solver and sensitivity lanes request 1000 output points by default; use +`--output-points` to measure trajectory scaling explicitly. `Build` covers model +processing and discretisation. `Prep` combines first-solve setup/compilation with +the first forced output materialisation, including lazy observation compilation. +`Cold` is the measured wall clock from the start of model build through that first +materialised result. `Solve` is the backend's internal timer, `Wall` covers the +complete warm `Simulation.solve` call, and `E2E` is a paired warm wall-clock sample +through forced output materialisation. JSON output retains the raw samples and +runtime metadata. Rust rows also report parent colours and dense-row count, entries, +and sweeps. Backend order is randomized reproducibly with `--backend-order-seed`, +including the CasADi baseline: pinning the reference first measures it on a +systematically colder machine than everything it is compared against. The shuffle key +covers model *and* protocol, so protocols of one model no longer share an order. Every +comparison is computed after execution, leaving the numbers independent of the order +things ran in. Repeats within a case are still contiguous, so this reduces +between-backend drift rather than eliminating it; interleaving individual repeats would +need every backend's simulation resident at once. Console and JSON results are reordered +deterministically for comparison: full-state rows then output-only rows, each ordered as +CasADi IDAKLU, CasADi AOT, Rust IDAKLU, and Rust diffsol. + +JSON metadata records the commit, whether the tree was dirty, and a digest of the +combined staged and unstaged diff, so two runs from different uncommitted work are +distinguishable. Untracked files fall outside the digest, as they fall outside the diff. + +Solver and sensitivity CasADi AOT rows use an isolated empty cache, so their main +`Prep` and `Cold` measurements always include genuine code generation, native +compilation, and library loading. The separate AOT profile verifies every fresh +cache miss, repeats the same case in a new Python process, and reports the +corresponding persistent-cache disk hits, phase timings, disk-cached `Prep`/`Cold`, +and generated library size. A failed compile or unexpected cache state aborts the +row rather than silently reporting the CasADi VM as AOT. This makes solver and +sensitivity runs slower by design, especially for DFN; the warm timing columns +remain the steady-state comparison. + +Console output uses the detected terminal width. The solver and sensitivity lanes +keep one full table when it fits, split timing and validation into compact tables on +laptop-sized terminals, and fall back to one wrapped block per backend below 101 +columns. JSON output retains every field regardless of the console layout. diff --git a/benchmarks/dfn_ab_harness.py b/benchmarks/dfn_ab_harness.py new file mode 100644 index 0000000000..452532f23f --- /dev/null +++ b/benchmarks/dfn_ab_harness.py @@ -0,0 +1,176 @@ +"""Shared DFN A/B harness — CasADi expression tree vs Rust DAG. + +Each benchmark goes through `build_dfn_ab(npts)` for one discretised DFN model +and the matched CasADi Function / CompiledModel pair. Both consume the same +`(t, y, p)`, so the A/B is fair: identical math, state size and inputs. + +Examples +-------- +From a benchmark file in this directory:: + + import os, sys + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from dfn_ab_harness import build_dfn_ab, sample_state + + ab = build_dfn_ab(npts=20) + y = sample_state(ab) + p = ab.inputs_array + res_casadi = ab.casadi_residual(0.0, y, p) + res_rust = ab.rust_model.rhs(0.0, y, p) +""" + +from __future__ import annotations + +import numbers +from dataclasses import dataclass + +import numpy as np + +import pybamm + + +@dataclass +class DFNAB: + """Container for the matched CasADi/Rust DFN evaluation pair. + + All callables operate on the same `(t, y, p)` tuple, where `p` is a + 1-D numpy array of input-parameter values in declaration order. + """ + + n_states: int + n_inputs: int + npts: int + y0: np.ndarray + inputs_dict: dict[str, float] + inputs_array: np.ndarray + # CasADi side — `casadi.Function` instances bound to the discretised model. + casadi_residual: object + casadi_jacobian: object + # Rust side — single CompiledModel exposing all needed evals. + rust_model: object + # Sparsity, useful for assembled-Jacobian benchmarks. + csc_colptrs: np.ndarray + csc_rowinds: np.ndarray + nnz: int + + +def _make_var_pts(model: pybamm.BaseModel, npts: int) -> dict: + """Override scalar var_pts with `npts`; preserve direction grids (y, z).""" + out = {} + for k, v in model.default_var_pts.items(): + if isinstance(v, numbers.Number) and v > 1 and k not in {"y", "z"}: + out[k] = npts + else: + out[k] = v + return out + + +def build_dfn_ab( + npts: int = 20, + *, + with_input_current: bool = True, + parameter_set: str = "Chen2020", +) -> DFNAB: + """Build a discretised DFN and matched CasADi/Rust evaluation pair. + + Parameters + ---------- + npts + Number of finite-volume points in each spatial direction (other than + the unused y/z directions). 5 -> 92 states, 10 -> 282, 20 -> 962 + (default), 30 -> 2042. Used to drive scale sweeps. + with_input_current + If True, parameterise `Current function [A]` as InputParameter "I" so + the benchmark exercises the input-vector path. Defaults to True. + parameter_set + Name of the PyBaMM parameter set. Defaults to "Chen2020". + """ + from pybamm.rust import CompiledModel, ExprGraph + + model = pybamm.lithium_ion.DFN() + model.events = [] # Rust path doesn't support root-finding events yet. + + param = pybamm.ParameterValues(parameter_set) + if with_input_current: + param["Current function [A]"] = pybamm.InputParameter("I") + inputs_dict = {"I": 0.5} + else: + inputs_dict = {} + + var_pts = _make_var_pts(model, npts) + + model.convert_to_format = "casadi" + sim = pybamm.Simulation( + model, + parameter_values=param, + var_pts=var_pts, + solver=pybamm.IDAKLUSolver(), + ) + sim.build() + built = sim.built_model + + # Populate `rhs_algebraic_eval`, `jac_rhs_algebraic_eval`, etc. on the model. + sim.solver.set_up(built, inputs=inputs_dict) + + n_states = built.len_rhs_and_alg + n_inputs = len(inputs_dict) + inputs_array = np.array( + [float(inputs_dict[k]) for k in inputs_dict], dtype=np.float64 + ) + y0 = np.asarray(built.y0.full() if hasattr(built.y0, "full") else built.y0).reshape( + -1 + ) + + # Convert the same Symbol `_set_up_rust` would and build a CompiledModel + # directly, without going through the solver-group wrapper. + if built.len_alg > 0: + full_sym = pybamm.numpy_concatenation( + built.concatenated_rhs, built.concatenated_algebraic + ) + else: + full_sym = built.concatenated_rhs + graph = ExprGraph() + rhs_expr = full_sym.to_rust(graph, {}) + + mass = built.mass_matrix.entries # scipy CSR + rust_model = CompiledModel.from_expr( + graph, + rhs_expr, + mass.data.astype(np.float64), + mass.indptr.astype(np.int64), + mass.indices.astype(np.int64), + n_inputs=n_inputs, + ) + + csc_colptrs, csc_rowinds = rust_model.csc_sparsity_pattern() + + return DFNAB( + n_states=n_states, + n_inputs=n_inputs, + npts=npts, + y0=y0, + inputs_dict=inputs_dict, + inputs_array=inputs_array, + casadi_residual=built.rhs_algebraic_eval, + casadi_jacobian=built.jac_rhs_algebraic_eval, + rust_model=rust_model, + csc_colptrs=np.asarray(csc_colptrs, dtype=np.int64), + csc_rowinds=np.asarray(csc_rowinds, dtype=np.int64), + nnz=int(rust_model.nnz), + ) + + +def sample_state( + ab: DFNAB, *, perturb: float = 0.0, seed: int | None = None +) -> np.ndarray: + """Return `ab.y0` optionally perturbed by `perturb * randn(n_states)`.""" + if perturb <= 0: + return ab.y0.copy() + rng = np.random.default_rng(seed) + return ab.y0 + perturb * rng.standard_normal(ab.n_states) + + +def casadi_jacobian_dense(ab: DFNAB, t: float, y: np.ndarray) -> np.ndarray: + """Convert CasADi sparse Jacobian to a dense numpy array (for parity).""" + J = ab.casadi_jacobian(t, y, ab.inputs_array) + return np.asarray(J) diff --git a/benchmarks/run_rust_observability.py b/benchmarks/run_rust_observability.py new file mode 100644 index 0000000000..1813b694b2 --- /dev/null +++ b/benchmarks/run_rust_observability.py @@ -0,0 +1,327 @@ +"""Entry point for the Rust-vs-CasADi observability benchmark suite. + +Runs one or more measurement lanes, prints a table per lane, and optionally writes +the same results as JSON. Runnable either as a module or as a script — the +``sys.path`` insert below covers the script case. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import platform +import subprocess +import sys +from pathlib import Path + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from benchmarks.rust_observability.registry import ( + DEFAULT_OUTPUT_POINTS, + DEFAULT_PROTOCOLS, + INFERENCE_INPUTS, + get_artifact_scenarios, + get_inference_scenarios, + get_solver_scenarios, +) +from benchmarks.rust_observability.report import ( + render_artifact_table, + render_inference_table, + render_sensitivity_table, + render_solver_table, + suite_to_jsonable, +) +from benchmarks.rust_observability.runners import ( + DEFAULT_REFERENCE_TOLERANCE, + SENSITIVITY_INPUTS, + run_artifact_lane, + run_inference_lane, + run_sensitivity_lane, + run_solver_lane, +) + +_DEFAULT_REPEATS = 10 +_DEFAULT_INFERENCE_REPEATS = 50 + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser for the suite.""" + parser = argparse.ArgumentParser( + description="Run the lightweight Rust-vs-CasADi observability suite." + ) + parser.add_argument( + "--lane", + choices=("artifact", "solver", "sensitivity", "inference", "all"), + default="all", + ) + parser.add_argument("--artifact-scenarios", nargs="*") + parser.add_argument("--models", nargs="*") + parser.add_argument( + "--protocols", + nargs="*", + default=list(DEFAULT_PROTOCOLS), + help="Operating protocols to run (default: %(default)s).", + ) + parser.add_argument( + "--repeats", + type=int, + default=None, + help=( + f"Timed repeats per case (default: {_DEFAULT_REPEATS}, " + f"{_DEFAULT_INFERENCE_REPEATS} on the inference lane)." + ), + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument( + "--output-points", + type=int, + default=DEFAULT_OUTPUT_POINTS, + help="Requested solver output points (default: %(default)s).", + ) + parser.add_argument( + "--backend-order-seed", + type=int, + default=0, + help="Seed used to randomize candidate backend execution order.", + ) + parser.add_argument( + "--aot", + choices=("solver", "all", "none"), + default="solver", + help="Which lanes run the CasADi AOT rows (default: %(default)s).", + ) + parser.add_argument( + "--reference-tolerance", + type=float, + default=DEFAULT_REFERENCE_TOLERANCE, + help=( + "Tolerance of the converged CasADi reference every row is judged " + "against (default: %(default)s). 0 drops the reference and compares " + "backends against each other at the scenario tolerance instead." + ), + ) + parser.add_argument( + "--inference-sensitivities", + action="store_true", + help="Request forward sensitivities in the inference lane.", + ) + parser.add_argument( + "--inference-seed", + type=int, + default=0, + help="Seed for the inference lane's input vectors (default: %(default)s).", + ) + parser.add_argument("--json", type=Path) + return parser + + +def include_aot_for(lane: str, aot: str) -> bool: + """Whether ``lane`` runs the AOT backend rows under the ``--aot`` setting.""" + if aot == "none": + return False + if aot == "all": + return True + return lane == "solver" + + +def _package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _git(*arguments: str) -> str: + return subprocess.run( + ["git", *arguments], check=True, capture_output=True, text=True + ).stdout + + +def _git_metadata() -> dict[str, str | bool | None]: + """Identify the source that produced a run. + + A dirty tree makes the revision alone ambiguous, so the combined staged and + unstaged diff is digested too and two different local implementations of the + same idea no longer share their metadata. Untracked files are outside the + digest, as they are outside the diff. + """ + try: + revision = _git("rev-parse", "HEAD").strip() + status = _git("status", "--porcelain") + diff = _git("diff", "HEAD") + digest = hashlib.sha256(diff.encode()).hexdigest()[:16] if diff else None + except (OSError, subprocess.CalledProcessError): + return {"git_revision": None, "git_dirty": None, "git_diff_digest": None} + return { + "git_revision": revision, + "git_dirty": bool(status), + "git_diff_digest": digest, + } + + +def _runtime_metadata() -> dict: + return { + **_git_metadata(), + "python": platform.python_version(), + "platform": platform.platform(), + "machine": platform.machine(), + "processor": platform.processor(), + "packages": { + name: _package_version(name) + for name in ("pybamm", "pybammsolvers", "casadi") + }, + } + + +def repeats_for(lane: str, requested: int | None) -> int: + """Timed repeats for one lane; ``None`` takes that lane's own default. + + The inference lane reports a p10-p90 spread, which needs more samples than + the paired medians the other lanes report. + """ + if requested is not None: + return requested + return _DEFAULT_INFERENCE_REPEATS if lane == "inference" else _DEFAULT_REPEATS + + +def _comparison_metadata(args, runtime_metadata: dict, repeats: int) -> dict: + """The metadata every comparison lane records, before its own additions.""" + return { + "repeats": repeats, + "warmup": args.warmup, + "requested_output_points": args.output_points, + "reference_tolerance": args.reference_tolerance or None, + "protocols": args.protocols, + "aot": args.aot, + "backend_order_seed": args.backend_order_seed, + **runtime_metadata, + } + + +def main(argv: list[str] | None = None) -> int: + """Run the selected lanes and print their tables. + + Returns + ------- + int + Process exit status; 0 once every requested lane has run. + """ + args = build_parser().parse_args(argv) + if args.repeats is not None and args.repeats < 1: + raise ValueError("repeats must be at least 1") + repeats = repeats_for("solver", args.repeats) + inference_repeats = repeats_for("inference", args.repeats) + if args.warmup < 0: + raise ValueError("warmup must be non-negative") + payloads = {} + runtime_metadata = _runtime_metadata() + + if args.lane in {"artifact", "all"}: + artifact_results = run_artifact_lane( + get_artifact_scenarios(args.artifact_scenarios), + repeats=repeats, + warmup=args.warmup, + ) + print("Artifact lane") + print(render_artifact_table(artifact_results)) + payloads["artifact"] = suite_to_jsonable( + "artifact", + artifact_results, + metadata={ + **runtime_metadata, + "repeats": repeats, + "warmup": args.warmup, + }, + ) + print() + + if args.lane in {"solver", "all"}: + solver_results = run_solver_lane( + get_solver_scenarios( + args.models, args.protocols, output_points=args.output_points + ), + repeats=repeats, + warmup=args.warmup, + include_aot=include_aot_for("solver", args.aot), + backend_order_seed=args.backend_order_seed, + reference_tolerance=args.reference_tolerance, + ) + print("Solver lane") + print(render_solver_table(solver_results)) + payloads["solver"] = suite_to_jsonable( + "solver", + solver_results, + metadata={ + **_comparison_metadata(args, runtime_metadata, repeats), + "diffsol_always_attempted": True, + "aot_profile": "isolated miss plus fresh-process disk reload", + }, + ) + print() + + if args.lane in {"sensitivity", "all"}: + sensitivity_results = run_sensitivity_lane( + get_solver_scenarios( + args.models, args.protocols, output_points=args.output_points + ), + repeats=repeats, + warmup=args.warmup, + include_aot=include_aot_for("sensitivity", args.aot), + backend_order_seed=args.backend_order_seed, + reference_tolerance=args.reference_tolerance, + ) + print("Sensitivity lane") + print(render_sensitivity_table(sensitivity_results)) + payloads["sensitivity"] = suite_to_jsonable( + "sensitivity", + sensitivity_results, + metadata={ + **_comparison_metadata(args, runtime_metadata, repeats), + # The superset asked for; a protocol can narrow it, so the + # parameters actually differentiated are recorded per row. + "sensitivity_parameters_requested": sorted(SENSITIVITY_INPUTS.values()), + "aot_profile": "isolated miss plus fresh-process disk reload", + }, + ) + + if args.lane in {"inference", "all"}: + inference_results = run_inference_lane( + get_inference_scenarios( + args.models, args.protocols, output_points=args.output_points + ), + repeats=inference_repeats, + warmup=args.warmup, + seed=args.inference_seed, + sensitivities=args.inference_sensitivities, + include_aot=include_aot_for("inference", args.aot), + backend_order_seed=args.backend_order_seed, + reference_tolerance=args.reference_tolerance, + ) + print() + print("Inference lane") + print(render_inference_table(inference_results)) + payloads["inference"] = suite_to_jsonable( + "inference", + inference_results, + metadata={ + **_comparison_metadata(args, runtime_metadata, inference_repeats), + "inference_parameters": sorted(INFERENCE_INPUTS.values()), + "inference_seed": args.inference_seed, + "inference_sensitivities": args.inference_sensitivities, + }, + ) + + if args.json: + json_payload = payloads + if args.lane in {"artifact", "solver", "sensitivity", "inference"}: + json_payload = payloads[args.lane] + args.json.write_text(json.dumps(json_payload, indent=2)) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/rust_observability/__init__.py b/benchmarks/rust_observability/__init__.py new file mode 100644 index 0000000000..92b621b3a0 --- /dev/null +++ b/benchmarks/rust_observability/__init__.py @@ -0,0 +1 @@ +"""Shared Rust observability benchmark helpers.""" diff --git a/benchmarks/rust_observability/registry.py b/benchmarks/rust_observability/registry.py new file mode 100644 index 0000000000..63cf7fb195 --- /dev/null +++ b/benchmarks/rust_observability/registry.py @@ -0,0 +1,454 @@ +"""Scenario definitions for the Rust-vs-CasADi observability suite. + +Scenarios are declared once here so the runners and the report agree on what was +measured; ``get_*_scenarios`` is the only way to obtain them, and it validates the +names a caller asked for. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +import numpy as np + +import pybamm + +DEFAULT_OUTPUT_POINTS = 1000 +DEFAULT_PROTOCOLS = ("cc_discharge",) + +# Chen2020 gives both particle diffusivities as scalars, so the inference lane +# can swap them for inputs like for like. +BASE_PARAMETER_SET = "Chen2020" + +CC_DURATION_S = 3600.0 +CHARGE_C_RATE = 2.0 +# 2C from empty tops out at ~1150 s; the margin keeps the event inside the grid +# as the inference lane perturbs the inputs. +CHARGE_DURATION_S = 1800.0 +TRIANGLE_AMPLITUDE_A = 5.0 +TRIANGLE_PERIOD_S = 600.0 +TRIANGLE_DURATION_S = 1800.0 +PULSE_AMPLITUDE_A = 5.0 +PULSE_ON_S = 60.0 +PULSE_REST_S = 120.0 +PULSE_RAMP_S = 2.0 +PULSE_DURATION_S = 1800.0 +EXPERIMENT_PERIOD = "10 seconds" + +# Every entry must be live on every model in the matrix. No maximum concentrations: +# initial concentration is absolute here, so fitting one moves the stoichiometry past 1. +INFERENCE_INPUTS = { + "Negative particle diffusivity [m2.s-1]": "D_n", + "Positive particle diffusivity [m2.s-1]": "D_p", + "Negative electrode active material volume fraction": "eps_n", + "Positive electrode active material volume fraction": "eps_p", +} + +# Active material and pore volume sum to 1 in the base set, so each porosity has +# to track its electrode's fitted fraction or the geometry goes infeasible. +INFERENCE_COMPLEMENTS = { + "Negative electrode porosity": "eps_n", + "Positive electrode porosity": "eps_p", +} + +# Sampling half-width per parameter: 20% on a volume fraction swings porosity +# 0.39 to 0.165 about a nominal 0.25, which stalls DFN and saturates SPMe on charge. +INFERENCE_SPREADS = { + "D_n": 0.2, + "D_p": 0.2, + "eps_n": 0.05, + "eps_p": 0.05, +} + + +@dataclass(frozen=True) +class ArtifactScenario: + """One model whose compiled artifacts are timed, operation by operation.""" + + name: str + operations: tuple[str, ...] + atol: float + rtol: float + + +@dataclass(frozen=True) +class SolvePlan: + """How one protocol issues its solve. + + Parameters + ---------- + t_interp : numpy.ndarray or None + Output grid. ``None`` when the protocol determines its own grid. + t_eval : numpy.ndarray or None + Breakpoints passed as ``t_eval``. ``None`` means ``[t0, tf]``. + experiment : pybamm.Experiment or None + When set, the solve goes through the experiment path. + """ + + t_interp: np.ndarray | None + t_eval: np.ndarray | None + experiment: Any | None = None + + @property + def requested_points(self) -> int: + """Output points asked for; ``0`` when the protocol decides the grid.""" + return 0 if self.t_interp is None else self.t_interp.size + + +@dataclass(frozen=True) +class Protocol: + """One operating condition, independent of which model runs it.""" + + build_parameter_values: Callable[[Any], Any] + initial_soc: float | None + build_plan: Callable[[int], SolvePlan] + + +def _base_parameter_values(): + """The shared base parameter set, at its own constant-current default.""" + return pybamm.ParameterValues(BASE_PARAMETER_SET) + + +def _grid_plan(duration: float) -> Callable[[int], SolvePlan]: + """A plan builder issuing ``output_points`` samples evenly over ``duration``.""" + + def build(output_points: int) -> SolvePlan: + return SolvePlan( + t_interp=np.linspace(0.0, duration, output_points), t_eval=None + ) + + return build + + +def _charge_parameter_values(): + parameter_values = _base_parameter_values() + # 1C never reaches the voltage cutoff within the window; 2C terminates on it + # for all three models. + parameter_values["Current function [A]"] = -CHARGE_C_RATE * float( + parameter_values["Nominal cell capacity [A.h]"] + ) + return parameter_values + + +def _triangle_breakpoints() -> tuple[np.ndarray, np.ndarray]: + """Vertices of the triangle wave, which linear interpolation reproduces exactly.""" + vertices = np.unique( + np.concatenate( + [ + [0.0], + np.arange( + TRIANGLE_PERIOD_S / 4.0, + TRIANGLE_DURATION_S, + TRIANGLE_PERIOD_S / 2.0, + ), + [TRIANGLE_DURATION_S], + ] + ) + ) + values = ( + TRIANGLE_AMPLITUDE_A + * (2.0 / np.pi) + * np.arcsin(np.sin(2.0 * np.pi * vertices / TRIANGLE_PERIOD_S)) + ) + return vertices, values + + +def _pulse_breakpoints() -> tuple[np.ndarray, np.ndarray]: + """Edges of a ramped pulse/rest train; finite ramps keep the profile Lipschitz.""" + times = [0.0] + values = [0.0] + start = 0.0 + period = PULSE_ON_S + PULSE_REST_S + while start < PULSE_DURATION_S: + times.extend( + [ + start + PULSE_RAMP_S, + start + PULSE_ON_S, + start + PULSE_ON_S + PULSE_RAMP_S, + start + period, + ] + ) + values.extend([PULSE_AMPLITUDE_A, PULSE_AMPLITUDE_A, 0.0, 0.0]) + start += period + times_arr = np.asarray(times, dtype=np.float64) + values_arr = np.asarray(values, dtype=np.float64) + keep = times_arr <= PULSE_DURATION_S + return times_arr[keep], values_arr[keep] + + +def _interpolant_parameter_values(breakpoints): + times, values = breakpoints + parameter_values = _base_parameter_values() + parameter_values["Current function [A]"] = pybamm.Interpolant( + times, values, pybamm.t, interpolator="linear" + ) + return parameter_values + + +def _drive_cycle_parameter_values(): + return _interpolant_parameter_values(_triangle_breakpoints()) + + +def _pulse_train_parameter_values(): + return _interpolant_parameter_values(_pulse_breakpoints()) + + +def _interpolant_plan(breakpoints, duration: float, output_points: int) -> SolvePlan: + times, _ = breakpoints + grid = np.linspace(0.0, duration, output_points) + # Every breakpoint must appear in t_eval or PyBaMM warns about resolution. + return SolvePlan(t_interp=grid, t_eval=times) + + +def _drive_cycle_plan(output_points: int) -> SolvePlan: + return _interpolant_plan( + _triangle_breakpoints(), TRIANGLE_DURATION_S, output_points + ) + + +def _pulse_train_plan(output_points: int) -> SolvePlan: + return _interpolant_plan(_pulse_breakpoints(), PULSE_DURATION_S, output_points) + + +def _experiment_plan(output_points: int) -> SolvePlan: + del output_points # The step period fixes the grid, identically on every backend. + experiment = pybamm.Experiment( + [ + "Discharge at 1C for 10 minutes", + "Rest for 5 minutes", + "Charge at 1C for 10 minutes", + "Rest for 5 minutes", + ], + period=EXPERIMENT_PERIOD, + ) + return SolvePlan(t_interp=None, t_eval=None, experiment=experiment) + + +_PROTOCOLS = { + "cc_discharge": Protocol( + build_parameter_values=_base_parameter_values, + initial_soc=None, + build_plan=_grid_plan(CC_DURATION_S), + ), + "cc_charge": Protocol( + build_parameter_values=_charge_parameter_values, + initial_soc=0.0, + build_plan=_grid_plan(CHARGE_DURATION_S), + ), + "drive_cycle": Protocol( + build_parameter_values=_drive_cycle_parameter_values, + initial_soc=0.5, + build_plan=_drive_cycle_plan, + ), + "pulse_train": Protocol( + build_parameter_values=_pulse_train_parameter_values, + initial_soc=0.5, + build_plan=_pulse_train_plan, + ), + "experiment": Protocol( + build_parameter_values=_base_parameter_values, + initial_soc=None, + build_plan=_experiment_plan, + ), +} + +_MODELS = { + "SPM": pybamm.lithium_ion.SPM, + "SPMe": pybamm.lithium_ion.SPMe, + "DFN": pybamm.lithium_ion.DFN, +} + + +@dataclass(frozen=True) +class SolverScenario: + """One model run under one protocol, plus the output variable to compare on.""" + + name: str + protocol: str + model_factory: Callable[..., Any] + model_options: dict[str, Any] + parameter_values_builder: Callable[[], Any] + initial_soc: float | None + plan: SolvePlan + observed_output: str + atol: float + rtol: float + + +def filter_names(available: list[str], selected: list[str] | None) -> list[str]: + """Narrow ``available`` to ``selected``, keeping registry order. + + Raises + ------ + ValueError + If a requested name is not registered, so a typo fails the run instead of + silently measuring less. + """ + if not selected: + return list(available) + available_set = set(available) + unknown = sorted(set(selected) - available_set) + if unknown: + raise ValueError(f"Unknown names requested: {', '.join(unknown)}") + return [name for name in available if name in set(selected)] + + +_ARTIFACT_SCENARIOS = { + "toy_expr": ArtifactScenario( + name="toy_expr", + operations=( + "eval", + "jacobian_y", + "jacobian_p", + "jvp", + "eval_trajectory", + ), + atol=1e-12, + rtol=1e-9, + ), + "spm_residual": ArtifactScenario( + name="spm_residual", + operations=( + "eval", + "jacobian_y", + "jacobian_p", + "jvp", + "eval_trajectory", + ), + atol=1e-9, + rtol=1e-7, + ), + "spme_residual": ArtifactScenario( + name="spme_residual", + operations=( + "eval", + "jacobian_y", + "jacobian_p", + "jvp", + "eval_trajectory", + ), + atol=1e-9, + rtol=1e-7, + ), + "dfn_residual": ArtifactScenario( + name="dfn_residual", + operations=( + "eval", + "jacobian_y", + "jacobian_p", + "jvp", + "eval_trajectory", + ), + atol=1e-9, + rtol=1e-7, + ), +} + + +def get_artifact_scenarios(selected: list[str] | None = None) -> list[ArtifactScenario]: + """The registered artifact scenarios, or just ``selected`` in registry order.""" + names = filter_names(list(_ARTIFACT_SCENARIOS), selected) + return [_ARTIFACT_SCENARIOS[name] for name in names] + + +def get_protocol_names() -> list[str]: + """Every registered protocol name, in registry order.""" + return list(_PROTOCOLS) + + +def _inference_parameter_values_for(protocol_builder): + """Wrap a protocol's builder so the fitted parameters come back symbolic. + + Wrapping rather than replacing keeps the protocol's own control law, which + is the whole reason a protocol row is worth timing. + """ + + def build(): + parameter_values = protocol_builder() + for pybamm_name, input_name in INFERENCE_INPUTS.items(): + parameter_values[pybamm_name] = pybamm.InputParameter(input_name) + for pybamm_name, input_name in INFERENCE_COMPLEMENTS.items(): + parameter_values[pybamm_name] = 1 - pybamm.InputParameter(input_name) + return parameter_values + + return build + + +def inference_nominal_values() -> dict[str, float]: + """Nominal value of each fitted parameter, read from the base set.""" + parameter_values = _base_parameter_values() + return { + input_name: float(parameter_values[pybamm_name]) + for pybamm_name, input_name in INFERENCE_INPUTS.items() + } + + +def get_inference_scenarios( + selected: list[str] | None = None, + protocols: list[str] | None = None, + *, + output_points: int = DEFAULT_OUTPUT_POINTS, +) -> list[SolverScenario]: + """Solver scenarios with the fitted parameters swapped for inputs.""" + return [ + replace( + scenario, + parameter_values_builder=_inference_parameter_values_for( + scenario.parameter_values_builder + ), + ) + for scenario in get_solver_scenarios( + selected, protocols, output_points=output_points + ) + ] + + +def get_solver_scenarios( + selected: list[str] | None = None, + protocols: list[str] | None = None, + *, + output_points: int = DEFAULT_OUTPUT_POINTS, +) -> list[SolverScenario]: + """The model × protocol cross product, model-major, in registry order. + + Parameters + ---------- + selected : list of str, optional + Model names; all registered models when omitted. + protocols : list of str, optional + Protocol names. ``None`` selects :data:`DEFAULT_PROTOCOLS`, keeping a + bare run identical to the constant-current baseline; an empty list + selects every protocol, as an empty ``selected`` does for models. + output_points : int + Output grid size for protocols that own their grid. + + Raises + ------ + ValueError + If ``output_points`` is below 2, or a model or protocol is unknown. + """ + if output_points < 2: + raise ValueError("output_points must be at least 2") + model_names = filter_names(list(_MODELS), selected) + protocol_names = filter_names( + list(_PROTOCOLS), + list(DEFAULT_PROTOCOLS) if protocols is None else protocols, + ) + return [ + SolverScenario( + name=model_name, + protocol=protocol_name, + model_factory=_MODELS[model_name], + model_options={}, + parameter_values_builder=_PROTOCOLS[protocol_name].build_parameter_values, + initial_soc=_PROTOCOLS[protocol_name].initial_soc, + plan=_PROTOCOLS[protocol_name].build_plan(output_points), + observed_output="Voltage [V]", + atol=1e-6, + rtol=1e-6, + ) + for model_name in model_names + for protocol_name in protocol_names + ] diff --git a/benchmarks/rust_observability/report.py b/benchmarks/rust_observability/report.py new file mode 100644 index 0000000000..3573f922f6 --- /dev/null +++ b/benchmarks/rust_observability/report.py @@ -0,0 +1,823 @@ +"""Terminal tables and JSON export for the observability suite's results. + +Rendering is kept out of the runners so a saved JSON run can be re-rendered, and +so table width can adapt to the terminal without touching measurement code. +""" + +from __future__ import annotations + +import math +import shutil +import textwrap +from dataclasses import asdict + +# Cap the free-text Reason column so a verbose backend error (e.g. the rust +# output-sensitivity SolverError) can't blow the table line width out. +_MAX_REASON_WIDTH = 60 +_DEFAULT_TABLE_WIDTH = 120 +# Each compact validation row's fixed width: its own columns and the separating +# spaces, with the identity, status and reason columns excluded. +_SOLVER_VALIDATION_COLUMNS = 80 +_SENSITIVITY_VALIDATION_COLUMNS = 74 +_INFERENCE_VALIDATION_COLUMNS = 63 +_BACKEND_COMPARISON_ORDER = { + "casadi_idaklu": 0, + "casadi_idaklu_aot": 1, + "rust_idaklu": 2, + "rust_diffsol": 3, +} + + +def _truncate(text: str, width: int) -> str: + return text if len(text) <= width else text[: width - 1] + "…" + + +def _fits(table: str, width: int) -> bool: + """Whether every line of a rendered table is inside ``width``. + + Measured rather than predicted from the column list, so adding a column + narrows the layout instead of overflowing it. + """ + return max(map(len, table.splitlines()), default=0) <= width + + +def _delta_text(comparison) -> str: + """A comparison's worst absolute difference, or ``-`` when it was not run.""" + return "-" if comparison is None else f"{comparison.max_abs_diff:.2e}" + + +def _points_text(requested_output_points: int) -> str: + """Requested output points, or ``-`` when the protocol owns its own grid.""" + return str(requested_output_points) if requested_output_points else "-" + + +def _reference_caption(results) -> str: + """One line naming what this lane's Δ columns were measured against.""" + tolerances = sorted( + {result.reference_tolerance for result in results if result.reference_tolerance} + ) + if not tolerances: + return ( + "Δ is the raw difference from casadi_idaklu at the scenario tolerance; " + "no converged reference was run." + ) + listed = ", ".join(f"{tolerance:g}" for tolerance in tolerances) + caption = ( + f"Δ is the error against a converged casadi_idaklu reference at " + f"atol=rtol={listed}. Base Δ is the raw difference from casadi_idaklu at " + "the scenario tolerance, gated on nothing." + ) + if any( + result.supported and result.reference_tolerance is None for result in results + ): + caption += ( + " Some rows had no reference available; for those, Δ already holds " + "that baseline difference and Base Δ is blank." + ) + return caption + + +def _wrap_caption(caption: str, width: int) -> str: + return "\n".join(textwrap.wrap(caption, width=max(width, 40))) + + +def _backend_comparison_key(backend: str) -> tuple[bool, int, str]: + output_only = backend.endswith("_out") + base_backend = backend.removesuffix("_out") + rank = _BACKEND_COMPARISON_ORDER.get(base_backend, len(_BACKEND_COMPARISON_ORDER)) + return output_only, rank, backend + + +def _ordered_backend_results(results) -> list: + results = list(results) + scenario_order: dict[str, int] = {} + protocol_order: dict[str, int] = {} + for result in results: + scenario_order.setdefault(result.scenario, len(scenario_order)) + protocol_order.setdefault(result.protocol, len(protocol_order)) + return sorted( + results, + key=lambda result: ( + scenario_order[result.scenario], + protocol_order[result.protocol], + *_backend_comparison_key(result.backend), + ), + ) + + +def render_artifact_table(results) -> str: + """Render the artifact lane as a fixed-width table.""" + # Run times are per-call kernel cost in microseconds; prep is one-time in ms. + headers = ( + f"{'Scenario':<16} {'Operation':<16} {'Rust Prep':>10} {'Rust Run µs':>11} " + f"{'CasADi Prep':>12} {'CasADi Run µs':>13} {'AOT Prep':>10} {'AOT Run µs':>11} " + f"{'Rust Spd':>9} {'AOT Spd':>8} {'Rust Abs':>10} {'AOT Abs':>10} " + f"{'Status':>8}" + ) + lines = [headers, "-" * len(headers)] + for result in results: + rust_run = result.candidate_timings.run_ms + casadi_run = result.baseline_timings.run_ms + rust_spd = casadi_run / rust_run if rust_run > 0 else float("inf") + aot_timings = result.aot_timings + aot_comparison = result.aot_comparison + aot_prep = aot_timings.prepare_ms if aot_timings else float("nan") + aot_run = aot_timings.run_ms if aot_timings else float("nan") + aot_spd = casadi_run / aot_run if aot_timings and aot_run > 0 else float("nan") + aot_abs = aot_comparison.max_abs_diff if aot_comparison else float("nan") + lines.append( + f"{result.scenario:<16} {result.operation:<16} " + f"{result.candidate_timings.prepare_ms:>10.3f} " + f"{result.candidate_timings.run_ms * 1000.0:>11.3f} " + f"{result.baseline_timings.prepare_ms:>12.3f} " + f"{result.baseline_timings.run_ms * 1000.0:>13.3f} " + f"{aot_prep:>10.3f} " + f"{aot_run * 1000.0:>11.3f} " + f"{rust_spd:>8.2f}x " + f"{aot_spd:>7.2f}x " + f"{result.comparison.max_abs_diff:>10.3e} " + f"{aot_abs:>10.3e} " + f"{result.status:>8}" + ) + return "\n".join(lines) + + +def _solver_cells(result) -> dict[str, str]: + state_abs = ( + result.state_comparison.max_abs_diff + if result.state_comparison + else float("nan") + ) + output_abs = ( + result.output_comparison.max_abs_diff + if result.output_comparison + else float("nan") + ) + trajectory = result.trajectory_comparison + telemetry = result.jacobian_telemetry + return { + "scenario": result.scenario, + "protocol": result.protocol, + "backend": result.backend, + "points": _points_text(result.requested_output_points), + "build": f"{result.timings.build_ms:.2f}", + "prepare": f"{result.timings.prepare_ms:.2f}", + "cold_startup": f"{result.timings.cold_startup_ms:.2f}", + "warm_set_up": f"{result.timings.warm_set_up_ms:.2f}", + "solve": f"{result.timings.solve_ms:.2f}", + "wall": f"{result.timings.wall_solve_ms:.2f}", + "integration": f"{result.timings.integration_ms:.2f}", + "observe": f"{result.timings.observe_ms:.2f}", + "e2e": f"{result.timings.e2e_ms:.2f}", + "coverage": f"{trajectory.coverage:.3f}" if trajectory else "-", + "final_time_diff": (f"{trajectory.final_time_diff:.2e}" if trajectory else "-"), + "colors": str(telemetry.n_colors) if telemetry else "-", + "dense_rows": str(telemetry.n_dense_rows) if telemetry else "-", + "dense_entries": str(telemetry.dense_row_entries) if telemetry else "-", + "dense_tape": str(telemetry.dense_row_tape_instructions) if telemetry else "-", + "state_abs": "-" if math.isnan(state_abs) else f"{state_abs:.2e}", + "output_abs": "-" if math.isnan(output_abs) else f"{output_abs:.2e}", + "base_delta": _delta_text(result.baseline_delta), + "status": result.status, + "reason": result.reason or "", + } + + +def _identity_widths( + rows: list[dict[str, str]], *, scenario_max: int, backend_max: int +) -> tuple[int, int, int]: + scenario_width = min( + scenario_max, max(len("Scenario"), *(len(row["scenario"]) for row in rows)) + ) + backend_width = min( + backend_max, max(len("Backend"), *(len(row["backend"]) for row in rows)) + ) + status_width = min(12, max(len("Status"), *(len(row["status"]) for row in rows))) + return scenario_width, backend_width, status_width + + +def _render_solver_wide(rows: list[dict[str, str]]) -> str: + scenario_width, backend_width, status_width = _identity_widths( + rows, scenario_max=18, backend_max=26 + ) + headers = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} " + f"{'Pts':>5} {'Build':>8} {'Prep':>8} {'Cold':>8} " + f"{'WarmSet':>8} {'Solve':>8} {'Wall':>8} {'Integr.':>8} " + f"{'Obs':>4} {'E2E':>8} {'Cover':>7} {'T End Δ':>9} {'Clr':>4} " + f"{'Dense':>5} {'DEntry':>6} {'Tape':>7} {'State Abs':>8} {'Output Abs':>8} " + f"{'Base Δ':>9} {'Status':>{status_width}} Reason" + ) + lines = [headers, "-" * len(headers)] + for row in rows: + reason = _truncate(row["reason"], _MAX_REASON_WIDTH) + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>5} {row['build']:>8} {row['prepare']:>8} " + f"{row['cold_startup']:>8} " + f"{row['warm_set_up']:>8} {row['solve']:>8} {row['wall']:>8} " + f"{row['integration']:>10} {row['observe']:>4} {row['e2e']:>8} " + f"{row['coverage']:>7} {row['final_time_diff']:>9} " + f"{row['colors']:>4} {row['dense_rows']:>5} " + f"{row['dense_entries']:>6} {row['dense_tape']:>7} " + f"{row['state_abs']:>8} {row['output_abs']:>8} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} " + f"{reason}" + ).rstrip() + ) + return "\n".join(lines) + + +def _compact_identity_widths( + rows: list[dict[str, str]], width: int, validation_columns: int +) -> tuple[int, int, int]: + """Identity widths that leave room for the timing *and* validation rows. + + ``validation_columns`` is the validation row's fixed width with the identity + columns, the status and the reason excluded. Budgeting on the timing row + alone let the validation row overflow whenever it was the wider of the two. + """ + scenario_width, backend_width, status_width = _identity_widths( + rows, scenario_max=18, backend_max=23 + ) + timing_fixed_width = max(77, 63 + status_width) + 14 + validation_fixed_width = validation_columns + status_width + len(" Reason") + fixed_width = max(timing_fixed_width, validation_fixed_width) + excess = scenario_width + backend_width + fixed_width - width + backend_reduction = min(max(excess, 0), backend_width - 16) + backend_width -= backend_reduction + excess -= backend_reduction + scenario_width -= min(max(excess, 0), scenario_width - 8) + return scenario_width, backend_width, status_width + + +def _render_solver_compact(rows: list[dict[str, str]], width: int) -> str: + scenario_width, backend_width, status_width = _compact_identity_widths( + rows, width, _SOLVER_VALIDATION_COLUMNS + ) + timing_header = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Pts':>4} " + f"{'Build':>7} {'Prep':>7} {'Cold':>7} {'Warm':>7} {'Solve':>7} " + f"{'Wall':>7} {'Int':>7} {'Obs':>6} {'E2E':>7}" + ) + lines = ["Timings (ms)", timing_header, "-" * len(timing_header)] + for row in rows: + lines.append( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>4} {row['build']:>7} {row['prepare']:>7} " + f"{row['cold_startup']:>7} {row['warm_set_up']:>7} " + f"{row['solve']:>7} {row['wall']:>7} " + f"{row['integration']:>7} {row['observe']:>6} {row['e2e']:>7}" + ) + + validation_prefix = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} " + f"{'Cover':>5} {'Δt':>8} {'Clr':>3} {'Rows':>4} {'Entry':>5} " + f"{'Tape':>7} {'State Δ':>8} {'Output Δ':>8} {'Base Δ':>9} " + f"{'Status':>{status_width}}" + ) + reason_width = max(6, min(_MAX_REASON_WIDTH, width - len(validation_prefix) - 1)) + validation_header = f"{validation_prefix} Reason" + lines.extend(["", "Validation", validation_header, "-" * len(validation_header)]) + for row in rows: + reason = _truncate(row["reason"], reason_width) + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['coverage']:>5} {row['final_time_diff']:>8} " + f"{row['colors']:>3} {row['dense_rows']:>4} " + f"{row['dense_entries']:>5} {row['dense_tape']:>7} " + f"{row['state_abs']:>8} {row['output_abs']:>8} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} {reason:<{reason_width}}" + ).rstrip() + ) + return "\n".join(lines) + + +def _render_solver_stacked(rows: list[dict[str, str]], width: int) -> str: + lines = ["One block per backend; timings are milliseconds."] + for row in rows: + if len(lines) > 1: + lines.append("") + heading = ( + f"{row['scenario']} | {row['protocol']} | {row['backend']} | " + f"{row['points']} pts | {row['status']}" + ) + lines.append(_truncate(heading, width)) + details = ( + "Timing: " + f"build {row['build']}, prepare {row['prepare']}, " + f"cold {row['cold_startup']}, " + f"warm {row['warm_set_up']}, solve {row['solve']}, wall {row['wall']}, " + f"integration {row['integration']}, observe {row['observe']}, " + f"e2e {row['e2e']}" + ) + validation = ( + "Validation: " + f"cover {row['coverage']}, Δt {row['final_time_diff']}, " + f"colors {row['colors']}, dense rows/entries/sweeps " + f"{row['dense_rows']}/{row['dense_entries']}/{row['dense_tape']}, " + f"state Δ {row['state_abs']}, output Δ {row['output_abs']}, " + f"base Δ {row['base_delta']}" + ) + lines.extend(textwrap.wrap(details, width=width, subsequent_indent=" ")) + lines.extend(textwrap.wrap(validation, width=width, subsequent_indent=" ")) + if row["reason"]: + lines.extend( + textwrap.wrap( + f"Reason: {row['reason']}", width=width, subsequent_indent=" " + ) + ) + return "\n".join(lines) + + +def _sensitivity_cells(result) -> dict[str, str]: + state_abs = ( + result.state_sens_comparison.max_abs_diff + if result.state_sens_comparison + else float("nan") + ) + output_abs = ( + result.output_sens_comparison.max_abs_diff + if result.output_sens_comparison + else float("nan") + ) + trajectory = result.trajectory_comparison + return { + "scenario": result.scenario, + "protocol": result.protocol, + "backend": result.backend, + "points": _points_text(result.requested_output_points), + "build": f"{result.timings.build_ms:.2f}", + "prepare": f"{result.timings.prepare_ms:.2f}", + "cold_startup": f"{result.timings.cold_startup_ms:.2f}", + "warm_set_up": f"{result.timings.warm_set_up_ms:.2f}", + "solve": f"{result.timings.solve_ms:.2f}", + "wall": f"{result.timings.wall_solve_ms:.2f}", + "integration": f"{result.timings.integration_ms:.2f}", + "observe": f"{result.timings.observe_ms:.2f}", + "e2e": f"{result.timings.e2e_ms:.2f}", + "coverage": f"{trajectory.coverage:.3f}" if trajectory else "-", + "final_time_diff": (f"{trajectory.final_time_diff:.2e}" if trajectory else "-"), + "sensitivity_parameters": ",".join(result.sensitivity_parameters) or "-", + "state_abs": "-" if math.isnan(state_abs) else f"{state_abs:.2e}", + "output_abs": "-" if math.isnan(output_abs) else f"{output_abs:.2e}", + "base_delta": _delta_text(result.baseline_delta), + "status": result.status, + "reason": result.reason or "", + } + + +def _render_sensitivity_wide(rows: list[dict[str, str]]) -> str: + scenario_width, backend_width, status_width = _identity_widths( + rows, scenario_max=18, backend_max=26 + ) + headers = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} " + f"{'Pts':>5} {'Build':>8} {'Prep':>8} {'Cold':>8} " + f"{'WarmSet':>8} {'Solve':>8} {'Wall':>8} {'Integr.':>10} " + f"{'Obs':>8} {'E2E':>8} {'Cover':>7} {'T End Δ':>9} {'Params':<12} " + f"{'State Sens':>11} {'Output Sens':>12} {'Base Δ':>9} " + f"{'Status':>{status_width}} Reason" + ) + lines = [headers, "-" * len(headers)] + for row in rows: + reason = _truncate(row["reason"], _MAX_REASON_WIDTH) + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>5} {row['build']:>8} {row['prepare']:>8} " + f"{row['cold_startup']:>8} {row['warm_set_up']:>8} " + f"{row['solve']:>8} {row['wall']:>8} " + f"{row['integration']:>10} {row['observe']:>8} {row['e2e']:>8} " + f"{row['coverage']:>7} {row['final_time_diff']:>9} " + f"{_truncate(row['sensitivity_parameters'], 12):<12} " + f"{row['state_abs']:>11} {row['output_abs']:>12} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} {reason}" + ).rstrip() + ) + return "\n".join(lines) + + +def _render_sensitivity_compact(rows: list[dict[str, str]], width: int) -> str: + scenario_width, backend_width, status_width = _compact_identity_widths( + rows, width, _SENSITIVITY_VALIDATION_COLUMNS + ) + timing_header = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Pts':>4} " + f"{'Build':>7} {'Prep':>7} {'Cold':>7} {'Warm':>7} {'Solve':>7} " + f"{'Wall':>7} {'Int':>7} {'Obs':>6} {'E2E':>7}" + ) + lines = ["Timings (ms)", timing_header, "-" * len(timing_header)] + for row in rows: + lines.append( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>4} {row['build']:>7} {row['prepare']:>7} " + f"{row['cold_startup']:>7} {row['warm_set_up']:>7} " + f"{row['solve']:>7} {row['wall']:>7} {row['integration']:>7} " + f"{row['observe']:>6} {row['e2e']:>7}" + ) + + validation_prefix = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Params':<12} " + f"{'Cover':>5} {'Δt':>8} {'State Δ':>9} {'Output Δ':>9} {'Base Δ':>9} " + f"{'Status':>{status_width}}" + ) + reason_width = max(6, min(_MAX_REASON_WIDTH, width - len(validation_prefix) - 1)) + validation_header = f"{validation_prefix} Reason" + lines.extend(["", "Validation", validation_header, "-" * len(validation_header)]) + for row in rows: + reason = _truncate(row["reason"], reason_width) + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{_truncate(row['sensitivity_parameters'], 12):<12} " + f"{row['coverage']:>5} {row['final_time_diff']:>8} " + f"{row['state_abs']:>9} {row['output_abs']:>9} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} {reason:<{reason_width}}" + ).rstrip() + ) + return "\n".join(lines) + + +def _render_sensitivity_stacked(rows: list[dict[str, str]], width: int) -> str: + lines = ["One block per backend; timings are milliseconds."] + for row in rows: + if len(lines) > 1: + lines.append("") + heading = ( + f"{row['scenario']} | {row['protocol']} | {row['backend']} | " + f"{row['points']} pts | {row['status']}" + ) + lines.append(_truncate(heading, width)) + details = ( + "Timing: " + f"build {row['build']}, prepare {row['prepare']}, " + f"cold {row['cold_startup']}, warm {row['warm_set_up']}, " + f"solve {row['solve']}, wall {row['wall']}, " + f"integration {row['integration']}, observe {row['observe']}, " + f"e2e {row['e2e']}" + ) + validation = ( + "Validation: " + f"cover {row['coverage']}, Δt {row['final_time_diff']}, " + f"parameters {row['sensitivity_parameters']}, " + f"state sensitivity Δ {row['state_abs']}, " + f"output sensitivity Δ {row['output_abs']}, " + f"base Δ {row['base_delta']}" + ) + lines.extend(textwrap.wrap(details, width=width, subsequent_indent=" ")) + lines.extend(textwrap.wrap(validation, width=width, subsequent_indent=" ")) + if row["reason"]: + lines.extend( + textwrap.wrap( + f"Reason: {row['reason']}", width=width, subsequent_indent=" " + ) + ) + return "\n".join(lines) + + +def _cache_status_text(statuses: tuple[str, ...]) -> str: + if not statuses: + return "-" + if len(set(statuses)) == 1: + suffix = f"x{len(statuses)}" if len(statuses) > 1 else "" + return f"{statuses[0]}{suffix}" + return "+".join(statuses) + + +def _render_aot_profile(results, width: int) -> str: + profiled = [result for result in results if result.aot_profile is not None] + if not profiled: + return "" + scenario_width = min( + 16, max(len("Scenario"), *(len(result.scenario) for result in profiled)) + ) + backend_width = min( + 23, max(len("Backend"), *(len(result.backend) for result in profiled)) + ) + fixed_width = 98 + excess = scenario_width + backend_width + fixed_width - width + backend_reduction = min(max(excess, 0), backend_width - 16) + backend_width -= backend_reduction + excess -= backend_reduction + scenario_width -= min(max(excess, 0), scenario_width - 8) + + title = "AOT profile (isolated cache; phase timings in ms)" + if scenario_width + backend_width + fixed_width > width: + lines = [title] + for result in profiled: + profile = result.aot_profile + heading = f"{result.scenario} | {result.protocol} | {result.backend}" + details = ( + f"fresh {_cache_status_text(profile.fresh_cache_statuses)}, " + f"disk {_cache_status_text(profile.disk_cache_statuses)}, " + f"codegen {profile.codegen_ms:.2f}, compiler {profile.compiler_ms:.2f}, " + f"fresh load {profile.fresh_load_ms:.2f}, " + f"disk load {profile.disk_load_ms:.2f}, " + f"disk prep {profile.disk_prepare_ms:.2f}, " + f"disk cold {profile.disk_cold_startup_ms:.2f}, " + f"library {profile.library_size_bytes / 1024**2:.2f} MiB, " + f"verified {'yes' if profile.verified else 'no'}" + ) + lines.extend(["", _truncate(heading, width)]) + lines.extend(textwrap.wrap(details, width=width, subsequent_indent=" ")) + return "\n".join(lines) + + header = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} " + f"{'Fresh':>7} {'Disk':>7} {'Gen':>8} {'Compiler':>9} " + f"{'FLoad':>7} {'DLoad':>7} {'DPrep':>8} {'DCold':>8} " + f"{'MiB':>6} {'Status':>6}" + ) + lines = [title, header, "-" * len(header)] + for result in profiled: + profile = result.aot_profile + lines.append( + f"{_truncate(result.scenario, scenario_width):<{scenario_width}} " + f"{_truncate(result.backend, backend_width):<{backend_width}} " + f"{_truncate(result.protocol, 13):<13} " + f"{_cache_status_text(profile.fresh_cache_statuses):>7} " + f"{_cache_status_text(profile.disk_cache_statuses):>7} " + f"{profile.codegen_ms:>8.2f} {profile.compiler_ms:>9.2f} " + f"{profile.fresh_load_ms:>7.2f} {profile.disk_load_ms:>7.2f} " + f"{profile.disk_prepare_ms:>8.2f} " + f"{profile.disk_cold_startup_ms:>8.2f} " + f"{profile.library_size_bytes / 1024**2:>6.2f} " + f"{'pass' if profile.verified else 'fail':>6}" + ) + return "\n".join(lines) + + +def _render_lane( + results, + width: int | None, + *, + cells, + wide, + compact, + stacked, + empty_message: str, + title: str = "", + with_aot: bool = True, +) -> str: + """Render one lane, narrowing to a compact then a stacked layout as needed. + + ``width`` defaults to the detected terminal width. ``title`` labels the wide + layout only; the narrower layouts carry their own headings. + """ + ordered = _ordered_backend_results(results) + rows = [cells(result) for result in ordered] + if not rows: + return empty_message + table_width = ( + width or shutil.get_terminal_size(fallback=(_DEFAULT_TABLE_WIDTH, 24)).columns + ) + wide_table = wide(rows) + compact_table = compact(rows, table_width) + if _fits(wide_table, table_width): + main_table = f"{title}\n{wide_table}" if title else wide_table + elif _fits(compact_table, table_width): + main_table = compact_table + else: + main_table = stacked(rows, max(table_width, 40)) + tables = [_wrap_caption(_reference_caption(ordered), table_width), main_table] + if with_aot and (aot_table := _render_aot_profile(ordered, max(table_width, 40))): + tables.append(aot_table) + return "\n\n".join(tables) + + +def render_solver_table(results, *, width: int | None = None) -> str: + """Render the solver lane, narrowing to a compact layout if ``width`` is tight. + + ``width`` defaults to the detected terminal width. + """ + return _render_lane( + results, + width, + cells=_solver_cells, + wide=_render_solver_wide, + compact=_render_solver_compact, + stacked=_render_solver_stacked, + empty_message="No solver results.", + ) + + +def render_sensitivity_table(results, *, width: int | None = None) -> str: + """Render the sensitivity lane, with the same width handling as the solver table. + + Solves differentiate the parameters in ``runners.SENSITIVITY_INPUTS``. + """ + # State/Output Sens are max abs diffs of the stacked "all" sensitivity block + # against the CasADi-IDAKLU baseline. + return _render_lane( + results, + width, + cells=_sensitivity_cells, + wide=_render_sensitivity_wide, + compact=_render_sensitivity_compact, + stacked=_render_sensitivity_stacked, + empty_message="No sensitivity results.", + ) + + +def _inference_cells(result) -> dict[str, str]: + if result.eval_samples_ms: + spread = f"{result.eval_p10_ms:.2f}-{result.eval_p90_ms:.2f}" + eval_p50 = f"{result.eval_median_ms:.2f}" + solve = f"{result.solve_median_ms:.2f}" + observe = f"{result.observe_median_ms:.2f}" + else: + spread = eval_p50 = solve = observe = "-" + trajectory = result.trajectory_comparison + return { + "scenario": result.scenario, + "protocol": result.protocol, + "backend": result.backend, + "points": _points_text(result.requested_output_points), + "build": f"{result.build_ms:.2f}", + "setup": f"{result.setup_ms:.2f}", + "cold_observe": f"{result.cold_observe_ms:.2f}", + "aot": result.aot_cache_status, + "eval_p50": eval_p50, + "spread": spread, + "solve": solve, + "observe": observe, + "coverage": f"{trajectory.coverage:.3f}" if trajectory else "-", + "final_time_diff": (f"{trajectory.final_time_diff:.2e}" if trajectory else "-"), + "output_abs": _delta_text(result.output_comparison), + "sensitivity_abs": _delta_text(result.sensitivity_comparison), + "base_delta": _delta_text(result.baseline_delta), + "status": result.status, + "reason": result.reason or "", + } + + +def _render_inference_wide(rows: list[dict[str, str]]) -> str: + scenario_width, backend_width, status_width = _identity_widths( + rows, scenario_max=18, backend_max=20 + ) + header = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Pts':>5} {'Build':>8} {'Setup':>8} {'ColdObs':>8} " + f"{'AOT':>7} {'Eval p50':>9} {'p10-p90':>17} {'Solve':>8} {'Obs':>8} " + f"{'Cover':>6} {'T End Δ':>9} {'Output Δ':>10} {'Sens Δ':>10} " + f"{'Base Δ':>9} {'Status':>{status_width}} Reason" + ) + lines = [header, "-" * len(header)] + for row in rows: + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>5} {row['build']:>8} {row['setup']:>8} " + f"{row['cold_observe']:>8} {row['aot']:>7} {row['eval_p50']:>9} " + f"{row['spread']:>17} {row['solve']:>8} {row['observe']:>8} " + f"{row['coverage']:>6} {row['final_time_diff']:>9} " + f"{row['output_abs']:>10} {row['sensitivity_abs']:>10} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} " + f"{_truncate(row['reason'], _MAX_REASON_WIDTH)}" + ).rstrip() + ) + return "\n".join(lines) + + +def _render_inference_compact(rows: list[dict[str, str]], width: int) -> str: + scenario_width, backend_width, status_width = _compact_identity_widths( + rows, width, _INFERENCE_VALIDATION_COLUMNS + ) + timing_header = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Pts':>4} {'Build':>7} {'Setup':>7} {'ColdObs':>7} " + f"{'AOT':>6} {'Eval p50':>8} {'p10-p90':>15} {'Solve':>7} {'Obs':>6}" + ) + lines = ["Timings (ms)", timing_header, "-" * len(timing_header)] + for row in rows: + lines.append( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['points']:>4} {row['build']:>7} {row['setup']:>7} " + f"{row['cold_observe']:>7} {row['aot']:>6} {row['eval_p50']:>8} " + f"{row['spread']:>15} {row['solve']:>7} {row['observe']:>6}" + ) + + validation_prefix = ( + f"{'Scenario':<{scenario_width}} {'Backend':<{backend_width}} " + f"{'Protocol':<13} {'Cover':>5} {'Δt':>8} {'Output Δ':>10} " + f"{'Sens Δ':>10} {'Base Δ':>9} {'Status':>{status_width}}" + ) + reason_width = max(6, min(_MAX_REASON_WIDTH, width - len(validation_prefix) - 1)) + validation_header = f"{validation_prefix} Reason" + lines.extend(["", "Validation", validation_header, "-" * len(validation_header)]) + for row in rows: + lines.append( + ( + f"{_truncate(row['scenario'], scenario_width):<{scenario_width}} " + f"{_truncate(row['backend'], backend_width):<{backend_width}} " + f"{_truncate(row['protocol'], 13):<13} " + f"{row['coverage']:>5} {row['final_time_diff']:>8} " + f"{row['output_abs']:>10} {row['sensitivity_abs']:>10} " + f"{row['base_delta']:>9} " + f"{row['status']:>{status_width}} " + f"{_truncate(row['reason'], reason_width):<{reason_width}}" + ).rstrip() + ) + return "\n".join(lines) + + +def _render_inference_stacked(rows: list[dict[str, str]], width: int) -> str: + lines = ["One block per backend; timings are milliseconds."] + for row in rows: + if len(lines) > 1: + lines.append("") + heading = ( + f"{row['scenario']} | {row['protocol']} | {row['backend']} | " + f"{row['points']} pts | {row['status']}" + ) + lines.append(_truncate(heading, width)) + details = ( + "Timing: " + f"build {row['build']}, setup {row['setup']}, " + f"cold observe {row['cold_observe']}, aot {row['aot']}, " + f"eval p50 {row['eval_p50']}, p10-p90 {row['spread']}, " + f"solve {row['solve']}, observe {row['observe']}" + ) + validation = ( + "Validation: " + f"cover {row['coverage']}, Δt {row['final_time_diff']}, " + f"output Δ {row['output_abs']}, sensitivity Δ {row['sensitivity_abs']}, " + f"base Δ {row['base_delta']}" + ) + lines.extend(textwrap.wrap(details, width=width, subsequent_indent=" ")) + lines.extend(textwrap.wrap(validation, width=width, subsequent_indent=" ")) + if row["reason"]: + lines.extend( + textwrap.wrap( + f"Reason: {row['reason']}", width=width, subsequent_indent=" " + ) + ) + return "\n".join(lines) + + +def render_inference_table(results, *, width: int | None = None) -> str: + """Render the inference lane: one-time costs beside per-evaluation costs. + + Same width handling as the solver and sensitivity lanes; ``width`` defaults + to the detected terminal width. + """ + return _render_lane( + results, + width, + cells=_inference_cells, + wide=_render_inference_wide, + compact=_render_inference_compact, + stacked=_render_inference_stacked, + empty_message="No inference results.", + title="Per-evaluation costs under changing inputs (ms)", + with_aot=False, + ) + + +def suite_to_jsonable(lane: str, results, *, metadata: dict | None = None) -> dict: + """Convert one lane's results to JSON-serialisable form for ``--json``. + + Solver and sensitivity results are emitted in backend-comparison order so a + diff between two saved runs lines up row for row. + """ + ordered_results = ( + _ordered_backend_results(results) + if lane in {"solver", "sensitivity", "inference"} + else list(results) + ) + return { + "lane": lane, + "metadata": metadata or {}, + "results": [asdict(result) for result in ordered_results], + } diff --git a/benchmarks/rust_observability/runners.py b/benchmarks/rust_observability/runners.py new file mode 100644 index 0000000000..ba43dc3787 --- /dev/null +++ b/benchmarks/rust_observability/runners.py @@ -0,0 +1,2626 @@ +"""Measurement lanes for the Rust-vs-CasADi observability suite. + +Each lane runs the same scenario across the backend matrix, samples the phase +timings, and pairs every candidate against the CasADi baseline so a regression +shows up as a comparison status rather than as a raw number a reader has to judge. +Backend order is shuffled per scenario, because a fixed order lets machine warm-up +flatter whichever backend runs last. +""" + +from __future__ import annotations + +import math +import multiprocessing +import os +import random +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass, replace +from time import perf_counter + +import numpy as np + +import pybamm +from benchmarks.rust_observability.registry import ( + DEFAULT_OUTPUT_POINTS, + INFERENCE_INPUTS, + INFERENCE_SPREADS, + ArtifactScenario, + SolverScenario, + get_solver_scenarios, + inference_nominal_values, +) + +# The row every candidate's raw difference is reported against. +_BASELINE_CASE = ("casadi_idaklu", False) + +# Two backends at one tolerance differ by their mutual error, not their accuracy; +# same-integrator rows also cancel the error they share. Hence a converged oracle. +_REFERENCE_BACKEND = "casadi_idaklu" +DEFAULT_REFERENCE_TOLERANCE = 1e-10 + +# Tolerance bounds one step's local error; global error accumulates it over +# thousands. Measured: casadi_idaklu itself lands ~18 tolerance units out on states. +_REFERENCE_ACCURACY_HEADROOM = 100.0 + +# A reference must stay clear of what it judges; two decades holds its own error +# to ~1% of the candidate's. +_MIN_REFERENCE_DECADES = 2 + +_BACKEND_CASES = ( + _BASELINE_CASE, + ("casadi_idaklu", True), + ("casadi_idaklu_aot", False), + ("casadi_idaklu_aot", True), + ("rust_idaklu", False), + ("rust_idaklu", True), + ("rust_diffsol", False), + ("rust_diffsol", True), +) + + +def backend_cases(include_aot: bool) -> tuple[tuple[str, bool], ...]: + """The candidate backend matrix, with the AOT rows optional.""" + return tuple( + case for case in _BACKEND_CASES if include_aot or case[0] != "casadi_idaklu_aot" + ) + + +def resolve_reference_tolerance(scenarios, tolerance: float | None) -> float | None: + """Validate the converged reference's tolerance, or ``None`` when disabled. + + Parameters + ---------- + scenarios : list of SolverScenario + The scenarios the reference will be solved for. + tolerance : float or None + Requested reference tolerance; ``0`` or ``None`` disables the reference + and reverts every lane to comparing against the baseline backend. + + Returns + ------- + float or None + The tolerance to solve the reference at, or ``None`` when disabled. + + Raises + ------ + ValueError + If the tolerance is negative, or sits within + :data:`_MIN_REFERENCE_DECADES` of the loosest scenario it would judge -- + a reference no better than the candidate measures the gap between two + approximations, which is the artifact the reference exists to remove. + """ + if not tolerance: + return None + if tolerance < 0.0: + raise ValueError("reference_tolerance must be non-negative") + loosest = min( + (min(scenario.atol, scenario.rtol) for scenario in scenarios), default=1.0 + ) + ceiling = loosest * 10.0**-_MIN_REFERENCE_DECADES + if tolerance > ceiling: + raise ValueError( + f"reference_tolerance {tolerance:g} is not at least " + f"{_MIN_REFERENCE_DECADES} decades tighter than the scenario " + f"tolerance {loosest:g}" + ) + return tolerance + + +def _reference_ladder(scenario, tolerance: float) -> list[float]: + """Reference tolerances to try, loosening a decade at a time. + + A converged solve is not reachable everywhere -- DFN under a ramping current + fails IDA's error test at ``t = 0`` below 1e-9 -- and the loosest reference + that still clears the scenario by :data:`_MIN_REFERENCE_DECADES` is worth + more than no reference at all. + """ + ceiling = min(scenario.atol, scenario.rtol) * 10.0**-_MIN_REFERENCE_DECADES + rungs = math.floor(math.log10(ceiling / tolerance) + 1e-9) + # Rebuilt from the exponent rather than multiplied up, so a decade stays a decade. + exponent = math.log10(tolerance) + return [10.0 ** (exponent + step) for step in range(max(int(rungs), 0) + 1)] + + +def _resolve_reference(scenario, tolerance: float | None, solve, *, what: str): + """Run ``solve(attempt)`` at the tightest reference tolerance that converges. + + Returns the tolerance used and the solve's value, or ``(None, None)`` when + nothing in the ladder converges, so the caller falls back to the baseline. + """ + if tolerance is None: + return None, None + failure = None + for attempt in _reference_ladder(scenario, tolerance): + try: + return attempt, solve(attempt) + except Exception as exc: # noqa: BLE001 + failure = exc + pybamm.logger.warning( + f"{scenario.name}/{scenario.protocol}: no converged {what} from " + f"{tolerance:g} upwards ({_short_reason(failure)}); comparing against " + "the baseline backend instead." + ) + return None, None + + +def _reference_tolerances(scenario) -> tuple[float, float]: + """Allowance for judging a value against the converged reference.""" + return ( + _REFERENCE_ACCURACY_HEADROOM * scenario.atol, + _REFERENCE_ACCURACY_HEADROOM * scenario.rtol, + ) + + +def _reference_solve( + scenario, + tolerance: float, + *, + inputs: dict[str, float] | None = None, + parameter_values=None, + sensitivity_parameters: list[str] | None = None, +): + """One converged solve of ``scenario``, the answer every row is judged against. + + Always full state and never ``output_variables``-restricted, so one solve + serves the state, output and sensitivity comparisons alike. + """ + solver = _make_solver( + _REFERENCE_BACKEND, atol=tolerance, rtol=tolerance, output_variables=None + ) + simulation = _build_simulation( + scenario, solver, _REFERENCE_BACKEND, parameter_values=parameter_values + ) + _build_and_time(simulation, scenario, inputs=inputs) + extra: dict = {} + if inputs is not None: + extra["inputs"] = inputs + if sensitivity_parameters: + extra["calculate_sensitivities"] = list(sensitivity_parameters) + return simulation.solve(**_solve_kwargs(scenario, extra)) + + +def _worst_comparison(summaries) -> ComparisonSummary | None: + """The summary that came closest to, or furthest past, its own tolerance.""" + present = [summary for summary in summaries if summary is not None] + if not present: + return None + return max(present, key=lambda summary: summary.max_normalized_error) + + +@dataclass(frozen=True) +class PhaseTiming: + """Milliseconds attributed to each phase of one measured run.""" + + build_ms: float = 0.0 + prepare_ms: float = 0.0 + cold_startup_ms: float = 0.0 + run_ms: float = 0.0 + set_up_ms: float = 0.0 + warm_set_up_ms: float = 0.0 + solve_ms: float = 0.0 + wall_solve_ms: float = 0.0 + integration_ms: float = 0.0 + observe_ms: float = 0.0 + e2e_ms: float = 0.0 + + +@dataclass(frozen=True) +class ComparisonSummary: + """Worst-case agreement between a candidate backend and the baseline. + + ``max_normalized_error`` is the absolute difference over the tolerance the + status is judged against, so it crosses 1.0 exactly when the status flips. + Rank comparisons by it rather than by ``max_abs_diff``, whose admissible + size varies with the baseline magnitude. + """ + + max_abs_diff: float + max_rel_diff: float + max_normalized_error: float + + @property + def status(self) -> str: + """``"pass"`` while the normalised error is inside the tolerance.""" + return "pass" if self.max_normalized_error <= 1.0 else "warn" + + +# Shapes that cannot be lined up at all: reported as an infinite miss rather than +# raised, so one unsupported backend does not abandon the rest of the suite. +_UNCOMPARABLE = ComparisonSummary( + max_abs_diff=float("inf"), + max_rel_diff=float("inf"), + max_normalized_error=float("inf"), +) + + +@dataclass(frozen=True) +class TimingSamples: + """Per-repeat samples behind the reported medians, kept for spread checks.""" + + warm_set_up_ms: tuple[float, ...] = () + solve_ms: tuple[float, ...] = () + wall_solve_ms: tuple[float, ...] = () + integration_ms: tuple[float, ...] = () + observe_ms: tuple[float, ...] = () + e2e_ms: tuple[float, ...] = () + + +@dataclass(frozen=True) +class AotProfile: + """What CasADi's ahead-of-time path cost, split by cold and warm disk cache.""" + + fresh_cache_statuses: tuple[str, ...] + disk_cache_statuses: tuple[str, ...] + codegen_ms: float + compiler_ms: float + fresh_load_ms: float + disk_load_ms: float + fresh_total_ms: float + disk_total_ms: float + disk_prepare_ms: float + disk_cold_startup_ms: float + library_size_bytes: int + verified: bool + + +@dataclass(frozen=True) +class TrajectorySummary: + """How far two backends agreed on the time grid they actually returned. + + Adaptive steppers stop at different points, so timings are only comparable + where the grids overlap; ``coverage`` is that overlap. + """ + + baseline_points: int + candidate_points: int + common_points: int + coverage: float + max_time_diff: float + final_time_diff: float + baseline_termination: str + candidate_termination: str + status: str + + +@dataclass(frozen=True) +class JacobianTelemetry: + """Compile-time Jacobian metrics read back from the Rust model.""" + + strategy: str + n_colors: int + nnz: int + n_dense_rows: int + dense_row_entries: int + dense_row_tape_instructions: int + split_eval_primal_instructions: int | None + split_eval_total_instructions: int | None + split_eval_raw_instructions: int | None + branch_block_lens: tuple[int, ...] + + +@dataclass(frozen=True) +class ArtifactResult: + """One artifact operation timed on baseline, candidate and the AOT backend.""" + + scenario: str + operation: str + baseline_backend: str + candidate_backend: str + baseline_timings: PhaseTiming + candidate_timings: PhaseTiming + comparison: ComparisonSummary + baseline_run_samples_ms: tuple[float, ...] = () + candidate_run_samples_ms: tuple[float, ...] = () + aot_run_samples_ms: tuple[float, ...] = () + # Third backend: AOT-compiled CasADi (same kernels as baseline, lowered to a + # native shared lib). ``aot_comparison`` is AOT vs the CasADi VM baseline. + aot_backend: str | None = None + aot_timings: PhaseTiming | None = None + aot_comparison: ComparisonSummary | None = None + + @property + def status(self) -> str: + """Agreement status of this operation against the baseline.""" + return self.comparison.status + + +def _comparison_status(supported: bool, summaries) -> str: + """Worst status across the comparisons that ran. + + ``"unsupported"`` when the backend could not run the scenario at all and + ``"baseline"`` when nothing was compared, so neither reads as a pass. + """ + if not supported: + return "unsupported" + statuses = [summary.status for summary in summaries if summary is not None] + if not statuses: + return "baseline" + return "pass" if all(status == "pass" for status in statuses) else "warn" + + +@dataclass(frozen=True) +class SolverResult: + """One scenario solved on every backend, with agreement on states and outputs. + + The comparisons are against the converged reference when + ``reference_tolerance`` is set, and against the baseline backend otherwise. + ``baseline_delta`` carries the raw cross-backend difference either way. + """ + + scenario: str + backend: str + timings: PhaseTiming + requested_output_points: int + protocol: str = "cc_discharge" + timing_samples: TimingSamples = TimingSamples() + state_comparison: ComparisonSummary | None = None + output_comparison: ComparisonSummary | None = None + trajectory_comparison: TrajectorySummary | None = None + reference_tolerance: float | None = None + baseline_delta: ComparisonSummary | None = None + jacobian_telemetry: JacobianTelemetry | None = None + aot_profile: AotProfile | None = None + supported: bool = True + reason: str | None = None + + @property + def status(self) -> str: + """Worst status across the state, output and trajectory comparisons.""" + return _comparison_status( + self.supported, + ( + self.state_comparison, + self.output_comparison, + self.trajectory_comparison, + ), + ) + + +@dataclass(frozen=True) +class SensitivityResult: + """One scenario solved with forward sensitivities on every backend. + + Same comparison target convention as :class:`SolverResult`. + """ + + scenario: str + backend: str + timings: PhaseTiming + requested_output_points: int + protocol: str = "cc_discharge" + timing_samples: TimingSamples = TimingSamples() + # The parameters actually differentiated, which a protocol can narrow. + sensitivity_parameters: tuple[str, ...] = () + state_sens_comparison: ComparisonSummary | None = None + output_sens_comparison: ComparisonSummary | None = None + trajectory_comparison: TrajectorySummary | None = None + reference_tolerance: float | None = None + baseline_delta: ComparisonSummary | None = None + aot_profile: AotProfile | None = None + supported: bool = True + reason: str | None = None + + @property + def status(self) -> str: + """Worst status across the sensitivity and trajectory comparisons.""" + return _comparison_status( + self.supported, + ( + self.state_sens_comparison, + self.output_sens_comparison, + self.trajectory_comparison, + ), + ) + + +@dataclass(frozen=True) +class InferenceResult: + """One backend's cost per likelihood evaluation under changing inputs. + + Same comparison target convention as :class:`SolverResult`. + """ + + scenario: str + protocol: str + backend: str + build_ms: float + setup_ms: float + aot_cache_status: str + eval_samples_ms: tuple[float, ...] + solve_samples_ms: tuple[float, ...] + observe_samples_ms: tuple[float, ...] + requested_output_points: int + cold_observe_ms: float = 0.0 + output_comparison: ComparisonSummary | None = None + sensitivity_comparison: ComparisonSummary | None = None + trajectory_comparison: TrajectorySummary | None = None + reference_tolerance: float | None = None + baseline_delta: ComparisonSummary | None = None + supported: bool = True + reason: str | None = None + + @property + def eval_median_ms(self) -> float: + """Median wall-clock cost of one likelihood evaluation.""" + return float(np.median(self.eval_samples_ms)) + + @property + def eval_p10_ms(self) -> float: + """10th-percentile evaluation cost, the fast end of the spread.""" + return float(np.percentile(self.eval_samples_ms, 10)) + + @property + def eval_p90_ms(self) -> float: + """90th-percentile evaluation cost, the slow end of the spread.""" + return float(np.percentile(self.eval_samples_ms, 90)) + + @property + def solve_median_ms(self) -> float: + """Median time inside ``Simulation.solve``.""" + return float(np.median(self.solve_samples_ms)) + + @property + def observe_median_ms(self) -> float: + """Median time materialising the observed output.""" + return float(np.median(self.observe_samples_ms)) + + @property + def status(self) -> str: + """Worst status across the value, gradient and trajectory comparisons.""" + return _comparison_status( + self.supported, + ( + self.output_comparison, + self.sensitivity_comparison, + self.trajectory_comparison, + ), + ) + + +@dataclass +class _PreparedArtifactCase: + baseline_backend: str + candidate_backend: str + baseline_timings: PhaseTiming + candidate_timings: PhaseTiming + # operation -> (baseline call, candidate call, aot call, materializer). Kernel + # calls return native output (timed); the materializer densifies it (untimed). + operations: dict[str, tuple[callable, callable, callable, callable]] + atol: float + rtol: float + aot_backend: str = "casadi_aot" + aot_timings: PhaseTiming = PhaseTiming() + + +def summarize_diff( + baseline, + candidate, + *, + atol: float, + rtol: float, +) -> ComparisonSummary: + """Compare two result arrays under the scenario's own tolerances. + + A shape mismatch is reported as infinite difference rather than raised, so one + unsupported backend does not abandon the rest of the suite. + """ + baseline_arr = _as_dense(baseline) + candidate_arr = _as_dense(candidate) + if baseline_arr.shape != candidate_arr.shape: + return _UNCOMPARABLE + abs_diff = np.abs(candidate_arr - baseline_arr) + denom = np.maximum(np.abs(baseline_arr), atol) + rel_diff = abs_diff / denom + # The np.allclose tolerance, so the reported error and the status agree. + normalized = abs_diff / (atol + rtol * np.abs(baseline_arr)) + max_normalized = float(normalized.max(initial=0.0)) + return ComparisonSummary( + max_abs_diff=float(np.round(abs_diff.max(initial=0.0), 15)), + max_rel_diff=float(np.round(rel_diff.max(initial=0.0), 15)), + max_normalized_error=max_normalized, + ) + + +def run_artifact_lane( + scenarios: list[ArtifactScenario], + *, + repeats: int, + warmup: int, +) -> list[ArtifactResult]: + """Time each scenario's compiled artifacts, one result per operation. + + Artifacts are prepared once per scenario and then called ``repeats`` times + after ``warmup`` discarded calls, so what is measured is steady-state + evaluation rather than compilation. + """ + _validate_counts(repeats, warmup) + results: list[ArtifactResult] = [] + for scenario in scenarios: + prepared = _prepare_artifact_case(scenario) + for operation in scenario.operations: + ( + baseline_callable, + candidate_callable, + aot_callable, + materialize, + ) = prepared.operations[operation] + baseline_run_ms, baseline_value, baseline_samples = _time_callable( + baseline_callable, repeats=repeats, warmup=warmup + ) + candidate_run_ms, candidate_value, candidate_samples = _time_callable( + candidate_callable, repeats=repeats, warmup=warmup + ) + aot_run_ms, aot_value, aot_samples = _time_callable( + aot_callable, repeats=repeats, warmup=warmup + ) + # Densify outside the timed region so run timings reflect kernel cost, + # not format marshalling. + comparison = summarize_diff( + materialize(baseline_value), + materialize(candidate_value), + atol=prepared.atol, + rtol=prepared.rtol, + ) + aot_comparison = summarize_diff( + materialize(baseline_value), + materialize(aot_value), + atol=prepared.atol, + rtol=prepared.rtol, + ) + results.append( + ArtifactResult( + scenario=scenario.name, + operation=operation, + baseline_backend=prepared.baseline_backend, + candidate_backend=prepared.candidate_backend, + baseline_timings=replace( + prepared.baseline_timings, + run_ms=baseline_run_ms, + ), + candidate_timings=replace( + prepared.candidate_timings, + run_ms=candidate_run_ms, + ), + comparison=comparison, + baseline_run_samples_ms=baseline_samples, + candidate_run_samples_ms=candidate_samples, + aot_run_samples_ms=aot_samples, + aot_backend=prepared.aot_backend, + aot_timings=replace(prepared.aot_timings, run_ms=aot_run_ms), + aot_comparison=aot_comparison, + ) + ) + return results + + +def _compare_solver_pair( + target_solution, + target_output, + solution, + output, + scenario, + *, + output_only: bool, + atol: float, + rtol: float, +): + """State, output and trajectory agreement of one solve against a target solve.""" + trajectory = summarize_trajectory( + target_solution, solution, atol=scenario.atol, rtol=scenario.rtol + ) + common_points = trajectory.common_points + state = None + if not output_only: + target_y, candidate_y = _align_time_axis( + target_solution.y, solution.y, common_points=common_points + ) + state = summarize_diff(target_y, candidate_y, atol=atol, rtol=rtol) + target_out, candidate_out = _align_time_axis( + target_output, output, common_points=common_points + ) + values = summarize_diff(target_out, candidate_out, atol=atol, rtol=rtol) + return state, values, trajectory + + +def _solver_reference(scenario, tolerance: float | None): + """The converged solve, its observed output, and the tolerance that produced them.""" + + def solve(attempt): + solution = _reference_solve(scenario, attempt) + return solution, _extract_output(solution, scenario.observed_output) + + return _resolve_reference(scenario, tolerance, solve, what="reference") + + +def run_solver_lane( + scenarios: list[SolverScenario], + *, + repeats: int, + warmup: int, + include_aot: bool = True, + backend_order_seed: int = 0, + reference_tolerance: float | None = DEFAULT_REFERENCE_TOLERANCE, +) -> list[SolverResult]: + """Solve each scenario on every backend and compare states and outputs. + + Every row, the baseline included, is judged against one converged + ``casadi_idaklu`` solve at ``reference_tolerance``; ``None`` falls back to + comparing against the baseline backend, leaving that row ungated. + + ``backend_order_seed`` seeds the per-scenario backend shuffle; any fixed + value gives a reproducible order. The baseline is shuffled in with the + candidates and every comparison is computed after execution, so no backend + is systematically measured on a colder machine. + """ + _validate_counts(repeats, warmup) + reference_tolerance = resolve_reference_tolerance(scenarios, reference_tolerance) + results: list[SolverResult] = [] + for scenario in scenarios: + measured = _execute_backend_cases( + scenario, + "solver", + repeats=repeats, + warmup=warmup, + include_aot=include_aot, + backend_order_seed=backend_order_seed, + ) + _, baseline_solution, baseline_output = measured[_BASELINE_CASE] + reference_used, reference = _solver_reference(scenario, reference_tolerance) + reference_atol, reference_rtol = _reference_tolerances(scenario) + + for case, (result, solution, output) in measured.items(): + output_only = case[1] + if not result.supported: + results.append(result) + continue + if reference is None and case == _BASELINE_CASE: + results.append(result) + continue + + if reference is None: + target, target_output = baseline_solution, baseline_output + atol, rtol = scenario.atol, scenario.rtol + else: + target, target_output = reference + atol, rtol = reference_atol, reference_rtol + state_comparison, output_comparison, trajectory_comparison = ( + _compare_solver_pair( + target, + target_output, + solution, + output, + scenario, + output_only=output_only, + atol=atol, + rtol=rtol, + ) + ) + baseline_delta = None + if reference is not None and case != _BASELINE_CASE: + baseline_state, baseline_values, _ = _compare_solver_pair( + baseline_solution, + baseline_output, + solution, + output, + scenario, + output_only=output_only, + atol=scenario.atol, + rtol=scenario.rtol, + ) + baseline_delta = _worst_comparison((baseline_state, baseline_values)) + results.append( + replace( + result, + state_comparison=state_comparison, + output_comparison=output_comparison, + trajectory_comparison=trajectory_comparison, + reference_tolerance=reference_used, + baseline_delta=baseline_delta, + ) + ) + return results + + +def _compare_sensitivity_pair( + target_solution, + target_state, + target_output, + solution, + state_sens, + output_sens, + scenario, + *, + output_only: bool, + atol: float | None, + rtol: float, +): + """State- and output-sensitivity agreement of one solve against a target solve. + + ``atol`` of ``None`` takes each block's near-zero floor from that block's own + peak instead: sensitivity magnitudes run from ``dV/dI`` at 1e-2 to ``dc/dp`` + at 1e4, so one shared absolute floor is either meaningless or dominant. + """ + trajectory = summarize_trajectory( + target_solution, solution, atol=scenario.atol, rtol=scenario.rtol + ) + align = { + "baseline_points": target_solution.t.size, + "candidate_points": solution.t.size, + "common_points": trajectory.common_points, + } + + def compare(target, candidate): + target_rows, candidate_rows = _align_rows(target, candidate, **align) + floor = atol + if floor is None: + floor = rtol * float(np.abs(_as_dense(target_rows)).max(initial=0.0)) + return summarize_diff(target_rows, candidate_rows, atol=floor, rtol=rtol) + + state = None + if not output_only and state_sens is not None and target_state is not None: + state = compare(target_state, state_sens) + return state, compare(target_output, output_sens), trajectory + + +def _sensitivity_reference(scenario, tolerance: float | None): + """The converged sensitivity solve and its blocks, with the tolerance used.""" + parameter_values, inputs = _build_sensitivity_parameters(scenario) + + def solve(attempt): + solution = _reference_solve( + scenario, + attempt, + inputs=inputs, + parameter_values=parameter_values, + sensitivity_parameters=sorted(inputs), + ) + state_sens, output_sens = _extract_sensitivities( + solution, scenario.observed_output, output_only=False + ) + return solution, state_sens, output_sens + + return _resolve_reference(scenario, tolerance, solve, what="sensitivity reference") + + +def run_sensitivity_lane( + scenarios: list[SolverScenario], + *, + repeats: int, + warmup: int, + include_aot: bool = False, + backend_order_seed: int = 0, + reference_tolerance: float | None = DEFAULT_REFERENCE_TOLERANCE, +) -> list[SensitivityResult]: + """Solve each scenario with forward sensitivities and compare them. + + Compares the state- and output-sensitivity blocks plus the trajectory, not the + state and output values the solver lane checks, since a backend can integrate + correctly and still get ``dy/dp`` wrong. Same converged-reference convention as + :func:`run_solver_lane`, judged at the looser gradient tolerance because + forward sensitivities are not error-controlled to the state tolerance. + """ + _validate_counts(repeats, warmup) + reference_tolerance = resolve_reference_tolerance(scenarios, reference_tolerance) + results: list[SensitivityResult] = [] + for scenario in scenarios: + measured = _execute_backend_cases( + scenario, + "sensitivity", + repeats=repeats, + warmup=warmup, + include_aot=include_aot, + backend_order_seed=backend_order_seed, + ) + _, baseline_state, baseline_output, baseline_solution = measured[_BASELINE_CASE] + reference_used, reference = _sensitivity_reference( + scenario, reference_tolerance + ) + _, gradient_rtol = _sensitivity_tolerances(scenario) + + for case, (result, state_sens, output_sens, solution) in measured.items(): + output_only = case[1] + if not result.supported: + results.append(result) + continue + if reference is None and case == _BASELINE_CASE: + results.append(result) + continue + + # A baseline fallback does not make gradients error-controlled. + target = ( + (baseline_solution, baseline_state, baseline_output) + if reference is None + else reference + ) + atol, rtol = None, gradient_rtol + state_sens_comparison, output_sens_comparison, trajectory_comparison = ( + _compare_sensitivity_pair( + *target, + solution, + state_sens, + output_sens, + scenario, + output_only=output_only, + atol=atol, + rtol=rtol, + ) + ) + baseline_delta = None + if reference is not None and case != _BASELINE_CASE: + delta_state, delta_output, _ = _compare_sensitivity_pair( + baseline_solution, + baseline_state, + baseline_output, + solution, + state_sens, + output_sens, + scenario, + output_only=output_only, + atol=scenario.atol, + rtol=scenario.rtol, + ) + baseline_delta = _worst_comparison((delta_state, delta_output)) + results.append( + replace( + result, + state_sens_comparison=state_sens_comparison, + output_sens_comparison=output_sens_comparison, + trajectory_comparison=trajectory_comparison, + reference_tolerance=reference_used, + baseline_delta=baseline_delta, + ) + ) + return results + + +def _run_backend_case( + lane: str, + scenario: SolverScenario, + *, + backend: str, + repeats: int, + warmup: int, + output_only: bool, +) -> tuple: + """Measure one ``(backend, output_only)`` case on one scenario.""" + if backend == "casadi_idaklu_aot": + return _run_profiled_aot_backend( + lane, + scenario, + repeats=repeats, + warmup=warmup, + output_only=output_only, + output_points=scenario.plan.requested_points or DEFAULT_OUTPUT_POINTS, + ) + runner = _run_solver_backend if lane == "solver" else _run_sensitivity_backend + return runner( + scenario, + backend=backend, + repeats=repeats, + warmup=warmup, + output_only=output_only, + ) + + +def _unsupported_lane_result( + lane: str, scenario: SolverScenario, backend: str, output_only: bool, reason: str +) -> tuple: + """A fixed-arity failure record, so the comparison pass unpacks uniformly.""" + factory = SolverResult if lane == "solver" else SensitivityResult + result = factory( + scenario=scenario.name, + protocol=scenario.protocol, + backend=_solver_row_name(backend, output_only), + timings=PhaseTiming(), + requested_output_points=scenario.plan.requested_points, + supported=False, + reason=reason, + ) + return (result, None, None) if lane == "solver" else (result, None, None, None) + + +def _execute_backend_cases( + scenario: SolverScenario, + lane: str, + *, + repeats: int, + warmup: int, + include_aot: bool, + backend_order_seed: int, +) -> dict[tuple[str, bool], tuple]: + """Measure every backend case for one scenario, in shuffled execution order. + + The baseline is shuffled in with the candidates rather than pinned first, so + machine warm-up cannot systematically favour one backend. The returned + mapping is re-keyed to start at the baseline, leaving the comparison pass + independent of the order things ran in. + """ + cases = _shuffled_backend_cases( + backend_order_seed, + f"{lane}:{scenario.name}:{scenario.protocol}", + include_aot=include_aot, + ) + measured: dict[tuple[str, bool], tuple] = {} + for backend, output_only in cases: + try: + measured[(backend, output_only)] = _run_backend_case( + lane, + scenario, + backend=backend, + repeats=repeats, + warmup=warmup, + output_only=output_only, + ) + except Exception as exc: + # Only the baseline is load-bearing; anything else degrades to a + # visible row rather than discarding a whole run's timings. + if (backend, output_only) == _BASELINE_CASE: + raise + measured[(backend, output_only)] = _unsupported_lane_result( + lane, scenario, backend, output_only, _short_reason(exc) + ) + return {_BASELINE_CASE: measured.pop(_BASELINE_CASE), **measured} + + +def sample_input_vectors( + nominal: dict[str, float], + count: int, + *, + seed: int, + spread: float | dict[str, float] = 0.2, +) -> list[dict[str, float]]: + """Log-uniform input vectors within ``spread`` of ``nominal``. + + Drawn once per scenario and shared across backends, so repeat *i* uses the + same inputs everywhere and cross-backend comparison stays exact. + + Parameters + ---------- + nominal : dict of str to float + Nominal value per input name. + count : int + Number of vectors to draw. + seed : int + Seed for the generator, making the sequence reproducible. + spread : float or dict of str to float + Fractional half-width of the sampling interval, either shared by every + parameter or given per parameter. One width cannot suit parameters of + different natures: what is routine for a diffusivity takes a bounded + volume fraction somewhere its model cannot be solved. + + Returns + ------- + list of dict + One input dictionary per draw, in draw order. + + Raises + ------ + KeyError + If a per-parameter mapping omits a name in ``nominal``, so a new fitted + parameter cannot silently inherit someone else's width. + """ + rng = np.random.default_rng(seed) + names = list(nominal) + widths = np.array( + [spread if isinstance(spread, float | int) else spread[name] for name in names], + dtype=np.float64, + ) + low = np.log(1.0 - widths) + high = np.log(1.0 + widths) + factors = np.exp(rng.uniform(low, high, size=(count, len(names)))) + return [ + { + name: float(nominal[name] * factor) + for name, factor in zip(names, row, strict=True) + } + for row in factors + ] + + +def run_inference_lane( + scenarios: list[SolverScenario], + *, + repeats: int, + warmup: int, + seed: int = 0, + sensitivities: bool = False, + include_aot: bool = False, + backend_order_seed: int = 0, + reference_tolerance: float | None = DEFAULT_REFERENCE_TOLERANCE, +) -> list[InferenceResult]: + """Time one likelihood evaluation per repeat, with inputs changing each time. + + Unlike the solver lane, the fitted parameters stay symbolic and every timed + repeat uses a different input vector, and observation goes through the + interpolating call interface rather than the raw entries. Same + converged-reference convention as :func:`run_solver_lane`, one reference + solve per measured draw. + """ + _validate_counts(repeats, warmup) + reference_tolerance = resolve_reference_tolerance(scenarios, reference_tolerance) + nominal = inference_nominal_values() + results: list[InferenceResult] = [] + for scenario in scenarios: + vectors = sample_input_vectors( + nominal, warmup + repeats, seed=seed, spread=INFERENCE_SPREADS + ) + cases = _shuffled_backend_cases( + backend_order_seed, + f"inference:{scenario.name}:{scenario.protocol}", + include_aot=include_aot, + ) + measured: dict[tuple[str, bool], tuple] = {} + failed_scenario: str | None = None + try: + grid = _shared_observation_grid(scenario, vectors[0]) + except Exception as exc: # noqa: BLE001 + failed_scenario = _short_reason(exc) + + for backend, output_only in cases: + if failed_scenario is not None: + measured[(backend, output_only)] = _unsupported_inference_case( + scenario, backend, output_only, failed_scenario + ) + continue + try: + measured[(backend, output_only)] = _run_inference_backend( + scenario, + backend=backend, + output_only=output_only, + vectors=vectors, + warmup=warmup, + sensitivities=sensitivities, + grid=grid, + ) + except Exception as exc: # noqa: BLE001 + reason = _short_reason(exc) + # Without a baseline there is nothing to compare against, so the + # scenario degrades whole rather than aborting the suite. + if (backend, output_only) == _BASELINE_CASE: + failed_scenario = reason + measured = { + case: _unsupported_inference_case(scenario, *case, reason) + for case in measured + } + measured[(backend, output_only)] = _unsupported_inference_case( + scenario, backend, output_only, reason + ) + + # Re-keyed to start at the baseline, so the reported order is independent + # of the shuffled execution order. + measured = {_BASELINE_CASE: measured.pop(_BASELINE_CASE), **measured} + _, baseline_repeats = measured[_BASELINE_CASE] + reference_used, reference_repeats = ( + (None, None) + if failed_scenario is not None + else _inference_reference( + scenario, + vectors, + warmup=warmup, + grid=grid, + sensitivities=sensitivities, + tolerance=reference_tolerance, + ) + ) + sensitivity_atol, sensitivity_rtol = _sensitivity_tolerances(scenario) + reference_atol, reference_rtol = _reference_tolerances(scenario) + for case, (result, candidate_repeats) in measured.items(): + if not result.supported: + results.append(result) + continue + if reference_repeats is None and case == _BASELINE_CASE: + results.append(result) + continue + + if reference_repeats is None: + target_repeats = baseline_repeats + value_atol, value_rtol = scenario.atol, scenario.rtol + else: + target_repeats = reference_repeats + value_atol, value_rtol = reference_atol, reference_rtol + baseline_delta = None + if reference_repeats is not None and case != _BASELINE_CASE: + baseline_delta = _worst_comparison( + ( + _worst_repeat_comparison( + baseline_repeats, + candidate_repeats, + select=_observed_values, + atol=scenario.atol, + rtol=scenario.rtol, + ), + _worst_repeat_comparison( + baseline_repeats, + candidate_repeats, + select=_observed_sensitivities, + atol=sensitivity_atol, + rtol=sensitivity_rtol, + ), + ) + ) + results.append( + replace( + result, + output_comparison=_worst_repeat_comparison( + target_repeats, + candidate_repeats, + select=_observed_values, + atol=value_atol, + rtol=value_rtol, + ), + sensitivity_comparison=_worst_repeat_comparison( + target_repeats, + candidate_repeats, + select=_observed_sensitivities, + atol=sensitivity_atol, + rtol=sensitivity_rtol, + ), + trajectory_comparison=_worst_repeat_trajectory( + target_repeats, candidate_repeats, scenario + ), + reference_tolerance=reference_used, + baseline_delta=baseline_delta, + ) + ) + return results + + +def _unsupported_inference_case( + scenario, backend: str, output_only: bool, reason: str +) -> tuple: + """An in-row failure paired with empty repeats, so the comparison pass unpacks + uniformly and one bad backend never abandons the rest of the suite.""" + return ( + InferenceResult( + scenario=scenario.name, + protocol=scenario.protocol, + backend=_solver_row_name(backend, output_only), + build_ms=0.0, + setup_ms=0.0, + aot_cache_status="-", + eval_samples_ms=(), + solve_samples_ms=(), + observe_samples_ms=(), + requested_output_points=scenario.plan.requested_points, + supported=False, + reason=reason, + ), + (), + ) + + +@dataclass(frozen=True) +class RepeatObservation: + """What one likelihood evaluation produced, for cross-backend comparison. + + ``times`` are the timestamps the values were read at, stopping at the + solution's own end, so a candidate that terminated early carries fewer. + Sensitivities come back on the solution's own grid rather than the + observation grid, hence the second time axis. + """ + + values: np.ndarray + times: np.ndarray + sensitivities: np.ndarray | None + sensitivity_times: np.ndarray | None + final_time: float + termination: str + + +def _observed_values(repeat: RepeatObservation): + return repeat.values, repeat.times + + +def _observed_sensitivities(repeat: RepeatObservation): + if repeat.sensitivities is None: + return None + return repeat.sensitivities, repeat.sensitivity_times + + +def _comparable_length(baseline_times, candidate_times, *, endpoint_gap: float) -> int: + """Points that both repeats can be judged on. + + The shared prefix, minus any point inside the window where the two solutions' + endpoints disagree: a moved event time leaves both trajectories racing to a + cutoff at different moments, so points there measure the endpoint gap rather + than trajectory agreement. Sizing the window from the measured gap keeps it at + zero when the terminations coincide, and independent of ``--output-points``. + """ + length = min(baseline_times.size, candidate_times.size) + if length == 0 or endpoint_gap <= 0.0: + return length + shared_times = baseline_times[:length] + cutoff = float(shared_times[-1]) - endpoint_gap + return int(np.searchsorted(shared_times, cutoff, side="right")) + + +def _trim_to_points(values: np.ndarray, times: np.ndarray, points: int) -> np.ndarray: + """Trim a time-outer block to ``points`` timepoints, whole rows at a time.""" + if times.size == 0 or values.shape[0] % times.size: + return values + return values[: points * (values.shape[0] // times.size)] + + +def _worst_repeat_comparison( + baseline_repeats, candidate_repeats, *, select, atol: float, rtol: float +): + """Worst agreement across every repeat, not just the last. + + Ranked by tolerance-normalised error, so a repeat that breached tolerance can + never be masked by one whose absolute difference is larger but still allowed. + Returns ``None`` when ``select`` yields nothing to compare. + """ + summaries = [] + for baseline, candidate in zip(baseline_repeats, candidate_repeats, strict=True): + selected = (select(baseline), select(candidate)) + if any(item is None for item in selected): + return None + (base_values, base_times), (cand_values, cand_times) = selected + length = _comparable_length( + base_times, + cand_times, + endpoint_gap=abs(baseline.final_time - candidate.final_time), + ) + if length <= 0: + summaries.append(_UNCOMPARABLE) + continue + summaries.append( + summarize_diff( + _trim_to_points(base_values, base_times, length), + _trim_to_points(cand_values, cand_times, length), + atol=atol, + rtol=rtol, + ) + ) + return _worst_comparison(summaries) or _UNCOMPARABLE + + +def _worst_repeat_trajectory(baseline_repeats, candidate_repeats, scenario): + """Worst trajectory agreement across repeats: coverage, endpoint, termination. + + Both backends are read on the same grid, truncated at their own solution's + end, so a shortfall in observed points *is* an early termination. + """ + summaries = [ + _repeat_trajectory(baseline, candidate, scenario) + for baseline, candidate in zip(baseline_repeats, candidate_repeats, strict=True) + ] + if not summaries: + return None + return min( + summaries, key=lambda summary: (summary.status == "pass", summary.coverage) + ) + + +def _repeat_trajectory( + baseline: RepeatObservation, candidate: RepeatObservation, scenario +) -> TrajectorySummary: + """Coverage and termination of one inference repeat against another. + + Both read the same shared observation grid, truncated at each solution's + own end, so the time axes agree exactly over their common prefix and only + the lengths and the final time carry information. + """ + summary = _summarize_time_axes( + np.asarray(baseline.times, dtype=np.float64).reshape(-1), + np.asarray(candidate.times, dtype=np.float64).reshape(-1), + baseline.termination, + candidate.termination, + atol=scenario.atol, + rtol=scenario.rtol, + ) + return replace( + summary, final_time_diff=abs(baseline.final_time - candidate.final_time) + ) + + +def _observation_grid(scenario) -> np.ndarray: + """Experimental timestamps to observe at, offset off the solver's own grid. + + The first interval is skipped. It straddles the initial transient, which no + two-point interpolant can represent from its endpoints alone -- a converged + reference needs hundreds of internal steps inside it at low output-point + counts. Reading there scores the output grid's resolution, not the backend: + a Hermite interpolant given correct data lands further out than a linear + chord, so the cell ranks backends by which interpolant they happen to use. + """ + plan = scenario.plan + if plan.t_interp is not None: + # Midpoints, so observation genuinely interpolates rather than hitting nodes. + return 0.5 * (plan.t_interp[1:-1] + plan.t_interp[2:]) + return np.array([], dtype=np.float64) + + +@contextmanager +def _aot_compile_events(backend: str): + """Capture the AOT cache telemetry a solve emits, for AOT backends only.""" + if backend != "casadi_idaklu_aot": + yield None + return + from pybamm.codegen.compilation import _capture_aot_compile_events + + with _capture_aot_compile_events() as events: + yield events + + +def _summarize_cache_statuses(events) -> str: + """Collapse per-kernel cache statuses into one cell. + + Reports what the compiler actually did (``miss``/``disk``/``memory``) so a + cheap warm ``Setup`` is never misread as a fresh compile. + """ + if not events: + return "-" + statuses = sorted({event.cache_status for event in events}) + return statuses[0] if len(statuses) == 1 else "+".join(statuses) + + +def _scaled_sensitivities(solution, name: str, *, inputs: dict[str, float]): + """``p . d(var)/dp`` for every fitted parameter, one column each. + + The fitted parameters span eighteen orders of magnitude (a diffusivity at + 1e-14 beside a concentration at 1e4), so raw ``d/dp`` columns are not + comparable under one tolerance. Scaling by the parameter gives the derivative + with respect to a fractional change: dimensionless, the same size as the + variable itself, and the quantity a log-space fitting loop actually consumes. + Columns are taken by name rather than from the stacked ``"all"`` block so the + scale factor cannot be paired with the wrong parameter. + """ + sensitivities = solution[name].sensitivities + return np.column_stack( + [ + np.asarray(sensitivities[input_name], dtype=np.float64).reshape(-1) + * inputs[input_name] + for input_name in sorted(inputs) + ] + ) + + +# A gate placed on the noise floor tracks it instead of discriminating against it. +_SENSITIVITY_NOISE_HEADROOM = 10.0 + + +def _sensitivity_tolerances(scenario) -> tuple[float, float]: + """Tolerances for judging a gradient, looser than for the state. + + A solver error-controls the state to ``(atol, rtol)``; the forward + sensitivities integrated alongside it are not, and degrade to roughly the + square root of it. Approaching an event the scaled gradient is also + near-singular -- 2e3 against order 1 mid-run -- and the worst measured + cross-backend agreement there is ~1e-3 relative, right at that square root. + The gate therefore sits a decade above it: still three orders below the + order-one relative miss a broken chain rule or an unseeded ``dy0/dp`` gives. + """ + return ( + _SENSITIVITY_NOISE_HEADROOM * math.sqrt(scenario.atol), + _SENSITIVITY_NOISE_HEADROOM * math.sqrt(scenario.rtol), + ) + + +def _shared_observation_grid(scenario, inputs: dict[str, float]) -> np.ndarray: + """The timestamps every backend in this scenario is read at. + + Declared grids come straight from the protocol. A period-driven protocol has + none, so one throwaway baseline solve establishes it; probing here rather than + inside a measured run keeps the grid independent of the shuffled backend order. + """ + declared = _observation_grid(scenario) + if declared.size: + return declared + solver = _make_solver( + "casadi_idaklu", atol=scenario.atol, rtol=scenario.rtol, output_variables=None + ) + simulation = _build_simulation(scenario, solver, "casadi_idaklu") + _build_and_time(simulation, scenario, inputs=inputs) + probe = simulation.solve(**_solve_kwargs(scenario, {"inputs": inputs})) + return np.asarray(probe.t, dtype=np.float64) + + +def _observe_inference( + solution, inputs: dict[str, float], *, grid, scenario, sensitivities: bool +) -> RepeatObservation: + """Read one solve on the shared grid, the way an inference loop would.""" + # The interpolating call interface, which is what an inference loop uses. + times = grid[grid <= solution.t[-1]] + values = np.asarray( + solution[scenario.observed_output](times), dtype=np.float64 + ).reshape(-1) + gradient = sensitivity_times = None + if sensitivities: + gradient = _scaled_sensitivities( + solution, scenario.observed_output, inputs=inputs + ) + sensitivity_times = np.asarray(solution.t, dtype=np.float64) + return RepeatObservation( + values=values, + times=times, + sensitivities=gradient, + sensitivity_times=sensitivity_times, + final_time=float(solution.t[-1]), + termination=str(solution.termination), + ) + + +def _reference_solve_kwargs(scenario, extra: dict | None = None) -> dict: + """Solve arguments for a reference, storing the solver's own steps. + + A reference restricted to the candidate's output grid carries that grid's + interpolation error, and the gate would charge it to any candidate whose + interpolant is *better* than the reference's -- ranking interpolants rather + than checking correctness. Dropping ``t_interp`` stores every internal step, + so an off-node read is judged against a trajectory dense enough to resolve it. + """ + kwargs = _solve_kwargs(scenario, extra) + if kwargs.get("t_interp") is not None: + kwargs["t_interp"] = None + return kwargs + + +def _inference_reference( + scenario, + vectors, + *, + warmup: int, + grid, + sensitivities: bool, + tolerance: float | None, +): + """One converged observation per measured draw, with the tolerance used. + + Solved outside every timed region, so the reference never lands in a sample. + Built from ``vectors[0]`` like every backend is: the lane holds ``y0`` at the + first draw, and a reference resolved from a different one starts the cell at + a different state of charge. A draw that will not converge fails the whole + ladder rung -- the comparison needs a reference for every draw, not most. + """ + measured = list(vectors[warmup:]) + # The draw that failed the last rung decides the next one too, so trying it + # first costs a doomed rung one converged solve instead of all of them. + decides_the_rung = 0 + + def solve(attempt): + nonlocal decides_the_rung + solver = _make_solver( + _REFERENCE_BACKEND, atol=attempt, rtol=attempt, output_variables=None + ) + simulation = _build_simulation(scenario, solver, _REFERENCE_BACKEND) + _build_and_time(simulation, scenario, inputs=vectors[0]) + observations: dict[int, object] = {} + order = sorted(range(len(measured)), key=lambda i: i != decides_the_rung) + for index in order: + inputs = measured[index] + extra: dict = {"inputs": inputs} + if sensitivities: + extra["calculate_sensitivities"] = sorted(INFERENCE_INPUTS.values()) + try: + solution = simulation.solve(**_reference_solve_kwargs(scenario, extra)) + except Exception: + decides_the_rung = index + raise + observations[index] = _observe_inference( + solution, + inputs, + grid=grid, + scenario=scenario, + sensitivities=sensitivities, + ) + return [observations[index] for index in range(len(measured))] + + return _resolve_reference(scenario, tolerance, solve, what="reference") + + +def _run_inference_backend( + scenario, + *, + backend: str, + output_only: bool, + vectors: list[dict[str, float]], + warmup: int, + sensitivities: bool, + grid: np.ndarray, +): + """Solve one scenario once per input vector and time each evaluation. + + ``grid`` is the shared observation grid, so every backend is read at the same + timestamps. Cold observation is forced and timed before the warmup loop, + keeping lazy variable compilation out of the per-evaluation samples whatever + ``warmup`` is set to. + """ + output_variables = [scenario.observed_output] if output_only else None + solver = _make_solver( + backend, + atol=scenario.atol, + rtol=scenario.rtol, + output_variables=output_variables, + ) + simulation = _build_simulation(scenario, solver, backend) + build_ms = _build_and_time(simulation, scenario, inputs=vectors[0]) + + def solve_extra(inputs): + extra: dict = {"inputs": inputs} + if sensitivities: + extra["calculate_sensitivities"] = sorted(INFERENCE_INPUTS.values()) + return extra + + with _aot_compile_events(backend) as events: + cold_solution = simulation.solve( + **_solve_kwargs(scenario, solve_extra(vectors[0])) + ) + cache_status = _summarize_cache_statuses(events) + setup_ms = _time_to_ms(cold_solution.set_up_time) + + def observe(solution, inputs) -> tuple[float, RepeatObservation]: + start = perf_counter() + observation = _observe_inference( + solution, inputs, grid=grid, scenario=scenario, sensitivities=sensitivities + ) + return (perf_counter() - start) * 1000.0, observation + + cold_observe_ms, _ = observe(cold_solution, vectors[0]) + + def evaluate(inputs): + solve_start = perf_counter() + solution = simulation.solve(**_solve_kwargs(scenario, solve_extra(inputs))) + solve_ms = (perf_counter() - solve_start) * 1000.0 + observe_ms, observation = observe(solution, inputs) + return solve_ms, observe_ms, observation + + for inputs in vectors[:warmup]: + evaluate(inputs) + + solve_samples: list[float] = [] + observe_samples: list[float] = [] + eval_samples: list[float] = [] + observations: list[RepeatObservation] = [] + for inputs in vectors[warmup:]: + eval_start = perf_counter() + solve_ms, observe_ms, observation = evaluate(inputs) + eval_samples.append((perf_counter() - eval_start) * 1000.0) + solve_samples.append(solve_ms) + observe_samples.append(observe_ms) + observations.append(observation) + + result = InferenceResult( + scenario=scenario.name, + protocol=scenario.protocol, + backend=_solver_row_name(backend, output_only), + build_ms=build_ms, + setup_ms=setup_ms, + cold_observe_ms=cold_observe_ms, + aot_cache_status=cache_status, + eval_samples_ms=tuple(eval_samples), + solve_samples_ms=tuple(solve_samples), + observe_samples_ms=tuple(observe_samples), + requested_output_points=scenario.plan.requested_points, + ) + return result, observations + + +def _prepare_artifact_case(scenario: ArtifactScenario) -> _PreparedArtifactCase: + if scenario.name == "toy_expr": + return _prepare_toy_expr_case(scenario) + if scenario.name in {"spm_residual", "spme_residual", "dfn_residual"}: + model_factory = { + "spm_residual": pybamm.lithium_ion.SPM, + "spme_residual": pybamm.lithium_ion.SPMe, + "dfn_residual": pybamm.lithium_ion.DFN, + }[scenario.name] + return _prepare_model_residual_case(scenario, model_factory) + raise ValueError(f"Unsupported artifact scenario: {scenario.name}") + + +def _prepare_aot_kernels(functions, n_traj_cols): + """AOT-compile the CasADi kernels to a native shared lib and return the + externals, the mapped eval, and the (compile-dominated) prep time. + + Raises if compilation silently fell back to the VM so the AOT row can never + be mislabelled as native code.""" + from pybamm.codegen.compilation import aot_compile + + start = perf_counter() + cf, cjy, cjp, cjvp = aot_compile(list(functions)) + cf_map = cf.map(n_traj_cols) + prepare_ms = (perf_counter() - start) * 1000.0 + for fn in (cf, cjy, cjp, cjvp): + if fn.class_name() != "External": + raise RuntimeError( + "AOT compilation fell back to the CasADi VM (compiler missing " + "or failed); refusing to report a mislabelled AOT row." + ) + return cf, cjy, cjp, cjvp, cf_map, prepare_ms + + +def _prepare_toy_expr_case(scenario: ArtifactScenario) -> _PreparedArtifactCase: + import casadi + + from pybamm.rust import ExprGraph + + build_start = perf_counter() + expr, n_states, input_names = _toy_expr() + t = 0.7 + y = np.array([0.3, 1.2]) + p = np.array([2.5, -0.8]) + v = np.array([0.6, -1.1]) + ts = np.linspace(0.0, 2.0, 100) + y_traj = np.vstack([np.linspace(0.1, 1.0, 100), np.linspace(-0.5, 1.5, 100)]) + build_ms = (perf_counter() - build_start) * 1000.0 + + baseline_prepare_start = perf_counter() + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", n_states) + y_dot_sym = casadi.MX.sym("y_dot", n_states) + p_syms = {name: casadi.MX.sym(name) for name in input_names} + cexpr = expr.to_casadi( + t_sym, + y_sym, + y_dot_sym, + p_syms, + {"t": t_sym, "y": y_sym, "y_dot": y_dot_sym, "inputs": p_syms}, + ) + p_stacked = casadi.vertcat(*p_syms.values()) + v_sym = casadi.MX.sym("v", n_states) + cf = casadi.Function("f", [t_sym, y_sym, p_stacked], [cexpr]) + cjy = casadi.Function( + "jy", + [t_sym, y_sym, p_stacked], + [casadi.jacobian(cexpr, y_sym)], + ) + cjp = casadi.Function( + "jp", [t_sym, y_sym, p_stacked], [casadi.jacobian(cexpr, p_stacked)] + ) + cjvp = casadi.Function( + "jvp", [t_sym, y_sym, p_stacked, v_sym], [casadi.jtimes(cexpr, y_sym, v_sym)] + ) + cf_map = cf.map(ts.size) + ts_row, p_tiled = _trajectory_inputs(ts, p) + baseline_prepare_ms = (perf_counter() - baseline_prepare_start) * 1000.0 + + candidate_prepare_start = perf_counter() + graph = ExprGraph() + rust_expr = expr.to_rust(graph, {}) + rust_function = graph.compile(rust_expr, name="toy_expr", n_states=n_states) + rust_jacobian_y = rust_function.jacobian() + rust_jacobian_p = rust_function.jacobian(wrt="p") + rust_function.jvp(t, y, p, v) + candidate_prepare_ms = (perf_counter() - candidate_prepare_start) * 1000.0 + + aot_cf, aot_cjy, aot_cjp, aot_cjvp, aot_cf_map, aot_prepare_ms = ( + _prepare_aot_kernels((cf, cjy, cjp, cjvp), ts.size) + ) + + operations = { + "eval": ( + lambda: cf(t, y, p), + lambda: rust_function(t, y, p), + lambda: aot_cf(t, y, p), + _as_vec, + ), + "jacobian_y": ( + lambda: cjy(t, y, p), + lambda: rust_jacobian_y(t, y, p), + lambda: aot_cjy(t, y, p), + _as_dense, + ), + "jacobian_p": ( + lambda: cjp(t, y, p), + lambda: rust_jacobian_p(t, y, p), + lambda: aot_cjp(t, y, p), + _as_dense, + ), + "jvp": ( + lambda: cjvp(t, y, p, v), + lambda: rust_function.jvp(t, y, p, v), + lambda: aot_cjvp(t, y, p, v), + _as_vec, + ), + "eval_trajectory": ( + lambda: cf_map(ts_row, y_traj, p_tiled), + lambda: rust_function.eval_trajectory(ts, y_traj, p), + lambda: aot_cf_map(ts_row, y_traj, p_tiled), + _as_dense, + ), + } + return _PreparedArtifactCase( + baseline_backend="casadi", + candidate_backend="rust", + baseline_timings=PhaseTiming( + build_ms=build_ms, + prepare_ms=baseline_prepare_ms, + ), + candidate_timings=PhaseTiming( + build_ms=build_ms, + prepare_ms=candidate_prepare_ms, + ), + operations=operations, + atol=scenario.atol, + rtol=scenario.rtol, + aot_timings=PhaseTiming(build_ms=build_ms, prepare_ms=aot_prepare_ms), + ) + + +def _prepare_model_residual_case( + scenario: ArtifactScenario, + model_factory, +) -> _PreparedArtifactCase: + import casadi + + from pybamm.rust import ExprGraph + + build_start = perf_counter() + model = model_factory() + model.events = [] + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Current function [A]"] = pybamm.InputParameter("I") + simulation = pybamm.Simulation( + model, + parameter_values=parameter_values, + var_pts=_make_var_pts(model, 10), + ) + simulation.build() + built = simulation.built_model + full_symbol = _full_residual_symbol(built) + y = np.asarray( + built.concatenated_initial_conditions.evaluate(), dtype=np.float64 + ).reshape(-1) + p = np.array([0.5], dtype=np.float64) + v = np.random.default_rng(0).standard_normal(y.size) + ts = np.linspace(0.0, 10.0, 100) + y_traj = np.tile(y[:, None], (1, ts.size)) * np.linspace(1.0, 1.01, ts.size) + build_ms = (perf_counter() - build_start) * 1000.0 + + baseline_prepare_start = perf_counter() + n_states = built.len_rhs_and_alg + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", n_states) + y_dot_sym = casadi.MX.sym("y_dot", n_states) + p_syms = {"I": casadi.MX.sym("I")} + cexpr = full_symbol.to_casadi( + t_sym, + y_sym, + y_dot_sym, + p_syms, + {"t": t_sym, "y": y_sym, "y_dot": y_dot_sym, "inputs": p_syms}, + ) + p_stacked = casadi.vertcat(*p_syms.values()) + v_sym = casadi.MX.sym("v", n_states) + cf = casadi.Function("f", [t_sym, y_sym, p_stacked], [cexpr]) + cjy = casadi.Function( + "jy", + [t_sym, y_sym, p_stacked], + [casadi.jacobian(cexpr, y_sym)], + ) + cjp = casadi.Function( + "jp", [t_sym, y_sym, p_stacked], [casadi.jacobian(cexpr, p_stacked)] + ) + cjvp = casadi.Function( + "jvp", [t_sym, y_sym, p_stacked, v_sym], [casadi.jtimes(cexpr, y_sym, v_sym)] + ) + cf_map = cf.map(ts.size) + ts_row, p_tiled = _trajectory_inputs(ts, p) + baseline_prepare_ms = (perf_counter() - baseline_prepare_start) * 1000.0 + + candidate_prepare_start = perf_counter() + graph = ExprGraph() + rust_expr = full_symbol.to_rust(graph, {}) + rust_function = graph.compile(rust_expr, name=scenario.name, n_states=n_states) + rust_jacobian_y = rust_function.jacobian() + rust_jacobian_p = rust_function.jacobian(wrt="p") + rust_function.jvp(0.0, y, p, v) + candidate_prepare_ms = (perf_counter() - candidate_prepare_start) * 1000.0 + + aot_cf, aot_cjy, aot_cjp, aot_cjvp, aot_cf_map, aot_prepare_ms = ( + _prepare_aot_kernels((cf, cjy, cjp, cjvp), ts.size) + ) + + operations = { + "eval": ( + lambda: cf(0.0, y, p), + lambda: rust_function(0.0, y, p), + lambda: aot_cf(0.0, y, p), + _as_vec, + ), + "jacobian_y": ( + lambda: cjy(0.0, y, p), + lambda: rust_jacobian_y(0.0, y, p), + lambda: aot_cjy(0.0, y, p), + _as_dense, + ), + "jacobian_p": ( + lambda: cjp(0.0, y, p), + lambda: rust_jacobian_p(0.0, y, p), + lambda: aot_cjp(0.0, y, p), + _as_dense, + ), + "jvp": ( + lambda: cjvp(0.0, y, p, v), + lambda: rust_function.jvp(0.0, y, p, v), + lambda: aot_cjvp(0.0, y, p, v), + _as_vec, + ), + "eval_trajectory": ( + lambda: cf_map(ts_row, y_traj, p_tiled), + lambda: rust_function.eval_trajectory(ts, y_traj, p), + lambda: aot_cf_map(ts_row, y_traj, p_tiled), + _as_dense, + ), + } + return _PreparedArtifactCase( + baseline_backend="casadi", + candidate_backend="rust", + baseline_timings=PhaseTiming( + build_ms=build_ms, + prepare_ms=baseline_prepare_ms, + ), + candidate_timings=PhaseTiming( + build_ms=build_ms, + prepare_ms=candidate_prepare_ms, + ), + operations=operations, + atol=scenario.atol, + rtol=scenario.rtol, + aot_timings=PhaseTiming(build_ms=build_ms, prepare_ms=aot_prepare_ms), + ) + + +def _toy_expr(): + y0 = pybamm.StateVector(slice(0, 1)) + y1 = pybamm.StateVector(slice(1, 2)) + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.NumpyConcatenation( + a * y0 * y1 + pybamm.t, + pybamm.sin(y0) * b + pybamm.exp(-y1), + ) + return expr, 2, ("a", "b") + + +def _build_simulation(scenario, solver, backend, *, parameter_values=None): + """Construct (but do not build) a simulation for one scenario.""" + model = scenario.model_factory(options=scenario.model_options) + model.convert_to_format = _backend_convert_to_format(backend) + if parameter_values is None: + parameter_values = scenario.parameter_values_builder() + kwargs = {"parameter_values": parameter_values, "solver": solver} + if scenario.plan.experiment is not None: + kwargs["experiment"] = scenario.plan.experiment + return pybamm.Simulation(model, **kwargs) + + +def _build_and_time(simulation, scenario, inputs: dict | None = None) -> float: + """Build the simulation once, outside the experiment path, and time it. + + Pre-building an experiment-attached simulation reparameterises the same + model that ``Simulation.solve`` then parameterises again per step, which + trips PyBaMM's reparameterised-model guard. The experiment path therefore + builds lazily on its first ``solve`` call instead, contributing 0 ms here. + + Parameters + ---------- + simulation : pybamm.Simulation + The constructed, unbuilt simulation. + scenario : SolverScenario + Supplies the protocol's ``initial_soc`` and solve plan. + inputs : dict, optional + Values for any symbolic parameters. Required when ``initial_soc`` is + set and the parameter set carries ``InputParameter``s, because mapping + SOC to concentrations runs an ElectrodeSOH solve that must evaluate + them. The initial state is fixed at these values for every repeat, so + the warm path stays warm. + """ + build_start = perf_counter() + if scenario.plan.experiment is None: + # initial_soc is applied here, once, never per solve: passing it to + # Simulation.solve re-runs set_initial_state on every call. + simulation.build(initial_soc=scenario.initial_soc, inputs=inputs) + return (perf_counter() - build_start) * 1000.0 + + +def _solve_kwargs(scenario, extra: dict | None = None) -> dict: + """Time arguments for one protocol's solve, plus any caller extras. + + The experiment path takes its times from the experiment itself, so it + contributes no time arguments at all. + """ + kwargs: dict = {} + plan = scenario.plan + if plan.experiment is None: + kwargs["t_eval"] = ( + [float(plan.t_interp[0]), float(plan.t_interp[-1])] + if plan.t_eval is None + else plan.t_eval + ) + kwargs["t_interp"] = plan.t_interp + kwargs.update(extra or {}) + return kwargs + + +def _run_solver_backend( + scenario: SolverScenario, + *, + backend: str, + repeats: int, + warmup: int, + output_only: bool, +): + output_variables = [scenario.observed_output] if output_only else None + solver = _make_solver( + backend, + atol=scenario.atol, + rtol=scenario.rtol, + output_variables=output_variables, + ) + simulation = _build_simulation(scenario, solver, backend) + + cold_startup_start = perf_counter() + build_ms = _build_and_time(simulation, scenario) + + solve_kwargs = _solve_kwargs(scenario) + + cold_solution = simulation.solve(**solve_kwargs) + cold_observe_start = perf_counter() + last_output = _extract_output(cold_solution, scenario.observed_output) + cold_observe_ms = (perf_counter() - cold_observe_start) * 1000.0 + cold_startup_ms = (perf_counter() - cold_startup_start) * 1000.0 + + for _ in range(warmup): + warm_solution = simulation.solve(**solve_kwargs) + _extract_output(warm_solution, scenario.observed_output) + + warm_set_up_samples = [] + solve_samples = [] + wall_solve_samples = [] + integration_samples = [] + observe_samples = [] + e2e_samples = [] + last_solution = cold_solution + + for _ in range(repeats): + e2e_start = perf_counter() + wall_solve_start = perf_counter() + solution = simulation.solve(**solve_kwargs) + wall_solve_samples.append((perf_counter() - wall_solve_start) * 1000.0) + warm_set_up_samples.append(_time_to_ms(solution.set_up_time)) + solve_samples.append(_time_to_ms(solution.solve_time)) + integration_samples.append(_time_to_ms(solution.integration_time)) + observe_start = perf_counter() + output = _extract_output(solution, scenario.observed_output) + observe_samples.append((perf_counter() - observe_start) * 1000.0) + e2e_samples.append((perf_counter() - e2e_start) * 1000.0) + last_solution = solution + last_output = output + + samples = TimingSamples( + warm_set_up_ms=tuple(warm_set_up_samples), + solve_ms=tuple(solve_samples), + wall_solve_ms=tuple(wall_solve_samples), + integration_ms=tuple(integration_samples), + observe_ms=tuple(observe_samples), + e2e_ms=tuple(e2e_samples), + ) + result = SolverResult( + scenario=scenario.name, + protocol=scenario.protocol, + backend=_solver_row_name(backend, output_only), + timings=_summarize_timing_samples( + samples, + build_ms=build_ms, + cold_set_up_ms=_time_to_ms(cold_solution.set_up_time), + cold_observe_ms=cold_observe_ms, + cold_startup_ms=cold_startup_ms, + ), + requested_output_points=scenario.plan.requested_points, + timing_samples=samples, + jacobian_telemetry=_get_jacobian_telemetry(solver, backend), + ) + return result, last_solution, last_output + + +SENSITIVITY_INPUTS = { + "Current function [A]": "I", + "Positive electrode active material volume fraction": "eps_p", +} + + +def _build_sensitivity_parameters(scenario): + """Swap the sensitivity parameters for input parameters on top of the protocol's values.""" + parameter_values = scenario.parameter_values_builder() + inputs = {} + for pybamm_name, input_name in SENSITIVITY_INPUTS.items(): + if ( + pybamm_name == "Current function [A]" + and scenario.plan.experiment is not None + ): + # An Experiment's own control law supersedes "Current function [A]", + # so it would be a dead, unreferenced sensitivity input. + continue + value = parameter_values[pybamm_name] + if isinstance(value, pybamm.Symbol): + # A protocol has made this parameter time-varying; it cannot be an input. + continue + inputs[input_name] = float(value) + parameter_values[pybamm_name] = pybamm.InputParameter(input_name) + return parameter_values, inputs + + +def _extract_sensitivities(solution, name: str, *, output_only: bool): + """Materialize the stacked ``"all"`` sensitivity blocks. + + Output sensitivities ``d(var)/dp`` are always available (chain rule); full + state sensitivities ``dy/dp`` only when the solve was not output-restricted. + """ + output_sens = np.asarray(solution[name].sensitivities["all"], dtype=np.float64) + if output_only: + return None, output_sens + state_sens = np.asarray(solution.sensitivities["all"], dtype=np.float64) + return state_sens, output_sens + + +def _run_sensitivity_backend( + scenario: SolverScenario, + *, + backend: str, + repeats: int, + warmup: int, + output_only: bool, +): + parameter_values, inputs = _build_sensitivity_parameters(scenario) + output_variables = [scenario.observed_output] if output_only else None + solver = _make_solver( + backend, + atol=scenario.atol, + rtol=scenario.rtol, + output_variables=output_variables, + ) + simulation = _build_simulation( + scenario, solver, backend, parameter_values=parameter_values + ) + + cold_startup_start = perf_counter() + build_ms = _build_and_time(simulation, scenario, inputs=inputs) + + solve_kwargs = _solve_kwargs( + scenario, + {"inputs": inputs, "calculate_sensitivities": sorted(inputs)}, + ) + + cold_solution = simulation.solve(**solve_kwargs) + cold_observe_start = perf_counter() + last_state_sens, last_output_sens = _extract_sensitivities( + cold_solution, scenario.observed_output, output_only=output_only + ) + cold_observe_ms = (perf_counter() - cold_observe_start) * 1000.0 + cold_startup_ms = (perf_counter() - cold_startup_start) * 1000.0 + + for _ in range(warmup): + warm_solution = simulation.solve(**solve_kwargs) + _extract_sensitivities( + warm_solution, scenario.observed_output, output_only=output_only + ) + + warm_set_up_samples = [] + solve_samples = [] + wall_solve_samples = [] + integration_samples = [] + observe_samples = [] + e2e_samples = [] + last_solution = cold_solution + for _ in range(repeats): + e2e_start = perf_counter() + wall_solve_start = perf_counter() + solution = simulation.solve(**solve_kwargs) + wall_solve_samples.append((perf_counter() - wall_solve_start) * 1000.0) + warm_set_up_samples.append(_time_to_ms(solution.set_up_time)) + solve_samples.append(_time_to_ms(solution.solve_time)) + integration_samples.append(_time_to_ms(solution.integration_time)) + observe_start = perf_counter() + state_sens, output_sens = _extract_sensitivities( + solution, scenario.observed_output, output_only=output_only + ) + observe_samples.append((perf_counter() - observe_start) * 1000.0) + e2e_samples.append((perf_counter() - e2e_start) * 1000.0) + last_state_sens = state_sens + last_output_sens = output_sens + last_solution = solution + + samples = TimingSamples( + warm_set_up_ms=tuple(warm_set_up_samples), + solve_ms=tuple(solve_samples), + wall_solve_ms=tuple(wall_solve_samples), + integration_ms=tuple(integration_samples), + observe_ms=tuple(observe_samples), + e2e_ms=tuple(e2e_samples), + ) + result = SensitivityResult( + scenario=scenario.name, + protocol=scenario.protocol, + backend=_solver_row_name(backend, output_only), + timings=_summarize_timing_samples( + samples, + build_ms=build_ms, + cold_set_up_ms=_time_to_ms(cold_solution.set_up_time), + cold_observe_ms=cold_observe_ms, + cold_startup_ms=cold_startup_ms, + ), + requested_output_points=scenario.plan.requested_points, + timing_samples=samples, + sensitivity_parameters=tuple(sorted(inputs)), + ) + return result, last_state_sens, last_output_sens, last_solution + + +@contextmanager +def _aot_cache_environment(cache_dir: str): + previous = os.environ.get("PYBAMM_CASADI_AOT_CACHE") + os.environ["PYBAMM_CASADI_AOT_CACHE"] = cache_dir + try: + yield + finally: + if previous is None: + os.environ.pop("PYBAMM_CASADI_AOT_CACHE", None) + else: + os.environ["PYBAMM_CASADI_AOT_CACHE"] = previous + + +def _run_aot_once( + lane: str, + scenario: SolverScenario, + *, + repeats: int, + warmup: int, + output_only: bool, + cache_dir: str, +): + from pybamm.codegen.compilation import _CACHE, _capture_aot_compile_events + + _CACHE.clear() + try: + with _aot_cache_environment(cache_dir), _capture_aot_compile_events() as events: + if lane == "solver": + values = _run_solver_backend( + scenario, + backend="casadi_idaklu_aot", + repeats=repeats, + warmup=warmup, + output_only=output_only, + ) + else: + values = _run_sensitivity_backend( + scenario, + backend="casadi_idaklu_aot", + repeats=repeats, + warmup=warmup, + output_only=output_only, + ) + return values, tuple(events) + finally: + _CACHE.clear() + + +def _aot_worker_payload( + lane, scenario, *, output_only, cache_dir, output_points +) -> tuple: + """Name-only payload for the AOT disk worker. + + A ``SolverScenario`` pickles fine today, but its protocol builders are + references to module-level functions, which would break silently if one + ever became a closure or lambda. Names are simpler and immune to that. + """ + return ( + lane, + scenario.name, + scenario.protocol, + output_points, + output_only, + cache_dir, + ) + + +def _run_aot_disk_worker(payload): + lane, model_name, protocol_name, output_points, output_only, cache_dir = payload + scenario = get_solver_scenarios( + [model_name], [protocol_name], output_points=output_points + )[0] + values, events = _run_aot_once( + lane, + scenario, + repeats=1, + warmup=0, + output_only=output_only, + cache_dir=cache_dir, + ) + return values[0].timings, events + + +def _summarize_aot_profile(fresh_events, disk_events, disk_timings) -> AotProfile: + fresh_statuses = tuple(event.cache_status for event in fresh_events) + disk_statuses = tuple(event.cache_status for event in disk_events) + fresh_keys = tuple(event.cache_key for event in fresh_events) + disk_keys = tuple(event.cache_key for event in disk_events) + verified = bool(fresh_events) and ( + all(status == "miss" for status in fresh_statuses) + and all(status == "disk" for status in disk_statuses) + and fresh_keys == disk_keys + ) + profile = AotProfile( + fresh_cache_statuses=fresh_statuses, + disk_cache_statuses=disk_statuses, + codegen_ms=sum(event.codegen_ms for event in fresh_events), + compiler_ms=sum(event.compiler_ms for event in fresh_events), + fresh_load_ms=sum(event.load_ms for event in fresh_events), + disk_load_ms=sum(event.load_ms for event in disk_events), + fresh_total_ms=sum(event.total_ms for event in fresh_events), + disk_total_ms=sum(event.total_ms for event in disk_events), + disk_prepare_ms=disk_timings.prepare_ms, + disk_cold_startup_ms=disk_timings.cold_startup_ms, + library_size_bytes=sum(event.library_size_bytes or 0 for event in fresh_events), + verified=verified, + ) + if not verified: + raise RuntimeError( + "AOT profile was not a verified cache miss followed by a fresh-process " + f"disk hit: fresh={fresh_statuses}, disk={disk_statuses}" + ) + return profile + + +def _run_profiled_aot_backend( + lane: str, + scenario: SolverScenario, + *, + repeats: int, + warmup: int, + output_only: bool, + output_points: int, +): + with tempfile.TemporaryDirectory(prefix="pybamm-aot-benchmark-") as cache_dir: + fresh_values, fresh_events = _run_aot_once( + lane, + scenario, + repeats=repeats, + warmup=warmup, + output_only=output_only, + cache_dir=cache_dir, + ) + context = multiprocessing.get_context("spawn") + with context.Pool(processes=1) as pool: + disk_timings, disk_events = pool.apply( + _run_aot_disk_worker, + ( + _aot_worker_payload( + lane, + scenario, + output_only=output_only, + cache_dir=cache_dir, + output_points=output_points, + ), + ), + ) + profile = _summarize_aot_profile( + fresh_events, + disk_events, + disk_timings, + ) + result = replace(fresh_values[0], aot_profile=profile) + return result, *fresh_values[1:] + + +def _backend_convert_to_format(backend: str) -> str: + """Map a bench backend name to the ``model.convert_to_format`` that selects + it.""" + if backend in ("casadi_idaklu", "casadi_idaklu_aot"): + return "casadi" + if backend in ("rust_idaklu", "rust_diffsol"): + return "rust" + raise ValueError(f"Unsupported solver backend: {backend}") + + +def _make_solver(backend: str, *, atol: float, rtol: float, output_variables=None): + if backend == "casadi_idaklu_aot": + return pybamm.IDAKLUSolver( + atol=atol, + rtol=rtol, + output_variables=output_variables, + options={"compile": True}, + ) + if backend in ("casadi_idaklu", "rust_idaklu"): + return pybamm.IDAKLUSolver( + atol=atol, + rtol=rtol, + output_variables=output_variables, + ) + if backend == "rust_diffsol": + return pybamm.DiffsolSolver( + atol=atol, + rtol=rtol, + output_variables=output_variables, + ) + raise ValueError(f"Unsupported solver backend: {backend}") + + +def _solver_row_name(backend: str, output_only: bool) -> str: + return f"{backend}_out" if output_only else backend + + +def _short_reason(exc: Exception) -> str: + reason = " ".join(str(exc).split()) + if not reason: + reason = exc.__class__.__name__ + return f"{exc.__class__.__name__}: {reason}" + + +def _make_var_pts(model, npts: int) -> dict: + var_pts = {} + for key, value in model.default_var_pts.items(): + if isinstance(value, (int, float)) and value > 1 and key not in {"y", "z"}: + var_pts[key] = npts + else: + var_pts[key] = value + return var_pts + + +def _full_residual_symbol(built_model): + if built_model.len_alg > 0: + return pybamm.numpy_concatenation( + built_model.concatenated_rhs, + built_model.concatenated_algebraic, + ) + return built_model.concatenated_rhs + + +def _extract_output(solution, name: str) -> np.ndarray: + """Read the observed variable off the solver's own stored grid. + + Deliberately not the interpolating call interface: the solver lane measures + what materialising the stored trajectory costs, and the inference lane + measures interpolated reads. Between them both doors are covered. + """ + return np.asarray(solution[name].data, dtype=np.float64) + + +def summarize_trajectory( + baseline_solution, + candidate_solution, + *, + atol: float, + rtol: float, +) -> TrajectorySummary: + """Compare time coverage and termination before numerical parity.""" + return _summarize_time_axes( + np.asarray(baseline_solution.t, dtype=np.float64).reshape(-1), + np.asarray(candidate_solution.t, dtype=np.float64).reshape(-1), + str(baseline_solution.termination), + str(candidate_solution.termination), + atol=atol, + rtol=rtol, + ) + + +def _summarize_time_axes( + baseline_t: np.ndarray, + candidate_t: np.ndarray, + baseline_termination: str, + candidate_termination: str, + *, + atol: float, + rtol: float, +) -> TrajectorySummary: + """The coverage and termination rules every lane classifies against. + + Shared so the solver and inference lanes cannot come to different verdicts + about the same divergence. + """ + baseline_points = baseline_t.size + candidate_points = candidate_t.size + minimum_points = min(baseline_points, candidate_points) + span = max( + abs(float(baseline_t[-1] - baseline_t[0])) if baseline_points else 0.0, + abs(float(candidate_t[-1] - candidate_t[0])) if candidate_points else 0.0, + ) + time_atol = max(atol, span * rtol) + + common_points = 0 + if minimum_points: + matches = np.isclose( + baseline_t[:minimum_points], + candidate_t[:minimum_points], + atol=time_atol, + rtol=0.0, + ) + mismatch = np.flatnonzero(~matches) + common_points = int(mismatch[0]) if mismatch.size else minimum_points + + if common_points: + max_time_diff = float( + np.max(np.abs(baseline_t[:common_points] - candidate_t[:common_points])) + ) + else: + max_time_diff = float("inf") + final_time_diff = ( + abs(float(baseline_t[-1] - candidate_t[-1])) + if baseline_points and candidate_points + else float("inf") + ) + coverage = common_points / max(baseline_points, candidate_points, 1) + only_terminal_sample_differs = ( + common_points >= minimum_points - 1 + and abs(baseline_points - candidate_points) <= 1 + ) + status = ( + "pass" + if baseline_termination == candidate_termination + and only_terminal_sample_differs + and final_time_diff <= time_atol + else "warn" + ) + return TrajectorySummary( + baseline_points=baseline_points, + candidate_points=candidate_points, + common_points=common_points, + coverage=coverage, + max_time_diff=max_time_diff, + final_time_diff=final_time_diff, + baseline_termination=baseline_termination, + candidate_termination=candidate_termination, + status=status, + ) + + +def _align_time_axis(baseline, candidate, *, common_points: int): + """Trim time-last arrays to a trajectory-validated common prefix.""" + base = _as_dense(baseline) + cand = _as_dense(candidate) + return base[..., :common_points], cand[..., :common_points] + + +def _align_rows( + baseline, + candidate, + *, + baseline_points: int, + candidate_points: int, + common_points: int, +): + """Trim fused time-outer sensitivity rows on timepoint boundaries.""" + base = _as_dense(baseline) + cand = _as_dense(candidate) + if baseline_points == 0 or candidate_points == 0: + return base[:0], cand[:0] + if base.shape[0] % baseline_points or cand.shape[0] % candidate_points: + return base, cand + base_rows_per_point = base.shape[0] // baseline_points + candidate_rows_per_point = cand.shape[0] // candidate_points + return ( + base[: common_points * base_rows_per_point], + cand[: common_points * candidate_rows_per_point], + ) + + +def _get_jacobian_telemetry(solver, backend: str) -> JacobianTelemetry | None: + if backend == "rust_idaklu": + # ``_setup`` is only assigned inside set_up(); the experiment path solves + # through a per-step copy of this solver, so it may never run here. + model = getattr(solver, "_setup", {}).get("rust_model") + elif backend == "rust_diffsol": + model = getattr(solver, "_rust_model", None) + else: + return None + if model is None: + return None + stats = model.jacobian_stats() + return JacobianTelemetry( + strategy=stats["strategy"], + n_colors=int(stats["n_colors"]), + nnz=int(stats["nnz"]), + n_dense_rows=int(stats.get("n_dense_rows", 0)), + dense_row_entries=int(stats.get("dense_row_entries", 0)), + dense_row_tape_instructions=int(stats.get("dense_row_tape_instructions", 0)), + split_eval_primal_instructions=stats["split_eval_primal_instructions"], + split_eval_total_instructions=stats["split_eval_total_instructions"], + split_eval_raw_instructions=stats.get("split_eval_raw_instructions"), + branch_block_lens=tuple(stats.get("branch_block_lens", ())), + ) + + +def _summarize_timing_samples( + samples: TimingSamples, + *, + build_ms: float, + cold_set_up_ms: float, + cold_observe_ms: float, + cold_startup_ms: float, +) -> PhaseTiming: + return PhaseTiming( + build_ms=build_ms, + prepare_ms=cold_set_up_ms + cold_observe_ms, + cold_startup_ms=cold_startup_ms, + set_up_ms=cold_set_up_ms, + warm_set_up_ms=float(np.median(samples.warm_set_up_ms)), + solve_ms=float(np.median(samples.solve_ms)), + wall_solve_ms=float(np.median(samples.wall_solve_ms)), + integration_ms=float(np.median(samples.integration_ms)), + observe_ms=float(np.median(samples.observe_ms)), + e2e_ms=float(np.median(samples.e2e_ms)), + ) + + +def _as_dense(value) -> np.ndarray: + """Materialize a backend's native output (CasADi DM, Rust sparse, ndarray) as dense float64.""" + if hasattr(value, "toarray"): + value = value.toarray() + return np.asarray(value, dtype=np.float64) + + +def _as_vec(value) -> np.ndarray: + """Materialize a vector-valued output as a flat float64 array.""" + return _as_dense(value).reshape(-1) + + +def _trajectory_inputs(ts: np.ndarray, p: np.ndarray): + """Shape the trajectory inputs for a CasADi mapped function call. + + ``Function.map(N)`` maps every input over N columns, so the scalar time + becomes a (1, N) row and the constant parameter vector is tiled to (n_p, N). + """ + ts_row = np.asarray(ts, dtype=np.float64).reshape(1, -1) + p_tiled = np.tile(np.asarray(p, dtype=np.float64).reshape(-1, 1), (1, ts.size)) + return ts_row, p_tiled + + +def _time_callable( + callable_, *, repeats: int, warmup: int, min_batch_seconds: float = 0.02 +): + """Median per-call time (ms) using auto-calibrated batching. + + A single perf_counter() span around one sub-microsecond call is dominated by + timer resolution and Python dispatch. Instead we size an inner batch so each + timed span exceeds ``min_batch_seconds``, then take the median per-call time + across ``repeats`` such batches. ``repeats`` therefore counts batches, not + individual calls. + """ + for _ in range(warmup): + callable_() + + batch = 1 + result = None + while True: + start = perf_counter() + for _ in range(batch): + result = callable_() + elapsed = perf_counter() - start + if elapsed >= min_batch_seconds or batch >= 1_000_000: + break + if elapsed <= 0.0: + batch *= 8 + else: + batch = max(batch * 2, int(batch * min_batch_seconds / elapsed) + 1) + + per_call_ms = [elapsed / batch * 1000.0] + for _ in range(repeats - 1): + start = perf_counter() + for _ in range(batch): + result = callable_() + per_call_ms.append((perf_counter() - start) / batch * 1000.0) + return float(np.median(per_call_ms)), result, tuple(per_call_ms) + + +def _time_to_ms(value) -> float: + if value is None: + return 0.0 + raw_value = value.value if hasattr(value, "value") else value + return float(raw_value) * 1000.0 + + +def _validate_counts(repeats: int, warmup: int) -> None: + if repeats < 1: + raise ValueError("repeats must be at least 1") + if warmup < 0: + raise ValueError("warmup must be non-negative") + + +def _shuffled_backend_cases( + seed: int, key: str, *, include_aot: bool +) -> list[tuple[str, bool]]: + cases = list(backend_cases(include_aot)) + random.Random(f"{seed}:{key}").shuffle(cases) + return cases diff --git a/docs/source/api/solvers/diffsol_solver.rst b/docs/source/api/solvers/diffsol_solver.rst new file mode 100644 index 0000000000..517845572e --- /dev/null +++ b/docs/source/api/solvers/diffsol_solver.rst @@ -0,0 +1,5 @@ +Diffsol Solver +============== + +.. autoclass:: pybamm.DiffsolSolver + :members: diff --git a/docs/source/api/solvers/index.rst b/docs/source/api/solvers/index.rst index 2bb15503bc..bd108708d3 100644 --- a/docs/source/api/solvers/index.rst +++ b/docs/source/api/solvers/index.rst @@ -4,6 +4,7 @@ Solvers .. toctree:: base_solver + diffsol_solver dummy_solver scipy_solver jax_solver diff --git a/docs/source/developer/rust_idaklu_build.md b/docs/source/developer/rust_idaklu_build.md new file mode 100644 index 0000000000..108ef2bf7d --- /dev/null +++ b/docs/source/developer/rust_idaklu_build.md @@ -0,0 +1,168 @@ +# Building the Rust Backend + +This guide covers building the Rust compute core, its Python bindings, and the +Rust-capable IDAKLU solver for development. Everything lives in this +repository's `uv` workspace; there are no side-by-side clones or environment +variables to configure. + +## Components + +| Component | Location | Build tool | Output | +|-----------|----------|------------|--------| +| `pybamm-core` | `packages/pybamm-rust/pybamm-core` | Cargo | Rust library (rlib + cdylib) | +| `pybamm-python` | `packages/pybamm-rust/pybamm-python` | maturin (via `uv sync`) | `pybamm.rust._core` extension module | +| pybammsolvers | `packages/pybammsolvers` | scikit-build-core / CMake | `idaklu` extension module | + +**How the pieces connect:** + +- `pybamm.rust` provides the Python API that lowers a discretised model to a + `CompiledModel` and drives the diffsol solver. +- pybammsolvers does **not** link against the Rust core (the package + dependency only flows the other way). The IDAKLU C++ consumer resolves the + `extern "C"` entry points at runtime — `dlsym(RTLD_DEFAULT)` on POSIX, a + loaded-module walk with `GetProcAddress` on Windows — once + `pybamm.rust._core` is imported in-process. +- The FFI contract's source of truth is the `extern "C"` exports in + `packages/pybamm-rust/pybamm-core/src/ffi.rs` (symbols carry a + `pybamm_rust_` prefix). The C++ consumer mirrors them as a function-pointer + table in + `packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h`. + A version handshake (`pybamm_rust_abi_version()` / + `PYBAMM_RUST_ABI_VERSION`) and the `test_ffi_abi_contract` drift test keep + the two sides in lockstep. + +## Prerequisites + +- **Rust 1.89+**: `curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh` +- **uv**: `curl -LsSf https://astral.sh/uv/install.sh | sh` +- **CMake 3.20+** and a C++17 compiler +- **macOS only**: Homebrew with `libomp` (`brew install libomp`) + +## Quick start + +One sync builds everything — the maturin hook compiles `pybamm.rust._core` in +release mode, and pybammsolvers builds SUNDIALS/SuiteSparse from its bundled +submodules on first build: + +```bash +git clone https://github.com/pybamm-team/PyBaMM.git +cd PyBaMM +git submodule update --init --recursive +uv sync --extra all --group dev +``` + +Verify the two halves see each other: + +```bash +uv run python -c " +from pybamm.rust import CompiledModel +from pybammsolvers import idaklu +print('Rust bindings:', CompiledModel) +print('Rust FFI available:', hasattr(idaklu, 'create_rust_solver_group')) +" +``` + +## Rebuilding after changes + +`uv`'s cache keys cover the Rust sources, so a plain `uv sync` rebuilds +`pybamm.rust._core` whenever a `.rs` file changes. + +pybammsolvers' cache key is only its `pyproject.toml`. After changing any of +its C++ sources — including `pybamm_rust_ffi.h` — force a rebuild: + +```bash +uv cache clean pybammsolvers +uv sync --extra all --group dev --reinstall-package pybammsolvers +``` + +## Changing the FFI surface + +The drift test parses both sides of the boundary, so the workflow is +mechanical: + +1. Edit the exports in `ffi.rs` and the consumer table in + `pybamm_rust_ffi.h` together. +2. Run `cargo test -p pybamm-core --test test_ffi_abi_contract` from + `packages/pybamm-rust/`. On any change to the exported surface it fails + and prints the new `EXPECTED_ABI_HASH`; update it in `ffi.rs` and bump + `RUST_ABI_VERSION` / `PYBAMM_RUST_ABI_VERSION` in lockstep. +3. Rebuild both extensions (previous section). A stale pybammsolvers binary + fails loudly at first use — either an unresolvable symbol or an ABI + version mismatch — rather than corrupting an evaluation. + +## Running tests + +```bash +# Rust unit + contract tests +cd packages/pybamm-rust && cargo test --workspace + +# Rust-IDAKLU integration tests (from the repo root) +uv run --group dev pytest packages/pybamm/tests/integration/test_rust_idaklu_parity.py +uv run --group dev pytest packages/pybamm/tests/integration/test_rust_idaklu_spm.py + +# Diffsol solver unit tests +uv run --group dev pytest packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py +``` + +## Verifying FFI symbols + +The entry points are exported by `pybamm.rust._core` and resolved from it at +runtime, so they should be *defined* there and appear nowhere in idaklu: + +```bash +# macOS (Linux: nm -D --defined-only) +nm -gU packages/pybamm/src/pybamm/rust/_core.abi3.so | grep pybamm_rust_ +``` + +Expected output lists symbols like `_pybamm_rust_eval_rhs` and +`_pybamm_rust_abi_version`. If a Rust-backed solve fails with an +unresolvable-symbol error instead, the loaded extension predates the current +FFI surface — rebuild both extensions from the same source tree. + +## Running benchmarks + +```bash +# Rust micro-benchmarks (Criterion) +cd packages/pybamm-rust/pybamm-core +cargo bench + +# Python end-to-end observability benchmark (from the repo root) +uv run python benchmarks/run_rust_observability.py +``` + +## Troubleshooting + +### Segfaults in SUNDIALS callbacks + +Crashes during `solve()` (especially in residual/Jacobian callbacks) usually +indicate a **SuiteSparse version mismatch**: the `.idaklu/` build directory +holds a stale build while the submodules (or a Homebrew install picked up via +`@rpath`) moved on. Rebuild the native dependencies from the current +submodules: + +```bash +cd packages/pybammsolvers +rm -rf .idaklu +uv run python install_KLU_Sundials.py +cd ../.. && uv sync --extra all --group dev --reinstall-package pybammsolvers +``` + +Do **not** point the build at Homebrew's suite-sparse; version mismatches +between it and the local SUNDIALS build cause runtime crashes. + +### `pybamm.rust._core` import fails + +```bash +uv sync --extra all --group dev --reinstall-package pybamm +``` + +### Other symbol errors + +An ABI mismatch between components — rebuild both extensions from the same +tree (see "Rebuilding after changes"). + +## Supported features + +Feature coverage (events, DAEs, forward sensitivities, output variables, +experiments) and the expression types the backend can lower are documented in +the [Rust backend user guide](../user_guide/rust_backend.rst). diff --git a/docs/source/examples/notebooks/models/jelly-roll-model.ipynb b/docs/source/examples/notebooks/models/jelly-roll-model.ipynb index 33ca0123d1..8aedbd7a51 100644 --- a/docs/source/examples/notebooks/models/jelly-roll-model.ipynb +++ b/docs/source/examples/notebooks/models/jelly-roll-model.ipynb @@ -212,15 +212,11 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "straight-anime", "metadata": {}, "outputs": [], - "source": [ - "# solver\n", - "solver = pybamm.CasadiAlgebraicSolver()\n", - "solution = solver.solve(model)" - ] + "source": "# solver\nsolver = pybamm.NonlinearSolver()\nsolution = solver.solve(model)" }, { "cell_type": "markdown", diff --git a/docs/source/examples/notebooks/parameterization/bpx.ipynb b/docs/source/examples/notebooks/parameterization/bpx.ipynb index 4b19e1bd8c..f6b7ecb10d 100644 --- a/docs/source/examples/notebooks/parameterization/bpx.ipynb +++ b/docs/source/examples/notebooks/parameterization/bpx.ipynb @@ -108,7 +108,7 @@ "Negative electrode active material volume fraction\t0.6860102133333333\n", "Negative electrode conductivity [S.m-1]\t0.222\n", "Negative electrode density [kg.m-3]\t1847.0\n", - "Negative electrode diffusivity [m2.s-1]\tfunctools.partial(._diffusivity at 0x17da31800>, D_ref=2.728e-14, Ea=30000.0, constant=True)\n", + "Negative particle diffusivity [m2.s-1]\tfunctools.partial(._diffusivity at 0x17da31800>, D_ref=2.728e-14, Ea=30000.0, constant=True)\n", "Negative electrode diffusivity activation energy [J.mol-1]\t30000.0\n", "Negative electrode exchange-current density [A.m-2]\tfunctools.partial(._exchange_current_density at 0x17da316c0>, k_ref=5.33563612557725e-07, Ea=55000.0)\n", "Negative electrode maximum stoichiometry\t0.75668\n", @@ -240,7 +240,7 @@ "ax[1, 0].plot(\n", " xLi_n.entries,\n", " parameter_values.evaluate(\n", - " parameter_values[\"Negative electrode diffusivity [m2.s-1]\"](xLi_n, T)\n", + " parameter_values[\"Negative particle diffusivity [m2.s-1]\"](xLi_n, T)\n", " )\n", " * np.ones_like(xLi_n.entries),\n", ")\n", @@ -248,7 +248,7 @@ "ax[1, 1].plot(\n", " xLi_p.entries,\n", " parameter_values.evaluate(\n", - " parameter_values[\"Positive electrode diffusivity [m2.s-1]\"](xLi_p, T)\n", + " parameter_values[\"Positive particle diffusivity [m2.s-1]\"](xLi_p, T)\n", " )\n", " * np.ones_like(xLi_p.entries),\n", ")\n", diff --git a/docs/source/examples/notebooks/performance/07-multithreading.ipynb b/docs/source/examples/notebooks/performance/07-multithreading.ipynb index 5d548372ac..b43593ec5b 100644 --- a/docs/source/examples/notebooks/performance/07-multithreading.ipynb +++ b/docs/source/examples/notebooks/performance/07-multithreading.ipynb @@ -183,6 +183,22 @@ "source": [ "So in this case the speed-up for using multiple threads to solve 1000 SPM simulations is much less than for the DFN simulations, and above 4 threads no speed-up is observed at all." ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## The same option on the other backends\n", + "\n", + "`num_threads` means the same thing wherever it appears: *this many input sets solve at once*.\n", + "\n", + "- `IDAKLUSolver` schedules them with OpenMP, one solver per thread, for both `convert_to_format=\"casadi\"` and `\"rust\"`. Work is handed out dynamically, so a sweep whose sets terminate at different times (a current sweep hitting a voltage cut-off, say) does not leave threads idle waiting on the slowest block.\n", + "- `DiffsolSolver` takes the same option and schedules with rayon: `pybamm.DiffsolSolver(options={\"num_threads\": 8})`. Pools are shared process-wide by thread count, so ten solvers asking for 8 threads share one pool of 8 rather than spawning eighty.\n", + "\n", + "Either way, leave `num_threads` at its default of 1 if something outside PyBaMM is already parallelising over your sweep — a `ThreadPoolExecutor`, joblib, or a parameter-estimation library — or the two thread counts multiply.\n", + "\n", + "The timings above are from one machine and one PyBaMM version; treat them as the shape of the curve rather than numbers to reproduce." + ] } ], "metadata": { diff --git a/docs/source/user_guide/index.md b/docs/source/user_guide/index.md index 6b700e15fc..108c585479 100644 --- a/docs/source/user_guide/index.md +++ b/docs/source/user_guide/index.md @@ -23,6 +23,7 @@ maxdepth: 2 fundamentals/index fundamentals/battery_models fundamentals/public_api +rust_backend ``` ```{toctree} @@ -31,6 +32,7 @@ caption: Contributing guide maxdepth: 1 --- contributing +../developer/rust_idaklu_build release ``` diff --git a/docs/source/user_guide/installation/index.rst b/docs/source/user_guide/installation/index.rst index 855be9887f..bb7f26fd18 100644 --- a/docs/source/user_guide/installation/index.rst +++ b/docs/source/user_guide/installation/index.rst @@ -118,6 +118,13 @@ For an introduction to virtual environments, see Install PyBaMM -------------- +PyBaMM publishes platform-specific wheels for macOS (arm64, x86_64), Linux +(``manylinux`` x86_64 and aarch64), and Windows (AMD64). One wheel per platform +covers Python 3.10 to 3.14, because the bundled Rust compute core is built +against CPython's stable ABI. On any other platform ``pip`` falls back to the +source distribution, which compiles that core and therefore needs a Rust +toolchain — see :doc:`install-from-source`. + .. tab:: uv PyBaMM can be installed via `uv `__ from `PyPI `__. diff --git a/docs/source/user_guide/installation/install-from-source.rst b/docs/source/user_guide/installation/install-from-source.rst index 9cdf6b3f5e..e9f8369e6c 100644 --- a/docs/source/user_guide/installation/install-from-source.rst +++ b/docs/source/user_guide/installation/install-from-source.rst @@ -36,9 +36,18 @@ To install PyBaMM, you will need: - A BLAS library (for instance `openblas `_). - A C compiler (ex: ``gcc``). - A Fortran compiler (ex: ``gfortran``). +- A Rust toolchain (``cargo`` 1.89 or newer), for the Rust compute core — install via `rustup `__. - ``graphviz`` (optional), if you wish to build the documentation locally. - ``pandoc`` (optional) to convert the example Jupyter notebooks when building the documentation. +.. note:: + + PyBaMM's wheels on PyPI bundle a prebuilt Rust extension, so ``pip install pybamm`` + needs no Rust toolchain. Building from source or from the source distribution + compiles it, which requires ``cargo`` 1.89 or newer — the floor comes from the + crate's dependency graph, not from Rust edition 2024, which alone needs only 1.85. + To force a prebuilt wheel, use ``pip install --only-binary pybamm pybamm``. + You can install the above with .. tab:: Ubuntu/Debian diff --git a/docs/source/user_guide/rust_backend.rst b/docs/source/user_guide/rust_backend.rst new file mode 100644 index 0000000000..77e79e0f29 --- /dev/null +++ b/docs/source/user_guide/rust_backend.rst @@ -0,0 +1,148 @@ +.. _rust-backend: + +Rust Backend +============ + +PyBaMM ships a Rust compute core (``pybamm.rust``) that compiles a discretised +model's expression trees: right-hand sides, algebraic residuals, Jacobians, +events, and output variables into a form evaluated entirely outside Python. +Two solvers use it: + +- :class:`pybamm.IDAKLUSolver` with ``model.convert_to_format = "rust"`` runs + the usual SUNDIALS IDA integrator with every callback evaluated by the Rust + core instead of CasADi. Solver options, forward sensitivities, + ``output_variables``, events, and experiments work as they do with the + CasADi backend. +- :class:`pybamm.DiffsolSolver` integrates with the pure-Rust BDF + implementation from the `diffsol `_ + crate, with no CasADi or SUNDIALS involvement. It supports forward + sensitivities under solver error control, ``output_variables``-only solves, + events, and ``t_interp``. + +Selecting the backend +--------------------- + +Which backend a solver uses follows the model's ``convert_to_format`` +attribute (documented on :class:`pybamm.BaseModel`): + +.. code-block:: python + + import pybamm + + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + + sim = pybamm.Simulation(model) # the default IDAKLUSolver picks up "rust" + sol = sim.solve([0, 3600]) + +:class:`pybamm.DiffsolSolver` always uses the Rust backend and converts the +model itself: + +.. code-block:: python + + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve([0, 3600]) + +Output times with ``DiffsolSolver`` +----------------------------------- + +Both solvers store the full state trajectory at every output time; they +differ in which times those are. ``IDAKLUSolver`` uses its internal +integrator steps as the output grid when ``t_interp`` is omitted, whereas +diffsol evaluates its error-controlled dense output at the requested times +alone. A bare span — ``solve([t0, tf])`` with no ``t_interp`` — is answered +on a uniform 100-point grid; pass ``t_interp`` (or a ``t_eval`` of three or +more points) to choose the output times exactly. + +Off-grid reads — ``sol["Voltage [V]"](t)`` at a ``t`` between output points — +interpolate with cubic Hermite, as they do for ``IDAKLUSolver``: the solver +stores the state time derivatives alongside the solution by default. Pass +``hermite_interpolation=False`` to drop them and halve trajectory memory, at +the cost of off-grid reads falling back to linear interpolation. IDAKLU's +Hermite knots are its internal steps, while diffsol's are the requested output +times, so on a coarse output grid IDAKLU's off-grid reads remain the more +accurate of the two. + +Supported expression types +-------------------------- + +The Rust backend supports the following expression types: + +**Leaf nodes:** + +- ``Scalar`` — constant scalar values +- ``Array`` / ``Vector`` / ``Matrix`` — dense arrays and matrices +- ``SparseMatrix`` — CSR sparse matrices +- ``StateVector`` / ``StateVectorDot`` — state variable slices +- ``InputParameter`` — named input parameters +- ``Time`` — simulation time + +**Binary operators:** + +- Arithmetic: ``+``, ``-``, ``*``, ``/``, ``**`` +- Matrix: ``@`` (matrix multiplication) +- Comparison: ``minimum``, ``maximum`` +- Other: ``modulo``, ``hypot``, ``EqualHeaviside``, ``NotEqualHeaviside`` + +**Unary operators:** + +- ``-`` (negation), ``abs`` +- Math functions: ``sqrt``, ``exp``, ``log``, ``sin``, ``cos``, ``tanh``, + ``sinh``, ``cosh``, ``arcsinh``, ``arctan``, ``erf``, ``sign``, ``floor``, + ``ceiling`` +- Reductions: ``max``, ``min`` (over arrays), differentiated to the + argmax/argmin subgradient + +**Structural:** + +- ``Index`` — array slicing +- ``Concatenation`` — combining arrays +- ``Conditional`` — branch selection +- ``VectorField`` — stacked components (read scalar components back with + ``pybamm.Component``) + +**Interpolation:** + +- 1D ``linear``, ``cubic``, and ``pchip`` interpolation +- 2D and 3D ``linear`` and ``cubic`` regular-grid interpolation + +A symbol the backend cannot convert raises an error naming the symbol at +solver set-up, so an unsupported model fails loudly rather than +mis-evaluating. + +Adding new expression types +--------------------------- + +To add support for a new expression type: + +1. **Rust side** (if a new node type is needed): + + - Add a variant to the ``Node`` enum in + ``packages/pybamm-rust/pybamm-core/src/node.rs`` + - Implement evaluation in ``packages/pybamm-rust/pybamm-core/src/eval.rs`` + - Add a PyO3 binding in ``packages/pybamm-rust/pybamm-python/src/expr.rs`` + +2. **Python side:** + + - Add a ``_to_rust(self, graph, rust_symbols)`` method to the expression + class + - Follow the pattern of existing implementations (e.g. + ``binary_operators.py``) + +3. **Testing:** + + - Add a unit test in + ``packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_rust.py`` + - Add a parity test in + ``packages/pybamm/tests/integration/test_rust_parity.py`` + +Example ``_to_rust`` implementation: + +.. code-block:: python + + def _to_rust(self, graph, rust_symbols): + """Convert to Rust expression graph.""" + # Convert children first + converted_children = self._children_to_rust(graph, rust_symbols) + # Call appropriate ExprGraph method + return graph.some_method(*converted_children) diff --git a/packages/pybamm-rust/.gitignore b/packages/pybamm-rust/.gitignore new file mode 100644 index 0000000000..fcf99788dd --- /dev/null +++ b/packages/pybamm-rust/.gitignore @@ -0,0 +1,7 @@ +/target/ +*/target/ +*.so +*.dylib +*.pyd +.venv/ +uv.lock diff --git a/packages/pybamm-rust/Cargo.lock b/packages/pybamm-rust/Cargo.lock new file mode 100644 index 0000000000..4adca55ea4 --- /dev/null +++ b/packages/pybamm-rust/Cargo.lock @@ -0,0 +1,2155 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +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" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "is-terminal", + "itertools", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +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.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b324c954a5178683516e35e06051f02e3e2471d22b8175d9a01b274bd81269d" +dependencies = [ + "diffsol-la", + "diffsol-nl", + "log", + "num-traits", + "petgraph", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "diffsol-la" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fd836b0b52b6459b1e16c4ebfe267599b39c07a2eb638426a4e72a279a4e384" +dependencies = [ + "faer", + "faer-traits", + "nalgebra", + "num-traits", + "thiserror 2.0.19", +] + +[[package]] +name = "diffsol-nl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b24ab3f56ff94c563ab85e7897ba57f884617d0dc83a081028d28ec800213a1" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "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]] +name = "fastrand" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "linux-raw-sys" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "memchr" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", + "rand 0.8.7", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +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" +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]] +name = "numpy" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778da78c64ddc928ebf5ad9df5edf0789410ff3bdbf3619aed51cd789a6af1e2" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "paste" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + +[[package]] +name = "portable-atomic" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand 0.9.5", + "rand_chacha", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + +[[package]] +name = "pulp" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "046aa45b989642ec2e4717c8e72d677b13edd831a4d3b6cf37d9a3e54912496a" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "paste", + "pulp-wasm-simd-flag", + "raw-cpuid", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp-wasm-simd-flag" +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", + "diffsol-la", + "dyn-stack", + "faer", + "proptest", + "rand 0.9.5", + "rayon", + "rustc-hash", + "serde", + "thiserror 2.0.19", +] + +[[package]] +name = "pybamm-python" +version = "0.1.0" +dependencies = [ + "bincode", + "numpy", + "pybamm-core", + "pyo3", + "rayon", + "serde", +] + +[[package]] +name = "pyo3" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +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 0.9.5", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "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]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "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]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + +[[package]] +name = "unicode-ident" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "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 0.61.2", +] + +[[package]] +name = "windows-link" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/packages/pybamm-rust/Cargo.toml b/packages/pybamm-rust/Cargo.toml new file mode 100644 index 0000000000..921b5e7812 --- /dev/null +++ b/packages/pybamm-rust/Cargo.toml @@ -0,0 +1,131 @@ +[workspace] +resolver = "3" +members = [ + "pybamm-core", + "pybamm-python", +] + +# Floor set by the dep graph (nalgebra 0.35 needs 1.89), not by edition 2024. +# Resolver 3 is Rust-version-aware and keeps deps at or below this floor. +[workspace.package] +rust-version = "1.89" + +[workspace.lints.rust] +unsafe_code = "warn" +missing_debug_implementations = "warn" +rust_2018_idioms = { level = "warn", priority = -1 } +trivial_casts = "warn" +trivial_numeric_casts = "warn" +unused_lifetimes = "warn" +unused_qualifications = "warn" + +[workspace.lints.clippy] +# Correctness - these catch real bugs +correctness = { level = "deny", priority = -1 } + +# Pedantic - stricter but valuable +pedantic = { level = "warn", priority = -1 } + +# Nursery - experimental but useful +nursery = { level = "warn", priority = -1 } + +# Perf - catch performance lints +perf = { level = "warn", priority = -1 } + +# Complexity - express complex code in a simple form +complexity = { level = "warn", priority = -1 } + +# Style - Idiomatic rust code +style = { level = "warn", priority = -1 } + +# Specific high-value lints +cast_possible_truncation = "warn" +cast_possible_wrap = "warn" +cast_sign_loss = "warn" +checked_conversions = "warn" +clone_on_ref_ptr = "warn" +debug_assert_with_mut_call = "warn" +empty_enums = "warn" +enum_glob_use = "warn" +expl_impl_clone_on_copy = "warn" +explicit_deref_methods = "warn" +explicit_into_iter_loop = "warn" +explicit_iter_loop = "warn" +filter_map_next = "warn" +flat_map_option = "warn" +float_cmp = "warn" +fn_params_excessive_bools = "warn" +implicit_clone = "warn" +inefficient_to_string = "warn" +invalid_upcast_comparisons = "warn" +large_digit_groups = "warn" +large_stack_arrays = "warn" +large_types_passed_by_value = "warn" +manual_ok_or = "warn" +map_flatten = "warn" +map_unwrap_or = "warn" +match_bool = "warn" +match_same_arms = "warn" +match_wild_err_arm = "warn" +match_wildcard_for_single_variants = "warn" +mut_mut = "warn" +needless_borrow = "warn" +needless_continue = "warn" +needless_for_each = "warn" +needless_pass_by_value = "warn" +option_option = "warn" +pub_underscore_fields = "warn" +range_minus_one = "warn" +range_plus_one = "warn" +redundant_closure_for_method_calls = "warn" +redundant_else = "warn" +ref_option_ref = "warn" +same_functions_in_if_condition = "warn" +semicolon_if_nothing_returned = "warn" +string_add_assign = "warn" +struct_excessive_bools = "warn" +trait_duplication_in_bounds = "warn" +trivially_copy_pass_by_ref = "warn" +uninlined_format_args = "warn" +unnested_or_patterns = "warn" +unreadable_literal = "warn" +unused_self = "warn" +used_underscore_binding = "warn" +zero_sized_map_values = "warn" + +# Allow some pedantic lints that are too noisy +missing_errors_doc = "allow" +missing_panics_doc = "allow" +module_name_repetitions = "allow" +must_use_candidate = "allow" +too_many_lines = "allow" +similar_names = "allow" +# Numeric code uses indices with bounds guaranteed by construction +indexing_slicing = "allow" +# Benchmark/test data generation +cast_precision_loss = "allow" +# Mathematical notation (x, y, t, a, b) +many_single_char_names = "allow" +# PyO3 convention for uniform error handling +unnecessary_wraps = "allow" +# We ship baseline x86-64 wheels, where mul_add lowers to a libm call rather +# than an FMA instruction, and its single rounding breaks CasADi parity. +suboptimal_flops = "allow" + +[profile.release] +lto = "thin" +codegen-units = 1 + +[profile.bench] +lto = "thin" +codegen-units = 1 + +[profile.release-baseline] +inherits = "release" +lto = false +codegen-units = 16 + +[profile.bench-baseline] +inherits = "bench" +lto = false +codegen-units = 16 diff --git a/packages/pybamm-rust/README.md b/packages/pybamm-rust/README.md new file mode 100644 index 0000000000..8df7474cb9 --- /dev/null +++ b/packages/pybamm-rust/README.md @@ -0,0 +1,128 @@ +# PyBaMM Rust Core + +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, differentiate and +solve. + +Models opt in with `model.convert_to_format = "rust"`. + +## Layout + +Two crates: + +- **`pybamm-core`** — the compiler and interpreter. No Python dependency, so it can be + tested and benchmarked on its own. +- **`pybamm-python`** — PyO3 bindings, built as the `_core` extension and re-exported + through `pybamm.rust`. Owns the boundary checks, so core code can assume array lengths + and index ranges are already valid. + +## How a model becomes executable + +Five stages, each a module in `pybamm-core`: + +1. **Build** — the bindings allocate nodes into an arena. A node is referenced by a + small integer id, so sharing a subexpression is repeating an id rather than cloning a + subtree. +2. **Rewrite** — constant folding and algebraic identities, CSE and DCE, plus a pass + that proves subtrees identically zero and folds them away. +3. **Differentiate** — forward mode emits a tangent DAG for JVPs; reverse mode fills a + single wide Jacobian row from one backward pass. +4. **Lower** — the DAG flattens to fixed-size instructions addressing slots in one + scratch buffer, with constants held in a side table. +5. **Evaluate** — the interpreter walks a tape against a caller-supplied buffer, either + one time point at a time or several lanes at once. + +Assembling a sparse Jacobian is where the design earns its keep: a column coloring +groups columns that never share a row, so assembly costs one primal pass whose result +every color reuses, plus one tangent sweep per color. Rows too wide to color cheaply are +split out and filled by reverse mode instead, so a single dense row cannot force one +color per column. + +## 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 + runtime with `dlsym` once that extension is loaded. IDAKLU therefore has no undefined + symbols and loads standalone, with the Rust path resolved lazily on first use. The two + sides agree on an ABI version and refuse to proceed on a mismatch; bumping the Rust + constant means bumping the C++ one in the same change. + +## Building + +Build through the workspace, not `cargo` or `maturin` directly: + +```bash +uv sync --extra all --group dev +``` + +The pybamm package declares `cache-keys` over both crates' sources and manifests, so +`uv sync` and `uv run` rebuild the extension whenever Rust code changes. Reaching for +maturin by hand produces a wheel the Python side will not pick up. + +For work that stays inside Rust, the usual cargo commands apply from this directory: + +```bash +cargo build --release +cargo test +cargo clippy --all-targets -- -D warnings +``` + +## Testing + +```bash +cargo test # rust unit + integration +cargo test --all-features # includes feature-gated tests +uv run --group dev pytest -m unit packages/pybamm/tests # python side, from repo root +``` + +Rust and Python tests cover different things. The Rust suite owns the compiler and +interpreter, including property tests that check AD against finite differences and the +split primal/tangent tape against a monolithic one. The Python suite owns parity: the +same model solved through Rust and through CasADi must agree. + +## Benchmarking + +```bash +cargo bench -p pybamm-core --all-features +uv run python benchmarks/run_rust_observability.py --lane all # from repo root +``` + +Pass `--all-features`, or cargo silently skips the benches whose `required-features` are +unmet rather than telling you they did not run. + +The cargo benches measure the compiler and interpreter in isolation. The observability +harness is the one to trust for backend comparisons: it runs the same scenarios across +Rust, CasADi and CasADi's ahead-of-time path, shuffles backend order to keep machine +warm-up from flattering whichever runs last, and reports agreement against a baseline +rather than raw timings alone. Quote numbers from a run on the machine in question, not +from documentation. + +## Finding your way around + +Every module carries a doc comment explaining what it owns and why it exists: + +```bash +cargo doc --no-deps --open +``` + +Start at the crate root for the pipeline overview, then read the module you need. The +doc comments are the reference; this README only orients you. + +## Conventions + +- **Lints are strict.** The workspace turns on clippy's pedantic and nursery groups and + denies correctness. Fix warnings at the source; reach for a scoped `allow` only for a + confirmed false positive, with a comment saying why. +- **MSRV and dependency versions** live in the workspace `Cargo.toml`. The Rust floor is + set by the dependency graph rather than by the edition, so check there before assuming + a newer language feature is available. +- **Feature flags**: `diffsol` (on by default) pulls in the integrator and its matrix + backend; `serialize` enables snapshot support that some tests and benches require; + `profile` adds FFI call counters. 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 new file mode 100644 index 0000000000..c7634ec183 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/Cargo.toml @@ -0,0 +1,58 @@ +[package] +name = "pybamm-core" +version = "0.1.0" +edition = "2024" +rust-version.workspace = true + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +rustc-hash = "2" +thiserror = "2" +serde = { version = "1", features = ["derive"], optional = true } +bincode = { version = "1", optional = true } +diffsol = { version = "0.16.2", optional = true, default-features = false, features = ["faer", "nalgebra"] } +diffsol-la = { version = "0.1.1", optional = true } +faer = { version = "0.24", optional = true } +# MemStack scratch API for ReusedFaerLu; must unify with faer's own dyn-stack. +dyn-stack = { version = "0.13", optional = true } +# Schedules the batch solve layer; only the `solver` module uses it. +rayon = { version = "1", optional = true } + +[dev-dependencies] +criterion = { version = "0.5", features = ["html_reports"] } +proptest = "1" +rand = "0.9" + +[features] +default = ["diffsol"] +profile = [] +serialize = ["serde", "bincode"] +diffsol = ["dep:diffsol", "dep:diffsol-la", "dep:faer", "dep:dyn-stack", "dep:rayon"] + +[[bench]] +name = "compile" +harness = false + +[[bench]] +name = "eval" +harness = false + +[[bench]] +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"] + +[lints] +workspace = true diff --git a/packages/pybamm-rust/pybamm-core/benches/compile.rs b/packages/pybamm-rust/pybamm-core/benches/compile.rs new file mode 100644 index 0000000000..cab107cf5c --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/compile.rs @@ -0,0 +1,296 @@ +mod helpers; + +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use helpers::{build_coupled, identity_mass_matrix}; +use pybamm_core::{Arena, ModelEvaluator, Node, NodeId, TypedIr}; + +fn build_linear_chain(arena: &mut Arena, depth: usize, vec_len: usize) -> NodeId { + let y = arena.alloc(Node::StateVector { + start: 0, + end: vec_len, + }); + let mut current = y; + for _ in 0..depth { + current = arena.alloc(Node::Sin(current)); + current = arena.alloc(Node::Neg(current)); + } + current +} + +fn build_wide_fanout(arena: &mut Arena, width: usize) -> NodeId { + let mut terms = Vec::with_capacity(width); + for i in 0..width { + let y = arena.alloc(Node::StateVector { + start: i * 10, + end: (i + 1) * 10, + }); + let a = arena.alloc(Node::Sin(y)); + let b = arena.alloc(Node::Exp(a)); + terms.push(b); + } + arena.alloc(Node::Concat(terms)) +} + +fn build_high_fanout(arena: &mut Arena, fanout: usize) -> NodeId { + let y = arena.alloc(Node::StateVector { start: 0, end: 50 }); + let mut terms = Vec::with_capacity(fanout); + for _ in 0..fanout { + let a = arena.alloc(Node::Sin(y)); + let b = arena.alloc(Node::Exp(a)); + terms.push(b); + } + arena.alloc(Node::Concat(terms)) +} + +fn build_diamond_dag(arena: &mut Arena, depth: usize) -> NodeId { + let y = arena.alloc(Node::StateVector { start: 0, end: 10 }); + let two = arena.alloc(Node::Scalar(2.0)); + + let mut layer = vec![y]; + for _ in 0..depth { + let mut next_layer = Vec::new(); + for &node in &layer { + let a = arena.alloc(Node::Sin(node)); + let b = arena.alloc(Node::Mul(node, two)); + next_layer.push(a); + next_layer.push(b); + } + let cap = next_layer.len().min(16); + next_layer.truncate(cap); + layer = next_layer; + } + + let mut acc = layer[0]; + for &node in &layer[1..] { + acc = arena.alloc(Node::Add(acc, node)); + } + acc +} + +fn bench_dag_to_ir(c: &mut Criterion) { + let mut group = c.benchmark_group("dag_to_ir"); + group.sample_size(20); + + for depth in [10, 50, 100, 500] { + let mut arena = Arena::new(); + let root = build_linear_chain(&mut arena, depth, 10); + group.bench_with_input(BenchmarkId::new("linear_chain", depth), &depth, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + for width in [10, 50, 100, 500] { + let mut arena = Arena::new(); + let root = build_wide_fanout(&mut arena, width); + group.bench_with_input(BenchmarkId::new("wide_fanout", width), &width, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + for depth in [4, 6, 8, 10] { + let mut arena = Arena::new(); + let root = build_diamond_dag(&mut arena, depth); + group.bench_with_input(BenchmarkId::new("diamond", depth), &depth, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + for fanout in [20, 100] { + let mut arena = Arena::new(); + let root = build_high_fanout(&mut arena, fanout); + group.bench_with_input(BenchmarkId::new("high_fanout", fanout), &fanout, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + for n in [50, 100, 200, 500] { + let mut arena = Arena::new(); + let root = build_coupled(&mut arena, n); + group.bench_with_input(BenchmarkId::new("coupled", n), &n, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + group.finish(); +} + +fn bench_model_compile(c: &mut Criterion) { + let mut group = c.benchmark_group("compiled_model_new"); + group.sample_size(10); + + for n in [50, 100, 200, 500] { + group.bench_with_input(BenchmarkId::new("coupled", n), &n, |b, &n| { + b.iter(|| { + let mut arena = Arena::new(); + let root = build_coupled(&mut arena, n); + let mass = identity_mass_matrix(n); + let _model = ModelEvaluator::new(black_box(&arena), root, mass, n, 0); + }); + }); + } + + group.finish(); +} + +struct Case { + label: &'static str, + param: usize, +} + +fn bench_slot_stats(c: &mut Criterion) { + let mut group = c.benchmark_group("slot_stats"); + group.sample_size(20); + + let cases = [ + Case { + label: "linear_chain", + param: 100, + }, + Case { + label: "linear_chain", + param: 500, + }, + Case { + label: "wide_fanout", + param: 50, + }, + Case { + label: "wide_fanout", + param: 200, + }, + Case { + label: "high_fanout", + param: 20, + }, + Case { + label: "high_fanout", + param: 100, + }, + Case { + label: "coupled", + param: 100, + }, + Case { + label: "coupled", + param: 500, + }, + ]; + + println!( + "\n{:>25} {:>8} {:>8} {:>8} {:>6}", + "topology", "naive", "actual", "saved", "ratio" + ); + println!("{:-<25} {:-<8} {:-<8} {:-<8} {:-<6}", "", "", "", "", ""); + + for case in &cases { + let mut arena = Arena::new(); + let root = match case.label { + "linear_chain" => build_linear_chain(&mut arena, case.param, 10), + "wide_fanout" => build_wide_fanout(&mut arena, case.param), + "high_fanout" => build_high_fanout(&mut arena, case.param), + "coupled" => build_coupled(&mut arena, case.param), + _ => unreachable!(), + }; + + let stats = TypedIr::slot_stats(&arena, root); + let saved = stats.naive_size.saturating_sub(stats.buffer_size); + let tag = format!("{}({})", case.label, case.param); + println!( + "{tag:>25} {:>8} {:>8} {:>8} {:>6.3}", + stats.naive_size, stats.buffer_size, saved, stats.reuse_ratio + ); + + group.bench_with_input(BenchmarkId::new("compile", &tag), &case.param, |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&arena), black_box(root))); + }); + } + + group.finish(); +} + +#[cfg(feature = "serialize")] +fn bench_real_models(c: &mut Criterion) { + use pybamm_core::DagSnapshot; + + let mut group = c.benchmark_group("real_model_compile"); + group.sample_size(20); + + println!( + "\n{:>6} {:>8} {:>8} {:>8} {:>8} {:>6}", + "model", "states", "nodes", "naive", "actual", "ratio" + ); + println!( + "{:-<6} {:-<8} {:-<8} {:-<8} {:-<8} {:-<6}", + "", "", "", "", "", "" + ); + + for fixture in helpers::FIXTURES { + let snap = DagSnapshot::from_bytes(fixture.bytes); + let stats = TypedIr::slot_stats(&snap.arena, snap.root); + let eval_order = snap.arena.topological_order(snap.root); + + println!( + "{:>6} {:>8} {:>8} {:>8} {:>8} {:>6.3}", + fixture.name, + snap.n_states, + eval_order.len(), + stats.naive_size, + stats.buffer_size, + stats.reuse_ratio, + ); + + group.bench_with_input( + BenchmarkId::new("dag_to_ir", fixture.name), + &fixture.name, + |b, _| { + b.iter(|| TypedIr::from_arena(black_box(&snap.arena), black_box(snap.root))); + }, + ); + + group.bench_with_input( + BenchmarkId::new("dag_to_split_ir", fixture.name), + &fixture.name, + |b, _| { + b.iter(|| { + TypedIr::from_arena_split_eval(black_box(&snap.arena), black_box(snap.root)) + }); + }, + ); + + // Surfaces any memory-driven regression in the whole of + // `ModelEvaluator::new` that the focused detect_sparsity bench misses. + group.bench_with_input( + BenchmarkId::new("compiled_model_new", fixture.name), + &fixture.name, + |b, _| { + b.iter(|| { + let mass = snap + .mass_matrix + .clone() + .unwrap_or_else(|| identity_mass_matrix(snap.n_states)); + ModelEvaluator::new( + black_box(&snap.arena), + black_box(snap.root), + mass, + snap.n_states, + snap.n_params, + ) + }); + }, + ); + } + + group.finish(); +} + +#[cfg(not(feature = "serialize"))] +const fn bench_real_models(_c: &mut Criterion) {} + +criterion_group!( + benches, + bench_dag_to_ir, + bench_model_compile, + bench_slot_stats, + bench_real_models, +); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/benches/eval.rs b/packages/pybamm-rust/pybamm-core/benches/eval.rs new file mode 100644 index 0000000000..1404553520 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/eval.rs @@ -0,0 +1,249 @@ +mod helpers; + +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use helpers::{build_coupled, identity_mass_matrix}; +use pybamm_core::{ + Arena, CompiledExpr, ModelEvaluator, Node, NodeId, SimplifyMode, TangentInputs, TypedIr, + simplify_with_mode, tangent_wrt_states, +}; + +fn build_simple_expression(arena: &mut Arena, n_states: usize) -> NodeId { + let y = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let y_sq = arena.alloc(Node::Pow(y, two)); + let sin_y = arena.alloc(Node::Sin(y)); + arena.alloc(Node::Add(y_sq, sin_y)) +} + +fn bench_primal_eval(c: &mut Criterion) { + let mut group = c.benchmark_group("primal_eval"); + + for n in [10, 50, 100, 200, 500, 1000] { + let mut arena = Arena::new(); + let root = build_coupled(&mut arena, n); + let y: Vec = (0..n).map(|i| (i as f64) * 0.01).collect(); + let compiled = CompiledExpr::new(&arena, root); + let mut scratch = vec![0.0; compiled.scratch_len()]; + + group.bench_with_input(BenchmarkId::new("coupled", n), &n, |b, _| { + b.iter(|| { + let result = compiled.eval( + black_box(&mut scratch), + black_box(0.0), + black_box(&y), + black_box(&[]), + black_box(&[]), + ); + black_box(result); + }); + }); + } + + group.finish(); +} + +fn bench_jvp(c: &mut Criterion) { + let mut group = c.benchmark_group("jvp"); + + for n in [50, 100, 200, 500, 1000, 2000] { + let mut arena = Arena::new(); + let root = build_coupled(&mut arena, n); + let y: Vec = (0..n).map(|i| (i as f64) * 0.01).collect(); + let v: Vec = vec![1.0; n]; + let compiled = CompiledExpr::new(&arena, root); + let mut scratch = vec![0.0; compiled.scratch_len()]; + + group.bench_with_input(BenchmarkId::new("primal", n), &n, |b, _| { + b.iter(|| { + let result = compiled.eval( + black_box(&mut scratch), + black_box(0.0), + black_box(&y), + black_box(&[]), + black_box(&[]), + ); + black_box(result); + }); + }); + + group.bench_with_input(BenchmarkId::new("forward", n), &n, |b, _| { + b.iter(|| { + let tangent = TangentInputs { + dy: Some(&v), + dp: None, + }; + let result = compiled.eval_with_tangent( + black_box(&mut scratch), + black_box(0.0), + black_box(&y), + black_box(&[]), + black_box(&[]), + black_box(&tangent), + ); + black_box(result); + }); + }); + } + + group.finish(); +} + +fn bench_jacobian_symbolic(c: &mut Criterion) { + let mut group = c.benchmark_group("jacobian_symbolic"); + + for n in [10, 50, 100, 200, 500] { + let mut arena = Arena::new(); + let root = build_simple_expression(&mut arena, n); + + group.bench_with_input(BenchmarkId::new("build", n), &n, |b, _| { + b.iter(|| { + let mut a = Arena::new(); + let r = build_simple_expression(&mut a, black_box(n)); + let jac = tangent_wrt_states(&mut a, r); + let _jac = simplify_with_mode(&mut a, jac, SimplifyMode::Aggressive); + }); + }); + + let jac_root = tangent_wrt_states(&mut arena, root); + let jac_root = simplify_with_mode(&mut arena, jac_root, SimplifyMode::Aggressive); + let ir = TypedIr::from_arena(&arena, jac_root); + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y: Vec = (0..n).map(|i| (i as f64) * 0.01).collect(); + let v: Vec = vec![1.0; n]; + + group.bench_with_input(BenchmarkId::new("eval", n), &n, |b, _| { + b.iter(|| { + let tangent = TangentInputs { + dy: Some(black_box(&v)), + dp: None, + }; + let result = expr.eval_with_tangent( + black_box(&mut scratch), + black_box(0.0), + black_box(&y), + &[], + &[], + &tangent, + ); + black_box(result.len()); + }); + }); + } + + group.finish(); +} + +fn bench_jacobian_assembly(c: &mut Criterion) { + let mut group = c.benchmark_group("jacobian_assembly"); + + for n in [50, 100, 200, 500] { + let mut arena = Arena::new(); + let root = build_coupled(&mut arena, n); + let mass = identity_mass_matrix(n); + let mut model = ModelEvaluator::new(&arena, root, mass, n, 0); + let y: Vec = (0..n).map(|i| (i as f64) * 0.01).collect(); + let mut jac_data = vec![0.0; model.nnz()]; + + group.bench_with_input(BenchmarkId::new("csc_into", n), &n, |b, _| { + b.iter(|| { + model.set_cj(1.0); + model.assemble_jacobian_csc_into(black_box(0.0), black_box(&y), &[], &mut jac_data); + }); + }); + } + + group.finish(); +} + +#[cfg(feature = "serialize")] +fn bench_real_model_eval(c: &mut Criterion) { + use pybamm_core::DagSnapshot; + + let mut group = c.benchmark_group("real_model_eval"); + group.sample_size(20); + + for fixture in helpers::FIXTURES { + let snap = DagSnapshot::from_bytes(fixture.bytes); + // Unsimplified on purpose: `CompiledModel::new` skips the rhs simplify + // pass, so this is the tape the solver actually walks. + let compiled = CompiledExpr::new(&snap.arena, snap.root); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let (y, inputs) = helpers::fixture_state(snap.n_states, snap.n_params); + + group.bench_with_input( + BenchmarkId::new("rhs", fixture.name), + &fixture.name, + |b, _| { + b.iter(|| { + let result = compiled.eval( + black_box(&mut scratch), + black_box(0.0), + black_box(&y), + black_box(&[]), + black_box(&inputs), + ); + black_box(result); + }); + }, + ); + } + + group.finish(); +} + +#[cfg(not(feature = "serialize"))] +const fn bench_real_model_eval(_c: &mut Criterion) {} + +/// Per-call assembly on the real fixtures, which is the number any change to +/// the batched sweep or the constant table has to move. +#[cfg(feature = "serialize")] +fn bench_real_model_assembly(c: &mut Criterion) { + use pybamm_core::DagSnapshot; + + let mut group = c.benchmark_group("real_model_assembly"); + group.sample_size(20); + + for fixture in helpers::FIXTURES { + let snap = DagSnapshot::from_bytes(fixture.bytes); + let mass = identity_mass_matrix(snap.n_states); + let mut model = ModelEvaluator::new(&snap.arena, snap.root, mass, snap.n_states, 0); + let (y, inputs) = helpers::fixture_state(snap.n_states, snap.n_params); + let mut jac_data = vec![0.0; model.nnz()]; + + group.bench_with_input( + BenchmarkId::new("csc_into", fixture.name), + &fixture.name, + |b, _| { + b.iter(|| { + model.set_cj(1.0); + model.assemble_jacobian_csc_into( + black_box(0.0), + black_box(&y), + black_box(&inputs), + &mut jac_data, + ); + }); + }, + ); + } + + group.finish(); +} + +#[cfg(not(feature = "serialize"))] +const fn bench_real_model_assembly(_c: &mut Criterion) {} + +criterion_group!( + benches, + bench_primal_eval, + bench_jvp, + bench_jacobian_symbolic, + bench_jacobian_assembly, + bench_real_model_eval, + bench_real_model_assembly, +); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/benches/fixtures/dfn.bin b/packages/pybamm-rust/pybamm-core/benches/fixtures/dfn.bin new file mode 100644 index 0000000000..70b96aaae3 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/benches/fixtures/dfn.bin differ diff --git a/packages/pybamm-rust/pybamm-core/benches/fixtures/spm.bin b/packages/pybamm-rust/pybamm-core/benches/fixtures/spm.bin new file mode 100644 index 0000000000..844e099b53 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/benches/fixtures/spm.bin differ diff --git a/packages/pybamm-rust/pybamm-core/benches/fixtures/spme.bin b/packages/pybamm-rust/pybamm-core/benches/fixtures/spme.bin new file mode 100644 index 0000000000..0916759d63 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/benches/fixtures/spme.bin differ diff --git a/packages/pybamm-rust/pybamm-core/benches/helpers.rs b/packages/pybamm-rust/pybamm-core/benches/helpers.rs new file mode 100644 index 0000000000..360e57c7ca --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/helpers.rs @@ -0,0 +1,74 @@ +//! Shared benchmark scaffolding; not every bench target uses every helper. +#![allow(dead_code)] + +use pybamm_core::node::{CsrData, Shape}; +use pybamm_core::{Arena, Node, NodeId}; + +pub fn build_coupled(arena: &mut Arena, n_states: usize) -> NodeId { + let mut terms = Vec::new(); + for i in 0..n_states { + let y_i = arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }); + let y_next = arena.alloc(Node::StateVector { + start: (i + 1) % n_states, + end: (i + 1) % n_states + 1, + }); + let y_prev = arena.alloc(Node::StateVector { + start: (i + n_states - 1) % n_states, + end: (i + n_states - 1) % n_states + 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let y_sq = arena.alloc(Node::Pow(y_i, two)); + let sin_next = arena.alloc(Node::Sin(y_next)); + let prod = arena.alloc(Node::Mul(y_i, y_prev)); + let sum1 = arena.alloc(Node::Add(y_sq, sin_next)); + let term = arena.alloc(Node::Add(sum1, prod)); + terms.push(term); + } + arena.alloc(Node::Concat(terms)) +} + +pub fn identity_mass_matrix(n: usize) -> CsrData { + CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("identity mass matrix is valid") +} + +#[cfg(feature = "serialize")] +#[derive(Debug)] +pub struct Fixture { + pub name: &'static str, + pub bytes: &'static [u8], +} + +/// The `(y, inputs)` the real-model benches evaluate at. Shared so the eval and +/// assembly numbers stay comparable. +#[cfg(feature = "serialize")] +pub fn fixture_state(n_states: usize, n_params: usize) -> (Vec, Vec) { + let y = (0..n_states) + .map(|i| 0.01f64.mul_add(i as f64 / n_states as f64, 0.5)) + .collect(); + (y, vec![0.0; n_params]) +} + +#[cfg(feature = "serialize")] +pub const FIXTURES: &[Fixture] = &[ + Fixture { + name: "SPM", + bytes: include_bytes!("fixtures/spm.bin"), + }, + Fixture { + name: "SPMe", + bytes: include_bytes!("fixtures/spme.bin"), + }, + Fixture { + name: "DFN", + bytes: include_bytes!("fixtures/dfn.bin"), + }, +]; 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..d1d852d9c4 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/sens_params.rs @@ -0,0 +1,86 @@ +//! Benchmark: the saving from integrating only the requested sensitivity columns. +//! +//! A forward sensitivity solve integrates one augmented state vector per column, +//! so narrowing an `n`-parameter model to `k` requested columns should scale the +//! solve cost with `k` rather than `n`. The fixture makes every state depend on +//! every parameter, so per-column cost is uniform and the `k = 1` versus +//! `k = n_params` ratio is the whole signal. + +#![cfg(feature = "diffsol")] + +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, CompiledModelOptions, ModelEvaluator, Node}; + +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: &[usize]) -> ModelEvaluator { + 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)); + + ModelEvaluator::new_with_options( + &arena, + rhs, + identity_mass_matrix(N_STATES), + N_STATES, + N_PARAMS, + CompiledModelOptions::new().with_sensitivities(sens), + ) +} + +fn bench_sens_params(c: &mut Criterion) { + let mut group = c.benchmark_group("sens_params"); + let all: Vec = (0..N_PARAMS).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_with_input(BenchmarkId::from_parameter(k), &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..8a28ff1595 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/solver_setup.rs @@ -0,0 +1,227 @@ +//! Benchmark: solver setup cost vs solve cost. +//! +//! Measures `PreparedSolver::new()` separately from `solve()` to confirm +//! that building a fresh diffsol solver per solve is cheap relative to the +//! solve itself. This justifies the immutable `PreparedSolver` design, which +//! constructs a fresh `Workspace` and BDF solver on every `solve` rather than +//! holding a mutable, reused solver across calls. + +#![cfg(feature = "serialize")] + +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::{DagSnapshot, ModelEvaluator}; + +const SPM: &[u8] = include_bytes!("fixtures/spm.bin"); +const SPME: &[u8] = include_bytes!("fixtures/spme.bin"); +const DFN: &[u8] = include_bytes!("fixtures/dfn.bin"); + +struct TestCase { + name: &'static str, + bytes: &'static [u8], +} + +const CASES: &[TestCase] = &[ + TestCase { + name: "SPM", + bytes: SPM, + }, + TestCase { + name: "SPMe", + bytes: SPME, + }, + TestCase { + name: "DFN", + bytes: DFN, + }, +]; + +fn build_compiled_model(snap: &DagSnapshot) -> ModelEvaluator { + let mass = snap + .mass_matrix + .clone() + .unwrap_or_else(|| identity_mass_matrix(snap.n_states)); + ModelEvaluator::new(&snap.arena, snap.root, mass, snap.n_states, snap.n_params) +} + +/// Benchmark A: `PreparedSolver::new()` only (BDF solver construction). +/// +/// Uses `iter_batched` so that `build_compiled_model()` runs per iteration +/// but is NOT included in the timing. This isolates the diffsol/BDF setup +/// cost from the expression compilation cost. +fn bench_solver_setup(c: &mut Criterion) { + let mut group = c.benchmark_group("solver_setup"); + group.sample_size(100); + + for case in CASES { + let snap = DagSnapshot::from_bytes(case.bytes); + + group.bench_with_input( + BenchmarkId::from_parameter(case.name), + &case.name, + |b, _| { + b.iter_batched( + || build_compiled_model(&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(); +} + +/// `ModelEvaluator::new()` only: differentiation, simplification, sparsity +/// detection, coloring and bytecode compilation. +fn bench_model_compilation(c: &mut Criterion) { + let mut group = c.benchmark_group("model_compilation"); + group.sample_size(50); + + for case in CASES { + let snap = DagSnapshot::from_bytes(case.bytes); + + group.bench_with_input( + BenchmarkId::from_parameter(case.name), + &case.name, + |b, _| { + b.iter(|| { + let model = build_compiled_model(&snap); + black_box(model); + }); + }, + ); + } + + 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 CASES { + let snap = DagSnapshot::from_bytes(case.bytes); + let model = build_compiled_model(&snap); + let n_states = model.n_states(); + let n_params = model.n_params(); + let atol = vec![1e-6; n_states]; + + let prepared = PreparedSolver::new(model, 1e-6, &atol).unwrap(); + + // Use a small perturbation around 0.5 as initial state, physically + // more reasonable than all-ones for battery models. + let y0: Vec = (0..n_states) + .map(|i| 0.01f64.mul_add(i as f64 / n_states as f64, 0.5)) + .collect(); + let inputs: Vec = vec![0.0; n_params]; + let t_eval: Vec = (0..100).map(|i| f64::from(i) * 36.0).collect(); + + // Warmup: skip this model if the dummy y0 doesn't converge + 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 CASES { + let snap = DagSnapshot::from_bytes(case.bytes); + let n_states = snap.n_states; + let n_params = snap.n_params; + + let y0: Vec = (0..n_states) + .map(|i| 0.01f64.mul_add(i as f64 / n_states as f64, 0.5)) + .collect(); + let inputs: Vec = vec![0.0; n_params]; + let t_eval: Vec = (0..100).map(|i| f64::from(i) * 36.0).collect(); + + // Skip models where dummy y0 doesn't converge + { + let model = build_compiled_model(&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( + || build_compiled_model(&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_model_compilation, + bench_solver_setup, + bench_solve_reuse, + bench_fresh_setup_and_solve, +); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/benches/sparsity.rs b/packages/pybamm-rust/pybamm-core/benches/sparsity.rs new file mode 100644 index 0000000000..47d8c44633 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/benches/sparsity.rs @@ -0,0 +1,37 @@ +//! Focused microbench for sparsity analysis on real model fixtures. + +#![cfg(feature = "serialize")] + +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use pybamm_core::{DagSnapshot, TypedIr, detect_sparsity_per_output}; + +const SPM: &[u8] = include_bytes!("fixtures/spm.bin"); +const SPME: &[u8] = include_bytes!("fixtures/spme.bin"); +const DFN: &[u8] = include_bytes!("fixtures/dfn.bin"); + +fn bench_detect_sparsity(c: &mut Criterion) { + let mut group = c.benchmark_group("detect_sparsity_per_output"); + group.sample_size(50); + + for (name, bytes) in [("SPM", SPM), ("SPMe", SPME), ("DFN", DFN)] { + let snap = DagSnapshot::from_bytes(bytes); + let ir = TypedIr::from_arena(&snap.arena, snap.root); + let n_outputs = ir.output_len(); + + group.bench_with_input(BenchmarkId::from_parameter(name), &name, |b, _| { + b.iter(|| { + detect_sparsity_per_output( + black_box(&snap.arena), + black_box(snap.root), + black_box(n_outputs), + black_box(snap.n_states), + ) + }); + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_detect_sparsity); +criterion_main!(benches); diff --git a/packages/pybamm-rust/pybamm-core/src/adjoint.rs b/packages/pybamm-rust/pybamm-core/src/adjoint.rs new file mode 100644 index 0000000000..2aec95a361 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/adjoint.rs @@ -0,0 +1,893 @@ +//! Reverse-mode (VJP) assembly of dense Jacobian rows. +//! +//! Wide rows split out of the column coloring are filled by backward passes +//! over their primal sub-expression instead of one forward JVP sweep per +//! column. The sub-expression is compiled with a no-reuse (SSA) slot layout +//! ([`CompiledExpr::new_pinned`]), so after one primal evaluation the scratch +//! buffer holds every intermediate; each instruction's operand slots are then +//! stable value-tape offsets the adjoint reads directly. The backward `match` +//! over [`Instruction`] is exhaustive: every op has an adjoint, so there is no +//! runtime fallback. +//! +//! A tape's root may be wider than one element, in which case seeding element +//! `r` recovers row `r`, so a group of rows shares one forward pass and one +//! compiled tape. See [`crate::row_extract`] for how many rows share one. +//! +//! The pinned layout forms branch blocks like every other layout, so the +//! backward walk jumps over the blocks of inactive conditional branches instead +//! of replaying adjoints that are all no-ops. + +// `NodeId` must come from `arena`: `node.rs` only re-imports it privately, so +// `crate::node::NodeId` is not reachable from another module. +use crate::arena::{Arena, NodeId}; +use crate::branch_regions::{active_branch, dispatch_span_end}; +use crate::eval::{ + CompiledExpr, interp_cubic_1d_deriv, interp_linear_1d_deriv, locate_nd_cell, sign, + tensor_horner_partial, +}; +use crate::ir::{BinaryOp, BroadcastKind, ConstPool, Instruction, TypedIr, UnaryOp}; + +/// Prepared adjoint (reverse-mode AD) tape for one expression, whose rows are +/// recovered one seeded backward pass at a time. +#[derive(Debug, Clone)] +pub struct AdjointTape { + expr: CompiledExpr, + n_states: usize, + /// For each instruction index, the index of the `Dispatch` whose block span + /// *ends* there, or `u32::MAX`. Precomputed once so the hot backward walk + /// allocates nothing. + span_owner: Vec, +} + +impl AdjointTape { + /// Compile `root` with a no-reuse (SSA) layout so the primal scratch is the + /// reverse value tape. Its width is how many rows the tape can recover. + /// + /// # Panics + /// Panics if `root` contains a derivative-only instruction (a tangent load + /// or `ReduceArgSelect`); a primal row lifted out of the residual never + /// does. + pub fn new(arena: &Arena, root: NodeId, n_states: usize) -> Self { + let expr = CompiledExpr::new_pinned(arena, root); + // Assert rather than assume: `ReduceArgSelect` and the tangent loads are + // built only by `tangent.rs`, never from Python, so no primal row has them. + for instr in expr.ir().instructions() { + assert!( + !matches!( + instr, + Instruction::LoadTangentState { .. } + | Instruction::LoadTangentParameter { .. } + | Instruction::ReduceArgSelect { .. } + ), + "adjoint tape requires a primal row: found derivative-only {instr:?}" + ); + } + let span_owner = build_span_owner(expr.ir()); + Self { + expr, + n_states, + span_owner, + } + } + + /// Tape length, so what this artifact costs in compiled memory. + #[inline] + pub fn instruction_count(&self) -> usize { + self.expr.ir().instructions().len() + } + + /// Scratch length for the value tape (and the parallel `bar` buffer). + #[inline] + pub const fn scratch_len(&self) -> usize { + self.expr.scratch_len() + } + + /// State dimension the gradient row spans. + #[inline] + pub const fn n_states(&self) -> usize { + self.n_states + } + + /// Rows this tape can recover, so its root's width. + #[inline] + pub const fn n_rows(&self) -> usize { + self.expr.output_len() + } + + /// Fill the value tape, which every row's backward pass then reads. + /// + /// Split out from [`assemble_row`](Self::assemble_row) so a block of rows + /// pays for the shared forward work once rather than per row. + pub fn eval_forward( + &self, + scratch: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) { + self.expr.eval(scratch, t, y, y_dot, inputs); + } + + /// Assemble row `row` into `grad[..n_states]` from an already-filled + /// `scratch`, as left by [`eval_forward`](Self::eval_forward), returning how + /// many instructions the backward walk touched. + /// + /// # Panics + /// Panics if `row` is outside the root's width. + pub fn assemble_row( + &self, + scratch: &[f64], + bar: &mut [f64], + grad: &mut [f64], + row: usize, + ) -> usize { + self.seed_row(bar, grad, row); + backward(self.expr.ir(), &self.span_owner, scratch, bar, grad) + } + + /// Clear the gradient and `bar` buffers and seed row `row`'s adjoint. + /// + /// # Panics + /// Panics if `row` is outside the root's width. + fn seed_row(&self, bar: &mut [f64], grad: &mut [f64], row: usize) { + let ir = self.expr.ir(); + let root = ir.root_slot(); + assert!( + row < root.len_usize(), + "row {row} is outside the tape's {} rows", + root.len_usize() + ); + grad[..self.n_states].fill(0.0); + bar[..ir.buffer_size()].fill(0.0); + bar[root.offset_usize() + row] = 1.0; + } + + /// Forward pass then row 0, so the whole gradient row `df/dy` in one call, + /// returning how many instructions the backward walk touched. + /// + /// `scratch` and `bar` must each be at least [`scratch_len`](Self::scratch_len); + /// `grad` at least `n_states`. All three are caller-provided and reset here, + /// no allocation occurs. Skipped blocks are excluded from the count, so a + /// test can assert an inactive branch's adjoint arms never ran. + #[allow(clippy::too_many_arguments)] + pub fn assemble( + &self, + scratch: &mut [f64], + bar: &mut [f64], + grad: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> usize { + // Forward record: fill the value tape (unique slot per node). + self.expr.eval(scratch, t, y, y_dot, inputs); + self.assemble_row(scratch, bar, grad, 0) + } + + #[cfg(test)] + pub(crate) const fn expr_for_test(&self) -> &CompiledExpr { + &self.expr + } +} + +/// Argmax (`is_max`) or argmin of `vals`, taking the earliest element under a +/// strict comparison to match the primal `MaxReduce`/`MinReduce` eval. +fn arg_reduce(vals: &[f64], is_max: bool) -> usize { + let mut best = 0; + for i in 1..vals.len() { + let better = if is_max { + vals[i] > vals[best] + } else { + vals[i] < vals[best] + }; + if better { + best = i; + } + } + best +} + +/// Reverse of `broadcast_apply`: accumulate operand adjoints given the two +/// per-element partials `pa = ∂f/∂x` and `pb = ∂f/∂y`. The scalar operand of a +/// broadcast has its `bar` sum-reduced (the transpose of broadcasting). A zero +/// output adjoint skips the local partial, which may be undefined on an inactive branch. +#[allow(clippy::too_many_arguments)] +#[inline] +fn reverse_broadcast( + scratch: &[f64], + bar: &mut [f64], + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, + pa: Fa, + pb: Fb, +) where + Fa: Fn(f64, f64) -> f64, + Fb: Fn(f64, f64) -> f64, +{ + match kind { + BroadcastKind::ScalarScalar => { + let bd = bar[dst]; + if bd == 0.0 { + return; + } + let (x, y) = (scratch[a], scratch[b]); + bar[a] += pa(x, y) * bd; + bar[b] += pb(x, y) * bd; + }, + BroadcastKind::ScalarVector => { + let x = scratch[a]; + let mut acc = 0.0; + for i in 0..len { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let y = scratch[b + i]; + acc += pa(x, y) * bd; + bar[b + i] += pb(x, y) * bd; + } + bar[a] += acc; + }, + BroadcastKind::VectorScalar => { + let y = scratch[b]; + let mut acc = 0.0; + for i in 0..len { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let x = scratch[a + i]; + bar[a + i] += pa(x, y) * bd; + acc += pb(x, y) * bd; + } + bar[b] += acc; + }, + BroadcastKind::VectorVector => { + for i in 0..len { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let (x, y) = (scratch[a + i], scratch[b + i]); + bar[a + i] += pa(x, y) * bd; + bar[b + i] += pb(x, y) * bd; + } + }, + } +} + +/// Dispatch the binary adjoint by op. Partials mirror `eval_binary_op`. +#[allow(clippy::suboptimal_flops, clippy::too_many_arguments)] +fn reverse_binary( + op: BinaryOp, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, + scratch: &[f64], + bar: &mut [f64], +) { + match op { + BinaryOp::Add => { + reverse_broadcast(scratch, bar, a, b, dst, len, kind, |_, _| 1.0, |_, _| 1.0); + }, + BinaryOp::Sub => { + reverse_broadcast(scratch, bar, a, b, dst, len, kind, |_, _| 1.0, |_, _| -1.0); + }, + BinaryOp::Mul => reverse_broadcast(scratch, bar, a, b, dst, len, kind, |_, y| y, |x, _| x), + BinaryOp::Div => { + reverse_broadcast( + scratch, + bar, + a, + b, + dst, + len, + kind, + |_, y| 1.0 / y, + |x, y| -x / (y * y), + ); + }, + BinaryOp::Pow => reverse_broadcast( + scratch, + bar, + a, + b, + dst, + len, + kind, + |x, y| y * x.powf(y - 1.0), + |x, y| if x > 0.0 { x.powf(y) * x.ln() } else { 0.0 }, + ), + BinaryOp::Minimum => reverse_broadcast( + scratch, + bar, + a, + b, + dst, + len, + kind, + |x, y| if x <= y { 1.0 } else { 0.0 }, + |x, y| if x <= y { 0.0 } else { 1.0 }, + ), + BinaryOp::Maximum => reverse_broadcast( + scratch, + bar, + a, + b, + dst, + len, + kind, + |x, y| if x >= y { 1.0 } else { 0.0 }, + |x, y| if x >= y { 0.0 } else { 1.0 }, + ), + BinaryOp::Modulo => { + reverse_broadcast(scratch, bar, a, b, dst, len, kind, |_, _| 1.0, |_, _| 0.0); + }, + BinaryOp::Hypot => reverse_broadcast( + scratch, + bar, + a, + b, + dst, + len, + kind, + |x, y| x / x.hypot(y), + |x, y| y / x.hypot(y), + ), + BinaryOp::EqualHeaviside | BinaryOp::NotEqualHeaviside | BinaryOp::Equality => { + reverse_broadcast(scratch, bar, a, b, dst, len, kind, |_, _| 0.0, |_, _| 0.0); + }, + } +} + +/// Elementwise unary adjoint: `bar[src+i] += f'(x) * bar[dst+i]`. Derivatives +/// mirror `eval_unary_op`; zero adjoints skip potentially undefined partials. +#[allow(clippy::suboptimal_flops)] +fn reverse_unary( + op: UnaryOp, + src: usize, + dst: usize, + len: usize, + scratch: &[f64], + bar: &mut [f64], +) { + for i in 0..len { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let x = scratch[src + i]; + let out = scratch[dst + i]; + let d = match op { + UnaryOp::Neg => -1.0, + UnaryOp::Abs => sign(x), + UnaryOp::Sqrt => 0.5 / out, + UnaryOp::Exp => out, + UnaryOp::Log => 1.0 / x, + UnaryOp::Sin => x.cos(), + UnaryOp::Cos => -x.sin(), + UnaryOp::Tanh => 1.0 - out * out, + UnaryOp::Sinh => x.cosh(), + UnaryOp::Cosh => x.sinh(), + UnaryOp::Arcsinh => 1.0 / (1.0 + x * x).sqrt(), + UnaryOp::Arctan => 1.0 / (1.0 + x * x), + UnaryOp::Erf => 2.0 / std::f64::consts::PI.sqrt() * (-x * x).exp(), + UnaryOp::Sign | UnaryOp::Floor | UnaryOp::Ceiling => 0.0, + }; + bar[src + i] += d * bd; + } +} + +/// Map each `Dispatch` span's last instruction index to the `Dispatch` itself. +/// +/// Blocks are flat and pairwise disjoint (a nested conditional is never +/// blocked), so one slot per index is enough. Both assertions ship rather than +/// being debug-only, for the same reason as `ir.rs`'s block-safety guards: an +/// overlapping span would silently drop one owner, and the failure mode is wrong +/// gradients, not a crash. This runs once per compile, off the hot path. +/// +/// # Panics +/// Panics if two `Dispatch` spans end at the same instruction. +fn build_span_owner(ir: &TypedIr) -> Vec { + let instrs = ir.instructions(); + let consts = ir.consts(); + let mut span_owner = vec![u32::MAX; instrs.len()]; + for (i, instr) in instrs.iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + let end = dispatch_span_end(consts, i, blocks_idx, blocks_len); + assert!(end > i, "a dispatch span must cover at least itself"); + assert_eq!( + span_owner[end - 1], + u32::MAX, + "overlapping dispatch spans: blocks must be flat and disjoint" + ); + span_owner[end - 1] = + u32::try_from(i).expect("instruction indices are u32 throughout the IR"); + } + } + span_owner +} + +/// Backward (adjoint) replay over a pinned primal instruction stream, jumping +/// over the blocks of inactive conditional branches. Returns the number of +/// instructions touched, and expects `bar[root] = 1` already seeded. +/// +/// Walking backwards reaches a block before its `Dispatch`, so `span_owner` says +/// which indices end a span; there the active branch is resolved from the +/// recorded primal selector and only that block is replayed. Skipping is sound +/// because the `Conditional` arm pushes the output adjoint into the matched +/// branch alone, so an inactive branch's slots carry `bar == 0`. +fn backward( + ir: &TypedIr, + span_owner: &[u32], + scratch: &[f64], + bar: &mut [f64], + grad: &mut [f64], +) -> usize { + let consts = ir.consts(); + let instructions = ir.instructions(); + let mut walked = 0_usize; + let mut i = instructions.len(); + while i > 0 { + i -= 1; + let owner = span_owner[i]; + if owner != u32::MAX { + let d = owner as usize; + let Instruction::Dispatch { + selector, + blocks_idx, + blocks_len, + } = instructions[d] + else { + unreachable!("span_owner points at a Dispatch") + }; + walked += 1; + if let Some(active) = active_branch(scratch[selector as usize], blocks_len as usize) { + let (rel, len) = consts.branch_blocks[blocks_idx as usize + active]; + let start = d + rel as usize; + // A block never contains a `Dispatch`, so replaying one needs no + // further span handling. + for k in (start..start + len as usize).rev() { + backward_instruction(consts, instructions[k], scratch, bar, grad); + } + walked += len as usize; + } + i = d; // the Dispatch itself has no adjoint + continue; + } + walked += 1; + backward_instruction(consts, instructions[i], scratch, bar, grad); + } + walked +} + +/// Accumulate one instruction's adjoint contribution into `bar` (and `grad` for +/// a state load). Exhaustive over `Instruction`, no fallback. +#[inline] +fn backward_instruction( + consts: &ConstPool, + instr: Instruction, + scratch: &[f64], + bar: &mut [f64], + grad: &mut [f64], +) { + match instr { + // Terminal instructions: their adjoint stops here. Interpolant + // derivatives are constant w.r.t. input (forward returns zero). + Instruction::LoadScalar { .. } + | Instruction::LoadTime { .. } + | Instruction::LoadArray { .. } + | Instruction::FillZero { .. } + | Instruction::LoadStateVectorDot { .. } + | Instruction::LoadInputParameter { .. } + | Instruction::Interp1DLinearDeriv { .. } + | Instruction::Interp1DCubicDeriv { .. } + | Instruction::InterpNdPartial { .. } => {}, + + // The gradient sink. + Instruction::LoadStateVector { start, end, dst } => { + let dst = dst as usize; + for (k, s) in (start as usize..end as usize).enumerate() { + grad[s] += bar[dst + k]; + } + }, + + // Derivative-only instructions are rejected in `AdjointTape::new`; this + // arm only keeps the `match` exhaustive with no runtime fallback. + Instruction::LoadTangentState { .. } + | Instruction::LoadTangentParameter { .. } + | Instruction::ReduceArgSelect { .. } => { + unreachable!("derivative-only instruction in primal adjoint tape") + }, + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => { + reverse_binary( + op, + a as usize, + b as usize, + dst as usize, + len as usize, + kind, + scratch, + bar, + ); + }, + Instruction::Unary { op, src, dst, len } => { + reverse_unary(op, src as usize, dst as usize, len as usize, scratch, bar); + }, + + Instruction::MaxReduce { src, src_len, dst } => { + let bd = bar[dst as usize]; + if bd == 0.0 { + return; + } + let src = src as usize; + let k = arg_reduce(&scratch[src..src + src_len as usize], true); + bar[src + k] += bd; + }, + Instruction::MinReduce { src, src_len, dst } => { + let bd = bar[dst as usize]; + if bd == 0.0 { + return; + } + let src = src as usize; + let k = arg_reduce(&scratch[src..src + src_len as usize], false); + bar[src + k] += bd; + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + let src = src as usize + start as usize; + let dst = dst as usize; + for i in 0..len as usize { + bar[src + i] += bar[dst + i]; + } + }, + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let dst = dst as usize; + let mut o = 0usize; + for s in 0..sources_len as usize { + let (off, slen) = consts.concat_sources[sources_idx as usize + s]; + for i in 0..slen as usize { + bar[off as usize + i] += bar[dst + o + i]; + } + o += slen as usize; + } + }, + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + // dst = A @ vec ; A constant, so bar_vec += Aᵀ @ bar_dst. + let mat = &consts.csr_data[csr_idx as usize]; + let vec_src = vec_src as usize; + let dst = dst as usize; + for row in 0..mat.shape.rows { + let bd = bar[dst + row]; + if bd == 0.0 { + continue; + } + for k in mat.indptr[row]..mat.indptr[row + 1] { + bar[vec_src + mat.indices[k]] += mat.data[k] * bd; + } + } + }, + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + dst, + } => { + // dst = A @ vec, A row-major constant ; bar_vec += Aᵀ @ bar_dst. + let mat_src = mat_src as usize; + let vec_src = vec_src as usize; + let dst = dst as usize; + let cols = cols as usize; + for row in 0..rows as usize { + let bd = bar[dst + row]; + if bd == 0.0 { + continue; + } + let base = mat_src + row * cols; + for j in 0..cols { + bar[vec_src + j] += scratch[base + j] * bd; + } + } + }, + + Instruction::Interp1DLinear { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let slope = + interp_linear_1d_deriv(&interp.x_data, &interp.y_data, scratch[src + i]); + bar[src + i] += slope * bd; + } + }, + Instruction::Interp1DCubic { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + let d = + interp_cubic_1d_deriv(&interp.breakpoints, &interp.coeffs, scratch[src + i]); + bar[src + i] += d * bd; + } + }, + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len, + } => { + let interp = &consts.nd_interpolants[interp_idx as usize]; + let ndim = interp.breakpoints.len(); + let order = interp.order as usize; + let dst = dst as usize; + let mut coords = [0.0_f64; 3]; + let mut dxs = [0.0_f64; 3]; + for i in 0..len as usize { + let bd = bar[dst + i]; + if bd == 0.0 { + continue; + } + for (a, coord) in coords.iter_mut().enumerate().take(ndim) { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + let j = if slen == 1 { 0 } else { i }; + *coord = scratch[off as usize + j]; + } + let cell = locate_nd_cell( + &interp.breakpoints, + &interp.coeffs, + order, + &coords[..ndim], + &mut dxs, + ); + for a in 0..ndim { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + let j = if slen == 1 { 0 } else { i }; + let partial = tensor_horner_partial(cell, &dxs[..ndim], order, a); + bar[off as usize + j] += partial * bd; + } + } + }, + + Instruction::Conditional { + selector, + branches_idx, + branches_len, + dst, + out_len, + } => { + let dst = dst as usize; + let out_len = out_len as usize; + if let Some(i) = active_branch(scratch[selector as usize], branches_len as usize) { + let (off, _) = consts.branch_offsets[branches_idx as usize + i]; + let off = off as usize; + for k in 0..out_len { + bar[off + k] += bar[dst + k]; + } + } + }, + + Instruction::Dispatch { .. } => unreachable!("handled via the span table"), + } +} + +#[cfg(test)] +mod tests { + use crate::adjoint::AdjointTape; + use crate::arena::Arena; + use crate::node::{CsrData, Node, Shape}; + + fn assemble(tape: &AdjointTape, y: &[f64]) -> Vec { + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; tape.n_states()]; + tape.assemble(&mut scratch, &mut bar, &mut grad, 0.0, y, &[], &[]); + grad + } + + #[test] + fn test_reverse_mul_scalar_scalar() { + // f = y0 * y1 ; grad = [y1, y0, 0]. y=[2,3,5] -> [3,2,0]. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y0 = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y, + start: 1, + end: 2, + }); + let f = arena.alloc(Node::Mul(y0, y1)); + let tape = AdjointTape::new(&arena, f, 3); + let grad = assemble(&tape, &[2.0, 3.0, 5.0]); + assert_eq!(grad, vec![3.0, 2.0, 0.0]); + } + + #[test] + fn test_reverse_matmul_and_broadcast_sum() { + // s = y0, v = [y1,y2,y3], p = s*v (ScalarVector broadcast), f = [1,1,1] @ p. + // grad = [y1+y2+y3, s, s, s], so y = [2,3,4,5] -> [12, 2, 2, 2]. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let s = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let v = arena.alloc(Node::Index { + child: y, + start: 1, + end: 4, + }); + let p = arena.alloc(Node::Mul(s, v)); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 3], + vec![0, 1, 2], + vec![1.0, 1.0, 1.0], + Shape::matrix(1, 3), + ) + .unwrap(), + ))); + let f = arena.alloc(Node::MatMul(ones, p)); + let tape = AdjointTape::new(&arena, f, 4); + let grad = assemble(&tape, &[2.0, 3.0, 4.0, 5.0]); + assert_eq!(grad, vec![12.0, 2.0, 2.0, 2.0]); + } + + #[test] + fn test_reverse_exp_index() { + // f = exp(y2) ; grad = [0, 0, exp(y2)]. y2=0 -> [0,0,1]. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y2 = arena.alloc(Node::Index { + child: y, + start: 2, + end: 3, + }); + let f = arena.alloc(Node::Exp(y2)); + let tape = AdjointTape::new(&arena, f, 3); + let grad = assemble(&tape, &[1.0, 1.0, 0.0]); + assert!((grad[2] - 1.0).abs() < 1e-14 && grad[0] == 0.0 && grad[1] == 0.0); + } + + #[test] + fn test_inactive_conditional_unary_does_not_poison_gradient() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let invalid = arena.alloc(Node::Sqrt(y)); + let selector = arena.alloc(Node::Scalar(1.0)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![y, invalid], + }); + + let tape = AdjointTape::new(&arena, root, 1); + let grad = assemble(&tape, &[-1.0]); + assert_eq!(grad, vec![1.0]); + } + + #[test] + fn test_inactive_conditional_binary_does_not_poison_gradient() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let invalid = arena.alloc(Node::Div(y, zero)); + let selector = arena.alloc(Node::Scalar(1.0)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![y, invalid], + }); + + let tape = AdjointTape::new(&arena, root, 1); + let grad = assemble(&tape, &[2.0]); + assert_eq!(grad, vec![1.0]); + } + + #[test] + fn test_pinned_reverse_tape_has_no_self_aliasing_binary() { + // Regression guard for the eval.rs split_dst_two_src alias panic: the + // no-reuse layout must never place a binary operand on top of its dst. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y, + start: 1, + end: 2, + }); + let f = arena.alloc(Node::Mul(y0, y1)); + let tape = AdjointTape::new(&arena, f, 2); + for instr in tape.expr_for_test().ir().instructions() { + if let crate::ir::Instruction::Binary { a, b, dst, len, .. } = *instr { + let (a, b, dst, len) = (a as usize, b as usize, dst as usize, len as usize); + assert!(a + len <= dst || dst + len <= a, "operand a aliases dst"); + assert!(b + len <= dst || dst + len <= b, "operand b aliases dst"); + } + } + } + + #[test] + #[should_panic(expected = "primal row")] + fn test_reverse_tape_rejects_derivative_only_node() { + // `ReduceArgSelect` is derivative-only: rejected at build time, not + // mis-differentiated later. Width-one operands let the reject fire first. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let basis = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let picker = arena.alloc(Node::Index { + child: y, + start: 1, + end: 2, + }); + let ras = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let _ = AdjointTape::new(&arena, ras, 2); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/arena.rs b/packages/pybamm-rust/pybamm-core/src/arena.rs new file mode 100644 index 0000000000..05a01bbe19 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/arena.rs @@ -0,0 +1,355 @@ +//! Flat storage for the expression DAG. +//! +//! Every [`Node`] of a model lives in one `Vec` inside an [`Arena`] and is named +//! by a `u32` [`NodeId`]. Ids are `Copy` and stable for the arena's lifetime, so +//! a shared subexpression is one node referenced twice rather than a cloned +//! subtree, and passes can annotate nodes out-of-band through [`NodeMap`] +//! instead of touching the DAG. +//! +//! Ids are only meaningful against the arena that issued them: rewriting passes +//! that build a new arena (`cse`, `dce`, `privatise_conditionals`) return that +//! arena with a new root, and mixing the two id spaces indexes the wrong node. +//! `simplify` instead appends to the arena it is given, so its root stays valid +//! against the same arena. + +use std::ops::Index; + +use crate::node::Node; + +/// Handle to a node in one [`Arena`], valid only against that arena. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct NodeId(u32); + +impl NodeId { + /// Position in the arena, for indexing side tables. + #[inline] + pub const fn index(self) -> usize { + self.0 as usize + } + + /// The underlying `u32`, as the instruction tape and FFI carry it. + #[inline] + pub const fn raw(self) -> u32 { + self.0 + } +} + +impl From for NodeId { + #[inline] + fn from(val: u32) -> Self { + Self(val) + } +} + +impl From for NodeId { + #[inline] + #[allow(clippy::cast_possible_truncation)] + fn from(val: usize) -> Self { + Self(val as u32) + } +} + +/// Owner of an expression DAG: nodes in allocation order, addressed by +/// [`NodeId`]. +/// +/// An arena only grows, so ids stay valid while it is alive. Because a node names +/// its children by id, a DAG is acyclic by construction as long as callers keep +/// allocating children before parents, which every builder here does. +#[derive(Clone, Debug, Default)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct Arena { + nodes: Vec, +} + +impl Arena { + /// An empty arena. + pub const fn new() -> Self { + Self { nodes: Vec::new() } + } + + /// Append `node` and return its id. + #[must_use] + #[allow(clippy::cast_possible_truncation)] + pub fn alloc(&mut self, node: Node) -> NodeId { + let id = NodeId(self.nodes.len() as u32); + self.nodes.push(node); + id + } + + /// Borrow a node. + /// + /// # Panics + /// + /// Panics if `id` came from a different arena and is out of range here. + #[inline] + pub fn get(&self, id: NodeId) -> &Node { + &self.nodes[id.index()] + } + + /// Nodes allocated so far, reachable or not. + pub const fn len(&self) -> usize { + self.nodes.len() + } + + /// Every allocated node, in allocation order, including nodes unreachable + /// from any particular root. For cheap whole-arena scans. + #[inline] + pub fn nodes(&self) -> &[Node] { + &self.nodes + } + + /// Whether nothing has been allocated yet. + pub const fn is_empty(&self) -> bool { + self.nodes.is_empty() + } + + /// Nodes reachable from `root`, children before parents and each visited once. + /// + /// This is the evaluation order every pass works in, and the reason a shared + /// subexpression is computed once rather than per parent. + pub fn topological_order(&self, root: NodeId) -> Vec { + // Iterative post-order DFS. Two frame variants avoid the per-node + // children buffer an explicit (id, flag) stack would need. + let n = self.len(); + let mut visited = vec![false; n]; + let mut order = Vec::with_capacity(n); + let mut stack: Vec = Vec::with_capacity(n); + stack.push(TopoFrame::Pending(root)); + while let Some(frame) = stack.pop() { + match frame { + TopoFrame::Visit(id) => { + if visited[id.index()] { + continue; + } + visited[id.index()] = true; + order.push(id); + }, + TopoFrame::Pending(id) => { + if visited[id.index()] { + continue; + } + stack.push(TopoFrame::Visit(id)); + self.get(id).for_each_child(|c| { + if !visited[c.index()] { + stack.push(TopoFrame::Pending(c)); + } + }); + }, + } + } + order + } +} + +/// State-input usage of the sub-DAG rooted at `root`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StateUsage { + /// 1 + max state index referenced across both `y` and `y_dot` + /// (0 if no state nodes are reachable). + pub n_states: usize, + /// Whether any `StateVectorDot` node is reachable. + pub uses_y_dot: bool, +} + +/// Scan the reachable sub-DAG for `StateVector`/`StateVectorDot` extents. +/// +/// Primal DAGs only: tangent nodes (`TangentStateVector` etc.) are +/// intentionally not counted. +pub fn scan_state_usage(arena: &Arena, root: NodeId) -> StateUsage { + let mut visited = vec![false; arena.len()]; + let mut stack = vec![root]; + let mut usage = StateUsage { + n_states: 0, + uses_y_dot: false, + }; + while let Some(id) = stack.pop() { + if std::mem::replace(&mut visited[id.index()], true) { + continue; + } + let node = arena.get(id); + match node { + Node::StateVector { end, .. } => usage.n_states = usage.n_states.max(*end), + Node::StateVectorDot { end, .. } => { + usage.uses_y_dot = true; + usage.n_states = usage.n_states.max(*end); + }, + _ => {}, + } + node.for_each_child(|c| stack.push(c)); + } + usage +} + +enum TopoFrame { + /// Expands into this node's `Visit` frame with its children stacked above. + Pending(NodeId), + /// Children are all emitted; emit this node now. + Visit(NodeId), +} + +impl Index for Arena { + type Output = Node; + + #[inline] + fn index(&self, id: NodeId) -> &Node { + &self.nodes[id.index()] + } +} + +/// Per-node side table, keyed by [`NodeId`]. +/// +/// A dense `Vec>` rather than a hash map: keys are arena positions, so +/// lookup is an index and a pass over every node touches memory in order. Ids are +/// dense and small, which is what makes the wasted `None` slots cheaper than +/// hashing. +#[derive(Clone, Debug)] +pub struct NodeMap { + slots: Vec>, +} + +impl NodeMap { + /// A table sized for an arena of `arena_len` nodes, all entries empty. + #[inline] + pub fn new(arena_len: usize) -> Self { + let mut slots = Vec::with_capacity(arena_len); + slots.resize_with(arena_len, || None); + Self { slots } + } + + /// The value for `id`, or `None` if unset or out of range. + #[inline] + pub fn get(&self, id: NodeId) -> Option<&T> { + self.slots.get(id.index()).and_then(|opt| opt.as_ref()) + } + + /// Set the value for `id`, returning what it replaced. Grows the table when + /// the arena has outgrown the size passed to [`new`](Self::new). + #[inline] + pub fn insert(&mut self, id: NodeId, value: T) -> Option { + if id.index() >= self.slots.len() { + self.slots.resize_with(id.index() + 1, || None); + } + self.slots[id.index()].replace(value) + } + + /// Whether `id` has a value set. + #[inline] + pub fn contains_key(&self, id: NodeId) -> bool { + self.slots.get(id.index()).is_some_and(Option::is_some) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + #[test] + fn test_alloc_and_get() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Scalar(42.5)); + match arena.get(id) { + Node::Scalar(v) => assert!((*v - 42.5).abs() < f64::EPSILON), + _ => panic!("expected Scalar"), + } + } + + #[test] + fn test_sequential_ids() { + let mut arena = Arena::new(); + let id0 = arena.alloc(Node::Scalar(1.0)); + let id1 = arena.alloc(Node::Scalar(2.0)); + assert_eq!(id0.index(), 0); + assert_eq!(id1.index(), 1); + } + + #[test] + fn test_node_id_copy() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Scalar(1.0)); + let id_copy = id; + assert_eq!(id, id_copy); + assert_eq!(arena.get(id), arena.get(id_copy)); + } + + #[test] + fn test_topological_order_visits_children_first() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Scalar(2.0)); + let add = arena.alloc(Node::Add(a, b)); + + let order = arena.topological_order(add); + + let pos_a = order.iter().position(|&n| n == a).unwrap(); + let pos_b = order.iter().position(|&n| n == b).unwrap(); + let pos_add = order.iter().position(|&n| n == add).unwrap(); + assert!(pos_a < pos_add); + assert!(pos_b < pos_add); + assert_eq!(order.len(), 3); + } + + #[test] + fn test_topological_order_deduplicates_shared_subexpression() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let neg_a = arena.alloc(Node::Neg(a)); + let combined = arena.alloc(Node::Add(neg_a, neg_a)); + + let order = arena.topological_order(combined); + assert_eq!(order.len(), 3); + assert_eq!(order.iter().filter(|&&n| n == a).count(), 1); + assert_eq!(order.iter().filter(|&&n| n == neg_a).count(), 1); + } + + #[test] + fn test_topological_order_skips_unreachable_nodes() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let _unreachable = arena.alloc(Node::Scalar(99.0)); + let neg_a = arena.alloc(Node::Neg(a)); + + let order = arena.topological_order(neg_a); + assert_eq!(order.len(), 2); + assert!(order.contains(&a)); + assert!(order.contains(&neg_a)); + } + + #[test] + fn scan_state_usage_finds_extents_and_ydot() { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 2, end: 5 }); + let svd = arena.alloc(Node::StateVectorDot { start: 0, end: 1 }); + let sum = arena.alloc(Node::Add(sv, svd)); + let usage = scan_state_usage(&arena, sum); + assert_eq!(usage.n_states, 5); + assert!(usage.uses_y_dot); + + let lone = arena.alloc(Node::Scalar(1.0)); + let usage = scan_state_usage(&arena, lone); + assert_eq!(usage.n_states, 0); + assert!(!usage.uses_y_dot); + } + + #[test] + fn test_node_map_basic_get_insert() { + let mut map: NodeMap = NodeMap::new(4); + let id0: NodeId = 0u32.into(); + let id3: NodeId = 3u32.into(); + + assert!(!map.contains_key(id0)); + assert_eq!(map.get(id0), None); + + assert_eq!(map.insert(id0, 42), None); + assert_eq!(map.get(id0), Some(&42)); + assert!(map.contains_key(id0)); + + assert_eq!(map.insert(id0, 99), Some(42)); + assert_eq!(map.get(id0), Some(&99)); + + assert!(!map.contains_key(id3)); + map.insert(id3, 7); + assert_eq!(map.get(id3), Some(&7)); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/branch_regions.rs b/packages/pybamm-rust/pybamm-core/src/branch_regions.rs new file mode 100644 index 0000000000..56f1809b62 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/branch_regions.rs @@ -0,0 +1,1142 @@ +//! Conditional branch region analysis. +//! +//! A `Node::Conditional` lowers every branch into one flat instruction tape. To +//! execute only the active branch, the compiler must first know which nodes +//! belong exclusively to which branch. This module answers that with a +//! consumer-propagation pass over the DAG, and uses the answer twice: to +//! privatise subgraphs `cse` made shared between a strict subset of branches +//! ([`privatise_conditionals`]), and to schedule each branch's nodes into one +//! contiguous instruction block. +//! +//! Every ambiguous case degrades to [`Ownership::Common`], "compute it +//! always", which costs performance but never correctness. + +use crate::arena::{Arena, NodeId}; +use crate::ir::ConstPool; +use crate::node::Node; + +/// One branch of one conditional: `branches[index]` of `cond`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BranchLabel { + /// The `Node::Conditional` this branch belongs to. + pub cond: NodeId, + /// Position of the branch within that conditional, 0-based. + pub index: u32, +} + +/// Which conditional branches exclusively need a node's value. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Ownership { + /// No consumer seen (unreachable from the root). + Unreached, + /// Needed only by the listed branches of one conditional. `indices` is + /// sorted and deduplicated. + Branches { cond: NodeId, indices: Vec }, + /// Needed regardless of which branch runs: reachable from outside the + /// conditional, needed by branches of more than one conditional, or the + /// selector's own cone. + Common, +} + +impl Ownership { + /// Join two labels in the lattice `Unreached < Branches < Common`. Two + /// `Branches` over the same conditional merge their index sets; over + /// different conditionals they degrade to `Common`. + fn join(&mut self, other: &Self) { + match (&mut *self, other) { + (_, Self::Unreached) | (Self::Common, _) => {}, + (Self::Unreached, _) => *self = other.clone(), + (_, Self::Common) => *self = Self::Common, + ( + Self::Branches { cond, indices }, + Self::Branches { + cond: other_cond, + indices: other_indices, + }, + ) => { + if cond == other_cond { + for &i in other_indices { + if let Err(pos) = indices.binary_search(&i) { + indices.insert(pos, i); + } + } + } else { + *self = Self::Common; + } + }, + } + } +} + +/// The single branch that exclusively owns this node, if exactly one does. +#[must_use] +pub fn sole_owner(ownership: &Ownership) -> Option { + match ownership { + Ownership::Branches { cond, indices } if indices.len() == 1 => Some(BranchLabel { + cond: *cond, + index: indices[0], + }), + _ => None, + } +} + +/// Whether `arena` holds any conditional that could become blocks at all. +/// +/// A linear, allocation-free scan that skips [`owner_sets`] on the overwhelming +/// majority of compiles, which carry no conditional. Answers over the whole arena +/// rather than one root's cone, so it can only over-approximate. +#[must_use] +pub fn has_multi_branch_conditional(arena: &Arena) -> bool { + arena + .nodes() + .iter() + .any(|node| matches!(node, Node::Conditional { branches, .. } if branches.len() >= 2)) +} + +/// Whether `cond`'s branches may become instruction blocks. +/// +/// True for a top-level conditional: one whose own value is needed regardless of +/// which branch of anything else runs. It is also true for a conditional nested +/// three or more deep, whose own label the nested-conditional pass has degraded to +/// [`Ownership::Common`]. That costs nothing and cannot nest blocks: the same pass +/// degraded every node its branches could own, so its branch runs come out empty +/// and scheduling drops the group. +#[must_use] +pub fn is_conditional_blockable(arena: &Arena, owners: &[Ownership], cond: NodeId) -> bool { + matches!(arena.get(cond), Node::Conditional { branches, .. } if branches.len() >= 2) + && owners[cond.index()] == Ownership::Common +} + +/// Label every node in `eval_order` with the branches that exclusively need it. +/// +/// Propagates labels from consumers to producers over the reverse of the +/// topological order. The root is `Common`; a `Conditional`'s branch edge +/// carries that branch's label, its selector edge carries `Common`, and every +/// other edge carries the consumer's own label. +/// +/// Returns a vector indexed by `NodeId::index()`, length `arena.len()`; nodes +/// outside `eval_order` stay [`Ownership::Unreached`]. +#[must_use] +// A branch index never approaches u32::MAX; `Vec::len()` is `usize` only by convention. +#[allow(clippy::cast_possible_truncation)] +pub fn owner_sets(arena: &Arena, eval_order: &[NodeId]) -> Vec { + let mut owners = vec![Ownership::Unreached; arena.len()]; + if let Some(&root) = eval_order.last() { + owners[root.index()] = Ownership::Common; + } + + for &id in eval_order.iter().rev() { + let label = owners[id.index()].clone(); + match arena.get(id) { + Node::Conditional { selector, branches } if branches.len() >= 2 => { + owners[selector.index()].join(&Ownership::Common); + for (i, &b) in branches.iter().enumerate() { + let edge = Ownership::Branches { + cond: id, + indices: vec![i as u32], + }; + owners[b.index()].join(&edge); + } + }, + node => node.for_each_child(|c| owners[c.index()].join(&label)), + } + } + + close_full_branch_coverage(arena, &mut owners); + force_nested_conditionals_common(arena, eval_order, &mut owners); + owners +} + +/// Degrade a node to `Common` once its `Branches` label already spans every +/// branch of its conditional. +/// +/// Consumer-propagation merges branch labels one edge at a time, so a node fed +/// by all of a conditional's branches only reveals that once every branch has +/// contributed its edge. Full coverage means "needed no matter which branch +/// runs", i.e. `Common`, not a `Branches` set naming all of them. +fn close_full_branch_coverage(arena: &Arena, owners: &mut [Ownership]) { + for owner in owners.iter_mut() { + if let Ownership::Branches { cond, indices } = owner + && let Node::Conditional { branches, .. } = arena.get(*cond) + && indices.len() >= branches.len() + { + *owner = Ownership::Common; + } + } +} + +/// Degrade the cone of every non-top-level conditional to `Common`. +/// +/// Keeps blocks flat and pairwise disjoint, which is what makes the evaluator's +/// jump-over-inactive-block correct without a nesting-aware skip table. The +/// rewrite is closed: a node labelled `Branches { cond: inner, .. }` can only be +/// consumed by nodes with that same label or `Common`, so no further propagation +/// is needed. +fn force_nested_conditionals_common( + arena: &Arena, + eval_order: &[NodeId], + owners: &mut [Ownership], +) { + let nested: Vec = eval_order + .iter() + .copied() + .filter(|&id| { + matches!(arena.get(id), Node::Conditional { branches, .. } if branches.len() >= 2) + && owners[id.index()] != Ownership::Common + }) + .collect(); + if nested.is_empty() { + return; + } + for &id in eval_order { + if let Ownership::Branches { cond, .. } = &owners[id.index()] + && nested.contains(cond) + { + owners[id.index()] = Ownership::Common; + } + } +} + +/// Ceiling on the privatised arena's node count, as a multiple of the input's. +/// +/// Cloning is bounded only when sharing is all-or-nothing. In chain-nested cones, +/// where branch *i*'s cone strictly contains branch *i-1*'s, every interior node is +/// shared with a strict *subset*, so the arena grows quadratically in branch count: +/// a `tanh` chain measures 1.4x at 4 branches, 2.9x at 8 and 14.3x at 32. +/// +/// 4x admits every shape `PyBaMM` builds today (widest measured 1.97x) while +/// still capping the chain shape from 16 branches up. Over budget costs only the +/// short-circuit, never correctness. +const CLONE_BUDGET_MULTIPLE: usize = 4; + +/// Clone subgraphs shared between a strict subset of one conditional's branches +/// so each sharing branch owns its own copy. +/// +/// Returns `None` when nothing needs cloning, meaning there is no conditional or +/// every shared cone node is `Common` or shared by *all* branches. Also returns +/// `None` when the projected clone count would exceed the clone budget. `None` is +/// only slower, never wrong: the caller lowers the unprivatised arena. +/// +/// Must run **after** `cse`, which created the sharing, and immediately before +/// lowering, which is why `IRBuilder` calls it. +#[must_use] +pub fn privatise_conditionals(arena: &Arena, root: NodeId) -> Option<(Arena, NodeId)> { + if !has_multi_branch_conditional(arena) { + return None; + } + let eval_order = arena.topological_order(root); + let owners = owner_sets(arena, &eval_order); + + // A node is privatisable iff it is owned by 2..n of one *blockable* + // conditional's branches. + let privatisable = |id: NodeId| -> Option<(NodeId, &[u32])> { + let Ownership::Branches { cond, indices } = &owners[id.index()] else { + return None; + }; + if indices.len() < 2 || !is_conditional_blockable(arena, &owners, *cond) { + return None; + } + let Node::Conditional { branches, .. } = arena.get(*cond) else { + return None; + }; + (indices.len() < branches.len()).then_some((*cond, indices.as_slice())) + }; + + // Projected new arena size, in the pass that already decides whether anything + // is privatisable: one copy per kept node, one per sharing branch otherwise. + let mut projected_nodes = 0_usize; + let mut any_privatisable = false; + for &id in &eval_order { + match privatisable(id) { + Some((_, indices)) => { + any_privatisable = true; + projected_nodes += indices.len(); + }, + None => projected_nodes += 1, + } + } + if !any_privatisable || projected_nodes > CLONE_BUDGET_MULTIPLE * arena.len() { + return None; + } + + // Rebuild in topological order: a common node is copied once, a privatised + // node once per sharing branch with children resolved to that branch's copy. + let mut new_arena = Arena::new(); + let mut common: Vec> = vec![None; arena.len()]; + let mut per_branch: Vec> = vec![Vec::new(); arena.len()]; + + for &id in &eval_order { + let node = arena.get(id); + match privatisable(id) { + None => { + let remapped = + remap_children(node, owner_branch_of(&owners, id), &common, &per_branch); + common[id.index()] = Some(new_arena.alloc(remapped)); + }, + Some((_, indices)) => { + for &i in indices { + let remapped = remap_children(node, Some(i), &common, &per_branch); + let clone = new_arena.alloc(remapped); + per_branch[id.index()].push((i, clone)); + } + }, + } + } + + let new_root = common[root.index()].expect("the root is never privatised"); + Some((new_arena, new_root)) +} + +/// Rebuild `node` against the privatised copies, reading each child from +/// `branch`'s copy where one exists. +/// +/// `Conditional` is the one node whose child *positions* carry per-branch +/// meaning: `cse` can alias two branch slots onto the same node, and each slot +/// must still resolve to its own branch's clone. Every other node's children +/// inherit the consumer's own branch, so one uniform closure suffices. +/// +/// A slot index is unambiguous even though [`resolve_child`] keys clones by index +/// alone: a node needed by two conditionals is `Ownership::Common`, which is not +/// privatisable. +// A branch index never approaches u32::MAX; `Vec::len()` is `usize` only by convention. +#[allow(clippy::cast_possible_truncation)] +fn remap_children( + node: &Node, + branch: Option, + common: &[Option], + per_branch: &[Vec<(u32, NodeId)>], +) -> Node { + match node { + Node::Conditional { selector, branches } => Node::Conditional { + selector: resolve_child(common, per_branch, *selector, branch), + branches: branches + .iter() + .enumerate() + .map(|(i, &b)| resolve_child(common, per_branch, b, Some(i as u32))) + .collect(), + }, + other => other.map_children(|c| resolve_child(common, per_branch, c, branch)), + } +} + +/// The branch index a node was assigned to, if exactly one owns it. Picks which +/// copy of a privatised child a consumer should read. +fn owner_branch_of(owners: &[Ownership], id: NodeId) -> Option { + sole_owner(&owners[id.index()]).map(|label| label.index) +} + +/// One conditional's branch nodes, as a contiguous range of the emission order. +#[derive(Clone, Debug)] +pub struct RegionGroup { + /// The conditional whose branches this group covers. + pub cond: NodeId, + /// Index in [`RegionSchedule::order`] where this group's block nodes start. + pub anchor: usize, + /// Node count per branch, in branch order. The branches occupy + /// `order[anchor .. anchor + branch_lens.iter().sum()]` back to back. + pub branch_lens: Vec, +} + +/// A block-contiguous emission order plus the groups to wrap in a `Dispatch`. +/// +/// `order` holds **every** node, group nodes sit at their block position, so +/// slot allocation can be driven by it with no allocator change. +#[derive(Clone, Debug)] +pub struct RegionSchedule { + /// Emission order for every reachable node. + pub order: Vec, + /// The conditionals whose branch runs came out contiguous, and so can be + /// guarded by a `Dispatch`. + pub groups: Vec, +} + +/// Reorder `base_order` so each blockable conditional's branch-owned nodes are +/// contiguous, grouped by branch, and immediately precede the `Conditional`. +/// +/// Valid because a branch-`i`-owned node can only depend on `Common` nodes +/// (all emitted before the `Conditional`, hence before the group) and on other +/// branch-`i`-owned nodes (kept in `base_order` order within the run). A +/// cross-branch dependency is impossible: if a branch-`i` node consumed a node, +/// that node's owner set would include `i`. +/// +/// `base_order` must be a topological order (children first). +#[must_use] +pub fn schedule_regions(arena: &Arena, base_order: &[NodeId]) -> RegionSchedule { + let unscheduled = || RegionSchedule { + order: base_order.to_vec(), + groups: Vec::new(), + }; + if !has_multi_branch_conditional(arena) { + return unscheduled(); + } + let owners = owner_sets(arena, base_order); + let Some(mut runs) = collect_branch_runs(arena, &owners, base_order, base_order) else { + return unscheduled(); + }; + + let mut order = Vec::with_capacity(base_order.len()); + let mut groups = Vec::with_capacity(runs.len()); + for &id in base_order { + if owned_by_a_run(&owners, &runs, id) { + continue; // deferred into its group + } + if let Some(pos) = runs.iter().position(|(c, _)| *c == id) { + let branch_runs = std::mem::take(&mut runs[pos].1); + flush_group(&mut order, &mut groups, id, &branch_runs); + } + order.push(id); + } + + // A dropped node leaves its slot at `(0, 0)`, which lowers to silently wrong + // values rather than a crash, so this ships rather than being debug-only. + assert_eq!( + order.len(), + base_order.len(), + "region scheduling dropped or duplicated a node" + ); + RegionSchedule { order, groups } +} + +/// [`schedule_regions`] for one half of a primal/tangent partition. +/// +/// `partition` is the subset of `base_order`, in order, that this half emits. A +/// group whose `Conditional` lands in the *other* half anchors at the end of this +/// partition: nothing here consumes it and its `Common` dependencies are all +/// earlier, which lets one branch yield two blocks with a valid split point. +/// +/// Such a group forms only when this partition holds the selector, since a +/// `Dispatch` reads the selector's slot. Otherwise it is dropped and its nodes go +/// inline, costing the short-circuit but never correctness. +#[must_use] +pub fn schedule_regions_partitioned( + arena: &Arena, + partition: &[NodeId], + base_order: &[NodeId], +) -> RegionSchedule { + let unscheduled = || RegionSchedule { + order: partition.to_vec(), + groups: Vec::new(), + }; + if !has_multi_branch_conditional(arena) { + return unscheduled(); + } + let owners = owner_sets(arena, base_order); + // Runs are collected from `partition` only, so a branch's primal nodes and + // its tangent nodes end up in different halves' groups. + let Some(mut runs) = collect_branch_runs(arena, &owners, base_order, partition) else { + return unscheduled(); + }; + + let mut in_partition = vec![false; arena.len()]; + for &id in partition { + in_partition[id.index()] = true; + } + runs.retain(|(cond, _)| { + in_partition[cond.index()] + || matches!(arena.get(*cond), Node::Conditional { selector, .. } + if in_partition[selector.index()]) + }); + + let mut order = Vec::with_capacity(partition.len()); + let mut groups = Vec::with_capacity(runs.len()); + let mut pending = vec![true; runs.len()]; + + for &id in partition { + if owned_by_a_run(&owners, &runs, id) { + continue; // deferred into its group + } + if let Some(pos) = runs.iter().position(|(c, _)| *c == id) { + let branch_runs = std::mem::take(&mut runs[pos].1); + flush_group(&mut order, &mut groups, id, &branch_runs); + pending[pos] = false; + } + order.push(id); + } + // Groups whose `Conditional` lives in the other half: nothing here consumes + // them, so they go last. + for (pos, run) in runs.iter_mut().enumerate() { + if pending[pos] { + let branch_runs = std::mem::take(&mut run.1); + flush_group(&mut order, &mut groups, run.0, &branch_runs); + } + } + + // Ships for the same reason as in `schedule_regions`: a dropped node leaves + // its slot at `(0, 0)`, which lowers to silently wrong values, not a crash. + assert_eq!( + order.len(), + partition.len(), + "partitioned region scheduling dropped or duplicated a node" + ); + RegionSchedule { order, groups } +} + +/// Per-branch runs for every blockable conditional, in `source` order. +/// `None` when there is no blockable conditional at all. +fn collect_branch_runs( + arena: &Arena, + owners: &[Ownership], + base_order: &[NodeId], + source: &[NodeId], +) -> Option>)>> { + let mut runs: Vec<(NodeId, Vec>)> = Vec::new(); + for &id in base_order { + if is_conditional_blockable(arena, owners, id) { + let Node::Conditional { branches, .. } = arena.get(id) else { + unreachable!("is_conditional_blockable checked the node kind") + }; + runs.push((id, vec![Vec::new(); branches.len()])); + } + } + if runs.is_empty() { + return None; + } + for &id in source { + if let Some(label) = sole_owner(&owners[id.index()]) + && let Some((_, branch_runs)) = runs.iter_mut().find(|(c, _)| *c == label.cond) + { + branch_runs[label.index as usize].push(id); + } + } + Some(runs) +} + +/// Whether `id` was deferred into one of `runs`' groups. +fn owned_by_a_run(owners: &[Ownership], runs: &[(NodeId, Vec>)], id: NodeId) -> bool { + sole_owner(&owners[id.index()]).is_some_and(|label| runs.iter().any(|(c, _)| *c == label.cond)) +} + +/// Append one group's branch runs to `order` and record the annotated range. +/// A group whose branches own nothing is dropped: a `Dispatch` guarding no +/// instructions is pure overhead, and dropping it keeps anchors distinct. +fn flush_group( + order: &mut Vec, + groups: &mut Vec, + cond: NodeId, + branch_runs: &[Vec], +) { + if branch_runs.iter().all(Vec::is_empty) { + return; + } + let anchor = order.len(); + let branch_lens = branch_runs.iter().map(Vec::len).collect(); + for run in branch_runs { + order.extend_from_slice(run); + } + groups.push(RegionGroup { + cond, + anchor, + branch_lens, + }); +} + +/// The active branch for `selector`, or `None` when nothing matches. +/// +/// 1-based round-to-nearest windows, first match wins: branch `i` is active iff +/// `selector > (i+1) - 0.5 && selector < (i+1) + 0.5`. The single home for the +/// window, shared by the scalar, batch and reverse evaluators so their branch +/// choices cannot drift apart. +#[inline] +#[must_use] +pub fn active_branch(selector: f64, n_branches: usize) -> Option { + for i in 0..n_branches { + let idx = (i + 1) as f64; + if selector > idx - 0.5 && selector < idx + 0.5 { + return Some(i); + } + } + None +} + +/// One past the last instruction guarded by the `Dispatch` at `pc`. +/// +/// The companion of [`active_branch`], and the single home for the *other* half +/// of the `Dispatch` contract: forward evaluation resumes here, and the backward +/// walk skips everything in `(pc, end)` bar the active block. Two copies of this +/// reduction could desynchronise the two directions silently. +/// +/// `blocks_len` of 0 falls back to `pc + 1`, so a span always covers at least +/// the `Dispatch` itself. +#[inline] +#[must_use] +pub(crate) fn dispatch_span_end( + consts: &ConstPool, + pc: usize, + blocks_idx: u32, + blocks_len: u32, +) -> usize { + (0..blocks_len as usize) + .map(|b| { + let (rel, len) = consts.branch_blocks[blocks_idx as usize + b]; + pc + rel as usize + len as usize + }) + .max() + .unwrap_or(pc + 1) +} + +/// Resolve `child` to the copy the consumer should read: its `branch` copy when +/// one was made, otherwise the single common copy. +fn resolve_child( + common: &[Option], + per_branch: &[Vec<(u32, NodeId)>], + child: NodeId, + branch: Option, +) -> NodeId { + if let Some(b) = branch + && let Some(&(_, id)) = per_branch[child.index()].iter().find(|&&(i, _)| i == b) + { + return id; + } + common[child.index()].unwrap_or_else(|| { + panic!("child {child:?} has no copy reachable from its consumer's branch") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + /// `cond(sel, [sin(y0), cos(y0)])` with `y0` shared: the two unary nodes are + /// branch-exclusive, `y0` and `sel` are common. + #[test] + fn exclusive_branch_cones_are_labelled_per_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!( + sole_owner(&owners[b1.index()]), + Some(BranchLabel { cond, index: 0 }) + ); + assert_eq!( + sole_owner(&owners[b2.index()]), + Some(BranchLabel { cond, index: 1 }) + ); + assert_eq!(owners[y.index()], Ownership::Common); + assert_eq!(owners[sel.index()], Ownership::Common); + assert_eq!(owners[cond.index()], Ownership::Common); + } + + /// A node consumed by two of three branches is `Branches` with both indices, + /// not `Common`, which is exactly the shape privatisation targets. + #[test] + fn strict_subset_sharing_records_every_sharing_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!( + owners[shared.index()], + Ownership::Branches { + cond, + indices: vec![0, 1] + } + ); + assert_eq!(sole_owner(&owners[shared.index()]), None); + } + + /// Reachable from outside the conditional => never owned by a branch. + #[test] + fn nodes_used_outside_the_conditional_stay_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + // `shared` also feeds the root outside the conditional. + let root = arena.alloc(Node::Add(cond, shared)); + + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[shared.index()], Ownership::Common); + assert_eq!( + sole_owner(&owners[b1.index()]), + Some(BranchLabel { cond, index: 0 }) + ); + } + + /// The selector edge carries `Common`: the dispatch reads it before the blocks. + #[test] + fn selector_cone_is_common_even_when_otherwise_exclusive() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Floor(y)); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[sel.index()], Ownership::Common); + } + + /// An inner conditional inside an outer branch cone is not blockable, and its + /// own cone is forced `Common`. + #[test] + fn nested_conditional_cone_is_forced_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let outer_sel = arena.alloc(Node::Scalar(1.0)); + let inner_sel = arena.alloc(Node::Scalar(2.0)); + let inner_b1 = arena.alloc(Node::Sin(y)); + let inner_b2 = arena.alloc(Node::Cos(y)); + let inner = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![inner_b1, inner_b2], + }); + let outer_b2 = arena.alloc(Node::Neg(y)); + let outer = arena.alloc(Node::Conditional { + selector: outer_sel, + branches: vec![inner, outer_b2], + }); + + let order = arena.topological_order(outer); + let owners = owner_sets(&arena, &order); + + assert!(is_conditional_blockable(&arena, &owners, outer)); + assert!(!is_conditional_blockable(&arena, &owners, inner)); + // The inner cone degraded to Common; the inner Conditional itself is + // still owned by the outer branch it sits in. + assert_eq!(owners[inner_b1.index()], Ownership::Common); + assert_eq!(owners[inner_b2.index()], Ownership::Common); + assert_eq!( + sole_owner(&owners[inner.index()]), + Some(BranchLabel { + cond: outer, + index: 0 + }) + ); + } + + /// A node needed by branches of two different conditionals degrades. + #[test] + fn sharing_across_two_conditionals_degrades_to_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let a1 = arena.alloc(Node::Sin(shared)); + let a2 = arena.alloc(Node::Neg(y)); + let c1 = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![a1, a2], + }); + let b1 = arena.alloc(Node::Cos(shared)); + let b2 = arena.alloc(Node::Abs(y)); + let c2 = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let root = arena.alloc(Node::Add(c1, c2)); + + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[shared.index()], Ownership::Common); + } + + /// No conditional at all: everything is `Common`, nothing panics. + #[test] + fn plain_expression_is_all_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let root = arena.alloc(Node::Sin(y)); + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + assert!(order.iter().all(|n| owners[n.index()] == Ownership::Common)); + } + + /// Nothing to privatise: no conditional, or no strict-subset sharing. + #[test] + fn privatise_is_a_no_op_without_strict_subset_sharing() { + let mut plain = Arena::new(); + let y = plain.alloc(Node::StateVector { start: 0, end: 1 }); + let root = plain.alloc(Node::Sin(y)); + assert!(privatise_conditionals(&plain, root).is_none()); + + // Shared by ALL branches: needed whichever branch runs, so cloning + // would be pure waste. + let mut all = Arena::new(); + let y = all.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = all.alloc(Node::Scalar(1.0)); + let shared = all.alloc(Node::Exp(y)); + let b1 = all.alloc(Node::Sin(shared)); + let b2 = all.alloc(Node::Cos(shared)); + let cond = all.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + assert!(privatise_conditionals(&all, cond).is_none()); + } + + /// Shared by 2 of 3 branches: cloned once per sharing branch, so each + /// branch's cone becomes exclusive. + #[test] + fn privatise_clones_a_strict_subset_cone_per_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let inner = arena.alloc(Node::Exp(y)); + let shared = arena.alloc(Node::Sqrt(inner)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let (new_arena, new_root) = + privatise_conditionals(&arena, cond).expect("strict subset must privatise"); + // The clone budget is a cap on pathological shapes, not on this one. + assert!(new_arena.len() <= CLONE_BUDGET_MULTIPLE * arena.len()); + + let order = new_arena.topological_order(new_root); + // Both nodes of the shared cone now exist twice. + let n_sqrt = order + .iter() + .filter(|&&id| matches!(new_arena.get(id), Node::Sqrt(_))) + .count(); + let n_exp = order + .iter() + .filter(|&&id| matches!(new_arena.get(id), Node::Exp(_))) + .count(); + assert_eq!((n_sqrt, n_exp), (2, 2)); + + // And every cone node is now exclusively owned by one branch. + let owners = owner_sets(&new_arena, &order); + for &id in &order { + if matches!(new_arena.get(id), Node::Sqrt(_) | Node::Exp(_)) { + assert!( + sole_owner(&owners[id.index()]).is_some(), + "cone node {id:?} is still shared after privatisation" + ); + } + } + } + + /// Privatisation is value-preserving: the cloned graph evaluates identically + /// for every selector, including no-match. + #[test] + fn privatise_preserves_values_for_every_selector() { + use crate::eval::CompiledExpr; + use crate::ir::TypedIr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let (new_arena, new_root) = privatise_conditionals(&arena, cond).expect("privatises"); + let before = CompiledExpr::from_ir(TypedIr::from_arena_raw(&arena, cond)); + let after = CompiledExpr::from_ir(TypedIr::from_arena_raw(&new_arena, new_root)); + + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 0.5, 1.5, f64::NAN] { + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.3], &[], &[sel_val]); + let b = after.eval(&mut s2, 0.0, &[0.3], &[], &[sel_val]); + assert_eq!(a[0].to_bits(), b[0].to_bits(), "selector {sel_val}"); + } + } + + /// Regression for a real unified-experiment-model panic: two branches of a + /// 3-way `Conditional` (a current-control step and a power-control step, + /// both with no extra termination condition) canonicalise to the literal + /// same `Scalar(1.0)` node under `cse`, aliasing `branches[0]` and + /// `branches[2]`. The slot stays privatisable, `remap_children` resolves + /// each branch position to its own clone, so the aliased node is cloned + /// once per sharing branch and the values are unchanged. + #[test] + fn aliased_branch_slot_privatises_per_position() { + use crate::eval::CompiledExpr; + use crate::ir::TypedIr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let shared = arena.alloc(Node::Scalar(1.0)); + let other = arena.alloc(Node::Sin(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + // branches[0] and branches[2] are the literal same NodeId. + branches: vec![shared, other, shared], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + assert_eq!( + owners[shared.index()], + Ownership::Branches { + cond, + indices: vec![0, 2] + }, + "an aliased branches slot is owned by exactly the positions that name it" + ); + + let before = CompiledExpr::from_ir(TypedIr::from_arena_raw(&arena, cond)); + // Must not panic: this call is the regression this test guards. + let (new_arena, new_root) = + privatise_conditionals(&arena, cond).expect("the aliased slot privatises"); + let Node::Conditional { branches, .. } = new_arena.get(new_root) else { + panic!("the root is still the conditional") + }; + assert_ne!( + branches[0], branches[2], + "each aliased position must resolve to its own clone" + ); + let after = CompiledExpr::from_ir(TypedIr::from_arena_raw(&new_arena, new_root)); + + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 0.5, 1.5, f64::NAN] { + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.3], &[], &[sel_val]); + let b = after.eval(&mut s2, 0.0, &[0.3], &[], &[sel_val]); + assert_eq!(a[0].to_bits(), b[0].to_bits(), "selector {sel_val}"); + } + } + + /// `cond(sel, [tanh(y), tanh²(y), ..., tanhⁿ(y)])`: branch `i`'s cone strictly + /// contains branch `i-1`'s, the shape `cse` makes of a progressive expression. + fn chain_nested_conditional(n: usize) -> (Arena, NodeId) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let mut node = y; + let mut branches = Vec::with_capacity(n); + for _ in 0..n { + node = arena.alloc(Node::Tanh(node)); + branches.push(node); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + (arena, cond) + } + + /// The semantics contract for [`chain_nested_conditional`], written out + /// independently of the evaluator. + fn expected_chain(y: f64, n: usize, selector: f64) -> f64 { + active_branch(selector, n).map_or(0.0, |i| { + let mut v = y; + for _ in 0..=i { + v = v.tanh(); + } + v + }) + } + + /// Chain-nested cones share every interior node with a strict *subset* of + /// branches, so privatising clones each once per sharing branch: quadratic in + /// the branch count (~14x the arena at 32 branches). The budget must bail out, + /// leaving the tape the size of the arena and the values untouched. + #[test] + fn a_chain_nested_conditional_bails_out_of_the_clone_budget() { + use crate::eval::CompiledExpr; + use crate::ir::TypedIr; + + const N: usize = 32; + let (arena, cond) = chain_nested_conditional(N); + assert!( + privatise_conditionals(&arena, cond).is_none(), + "the projected clone count must exceed the budget" + ); + + let ir = TypedIr::from_arena(&arena, cond); + let budget = CLONE_BUDGET_MULTIPLE * arena.len(); + assert!( + ir.instructions().len() <= budget, + "raw tape of {} instructions exceeds the {budget}-instruction budget", + ir.instructions().len() + ); + + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.4_f64; + for sel_val in [ + 0.0_f64, + 1.0, + 2.0, + 17.0, + 32.0, + 33.0, + 0.5, + 1.5, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = expr.eval(&mut scratch, 0.0, &[y0], &[], &[sel_val])[0]; + let want = expected_chain(y0, N, sel_val); + assert_eq!(got.to_bits(), want.to_bits(), "selector {sel_val}"); + } + } + + /// The same shape within budget still privatises, and agrees with the same + /// oracle: the cap changes how a chain is scheduled, never what it computes. + #[test] + fn a_short_chain_still_privatises_within_the_budget() { + use crate::eval::CompiledExpr; + use crate::ir::TypedIr; + + const N: usize = 4; + let (arena, cond) = chain_nested_conditional(N); + let (new_arena, _) = + privatise_conditionals(&arena, cond).expect("a 4-branch chain fits the budget"); + assert!(new_arena.len() > arena.len(), "cones must have been cloned"); + + let expr = CompiledExpr::from_ir(TypedIr::from_arena(&arena, cond)); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.4_f64; + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 1.5, f64::NAN] { + let got = expr.eval(&mut scratch, 0.0, &[y0], &[], &[sel_val])[0]; + let want = expected_chain(y0, N, sel_val); + assert_eq!(got.to_bits(), want.to_bits(), "selector {sel_val}"); + } + } + + /// A conditional can carry BOTH an aliased slot pair AND a genuinely shared, + /// disjoint subgraph at the same time: `branches[0] == branches[2] == x` + /// (index-set `{0,2}`) while `branches[1]` and `branches[3]` independently + /// share `q` (index-set `{1,3}`). Both privatise, independently: the two + /// index-sets never interfere. + #[test] + fn aliased_and_disjoint_shared_cones_privatise_independently() { + use crate::eval::CompiledExpr; + use crate::ir::TypedIr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let x = arena.alloc(Node::Scalar(1.0)); + let inner = arena.alloc(Node::Exp(y)); + let q = arena.alloc(Node::Sqrt(inner)); + let b1 = arena.alloc(Node::Sin(q)); + let b3 = arena.alloc(Node::Cos(q)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + // branches[0] and branches[2] alias the same node `x`; + // branches[1] and branches[3] independently share `q`. + branches: vec![x, b1, x, b3], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + assert_eq!( + owners[x.index()], + Ownership::Branches { + cond, + indices: vec![0, 2] + }, + "the aliased slot is owned by exactly the positions that name it" + ); + assert_eq!( + owners[q.index()], + Ownership::Branches { + cond, + indices: vec![1, 3] + }, + "a disjoint shared cone privatises on its own index-set" + ); + + let (new_arena, new_root) = + privatise_conditionals(&arena, cond).expect("the disjoint shared cone must privatise"); + + let new_order = new_arena.topological_order(new_root); + let n_sqrt = new_order + .iter() + .filter(|&&id| matches!(new_arena.get(id), Node::Sqrt(_))) + .count(); + let n_exp = new_order + .iter() + .filter(|&&id| matches!(new_arena.get(id), Node::Exp(_))) + .count(); + assert_eq!( + (n_sqrt, n_exp), + (2, 2), + "q and inner must be cloned once per sharing branch" + ); + let n_aliased_scalar = new_order + .iter() + .filter(|&&id| matches!(new_arena.get(id), Node::Scalar(v) if v.to_bits() == 1.0_f64.to_bits())) + .count(); + assert_eq!( + n_aliased_scalar, 2, + "the aliased slot is cloned once per position that names it" + ); + + let before = CompiledExpr::from_ir(TypedIr::from_arena_raw(&arena, cond)); + let after = CompiledExpr::from_ir(TypedIr::from_arena_raw(&new_arena, new_root)); + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 1.5, f64::NAN] { + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.3], &[], &[sel_val]); + let b = after.eval(&mut s2, 0.0, &[0.3], &[], &[sel_val]); + assert_eq!(a[0].to_bits(), b[0].to_bits(), "selector {sel_val}"); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/coloring.rs b/packages/pybamm-rust/pybamm-core/src/coloring.rs new file mode 100644 index 0000000000..be9025ab61 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/coloring.rs @@ -0,0 +1,639 @@ +//! Column grouping that compresses a sparse Jacobian into few JVP passes. +//! +//! Two columns that never share a row can be evaluated together: seed the +//! tangent vector with 1 in both and one forward pass returns both columns' +//! entries, unmixed. Grouping the columns of a [`SparsityPattern`] that way +//! turns `ncols` JVP passes into `n_colors`, which for a banded `PyBaMM` Jacobian +//! is a handful regardless of mesh size. +//! +//! Only the pattern is consulted, never the values, so a coloring computed once +//! at compile time stays valid for every later evaluation. +//! +//! Entries a compile pass already knows are exempt: [`color_columns_masked`] +//! seeds only the columns a sweep must still recover, and lets two of them +//! share a color unless a row makes one of them read the other. + +use crate::sparsity::SparsityPattern; + +/// A column no sweep produces, so it never receives a color. +pub const UNSEEDED: usize = usize::MAX; + +/// Result of graph coloring, groups columns that can be computed together. +#[derive(Debug, Clone)] +pub struct ColumnColoring { + /// Color assigned to each column (0-indexed), or [`UNSEEDED`]. + pub colors: Vec, + /// Total number of colors (= number of JVP calls needed). + pub n_colors: usize, + /// For each color, the columns assigned that color, in ascending order. + pub color_to_columns: Vec>, +} + +impl ColumnColoring { + /// Columns assigned a given color (slice into the precomputed table). + /// + /// Every seeded column is colored, including structurally-zero ones, so + /// consumers must scatter results via a pattern-restricted entry table + /// (e.g. `color_to_csc_entries`), never by writing all same-color columns. + #[inline] + pub fn columns_with_color(&self, color: usize) -> &[usize] { + &self.color_to_columns[color] + } + + /// Columns a sweep produces, i.e. those that received a color. + #[inline] + pub fn n_seeded_columns(&self) -> usize { + self.color_to_columns.iter().map(Vec::len).sum() + } +} + +/// Column-adjacency graph in CSR form. +pub(crate) struct ColumnAdjacency { + pub indptr: Vec, + pub indices: Vec, +} + +/// Column adjacency restricted to the entries a sweep must recover. +/// +/// `swept` is indexed by CSR entry (one flag per entry, always) and `seeded` by +/// column. An edge is emitted only where a swept entry would pick a seeded +/// column up: within a row, every swept column conflicts with every other +/// seeded column present, while two columns that are both merely constant +/// there never do. +fn build_column_adjacency_masked( + pattern: &SparsityPattern, + swept: &[bool], + seeded: &[bool], +) -> ColumnAdjacency { + build_column_adjacency_masked_counted(pattern, swept, seeded).0 +} + +/// [`build_column_adjacency_masked`], also returning how many neighbours were +/// appended. +/// +/// A test can assert that equals the graph's edge count, which is what pins the +/// dedup to insertion time: buffering duplicates instead costs one append per +/// (row, column) pair, so a structurally dense row would push `ncols` per column. +fn build_column_adjacency_masked_counted( + pattern: &SparsityPattern, + swept: &[bool], + seeded: &[bool], +) -> (ColumnAdjacency, usize) { + let ncols = pattern.ncols; + // Column-major view of the pattern, carrying the CSR entry index so the walk + // below can ask whether this column is swept in that row without rescanning. + let mut col_start = vec![0usize; ncols + 1]; + for &col in &pattern.indices { + col_start[col + 1] += 1; + } + for col in 0..ncols { + col_start[col + 1] += col_start[col]; + } + let mut col_entries = vec![(0usize, 0usize); pattern.indices.len()]; + let mut fill = col_start.clone(); + for row in 0..pattern.nrows { + for entry in pattern.indptr[row]..pattern.indptr[row + 1] { + let col = pattern.indices[entry]; + col_entries[fill[col]] = (row, entry); + fill[col] += 1; + } + } + // A row with no swept entry conflicts nothing, so hoist the test out of the + // per-column walk that would otherwise repeat it once per column in the row. + let row_has_swept: Vec = (0..pattern.nrows) + .map(|row| { + swept[pattern.indptr[row]..pattern.indptr[row + 1]] + .iter() + .any(|&s| s) + }) + .collect(); + + // Dedup against a per-source stamp rather than by sorting a multiset: each + // structurally dense row otherwise pushes `ncols` entries into every one of + // its columns, so d dense rows cost d * ncols^2 live before the dedup. + let mut sets: Vec> = vec![Vec::new(); ncols]; + let mut stamp = vec![usize::MAX; ncols]; + let mut appended = 0; + for col in 0..ncols { + stamp[col] = col; + for &(row, entry) in &col_entries[col_start[col]..col_start[col + 1]] { + let col_swept = swept[entry]; + if !row_has_swept[row] || (!col_swept && !seeded[col]) { + continue; + } + let span = pattern.indptr[row]..pattern.indptr[row + 1]; + for (&other, &other_swept) in pattern.indices[span.clone()].iter().zip(&swept[span]) { + let conflicts = if col_swept { + other_swept || seeded[other] + } else { + other_swept + }; + if conflicts && stamp[other] != col { + stamp[other] = col; + sets[col].push(other); + appended += 1; + } + } + } + sets[col].sort_unstable(); + } + + let mut indptr = Vec::with_capacity(ncols + 1); + let mut indices = Vec::new(); + indptr.push(0); + for set in &sets { + indices.extend_from_slice(set); + indptr.push(indices.len()); + } + (ColumnAdjacency { indptr, indices }, appended) +} + +/// DSATUR (maximum saturation degree) graph coloring for Jacobian column +/// grouping. +/// +/// At each step the uncolored column with the most distinctly-colored +/// neighbours is coloured next (ties broken by higher graph degree, then lower +/// index for determinism) and given the smallest colour no neighbour uses. +pub fn color_columns(pattern: &SparsityPattern) -> ColumnColoring { + // Every column is seeded here, structurally-empty ones included, which is + // what consumers that index `colors` unconditionally rely on. + color_with_seeds( + pattern, + &vec![true; pattern.nnz()], + &vec![true; pattern.ncols], + ) +} + +/// As [`color_columns`], but told which entries a sweep must recover. +/// +/// `swept` carries one flag per CSR entry. Two columns conflict on a shared row +/// when at least one is swept there; a column with no swept entry produces +/// nothing and gets [`UNSEEDED`] in `colors`. +pub fn color_columns_masked(pattern: &SparsityPattern, swept: &[bool]) -> ColumnColoring { + let mut seeded = vec![false; pattern.ncols]; + for (&col, &is_swept) in pattern.indices.iter().zip(swept) { + if is_swept { + seeded[col] = true; + } + } + color_with_seeds(pattern, swept, &seeded) +} + +/// DSATUR over the seeded columns, with `swept` deciding which entries create a +/// conflict. Shared by the masked and unmasked entry points. +fn color_with_seeds(pattern: &SparsityPattern, swept: &[bool], seeded: &[bool]) -> ColumnColoring { + let ncols = pattern.ncols; + if ncols == 0 { + return ColumnColoring { + colors: vec![], + n_colors: 0, + color_to_columns: vec![], + }; + } + + // Selection scans this list, not every column, which is what keeps the + // O(n^2) loop proportional to the columns a sweep actually produces. + let seeded_cols: Vec = (0..ncols).filter(|&col| seeded[col]).collect(); + + let adj = build_column_adjacency_masked(pattern, swept, seeded); + let degree = |v: usize| adj.indptr[v + 1] - adj.indptr[v]; + + let mut colors = vec![UNSEEDED; ncols]; + // Distinct neighbour colours seen by each still-uncolored column; its + // length is the column's saturation degree. + let mut neighbour_colors: Vec> = + vec![std::collections::BTreeSet::new(); ncols]; + let mut forbidden = vec![false; ncols]; + let mut n_colors: usize = 0; + + for _ in 0..seeded_cols.len() { + // Select the uncolored column with maximum saturation degree. + let mut best: Option = None; + for &v in &seeded_cols { + if colors[v] != UNSEEDED { + continue; + } + best = Some(best.map_or(v, |b| { + let (sv, sb) = (neighbour_colors[v].len(), neighbour_colors[b].len()); + if sv > sb || (sv == sb && degree(v) > degree(b)) { + v + } else { + b + } + })); + } + let col = best.expect("an uncolored seeded column remains while iterating"); + + let neighbours = &adj.indices[adj.indptr[col]..adj.indptr[col + 1]]; + + // Smallest colour not used by any neighbour. + for &nbr in neighbours { + if colors[nbr] != UNSEEDED { + forbidden[colors[nbr]] = true; + } + } + let mut c = 0; + while c < forbidden.len() && forbidden[c] { + c += 1; + } + for &nbr in neighbours { + if colors[nbr] != UNSEEDED { + forbidden[colors[nbr]] = false; + } + } + + colors[col] = c; + n_colors = n_colors.max(c + 1); + + // Propagate the new colour into uncolored neighbours' saturation sets. + for &nbr in neighbours { + if colors[nbr] == UNSEEDED { + neighbour_colors[nbr].insert(c); + } + } + } + + let mut color_to_columns: Vec> = vec![Vec::new(); n_colors]; + for (col, &c) in colors.iter().enumerate() { + if c != UNSEEDED { + color_to_columns[c].push(col); + } + } + + ColumnColoring { + colors, + n_colors, + color_to_columns, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + fn make_dense_pattern(nrows: usize, ncols: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(nrows, ncols); + for row in 0..nrows { + pattern.indptr[row] = row * ncols; + for col in 0..ncols { + pattern.indices.push(col); + } + } + pattern.indptr[nrows] = nrows * ncols; + pattern + } + + fn make_diagonal_pattern(n: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(n, n); + for i in 0..n { + pattern.indptr[i] = i; + pattern.indices.push(i); + } + pattern.indptr[n] = n; + pattern + } + + fn make_tridiagonal_pattern(n: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(n, n); + let mut idx = 0; + for row in 0..n { + pattern.indptr[row] = idx; + if row > 0 { + pattern.indices.push(row - 1); + idx += 1; + } + pattern.indices.push(row); + idx += 1; + if row < n - 1 { + pattern.indices.push(row + 1); + idx += 1; + } + } + pattern.indptr[n] = idx; + pattern + } + + #[test] + fn test_color_dense_matrix() { + let pattern = make_dense_pattern(3, 4); + let coloring = color_columns(&pattern); + assert_eq!(coloring.n_colors, 4); + let unique_colors: HashSet<_> = coloring.colors.iter().copied().collect(); + assert_eq!(unique_colors.len(), 4); + } + + #[test] + fn test_color_diagonal_matrix() { + let pattern = make_diagonal_pattern(5); + let coloring = color_columns(&pattern); + assert_eq!(coloring.n_colors, 1); + assert!(coloring.colors.iter().all(|&c| c == 0)); + } + + #[test] + fn test_color_tridiagonal_matrix() { + let pattern = make_tridiagonal_pattern(6); + let coloring = color_columns(&pattern); + assert!(coloring.n_colors <= 3); + for col in 0..5 { + assert_ne!(coloring.colors[col], coloring.colors[col + 1]); + } + } + + #[test] + fn test_columns_with_color() { + let pattern = make_diagonal_pattern(4); + let coloring = color_columns(&pattern); + let cols = coloring.columns_with_color(0); + assert_eq!(cols, vec![0, 1, 2, 3]); + } + + #[test] + fn test_color_to_columns_matches_columns_with_color() { + let pattern = make_tridiagonal_pattern(8); + let coloring = color_columns(&pattern); + for color in 0..coloring.n_colors { + let direct: Vec = coloring + .colors + .iter() + .enumerate() + .filter_map(|(i, &c)| (c == color).then_some(i)) + .collect(); + assert_eq!(coloring.color_to_columns[color], direct); + } + } + + /// Validity check that a colouring respects column-adjacency (distance-1). + fn assert_valid_coloring(pattern: &SparsityPattern, coloring: &ColumnColoring) { + assert_valid_masked_coloring(pattern, &vec![true; pattern.nnz()], coloring); + } + + /// As [`assert_valid_coloring`], under the mask the colouring was given. + fn assert_valid_masked_coloring( + pattern: &SparsityPattern, + swept: &[bool], + coloring: &ColumnColoring, + ) { + let seeded: Vec = (0..pattern.ncols) + .map(|col| coloring.colors[col] != UNSEEDED) + .collect(); + let adj = build_column_adjacency_masked(pattern, swept, &seeded); + for (col, _) in seeded.iter().enumerate().filter(|&(_, &s)| s) { + for &nbr in &adj.indices[adj.indptr[col]..adj.indptr[col + 1]] { + assert_ne!( + coloring.colors[col], coloring.colors[nbr], + "adjacent columns {col} and {nbr} share a colour" + ); + } + } + } + + /// Build a sparsity pattern whose column-adjacency graph is exactly `edges` + /// (one matrix row per edge). + fn make_graph_from_edges(ncols: usize, edges: &[(usize, usize)]) -> SparsityPattern { + let mut pattern = SparsityPattern::new(edges.len(), ncols); + let mut idx = 0; + for (row, &(i, j)) in edges.iter().enumerate() { + assert_ne!(i, j, "self-loops are not valid graph edges"); + let (lo, hi) = if i < j { (i, j) } else { (j, i) }; + pattern.indptr[row] = idx; + pattern.indices.push(lo); + pattern.indices.push(hi); + idx += 2; + } + pattern.indptr[edges.len()] = idx; + pattern + } + + /// Bidiagonal pattern: row k holds columns {k, k+1}. The resulting column + /// graph is a simple path 0-1-...-(n-1), which is bipartite. + fn make_bidiagonal_pattern(n: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(n - 1, n); + let mut idx = 0; + for row in 0..n - 1 { + pattern.indptr[row] = idx; + pattern.indices.push(row); + pattern.indices.push(row + 1); + idx += 2; + } + pattern.indptr[n - 1] = idx; + pattern + } + + #[test] + fn test_dsatur_reaches_optimum_on_path() { + // A path column-graph is bipartite, so 2 colours is optimal; DSATUR also + // hits the structural lower bound on the battery Jacobians. + let pattern = make_bidiagonal_pattern(12); + let coloring = color_columns(&pattern); + assert_valid_coloring(&pattern, &coloring); + assert_eq!(coloring.n_colors, 2, "DSATUR should 2-colour a path graph"); + } + + #[test] + fn test_dsatur_valid_on_dense_and_tridiagonal() { + for pattern in [make_dense_pattern(5, 5), make_tridiagonal_pattern(20)] { + let coloring = color_columns(&pattern); + assert_valid_coloring(&pattern, &coloring); + } + } + + #[test] + fn test_dsatur_matches_known_chromatic_number() { + // DSATUR must use exactly the proven chromatic number on each graph, built + // as a column-adjacency graph with one matrix row per edge. + + // (name, n_columns, edges, expected chromatic number) + type ChromaticCase = (&'static str, usize, Vec<(usize, usize)>, usize); + let cases: [ChromaticCase; 4] = [ + // K3 triangle: a 3-clique forces χ ≥ 3, and 3 suffices. + ("K3", 3, vec![(0, 1), (1, 2), (2, 0)], 3), + // C5 odd cycle: not bipartite, so χ = 3 though its largest clique is 2, + // a stronger bound than clique size, and the likeliest to mis-colour. + ("C5", 5, vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)], 3), + // C6 even cycle: bipartite, so χ = 2. + ( + "C6", + 6, + vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 0)], + 2, + ), + // K4 complete graph: every pair adjacent, so χ = 4. + ( + "K4", + 4, + vec![(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)], + 4, + ), + ]; + for (name, ncols, edges, chromatic) in &cases { + let pattern = make_graph_from_edges(*ncols, edges); + let coloring = color_columns(&pattern); + assert_valid_coloring(&pattern, &coloring); + assert_eq!( + coloring.n_colors, *chromatic, + "{name}: DSATUR should use exactly {chromatic} colours, got {}", + coloring.n_colors + ); + } + } + + #[test] + fn test_column_adjacency_for_dense_matrix() { + let pattern = make_dense_pattern(4, 4); + let adj = build_column_adjacency_masked(&pattern, &[true; 16], &[true; 4]); + for col in 0..4 { + let neighbours = &adj.indices[adj.indptr[col]..adj.indptr[col + 1]]; + let expected: Vec = (0..4).filter(|&c| c != col).collect(); + assert_eq!(neighbours, expected.as_slice()); + } + } + + /// Brute-force restatement of the edge rule this module documents, used to + /// pin the incremental build against a version with no dedup subtlety. + fn reference_adjacency( + pattern: &SparsityPattern, + swept: &[bool], + seeded: &[bool], + ) -> Vec> { + let mut sets = vec![std::collections::BTreeSet::new(); pattern.ncols]; + for row in 0..pattern.nrows { + let (start, end) = (pattern.indptr[row], pattern.indptr[row + 1]); + let cols = &pattern.indices[start..end]; + let (swept_cols, other_cols): (Vec, Vec) = ( + cols.iter() + .zip(&swept[start..end]) + .filter(|&(_, &s)| s) + .map(|(&c, _)| c) + .collect(), + cols.iter() + .zip(&swept[start..end]) + .filter(|&(&c, &s)| !s && seeded[c]) + .map(|(&c, _)| c) + .collect(), + ); + if swept_cols.is_empty() { + continue; + } + for &col in &swept_cols { + sets[col].extend(swept_cols.iter().filter(|&&o| o != col)); + sets[col].extend(other_cols.iter().filter(|&&o| o != col)); + } + for &col in &other_cols { + sets[col].extend(swept_cols.iter().filter(|&&o| o != col)); + } + } + sets.into_iter() + .map(|set| set.into_iter().collect()) + .collect() + } + + fn assert_matches_reference(pattern: &SparsityPattern, swept: &[bool], seeded: &[bool]) { + let adj = build_column_adjacency_masked(pattern, swept, seeded); + for (col, expected) in reference_adjacency(pattern, swept, seeded) + .iter() + .enumerate() + { + assert_eq!( + &adj.indices[adj.indptr[col]..adj.indptr[col + 1]], + expected.as_slice(), + "column {col} neighbours differ from the reference rule" + ); + } + } + + #[test] + fn column_adjacency_matches_the_reference_rule() { + // Dense, banded and mixed-mask shapes: the last is the one where a + // column is present in a row without conflicting there. + let dense = make_dense_pattern(6, 5); + assert_matches_reference(&dense, &vec![true; dense.nnz()], &[true; 5]); + + let banded = make_tridiagonal_pattern(9); + assert_matches_reference(&banded, &vec![true; banded.nnz()], &[true; 9]); + + let mut mask = vec![true; dense.nnz()]; + for (entry, flag) in mask.iter_mut().enumerate() { + *flag = entry % 3 != 0; + } + let mut seeded = vec![false; 5]; + for (&col, &is_swept) in dense.indices.iter().zip(&mask) { + seeded[col] |= is_swept; + } + assert_matches_reference(&dense, &mask, &seeded); + } + + #[test] + fn dense_rows_never_buffer_duplicate_neighbours() { + // Each of these rows alone makes every column adjacent to every other, + // so appending before deduping would cost `n_dense * ncols` per column + // instead of the ncols - 1 edges that survive. + let (ncols, n_dense) = (24, 8); + let mut pattern = SparsityPattern::new(n_dense, ncols); + for row in 0..n_dense { + pattern.indices.extend(0..ncols); + pattern.indptr[row + 1] = pattern.indices.len(); + } + let (adj, appended) = build_column_adjacency_masked_counted( + &pattern, + &vec![true; pattern.nnz()], + &vec![true; ncols], + ); + assert_eq!(adj.indices.len(), ncols * (ncols - 1)); + assert_eq!( + appended, + adj.indices.len(), + "every append must survive the dedup, not be sorted out afterwards" + ); + } + + #[test] + fn masked_coloring_ignores_wholly_constant_rows() { + // A dense pattern whose every entry is known needs no sweep at all. + let pattern = make_dense_pattern(4, 4); + let coloring = color_columns_masked(&pattern, &vec![false; pattern.nnz()]); + assert_eq!(coloring.n_colors, 0); + assert!(coloring.colors.iter().all(|&c| c == UNSEEDED)); + assert_eq!(coloring.n_seeded_columns(), 0); + } + + #[test] + fn masked_coloring_keeps_swept_entries_unpolluted() { + // Row 0 is dense but only column 0 varies there, so column 0 must not + // share a colour with any of the constants it would otherwise read. + let pattern = make_dense_pattern(1, 5); + let mut swept = vec![false; pattern.nnz()]; + swept[0] = true; + let coloring = color_columns_masked(&pattern, &swept); + assert_valid_masked_coloring(&pattern, &swept, &coloring); + assert_eq!(coloring.n_colors, 1, "only one column is ever seeded"); + assert_eq!(coloring.colors[0], 0); + assert!(coloring.colors[1..].iter().all(|&c| c == UNSEEDED)); + } + + #[test] + fn masked_coloring_beats_the_full_rule_on_a_mixed_row() { + // Column 1 is constant wherever it appears, so it costs a colour under + // the full rule and none here. + let mut pattern = SparsityPattern::new(2, 3); + pattern.indptr = vec![0, 3, 4]; + pattern.indices = vec![0, 1, 2, 0]; + let swept = vec![true, false, true, true]; + assert_eq!(color_columns(&pattern).n_colors, 3); + let coloring = color_columns_masked(&pattern, &swept); + assert_valid_masked_coloring(&pattern, &swept, &coloring); + assert_eq!(coloring.n_colors, 2); + assert_ne!(coloring.colors[0], coloring.colors[2]); + assert_eq!(coloring.colors[1], UNSEEDED); + } + + #[test] + fn an_all_swept_mask_matches_the_unmasked_coloring() { + for pattern in [make_dense_pattern(5, 5), make_tridiagonal_pattern(20)] { + let full = color_columns(&pattern); + let masked = color_columns_masked(&pattern, &vec![true; pattern.nnz()]); + assert_eq!(full.colors, masked.colors); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/const_entries.rs b/packages/pybamm-rust/pybamm-core/src/const_entries.rs new file mode 100644 index 0000000000..77e5981d93 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/const_entries.rs @@ -0,0 +1,1023 @@ +//! Jacobian entries a compile pass proves independent of `(t, y, y_dot, inputs)`. +//! +//! The tangent expression `T(y, s) = J(y) · s` is linear in the seed `s`, so the +//! coefficient of every seed in every output element falls out of one bottom-up +//! pass over the tape. A coefficient that folds to a literal is a Jacobian entry +//! no sweep has to produce: assembly writes it from a table, and the column +//! coloring only has to keep the remaining entries exact. +//! +//! The pass is conservative in one direction only. A coefficient it cannot +//! resolve is reported as varying, never as constant, and a node kind it does +//! not model marks the elements it reaches *incomplete*, after which a seed's +//! absence proves nothing and every entry of the affected row is swept. +//! +//! Classifying here — on the simplified tangent tape rather than on the primal +//! graph — is what makes a folded value bit-identical to the one the sweep +//! would have produced: the fold walks the same operators in the same order, so +//! a `MatMul` accumulates over its columns in the same sequence. The single +//! difference it allows itself is the sign of a zero, which dropping an exactly +//! zero term can flip and which `simplify` already documents as acceptable. +//! +//! One boundary sits outside that guarantee. Where the tape overflows, a term +//! the fold drops as exactly zero evaluates to `inf * 0.0 = NaN` and poisons +//! the sweep, so a folded entry can be finite where the swept one is not. The +//! rest of that row still comes out `NaN` and the step is rejected either way, +//! which is why nothing here tries to reproduce the poisoning. + +use std::cmp::Ordering; + +use crate::arena::{Arena, NodeId}; +use crate::ir::infer_sizes; +use crate::node::Node; +use crate::sparsity::SparsityPattern; + +/// Coefficients per element beyond which an element is abandoned rather than +/// grown. Bounds compile time and memory on a near-dense, fully linear model, +/// where the propagation would otherwise compute the whole Jacobian +/// symbolically; a row wider than this is swept, exactly as it is today. +const MAX_CONSTANT_DEGREE: usize = 64; + +/// One seed's coefficient in one element of the tangent expression. +#[derive(Clone, Copy, Debug, PartialEq)] +struct Coeff { + /// Seed index, which is the Jacobian column. + seed: usize, + /// Folded value, or `None` when the coefficient depends on the state. + value: Option, +} + +impl Coeff { + const fn unknown(seed: usize) -> Self { + Self { seed, value: None } + } + + /// Same seed, value transformed wherever it is resolved. + fn map(self, f: impl Fn(f64) -> f64) -> Self { + Self { + seed: self.seed, + value: self.value.map(f), + } + } + + fn negated(self) -> Self { + self.map(|v| -v) + } +} + +/// One element's resolved coefficients, ascending by seed. +#[derive(Clone, Copy, Debug)] +struct Element<'a> { + coeffs: &'a [Coeff], + /// False once a node kind the pass cannot model has been reached, after + /// which an absent seed proves nothing. + complete: bool, +} + +impl Element<'_> { + /// An element no seed reaches: every coefficient is exactly zero. + const SEED_FREE: Self = Self { + coeffs: &[], + complete: true, + }; +} + +/// Per-element seed coefficients of one node, flattened CSR-style. +/// +/// A single element means "broadcast": every element of the node reads it. An +/// element flagged incomplete carries no entries, since they could no longer +/// prove anything about the seeds they omit. +#[derive(Clone, Debug)] +struct Coeffs { + /// Element start offsets into `entries`, length `n_elems + 1`. + offsets: Vec, + /// Coefficients of every element, ascending by seed within an element. + entries: Vec, + /// Per element, parallel to the `offsets` windows. + complete: Vec, +} + +impl Coeffs { + /// Empty, sized for `n_elems` pushes. At least one coefficient per element + /// is the common shape, so `entries` is reserved for that. + fn with_capacity(n_elems: usize) -> Self { + let mut offsets = Vec::with_capacity(n_elems + 1); + offsets.push(0); + Self { + offsets, + entries: Vec::with_capacity(n_elems), + complete: Vec::with_capacity(n_elems), + } + } + + /// `n` elements the pass could not model at all. + fn unresolved(n: usize) -> Self { + Self { + offsets: vec![0; n + 1], + entries: Vec::new(), + complete: vec![false; n], + } + } + + /// Close one element. Over the cap it is recorded incomplete and its + /// entries dropped, which is what bounds the pass on a dense linear model. + fn push(&mut self, coeffs: &[Coeff], complete: bool) { + let keep = complete && coeffs.len() <= MAX_CONSTANT_DEGREE; + if keep { + self.entries.extend_from_slice(coeffs); + } + self.complete.push(keep); + self.offsets.push(self.entries.len()); + } + + fn push_element(&mut self, element: Element<'_>) { + self.push(element.coeffs, element.complete); + } + + const fn n_elems(&self) -> usize { + self.complete.len() + } + + fn elem(&self, i: usize) -> Element<'_> { + let i = i.min(self.n_elems() - 1); + Element { + coeffs: &self.entries[self.offsets[i]..self.offsets[i + 1]], + complete: self.complete[i], + } + } +} + +/// A node's value where the pass folded it to literals; `Scalar` broadcasts. +#[derive(Clone, Debug)] +enum Literal { + Scalar(f64), + Vector(Vec), +} + +impl Literal { + fn at(&self, i: usize) -> f64 { + match self { + Self::Scalar(v) => *v, + Self::Vector(v) => v[i.min(v.len() - 1)], + } + } + + fn negated(&self) -> Self { + match self { + Self::Scalar(v) => Self::Scalar(-v), + Self::Vector(v) => Self::Vector(v.iter().map(|x| -x).collect()), + } + } +} + +/// What the pass knows about one node's value. +#[derive(Debug)] +enum NodeInfo { + /// Free of every seed and not folded: an ordinary primal subexpression. + Primal, + /// Folds to a literal, which is what lets a `Mul`/`Div` scale a coefficient. + Literal(Literal), + /// Carries seeds; the general case. + Seeded(Coeffs), +} + +impl NodeInfo { + fn elem(&self, i: usize) -> Element<'_> { + match self { + Self::Seeded(coeffs) => coeffs.elem(i), + Self::Primal | Self::Literal(_) => Element::SEED_FREE, + } + } + + const fn n_elems(&self) -> usize { + match self { + Self::Primal | Self::Literal(Literal::Scalar(_)) => 1, + Self::Literal(Literal::Vector(v)) => v.len(), + Self::Seeded(coeffs) => coeffs.n_elems(), + } + } + + const fn literal(&self) -> Option<&Literal> { + match self { + Self::Literal(literal) => Some(literal), + Self::Primal | Self::Seeded(_) => None, + } + } + + const fn is_seeded(&self) -> bool { + matches!(self, Self::Seeded(_)) + } +} + +/// Merge buffers reused across nodes, so the pass allocates per node rather +/// than per element. +#[derive(Debug, Default)] +struct Scratch { + out: Vec, + acc: Vec, + children: Vec, +} + +/// `out = map_a(a) + map_b(b)`, a sorted merge over seed indices. Either side +/// unresolved leaves that seed's sum unresolved. +/// +/// The three uses -- sum, difference, scaled accumulation -- differ only in +/// those two maps, and each is exact: `x - y == x + (-y)` and `factor * y == +/// y * factor` bit for bit, so the fold still matches the tape. +fn merge( + a: &[Coeff], + b: &[Coeff], + out: &mut Vec, + map_a: impl Fn(Coeff) -> Coeff, + map_b: impl Fn(Coeff) -> Coeff, +) { + out.clear(); + let (mut i, mut j) = (0, 0); + while i < a.len() && j < b.len() { + match a[i].seed.cmp(&b[j].seed) { + Ordering::Less => { + out.push(map_a(a[i])); + i += 1; + }, + Ordering::Greater => { + out.push(map_b(b[j])); + j += 1; + }, + Ordering::Equal => { + let (left, right) = (map_a(a[i]), map_b(b[j])); + out.push(Coeff { + seed: left.seed, + value: left.value.zip(right.value).map(|(x, y)| x + y), + }); + i += 1; + j += 1; + }, + } + } + out.extend(a[i..].iter().copied().map(&map_a)); + out.extend(b[j..].iter().copied().map(&map_b)); +} + +/// `out = a + b`, or `a - b` when `subtract`. +fn merge_sum(a: &[Coeff], b: &[Coeff], subtract: bool, out: &mut Vec) { + merge(a, b, out, |c| c, |c| if subtract { c.negated() } else { c }); +} + +/// `out = acc + factor * coeffs`, the accumulation one `MatMul` row performs. +fn merge_scaled(acc: &[Coeff], coeffs: &[Coeff], factor: f64, out: &mut Vec) { + merge(acc, coeffs, out, |c| c, |c| c.map(|v| v * factor)); +} + +/// `out` = the seeds of `a` and `b`, every coefficient unresolved: all a node +/// kind the pass cannot model can still say is which columns it touches. +fn merge_unknown(a: &[Coeff], b: &[Coeff], out: &mut Vec) { + let unknown = |c: Coeff| Coeff::unknown(c.seed); + merge(a, b, out, unknown, unknown); +} + +/// Union the seeds of `elements` into `acc`, every coefficient unresolved. +/// +/// Returns false at the first incomplete element or once the degree cap is +/// passed, after which the accumulated seeds prove nothing about the ones they +/// omit. Seed-free elements contribute nothing and are skipped. +fn union_seeds<'a>( + elements: impl Iterator>, + acc: &mut Vec, + out: &mut Vec, +) -> bool { + acc.clear(); + for element in elements { + if !element.complete { + return false; + } + if element.coeffs.is_empty() { + continue; + } + merge_unknown(acc, element.coeffs, out); + std::mem::swap(acc, out); + if acc.len() > MAX_CONSTANT_DEGREE { + return false; + } + } + true +} + +/// Classify every entry of `pattern` against a tangent expression: a per-CSR +/// -entry mask of the ones a sweep must still produce, and `(csr_idx, value)` +/// for the rest. +/// +/// The tangent expression's linearity in the seed makes the coefficient of each +/// seed recoverable in one bottom-up pass; a coefficient that folds to a literal +/// is an entry no sweep has to produce. `tangent_root` must be the *simplified* +/// tangent expression, so that a folded value follows the same operator order +/// the tape executes. +/// +/// A pattern entry the tangent proves absent is a constant zero rather than an +/// omission, so the two halves together cover every entry of `pattern`. +#[must_use] +pub fn classify_constant_entries( + arena: &Arena, + tangent_root: NodeId, + pattern: &SparsityPattern, +) -> (Vec, Vec<(usize, f64)>) { + let rows = classify_rows(arena, tangent_root, pattern.nrows); + + let mut varying = vec![false; pattern.nnz()]; + let mut constants = Vec::new(); + for row in 0..pattern.nrows { + let (row_start, row_end) = (pattern.indptr[row], pattern.indptr[row + 1]); + let element = rows.elem(row); + if !element.complete { + varying[row_start..row_end].fill(true); + continue; + } + let mut pos = 0; + for (offset, &col) in pattern.indices[row_start..row_end].iter().enumerate() { + while pos < element.coeffs.len() && element.coeffs[pos].seed < col { + pos += 1; + } + let csr_idx = row_start + offset; + match element.coeffs.get(pos) { + Some(coeff) if coeff.seed == col => match coeff.value { + Some(value) => constants.push((csr_idx, value)), + None => varying[csr_idx] = true, + }, + // Absent from a complete element: exactly zero, forever. + _ => constants.push((csr_idx, 0.0)), + } + } + } + (varying, constants) +} + +/// Seed coefficients of each of `n_rows` output rows, never broadcast. +fn classify_rows(arena: &Arena, tangent_root: NodeId, n_rows: usize) -> Coeffs { + let order = arena.topological_order(tangent_root); + let sizes = infer_sizes(arena, &order); + + // Freeing at the last read keeps the live set a frontier over the + // tape rather than the whole graph. + let mut remaining = vec![0u32; arena.len()]; + for &id in &order { + arena.get(id).for_each_child(|child| { + remaining[child.index()] += 1; + }); + } + + let mut info: Vec> = (0..arena.len()).map(|_| None).collect(); + let mut scratch = Scratch::default(); + for &id in &order { + let computed = classify_node(arena, id, &sizes, &info, &mut scratch); + arena.get(id).for_each_child(|child| { + remaining[child.index()] -= 1; + if remaining[child.index()] == 0 { + info[child.index()] = None; + } + }); + info[id.index()] = Some(computed); + } + + let root = info[tangent_root.index()] + .take() + .expect("root is classified last"); + match root { + // Already one entry per row, and every element in it has been through + // the same cap check, so re-pushing it would only duplicate the table. + NodeInfo::Seeded(coeffs) if coeffs.n_elems() == n_rows => coeffs, + other => { + let mut rows = Coeffs::with_capacity(n_rows); + for row in 0..n_rows { + rows.push_element(other.elem(row)); + } + rows + }, + } +} + +fn child_info(info: &[Option], id: NodeId) -> &NodeInfo { + info[id.index()] + .as_ref() + .expect("children are classified before their parents") +} + +/// Elements to materialise: one when every operand broadcasts, else the node's +/// full width. +const fn combined_elems(a: &NodeInfo, b: &NodeInfo, len: usize) -> usize { + if a.n_elems() == 1 && b.n_elems() == 1 { + 1 + } else { + len + } +} + +/// Fold an element-wise binary operator over two literal operands. Anything +/// else is seed-free but unfoldable, which still blocks a `Mul` from scaling. +fn fold_literals(a: &NodeInfo, b: &NodeInfo, len: usize, op: impl Fn(f64, f64) -> f64) -> NodeInfo { + let (Some(left), Some(right)) = (a.literal(), b.literal()) else { + return NodeInfo::Primal; + }; + if let (Literal::Scalar(x), Literal::Scalar(y)) = (left, right) { + return NodeInfo::Literal(Literal::Scalar(op(*x, *y))); + } + NodeInfo::Literal(Literal::Vector( + (0..len).map(|i| op(left.at(i), right.at(i))).collect(), + )) +} + +/// Coefficients of `seeded` scaled by `factor` element-wise, or divided by it. +/// A non-literal factor leaves every coefficient unresolved: the tangent still +/// reaches those columns, but their values move with the state. +fn scale_seeded( + seeded: &NodeInfo, + factor: &NodeInfo, + len: usize, + divide: bool, + scratch: &mut Scratch, +) -> NodeInfo { + let n_elems = combined_elems(seeded, factor, len); + match factor.literal() { + Some(literal) if divide => map_coeffs(seeded, n_elems, scratch, |c, i| { + c.map(|v| v / literal.at(i)) + }), + Some(literal) => map_coeffs(seeded, n_elems, scratch, |c, i| { + c.map(|v| v * literal.at(i)) + }), + None => map_coeffs(seeded, n_elems, scratch, |c, _| Coeff::unknown(c.seed)), + } +} + +/// Every element's coefficients mapped through `f`, completeness preserved. +/// `f` is given the element index, which is what lets a per-element literal +/// factor scale each row by its own value. +fn map_coeffs( + seeded: &NodeInfo, + n_elems: usize, + scratch: &mut Scratch, + f: impl Fn(Coeff, usize) -> Coeff, +) -> NodeInfo { + let mut built = Coeffs::with_capacity(n_elems); + for i in 0..n_elems { + let element = seeded.elem(i); + scratch.out.clear(); + scratch.out.extend(element.coeffs.iter().map(|&c| f(c, i))); + built.push(&scratch.out, element.complete); + } + NodeInfo::Seeded(built) +} + +/// Seeds of both operands, every coefficient unresolved. The shape a product +/// or quotient of two seeded operands takes; it cannot arise from a tangent +/// transform, which is linear in the seed, but costs nothing to survive. +fn unknown_pair(a: &NodeInfo, b: &NodeInfo, len: usize, scratch: &mut Scratch) -> NodeInfo { + let n_elems = combined_elems(a, b, len); + let mut built = Coeffs::with_capacity(n_elems); + for i in 0..n_elems { + let (left, right) = (a.elem(i), b.elem(i)); + merge_unknown(left.coeffs, right.coeffs, &mut scratch.out); + built.push(&scratch.out, left.complete && right.complete); + } + NodeInfo::Seeded(built) +} + +/// One scalar element carrying every seed any child element touches: what a +/// reduction over an unknown position can promise. +fn reduce_all(children: &[NodeId], info: &[Option], scratch: &mut Scratch) -> NodeInfo { + let infos: Vec<&NodeInfo> = children.iter().map(|&c| child_info(info, c)).collect(); + if !infos.iter().any(|child| child.is_seeded()) { + return NodeInfo::Primal; + } + let Scratch { out, acc, .. } = scratch; + let complete = union_seeds( + infos + .iter() + .flat_map(|child| (0..child.n_elems()).map(move |i| child.elem(i))), + acc, + out, + ); + let mut built = Coeffs::with_capacity(1); + built.push(acc, complete); + NodeInfo::Seeded(built) +} + +/// Element-wise fallback for every node kind the pass does not model: the union +/// of the children's seeds, all unresolved. Sound because a node's value can +/// only depend on seeds its children depend on. +fn unknown_elementwise( + node: &Node, + sizes: &[usize], + info: &[Option], + len: usize, + scratch: &mut Scratch, +) -> NodeInfo { + scratch.children.clear(); + node.for_each_child(|child| scratch.children.push(child)); + let Scratch { out, acc, children } = scratch; + // Resolved once per node, not once per element. + let infos: Vec<(&NodeInfo, usize)> = children + .iter() + .map(|&c| (child_info(info, c), sizes[c.index()])) + .collect(); + if !infos.iter().any(|(child, _)| child.is_seeded()) { + return NodeInfo::Primal; + } + let broadcast = infos + .iter() + .all(|&(child, size)| child.n_elems() == 1 && size <= 1); + let n_elems = if broadcast { 1 } else { len }; + let mut built = Coeffs::with_capacity(n_elems); + for i in 0..n_elems { + let complete = union_seeds(infos.iter().map(|(child, _)| child.elem(i)), acc, out); + built.push(acc, complete); + } + NodeInfo::Seeded(built) +} + +/// Accumulate one matrix row's contributions in the tape's own column order, +/// so a folded value matches the sweep bit for bit. +fn matmul_row( + columns: impl Iterator, + vector: &NodeInfo, + scratch: &mut Scratch, +) -> bool { + let Scratch { out, acc, .. } = scratch; + acc.clear(); + for (col, value) in columns { + let element = vector.elem(col); + // A non-finite matrix entry makes the tape's `value * 0.0` term NaN + // where the fold drops it, so this row promises nothing. + if !element.complete || !value.is_finite() { + return false; + } + if element.coeffs.is_empty() { + continue; + } + merge_scaled(acc, element.coeffs, value, out); + std::mem::swap(acc, out); + if acc.len() > MAX_CONSTANT_DEGREE { + return false; + } + } + true +} + +/// `Add`/`Sub`: coefficients merge seed-wise; two seed-free operands fold. +fn classify_add_sub( + left: &NodeInfo, + right: &NodeInfo, + subtract: bool, + len: usize, + scratch: &mut Scratch, +) -> NodeInfo { + if !left.is_seeded() && !right.is_seeded() { + return fold_literals( + left, + right, + len, + |x, y| if subtract { x - y } else { x + y }, + ); + } + let n_elems = combined_elems(left, right, len); + let mut built = Coeffs::with_capacity(n_elems); + for i in 0..n_elems { + let (le, re) = (left.elem(i), right.elem(i)); + merge_sum(le.coeffs, re.coeffs, subtract, &mut scratch.out); + built.push(&scratch.out, le.complete && re.complete); + } + NodeInfo::Seeded(built) +} + +/// `MatMul` against a literal matrix: each row accumulates in the tape's own +/// column order. A non-literal matrix resolves nothing. +fn classify_matmul( + arena: &Arena, + matrix: NodeId, + vector: &NodeInfo, + len: usize, + scratch: &mut Scratch, +) -> NodeInfo { + if !vector.is_seeded() { + return NodeInfo::Primal; + } + match arena.get(matrix) { + Node::SparseMatrix(csr) => { + let mut built = Coeffs::with_capacity(csr.shape.rows); + for row in 0..csr.shape.rows { + let span = csr.indptr[row]..csr.indptr[row + 1]; + let columns = csr.indices[span.clone()] + .iter() + .copied() + .zip(csr.data[span].iter().copied()); + let complete = matmul_row(columns, vector, scratch); + built.push(&scratch.acc, complete); + } + NodeInfo::Seeded(built) + }, + Node::Array(array) => { + let (rows, cols) = (array.shape.rows, array.shape.cols); + let mut built = Coeffs::with_capacity(rows); + for row in 0..rows { + let columns = array.data[row * cols..(row + 1) * cols] + .iter() + .copied() + .enumerate(); + let complete = matmul_row(columns, vector, scratch); + built.push(&scratch.acc, complete); + } + NodeInfo::Seeded(built) + }, + _ => NodeInfo::Seeded(Coeffs::unresolved(len)), + } +} + +/// `Concat`: children's elements laid end to end, folding when all are literal. +fn classify_concat( + children: &[NodeId], + sizes: &[usize], + info: &[Option], + len: usize, +) -> NodeInfo { + let infos: Vec<&NodeInfo> = children.iter().map(|&c| child_info(info, c)).collect(); + let widths = children.iter().map(|&c| sizes[c.index()]); + + if infos.iter().any(|child| child.is_seeded()) { + let mut built = Coeffs::with_capacity(len); + for (child, width) in infos.iter().zip(widths) { + for i in 0..width { + built.push_element(child.elem(i)); + } + } + return NodeInfo::Seeded(built); + } + let Some(literals) = infos + .iter() + .map(|child| child.literal()) + .collect::>>() + else { + return NodeInfo::Primal; + }; + let mut values = Vec::with_capacity(len); + for (literal, width) in literals.iter().zip(widths) { + values.extend((0..width).map(|i| literal.at(i))); + } + NodeInfo::Literal(Literal::Vector(values)) +} + +/// `Index`: a window into the child's elements, or the child itself when it +/// broadcasts. +fn classify_index(child: &NodeInfo, start: usize, end: usize) -> NodeInfo { + match child { + NodeInfo::Primal => NodeInfo::Primal, + NodeInfo::Literal(Literal::Scalar(value)) => NodeInfo::Literal(Literal::Scalar(*value)), + NodeInfo::Literal(Literal::Vector(values)) => { + NodeInfo::Literal(Literal::Vector(values[start..end].to_vec())) + }, + NodeInfo::Seeded(coeffs) if coeffs.n_elems() == 1 => NodeInfo::Seeded(coeffs.clone()), + NodeInfo::Seeded(coeffs) => { + let mut built = Coeffs::with_capacity(end - start); + for i in start..end { + built.push_element(coeffs.elem(i)); + } + NodeInfo::Seeded(built) + }, + } +} + +fn classify_node( + arena: &Arena, + id: NodeId, + sizes: &[usize], + info: &[Option], + scratch: &mut Scratch, +) -> NodeInfo { + let len = sizes[id.index()]; + if len == 0 { + return NodeInfo::Primal; + } + match arena.get(id) { + Node::Scalar(value) => NodeInfo::Literal(Literal::Scalar(*value)), + Node::ZeroVector { .. } => NodeInfo::Literal(Literal::Scalar(0.0)), + Node::Array(array) => NodeInfo::Literal(Literal::Vector(array.data.clone())), + Node::SparseMatrix(_) + | Node::Time + | Node::InputParameter { .. } + | Node::StateVector { .. } + | Node::StateVectorDot { .. } => NodeInfo::Primal, + + Node::TangentStateVector { start, end } => { + let mut built = Coeffs::with_capacity(end - start); + for seed in *start..*end { + built.push( + &[Coeff { + seed, + value: Some(1.0), + }], + true, + ); + } + NodeInfo::Seeded(built) + }, + + // A parameter tangent is a seed outside this index space, so nothing + // downstream of it can be resolved against the state seeds. + Node::TangentParameter { .. } => NodeInfo::Seeded(Coeffs::unresolved(1)), + + Node::Neg(a) => match child_info(info, *a) { + NodeInfo::Primal => NodeInfo::Primal, + NodeInfo::Literal(literal) => NodeInfo::Literal(literal.negated()), + seeded @ NodeInfo::Seeded(_) => { + map_coeffs(seeded, seeded.n_elems(), scratch, |c, _| c.negated()) + }, + }, + + Node::Add(a, b) => classify_add_sub( + child_info(info, *a), + child_info(info, *b), + false, + len, + scratch, + ), + Node::Sub(a, b) => classify_add_sub( + child_info(info, *a), + child_info(info, *b), + true, + len, + scratch, + ), + + Node::Mul(a, b) => { + let (left, right) = (child_info(info, *a), child_info(info, *b)); + // Scaling is order-free in IEEE arithmetic, so one rule serves + // both `literal * tangent` and `tangent * literal`. + match (left.is_seeded(), right.is_seeded()) { + (false, false) => fold_literals(left, right, len, |x, y| x * y), + (false, true) => scale_seeded(right, left, len, false, scratch), + (true, false) => scale_seeded(left, right, len, false, scratch), + (true, true) => unknown_pair(left, right, len, scratch), + } + }, + + Node::Div(a, b) => { + let (left, right) = (child_info(info, *a), child_info(info, *b)); + if right.is_seeded() { + unknown_pair(left, right, len, scratch) + } else if left.is_seeded() { + scale_seeded(left, right, len, true, scratch) + } else { + fold_literals(left, right, len, |x, y| x / y) + } + }, + + Node::MatMul(matrix, vector) => { + classify_matmul(arena, *matrix, child_info(info, *vector), len, scratch) + }, + + Node::Index { child, start, end } => classify_index(child_info(info, *child), *start, *end), + + Node::Concat(children) => classify_concat(children, sizes, info, len), + + Node::MaxReduce(a) | Node::MinReduce(a) => reduce_all(&[*a], info, scratch), + Node::ReduceArgSelect { basis, picker, .. } => { + reduce_all(&[*basis, *picker], info, scratch) + }, + + other => unknown_elementwise(other, sizes, info, len, scratch), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::{CompiledExpr, TangentInputs}; + use crate::ir::TypedIr; + use crate::node::{ArrayData, CsrData, Shape}; + use crate::simplify::simplify_pipeline; + use crate::sparsity::detect_sparsity_per_output; + use crate::tangent::tangent_wrt_states; + + /// Classify `root`'s state Jacobian the way `JacobianData` does, and return + /// the pattern alongside the split. + fn classify_states( + arena: &Arena, + root: NodeId, + n_rows: usize, + n_states: usize, + ) -> (SparsityPattern, Vec, Vec<(usize, f64)>) { + let mut diff_arena = arena.clone(); + let tangent_root = tangent_wrt_states(&mut diff_arena, root); + let (diff_arena, tangent_root) = simplify_pipeline(diff_arena, tangent_root); + let pattern = detect_sparsity_per_output(arena, root, n_rows, n_states); + let (varying, entries) = classify_constant_entries(&diff_arena, tangent_root, &pattern); + (pattern, varying, entries) + } + + /// Dense `(row, col) -> value` view of a split's constant table. + fn constant_map( + pattern: &SparsityPattern, + entries: &[(usize, f64)], + ) -> std::collections::HashMap<(usize, usize), f64> { + let rows = pattern.entry_rows(); + entries + .iter() + .map(|&(csr_idx, value)| ((rows[csr_idx], pattern.indices[csr_idx]), value)) + .collect() + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins the fold + fn linear_row_folds_to_exact_coefficients() { + // f(y) = [2*y0 - y1, y1] over 2 states: row 0 is fully constant. + 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 two = arena.alloc(Node::Scalar(2.0)); + let scaled = arena.alloc(Node::Mul(two, y0)); + let row0 = arena.alloc(Node::Sub(scaled, y1)); + let root = arena.alloc(Node::Concat(vec![row0, y1])); + + let (pattern, varying, entries) = classify_states(&arena, root, 2, 2); + assert!(!varying.iter().any(|&v| v), "a linear model sweeps nothing"); + let constants = constant_map(&pattern, &entries); + assert_eq!(constants[&(0, 0)], 2.0); + assert_eq!(constants[&(0, 1)], -1.0); + assert_eq!(constants[&(1, 1)], 1.0); + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins the fold + fn array_factor_scales_each_element() { + // f(y) = a * y with a literal vector a, so df/dy is diag(a). + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let a = arena.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![1.5, -2.5, 4.0], Shape::vector(3)).expect("valid array"), + ))); + let root = arena.alloc(Node::Mul(a, y)); + + let (pattern, varying, entries) = classify_states(&arena, root, 3, 3); + assert!(!varying.iter().any(|&v| v)); + let constants = constant_map(&pattern, &entries); + for (row, expected) in [1.5, -2.5, 4.0].into_iter().enumerate() { + assert_eq!(constants[&(row, row)], expected); + } + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins the fold + fn sparse_matmul_folds_matrix_entries() { + // f(y) = A @ y for a literal tridiagonal A: df/dy is A itself. + let n = 4; + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let (mut indptr, mut indices, mut data) = (vec![0usize], Vec::new(), Vec::new()); + for row in 0..n { + for col in row.saturating_sub(1)..(row + 2).min(n) { + indices.push(col); + data.push((row * n + col) as f64 * 0.25); + } + indptr.push(indices.len()); + } + let mut expected: Vec<(usize, usize, f64)> = Vec::new(); + for row in 0..n { + expected.extend((indptr[row]..indptr[row + 1]).map(|k| (row, indices[k], data[k]))); + } + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(indptr, indices, data, Shape::matrix(n, n)).expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, y)); + + let (pattern, varying, entries) = classify_states(&arena, root, n, n); + assert!(!varying.iter().any(|&v| v)); + let constants = constant_map(&pattern, &entries); + for (row, col, value) in expected { + assert_eq!(constants[&(row, col)], value, "entry ({row}, {col})"); + } + } + + #[test] + fn state_dependent_factor_stays_unproven() { + // f(y) = exp(y) * y: the tangent carries exp(y) as a factor, which is + // seed-free but not literal, so the entry has to be swept. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let factor = arena.alloc(Node::Exp(y)); + let root = arena.alloc(Node::Mul(factor, y)); + + let (_, varying, entries) = classify_states(&arena, root, 2, 2); + assert!(varying.iter().all(|&v| v), "every entry must be swept"); + assert!(entries.is_empty()); + } + + #[test] + fn unmodelled_node_marks_its_row_incomplete() { + // A reduction over the whole state is not modelled element-wise, so its + // row is swept even though the summed rows are linear. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let reduced = arena.alloc(Node::MaxReduce(y)); + let root = arena.alloc(Node::Concat(vec![reduced, y])); + + let (pattern, varying, _) = classify_states(&arena, root, 4, 3); + let row0 = pattern.indptr[0]..pattern.indptr[1]; + assert!( + varying[row0.clone()].iter().all(|&v| v), + "the reduced row must sweep" + ); + assert!( + !varying[row0.end..].iter().any(|&v| v), + "the linear rows must not" + ); + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins the fold + fn structurally_present_but_absent_entries_classify_as_zero() { + // `sign` keeps its argument in the pattern but differentiates to + // nothing, so entry (0, 0) is a constant zero, not an omission. + 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 step = arena.alloc(Node::Sign(y0)); + let root = arena.alloc(Node::Concat(vec![step, y1])); + + let (pattern, varying, entries) = classify_states(&arena, root, 2, 2); + assert_eq!(pattern.indices[pattern.indptr[0]..pattern.indptr[1]], [0]); + assert!(!varying.iter().any(|&v| v)); + let constants = constant_map(&pattern, &entries); + assert_eq!(constants[&(0, 0)], 0.0); + assert_eq!(constants[&(1, 1)], 1.0); + } + + #[test] + fn folded_values_match_the_tape() { + // The soundness property in miniature: every folded entry equals what + // the tangent tape produces for that column, bit for bit. + let n = 3; + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 3, 5], + vec![0, 1, 1, 1, 2], + vec![0.25, -1.75, 3.5, 0.125, -0.5], + Shape::matrix(n, n), + ) + .expect("valid matrix"), + ))); + let linear = arena.alloc(Node::MatMul(matrix, y)); + let nonlinear = arena.alloc(Node::Sin(y)); + let root = arena.alloc(Node::Add(linear, nonlinear)); + + let mut diff_arena = arena.clone(); + let tangent_root = tangent_wrt_states(&mut diff_arena, root); + let (diff_arena, tangent_root) = simplify_pipeline(diff_arena, tangent_root); + let pattern = detect_sparsity_per_output(&arena, root, n, n); + let (_, entries) = classify_constant_entries(&diff_arena, tangent_root, &pattern); + + let expr = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&diff_arena, tangent_root)); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y_values = [0.3, -1.2, 2.4]; + let table = constant_map(&pattern, &entries); + let mut cache = expr.eval_primal(&mut scratch, 0.0, &y_values, &[], &[]); + for col in 0..n { + let mut seed = vec![0.0; n]; + seed[col] = 1.0; + let swept = cache + .eval_tangent(&TangentInputs { + dy: Some(&seed), + dp: None, + }) + .to_vec(); + for (row, &tape) in swept.iter().enumerate() { + if let Some(&value) = table.get(&(row, col)) { + assert_eq!( + value.to_bits(), + tape.to_bits(), + "entry ({row}, {col}) folded to {value}, tape gave {tape}" + ); + } + } + } + } + + #[test] + fn wide_linear_rows_hit_the_degree_cap() { + // Past MAX_CONSTANT_DEGREE columns an element is abandoned rather than + // grown, so the row falls back to being swept. + let n = MAX_CONSTANT_DEGREE + 8; + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .expect("valid row matrix"), + ))); + let root = arena.alloc(Node::MatMul(ones, y)); + + let (_, varying, entries) = classify_states(&arena, root, 1, n); + assert!(varying.iter().all(|&v| v)); + assert!(entries.is_empty()); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/error.rs b/packages/pybamm-rust/pybamm-core/src/error.rs new file mode 100644 index 0000000000..6ed193ebfd --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/error.rs @@ -0,0 +1,92 @@ +//! Crate-wide error type for recoverable failures that should surface as +//! ordinary `Result::Err` values (and, at the Python boundary, as `ValueError`) +//! rather than panics. +//! +//! Panics remain reserved for broken internal invariants; anything driven by +//! caller-supplied arguments or externally-constructed data is reported here. + +/// Errors returned by fallible `pybamm-core` entry points. +#[derive(Debug, thiserror::Error)] +pub enum CoreError { + /// The evaluation time grid was empty; there is nothing to integrate. + #[error("t_eval must not be empty")] + EmptyTimePoints, + + /// The evaluation time grid decreased at some point, or held a NaN. + #[error("t_eval must increase monotonically, but t_eval[{index}] = {got} follows {previous}")] + UnsortedTimePoints { + index: usize, + got: f64, + previous: f64, + }, + + /// The initial state vector length did not match the model's state count. + #[error("y0 has length {got} but the model has {expected} states")] + Y0Length { got: usize, expected: usize }, + + /// The packed input array length did not match the model's parameter width. + #[error("inputs has length {got} but the model expects {expected} packed parameter value(s)")] + InputsLength { got: usize, expected: usize }, + + /// The absolute-tolerance vector length did not match the model's state count. + #[error("atol has length {got} but the model has {expected} states")] + AtolLength { got: usize, expected: usize }, + + /// A sensitivity solve was requested of a model compiled without any. + #[error("no sensitivity parameters were requested when the model was compiled")] + NoSensitivityParams, + + /// An output-variable solve was requested of a model carrying none. + #[error("no output variables were registered when the model was compiled")] + NoOutputVariables, + + /// The `dy0/dp` seed was neither empty nor `n_states x n_sens_params`. + #[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 }, + + /// A batch supplied a different number of initial states and input vectors. + #[error( + "a batch needs one entry per input set, but got {y0} initial state(s) and {inputs} input vector(s)" + )] + BatchWidths { y0: usize, inputs: usize }, + + /// A batch supplied a different number of `dy0/dp` seeds than input sets. + #[error("a batch of {expected} input set(s) needs {expected} dy0/dp seed(s), but got {got}")] + BatchSensWidth { got: usize, expected: usize }, + + /// The sensitivity atol factor was not a finite, strictly positive number. + #[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 finite and > 0, got {got}")] + SolverOption { name: String, 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 }, + + /// A CSR matrix supplied at the boundary violated its structural + /// invariants (indptr length/monotonicity, indices/data lengths, or a + /// column index out of range). + #[error("invalid CSR matrix: {0}")] + Csr(String), + + /// A dense array's data length did not match its declared shape. + #[error("invalid array: {0}")] + Array(String), + + /// An interpolation table violated its invariants (non-empty, matching + /// lengths, strictly increasing finite knots, coefficient counts). + #[error("invalid interpolant: {0}")] + Interpolant(String), + + /// An underlying diffsol integration error. + #[cfg(feature = "diffsol")] + #[error(transparent)] + Diffsol(#[from] diffsol::DiffsolError), +} diff --git a/packages/pybamm-rust/pybamm-core/src/eval.rs b/packages/pybamm-rust/pybamm-core/src/eval.rs new file mode 100644 index 0000000000..9bfcb3e273 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/eval.rs @@ -0,0 +1,1872 @@ +// Intentional u32 usage for compact instruction storage - expression graphs +// won't exceed 4B nodes in practice +#![allow(clippy::cast_possible_truncation)] + +//! Expression evaluator using `TypedIr`. +//! +//! `CompiledExpr` wraps a `TypedIr` and provides efficient evaluation against a +//! caller-supplied scratch buffer. It also supports evaluation with tangent +//! inputs for forward-mode automatic differentiation. + +use crate::arena::{Arena, NodeId}; +use crate::branch_regions::{active_branch, dispatch_span_end}; +use crate::ir::{BinaryOp, BroadcastKind, ConstPool, Instruction, TypedIr, UnaryOp}; + +/// Linear interpolation with binary search. Extends the boundary segment +/// linearly outside the data domain (no flat clamp). +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn interp_linear_1d(x_data: &[f64], y_data: &[f64], x: f64) -> f64 { + let n = x_data.len(); + if n == 1 { + return y_data[0]; + } + // Select the segment [lo, lo+1]: clamp to the first/last segment for + // out-of-domain x so the boundary line is extended. + let lo = if x <= x_data[0] { + 0 + } else if x >= x_data[n - 1] { + n - 2 + } else { + let mut lo = 0; + let mut hi = n - 1; + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if x_data[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + lo + }; + let t = (x - x_data[lo]) / (x_data[lo + 1] - x_data[lo]); + t.mul_add(y_data[lo + 1] - y_data[lo], y_data[lo]) +} + +/// Pre-computed slope lookup for linear-interpolation derivative (extends +/// boundary segment outside data domain, matching the value function). +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn interp_linear_1d_slope_lookup(x_data: &[f64], slopes: &[f64], x: f64) -> f64 { + let n = x_data.len(); + if n < 2 || slopes.is_empty() { + return 0.0; + } + if x <= x_data[0] { + return slopes[0]; + } + if x >= x_data[n - 1] { + return slopes[n - 2]; + } + let mut lo = 0; + let mut hi = n - 1; + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if x_data[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + slopes[lo] +} + +/// Local slope of a piecewise-linear interpolant at `x`, computed from the +/// primal knot values. Matches `interp_linear_1d_slope_lookup`'s segment +/// choice and `compute_interpolant_slopes`' value, so the reverse-AD adjoint +/// equals the forward derivative path exactly. +pub(crate) fn interp_linear_1d_deriv(x_data: &[f64], y_data: &[f64], x: f64) -> f64 { + let n = x_data.len(); + if n < 2 { + return 0.0; + } + let seg = if x <= x_data[0] { + 0 + } else if x >= x_data[n - 1] { + n - 2 + } else { + let mut lo = 0; + let mut hi = n - 1; + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if x_data[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + lo + }; + let dx = x_data[seg + 1] - x_data[seg]; + if dx.abs() > f64::EPSILON { + (y_data[seg + 1] - y_data[seg]) / dx + } else { + 0.0 + } +} + +/// Interval index for `x`, clamped to `[0, nseg-1]` (extends boundary polynomial). +/// `breakpoints.len() >= 2` is guaranteed by Python lowering. +#[allow(clippy::inline_always)] +#[inline(always)] +fn locate_cubic_interval(breakpoints: &[f64], x: f64) -> usize { + debug_assert!(breakpoints.len() >= 2); + let nseg = breakpoints.len() - 1; + if x <= breakpoints[0] { + return 0; + } + if x >= breakpoints[nseg] { + return nseg - 1; + } + let mut lo = 0; + let mut hi = nseg; + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if breakpoints[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + lo +} + +/// Piecewise-cubic interpolation (cubic spline / pchip), power-basis Horner. +#[allow(clippy::inline_always, clippy::suboptimal_flops)] +#[inline(always)] +pub(crate) fn interp_cubic_1d(breakpoints: &[f64], coeffs: &[[f64; 4]], x: f64) -> f64 { + let i = locate_cubic_interval(breakpoints, x); + let dx = x - breakpoints[i]; + let [c0, c1, c2, c3] = coeffs[i]; + c0 + (c1 + (c2 + c3 * dx) * dx) * dx +} + +/// Derivative of the piecewise-cubic interpolant: `p'(dx) = c1 + 2*c2*dx + 3*c3*dx^2`. +#[allow(clippy::inline_always, clippy::suboptimal_flops)] +#[inline(always)] +pub(crate) fn interp_cubic_1d_deriv(breakpoints: &[f64], coeffs: &[[f64; 4]], x: f64) -> f64 { + let i = locate_cubic_interval(breakpoints, x); + let dx = x - breakpoints[i]; + let [_c0, c1, c2, c3] = coeffs[i]; + c1 + (2.0 * c2 + 3.0 * c3 * dx) * dx +} + +/// Locate the N-D cell for `coords`, writing per-axis offsets into `dxs` and +/// returning the cell's `order^ndim` power-basis coefficients (clamped per +/// axis, matching scipy `RegularGridInterpolator` with `fill_value=None`). +pub(crate) fn locate_nd_cell<'a>( + breakpoints: &[Vec], + coeffs: &'a [f64], + order: usize, + coords: &[f64], + dxs: &mut [f64; 3], +) -> &'a [f64] { + let ndim = breakpoints.len(); + debug_assert!((2..=3).contains(&ndim) && coords.len() == ndim); + let mut cell = 0; + for a in 0..ndim { + let knots = &breakpoints[a]; + let i = locate_cubic_interval(knots, coords[a]); + dxs[a] = coords[a] - knots[i]; + cell = cell * (knots.len() - 1) + i; + } + let csize = order.pow(ndim as u32); + &coeffs[cell * csize..(cell + 1) * csize] +} + +/// Evaluate an N-D tensor-product polynomial in nested Horner form. +/// `coeffs` has `order^(dxs.len())` entries, axis-0 power slowest, ascending. +#[allow(clippy::suboptimal_flops)] +pub(crate) fn tensor_horner(coeffs: &[f64], dxs: &[f64], order: usize) -> f64 { + debug_assert!(order >= 2); + if dxs.len() == 1 { + let mut acc = coeffs[order - 1]; + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + coeffs[a]; + } + return acc; + } + let stride = coeffs.len() / order; + let mut acc = tensor_horner(&coeffs[(order - 1) * stride..], &dxs[1..], order); + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + tensor_horner(&coeffs[a * stride..(a + 1) * stride], &dxs[1..], order); + } + acc +} + +/// Partial derivative of the N-D tensor-product polynomial along `axis`. +#[allow(clippy::suboptimal_flops)] +pub(crate) fn tensor_horner_partial(coeffs: &[f64], dxs: &[f64], order: usize, axis: usize) -> f64 { + debug_assert!(order >= 2); + debug_assert!(axis < dxs.len()); + if axis == 0 { + if dxs.len() == 1 { + let mut acc = (order - 1) as f64 * coeffs[order - 1]; + for a in (1..order - 1).rev() { + acc = acc * dxs[0] + a as f64 * coeffs[a]; + } + return acc; + } + let stride = coeffs.len() / order; + let mut acc = + (order - 1) as f64 * tensor_horner(&coeffs[(order - 1) * stride..], &dxs[1..], order); + for a in (1..order - 1).rev() { + acc = acc * dxs[0] + + a as f64 * tensor_horner(&coeffs[a * stride..(a + 1) * stride], &dxs[1..], order); + } + return acc; + } + let stride = coeffs.len() / order; + let mut acc = + tensor_horner_partial(&coeffs[(order - 1) * stride..], &dxs[1..], order, axis - 1); + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + + tensor_horner_partial( + &coeffs[a * stride..(a + 1) * stride], + &dxs[1..], + order, + axis - 1, + ); + } + acc +} + +/// Approximation of the error function +/// Tolerance `BinaryOp::Equality` compares within, shared with the batched +/// tangent sweep so both paths agree. +pub(crate) const EQUALITY_EPS: f64 = 1e-14; + +/// +/// Shared with constant folding in `simplify` so folded values match +/// runtime evaluation exactly. +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn erf_approx(x: f64) -> f64 { + if x == 0.0 { + return 0.0; + } + let sign = x.signum(); + let x = x.abs(); + let t = 1.0 / 0.327_591_1f64.mul_add(x, 1.0); + let poly = t * t.mul_add( + t.mul_add( + t.mul_add(t.mul_add(1.061_405_429, -1.453_152_027), 1.421_413_741), + -0.284_496_736, + ), + 0.254_829_592, + ); + sign * poly.mul_add(-(-x * x).exp(), 1.0) +} + +/// Sign function with `sign(0) = 0` (unlike `f64::signum`). +/// +/// Shared with constant folding in `simplify` so folded values match +/// runtime evaluation exactly. +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn sign(x: f64) -> f64 { + if x == 0.0 { 0.0 } else { x.signum() } +} + +/// Carve a read window `[src, src+len)` and a write window `[dst, dst+len)` +/// out of `buf` as disjoint slices, which the IR allocator guarantees. +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn split_src_dst( + buf: &mut [f64], + src: usize, + dst: usize, + len: usize, +) -> (&[f64], &mut [f64]) { + debug_assert!( + src + len <= dst || dst + len <= src, + "src/dst windows overlap (src={src}, dst={dst}, len={len})" + ); + if dst < src { + let (left, right) = buf.split_at_mut(src); + (&right[..len], &mut left[dst..dst + len]) + } else { + let (left, right) = buf.split_at_mut(dst); + (&left[src..src + len], &mut right[..len]) + } +} + +/// Carve two read windows `a` and `b` and a write window `dst`, all length +/// `len`, out of `buf`. `dst` must be disjoint from both `a` and `b` (the IR +/// allocator guarantees this); `a` and `b` may overlap each other (read-only). +#[allow(clippy::inline_always)] +#[inline(always)] +pub(crate) fn split_dst_two_src( + buf: &mut [f64], + a: usize, + b: usize, + dst: usize, + len: usize, +) -> (&[f64], &[f64], &mut [f64]) { + debug_assert!( + a + len <= dst || dst + len <= a, + "operand a overlaps dst (a={a}, dst={dst}, len={len})" + ); + debug_assert!( + b + len <= dst || dst + len <= b, + "operand b overlaps dst (b={b}, dst={dst}, len={len})" + ); + let (left, rest) = buf.split_at_mut(dst); + let (dst_s, right) = rest.split_at_mut(len); + let a_s = if a + len <= dst { + &left[a..a + len] + } else { + let off = a - (dst + len); + &right[off..off + len] + }; + let b_s = if b + len <= dst { + &left[b..b + len] + } else { + let off = b - (dst + len); + &right[off..off + len] + }; + (a_s, b_s, dst_s) +} + +/// Apply a binary operation with broadcasting. +#[allow(clippy::inline_always)] +#[inline(always)] +fn broadcast_apply f64>( + buf: &mut [f64], + f: F, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, +) { + match kind { + BroadcastKind::ScalarScalar => { + buf[dst] = f(buf[a], buf[b]); + }, + BroadcastKind::ScalarVector => { + let scalar = buf[a]; + let (b_s, d_s) = split_src_dst(buf, b, dst, len); + for (o, &y) in d_s.iter_mut().zip(b_s) { + *o = f(scalar, y); + } + }, + BroadcastKind::VectorScalar => { + let scalar = buf[b]; + let (a_s, d_s) = split_src_dst(buf, a, dst, len); + for (o, &x) in d_s.iter_mut().zip(a_s) { + *o = f(x, scalar); + } + }, + BroadcastKind::VectorVector => { + let (a_s, b_s, d_s) = split_dst_two_src(buf, a, b, dst, len); + for ((o, &x), &y) in d_s.iter_mut().zip(a_s).zip(b_s) { + *o = f(x, y); + } + }, + } +} + +/// Evaluate a binary operation. +#[allow(clippy::inline_always)] +#[inline(always)] +fn eval_binary_op( + buf: &mut [f64], + op: BinaryOp, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, +) { + match op { + BinaryOp::Add => broadcast_apply(buf, |x, y| x + y, a, b, dst, len, kind), + BinaryOp::Sub => broadcast_apply(buf, |x, y| x - y, a, b, dst, len, kind), + BinaryOp::Mul => broadcast_apply(buf, |x, y| x * y, a, b, dst, len, kind), + BinaryOp::Div => broadcast_apply(buf, |x, y| x / y, a, b, dst, len, kind), + BinaryOp::Pow => broadcast_apply(buf, f64::powf, a, b, dst, len, kind), + BinaryOp::Minimum => broadcast_apply(buf, f64::min, a, b, dst, len, kind), + BinaryOp::Maximum => broadcast_apply(buf, f64::max, a, b, dst, len, kind), + BinaryOp::Modulo => broadcast_apply(buf, |x, y| x % y, a, b, dst, len, kind), + BinaryOp::Hypot => broadcast_apply(buf, f64::hypot, a, b, dst, len, kind), + BinaryOp::EqualHeaviside => broadcast_apply( + buf, + |x, y| if x <= y { 1.0 } else { 0.0 }, + a, + b, + dst, + len, + kind, + ), + BinaryOp::NotEqualHeaviside => broadcast_apply( + buf, + |x, y| if x < y { 1.0 } else { 0.0 }, + a, + b, + dst, + len, + kind, + ), + BinaryOp::Equality => { + broadcast_apply( + buf, + |x, y| { + if (x - y).abs() < EQUALITY_EPS { + 1.0 + } else { + 0.0 + } + }, + a, + b, + dst, + len, + kind, + ); + }, + } +} + +/// Apply a unary function. +#[allow(clippy::inline_always)] +#[inline(always)] +fn unary_apply f64>(buf: &mut [f64], f: F, src: usize, dst: usize, len: usize) { + if src == dst { + for x in &mut buf[dst..dst + len] { + *x = f(*x); + } + return; + } + let (src_s, dst_s) = split_src_dst(buf, src, dst, len); + for (o, &x) in dst_s.iter_mut().zip(src_s) { + *o = f(x); + } +} + +/// Evaluate a unary operation. +#[allow(clippy::inline_always)] +#[inline(always)] +fn eval_unary_op(buf: &mut [f64], op: UnaryOp, src: usize, dst: usize, len: usize) { + match op { + UnaryOp::Neg => unary_apply(buf, |x| -x, src, dst, len), + UnaryOp::Abs => unary_apply(buf, f64::abs, src, dst, len), + UnaryOp::Sqrt => unary_apply(buf, f64::sqrt, src, dst, len), + UnaryOp::Exp => unary_apply(buf, f64::exp, src, dst, len), + UnaryOp::Log => unary_apply(buf, f64::ln, src, dst, len), + UnaryOp::Sin => unary_apply(buf, f64::sin, src, dst, len), + UnaryOp::Cos => unary_apply(buf, f64::cos, src, dst, len), + UnaryOp::Tanh => unary_apply(buf, f64::tanh, src, dst, len), + UnaryOp::Sinh => unary_apply(buf, f64::sinh, src, dst, len), + UnaryOp::Cosh => unary_apply(buf, f64::cosh, src, dst, len), + UnaryOp::Arcsinh => unary_apply(buf, f64::asinh, src, dst, len), + UnaryOp::Arctan => unary_apply(buf, f64::atan, src, dst, len), + UnaryOp::Erf => unary_apply(buf, erf_approx, src, dst, len), + UnaryOp::Sign => unary_apply(buf, sign, src, dst, len), + UnaryOp::Floor => unary_apply(buf, f64::floor, src, dst, len), + UnaryOp::Ceiling => unary_apply(buf, f64::ceil, src, dst, len), + } +} + +/// Seed vectors for tangent evaluation (forward-mode AD). +/// +/// A seed is the direction the JVP is taken in, so one entry per state or per +/// parameter; `None` leaves that family of tangent reads at zero. Colored Jacobian +/// assembly seeds `dy` with a 1 in every column of the current color. +#[derive(Debug, Default)] +pub struct TangentInputs<'a> { + /// Tangent of the state vector, indexed like `y`. + pub dy: Option<&'a [f64]>, + /// Tangent of the parameters, one entry per parameter regardless of width: + /// indexed by `InputParameter::index`, *not* by offset into `inputs`. + pub dp: Option<&'a [f64]>, +} + +#[derive(Debug, Clone, Copy)] +struct PrimalEvalInputs<'a> { + t: f64, + y: &'a [f64], + y_dot: &'a [f64], + inputs: &'a [f64], +} + +#[derive(Debug, Clone, Copy)] +struct EvalContext<'a> { + primal: Option>, + tangent: Option<&'a TangentInputs<'a>>, +} + +impl<'a> EvalContext<'a> { + const fn with_primal( + t: f64, + y: &'a [f64], + y_dot: &'a [f64], + inputs: &'a [f64], + tangent: Option<&'a TangentInputs<'a>>, + ) -> Self { + Self { + primal: Some(PrimalEvalInputs { + t, + y, + y_dot, + inputs, + }), + tangent, + } + } + + const fn tangent_only(tangent: &'a TangentInputs<'a>) -> Self { + Self { + primal: None, + tangent: Some(tangent), + } + } + + const fn primal(self) -> PrimalEvalInputs<'a> { + self.primal + .expect("primal instructions require full evaluation inputs") + } + + fn tangent_state(self) -> Option<&'a [f64]> { + self.tangent.and_then(|tangent| tangent.dy) + } + + fn tangent_parameter(self, index: usize) -> f64 { + self.tangent + .and_then(|tangent| tangent.dp) + .map_or(0.0, |dp| dp[index]) + } +} + +/// Execute a slice of instructions against a shared buffer, returning the number +/// of instructions actually executed. +/// +/// `Dispatch` advances past the blocks of every inactive branch, so the count is +/// the honest cost of this evaluation rather than the tape length. The +/// accumulator costs one register; callers that ignore it pay nothing. +#[inline] +fn eval_instructions( + buf: &mut [f64], + instructions: &[Instruction], + consts: &ConstPool, + ctx: EvalContext<'_>, +) -> usize { + let mut executed = 0_usize; + // Named `pc`, not `i`: many arms bind their own loop-local `i`. + let mut pc = 0_usize; + while pc < instructions.len() { + if let Instruction::Dispatch { + selector, + blocks_idx, + blocks_len, + } = instructions[pc] + { + let base = blocks_idx as usize; + let n = blocks_len as usize; + let span_end = dispatch_span_end(consts, pc, blocks_idx, blocks_len); + executed += 1; + if let Some(active) = active_branch(buf[selector as usize], n) { + let (rel, len) = consts.branch_blocks[base + active]; + let start = pc + rel as usize; + executed += + eval_instructions(buf, &instructions[start..start + len as usize], consts, ctx); + } + pc = span_end; + continue; + } + + match instructions[pc] { + Instruction::LoadScalar { value, dst } => { + buf[dst as usize] = value; + }, + Instruction::LoadTime { dst } => { + buf[dst as usize] = ctx.primal().t; + }, + Instruction::LoadArray { data_idx, len, dst } => { + let src = consts.get_array(data_idx, len); + buf[dst as usize..dst as usize + len as usize].copy_from_slice(src); + }, + Instruction::FillZero { dst, len } => { + buf[dst as usize..dst as usize + len as usize].fill(0.0); + }, + Instruction::LoadStateVector { start, end, dst } => { + let primal = ctx.primal(); + buf[dst as usize..dst as usize + (end - start) as usize] + .copy_from_slice(&primal.y[start as usize..end as usize]); + }, + Instruction::LoadStateVectorDot { start, end, dst } => { + let primal = ctx.primal(); + buf[dst as usize..dst as usize + (end - start) as usize] + .copy_from_slice(&primal.y_dot[start as usize..end as usize]); + }, + Instruction::LoadInputParameter { offset, width, dst } => { + let primal = ctx.primal(); + let offset = offset as usize; + let width = width as usize; + buf[dst as usize..dst as usize + width] + .copy_from_slice(&primal.inputs[offset..offset + width]); + }, + Instruction::LoadTangentState { start, end, dst } => { + if let Some(dy) = ctx.tangent_state() { + buf[dst as usize..dst as usize + (end - start) as usize] + .copy_from_slice(&dy[start as usize..end as usize]); + } else { + buf[dst as usize..dst as usize + (end - start) as usize].fill(0.0); + } + }, + Instruction::LoadTangentParameter { index, dst } => { + buf[dst as usize] = ctx.tangent_parameter(index as usize); + }, + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => { + eval_binary_op( + buf, + op, + a as usize, + b as usize, + dst as usize, + len as usize, + kind, + ); + }, + + Instruction::Unary { op, src, dst, len } => { + eval_unary_op(buf, op, src as usize, dst as usize, len as usize); + }, + + Instruction::MaxReduce { src, src_len, dst } => { + let src = src as usize; + let src_len = src_len as usize; + let mut max_val = f64::NEG_INFINITY; + for i in 0..src_len { + let v = buf[src + i]; + if v > max_val { + max_val = v; + } + } + buf[dst as usize] = max_val; + }, + + Instruction::MinReduce { src, src_len, dst } => { + let src = src as usize; + let src_len = src_len as usize; + let mut min_val = f64::INFINITY; + for i in 0..src_len { + let v = buf[src + i]; + if v < min_val { + min_val = v; + } + } + buf[dst as usize] = min_val; + }, + + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + is_max, + dst, + } => { + let basis = basis_src as usize; + let picker = picker_src as usize; + let len = len as usize; + // Argmax/argmin takes the earliest element under a strict + // comparison, matching the primal MaxReduce/MinReduce eval. + let mut best_idx = 0; + let mut best_val = buf[picker]; + for i in 1..len { + let v = buf[picker + i]; + if (is_max && v > best_val) || (!is_max && v < best_val) { + best_val = v; + best_idx = i; + } + } + buf[dst as usize] = buf[basis + best_idx]; + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + buf.copy_within( + src as usize + start as usize..src as usize + start as usize + len as usize, + dst as usize, + ); + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write_pos = dst as usize; + for i in 0..sources_len as usize { + let (src_off, src_len) = consts.concat_sources[sources_idx as usize + i]; + buf.copy_within( + src_off as usize..src_off as usize + src_len as usize, + write_pos, + ); + write_pos += src_len as usize; + } + }, + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx as usize]; + let vec_src = vec_src as usize; + let dst = dst as usize; + for row in 0..csr.shape.rows { + let start = csr.indptr[row]; + let end = csr.indptr[row + 1]; + let cols = &csr.indices[start..end]; + let vals = &csr.data[start..end]; + let mut sum = 0.0; + // The data-dependent gather `buf[vec_src + col]` stays + // bounds-checked; the row's value/index reads do not. + for (&col, &val) in cols.iter().zip(vals) { + sum += val * buf[vec_src + col]; + } + buf[dst + row] = sum; + } + }, + + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + dst, + } => { + let (mat_src, vec_src, dst) = (mat_src as usize, vec_src as usize, dst as usize); + for row in 0..rows as usize { + let mut sum = 0.0; + for col in 0..cols as usize { + sum += buf[mat_src + row * cols as usize + col] * buf[vec_src + col]; + } + buf[dst + row] = sum; + } + }, + + Instruction::Interp1DLinear { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + buf[dst + i] = interp_linear_1d(&interp.x_data, &interp.y_data, buf[src + i]); + } + }, + + Instruction::Interp1DLinearDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + buf[dst + i] = + interp_linear_1d_slope_lookup(&interp.x_data, &interp.y_data, buf[src + i]); + } + }, + + Instruction::Interp1DCubic { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + buf[dst + i] = + interp_cubic_1d(&interp.breakpoints, &interp.coeffs, buf[src + i]); + } + }, + Instruction::Interp1DCubicDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx as usize]; + let src = src as usize; + let dst = dst as usize; + for i in 0..len as usize { + buf[dst + i] = + interp_cubic_1d_deriv(&interp.breakpoints, &interp.coeffs, buf[src + i]); + } + }, + + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len, + } => { + let interp = &consts.nd_interpolants[interp_idx as usize]; + let ndim = interp.breakpoints.len(); + let order = interp.order as usize; + let dst = dst as usize; + let mut coords = [0.0_f64; 3]; + let mut dxs = [0.0_f64; 3]; + for i in 0..len as usize { + for (a, coord) in coords.iter_mut().enumerate().take(ndim) { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + // Length-1 children broadcast over the output length. + let j = if slen == 1 { 0 } else { i }; + *coord = buf[off as usize + j]; + } + let cell = locate_nd_cell( + &interp.breakpoints, + &interp.coeffs, + order, + &coords[..ndim], + &mut dxs, + ); + buf[dst + i] = tensor_horner(cell, &dxs[..ndim], order); + } + }, + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis, + dst, + len, + } => { + let interp = &consts.nd_interpolants[interp_idx as usize]; + let ndim = interp.breakpoints.len(); + let order = interp.order as usize; + let dst = dst as usize; + let mut coords = [0.0_f64; 3]; + let mut dxs = [0.0_f64; 3]; + for i in 0..len as usize { + for (a, coord) in coords.iter_mut().enumerate().take(ndim) { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + let j = if slen == 1 { 0 } else { i }; + *coord = buf[off as usize + j]; + } + let cell = locate_nd_cell( + &interp.breakpoints, + &interp.coeffs, + order, + &coords[..ndim], + &mut dxs, + ); + buf[dst + i] = tensor_horner_partial(cell, &dxs[..ndim], order, axis as usize); + } + }, + + Instruction::Conditional { + selector, + branches_idx, + branches_len, + dst, + out_len, + } => { + let dst = dst as usize; + let out_len = out_len as usize; + match active_branch(buf[selector as usize], branches_len as usize) { + Some(i) => { + let (branch_off, _) = consts.branch_offsets[branches_idx as usize + i]; + buf.copy_within(branch_off as usize..branch_off as usize + out_len, dst); + }, + None => buf[dst..dst + out_len].fill(0.0), + } + }, + Instruction::Dispatch { .. } => unreachable!("handled above"), + } + executed += 1; + pc += 1; + } + executed +} + +/// Zero-allocation expression evaluator. +/// +/// Wraps a `TypedIr` and provides efficient evaluation. Callers must supply +/// an external scratch buffer of length `scratch_len()` so that the same +/// `CompiledExpr` can be shared across concurrent or re-entrant solves without +/// interior-mutable state. +#[derive(Debug, Clone)] +pub struct CompiledExpr { + ir: TypedIr, +} + +impl CompiledExpr { + /// Compile an expression DAG into a new `CompiledExpr`. + pub fn new(arena: &Arena, root: NodeId) -> Self { + Self::from_ir(TypedIr::from_arena(arena, root)) + } + + /// Compile with a no-reuse (SSA) slot layout for reverse-mode AD. After + /// [`eval`](Self::eval) the scratch holds every intermediate value, so the + /// reverse backward pass reads operands by their stable slots. + pub fn new_pinned(arena: &Arena, root: NodeId) -> Self { + Self::from_ir(TypedIr::from_arena_pinned(arena, root)) + } + + /// Wrap a pre-built `TypedIr`. + pub const fn from_ir(ir: TypedIr) -> Self { + Self { ir } + } + + /// Reference to the underlying `TypedIr`. + #[inline] + pub const fn ir(&self) -> &TypedIr { + &self.ir + } + + /// Length of the scratch buffer this expression needs for evaluation. + #[inline] + pub const fn scratch_len(&self) -> usize { + self.ir.buffer_size() + } + + /// Output length of the root expression. + #[inline] + pub const fn output_len(&self) -> usize { + self.ir.output_len() + } + + /// Evaluate the expression into `scratch` (length `>= scratch_len()`), + /// returning the slice holding the root result. + pub fn eval<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> &'s [f64] { + self.eval_internal(scratch, t, y, y_dot, inputs, None).1 + } + + /// [`eval`](Self::eval), also returning how many instructions ran. + /// + /// The direct test of only-active-branch execution: a reported tape length + /// cannot distinguish work that was skipped from work that was not counted. + pub fn eval_counted<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> (usize, &'s [f64]) { + self.eval_internal(scratch, t, y, y_dot, inputs, None) + } + + /// Evaluate with tangent inputs for forward-mode AD (JVP). + pub fn eval_with_tangent<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + ) -> &'s [f64] { + self.eval_internal(scratch, t, y, y_dot, inputs, Some(tangent)) + .1 + } + + /// [`eval_with_tangent`](Self::eval_with_tangent), also returning how many + /// instructions ran. + /// + /// The direct test of only-active-branch execution on a tape that carries + /// tangent work, where the split layout puts one branch in two blocks. + pub fn eval_counted_with_tangent<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + ) -> (usize, &'s [f64]) { + self.eval_internal(scratch, t, y, y_dot, inputs, Some(tangent)) + } + + fn eval_internal<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: Option<&TangentInputs<'_>>, + ) -> (usize, &'s [f64]) { + let executed = eval_instructions( + scratch, + self.ir.instructions(), + self.ir.consts(), + EvalContext::with_primal(t, y, y_dot, inputs, tangent), + ); + + let root = self.ir.root_slot(); + ( + executed, + &scratch[root.offset_usize()..root.offset_usize() + root.len_usize()], + ) + } + + /// Evaluate the primal section of a split-eval expression into `scratch`, + /// returning a [`PrimalCache`] that owns the buffer. The tangent sweep runs + /// through the returned cache, which encodes the "primal first, same + /// scratch" contract that a bare tangent call cannot express. + pub fn eval_primal<'a>( + &'a self, + scratch: &'a mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> PrimalCache<'a> { + self.run_primal_section(scratch, t, y, y_dot, inputs); + PrimalCache { + expr: self, + scratch, + } + } + + /// Raw primal-section evaluation into `scratch` (no cache). Primitive for + /// the sensitivity and batched-tangent paths, whose primal pass and tangent + /// sweeps write separate buffers and so cannot share a borrow-based cache; + /// prefer [`eval_primal`](Self::eval_primal) elsewhere. + pub fn run_primal_section( + &self, + scratch: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) { + let primal_end = self + .ir + .split_eval_info() + .map_or_else(|| self.ir.instructions().len(), |s| s.primal_end); + + eval_instructions( + scratch, + &self.ir.instructions()[..primal_end], + self.ir.consts(), + EvalContext::with_primal(t, y, y_dot, inputs, None), + ); + } + + /// Raw tangent-section evaluation, reusing the primal region a prior + /// [`run_primal_section`](Self::run_primal_section) left in `scratch`. + /// Internal primitive for the sensitivity path; prefer + /// [`PrimalCache::eval_tangent`] elsewhere. + pub(crate) fn run_tangent_section<'s>( + &self, + scratch: &'s mut [f64], + tangent: &TangentInputs<'_>, + ) -> &'s [f64] { + debug_assert!( + self.ir.split_eval_info().is_some(), + "run_tangent_section requires a split-eval IR" + ); + let primal_end = self + .ir + .split_eval_info() + .map_or_else(|| self.ir.instructions().len(), |s| s.primal_end); + + eval_instructions( + scratch, + &self.ir.instructions()[primal_end..], + self.ir.consts(), + EvalContext::tangent_only(tangent), + ); + + let root = self.ir.root_slot(); + &scratch[root.offset_usize()..root.offset_usize() + root.len_usize()] + } + + /// True when this expression was compiled with `from_arena_split_eval`, + /// i.e. split primal/tangent evaluation via [`eval_primal`](Self::eval_primal) + /// is available. + #[inline] + pub const fn has_split_eval(&self) -> bool { + self.ir.split_eval_info().is_some() + } + + /// Evaluate and copy the result into `out`, using `scratch` for working space. + #[inline] + pub fn eval_into( + &self, + scratch: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + let result = self.eval(scratch, t, y, y_dot, inputs); + out[..result.len()].copy_from_slice(result); + } +} + +/// Borrow-based typestate for split evaluation. +/// +/// Produced only by [`CompiledExpr::eval_primal`], it owns the scratch buffer +/// whose primal region was just filled. Because a cache is the only way to +/// reach the tangent sweep, calling tangent before primal or on a different +/// buffer is a compile error rather than a silent stale-buffer read. +pub struct PrimalCache<'a> { + expr: &'a CompiledExpr, + scratch: &'a mut [f64], +} + +impl std::fmt::Debug for PrimalCache<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrimalCache").finish_non_exhaustive() + } +} + +impl PrimalCache<'_> { + /// Evaluate the tangent section for `tangent`, reusing the primal region + /// filled by [`CompiledExpr::eval_primal`], and return the root slice. + pub fn eval_tangent(&mut self, tangent: &TangentInputs<'_>) -> &[f64] { + self.expr.run_tangent_section(self.scratch, tangent) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::{ + ArrayData, CsrData, CubicInterpolantData, InterpolantData, NdInterpolantData, Node, Shape, + }; + + #[test] + fn external_scratch_eval_matches_len() { + // dy/dt expression: 2.0 * y[0] + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr_node = arena.alloc(Node::Mul(two, sv)); + let expr = CompiledExpr::new(&arena, expr_node); + + let mut scratch = vec![0.0; expr.scratch_len()]; + let result = expr.eval(&mut scratch, 0.0, &[3.0], &[], &[]); + assert_eq!(result, &[6.0]); + } + + #[test] + fn size_check() { + use std::mem::size_of; + // These live in hot per-node arrays, so guard against accidental bloat; + // bump a bound deliberately if a type must grow. + assert!( + size_of::() <= 32, + "Instruction grew to {} bytes", + size_of::() + ); + assert!( + size_of::() <= 48, + "Node grew to {} bytes", + size_of::() + ); + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 1); + assert_eq!(size_of::(), 1); + assert_eq!(size_of::(), 1); + } + + #[test] + fn test_eval_scalar() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Scalar(42.0)); + let compiled = CompiledExpr::new(&arena, id); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[42.0]); + } + + #[test] + fn test_eval_time() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Time); + let compiled = CompiledExpr::new(&arena, id); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 3.5, &[], &[], &[]); + assert_eq!(result, &[3.5]); + } + + #[test] + fn test_eval_state_vector() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::StateVector { start: 1, end: 3 }); + let compiled = CompiledExpr::new(&arena, id); + let y = [10.0, 20.0, 30.0, 40.0]; + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &y, &[], &[]); + assert_eq!(result, &[20.0, 30.0]); + } + + #[test] + fn test_eval_array() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: Shape::vector(3), + }))); + let compiled = CompiledExpr::new(&arena, id); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[1.0, 2.0, 3.0]); + } + + #[test] + fn test_eval_add_scalars() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(2.0)); + let b = arena.alloc(Node::Scalar(3.0)); + let sum = arena.alloc(Node::Add(a, b)); + let compiled = CompiledExpr::new(&arena, sum); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[5.0]); + } + + #[test] + fn test_eval_add_vectors() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: Shape::vector(3), + }))); + let b = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0], + shape: Shape::vector(3), + }))); + let sum = arena.alloc(Node::Add(a, b)); + let compiled = CompiledExpr::new(&arena, sum); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[11.0, 22.0, 33.0]); + } + + #[test] + fn test_eval_scalar_times_vector() { + let mut arena = Arena::new(); + let s_node = arena.alloc(Node::Scalar(2.0)); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: Shape::vector(3), + }))); + let prod = arena.alloc(Node::Mul(s_node, v)); + let compiled = CompiledExpr::new(&arena, prod); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[2.0, 4.0, 6.0]); + } + + #[test] + fn test_eval_nested_arithmetic() { + let mut arena = Arena::new(); + let two = arena.alloc(Node::Scalar(2.0)); + let three = arena.alloc(Node::Scalar(3.0)); + let four = arena.alloc(Node::Scalar(4.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let sum = arena.alloc(Node::Add(two, three)); + let prod = arena.alloc(Node::Mul(sum, four)); + let result_node = arena.alloc(Node::Sub(prod, one)); + let compiled = CompiledExpr::new(&arena, result_node); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[19.0]); + } + + #[test] + fn test_eval_neg() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(5.0)); + let neg = arena.alloc(Node::Neg(a)); + let compiled = CompiledExpr::new(&arena, neg); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[-5.0]); + } + + #[test] + fn test_eval_sqrt_vector() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 4.0, 9.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::Sqrt(a)); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[1.0, 2.0, 3.0]); + } + + #[test] + fn test_eval_sin_cos() { + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let sin = arena.alloc(Node::Sin(zero)); + let cos = arena.alloc(Node::Cos(zero)); + let sum = arena.alloc(Node::Add(sin, cos)); + let compiled = CompiledExpr::new(&arena, sum); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[1.0]); + } + + #[test] + fn test_eval_max_reduce() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 5.0, 3.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::MaxReduce(a)); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[5.0]); + } + + #[test] + fn test_eval_reduce_arg_select_max() { + // basis[argmax(picker)]; picker=[1,5,3] -> k=1 -> basis[1]=20 + let mut arena = Arena::new(); + let picker = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 5.0, 3.0], + shape: Shape::vector(3), + }))); + let basis = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[20.0]); + } + + #[test] + fn test_eval_reduce_arg_select_min() { + // picker=[1,5,3] -> argmin k=0 -> basis[0]=10 + let mut arena = Arena::new(); + let picker = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 5.0, 3.0], + shape: Shape::vector(3), + }))); + let basis = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: false, + }); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[10.0]); + } + + #[test] + fn test_eval_reduce_arg_select_tie_first_wins() { + // picker=[5,5,3] max ties at 0 and 1 -> first index 0 -> basis[0]=10 + let mut arena = Arena::new(); + let picker = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![5.0, 5.0, 3.0], + shape: Shape::vector(3), + }))); + let basis = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[10.0]); + } + + #[test] + fn test_eval_reduce_arg_select_min_tie_first_wins() { + // picker=[5,3,3] min ties at 1 and 2 -> first index 1 -> basis[1]=20 + let mut arena = Arena::new(); + let picker = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![5.0, 3.0, 3.0], + shape: Shape::vector(3), + }))); + let basis = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0], + shape: Shape::vector(3), + }))); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: false, + }); + let compiled = CompiledExpr::new(&arena, node); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[20.0]); + } + + #[test] + fn test_eval_index() { + let mut arena = Arena::new(); + let arr = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![10.0, 20.0, 30.0, 40.0, 50.0], + shape: Shape::vector(5), + }))); + let idx = arena.alloc(Node::Index { + child: arr, + start: 1, + end: 4, + }); + let compiled = CompiledExpr::new(&arena, idx); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[20.0, 30.0, 40.0]); + } + + #[test] + fn test_eval_concat() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![2.0, 3.0], + shape: Shape::vector(2), + }))); + let c = arena.alloc(Node::Scalar(4.0)); + let concat = arena.alloc(Node::Concat(vec![a, b, c])); + let compiled = CompiledExpr::new(&arena, concat); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[1.0, 2.0, 3.0, 4.0]); + } + + #[test] + fn test_eval_matmul() { + let mut arena = Arena::new(); + let sparse = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: Shape::vector(3), + }))); + let matmul = arena.alloc(Node::MatMul(sparse, v)); + let compiled = CompiledExpr::new(&arena, matmul); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[2.0, 6.0]); + } + + #[test] + fn test_eval_dense_matmul() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // row-major 2x3 + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: Shape::vector(3), + }))); + let matmul = arena.alloc(Node::MatMul(a, v)); + let compiled = CompiledExpr::new(&arena, matmul); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[14.0, 32.0]); + } + + #[test] + fn test_eval_interpolant() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::Scalar(1.5)); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(InterpolantData { + x_data: vec![0.0, 1.0, 2.0], + y_data: vec![0.0, 10.0, 20.0], + }), + child: x, + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[15.0]); + } + + #[test] + fn test_eval_conditional() { + let mut arena = Arena::new(); + let selector = arena.alloc(Node::Scalar(2.0)); + let branch1 = arena.alloc(Node::Scalar(100.0)); + let branch2 = arena.alloc(Node::Scalar(200.0)); + let branch3 = arena.alloc(Node::Scalar(300.0)); + let cond = arena.alloc(Node::Conditional { + selector, + branches: vec![branch1, branch2, branch3], + }); + let compiled = CompiledExpr::new(&arena, cond); + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &[], &[], &[]); + assert_eq!(result, &[200.0]); + } + + #[test] + fn test_from_ir() { + use crate::ir::TypedIr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, y)); + + let ir = TypedIr::from_arena(&arena, expr); + let compiled = CompiledExpr::from_ir(ir); + + let y = [1.0, 2.0, 3.0]; + let mut s = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut s, 0.0, &y, &[], &[]); + assert_eq!(result, &[2.0, 4.0, 6.0]); + } + + #[test] + fn test_ir_accessor() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::Scalar(1.0)); + let compiled = CompiledExpr::new(&arena, x); + + // Verify we can access the IR + assert_eq!(compiled.ir().output_len(), 1); + assert_eq!(compiled.ir().instructions().len(), 1); + } + + #[test] + fn test_eval_interpolant_breakpoint_exactness() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(InterpolantData { + x_data: vec![0.0, 1.0, 3.0, 7.0], + y_data: vec![10.0, 20.0, 60.0, -5.0], + }), + child: x, + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + for (xv, expected) in [(0.0, 10.0), (1.0, 20.0), (3.0, 60.0), (7.0, -5.0)] { + let result = compiled.eval(&mut s, 0.0, &[xv], &[], &[]); + assert!( + (result[0] - expected).abs() < 1e-14, + "at x={xv}: expected {expected}, got {}", + result[0] + ); + } + } + + #[test] + fn test_eval_interpolant_linear_extrapolation() { + // Outside the data domain, linear interp extends the boundary segment + // (matches scipy interp1d(fill_value="extrapolate") and casadi interpn_linear). + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(InterpolantData { + x_data: vec![1.0, 2.0, 3.0], + y_data: vec![100.0, 200.0, 300.0], // slope 100 everywhere + }), + child: x, + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + + // Below range: extend first segment (slope 100): y(-5) = 100 + 100*(-5-1) = -500 + let below = compiled.eval(&mut s, 0.0, &[-5.0], &[], &[]); + assert_eq!(below, &[-500.0]); + + // Above range: extend last segment: y(5) = 300 + 100*(5-3) = 500 + let above = compiled.eval(&mut s, 0.0, &[5.0], &[], &[]); + assert_eq!(above, &[500.0]); + } + + #[test] + fn test_eval_conditional_out_of_range_fills_zero() { + let mut arena = Arena::new(); + let branch1 = arena.alloc(Node::Scalar(100.0)); + let branch2 = arena.alloc(Node::Scalar(200.0)); + let branches = vec![branch1, branch2]; + + // selector = 0 is below range (branches are 1-indexed) + let sel0 = arena.alloc(Node::Scalar(0.0)); + let cond = arena.alloc(Node::Conditional { + selector: sel0, + branches: branches.clone(), + }); + let compiled = CompiledExpr::new(&arena, cond); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[0.0]); + + // selector = 10 is above range + let sel10 = arena.alloc(Node::Scalar(10.0)); + let cond2 = arena.alloc(Node::Conditional { + selector: sel10, + branches: branches.clone(), + }); + let compiled2 = CompiledExpr::new(&arena, cond2); + let mut s2 = vec![0.0; compiled2.scratch_len()]; + assert_eq!(compiled2.eval(&mut s2, 0.0, &[], &[], &[]), &[0.0]); + + // selector = -1 is below range + let sel_neg = arena.alloc(Node::Scalar(-1.0)); + let cond3 = arena.alloc(Node::Conditional { + selector: sel_neg, + branches, + }); + let compiled3 = CompiledExpr::new(&arena, cond3); + let mut s3 = vec![0.0; compiled3.scratch_len()]; + assert_eq!(compiled3.eval(&mut s3, 0.0, &[], &[], &[]), &[0.0]); + } + + #[test] + fn test_eval_conditional_nan_fills_zero() { + let mut arena = Arena::new(); + let sel = arena.alloc(Node::Scalar(f64::NAN)); + let branch = arena.alloc(Node::Scalar(999.0)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![branch], + }); + let compiled = CompiledExpr::new(&arena, cond); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[0.0]); + } + + #[test] + fn test_eval_interpolant_cubic() { + // One interval [0, 2] with p(dx) = 1 + 2*dx + 3*dx^2 + 4*dx^3. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 2.0], + coeffs: vec![[1.0, 2.0, 3.0, 4.0]], + }), + child: x, + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + + // In-interval at x=1 (dx=1): 1+2+3+4 = 10 + assert_eq!(compiled.eval(&mut s, 0.0, &[1.0], &[], &[]), &[10.0]); + // At left breakpoint x=0 (dx=0): 1 + assert_eq!(compiled.eval(&mut s, 0.0, &[0.0], &[], &[]), &[1.0]); + // Extrapolate right at x=3 (clamp to interval 0, dx=3): 1+6+27+108 = 142 + assert_eq!(compiled.eval(&mut s, 0.0, &[3.0], &[], &[]), &[142.0]); + } + + #[test] + fn test_eval_interpolant_cubic_breakpoints() { + // Discontinuous constants pin interval selection at interior points, at + // breakpoints (right-continuous like scipy PPoly) and past the right edge. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 1.0, 3.0], + // interval 0: p(dx)=5 (constant); interval 1: p(dx)=7 (constant) + coeffs: vec![[5.0, 0.0, 0.0, 0.0], [7.0, 0.0, 0.0, 0.0]], + }), + child: x, + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + // Interior points + assert_eq!(compiled.eval(&mut s, 0.0, &[0.5], &[], &[]), &[5.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[2.0], &[], &[]), &[7.0]); + // Breakpoints: left edge, interior knot (right interval), right edge + assert_eq!(compiled.eval(&mut s, 0.0, &[0.0], &[], &[]), &[5.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[1.0], &[], &[]), &[7.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[3.0], &[], &[]), &[7.0]); + } + + #[test] + fn test_eval_interpolant_nd_bilinear() { + // One cell [0,2]x[0,2] with p(dx0,dx1) = 1 + 3*dx1 + 2*dx0 + 4*dx0*dx1. + // Coeff layout: index = a0*order + a1 -> [c00, c01, c10, c11]. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 2.0], vec![0.0, 2.0]], + coeffs: vec![1.0, 3.0, 2.0, 4.0], + order: 2, + }), + children: vec![x0, x1], + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + + // In-cell at (1,1): 1 + 3 + 2 + 4 = 10 + assert_eq!(compiled.eval(&mut s, 0.0, &[1.0, 1.0], &[], &[]), &[10.0]); + // Corner (0,0): constant term only + assert_eq!(compiled.eval(&mut s, 0.0, &[0.0, 0.0], &[], &[]), &[1.0]); + // Extrapolate both axes at (3,4): 1 + 3*4 + 2*3 + 4*3*4 = 67 + assert_eq!(compiled.eval(&mut s, 0.0, &[3.0, 4.0], &[], &[]), &[67.0]); + } + + #[test] + fn test_eval_interpolant_nd_cell_selection() { + // Two cells along axis 0 x one along axis 1; discontinuous constants pin + // cell choice, right-continuous like scipy and clamped at both edges. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 1.0, 3.0], vec![0.0, 1.0]], + // cell (0,0): p=5; cell (1,0): p=7 + coeffs: vec![5.0, 0.0, 0.0, 0.0, 7.0, 0.0, 0.0, 0.0], + order: 2, + }), + children: vec![x0, x1], + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[0.5, 0.5], &[], &[]), &[5.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[2.0, 0.5], &[], &[]), &[7.0]); + // Interior knot x0=1: right cell; edges clamp to boundary cells. + assert_eq!(compiled.eval(&mut s, 0.0, &[1.0, 0.5], &[], &[]), &[7.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[0.0, 0.5], &[], &[]), &[5.0]); + assert_eq!(compiled.eval(&mut s, 0.0, &[3.0, 0.5], &[], &[]), &[7.0]); + } + + #[test] + fn test_eval_interpolant_nd_tricubic() { + // One cell [0,2]^3, order 4: p = 3 + dx0^3 + 2*dx1^2 + 5*dx2. + // Coeff index = (a0*4 + a1)*4 + a2. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let x2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let mut coeffs = vec![0.0; 64]; + coeffs[0] = 3.0; // (0,0,0) + coeffs[(3 * 4) * 4] = 1.0; // dx0^3 + coeffs[2 * 4] = 2.0; // dx1^2 + coeffs[1] = 5.0; // dx2^1 + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 2.0], vec![0.0, 2.0], vec![0.0, 2.0]], + coeffs, + order: 4, + }), + children: vec![x0, x1, x2], + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + // (1,1,1): 3 + 1 + 2 + 5 = 11 + assert_eq!( + compiled.eval(&mut s, 0.0, &[1.0, 1.0, 1.0], &[], &[]), + &[11.0] + ); + // Extrapolate axis 0 at (3,0,0): 3 + 27 = 30 + assert_eq!( + compiled.eval(&mut s, 0.0, &[3.0, 0.0, 0.0], &[], &[]), + &[30.0] + ); + } + + #[test] + fn test_eval_interpolant_nd_vector_children_broadcast() { + // p(dx0,dx1) = dx0 + 10*dx1 on one cell [0,10]x[0,10]; vector child + // (len 3) on axis 0, length-1 (broadcast) child on axis 1. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let x1 = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 10.0], vec![0.0, 10.0]], + coeffs: vec![0.0, 10.0, 1.0, 0.0], + order: 2, + }), + children: vec![x0, x1], + }); + let compiled = CompiledExpr::new(&arena, interp); + let mut s = vec![0.0; compiled.scratch_len()]; + let r = compiled.eval(&mut s, 0.0, &[1.0, 2.0, 3.0, 0.5], &[], &[]); + assert_eq!(r, &[6.0, 7.0, 8.0]); + } + + #[test] + fn test_eval_conditional_half_open_window() { + // Selector matching uses (branch_index - 0.5, branch_index + 0.5) + // branch_index is 1-indexed: branch 0 matches selector ∈ (0.5, 1.5) + let mut arena = Arena::new(); + let branch1 = arena.alloc(Node::Scalar(100.0)); + let branch2 = arena.alloc(Node::Scalar(200.0)); + let branches = vec![branch1, branch2]; + + // 0.6 is in (0.5, 1.5) → selects branch 0 + let sel = arena.alloc(Node::Scalar(0.6)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: branches.clone(), + }); + let compiled = CompiledExpr::new(&arena, cond); + let mut s = vec![0.0; compiled.scratch_len()]; + assert_eq!(compiled.eval(&mut s, 0.0, &[], &[], &[]), &[100.0]); + + // 1.4 is in (0.5, 1.5) → selects branch 0 + let sel2 = arena.alloc(Node::Scalar(1.4)); + let cond2 = arena.alloc(Node::Conditional { + selector: sel2, + branches: branches.clone(), + }); + let compiled2 = CompiledExpr::new(&arena, cond2); + let mut s2 = vec![0.0; compiled2.scratch_len()]; + assert_eq!(compiled2.eval(&mut s2, 0.0, &[], &[], &[]), &[100.0]); + + // Exactly 0.5 is NOT in the open interval (0.5, 1.5) → fills zero + let sel3 = arena.alloc(Node::Scalar(0.5)); + let cond3 = arena.alloc(Node::Conditional { + selector: sel3, + branches, + }); + let compiled3 = CompiledExpr::new(&arena, cond3); + let mut s3 = vec![0.0; compiled3.scratch_len()]; + assert_eq!(compiled3.eval(&mut s3, 0.0, &[], &[], &[]), &[0.0]); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/eval_batch.rs b/packages/pybamm-rust/pybamm-core/src/eval_batch.rs new file mode 100644 index 0000000000..bf53792154 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/eval_batch.rs @@ -0,0 +1,1183 @@ +// Intentional u32 usage for compact instruction storage - expression graphs +// won't exceed 4B nodes in practice +#![allow(clippy::cast_possible_truncation)] + +//! Lane-batched (K-wide) evaluator for the primal observe path. +//! +//! [`CompiledExpr::eval_batch`](crate::CompiledExpr::eval_batch) interprets a +//! tape for `k` time points ("lanes") in one pass. The scratch layout is the +//! scalar layout scaled by `k`: scalar buffer index `i` maps to the contiguous +//! lane block `[i*k, (i+1)*k)`. Every handler is then its scalar counterpart with +//! slot offsets scaled by `k` and an inner contiguous loop over lanes, so +//! per-element operations and their order are unchanged and results are +//! **bitwise identical** to `k` independent `eval` calls. +//! +//! Primal instructions only: tangent and state-derivative loads return +//! [`BatchEvalError`] rather than guessing. + +use crate::branch_regions::{active_branch, dispatch_span_end}; +use crate::eval::{ + CompiledExpr, erf_approx, interp_cubic_1d, interp_cubic_1d_deriv, interp_linear_1d, + interp_linear_1d_slope_lookup, locate_nd_cell, sign, split_dst_two_src, split_src_dst, + tensor_horner, tensor_horner_partial, +}; +use crate::ir::{BinaryOp, BroadcastKind, ConstPool, Instruction, UnaryOp}; + +/// Error returned when a tape cannot be evaluated by the primal batch path. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum BatchEvalError { + /// The tape contains a tangent load; `eval_batch` is primal-only. + #[error("eval_batch is primal-only but the tape contains a tangent load")] + NonPrimalInstruction, + /// The tape references state derivatives, which the batch path does not supply. + #[error("eval_batch does not supply y_dot but the tape references state derivatives")] + StateDotUnsupported, +} + +/// Apply a binary broadcast op over `k` lanes with a monomorphic closure. +/// +/// Slot offsets are scaled by `k`; each broadcast kind resolves to contiguous +/// `k`-wide slices iterated with a bounds-check-free `zip`, so the closure +/// vectorises exactly as the scalar `broadcast_apply` does. `lane_tmp` (length +/// `k`) holds the broadcast scalar's per-lane values for the mixed kinds, whose +/// operand aliases the destination-disjoint region the borrow checker cannot +/// otherwise split. +#[allow(clippy::inline_always, clippy::too_many_arguments)] +#[inline(always)] +fn batch_broadcast_apply f64>( + buf: &mut [f64], + k: usize, + f: F, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, + lane_tmp: &mut [f64], +) { + match kind { + BroadcastKind::ScalarScalar => { + let (a_s, b_s, d_s) = split_dst_two_src(buf, a * k, b * k, dst * k, k); + for ((o, &x), &y) in d_s.iter_mut().zip(a_s).zip(b_s) { + *o = f(x, y); + } + }, + BroadcastKind::VectorVector => { + let n = len * k; + let (a_s, b_s, d_s) = split_dst_two_src(buf, a * k, b * k, dst * k, n); + for ((o, &x), &y) in d_s.iter_mut().zip(a_s).zip(b_s) { + *o = f(x, y); + } + }, + BroadcastKind::ScalarVector => { + lane_tmp.copy_from_slice(&buf[a * k..a * k + k]); + let (b_s, d_s) = split_src_dst(buf, b * k, dst * k, len * k); + for (d_chunk, b_chunk) in d_s.chunks_exact_mut(k).zip(b_s.chunks_exact(k)) { + for ((o, &s), &y) in d_chunk.iter_mut().zip(lane_tmp.iter()).zip(b_chunk) { + *o = f(s, y); + } + } + }, + BroadcastKind::VectorScalar => { + lane_tmp.copy_from_slice(&buf[b * k..b * k + k]); + let (a_s, d_s) = split_src_dst(buf, a * k, dst * k, len * k); + for (d_chunk, a_chunk) in d_s.chunks_exact_mut(k).zip(a_s.chunks_exact(k)) { + for ((o, &x), &s) in d_chunk.iter_mut().zip(a_chunk).zip(lane_tmp.iter()) { + *o = f(x, s); + } + } + }, + } +} + +/// Dispatch a binary op to `batch_broadcast_apply` with a monomorphic closure. +/// Must mirror `eval_binary_op` in `eval.rs` closure-for-closure so results +/// stay bitwise identical to the scalar path. +#[allow(clippy::inline_always, clippy::too_many_arguments)] +#[inline(always)] +fn batch_binary( + buf: &mut [f64], + k: usize, + op: BinaryOp, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, + lane_tmp: &mut [f64], +) { + match op { + BinaryOp::Add => { + batch_broadcast_apply(buf, k, |x, y| x + y, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Sub => { + batch_broadcast_apply(buf, k, |x, y| x - y, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Mul => { + batch_broadcast_apply(buf, k, |x, y| x * y, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Div => { + batch_broadcast_apply(buf, k, |x, y| x / y, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Pow => batch_broadcast_apply(buf, k, f64::powf, a, b, dst, len, kind, lane_tmp), + BinaryOp::Minimum => { + batch_broadcast_apply(buf, k, f64::min, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Maximum => { + batch_broadcast_apply(buf, k, f64::max, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Modulo => { + batch_broadcast_apply(buf, k, |x, y| x % y, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::Hypot => { + batch_broadcast_apply(buf, k, f64::hypot, a, b, dst, len, kind, lane_tmp); + }, + BinaryOp::EqualHeaviside => batch_broadcast_apply( + buf, + k, + |x, y| if x <= y { 1.0 } else { 0.0 }, + a, + b, + dst, + len, + kind, + lane_tmp, + ), + BinaryOp::NotEqualHeaviside => batch_broadcast_apply( + buf, + k, + |x, y| if x < y { 1.0 } else { 0.0 }, + a, + b, + dst, + len, + kind, + lane_tmp, + ), + BinaryOp::Equality => { + const EPS: f64 = 1e-14; + batch_broadcast_apply( + buf, + k, + |x, y| if (x - y).abs() < EPS { 1.0 } else { 0.0 }, + a, + b, + dst, + len, + kind, + lane_tmp, + ); + }, + } +} + +/// Apply a unary op over `len * k` contiguous lane values with a monomorphic +/// closure (vectorises like `unary_apply`); handles the in-place `src == dst`. +#[allow(clippy::inline_always)] +#[inline(always)] +fn batch_unary_apply f64>( + buf: &mut [f64], + k: usize, + f: F, + src: usize, + dst: usize, + len: usize, +) { + let n = len * k; + if src == dst { + for x in &mut buf[dst * k..dst * k + n] { + *x = f(*x); + } + return; + } + let (s_s, d_s) = split_src_dst(buf, src * k, dst * k, n); + for (o, &x) in d_s.iter_mut().zip(s_s) { + *o = f(x); + } +} + +/// Dispatch a unary op to `batch_unary_apply`. Must mirror `eval_unary_op` in +/// `eval.rs` closure-for-closure so results stay bitwise identical. +#[allow(clippy::inline_always)] +#[inline(always)] +fn batch_unary(buf: &mut [f64], k: usize, op: UnaryOp, src: usize, dst: usize, len: usize) { + match op { + UnaryOp::Neg => batch_unary_apply(buf, k, |x| -x, src, dst, len), + UnaryOp::Abs => batch_unary_apply(buf, k, f64::abs, src, dst, len), + UnaryOp::Sqrt => batch_unary_apply(buf, k, f64::sqrt, src, dst, len), + UnaryOp::Exp => batch_unary_apply(buf, k, f64::exp, src, dst, len), + UnaryOp::Log => batch_unary_apply(buf, k, f64::ln, src, dst, len), + UnaryOp::Sin => batch_unary_apply(buf, k, f64::sin, src, dst, len), + UnaryOp::Cos => batch_unary_apply(buf, k, f64::cos, src, dst, len), + UnaryOp::Tanh => batch_unary_apply(buf, k, f64::tanh, src, dst, len), + UnaryOp::Sinh => batch_unary_apply(buf, k, f64::sinh, src, dst, len), + UnaryOp::Cosh => batch_unary_apply(buf, k, f64::cosh, src, dst, len), + UnaryOp::Arcsinh => batch_unary_apply(buf, k, f64::asinh, src, dst, len), + UnaryOp::Arctan => batch_unary_apply(buf, k, f64::atan, src, dst, len), + UnaryOp::Erf => batch_unary_apply(buf, k, erf_approx, src, dst, len), + UnaryOp::Sign => batch_unary_apply(buf, k, sign, src, dst, len), + UnaryOp::Floor => batch_unary_apply(buf, k, f64::floor, src, dst, len), + UnaryOp::Ceiling => batch_unary_apply(buf, k, f64::ceil, src, dst, len), + } +} + +/// Execute a primal instruction slice against a lane-batched buffer. +/// +/// `Dispatch` runs the union of blocks any lane's selector picks, skipping a +/// block only when no lane needs it, `Conditional` then still selects the +/// output per lane from within that union. +#[allow(clippy::too_many_arguments, clippy::too_many_lines)] +fn eval_batch_instructions( + buf: &mut [f64], + k: usize, + instructions: &[Instruction], + consts: &ConstPool, + ts: &[f64], + y_cols: &[f64], + n_states: usize, + inputs: &[f64], +) -> Result<(), BatchEvalError> { + // Reused per-lane scratch: broadcast-scalar operands and matmul row sums. + let mut lane_tmp = vec![0.0_f64; k]; + let mut pc = 0_usize; + while pc < instructions.len() { + if let Instruction::Dispatch { + selector, + blocks_idx, + blocks_len, + } = instructions[pc] + { + let base = blocks_idx as usize; + let n = blocks_len as usize; + let span_end = dispatch_span_end(consts, pc, blocks_idx, blocks_len); + let sel_base = selector as usize * k; + // Union over lanes: a block runs if any lane selects it, as the + // un-blocked tape did, and branch roots are never recycled. + let mut needed = vec![false; n]; + for l in 0..k { + if let Some(active) = active_branch(buf[sel_base + l], n) { + needed[active] = true; + } + } + for (b, &need) in needed.iter().enumerate() { + if need { + let (rel, len) = consts.branch_blocks[base + b]; + let start = pc + rel as usize; + let end = start + len as usize; + eval_batch_instructions( + buf, + k, + &instructions[start..end], + consts, + ts, + y_cols, + n_states, + inputs, + )?; + } + } + pc = span_end; + continue; + } + + match instructions[pc] { + Instruction::LoadScalar { value, dst } => { + let base = dst as usize * k; + buf[base..base + k].fill(value); + }, + Instruction::LoadTime { dst } => { + let base = dst as usize * k; + buf[base..base + k].copy_from_slice(&ts[..k]); + }, + Instruction::LoadArray { data_idx, len, dst } => { + let src = consts.get_array(data_idx, len); + let base = dst as usize * k; + for (e, &v) in src.iter().enumerate() { + let o = base + e * k; + buf[o..o + k].fill(v); + } + }, + Instruction::FillZero { dst, len } => { + let base = dst as usize * k; + buf[base..base + len as usize * k].fill(0.0); + }, + Instruction::LoadStateVector { start, end, dst } => { + let base = dst as usize * k; + let slen = (end - start) as usize; + // Transposing gather: element e, lane l reads column l of the + // (n_states, k) F-contiguous state matrix. + for e in 0..slen { + let o = base + e * k; + let row = start as usize + e; + for l in 0..k { + buf[o + l] = y_cols[l * n_states + row]; + } + } + }, + Instruction::LoadStateVectorDot { .. } => { + return Err(BatchEvalError::StateDotUnsupported); + }, + Instruction::LoadInputParameter { offset, width, dst } => { + let base = dst as usize * k; + let off = offset as usize; + for e in 0..width as usize { + let v = inputs[off + e]; + let o = base + e * k; + buf[o..o + k].fill(v); + } + }, + Instruction::LoadTangentState { .. } | Instruction::LoadTangentParameter { .. } => { + return Err(BatchEvalError::NonPrimalInstruction); + }, + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => { + batch_binary( + buf, + k, + op, + a as usize, + b as usize, + dst as usize, + len as usize, + kind, + &mut lane_tmp, + ); + }, + + Instruction::Unary { op, src, dst, len } => { + batch_unary(buf, k, op, src as usize, dst as usize, len as usize); + }, + + Instruction::MaxReduce { src, src_len, dst } => { + let (src, src_len) = (src as usize, src_len as usize); + let dbase = dst as usize * k; + for l in 0..k { + let mut max_val = f64::NEG_INFINITY; + for e in 0..src_len { + let v = buf[(src + e) * k + l]; + if v > max_val { + max_val = v; + } + } + buf[dbase + l] = max_val; + } + }, + + Instruction::MinReduce { src, src_len, dst } => { + let (src, src_len) = (src as usize, src_len as usize); + let dbase = dst as usize * k; + for l in 0..k { + let mut min_val = f64::INFINITY; + for e in 0..src_len { + let v = buf[(src + e) * k + l]; + if v < min_val { + min_val = v; + } + } + buf[dbase + l] = min_val; + } + }, + + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + is_max, + dst, + } => { + let (basis, picker, len) = (basis_src as usize, picker_src as usize, len as usize); + let dbase = dst as usize * k; + // Per lane: first-occurrence argmax/argmin with a strict + // comparison from element 0, matching the scalar eval. + for l in 0..k { + let mut best_idx = 0; + let mut best_val = buf[picker * k + l]; + for e in 1..len { + let v = buf[(picker + e) * k + l]; + if (is_max && v > best_val) || (!is_max && v < best_val) { + best_val = v; + best_idx = e; + } + } + buf[dbase + l] = buf[(basis + best_idx) * k + l]; + } + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + // Slot elements [start, start+len) map to contiguous lane blocks, + // so the whole window is one contiguous copy (memmove-safe). + let s = (src as usize + start as usize) * k; + let n = len as usize * k; + buf.copy_within(s..s + n, dst as usize * k); + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write_pos = dst as usize * k; + for i in 0..sources_len as usize { + let (src_off, src_len) = consts.concat_sources[sources_idx as usize + i]; + let n = src_len as usize * k; + buf.copy_within(src_off as usize * k..src_off as usize * k + n, write_pos); + write_pos += n; + } + }, + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx as usize]; + let (vec_src, dst) = (vec_src as usize, dst as usize); + for row in 0..csr.shape.rows { + // `lane_tmp` is disjoint from `buf` so the inner lane loop + // vectorises; entries accumulate in scalar SpMV order. + lane_tmp.fill(0.0); + let (start, end) = (csr.indptr[row], csr.indptr[row + 1]); + for (&col, &val) in csr.indices[start..end].iter().zip(&csr.data[start..end]) { + let cbase = (vec_src + col) * k; + let x = &buf[cbase..cbase + k]; + for (acc, &xi) in lane_tmp.iter_mut().zip(x) { + *acc += val * xi; + } + } + let rbase = (dst + row) * k; + buf[rbase..rbase + k].copy_from_slice(&lane_tmp); + } + }, + + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + dst, + } => { + let (mat_src, vec_src, dst) = (mat_src as usize, vec_src as usize, dst as usize); + let cols = cols as usize; + for row in 0..rows as usize { + lane_tmp.fill(0.0); + for col in 0..cols { + let mbase = (mat_src + row * cols + col) * k; + let vbase = (vec_src + col) * k; + let m = &buf[mbase..mbase + k]; + let v = &buf[vbase..vbase + k]; + for ((acc, &mi), &vi) in lane_tmp.iter_mut().zip(m).zip(v) { + *acc += mi * vi; + } + } + let rbase = (dst + row) * k; + buf[rbase..rbase + k].copy_from_slice(&lane_tmp); + } + }, + + Instruction::Interp1DLinear { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx as usize]; + let (src, dst) = (src as usize, dst as usize); + for e in 0..len as usize { + let (sb, db) = ((src + e) * k, (dst + e) * k); + for l in 0..k { + buf[db + l] = interp_linear_1d(&interp.x_data, &interp.y_data, buf[sb + l]); + } + } + }, + + Instruction::Interp1DLinearDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx as usize]; + let (src, dst) = (src as usize, dst as usize); + for e in 0..len as usize { + let (sb, db) = ((src + e) * k, (dst + e) * k); + for l in 0..k { + buf[db + l] = interp_linear_1d_slope_lookup( + &interp.x_data, + &interp.y_data, + buf[sb + l], + ); + } + } + }, + + Instruction::Interp1DCubic { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx as usize]; + let (src, dst) = (src as usize, dst as usize); + for e in 0..len as usize { + let (sb, db) = ((src + e) * k, (dst + e) * k); + for l in 0..k { + buf[db + l] = + interp_cubic_1d(&interp.breakpoints, &interp.coeffs, buf[sb + l]); + } + } + }, + + Instruction::Interp1DCubicDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx as usize]; + let (src, dst) = (src as usize, dst as usize); + for e in 0..len as usize { + let (sb, db) = ((src + e) * k, (dst + e) * k); + for l in 0..k { + buf[db + l] = + interp_cubic_1d_deriv(&interp.breakpoints, &interp.coeffs, buf[sb + l]); + } + } + }, + + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len, + } => { + let interp = &consts.nd_interpolants[interp_idx as usize]; + let ndim = interp.breakpoints.len(); + let order = interp.order as usize; + let dst = dst as usize; + let mut coords = [0.0_f64; 3]; + let mut dxs = [0.0_f64; 3]; + for e in 0..len as usize { + for l in 0..k { + for (a, coord) in coords.iter_mut().enumerate().take(ndim) { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + let j = if slen == 1 { 0 } else { e }; + *coord = buf[(off as usize + j) * k + l]; + } + let cell = locate_nd_cell( + &interp.breakpoints, + &interp.coeffs, + order, + &coords[..ndim], + &mut dxs, + ); + buf[(dst + e) * k + l] = tensor_horner(cell, &dxs[..ndim], order); + } + } + }, + + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis, + dst, + len, + } => { + let interp = &consts.nd_interpolants[interp_idx as usize]; + let ndim = interp.breakpoints.len(); + let order = interp.order as usize; + let dst = dst as usize; + let mut coords = [0.0_f64; 3]; + let mut dxs = [0.0_f64; 3]; + for e in 0..len as usize { + for l in 0..k { + for (a, coord) in coords.iter_mut().enumerate().take(ndim) { + let (off, slen) = consts.interp_nd_sources[sources_idx as usize + a]; + let j = if slen == 1 { 0 } else { e }; + *coord = buf[(off as usize + j) * k + l]; + } + let cell = locate_nd_cell( + &interp.breakpoints, + &interp.coeffs, + order, + &coords[..ndim], + &mut dxs, + ); + buf[(dst + e) * k + l] = + tensor_horner_partial(cell, &dxs[..ndim], order, axis as usize); + } + } + }, + + Instruction::Conditional { + selector, + branches_idx, + branches_len, + dst, + out_len, + } => { + let sel_base = selector as usize * k; + let dst = dst as usize; + let out_len = out_len as usize; + // Each lane selects its branch independently; branch offsets are + // slot offsets, scaled by k at the copy site. + for l in 0..k { + match active_branch(buf[sel_base + l], branches_len as usize) { + Some(i) => { + let (branch_off, _) = consts.branch_offsets[branches_idx as usize + i]; + let bbase = branch_off as usize; + for e in 0..out_len { + buf[(dst + e) * k + l] = buf[(bbase + e) * k + l]; + } + }, + None => { + for e in 0..out_len { + buf[(dst + e) * k + l] = 0.0; + } + }, + } + } + }, + + Instruction::Dispatch { .. } => unreachable!("handled above"), + } + pc += 1; + } + Ok(()) +} + +impl CompiledExpr { + /// Evaluate the tape for `k` lanes (time points) at once. + /// + /// `scratch` must hold at least `scratch_len() * k` elements. Slot `s` occupies + /// `[s*k, (s + slot_len)*k)`, and element `e` of it holds its `k` lane values + /// contiguously at `[(s+e)*k, (s+e+1)*k)`. `ts` supplies the `k` time values, + /// `y_cols` is the `(n_states, k)` F-contiguous state matrix, and `inputs` is + /// shared across lanes. Returns the root slot as `(out_len, k)` lane-minor: + /// element `e`, lane `l` at relative index `e*k + l`. + /// + /// Results are bitwise identical to `k` independent [`eval`](Self::eval) calls. + /// Primal-only: a tangent or state-derivative load returns [`BatchEvalError`]. + pub fn eval_batch<'s>( + &self, + scratch: &'s mut [f64], + k: usize, + ts: &[f64], + y_cols: &[f64], + inputs: &[f64], + ) -> Result<&'s [f64], BatchEvalError> { + let ir = self.ir(); + debug_assert!( + ir.split_eval_info().is_none(), + "eval_batch requires a primal (non-split-eval) tape" + ); + debug_assert!(k >= 1, "eval_batch needs at least one lane"); + debug_assert!(ts.len() >= k, "ts must have at least k time values"); + debug_assert!( + scratch.len() >= ir.buffer_size() * k, + "scratch too small for k lanes" + ); + let n_states = y_cols.len() / k; + eval_batch_instructions( + scratch, + k, + ir.instructions(), + ir.consts(), + ts, + y_cols, + n_states, + inputs, + )?; + let root = ir.root_slot(); + Ok(&scratch[root.offset_usize() * k..(root.offset_usize() + root.len_usize()) * k]) + } +} + +#[cfg(test)] +mod tests { + // Test-fixture arithmetic favours readable data generation over FMA. + #![allow(clippy::suboptimal_flops)] + use super::*; + use crate::arena::{Arena, NodeId}; + use crate::node::{CsrData, InterpolantData, Node, Shape}; + + /// Evaluate `k` lanes of random `(t, y)` through both the scalar `eval` + /// (once per lane) and `eval_batch`, asserting bitwise equality. + fn assert_batch_matches_scalar( + arena: &Arena, + root: NodeId, + n_states: usize, + k: usize, + ts: &[f64], + y_cols: &[f64], + inputs: &[f64], + ) { + let expr = CompiledExpr::new(arena, root); + let out_len = expr.output_len(); + + // Scalar reference: one eval per lane. + let mut scalar = vec![0.0_f64; out_len * k]; + let mut s = vec![0.0_f64; expr.scratch_len()]; + for l in 0..k { + let y = &y_cols[l * n_states..(l + 1) * n_states]; + let res = expr.eval(&mut s, ts[l], y, &[], inputs); + scalar[l * out_len..(l + 1) * out_len].copy_from_slice(res); + } + + // Batched: one eval_batch over all lanes. + let mut batch_scratch = vec![0.0_f64; expr.scratch_len() * k]; + let root_slice = expr + .eval_batch(&mut batch_scratch, k, ts, y_cols, inputs) + .expect("primal tape must batch-evaluate"); + for l in 0..k { + for e in 0..out_len { + let got = root_slice[e * k + l]; + let want = scalar[l * out_len + e]; + assert_eq!( + got.to_bits(), + want.to_bits(), + "lane {l}, elem {e}: batch {got} != scalar {want}" + ); + } + } + } + + /// Build a column-major `(n_states, k)` state matrix from a closure. + fn state_cols(n_states: usize, k: usize, f: impl Fn(usize, usize) -> f64) -> Vec { + let mut cols = vec![0.0_f64; n_states * k]; + for l in 0..k { + for i in 0..n_states { + cols[l * n_states + i] = f(i, l); + } + } + cols + } + + #[test] + fn all_binary_ops_match_scalar() { + // Every BinaryOp variant, guarding drift in the duplicated op table. + let ops = [ + BinaryOp::Add, + BinaryOp::Sub, + BinaryOp::Mul, + BinaryOp::Div, + BinaryOp::Pow, + BinaryOp::Minimum, + BinaryOp::Maximum, + BinaryOp::Modulo, + BinaryOp::Hypot, + BinaryOp::EqualHeaviside, + BinaryOp::NotEqualHeaviside, + BinaryOp::Equality, + ]; + let k = 7; + let ts: Vec = (0..k).map(|l| l as f64 * 0.3).collect(); + let y_cols = state_cols(2, k, |i, l| 0.5 + i as f64 + l as f64 * 0.11); + for op in ops { + let mut arena = Arena::new(); + let a = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = alloc_binary(&mut arena, op, a, b); + assert_batch_matches_scalar(&arena, root, 2, k, &ts, &y_cols, &[]); + } + } + + #[test] + fn all_unary_ops_match_scalar() { + let ops = [ + UnaryOp::Neg, + UnaryOp::Abs, + UnaryOp::Sqrt, + UnaryOp::Exp, + UnaryOp::Log, + UnaryOp::Sin, + UnaryOp::Cos, + UnaryOp::Tanh, + UnaryOp::Sinh, + UnaryOp::Cosh, + UnaryOp::Arcsinh, + UnaryOp::Arctan, + UnaryOp::Erf, + UnaryOp::Sign, + UnaryOp::Floor, + UnaryOp::Ceiling, + ]; + let k = 5; + let ts: Vec = (0..k).map(|l| l as f64 * 0.2).collect(); + // Positive inputs so Sqrt/Log stay in-domain; varied per lane. + let y_cols = state_cols(1, k, |_, l| 0.3 + l as f64 * 0.37); + for op in ops { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = alloc_unary(&mut arena, op, x); + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[]); + } + } + + #[test] + fn broadcast_scalar_vector_per_lane() { + // (scalar y0) * (vector y1..y4): ScalarVector with a per-lane scalar. + let mut arena = Arena::new(); + let s = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let v = arena.alloc(Node::StateVector { start: 1, end: 4 }); + let root = arena.alloc(Node::Mul(s, v)); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(4, k, |i, l| 1.0 + i as f64 * 0.5 + l as f64); + assert_batch_matches_scalar(&arena, root, 4, k, &ts, &y_cols, &[]); + } + + #[test] + fn conditional_selects_different_branches_per_lane() { + // selector = y0 (1.0 or 2.0 per lane) picks branch 1 or 2. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let three = arena.alloc(Node::Scalar(3.0)); + let b1 = arena.alloc(Node::Add(y1, two)); + let b2 = arena.alloc(Node::Mul(y2, three)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + // Alternate selector 1.0 / 2.0 so lanes take different branches, plus a + // lane whose selector matches nothing (0.0 -> zero-filled). + let y_cols = state_cols(3, k, |i, l| match i { + 0 => { + if l == 3 { + 0.0 + } else { + (l % 2 + 1) as f64 + } + }, + _ => i as f64 + l as f64 * 0.25, + }); + assert_batch_matches_scalar(&arena, root, 3, k, &ts, &y_cols, &[]); + } + + #[test] + fn reduce_argselect_tie_breaking_per_lane() { + // pybamm.max subgradient: basis[argmax(picker)] with first-occurrence + // ties, exercised with per-lane tie positions. + let mut arena = Arena::new(); + let picker = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let root = arena.alloc(Node::MaxReduce(picker)); + let k = 3; + let ts = vec![0.0; k]; + // Lane 0: [5,5,3] (tie at 0), lane 1: [1,5,5] (tie at 1), lane 2: strict. + let picks = [[5.0, 5.0, 3.0], [1.0, 5.0, 5.0], [1.0, 2.0, 3.0]]; + let y_cols = state_cols(3, k, |i, l| picks[l][i]); + assert_batch_matches_scalar(&arena, root, 3, k, &ts, &y_cols, &[]); + } + + #[test] + fn sparse_matmul_matches_scalar() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 4], + vec![0, 1, 2, 3], + vec![1.0, -1.0, 2.0, 3.0], + Shape::matrix(2, 4), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(mat, y)); + let k = 6; + let ts = vec![0.0; k]; + let y_cols = state_cols(4, k, |i, l| 0.5 + i as f64 + l as f64 * 0.13); + assert_batch_matches_scalar(&arena, root, 4, k, &ts, &y_cols, &[]); + } + + #[test] + fn interpolant_matches_scalar() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 25.0]) + .expect("valid interpolant"), + ), + child: x, + }); + let k = 5; + let ts = vec![0.0; k]; + // Spread across, below, and above the knot range to hit extrapolation. + let y_cols = state_cols(1, k, |_, l| -0.5 + l as f64 * 0.7); + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[]); + } + + #[test] + fn ragged_tail_k1_matches_scalar() { + // k == 1 is the ragged-tail degenerate case (compact stride 1). + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let root = arena.alloc(Node::Mul(y, two)); + let ts = [0.4]; + let y_cols = [1.5, -2.5]; + assert_batch_matches_scalar(&arena, root, 2, 1, &ts, &y_cols, &[]); + } + + #[test] + fn time_and_concat_match_scalar() { + // Concat([t, y0]) exercises LoadTime + Concat over lanes. + let mut arena = Arena::new(); + let t = arena.alloc(Node::Time); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![t, y0])); + let k = 4; + let ts: Vec = (0..k).map(|l| 0.1 + l as f64 * 0.9).collect(); + let y_cols = state_cols(1, k, |_, l| l as f64 * 2.0 - 1.0); + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[]); + } + + #[test] + fn tangent_load_returns_error() { + // A split-eval (tangent) tape must be rejected, not silently mis-evaluated. + use crate::tangent::tangent_wrt_states; + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let root = arena.alloc(Node::Mul(y, two)); + let troot = tangent_wrt_states(&mut arena, root); + let expr = CompiledExpr::new(&arena, troot); + let mut scratch = vec![0.0_f64; expr.scratch_len() * 2]; + let err = expr.eval_batch(&mut scratch, 2, &[0.0, 0.0], &[1.0, 2.0], &[]); + assert_eq!(err, Err(BatchEvalError::NonPrimalInstruction)); + } + + #[test] + fn dispatch_skips_blocks_no_lane_selects() { + // Selector is an InputParameter, so all lanes agree: only branch 2's + // block may run, and the result must match the scalar path bitwise. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let mut b2 = y; + for _ in 0..6 { + b2 = arena.alloc(Node::Exp(b2)); + } + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + for sel in [1.0_f64, 2.0, 0.0] { + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[sel]); + } + } + + #[test] + fn dispatch_falls_back_to_the_union_on_lane_divergence() { + // y-derived selector: lanes pick different branches, so both blocks must + // run. Bitwise parity with per-lane scalar eval is the assertion. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let b1 = arena.alloc(Node::Sin(y1)); + let b2 = arena.alloc(Node::Exp(y2)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + // Lanes 0..2 alternate selector 1/2; lane 3 matches nothing. + let y_cols = state_cols(3, k, |i, l| match i { + 0 => { + if l == 3 { + 0.0 + } else { + (l % 2 + 1) as f64 + } + }, + _ => i as f64 + l as f64 * 0.25, + }); + assert_batch_matches_scalar(&arena, root, 3, k, &ts, &y_cols, &[]); + } + + #[test] + fn dispatch_skips_unneeded_block_even_when_it_would_error() { + // Branch 2 loads a state derivative, which eval_batch rejects whenever it + // is evaluated, so this passes only if branch 2's block is truly skipped. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let y_dot = arena.alloc(Node::StateVectorDot { start: 0, end: 1 }); + let b2 = arena.alloc(Node::Exp(y_dot)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let expr = CompiledExpr::new(&arena, root); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + let mut scratch = vec![0.0_f64; expr.scratch_len() * k]; + let result = expr.eval_batch(&mut scratch, k, &ts, &y_cols, &[1.0]); + assert!( + result.is_ok(), + "branch 2 must never run when no lane selects it: {result:?}" + ); + } + + #[test] + fn two_dispatches_in_one_tape_match_scalar() { + // Two independent Conditionals lower to two Dispatch instructions, so the + // union-dispatch loop must resume after the first to find the second. + let mut arena = Arena::new(); + let sel1 = arena.alloc(Node::InputParameter { + name: "s1".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let sel2 = arena.alloc(Node::InputParameter { + name: "s2".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y2 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a1 = arena.alloc(Node::Sin(y1)); + let a2 = arena.alloc(Node::Cos(y1)); + let cond1 = arena.alloc(Node::Conditional { + selector: sel1, + branches: vec![a1, a2], + }); + let b1 = arena.alloc(Node::Exp(y2)); + let b2 = arena.alloc(Node::Sqrt(y2)); + let cond2 = arena.alloc(Node::Conditional { + selector: sel2, + branches: vec![b1, b2], + }); + let root = arena.alloc(Node::Add(cond1, cond2)); + let expr = CompiledExpr::new(&arena, root); + assert_eq!( + expr.ir().branch_block_lens().len(), + 4, + "two two-branch dispatches, or this test covers nothing" + ); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(2, k, |i, l| match i { + 0 => 0.2 + l as f64 * 0.1, + _ => 0.5 + l as f64 * 0.2, + }); + for sels in [[1.0, 1.0], [2.0, 2.0], [1.0, 2.0], [0.0, 1.0]] { + assert_batch_matches_scalar(&arena, root, 2, k, &ts, &y_cols, &sels); + } + } + + #[test] + fn dispatch_on_boundary_and_nan_selectors_matches_scalar() { + // The semantics contract's edge cases: half-integer window boundaries, + // NaN, negative and infinite selectors all match no branch. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + for sel in [0.5_f64, 1.5, 2.5, -1.0, f64::NAN, f64::INFINITY] { + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[sel]); + } + + // And per-lane: one weird selector per lane, so the union is empty for + // some lanes while others still pick a branch. + let lane_selector = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let lb1 = arena.alloc(Node::Sin(y1)); + let lb2 = arena.alloc(Node::Cos(y1)); + let lane_root = arena.alloc(Node::Conditional { + selector: lane_selector, + branches: vec![lb1, lb2], + }); + let weird = [f64::NAN, 0.5, 1.5, 2.5, -1.0, f64::INFINITY, 1.0, 2.0]; + let kw = weird.len(); + let y_cols = state_cols(2, kw, |i, l| if i == 0 { weird[l] } else { 0.3 }); + assert_batch_matches_scalar(&arena, lane_root, 2, kw, &vec![0.0; kw], &y_cols, &[]); + } + + fn alloc_binary(arena: &mut Arena, op: BinaryOp, a: NodeId, b: NodeId) -> NodeId { + match op { + BinaryOp::Add => arena.alloc(Node::Add(a, b)), + BinaryOp::Sub => arena.alloc(Node::Sub(a, b)), + BinaryOp::Mul => arena.alloc(Node::Mul(a, b)), + BinaryOp::Div => arena.alloc(Node::Div(a, b)), + BinaryOp::Pow => arena.alloc(Node::Pow(a, b)), + BinaryOp::Minimum => arena.alloc(Node::Minimum(a, b)), + BinaryOp::Maximum => arena.alloc(Node::Maximum(a, b)), + BinaryOp::Modulo => arena.alloc(Node::Modulo(a, b)), + BinaryOp::Hypot => arena.alloc(Node::Hypot(a, b)), + BinaryOp::EqualHeaviside => arena.alloc(Node::EqualHeaviside(a, b)), + BinaryOp::NotEqualHeaviside => arena.alloc(Node::NotEqualHeaviside(a, b)), + BinaryOp::Equality => arena.alloc(Node::Equality(a, b)), + } + } + + fn alloc_unary(arena: &mut Arena, op: UnaryOp, x: NodeId) -> NodeId { + match op { + UnaryOp::Neg => arena.alloc(Node::Neg(x)), + UnaryOp::Abs => arena.alloc(Node::Abs(x)), + UnaryOp::Sqrt => arena.alloc(Node::Sqrt(x)), + UnaryOp::Exp => arena.alloc(Node::Exp(x)), + UnaryOp::Log => arena.alloc(Node::Log(x)), + UnaryOp::Sin => arena.alloc(Node::Sin(x)), + UnaryOp::Cos => arena.alloc(Node::Cos(x)), + UnaryOp::Tanh => arena.alloc(Node::Tanh(x)), + UnaryOp::Sinh => arena.alloc(Node::Sinh(x)), + UnaryOp::Cosh => arena.alloc(Node::Cosh(x)), + UnaryOp::Arcsinh => arena.alloc(Node::Arcsinh(x)), + UnaryOp::Arctan => arena.alloc(Node::Arctan(x)), + UnaryOp::Erf => arena.alloc(Node::Erf(x)), + UnaryOp::Sign => arena.alloc(Node::Sign(x)), + UnaryOp::Floor => arena.alloc(Node::Floor(x)), + UnaryOp::Ceiling => arena.alloc(Node::Ceiling(x)), + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/ffi.rs b/packages/pybamm-rust/pybamm-core/src/ffi.rs new file mode 100644 index 0000000000..258b252643 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/ffi.rs @@ -0,0 +1,2608 @@ +//! FFI (Foreign Function Interface) for IDAKLU solver integration. +//! +//! This module provides C ABI wrappers for the `ModelEvaluator` methods, +//! allowing the Rust expression evaluator to be called from C/C++ code +//! (specifically the IDAKLU DAE solver used by `PyBaMM`). +//! +//! # Safety +//! +//! All functions in this module are `unsafe extern "C"` and require: +//! - Valid, non-null pointers for all array arguments +//! - Correctly sized arrays (`n_states` elements) +//! - Valid `user_data` pointer to a `ModelEvaluator` instance +//! +//! Every entry point enters through [`with_model`] (a query) or +//! [`with_model_and_inputs`] (an evaluation), which own the whole boundary +//! contract — null rejection, panic containment, the cast of `user_data`, and +//! the input-parameter buffer — so an entry point body is only its own work. + +// FFI code requires unsafe - this module is intentionally using unsafe patterns +// for C interop with proper safety checks at the boundaries. +#![allow(unsafe_code)] +#![allow(clippy::not_unsafe_ptr_arg_deref)] // We check null before deref +#![allow(clippy::missing_panics_doc)] // Panics are caught at boundary +#![allow(clippy::missing_const_for_fn)] // FFI functions cannot be const + +use crate::model::ModelEvaluator; +use crate::observable::ObservableKind; +use std::ffi::c_void; +use std::os::raw::c_int; + +/// Success return code. +pub const SUCCESS: c_int = 0; + +/// Error: null pointer passed to function. +pub const ERROR_NULL_POINTER: c_int = -1; + +/// Error: panic occurred during execution. +pub const ERROR_PANIC: c_int = -2; + +/// Error: invalid parameter index supplied to a sensitivity-related FFI fn. +pub const ERROR_INVALID_PARAM: c_int = -3; + +/// Error: invalid output-variable index supplied to an output-related FFI fn. +pub const ERROR_INVALID_OUTPUT: c_int = -4; + +/// Error: caller invoked a sensitivity FFI on a model with no sensitivities. +pub const ERROR_NO_SENS: c_int = -6; + +/// Error: caller invoked an output FFI on a model with no output variables. +pub const ERROR_NO_OUTPUTS: c_int = -7; + +/// Error: caller invoked an algebraic FFI on a model with no algebraic block. +pub const ERROR_NO_ALG: c_int = -8; + +/// Error: caller invoked an event FFI on a model with no events. +pub const ERROR_NO_EVENTS: c_int = -9; + +/// ABI contract version for the Rust FFI surface. +/// +/// Bump by 1 whenever ANY exported signature changes, +/// arg/return types, or a function added, removed, or reordered. The C++ +/// consumer pins the expected value in `PYBAMM_RUST_ABI_VERSION`, and the +/// drift test asserts the two are equal. +/// +/// `pybamm_rust_abi_version` itself must forever keep the signature `-> u32` with no +/// arguments: it is the probe the C++ consumer calls to read this version, so it +/// cannot follow the bump rule it enforces. +pub const RUST_ABI_VERSION: u32 = 1; + +/// Return the FFI ABI contract version. +/// +/// # Safety +/// +/// Takes no pointer arguments and is always safe to call. +#[unsafe(no_mangle)] +pub extern "C" fn pybamm_rust_abi_version() -> u32 { + RUST_ABI_VERSION +} + +/// Golden hash of the normalized exported symbol surface, pinned by the +/// `test_ffi_abi_contract` drift test. +/// +/// Changing any exported signature changes this hash and fails that test, +/// which then instructs you to update this value AND bump `RUST_ABI_VERSION` +/// / `PYBAMM_RUST_ABI_VERSION` in lockstep. This makes the version bump +/// enforceable rather than a manual convention. +pub const EXPECTED_ABI_HASH: u64 = 0x6584_760b_149a_9779; + +/// Run `body` with the model borrowed immutably from `user_data`. +/// +/// Returns `ERROR_NULL_POINTER` if `user_data` or any pointer in `required` is +/// null, and `ERROR_PANIC` if `body` unwinds; otherwise `body`'s own status. +/// +/// # Safety +/// +/// `user_data` must point to a live `ModelEvaluator` that stays valid, and is +/// not mutably aliased, for the call. +#[inline] +unsafe fn with_model( + user_data: *const c_void, + required: [*const c_void; N], + body: F, +) -> c_int +where + F: FnOnce(&ModelEvaluator) -> c_int, +{ + if user_data.is_null() || any_null(required) { + return ERROR_NULL_POINTER; + } + // SAFETY: null-checked above; the caller guarantees a live `ModelEvaluator`. + let model = unsafe { &*user_data.cast::() }; + caught(|| body(model)) +} + +/// Run `body` with the model borrowed mutably from `user_data` and its input +/// parameters borrowed from `inputs`. +/// +/// Rejects nulls and contains panics as [`with_model`] does, and holds `inputs` +/// to the same rule: null is accepted only when the model declares no input +/// parameters. A body handed an empty slice for a model that has parameters +/// would read past it and panic, reporting a caller's null as a Rust fault. +/// +/// # Safety +/// +/// `user_data` must point to a live `ModelEvaluator` that stays valid, and is +/// not otherwise aliased, for the call. When the model declares input +/// parameters, `inputs` must point to at least that many `f64`. +#[inline] +unsafe fn with_model_and_inputs( + user_data: *mut c_void, + required: [*const c_void; N], + inputs: *const f64, + body: F, +) -> c_int +where + F: FnOnce(&mut ModelEvaluator, &[f64]) -> c_int, +{ + if user_data.is_null() || any_null(required) { + return ERROR_NULL_POINTER; + } + // SAFETY: null-checked above; the caller guarantees a live `ModelEvaluator`. + let model = unsafe { &mut *user_data.cast::() }; + let n_params = model.n_params(); + if n_params != 0 && inputs.is_null() { + return ERROR_NULL_POINTER; + } + let inputs: &[f64] = if n_params == 0 { + &[] + } else { + // SAFETY: non-null per the check above, and the caller guarantees + // `n_params` elements. + unsafe { borrow_slice(inputs, n_params) } + }; + caught(|| body(model, inputs)) +} + +/// Whether any caller-supplied buffer pointer is null. +#[inline] +fn any_null(pointers: [*const c_void; N]) -> bool { + pointers.iter().any(|pointer| pointer.is_null()) +} + +/// Run `body`, turning an unwind into `ERROR_PANIC` rather than letting it +/// cross back into C, which would be undefined behaviour. +#[inline] +fn caught c_int>(body: F) -> c_int { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)).unwrap_or(ERROR_PANIC) +} + +/// Borrow `len` elements of a caller-provided buffer. +/// +/// # Safety +/// +/// `pointer` must be non-null and point to at least `len` initialized `T`, +/// unaliased for the borrow. +#[inline] +unsafe fn borrow_slice<'a, T>(pointer: *const T, len: usize) -> &'a [T] { + // SAFETY: guaranteed by this function's own contract. + unsafe { std::slice::from_raw_parts(pointer, len) } +} + +/// Borrow `len` writable elements of a caller-provided buffer. +/// +/// # Safety +/// +/// `pointer` must be non-null and point to at least `len` writable `T`, +/// unaliased for the borrow. +#[inline] +unsafe fn borrow_slice_mut<'a, T>(pointer: *mut T, len: usize) -> &'a mut [T] { + // SAFETY: guaranteed by this function's own contract. + unsafe { std::slice::from_raw_parts_mut(pointer, len) } +} + +/// Time the enclosing scope into `$counter`, or expand to nothing when the +/// `profile` feature is off. +macro_rules! profile_scope { + ($counter:ident) => { + #[cfg(feature = "profile")] + let _profile_scope = profile::$counter.scope(); + }; +} + +/// Compile-time FFI profiling, absent entirely when the feature is disabled. +#[cfg(feature = "profile")] +mod profile { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Instant; + + /// Calls and accumulated wall time for one FFI entry point. + #[derive(Debug)] + pub struct Counter { + calls: AtomicU64, + nanos: AtomicU64, + } + + impl Counter { + const fn new() -> Self { + Self { + calls: AtomicU64::new(0), + nanos: AtomicU64::new(0), + } + } + + /// Start timing a call; the returned scope records it when dropped. + pub fn scope(&'static self) -> Scope { + Scope { + counter: self, + start: Instant::now(), + } + } + + /// Read and reset the accumulated `(calls, nanos)`. + pub fn take(&self) -> (u64, u64) { + ( + self.calls.swap(0, Ordering::Relaxed), + self.nanos.swap(0, Ordering::Relaxed), + ) + } + } + + /// Records one call into its counter on drop, including on unwind. + #[derive(Debug)] + pub struct Scope { + counter: &'static Counter, + start: Instant, + } + + impl Drop for Scope { + fn drop(&mut self) { + self.counter.calls.fetch_add(1, Ordering::Relaxed); + self.counter.nanos.fetch_add( + u64::try_from(self.start.elapsed().as_nanos()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + } + + pub static RESIDUAL: Counter = Counter::new(); + pub static JAC_ASSEMBLE: Counter = Counter::new(); + pub static JAC_MUL: Counter = Counter::new(); + pub static RHS_EVAL: Counter = Counter::new(); + + /// Every profiled entry point, in report order. + pub const ALL: [(&str, &Counter); 4] = [ + ("residual", &RESIDUAL), + ("jac_assemble", &JAC_ASSEMBLE), + ("jac_mul", &JAC_MUL), + ("rhs_eval", &RHS_EVAL), + ]; +} + +/// Print accumulated FFI profiling statistics and reset all counters. +/// +/// Only available when compiled with `--features profile`. +/// +/// # Safety +/// +/// This function has no pointer arguments and is always safe to call. +#[cfg(feature = "profile")] +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_profile_report() { + eprintln!("=== pybamm-core FFI profile ==="); + for (name, counter) in profile::ALL { + let (calls, nanos) = counter.take(); + eprintln!( + " {name:<13} {calls:>8} calls, {:.3} ms total", + nanos as f64 / 1_000_000.0 + ); + } + eprintln!("==============================="); +} + +/// Evaluate residual: r = M*y' - f(t,y) +/// +/// This computes the DAE residual function for IDAKLU. For a system +/// M*y' = f(t,y), the residual is M*y' - f(t,y). +/// +/// # Safety +/// +/// - `y`, `yp`, `r`, and `user_data` must be valid and non-null +/// - `y` and `yp` must point to arrays of at least `model.n_states()` elements +/// - `r` must point to an array of at least `model.output_len()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_residual( + t: f64, + y: *const f64, + yp: *const f64, + inputs: *const f64, + r: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), yp.cast(), r.cast()], + inputs, + |model, inputs| { + profile_scope!(RESIDUAL); + let (n_states, n_out) = (model.n_states(), model.output_len()); + model.eval_residual( + t, + borrow_slice(y, n_states), + borrow_slice(yp, n_states), + inputs, + borrow_slice_mut(r, n_out), + ); + SUCCESS + }, + ) + } +} + +/// Compute Jacobian-vector product: (df/dy - cj*M) @ v +/// +/// This computes the matrix-vector product needed for Newton iteration +/// in the DAE solver. The Jacobian is J = df/dy - cj*M where cj is a +/// scalar coefficient provided by the solver. +/// +/// # Safety +/// +/// - `y`, `v`, `jv`, and `user_data` must be valid and non-null +/// - `y`, `v`, and `jv` must point to arrays of at least `model.n_states()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_jac_mul( + t: f64, + y: *const f64, + inputs: *const f64, + cj: f64, + v: *const f64, + jv: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), v.cast(), jv.cast()], + inputs, + |model, inputs| { + profile_scope!(JAC_MUL); + let n_states = model.n_states(); + model.set_cj(cj); + model.jac_mul( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice(v, n_states), + borrow_slice_mut(jv, n_states), + ); + SUCCESS + }, + ) + } +} + +/// Evaluate the right-hand side f(t, y). +/// +/// This computes f(t, y) for the DAE system M*y' = f(t, y). +/// +/// # Safety +/// +/// - `y`, `f_out`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `f_out` must point to an array of at least `model.output_len()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_eval_rhs( + t: f64, + y: *const f64, + inputs: *const f64, + f_out: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), f_out.cast()], + inputs, + |model, inputs| { + profile_scope!(RHS_EVAL); + let (n_states, n_out) = (model.n_states(), model.output_len()); + model.eval_rhs( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(f_out, n_out), + ); + SUCCESS + }, + ) + } +} + +/// Get the number of states in the model. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of states on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_n_states(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.n_states() as c_int) } +} + +/// Get the number of input parameters in the model. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of input parameters on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_n_inputs(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.n_params() as c_int) } +} + +/// Write algebraic-state IDs into the provided buffer using IDA's convention: +/// `1.0` for differential, `0.0` for algebraic. +/// +/// # Safety +/// +/// - `ids_out` must point to an array of at least `model.n_states()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_algebraic_ids( + ids_out: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model(user_data, [ids_out.cast()], |model| { + model.algebraic_ids_f64(borrow_slice_mut(ids_out, model.n_states())); + SUCCESS + }) + } +} + +/// Assemble the Jacobian matrix into a pre-allocated CSC data buffer. +/// +/// This is the zero-allocation callback for IDAKLU integration. The Jacobian +/// is computed as `J = df/dy - cj * M` and stored in CSC format. +/// +/// # Safety +/// +/// - `y`, `jac_data`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `jac_data` must point to an array of at least `model.nnz()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_jac_assemble( + t: f64, + y: *const f64, + inputs: *const f64, + cj: f64, + jac_data: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), jac_data.cast()], + inputs, + |model, inputs| { + profile_scope!(JAC_ASSEMBLE); + let (n_states, nnz) = (model.n_states(), model.nnz()); + model.set_cj(cj); + model.assemble_jacobian_csc_into( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(jac_data, nnz), + ); + SUCCESS + }, + ) + } +} + +/// Get the number of non-zeros in the Jacobian. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of non-zeros on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_jac_nnz(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.nnz() as c_int) } +} + +/// Compute pure Jacobian-vector product: df/dy @ v (no mass term) +/// +/// This matches pybammsolvers ABI where mass is subtracted separately. +/// +/// # Safety +/// +/// - `y`, `v`, `jv`, and `user_data` must be valid and non-null +/// - `y`, `v`, and `jv` must point to arrays of at least `model.n_states()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_jac_action( + t: f64, + y: *const f64, + inputs: *const f64, + v: *const f64, + jv: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), v.cast(), jv.cast()], + inputs, + |model, inputs| { + let n_states = model.n_states(); + model.jac_action( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice(v, n_states), + borrow_slice_mut(jv, n_states), + ); + SUCCESS + }, + ) + } +} + +/// Compute mass matrix action: M @ v +/// +/// For identity mass matrix (ODE prototype), this copies v to mv. +/// +/// # Safety +/// +/// - All pointers must be valid and non-null +/// - `v` and `mv` must point to arrays of at least `model.n_states()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_mass_action( + v: *const f64, + mv: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model(user_data, [v.cast(), mv.cast()], |model| { + let n_states = model.n_states(); + model.mass_action(borrow_slice(v, n_states), borrow_slice_mut(mv, n_states)); + SUCCESS + }) + } +} + +/// Copy the CSC column pointers to a pre-allocated buffer. +/// +/// The buffer must have length `n_states + 1`. +/// +/// # Safety +/// +/// - `colptr` must point to an array of at least `model.n_states() + 1` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_jac_csc_colptr( + colptr: *mut i64, + user_data: *const c_void, +) -> c_int { + unsafe { + with_model(user_data, [colptr.cast()], |model| { + let csc = model.csc_sparsity(); + let out = borrow_slice_mut(colptr, csc.ncols + 1); + for (i, &val) in csc.colptr.iter().enumerate() { + out[i] = val as i64; + } + SUCCESS + }) + } +} + +/// Copy the CSC row indices to a pre-allocated buffer. +/// +/// The buffer must have length `nnz`. +/// +/// # Safety +/// +/// - `rowind` must point to an array of at least `model.nnz()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any pointer is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_jac_csc_rowind( + rowind: *mut i64, + user_data: *const c_void, +) -> c_int { + unsafe { + with_model(user_data, [rowind.cast()], |model| { + let csc = model.csc_sparsity(); + let out = borrow_slice_mut(rowind, csc.nnz()); + for (i, &val) in csc.rowind.iter().enumerate() { + out[i] = val as i64; + } + SUCCESS + }) + } +} + +/// Get the number of forward-sensitivity parameters configured on the model. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of sensitivity parameters on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_n_sens_params(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.n_sens_params() as c_int) } +} + +/// Evaluate `df/dp_i` for a single sensitivity parameter into `df_dp`. +/// +/// # Safety +/// +/// - `y`, `df_dp`, `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `df_dp` must point to an array of at least `model.n_states()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_INVALID_PARAM` (-3) if `param_idx` is out of range +/// - `ERROR_NO_SENS` (-6) if the model has no sensitivities +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_sign_loss)] // param_idx >= 0 is checked before the cast +pub unsafe extern "C" fn pybamm_rust_sens_eval( + t: f64, + y: *const f64, + inputs: *const f64, + param_idx: c_int, + df_dp: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), df_dp.cast()], + inputs, + |model, inputs| { + if model.n_sens_params() == 0 { + return ERROR_NO_SENS; + } + if param_idx < 0 || (param_idx as usize) >= model.n_sens_params() { + return ERROR_INVALID_PARAM; + } + let n_states = model.n_states(); + model.eval_sens( + t, + borrow_slice(y, n_states), + inputs, + param_idx as usize, + borrow_slice_mut(df_dp, n_states), + ); + SUCCESS + }, + ) + } +} + +/// Evaluate `df/dp` for all configured sensitivity parameters at once. +/// +/// Layout of `df_dp_out`: `df_dp_out[i*n_states + j] = ∂f_j/∂p_i`. +/// +/// # Safety +/// +/// Same buffer-size and pointer requirements as [`pybamm_rust_sens_eval`], with +/// `df_dp_out` length at least `n_sens_params * n_states`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_sens_eval_all( + t: f64, + y: *const f64, + inputs: *const f64, + df_dp_out: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), df_dp_out.cast()], + inputs, + |model, inputs| { + if model.n_sens_params() == 0 { + return ERROR_NO_SENS; + } + let n_states = model.n_states(); + let n_sens = model.n_sens_params(); + model.eval_sens_all( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(df_dp_out, n_sens * n_states), + ); + SUCCESS + }, + ) + } +} + +/// Project state sensitivities onto output sensitivities for all outputs and +/// all configured sensitivity parameters. +/// +/// `dvar/dp_k = dH/dp . e_k + dH/dy . y_sens_k`. +/// +/// Layout of `y_sens`: `y_sens[k*n_states + j]` (len `n_sens_params * n_states`). +/// Layout of `out`: `out[k*n_out + o]` (len `n_sens_params * total_output_len`). +/// +/// # Safety +/// +/// - `y`, `y_sens`, `out`, `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `y_sens` must point to an array of at least `n_sens_params * n_states` elements +/// - `out` must point to an array of at least `n_sens_params * total_output_len` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_NO_OUTPUTS` (-7) if the model has no output variables +/// - `ERROR_NO_SENS` (-6) if the model has no sensitivities +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_output_sens_project( + t: f64, + y: *const f64, + inputs: *const f64, + y_sens: *const f64, + out: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), y_sens.cast(), out.cast()], + inputs, + |model, inputs| { + if model.n_outputs() == 0 { + return ERROR_NO_OUTPUTS; + } + if model.n_sens_params() == 0 { + return ERROR_NO_SENS; + } + let n_states = model.n_states(); + let (n_sens, n_out) = (model.n_sens_params(), model.total_output_len()); + model.output_sens_project( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice(y_sens, n_sens * n_states), + borrow_slice_mut(out, n_sens * n_out), + ); + SUCCESS + }, + ) + } +} + +/// Get the number of compiled output-variable expressions on the model. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of output variables on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_n_outputs(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.n_outputs() as c_int) } +} + +/// Get the length of output variable `var_idx`. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Output length on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_INVALID_OUTPUT` (-4) if `var_idx` is out of range +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow( + clippy::cast_possible_truncation, + clippy::cast_possible_wrap, + clippy::cast_sign_loss +)] +pub unsafe extern "C" fn pybamm_rust_output_len(user_data: *const c_void, var_idx: c_int) -> c_int { + unsafe { + with_model(user_data, [], |model| { + if var_idx < 0 || (var_idx as usize) >= model.n_outputs() { + return ERROR_INVALID_OUTPUT; + } + model.output_len_at(var_idx as usize) as c_int + }) + } +} + +/// Evaluate output variable `var_idx` into `out`. +/// +/// # Safety +/// +/// - `y`, `out`, `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `out` must point to an array of at least `output_len_at(var_idx)` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `out_len` may be null; if non-null it receives the count of values written +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_INVALID_OUTPUT` (-4) if `var_idx` is out of range +/// - `ERROR_NO_OUTPUTS` (-7) if the model has no output variables +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow( + clippy::cast_sign_loss, + clippy::cast_possible_truncation, + clippy::cast_possible_wrap +)] +pub unsafe extern "C" fn pybamm_rust_output_eval( + t: f64, + y: *const f64, + inputs: *const f64, + var_idx: c_int, + out: *mut f64, + out_len: *mut c_int, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), out.cast()], + inputs, + |model, inputs| { + if model.n_outputs() == 0 { + return ERROR_NO_OUTPUTS; + } + if var_idx < 0 || (var_idx as usize) >= model.n_outputs() { + return ERROR_INVALID_OUTPUT; + } + let var = var_idx as usize; + let n_states = model.n_states(); + let len = model.output_len_at(var); + let written = model.eval_output( + t, + borrow_slice(y, n_states), + inputs, + var, + borrow_slice_mut(out, len), + ); + if !out_len.is_null() { + *out_len = written as c_int; + } + SUCCESS + }, + ) + } +} + +/// Batch-evaluate every output variable over `n_points` trajectory points. +/// +/// Amortises interpreter dispatch across the batch instead of walking every +/// output tape once per point; results match per-point evaluation bitwise. +/// +/// # Safety +/// +/// - `ts`, `ys`, `out`, `user_data` must be valid and non-null +/// - `ts` must point to `n_points` times +/// - `ys` must point to `n_points * n_states` states, each point contiguous +/// (`(n_states, n_points)` F-contiguous) +/// - `out` must point to `n_points * total_output_len` elements and is written +/// with each point's stacked outputs contiguous +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_INVALID_PARAM` (-3) if `n_points` is not positive +/// - `ERROR_NO_OUTPUTS` (-7) if the model has no output variables +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_sign_loss)] +pub unsafe extern "C" fn pybamm_rust_output_eval_batch( + ts: *const f64, + ys: *const f64, + n_points: c_int, + inputs: *const f64, + out: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [ts.cast(), ys.cast(), out.cast()], + inputs, + |model, inputs| { + if n_points <= 0 { + return ERROR_INVALID_PARAM; + } + if model.n_outputs() == 0 { + return ERROR_NO_OUTPUTS; + } + let points = n_points as usize; + let n_states = model.n_states(); + let total = model.total_output_len(); + model.eval_outputs_batch( + points, + borrow_slice(ts, points), + borrow_slice(ys, points * n_states), + inputs, + borrow_slice_mut(out, points * total), + ); + SUCCESS + }, + ) + } +} + +/// Evaluate algebraic residual g(t, y). +/// +/// # Safety +/// +/// - `y`, `output`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `output` must point to an array of at least `model.n_algebraic()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_NO_ALG` (-8) if the model has no algebraic block +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_alg_res( + t: f64, + y: *const f64, + inputs: *const f64, + output: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), output.cast()], + inputs, + |model, inputs| { + if !model.has_algebraic() { + return ERROR_NO_ALG; + } + let n_states = model.n_states(); + let n_algebraic = model.n_algebraic(); + model.eval_algebraic_residual( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(output, n_algebraic), + ); + SUCCESS + }, + ) + } +} + +/// Assemble the algebraic Jacobian `dg/dy_alg`. +/// +/// # Safety +/// +/// - `y`, `output`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `output` must point to an array of at least `model.algebraic_jacobian_nnz()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_NO_ALG` (-8) if the model has no algebraic block +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_alg_jac_assemble( + t: f64, + y: *const f64, + inputs: *const f64, + output: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), output.cast()], + inputs, + |model, inputs| { + if !model.has_algebraic() { + return ERROR_NO_ALG; + } + let n_states = model.n_states(); + let n_alg_jac = model.algebraic_jacobian_nnz(); + model.assemble_algebraic_jacobian_into( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(output, n_alg_jac), + ); + SUCCESS + }, + ) + } +} + +/// Compute algebraic Jacobian-vector product `(dg/dy_alg) @ v`. +/// +/// # Safety +/// +/// - `y`, `v`, `jv`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `v` and `jv` must point to arrays of at least `model.n_algebraic()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_NO_ALG` (-8) if the model has no algebraic block +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_alg_jac_action( + t: f64, + y: *const f64, + inputs: *const f64, + v: *const f64, + jv: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), v.cast(), jv.cast()], + inputs, + |model, inputs| { + if !model.has_algebraic() { + return ERROR_NO_ALG; + } + let n_states = model.n_states(); + let n_algebraic = model.n_algebraic(); + model.eval_algebraic_jacobian_action( + t, + borrow_slice(y, n_states), + inputs, + borrow_slice(v, n_algebraic), + borrow_slice_mut(jv, n_algebraic), + ); + SUCCESS + }, + ) + } +} + +/// Get the number of compiled event expressions on the model. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Number of events on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_n_events(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.n_events() as c_int) } +} + +/// Get the total length of all events concatenated. +/// +/// # Safety +/// +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - Total event length on success (>= 0) +/// - `ERROR_NULL_POINTER` (-1) if `user_data` is null +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +pub unsafe extern "C" fn pybamm_rust_total_event_len(user_data: *const c_void) -> c_int { + unsafe { with_model(user_data, [], |model| model.total_event_len() as c_int) } +} + +/// Evaluate all events and write concatenated results into `output`. +/// +/// This is the primary interface for DAE solvers that need all event values +/// at once for root-finding. Events are evaluated in order and their results +/// are concatenated into the output buffer. +/// +/// # Safety +/// +/// - `y`, `output`, and `user_data` must be valid and non-null +/// - `y` must point to an array of at least `model.n_states()` elements +/// - `output` must point to an array of at least `model.total_event_len()` elements +/// - `inputs` may be null when the model has no input parameters; otherwise it +/// must point to an array of at least `model.n_params()` elements +/// - `user_data` must point to a valid `ModelEvaluator` instance +/// +/// # Returns +/// +/// - `SUCCESS` (0) on success +/// - `ERROR_NULL_POINTER` (-1) if any required pointer is null +/// - `ERROR_NO_EVENTS` (-9) if the model has no events +/// - `ERROR_PANIC` (-2) if a Rust panic occurred +#[unsafe(no_mangle)] +pub unsafe extern "C" fn pybamm_rust_events_eval( + t: f64, + y: *const f64, + inputs: *const f64, + output: *mut f64, + user_data: *mut c_void, +) -> c_int { + unsafe { + with_model_and_inputs( + user_data, + [y.cast(), output.cast()], + inputs, + |model, inputs| { + if model.n_events() == 0 { + return ERROR_NO_EVENTS; + } + let n_states = model.n_states(); + let total_len = model.total_event_len(); + model.eval_observables( + ObservableKind::Events, + t, + borrow_slice(y, n_states), + inputs, + borrow_slice_mut(output, total_len), + ); + SUCCESS + }, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::Arena; + use crate::node::{CsrData, Node, Shape}; + + #[test] + fn abi_version_is_stable_and_positive() { + // Bumping the const is deliberate, paired with a PYBAMM_RUST_ABI_VERSION + // bump in the C++ consumer. + assert_eq!(pybamm_rust_abi_version(), RUST_ABI_VERSION); + assert_eq!(RUST_ABI_VERSION, 1); + } + + /// Create an identity mass matrix of size n. + fn identity_mass_matrix(n: usize) -> CsrData { + CsrData { + shape: Shape::matrix(n, n), + indptr: (0..=n).collect(), + indices: (0..n).collect(), + data: vec![1.0; n], + } + } + + /// Helper to create a simple test model: f(y) = 2*y with identity mass matrix + fn create_test_model() -> Box { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); // f(y) = 2*y + + let mass = identity_mass_matrix(2); + Box::new(ModelEvaluator::new(&arena, rhs, mass, 2, 0)) + } + + fn create_algebraic_test_model() -> Box { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let y2 = arena.alloc(Node::Index { + child: y_full, + start: 2, + end: 3, + }); + let rhs = y_full; + let alg0_mul = arena.alloc(Node::Mul(y1, y2)); + let alg0 = arena.alloc(Node::Add(y0, alg0_mul)); + let two = arena.alloc(Node::Scalar(2.0)); + let y2_sq = arena.alloc(Node::Pow(y2, two)); + let alg1 = arena.alloc(Node::Add(y1, y2_sq)); + let alg = arena.alloc(Node::Concat(vec![alg0, alg1])); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 1], + indices: vec![0], + data: vec![1.0], + }; + + Box::new(ModelEvaluator::new_with_algebraic( + &arena, + rhs, + mass, + 3, + 0, + Some(alg), + &[1, 2], + )) + } + + #[test] + fn test_rust_residual_null_check() { + // All nulls + let result = unsafe { + pybamm_rust_residual( + 0.0, + std::ptr::null(), + std::ptr::null(), + std::ptr::null::(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + + // y is null + let yp = [1.0, 2.0]; + let mut r = [0.0, 0.0]; + let result = unsafe { + pybamm_rust_residual( + 0.0, + std::ptr::null(), + yp.as_ptr(), + std::ptr::null::(), + r.as_mut_ptr(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + + // user_data is null (but other pointers valid) + let y = [1.0, 2.0]; + let result = unsafe { + pybamm_rust_residual( + 0.0, + y.as_ptr(), + yp.as_ptr(), + std::ptr::null::(), + r.as_mut_ptr(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_residual_success() { + let mut model = create_test_model(); + + // f(y) = 2*y, M = I: residual = M*y' - f(y) = [3, 4] - [2, 4] = [1, 0] + let y = [1.0, 2.0]; + let yp = [3.0, 4.0]; + let mut r = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_residual( + 0.0, + y.as_ptr(), + yp.as_ptr(), + std::ptr::null::(), + r.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!( + (r[0] - 1.0).abs() < 1e-14, + "Expected r[0] = 1.0, got {}", + r[0] + ); + assert!(r[1].abs() < 1e-14, "Expected r[1] = 0.0, got {}", r[1]); + } + + #[test] + fn test_rust_jac_mul_null_check() { + let result = unsafe { + pybamm_rust_jac_mul( + 0.0, + std::ptr::null(), + std::ptr::null::(), + 1.0, + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_jac_mul_success() { + let mut model = create_test_model(); + + // f(y) = 2*y, so df/dy = 2*I + // With cj=0.5 and M=I: (df/dy - cj*M) @ v = (2*I - 0.5*I) @ v = 1.5*v + let y = [1.0, 2.0]; + let v = [1.0, 0.0]; + let mut jv = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_jac_mul( + 0.0, + y.as_ptr(), + std::ptr::null::(), + 0.5, + v.as_ptr(), + jv.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + // (2 - 0.5) * [1, 0] = [1.5, 0] + assert!( + (jv[0] - 1.5).abs() < 1e-14, + "Expected jv[0] = 1.5, got {}", + jv[0] + ); + assert!(jv[1].abs() < 1e-14, "Expected jv[1] = 0.0, got {}", jv[1]); + } + + #[test] + fn test_rust_jac_mul_uses_inputs() { + // f(y)=k*y (k an input); (df/dy - cj*M)@v = (k-cj)*v. k=7, cj=0.5, v=[1,0] -> [6.5, 0]. + // If inputs were dropped (k read as 0), the result would instead be [-0.5, 0]. + let mut model = create_sens_test_model(); + let y = [3.0, 4.0]; + let inputs = [7.0]; + let v = [1.0, 0.0]; + let mut jv = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_jac_mul( + 0.0, + y.as_ptr(), + inputs.as_ptr(), + 0.5, + v.as_ptr(), + jv.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!( + (jv[0] - 6.5).abs() < 1e-12, + "expected jv[0]=6.5, got {}", + jv[0] + ); + assert!(jv[1].abs() < 1e-12, "expected jv[1]=0.0, got {}", jv[1]); + } + + #[test] + fn test_rust_eval_rhs_null_check() { + let result = unsafe { + pybamm_rust_eval_rhs( + 0.0, + std::ptr::null(), + std::ptr::null::(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_eval_rhs_success() { + let mut model = create_test_model(); + + // f(y) = 2*y + // y = [1, 2] -> f(y) = [2, 4] + let y = [1.0, 2.0]; + let mut f_out = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_eval_rhs( + 0.0, + y.as_ptr(), + std::ptr::null::(), + f_out.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!( + (f_out[0] - 2.0).abs() < 1e-14, + "Expected f_out[0] = 2.0, got {}", + f_out[0] + ); + assert!( + (f_out[1] - 4.0).abs() < 1e-14, + "Expected f_out[1] = 4.0, got {}", + f_out[1] + ); + } + + #[test] + fn test_rust_alg_res_no_algebraic_block() { + let mut model = create_test_model(); + let y = [1.0, 2.0]; + let mut out = [0.0, 0.0]; + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_alg_res( + 0.0, + y.as_ptr(), + std::ptr::null::(), + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, ERROR_NO_ALG); + } + + #[test] + fn test_rust_alg_res_success() { + let mut model = create_algebraic_test_model(); + let y = [10.0, 2.0, 3.0]; + let mut out = [0.0, 0.0]; + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_alg_res( + 0.0, + y.as_ptr(), + std::ptr::null::(), + out.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!((out[0] - 16.0).abs() < 1e-12, "expected 16, got {}", out[0]); + assert!((out[1] - 11.0).abs() < 1e-12, "expected 11, got {}", out[1]); + } + + #[test] + fn test_rust_alg_jac_assemble_success() { + let mut model = create_algebraic_test_model(); + let y = [10.0, 2.0, 3.0]; + let mut jac = vec![0.0; model.algebraic_jacobian_nnz()]; + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_alg_jac_assemble( + 0.0, + y.as_ptr(), + std::ptr::null::(), + jac.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + // CSC order, matching the (rows, cols) the model publishes beside it. + assert_eq!(jac, vec![3.0, 1.0, 2.0, 6.0]); + } + + #[test] + fn test_rust_alg_jac_action_success() { + let mut model = create_algebraic_test_model(); + let y = [10.0, 2.0, 3.0]; + let v = [5.0, 7.0]; + let mut jv = [0.0, 0.0]; + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_alg_jac_action( + 0.0, + y.as_ptr(), + std::ptr::null::(), + v.as_ptr(), + jv.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!((jv[0] - 29.0).abs() < 1e-12, "expected 29, got {}", jv[0]); + assert!((jv[1] - 47.0).abs() < 1e-12, "expected 47, got {}", jv[1]); + } + + #[test] + fn test_rust_n_states_null_check() { + let result = unsafe { pybamm_rust_n_states(std::ptr::null()) }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_n_states_success() { + let model = create_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_n_states(user_data) }; + assert_eq!(result, 2); + } + + #[test] + fn test_rust_jac_assemble_null_check() { + let result = unsafe { + pybamm_rust_jac_assemble( + 0.0, + std::ptr::null(), + std::ptr::null::(), + 1.0, + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_jac_assemble_success() { + let mut model = create_test_model(); + + // f(y) = 2*y, df/dy = 2*I, M = I + // With cj=0.5: J = df/dy - cj*M = 2*I - 0.5*I = 1.5*I + let y = [1.0, 2.0]; + let nnz = model.nnz(); + let mut jac_data = vec![0.0; nnz]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_jac_assemble( + 0.0, + y.as_ptr(), + std::ptr::null::(), + 0.5, + jac_data.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + // Diagonal entries should be 1.5 + for &val in &jac_data { + assert!((val - 1.5).abs() < 1e-14, "Expected 1.5, got {val}"); + } + } + + #[test] + fn test_rust_jac_nnz_null_check() { + let result = unsafe { pybamm_rust_jac_nnz(std::ptr::null()) }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_jac_nnz_success() { + let model = create_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_jac_nnz(user_data) }; + // f(y) = 2*y has diagonal Jacobian, so nnz = n_states = 2 + assert_eq!(result, 2); + } + + #[test] + fn test_rust_jac_csc_colptr_success() { + let model = create_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + + // For 2x2 diagonal matrix, colptr should be [0, 1, 2] + let mut colptr = [0i64; 3]; + let result = unsafe { pybamm_rust_jac_csc_colptr(colptr.as_mut_ptr(), user_data) }; + + assert_eq!(result, SUCCESS); + assert_eq!(colptr, [0, 1, 2]); + } + + #[test] + fn test_rust_jac_csc_rowind_success() { + let model = create_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + + // For 2x2 diagonal matrix, rowind should be [0, 1] + let mut rowind = [0i64; 2]; + let result = unsafe { pybamm_rust_jac_csc_rowind(rowind.as_mut_ptr(), user_data) }; + + assert_eq!(result, SUCCESS); + assert_eq!(rowind, [0, 1]); + } + + #[test] + fn test_rust_jac_action_null_check() { + let result = unsafe { + pybamm_rust_jac_action( + 0.0, + std::ptr::null(), + std::ptr::null::(), + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_jac_action_success() { + let mut model = create_test_model(); + + // f(y) = 2*y, so df/dy = 2*I + // jac_action(v) = df/dy @ v = 2*v (no mass term) + let y = [1.0, 2.0]; + let v = [1.0, 0.0]; + let mut jv = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_jac_action( + 0.0, + y.as_ptr(), + std::ptr::null::(), + v.as_ptr(), + jv.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!( + (jv[0] - 2.0).abs() < 1e-14, + "Expected jv[0] = 2.0, got {}", + jv[0] + ); + assert!(jv[1].abs() < 1e-14, "Expected jv[1] = 0.0, got {}", jv[1]); + } + + #[test] + fn test_rust_mass_action_null_check() { + let result = unsafe { + pybamm_rust_mass_action(std::ptr::null(), std::ptr::null_mut(), std::ptr::null_mut()) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_mass_action_success() { + let model = create_test_model(); + + // Identity mass: M @ v = v + let v = [3.0, 4.0]; + let mut mv = [0.0, 0.0]; + + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = + unsafe { pybamm_rust_mass_action(v.as_ptr(), mv.as_mut_ptr(), user_data.cast_mut()) }; + + assert_eq!(result, SUCCESS); + // Use approx comparison for floats instead of assert_eq! + assert!( + (mv[0] - v[0]).abs() < 1e-14, + "Expected mv[0] = {}, got {}", + v[0], + mv[0] + ); + assert!( + (mv[1] - v[1]).abs() < 1e-14, + "Expected mv[1] = {}, got {}", + v[1], + mv[1] + ); + } + + #[test] + fn test_rust_jac_action_preserves_cj() { + // Verify that jac_action doesn't permanently change cj + let mut model = create_test_model(); + + // Set a non-zero cj + model.set_cj(0.5); + + let y = [1.0, 2.0]; + let v = [1.0, 0.0]; + let mut jv = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_jac_action( + 0.0, + y.as_ptr(), + std::ptr::null::(), + v.as_ptr(), + jv.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + // jac_action should still give df/dy @ v = 2*v (no mass) + assert!( + (jv[0] - 2.0).abs() < 1e-14, + "Expected jv[0] = 2.0, got {}", + jv[0] + ); + + // cj should be restored to 0.5 + assert!( + (model.cj() - 0.5).abs() < 1e-14, + "Expected cj = 0.5, got {}", + model.cj() + ); + } + + #[test] + fn test_rust_eval_rhs_with_inputs() { + // f(y) = k*y where k is an input parameter at index 0 + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y)); + let mass = identity_mass_matrix(2); + let mut model = Box::new(ModelEvaluator::new(&arena, rhs, mass, 2, 1)); + + // y = [1, 2], k = 3 -> f = [3, 6] + let y = [1.0, 2.0]; + let inputs = [3.0]; + let mut f_out = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_eval_rhs( + 0.0, + y.as_ptr(), + inputs.as_ptr(), + f_out.as_mut_ptr(), + user_data, + ) + }; + + assert_eq!(result, SUCCESS); + assert!((f_out[0] - 3.0).abs() < 1e-14); + assert!((f_out[1] - 6.0).abs() < 1e-14); + } + + #[test] + fn test_rust_eval_rhs_null_inputs_when_no_params() { + // When n_inputs == 0, inputs may be null + let mut model = create_test_model(); + let y = [1.0, 2.0]; + let mut f_out = [0.0, 0.0]; + + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let result = unsafe { + pybamm_rust_eval_rhs( + 0.0, + y.as_ptr(), + std::ptr::null(), + f_out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!((f_out[0] - 2.0).abs() < 1e-14); + } + + #[test] + fn test_rust_n_inputs() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y)); + let mass = identity_mass_matrix(2); + let model = Box::new(ModelEvaluator::new(&arena, rhs, mass, 2, 1)); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_n_inputs(user_data) }; + assert_eq!(result, 1); + } + + #[test] + fn test_rust_algebraic_ids_dae() { + // Mass with row 1 missing diagonal -> row 1 is algebraic. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 2], + indices: vec![0, 2], + data: vec![1.0, 1.0], + }; + let mut model = Box::new(ModelEvaluator::new(&arena, y, mass, 3, 0)); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let mut ids = [0.0f64; 3]; + let result = unsafe { pybamm_rust_algebraic_ids(ids.as_mut_ptr(), user_data) }; + assert_eq!(result, SUCCESS); + // 1.0 = differential, 0.0 = algebraic; values are exact (no FP arithmetic). + assert!((ids[0] - 1.0).abs() < f64::EPSILON); + assert!(ids[1].abs() < f64::EPSILON); + assert!((ids[2] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_rust_algebraic_ids_null_check() { + let result = + unsafe { pybamm_rust_algebraic_ids(std::ptr::null_mut(), std::ptr::null_mut()) }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + fn create_sens_test_model() -> Box { + // f(y) = k * y with k as InputParameter index 0; sensitivity w.r.t. k. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y)); + let mass = identity_mass_matrix(2); + Box::new(ModelEvaluator::new_with_sens(&arena, rhs, mass, 2, 1, &[0])) + } + + #[test] + fn test_rust_n_sens_params() { + let model = create_sens_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + assert_eq!(unsafe { pybamm_rust_n_sens_params(user_data) }, 1); + } + + #[test] + fn test_rust_sens_eval_basic() { + // f(y) = k*y, ∂f/∂k = y. At y=[3, 4] -> [3, 4]. + let mut model = create_sens_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [3.0, 4.0]; + let inputs = [7.0]; + let mut out = [0.0; 2]; + let result = unsafe { + pybamm_rust_sens_eval( + 0.0, + y.as_ptr(), + inputs.as_ptr(), + 0, + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!((out[0] - 3.0).abs() < 1e-12); + assert!((out[1] - 4.0).abs() < 1e-12); + } + + #[test] + fn test_rust_sens_eval_no_sensitivities() { + // Model has 0 sens params -> any param_idx returns ERROR_NO_SENS. + let mut model = create_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [1.0, 2.0]; + let mut out = [0.0; 2]; + let result = unsafe { + pybamm_rust_sens_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + 0, + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, ERROR_NO_SENS); + } + + #[test] + fn test_rust_sens_eval_invalid_index() { + let mut model = create_sens_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [3.0, 4.0]; + let inputs = [7.0]; + let mut out = [0.0; 2]; + let result = unsafe { + pybamm_rust_sens_eval( + 0.0, + y.as_ptr(), + inputs.as_ptr(), + 5, // out of range + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, ERROR_INVALID_PARAM); + } + + #[test] + fn test_rust_sens_eval_all_basic() { + let mut model = create_sens_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [3.0, 4.0]; + let inputs = [7.0]; + let mut out = [0.0; 2]; // n_sens_params=1, n_states=2 -> 2 entries + let result = unsafe { + pybamm_rust_sens_eval_all( + 0.0, + y.as_ptr(), + inputs.as_ptr(), + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!((out[0] - 3.0).abs() < 1e-12); + assert!((out[1] - 4.0).abs() < 1e-12); + } + + fn create_output_test_model() -> Box { + // rhs: f(y) = y; output: 2 * y[0] + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let var0 = arena.alloc(Node::Mul(two, y0)); + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_output(&arena, var0); + Box::new(model) + } + + #[test] + fn test_rust_n_outputs() { + let model = create_output_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + assert_eq!(unsafe { pybamm_rust_n_outputs(user_data) }, 1); + } + + #[test] + fn test_rust_output_len() { + let model = create_output_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + assert_eq!(unsafe { pybamm_rust_output_len(user_data, 0) }, 1); + assert_eq!( + unsafe { pybamm_rust_output_len(user_data, 5) }, + ERROR_INVALID_OUTPUT + ); + } + + #[test] + fn test_rust_output_eval_basic() { + let mut model = create_output_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [3.0, 4.0]; + let mut out = [0.0; 1]; + let mut out_len: c_int = 0; + let result = unsafe { + pybamm_rust_output_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + 0, + out.as_mut_ptr(), + &raw mut out_len, + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert_eq!(out_len, 1); + assert!((out[0] - 6.0).abs() < 1e-12); + } + + #[test] + fn test_rust_output_eval_no_outputs() { + let mut model = create_test_model(); // no outputs configured + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [1.0, 2.0]; + let mut out = [0.0]; + let mut out_len: c_int = 0; + let result = unsafe { + pybamm_rust_output_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + 0, + out.as_mut_ptr(), + &raw mut out_len, + user_data, + ) + }; + assert_eq!(result, ERROR_NO_OUTPUTS); + } + + fn create_event_test_model() -> Box { + // rhs: f(y) = y; event: y[0] - 0.5 + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let thresh = arena.alloc(Node::Scalar(0.5)); + let event_expr = arena.alloc(Node::Sub(y0, thresh)); + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_event(&arena, event_expr); + Box::new(model) + } + + #[test] + fn test_rust_n_events_null_check() { + let result = unsafe { pybamm_rust_n_events(std::ptr::null()) }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_n_events_success() { + let model = create_event_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_n_events(user_data) }; + assert_eq!(result, 1); + } + + #[test] + fn test_rust_n_events_zero_when_no_events() { + let model = create_test_model(); // no events + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_n_events(user_data) }; + assert_eq!(result, 0); + } + + #[test] + fn test_rust_total_event_len_null_check() { + let result = unsafe { pybamm_rust_total_event_len(std::ptr::null()) }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_total_event_len_success() { + let model = create_event_test_model(); + let user_data: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let result = unsafe { pybamm_rust_total_event_len(user_data) }; + assert_eq!(result, 1); + } + + #[test] + fn test_rust_events_eval_null_check() { + let result = unsafe { + pybamm_rust_events_eval( + 0.0, + std::ptr::null(), + std::ptr::null(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + assert_eq!(result, ERROR_NULL_POINTER); + } + + #[test] + fn test_rust_events_eval_no_events() { + let mut model = create_test_model(); // no events + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [1.0, 2.0]; + let mut out = [0.0]; + let result = unsafe { + pybamm_rust_events_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, ERROR_NO_EVENTS); + } + + #[test] + fn test_rust_events_eval_success() { + let mut model = create_event_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + + // event = y[0] - 0.5; y = [0.7, 1.0] => 0.7 - 0.5 = 0.2 + let y = [0.7, 1.0]; + let mut out = [0.0]; + let result = unsafe { + pybamm_rust_events_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!((out[0] - 0.2).abs() < 1e-12, "expected 0.2, got {}", out[0]); + } + + #[test] + fn test_rust_events_eval_at_threshold() { + let mut model = create_event_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + + // event = y[0] - 0.5; y = [0.5, 1.0] => 0.0 (at threshold) + let y = [0.5, 1.0]; + let mut out = [0.0]; + let result = unsafe { + pybamm_rust_events_eval( + 0.0, + y.as_ptr(), + std::ptr::null(), + out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!(out[0].abs() < 1e-12, "expected 0.0, got {}", out[0]); + } + + #[test] + fn guard_maps_a_panicking_body_to_error_panic() { + let model = create_test_model(); + let shared: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + assert_eq!( + unsafe { with_model(shared, [], |_| panic!("boom")) }, + ERROR_PANIC + ); + + let mut model = create_test_model(); + let owned: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + assert_eq!( + unsafe { with_model_and_inputs(owned, [], std::ptr::null(), |_, _| panic!("boom")) }, + ERROR_PANIC + ); + } + + #[test] + fn guard_rejects_nulls_without_running_the_body() { + let model = create_test_model(); + let shared: *const c_void = std::ptr::from_ref(model.as_ref()).cast(); + let unreachable = |_: &ModelEvaluator| unreachable!("guard let a null through"); + + assert_eq!( + unsafe { with_model(std::ptr::null(), [], unreachable) }, + ERROR_NULL_POINTER + ); + assert_eq!( + unsafe { with_model(shared, [std::ptr::null()], unreachable) }, + ERROR_NULL_POINTER + ); + assert_eq!( + unsafe { with_model(shared, [shared, std::ptr::null(), shared], unreachable) }, + ERROR_NULL_POINTER + ); + } + + /// `inputs` obeys the same null rule as every other buffer, decided against + /// the model rather than the signature: null is the caller saying "no + /// inputs", which only a model without input parameters can be told. + #[test] + fn null_inputs_are_rejected_for_a_model_that_takes_parameters() { + let mut model = create_sens_test_model(); // f(y) = k*y, one input + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [3.0, 4.0]; + let mut f_out = [0.0; 2]; + + let mut call = |inputs| unsafe { + pybamm_rust_eval_rhs(0.0, y.as_ptr(), inputs, f_out.as_mut_ptr(), user_data) + }; + assert_eq!(call(std::ptr::null()), ERROR_NULL_POINTER); + + let inputs = [3.0]; + assert_eq!(call(inputs.as_ptr()), SUCCESS); + assert!((f_out[0] - 9.0).abs() < 1e-12, "got {}", f_out[0]); + } + + #[test] + fn a_supplied_inputs_buffer_is_ignored_when_the_model_takes_none() { + let mut model = create_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let y = [1.0, 2.0]; + let stray = [99.0]; + let mut f_out = [0.0; 2]; + + let result = unsafe { + pybamm_rust_eval_rhs( + 0.0, + y.as_ptr(), + stray.as_ptr(), + f_out.as_mut_ptr(), + user_data, + ) + }; + assert_eq!(result, SUCCESS); + assert!((f_out[0] - 2.0).abs() < 1e-14, "got {}", f_out[0]); + } + + /// Null `user_data` is the one input every entry point shares, so a missing + /// guard shows up here as a segfault or a stale value instead of the code. + #[test] + fn every_entry_point_rejects_a_null_user_data() { + let nil = std::ptr::null::(); + let nil_mut = std::ptr::null_mut::(); + let nil_data = std::ptr::null_mut::(); + let checks: [(&str, c_int); 26] = [ + ("pybamm_rust_residual", unsafe { + pybamm_rust_residual(0.0, nil, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_jac_mul", unsafe { + pybamm_rust_jac_mul(0.0, nil, nil, 1.0, nil, nil_mut, nil_data) + }), + ("pybamm_rust_eval_rhs", unsafe { + pybamm_rust_eval_rhs(0.0, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_n_states", unsafe { + pybamm_rust_n_states(std::ptr::null()) + }), + ("pybamm_rust_n_inputs", unsafe { + pybamm_rust_n_inputs(std::ptr::null()) + }), + ("pybamm_rust_algebraic_ids", unsafe { + pybamm_rust_algebraic_ids(nil_mut, nil_data) + }), + ("pybamm_rust_jac_assemble", unsafe { + pybamm_rust_jac_assemble(0.0, nil, nil, 1.0, nil_mut, nil_data) + }), + ("pybamm_rust_jac_nnz", unsafe { + pybamm_rust_jac_nnz(std::ptr::null()) + }), + ("pybamm_rust_jac_action", unsafe { + pybamm_rust_jac_action(0.0, nil, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_mass_action", unsafe { + pybamm_rust_mass_action(nil, nil_mut, nil_data) + }), + ("pybamm_rust_jac_csc_colptr", unsafe { + pybamm_rust_jac_csc_colptr(std::ptr::null_mut(), std::ptr::null()) + }), + ("pybamm_rust_jac_csc_rowind", unsafe { + pybamm_rust_jac_csc_rowind(std::ptr::null_mut(), std::ptr::null()) + }), + ("pybamm_rust_n_sens_params", unsafe { + pybamm_rust_n_sens_params(std::ptr::null()) + }), + ("pybamm_rust_sens_eval", unsafe { + pybamm_rust_sens_eval(0.0, nil, nil, 0, nil_mut, nil_data) + }), + ("pybamm_rust_sens_eval_all", unsafe { + pybamm_rust_sens_eval_all(0.0, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_output_sens_project", unsafe { + pybamm_rust_output_sens_project(0.0, nil, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_n_outputs", unsafe { + pybamm_rust_n_outputs(std::ptr::null()) + }), + ("pybamm_rust_output_len", unsafe { + pybamm_rust_output_len(std::ptr::null(), 0) + }), + ("pybamm_rust_output_eval", unsafe { + pybamm_rust_output_eval(0.0, nil, nil, 0, nil_mut, std::ptr::null_mut(), nil_data) + }), + // Also pins the order: nulls outrank the `n_points > 0` rule. + ("pybamm_rust_output_eval_batch", unsafe { + pybamm_rust_output_eval_batch(nil, nil, 0, nil, nil_mut, nil_data) + }), + ("pybamm_rust_alg_res", unsafe { + pybamm_rust_alg_res(0.0, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_alg_jac_assemble", unsafe { + pybamm_rust_alg_jac_assemble(0.0, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_alg_jac_action", unsafe { + pybamm_rust_alg_jac_action(0.0, nil, nil, nil, nil_mut, nil_data) + }), + ("pybamm_rust_n_events", unsafe { + pybamm_rust_n_events(std::ptr::null()) + }), + ("pybamm_rust_total_event_len", unsafe { + pybamm_rust_total_event_len(std::ptr::null()) + }), + ("pybamm_rust_events_eval", unsafe { + pybamm_rust_events_eval(0.0, nil, nil, nil_mut, nil_data) + }), + ]; + for (name, status) in checks { + assert_eq!( + status, ERROR_NULL_POINTER, + "{name} accepted a null user_data" + ); + } + } + + /// Every buffer an entry point dereferences must be named in its guard, and + /// the model here has no sensitivities, outputs, events or algebraic block — + /// so these also pin that a null outranks the model-state error codes. + #[test] + fn every_required_buffer_of_the_solver_callbacks_is_null_checked() { + let mut model = create_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let mut scratch = [0.0f64; 4]; + let ok = scratch.as_mut_ptr(); + let ok_const: *const f64 = ok; + let nil = std::ptr::null::(); + let nil_mut = std::ptr::null_mut::(); + + let residual = |y, yp, r| unsafe { pybamm_rust_residual(0.0, y, yp, nil, r, user_data) }; + assert_eq!(residual(nil, ok_const, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(residual(ok_const, nil, ok), ERROR_NULL_POINTER, "yp"); + assert_eq!( + residual(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "r" + ); + + let jac_mul = |y, v, jv| unsafe { pybamm_rust_jac_mul(0.0, y, nil, 1.0, v, jv, user_data) }; + assert_eq!(jac_mul(nil, ok_const, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(jac_mul(ok_const, nil, ok), ERROR_NULL_POINTER, "v"); + assert_eq!( + jac_mul(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "jv" + ); + + let eval_rhs = |y, f_out| unsafe { pybamm_rust_eval_rhs(0.0, y, nil, f_out, user_data) }; + assert_eq!(eval_rhs(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(eval_rhs(ok_const, nil_mut), ERROR_NULL_POINTER, "f_out"); + + let jac_assemble = + |y, jac| unsafe { pybamm_rust_jac_assemble(0.0, y, nil, 1.0, jac, user_data) }; + assert_eq!(jac_assemble(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!( + jac_assemble(ok_const, nil_mut), + ERROR_NULL_POINTER, + "jac_data" + ); + + let jac_action = + |y, v, jv| unsafe { pybamm_rust_jac_action(0.0, y, nil, v, jv, user_data) }; + assert_eq!(jac_action(nil, ok_const, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(jac_action(ok_const, nil, ok), ERROR_NULL_POINTER, "v"); + assert_eq!( + jac_action(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "jv" + ); + + let mass_action = |v, mv| unsafe { pybamm_rust_mass_action(v, mv, user_data) }; + assert_eq!(mass_action(nil, ok), ERROR_NULL_POINTER, "v"); + assert_eq!(mass_action(ok_const, nil_mut), ERROR_NULL_POINTER, "mv"); + + assert_eq!( + unsafe { pybamm_rust_algebraic_ids(nil_mut, user_data) }, + ERROR_NULL_POINTER, + "ids_out" + ); + assert_eq!( + unsafe { pybamm_rust_jac_csc_colptr(std::ptr::null_mut(), user_data) }, + ERROR_NULL_POINTER, + "colptr" + ); + assert_eq!( + unsafe { pybamm_rust_jac_csc_rowind(std::ptr::null_mut(), user_data) }, + ERROR_NULL_POINTER, + "rowind" + ); + } + + #[test] + fn every_required_buffer_of_the_sensitivity_and_output_calls_is_null_checked() { + let mut model = create_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let mut scratch = [0.0f64; 4]; + let ok = scratch.as_mut_ptr(); + let ok_const: *const f64 = ok; + let nil = std::ptr::null::(); + let nil_mut = std::ptr::null_mut::(); + + let sens_eval = + |y, df_dp| unsafe { pybamm_rust_sens_eval(0.0, y, nil, 0, df_dp, user_data) }; + assert_eq!(sens_eval(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(sens_eval(ok_const, nil_mut), ERROR_NULL_POINTER, "df_dp"); + + let sens_all = |y, out| unsafe { pybamm_rust_sens_eval_all(0.0, y, nil, out, user_data) }; + assert_eq!(sens_all(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(sens_all(ok_const, nil_mut), ERROR_NULL_POINTER, "df_dp_out"); + + let project = |y, y_sens, out| unsafe { + pybamm_rust_output_sens_project(0.0, y, nil, y_sens, out, user_data) + }; + assert_eq!(project(nil, ok_const, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(project(ok_const, nil, ok), ERROR_NULL_POINTER, "y_sens"); + assert_eq!( + project(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "out" + ); + + // `out_len` is the one optional buffer: null there is not an error. + let output_eval = |y, out| unsafe { + pybamm_rust_output_eval(0.0, y, nil, 0, out, std::ptr::null_mut(), user_data) + }; + assert_eq!(output_eval(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(output_eval(ok_const, nil_mut), ERROR_NULL_POINTER, "out"); + + let batch = + |ts, ys, out| unsafe { pybamm_rust_output_eval_batch(ts, ys, 1, nil, out, user_data) }; + assert_eq!(batch(nil, ok_const, ok), ERROR_NULL_POINTER, "ts"); + assert_eq!(batch(ok_const, nil, ok), ERROR_NULL_POINTER, "ys"); + assert_eq!( + batch(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "out" + ); + } + + #[test] + fn every_required_buffer_of_the_algebraic_and_event_calls_is_null_checked() { + let mut model = create_test_model(); + let user_data: *mut c_void = std::ptr::from_mut(model.as_mut()).cast(); + let mut scratch = [0.0f64; 4]; + let ok = scratch.as_mut_ptr(); + let ok_const: *const f64 = ok; + let nil = std::ptr::null::(); + let nil_mut = std::ptr::null_mut::(); + + let alg_res = |y, out| unsafe { pybamm_rust_alg_res(0.0, y, nil, out, user_data) }; + assert_eq!(alg_res(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(alg_res(ok_const, nil_mut), ERROR_NULL_POINTER, "output"); + + let alg_jac = |y, out| unsafe { pybamm_rust_alg_jac_assemble(0.0, y, nil, out, user_data) }; + assert_eq!(alg_jac(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(alg_jac(ok_const, nil_mut), ERROR_NULL_POINTER, "output"); + + let alg_action = + |y, v, jv| unsafe { pybamm_rust_alg_jac_action(0.0, y, nil, v, jv, user_data) }; + assert_eq!(alg_action(nil, ok_const, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(alg_action(ok_const, nil, ok), ERROR_NULL_POINTER, "v"); + assert_eq!( + alg_action(ok_const, ok_const, nil_mut), + ERROR_NULL_POINTER, + "jv" + ); + + let events = |y, out| unsafe { pybamm_rust_events_eval(0.0, y, nil, out, user_data) }; + assert_eq!(events(nil, ok), ERROR_NULL_POINTER, "y"); + assert_eq!(events(ok_const, nil_mut), ERROR_NULL_POINTER, "output"); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/ir.rs b/packages/pybamm-rust/pybamm-core/src/ir.rs new file mode 100644 index 0000000000..f7d0d1ffcc --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/ir.rs @@ -0,0 +1,2373 @@ +//! Lowering from the expression DAG to a flat, executable instruction tape. +//! +//! [`TypedIr`] is what actually runs: one [`Instruction`] per node in topological +//! order, operands and results addressed as [`Slot`]s into a single scratch +//! buffer, and arrays, CSR matrices and interpolant tables held in a +//! [`ConstPool`] and addressed by index. Element counts are resolved here, not at +//! evaluation. +//! +//! The slot layout is the caller's one choice. [`from_arena`](TypedIr::from_arena) +//! reuses a slot once its value is dead; `from_arena_split_eval` separates primal +//! from tangent work so one primal pass serves every color; `from_arena_pinned` +//! keeps every intermediate, leaving the buffer standing as reverse-AD's value +//! tape. +//! +//! All three privatise conditional branch cones and emit each branch as a +//! contiguous block, so evaluation skips the branches it did not select. + +// Intentional u32 usage for compact instruction storage - expression graphs +// won't exceed 4B nodes in practice +#![allow(clippy::cast_possible_truncation)] + +use crate::arena::{Arena, NodeId}; +use crate::branch_regions::{ + RegionGroup, privatise_conditionals, schedule_regions, schedule_regions_partitioned, +}; +use crate::node::{CsrData, Node}; + +/// Metrics for buffer slot reuse quality. +#[derive(Clone, Debug)] +pub struct SlotStats { + /// Elements the lowered tape actually needs. + pub buffer_size: usize, + /// Elements a slot-per-node layout would need, the no-reuse baseline. + pub naive_size: usize, + /// Instructions emitted. + pub num_instructions: usize, + /// `buffer_size / naive_size`; 1.0 means no slot was ever reused. + pub reuse_ratio: f64, +} + +/// Buffer slot within the evaluation arena: `(offset, len)`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Slot { + pub offset: u32, + pub len: u32, +} + +impl Slot { + /// A slot of `len` elements starting at `offset` in the evaluation buffer. + #[inline] + pub const fn new(offset: u32, len: u32) -> Self { + Self { offset, len } + } + + /// Start offset, widened for indexing. + #[inline] + pub const fn offset_usize(self) -> usize { + self.offset as usize + } + + /// Element count, widened for indexing. + #[inline] + pub const fn len_usize(self) -> usize { + self.len as usize + } + + /// Whether this slot holds a single element, and so broadcasts against any + /// other operand. + #[inline] + pub const fn is_scalar(self) -> bool { + self.len == 1 + } +} + +/// Compile-time broadcast pattern for binary operations. +/// +/// Resolving which side is scalar once at lowering lets evaluation pick a +/// specialised loop per variant instead of testing operand widths per element. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum BroadcastKind { + /// Both operands are single elements. + ScalarScalar = 0, + /// Left operand is a single element, broadcast across the right. + ScalarVector = 1, + /// Right operand is a single element, broadcast across the left. + VectorScalar = 2, + /// Operands are equal-length vectors, paired element-wise. + VectorVector = 3, +} + +impl BroadcastKind { + /// Classify a pair of operand widths. + /// + /// # Panics + /// + /// Panics if the widths neither match nor have a scalar side. [`first_invalid`] + /// rejects that combination before lowering, so a panic here means the DAG + /// reached the builder unvalidated. + #[inline] + pub fn from_lens(a_len: usize, b_len: usize) -> Self { + assert!( + broadcast_widths_compatible(a_len, b_len), + "binary operand widths are incompatible: left={a_len}, right={b_len}" + ); + match (a_len, b_len) { + (1, 1) => Self::ScalarScalar, + (1, _) => Self::ScalarVector, + (_, 1) => Self::VectorScalar, + _ => Self::VectorVector, + } + } +} + +#[inline] +const fn broadcast_widths_compatible(a_len: usize, b_len: usize) -> bool { + a_len == b_len || (a_len == 1 && b_len > 1) || (b_len == 1 && a_len > 1) +} + +/// Element-wise binary operation. +/// +/// Every variant applies one `f64` operation across the broadcast operands, with +/// no domain guards: division by zero yields an infinity and out-of-domain powers +/// a NaN, exactly as the hardware would. The comparison variants return 1.0 or +/// 0.0 so they can be multiplied into expressions. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum BinaryOp { + Add, + Sub, + Mul, + Div, + /// `f64::powf`. `simplify` rewrites the exponents in `{2, 3, 4, -1, -2}` to + /// multiply/divide chains and folds away `1` and (scalar-base) `0`, so none of + /// those reach here; every other exponent, integer or not, does. + Pow, + /// `f64::min`, which returns the non-NaN operand if exactly one is NaN. + Minimum, + /// `f64::max`, which returns the non-NaN operand if exactly one is NaN. + Maximum, + /// Truncated remainder (`%`), so the sign follows the dividend rather than + /// the divisor as Python's `%` does. + Modulo, + Hypot, + /// `a <= b`, the branch-inclusive Heaviside: 1.0 when the operands are equal. + EqualHeaviside, + /// `a < b`, the strict Heaviside: 0.0 when the operands are equal. + NotEqualHeaviside, + /// `|a - b| < 1e-14`, a tolerance test rather than an exact comparison, so + /// values that differ only by rounding still compare equal. + Equality, +} + +/// Element-wise unary operation. +/// +/// Each variant is the `f64` function of the same name applied per element, again +/// without domain guards: `sqrt(-1)` is NaN and `log(0)` is `-inf` rather than an +/// error. Guarding is the model's job, through a smoothed operator in the +/// expression itself. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum UnaryOp { + Neg, + Abs, + Sqrt, + Exp, + Log, + Sin, + Cos, + Tanh, + Sinh, + Cosh, + Arcsinh, + Arctan, + /// Abramowitz & Stegun 7.1.26 approximation (error below 1.5e-7), not a libm + /// call, and shared with constant folding so a folded value matches a runtime + /// one bit for bit. + Erf, + /// `sign(0) = 0`, unlike `f64::signum`. + Sign, + Floor, + Ceiling, +} + +/// Single evaluation step. +/// +/// The field names are a fixed vocabulary across variants: `dst`, `src`, `a` and +/// `b` are element offsets into the evaluation buffer (not [`Slot`]s, since the +/// length is already fixed at lowering), `len` counts `f64` elements, and any `*_idx` +/// addresses the [`ConstPool`] rather than the buffer. Variants are `Copy` and +/// fixed-size, so a tape is one flat `Vec` an interpreter can walk without +/// chasing pointers. +/// +/// Write windows never overlap their read windows; the slot allocator guarantees +/// it, and evaluation relies on it to split the buffer into disjoint slices. +#[derive(Clone, Copy, Debug)] +#[repr(C)] +pub enum Instruction { + /// Load a scalar constant into `dst`. + LoadScalar { value: f64, dst: u32 }, + /// Load the current time value into `dst`. + LoadTime { dst: u32 }, + /// Load array data from `ConstPool` into `dst`. + LoadArray { data_idx: u32, len: u32, dst: u32 }, + /// Fill `dst` with zeros (length `len`). + FillZero { dst: u32, len: u32 }, + /// Load state vector slice `[start, end)` into `dst`. + LoadStateVector { start: u32, end: u32, dst: u32 }, + /// Load state derivative slice `[start, end)` into `dst`. + LoadStateVectorDot { start: u32, end: u32, dst: u32 }, + /// Load `width` consecutive packed input values starting at `offset` into `dst`. + LoadInputParameter { offset: u32, width: u32, dst: u32 }, + /// Load tangent state vector slice `[start, end)` into `dst`. + LoadTangentState { start: u32, end: u32, dst: u32 }, + /// Load the parameter seed `dp[index]` into `dst`. + LoadTangentParameter { index: u32, dst: u32 }, + /// Element-wise binary operation: `dst = a b` with broadcast. + Binary { + op: BinaryOp, + a: u32, + b: u32, + dst: u32, + len: u32, + kind: BroadcastKind, + }, + /// Element-wise unary operation: `dst = op(src)`. + Unary { + op: UnaryOp, + src: u32, + dst: u32, + len: u32, + }, + /// Reduce to a scalar: `dst = max(src[..src_len])`. + MaxReduce { src: u32, src_len: u32, dst: u32 }, + /// Reduce to a scalar: `dst = min(src[..src_len])`. + MinReduce { src: u32, src_len: u32, dst: u32 }, + /// Reduce subgradient: `dst = basis[k]`, `k` = argmax/argmin of + /// `picker[..len]` (first occurrence, strict comparison). + ReduceArgSelect { + basis_src: u32, + picker_src: u32, + len: u32, + is_max: bool, + dst: u32, + }, + /// Slice: `dst = src[start..start+len]`. + Index { + src: u32, + start: u32, + dst: u32, + len: u32, + }, + /// Concatenate `sources_len` slot ranges into `dst`. + Concat { + sources_idx: u32, + sources_len: u32, + dst: u32, + }, + /// Sparse matrix-vector multiply: `dst = CSR(csr_idx) @ vec(vec_src)`. + MatMul { + csr_idx: u32, + vec_src: u32, + dst: u32, + }, + /// Dense matrix-vector multiply: `dst = A @ v`, A row-major at slot `mat_src` (rows × cols). + DenseMatMul { + mat_src: u32, + rows: u32, + cols: u32, + vec_src: u32, + dst: u32, + }, + /// 1-D linear interpolation: `dst[i] = interp(src[i])`. + Interp1DLinear { + interp_idx: u32, + src: u32, + dst: u32, + len: u32, + }, + /// 1-D linear interpolation derivative. + Interp1DLinearDeriv { + interp_idx: u32, + src: u32, + dst: u32, + len: u32, + }, + /// 1-D cubic (pchip) interpolation. + Interp1DCubic { + interp_idx: u32, + src: u32, + dst: u32, + len: u32, + }, + /// 1-D cubic (pchip) interpolation derivative. + Interp1DCubicDeriv { + interp_idx: u32, + src: u32, + dst: u32, + len: u32, + }, + /// N-D tensor-product interpolation. + InterpNd { + interp_idx: u32, + sources_idx: u32, + dst: u32, + len: u32, + }, + /// N-D tensor-product interpolation partial derivative along `axis`. + InterpNdPartial { + interp_idx: u32, + sources_idx: u32, + axis: u32, + dst: u32, + len: u32, + }, + /// Piecewise conditional: select branches by `selector` value. + Conditional { + selector: u32, + branches_idx: u32, + branches_len: u32, + dst: u32, + out_len: u32, + }, + /// Skip all but the active branch's instruction block. + /// + /// Sits immediately before the blocks it guards. `blocks_idx` indexes + /// [`ConstPool::branch_blocks`], holding `(rel_start, len)` per branch where + /// `rel_start` is an offset **from this instruction's own index**, the tape + /// is executed as sub-slices (split-eval), so absolute indices would be + /// wrong. A `len` of 0 means that branch owns no instructions here. + Dispatch { + selector: u32, + blocks_idx: u32, + blocks_len: u32, + }, +} + +/// Interned table for a 1-D linear interpolant. +/// +/// Evaluation binary-searches `x_data`, so the knots must be strictly increasing; +/// outside the data range the boundary segment is extended linearly rather than +/// clamped flat. +#[derive(Debug, Clone)] +pub struct InterpolantEntry { + /// Knot vector, strictly increasing. + pub x_data: Vec, + /// Function values at knots. + pub y_data: Vec, +} + +/// Interned table for a 1-D cubic interpolant, pre-divided into per-interval +/// power-basis coefficients so evaluation is a search plus a Horner step. +#[derive(Debug, Clone)] +pub struct CubicInterpolantEntry { + /// Interval breakpoints, length `nseg + 1`. + pub breakpoints: Vec, + /// Per-interval power-basis coefficients `[c0, c1, c2, c3]`, length `nseg`. + /// `p(dx) = c0 + c1·dx + c2·dx² + c3·dx³`, `dx = x - breakpoints[i]`. + pub coeffs: Vec<[f64; 4]>, +} + +/// Interned table for a 2-D or 3-D tensor-product interpolant. +/// +/// One flat coefficient block per cell, so a lookup is one cell index per axis +/// followed by a tensor Horner evaluation; `order` says whether the cells are +/// multilinear or cubic in every axis, since the two are not mixed. +#[derive(Debug, Clone)] +pub struct NdInterpolantEntry { + /// Per-axis knot vectors (2 or 3 axes), each of length `nseg_a + 1`. + pub breakpoints: Vec>, + /// Flat per-cell power-basis tensors: cell-major, `order^ndim` coeffs per cell. + pub coeffs: Vec, + /// Per-axis polynomial order: 2 = multilinear, 4 = tensor cubic. + pub order: u32, +} + +/// Side tables an instruction addresses by index instead of carrying inline. +/// +/// Instructions are fixed-size and `Copy`, so anything variable-length lives here +/// and is referenced by a `u32`: dense arrays, CSR matrices, interpolant tables, +/// the operand lists of `Concat` and the branch ranges of conditionals. Adding +/// an entry deduplicates nothing, so a tape may hold the same array twice; CSE on +/// the DAG is what keeps that rare. +#[derive(Debug, Default, Clone)] +pub struct ConstPool { + /// Flattened array data; `array_offsets[i]` gives the start of slot `i`. + pub array_data: Vec, + /// Start offset of each array in `array_data`. + pub array_offsets: Vec, + + /// CSR sparse matrices for matmul instructions. + pub csr_data: Vec, + + /// 1-D linear interpolation tables. + pub interpolants: Vec, + + /// 1-D cubic (pchip) interpolation tables. + pub cubic_interpolants: Vec, + + /// N-D tensor-product interpolation tables. + pub nd_interpolants: Vec, + + /// Concat source slot ranges: `(offset, len)` pairs. + pub concat_sources: Vec<(u32, u32)>, + + /// Conditional branch slot ranges: `(offset, len)` pairs. + pub branch_offsets: Vec<(u32, u32)>, + + /// `Dispatch` branch block ranges: `(rel_start, len)` instruction offsets + /// relative to the owning `Dispatch`. + pub branch_blocks: Vec<(u32, u32)>, + + /// N-D interpolant child slot ranges: `(offset, len)` per input. + pub interp_nd_sources: Vec<(u32, u32)>, +} + +impl ConstPool { + /// An empty pool. + pub fn new() -> Self { + Self::default() + } + + /// Insert array data, returning its index. + pub fn add_array(&mut self, data: &[f64]) -> u32 { + let idx = self.array_offsets.len() as u32; + self.array_offsets.push(self.array_data.len()); + self.array_data.extend_from_slice(data); + idx + } + + /// Retrieve a slice of array data by index and length. + pub fn get_array(&self, idx: u32, len: u32) -> &[f64] { + let start = self.array_offsets[idx as usize]; + &self.array_data[start..start + len as usize] + } + + /// Insert a CSR matrix, returning its index. + pub fn add_csr(&mut self, csr: CsrData) -> u32 { + let idx = self.csr_data.len() as u32; + self.csr_data.push(csr); + idx + } + + /// Insert a 1-D linear interpolation table, returning its index. + pub fn add_interpolant(&mut self, x_data: Vec, y_data: Vec) -> u32 { + let idx = self.interpolants.len() as u32; + self.interpolants.push(InterpolantEntry { x_data, y_data }); + idx + } + + /// Insert a 1-D cubic interpolation table, returning its index. + pub fn add_cubic_interpolant(&mut self, breakpoints: Vec, coeffs: Vec<[f64; 4]>) -> u32 { + let idx = self.cubic_interpolants.len() as u32; + self.cubic_interpolants.push(CubicInterpolantEntry { + breakpoints, + coeffs, + }); + idx + } + + /// Insert an N-D interpolation table, returning its index. + pub fn add_nd_interpolant( + &mut self, + breakpoints: Vec>, + coeffs: Vec, + order: u32, + ) -> u32 { + let idx = self.nd_interpolants.len() as u32; + self.nd_interpolants.push(NdInterpolantEntry { + breakpoints, + coeffs, + order, + }); + idx + } +} + +/// Metadata for a partitioned primal/tangent instruction stream. +/// +/// When present, the instruction stream is ordered so all primal instructions +/// precede all tangent instructions, and buffer slots are partitioned into +/// two non-overlapping pools. This allows evaluating the primal section once +/// and re-running only the tangent section with different seed vectors. +#[derive(Debug, Clone, Copy)] +pub struct SplitEvalInfo { + /// Index into the instruction stream where tangent instructions begin. + pub primal_end: usize, + /// Buffer region `[0, primal_buffer_size)` holds primal values. + pub primal_buffer_size: usize, +} + +/// Typed Intermediate Representation +/// +/// A compiled representation of an expression DAG with explicit shapes +/// for all intermediate values. This IR can be: +/// - Interpreted directly by `CompiledExpr` +/// - Transformed by symbolic differentiation +#[derive(Debug, Clone)] +pub struct TypedIr { + /// Linear sequence of instructions in topological order + instructions: Vec, + + /// Root output slot + root_slot: Slot, + + /// Total buffer size needed for evaluation + buffer_size: usize, + + /// Constant pool for large data + consts: ConstPool, + + /// Metadata for differentiation + n_states: usize, + n_params: usize, + uses_state_dot: bool, + + /// Partitioning for split evaluation, `None` for standard IR. + split_eval_info: Option, +} + +impl TypedIr { + /// Compile an expression DAG into the IR. + /// + /// Runs [`privatise_conditionals`] first so branch cones `cse` made shared + /// become exclusive and can be short-circuited. Every `from_arena*` entry + /// point does this, which is why no caller has to. + pub fn from_arena(arena: &Arena, root: NodeId) -> Self { + Self::privatise_then(arena, root, IRBuilder::build) + } + + /// [`from_arena`](Self::from_arena) without branch privatisation. Test-only + /// escape hatch for comparing against an unprivatised tape. + #[cfg(test)] + pub fn from_arena_raw(arena: &Arena, root: NodeId) -> Self { + IRBuilder::build(arena, root) + } + + /// Run [`privatise_conditionals`] on `arena`/`root`, then hand the (possibly + /// rewritten) arena to `build`. Shared by every `from_arena*` constructor so + /// privatisation always happens exactly once, right before lowering. + fn privatise_then( + arena: &Arena, + root: NodeId, + build: impl FnOnce(&Arena, NodeId) -> Self, + ) -> Self { + match privatise_conditionals(arena, root) { + Some((owned, owned_root)) => build(&owned, owned_root), + None => build(arena, root), + } + } + + /// Output length of the root expression. + #[inline] + pub const fn output_len(&self) -> usize { + self.root_slot.len as usize + } + + /// Instruction stream in evaluation order. + #[inline] + pub fn instructions(&self) -> &[Instruction] { + &self.instructions + } + + /// Per-branch block lengths in tape order, one entry per branch of every + /// `Dispatch`. Empty when no conditional was short-circuited. + #[must_use] + pub fn branch_block_lens(&self) -> Vec { + let mut lens = Vec::new(); + for instr in &self.instructions { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len as usize { + lens.push(self.consts.branch_blocks[blocks_idx as usize + b].1); + } + } + } + lens + } + + /// How many `Dispatch` instructions the tape carries, one per + /// short-circuited conditional, in each half of a split-eval tape. + /// + /// The part of [`common_instruction_count`](Self::common_instruction_count) + /// that exists only because a conditional was short-circuited, which + /// `branch_block_lens` cannot recover: it is flattened across dispatches. + #[must_use] + pub fn dispatch_count(&self) -> usize { + self.instructions + .iter() + .filter(|instr| matches!(instr, Instruction::Dispatch { .. })) + .count() + } + + /// Instruction count excluding conditional branch blocks: the common tape + /// plus one `Dispatch` per short-circuited conditional. + /// + /// This is the quantity `casadi.Function.n_instructions()` reports for a + /// Switch-lowered node, so the two backends are comparable. Use + /// [instructions](Self::instructions).len() for the raw tape + /// length. + #[must_use] + pub fn common_instruction_count(&self) -> usize { + self.instructions.len() - self.branch_block_lens().iter().sum::() as usize + } + + /// Output slot of the root expression. + #[inline] + pub const fn root_slot(&self) -> Slot { + self.root_slot + } + + /// Total evaluation buffer size in elements. + #[inline] + pub const fn buffer_size(&self) -> usize { + self.buffer_size + } + + /// Reference to the constant pool. + #[inline] + pub const fn consts(&self) -> &ConstPool { + &self.consts + } + + /// Number of state variables referenced. + #[inline] + pub const fn n_states(&self) -> usize { + self.n_states + } + + /// Number of input parameters referenced. + #[inline] + pub const fn n_params(&self) -> usize { + self.n_params + } + + /// Whether the expression references state derivatives. + #[inline] + pub const fn uses_state_dot(&self) -> bool { + self.uses_state_dot + } + + /// Partitioning for split evaluation, `None` for standard IR. + #[inline] + pub const fn split_eval_info(&self) -> Option { + self.split_eval_info + } + + /// Compute buffer slot allocation quality metrics. + pub fn slot_stats(arena: &Arena, root: NodeId) -> SlotStats { + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(arena, &eval_order); + let naive_size: usize = eval_order.iter().map(|n| sizes[n.index()]).sum(); + let ir = Self::from_arena(arena, root); + let reuse_ratio = if naive_size == 0 { + 1.0 + } else { + ir.buffer_size as f64 / naive_size as f64 + }; + SlotStats { + buffer_size: ir.buffer_size, + naive_size, + num_instructions: ir.instructions.len(), + reuse_ratio, + } + } + + /// Compile with primal/tangent split for efficient Jacobian assembly. + /// + /// The instruction stream is partitioned so primal instructions precede + /// tangent instructions, with buffer slots in two disjoint pools. + /// Primal work is evaluated once; only the tangent section is re-run + /// per color during coloring-based Jacobian assembly. + pub fn from_arena_split_eval(arena: &Arena, root: NodeId) -> Self { + Self::privatise_then(arena, root, IRBuilder::build_split_eval) + } + + /// Compile with a no-reuse (SSA) slot layout for reverse-mode AD. + /// Every intermediate keeps a unique slot, so after one primal eval the + /// scratch buffer is the reverse value tape. + pub fn from_arena_pinned(arena: &Arena, root: NodeId) -> Self { + Self::privatise_then(arena, root, IRBuilder::build_pinned) + } +} + +/// Builder for constructing `TypedIr` from an expression DAG. +struct IRBuilder { + /// Emitted instruction stream in evaluation order. + instructions: Vec, + consts: ConstPool, + /// `(offset, len)` for each `NodeId` in the arena. + slots: Vec<(usize, usize)>, + /// Maximum state index referenced. + n_states: usize, + /// Number of input parameters referenced. + n_params: usize, + /// Whether the expression uses state derivatives. + uses_state_dot: bool, +} + +impl IRBuilder { + fn build(arena: &Arena, root: NodeId) -> TypedIr { + let base_order = arena.topological_order(root); + let schedule = schedule_regions(arena, &base_order); + let sizes = infer_sizes(arena, &schedule.order); + let (slots, total_size) = assign_slots(arena, &sizes, &schedule.order, &[root]); + + let mut builder = Self { + instructions: Vec::with_capacity(schedule.order.len()), + consts: ConstPool::new(), + slots, + n_states: 0, + n_params: 0, + uses_state_dot: false, + }; + builder.emit_scheduled(arena, &schedule.order, &schedule.groups); + builder.assert_block_slots_private(total_size); + + let (root_offset, root_len) = builder.slots[root.index()]; + + TypedIr { + instructions: builder.instructions, + root_slot: Slot::new(root_offset as u32, root_len as u32), + buffer_size: total_size, + consts: builder.consts, + n_states: builder.n_states, + n_params: builder.n_params, + uses_state_dot: builder.uses_state_dot, + split_eval_info: None, + } + } + + /// Build a primal `TypedIr` with a no-reuse (SSA) slot layout. See + /// [`assign_slots_pinned`]. Mirrors [`Self::build`] but never recycles + /// a slot, so the reverse-AD backward pass can read every intermediate. + fn build_pinned(arena: &Arena, root: NodeId) -> TypedIr { + let base_order = arena.topological_order(root); + let schedule = schedule_regions(arena, &base_order); + let sizes = infer_sizes(arena, &schedule.order); + let (slots, total_size) = assign_slots_pinned(&sizes, &schedule.order); + + let mut builder = Self { + instructions: Vec::with_capacity(schedule.order.len()), + consts: ConstPool::new(), + slots, + n_states: 0, + n_params: 0, + uses_state_dot: false, + }; + builder.emit_scheduled(arena, &schedule.order, &schedule.groups); + builder.assert_block_slots_private(total_size); + + let (root_offset, root_len) = builder.slots[root.index()]; + + TypedIr { + instructions: builder.instructions, + root_slot: Slot::new(root_offset as u32, root_len as u32), + buffer_size: total_size, + consts: builder.consts, + n_states: builder.n_states, + n_params: builder.n_params, + uses_state_dot: builder.uses_state_dot, + split_eval_info: None, + } + } + + fn build_split_eval(arena: &Arena, root: NodeId) -> TypedIr { + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(arena, &eval_order); + let is_tangent = classify_tangent_nodes(arena, &eval_order); + + // Stable-partition: primal nodes first, tangent nodes second + let mut primal_order = Vec::new(); + let mut tangent_order = Vec::new(); + for &nid in &eval_order { + if is_tangent[nid.index()] { + tangent_order.push(nid); + } else { + primal_order.push(nid); + } + } + // Each half schedules on its own, so a `Dispatch` and the blocks it + // guards land in one half; a cone spanning the split yields one each. + let primal_schedule = schedule_regions_partitioned(arena, &primal_order, &eval_order); + let tangent_schedule = schedule_regions_partitioned(arena, &tangent_order, &eval_order); + let partitioned_order: Vec = primal_schedule + .order + .iter() + .chain(tangent_schedule.order.iter()) + .copied() + .collect(); + + let n = sizes.len(); + let mut last_use = vec![0_usize; n]; + for (pos, &nid) in partitioned_order.iter().enumerate() { + arena.get(nid).for_each_child(|c| { + if pos > last_use[c.index()] { + last_use[c.index()] = pos; + } + }); + } + // Pin root + last_use[root.index()] = usize::MAX; + + // Pass 1: allocate primal slots in pool [0, primal_hw) + let (primal_slots, primal_hw) = + assign_slots_configured(&sizes, &primal_schedule.order, &last_use, 0); + + // Pass 2: allocate tangent slots in pool [primal_hw, total_hw) + let (tangent_slots, total_hw) = + assign_slots_configured(&sizes, &tangent_schedule.order, &last_use, primal_hw); + + // Merge slot maps + let mut slots = vec![(0_usize, 0_usize); n]; + for &nid in &primal_schedule.order { + slots[nid.index()] = primal_slots[nid.index()]; + } + for &nid in &tangent_schedule.order { + slots[nid.index()] = tangent_slots[nid.index()]; + } + + let mut builder = Self { + instructions: Vec::with_capacity(partitioned_order.len()), + consts: ConstPool::new(), + slots, + n_states: 0, + n_params: 0, + uses_state_dot: false, + }; + + builder.emit_scheduled(arena, &primal_schedule.order, &primal_schedule.groups); + let primal_end = builder.instructions.len(); + builder.emit_scheduled(arena, &tangent_schedule.order, &tangent_schedule.groups); + // After both halves, so the tangent half's blocks are checked too. + builder.assert_no_block_straddles(primal_end); + builder.assert_block_slots_private(total_hw); + + let (root_offset, root_len) = builder.slots[root.index()]; + + TypedIr { + instructions: builder.instructions, + root_slot: Slot::new(root_offset as u32, root_len as u32), + buffer_size: total_hw, + consts: builder.consts, + n_states: builder.n_states, + n_params: builder.n_params, + uses_state_dot: builder.uses_state_dot, + split_eval_info: Some(SplitEvalInfo { + primal_end, + primal_buffer_size: primal_hw, + }), + } + } + + fn slot_for(&self, id: NodeId) -> (u32, u32) { + let (off, len) = self.slots[id.index()]; + (off as u32, len as u32) + } + + /// Emit `order`, wrapping each annotated group range in a `Dispatch`. + /// + /// Group anchors are distinct (an empty group is never recorded) and their + /// ranges are disjoint, so one lookup slot per position suffices. + /// Emit `order`, wrapping each annotated group range in a `Dispatch`. + /// + /// Group anchors are distinct (an empty group is never recorded) and their + /// ranges are disjoint, so one lookup slot per position suffices. + fn emit_scheduled(&mut self, arena: &Arena, order: &[NodeId], groups: &[RegionGroup]) { + let mut group_at: Vec> = vec![None; order.len() + 1]; + for g in groups { + assert!( + group_at[g.anchor].is_none(), + "two region groups share anchor {}", + g.anchor + ); + group_at[g.anchor] = Some(g); + } + + let mut pos = 0; + while pos < order.len() { + if let Some(group) = group_at[pos] { + let total: usize = group.branch_lens.iter().sum(); + assert!(total > 0, "an empty region group must not be recorded"); + self.emit_dispatch(arena, group, &order[pos..pos + total]); + pos += total; + continue; + } + self.emit_node(arena, order[pos]); + pos += 1; + } + } + + /// Emit a `Dispatch` followed by one contiguous block per branch, recording + /// each block's `(rel_start, len)` relative to the `Dispatch`'s own index. + /// + /// `nodes` is the group's slice of the emission order, branch runs back to + /// back. Block lengths are measured in *instructions*, not nodes, because + /// `Node::SparseMatrix` emits none. + fn emit_dispatch(&mut self, arena: &Arena, group: &RegionGroup, nodes: &[NodeId]) { + let Node::Conditional { selector, .. } = arena.get(group.cond) else { + unreachable!("a region group is always anchored on a Conditional") + }; + let selector_slot = self.slot_for(*selector).0; + let dispatch_at = self.instructions.len(); + let blocks_idx = self.consts.branch_blocks.len() as u32; + // Reserve the table so the blocks can be emitted before it is filled. + self.consts + .branch_blocks + .resize(blocks_idx as usize + group.branch_lens.len(), (0, 0)); + self.instructions.push(Instruction::Dispatch { + selector: selector_slot, + blocks_idx, + blocks_len: group.branch_lens.len() as u32, + }); + + let mut taken = 0; + for (i, &n_nodes) in group.branch_lens.iter().enumerate() { + let start = self.instructions.len(); + for &node_id in &nodes[taken..taken + n_nodes] { + self.emit_node(arena, node_id); + } + taken += n_nodes; + let len = self.instructions.len() - start; + self.consts.branch_blocks[blocks_idx as usize + i] = + ((start - dispatch_at) as u32, len as u32); + } + assert_eq!(taken, nodes.len(), "group range and branch_lens disagree"); + } + + /// Assert that every `Dispatch` span lies wholly on one side of `split`. + /// Otherwise `run_tangent_section`, which executes `instructions[split..]`, + /// would enter a block without its `Dispatch` and run it unconditionally. + /// Ships alongside [`Self::assert_block_slots_private`], for the same reason: + /// the failure is silently wrong values, not a crash. + /// + /// # Panics + /// Panics if a block starts before `split` and ends after it. + fn assert_no_block_straddles(&self, split: usize) { + for (pc, instr) in self.instructions.iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len as usize { + let (rel, len) = self.consts.branch_blocks[blocks_idx as usize + b]; + let start = pc + rel as usize; + let end = start + len as usize; + assert!( + split <= start || split >= end, + "split point {split} straddles block [{start}, {end})" + ); + } + } + } + } + + /// Which instructions sit inside a branch block. + /// + /// # Panics + /// Panics if a block contains a `Dispatch`. Every evaluator skips one level + /// of block only, and `reverse.rs`'s span table keys one owner per span end, + /// so a nested span would be silently mis-walked; `build_span_owner`'s + /// overlap check cannot see a span fully contained in another. + fn block_instruction_mask(&self) -> Vec { + let mut in_block = vec![false; self.instructions.len()]; + for (pc, instr) in self.instructions.iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len as usize { + let (rel, len) = self.consts.branch_blocks[blocks_idx as usize + b]; + for k in 0..len as usize { + let at = pc + rel as usize + k; + assert!( + !matches!(self.instructions[at], Instruction::Dispatch { .. }), + "instruction {at} is a Dispatch inside the block of the Dispatch at \ + {pc}: branch blocks must be flat" + ); + in_block[at] = true; + } + } + } + } + in_block + } + + /// Assert that no instruction outside a branch block reads a value a block + /// defined. This is the one miscompilation this plan produces silently, where + /// the reader sees whatever a previous solve left in the recycled slot. Runs + /// once per compile, in release too. + /// + /// Tracks the *last writer* of each buffer element, not set membership: + /// recycling makes "something outside also writes here" routinely true. + /// Extents are element-wise, so a partial overlap still trips. A `Dispatch` + /// selector that is never written has no writer to blame, so its definition + /// is asserted directly. + /// + /// # Panics + /// Panics if an outside instruction's read resolves to a block-owned + /// definition, naming the reader, the element and the defining instruction, + /// or if a `Dispatch`'s selector has no earlier definition. + fn assert_block_slots_private(&self, buffer_size: usize) { + // A tape with no blocks cannot violate this, and most tapes have none. + if self.consts.branch_blocks.is_empty() { + return; + } + let in_block = self.block_instruction_mask(); + let mut last_writer: Vec> = vec![None; buffer_size]; + + for (pc, instr) in self.instructions.iter().enumerate() { + if !in_block[pc] { + // The selector decides which block runs, so an undefined one is + // a miscompilation the last-writer scan below cannot name. + if let Instruction::Dispatch { selector, .. } = *instr { + assert!( + last_writer[selector as usize].is_some(), + "Dispatch at {pc} selects on buffer element {selector}, which no earlier \ + instruction defines" + ); + } + for src in instruction_src_extents(instr, &self.consts) { + // A `Conditional`'s branch slots are the one legitimate + // outside read of a block-owned value. + if src.is_branch_slot { + continue; + } + let start = src.offset as usize; + let end = start + src.len as usize; + for (index, writer) in last_writer[start..end].iter().enumerate() { + if let Some(writer) = *writer + && in_block[writer] + { + panic!( + "instruction {pc} ({instr:?}) outside a branch block reads buffer \ + element {}, defined by block-private instruction {writer} ({:?})", + start + index, + self.instructions[writer] + ); + } + } + } + } + if let Some((offset, len)) = instruction_dst_extent(instr, &self.consts) { + let start = offset as usize; + for writer in &mut last_writer[start..start + len as usize] { + *writer = Some(pc); + } + } + } + } + + fn emit_node(&mut self, arena: &Arena, node_id: NodeId) { + let (dst, out_len) = self.slots[node_id.index()]; + let dst = dst as u32; + let out_len_u32 = out_len as u32; + + let instr = match &arena[node_id] { + Node::Scalar(v) => Instruction::LoadScalar { value: *v, dst }, + Node::Time => Instruction::LoadTime { dst }, + Node::Array(arr) => { + let data_idx = self.consts.add_array(&arr.data); + Instruction::LoadArray { + data_idx, + len: arr.data.len() as u32, + dst, + } + }, + Node::ZeroVector { len } => Instruction::FillZero { + dst, + len: *len as u32, + }, + Node::StateVector { start, end } => { + self.n_states = self.n_states.max(*end); + Instruction::LoadStateVector { + start: *start as u32, + end: *end as u32, + dst, + } + }, + Node::StateVectorDot { start, end } => { + self.uses_state_dot = true; + self.n_states = self.n_states.max(*end); + Instruction::LoadStateVectorDot { + start: *start as u32, + end: *end as u32, + dst, + } + }, + Node::InputParameter { + index, + offset, + width, + .. + } => { + self.n_params = self.n_params.max(*index + 1); + Instruction::LoadInputParameter { + offset: *offset as u32, + width: *width as u32, + dst, + } + }, + Node::TangentStateVector { start, end } => { + self.n_states = self.n_states.max(*end); + Instruction::LoadTangentState { + start: *start as u32, + end: *end as u32, + dst, + } + }, + Node::TangentParameter { index } => { + self.n_params = self.n_params.max(*index + 1); + Instruction::LoadTangentParameter { + index: *index as u32, + dst, + } + }, + Node::SparseMatrix(_) => return, // Not emitted as instruction + + // Binary operations + Node::Add(a, b) => self.emit_binary(BinaryOp::Add, *a, *b, dst, out_len_u32), + Node::Sub(a, b) => self.emit_binary(BinaryOp::Sub, *a, *b, dst, out_len_u32), + Node::Mul(a, b) => self.emit_binary(BinaryOp::Mul, *a, *b, dst, out_len_u32), + Node::Div(a, b) => self.emit_binary(BinaryOp::Div, *a, *b, dst, out_len_u32), + Node::Pow(a, b) => self.emit_binary(BinaryOp::Pow, *a, *b, dst, out_len_u32), + Node::Minimum(a, b) => self.emit_binary(BinaryOp::Minimum, *a, *b, dst, out_len_u32), + Node::Maximum(a, b) => self.emit_binary(BinaryOp::Maximum, *a, *b, dst, out_len_u32), + Node::Modulo(a, b) => self.emit_binary(BinaryOp::Modulo, *a, *b, dst, out_len_u32), + Node::Hypot(a, b) => self.emit_binary(BinaryOp::Hypot, *a, *b, dst, out_len_u32), + Node::EqualHeaviside(a, b) => { + self.emit_binary(BinaryOp::EqualHeaviside, *a, *b, dst, out_len_u32) + }, + Node::NotEqualHeaviside(a, b) => { + self.emit_binary(BinaryOp::NotEqualHeaviside, *a, *b, dst, out_len_u32) + }, + Node::Equality(a, b) => self.emit_binary(BinaryOp::Equality, *a, *b, dst, out_len_u32), + + // Unary operations + Node::Neg(a) => self.emit_unary(UnaryOp::Neg, *a, dst, out_len_u32), + Node::Abs(a) => self.emit_unary(UnaryOp::Abs, *a, dst, out_len_u32), + Node::Sqrt(a) => self.emit_unary(UnaryOp::Sqrt, *a, dst, out_len_u32), + Node::Exp(a) => self.emit_unary(UnaryOp::Exp, *a, dst, out_len_u32), + Node::Log(a) => self.emit_unary(UnaryOp::Log, *a, dst, out_len_u32), + Node::Sin(a) => self.emit_unary(UnaryOp::Sin, *a, dst, out_len_u32), + Node::Cos(a) => self.emit_unary(UnaryOp::Cos, *a, dst, out_len_u32), + Node::Tanh(a) => self.emit_unary(UnaryOp::Tanh, *a, dst, out_len_u32), + Node::Sinh(a) => self.emit_unary(UnaryOp::Sinh, *a, dst, out_len_u32), + Node::Cosh(a) => self.emit_unary(UnaryOp::Cosh, *a, dst, out_len_u32), + Node::Arcsinh(a) => self.emit_unary(UnaryOp::Arcsinh, *a, dst, out_len_u32), + Node::Arctan(a) => self.emit_unary(UnaryOp::Arctan, *a, dst, out_len_u32), + Node::Erf(a) => self.emit_unary(UnaryOp::Erf, *a, dst, out_len_u32), + Node::Sign(a) => self.emit_unary(UnaryOp::Sign, *a, dst, out_len_u32), + Node::Floor(a) => self.emit_unary(UnaryOp::Floor, *a, dst, out_len_u32), + Node::Ceiling(a) => self.emit_unary(UnaryOp::Ceiling, *a, dst, out_len_u32), + + // Reduction ops + Node::MaxReduce(a) => { + let (src, src_len) = self.slot_for(*a); + Instruction::MaxReduce { src, src_len, dst } + }, + Node::MinReduce(a) => { + let (src, src_len) = self.slot_for(*a); + Instruction::MinReduce { src, src_len, dst } + }, + Node::ReduceArgSelect { + basis, + picker, + is_max, + } => { + let (basis_src, len) = self.slot_for(*basis); + let (picker_src, picker_len) = self.slot_for(*picker); + // basis (a tangent) and its primal picker always share width; the + // scan reads picker[0..len], so a mismatch would read the wrong slot. + debug_assert_eq!( + len, picker_len, + "ReduceArgSelect basis/picker width mismatch" + ); + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + is_max: *is_max, + dst, + } + }, + + // Structural nodes + Node::Index { child, start, .. } => Instruction::Index { + src: self.slot_for(*child).0, + start: *start as u32, + dst, + len: out_len_u32, + }, + Node::Concat(children) => { + let sources_idx = self.consts.concat_sources.len() as u32; + for c in children { + let (off, len) = self.slot_for(*c); + self.consts.concat_sources.push((off, len)); + } + Instruction::Concat { + sources_idx, + sources_len: children.len() as u32, + dst, + } + }, + + // Matrix operations + Node::MatMul(a, b) => match &arena[*a] { + Node::SparseMatrix(csr) => { + let csr_idx = self.consts.add_csr(csr.as_ref().clone()); + Instruction::MatMul { + csr_idx, + vec_src: self.slot_for(*b).0, + dst, + } + }, + Node::Array(arr) => Instruction::DenseMatMul { + mat_src: self.slot_for(*a).0, + rows: arr.shape.rows as u32, + cols: arr.shape.cols as u32, + vec_src: self.slot_for(*b).0, + dst, + }, + _ => panic!("MatMul requires a constant matrix on the left"), + }, + + // Interpolation nodes + Node::Interpolant1DLinear { data, child } => { + let interp_idx = self + .consts + .add_interpolant(data.x_data.clone(), data.y_data.clone()); + Instruction::Interp1DLinear { + interp_idx, + src: self.slot_for(*child).0, + dst, + len: out_len_u32, + } + }, + Node::Interpolant1DLinearDeriv { + slopes, + x_data, + child, + } => { + // Store x_data as breakpoints and slopes as y_data (the derivative values) + let interp_idx = self + .consts + .add_interpolant(x_data.to_vec(), slopes.to_vec()); + Instruction::Interp1DLinearDeriv { + interp_idx, + src: self.slot_for(*child).0, + dst, + len: out_len_u32, + } + }, + Node::Interpolant1DCubic { data, child } => { + let interp_idx = self + .consts + .add_cubic_interpolant(data.breakpoints.clone(), data.coeffs.clone()); + Instruction::Interp1DCubic { + interp_idx, + src: self.slot_for(*child).0, + dst, + len: out_len_u32, + } + }, + Node::Interpolant1DCubicDeriv { data, child } => { + let interp_idx = self + .consts + .add_cubic_interpolant(data.breakpoints.clone(), data.coeffs.clone()); + Instruction::Interp1DCubicDeriv { + interp_idx, + src: self.slot_for(*child).0, + dst, + len: out_len_u32, + } + }, + Node::InterpolantNd { data, children } => { + let interp_idx = self.consts.add_nd_interpolant( + data.breakpoints.clone(), + data.coeffs.clone(), + data.order, + ); + let sources_idx = self.consts.interp_nd_sources.len() as u32; + for c in children { + let (off, len) = self.slot_for(*c); + self.consts.interp_nd_sources.push((off, len)); + } + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len: out_len_u32, + } + }, + Node::InterpolantNdPartial { + data, + children, + axis, + } => { + let interp_idx = self.consts.add_nd_interpolant( + data.breakpoints.clone(), + data.coeffs.clone(), + data.order, + ); + let sources_idx = self.consts.interp_nd_sources.len() as u32; + for c in children { + let (off, len) = self.slot_for(*c); + self.consts.interp_nd_sources.push((off, len)); + } + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis: *axis, + dst, + len: out_len_u32, + } + }, + + Node::Conditional { selector, branches } => { + let branches_idx = self.consts.branch_offsets.len() as u32; + for b in branches { + let (off, len) = self.slot_for(*b); + self.consts.branch_offsets.push((off, len)); + } + Instruction::Conditional { + selector: self.slot_for(*selector).0, + branches_idx, + branches_len: branches.len() as u32, + dst, + out_len: out_len_u32, + } + }, + }; + + self.instructions.push(instr); + } + + fn emit_binary(&self, op: BinaryOp, a: NodeId, b: NodeId, dst: u32, len: u32) -> Instruction { + let (a_off, a_len) = self.slot_for(a); + let (b_off, b_len) = self.slot_for(b); + Instruction::Binary { + op, + a: a_off, + b: b_off, + dst, + len, + kind: BroadcastKind::from_lens(a_len as usize, b_len as usize), + } + } + + fn emit_unary(&self, op: UnaryOp, a: NodeId, dst: u32, len: u32) -> Instruction { + let (src, _) = self.slot_for(a); + Instruction::Unary { op, src, dst, len } + } +} + +#[cfg(test)] +impl IRBuilder { + /// Hand-build a tape that violates `assert_block_slots_private`: an + /// instruction outside the one `Dispatch` block reads the slot that + /// block's only instruction wrote. No real scheduler produces this; it + /// exists to exercise the guard, which otherwise ships untested. + fn test_tape_reading_across_block_boundary() -> Self { + let mut consts = ConstPool::new(); + consts.branch_blocks.push((1, 1)); // one block: 1 instruction, starting right after the Dispatch + + let instructions = vec![ + Instruction::LoadScalar { value: 1.0, dst: 0 }, // defines the selector + Instruction::Dispatch { + selector: 0, + blocks_idx: 0, + blocks_len: 1, + }, + Instruction::LoadScalar { value: 1.0, dst: 5 }, // block-private write to slot 5 + Instruction::Unary { + op: UnaryOp::Neg, + src: 5, // outside the block, illegally reads the block's slot + dst: 6, + len: 1, + }, + ]; + + Self { + instructions, + consts, + slots: Vec::new(), + n_states: 0, + n_params: 0, + uses_state_dot: false, + } + } + + /// Hand-build a tape whose `Dispatch` selects on a slot no instruction ever + /// writes. The last-writer scan has no writer to blame here, so the guard + /// asserts the definition itself; no real scheduler produces this either. + fn test_tape_dispatching_on_an_undefined_selector() -> Self { + let mut consts = ConstPool::new(); + consts.branch_blocks.push((1, 1)); + + let instructions = vec![ + Instruction::Dispatch { + selector: 0, // never written + blocks_idx: 0, + blocks_len: 1, + }, + Instruction::LoadScalar { value: 1.0, dst: 5 }, + ]; + + Self { + instructions, + consts, + slots: Vec::new(), + n_states: 0, + n_params: 0, + uses_state_dot: false, + } + } +} + +/// The earliest lowering blocker reachable from `root`, if any. +/// +/// Checked before lowering so unsupported inputs surface as a Python error +/// instead of an FFI panic. Currently flags a `MatMul` whose left operand is +/// not a constant (`SparseMatrix` or `Array`). +pub fn first_unsupported(arena: &Arena, root: NodeId) -> Option { + for id in arena.topological_order(root) { + if let Node::MatMul(a, _) = &arena[id] { + match &arena[*a] { + Node::SparseMatrix(_) | Node::Array(_) => {}, + other => { + return Some(format!( + "MatMul left operand must be a constant matrix, got {other:?}" + )); + }, + } + } + } + None +} + +/// The earliest invalid evaluator shape relationship reachable from `root`. +/// +/// Checked at Python entry points before lowering so malformed expression +/// graphs raise `ValueError` instead of reading adjacent evaluator scratch. +pub fn first_invalid(arena: &Arena, root: NodeId) -> Option { + let eval_order = arena.topological_order(root); + if eval_order.iter().any(|&id| { + matches!( + &arena[id], + Node::MatMul(a, _) + if !matches!(&arena[*a], Node::SparseMatrix(_) | Node::Array(_)) + ) + }) { + return None; + } + infer_sizes_checked(arena, &eval_order).err() +} + +fn infer_sizes_checked(arena: &Arena, eval_order: &[NodeId]) -> Result, String> { + let mut sizes = vec![0usize; arena.len()]; + + for &node_id in eval_order { + let size = match &arena[node_id] { + Node::Scalar(_) + | Node::Time + | Node::TangentParameter { .. } + | Node::MaxReduce(_) + | Node::MinReduce(_) + | Node::ReduceArgSelect { .. } => 1, + Node::InputParameter { width, .. } => *width, + Node::Array(arr) => arr.data.len(), + Node::ZeroVector { len } => *len, + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } => end.checked_sub(*start).ok_or_else(|| { + format!("node {node_id:?} has an inverted extent: start={start}, end={end}") + })?, + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => { + let (a_len, b_len) = (sizes[a.index()], sizes[b.index()]); + if !broadcast_widths_compatible(a_len, b_len) { + return Err(format!( + "binary node {node_id:?} has incompatible operand widths {a_len} and \ + {b_len}; widths must match or one operand must be scalar" + )); + } + a_len.max(b_len) + }, + Node::MatMul(a, b) => { + let (rows, cols) = match &arena[*a] { + Node::SparseMatrix(csr) => (csr.shape.rows, csr.shape.cols), + Node::Array(arr) => (arr.shape.rows, arr.shape.cols), + _ => { + return Err("MatMul requires a constant matrix on the left".to_string()); + }, + }; + let b_len = sizes[b.index()]; + if b_len != cols { + return Err(format!( + "MatMul node {node_id:?} has {cols} columns but its vector operand has \ + width {b_len}" + )); + } + rows + }, + Node::Concat(children) => children.iter().map(|c| sizes[c.index()]).sum(), + Node::InterpolantNd { data, children } + | Node::InterpolantNdPartial { data, children, .. } => { + if children.len() != data.breakpoints.len() { + return Err(format!( + "N-D interpolant node {node_id:?} has {} children for {} axes", + children.len(), + data.breakpoints.len() + )); + } + let out_len = children.iter().map(|c| sizes[c.index()]).max().unwrap_or(1); + if let Some(child) = children + .iter() + .find(|c| !matches!(sizes[c.index()], 1) && sizes[c.index()] != out_len) + { + return Err(format!( + "N-D interpolant node {node_id:?} has child widths that cannot broadcast: \ + {} and {out_len}", + sizes[child.index()] + )); + } + out_len + }, + Node::Neg(a) + | Node::Abs(a) + | Node::Sqrt(a) + | Node::Exp(a) + | Node::Log(a) + | Node::Sin(a) + | Node::Cos(a) + | Node::Tanh(a) + | Node::Sinh(a) + | Node::Cosh(a) + | Node::Arcsinh(a) + | Node::Arctan(a) + | Node::Erf(a) + | Node::Sign(a) + | Node::Floor(a) + | Node::Ceiling(a) => sizes[a.index()], + Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => sizes[child.index()], + Node::Conditional { selector, branches } => { + let selector_len = sizes[selector.index()]; + if selector_len != 1 { + return Err(format!( + "Conditional node {node_id:?} requires a scalar selector, got width \ + {selector_len}" + )); + } + let out_len = branches.first().map_or(1, |branch| sizes[branch.index()]); + if let Some(branch) = branches + .iter() + .find(|branch| sizes[branch.index()] != out_len) + { + return Err(format!( + "Conditional node {node_id:?} has branch widths {out_len} and {}", + sizes[branch.index()] + )); + } + out_len + }, + Node::SparseMatrix(_) => 0, + }; + + match &arena[node_id] { + Node::Index { child, end, .. } if *end > sizes[child.index()] => { + return Err(format!( + "Index node {node_id:?} ends at {end}, beyond child width {}", + sizes[child.index()] + )); + }, + Node::ReduceArgSelect { basis, picker, .. } => { + let (basis_len, picker_len) = (sizes[basis.index()], sizes[picker.index()]); + if basis_len == 0 { + return Err(format!( + "ReduceArgSelect node {node_id:?} requires non-empty operands" + )); + } + if basis_len != picker_len { + return Err(format!( + "ReduceArgSelect node {node_id:?} has basis width {basis_len} and picker \ + width {picker_len}" + )); + } + }, + Node::InterpolantNdPartial { data, axis, .. } + if *axis as usize >= data.breakpoints.len() => + { + return Err(format!( + "N-D interpolant partial node {node_id:?} has axis {axis} for {} axes", + data.breakpoints.len() + )); + }, + _ => {}, + } + sizes[node_id.index()] = size; + } + Ok(sizes) +} + +/// Infer the output size of each node in topological evaluation order. +pub fn infer_sizes(arena: &Arena, eval_order: &[NodeId]) -> Vec { + infer_sizes_checked(arena, eval_order) + .unwrap_or_else(|message| panic!("invalid expression graph: {message}")) +} + +/// One slot range an instruction reads. +#[derive(Clone, Copy, Debug)] +struct SrcExtent { + offset: u32, + len: u32, + /// A `Conditional`'s branch slot: the only read that may legitimately + /// resolve to a definition owned by a branch block. + is_branch_slot: bool, +} + +impl SrcExtent { + const fn read(offset: u32, len: u32) -> Self { + Self { + offset, + len, + is_branch_slot: false, + } + } + + const fn branch_slot(offset: u32, len: u32) -> Self { + Self { + offset, + len, + is_branch_slot: true, + } + } +} + +/// Destination slot range `(offset, len)` an instruction writes, or `None` for +/// `Dispatch`, which writes nothing. +/// +/// Exact, not an over-approximation: [`IRBuilder::assert_block_slots_private`] +/// tracks per-element last writers, so a short extent would leave stale entries +/// and a long one would mask a real violation. +fn instruction_dst_extent(instr: &Instruction, consts: &ConstPool) -> Option<(u32, u32)> { + let extent = match *instr { + Instruction::LoadScalar { dst, .. } + | Instruction::LoadTime { dst } + | Instruction::LoadTangentParameter { dst, .. } + | Instruction::MaxReduce { dst, .. } + | Instruction::MinReduce { dst, .. } + | Instruction::ReduceArgSelect { dst, .. } => (dst, 1), + Instruction::LoadArray { dst, len, .. } + | Instruction::FillZero { dst, len } + | Instruction::Binary { dst, len, .. } + | Instruction::Unary { dst, len, .. } + | Instruction::Index { dst, len, .. } + | Instruction::Interp1DLinear { dst, len, .. } + | Instruction::Interp1DLinearDeriv { dst, len, .. } + | Instruction::Interp1DCubic { dst, len, .. } + | Instruction::Interp1DCubicDeriv { dst, len, .. } + | Instruction::InterpNd { dst, len, .. } + | Instruction::InterpNdPartial { dst, len, .. } => (dst, len), + Instruction::LoadStateVector { start, end, dst } + | Instruction::LoadStateVectorDot { start, end, dst } + | Instruction::LoadTangentState { start, end, dst } => (dst, end - start), + Instruction::LoadInputParameter { width, dst, .. } => (dst, width), + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let width = (0..sources_len as usize) + .map(|i| consts.concat_sources[sources_idx as usize + i].1) + .sum(); + (dst, width) + }, + Instruction::MatMul { csr_idx, dst, .. } => { + (dst, consts.csr_data[csr_idx as usize].shape.rows as u32) + }, + Instruction::DenseMatMul { rows, dst, .. } => (dst, rows), + Instruction::Conditional { dst, out_len, .. } => (dst, out_len), + Instruction::Dispatch { .. } => return None, + }; + Some(extent) +} + +/// Slot ranges an instruction reads, with exact widths. +fn instruction_src_extents(instr: &Instruction, consts: &ConstPool) -> Vec { + match *instr { + Instruction::LoadScalar { .. } + | Instruction::LoadTime { .. } + | Instruction::LoadArray { .. } + | Instruction::FillZero { .. } + | Instruction::LoadStateVector { .. } + | Instruction::LoadStateVectorDot { .. } + | Instruction::LoadInputParameter { .. } + | Instruction::LoadTangentState { .. } + | Instruction::LoadTangentParameter { .. } => Vec::new(), + Instruction::Binary { + a, b, len, kind, .. + } => { + // The broadcast pattern decides which operand is the scalar one. + let (a_len, b_len) = match kind { + BroadcastKind::ScalarScalar => (1, 1), + BroadcastKind::ScalarVector => (1, len), + BroadcastKind::VectorScalar => (len, 1), + BroadcastKind::VectorVector => (len, len), + }; + vec![SrcExtent::read(a, a_len), SrcExtent::read(b, b_len)] + }, + Instruction::Unary { src, len, .. } + | Instruction::Interp1DLinear { src, len, .. } + | Instruction::Interp1DLinearDeriv { src, len, .. } + | Instruction::Interp1DCubic { src, len, .. } + | Instruction::Interp1DCubicDeriv { src, len, .. } => vec![SrcExtent::read(src, len)], + Instruction::MaxReduce { src, src_len, .. } + | Instruction::MinReduce { src, src_len, .. } => { + vec![SrcExtent::read(src, src_len)] + }, + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + .. + } => vec![ + SrcExtent::read(basis_src, len), + SrcExtent::read(picker_src, len), + ], + Instruction::Index { + src, start, len, .. + } => vec![SrcExtent::read(src + start, len)], + Instruction::Concat { + sources_idx, + sources_len, + .. + } => (0..sources_len as usize) + .map(|i| { + let (off, len) = consts.concat_sources[sources_idx as usize + i]; + SrcExtent::read(off, len) + }) + .collect(), + Instruction::MatMul { + csr_idx, vec_src, .. + } => vec![SrcExtent::read( + vec_src, + consts.csr_data[csr_idx as usize].shape.cols as u32, + )], + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + .. + } => vec![ + SrcExtent::read(mat_src, rows * cols), + SrcExtent::read(vec_src, cols), + ], + Instruction::InterpNd { + interp_idx, + sources_idx, + .. + } + | Instruction::InterpNdPartial { + interp_idx, + sources_idx, + .. + } => { + // The axis count lives on the interpolant table, not the instruction. + let ndim = consts.nd_interpolants[interp_idx as usize] + .breakpoints + .len(); + (0..ndim) + .map(|a| { + let (off, len) = consts.interp_nd_sources[sources_idx as usize + a]; + SrcExtent::read(off, len) + }) + .collect() + }, + Instruction::Conditional { + selector, + branches_idx, + branches_len, + out_len, + .. + } => { + let mut srcs = vec![SrcExtent::read(selector, 1)]; + srcs.extend((0..branches_len as usize).map(|i| { + let (off, _) = consts.branch_offsets[branches_idx as usize + i]; + SrcExtent::branch_slot(off, out_len) + })); + srcs + }, + Instruction::Dispatch { selector, .. } => vec![SrcExtent::read(selector, 1)], + } +} + +/// Classify nodes as primal or tangent via forward taint propagation. +fn classify_tangent_nodes(arena: &Arena, eval_order: &[NodeId]) -> Vec { + let mut is_tangent = vec![false; arena.len()]; + + for &nid in eval_order { + let node = arena.get(nid); + let tainted = match node { + Node::TangentStateVector { .. } | Node::TangentParameter { .. } => true, + _ => { + let mut any_child_tangent = false; + node.for_each_child(|c| { + if is_tangent[c.index()] { + any_child_tangent = true; + } + }); + any_child_tangent + }, + }; + is_tangent[nid.index()] = tainted; + } + + is_tangent +} + +/// Assign buffer slots, computing last-use positions from the evaluation order. +fn assign_slots( + arena: &Arena, + sizes: &[usize], + eval_order: &[NodeId], + roots: &[NodeId], +) -> (Vec<(usize, usize)>, usize) { + let n = sizes.len(); + + // Compute last_use[node] from the eval_order + let mut last_use = vec![0_usize; n]; + for (pos, &node_id) in eval_order.iter().enumerate() { + arena.get(node_id).for_each_child(|c| { + let cur = last_use[c.index()]; + if pos > cur { + last_use[c.index()] = pos; + } + }); + } + for &r in roots { + last_use[r.index()] = usize::MAX; + } + + assign_slots_configured(sizes, eval_order, &last_use, 0) +} + +/// Slot allocation with **no reuse**: every node's slot is pinned for the +/// whole evaluation, so the primal scratch preserves every intermediate. +/// Required by the reverse-AD value tape, whose backward pass reads each +/// operand's recorded value by its (now stable) slot. +fn assign_slots_pinned(sizes: &[usize], eval_order: &[NodeId]) -> (Vec<(usize, usize)>, usize) { + let last_use = vec![usize::MAX; sizes.len()]; + assign_slots_configured(sizes, eval_order, &last_use, 0) +} + +/// Sweep-line buffer slot allocator with externalized lifetime control. +/// +/// Assigns buffer offsets in `eval_order`, reusing freed regions. +/// `last_use[node_id]` is the last position that reads from the node; +/// set to `usize::MAX` to pin a slot permanently (e.g. roots). +/// `initial_high_water` offsets all allocations (for dual-pool partitioning). +fn assign_slots_configured( + sizes: &[usize], + eval_order: &[NodeId], + last_use: &[usize], + initial_high_water: usize, +) -> (Vec<(usize, usize)>, usize) { + use std::collections::{BTreeMap, BTreeSet}; + + let n = sizes.len(); + let mut slots = vec![(0_usize, 0_usize); n]; + + let mut by_offset: BTreeMap = BTreeMap::new(); + let mut by_size: BTreeSet<(usize, usize)> = BTreeSet::new(); + let mut high_water: usize = initial_high_water; + + let mut release_at: Vec> = vec![Vec::new(); eval_order.len() + 1]; + for &node_id in eval_order { + let lu = last_use[node_id.index()]; + if lu != usize::MAX && lu < eval_order.len() { + release_at[lu + 1].push(node_id); + } + } + + for (pos, &node_id) in eval_order.iter().enumerate() { + // Release phase + for &n_to_free in &release_at[pos] { + let (off, len) = slots[n_to_free.index()]; + if len == 0 { + continue; + } + + let mut new_off = off; + let mut new_len = len; + if let Some(&right_len) = by_offset.get(&(off + len)) { + by_offset.remove(&(off + len)); + by_size.remove(&(right_len, off + len)); + new_len += right_len; + } + if let Some((&left_off, &left_len)) = by_offset.range(..off).next_back() + && left_off + left_len == off + { + by_offset.remove(&left_off); + by_size.remove(&(left_len, left_off)); + new_off = left_off; + new_len += left_len; + } + by_offset.insert(new_off, new_len); + by_size.insert((new_len, new_off)); + } + + // Allocate phase + let len = sizes[node_id.index()]; + if len == 0 { + slots[node_id.index()] = (0, 0); + continue; + } + + let chosen = by_size.range((len, 0)..).next().copied(); + let off = if let Some((region_len, region_off)) = chosen { + by_size.remove(&(region_len, region_off)); + by_offset.remove(®ion_off); + if region_len > len { + let rem_off = region_off + len; + let rem_len = region_len - len; + by_offset.insert(rem_off, rem_len); + by_size.insert((rem_len, rem_off)); + } + region_off + } else { + let off = high_water; + high_water += len; + off + }; + slots[node_id.index()] = (off, len); + } + + (slots, high_water) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::CompiledExpr; + use crate::node::{ArrayData, NdInterpolantData, Shape}; + + #[test] + #[should_panic(expected = "outside a branch block reads buffer element")] + fn assert_block_slots_private_catches_a_cross_block_read() { + let builder = IRBuilder::test_tape_reading_across_block_boundary(); + builder.assert_block_slots_private(7); + } + + #[test] + #[should_panic(expected = "which no earlier instruction defines")] + fn assert_block_slots_private_catches_an_undefined_selector() { + let builder = IRBuilder::test_tape_dispatching_on_an_undefined_selector(); + builder.assert_block_slots_private(7); + } + + #[test] + fn test_typed_ir_from_expression() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, y)); + + let ir = TypedIr::from_arena(&arena, expr); + + assert_eq!(ir.output_len(), 3); + assert_eq!(ir.instructions().len(), 3); // Scalar, StateVector, Mul + } + + #[test] + fn test_slot_struct() { + let slot = Slot::new(10, 5); + assert_eq!(slot.offset, 10); + assert_eq!(slot.len, 5); + assert_eq!(slot.offset_usize(), 10); + assert_eq!(slot.len_usize(), 5); + assert!(!slot.is_scalar()); + + let scalar_slot = Slot::new(0, 1); + assert!(scalar_slot.is_scalar()); + } + + #[test] + fn test_broadcast_kind_inference() { + assert_eq!(BroadcastKind::from_lens(1, 1), BroadcastKind::ScalarScalar); + assert_eq!(BroadcastKind::from_lens(1, 5), BroadcastKind::ScalarVector); + assert_eq!(BroadcastKind::from_lens(5, 1), BroadcastKind::VectorScalar); + assert_eq!(BroadcastKind::from_lens(5, 5), BroadcastKind::VectorVector); + } + + #[test] + fn test_ir_metadata() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 10 }); + let p = arena.alloc(Node::InputParameter { + name: "param".to_string(), + index: 2, + offset: 2, + width: 1, + }); + let expr = arena.alloc(Node::Mul(y, p)); + + let ir = TypedIr::from_arena(&arena, expr); + + assert_eq!(ir.n_states(), 10); + assert_eq!(ir.n_params(), 3); // index 2 means 3 params (0, 1, 2) + assert!(!ir.uses_state_dot()); + } + + #[test] + fn test_ir_with_state_dot() { + let mut arena = Arena::new(); + let y_dot = arena.alloc(Node::StateVectorDot { start: 0, end: 5 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, y_dot)); + + let ir = TypedIr::from_arena(&arena, expr); + + assert!(ir.uses_state_dot()); + assert_eq!(ir.n_states(), 5); + } + + #[test] + fn test_ir_nested_expression() { + let mut arena = Arena::new(); + // Build: sin(x * 2 + 1) where x is a state vector + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let mul = arena.alloc(Node::Mul(x, two)); + let add = arena.alloc(Node::Add(mul, one)); + let sin = arena.alloc(Node::Sin(add)); + + let ir = TypedIr::from_arena(&arena, sin); + + assert_eq!(ir.output_len(), 3); + // Instructions: StateVector, Scalar(2), Scalar(1), Mul, Add, Sin + assert_eq!(ir.instructions().len(), 6); + } + + #[test] + fn test_ir_concat() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![2.0, 3.0], + shape: Shape::vector(2), + }))); + let c = arena.alloc(Node::Scalar(4.0)); + let concat = arena.alloc(Node::Concat(vec![a, b, c])); + + let ir = TypedIr::from_arena(&arena, concat); + + assert_eq!(ir.output_len(), 4); + assert_eq!(ir.consts().concat_sources.len(), 3); + } + + #[test] + fn test_first_unsupported_flags_non_constant_matmul_lhs() { + let mut arena = Arena::new(); + let s = arena.alloc(Node::Scalar(2.0)); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0], + shape: Shape::vector(1), + }))); + let mm = arena.alloc(Node::MatMul(s, v)); + assert!(first_unsupported(&arena, mm).is_some()); + assert!(first_unsupported(&arena, v).is_none()); + } + + #[test] + fn test_symbolic_jacobian_width_mismatch_rejected_before_eval_overlap() { + let mut arena = Arena::new(); + let short = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = arena.alloc(Node::StateVector { start: 2, end: 5 }); + let add = arena.alloc(Node::Add(short, long)); + let scalar_entry = arena.alloc(Node::Scalar(1.0)); + let symbolic_jacobian = arena.alloc(Node::Concat(vec![add, scalar_entry])); + + // This layout previously reached split_dst_two_src with len=3 for the + // width-2 operand, making its fictitious read window overlap dst. + let message = + first_invalid(&arena, symbolic_jacobian).expect("binary widths must be rejected"); + assert!(message.contains("incompatible operand widths 2 and 3")); + } + + #[test] + fn test_first_invalid_flags_other_evaluator_width_hazards() { + let mut matmul_arena = Arena::new(); + let matrix = matmul_arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0; 6], + shape: Shape::matrix(2, 3), + }))); + let short = matmul_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let matmul = matmul_arena.alloc(Node::MatMul(matrix, short)); + assert!( + first_invalid(&matmul_arena, matmul) + .expect("MatMul width must be rejected") + .contains("has 3 columns but its vector operand has width 2") + ); + + let mut index_arena = Arena::new(); + let child = index_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let index = index_arena.alloc(Node::Index { + child, + start: 1, + end: 3, + }); + assert!( + first_invalid(&index_arena, index) + .expect("Index bounds must be rejected") + .contains("beyond child width 2") + ); + + let mut conditional_arena = Arena::new(); + let selector = conditional_arena.alloc(Node::Scalar(1.0)); + let short = conditional_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = conditional_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let conditional = conditional_arena.alloc(Node::Conditional { + selector, + branches: vec![short, long], + }); + assert!( + first_invalid(&conditional_arena, conditional) + .expect("Conditional widths must be rejected") + .contains("branch widths 2 and 3") + ); + + let mut interpolant_arena = Arena::new(); + let short = interpolant_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = interpolant_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let interpolant = interpolant_arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 1.0], vec![0.0, 1.0]], + coeffs: vec![0.0; 4], + order: 2, + }), + children: vec![short, long], + }); + assert!( + first_invalid(&interpolant_arena, interpolant) + .expect("N-D interpolant widths must be rejected") + .contains("cannot broadcast: 2 and 3") + ); + + let mut reduce_arena = Arena::new(); + let basis = reduce_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let picker = reduce_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let reduce = reduce_arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + assert!( + first_invalid(&reduce_arena, reduce) + .expect("ReduceArgSelect widths must be rejected") + .contains("basis width 2 and picker width 3") + ); + + let mut empty_reduce_arena = Arena::new(); + let basis = empty_reduce_arena.alloc(Node::ZeroVector { len: 0 }); + let picker = empty_reduce_arena.alloc(Node::ZeroVector { len: 0 }); + let reduce = empty_reduce_arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + assert!( + first_invalid(&empty_reduce_arena, reduce) + .expect("empty ReduceArgSelect must be rejected") + .contains("requires non-empty operands") + ); + } + + #[test] + fn test_ir_dense_matmul_output_len() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // row-major 2x3 + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 1.0, 1.0], + shape: Shape::vector(3), + }))); + let matmul = arena.alloc(Node::MatMul(a, v)); + let ir = TypedIr::from_arena(&arena, matmul); + assert_eq!(ir.output_len(), 2); + } + + #[test] + fn test_ir_interpolant() { + use crate::node::InterpolantData; + + let mut arena = Arena::new(); + let x = arena.alloc(Node::Scalar(1.5)); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(InterpolantData { + x_data: vec![0.0, 1.0, 2.0], + y_data: vec![0.0, 10.0, 20.0], + }), + child: x, + }); + + let ir = TypedIr::from_arena(&arena, interp); + + assert_eq!(ir.output_len(), 1); + assert_eq!(ir.consts().interpolants.len(), 1); + assert_eq!(ir.consts().interpolants[0].x_data, vec![0.0, 1.0, 2.0]); + } + + #[test] + fn test_const_pool() { + let mut pool = ConstPool::new(); + + // Test array storage + let idx1 = pool.add_array(&[1.0, 2.0, 3.0]); + let idx2 = pool.add_array(&[4.0, 5.0]); + assert_eq!(idx1, 0); + assert_eq!(idx2, 1); + assert_eq!(pool.get_array(idx1, 3), &[1.0, 2.0, 3.0]); + assert_eq!(pool.get_array(idx2, 2), &[4.0, 5.0]); + + // Test interpolant storage + let interp_idx = pool.add_interpolant(vec![0.0, 1.0], vec![0.0, 10.0]); + assert_eq!(interp_idx, 0); + assert_eq!(pool.interpolants[0].x_data, vec![0.0, 1.0]); + } + + #[test] + fn test_assign_slots_reuses_freed_slot_in_simple_chain() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let a = arena.alloc(Node::Sin(y)); + let b = arena.alloc(Node::Cos(a)); + let root = arena.alloc(Node::Neg(b)); + + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(&arena, &eval_order); + let (slots, total) = assign_slots(&arena, &sizes, &eval_order, &[root]); + + assert!( + total <= 6, + "expected slot reuse to reduce buffer below 6, got {total}" + ); + + for &(off, len) in &slots { + assert!( + off + len <= total, + "invalid slot ({off}, {len}) > total {total}" + ); + } + } + + #[test] + fn test_assign_slots_coalesces_adjacent_free_regions() { + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y1 = arena.alloc(Node::StateVector { start: 2, end: 4 }); + let s0 = arena.alloc(Node::Sin(y0)); + let s1 = arena.alloc(Node::Sin(y1)); + let sum = arena.alloc(Node::Add(s0, s1)); + let cat = arena.alloc(Node::Concat(vec![y0, y1])); + let big = arena.alloc(Node::Sin(cat)); + let repeated_sum = arena.alloc(Node::Concat(vec![sum, sum])); + let root = arena.alloc(Node::Add(big, repeated_sum)); + + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(&arena, &eval_order); + let (_slots, total) = assign_slots(&arena, &sizes, &eval_order, &[root]); + + let no_reuse: usize = eval_order.iter().map(|n| sizes[n.index()]).sum(); + assert!( + total < no_reuse, + "expected slot reuse total ({total}) < sequential sum ({no_reuse})" + ); + } + + #[test] + fn test_pinned_layout_disables_reuse_and_preserves_intermediates() { + // f = (y0*y1) + y2 ; the product `a` is dead after the add, so the + // reuse allocator recycles its slot, but the pinned layout must not. + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let a = arena.alloc(Node::Mul(y0, y1)); + let root = arena.alloc(Node::Add(a, y2)); + + let reuse = TypedIr::from_arena(&arena, root); + let pinned = TypedIr::from_arena_pinned(&arena, root); + + // No reuse => strictly larger (or equal) buffer; here strictly larger. + assert!(pinned.buffer_size() > reuse.buffer_size()); + + // After a pinned eval, the intermediate product a = y0*y1 = 2*3 = 6 is + // still live at its own slot (reuse layout would have overwritten it). + let compiled = CompiledExpr::from_ir(pinned); + let mut s = vec![0.0; compiled.scratch_len()]; + let out = compiled.eval(&mut s, 0.0, &[2.0, 3.0, 5.0], &[], &[]); + assert_eq!(out, &[11.0]); + // Every node has a unique slot: the number of distinct slot offsets in + // the instruction stream equals the instruction count (SSA). + let ir = compiled.ir(); + let mut offsets: Vec = Vec::new(); + for instr in ir.instructions() { + let dst = instruction_dst(instr); + if dst != usize::MAX { + offsets.push(dst); + } + } + let unique: std::collections::BTreeSet = offsets.iter().copied().collect(); + assert_eq!(unique.len(), offsets.len(), "pinned layout reused a slot"); + } + + // Small test-only helper: the dst offset of any instruction. + fn instruction_dst(instr: &Instruction) -> usize { + match *instr { + Instruction::LoadScalar { dst, .. } + | Instruction::LoadTime { dst } + | Instruction::LoadArray { dst, .. } + | Instruction::FillZero { dst, .. } + | Instruction::LoadStateVector { dst, .. } + | Instruction::LoadStateVectorDot { dst, .. } + | Instruction::LoadInputParameter { dst, .. } + | Instruction::LoadTangentState { dst, .. } + | Instruction::LoadTangentParameter { dst, .. } + | Instruction::Binary { dst, .. } + | Instruction::Unary { dst, .. } + | Instruction::MaxReduce { dst, .. } + | Instruction::MinReduce { dst, .. } + | Instruction::ReduceArgSelect { dst, .. } + | Instruction::Index { dst, .. } + | Instruction::Concat { dst, .. } + | Instruction::MatMul { dst, .. } + | Instruction::DenseMatMul { dst, .. } + | Instruction::Interp1DLinear { dst, .. } + | Instruction::Interp1DLinearDeriv { dst, .. } + | Instruction::Interp1DCubic { dst, .. } + | Instruction::Interp1DCubicDeriv { dst, .. } + | Instruction::InterpNd { dst, .. } + | Instruction::InterpNdPartial { dst, .. } + | Instruction::Conditional { dst, .. } => dst as usize, + Instruction::Dispatch { .. } => usize::MAX, + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/jacobian.rs b/packages/pybamm-rust/pybamm-core/src/jacobian.rs new file mode 100644 index 0000000000..f907e44f01 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/jacobian.rs @@ -0,0 +1,2015 @@ +//! Jacobian construction: symbolic derivative, sparsity, coloring, assembly. +//! +//! [`JacobianData`] is built once per model and holds everything an evaluation +//! needs: the derivative DAG lowered to split-eval tapes, the CSR sparsity, a +//! column coloring, and the scatter tables that place each color's results at +//! their CSC positions. Assembly then costs a *single* primal pass, whose cached +//! result every color reuses, plus one tangent sweep per color. That sharing is +//! the whole point of lowering to split-eval tapes. +//! +//! Rows too wide to color cheaply are split out and filled by reverse mode +//! instead, so a single dense row cannot force one color per column. Each such +//! row adds one backward pass over that same primal scratch. +//! +//! Entries a compile pass proves constant are lifted out of the sweep the same +//! way: they are written from a table and, being nobody's sweep result, stop +//! constraining the coloring of the columns that share their rows. +//! +//! [`JacobianData::assemble_into`] is the *only* implementation of that sweep. +//! A consumer supplies a [`JacobianLayout`] saying where this artifact's entries +//! land in its own value buffer -- the artifact's CSC, or the wider merged CSC a +//! model shares with its mass matrix -- and a [`JacobianScratch`] carrying the +//! buffers and the lane width. Batching, the constant table, the seed-lane +//! hygiene and the dense-row passes therefore exist once, whatever the consumer. + +use std::collections::HashSet; +use std::sync::{Arc, OnceLock}; + +use crate::adjoint::AdjointTape; +use crate::arena::{Arena, NodeId, NodeMap}; +use crate::coloring::{ColumnColoring, color_columns_masked}; +use crate::const_entries::classify_constant_entries; +use crate::eval::{CompiledExpr, TangentInputs}; +use crate::ir::TypedIr; +use crate::node::Node; +use crate::row_extract::extract_scalar_rows; +use crate::simplify::{node_len, simplify_pipeline}; +use crate::sparsity::{SparsityPattern, detect_sparsity_per_output}; +use crate::tangent::{tangent_wrt_params, tangent_wrt_states, tangent_wrt_subset}; +use crate::tangent_batch; + +/// Scratch ceiling for the batched tangent region, per assembly. Wider lanes +/// stop paying once the region leaves cache, and an unbounded width would let a +/// large model allocate a lane buffer far bigger than its own tape. +const LANE_SCRATCH_BUDGET: usize = 32 << 20; + +/// Colours one batched walk of `ir` should carry, or 1 for the scalar path. +/// +/// The widest monomorphised width whose lane region fits [`LANE_SCRATCH_BUDGET`]; +/// tapes with too few colours or unbatchable instructions stay scalar. +fn decide_lane_width(ir: &TypedIr, n_colors: usize) -> usize { + if n_colors < 2 || !tangent_batch::is_batchable(ir) { + return 1; + } + tangent_batch::SUPPORTED_LANES + .into_iter() + .find(|&lanes| { + lanes <= n_colors + && tangent_batch::tangent_scratch_len(ir, lanes) * size_of::() + <= LANE_SCRATCH_BUDGET + }) + .unwrap_or(1) +} + +/// Which inputs a Jacobian differentiates against; the same two cases the +/// tangent transform seeds, so the enum is shared rather than duplicated. +pub use crate::tangent::DiffTarget; + +/// CSC (Compressed Sparse Column) sparsity pattern for KLU compatibility. +/// +/// KLU and other direct sparse solvers expect CSC format, while internal +/// representations use CSR. This struct stores the CSC pattern for zero-allocation +/// Jacobian assembly. +#[derive(Debug, Clone)] +pub struct CscPattern { + /// Number of rows + pub nrows: usize, + /// Number of columns + pub ncols: usize, + /// Column pointers (length ncols + 1) + pub colptr: Vec, + /// Row indices for each non-zero (length nnz) + pub rowind: Vec, + /// Mapping from CSC index to (row, col) for assembly. + pub csc_to_csr_map: Vec<(usize, usize)>, +} + +impl CscPattern { + /// Convert CSR sparsity pattern to CSC. + pub fn from_csr(csr: &SparsityPattern) -> Self { + let nrows = csr.nrows; + let ncols = csr.ncols; + let nnz = csr.nnz(); + + // Count entries per column. + let mut col_counts = vec![0usize; ncols]; + for &col in &csr.indices { + col_counts[col] += 1; + } + + // Build column pointers. + let mut colptr = vec![0usize; ncols + 1]; + for (i, &count) in col_counts.iter().enumerate() { + colptr[i + 1] = colptr[i] + count; + } + + // Build row indices and CSC-to-CSR mapping. + let mut rowind = vec![0usize; nnz]; + let mut csc_to_csr_map = vec![(0usize, 0usize); nnz]; + let mut col_pos = colptr.clone(); + + for row in 0..nrows { + let row_start = csr.indptr[row]; + let row_end = csr.indptr[row + 1]; + for csr_idx in row_start..row_end { + let col = csr.indices[csr_idx]; + let csc_idx = col_pos[col]; + rowind[csc_idx] = row; + csc_to_csr_map[csc_idx] = (row, col); + col_pos[col] += 1; + } + } + + Self { + nrows, + ncols, + colptr, + rowind, + csc_to_csr_map, + } + } + + /// Number of non-zeros. + #[inline] + pub const fn nnz(&self) -> usize { + self.rowind.len() + } +} + +/// Where one entry of a color's JVP result belongs in the CSC value buffer. +/// +/// Precomputing the destination is what makes assembly a copy: after a sweep, each +/// entry reads `row` of the result vector and writes `csc_idx`, with no pattern +/// search per call. +#[derive(Debug, Clone, Copy)] +pub struct ColorScatterEntry { + /// Position in the CSC value buffer to write. + pub csc_idx: usize, + /// Row of the JVP result to read. + pub row: usize, +} + +/// The split-out dense rows, filled by one reverse (VJP) pass each instead of +/// one forward JVP sweep per column. +/// +/// A dense row (nnz ≥ `DENSE_ROW_MIN_NNZ`) would otherwise force the whole +/// matrix coloring toward one color per touched column. Splitting the rows out +/// lets the sparse remainder colour cheaply. See [`crate::row_extract`] for why +/// the rows share a tape rather than getting one each. +#[derive(Debug, Clone)] +pub struct AdjointDenseRows { + /// Parent output rows these gradients reconstruct, in tape-element order. + pub rows: Vec, + /// Reverse-AD tape whose element `i` is row `rows[i]`. + pub tape: AdjointTape, + /// `(column, csc_idx)` scatter targets in the parent CSC, per row. + pub entries: Vec>, +} + +impl AdjointDenseRows { + /// Group `rows` onto `tape`, whose element `i` must be row `rows[i]`. + /// + /// # Panics + /// Panics unless the tape recovers exactly one element per row, which is + /// the correspondence every other method here relies on. + pub fn new(rows: Vec, tape: AdjointTape, entries: Vec>) -> Self { + assert_eq!( + tape.n_rows(), + rows.len(), + "the tape must hold one element per split row" + ); + assert_eq!(entries.len(), rows.len(), "one scatter list per split row"); + Self { + rows, + tape, + entries, + } + } + + /// Rows filled by reverse mode, so reverse passes per assembly. + #[inline] + pub const fn n_rows(&self) -> usize { + self.rows.len() + } + + /// Jacobian entries these rows account for. + pub fn n_entries(&self) -> usize { + self.entries.iter().map(Vec::len).sum() + } + + /// Fill every row's `scatter` targets in `out` from one shared forward pass. + /// + /// `scatter` is indexed by tape element, so it is [`Self::entries`] or the + /// same lists remapped to another CSC. `grad` doubles as the gradient + /// buffer and must span the tape's state dimension. + #[allow(clippy::too_many_arguments)] + pub fn assemble_into( + &self, + scratch: &mut [f64], + bar: &mut [f64], + grad: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + scatter: &[Vec<(usize, usize)>], + out: &mut [f64], + ) { + self.tape.eval_forward(scratch, t, y, y_dot, inputs); + for (row, targets) in scatter.iter().enumerate() { + self.tape.assemble_row(scratch, bar, grad, row); + for &(col, csc_idx) in targets { + out[csc_idx] = grad[col]; + } + } + } +} + +/// Where one [`JacobianData`]'s entries land in a consumer's value buffer. +/// +/// Two consumers assemble the same artifact into differently indexed buffers: a +/// standalone one into the artifact's own CSC ([`JacobianData::own_layout`]), and +/// one whose buffer merges this pattern with another -- `CompiledModel`, which +/// carries `df/dy` and the mass matrix in a single array -- into the merged slots +/// ([`JacobianData::layout_in`]). Both are the same three scatter tables against +/// different slot indices, so the sweep takes a layout rather than each consumer +/// re-deriving the tables from the artifact's public parts. +#[derive(Debug, Clone)] +pub struct JacobianLayout { + /// Per colour, the `(slot, row)` pairs that colour's sweep recovers. + color_to_slots: Vec>, + /// Per dense-row group, per row of that group, `(column, slot)` targets. + dense_row_slots: Vec>>, + /// The starting buffer for one assembly: constants in place, zero everywhere + /// else. One `copy_from_slice` replaces a memset plus a scattered rewrite of + /// most of what it just wrote. + template: Vec, + /// `(slot, value)` for every entry a compile pass proved constant, for + /// consumers that report the table as well as writing it. + constant_slots: Vec<(usize, f64)>, +} + +impl JacobianLayout { + /// Slots the target buffer must hold, which is its consumer's `nnz`. + #[inline] + pub const fn n_slots(&self) -> usize { + self.template.len() + } + + /// `(slot, value)` for every entry a compile pass proved constant. Split + /// dense rows are absent: their own tape recomputes them. + #[inline] + pub fn constant_slots(&self) -> &[(usize, f64)] { + &self.constant_slots + } +} + +/// Per-solve mutable buffers for one [`JacobianData`]'s assembly and actions. +/// +/// Sized from the artifact, and carrying the lane width its tangent tape +/// supports, so a caller can neither mis-size a buffer nor drive the batched +/// sweep with a scalar one. Create one per solve; two concurrent assemblies of +/// the same artifact need two. +#[derive(Debug, Clone)] +pub struct JacobianScratch { + /// The parent tangent tape, then in sequence every dense-row sub-tape, which + /// reuse this buffer after the parent sweep. + tape: Vec, + /// Scalar colour seed, spanning [`JacobianData::seed_len`]. + seed: Vec, + /// Adjoint `bar` buffer, parallel to the widest dense-row value tape. + bar: Vec, + /// Colours per batched tangent walk, or 1 for the scalar per-colour path. + lanes: usize, + /// Lane-minor tangent region; empty when `lanes == 1`. + tangent_lanes: Vec, + /// Lane-minor colour seeds; empty when `lanes == 1`. + seed_lanes: Vec, +} + +impl JacobianScratch { + /// Buffers for `jac` at the widest lane width its tangent tape supports. + pub fn new(jac: &JacobianData) -> Self { + Self::with_lanes(jac, jac.lane_width()) + } + + /// Buffers pinned to the scalar one-colour-per-walk path: the reference the + /// batched sweep is compared against, and what a tape that cannot batch gets + /// from [`Self::new`] anyway. + pub fn scalar(jac: &JacobianData) -> Self { + Self::with_lanes(jac, 1) + } + + fn with_lanes(jac: &JacobianData, lanes: usize) -> Self { + let batched = lanes > 1; + Self { + tape: vec![0.0; jac.max_scratch_len()], + seed: vec![0.0; jac.seed_len().max(1)], + bar: vec![0.0; jac.max_adjoint_tape_len().max(1)], + lanes, + tangent_lanes: vec![ + 0.0; + if batched { + tangent_batch::tangent_scratch_len(jac.assembly_tape().ir(), lanes) + } else { + 0 + } + ], + seed_lanes: vec![0.0; if batched { jac.seed_dim * lanes } else { 0 }], + } + } + + /// Colours this scratch sweeps per walk of the tangent tape. + #[inline] + pub const fn lane_width(&self) -> usize { + self.lanes + } +} + +/// Everything `finish_with_options` switches on, so the three constructors +/// differ by data rather than by a run of positional flags. +#[derive(Debug, Clone, Copy)] +struct BuildOptions { + /// Fill wide scalar rows by reverse mode rather than coloring against + /// them. Only a full-state build can, since it owns every column. + split_dense_rows: bool, + /// Sweep only the entries a compile pass could not prove constant, writing + /// the rest from a table. Off is the reference path the exactness tests + /// compare against. + split_constants: bool, +} + +/// What [`JacobianData::decide_coloring`] settled on: the coloring driving the +/// sweep, the entry mask it was chosen under, and the rows it handed to reverse +/// mode instead. +struct ColoringDecision { + coloring: ColumnColoring, + swept: Vec, + dense_rows: Vec, + /// Rows `detect_dense_rows` nominated, adopted or not. + n_candidate_rows: usize, +} + +/// Prepared derivative artifact for one (expression, wrt) pair. +/// +/// Holds the tangent-transformed expression, sparsity pattern, column coloring +/// and dense-row split. Immutable after build, so it is shared via `Arc` and read +/// by any number of concurrent assemblies; each supplies its own +/// [`JacobianScratch`] and a [`JacobianLayout`] naming its output slots. +/// +/// The fields are private on purpose: every one of them is an input to +/// [`Self::assemble_into`], and the whole point of this type is that the sweep +/// they describe has one implementation. Read them through the accessors, add a +/// [`JacobianLayout`] for a new output ordering, and never rebuild the tables. +#[derive(Debug)] +pub struct JacobianData { + /// Complete tangent-transformed expression with split-eval partitioning. + /// Matrix-free Jacobian actions use this tape, including split dense rows. + jvp_expr: Arc, + /// Parent tape used by colored assembly. When dense rows are split out, + /// their outputs are zeroed so only the adjoint tapes evaluate them. + assembly_jvp_expr: Arc, + /// CSR sparsity pattern of the derivative (no mass, no `cj`, pure `df/d·`). + sparsity: SparsityPattern, + /// CSC sparsity pattern (KLU/scipy ordering) of the same entries. + csc: CscPattern, + /// Column coloring that drives the per-call JVP sweep count. + coloring: ColumnColoring, + /// Rows of the derivative, matching the primal expression's output length. + n_rows: usize, + /// Columns of the derivative: states, or parameters, or the subset size. + n_cols: usize, + /// What this Jacobian differentiates with respect to. + wrt: DiffTarget, + /// For `States` with a column subset (algebraic blocks): global state + /// index per local column. Empty means identity (column `i` = state `i`). + col_to_global: Vec, + /// Length of the seed buffer callers must supply. Full-state/params: + /// `n_cols`. Subset: the full state dimension, tangent nodes index + /// global state positions, so the `dy` slice must span all of them. + seed_dim: usize, + /// Dense rows split out of the column coloring, all on one shared tape. + /// Empty unless a `new_wrt_states` build detected splittable dense rows AND + /// the split strictly lowered the colour count. + dense_rows: Vec, + /// Rows `detect_dense_rows` nominated, whether or not they were adopted. + /// `n_candidate_rows` above `n_dense_rows()` is a declined split, which is + /// correct but slower, and would otherwise be invisible. + n_candidate_rows: usize, + /// Per CSR entry of `sparsity`: whether a sweep must recover it. Always + /// one flag per entry, all true when the build split nothing. + swept_entries: Vec, + /// `(csr_idx, value)` for every entry a compile pass proved constant, + /// excluding split dense rows, whose own tape fills them. Kept in CSR order, + /// since a layout is what maps them to one consumer's slots. + constant_csr_entries: Vec<(usize, f64)>, + /// Colours a batched walk of the tangent tape carries, or 1 for the scalar + /// path. Decided once here rather than re-scanning the tape per scratch. + lane_width: usize, + /// This artifact's own-CSC layout, built on first ask and shared thereafter. + /// Identical by construction for every standalone consumer, so minting one + /// each would be the multi-copy-of-derived-tables shape this module exists to + /// remove; lazy because a consumer that only assembles into a merged buffer + /// never needs it. + own_layout: OnceLock, + /// The COO triplet for that layout, on the same terms. + own_coo: OnceLock<(Vec, Vec)>, +} + +impl JacobianData { + /// `df/dy` over the full state vector. + /// + /// Enables dense-row splitting: wide scalar rows are extracted into 1×n + /// sub-Jacobians so the sparse remainder colours cheaply. + pub fn new_wrt_states(arena: &Arena, root: NodeId, n_rows: usize, n_states: usize) -> Self { + Self::new_wrt_states_inner(arena, root, n_rows, n_states, true) + } + + /// As [`Self::new_wrt_states`], sweeping every entry rather than lifting the + /// constant ones out. The reference path the exactness tests compare against. + pub fn new_wrt_states_unsplit( + arena: &Arena, + root: NodeId, + n_rows: usize, + n_states: usize, + ) -> Self { + Self::new_wrt_states_inner(arena, root, n_rows, n_states, false) + } + + fn new_wrt_states_inner( + arena: &Arena, + root: NodeId, + n_rows: usize, + n_states: usize, + split_constants: bool, + ) -> Self { + let mut diff_arena = arena.clone(); + let tangent_root = tangent_wrt_states(&mut diff_arena, root); + let sparsity = detect_sparsity_per_output(arena, root, n_rows, n_states); + Self::finish_with_options( + arena, + root, + diff_arena, + tangent_root, + sparsity, + DiffTarget::States, + Vec::new(), + n_states, + BuildOptions { + split_dense_rows: true, + split_constants, + }, + ) + } + + /// `df/dp` over the full parameter vector. Sparsity is taken as dense + /// (parameters are scalars; coloring degenerates to one sweep per + /// column, matching unit-seed cost). Parameter sparsity detection is + /// a later optimisation with no API impact. + pub fn new_wrt_params(arena: &Arena, root: NodeId, n_rows: usize, n_params: usize) -> Self { + let mut diff_arena = arena.clone(); + let tangent_root = tangent_wrt_params(&mut diff_arena, root); + let sparsity = SparsityPattern::dense(n_rows, n_params); + Self::finish_with_options( + arena, + root, + diff_arena, + tangent_root, + sparsity, + DiffTarget::Params, + Vec::new(), + n_params, + BuildOptions { + split_dense_rows: false, + split_constants: false, + }, + ) + } + + /// `dg/dy_subset` for an algebraic block: differentiate w.r.t. the given + /// global state indices only; columns are local (`0..subset.len()`). + /// + /// `subset` must be strictly ascending (sorted, no duplicates). The + /// global-to-local column remap and `build_csr_to_csc_map` both rely on + /// ascending CSR column order; an unsorted subset silently mis-maps entries. + pub fn new_wrt_state_subset( + arena: &Arena, + root: NodeId, + n_rows: usize, + n_states: usize, + subset: &[usize], + ) -> Self { + assert!( + subset.windows(2).all(|w| w[0] < w[1]), + "subset must be strictly ascending" + ); + let active: HashSet = subset.iter().copied().collect(); + let mut diff_arena = arena.clone(); + let tangent_root = tangent_wrt_subset(&mut diff_arena, root, &active); + let full = detect_sparsity_per_output(arena, root, n_rows, n_states); + let filtered = filter_sparsity_columns(&full, subset); + // Remap global columns to local positions so the pattern is (n_rows, subset.len()). + let mut global_to_local = vec![usize::MAX; n_states]; + for (local, &g) in subset.iter().enumerate() { + global_to_local[g] = local; + } + let mut local = SparsityPattern::new(n_rows, subset.len()); + local.indptr.clone_from(&filtered.indptr); + local.indices = filtered + .indices + .iter() + .map(|&c| global_to_local[c]) + .collect(); + Self::finish_with_options( + arena, + root, + diff_arena, + tangent_root, + local, + DiffTarget::States, + subset.to_vec(), + n_states, + BuildOptions { + // A dense row's adjoint tape reconstructs gradients over every + // column, which a subset build does not own. + split_dense_rows: false, + // `classify_constant_entries` keys each coefficient by seed + // index, i.e. global state position, while this pattern's columns + // are local -- so lifting constants here needs the pattern + // restated in seed space first. Not a reference-path choice. + split_constants: false, + }, + ) + } + + /// Assemble the artifact. `arena`/`root` are the ORIGINAL (pre-tangent) + /// primal graph, dense-row splitting resolves and re-differentiates rows + /// over them, so they must be threaded through even though the parent tape + /// evaluates `tangent_root` in `diff_arena`. + #[allow(clippy::too_many_arguments)] + fn finish_with_options( + arena: &Arena, + root: NodeId, + diff_arena: Arena, + tangent_root: NodeId, + sparsity: SparsityPattern, + wrt: DiffTarget, + col_to_global: Vec, + seed_dim: usize, + options: BuildOptions, + ) -> Self { + let (mut diff_arena, tangent_root) = simplify_pipeline(diff_arena, tangent_root); + + let csc = CscPattern::from_csr(&sparsity); + let row_to_csc = build_row_to_csc_entries(&csc); + + // Classify on the simplified tangent tape, so a folded value follows + // the same operator order the sweep would have executed. + let (swept, mut constant_csr_entries) = if options.split_constants { + classify_constant_entries(&diff_arena, tangent_root, &sparsity) + } else { + (vec![true; sparsity.nnz()], Vec::new()) + }; + + // Adopt the reduced coloring only if every dense row extracts AND it strictly + // beats the full coloring; else fall back. The sparsity is never altered. + let ColoringDecision { + coloring, + swept: swept_entries, + dense_rows, + n_candidate_rows, + } = Self::decide_coloring( + arena, + root, + &sparsity, + &row_to_csc, + swept, + options.split_dense_rows, + ); + let skip_rows: Vec = dense_rows + .iter() + .flat_map(|split| split.rows.iter().copied()) + .collect(); + retain_outside_rows(&sparsity, &skip_rows, &mut constant_csr_entries); + + let jvp_ir = TypedIr::from_arena_split_eval(&diff_arena, tangent_root); + let jvp_expr = Arc::new(CompiledExpr::from_ir(jvp_ir)); + // Masking is pruning, not correctness: nothing reads a split row's sweep + // output, so an unmaskable row only leaves a value the assembly discards. + let assembly_jvp_expr = match mask_scalar_rows(&mut diff_arena, tangent_root, &skip_rows) { + Some(root) if root != tangent_root => { + let masked_ir = TypedIr::from_arena_split_eval(&diff_arena, root); + Arc::new(CompiledExpr::from_ir(masked_ir)) + }, + _ => Arc::clone(&jvp_expr), + }; + let (n_rows, n_cols) = (sparsity.nrows, sparsity.ncols); + let lane_width = decide_lane_width(assembly_jvp_expr.ir(), coloring.n_colors); + Self { + jvp_expr, + assembly_jvp_expr, + sparsity, + csc, + coloring, + n_rows, + n_cols, + wrt, + col_to_global, + seed_dim, + dense_rows, + n_candidate_rows, + swept_entries, + constant_csr_entries, + lane_width, + own_layout: OnceLock::new(), + own_coo: OnceLock::new(), + } + } + + /// Choose the column coloring, and with it which rows leave the sweep. + /// + /// `swept` is the mask both decisions read; an adopted split returns it + /// narrowed, since a split row is nobody's sweep result and clearing its + /// entries is what takes it out of the coloring. + fn decide_coloring( + arena: &Arena, + root: NodeId, + sparsity: &SparsityPattern, + row_to_csc: &[Vec<(usize, usize)>], + swept: Vec, + split_dense_rows: bool, + ) -> ColoringDecision { + let no_candidates = |swept: Vec| ColoringDecision { + coloring: color_columns_masked(sparsity, &swept), + swept, + dense_rows: Vec::new(), + n_candidate_rows: 0, + }; + if !split_dense_rows { + return no_candidates(swept); + } + // Measure candidates by swept width: a row of known entries costs no + // colours, so filling it by reverse mode would waste a pass. + let widths = (0..sparsity.nrows) + .map(|row| { + swept[sparsity.indptr[row]..sparsity.indptr[row + 1]] + .iter() + .filter(|&&s| s) + .count() + }) + .collect::>(); + let candidates = detect_dense_rows(&widths); + if candidates.is_empty() { + return no_candidates(swept); + } + + let mut narrowed = swept.clone(); + for &row in &candidates { + narrowed[sparsity.indptr[row]..sparsity.indptr[row + 1]].fill(false); + } + let reduced = color_columns_masked(sparsity, &narrowed); + let full = color_columns_masked(sparsity, &swept); + let improves = reduced.n_colors < full.n_colors; + // Built once and moved out on either decline path, since the adopt path + // returns `narrowed` and `reduced` instead. + let declined = ColoringDecision { + coloring: full, + swept, + dense_rows: Vec::new(), + n_candidate_rows: candidates.len(), + }; + if !improves { + return declined; + } + + // Extraction is the expensive half of the decision, so it runs only once + // the colouring is known to improve. + let mut dense_rows = Vec::with_capacity(candidates.len().div_ceil(ROWS_PER_TAPE)); + for group in candidates.chunks(ROWS_PER_TAPE) { + let Some(block) = extract_scalar_rows(arena, root, group) else { + return declined; + }; + let entries = block + .rows + .iter() + .map(|&row| row_to_csc[row].clone()) + .collect(); + let tape = AdjointTape::new(&block.arena, block.root, sparsity.ncols); + dense_rows.push(AdjointDenseRows::new(block.rows, tape, entries)); + } + ColoringDecision { + coloring: reduced, + swept: narrowed, + n_candidate_rows: candidates.len(), + dense_rows, + } + } + + /// This artifact's entries laid out on its own CSC, which is what a standalone + /// consumer (a `scipy` matrix, a Newton sub-block) assembles into. + /// + /// Built on first ask and shared thereafter: it is the same tables for every + /// such consumer, so handing each one its own copy would be the duplication + /// this module exists to remove. + pub fn layout(&self) -> &JacobianLayout { + self.own_layout.get_or_init(|| self.layout_in(&self.csc)) + } + + /// This artifact's entries laid out on `buffer`, a pattern whose slots it + /// shares with another -- the merged `df/dy` + mass CSC a model assembles + /// into. Slots this artifact has no entry for keep the zero the template + /// writes, so the consumer's own pass can fold into them afterwards. + /// + /// # Panics + /// Panics unless `buffer` is a superset of [`Self::sparsity`], since every + /// entry of this artifact must have a slot to land in. + #[must_use] + pub fn layout_in(&self, buffer: &CscPattern) -> JacobianLayout { + let row_to_slots = &build_row_to_csc_entries(buffer); + let n_slots = buffer.nnz(); + let csr_to_slot = build_csr_to_csc_map(&self.sparsity, row_to_slots); + let color_to_slots = build_color_scatter_entries( + &self.sparsity, + &csr_to_slot, + &self.coloring, + &self.swept_entries, + ); + // Each dense row scatters grad[col] -> slot; a merged buffer numbers its + // slots differently from this artifact's own CSC, so take its row map. + let dense_row_slots = self + .dense_rows + .iter() + .map(|split| { + split + .rows + .iter() + .map(|&row| row_to_slots[row].clone()) + .collect() + }) + .collect(); + let constant_slots = map_entries_to_csc(&self.constant_csr_entries, &csr_to_slot); + let mut template = vec![0.0; n_slots]; + for &(slot, value) in &constant_slots { + template[slot] = value; + } + // A split row is nobody's sweep result: several same-coloured columns land + // on it, so bucketing one would alias their sums. `decide_coloring` clears + // those entries from `swept_entries`, and this is that holding. + debug_assert!( + { + let split: HashSet = self + .dense_rows + .iter() + .flat_map(|group| group.rows.iter().copied()) + .collect(); + color_to_slots + .iter() + .flatten() + .all(|entry| !split.contains(&entry.row)) + }, + "a split dense row must not be scattered by any colour" + ); + JacobianLayout { + color_to_slots, + dense_row_slots, + template, + constant_slots, + } + } + + /// COO `(rows, cols)` in the order [`Self::assemble_into`] writes a + /// [`layout`](Self::layout) buffer, with each local column reported as the + /// global state index it stands for. + /// + /// The pair a consumer needs when it wants the triplet form rather than this + /// artifact's CSC. Built on first ask and shared thereafter. + pub fn coo_global_indices(&self) -> (&[usize], &[usize]) { + let (rows, cols) = self.own_coo.get_or_init(|| { + self.csc + .csc_to_csr_map + .iter() + .map(|&(row, col)| (row, self.global_column(col))) + .unzip() + }); + (rows, cols) + } + + /// Global state index of a local column; `col_to_global` empty is identity. + #[inline] + fn global_column(&self, col: usize) -> usize { + self.col_to_global.get(col).copied().unwrap_or(col) + } + + /// CSR pattern of the derivative alone: no mass term, no `cj`, pure `df/d·`. + #[inline] + pub const fn sparsity(&self) -> &SparsityPattern { + &self.sparsity + } + + /// The same entries in CSC (KLU/scipy) ordering. + #[inline] + pub const fn csc(&self) -> &CscPattern { + &self.csc + } + + /// Column coloring that decides the per-call sweep count. With a dense-row + /// split adopted it is the reduced coloring, so it does not cover every row. + #[inline] + pub const fn coloring(&self) -> &ColumnColoring { + &self.coloring + } + + /// Sweeps of the tangent tape one assembly costs, before batching. + #[inline] + pub const fn n_colors(&self) -> usize { + self.coloring.n_colors + } + + /// What this Jacobian differentiates with respect to. + #[inline] + pub const fn wrt(&self) -> DiffTarget { + self.wrt + } + + /// Rows of the derivative, matching the primal expression's output length. + #[inline] + pub const fn n_rows(&self) -> usize { + self.n_rows + } + + /// Columns of the derivative: states, parameters, or the subset size. + #[inline] + pub const fn n_cols(&self) -> usize { + self.n_cols + } + + /// `(csr_idx, value)` for every entry a compile pass proved constant, in this + /// artifact's CSR order. Split dense rows are absent. + #[inline] + pub fn constant_csr_entries(&self) -> &[(usize, f64)] { + &self.constant_csr_entries + } + + /// The rows filled by reverse mode, grouped by shared tape. + #[inline] + pub fn dense_rows(&self) -> &[AdjointDenseRows] { + &self.dense_rows + } + + /// Rows `detect_dense_rows` nominated, adopted or not. + #[inline] + pub const fn n_candidate_rows(&self) -> usize { + self.n_candidate_rows + } + + /// The unpruned tangent tape, which every matrix-free action walks. + #[inline] + pub const fn action_tape(&self) -> &Arc { + &self.jvp_expr + } + + /// The tape colored assembly walks: [`action_tape`](Self::action_tape) with + /// the split dense rows' outputs masked off, or the same tape when no row + /// could be masked. + #[inline] + pub const fn assembly_tape(&self) -> &Arc { + &self.assembly_jvp_expr + } + + /// Colours a batched walk of the tangent tape carries, or 1 for the scalar + /// per-colour path. Decided once at build time. + #[inline] + pub const fn lane_width(&self) -> usize { + self.lane_width + } + + /// Length of the seed buffer [`JacobianScratch`] allocates: `n_cols` for a + /// full build, the whole state dimension for a subset, whose tangent nodes + /// index global state positions. + const fn seed_len(&self) -> usize { + self.seed_dim + } + + /// Scratch length covering the parent tape and every dense-row value tape. + /// The tapes run in sequence after the parent sweep and reuse the same + /// buffer, so [`JacobianScratch`] sizes to this maximum. + fn max_scratch_len(&self) -> usize { + self.jvp_expr + .scratch_len() + .max(self.assembly_jvp_expr.scratch_len()) + .max(self.max_adjoint_tape_len()) + } + + /// Length of the adjoint `bar` buffer: the largest dense-row value tape, + /// 0 when no rows were split. + pub(crate) fn max_adjoint_tape_len(&self) -> usize { + self.dense_rows + .iter() + .map(|split| split.tape.scratch_len()) + .max() + .unwrap_or(0) + } + + /// Rows filled by reverse mode rather than by the colour sweep. + pub fn n_dense_rows(&self) -> usize { + self.dense_rows.iter().map(AdjointDenseRows::n_rows).sum() + } + + /// Jacobian entries those rows account for. + pub fn dense_row_entries(&self) -> usize { + self.dense_rows + .iter() + .map(AdjointDenseRows::n_entries) + .sum() + } + + /// Instructions across the split rows' tapes: what the split costs in + /// compiled memory, against the sweeps it removes. + pub fn dense_row_tape_instructions(&self) -> usize { + self.dense_rows + .iter() + .map(|split| split.tape.instruction_count()) + .sum() + } + + /// Number of non-zeros in the derivative sparsity pattern. + pub const fn nnz(&self) -> usize { + self.sparsity.nnz() + } + + /// Assemble the derivative values into `out` at `layout`'s slots. + /// + /// The constant template, then the colour sweep at `scratch`'s lane width, + /// then one reverse pass per split dense row. Every slot the layout names is + /// written -- from the template, one colour's scatter, or a dense row -- so + /// callers need not zero `out`; nothing outside those slots is touched, which + /// is what lets a merged buffer carry another pattern's entries through the + /// call and fold its own term in afterwards. + /// + /// No mass term and no `cj`: those are the solver's, and keeping them out is + /// why the three consumers can share this one driver. + /// + /// # Panics + /// Panics if `out` is shorter than [`JacobianLayout::n_slots`]. + #[allow(clippy::too_many_arguments)] + pub fn assemble_into( + &self, + scratch: &mut JacobianScratch, + layout: &JacobianLayout, + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + // A layout built from another artifact would scatter this one's rows into + // the wrong slots, which the sweep itself cannot notice. + debug_assert_eq!( + layout.color_to_slots.len(), + self.coloring.n_colors, + "layout was built for a different jacobian" + ); + debug_assert_eq!( + layout.dense_row_slots.len(), + self.dense_rows.len(), + "layout was built for a different jacobian" + ); + let n_slots = layout.n_slots(); + assert!( + out.len() >= n_slots, + "jacobian buffer too small: need {n_slots}, got {}", + out.len() + ); + out[..n_slots].copy_from_slice(&layout.template); + + if scratch.lanes > 1 { + self.sweep_colors_batched(scratch, layout, t, y, y_dot, inputs, out); + } else { + self.sweep_colors_scalar(scratch, layout, t, y, y_dot, inputs, out); + } + + // Runs after the parent sweep, so `tape` is free as the value tape. Lands + // before any consumer postpass: a mass term subtracts on diagonal slots, + // so it must follow ALL derivative fills. + for (split, slots) in self.dense_rows.iter().zip(&layout.dense_row_slots) { + let JacobianScratch { + tape, bar, seed, .. + } = scratch; + split.assemble_into(tape, bar, seed, t, y, y_dot, inputs, slots, out); + } + } + + /// One tangent-tape walk per colour. Nothing reads the primal section when no + /// colour sweeps, which is the shape of a Jacobian that folded entirely. + #[allow(clippy::too_many_arguments)] + fn sweep_colors_scalar( + &self, + scratch: &mut JacobianScratch, + layout: &JacobianLayout, + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + if self.coloring.n_colors == 0 { + return; + } + self.assembly_jvp_expr + .run_primal_section(&mut scratch.tape, t, y, y_dot, inputs); + for color in 0..self.coloring.n_colors { + self.sweep_one_color_scalar(scratch, layout, color, out); + } + } + + /// Sweep colours in blocks, one tangent-tape walk per block. Lane `l` carries + /// colour `block_start + l`. + /// + /// Block widths come from [`Self::block_width`], so a tail shorter than the + /// scratch's lane count costs a narrower walk rather than a padded full-width + /// one. Per-lane arithmetic does not depend on the width, so which widths a + /// sweep happens to use cannot change the assembled values. + #[allow(clippy::too_many_arguments)] + fn sweep_colors_batched( + &self, + scratch: &mut JacobianScratch, + layout: &JacobianLayout, + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + let n_colors = self.coloring.n_colors; + if n_colors == 0 { + return; + } + self.assembly_jvp_expr + .run_primal_section(&mut scratch.tape, t, y, y_dot, inputs); + + // Undoing a block's seeds beats memsetting the lane region: each column + // is seeded once per assembly, so this is O(n_states), not O(n * blocks). + let mut block = 0; + while block < n_colors { + let stride = Self::block_width(scratch.lanes, n_colors - block); + if stride == 1 { + // One colour left: a scalar walk beats a vector one carrying it. + self.sweep_one_color_scalar(scratch, layout, block, out); + block += 1; + continue; + } + // A tail block may carry fewer colours than its stride; the surplus + // lanes are never seeded, so they stay zero and scatter nothing. + let carried = stride.min(n_colors - block); + self.write_seed_lanes(scratch, block, carried, stride, 1.0); + + let result = self.run_tangent_lanes(scratch, stride); + + for lane in 0..carried { + for entry in &layout.color_to_slots[block + lane] { + out[entry.csc_idx] = result[entry.row * stride + lane]; + } + } + self.write_seed_lanes(scratch, block, carried, stride, 0.0); + block += carried; + } + } + + /// Colours a tail must carry before a part-empty vector block beats that many + /// scalar walks. + /// + /// A `K`-lane walk costs `ratio` scalar walks, where `ratio < K` because the + /// operator gather is paid once and the lanes go through SIMD. Padding `r` + /// colours into one block therefore wins exactly when `ratio < r`. Measured + /// over deliberately different tangent tapes -- sparse-matmul (FV-like), + /// shallow and deep elementwise, and the SPMe/DFN fixtures -- `ratio` at + /// `K = 4` spans 1.53 to 2.33, so `r >= 3` wins everywhere while `r == 2` is + /// a loss on the gather-light half. Tapes that would amortise worse still + /// cannot reach here: `is_batchable` keeps them on the scalar path entirely. + const MIN_PADDED_TAIL: usize = 3; + + /// Width of the next block: the widest monomorphised width the remaining + /// colours fill outright, else the narrowest vector width once the tail is + /// long enough to pay for the empty lanes, else a scalar walk. Never exceeds + /// `lanes`, which is what the scratch buffers were sized for. + fn block_width(lanes: usize, remaining: usize) -> usize { + let usable = || { + tangent_batch::SUPPORTED_LANES + .into_iter() + .filter(|&w| w <= lanes) + }; + if let Some(width) = usable().find(|&width| width <= remaining) { + return width; + } + match usable().min() { + Some(narrowest) if remaining >= Self::MIN_PADDED_TAIL => narrowest, + _ => 1, + } + } + + /// One colour through the scalar tangent section, reusing the primal region + /// the batched sweep already filled. Only ever the tail of the block loop, so + /// no later block reads the tangent slots this overwrites. + fn sweep_one_color_scalar( + &self, + scratch: &mut JacobianScratch, + layout: &JacobianLayout, + color: usize, + out: &mut [f64], + ) { + let JacobianScratch { tape, seed, .. } = scratch; + seed.fill(0.0); + for &col in self.coloring.columns_with_color(color) { + seed[self.global_column(col)] = 1.0; + } + let result = self + .assembly_jvp_expr + .run_tangent_section(tape, &self.seed_as_tangent(seed)); + for entry in &layout.color_to_slots[color] { + out[entry.csc_idx] = result[entry.row]; + } + } + + /// Write `value` into the seed positions of the `carried` colours starting at + /// `block`, at that block's lane `stride`, leaving every other slot untouched. + /// Called with 1.0 before the sweep and 0.0 after, which is what keeps + /// `seed_lanes` zero between blocks without a memset. + fn write_seed_lanes( + &self, + scratch: &mut JacobianScratch, + block: usize, + carried: usize, + stride: usize, + value: f64, + ) { + for lane in 0..carried { + for &col in self.coloring.columns_with_color(block + lane) { + scratch.seed_lanes[self.global_column(col) * stride + lane] = value; + } + } + } + + /// Dispatch the batched tangent sweep to its monomorphised lane width. + fn run_tangent_lanes<'s>(&self, scratch: &'s mut JacobianScratch, lanes: usize) -> &'s [f64] { + let ir = self.assembly_jvp_expr.ir(); + match lanes { + 8 => tangent_batch::run_tangent_batch::<8>( + ir, + &scratch.tape, + &mut scratch.tangent_lanes, + &scratch.seed_lanes, + ), + 4 => tangent_batch::run_tangent_batch::<4>( + ir, + &scratch.tape, + &mut scratch.tangent_lanes, + &scratch.seed_lanes, + ), + other => unreachable!("unsupported jacobian lane width {other}"), + } + } + + /// Point the seed buffer at whichever tangent input this artifact seeds, the + /// one place [`DiffTarget`] decides anything at assembly time. + #[inline] + const fn seed_as_tangent<'s>(&self, seed: &'s [f64]) -> TangentInputs<'s> { + match self.wrt { + DiffTarget::States => TangentInputs { + dy: Some(seed), + dp: None, + }, + DiffTarget::Params => TangentInputs { + dy: None, + dp: Some(seed), + }, + } + } + + /// Matrix-free `J . v` into `out`, through the unpruned tangent tape. + /// + /// `v` is one value per column of this Jacobian. A subset build's columns are + /// local while its tangent nodes index global state positions, so the seed is + /// scattered through `col_to_global` first; an artifact owning every column + /// seeds from `v` directly, with no copy. + /// + /// # Panics + /// Panics if `v` is shorter than [`Self::n_cols`]. + #[allow(clippy::too_many_arguments)] + pub fn action_into( + &self, + scratch: &mut JacobianScratch, + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + v: &[f64], + out: &mut [f64], + ) { + assert!( + v.len() >= self.n_cols, + "jacobian tangent buffer too small: need {}, got {}", + self.n_cols, + v.len() + ); + if !self.col_to_global.is_empty() { + scratch.seed.fill(0.0); + for (col, &seed) in v.iter().enumerate().take(self.n_cols) { + scratch.seed[self.global_column(col)] = seed; + } + } + let JacobianScratch { tape, seed, .. } = scratch; + let seed = if self.col_to_global.is_empty() { + v + } else { + seed + }; + let result = + self.jvp_expr + .eval_with_tangent(tape, t, y, y_dot, inputs, &self.seed_as_tangent(seed)); + out[..result.len()].copy_from_slice(result); + } +} + +/// Build a row→[(col, `csc_idx`)] lookup table from a `CscPattern`. +pub fn build_row_to_csc_entries(csc: &CscPattern) -> Vec> { + let mut row_to_csc_entries = vec![Vec::new(); csc.nrows]; + for (csc_idx, &(row, col)) in csc.csc_to_csr_map.iter().enumerate() { + row_to_csc_entries[row].push((col, csc_idx)); + } + row_to_csc_entries +} + +/// Map each CSR non-zero to its CSC slot index. +pub fn build_csr_to_csc_map( + csr: &SparsityPattern, + row_to_csc_entries: &[Vec<(usize, usize)>], +) -> Vec { + let mut map = vec![0usize; csr.nnz()]; + for (row, csc_entries) in row_to_csc_entries.iter().enumerate().take(csr.nrows) { + let row_start = csr.indptr[row]; + let row_end = csr.indptr[row + 1]; + let mut csc_pos = 0usize; + for (csr_idx, map_entry) in map.iter_mut().enumerate().take(row_end).skip(row_start) { + let col = csr.indices[csr_idx]; + while csc_entries[csc_pos].0 < col { + csc_pos += 1; + } + assert_eq!( + csc_entries[csc_pos].0, col, + "missing CSC slot for CSR entry ({row}, {col})" + ); + *map_entry = csc_entries[csc_pos].1; + } + } + map +} + +/// Build per-color scatter entries from a sparsity pattern, CSR→CSC map, and coloring. +/// +/// Only entries `swept` marks are bucketed. An entry a sweep does not recover +/// MUST be left out: for a split dense row several same-colour columns land on +/// it and bucketing would alias their sums, and for a constant entry the sweep +/// value is polluted by design. +fn build_color_scatter_entries( + jac_y_sparsity: &SparsityPattern, + jac_y_csr_to_csc_map: &[usize], + coloring: &ColumnColoring, + swept: &[bool], +) -> Vec> { + let mut color_to_csc_entries = vec![Vec::new(); coloring.n_colors]; + for (csr_idx, row) in jac_y_sparsity.entry_rows().into_iter().enumerate() { + if !swept[csr_idx] { + continue; + } + let color = coloring.colors[jac_y_sparsity.indices[csr_idx]]; + color_to_csc_entries[color].push(ColorScatterEntry { + csc_idx: jac_y_csr_to_csc_map[csr_idx], + row, + }); + } + color_to_csc_entries +} + +/// Remap `(csr_idx, value)` entries into a CSC ordering. +fn map_entries_to_csc(entries: &[(usize, f64)], csr_to_csc: &[usize]) -> Vec<(usize, f64)> { + entries + .iter() + .map(|&(csr_idx, value)| (csr_to_csc[csr_idx], value)) + .collect() +} + +/// Keep only the entries outside `rows`. A split dense row is filled wholesale +/// by its own tape, so a table write there would be overwritten anyway. +fn retain_outside_rows(pattern: &SparsityPattern, rows: &[usize], entries: &mut Vec<(usize, f64)>) { + if rows.is_empty() || entries.is_empty() { + return; + } + let mut dropped = vec![false; pattern.nnz()]; + for &row in rows { + dropped[pattern.indptr[row]..pattern.indptr[row + 1]].fill(true); + } + entries.retain(|&(csr_idx, _)| !dropped[csr_idx]); +} + +/// Filter a sparsity pattern to retain only the columns in `allowed_columns`. +fn filter_sparsity_columns( + pattern: &SparsityPattern, + allowed_columns: &[usize], +) -> SparsityPattern { + let mut filtered = SparsityPattern::new(pattern.nrows, pattern.ncols); + let mut allowed = vec![false; pattern.ncols]; + for &col in allowed_columns { + allowed[col] = true; + } + + for row in 0..pattern.nrows { + let row_start = pattern.indptr[row]; + let row_end = pattern.indptr[row + 1]; + filtered.indptr[row] = filtered.indices.len(); + for &col in &pattern.indices[row_start..row_end] { + if allowed[col] { + filtered.indices.push(col); + } + } + } + filtered.indptr[pattern.nrows] = filtered.indices.len(); + filtered +} + +/// Rows this wide force column coloring toward one color per touched column +/// (max-row-nnz is the coloring lower bound), so they are split out. +pub(crate) const DENSE_ROW_MIN_NNZ: usize = 16; +/// Safety cap on reverse passes, so a pathological pattern cannot make the build +/// allocate an unbounded number of adjoint tapes. +pub(crate) const MAX_DENSE_ROWS: usize = 1024; +/// Rows per shared adjoint tape. +/// +/// Every row of a group walks its group's whole tape backwards, so a group +/// trades assembly time for compiled memory. Measured on the 12x12 pouch cell +/// (one assembly / total tape): 1 row 66 ms / 7.8M instrs, 4 rows 58 ms / 2.4M, +/// 16 rows 85 ms / 1.0M, all 144 rows 215 ms / 283k. Four is the knee. +pub(crate) const ROWS_PER_TAPE: usize = 4; + +/// The `k` widest rows, with `k` chosen to minimise what the Jacobian costs. +/// +/// Splitting the `k` widest leaves a colouring bound of `nnz(wide[k])` and buys it +/// for `k` reverse passes, so the sweep count is `k + nnz(wide[k])`, against +/// `nnz(wide[0])` unsplit. Taking the minimum covers both a lone outlier row and a +/// plateau of equally wide rows, where splitting part of it lowers nothing at all +/// and only splitting the whole plateau pays -- a 2D current collector puts one +/// dense row per collector node, so the plateau is the common case, not the +/// pathological one. Past `nnz(wide[0])` passes a split costs more than colouring +/// everything, so the objective caps itself and `MAX_DENSE_ROWS` only bounds how +/// many tapes a build will construct. +/// +/// `row_widths` counts only the entries a sweep must recover, so a row of +/// known entries costs no colours and is never worth a reverse pass. +pub(crate) fn detect_dense_rows(row_widths: &[usize]) -> Vec { + let nnz = |row: usize| row_widths[row]; + let mut wide: Vec = (0..row_widths.len()) + .filter(|&row| nnz(row) >= DENSE_ROW_MIN_NNZ) + .collect(); + wide.sort_unstable_by_key(|&row| std::cmp::Reverse(nnz(row))); + + let bound_after = |k: usize| wide.get(k).map_or(0, |&row| nnz(row)); + let unsplit = bound_after(0); + let cap = wide + .len() + .min(MAX_DENSE_ROWS) + .min(unsplit.saturating_sub(1)); + let mut best_cost = unsplit; + let mut best_k = 0; + for k in 1..=cap { + let cost = k + bound_after(k); + if cost < best_cost { + best_cost = cost; + best_k = k; + } + } + wide.truncate(best_k); + wide.sort_unstable(); + wide +} + +/// Replace selected scalar output rows with zeros while preserving output width. +fn mask_scalar_rows(arena: &mut Arena, root: NodeId, rows: &[usize]) -> Option { + if rows.is_empty() { + return Some(root); + } + if !rows.windows(2).all(|pair| pair[0] < pair[1]) { + return None; + } + let mut lens = NodeMap::new(arena.len()); + mask_scalar_rows_inner(arena, root, rows, &mut lens) +} + +fn mask_scalar_rows_inner( + arena: &mut Arena, + id: NodeId, + rows: &[usize], + lens: &mut NodeMap, +) -> Option { + if rows.is_empty() { + return Some(id); + } + match arena.get(id).clone() { + Node::Concat(children) => { + let mut offset = 0; + let mut masked = Vec::with_capacity(children.len()); + for child in children { + let len = node_len(arena, child, lens); + let start = rows.partition_point(|&row| row < offset); + let end = rows.partition_point(|&row| row < offset + len); + let local_rows: Vec = + rows[start..end].iter().map(|&row| row - offset).collect(); + masked.push(mask_scalar_rows_inner(arena, child, &local_rows, lens)?); + offset += len; + } + if rows.iter().any(|&row| row >= offset) { + return None; + } + Some(arena.alloc(Node::Concat(masked))) + }, + _ if rows == [0] && node_len(arena, id, lens) == 1 => Some(arena.alloc(Node::Scalar(0.0))), + _ => None, + } +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use crate::arena::Arena; + use crate::coloring::color_columns; + use crate::node::Node; + + /// f(y, p) = [y0*y1 + p0, sin(y0) * p1], mixed state and parameter dependencies. + fn toy(arena: &mut Arena) -> NodeId { + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let p0 = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let p1 = arena.alloc(Node::InputParameter { + name: "b".into(), + index: 1, + offset: 1, + width: 1, + }); + let prod = arena.alloc(Node::Mul(y0, y1)); + let r0 = arena.alloc(Node::Add(prod, p0)); + let s = arena.alloc(Node::Sin(y0)); + let r1 = arena.alloc(Node::Mul(s, p1)); + arena.alloc(Node::Concat(vec![r0, r1])) + } + + fn assemble_full(jac: &JacobianData, t: f64, y: &[f64], p: &[f64]) -> Vec> { + let layout = jac.layout(); + let mut scratch = JacobianScratch::new(jac); + let mut data = vec![0.0; layout.n_slots()]; + jac.assemble_into(&mut scratch, layout, t, y, &[], p, &mut data); + // Expand CSC data into a dense matrix for assertion. + let mut dense = vec![vec![0.0; jac.n_cols()]; jac.n_rows()]; + for (col, span) in jac.csc().colptr.windows(2).enumerate() { + for k in span[0]..span[1] { + dense[jac.csc().rowind[k]][col] = data[k]; + } + } + dense + } + + /// Nonlinear banded expression over `n` states with half-width `w`: enough + /// colours to batch, and a tangent tape of nothing but batchable + /// instructions. Shared with `model`'s tests, whose batching fixtures must be + /// the same shape as these to be comparable. + pub fn banded_nonlinear(n: usize, w: usize) -> (Arena, NodeId) { + let mut arena = Arena::new(); + let states: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + let rows: Vec<_> = (0..n) + .map(|i| { + let mut acc = arena.alloc(Node::Exp(states[i])); + for offset in 1..=w { + let left = arena.alloc(Node::Sin(states[i.saturating_sub(offset)])); + let right = arena.alloc(Node::Cos(states[(i + offset).min(n - 1)])); + let pair = arena.alloc(Node::Mul(left, right)); + acc = arena.alloc(Node::Add(acc, pair)); + } + acc + }) + .collect(); + let root = arena.alloc(Node::Concat(rows)); + (arena, root) + } + + fn assemble_with(jac: &JacobianData, scratch: &mut JacobianScratch, y: &[f64]) -> Vec { + let layout = jac.layout(); + let mut out = vec![f64::NAN; layout.n_slots()]; + jac.assemble_into(scratch, layout, 0.5, y, &[], &[], &mut out); + out + } + + /// The batched sweep is the scalar one's arithmetic re-laned, so it must agree + /// bit for bit. With one driver this pins every consumer at once, not just the + /// one a test picks. + #[allow(clippy::float_cmp)] // exact equality is the point + fn assert_batched_matches_scalar(jac: &JacobianData, y: &[f64]) { + assert!(jac.lane_width() > 1, "the fixture must batch"); + let batched = assemble_with(jac, &mut JacobianScratch::new(jac), y); + let scalar = assemble_with(jac, &mut JacobianScratch::scalar(jac), y); + assert_eq!(batched, scalar); + } + + fn batching_probe_states(n: usize) -> Vec { + (0..n).map(|i| (i as f64).mul_add(0.037, 0.41)).collect() + } + + #[test] + fn a_batched_sweep_matches_the_scalar_one_bit_for_bit() { + let (n, w) = (64, 3); + let (arena, root) = banded_nonlinear(n, w); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n); + assert_batched_matches_scalar(&jac, &batching_probe_states(n)); + } + + /// A layout onto a wider buffer is a remap, not a second derivation: the + /// derivative values land unchanged, and the slots this pattern has no entry + /// for are left for their own owner. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point + fn a_merged_layout_only_renumbers_the_slots() { + let (n, w) = (32, 2); + let (arena, root) = banded_nonlinear(n, w); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n); + + // A dense row per state, as a mass matrix that couples everything would + // give: every merged row is wider than the derivative's, so no slot + // number survives by luck. + let mut extra = SparsityPattern::new(n_rows, n); + extra.indptr = (0..=n_rows).map(|row| row * n).collect(); + extra.indices = (0..n_rows).flat_map(|_| 0..n).collect(); + let mut merged = jac.sparsity().clone(); + merged.merge_with(&extra); + let merged_csc = CscPattern::from_csr(&merged); + let merged_layout = jac.layout_in(&merged_csc); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.041, 0.23)).collect(); + let own = assemble_with(&jac, &mut JacobianScratch::new(&jac), &y); + let mut wide = vec![f64::NAN; merged_layout.n_slots()]; + jac.assemble_into( + &mut JacobianScratch::new(&jac), + &merged_layout, + 0.5, + &y, + &[], + &[], + &mut wide, + ); + + // Every derivative entry, matched (row, col) to (row, col) across the two + // slot numberings. + let mut own_by_entry = std::collections::HashMap::new(); + for (slot, &(row, col)) in jac.csc().csc_to_csr_map.iter().enumerate() { + own_by_entry.insert((row, col), own[slot]); + } + let mut covered = 0; + for (slot, &(row, col)) in merged_csc.csc_to_csr_map.iter().enumerate() { + match own_by_entry.get(&(row, col)) { + Some(&value) => { + assert_eq!(wide[slot], value, "entry ({row}, {col}) moved value"); + covered += 1; + }, + // A slot only the other pattern owns: zeroed, never written. + None => assert_eq!(wide[slot], 0.0, "slot ({row}, {col}) is not ours to write"), + } + } + assert_eq!(covered, jac.nnz(), "every derivative entry must be placed"); + } + + /// A subset build seeds global state positions through `col_to_global`, so its + /// batched lane seeds must land there too -- the one place the two sweeps + /// could disagree about which column a lane carries. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point + fn a_subset_build_batches_onto_the_same_values() { + let (n, w) = (64, 3); + let (arena, root) = banded_nonlinear(n, w); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let subset: Vec = (0..n).collect(); + let jac = JacobianData::new_wrt_state_subset(&arena, root, n_rows, n, &subset); + let y = batching_probe_states(n); + assert_batched_matches_scalar(&jac, &y); + + // And the same derivative the full-state build produces. + let subset_values = assemble_with(&jac, &mut JacobianScratch::new(&jac), &y); + let full = JacobianData::new_wrt_states(&arena, root, n_rows, n); + let reference = assemble_with(&full, &mut JacobianScratch::new(&full), &y); + assert_eq!(jac.csc().csc_to_csr_map, full.csc().csc_to_csr_map); + for (k, (&got, &want)) in subset_values.iter().zip(&reference).enumerate() { + assert!( + (got - want).abs() <= 1e-12 * want.abs().mul_add(1.0, 1.0), + "slot {k}: subset {got} vs full {want}" + ); + } + } + + #[test] + fn jacobian_wrt_states_matches_analytic() { + let mut arena = Arena::new(); + let root = toy(&mut arena); + let jac = JacobianData::new_wrt_states(&arena, root, 2, 2); + let (y, p) = ([0.5, 2.0], [3.0, 4.0]); + let dense = assemble_full(&jac, 0.0, &y, &p); + // df0/dy = [y1, y0]; df1/dy = [cos(y0)*p1, 0] + assert!((dense[0][0] - 2.0).abs() < 1e-12); + assert!((dense[0][1] - 0.5).abs() < 1e-12); + assert!(0.5f64.cos().mul_add(-4.0, dense[1][0]).abs() < 1e-12); + assert!(dense[1][1].abs() < 1e-12); + } + + #[test] + fn jacobian_wrt_params_matches_analytic() { + let mut arena = Arena::new(); + let root = toy(&mut arena); + let jac = JacobianData::new_wrt_params(&arena, root, 2, 2); + let (y, p) = ([0.5, 2.0], [3.0, 4.0]); + let dense = assemble_full(&jac, 0.0, &y, &p); + // df0/dp = [1, 0]; df1/dp = [0, sin(y0)] + assert!((dense[0][0] - 1.0).abs() < 1e-12); + assert!(dense[0][1].abs() < 1e-12); + assert!(dense[1][0].abs() < 1e-12); + assert!((dense[1][1] - 0.5f64.sin()).abs() < 1e-12); + } + + #[test] + fn jacobian_wrt_states_is_rectangular_for_partial_group() { + // Single output row over 2 states: shape (1, 2) + 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 root = arena.alloc(Node::Mul(y0, y1)); + let jac = JacobianData::new_wrt_states(&arena, root, 1, 2); + assert_eq!((jac.n_rows(), jac.n_cols()), (1, 2)); + } + + #[test] + fn jacobian_wrt_state_subset_analytic() { + // f(y) = [y0*y2 + y1], so df/dy = [y2, 1, y0]. subset = [0, 2] maps local + // cols 0,1 to global 0,2; at y=[2,3,5] df/dy1 = 1 must not appear. + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let prod = arena.alloc(Node::Mul(y0, y2)); + let root = arena.alloc(Node::Add(prod, y1)); + + let jac = JacobianData::new_wrt_state_subset(&arena, root, 1, 3, &[0, 2]); + + // Shape checks. + assert_eq!((jac.n_rows(), jac.n_cols()), (1, 2)); + // seed_len must equal full state dimension, not subset size. + assert_eq!(jac.seed_len(), 3); + + let dense = assemble_full(&jac, 0.0, &[2.0, 3.0, 5.0], &[]); + // col 0 (global 0): df/dy0 = y2 = 5 + assert!( + (dense[0][0] - 5.0).abs() < 1e-12, + "df/dy0 wrong: {}", + dense[0][0] + ); + // col 1 (global 2): df/dy2 = y0 = 2 + assert!( + (dense[0][1] - 2.0).abs() < 1e-12, + "df/dy2 wrong: {}", + dense[0][1] + ); + // df/dy1 = 1 must not appear in the 2-column result. + assert_eq!( + jac.n_cols(), + 2, + "subset Jacobian must have exactly 2 columns" + ); + } + + #[test] + fn jacobian_wrt_params_param_independent_row() { + // f = [y0*y1, y0*p0]: row 0 depends on no parameter, so df/dp must give + // (0,0) == 0 and (1,0) == y0 with no read past a collapsed tangent tape. + 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 p0 = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let r0 = arena.alloc(Node::Mul(y0, y1)); + let r1 = arena.alloc(Node::Mul(y0, p0)); + let root = arena.alloc(Node::Concat(vec![r0, r1])); + + let jac = JacobianData::new_wrt_params(&arena, root, 2, 1); + let dense = assemble_full(&jac, 0.0, &[3.0, 5.0], &[7.0]); + assert!( + dense[0][0].abs() < 1e-12, + "df0/dp0 must be 0, got {}", + dense[0][0] + ); + assert!( + (dense[1][0] - 3.0).abs() < 1e-12, + "df1/dp0 must be y0=3, got {}", + dense[1][0] + ); + } + + #[test] + fn jacobian_wrt_states_state_independent_leading_row() { + // f = concat(const_vec[1,2], y0): when the leading constant child collapses + // in the tangent tape, y0's derivative must stay in row 2, not shift to 0. + use crate::node::{ArrayData, Shape}; + let mut arena = Arena::new(); + let cv = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0], + shape: Shape::vector(2), + }))); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![cv, y0])); + + let jac = JacobianData::new_wrt_states(&arena, root, 3, 1); + let dense = assemble_full(&jac, 0.0, &[9.0], &[]); + assert!( + dense[0][0].abs() < 1e-12, + "row 0 must be 0, got {}", + dense[0][0] + ); + assert!( + dense[1][0].abs() < 1e-12, + "row 1 must be 0, got {}", + dense[1][0] + ); + assert!( + (dense[2][0] - 1.0).abs() < 1e-12, + "df2/dy0 must be 1, got {}", + dense[2][0] + ); + } + + #[test] + #[should_panic(expected = "subset must be strictly ascending")] + fn jacobian_wrt_state_subset_rejects_unsorted() { + 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 root = arena.alloc(Node::Add(y0, y1)); + // Reversed subset, must panic. + let _ = JacobianData::new_wrt_state_subset(&arena, root, 1, 2, &[1, 0]); + } + + /// `nrows` x `ncols` pattern: row 0 fully dense (`ncols` nnz), every + /// other row carries a single diagonal-ish nonzero. + fn make_pattern_with_dense_row(ncols: usize, nrows: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(nrows, ncols); + let mut idx = 0; + for row in 0..nrows { + pattern.indptr[row] = idx; + if row == 0 { + pattern.indices.extend(0..ncols); + idx += ncols; + } else { + pattern.indices.push((row - 1) % ncols); + idx += 1; + } + } + pattern.indptr[nrows] = idx; + pattern + } + + /// `n` x `n` identity-like pattern: one nonzero per row, on the diagonal. + fn make_diagonal_pattern_local(n: usize) -> SparsityPattern { + let mut pattern = SparsityPattern::new(n, n); + for row in 0..n { + pattern.indptr[row] = row; + pattern.indices.push(row); + } + pattern.indptr[n] = n; + pattern + } + + /// `ncols` x `ncols` pattern with the first `n_dense` rows fully dense + /// and the remainder carrying a single nonzero each. + fn make_pattern_with_n_dense_rows(ncols: usize, n_dense: usize) -> SparsityPattern { + let nrows = ncols; + let mut pattern = SparsityPattern::new(nrows, ncols); + let mut idx = 0; + for row in 0..nrows { + pattern.indptr[row] = idx; + if row < n_dense { + pattern.indices.extend(0..ncols); + idx += ncols; + } else { + pattern.indices.push(row % ncols); + idx += 1; + } + } + pattern.indptr[nrows] = idx; + pattern + } + + #[test] + fn detect_dense_rows_flags_wide_rows_only() { + // 20 cols; row 0 dense (20 nnz), rows 1..19 diagonal-ish (1 nnz) + let pattern = make_pattern_with_dense_row(20, 20); + assert_eq!(detect_dense_rows(&pattern.row_widths()), vec![0]); + // all-sparse pattern → no detection + let diag = make_diagonal_pattern_local(20); + assert!(detect_dense_rows(&diag.row_widths()).is_empty()); + } + + /// A 2D current collector puts one dense row per collector node, so the + /// widest rows are a plateau of equal widths rather than a lone outlier. + /// Splitting part of a plateau lowers nothing, so a heuristic that stops at a + /// fixed handful splits none of it and leaves colouring at one colour per + /// column -- the shape that took a 2818-state pouch model to 2816 colours. + #[test] + fn detect_dense_rows_splits_a_whole_plateau() { + for n_dense in [5, 8, 40] { + let pattern = make_pattern_with_n_dense_rows(64, n_dense); + let widths = pattern.row_widths(); + let split = detect_dense_rows(&widths); + assert_eq!( + split, + (0..n_dense).collect::>(), + "the whole plateau of {n_dense} dense rows must be split" + ); + // and the split has to actually pay for itself + let unsplit = widths.iter().copied().max().unwrap(); + let remaining = (n_dense..pattern.nrows) + .map(|row| widths[row]) + .max() + .unwrap_or(0); + assert!( + split.len() + remaining < unsplit, + "{n_dense} reverse passes + {remaining} colours must beat {unsplit}" + ); + } + } + + /// Splitting cannot pay when every row is equally dense: `k` reverse passes + /// leave the bound untouched, so plain colouring stays the cheaper option. + #[test] + fn detect_dense_rows_declines_when_no_split_pays() { + let all_dense = make_pattern_with_n_dense_rows(20, 20); + assert!(detect_dense_rows(&all_dense.row_widths()).is_empty()); + } + + /// The number of reverse passes stays bounded however pathological the + /// pattern, so a build cannot be made to allocate unbounded adjoint tapes. + #[test] + fn detect_dense_rows_respects_the_tape_cap() { + let ncols = MAX_DENSE_ROWS + 64; + let pattern = make_pattern_with_n_dense_rows(ncols, MAX_DENSE_ROWS + 32); + assert!(detect_dense_rows(&pattern.row_widths()).len() <= MAX_DENSE_ROWS); + } + + #[test] + fn detect_dense_rows_splits_outlier_among_many_wide_rows() { + // The outlier alone sets the coloring bound, so it must still be split + // when the count of merely-wide rows runs far past MAX_DENSE_ROWS. + let ncols = 64; + let nrows = 40; + let mut pattern = SparsityPattern::new(nrows, ncols); + let mut idx = 0; + for row in 0..nrows { + pattern.indptr[row] = idx; + let width = if row == 7 { ncols } else { DENSE_ROW_MIN_NNZ }; + pattern + .indices + .extend((0..width).map(|c| (row + c) % ncols)); + idx += width; + } + pattern.indptr[nrows] = idx; + + assert_eq!(detect_dense_rows(&pattern.row_widths()), vec![7]); + // and dropping it genuinely lowers the coloring bound + let full = color_columns(&pattern); + let mut swept = vec![true; pattern.nnz()]; + swept[pattern.indptr[7]..pattern.indptr[8]].fill(false); + let reduced = color_columns_masked(&pattern, &swept); + assert!( + reduced.n_colors < full.n_colors, + "reduced {} should beat full {}", + reduced.n_colors, + full.n_colors + ); + } + + #[test] + fn mask_scalar_rows_preserves_nested_output_shape() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let b = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let c = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let inner = arena.alloc(Node::Concat(vec![a, b])); + let root = arena.alloc(Node::Concat(vec![inner, c])); + + let masked = mask_scalar_rows(&mut arena, root, &[2, 3]).expect("scalar rows"); + let ir = TypedIr::from_arena(&arena, masked); + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let result = expr.eval(&mut scratch, 0.0, &[1.0, 2.0, 3.0, 4.0], &[], &[]); + assert_eq!(result, &[1.0, 2.0, 0.0, 0.0]); + } + + #[test] + fn mask_scalar_rows_rejects_vector_interior_and_unsorted_rows() { + let mut arena = Arena::new(); + let vector = arena.alloc(Node::StateVector { start: 0, end: 3 }); + assert!(mask_scalar_rows(&mut arena, vector, &[1]).is_none()); + assert!(mask_scalar_rows(&mut arena, vector, &[2, 1]).is_none()); + } + + #[test] + fn dense_rows_are_pruned_only_from_assembly_jvp() { + let mut arena = Arena::new(); + let n = 20; + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let square = arena.alloc(Node::Mul(y, y)); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + crate::node::CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + crate::node::Shape::matrix(1, n), + ) + .expect("valid row matrix"), + ))); + let dense = arena.alloc(Node::MatMul(ones, square)); + let mut rows = vec![dense]; + for i in 1..n { + rows.push(arena.alloc(Node::Index { + child: y, + start: i, + end: i + 1, + })); + } + let root = arena.alloc(Node::Concat(rows)); + let jac = JacobianData::new_wrt_states(&arena, root, n, n); + assert_eq!(jac.n_dense_rows(), 1); + // This fixture has no Conditional, so raw and common tape lengths + // coincide; the assertion is about dense-row pruning, not branches. + assert!( + jac.assembly_tape().ir().instructions().len() + < jac.action_tape().ir().instructions().len(), + "dense-row pruning must shorten the raw assembly tape" + ); + + let values: Vec = (1..=n).map(|value| value as f64).collect(); + let seed = vec![1.0; n]; + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let mut full_scratch = vec![0.0; jac.action_tape().scratch_len()]; + let full = jac + .jvp_expr + .eval_with_tangent(&mut full_scratch, 0.0, &values, &[], &[], &tangent) + .to_vec(); + let mut assembly_scratch = vec![0.0; jac.assembly_tape().scratch_len()]; + let assembly = jac.assembly_tape().eval_with_tangent( + &mut assembly_scratch, + 0.0, + &values, + &[], + &[], + &tangent, + ); + + let expected_dense = 2.0 * values.iter().sum::(); + assert_eq!(full[0].to_bits(), expected_dense.to_bits()); + assert_eq!(assembly[0].to_bits(), 0.0f64.to_bits()); + assert!( + assembly[1..] + .iter() + .zip(&full[1..]) + .all(|(&left, &right)| left.to_bits() == right.to_bits()) + ); + } + + #[test] + fn unsplit_jacobian_shares_full_and_assembly_jvp() { + let mut arena = Arena::new(); + let root = toy(&mut arena); + let jac = JacobianData::new_wrt_states(&arena, root, 2, 2); + assert!(Arc::ptr_eq(jac.action_tape(), jac.assembly_tape())); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/lib.rs b/packages/pybamm-rust/pybamm-core/src/lib.rs new file mode 100644 index 0000000000..72990838a5 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/lib.rs @@ -0,0 +1,97 @@ +//! 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 and this crate +//! turns it into flat instruction tapes it can evaluate, differentiate and +//! solve. The stages, each a module below: +//! +//! 1. **Build**: the bindings allocate [`Node`]s into an [`Arena`]; a node is +//! referenced by a `u32` [`NodeId`], so sharing a subexpression is repeating +//! an id. +//! 2. **Rewrite**: `simplify` folds identities and runs CSE/DCE, +//! `zero_propagate` proves subtrees identically zero. +//! 3. **Differentiate**: `tangent` emits forward-mode JVP DAGs; `adjoint` +//! fills one wide Jacobian row from a single backward pass instead. +//! 4. **Lower**: [`TypedIr::from_arena`] flattens a DAG into slot-addressed +//! [`Instruction`]s with constants interned in a [`ConstPool`]. +//! 5. **Evaluate**: [`CompiledExpr`] interprets a tape against a +//! caller-supplied scratch buffer, one time point at a time or `k` lanes +//! at once (`eval`, `eval_batch`). +//! +//! [`CompiledModel`] assembles those pieces into a DAE `M y' = f(t, y; p)` +//! carrying a symbolic `df/dy`, its sparsity pattern and a column coloring, and +//! two consumers drive it: `solver` runs diffsol in-process, while `ffi` +//! exposes a C ABI for the IDAKLU solver to call. +//! +//! Nothing here is thread-confined by construction, but evaluation writes into +//! scratch buffers, so a [`Workspace`] belongs to one solve at a time. +//! [`CompiledModel`] is immutable and shared via `Arc`; [`ModelEvaluator`] +//! pairs one with an owned [`Workspace`] for callers that want a single +//! `&mut self` handle instead, adding only what owning the workspace buys and +//! dereferencing to the model for the rest. + +#![deny(unused_must_use)] + +pub mod adjoint; +pub mod arena; +pub mod branch_regions; +pub mod coloring; +pub mod const_entries; +pub mod error; +pub mod eval; +pub mod eval_batch; +pub mod ffi; +pub mod ir; +pub mod jacobian; +pub mod model; +pub mod node; +pub mod observable; +mod row_extract; +pub mod simplify; +#[cfg(feature = "serialize")] +pub mod snapshot; +pub mod sparsity; +pub mod tangent; +pub mod tangent_batch; +pub mod zero_propagate; + +#[cfg(feature = "diffsol")] +pub mod solver; + +pub use arena::{Arena, NodeId, NodeMap, StateUsage, scan_state_usage}; +pub use branch_regions::{ + BranchLabel, Ownership, RegionGroup, RegionSchedule, active_branch, owner_sets, + privatise_conditionals, schedule_regions, schedule_regions_partitioned, +}; +pub use coloring::{ColumnColoring, color_columns}; +pub use error::CoreError; +pub use eval::{CompiledExpr, PrimalCache, TangentInputs}; +pub use eval_batch::BatchEvalError; +#[cfg(feature = "profile")] +pub use ffi::pybamm_rust_profile_report; +pub use ffi::{ + ERROR_NULL_POINTER, ERROR_PANIC, RUST_ABI_VERSION, SUCCESS, pybamm_rust_abi_version, + pybamm_rust_eval_rhs, pybamm_rust_jac_mul, pybamm_rust_n_inputs, pybamm_rust_n_states, + pybamm_rust_residual, +}; +pub use ir::{ + BinaryOp, BroadcastKind, ConstPool, Instruction, Slot, SlotStats, SplitEvalInfo, TypedIr, + UnaryOp, first_invalid, first_unsupported, infer_sizes, +}; +pub use jacobian::{CscPattern, JacobianData, JacobianLayout, JacobianScratch}; +pub use model::{ + CompiledModel, CompiledModelAlgebraicBlock, CompiledModelOptions, JacobianStats, + JacobianStrategy, ModelEvaluator, Workspace, +}; +pub use node::{ + ArrayData, CsrData, CubicInterpolantData, InterpolantData, NdInterpolantData, Node, Shape, + structural_hash, +}; +pub use observable::{CompiledObservable, ObservableKind, ObservableScratch, ObservableSet}; +pub use row_extract::{ScalarRowBlock, extract_scalar_rows}; +pub use simplify::{SimplifyMode, cse, dce, simplify, simplify_pipeline, simplify_with_mode}; +#[cfg(feature = "serialize")] +pub use snapshot::DagSnapshot; +pub use sparsity::{SparsityPattern, detect_sparsity_per_output}; +pub use tangent::{DiffTarget, tangent_wrt_params, tangent_wrt_states, tangent_wrt_subset}; +pub use zero_propagate::{ShapeInfo, ZeroStatus, zero_propagate}; diff --git a/packages/pybamm-rust/pybamm-core/src/model.rs b/packages/pybamm-rust/pybamm-core/src/model.rs new file mode 100644 index 0000000000..82c914881a --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/model.rs @@ -0,0 +1,4631 @@ +//! `CompiledModel` for DAE system evaluation. +//! +//! `PyBaMM` models are DAE systems of the form `M * y' = f(t, y)` where `M` is +//! a constant mass matrix (often singular), `f(t, y)` is the right-hand side, +//! and `y` is the state vector. Newton iteration computes `J = df/dy - cj * M`. +//! +//! `CompiledModel` holds the primal evaluator, symbolic `df/dy`, mass matrix, +//! and sparsity/coloring information for efficient Jacobian assembly. It is +//! immutable, so evaluation needs a `Workspace` alongside it; `ModelEvaluator` +//! pairs the two into one `&mut self` handle for callers that want that. + +use std::sync::Arc; + +use crate::arena::{Arena, NodeId}; +use crate::coloring::ColumnColoring; +use crate::eval::{CompiledExpr, TangentInputs}; +use crate::ir::TypedIr; +use crate::jacobian::{ + CscPattern, JacobianData, JacobianLayout, JacobianScratch, build_csr_to_csc_map, + build_row_to_csc_entries, +}; +use crate::node::CsrData; +use crate::observable::{ObservableKind, ObservableScratch, ObservableSet, seed_param_tangent}; +use crate::simplify::simplify_pipeline; +use crate::sparsity::SparsityPattern; +use crate::tangent::tangent_wrt_params; + +/// Classification of the mass matrix for dispatch in hot-path operations. +/// +/// Detected once at `CompiledModel::new()` time. Identity and `DiagonalSelector` +/// paths replace the general sparse matvec with O(n) operations. +#[derive(Clone, Debug)] +enum MassKind { + /// M = I (pure ODE, all states differential). + Identity, + /// M = diag(mask): 1 for differential states, 0 for algebraic. + DiagonalSelector(Vec), + /// Arbitrary sparse M, fallback to CSR matvec. + General, +} + +/// Classify a CSR mass matrix into the fastest applicable [`MassKind`]. +/// +/// Returns `Identity` when `M = I`, `DiagonalSelector` when every row has at +/// most one non-zero on the diagonal with value 0 or 1, and `General` +/// otherwise. +// Exact 0.0/1.0 compares classify structural mass entries; an epsilon +// tolerance would misclassify a near-identity matrix as `Identity`. +#[allow(clippy::float_cmp)] +fn classify_mass_matrix(mass: &CsrData) -> MassKind { + let n = mass.shape.rows; + if n != mass.shape.cols { + return MassKind::General; + } + + for row in 0..n { + let start = mass.indptr[row]; + let end = mass.indptr[row + 1]; + let row_nnz = end - start; + + if row_nnz > 1 { + return MassKind::General; + } + if row_nnz == 1 { + if mass.indices[start] != row { + return MassKind::General; + } + let val = mass.data[start]; + if val != 1.0 && val != 0.0 { + return MassKind::General; + } + } + } + + let nnz = mass.indptr[n]; + if nnz == n && mass.data.iter().all(|&v| v == 1.0) { + return MassKind::Identity; + } + + let mut mask = vec![false; n]; + for (row, m) in mask.iter_mut().enumerate() { + let start = mass.indptr[row]; + let end = mass.indptr[row + 1]; + if start < end && mass.data[start] == 1.0 { + *m = true; + } + } + MassKind::DiagonalSelector(mask) +} + +/// Jacobian assembly policy used by [`CompiledModel`]. +/// +/// One variant today, and every model compiles to it. The enum exists so a +/// second policy (a dense or reverse-only assembly) can be added without +/// changing the reporting surface; until one lands there is nothing to select +/// between, so nothing reports a request separately from an outcome. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JacobianStrategy { + /// Assemble through colored JVP passes. + Coloring, +} + +impl JacobianStrategy { + /// Stable name for stats and Python-side reporting. + #[inline] + pub const fn as_str(self) -> &'static str { + match self { + Self::Coloring => "coloring", + } + } +} + +/// Lightweight Jacobian-assembly stats exposed to benchmarks. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct JacobianStats { + /// Strategy this model compiled to. `JacobianStrategy` has a single variant, + /// so there is no request/selection to report separately. + pub strategy: JacobianStrategy, + /// Colors in the adopted coloring, so tangent sweeps per assembly. + pub n_colors: usize, + /// Non-zeros in the assembled Jacobian. + pub nnz: usize, + /// Rows taken out of the coloring and filled by reverse mode instead, so + /// also the reverse passes one assembly adds. + pub n_dense_rows: usize, + /// Rows `detect_dense_rows` nominated. Above `n_dense_rows` means the + /// split was declined — correct, but the colouring stayed wide. + pub n_dense_row_candidates: usize, + /// Entries a compile pass proved constant, written from a table rather + /// than swept. + pub n_constant_entries: usize, + /// Columns a sweep still has to produce, so those the coloring covers. + pub n_swept_columns: usize, + /// Jacobian entries those dense rows account for. + pub dense_row_entries: usize, + /// Instructions across those tapes: the split's compiled-memory cost. + pub dense_row_tape_instructions: usize, + /// Instructions in the shared primal section of the split-eval tape, `None` + /// when the Jacobian tape is not split. + pub split_eval_primal_instructions: Option, + /// Common tape + dispatch: what actually runs regardless of branch. + pub split_eval_total_instructions: Option, + /// Raw tape length, branch blocks included. + pub split_eval_raw_instructions: Option, + /// How many dispatches `split_eval_total_instructions` includes, one per + /// short-circuited conditional, in each of the primal and tangent halves. + pub split_eval_dispatch_count: usize, + /// Per-branch block lengths, in tape order. + pub branch_block_lens: Vec, + /// Colours per batched tangent sweep, or 1 when the tape runs the scalar + /// per-colour path. + pub jac_lane_width: usize, +} + +/// The compiled algebraic sub-block: residual `g(t, y)` and its Jacobian +/// `dg/dy_alg`, for the Newton solver that consistently initialises the +/// algebraic states. +/// +/// One of three [`JacobianData`] adapters, so the block gets the colour sweep, +/// its lane batching and its dense-row handling from the assembly module rather +/// than restating them; what is left here is the block's own descriptor. +#[derive(Debug, Clone)] +struct AlgebraicBlock { + /// `g(t, y)`, compiled through the same simplify pipeline as the residual. + residual: Arc, + /// `dg/dy_alg` over the algebraic columns only. Its layout, its COO triplet + /// and the block's width are all facts about this artifact, so they are read + /// from it rather than cached here where they could drift out of step. + jac: Arc, +} + +/// The algebraic sub-block to compile alongside the right-hand side, for the +/// Newton solver that consistently initialises algebraic states. +#[derive(Debug, Clone, Copy)] +pub struct CompiledModelAlgebraicBlock<'a> { + rhs: NodeId, + var_indices: &'a [usize], +} + +impl<'a> CompiledModelAlgebraicBlock<'a> { + /// `rhs` is the algebraic residual `g(t, y)`, and `var_indices` names the + /// algebraic states by global state index. + #[inline] + pub const fn new(rhs: NodeId, var_indices: &'a [usize]) -> Self { + Self { rhs, var_indices } + } +} + +/// Optional artifacts to compile with a model, so the many-argument constructors +/// do not multiply with every combination. +#[derive(Debug, Clone, Copy, Default)] +#[must_use] +pub struct CompiledModelOptions<'a> { + sens_param_indices: &'a [usize], + algebraic_block: Option>, +} + +impl<'a> CompiledModelOptions<'a> { + /// No sensitivities and no algebraic block, matching a plain + /// [`CompiledModel::new`]. + #[inline] + pub const fn new() -> Self { + Self { + sens_param_indices: &[], + algebraic_block: None, + } + } + + /// Compile `df/dp` for these parameters, named by global parameter index. The + /// order given becomes the sensitivity column order for the whole solve. + #[inline] + pub const fn with_sensitivities(mut self, sens_param_indices: &'a [usize]) -> Self { + self.sens_param_indices = sens_param_indices; + self + } + + /// Compile the algebraic residual and its Jacobian as well as the RHS. + #[inline] + pub const fn with_algebraic( + mut self, + algebraic_block: CompiledModelAlgebraicBlock<'a>, + ) -> Self { + self.algebraic_block = Some(algebraic_block); + self + } + + #[inline] + const fn sens_param_indices(self) -> &'a [usize] { + self.sens_param_indices + } + + #[inline] + const fn algebraic_block(self) -> Option> { + self.algebraic_block + } +} + +/// Per-solve mutable scratch for evaluating a [`CompiledModel`]. +/// +/// One evaluation buffer per `CompiledExpr` plus auxiliary buffers, created +/// fresh per solve so repeated solves never share mutable state. +#[derive(Debug, Clone)] +pub struct Workspace { + pub(crate) primal_scratch: Vec, + /// Every buffer the `df/dy` assembly and its matrix-free actions need, + /// sized and lane-widthed by the assembly module. + pub(crate) jac_scratch: JacobianScratch, + pub(crate) sens_scratch: Option>, + /// Buffers for the two observable families, each sized by its own set. + pub(crate) output_scratch: ObservableScratch, + pub(crate) event_scratch: ObservableScratch, + pub(crate) algebraic_scratch: Option>, + /// The algebraic block's own assembly buffers; `None` without a block. + pub(crate) algebraic_jac_scratch: Option, + pub(crate) cj: f64, + pub(crate) mv_buffer: Vec, + pub(crate) sens_dp_buffer: Vec, + /// Counts `sens_primal_pass` calls so tests can pin the batching mechanism. + #[cfg(test)] + pub(crate) sens_primal_passes: usize, +} + +impl Workspace { + /// Set the cj coefficient used by Jacobian/residual assembly. + #[inline] + pub const fn set_cj(&mut self, cj: f64) { + self.cj = cj; + } + + /// Current cj coefficient. + #[inline] + pub const fn cj(&self) -> f64 { + self.cj + } + + /// The buffers paired with the observables of `kind`, alongside the + /// parameter-tangent buffer their sens action seeds. + /// + /// Handed out together because the borrow checker cannot see through two + /// separate accessors that both take `&mut self`. + #[inline] + fn observable_tangent(&mut self, kind: ObservableKind) -> (&mut ObservableScratch, &mut [f64]) { + let scratch = match kind { + ObservableKind::Outputs => &mut self.output_scratch, + ObservableKind::Events => &mut self.event_scratch, + }; + (scratch, &mut self.sens_dp_buffer) + } + + /// The buffers paired with the observables of `kind`. + #[inline] + fn observable_scratch(&mut self, kind: ObservableKind) -> &mut ObservableScratch { + self.observable_tangent(kind).0 + } +} + +/// A compiled DAE system: primal expression, symbolic Jacobian, mass matrix. +/// +/// Immutable after construction, so it is shared via `Arc` and read by any +/// number of concurrent solves. Mutable scratch lives in a [`Workspace`] created +/// via [`Self::create_workspace`]; [`ModelEvaluator`] bundles the two for +/// callers that would rather hold one `&mut self` handle. +#[derive(Clone)] +pub struct CompiledModel { + primal_expr: Arc, + jac: Arc, + /// Mass matrix M in CSR format. + mass_matrix: CsrData, + mass_kind: MassKind, + sparsity: SparsityPattern, + csc_sparsity: CscPattern, + n_states: usize, + n_params: usize, + /// `jac`'s entries on the MERGED CSC this model assembles into, which carries + /// the mass pattern's slots too. The remap is all this model adds to the + /// assembly module's tables. + jac_layout: JacobianLayout, + /// Mass-matrix CSR entry to CSC slot mapping. + mass_to_csc_map: Vec, + algebraic_ids: Vec, + sens_expr: Option>, + /// Global parameter index for each sensitivity (range `0..n_params`). + sens_param_indices: Vec, + /// Output variables H(t, y; p) and their tangent tapes. + outputs: ObservableSet, + /// Event functions g(t, y; p) and their tangent tapes. Fused into one tape + /// once built, since both hot loops evaluate every event each step. + events: ObservableSet, + algebraic: Option, +} + +impl CompiledModel { + /// Create a new `CompiledModel` from a primal expression and mass matrix. + /// + /// Builds the symbolic derivative df/dy, detects sparsity, and computes + /// coloring for efficient Jacobian assembly. + pub fn new( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + ) -> Self { + // The rhs is compiled raw. Any value-preserving-but-ULP-shifting fold + // stalls IDA's Newton; `simplify`'s int-pow lowering alone moves the + // DFN residual 4096 ULP. `residual_is_compiled_bit_exactly` pins this. + let primal_ir = TypedIr::from_arena(arena, rhs); + let n_outputs = primal_ir.output_len(); + + let jac = Arc::new(JacobianData::new_wrt_states( + arena, rhs, n_outputs, n_states, + )); + + let mass_sparsity = SparsityPattern::from_csr_data(&mass_matrix); + let mut sparsity = jac.sparsity().clone(); + sparsity.merge_with(&mass_sparsity); + + // Convert CSR sparsity to CSC for KLU compatibility + let csc_sparsity = CscPattern::from_csr(&sparsity); + + let row_to_csc_entries = build_row_to_csc_entries(&csc_sparsity); + let mass_to_csc_map = build_csr_to_csc_map(&mass_sparsity, &row_to_csc_entries); + // The one thing this model adds to the assembly module: its own slot + // numbering, since a merged buffer interleaves the mass pattern's entries. + let jac_layout = jac.layout_in(&csc_sparsity); + + // Compile primal expression + let primal_expr = Arc::new(CompiledExpr::from_ir(primal_ir)); + + // Detect algebraic states: a state i is algebraic iff row i of the + // mass matrix has no diagonal entry. Used by IDA's `id` vector. + let algebraic_ids: Vec = (0..n_states) + .map(|i| { + let row_start = mass_matrix.indptr[i]; + let row_end = mass_matrix.indptr[i + 1]; + !mass_matrix.indices[row_start..row_end].contains(&i) + }) + .collect(); + + // Classify mass matrix once for fast-path dispatch in hot-path methods. + let mass_kind = classify_mass_matrix(&mass_matrix); + + Self { + primal_expr, + jac, + mass_matrix, + mass_kind, + sparsity, + csc_sparsity, + n_states, + n_params, + jac_layout, + mass_to_csc_map, + algebraic_ids, + sens_expr: None, + sens_param_indices: Vec::new(), + outputs: ObservableSet::new(), + events: ObservableSet::new(), + algebraic: None, + } + } + + /// Create a `CompiledModel` with forward-sensitivity expressions for + /// the given parameter indices (each must be `< n_params`). + /// + /// All sensitivities reuse a single compiled tangent expression. + /// Panics if any entry in `sens_param_indices` is out of range. + pub fn new_with_sens( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + sens_param_indices: &[usize], + ) -> Self { + Self::new_with_options( + arena, + rhs, + mass_matrix, + n_states, + n_params, + CompiledModelOptions::new().with_sensitivities(sens_param_indices), + ) + } + + /// Build a `CompiledModel` with optional algebraic sub-block expressions. + /// + /// `algebraic_rhs` is the algebraic residual expression g(t, y). + /// `algebraic_variable_indices` lists which state indices are algebraic. + /// + /// When `algebraic_rhs` is `None`, this is equivalent to calling [`Self::new`]. + pub fn new_with_algebraic( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + algebraic_rhs: Option, + algebraic_variable_indices: &[usize], + ) -> Self { + let options = algebraic_rhs.map_or_else(CompiledModelOptions::new, |algebraic_rhs| { + CompiledModelOptions::new().with_algebraic(CompiledModelAlgebraicBlock::new( + algebraic_rhs, + algebraic_variable_indices, + )) + }); + + Self::new_with_options(arena, rhs, mass_matrix, n_states, n_params, options) + } + + /// Build a `CompiledModel` with any combination of sensitivities and an + /// algebraic sub-block. + pub fn new_with_options( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + options: CompiledModelOptions<'_>, + ) -> Self { + let mut model = Self::new(arena, rhs, mass_matrix, n_states, n_params); + model.apply_options(arena, rhs, options); + model + } + + fn apply_options(&mut self, arena: &Arena, rhs: NodeId, options: CompiledModelOptions<'_>) { + self.compile_sensitivities(arena, rhs, options.sens_param_indices()); + if let Some(algebraic_block) = options.algebraic_block() { + self.compile_algebraic_block(arena, algebraic_block.rhs, algebraic_block.var_indices); + } + } + + fn compile_sensitivities(&mut self, arena: &Arena, rhs: NodeId, sens_param_indices: &[usize]) { + if sens_param_indices.is_empty() { + return; + } + let mut seen = vec![false; self.n_params]; + for &idx in sens_param_indices { + assert!( + idx < self.n_params, + "sens_param_indices[{idx}] is out of range for n_params={}", + self.n_params, + ); + assert!( + !std::mem::replace(&mut seen[idx], true), + "sens_param_indices contains a repeated index: {idx}", + ); + } + + let mut diff_arena = arena.clone(); + let df_dp = tangent_wrt_params(&mut diff_arena, rhs); + let (final_arena, df_dp) = simplify_pipeline(diff_arena, df_dp); + let ir = TypedIr::from_arena_split_eval(&final_arena, df_dp); + self.sens_expr = Some(Arc::new(CompiledExpr::from_ir(ir))); + self.sens_param_indices = sens_param_indices.to_vec(); + } + + /// Compile the algebraic residual and its Jacobian, the third adapter onto + /// the assembly module. + /// + /// The Jacobian is a plain [`JacobianData`] over the algebraic columns, so + /// this block gets the colour sweep, its lane batching and the seed-lane + /// hygiene from there; what is compiled here is the residual and the block's + /// COO descriptor. + fn compile_algebraic_block( + &mut self, + arena: &Arena, + algebraic: NodeId, + algebraic_variable_indices: &[usize], + ) { + let (algebraic_arena, algebraic_simplified) = simplify_pipeline(arena.clone(), algebraic); + let algebraic_ir = TypedIr::from_arena(&algebraic_arena, algebraic_simplified); + let n_rows = algebraic_ir.output_len(); + let residual = Arc::new(CompiledExpr::from_ir(algebraic_ir)); + + let jac = Arc::new(JacobianData::new_wrt_state_subset( + arena, + algebraic, + n_rows, + self.n_states, + algebraic_variable_indices, + )); + self.algebraic = Some(AlgebraicBlock { residual, jac }); + } + + /// Scratch length required by the primal `f(t, y)` expression. + #[cfg(test)] + pub fn primal_scratch_len(&self) -> usize { + self.primal_expr.scratch_len() + } + + /// Allocate a fresh [`Workspace`] sized for this model's expressions. + pub fn create_workspace(&self) -> Workspace { + Workspace { + primal_scratch: vec![0.0; self.primal_expr.scratch_len()], + jac_scratch: JacobianScratch::new(&self.jac), + sens_scratch: self.sens_expr.as_ref().map(|e| vec![0.0; e.scratch_len()]), + output_scratch: self.outputs.create_scratch(), + event_scratch: self.events.create_scratch(), + algebraic_scratch: self + .algebraic + .as_ref() + .map(|block| vec![0.0; block.residual.scratch_len()]), + algebraic_jac_scratch: self + .algebraic + .as_ref() + .map(|block| JacobianScratch::new(&block.jac)), + cj: 0.0, + mv_buffer: vec![0.0; self.n_states], + sens_dp_buffer: vec![0.0; self.n_params.max(1)], + #[cfg(test)] + sens_primal_passes: 0, + } + } + + /// Shared handle to the primal `f(t, y)` expression (no recompilation). + #[inline] + pub fn primal_expr_arc(&self) -> Arc { + Arc::clone(&self.primal_expr) + } + + /// Shared handle to the prepared df/dy artifact (no recompilation). + #[inline] + pub fn jacobian_data(&self) -> Arc { + Arc::clone(&self.jac) + } + + /// The observables of `kind`: their tapes, layout and scratch sizing. + #[inline] + pub const fn observables(&self, kind: ObservableKind) -> &ObservableSet { + match kind { + ObservableKind::Outputs => &self.outputs, + ObservableKind::Events => &self.events, + } + } + + /// Shared handle to the algebraic residual expression, if present. + #[inline] + pub fn algebraic_expr_arc(&self) -> Option> { + self.algebraic + .as_ref() + .map(|block| Arc::clone(&block.residual)) + } + + /// Shared handle to the prepared `dg/dy_alg` artifact, if present. The + /// binding's standalone algebraic Jacobian is a view onto this, never a + /// second tangent transform of the same expression. + #[inline] + pub fn algebraic_jacobian_data(&self) -> Option> { + self.algebraic.as_ref().map(|block| Arc::clone(&block.jac)) + } + + /// Length of the state vector this model was compiled against. + #[inline] + pub const fn n_states(&self) -> usize { + self.n_states + } + + /// Number of input parameters declared at compile time. + #[inline] + pub const fn n_params(&self) -> usize { + self.n_params + } + + /// Length of `f(t, y)`, which is `n_states` for a well-formed model. + #[inline] + #[allow(clippy::missing_const_for_fn)] // Calls non-const method + pub fn output_len(&self) -> usize { + self.primal_expr.output_len() + } + + /// Evaluate f(t, y) into `output`. + #[inline] + pub fn eval_rhs( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + let result = self + .primal_expr + .eval(&mut ws.primal_scratch, t, y, &[], inputs); + output[..result.len()].copy_from_slice(result); + } + + /// Compute (df/dy - cj*M) @ v via forward-mode AD and sparse matvec. + pub fn jac_mul( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + output: &mut [f64], + ) { + self.jac_mul_primal(ws, t, y, inputs, v, output); + + // Subtract cj * (M @ v), dispatching on mass kind for fast paths. + if ws.cj != 0.0 { + let n = self.n_states.min(output.len()); + match &self.mass_kind { + MassKind::Identity => { + for i in 0..n { + output[i] -= ws.cj * v[i]; + } + }, + MassKind::DiagonalSelector(mask) => { + for (i, &is_differential) in mask.iter().take(n).enumerate() { + if is_differential { + output[i] -= ws.cj * v[i]; + } + } + }, + MassKind::General => { + csr_matvec(&self.mass_matrix, v, &mut ws.mv_buffer); + let cj = ws.cj; + for (out, &mv) in output.iter_mut().zip(&ws.mv_buffer).take(n) { + *out -= cj * mv; + } + }, + } + } + } + + /// Compute DAE residual r = M*y' - f(t, y). + /// + /// Used by IDAKLU and other DAE solvers. + pub fn eval_residual( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + yp: &[f64], + inputs: &[f64], + r: &mut [f64], + ) { + // Compute f(t, y) first so we can fuse M*y' - f in a single pass. + let f = self + .primal_expr + .eval(&mut ws.primal_scratch, t, y, &[], inputs); + let n = self.n_states.min(r.len()); + match &self.mass_kind { + MassKind::Identity => { + for i in 0..n { + r[i] = yp[i] - f[i]; + } + }, + MassKind::DiagonalSelector(mask) => { + for (i, &is_differential) in mask.iter().take(n).enumerate() { + r[i] = if is_differential { yp[i] } else { 0.0 } - f[i]; + } + }, + MassKind::General => { + csr_matvec(&self.mass_matrix, yp, &mut ws.mv_buffer); + for i in 0..n { + r[i] = ws.mv_buffer[i] - f[i]; + } + }, + } + } + + /// Compute pure Jacobian-vector product: df/dy @ v (no mass term) + /// + /// This is for pybammsolvers ABI which subtracts cj*M@v separately. + pub fn jac_action( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + jv: &mut [f64], + ) { + self.jac_mul_primal(ws, t, y, inputs, v, jv); + } + + /// Compute mass matrix action: M @ v (no beta term). + /// + /// Dispatches to an O(n) fast path for `Identity` and `DiagonalSelector` + /// mass matrices, falling back to a general CSR matvec for `General`. + pub fn mass_action(&self, v: &[f64], mv: &mut [f64]) { + self.mass_action_into(v, mv); + } + + /// Helper: write M @ x into `mv` (no beta). Dispatches on mass kind. + fn mass_action_into(&self, x: &[f64], mv: &mut [f64]) { + match &self.mass_kind { + MassKind::Identity => { + mv[..self.n_states].copy_from_slice(&x[..self.n_states]); + }, + MassKind::DiagonalSelector(mask) => { + for i in 0..self.n_states { + mv[i] = if mask[i] { x[i] } else { 0.0 }; + } + }, + MassKind::General => { + csr_matvec(&self.mass_matrix, x, mv); + }, + } + } + + /// CSR pattern of the assembled Jacobian `df/dy - cj*M`, which is the union + /// of the `df/dy` and mass patterns. + #[inline] + pub const fn sparsity(&self) -> &SparsityPattern { + &self.sparsity + } + + /// CSR pattern of `df/dy` alone, without the mass entries. + #[inline] + pub fn jac_y_sparsity(&self) -> &SparsityPattern { + self.jac.sparsity() + } + + /// Coloring that drives the JVP sweep count, decided once at compile time. + /// With a dense-row split adopted it is the reduced coloring, so it does not + /// cover every row of [`jac_y_sparsity`](Self::jac_y_sparsity). + #[inline] + pub fn coloring(&self) -> &ColumnColoring { + self.jac.coloring() + } + + /// `(csc_idx, value)` in the assembled CSC for every df/dy entry a compile + /// pass proved constant. Split dense rows are absent: their own tape + /// recomputes them. + #[inline] + pub fn constant_jacobian_entries(&self) -> &[(usize, f64)] { + self.jac_layout.constant_slots() + } + + /// The mass matrix `M` in CSR, as Python supplied it. + #[inline] + pub const fn mass_matrix(&self) -> &CsrData { + &self.mass_matrix + } + + /// Get the algebraic-state mask: `true` = algebraic, `false` = differential. + /// + /// A state `i` is classified algebraic when row `i` of the mass matrix has + /// no diagonal entry, which matches `PyBaMM`'s convention. + #[inline] + pub fn algebraic_ids(&self) -> &[bool] { + &self.algebraic_ids + } + + /// The same mask in IDA's numeric convention: 1.0 differential, 0.0 + /// algebraic, the inverse polarity of [`algebraic_ids`](Self::algebraic_ids). + pub fn algebraic_ids_f64(&self, output: &mut [f64]) { + for (i, &is_alg) in self.algebraic_ids.iter().enumerate() { + output[i] = if is_alg { 0.0 } else { 1.0 }; + } + } + + /// Number of algebraic states, or 0 without an algebraic sub-block. + #[inline] + pub fn n_algebraic(&self) -> usize { + self.algebraic + .as_ref() + .map_or(0, |block| block.jac.n_cols()) + } + + /// Whether an algebraic sub-block was compiled for the Newton solver. + #[inline] + pub const fn has_algebraic(&self) -> bool { + self.algebraic.is_some() + } + + /// Non-zeros in the assembled algebraic Jacobian `dg/dy_alg`. + pub fn algebraic_jacobian_nnz(&self) -> usize { + self.algebraic.as_ref().map_or(0, |block| block.jac.nnz()) + } + + /// COO row indices of the algebraic Jacobian, in the order + /// [`assemble_algebraic_jacobian_into`](Self::assemble_algebraic_jacobian_into) writes. + #[inline] + pub fn algebraic_jacobian_row_indices(&self) -> &[usize] { + self.algebraic + .as_ref() + .map_or(&[], |block| block.jac.coo_global_indices().0) + } + + /// COO column indices matching + /// [`algebraic_jacobian_row_indices`](Self::algebraic_jacobian_row_indices), + /// as global state indices. + #[inline] + pub fn algebraic_jacobian_col_indices(&self) -> &[usize] { + self.algebraic + .as_ref() + .map_or(&[], |block| block.jac.coo_global_indices().1) + } + + /// Evaluate the algebraic residual g(t, y) into `output`. + /// + /// No-op when no algebraic sub-block was compiled. + pub fn eval_algebraic_residual( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + if let (Some(block), Some(scratch)) = + (self.algebraic.as_ref(), ws.algebraic_scratch.as_mut()) + { + let result = block.residual.eval(scratch, t, y, &[], inputs); + output[..result.len()].copy_from_slice(result); + } + } + + /// Compute (`dg/dy_alg`) @ v for the algebraic Jacobian-vector product. + /// + /// No-op when no algebraic sub-block was compiled. + pub fn eval_algebraic_jacobian_action( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + output: &mut [f64], + ) { + if let (Some(block), Some(scratch)) = + (self.algebraic.as_ref(), ws.algebraic_jac_scratch.as_mut()) + { + let n_algebraic = block.jac.n_cols(); + assert!( + output.len() >= n_algebraic, + "algebraic output buffer too small: need {n_algebraic}, got {}", + output.len() + ); + block.jac.action_into(scratch, t, y, &[], inputs, v, output); + } + } + + /// Assemble the algebraic Jacobian `dg/dy_alg` into `jac_data`. + /// + /// The output order matches [`Self::algebraic_jacobian_row_indices`] and + /// [`Self::algebraic_jacobian_col_indices`]. + pub fn assemble_algebraic_jacobian_into( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + let Some(block) = self.algebraic.as_ref() else { + return; + }; + let scratch = ws + .algebraic_jac_scratch + .as_mut() + .expect("an algebraic block implies its workspace scratch"); + block + .jac + .assemble_into(scratch, block.jac.layout(), t, y, &[], inputs, jac_data); + } + + /// Number of parameters sensitivities were requested for, which may be a + /// subset of [`n_params`](Self::n_params). Zero without sensitivities. + #[inline] + pub const fn n_sens_params(&self) -> usize { + self.sens_param_indices.len() + } + + /// Evaluate ∂f/∂p for the sensitivity indexed by `sens_idx` into `output`. + /// + /// `sens_idx` indexes into `sens_param_indices` (range `0..n_sens_params()`). + /// Panics if `sens_idx >= n_sens_params()` or if no sensitivities were + /// configured at construction time. + pub fn eval_sens( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + sens_idx: usize, + output: &mut [f64], + ) { + let expr = self + .sens_expr + .as_ref() + .expect("eval_sens called on a model without sensitivities"); + let scratch = ws + .sens_scratch + .as_mut() + .expect("sens_scratch missing from workspace"); + let global_idx = self.sens_param_indices[sens_idx]; + + // Build the unit dp vector in the pre-allocated buffer + debug_assert!(global_idx < ws.sens_dp_buffer.len()); + ws.sens_dp_buffer.fill(0.0); + ws.sens_dp_buffer[global_idx] = 1.0; + let tangent = TangentInputs { + dy: None, + dp: Some(&ws.sens_dp_buffer), + }; + + let result = expr.eval_with_tangent(scratch, t, y, &[], inputs, &tangent); + output[..result.len()].copy_from_slice(result); + } + + /// Run the primal section of `sens_expr` once for (t, y, inputs). The primal + /// slots in `sens_scratch` then serve every subsequent tangent column. + /// Only callable on models compiled with sensitivities; panics otherwise. + pub fn sens_primal_pass(&self, ws: &mut Workspace, t: f64, y: &[f64], inputs: &[f64]) { + #[cfg(test)] + { + ws.sens_primal_passes += 1; + } + let expr = self.sens_expr.as_ref().expect("no sensitivity expression"); + let scratch = ws + .sens_scratch + .as_mut() + .expect("sens_scratch missing from workspace"); + expr.run_primal_section(scratch, t, y, &[], inputs); + } + + /// Tangent-only sweep for the column of parameter `param_idx`. Requires a prior + /// `sens_primal_pass` at the same (t, y, inputs) on this workspace. + /// Only callable on models compiled with sensitivities; panics otherwise. + pub fn sens_tangent_column(&self, ws: &mut Workspace, param_idx: usize, output: &mut [f64]) { + let expr = self.sens_expr.as_ref().expect("no sensitivity expression"); + debug_assert!(param_idx < ws.sens_dp_buffer.len()); + ws.sens_dp_buffer.fill(0.0); + ws.sens_dp_buffer[param_idx] = 1.0; + let scratch = ws + .sens_scratch + .as_mut() + .expect("sens_scratch missing from workspace"); + let tangent = TangentInputs { + dy: None, + dp: Some(&ws.sens_dp_buffer), + }; + let result = expr.run_tangent_section(scratch, &tangent); + output[..result.len()].copy_from_slice(result); + } + + /// Evaluate ∂f/∂p for all configured sensitivity parameters into `output`. + /// + /// Layout: `output[i*n_states + j] = ∂f_j/∂p_i`. The buffer must have + /// length at least `n_sens_params() * n_states`. Runs the shared primal + /// section once, then one tangent-only sweep per parameter. + pub fn eval_sens_all( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + // No configured sensitivities: keep the historical silent no-op. + if self.n_sens_params() == 0 { + return; + } + let n_states = self.n_states; + self.sens_primal_pass(ws, t, y, inputs); + for (i, ¶m_idx) in self.sens_param_indices.iter().enumerate() { + let start = i * n_states; + self.sens_tangent_column(ws, param_idx, &mut output[start..start + n_states]); + } + } + + /// Compile and append an output-variable expression to this model. + /// + /// `node` must already exist in `arena`. Its length is captured at compile + /// time, along with the dH/dp and dH/dy tangent graphs. + pub fn add_output(&mut self, arena: &Arena, node: NodeId) { + self.outputs.push(arena, node); + } + + /// Compile and append an event expression to this model. + /// + /// Events are used for root-finding during integration. When any event + /// expression crosses zero, the solver can terminate or take action. + /// + /// `node` must already exist in `arena`. + pub fn add_event(&mut self, arena: &Arena, node: NodeId) { + self.events.push(arena, node); + } + + /// Fuse the events into one tape (see [`ObservableSet::fuse`]). + /// + /// Call once, after every event is added. `event_roots` must be the same + /// nodes, in the same order, passed to [`Self::add_event`], from `arena`. + pub fn fuse_events(&mut self, arena: &mut Arena, event_roots: &[NodeId]) { + self.events.fuse(arena, event_roots); + } + + /// Number of compiled output-variable expressions. + #[inline] + pub const fn n_outputs(&self) -> usize { + self.outputs.count() + } + + /// Length of output variable `var_idx`. + /// + /// Panics if `var_idx >= n_outputs()`. + #[inline] + pub fn output_len_at(&self, var_idx: usize) -> usize { + self.outputs.len_at(var_idx) + } + + /// Length of all output variables concatenated, the buffer size + /// [`eval_observables`](Self::eval_observables) needs for + /// [`ObservableKind::Outputs`]. + #[inline] + pub const fn total_output_len(&self) -> usize { + self.outputs.total_len() + } + + /// Number of compiled event expressions. + #[inline] + pub const fn n_events(&self) -> usize { + self.events.count() + } + + /// Length of all events concatenated, the buffer size + /// [`eval_observables`](Self::eval_observables) needs for + /// [`ObservableKind::Events`]. + #[inline] + pub const fn total_event_len(&self) -> usize { + self.events.total_len() + } + + /// Evaluate output variable `var_idx` into `output`, returning the count written. + /// + /// Panics if `var_idx >= n_outputs()` or `output.len() < output_len_at(var_idx)`. + pub fn eval_output( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + var_idx: usize, + output: &mut [f64], + ) -> usize { + self.outputs.eval_at( + ws.observable_scratch(ObservableKind::Outputs), + t, + y, + inputs, + var_idx, + output, + ) + } + + /// Evaluate event `event_idx` into `output`, returning the count written. + /// + /// Panics if `event_idx >= n_events()` or `output.len() < the event's length`. + pub fn eval_event( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + event_idx: usize, + output: &mut [f64], + ) -> usize { + self.events.eval_at( + ws.observable_scratch(ObservableKind::Events), + t, + y, + inputs, + event_idx, + output, + ) + } + + /// Evaluate every observable of `kind` into `output`, concatenated. + /// + /// Panics if `output.len()` is under the family's total length. + pub fn eval_observables( + &self, + ws: &mut Workspace, + kind: ObservableKind, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + self.observables(kind) + .eval_all(ws.observable_scratch(kind), t, y, inputs, output); + } + + /// Assemble `J = df/dy - cj * M` in COO (row, col, value) format, `cj` + /// taken from the workspace. + /// + /// Allocates on every call and has no binding, so it is a convenience for + /// tests and callers holding a `CompiledModel` directly; solver hot paths + /// want [`Self::assemble_jacobian_csc_into`]. Expanded from the CSC + /// assembly rather than re-deriving the colour loop, so the two cannot + /// drift over which route fills a given slot. + pub fn assemble_jacobian( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + ) -> (Vec, Vec, Vec) { + let nnz = self.sparsity.nnz(); + let mut values = vec![0.0; nnz]; + self.assemble_jacobian_csc_into_coloring(ws, t, y, inputs, &mut values); + + let (row_indices, col_indices) = self.csc_sparsity.csc_to_csr_map.iter().copied().unzip(); + (row_indices, col_indices, values) + } + + /// CSC pattern of the assembled Jacobian `df/dy - cj*M`, the layout KLU and + /// diffsol expect; not `df/dy` alone, which is + /// [`jac_y_sparsity`](Self::jac_y_sparsity). + #[inline] + pub const fn csc_sparsity(&self) -> &CscPattern { + &self.csc_sparsity + } + + /// Non-zeros in the assembled Jacobian, the buffer size every + /// `assemble_jacobian_csc_*` call needs. + #[inline] + pub const fn nnz(&self) -> usize { + self.sparsity.nnz() + } + + /// Jacobian assembly strategy this model compiled to. + #[inline] + pub const fn jacobian_strategy(&self) -> JacobianStrategy { + JacobianStrategy::Coloring + } + + /// Compile-time Jacobian metrics: colors, non-zeros and dense-row counts, + /// for benchmark attribution. + pub fn jacobian_stats(&self) -> JacobianStats { + let ir = self.jac.assembly_tape().ir(); + let split_info = ir.split_eval_info(); + JacobianStats { + strategy: JacobianStrategy::Coloring, + n_colors: self.jac.n_colors(), + nnz: self.nnz(), + n_dense_rows: self.jac.n_dense_rows(), + n_dense_row_candidates: self.jac.n_candidate_rows(), + n_constant_entries: self.jac.constant_csr_entries().len(), + n_swept_columns: self.jac.coloring().n_seeded_columns(), + dense_row_entries: self.jac.dense_row_entries(), + dense_row_tape_instructions: self.jac.dense_row_tape_instructions(), + split_eval_primal_instructions: split_info.map(|s| s.primal_end), + split_eval_total_instructions: Some(ir.common_instruction_count()), + split_eval_raw_instructions: Some(ir.instructions().len()), + split_eval_dispatch_count: ir.dispatch_count(), + branch_block_lens: ir.branch_block_lens(), + jac_lane_width: self.jac.lane_width(), + } + } + + /// Assemble the Jacobian into a pre-allocated CSC data buffer. + /// + /// Zero-allocation version for FFI/IDAKLU integration. The Jacobian + /// is `J = df/dy - cj * M` where `cj` lives in `ws`. + /// Panics if `jac_data.len() < nnz()`. + pub fn assemble_jacobian_csc_into( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.assemble_jacobian_csc_into_coloring(ws, t, y, inputs, jac_data); + } + + /// Assemble the Jacobian into a pre-allocated CSC data buffer using the + /// coloring-based JVP approach. + /// + /// Canonical Jacobian assembly method: one primal evaluation, then + /// `n_colors` tangent sweeps with precomputed scatter. + /// Panics if `jac_data.len() < nnz()`. + pub fn assemble_jacobian_csc_into_coloring( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + let nnz = self.sparsity.nnz(); + assert!( + jac_data.len() >= nnz, + "jac_data buffer too small: need {nnz}, got {}", + jac_data.len() + ); + + self.assemble_dfdy_into(ws, t, y, inputs, jac_data); + self.apply_mass_postpass(ws, jac_data); + } + + /// `df/dy` into the merged CSC, by handing the assembly module this model's + /// slot layout. Shared by both assembly entry points, which differ only in + /// the mass term, so a future stage cannot reach one and miss the other. + fn assemble_dfdy_into( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.jac.assemble_into( + &mut ws.jac_scratch, + &self.jac_layout, + t, + y, + &[], + inputs, + jac_data, + ); + } + + fn jac_mul_primal( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + output: &mut [f64], + ) { + self.jac + .action_into(&mut ws.jac_scratch, t, y, &[], inputs, v, output); + } + + /// Fold `-cj*M` into the merged CSC, after every derivative fill: the + /// diagonal slots it touches are ones a colour sweep or a dense row has + /// already written. + fn apply_mass_postpass(&self, ws: &Workspace, jac_data: &mut [f64]) { + if ws.cj == 0.0 { + return; + } + for (mass_idx, &csc_idx) in self.mass_to_csc_map.iter().enumerate() { + jac_data[csc_idx] -= ws.cj * self.mass_matrix.data[mass_idx]; + } + } + + /// Assemble df/dy in CSC format using graph coloring, without applying + /// the mass matrix post-pass (diffsol handles mass separately). + /// + /// Panics if `jac_data.len() < nnz`. + pub fn assemble_jacobian_csc_no_mass( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + let nnz = self.csc_sparsity.rowind.len(); + assert!( + jac_data.len() >= nnz, + "jac_data buffer too small: need {nnz}, got {}", + jac_data.len() + ); + self.assemble_dfdy_into(ws, t, y, inputs, jac_data); + // No apply_mass_postpass, diffsol handles mass matrix separately + } + + /// Batch-evaluate every output tape over `k` trajectory points; see + /// [`ObservableSet::eval_batch`] for the buffer layouts. + pub fn eval_outputs_batch( + &self, + ws: &mut Workspace, + k: usize, + ts: &[f64], + y_cols: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + self.outputs.eval_batch( + ws.observable_scratch(ObservableKind::Outputs), + k, + ts, + y_cols, + self.n_states, + inputs, + out, + ); + } + + /// Compute the sensitivity action df/dp @ v. + /// + /// `v` is parameter-space; `sens_params[i]` is the global index of `v[i]`. + /// Panics if no sensitivity expression was compiled. + // Evaluation point, tangent, mapping and output are all distinct arguments. + #[allow(clippy::too_many_arguments)] + pub fn sens_action( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + sens_params: &[usize], + v: &[f64], + output: &mut [f64], + ) { + let sens_expr = self.sens_expr.as_ref().expect("no sensitivity expression"); + seed_param_tangent(&mut ws.sens_dp_buffer, sens_params, v); + let scratch = ws.sens_scratch.as_mut().expect("no sensitivity scratch"); + let tangent = TangentInputs { + dy: None, + dp: Some(&ws.sens_dp_buffer), + }; + let result = sens_expr.eval_with_tangent(scratch, t, y, &[], inputs, &tangent); + output[..result.len()].copy_from_slice(result); + } + + /// Whether `df/dp` was compiled, which requires sensitivities at build time. + pub const fn has_sens(&self) -> bool { + self.sens_expr.is_some() + } + + /// Global parameter index for each configured sensitivity parameter. + #[inline] + pub fn sens_param_indices(&self) -> &[usize] { + &self.sens_param_indices + } + + /// Compute dH/dp · v over every observable of `kind`, into `output`. + /// + /// `v` is parameter-space; `sens_params[i]` is the global index of `v[i]`. + // Evaluation point, family, tangent, mapping and output are all distinct. + #[allow(clippy::too_many_arguments)] + pub fn observable_sens_action( + &self, + ws: &mut Workspace, + kind: ObservableKind, + t: f64, + y: &[f64], + inputs: &[f64], + sens_params: &[usize], + v: &[f64], + output: &mut [f64], + ) { + let (scratch, dp) = ws.observable_tangent(kind); + seed_param_tangent(dp, sens_params, v); + self.observables(kind) + .sens_action(scratch, t, y, inputs, dp, output); + } + + /// Compute dH/dy · v over every observable of `kind`, into `output`. + /// + /// `v` is state-space, of length `n_states`. + // Evaluation point, family, tangent and output are all distinct arguments. + #[allow(clippy::too_many_arguments)] + pub fn observable_jac_action( + &self, + ws: &mut Workspace, + kind: ObservableKind, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + output: &mut [f64], + ) { + self.observables(kind) + .jac_action(ws.observable_scratch(kind), t, y, inputs, v, output); + } + + /// Project state sensitivities onto output-variable sensitivities; see + /// [`ObservableSet::sens_project`] for the buffer layouts. + pub fn output_sens_project( + &self, + ws: &mut Workspace, + t: f64, + y: &[f64], + inputs: &[f64], + y_sens: &[f64], + out: &mut [f64], + ) { + let (scratch, dp) = ws.observable_tangent(ObservableKind::Outputs); + self.outputs.sens_project( + scratch, + dp, + t, + y, + inputs, + &self.sens_param_indices, + y_sens, + self.n_states, + out, + ); + } +} + +impl std::fmt::Debug for CompiledModel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledModel") + .field("n_states", &self.n_states) + .field("n_params", &self.n_params) + .field("output_len", &self.primal_expr.output_len()) + .field("sparsity_nnz", &self.sparsity.nnz()) + .field("n_colors", &self.jac.n_colors()) + .field("jacobian_strategy", &JacobianStrategy::Coloring.as_str()) + .finish_non_exhaustive() + } +} + +/// One [`CompiledModel`] paired with an owned mutable [`Workspace`], so a +/// caller can evaluate through `&mut self` without managing scratch by hand. +/// +/// Owning the workspace is what makes this per-solve rather than shareable: the +/// FFI/pybammsolvers ABI path, the Python `CompiledModel` binding, and one-shot +/// eval/tests each hold one. Solves that share a model concurrently pass +/// `Arc` and their own [`Workspace`] instead. +/// +/// Cloning shares the immutable [`CompiledModel`] via `Arc` and gives the clone +/// its own fresh [`Workspace`], so it does not deep-copy the mass matrix, +/// sparsity, scatter tables, or workspace buffers. +/// +/// Its interface is only what the workspace buys; everything else is reached +/// through [`Deref`](std::ops::Deref). +pub struct ModelEvaluator { + compiled: Arc, + workspace: Workspace, +} + +impl Clone for ModelEvaluator { + fn clone(&self) -> Self { + let mut workspace = self.compiled.create_workspace(); + workspace.set_cj(self.workspace.cj()); + Self { + compiled: Arc::clone(&self.compiled), + workspace, + } + } +} + +impl std::fmt::Debug for ModelEvaluator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ModelEvaluator") + .field("n_states", &self.compiled.n_states) + .field("n_params", &self.compiled.n_params) + .field("output_len", &self.compiled.primal_expr.output_len()) + .field("cj", &self.workspace.cj) + .field("sparsity_nnz", &self.compiled.sparsity.nnz()) + .field("n_colors", &self.compiled.jac.n_colors()) + .field("jacobian_strategy", &JacobianStrategy::Coloring.as_str()) + .finish_non_exhaustive() + } +} + +/// Lends the wrapped [`CompiledModel`]'s read-only interface rather than +/// restating it. +/// +/// No `DerefMut`: the artifact is shared behind an `Arc`, and only +/// `add_output`/`add_event`/`fuse_events` may mutate it. +impl std::ops::Deref for ModelEvaluator { + type Target = CompiledModel; + + #[inline] + fn deref(&self) -> &CompiledModel { + &self.compiled + } +} + +/// Drops the workspace, so an `impl Into>` parameter takes +/// an evaluator as readily as an `Arc`. +impl From for Arc { + fn from(evaluator: ModelEvaluator) -> Self { + evaluator.into_compiled() + } +} + +impl ModelEvaluator { + /// Create a new `ModelEvaluator` wrapper from a primal expression and mass matrix. + pub fn new( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + ) -> Self { + let compiled = CompiledModel::new(arena, rhs, mass_matrix, n_states, n_params); + let workspace = compiled.create_workspace(); + Self { + compiled: Arc::new(compiled), + workspace, + } + } + + /// Create a `ModelEvaluator` with forward-sensitivity expressions. + pub fn new_with_sens( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + sens_param_indices: &[usize], + ) -> Self { + let compiled = CompiledModel::new_with_sens( + arena, + rhs, + mass_matrix, + n_states, + n_params, + sens_param_indices, + ); + let workspace = compiled.create_workspace(); + Self { + compiled: Arc::new(compiled), + workspace, + } + } + + /// Create a `ModelEvaluator` with optional algebraic sub-block expressions. + pub fn new_with_algebraic( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + algebraic_rhs: Option, + algebraic_variable_indices: &[usize], + ) -> Self { + let compiled = CompiledModel::new_with_algebraic( + arena, + rhs, + mass_matrix, + n_states, + n_params, + algebraic_rhs, + algebraic_variable_indices, + ); + let workspace = compiled.create_workspace(); + Self { + compiled: Arc::new(compiled), + workspace, + } + } + + /// Create a `ModelEvaluator` with any combination of sensitivities and an algebraic sub-block. + pub fn new_with_options( + arena: &Arena, + rhs: NodeId, + mass_matrix: CsrData, + n_states: usize, + n_params: usize, + options: CompiledModelOptions<'_>, + ) -> Self { + let compiled = + CompiledModel::new_with_options(arena, rhs, mass_matrix, n_states, n_params, options); + let workspace = compiled.create_workspace(); + Self { + compiled: Arc::new(compiled), + workspace, + } + } + + /// Consume the wrapper, returning the shared immutable compiled model + /// (drops the workspace). The `Arc` is moved out without deep-copying. + pub fn into_compiled(self) -> Arc { + self.compiled + } + + /// Borrow the immutable compiled model. + pub fn compiled(&self) -> &CompiledModel { + &self.compiled + } + + /// Bind a fresh [`Workspace`] to an already-compiled model, the inverse of + /// [`into_compiled`](Self::into_compiled). Nothing is recompiled, so N + /// evaluators cost N workspaces and no lowering. + pub fn from_compiled(compiled: Arc) -> Self { + let workspace = compiled.create_workspace(); + Self { + compiled, + workspace, + } + } + + /// Set the cj coefficient for Jacobian computation. + #[inline] + pub const fn set_cj(&mut self, cj: f64) { + self.workspace.set_cj(cj); + } + + /// Get the current cj coefficient. + #[inline] + pub const fn cj(&self) -> f64 { + self.workspace.cj() + } + + /// Evaluate `f(t, y; p)` into `out`, using this model's own workspace. + #[inline] + pub fn eval_rhs(&mut self, t: f64, y: &[f64], inputs: &[f64], out: &mut [f64]) { + self.compiled + .eval_rhs(&mut self.workspace, t, y, inputs, out); + } + + /// `(df/dy - cj*M) @ v`, the Newton-iteration matrix action, with `cj` taken + /// from [`set_cj`](Self::set_cj). + pub fn jac_mul(&mut self, t: f64, y: &[f64], inputs: &[f64], v: &[f64], out: &mut [f64]) { + self.compiled + .jac_mul(&mut self.workspace, t, y, inputs, v, out); + } + + /// `df/dy @ v` with no mass term, for callers that subtract `cj*M @ v` + /// themselves. + pub fn jac_action(&mut self, t: f64, y: &[f64], inputs: &[f64], v: &[f64], jv: &mut [f64]) { + self.compiled + .jac_action(&mut self.workspace, t, y, inputs, v, jv); + } + + /// DAE residual `r = M*y' - f(t, y; p)`, the form IDAKLU solves. + pub fn eval_residual(&mut self, t: f64, y: &[f64], yp: &[f64], inputs: &[f64], r: &mut [f64]) { + self.compiled + .eval_residual(&mut self.workspace, t, y, yp, inputs, r); + } + + /// Algebraic residual `g(t, y; p)`, or nothing when the model has no + /// algebraic sub-block. + pub fn eval_algebraic_residual( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + self.compiled + .eval_algebraic_residual(&mut self.workspace, t, y, inputs, output); + } + + /// `dg/dy_alg @ v` over the algebraic states only, again a no-op without an + /// algebraic sub-block. + pub fn eval_algebraic_jacobian_action( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + output: &mut [f64], + ) { + self.compiled + .eval_algebraic_jacobian_action(&mut self.workspace, t, y, inputs, v, output); + } + + /// Assemble `dg/dy_alg` into `jac_data`, ordered to match + /// [`algebraic_jacobian_row_indices`](CompiledModel::algebraic_jacobian_row_indices). + pub fn assemble_algebraic_jacobian_into( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.compiled + .assemble_algebraic_jacobian_into(&mut self.workspace, t, y, inputs, jac_data); + } + + /// `df/dp` for one sensitivity parameter, where `sens_idx` is a position in + /// [`sens_param_indices`](CompiledModel::sens_param_indices), not a global parameter + /// index. + pub fn eval_sens( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + sens_idx: usize, + output: &mut [f64], + ) { + self.compiled + .eval_sens(&mut self.workspace, t, y, inputs, sens_idx, output); + } + + /// `df/dp` for every configured sensitivity parameter, laid out as + /// `output[i*n_states + j] = df_j/dp_i`, sharing one primal pass across all + /// of them. + pub fn eval_sens_all(&mut self, t: f64, y: &[f64], inputs: &[f64], output: &mut [f64]) { + self.compiled + .eval_sens_all(&mut self.workspace, t, y, inputs, output); + } + + /// Project state sensitivities onto output-variable sensitivities. See + /// `CompiledModel::output_sens_project`. + pub fn output_sens_project( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + y_sens: &[f64], + out: &mut [f64], + ) { + self.compiled + .output_sens_project(&mut self.workspace, t, y, inputs, y_sens, out); + } + + /// Evaluate output variable `var_idx`, returning how many elements it wrote. + pub fn eval_output( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + var_idx: usize, + output: &mut [f64], + ) -> usize { + self.compiled + .eval_output(&mut self.workspace, t, y, inputs, var_idx, output) + } + + /// Evaluate event `event_idx`, returning how many elements it wrote. + pub fn eval_event( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + event_idx: usize, + output: &mut [f64], + ) -> usize { + self.compiled + .eval_event(&mut self.workspace, t, y, inputs, event_idx, output) + } + + /// Evaluate every observable of `kind` into one concatenated buffer; see + /// [`CompiledModel::eval_observables`]. + pub fn eval_observables( + &mut self, + kind: ObservableKind, + t: f64, + y: &[f64], + inputs: &[f64], + output: &mut [f64], + ) { + self.compiled + .eval_observables(&mut self.workspace, kind, t, y, inputs, output); + } + + /// Batch-evaluate every output variable over `k` trajectory points; see + /// [`CompiledModel::eval_outputs_batch`] for the buffer layouts. + pub fn eval_outputs_batch( + &mut self, + k: usize, + ts: &[f64], + y_cols: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + self.compiled + .eval_outputs_batch(&mut self.workspace, k, ts, y_cols, inputs, out); + } + + /// Assemble `df/dy - cj*M` as fresh COO triples. + /// + /// Allocates on every call, so it suits the Python bindings rather than a + /// solver loop; the `_csc_into` forms below write into a caller's buffer. + pub fn assemble_jacobian( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + ) -> (Vec, Vec, Vec) { + self.compiled + .assemble_jacobian(&mut self.workspace, t, y, inputs) + } + + /// Assemble `df/dy - cj*M` into a caller-owned CSC value buffer of + /// [`nnz`](CompiledModel::nnz) elements. + pub fn assemble_jacobian_csc_into( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.compiled + .assemble_jacobian_csc_into(&mut self.workspace, t, y, inputs, jac_data); + } + + /// The same assembly driven by the column coloring: one primal pass then one + /// tangent sweep per color, which is the fast path solvers use. + pub fn assemble_jacobian_csc_into_coloring( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.compiled.assemble_jacobian_csc_into_coloring( + &mut self.workspace, + t, + y, + inputs, + jac_data, + ); + } + + /// Assemble `df/dy` without folding in the mass matrix, for diffsol, which + /// applies `M` itself. + pub fn assemble_jacobian_csc_no_mass( + &mut self, + t: f64, + y: &[f64], + inputs: &[f64], + jac_data: &mut [f64], + ) { + self.compiled + .assemble_jacobian_csc_no_mass(&mut self.workspace, t, y, inputs, jac_data); + } + + /// Compile and append an output-variable expression to this model. + /// Rebuilds the workspace after mutation to include a scratch buffer for the new expr. + pub fn add_output(&mut self, arena: &Arena, node: NodeId) { + // Construction-time mutation before the model is shared, so `make_mut` + // takes the unique `Arc` in place (copy-on-write only if ever shared). + Arc::make_mut(&mut self.compiled).add_output(arena, node); + self.workspace = self.compiled.create_workspace(); + } + + /// Compile and append an event expression to this model. + /// Rebuilds the workspace after mutation to include a scratch buffer for the new expr. + pub fn add_event(&mut self, arena: &Arena, node: NodeId) { + Arc::make_mut(&mut self.compiled).add_event(arena, node); + self.workspace = self.compiled.create_workspace(); + } + + /// Fuse the events (see [`CompiledModel::fuse_events`]) and resize the + /// workspace scratch to match. Call once, after all events are added. + pub fn fuse_events(&mut self, arena: &mut Arena, event_roots: &[NodeId]) { + Arc::make_mut(&mut self.compiled).fuse_events(arena, event_roots); + self.workspace = self.compiled.create_workspace(); + } +} + +/// CSR matrix-vector product: y = A @ x +/// +/// Computes the sparse matrix-vector product for the mass matrix. +#[inline] +fn csr_matvec(csr: &CsrData, x: &[f64], y: &mut [f64]) { + let rows = csr.shape.rows; + for (row, y_elem) in y.iter_mut().enumerate().take(rows) { + let start = csr.indptr[row]; + let end = csr.indptr[row + 1]; + let mut sum = 0.0; + for idx in start..end { + sum += csr.data[idx] * x[csr.indices[idx]]; + } + *y_elem = sum; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::jacobian::tests::banded_nonlinear; + use crate::node::{Node, Shape}; + + #[test] + fn data_workspace_roundtrip_eval_rhs() { + // 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 r0 = arena.alloc(Node::Mul(neg_one, sv0)); + let r1 = arena.alloc(Node::Mul(neg_two, sv1)); + let rhs = arena.alloc(Node::Concat(vec![r0, r1])); + let mass = CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![1.0, 1.0], + shape: Shape { rows: 2, cols: 2 }, + }; + let compiled = ModelEvaluator::new(&arena, rhs, mass, 2, 0).into_compiled(); + + let mut ws = compiled.create_workspace(); + let mut out = vec![0.0; 2]; + compiled.eval_rhs(&mut ws, 0.0, &[3.0, 5.0], &[], &mut out); + assert!((out[0] - (-3.0)).abs() < 1e-12); + assert!((out[1] - (-10.0)).abs() < 1e-12); + + // Workspace scratch lengths match the compiled model's expressions. + assert_eq!(ws.primal_scratch.len(), compiled.primal_scratch_len()); + } + + #[test] + fn clone_shares_data_and_is_workspace_independent() { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs = arena.alloc(Node::Mul(neg, sv)); + let model = ModelEvaluator::new(&arena, rhs, identity_mass_matrix(1), 1, 0); + + let mut clone = model.clone(); + // Immutable compiled data is shared via Arc, not deep-copied. + assert!(Arc::ptr_eq(&model.compiled, &clone.compiled)); + // Workspaces are independent: mutating the clone leaves the original. + clone.workspace.set_cj(7.0); + assert!((clone.workspace.cj() - 7.0).abs() < 1e-12); + assert!(model.workspace.cj().abs() < 1e-12); + } + + #[test] + fn deref_lends_the_artifact_rather_than_copying_it() { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs = arena.alloc(Node::Mul(neg, sv)); + let model = ModelEvaluator::new(&arena, rhs, identity_mass_matrix(1), 1, 0); + + // The lent reference IS the shared artifact, so an accessor reached + // through `Deref` cannot drift from one reached through `compiled()`. + assert!(std::ptr::eq(&raw const *model, model.compiled())); + assert!(std::ptr::eq( + &raw const *model, + Arc::as_ptr(&model.compiled) + )); + assert_eq!(model.n_states(), model.compiled().n_states()); + } + + #[test] + fn inherent_workspace_methods_win_over_the_lent_ones() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + let mut model = ModelEvaluator::new(&arena, rhs, identity_mass_matrix(2), 2, 0); + + // Same name on both types: the evaluator's own `&mut self` form is the + // one that resolves, and it agrees with driving the artifact by hand. + let mut through_wrapper = [0.0; 2]; + model.eval_rhs(0.0, &[1.0, 2.0], &[], &mut through_wrapper); + + let compiled = Arc::clone(&model.compiled); + let mut ws = compiled.create_workspace(); + let mut through_artifact = [0.0; 2]; + compiled.eval_rhs(&mut ws, 0.0, &[1.0, 2.0], &[], &mut through_artifact); + + // Bitwise: the two paths must be the same evaluation, not merely close. + assert_eq!( + through_wrapper.map(f64::to_bits), + through_artifact.map(f64::to_bits) + ); + } + + #[test] + fn mutators_repoint_the_lend() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + let mut model = ModelEvaluator::new(&arena, rhs, identity_mass_matrix(2), 2, 0); + assert_eq!(model.n_outputs(), 0); + + model.add_output(&arena, y); + + // Read through the lend, so `add_output`'s `Arc::make_mut` swap has to + // be visible there and not just through a stale forwarded copy. + assert_eq!(model.n_outputs(), 1); + assert_eq!(model.total_output_len(), 2); + } + + /// Create an identity mass matrix of size n. + fn identity_mass_matrix(n: usize) -> CsrData { + CsrData { + shape: Shape::matrix(n, n), + indptr: (0..=n).collect(), + indices: (0..n).collect(), + data: vec![1.0; n], + } + } + + /// Assemble `jac` directly, through its own canonical layout: what a caller + /// holding the shared artifact rather than the model would do. + fn assemble_through( + jac: &JacobianData, + mut scratch: JacobianScratch, + t: f64, + y: &[f64], + ) -> Vec { + let mut out = vec![f64::NAN; jac.layout().n_slots()]; + jac.assemble_into(&mut scratch, jac.layout(), t, y, &[], &[], &mut out); + out + } + + /// A mass matrix with no diagonal at all: every state is algebraic, so the + /// algebraic sub-block spans the whole system. + fn all_algebraic_mass_matrix(n: usize) -> CsrData { + CsrData { + shape: Shape::matrix(n, n), + indptr: vec![0; n + 1], + indices: Vec::new(), + data: Vec::new(), + } + } + + /// A fully algebraic banded model, whose `dg/dy_alg` is wide enough to batch. + fn build_batched_algebraic_model(n: usize, half_width: usize) -> ModelEvaluator { + let (arena, root) = banded_nonlinear(n, half_width); + let indices: Vec = (0..n).collect(); + ModelEvaluator::new_with_algebraic( + &arena, + root, + all_algebraic_mass_matrix(n), + n, + 0, + Some(root), + &indices, + ) + } + + /// The algebraic block is a third adapter onto the assembly module, so it gets + /// the lane batching the `df/dy` path has without restating it -- and the + /// values must not depend on which width ran. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point + fn the_algebraic_block_inherits_lane_batching() { + let (n, half_width) = (64usize, 3usize); + let mut model = build_batched_algebraic_model(n, half_width); + let jac = model + .algebraic_jacobian_data() + .expect("the fixture compiles an algebraic block"); + assert!( + JacobianScratch::new(&jac).lane_width() > 1, + "the algebraic fixture must batch" + ); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.037, 0.41)).collect(); + let mut batched = vec![f64::NAN; model.algebraic_jacobian_nnz()]; + model.assemble_algebraic_jacobian_into(0.5, &y, &[], &mut batched); + + // The same artifact swept one colour per walk: the reference the batched + // sweep has to reproduce. + let scalar = assemble_through(&jac, JacobianScratch::scalar(&jac), 0.5, &y); + assert_eq!(batched, scalar); + } + + /// The standalone `dg/dy_alg` handle is the artifact the model compiled, not a + /// second tangent transform of the same expression, so the two cannot drift + /// in pattern, colouring or value. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point + fn the_standalone_algebraic_jacobian_is_the_compiled_artifact() { + let n = 24usize; + let mut model = build_batched_algebraic_model(n, 2); + let first = model.algebraic_jacobian_data().expect("algebraic block"); + let second = model.algebraic_jacobian_data().expect("algebraic block"); + assert!( + Arc::ptr_eq(&first, &second), + "the handle must lend the compiled artifact, never rebuild it" + ); + + let (rows, cols) = first.coo_global_indices(); + assert_eq!(rows, model.algebraic_jacobian_row_indices()); + assert_eq!(cols, model.algebraic_jacobian_col_indices()); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.13, 0.7)).collect(); + let mut through_model = vec![f64::NAN; model.algebraic_jacobian_nnz()]; + model.assemble_algebraic_jacobian_into(0.25, &y, &[], &mut through_model); + + let through_handle = assemble_through(&first, JacobianScratch::new(&first), 0.25, &y); + assert_eq!(through_model, through_handle); + } + + #[test] + fn test_compiled_model_new() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); // f(y) = 2*y + + let mass = identity_mass_matrix(2); + + let model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + assert_eq!(model.n_states(), 2); + assert_eq!(model.n_params(), 0); + assert_eq!(model.output_len(), 2); + assert!(model.cj().abs() < f64::EPSILON); + } + + #[test] + fn test_compiled_model_eval_rhs() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); // f(y) = 2*y + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + let y_vals = [1.0, 2.0]; + let mut output = [0.0, 0.0]; + model.eval_rhs(0.0, &y_vals, &[], &mut output); + + // 2 * [1, 2] = [2, 4] + assert!((output[0] - 2.0).abs() < 1e-14); + assert!((output[1] - 4.0).abs() < 1e-14); + } + + #[test] + fn test_compiled_model_jac_mul_no_mass() { + // f(y) = 2*y, so df/dy = 2*I + // With cj=0: (df/dy - 0*M) @ v = 2*I @ v = 2*v + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + // cj = 0, so mass matrix doesn't contribute + model.set_cj(0.0); + + let y_vals = [1.0, 2.0]; + let v = [1.0, 0.0]; + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v, &mut output); + + // df/dy @ [1, 0] = 2 * [1, 0] = [2, 0] + assert!( + (output[0] - 2.0).abs() < 1e-14, + "Expected 2.0, got {}", + output[0] + ); + assert!(output[1].abs() < 1e-14, "Expected 0.0, got {}", output[1]); + } + + #[test] + fn test_compiled_model_jac_mul_with_mass() { + // f(y) = 2*y, so df/dy = 2*I + // With cj=0.5 and M=I: (df/dy - cj*M) = 2*I - 0.5*I = 1.5*I + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + model.set_cj(0.5); + + let y_vals = [1.0, 2.0]; + let v = [1.0, 0.0]; + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v, &mut output); + + // (2 - 0.5) * [1, 0] = [1.5, 0] + assert!( + (output[0] - 1.5).abs() < 1e-14, + "Expected 1.5, got {}", + output[0] + ); + assert!(output[1].abs() < 1e-14, "Expected 0.0, got {}", output[1]); + } + + #[test] + fn test_compiled_model_jac_mul_second_direction() { + // Same model, but test v = [0, 1] + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + model.set_cj(0.5); + + let y_vals = [1.0, 2.0]; + let v = [0.0, 1.0]; + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v, &mut output); + + // (2 - 0.5) * [0, 1] = [0, 1.5] + assert!(output[0].abs() < 1e-14, "Expected 0.0, got {}", output[0]); + assert!( + (output[1] - 1.5).abs() < 1e-14, + "Expected 1.5, got {}", + output[1] + ); + } + + #[test] + fn test_compiled_model_sparsity() { + // f(y) = [sin(y0), sin(y1)]: diagonal but state-dependent, so the + // coloring is what has to make assembly cheap here. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let rhs = arena.alloc(Node::Sin(y)); + + let mass = identity_mass_matrix(2); + let model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + // Diagonal Jacobian should only need 1 color + assert_eq!( + model.coloring().n_colors, + 1, + "Diagonal should need 1 color, got {}", + model.coloring().n_colors + ); + } + + /// Nonlinear banded ODE whose 5-wide stencil needs enough colours to + /// engage the batched sweep. + fn build_batched_model(n: usize) -> ModelEvaluator { + build_banded_model(n, 2) + } + + /// Nonlinear ODE with a one-sided stencil, whose colour count is `reach + 1` + /// -- so unlike the symmetric builder it can land on an even count, which is + /// what puts a tail of exactly two colours under test. + fn build_one_sided_model(n: usize, reach: usize) -> ModelEvaluator { + let mut arena = Arena::new(); + let states: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + let rows: Vec<_> = (0..n) + .map(|i| { + let mut acc = arena.alloc(Node::Exp(states[i])); + for offset in 1..=reach { + let term = arena.alloc(Node::Sin(states[(i + offset).min(n - 1)])); + acc = arena.alloc(Node::Add(acc, term)); + } + acc + }) + .collect(); + let rhs = arena.alloc(Node::Concat(rows)); + ModelEvaluator::new(&arena, rhs, identity_mass_matrix(n), n, 0) + } + + /// Nonlinear banded ODE of a chosen half-bandwidth, so the colour count -- + /// and with it the block/tail split -- can be dialled. + fn build_banded_model(n: usize, half_width: usize) -> ModelEvaluator { + let (arena, rhs) = banded_nonlinear(n, half_width); + ModelEvaluator::new(&arena, rhs, identity_mass_matrix(n), n, 0) + } + + /// A tail narrower than the lane width takes its own narrower walk, so the + /// block loop must still reproduce the unbatched tape column for column. + #[test] + fn a_narrow_tail_block_matches_the_unbatched_tape() { + // Symmetric half-widths give odd colour counts, one-sided reaches give + // even ones, so between them the tails cover every shape: scalar tails + // under lane 4 and lane 8, tails of two and three (which straddle + // MIN_PADDED_TAIL), a narrower vector block, and a vector block followed + // by several scalar walks. + let fixtures = [ + (64usize, 2usize, true), + (64, 4, true), + (64, 6, true), + (64, 7, true), + ] + .into_iter() + .chain([ + (64, 5, false), + (64, 9, false), + (64, 11, false), + (64, 13, false), + ]); + for (n, reach, symmetric) in fixtures { + let half_width = reach; + let mut model = if symmetric { + build_banded_model(n, half_width) + } else { + build_one_sided_model(n, reach) + }; + let stats = model.jacobian_stats(); + let (colors, lanes) = (stats.n_colors, stats.jac_lane_width); + assert!(lanes > 1, "n={n} w={half_width} must batch"); + assert_ne!( + colors % lanes, + 0, + "n={n} w={half_width} must leave a partial tail block" + ); + + // cj = 0 keeps this a pure df/dy comparison: no mass postpass. + model.set_cj(0.0); + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.037, 0.41)).collect(); + let mut assembled = vec![f64::NAN; model.nnz()]; + model.assemble_jacobian_csc_into(0.0, &y, &[], &mut assembled); + + let (colptr, rowind) = { + let csc = model.csc_sparsity(); + (csc.colptr.clone(), csc.rowind.clone()) + }; + let mut column = vec![0.0; n]; + let mut seed = vec![0.0; n]; + for col in 0..n { + seed.fill(0.0); + seed[col] = 1.0; + model.jac_action(0.0, &y, &[], &seed, &mut column); + for k in colptr[col]..colptr[col + 1] { + let (row, got) = (rowind[k], assembled[k]); + let want = column[row]; + assert!( + got.to_bits() == want.to_bits() || (got == 0.0 && want == 0.0), + "n={n} reach={reach} entry ({row}, {col}): batched {got}, tape {want}" + ); + } + } + } + } + + #[test] + fn batched_seeds_are_restored_between_assemblies() { + // The batched sweep clears only the lanes it set, so a stale 1.0 left + // behind would silently add another column into a later block. + let n = 24; + let mut model = build_batched_model(n); + assert!(model.jacobian_stats().jac_lane_width > 1, "needs batching"); + model.set_cj(0.25); + + let states: [Vec; 2] = [ + (0..n).map(|i| (i as f64).mul_add(0.11, 0.3)).collect(), + (0..n).map(|i| (i as f64).mul_add(-0.07, 0.9)).collect(), + ]; + let mut repeated = vec![0.0; model.nnz()]; + for _ in 0..3 { + for y in &states { + model.assemble_jacobian_csc_into(0.0, y, &[], &mut repeated); + } + } + + let mut fresh_model = build_batched_model(n); + fresh_model.set_cj(0.25); + let mut fresh = vec![0.0; fresh_model.nnz()]; + fresh_model.assemble_jacobian_csc_into(0.0, &states[1], &[], &mut fresh); + assert!( + repeated + .iter() + .zip(&fresh) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "a repeated assembly must match a first one bit for bit" + ); + } + + #[test] + fn a_linear_model_assembles_without_sweeping() { + // f(y) = [y0, y1]: every entry folds at compile time, so the coloring + // seeds nothing and assembly is a table write. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let mut model = ModelEvaluator::new(&arena, y, identity_mass_matrix(2), 2, 0); + + let stats = model.jacobian_stats(); + assert_eq!(stats.n_colors, 0); + assert_eq!(stats.n_swept_columns, 0); + assert_eq!(stats.n_constant_entries, 2); + + model.set_cj(0.25); + let mut data = vec![f64::NAN; model.nnz()]; + model.assemble_jacobian_csc_into_coloring(0.0, &[7.0, -3.0], &[], &mut data); + assert_eq!(data, vec![0.75, 0.75]); + } + + #[test] + fn test_compiled_model_nonlinear() { + // f(y) = y^2, so df/dy = 2*y (diagonal) + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Pow(y, two)); // y^2 + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + model.set_cj(0.0); + + // At y = [3, 4], df/dy = diag([6, 8]) + let y_vals = [3.0, 4.0]; + let v = [1.0, 1.0]; + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v, &mut output); + + // df/dy @ [1, 1] = [6, 8] + assert!( + (output[0] - 6.0).abs() < 1e-12, + "Expected 6.0, got {}", + output[0] + ); + assert!( + (output[1] - 8.0).abs() < 1e-12, + "Expected 8.0, got {}", + output[1] + ); + } + + #[test] + fn test_compiled_model_with_params() { + // f(y) = k * y where k is a parameter + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 1); + + assert_eq!(model.n_params(), 1); + + // Test eval_rhs with k=3 + let y_vals = [1.0, 2.0]; + let inputs = [3.0]; + let mut output = [0.0, 0.0]; + model.eval_rhs(0.0, &y_vals, &inputs, &mut output); + + // 3 * [1, 2] = [3, 6] + assert!((output[0] - 3.0).abs() < 1e-14); + assert!((output[1] - 6.0).abs() < 1e-14); + + // Test jac_mul: df/dy = k*I = 3*I + model.set_cj(1.0); + let v = [1.0, 0.0]; + model.jac_mul(0.0, &y_vals, &inputs, &v, &mut output); + + // (3 - 1) * [1, 0] = [2, 0] + assert!( + (output[0] - 2.0).abs() < 1e-14, + "Expected 2.0, got {}", + output[0] + ); + } + + #[test] + fn test_csr_matvec() { + // Test sparse matrix-vector product + // M = [[2, 0], [0, 3]] + let csr = CsrData { + shape: Shape::matrix(2, 2), + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], + }; + + let x = [1.0, 2.0]; + let mut y = [0.0, 0.0]; + csr_matvec(&csr, &x, &mut y); + + // [2*1, 3*2] = [2, 6] + assert!((y[0] - 2.0).abs() < 1e-14); + assert!((y[1] - 6.0).abs() < 1e-14); + } + + #[test] + fn test_csr_matvec_tridiagonal() { + // M = [[1, 2, 0], [3, 4, 5], [0, 6, 7]] + let csr = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 2, 5, 7], + indices: vec![0, 1, 0, 1, 2, 1, 2], + data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + }; + + let x = [1.0, 2.0, 3.0]; + let mut y = [0.0, 0.0, 0.0]; + csr_matvec(&csr, &x, &mut y); + + // y = [1*1 + 2*2, 3*1 + 4*2 + 5*3, 6*2 + 7*3] + assert!((y[0] - 5.0).abs() < 1e-14); + assert!((y[1] - 26.0).abs() < 1e-14); + assert!((y[2] - 33.0).abs() < 1e-14); + } + + #[test] + fn test_compiled_model_debug() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + let debug_str = format!("{model:?}"); + assert!(debug_str.contains("ModelEvaluator")); + assert!(debug_str.contains("n_states: 2")); + } + + #[test] + fn test_compiled_model_residual() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); // f(y) = 2*y + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + let y_vals = [1.0, 2.0]; + let yp = [3.0, 4.0]; + let mut r = [0.0, 0.0]; + model.eval_residual(0.0, &y_vals, &yp, &[], &mut r); + + // M*y' - f(y) = I*[3, 4] - 2*[1, 2] = [3, 4] - [2, 4] = [1, 0] + assert!((r[0] - 1.0).abs() < 1e-14, "Expected 1.0, got {}", r[0]); + assert!(r[1].abs() < 1e-14, "Expected 0.0, got {}", r[1]); + } + + #[test] + fn test_compiled_model_with_zero_tangent_direction() { + // When v = 0, zero propagation should optimize away most of the JVP computation + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.5); + + let y_vals = [1.0, 2.0]; + let v = [0.0, 0.0]; // Zero tangent direction + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v, &mut output); + + // Result should be [0, 0] + assert!(output[0].abs() < 1e-14); + assert!(output[1].abs() < 1e-14); + } + + #[test] + fn test_compiled_model_residual_nonidentity_mass() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + + // Non-identity mass matrix: M = diag([2, 3]) + let mass = CsrData { + shape: Shape::matrix(2, 2), + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], + }; + + let mut model = ModelEvaluator::new(&arena, y, mass, 2, 0); + + let y_vals = [1.0, 2.0]; + let yp = [3.0, 4.0]; + let mut r = [0.0, 0.0]; + model.eval_residual(0.0, &y_vals, &yp, &[], &mut r); + + // M*y' - f(y) = [2*3, 3*4] - [1, 2] = [6, 12] - [1, 2] = [5, 10] + assert!((r[0] - 5.0).abs() < 1e-14, "Expected 5.0, got {}", r[0]); + assert!((r[1] - 10.0).abs() < 1e-14, "Expected 10.0, got {}", r[1]); + } + + #[test] + fn test_algebraic_ids_inferred_from_mass() { + // Mass matrix: row 0 has diag, row 1 missing diag, row 2 has diag. + // Row 1 is therefore algebraic. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 2], + indices: vec![0, 2], + data: vec![1.0, 1.0], + }; + let model = ModelEvaluator::new(&arena, y, mass, 3, 0); + assert_eq!(model.algebraic_ids(), &[false, true, false]); + } + + #[test] + fn test_algebraic_ids_f64_for_ida() { + // 1.0 = differential, 0.0 = algebraic (matches IDA's id vector format). + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 2], + indices: vec![0, 2], + data: vec![1.0, 1.0], + }; + let model = ModelEvaluator::new(&arena, y, mass, 3, 0); + let mut ids = vec![0.0; 3]; + model.algebraic_ids_f64(&mut ids); + assert_eq!(ids, vec![1.0, 0.0, 1.0]); + } + + #[test] + fn test_algebraic_ids_identity_mass_all_differential() { + // Identity mass = pure ODE: every state is differential. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let model = ModelEvaluator::new(&arena, y, identity_mass_matrix(4), 4, 0); + assert_eq!(model.algebraic_ids(), &[false, false, false, false]); + } + + /// Shared fixture: f(y) = k * y^2 with k as `InputParameter` index 0. + fn build_sens_test_model() -> (ModelEvaluator, [f64; 2]) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let y_sq = arena.alloc(Node::Pow(y, two)); + let rhs = arena.alloc(Node::Mul(k, y_sq)); + + let mass = identity_mass_matrix(2); + let model = ModelEvaluator::new_with_sens(&arena, rhs, mass, 2, 1, &[0]); + (model, [3.0, 4.0]) + } + + /// Shared fixture: f(y) = k1 * y; sensitivities for both k1 (idx 0) and k2 (idx 1). + fn build_two_param_sens_model() -> (ModelEvaluator, [f64; 2]) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k1 = arena.alloc(Node::InputParameter { + name: "k1".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let _k2 = arena.alloc(Node::InputParameter { + name: "k2".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k1, y)); + + let mass = identity_mass_matrix(2); + let model = ModelEvaluator::new_with_sens(&arena, rhs, mass, 2, 2, &[0, 1]); + (model, [3.0, 4.0]) + } + + #[test] + fn test_sens_expr_compiled_with_split_eval() { + // Same fixture as test_eval_sens_matches_analytical (f = k*y^2, sens wrt k). + let (model, _y) = build_sens_test_model(); + let sens_expr = model.compiled.sens_expr.as_ref().expect("sens expr"); + assert!( + sens_expr.has_split_eval(), + "sens_expr must be split-eval compiled so the primal runs once per (t,y)" + ); + } + + #[test] + fn test_eval_sens_matches_analytical() { + // f(y) = k * y^2 with k as InputParameter index 0. ∂f/∂k = y^2. + let (mut model, y_vals) = build_sens_test_model(); + assert_eq!(model.n_sens_params(), 1); + + // y = [3, 4], k = 7 -> ∂f/∂k = y^2 = [9, 16] + let inputs = [7.0]; + let mut out = [0.0; 2]; + model.eval_sens(0.0, &y_vals, &inputs, 0, &mut out); + assert!((out[0] - 9.0).abs() < 1e-10, "expected 9, got {}", out[0]); + assert!((out[1] - 16.0).abs() < 1e-10, "expected 16, got {}", out[1]); + } + + #[test] + fn test_eval_sens_all_layout() { + // f(y) = k1 * y; sensitivities for both k1 (idx 0) and k2 (idx 1). + // ∂f/∂k1 = y, ∂f/∂k2 = 0. + let (mut model, y_vals) = build_two_param_sens_model(); + + let inputs = [5.0, 9.0]; + let mut out = vec![-1.0; 4]; + model.eval_sens_all(0.0, &y_vals, &inputs, &mut out); + + // Layout: out[i*n_states + j] = ∂f_j/∂p_i + // i=0 (k1): [3, 4]; i=1 (k2): [0, 0] + assert!((out[0] - 3.0).abs() < 1e-10); + assert!((out[1] - 4.0).abs() < 1e-10); + assert!(out[2].abs() < 1e-10); + assert!(out[3].abs() < 1e-10); + } + + #[test] + fn test_eval_sens_all_matches_single_column_eval() { + // eval_sens_all (primal-once batch) must agree with per-column eval_sens + // (full-stream) exactly, same tape, same arithmetic. + let (mut model, y) = build_two_param_sens_model(); + let n = model.compiled.n_states; + let k = model.compiled.n_sens_params(); + let mut batched = vec![0.0; n * k]; + model.eval_sens_all(0.0, &y, &[1.5, 2.5], &mut batched); + for i in 0..k { + let mut single = vec![0.0; n]; + model.eval_sens(0.0, &y, &[1.5, 2.5], i, &mut single); + assert_eq!( + &batched[i * n..(i + 1) * n], + &single[..], + "column {i} diverged" + ); + } + } + + /// Shared fixture whose requested subset is not a leading prefix. + /// + /// `f(y) = k0 * y + k1 * y^2` and `H = k0 * y[0] + k1 * y[1]`, with sensitivities + /// requested for k1 (index 1) alone. Every parameter derivative is distinct and + /// non-zero, so a seed landing on the wrong parameter shows up as a wrong value + /// rather than a zero. + fn build_trailing_param_sens_model() -> (CompiledModel, [f64; 2]) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y, + start: 1, + end: 2, + }); + let k0 = arena.alloc(Node::InputParameter { + name: "k0".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let k1 = arena.alloc(Node::InputParameter { + name: "k1".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let y_squared = arena.alloc(Node::Pow(y, two)); + let linear = arena.alloc(Node::Mul(k0, y)); + let quadratic = arena.alloc(Node::Mul(k1, y_squared)); + let rhs = arena.alloc(Node::Add(linear, quadratic)); + + let out_left = arena.alloc(Node::Mul(k0, y0)); + let out_right = arena.alloc(Node::Mul(k1, y1)); + let out_node = arena.alloc(Node::Add(out_left, out_right)); + + let mut data = + CompiledModel::new_with_sens(&arena, rhs, identity_mass_matrix(2), 2, 2, &[1]); + data.add_output(&arena, out_node); + (data, [3.0, 4.0]) + } + + #[test] + fn sens_action_seeds_in_parameter_space() { + // df/dk1 = y^2. The seed is indexed by parameter, so its 1.0 sits at index 1; + // reading it as a subset-space seed would pick up the 0.0 at index 0 instead. + let (model, y) = build_trailing_param_sens_model(); + let mut ws = model.create_workspace(); + let mut got = [0.0; 2]; + model.sens_action( + &mut ws, + 0.0, + &y, + &[5.0, 9.0], + &[0, 1], + &[0.0, 1.0], + &mut got, + ); + assert!((got[0] - 9.0).abs() < 1e-10, "expected 9, got {}", got[0]); + assert!((got[1] - 16.0).abs() < 1e-10, "expected 16, got {}", got[1]); + } + + #[test] + fn sens_action_maps_a_subset_seed_to_its_global_parameter() { + // A one-entry subset seed naming k1 must differentiate k1, not whatever + // parameter sits at position 0 of the seed. + let (model, y) = build_trailing_param_sens_model(); + let mut ws = model.create_workspace(); + let inputs = [5.0, 9.0]; + + let mut subset = [0.0; 2]; + model.sens_action(&mut ws, 0.0, &y, &inputs, &[1], &[1.0], &mut subset); + + let mut global = [0.0; 2]; + model.sens_action(&mut ws, 0.0, &y, &inputs, &[0, 1], &[0.0, 1.0], &mut global); + + assert_eq!(subset.as_slice(), global.as_slice()); + // df/dk1 = y^2 = [9, 16] + assert!((subset[0] - 9.0).abs() < 1e-10, "got {}", subset[0]); + assert!((subset[1] - 16.0).abs() < 1e-10, "got {}", subset[1]); + } + + #[test] + fn an_output_sens_action_maps_a_subset_seed_to_its_global_parameter() { + // H = k0*y0 + k1*y1, so dH/dk1 = y1 = 4.0. Reading the seed as a prefix + // would give dH/dk0 = y0 = 3.0 instead. + let (model, y) = build_trailing_param_sens_model(); + let mut ws = model.create_workspace(); + let mut got = vec![0.0; model.total_output_len()]; + model.observable_sens_action( + &mut ws, + ObservableKind::Outputs, + 0.0, + &y, + &[5.0, 9.0], + &[1], + &[1.0], + &mut got, + ); + assert!((got[0] - 4.0).abs() < 1e-10, "expected 4, got {}", got[0]); + } + + #[test] + fn sens_tangent_column_is_named_by_parameter_index() { + // Column 0 is k0, whose df/dk0 = y, even though only k1 was registered: the + // compiled tangent is generic over the parameters and the seed picks the column. + let (model, y) = build_trailing_param_sens_model(); + let mut ws = model.create_workspace(); + model.sens_primal_pass(&mut ws, 0.0, &y, &[5.0, 9.0]); + let mut got = [0.0; 2]; + model.sens_tangent_column(&mut ws, 0, &mut got); + assert!((got[0] - 3.0).abs() < 1e-10, "expected 3, got {}", got[0]); + assert!((got[1] - 4.0).abs() < 1e-10, "expected 4, got {}", got[1]); + } + + #[test] + fn output_sens_project_seeds_the_requested_parameter() { + // dH/dk1 = y[1] = 4 and dH/dk0 = y[0] = 3, so projecting the one requested + // column must give 4: the subset-to-parameter mapping belongs to the caller. + let (model, y) = build_trailing_param_sens_model(); + let mut ws = model.create_workspace(); + let y_sens = vec![0.0; model.n_states]; + let mut got = vec![0.0; model.total_output_len()]; + model.output_sens_project(&mut ws, 0.0, &y, &[5.0, 9.0], &y_sens, &mut got); + assert_eq!(got.len(), 1); + assert!((got[0] - 4.0).abs() < 1e-10, "expected 4, got {}", got[0]); + } + + #[test] + fn test_n_sens_params_zero_when_no_sensitivities() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let model = ModelEvaluator::new(&arena, y, identity_mass_matrix(2), 2, 0); + assert_eq!(model.n_sens_params(), 0); + } + + #[test] + fn test_eval_sens_all_noop_without_sensitivities() { + // Zero configured sensitivities: eval_sens_all must return without + // touching sens_expr (absent), matching the old 0-column loop. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let mut model = ModelEvaluator::new(&arena, y, identity_mass_matrix(2), 2, 0); + let mut out: [f64; 0] = []; + model.eval_sens_all(0.0, &[3.0, 4.0], &[], &mut out); + } + + #[test] + fn test_add_output_compiles_and_evaluates() { + // f(y) = y is the rhs; output_var = 2 * y[0]. + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let var0 = arena.alloc(Node::Mul(two, y0)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_output(&arena, var0); + + assert_eq!(model.n_outputs(), 1); + assert_eq!(model.output_len_at(0), 1); + + // var0(y=[3, 4]) = 2 * 3 = 6 + let mut out = [0.0; 1]; + let written = model.eval_output(0.0, &[3.0, 4.0], &[], 0, &mut out); + assert_eq!(written, 1); + assert!((out[0] - 6.0).abs() < 1e-12); + } + + #[test] + fn test_n_outputs_zero_initially() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let model = ModelEvaluator::new(&arena, y, identity_mass_matrix(2), 2, 0); + assert_eq!(model.n_outputs(), 0); + } + + #[test] + fn test_outputs_compose_with_sensitivities() { + // Verify add_output works on a model that already has sensitivities. + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y_full)); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new_with_sens(&arena, rhs, mass, 2, 1, &[0]); + model.add_output(&arena, y0); + + assert_eq!(model.n_sens_params(), 1); + assert_eq!(model.n_outputs(), 1); + assert_eq!(model.output_len_at(0), 1); + + let y_vals = [3.0, 4.0]; + let inputs = [2.0]; + + // ∂f/∂k = y -> [3, 4] + let mut sens_out = [0.0; 2]; + model.eval_sens(0.0, &y_vals, &inputs, 0, &mut sens_out); + assert!((sens_out[0] - 3.0).abs() < 1e-12); + assert!((sens_out[1] - 4.0).abs() < 1e-12); + + // var0 = y[0] = 3 + let mut out = [0.0; 1]; + model.eval_output(0.0, &y_vals, &inputs, 0, &mut out); + assert!((out[0] - 3.0).abs() < 1e-12); + } + + #[test] + fn test_new_with_options_compiles_algebraic_and_sensitivities() { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let y2 = arena.alloc(Node::Index { + child: y_full, + start: 2, + end: 3, + }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y_full)); + let algebraic_sum = arena.alloc(Node::Add(y1, y2)); + let algebraic_prod = arena.alloc(Node::Mul(y1, y2)); + let alg = arena.alloc(Node::Concat(vec![algebraic_sum, algebraic_prod])); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 1], + indices: vec![0], + data: vec![1.0], + }; + + let options = CompiledModelOptions::new() + .with_sensitivities(&[0]) + .with_algebraic(CompiledModelAlgebraicBlock::new(alg, &[1, 2])); + let mut model = ModelEvaluator::new_with_options(&arena, rhs, mass, 3, 1, options); + + assert_eq!(model.n_sens_params(), 1); + assert!(model.has_algebraic()); + assert_eq!(model.n_algebraic(), 2); + + let y_vals = [3.0, 4.0, 5.0]; + let inputs = [2.0]; + + let mut sens_out = [0.0; 3]; + model.eval_sens(0.0, &y_vals, &inputs, 0, &mut sens_out); + for (actual, expected) in sens_out.iter().zip(y_vals.iter()) { + assert!((actual - expected).abs() < 1e-12); + } + + let mut algebraic_out = [0.0; 2]; + model.eval_algebraic_residual(0.0, &y_vals, &inputs, &mut algebraic_out); + assert!((algebraic_out[0] - 9.0).abs() < 1e-12); + assert!((algebraic_out[1] - 20.0).abs() < 1e-12); + } + + #[test] + fn test_algebraic_jac_action_uses_algebraic_index_mapping() { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let y2 = arena.alloc(Node::Index { + child: y_full, + start: 2, + end: 3, + }); + let alg0_mul = arena.alloc(Node::Mul(y1, y2)); + let alg0 = arena.alloc(Node::Add(y0, alg0_mul)); + let two = arena.alloc(Node::Scalar(2.0)); + let y2_sq = arena.alloc(Node::Pow(y2, two)); + let alg1 = arena.alloc(Node::Add(y1, y2_sq)); + let alg = arena.alloc(Node::Concat(vec![alg0, alg1])); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 1], + indices: vec![0], + data: vec![1.0], + }; + + let mut model = + ModelEvaluator::new_with_algebraic(&arena, y_full, mass, 3, 0, Some(alg), &[1, 2]); + let y_vals = [10.0, 2.0, 3.0]; + let v_alg = [5.0, 7.0]; + let mut jv = [0.0; 2]; + model.eval_algebraic_jacobian_action(0.0, &y_vals, &[], &v_alg, &mut jv); + + assert!((jv[0] - 29.0).abs() < 1e-12, "expected 29, got {}", jv[0]); + assert!((jv[1] - 47.0).abs() < 1e-12, "expected 47, got {}", jv[1]); + } + + #[test] + fn test_assemble_alg_jacobian_into_matches_expected_sparse_values() { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let y2 = arena.alloc(Node::Index { + child: y_full, + start: 2, + end: 3, + }); + let alg0_mul = arena.alloc(Node::Mul(y1, y2)); + let alg0 = arena.alloc(Node::Add(y0, alg0_mul)); + let two = arena.alloc(Node::Scalar(2.0)); + let y2_sq = arena.alloc(Node::Pow(y2, two)); + let alg1 = arena.alloc(Node::Add(y1, y2_sq)); + let alg = arena.alloc(Node::Concat(vec![alg0, alg1])); + let mass = CsrData { + shape: Shape::matrix(3, 3), + indptr: vec![0, 1, 1, 1], + indices: vec![0], + data: vec![1.0], + }; + + let mut model = + ModelEvaluator::new_with_algebraic(&arena, y_full, mass, 3, 0, Some(alg), &[1, 2]); + + // The block shares `JacobianData`'s CSC ordering, so the triplet runs + // column-major: global column 1, then 2. Both C++ Newton drivers + // bucket-sort it, so the order is ours to choose -- but it is published + // alongside the values and must stay aligned with them. + assert_eq!(model.algebraic_jacobian_row_indices(), &[0, 1, 0, 1]); + assert_eq!(model.algebraic_jacobian_col_indices(), &[1, 1, 2, 2]); + + let mut jac = vec![0.0; model.algebraic_jacobian_nnz()]; + model.assemble_algebraic_jacobian_into(0.0, &[10.0, 2.0, 3.0], &[], &mut jac); + // dg/d(y1, y2) = [[y2, y1], [1, 2*y2]] at (y1, y2) = (2, 3). + assert_eq!(jac, vec![3.0, 1.0, 2.0, 6.0]); + } + + #[test] + fn test_jacobian_with_zero_elimination() { + // Build a model where some Jacobian columns are structurally zero + // Verify that zero propagation eliminates dead computation + let mut arena = Arena::new(); + + // f(y) = [y0 * y1, y0 + c], so df/dy = [[y1, y0], [1, 0]] and + // df1/dy1 is a structural zero. + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let c = arena.alloc(Node::Scalar(5.0)); + let prod = arena.alloc(Node::Mul(y0, y1)); + let sum = arena.alloc(Node::Add(y0, c)); + let rhs = arena.alloc(Node::Concat(vec![prod, sum])); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.0); + + // Test jac_mul with [1, 0] - should give [y1, 1] + let y_vals = [2.0, 3.0]; + let v1 = [1.0, 0.0]; + let mut output = [0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &v1, &mut output); + + assert!( + (output[0] - 3.0).abs() < 1e-12, + "df1/dy0 = y1 = 3, got {}", + output[0] + ); + assert!( + (output[1] - 1.0).abs() < 1e-12, + "df2/dy0 = 1, got {}", + output[1] + ); + + // Test jac_mul with [0, 1] - should give [y0, 0] + let v2 = [0.0, 1.0]; + model.jac_mul(0.0, &y_vals, &[], &v2, &mut output); + + assert!( + (output[0] - 2.0).abs() < 1e-12, + "df1/dy1 = y0 = 2, got {}", + output[0] + ); + assert!(output[1].abs() < 1e-12, "df2/dy1 = 0, got {}", output[1]); + } + + #[test] + fn test_assemble_jacobian_linear() { + // f(y) = 2*y, so df/dy = 2*I (diagonal) + // With cj=0: J = df/dy = 2*I + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.0); + + let y_vals = [1.0, 2.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Should have 2 non-zeros (diagonal) + assert_eq!(rows.len(), 2); + assert_eq!(cols.len(), 2); + assert_eq!(vals.len(), 2); + + // Build dense matrix to verify + let mut dense = [[0.0; 2]; 2]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + // J = 2*I + assert!( + (dense[0][0] - 2.0).abs() < 1e-12, + "J[0,0] = {}", + dense[0][0] + ); + assert!( + (dense[1][1] - 2.0).abs() < 1e-12, + "J[1,1] = {}", + dense[1][1] + ); + assert!(dense[0][1].abs() < 1e-12, "J[0,1] = {}", dense[0][1]); + assert!(dense[1][0].abs() < 1e-12, "J[1,0] = {}", dense[1][0]); + } + + #[test] + fn test_assemble_jacobian_with_mass() { + // f(y) = 2*y, M = I + // With cj=0.5: J = df/dy - cj*M = 2*I - 0.5*I = 1.5*I + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.5); + + let y_vals = [1.0, 2.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Build dense matrix to verify + let mut dense = [[0.0; 2]; 2]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + // J = 1.5*I + assert!( + (dense[0][0] - 1.5).abs() < 1e-12, + "J[0,0] = {}", + dense[0][0] + ); + assert!( + (dense[1][1] - 1.5).abs() < 1e-12, + "J[1,1] = {}", + dense[1][1] + ); + } + + /// Square DAE whose last row is a dense algebraic equation: + /// `f_i = sin(y_i)` for `i in 0..n-1` (differential, mass diagonal 1) and + /// `f_{n-1} = sum_j y_j^2` (algebraic, mass diagonal 0), depending on every + /// state. The dense row's `>= 16` columns also live on the diagonal rows + /// (the aliasing trap), so the reduced coloring must not sum them. + fn build_dense_row_dae(n: usize) -> (ModelEvaluator, Vec) { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: n }); + // Differential block: sin over the first n-1 states (diagonal). + let y_head = arena.alloc(Node::StateVector { + start: 0, + end: n - 1, + }); + let diff = arena.alloc(Node::Sin(y_head)); + // Dense algebraic row: ones(1 x n) @ (y*y) = sum_j y_j^2. + let gy = arena.alloc(Node::Mul(y_full, y_full)); + let ones = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, n], + indices: (0..n).collect(), + data: vec![1.0; n], + shape: Shape::matrix(1, n), + }))); + let dense = arena.alloc(Node::MatMul(ones, gy)); + let rhs = arena.alloc(Node::Concat(vec![diff, dense])); + // Mass: diag(1 for the n-1 differential rows, 0 for the algebraic row). + let mass = CsrData { + shape: Shape::matrix(n, n), + indptr: (0..n).chain(std::iter::once(n - 1)).collect(), + indices: (0..n - 1).collect(), + data: vec![1.0; n - 1], + }; + let model = ModelEvaluator::new(&arena, rhs, mass, n, 0); + let y: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.1, 0.3).sin() + 0.7) + .collect(); + (model, y) + } + + /// Differential dense row whose `df/dy` excludes its own column + /// (`f_0 = sum_{j>=1} y_j^2`), so with an identity mass the merged row 0 + /// carries column 0 as a mass-only slot the sub-Jacobian lacks, the + /// "merged has more columns than the sub-row" remap case. Rows `1..n-1` are + /// diagonal `sin(y_i)`. + fn build_dense_row_dae_mass_only_column(n: usize) -> (ModelEvaluator, Vec) { + let mut arena = Arena::new(); + let y_tail = arena.alloc(Node::StateVector { start: 1, end: n }); + let gy = arena.alloc(Node::Mul(y_tail, y_tail)); + let ones = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, n - 1], + indices: (0..n - 1).collect(), + data: vec![1.0; n - 1], + shape: Shape::matrix(1, n - 1), + }))); + let dense = arena.alloc(Node::MatMul(ones, gy)); + let diag = arena.alloc(Node::Sin(y_tail)); + let rhs = arena.alloc(Node::Concat(vec![dense, diag])); + let mass = identity_mass_matrix(n); + let model = ModelEvaluator::new(&arena, rhs, mass, n, 0); + let y: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.1, 0.4).cos() + 0.6) + .collect(); + (model, y) + } + + /// Dense df/dy via central finite differences of `eval_rhs`. + fn finite_difference_jacobian( + model: &mut ModelEvaluator, + t: f64, + y: &[f64], + inputs: &[f64], + eps: f64, + ) -> Vec> { + let n = model.n_states(); + let mut fp = vec![0.0; n]; + let mut fm = vec![0.0; n]; + let mut dense = vec![vec![0.0; n]; n]; + for col in 0..n { + let mut yp = y.to_vec(); + let mut ym = y.to_vec(); + yp[col] += eps; + ym[col] -= eps; + model.eval_rhs(t, &yp, inputs, &mut fp); + model.eval_rhs(t, &ym, inputs, &mut fm); + for row in 0..n { + dense[row][col] = (fp[row] - fm[row]) / (2.0 * eps); + } + } + dense + } + + fn assert_matrices_match(assembled: &[Vec], expected: &[Vec], tol: f64, ctx: &str) { + let n = expected.len(); + for row in 0..n { + for col in 0..n { + let err = (assembled[row][col] - expected[row][col]).abs(); + let scale = assembled[row][col] + .abs() + .max(expected[row][col].abs()) + .max(1e-12); + assert!( + err / scale < tol || err < tol, + "{ctx} J({row},{col}): assembled={}, expected={}", + assembled[row][col], + expected[row][col] + ); + } + } + } + + fn coo_to_dense(rows: &[usize], cols: &[usize], vals: &[f64], n: usize) -> Vec> { + let mut dense = vec![vec![0.0; n]; n]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + dense + } + + fn csc_to_dense(model: &ModelEvaluator, csc_vals: &[f64], n: usize) -> Vec> { + let csc = model.csc_sparsity(); + let mut dense = vec![vec![0.0; n]; n]; + for (col, span) in csc.colptr.windows(2).enumerate() { + for k in span[0]..span[1] { + dense[csc.rowind[k]][col] = csc_vals[k]; + } + } + dense + } + + #[test] + fn test_all_assembly_paths_agree_with_dense_row_split() { + let n = 20; + let (mut model, y) = build_dense_row_dae(n); + assert_eq!( + model.compiled.jac.n_dense_rows(), + 1, + "dense-row split must be active" + ); + assert_eq!(model.compiled.jac.dense_rows()[0].rows[0], n - 1); + let telemetry = model.jacobian_stats(); + assert_eq!(telemetry.n_dense_rows, 1); + assert_eq!(telemetry.dense_row_entries, n); + assert!(telemetry.dense_row_tape_instructions > 0); + + let dense_fd = finite_difference_jacobian(&mut model, 0.0, &y, &[], 1e-6); + + // Loop B (COO), cj = 0 -> J = df/dy. + model.set_cj(0.0); + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y, &[]); + assert_matrices_match( + &coo_to_dense(&rows, &cols, &vals, n), + &dense_fd, + 1e-5, + "COO", + ); + + // Loop D (CSC, no mass) -> pure df/dy. + let mut csc_no_mass = vec![0.0; model.nnz()]; + model.assemble_jacobian_csc_no_mass(0.0, &y, &[], &mut csc_no_mass); + assert_matrices_match( + &csc_to_dense(&model, &csc_no_mass, n), + &dense_fd, + 1e-5, + "CSC-no-mass", + ); + + // Loop C (CSC, with mass postpass), cj != 0 -> J = df/dy - cj*M. Mass is + // diag(1 on the n-1 differential rows, 0 on the algebraic dense row). + let cj = 0.7; + model.set_cj(cj); + let mut csc_mass = vec![0.0; model.nnz()]; + model.assemble_jacobian_csc_into_coloring(0.0, &y, &[], &mut csc_mass); + let mut expected = dense_fd; + for (i, row) in expected.iter_mut().enumerate().take(n - 1) { + row[i] -= cj; + } + assert_matrices_match( + &csc_to_dense(&model, &csc_mass, n), + &expected, + 1e-5, + "CSC-mass", + ); + } + + #[test] + fn test_dense_row_reverse_scales_linearly_not_quadratically() { + // One reverse pass per dense row at any mesh size, with the value tape + // growing ~linearly in state count rather than quadratically. + let sizes = [20usize, 40, 80, 160]; + let mut tape_lens = Vec::new(); + for &n in &sizes { + let (model, _y) = build_dense_row_dae(n); + assert_eq!( + model.jacobian_stats().n_dense_rows, + 1, + "n={n}: reverse pass count must stay 1" + ); + tape_lens.push(model.compiled.jac.max_adjoint_tape_len()); + } + // Doubling n at most doubles a linear tape (b > 0 keeps every ratio < 2); + // a quadratic tape would ~quadruple. `< 3.0` cleanly discriminates. + for pair in tape_lens.windows(2) { + let ratio = pair[1] as f64 / pair[0] as f64; + assert!( + ratio < 3.0, + "adjoint tape grew {ratio:.1}x on doubling n (expected ~linear)" + ); + } + } + + #[test] + fn test_dense_row_merged_remap_tolerates_mass_only_column() { + let n = 20; + let (mut model, y) = build_dense_row_dae_mass_only_column(n); + assert_eq!(model.compiled.jac.n_dense_rows(), 1, "split must be active"); + assert_eq!(model.compiled.jac.dense_rows()[0].rows[0], 0); + + let dense_fd = finite_difference_jacobian(&mut model, 0.0, &y, &[], 1e-6); + + // Loop C with identity mass: J = df/dy - cj*I. Column 0 of row 0 is a + // mass-only merged slot (df/dy has no entry) -> must become -cj. + let cj = 0.5; + model.set_cj(cj); + let mut csc_mass = vec![0.0; model.nnz()]; + model.assemble_jacobian_csc_into_coloring(0.0, &y, &[], &mut csc_mass); + let mut expected = dense_fd.clone(); + for (i, row) in expected.iter_mut().enumerate() { + row[i] -= cj; + } + assert_matrices_match( + &csc_to_dense(&model, &csc_mass, n), + &expected, + 1e-5, + "CSC-mass", + ); + + // Loop D (no mass): the mass-only slot (0,0) stays 0 (pure df/dy). + let mut csc_no_mass = vec![0.0; model.nnz()]; + model.assemble_jacobian_csc_no_mass(0.0, &y, &[], &mut csc_no_mass); + assert_matrices_match( + &csc_to_dense(&model, &csc_no_mass, n), + &dense_fd, + 1e-5, + "CSC-no-mass", + ); + + // Loop B (COO) at cj != 0: the dense row's triplets must fold -cj*M, + // emitting one triplet per merged nnz incl. the mass-only slot (0,0) = -cj. + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y, &[]); + assert_eq!(rows.len(), model.nnz(), "one COO triplet per merged nnz"); + let coo = coo_to_dense(&rows, &cols, &vals, n); + assert!( + (coo[0][0] - (-cj)).abs() < 1e-12, + "mass-only COO slot (0,0) must be -cj, got {}", + coo[0][0] + ); + assert_matrices_match(&coo, &expected, 1e-5, "COO-mass"); + } + + /// Differential dense row whose `df/dy` INCLUDES its own diagonal + /// (`f_0 = sum_j y_j^2` over all states, identity mass), so merged slot + /// (0,0) carries both a df/dy value and a mass entry. Rows `1..n-1` are + /// diagonal `sin(y_i)`. + fn build_dense_row_dae_shared_slot(n: usize) -> (ModelEvaluator, Vec) { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: n }); + let gy = arena.alloc(Node::Mul(y_full, y_full)); + let ones = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, n], + indices: (0..n).collect(), + data: vec![1.0; n], + shape: Shape::matrix(1, n), + }))); + let dense = arena.alloc(Node::MatMul(ones, gy)); + let y_tail = arena.alloc(Node::StateVector { start: 1, end: n }); + let diag = arena.alloc(Node::Sin(y_tail)); + let rhs = arena.alloc(Node::Concat(vec![dense, diag])); + let mass = identity_mass_matrix(n); + let model = ModelEvaluator::new(&arena, rhs, mass, n, 0); + let y: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.2, 0.5).sin() + 0.8) + .collect(); + (model, y) + } + + #[test] + fn test_dense_row_shared_dfdy_mass_slot_survives_postpass() { + let n = 20; + let (mut model, y) = build_dense_row_dae_shared_slot(n); + assert_eq!(model.compiled.jac.n_dense_rows(), 1, "split must be active"); + assert_eq!(model.compiled.jac.dense_rows()[0].rows[0], 0); + + let dense_fd = finite_difference_jacobian(&mut model, 0.0, &y, &[], 1e-6); + + // Loop C at cj != 0: slot (0,0) carries BOTH the dense fill (2*y0) and + // the mass postpass (-cj); a postpass-before-fill reorder clobbers -cj. + let cj = 0.5; + model.set_cj(cj); + let mut csc_mass = vec![0.0; model.nnz()]; + model.assemble_jacobian_csc_into_coloring(0.0, &y, &[], &mut csc_mass); + let assembled = csc_to_dense(&model, &csc_mass, n); + let analytic = 2.0f64.mul_add(y[0], -cj); + assert!( + (assembled[0][0] - analytic).abs() < 1e-12, + "shared slot (0,0) must be 2*y0 - cj = {analytic}, got {}", + assembled[0][0] + ); + let mut expected = dense_fd; + for (i, row) in expected.iter_mut().enumerate() { + row[i] -= cj; + } + assert_matrices_match(&assembled, &expected, 1e-5, "CSC-mass-shared"); + } + + /// Loop-B sparse scan: a differential row whose rhs does not depend on its + /// own state has a mass-only merged slot `(r,r)`. Nothing in the df/dy + /// sweep may reach that slot, which must carry a pure `-cj*M[r,r]` term. + #[test] + fn test_loop_b_sparse_scan_mass_only_slot_not_aliased() { + // f = [sin(y1), sin(y1)]: col 0 is structurally absent from df/dy, and + // identity mass adds slot (0,0). + let mut arena = Arena::new(); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let f0 = arena.alloc(Node::Sin(y1)); + let f1 = arena.alloc(Node::Sin(y1)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1])); + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + assert!( + model.compiled.jac.dense_rows().is_empty(), + "fixture must exercise the sparse scan, not the dense-row path" + ); + // Precondition: col 0 has no df/dy entry, so no sweep ever writes it. + assert_eq!( + model.compiled.jac.coloring().colors[0], + crate::coloring::UNSEEDED, + "fixture requires col 0 to be absent from df/dy" + ); + + let y = [0.3, 0.5]; + let dcos = 0.5_f64.cos(); + + // cj = 0 -> J = df/dy. The mass-only slot (0,0) has no df/dy entry, so it + // must be exactly 0 (RED: old scan emitted df/dy(0,1) = cos(y1)). + model.set_cj(0.0); + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y, &[]); + assert_eq!(rows.len(), model.nnz(), "one triplet per merged nnz"); + let coo = coo_to_dense(&rows, &cols, &vals, 2); + assert!( + coo[0][0].abs() < 1e-12, + "mass-only slot (0,0) must be 0 at cj=0, got {}", + coo[0][0] + ); + assert!( + (coo[0][1] - dcos).abs() < 1e-12, + "df/dy slot (0,1) must be cos(y1), got {}", + coo[0][1] + ); + + // cj != 0 -> slot (0,0) = -cj*M[0,0] = -cj (identity mass). + let cj = 0.5; + model.set_cj(cj); + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y, &[]); + let coo = coo_to_dense(&rows, &cols, &vals, 2); + assert!( + (coo[0][0] - (-cj)).abs() < 1e-12, + "mass-only slot (0,0) must be -cj, got {}", + coo[0][0] + ); + assert!( + (coo[0][1] - dcos).abs() < 1e-12, + "df/dy slot (0,1) must stay cos(y1) at cj!=0, got {}", + coo[0][1] + ); + } + + #[test] + fn test_assemble_jacobian_nonlinear() { + // f(y) = y^2, so df/dy = diag([2*y0, 2*y1]) + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Pow(y, two)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.0); + + // At y = [3, 4], df/dy = diag([6, 8]) + let y_vals = [3.0, 4.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Build dense matrix + let mut dense = [[0.0; 2]; 2]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + assert!( + (dense[0][0] - 6.0).abs() < 1e-12, + "J[0,0] = {}", + dense[0][0] + ); + assert!( + (dense[1][1] - 8.0).abs() < 1e-12, + "J[1,1] = {}", + dense[1][1] + ); + } + + #[test] + fn test_assemble_jacobian_coupled() { + // f(y) = [y0 * y1, y0 + y1] + // df/dy = [[y1, y0], [1, 1]] + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let prod = arena.alloc(Node::Mul(y0, y1)); + let sum = arena.alloc(Node::Add(y0, y1)); + let rhs = arena.alloc(Node::Concat(vec![prod, sum])); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + model.set_cj(0.0); + + // At y = [2, 3] + // df/dy = [[3, 2], [1, 1]] + let y_vals = [2.0, 3.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Should have 4 non-zeros (dense 2x2) + assert_eq!(rows.len(), 4, "Expected 4 non-zeros"); + + // Build dense matrix + let mut dense = [[0.0; 2]; 2]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + assert!( + (dense[0][0] - 3.0).abs() < 1e-12, + "J[0,0] = {}", + dense[0][0] + ); + assert!( + (dense[0][1] - 2.0).abs() < 1e-12, + "J[0,1] = {}", + dense[0][1] + ); + assert!( + (dense[1][0] - 1.0).abs() < 1e-12, + "J[1,0] = {}", + dense[1][0] + ); + assert!( + (dense[1][1] - 1.0).abs() < 1e-12, + "J[1,1] = {}", + dense[1][1] + ); + } + + #[test] + fn test_assemble_jacobian_with_params() { + // f(y) = k * y where k is a parameter + // df/dy = k * I + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(k, y)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 2, 1); + model.set_cj(1.0); // J = k*I - 1*I = (k-1)*I + + let y_vals = [1.0, 2.0]; + let inputs = [3.0]; // k = 3 + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &inputs); + + // Build dense matrix + let mut dense = [[0.0; 2]; 2]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + // J = (3 - 1) * I = 2 * I + assert!( + (dense[0][0] - 2.0).abs() < 1e-12, + "J[0,0] = {}", + dense[0][0] + ); + assert!( + (dense[1][1] - 2.0).abs() < 1e-12, + "J[1,1] = {}", + dense[1][1] + ); + } + + #[test] + fn test_assemble_jacobian_sparse_pattern() { + // f(y) = [y0, y1, y2] (identity) + // df/dy = I (diagonal sparsity) + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let rhs = arena.alloc(Node::Concat(vec![y0, y1, y2])); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + model.set_cj(0.0); + + // Identity is fully constant, so the whole diagonal comes from the table. + assert_eq!(model.coloring().n_colors, 0); + + let y_vals = [1.0, 2.0, 3.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Should have exactly 3 non-zeros (diagonal) + assert_eq!(rows.len(), 3, "Expected 3 non-zeros for diagonal"); + + // All diagonal entries should be 1.0 + for i in 0..3 { + assert_eq!(rows[i], cols[i], "Should be diagonal entry"); + assert!( + (vals[i] - 1.0).abs() < 1e-12, + "Diagonal entry {} = {}", + i, + vals[i] + ); + } + } + + #[test] + fn test_assemble_jacobian_tridiagonal() { + // Build tridiagonal-like structure + // f(y) = [y0+y1, y0+y1+y2, y1+y2] + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + + let f0 = arena.alloc(Node::Add(y0, y1)); + let f1_partial = arena.alloc(Node::Add(y0, y1)); + let f1 = arena.alloc(Node::Add(f1_partial, y2)); + let f2 = arena.alloc(Node::Add(y1, y2)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1, f2])); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + model.set_cj(0.0); + + // Should use <= 3 colors for tridiagonal + assert!( + model.coloring().n_colors <= 3, + "Expected <= 3 colors, got {}", + model.coloring().n_colors + ); + + let y_vals = [1.0, 2.0, 3.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Build dense matrix + let mut dense = [[0.0; 3]; 3]; + for i in 0..rows.len() { + dense[rows[i]][cols[i]] = vals[i]; + } + + assert!((dense[0][0] - 1.0).abs() < 1e-12); + assert!((dense[0][1] - 1.0).abs() < 1e-12); + assert!(dense[0][2].abs() < 1e-12); + assert!((dense[1][0] - 1.0).abs() < 1e-12); + assert!((dense[1][1] - 1.0).abs() < 1e-12); + assert!((dense[1][2] - 1.0).abs() < 1e-12); + assert!(dense[2][0].abs() < 1e-12); + assert!((dense[2][1] - 1.0).abs() < 1e-12); + assert!((dense[2][2] - 1.0).abs() < 1e-12); + } + + #[test] + fn test_assemble_jacobian_consistency_with_jac_mul() { + // J @ e_i == jac_mul(e_i) for each unit vector. + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let y2 = arena.alloc(Node::Index { + child: y_full, + start: 2, + end: 3, + }); + // f(y) = [y0*y1, y1*y2, y0+y2] + let f0 = arena.alloc(Node::Mul(y0, y1)); + let f1 = arena.alloc(Node::Mul(y1, y2)); + let f2 = arena.alloc(Node::Add(y0, y2)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1, f2])); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + model.set_cj(0.5); + + let y_vals = [2.0, 3.0, 4.0]; + let (rows, cols, vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Build dense Jacobian + let mut jac = [[0.0; 3]; 3]; + for i in 0..rows.len() { + jac[rows[i]][cols[i]] = vals[i]; + } + + // Verify J @ e_i == jac_mul(e_i) + for col in 0..3 { + let mut e_i = [0.0, 0.0, 0.0]; + e_i[col] = 1.0; + + let mut jac_mul_result = [0.0, 0.0, 0.0]; + model.jac_mul(0.0, &y_vals, &[], &e_i, &mut jac_mul_result); + + for row in 0..3 { + let expected = jac[row][col]; + let actual = jac_mul_result[row]; + assert!( + (expected - actual).abs() < 1e-12, + "Mismatch at ({row}, {col}): J={expected}, jac_mul={actual}" + ); + } + } + } + + #[test] + fn test_csc_pattern_from_csr() { + // Create a simple CSR pattern for a 3x3 matrix with entries at (0,0), (1,1), (2,1), (2,2) + let csr = SparsityPattern { + nrows: 3, + ncols: 3, + indptr: vec![0, 1, 2, 4], // row 0: 1 entry, row 1: 1 entry, row 2: 2 entries + indices: vec![0, 1, 1, 2], // col indices + }; + + let csc = CscPattern::from_csr(&csr); + + // Verify CSC structure + assert_eq!(csc.nrows, 3); + assert_eq!(csc.ncols, 3); + assert_eq!(csc.nnz(), 4); + + // Column 0 has 1 entry (row 0), column 1 has 2 entries (rows 1, 2), column 2 has 1 entry (row 2) + assert_eq!(csc.colptr, vec![0, 1, 3, 4]); + // Row indices sorted by column + assert_eq!(csc.rowind, vec![0, 1, 2, 2]); + } + + #[test] + fn test_csc_pattern_diagonal() { + // Diagonal matrix: entries at (0,0), (1,1), (2,2) + let csr = SparsityPattern { + nrows: 3, + ncols: 3, + indptr: vec![0, 1, 2, 3], + indices: vec![0, 1, 2], + }; + + let csc = CscPattern::from_csr(&csr); + + assert_eq!(csc.colptr, vec![0, 1, 2, 3]); + assert_eq!(csc.rowind, vec![0, 1, 2]); + } + + #[test] + fn test_assemble_jacobian_csc_into_diagonal() { + // f(y) = 2*y, so df/dy = 2*I (diagonal) + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + + let y_vals = [1.0, 2.0, 3.0]; + let nnz = model.nnz(); + + // Allocate CSC data buffer + let mut jac_data = vec![0.0; nnz]; + + // Assemble Jacobian with cj=0 (no mass term) + model.set_cj(0.0); + model.assemble_jacobian_csc_into(0.0, &y_vals, &[], &mut jac_data); + + // df/dy = 2*I, so all diagonal entries should be 2.0 + for &val in &jac_data { + assert!((val - 2.0).abs() < 1e-14, "Expected 2.0, got {val}"); + } + } + + #[test] + fn test_assemble_jacobian_csc_into_with_cj() { + // f(y) = 2*y, df/dy = 2*I, M = I + // J = df/dy - cj*M = 2*I - 0.5*I = 1.5*I + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + + let y_vals = [1.0, 2.0, 3.0]; + let nnz = model.nnz(); + + let mut jac_data = vec![0.0; nnz]; + + model.set_cj(0.5); + model.assemble_jacobian_csc_into(0.0, &y_vals, &[], &mut jac_data); + + // J = 2*I - 0.5*I = 1.5*I + for &val in &jac_data { + assert!((val - 1.5).abs() < 1e-14, "Expected 1.5, got {val}"); + } + } + + #[test] + fn test_assemble_jacobian_csc_into_matches_coo() { + // Compare CSC-into with the COO-returning version + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + + // f(y) = [y0*y1, y1*y2, y0+y2] - coupled system + let f0 = arena.alloc(Node::Mul(y0, y1)); + let f1 = arena.alloc(Node::Mul(y1, y2)); + let f2 = arena.alloc(Node::Add(y0, y2)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1, f2])); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + model.set_cj(0.5); + + let y_vals = [2.0, 3.0, 4.0]; + + // Get COO result + let (coo_rows, coo_cols, coo_vals) = model.assemble_jacobian(0.0, &y_vals, &[]); + + // Get CSC result + let nnz = model.nnz(); + let mut csc_data = vec![0.0; nnz]; + model.assemble_jacobian_csc_into(0.0, &y_vals, &[], &mut csc_data); + + // Build dense matrices from both and compare + let mut coo_dense = [[0.0; 3]; 3]; + for i in 0..coo_rows.len() { + coo_dense[coo_rows[i]][coo_cols[i]] = coo_vals[i]; + } + + let csc = model.csc_sparsity(); + let mut csc_dense = [[0.0; 3]; 3]; + for (csc_idx, &val) in csc_data.iter().enumerate().take(nnz) { + let (row, col) = csc.csc_to_csr_map[csc_idx]; + csc_dense[row][col] = val; + } + + for row in 0..3 { + for col in 0..3 { + assert!( + (coo_dense[row][col] - csc_dense[row][col]).abs() < 1e-12, + "Mismatch at ({row}, {col}): COO={}, CSC={}", + coo_dense[row][col], + csc_dense[row][col] + ); + } + } + } + + #[test] + fn test_assemble_jacobian_csc_into_zero_allocation() { + // Verify that repeated calls don't allocate (by checking buffer reuse) + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let rhs = arena.alloc(Node::Mul(two, y)); + + let mass = identity_mass_matrix(3); + let mut model = ModelEvaluator::new(&arena, rhs, mass, 3, 0); + + let nnz = model.nnz(); + let mut jac_data = vec![0.0; nnz]; + + // Call multiple times - should not allocate + for i in 0..10_i32 { + let y_vals = [1.0 + f64::from(i), 2.0, 3.0]; + model.assemble_jacobian_csc_into(0.0, &y_vals, &[], &mut jac_data); + // All values should be 2.0 regardless of y (linear function) + for &val in &jac_data { + assert!((val - 2.0).abs() < 1e-14); + } + } + } + + #[test] + fn test_jacobian_stats_report_the_compiled_strategy() { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let f0 = arena.alloc(Node::Mul(x0, x1)); + let f1 = arena.alloc(Node::Add(x0, x1)); + let rhs = arena.alloc(Node::Concat(vec![f0, f1])); + + let mass = identity_mass_matrix(2); + let model = ModelEvaluator::new(&arena, rhs, mass, 2, 0); + + let stats = model.jacobian_stats(); + assert_eq!(stats.strategy, JacobianStrategy::Coloring); + assert_eq!(stats.n_colors, model.coloring().n_colors); + assert_eq!(stats.nnz, model.nnz()); + assert_eq!(stats.n_dense_rows, 0); + assert_eq!(stats.dense_row_entries, 0); + assert_eq!(stats.dense_row_tape_instructions, 0); + } + + #[test] + fn test_n_events_zero_initially() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let model = ModelEvaluator::new(&arena, y, identity_mass_matrix(2), 2, 0); + assert_eq!(model.n_events(), 0); + assert_eq!(model.total_event_len(), 0); + } + + #[test] + fn test_add_event_compiles_and_evaluates() { + // f(y) = y is the rhs; event = y[0] - 0.5 (triggers when y[0] crosses 0.5) + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let threshold = arena.alloc(Node::Scalar(0.5)); + let event_expr = arena.alloc(Node::Sub(y0, threshold)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_event(&arena, event_expr); + + assert_eq!(model.n_events(), 1); + assert_eq!(model.observables(ObservableKind::Events).len_at(0), 1); + assert_eq!(model.total_event_len(), 1); + + // event(y=[0.7, 1.0]) = 0.7 - 0.5 = 0.2 + let mut out = [0.0; 1]; + let written = model.eval_event(0.0, &[0.7, 1.0], &[], 0, &mut out); + assert_eq!(written, 1); + assert!((out[0] - 0.2).abs() < 1e-12); + + // event(y=[0.5, 1.0]) = 0.5 - 0.5 = 0.0 (at threshold) + let written = model.eval_event(0.0, &[0.5, 1.0], &[], 0, &mut out); + assert_eq!(written, 1); + assert!(out[0].abs() < 1e-12); + + // event(y=[0.3, 1.0]) = 0.3 - 0.5 = -0.2 (crossed threshold) + let written = model.eval_event(0.0, &[0.3, 1.0], &[], 0, &mut out); + assert_eq!(written, 1); + assert!((out[0] + 0.2).abs() < 1e-12); + } + + #[test] + fn test_eval_events_multiple_events() { + // Two events: e0 = y[0] - 0.5, e1 = y[1] - 1.0 + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let thresh0 = arena.alloc(Node::Scalar(0.5)); + let thresh1 = arena.alloc(Node::Scalar(1.0)); + let event0 = arena.alloc(Node::Sub(y0, thresh0)); + let event1 = arena.alloc(Node::Sub(y1, thresh1)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_event(&arena, event0); + model.add_event(&arena, event1); + + assert_eq!(model.n_events(), 2); + assert_eq!(model.total_event_len(), 2); + + // y = [0.7, 1.5] => e0 = 0.2, e1 = 0.5 + let mut out = [0.0; 2]; + model.eval_observables(ObservableKind::Events, 0.0, &[0.7, 1.5], &[], &mut out); + assert!((out[0] - 0.2).abs() < 1e-12); + assert!((out[1] - 0.5).abs() < 1e-12); + } + + #[test] + fn test_events_compose_with_outputs() { + // Verify events work alongside outputs on the same model + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let output_expr = arena.alloc(Node::Mul(two, y0)); + let thresh = arena.alloc(Node::Scalar(0.5)); + let event_expr = arena.alloc(Node::Sub(y0, thresh)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_output(&arena, output_expr); + model.add_event(&arena, event_expr); + + assert_eq!(model.n_outputs(), 1); + assert_eq!(model.n_events(), 1); + + // y = [3.0, 4.0] => output = 6.0, event = 2.5 + let mut out_val = [0.0; 1]; + let mut event_val = [0.0; 1]; + model.eval_output(0.0, &[3.0, 4.0], &[], 0, &mut out_val); + model.eval_event(0.0, &[3.0, 4.0], &[], 0, &mut event_val); + + assert!((out_val[0] - 6.0).abs() < 1e-12); + assert!((event_val[0] - 2.5).abs() < 1e-12); + } + + #[test] + fn test_fuse_events_matches_the_per_event_tapes() { + // Two events over a shared subexpression `shared = y[0]*y[1] + y[0]`: + // e0 = shared - 0.5, e1 = shared*2 - 1 + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y_full, + start: 1, + end: 2, + }); + let prod = arena.alloc(Node::Mul(y0, y1)); + let shared = arena.alloc(Node::Add(prod, y0)); + let half = arena.alloc(Node::Scalar(0.5)); + let event0 = arena.alloc(Node::Sub(shared, half)); + let two = arena.alloc(Node::Scalar(2.0)); + let scaled = arena.alloc(Node::Mul(shared, two)); + let one = arena.alloc(Node::Scalar(1.0)); + let event1 = arena.alloc(Node::Sub(scaled, one)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_event(&arena, event0); + model.add_event(&arena, event1); + model.fuse_events(&mut arena, &[event0, event1]); + + // The fused tape is the one the hot loop now uses. + assert!( + model + .observables(ObservableKind::Events) + .fused_expr() + .is_some() + ); + assert_eq!(model.total_event_len(), 2); + + // The fused tape must match the per-event tapes BITWISE: same + // nodes, same instruction semantics. + for &y in &[[0.7, 1.5], [0.3, -2.0], [1.25, 0.0]] { + let mut e0 = [0.0; 1]; + let mut e1 = [0.0; 1]; + model.eval_event(0.0, &y, &[], 0, &mut e0); + model.eval_event(0.0, &y, &[], 1, &mut e1); + + let mut fused = [0.0; 2]; + model.eval_observables(ObservableKind::Events, 0.0, &y, &[], &mut fused); + // Compare raw bits: the fused tape must reproduce the per-event + // values exactly (identical nodes and instruction semantics). + assert_eq!(fused[0].to_bits(), e0[0].to_bits()); + assert_eq!(fused[1].to_bits(), e1[0].to_bits()); + + // `eval_events` (diffsol path) shares the same fused tape. + let mut fused_diffsol = [0.0; 2]; + model.eval_observables(ObservableKind::Events, 0.0, &y, &[], &mut fused_diffsol); + assert_eq!(fused_diffsol.map(f64::to_bits), fused.map(f64::to_bits)); + } + } + + #[test] + fn test_fuse_events_noop_for_single_event() { + // One event: fusion is a no-op and the per-event path is used unchanged. + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let half = arena.alloc(Node::Scalar(0.5)); + let event0 = arena.alloc(Node::Sub(y0, half)); + + let mass = identity_mass_matrix(2); + let mut model = ModelEvaluator::new(&arena, y_full, mass, 2, 0); + model.add_event(&arena, event0); + model.fuse_events(&mut arena, &[event0]); + + assert!( + model + .observables(ObservableKind::Events) + .fused_expr() + .is_none() + ); + let mut out = [0.0; 1]; + model.eval_observables(ObservableKind::Events, 0.0, &[0.7, 1.0], &[], &mut out); + assert!((out[0] - 0.2).abs() < 1e-12); + } + + #[test] + fn mass_kind_identity() { + let mass = CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![1.0, 1.0], + shape: Shape::matrix(2, 2), + }; + assert!(matches!(classify_mass_matrix(&mass), MassKind::Identity)); + } + + #[test] + fn mass_kind_diagonal_selector() { + let mass = CsrData { + indptr: vec![0, 1, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(2, 2), + }; + match classify_mass_matrix(&mass) { + MassKind::DiagonalSelector(mask) => assert_eq!(mask, vec![true, false]), + other => panic!("Expected DiagonalSelector, got {other:?}"), + } + } + + #[test] + fn mass_kind_general_offdiag() { + let mass = CsrData { + indptr: vec![0, 2, 3], + indices: vec![0, 1, 1], + data: vec![1.0, 0.5, 1.0], + shape: Shape::matrix(2, 2), + }; + assert!(matches!(classify_mass_matrix(&mass), MassKind::General)); + } + + /// Build a `CompiledModel` with one output and one sensitivity parameter. + /// + /// Model: f(y) = k * y (RHS), H(y, p) = k * y[0]^2 (output). + /// `n_states` = 2, `n_params` = 1, sens on param 0. + fn build_test_model_with_output_and_sens() -> CompiledModel { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + // RHS: k * y (2-vector) + let rhs = arena.alloc(Node::Mul(k, y_full)); + // Output: k * y[0]^2 (scalar) + let y0_sq = arena.alloc(Node::Pow(y0, two)); + let out_node = arena.alloc(Node::Mul(k, y0_sq)); + + let mass = identity_mass_matrix(2); + let mut data = CompiledModel::new_with_sens(&arena, rhs, mass, 2, 1, &[0]); + data.add_output(&arena, out_node); + data + } + + /// State vector for `build_test_model_with_output_and_sens`. + fn test_state(model: &CompiledModel) -> Vec { + vec![3.0; model.n_states] + } + + /// Input vector for `build_test_model_with_output_and_sens`. + fn test_inputs(_model: &CompiledModel) -> Vec { + vec![2.0] // k = 2 + } + + #[test] + fn an_output_sens_action_matches_finite_difference() { + // Build a compiled model with >=1 output and >=1 sensitivity parameter. + let model = build_test_model_with_output_and_sens(); + let mut ws = model.create_workspace(); + let t = 0.3; + let y = test_state(&model); + let inputs = test_inputs(&model); + + // Analytic dH/dp . e_k for the first sensitivity parameter. + let mut v = vec![0.0; model.n_sens_params()]; + v[0] = 1.0; + let mut got = vec![0.0; model.total_output_len()]; + model.observable_sens_action( + &mut ws, + ObservableKind::Outputs, + t, + &y, + &inputs, + model.sens_param_indices(), + &v, + &mut got, + ); + + // Finite-difference reference: perturb that parameter, re-eval outputs. + let h = 1e-6; + let pidx = model.sens_param_indices()[0]; + let mut ip = inputs.clone(); + let mut o_plus = vec![0.0; model.total_output_len()]; + let mut o_minus = vec![0.0; model.total_output_len()]; + ip[pidx] += h; + model.eval_observables(&mut ws, ObservableKind::Outputs, t, &y, &ip, &mut o_plus); + ip[pidx] -= 2.0 * h; + model.eval_observables(&mut ws, ObservableKind::Outputs, t, &y, &ip, &mut o_minus); + for j in 0..got.len() { + let fd = (o_plus[j] - o_minus[j]) / (2.0 * h); + assert!( + (got[j] - fd).abs() < 1e-5, + "out[{j}]: got {} fd {}", + got[j], + fd + ); + } + } + + #[test] + fn an_output_jac_action_matches_finite_difference() { + let model = build_test_model_with_output_and_sens(); + let mut ws = model.create_workspace(); + let (t, y, inputs) = (0.3, test_state(&model), test_inputs(&model)); + let mut v = vec![0.0; model.n_states]; + v[0] = 1.0; // dH/dy . e_0 + let mut got = vec![0.0; model.total_output_len()]; + model.observable_jac_action( + &mut ws, + ObservableKind::Outputs, + t, + &y, + &inputs, + &v, + &mut got, + ); + let h = 1e-6; + let (mut yp, mut ym) = (y.clone(), y.clone()); + yp[0] += h; + ym[0] -= h; + let mut op = vec![0.0; model.total_output_len()]; + let mut om = vec![0.0; model.total_output_len()]; + model.eval_observables(&mut ws, ObservableKind::Outputs, t, &yp, &inputs, &mut op); + model.eval_observables(&mut ws, ObservableKind::Outputs, t, &ym, &inputs, &mut om); + for j in 0..got.len() { + let fd = (op[j] - om[j]) / (2.0 * h); + assert!( + (got[j] - fd).abs() < 1e-5, + "jac[{j}]: got {} fd {}", + got[j], + fd + ); + } + } + + #[test] + fn output_sens_project_equals_action_sum() { + // projection[k] must equal dH/dp . e_k + dH/dy . y_sens_k + let model = build_test_model_with_output_and_sens(); + let mut ws = model.create_workspace(); + let (t, y, inputs) = (0.3, test_state(&model), test_inputs(&model)); + let n_s = model.n_sens_params(); + let n_states = model.n_states; + let n_out = model.total_output_len(); + + let mut y_sens = vec![0.0; n_s * n_states]; + for (i, v) in y_sens.iter_mut().enumerate() { + *v = 0.1 * (i as f64 + 1.0); + } + let mut got = vec![0.0; n_s * n_out]; + model.output_sens_project(&mut ws, t, &y, &inputs, &y_sens, &mut got); + + for k in 0..n_s { + let mut e_k = vec![0.0; n_s]; + e_k[k] = 1.0; + let mut sens_term = vec![0.0; n_out]; + model.observable_sens_action( + &mut ws, + ObservableKind::Outputs, + t, + &y, + &inputs, + model.sens_param_indices(), + &e_k, + &mut sens_term, + ); + let mut jac_term = vec![0.0; n_out]; + model.observable_jac_action( + &mut ws, + ObservableKind::Outputs, + t, + &y, + &inputs, + &y_sens[k * n_states..(k + 1) * n_states], + &mut jac_term, + ); + for o in 0..n_out { + let expected = sens_term[o] + jac_term[o]; + assert!( + (got[k * n_out + o] - expected).abs() < 1e-12, + "k={k} o={o}: got {} want {}", + got[k * n_out + o], + expected + ); + } + } + } + + /// Build a `CompiledModel` with one event and one sensitivity parameter. + /// + /// Model: f(y) = k * y (RHS), g(y, p) = k * y[0]^2 - 1.0 (event). + /// `n_states` = 2, `n_params` = 1, sens on param 0. + fn build_test_model_with_event_and_sens() -> CompiledModel { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y_full, + start: 0, + end: 1, + }); + let k = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let one = arena.alloc(Node::Scalar(1.0)); + // RHS: k * y (2-vector) + let rhs = arena.alloc(Node::Mul(k, y_full)); + // Event: k * y[0]^2 - 1.0 (scalar) + let y0_sq = arena.alloc(Node::Pow(y0, two)); + let ky0_sq = arena.alloc(Node::Mul(k, y0_sq)); + let event_node = arena.alloc(Node::Sub(ky0_sq, one)); + + let mass = identity_mass_matrix(2); + let mut data = CompiledModel::new_with_sens(&arena, rhs, mass, 2, 1, &[0]); + data.add_event(&arena, event_node); + data + } + + #[test] + fn event_sens_and_jac_actions_match_finite_difference() { + let model = build_test_model_with_event_and_sens(); + let mut ws = model.create_workspace(); + let (t, y, inputs) = (0.3, test_state(&model), test_inputs(&model)); + + // dg/dp . e_0 + let mut vp = vec![0.0; model.n_sens_params()]; + vp[0] = 1.0; + let mut got_p = vec![0.0; model.total_event_len()]; + model.observable_sens_action( + &mut ws, + ObservableKind::Events, + t, + &y, + &inputs, + model.sens_param_indices(), + &vp, + &mut got_p, + ); + let h = 1e-6; + let pidx = model.sens_param_indices()[0]; + let (mut ip_p, mut ip_m) = (inputs.clone(), inputs.clone()); + ip_p[pidx] += h; + ip_m[pidx] -= h; + let mut gp = vec![0.0; model.total_event_len()]; + let mut gm = vec![0.0; model.total_event_len()]; + model.eval_observables(&mut ws, ObservableKind::Events, t, &y, &ip_p, &mut gp); + model.eval_observables(&mut ws, ObservableKind::Events, t, &y, &ip_m, &mut gm); + for j in 0..got_p.len() { + let fd = (gp[j] - gm[j]) / (2.0 * h); + assert!((got_p[j] - fd).abs() < 1e-5); + } + + // dg/dy . e_0 + let mut vy = vec![0.0; model.n_states]; + vy[0] = 1.0; + let mut got_y = vec![0.0; model.total_event_len()]; + model.observable_jac_action( + &mut ws, + ObservableKind::Events, + t, + &y, + &inputs, + &vy, + &mut got_y, + ); + let (mut yp, mut ym) = (y.clone(), y.clone()); + yp[0] += h; + ym[0] -= h; + let mut gyp = vec![0.0; model.total_event_len()]; + let mut gym = vec![0.0; model.total_event_len()]; + model.eval_observables(&mut ws, ObservableKind::Events, t, &yp, &inputs, &mut gyp); + model.eval_observables(&mut ws, ObservableKind::Events, t, &ym, &inputs, &mut gym); + for j in 0..got_y.len() { + let fd = (gyp[j] - gym[j]) / (2.0 * h); + assert!((got_y[j] - fd).abs() < 1e-5); + } + } + + #[test] + #[should_panic(expected = "sens_param_indices contains a repeated index")] + fn duplicate_sens_param_indices_are_rejected() { + // A repeated index would turn the subset scatter into an accumulate and + // break the set_params/get_params round trip. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let rhs = arena.alloc(Node::Mul(a, y)); + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(1, 1), + }; + ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 1, + 1, + CompiledModelOptions::new().with_sensitivities(&[0, 0]), + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/node.rs b/packages/pybamm-rust/pybamm-core/src/node.rs new file mode 100644 index 0000000000..bf3722549f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/node.rs @@ -0,0 +1,1150 @@ +//! The expression vocabulary Python hands over. +//! +//! [`Node`] is one DAG node: an operator referencing its children by +//! [`NodeId`], a leaf reading state or parameters, or a literal +//! carrying its own data (dense arrays, CSR matrices, 1-D and N-D interpolant +//! tables). Everything downstream matches on this enum, from simplification +//! through differentiation to lowering, so support for a new `PyBaMM` operator +//! starts with a variant here. +//! +//! Nodes are shape-carrying but not shape-checked on construction; +//! [`first_invalid`](crate::first_invalid) and +//! [`first_unsupported`](crate::first_unsupported) report what lowering will +//! reject. + +use crate::arena::NodeId; +use crate::error::CoreError; + +/// Declared `rows × cols` of a literal. +/// +/// Carried for matmul dimension checks and for reporting. Elsewhere evaluation +/// works on the flat element count, so a `1 × n` and an `n × 1` array are alike. +/// As a [`MatMul`](Node::MatMul) left operand they are not, since `DenseMatMul` +/// reads `rows`/`cols` at evaluation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct Shape { + pub rows: usize, + pub cols: usize, +} + +/// Knots and values of a 1-D linear interpolant. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct InterpolantData { + pub(crate) x_data: Vec, + pub(crate) y_data: Vec, +} + +/// Breakpoints and per-interval coefficients of a 1-D cubic interpolant. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct CubicInterpolantData { + /// Interval breakpoints, length nseg + 1. + pub(crate) breakpoints: Vec, + /// Per-interval power-basis coeffs `[c0, c1, c2, c3]`, length nseg. + /// `p(dx) = c0 + c1*dx + c2*dx^2 + c3*dx^3`, `dx = x - breakpoints[i]`. + pub(crate) coeffs: Vec<[f64; 4]>, +} + +/// Per-axis breakpoints and per-cell coefficients of a 2-D or 3-D +/// tensor-product interpolant. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct NdInterpolantData { + /// Per-axis knot vectors (2 or 3 axes), each of length `nseg_a + 1`. + pub(crate) breakpoints: Vec>, + /// Flat per-cell power-basis tensors: cell-major (axis-0 segment + /// slowest), `order^ndim` coeffs per cell (axis-0 power slowest, + /// ascending powers): `p(dx) = Σ c[a0,..] · Π dx_i^a_i`. + pub(crate) coeffs: Vec, + /// Per-axis polynomial order: 2 = multilinear, 4 = tensor cubic. + pub(crate) order: u32, +} + +impl Shape { + /// `1 × 1`. + pub const fn scalar() -> Self { + Self { rows: 1, cols: 1 } + } + + /// `len × 1`, the column-vector orientation `PyBaMM` discretises into. + pub const fn vector(len: usize) -> Self { + Self { rows: len, cols: 1 } + } + + pub const fn matrix(rows: usize, cols: usize) -> Self { + Self { rows, cols } + } +} + +/// A constant sparse matrix, in the CSR form `PyBaMM`'s discretisation produces. +/// +/// Spatial operators arrive as matrices, so most of a discretised model's constant +/// data is one of these. CSR is kept as given and converted at the boundaries that +/// need column-major, namely the solver's Jacobian and mass sparsity. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct CsrData { + pub(crate) indptr: Vec, + pub(crate) indices: Vec, + pub(crate) data: Vec, + pub(crate) shape: Shape, +} + +/// A constant dense array, stored row-major and flat. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct ArrayData { + pub(crate) data: Vec, + pub(crate) shape: Shape, +} + +impl CsrData { + /// Build a CSR matrix, validating its structural invariants: `indptr` + /// length (`rows + 1`), leading zero, non-decreasing offsets whose tail + /// equals the nnz, matching `indices`/`data` lengths, and in-range column + /// indices. Internal transform passes that already hold the invariant + /// still construct via the (crate-visible) struct literal. + pub fn try_new( + indptr: Vec, + indices: Vec, + data: Vec, + shape: Shape, + ) -> Result { + let (rows, cols) = (shape.rows, shape.cols); + if indptr.len() != rows + 1 { + return Err(CoreError::Csr(format!( + "indptr length {} must equal rows + 1 = {}", + indptr.len(), + rows + 1 + ))); + } + if indptr[0] != 0 { + return Err(CoreError::Csr(format!( + "indptr[0] must be 0, got {}", + indptr[0] + ))); + } + if indptr.windows(2).any(|w| w[1] < w[0]) { + return Err(CoreError::Csr("indptr must be non-decreasing".to_string())); + } + if indices.len() != data.len() { + return Err(CoreError::Csr(format!( + "indices length {} must equal data length {}", + indices.len(), + data.len() + ))); + } + let nnz = indptr[rows]; // in bounds: len == rows + 1 + if nnz != data.len() { + return Err(CoreError::Csr(format!( + "indptr tail {nnz} must equal nnz {}", + data.len() + ))); + } + if let Some(&max_col) = indices.iter().max() + && max_col >= cols + { + return Err(CoreError::Csr(format!( + "column index {max_col} out of range for {cols} columns" + ))); + } + Ok(Self { + indptr, + indices, + data, + shape, + }) + } + + /// Row-offset array, length `shape.rows + 1`. + pub fn indptr(&self) -> &[usize] { + &self.indptr + } + + /// Column index for each stored entry. + pub fn indices(&self) -> &[usize] { + &self.indices + } + + /// Stored values, parallel to [`indices`](Self::indices). + pub fn data(&self) -> &[f64] { + &self.data + } + + /// Matrix shape. + pub const fn shape(&self) -> Shape { + self.shape + } +} + +impl ArrayData { + /// Build a dense array, validating that `data.len()` equals `rows * cols`. + pub fn try_new(data: Vec, shape: Shape) -> Result { + let expected = shape.rows.checked_mul(shape.cols).ok_or_else(|| { + CoreError::Array(format!( + "shape {}x{} overflows usize", + shape.rows, shape.cols + )) + })?; + if data.len() != expected { + return Err(CoreError::Array(format!( + "data length {} must equal rows*cols = {expected}", + data.len() + ))); + } + Ok(Self { data, shape }) + } + + /// Row-major dense values. + pub fn data(&self) -> &[f64] { + &self.data + } + + /// Array shape. + pub const fn shape(&self) -> Shape { + self.shape + } +} + +impl InterpolantData { + /// Build a 1D linear interpolation table, validating that the grid is + /// non-empty, `x`/`y` lengths match, and the knots are finite and strictly + /// increasing (the segment lookup and slope division rely on both). + pub fn try_new(x_data: Vec, y_data: Vec) -> Result { + if x_data.is_empty() { + return Err(CoreError::Interpolant( + "x_data must be non-empty".to_string(), + )); + } + if x_data.len() != y_data.len() { + return Err(CoreError::Interpolant(format!( + "x_data length {} must equal y_data length {}", + x_data.len(), + y_data.len() + ))); + } + if !x_data.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant("x_data must be finite".to_string())); + } + if x_data.windows(2).any(|w| w[1] <= w[0]) { + return Err(CoreError::Interpolant( + "x_data must be strictly increasing".to_string(), + )); + } + Ok(Self { x_data, y_data }) + } +} + +impl CubicInterpolantData { + /// Build a 1D cubic interpolation table, validating at least two finite, + /// strictly increasing breakpoints and one coefficient tuple per segment. + pub fn try_new(breakpoints: Vec, coeffs: Vec<[f64; 4]>) -> Result { + if breakpoints.len() < 2 { + return Err(CoreError::Interpolant(format!( + "cubic interpolant needs at least 2 breakpoints, got {}", + breakpoints.len() + ))); + } + let nseg = breakpoints.len() - 1; + if coeffs.len() != nseg { + return Err(CoreError::Interpolant(format!( + "cubic interpolant needs {nseg} coefficient tuples (breakpoints - 1), got {}", + coeffs.len() + ))); + } + if !breakpoints.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant( + "cubic interpolant breakpoints must be finite".to_string(), + )); + } + if breakpoints.windows(2).any(|w| w[1] <= w[0]) { + return Err(CoreError::Interpolant( + "cubic interpolant breakpoints must be strictly increasing".to_string(), + )); + } + Ok(Self { + breakpoints, + coeffs, + }) + } +} + +impl NdInterpolantData { + /// Build an N-D (2 or 3 axis) tensor interpolation table, validating the + /// axis count, order (2 or 4), per-axis finite strictly-increasing knots + /// (at least two each), and the `ncells * order^ndim` coefficient count. + pub fn try_new( + breakpoints: Vec>, + coeffs: Vec, + order: u32, + ) -> Result { + let ndim = breakpoints.len(); + if !(2..=3).contains(&ndim) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant supports 2 or 3 axes, got {ndim}" + ))); + } + if order != 2 && order != 4 { + return Err(CoreError::Interpolant(format!( + "N-D interpolant order must be 2 or 4, got {order}" + ))); + } + let mut ncells = 1usize; + for (axis, knots) in breakpoints.iter().enumerate() { + if knots.len() < 2 { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} needs at least 2 breakpoints, got {}", + knots.len() + ))); + } + if !knots.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} breakpoints must be finite" + ))); + } + if knots.windows(2).any(|w| w[1] <= w[0]) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} breakpoints must be strictly increasing" + ))); + } + ncells *= knots.len() - 1; + } + let order_usize = order as usize; + let per_cell = match ndim { + 2 => order_usize * order_usize, + 3 => order_usize * order_usize * order_usize, + _ => unreachable!("ndim validated to 2..=3"), + }; + let expected = ncells * per_cell; + if coeffs.len() != expected { + return Err(CoreError::Interpolant(format!( + "N-D interpolant needs {expected} coefficients (ncells * order^ndim), got {}", + coeffs.len() + ))); + } + Ok(Self { + breakpoints, + coeffs, + order, + }) + } +} + +/// One node of the expression DAG. +/// +/// Children are [`NodeId`]s into the arena that owns this node, so a node is +/// meaningless on its own and cheap to share. Values are `f64` vectors: the +/// arithmetic and unary variants are element-wise with a scalar operand +/// broadcasting against a vector, matching [`BinaryOp`](crate::BinaryOp) and +/// [`UnaryOp`](crate::UnaryOp), which they lower to one-for-one. +/// +/// Variants marked internal are produced by differentiation rather than by +/// Python, and a few carry `Box`ed payloads to keep the enum small enough that a +/// DAG of them stays compact. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub enum Node { + // Leaf nodes + Scalar(f64), + Array(Box), + /// Zero vector of specified length (first-class, not dense Array) + ZeroVector { + len: usize, + }, + SparseMatrix(Box), + /// Reads `y[start..end]`, a half-open range in the solver's global state + /// vector rather than any per-equation numbering. + StateVector { + start: usize, + end: usize, + }, + /// Reads `y'[start..end]` from the state derivative, in the same index space + /// as [`StateVector`](Self::StateVector). Only a residual formulation + /// supplies it. + StateVectorDot { + start: usize, + end: usize, + }, + InputParameter { + name: String, + /// Registration order among distinct names (0-based); used by + /// `TangentParameter`/sensitivity indexing, unaffected by width. + index: usize, + /// Cumulative offset into the packed `p` values array. + offset: usize, + /// Number of packed values this parameter occupies (>1 for vector inputs). + width: usize, + }, + Time, + + // Binary operations + Add(NodeId, NodeId), + Sub(NodeId, NodeId), + Mul(NodeId, NodeId), + Div(NodeId, NodeId), + Pow(NodeId, NodeId), + /// Matrix-vector product. The left child must be a constant + /// [`SparseMatrix`](Self::SparseMatrix) or [`Array`](Self::Array), since there + /// is no general matrix-matrix product, and its column count must equal the + /// right child's width. + MatMul(NodeId, NodeId), + Minimum(NodeId, NodeId), + Maximum(NodeId, NodeId), + Modulo(NodeId, NodeId), + Hypot(NodeId, NodeId), + EqualHeaviside(NodeId, NodeId), + NotEqualHeaviside(NodeId, NodeId), + Equality(NodeId, NodeId), + + // Structural nodes + /// Half-open slice `child[start..end]` of the child's own value, unrelated to + /// state-vector indices. + Index { + child: NodeId, + start: usize, + end: usize, + }, + /// Children joined end to end; the order here is the layout of the result, and + /// for a model's right-hand side it is the equation ordering the solver sees. + Concat(Vec), + + // Interpolation (boxed to reduce enum size) + Interpolant1DLinear { + data: Box, + child: NodeId, + }, + /// 1D cubic/pchip interpolation (piecewise cubic, power-basis coeffs). + Interpolant1DCubic { + data: Box, + child: NodeId, + }, + /// N-D (2D/3D) tensor-product interpolation: one child per axis, + /// evaluated element-wise over equal-length children. + InterpolantNd { + data: Box, + children: Vec, + }, + + // Unary operations + Neg(NodeId), + Abs(NodeId), + Sqrt(NodeId), + Exp(NodeId), + Log(NodeId), + Sin(NodeId), + Cos(NodeId), + Tanh(NodeId), + Sinh(NodeId), + Cosh(NodeId), + Arcsinh(NodeId), + Arctan(NodeId), + Erf(NodeId), + Sign(NodeId), + Floor(NodeId), + Ceiling(NodeId), + MaxReduce(NodeId), + MinReduce(NodeId), + + /// Internal: `basis[k]` where `k` is the first-occurrence argmax + /// (`is_max = true`) or argmin of `picker`. Created by `differentiate` + /// as the subgradient of `MaxReduce`/`MinReduce`; never built from Python. + ReduceArgSelect { + basis: NodeId, + picker: NodeId, + is_max: bool, + }, + + // Conditional branching + /// Selects `branches[i]` when `selector` falls in the open window + /// `(i + 0.5, i + 1.5)`, so 2.0 and 1.6 both pick `branches[1]` but an exact + /// half-way value matches nothing and evaluates to zeros. Branches are ordinary + /// subexpressions, and the compiler arranges for only the selected one to be + /// evaluated where it can prove ownership. + Conditional { + selector: NodeId, + branches: Vec, + }, + + // Tangent nodes for forward-mode AD + /// Reads a slice of the tangent state vector (corresponds to `LoadTangentState` in `TypedIr`) + TangentStateVector { + start: usize, + end: usize, + }, + /// Reads a single tangent parameter (corresponds to `LoadTangentParameter` in `TypedIr`) + TangentParameter { + index: usize, + }, + /// Derivative of linear interpolation (used during symbolic differentiation) + Interpolant1DLinearDeriv { + slopes: Box<[f64]>, + x_data: Box<[f64]>, + child: NodeId, + }, + /// Derivative of 1D cubic/pchip interpolation (used during AD). + Interpolant1DCubicDeriv { + data: Box, + child: NodeId, + }, + /// Partial derivative of N-D interpolation along `axis` (used during AD). + InterpolantNdPartial { + data: Box, + children: Vec, + axis: u32, + }, +} + +impl Node { + /// Visit each child `NodeId` of this node, allocation-free. + pub fn for_each_child(&self, mut f: F) { + match self { + // Leaves have no children + Self::Scalar(_) + | Self::Array(_) + | Self::ZeroVector { .. } + | Self::SparseMatrix(_) + | Self::StateVector { .. } + | Self::StateVectorDot { .. } + | Self::InputParameter { .. } + | Self::Time + | Self::TangentStateVector { .. } + | Self::TangentParameter { .. } => {}, + + // Binary operations + Self::Add(l, r) + | Self::Sub(l, r) + | Self::Mul(l, r) + | Self::Div(l, r) + | Self::Pow(l, r) + | Self::MatMul(l, r) + | Self::Minimum(l, r) + | Self::Maximum(l, r) + | Self::Modulo(l, r) + | Self::Hypot(l, r) + | Self::EqualHeaviside(l, r) + | Self::NotEqualHeaviside(l, r) + | Self::Equality(l, r) => { + f(*l); + f(*r); + }, + + // Unary operations + Self::Neg(c) + | Self::Abs(c) + | Self::Sqrt(c) + | Self::Exp(c) + | Self::Log(c) + | Self::Sin(c) + | Self::Cos(c) + | Self::Tanh(c) + | Self::Sinh(c) + | Self::Cosh(c) + | Self::Arcsinh(c) + | Self::Arctan(c) + | Self::Erf(c) + | Self::Sign(c) + | Self::Floor(c) + | Self::Ceiling(c) + | Self::MaxReduce(c) + | Self::MinReduce(c) => { + f(*c); + }, + + // Structural / Interpolation with single child + Self::Index { child, .. } + | Self::Interpolant1DLinear { child, .. } + | Self::Interpolant1DLinearDeriv { child, .. } + | Self::Interpolant1DCubic { child, .. } + | Self::Interpolant1DCubicDeriv { child, .. } => { + f(*child); + }, + Self::Concat(children) + | Self::InterpolantNd { children, .. } + | Self::InterpolantNdPartial { children, .. } => { + for &c in children { + f(c); + } + }, + + Self::Conditional { selector, branches } => { + f(*selector); + for &b in branches { + f(b); + } + }, + + Self::ReduceArgSelect { basis, picker, .. } => { + f(*basis); + f(*picker); + }, + } + } + + /// Returns a new `Node` with a transformation applied to each child `NodeId`. + /// + /// Leaves (scalars, arrays, time, etc.) are cloned unchanged. For nodes with + /// children (binary ops, unary ops, structural, interpolants, conditionals), + /// the closure `f` is called on every child `NodeId`, and a new node is + /// constructed with the transformed children. + /// + /// Common uses: + /// - Renumbering node IDs after deduplication + /// - Substituting nodes (e.g. replacing `InputParameter` with constants) + /// - Pruning unused branches (replacing with a zero node) + #[must_use] + pub fn map_children NodeId>(&self, mut f: F) -> Self { + match self { + // Leaves: clone unchanged. + Self::Scalar(_) + | Self::Array(_) + | Self::ZeroVector { .. } + | Self::SparseMatrix(_) + | Self::StateVector { .. } + | Self::StateVectorDot { .. } + | Self::InputParameter { .. } + | Self::Time + | Self::TangentStateVector { .. } + | Self::TangentParameter { .. } => self.clone(), + + // Binary operations + Self::Add(l, r) => Self::Add(f(*l), f(*r)), + Self::Sub(l, r) => Self::Sub(f(*l), f(*r)), + Self::Mul(l, r) => Self::Mul(f(*l), f(*r)), + Self::Div(l, r) => Self::Div(f(*l), f(*r)), + Self::Pow(l, r) => Self::Pow(f(*l), f(*r)), + Self::MatMul(l, r) => Self::MatMul(f(*l), f(*r)), + Self::Minimum(l, r) => Self::Minimum(f(*l), f(*r)), + Self::Maximum(l, r) => Self::Maximum(f(*l), f(*r)), + Self::Modulo(l, r) => Self::Modulo(f(*l), f(*r)), + Self::Hypot(l, r) => Self::Hypot(f(*l), f(*r)), + Self::EqualHeaviside(l, r) => Self::EqualHeaviside(f(*l), f(*r)), + Self::NotEqualHeaviside(l, r) => Self::NotEqualHeaviside(f(*l), f(*r)), + Self::Equality(l, r) => Self::Equality(f(*l), f(*r)), + + // Unary operations + Self::Neg(c) => Self::Neg(f(*c)), + Self::Abs(c) => Self::Abs(f(*c)), + Self::Sqrt(c) => Self::Sqrt(f(*c)), + Self::Exp(c) => Self::Exp(f(*c)), + Self::Log(c) => Self::Log(f(*c)), + Self::Sin(c) => Self::Sin(f(*c)), + Self::Cos(c) => Self::Cos(f(*c)), + Self::Tanh(c) => Self::Tanh(f(*c)), + Self::Sinh(c) => Self::Sinh(f(*c)), + Self::Cosh(c) => Self::Cosh(f(*c)), + Self::Arcsinh(c) => Self::Arcsinh(f(*c)), + Self::Arctan(c) => Self::Arctan(f(*c)), + Self::Erf(c) => Self::Erf(f(*c)), + Self::Sign(c) => Self::Sign(f(*c)), + Self::Floor(c) => Self::Floor(f(*c)), + Self::Ceiling(c) => Self::Ceiling(f(*c)), + Self::MaxReduce(c) => Self::MaxReduce(f(*c)), + Self::MinReduce(c) => Self::MinReduce(f(*c)), + + // Structural nodes + Self::Index { child, start, end } => Self::Index { + child: f(*child), + start: *start, + end: *end, + }, + Self::Concat(children) => Self::Concat(children.iter().map(|&c| f(c)).collect()), + + // Interpolation nodes + Self::Interpolant1DLinear { data, child } => Self::Interpolant1DLinear { + data: data.clone(), + child: f(*child), + }, + Self::Interpolant1DLinearDeriv { + slopes, + x_data, + child, + } => Self::Interpolant1DLinearDeriv { + slopes: slopes.clone(), + x_data: x_data.clone(), + child: f(*child), + }, + Self::Interpolant1DCubic { data, child } => Self::Interpolant1DCubic { + data: data.clone(), + child: f(*child), + }, + Self::Interpolant1DCubicDeriv { data, child } => Self::Interpolant1DCubicDeriv { + data: data.clone(), + child: f(*child), + }, + Self::InterpolantNd { data, children } => Self::InterpolantNd { + data: data.clone(), + children: children.iter().map(|&c| f(c)).collect(), + }, + Self::InterpolantNdPartial { + data, + children, + axis, + } => Self::InterpolantNdPartial { + data: data.clone(), + children: children.iter().map(|&c| f(c)).collect(), + axis: *axis, + }, + + Self::Conditional { selector, branches } => Self::Conditional { + selector: f(*selector), + branches: branches.iter().map(|&b| f(b)).collect(), + }, + + Self::ReduceArgSelect { + basis, + picker, + is_max, + } => Self::ReduceArgSelect { + basis: f(*basis), + picker: f(*picker), + is_max: *is_max, + }, + } + } +} + +fn hash_nd_interpolant_data(data: &NdInterpolantData, hasher: &mut H) { + use std::hash::Hash; + data.order.hash(hasher); + for knots in &data.breakpoints { + knots.len().hash(hasher); + for v in knots { + v.to_bits().hash(hasher); + } + } + for v in &data.coeffs { + v.to_bits().hash(hasher); + } +} + +/// Hash a node by structure, for the CSE key that finds duplicate subexpressions. +/// +/// Children are hashed through `child_remap`, so two nodes collide only when their +/// operands are already known to be the same value, which is what lets CSE work +/// bottom-up in one pass. Floats are hashed by bit pattern, which keeps `0.0` and +/// `-0.0` distinct rather than merging expressions that differ in sign of zero. +pub fn structural_hash( + node: &Node, + hasher: &mut H, + mut child_remap: impl FnMut(NodeId) -> NodeId, +) { + use std::hash::Hash; + std::mem::discriminant(node).hash(hasher); + match node { + Node::Scalar(v) => v.to_bits().hash(hasher), + Node::Array(arr) => { + for v in &arr.data { + v.to_bits().hash(hasher); + } + arr.shape.rows.hash(hasher); + arr.shape.cols.hash(hasher); + }, + Node::ZeroVector { len } => len.hash(hasher), + Node::SparseMatrix(csr) => { + csr.indptr.hash(hasher); + csr.indices.hash(hasher); + for v in &csr.data { + v.to_bits().hash(hasher); + } + csr.shape.rows.hash(hasher); + csr.shape.cols.hash(hasher); + }, + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } => { + start.hash(hasher); + end.hash(hasher); + }, + Node::InputParameter { + name, + index, + offset, + width, + } => { + name.hash(hasher); + index.hash(hasher); + offset.hash(hasher); + width.hash(hasher); + }, + Node::TangentParameter { index } => index.hash(hasher), + Node::Concat(children) => children.len().hash(hasher), + Node::Conditional { branches, .. } => branches.len().hash(hasher), + Node::Interpolant1DLinear { data, .. } => { + for v in &data.x_data { + v.to_bits().hash(hasher); + } + for v in &data.y_data { + v.to_bits().hash(hasher); + } + }, + Node::Interpolant1DLinearDeriv { slopes, x_data, .. } => { + for v in slopes { + v.to_bits().hash(hasher); + } + for v in x_data { + v.to_bits().hash(hasher); + } + }, + Node::Interpolant1DCubic { data, .. } | Node::Interpolant1DCubicDeriv { data, .. } => { + for v in &data.breakpoints { + v.to_bits().hash(hasher); + } + for c in &data.coeffs { + for v in c { + v.to_bits().hash(hasher); + } + } + }, + Node::InterpolantNd { data, .. } => hash_nd_interpolant_data(data, hasher), + Node::InterpolantNdPartial { data, axis, .. } => { + hash_nd_interpolant_data(data, hasher); + axis.hash(hasher); + }, + Node::Time + | Node::Add(_, _) + | Node::Sub(_, _) + | Node::Mul(_, _) + | Node::Div(_, _) + | Node::Pow(_, _) + | Node::MatMul(_, _) + | Node::Minimum(_, _) + | Node::Maximum(_, _) + | Node::Modulo(_, _) + | Node::Hypot(_, _) + | Node::EqualHeaviside(_, _) + | Node::NotEqualHeaviside(_, _) + | Node::Equality(_, _) + | Node::Neg(_) + | Node::Abs(_) + | Node::Sqrt(_) + | Node::Exp(_) + | Node::Log(_) + | Node::Sin(_) + | Node::Cos(_) + | Node::Tanh(_) + | Node::Sinh(_) + | Node::Cosh(_) + | Node::Arcsinh(_) + | Node::Arctan(_) + | Node::Erf(_) + | Node::Sign(_) + | Node::Floor(_) + | Node::Ceiling(_) + | Node::MaxReduce(_) + | Node::MinReduce(_) => {}, + Node::ReduceArgSelect { is_max, .. } => is_max.hash(hasher), + } + node.for_each_child(|c| { + use std::hash::Hash; + child_remap(c).hash(hasher); + }); +} + +const _: () = { + assert!(size_of::() <= 48); +}; + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::Arena; + + // ---- F3: validated data-structure constructors ---- + + #[test] + fn csr_try_new_accepts_valid_matrix() { + let m = CsrData::try_new( + vec![0, 1, 2], + vec![0, 1], + vec![1.0, 2.0], + Shape::matrix(2, 2), + ); + assert!(m.is_ok()); + } + + #[test] + fn csr_try_new_rejects_wrong_indptr_length() { + let m = CsrData::try_new(vec![0, 2], vec![0, 1], vec![1.0, 2.0], Shape::matrix(2, 2)); + assert!(m.is_err()); + } + + #[test] + fn csr_try_new_rejects_non_monotonic_indptr() { + let m = CsrData::try_new( + vec![0, 2, 1], + vec![0, 1], + vec![1.0, 2.0], + Shape::matrix(2, 2), + ); + assert!(m.is_err()); + } + + #[test] + fn csr_try_new_rejects_out_of_range_column() { + let m = CsrData::try_new( + vec![0, 1, 2], + vec![0, 5], + vec![1.0, 2.0], + Shape::matrix(2, 2), + ); + assert!(m.is_err()); + } + + #[test] + fn csr_try_new_rejects_indices_data_length_mismatch() { + let m = CsrData::try_new(vec![0, 1, 2], vec![0, 1], vec![1.0], Shape::matrix(2, 2)); + assert!(m.is_err()); + } + + #[test] + fn csr_try_new_rejects_indptr_tail_mismatch() { + let m = CsrData::try_new(vec![0, 1, 2], vec![0], vec![1.0], Shape::matrix(2, 2)); + assert!(m.is_err()); + } + + #[test] + fn array_try_new_accepts_matching_length() { + assert!(ArrayData::try_new(vec![1.0, 2.0, 3.0, 4.0], Shape::matrix(2, 2)).is_ok()); + } + + #[test] + fn array_try_new_rejects_length_mismatch() { + assert!(ArrayData::try_new(vec![1.0, 2.0, 3.0], Shape::matrix(2, 2)).is_err()); + } + + #[test] + fn interpolant_try_new_accepts_increasing_grid() { + assert!(InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![10.0, 20.0, 30.0]).is_ok()); + } + + #[test] + fn interpolant_try_new_rejects_empty_grid() { + assert!(InterpolantData::try_new(vec![], vec![]).is_err()); + } + + #[test] + fn interpolant_try_new_rejects_length_mismatch() { + assert!(InterpolantData::try_new(vec![0.0, 1.0], vec![10.0]).is_err()); + } + + #[test] + fn interpolant_try_new_rejects_non_increasing_knots() { + assert!(InterpolantData::try_new(vec![0.0, 2.0, 1.0], vec![1.0, 2.0, 3.0]).is_err()); + } + + #[test] + fn interpolant_try_new_rejects_non_finite_knot() { + assert!(InterpolantData::try_new(vec![0.0, f64::NAN], vec![1.0, 2.0]).is_err()); + } + + #[test] + fn cubic_try_new_accepts_valid() { + assert!(CubicInterpolantData::try_new(vec![0.0, 1.0], vec![[1.0, 2.0, 3.0, 4.0]]).is_ok()); + } + + #[test] + fn cubic_try_new_rejects_too_few_breakpoints() { + assert!(CubicInterpolantData::try_new(vec![0.0], vec![]).is_err()); + } + + #[test] + fn cubic_try_new_rejects_coeff_count_mismatch() { + let c = CubicInterpolantData::try_new( + vec![0.0, 1.0], + vec![[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + ); + assert!(c.is_err()); + } + + #[test] + fn cubic_try_new_rejects_non_increasing_breakpoints() { + assert!(CubicInterpolantData::try_new(vec![1.0, 0.0], vec![[1.0, 2.0, 3.0, 4.0]]).is_err()); + } + + #[test] + fn nd_try_new_accepts_valid_2d_linear() { + let nd = NdInterpolantData::try_new( + vec![vec![0.0, 1.0], vec![0.0, 1.0]], + vec![1.0, 2.0, 3.0, 4.0], + 2, + ); + assert!(nd.is_ok()); + } + + #[test] + fn nd_try_new_rejects_bad_order() { + let nd = NdInterpolantData::try_new( + vec![vec![0.0, 1.0], vec![0.0, 1.0]], + vec![1.0, 2.0, 3.0, 4.0], + 3, + ); + assert!(nd.is_err()); + } + + #[test] + fn nd_try_new_rejects_coeff_count_mismatch() { + let nd = NdInterpolantData::try_new( + vec![vec![0.0, 1.0], vec![0.0, 1.0]], + vec![1.0, 2.0, 3.0], + 2, + ); + assert!(nd.is_err()); + } + + #[test] + fn nd_try_new_rejects_axis_with_too_few_knots() { + let nd = NdInterpolantData::try_new(vec![vec![0.0], vec![0.0, 1.0]], vec![1.0, 2.0], 2); + assert!(nd.is_err()); + } + + #[test] + fn test_tangent_state_vector_node() { + let node = Node::TangentStateVector { start: 0, end: 3 }; + match node { + Node::TangentStateVector { start, end } => { + assert_eq!(start, 0); + assert_eq!(end, 3); + }, + _ => panic!("Expected TangentStateVector"), + } + } + + #[test] + fn test_tangent_parameter_node() { + let node = Node::TangentParameter { index: 2 }; + match node { + Node::TangentParameter { index } => { + assert_eq!(index, 2); + }, + _ => panic!("Expected TangentParameter"), + } + } + + #[test] + fn test_for_each_child_collects_all_children() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Scalar(2.0)); + let add = Node::Add(a, b); + + let mut collected = Vec::new(); + add.for_each_child(|c| collected.push(c)); + assert_eq!(collected, vec![a, b]); + + let leaf = Node::Scalar(2.72); + let mut leaf_children = Vec::new(); + leaf.for_each_child(|c| leaf_children.push(c)); + assert!(leaf_children.is_empty()); + + let concat = Node::Concat(vec![a, b, a]); + let mut concat_children = Vec::new(); + concat.for_each_child(|c| concat_children.push(c)); + assert_eq!(concat_children, vec![a, b, a]); + } + + #[test] + fn test_map_children_remaps_binary() { + let id_a: NodeId = 0u32.into(); + let id_b: NodeId = 1u32.into(); + let id_x: NodeId = 10u32.into(); + + let add = Node::Add(id_a, id_b); + let mapped = add.map_children(|c| if c == id_a { id_x } else { c }); + match mapped { + Node::Add(l, r) => { + assert_eq!(l, id_x); + assert_eq!(r, id_b); + }, + _ => panic!("expected Add"), + } + } + + #[test] + fn test_map_children_preserves_leaves() { + let s = Node::Scalar(3.15); + let mapped = s.map_children(|c| c); + assert_eq!(s, mapped); + + let zv = Node::ZeroVector { len: 4 }; + let mapped_zv = zv.map_children(|c| c); + assert_eq!(zv, mapped_zv); + } + + #[test] + fn test_map_children_preserves_concat_arity_and_order() { + let a: NodeId = 0u32.into(); + let b: NodeId = 1u32.into(); + let c: NodeId = 2u32.into(); + let concat = Node::Concat(vec![a, b, c]); + let mapped = concat.map_children(|x| NodeId::from(x.raw() + 100)); + match mapped { + Node::Concat(children) => { + assert_eq!(children, vec![100u32.into(), 101u32.into(), 102u32.into()]); + }, + _ => panic!("expected Concat"), + } + } + + #[test] + fn test_interpolant1d_linear_deriv_node() { + let node = Node::Interpolant1DLinearDeriv { + slopes: vec![10.0, 10.0].into_boxed_slice(), + x_data: vec![0.0, 1.0, 2.0].into_boxed_slice(), + child: { + let mut arena = Arena::new(); + arena.alloc(Node::Scalar(1.5)) + }, + }; + match node { + Node::Interpolant1DLinearDeriv { slopes, x_data, .. } => { + assert_eq!(slopes.len(), 2); + assert_eq!(x_data.len(), 3); + }, + _ => panic!("Expected Interpolant1DLinearDeriv"), + } + } + + #[test] + fn test_structural_hash_distinguishes_payload_only_differences() { + use rustc_hash::FxHasher; + use std::hash::Hasher; + + let s1 = Node::StateVector { start: 0, end: 5 }; + let s2 = Node::StateVector { start: 1, end: 6 }; + + let mut h1 = FxHasher::default(); + let mut h2 = FxHasher::default(); + structural_hash(&s1, &mut h1, |c| c); + structural_hash(&s2, &mut h2, |c| c); + assert_ne!(h1.finish(), h2.finish()); + + let zv1 = Node::ZeroVector { len: 3 }; + let zv2 = Node::ZeroVector { len: 7 }; + let mut h3 = FxHasher::default(); + let mut h4 = FxHasher::default(); + structural_hash(&zv1, &mut h3, |c| c); + structural_hash(&zv2, &mut h4, |c| c); + assert_ne!(h3.finish(), h4.finish()); + + let ip1 = Node::InputParameter { + name: "p1".into(), + index: 0, + offset: 0, + width: 1, + }; + let ip2 = Node::InputParameter { + name: "p1".into(), + index: 1, + offset: 1, + width: 1, + }; + let mut h5 = FxHasher::default(); + let mut h6 = FxHasher::default(); + structural_hash(&ip1, &mut h5, |c| c); + structural_hash(&ip2, &mut h6, |c| c); + assert_ne!(h5.finish(), h6.finish()); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/observable.rs b/packages/pybamm-rust/pybamm-core/src/observable.rs new file mode 100644 index 0000000000..5a78e4614f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/observable.rs @@ -0,0 +1,719 @@ +//! Observables: what a solve reports rather than integrates. +//! +//! Output variables and events are one structure — a primal tape, a `dH/dp` +//! tape, a `dH/dy` tape, a component count, a scratch buffer per tape — so it is +//! written once here and held twice. [`ObservableSet`] owns a family's +//! concatenated layout, its optional fused primal tape and its scratch sizing; +//! [`ObservableKind`] names a family, so a `CompiledModel` method takes one as +//! an argument instead of existing per family. + +use std::sync::Arc; + +use crate::arena::{Arena, NodeId}; +use crate::eval::{CompiledExpr, TangentInputs}; +use crate::ir::TypedIr; +use crate::node::Node; +use crate::simplify::simplify_pipeline; +use crate::tangent::{tangent_wrt_params, tangent_wrt_states}; + +/// Which family of observables a caller means. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ObservableKind { + /// Output variables `H(t, y; p)`, reported to the caller with their + /// sensitivities. + Outputs, + /// Event functions `g(t, y; p)`, which the integrator roots on. + Events, +} + +/// Write a parameter-space tangent into a global `dp` buffer. +/// +/// `sens_params[i]` is the global index of `v[i]`. The compiled tangent tapes are +/// generic over every parameter, so a seed at any global index is valid; this is +/// the one home of that convention, shared by the rhs and observable actions. +pub fn seed_param_tangent(dp: &mut [f64], sens_params: &[usize], v: &[f64]) { + debug_assert_eq!( + v.len(), + sens_params.len(), + "tangent length must match the parameter mapping", + ); + dp.fill(0.0); + for (&global, &vi) in sens_params.iter().zip(v) { + dp[global] = vi; + } +} + +/// Seed the unit tangent for sensitivity column `column`. +/// +/// [`seed_param_tangent`] with `v = e_column`, without a buffer to hold `e`. +pub fn seed_param_tangent_unit(dp: &mut [f64], sens_params: &[usize], column: usize) { + dp.fill(0.0); + dp[sens_params[column]] = 1.0; +} + +/// One compiled observable: the primal tape plus its two tangent tapes. +#[derive(Debug, Clone)] +pub struct CompiledObservable { + expr: Arc, + /// Components this observable writes, captured at compile time. + len: usize, + /// `dH/dp`, the param-tangent tape. + sens_expr: Arc, + /// `dH/dy`, the state-tangent tape. + jac_expr: Arc, +} + +impl CompiledObservable { + /// Compile `node` and its two tangent graphs off `arena`. + /// + /// `node` must already exist in `arena`. + pub fn new(arena: &Arena, node: NodeId) -> Self { + let ir = TypedIr::from_arena(arena, node); + Self { + len: ir.output_len(), + expr: Arc::new(CompiledExpr::from_ir(ir)), + sens_expr: compile_tangent(arena, node, TangentTarget::Params), + jac_expr: compile_tangent(arena, node, TangentTarget::States), + } + } + + /// Components this observable writes. + #[inline] + pub const fn output_len(&self) -> usize { + self.len + } + + /// Shared handle to the primal tape (no recompilation). + #[inline] + pub const fn expr(&self) -> &Arc { + &self.expr + } + + /// The `dH/dp` tape, as a selector for [`ObservableSet::sens_action`]. + #[inline] + const fn sens_tape(&self) -> &Arc { + &self.sens_expr + } + + /// The `dH/dy` tape, as a selector for [`ObservableSet::jac_action`]. + #[inline] + const fn jac_tape(&self) -> &Arc { + &self.jac_expr + } +} + +/// Which variable a tangent graph differentiates with respect to. +#[derive(Debug, Clone, Copy)] +enum TangentTarget { + Params, + States, +} + +/// Compile a tangent graph for `node` over a clone of the model's arena. +/// +/// Derivative tapes are simplified; the primal is not (see +/// `CompiledModel::new`), so this entry point is derivative-only. +fn compile_tangent(arena: &Arena, node: NodeId, target: TangentTarget) -> Arc { + let mut diff_arena = arena.clone(); + let root = match target { + TangentTarget::Params => tangent_wrt_params(&mut diff_arena, node), + TangentTarget::States => tangent_wrt_states(&mut diff_arena, node), + }; + let (final_arena, root) = simplify_pipeline(diff_arena, root); + Arc::new(CompiledExpr::from_ir(TypedIr::from_arena( + &final_arena, + root, + ))) +} + +/// A family of observables evaluated together, and the concatenated layout they +/// are reported in. +/// +/// Element `i` occupies `[sum(len_at(..i)), sum(len_at(..=i)))` of every buffer +/// the family-wide methods write, which is the layout `PyBaMM` reads output +/// variables and event values back in. +#[derive(Debug, Clone, Default)] +pub struct ObservableSet { + items: Vec, + /// One `Concat` tape over every root, evaluated by + /// [`eval_all`](Self::eval_all) when present. `None` until + /// [`fuse`](Self::fuse) is called with at least two roots. + fused: Option>, + /// `items`' component counts summed, maintained on push. + total_len: usize, +} + +impl ObservableSet { + /// An empty set. + #[inline] + pub const fn new() -> Self { + Self { + items: Vec::new(), + fused: None, + total_len: 0, + } + } + + /// Compile `node` and append it to the family. + /// + /// `node` must already exist in `arena`. Invalidates any fused tape, so + /// [`fuse`](Self::fuse) belongs after the last push. + pub fn push(&mut self, arena: &Arena, node: NodeId) { + let observable = CompiledObservable::new(arena, node); + self.total_len += observable.len; + self.items.push(observable); + self.fused = None; + } + + /// Build one fused primal tape from a `Concat` of `roots`. + /// + /// Per-observable tapes re-evaluate whatever subexpressions the family + /// shares; lowering a `Concat` of the roots emits each shared arena node + /// once. No-op for fewer than two observables, which have nothing to share. + /// `roots` must be the same nodes, in the same order, passed to + /// [`push`](Self::push), from `arena`. + pub fn fuse(&mut self, arena: &mut Arena, roots: &[NodeId]) { + if self.items.len() < 2 { + return; + } + debug_assert_eq!( + roots.len(), + self.items.len(), + "fuse: roots must match the pushed observables", + ); + let concat = arena.alloc(Node::Concat(roots.to_vec())); + // Same lowering entry point as `push`, so simplification and instruction + // semantics are identical to the per-observable tapes. + let ir = TypedIr::from_arena(arena, concat); + self.fused = Some(Arc::new(CompiledExpr::from_ir(ir))); + } + + /// Observables in the family. + #[inline] + pub const fn count(&self) -> usize { + self.items.len() + } + + /// Whether the family is empty. + #[inline] + pub const fn is_empty(&self) -> bool { + self.items.is_empty() + } + + /// Components observable `i` writes. + /// + /// Panics if `i >= count()`. + #[inline] + pub fn len_at(&self, i: usize) -> usize { + self.items[i].len + } + + /// Components the whole family writes, which is the buffer size every + /// family-wide method needs. + #[inline] + pub const fn total_len(&self) -> usize { + self.total_len + } + + /// Shared handle to observable `i`'s primal tape (no recompilation). + /// + /// Panics if `i >= count()`. + #[inline] + pub fn expr_arc(&self, i: usize) -> Arc { + Arc::clone(&self.items[i].expr) + } + + /// The fused primal tape, if [`fuse`](Self::fuse) built one. + #[inline] + pub const fn fused_expr(&self) -> Option<&Arc> { + self.fused.as_ref() + } + + /// Allocate per-solve buffers sized for this family. + pub fn create_scratch(&self) -> ObservableScratch { + ObservableScratch { + primal: self.scratches(CompiledObservable::expr), + sens: self.scratches(CompiledObservable::sens_tape), + jac: self.scratches(CompiledObservable::jac_tape), + fused: self + .fused + .as_ref() + .map_or_else(Vec::new, |e| vec![0.0; e.scratch_len()]), + batch: Vec::new(), + jac_term: vec![0.0; self.total_len], + } + } + + /// One zeroed buffer per observable, sized for the tape `tape` selects. + fn scratches(&self, tape: fn(&CompiledObservable) -> &Arc) -> Vec> { + self.items + .iter() + .map(|item| vec![0.0; tape(item).scratch_len()]) + .collect() + } + + /// Evaluate observable `i` into `out`, returning the count written. + /// + /// Panics if `i >= count()` or `out.len() < len_at(i)`. + pub fn eval_at( + &self, + scratch: &mut ObservableScratch, + t: f64, + y: &[f64], + inputs: &[f64], + i: usize, + out: &mut [f64], + ) -> usize { + let result = self.items[i] + .expr + .eval(&mut scratch.primal[i], t, y, &[], inputs); + out[..result.len()].copy_from_slice(result); + result.len() + } + + /// Evaluate the whole family into `out`, in the concatenated layout. + /// + /// Takes the fused tape when the family carries one. Panics if + /// `out.len() < total_len()`. + pub fn eval_all( + &self, + scratch: &mut ObservableScratch, + t: f64, + y: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + if let Some(fused) = &self.fused { + let result = fused.eval(&mut scratch.fused, t, y, &[], inputs); + out[..self.total_len].copy_from_slice(&result[..self.total_len]); + return; + } + let mut offset = 0; + for (item, buffer) in self.items.iter().zip(&mut scratch.primal) { + let result = item.expr.eval(buffer, t, y, &[], inputs); + out[offset..offset + item.len].copy_from_slice(&result[..item.len]); + offset += item.len; + } + } + + /// Batch-evaluate the whole family over `k` trajectory points. + /// + /// `ts` holds the `k` times, `y_cols` the `(n_states, k)` F-contiguous state + /// matrix, and `out` receives the `(total_len, k)` F-contiguous matrix. Tapes + /// the batch evaluator rejects (e.g. a `y_dot` reference) fall back to + /// per-point scalar evaluation, so results always match `k` + /// [`eval_all`](Self::eval_all) calls bitwise. + // Times, states, inputs and the output matrix are all distinct arguments. + #[allow(clippy::too_many_arguments)] + pub fn eval_batch( + &self, + scratch: &mut ObservableScratch, + k: usize, + ts: &[f64], + y_cols: &[f64], + n_states: usize, + inputs: &[f64], + out: &mut [f64], + ) { + let ObservableScratch { primal, batch, .. } = scratch; + let n_total = self.total_len; + let mut offset = 0; + for (item, fallback) in self.items.iter().zip(primal.iter_mut()) { + let (expr, len) = (&item.expr, item.len); + let needed = expr.scratch_len() * k; + if batch.len() < needed { + batch.resize(needed, 0.0); + } + match expr.eval_batch(&mut batch[..needed], k, ts, y_cols, inputs) { + Ok(result) => { + // result is lane-minor (`result[row * k + lane]`); out is + // F-contiguous. + for row in 0..len { + for (lane, &v) in result[row * k..(row + 1) * k].iter().enumerate() { + out[lane * n_total + offset + row] = v; + } + } + }, + Err(_) => { + for lane in 0..k { + let y = &y_cols[lane * n_states..(lane + 1) * n_states]; + let result = expr.eval(fallback, ts[lane], y, &[], inputs); + out[lane * n_total + offset..lane * n_total + offset + len] + .copy_from_slice(&result[..len]); + } + }, + } + offset += len; + } + } + + /// Compute `dH/dp . v` over the whole family into `out`. + /// + /// `dp` is the global parameter-space tangent, seeded by + /// [`seed_param_tangent`]. Panics if `out.len() < total_len()`. + pub fn sens_action( + &self, + scratch: &mut ObservableScratch, + t: f64, + y: &[f64], + inputs: &[f64], + dp: &[f64], + out: &mut [f64], + ) { + let tangent = TangentInputs { + dy: None, + dp: Some(dp), + }; + Self::tangent_action( + &self.items, + CompiledObservable::sens_tape, + &mut scratch.sens, + t, + y, + inputs, + &tangent, + out, + ); + } + + /// Compute `dH/dy . v` over the whole family into `out`. + /// + /// `v` is a state-space tangent of length `n_states`. Panics if + /// `out.len() < total_len()`. + pub fn jac_action( + &self, + scratch: &mut ObservableScratch, + t: f64, + y: &[f64], + inputs: &[f64], + v: &[f64], + out: &mut [f64], + ) { + let tangent = TangentInputs { + dy: Some(v), + dp: None, + }; + Self::tangent_action( + &self.items, + CompiledObservable::jac_tape, + &mut scratch.jac, + t, + y, + inputs, + &tangent, + out, + ); + } + + /// Project state sensitivities onto observable sensitivities. + /// + /// `dH/dp_k = sens_action(e_k) + jac_action(y_sens_k)` for every sensitivity + /// column, where `sens_params[k]` is column `k`'s global parameter index. + /// `y_sens` layout: `y_sens[k * n_states + j]` + /// `out` layout: `out[k * total_len() + o]` + /// `dp` is the caller's global tangent buffer, reseeded per column. + // Evaluation point, tangent buffers, mapping and output are all distinct. + #[allow(clippy::too_many_arguments)] + pub fn sens_project( + &self, + scratch: &mut ObservableScratch, + dp: &mut [f64], + t: f64, + y: &[f64], + inputs: &[f64], + sens_params: &[usize], + y_sens: &[f64], + n_states: usize, + out: &mut [f64], + ) { + let n_out = self.total_len; + if n_out == 0 || sens_params.is_empty() { + return; + } + // Destructured so the two tangent passes and the accumulator can borrow + // disjoint buffers of one scratch. + let ObservableScratch { + sens, + jac, + jac_term, + .. + } = scratch; + for k in 0..sens_params.len() { + seed_param_tangent_unit(dp, sens_params, k); + let dst = &mut out[k * n_out..(k + 1) * n_out]; + Self::tangent_action( + &self.items, + CompiledObservable::sens_tape, + sens, + t, + y, + inputs, + &TangentInputs { + dy: None, + dp: Some(dp), + }, + dst, + ); + Self::tangent_action( + &self.items, + CompiledObservable::jac_tape, + jac, + t, + y, + inputs, + &TangentInputs { + dy: Some(&y_sens[k * n_states..(k + 1) * n_states]), + dp: None, + }, + jac_term, + ); + for (dst_o, &term) in dst.iter_mut().zip(jac_term.iter()) { + *dst_o += term; // + dH/dy . y_sens_k + } + } + } + + /// The one tangent-action loop, over whichever tape `tape` selects. + /// + /// Walks observables and their buffers in lockstep and writes each result + /// into its slice of the concatenated layout. + // Tapes, buffers, evaluation point, tangent and output are distinct groups. + #[allow(clippy::too_many_arguments)] + fn tangent_action( + items: &[CompiledObservable], + tape: fn(&CompiledObservable) -> &Arc, + scratches: &mut [Vec], + t: f64, + y: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + out: &mut [f64], + ) { + let mut offset = 0; + for (item, buffer) in items.iter().zip(scratches.iter_mut()) { + let result = tape(item).eval_with_tangent(buffer, t, y, &[], inputs, tangent); + out[offset..offset + item.len].copy_from_slice(&result[..item.len]); + offset += item.len; + } + } +} + +/// Per-solve mutable buffers for one [`ObservableSet`]. +/// +/// Sized from the set, one buffer per tape, so a caller can neither mis-size a +/// buffer nor pair a scratch with the wrong family. Create one per solve; two +/// concurrent evaluations of the same set need two. +#[derive(Debug, Clone, Default)] +pub struct ObservableScratch { + /// One buffer per observable's primal tape. + primal: Vec>, + /// One buffer per observable's `dH/dp` tape. + sens: Vec>, + /// One buffer per observable's `dH/dy` tape. + jac: Vec>, + /// The fused tape's buffer; empty unless the set carries one. + fused: Vec, + /// Lane-scaled buffer for [`ObservableSet::eval_batch`], grown on demand and + /// reused across windows. + batch: Vec, + /// `dH/dy . y_sens_k` for one column of + /// [`ObservableSet::sens_project`]. + jac_term: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + /// Assert values, at a tolerance these small-integer fixtures never need. + fn assert_close(got: &[f64], want: &[f64]) { + for (g, w) in got.iter().zip(want) { + assert!((g - w).abs() < 1e-12, "got {got:?}, want {want:?}"); + } + } + + /// `H = [y0 * y1 + p0, y0 * y1 - 1]` as two observables over one shared + /// product, so fusing has something to share and both tangents are nonzero. + fn build_pair() -> (Arena, ObservableSet, Vec) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y0 = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let y1 = arena.alloc(Node::Index { + child: y, + start: 1, + end: 2, + }); + let product = arena.alloc(Node::Mul(y0, y1)); + let p0 = arena.alloc(Node::InputParameter { + name: "p0".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let first = arena.alloc(Node::Add(product, p0)); + let one = arena.alloc(Node::Scalar(1.0)); + let second = arena.alloc(Node::Sub(product, one)); + + let mut set = ObservableSet::new(); + set.push(&arena, first); + set.push(&arena, second); + (arena, set, vec![first, second]) + } + + #[test] + fn a_set_reports_its_concatenated_layout() { + let (_arena, set, _roots) = build_pair(); + assert_eq!(set.count(), 2); + assert_eq!(set.len_at(0), 1); + assert_eq!(set.len_at(1), 1); + assert_eq!(set.total_len(), 2); + assert!(!set.is_empty()); + assert!(ObservableSet::new().is_empty()); + } + + #[test] + fn eval_all_concatenates_what_eval_at_writes() { + let (_arena, set, _roots) = build_pair(); + let mut scratch = set.create_scratch(); + let (y, inputs) = ([3.0, 5.0], [7.0]); + + let mut all = [0.0; 2]; + set.eval_all(&mut scratch, 0.0, &y, &inputs, &mut all); + assert_close(&all, &[22.0, 14.0]); + + for (i, &concatenated) in all.iter().enumerate() { + let mut one = [0.0; 1]; + assert_eq!(set.eval_at(&mut scratch, 0.0, &y, &inputs, i, &mut one), 1); + assert_eq!(one[0].to_bits(), concatenated.to_bits()); + } + } + + #[test] + fn a_fused_set_evaluates_bitwise_what_the_per_item_tapes_do() { + let (mut arena, mut set, roots) = build_pair(); + let mut per_item_scratch = set.create_scratch(); + let (y, inputs) = ([0.75, -1.5], [0.25]); + let mut per_item = [0.0; 2]; + set.eval_all(&mut per_item_scratch, 0.0, &y, &inputs, &mut per_item); + + set.fuse(&mut arena, &roots); + assert!(set.fused_expr().is_some()); + let mut scratch = set.create_scratch(); + let mut fused = [0.0; 2]; + set.eval_all(&mut scratch, 0.0, &y, &inputs, &mut fused); + + assert_eq!(fused.map(f64::to_bits), per_item.map(f64::to_bits)); + // CSE across the roots: the shared product is emitted once. + let fused_instructions = set.fused_expr().unwrap().ir().instructions().len(); + let per_item_instructions: usize = (0..set.count()) + .map(|i| set.expr_arc(i).ir().instructions().len()) + .sum(); + assert!( + fused_instructions < per_item_instructions, + "fused tape ({fused_instructions}) should be shorter than the \ + per-item sum ({per_item_instructions})", + ); + } + + #[test] + fn fusing_fewer_than_two_observables_is_a_noop() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let mut set = ObservableSet::new(); + set.push(&arena, y); + set.fuse(&mut arena, &[y]); + assert!(set.fused_expr().is_none()); + } + + #[test] + fn a_push_after_fusing_drops_the_stale_tape() { + let (mut arena, mut set, roots) = build_pair(); + set.fuse(&mut arena, &roots); + assert!(set.fused_expr().is_some()); + set.push(&arena, roots[0]); + assert!( + set.fused_expr().is_none(), + "a fused tape that predates an observable would under-report the family", + ); + } + + #[test] + fn the_tangent_actions_differentiate_the_family() { + let (_arena, set, _roots) = build_pair(); + let mut scratch = set.create_scratch(); + let (y, inputs) = ([3.0, 5.0], [7.0]); + + // dH/dy . v with v = [1, 0] is [y1, y1]. + let mut jac = [0.0; 2]; + set.jac_action(&mut scratch, 0.0, &y, &inputs, &[1.0, 0.0], &mut jac); + assert_close(&jac, &[5.0, 5.0]); + + // dH/dp . e_0 is [1, 0]: only the first observable reads p0. + let mut dp = [0.0]; + seed_param_tangent(&mut dp, &[0], &[1.0]); + let mut sens = [0.0; 2]; + set.sens_action(&mut scratch, 0.0, &y, &inputs, &dp, &mut sens); + assert_close(&sens, &[1.0, 0.0]); + } + + #[test] + fn sens_project_sums_the_two_actions() { + let (_arena, set, _roots) = build_pair(); + let mut scratch = set.create_scratch(); + let (y, inputs) = ([3.0, 5.0], [7.0]); + let sens_params = [0]; + let y_sens = [1.0, 0.0]; // dy/dp0 + + let mut projected = [0.0; 2]; + let mut dp = [0.0]; + set.sens_project( + &mut scratch, + &mut dp, + 0.0, + &y, + &inputs, + &sens_params, + &y_sens, + 2, + &mut projected, + ); + + // dH/dp0 + dH/dy . dy/dp0 = [1, 0] + [y1, y1]. + assert_close(&projected, &[6.0, 5.0]); + } + + #[test] + fn eval_batch_matches_a_loop_of_eval_all() { + let (_arena, set, _roots) = build_pair(); + let mut scratch = set.create_scratch(); + let inputs = [7.0]; + let ts = [0.0, 1.0, 2.0]; + let y_cols = [3.0, 5.0, 1.0, 2.0, -4.0, 0.5]; + + let mut batched = [0.0; 6]; + set.eval_batch(&mut scratch, 3, &ts, &y_cols, 2, &inputs, &mut batched); + + for lane in 0..3 { + let mut expected = [0.0; 2]; + set.eval_all( + &mut scratch, + ts[lane], + &y_cols[lane * 2..(lane + 1) * 2], + &inputs, + &mut expected, + ); + assert_eq!( + batched[lane * 2..(lane + 1) * 2] + .iter() + .map(|v| v.to_bits()) + .collect::>(), + expected.map(f64::to_bits).to_vec(), + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/row_extract.rs b/packages/pybamm-rust/pybamm-core/src/row_extract.rs new file mode 100644 index 0000000000..4dfec6de28 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/row_extract.rs @@ -0,0 +1,612 @@ +//! Lift wide output rows of an expression out as standalone scalar expressions. +//! +//! A row wide enough to dominate the Jacobian colouring is filled by a reverse +//! pass over its own sub-expression instead ([`crate::adjoint`]). That needs the +//! row as a width-1 root, and a discretised model rarely offers one: its rhs is +//! a `Concat` of *blocks*, and a 2-D current collector puts one dense row per +//! collector node inside a single vector-valued block. +//! +//! So the row is synthesised rather than found. Extraction pushes an element +//! index down the DAG — `Concat` picks a child, `StateVector` narrows to one +//! entry, elementwise nodes index their operands, and `MatMul` against a +//! constant expands row `r` into its dot product. Nodes with no cheap indexed +//! form (interpolants, whose tables would be copied per element) stop the walk, +//! and the caller falls back to colouring the row like any other. +//! +//! Rows are extracted in groups, because they overlap almost entirely: a dense +//! operator makes every row of a block read every lane of the same upstream +//! expression. On a 12x12 pouch cell one row alone is 50k instructions of which +//! 49k is that shared upstream, so a tape per row costs 7.8M instructions where +//! four rows per tape costs 2.4M and one tape for all 144 costs 283k. +//! +//! Grouping is not free in the other direction: every row walks its group's +//! whole tape backwards, so one assembly costs 66 ms at one row per tape, 58 ms +//! at four, and 215 ms at all 144 -- worse than not splitting at all. +//! [`crate::jacobian::ROWS_PER_TAPE`] picks the knee. + +use rustc_hash::FxHashMap; + +use crate::arena::{Arena, NodeId, NodeMap}; +use crate::node::Node; +use crate::simplify::{dce, node_len}; + +/// A group of output rows lifted out as one expression owning its own arena. +/// +/// Element `i` of [`root`](Self::root) is element `i` of [`rows`](Self::rows). +#[derive(Debug)] +pub struct ScalarRowBlock { + /// Parent output rows this block holds, in element order. + pub rows: Vec, + /// Arena holding only the nodes those rows reach. + pub arena: Arena, + /// `Concat` of the extracted rows, one element each. + pub root: NodeId, +} + +/// Extract `rows` of `root` as one block: a `Concat` whose element `i` is +/// `rows[i]`. +/// +/// The rows share one tape and one forward pass. Each row's backward pass walks +/// only the shared upstream plus its own cone, so holding every row on one tape +/// costs nothing per row and carries the shared expression exactly once. +/// +/// `None` means the caller must colour every row as before: extraction is +/// all-or-nothing, since the colouring that adopted `rows` assumed all of them +/// leave the sweep. +pub fn extract_scalar_rows(arena: &Arena, root: NodeId, rows: &[usize]) -> Option { + debug_assert!(!rows.is_empty(), "the caller declines before extracting"); + let mut extractor = Extractor::new(arena); + let mut extracted = Vec::with_capacity(rows.len()); + for &row in rows { + extracted.push(extractor.push(root, row)?); + } + let concat = extractor.work.alloc(Node::Concat(extracted)); + // Compacted because tape construction allocates per node of the arena it is + // handed, and the extraction arena still carries the whole model. + let (arena, root) = dce(&extractor.work, concat); + Some(ScalarRowBlock { + rows: rows.to_vec(), + arena, + root, + }) +} + +/// Index-push walker over one primal DAG. +struct Extractor<'a> { + /// The graph being indexed. Its ids stay valid in `work`, which starts as + /// its clone, so widths and children can be read straight from here. + source: &'a Arena, + /// `source` plus every synthesised scalar node. + work: Arena, + /// `node_len` memo over `source`; synthesised nodes are all width 1 and are + /// never queried. + widths: NodeMap, + /// Element `index` of node `id`, shared across every requested row. + cache: FxHashMap<(NodeId, usize), NodeId>, + /// Synthesised `Scalar` nodes by bit pattern. A discretisation matrix has a + /// handful of distinct coefficients across hundreds of thousands of entries. + scalars: FxHashMap, +} + +impl<'a> Extractor<'a> { + fn new(source: &'a Arena) -> Self { + Self { + source, + work: source.clone(), + widths: NodeMap::new(source.len()), + cache: FxHashMap::default(), + scalars: FxHashMap::default(), + } + } + + fn alloc(&mut self, node: Node) -> NodeId { + self.work.alloc(node) + } + + fn width(&mut self, id: NodeId) -> usize { + node_len(self.source, id, &mut self.widths) + } + + /// A `Scalar` node for `value`, interned so equal coefficients share one. + fn scalar(&mut self, value: f64) -> NodeId { + if let Some(&hit) = self.scalars.get(&value.to_bits()) { + return hit; + } + let id = self.alloc(Node::Scalar(value)); + self.scalars.insert(value.to_bits(), id); + id + } + + /// Element `index` of `id` as a width-1 node in `work`. + fn push(&mut self, id: NodeId, index: usize) -> Option { + let width = self.width(id); + if index >= width { + return None; + } + // Reused whole: a width-1 node already is its own element 0. + if width == 1 { + return Some(id); + } + if let Some(&hit) = self.cache.get(&(id, index)) { + return Some(hit); + } + let pushed = self.push_uncached(id, index)?; + self.cache.insert((id, index), pushed); + Some(pushed) + } + + fn push_uncached(&mut self, id: NodeId, index: usize) -> Option { + // Copied out so the match borrows the source graph, not `self`. + let source = self.source; + match source.get(id) { + Node::Concat(children) => { + let mut offset = 0; + for &child in children { + let len = self.width(child); + if index < offset + len { + return self.push(child, index - offset); + } + offset += len; + } + None + }, + Node::Index { child, start, .. } => self.push(*child, start + index), + Node::StateVector { start, .. } => { + let at = start + index; + Some(self.alloc(Node::StateVector { + start: at, + end: at + 1, + })) + }, + Node::StateVectorDot { start, .. } => { + let at = start + index; + Some(self.alloc(Node::StateVectorDot { + start: at, + end: at + 1, + })) + }, + Node::Array(array) => { + let value = array.data()[index]; + Some(self.scalar(value)) + }, + Node::ZeroVector { .. } => Some(self.scalar(0.0)), + // A vector input has no narrower load, but slicing a leaf costs + // nothing: the tape reads the packed values either way. + Node::InputParameter { .. } => Some(self.alloc(Node::Index { + child: id, + start: index, + end: index + 1, + })), + Node::MatMul(matrix, vector) => self.push_matmul_row(*matrix, *vector, index), + node if is_elementwise(node) => self.push_elementwise(node, index), + // Interpolants are the notable omission: `map_children` clones their + // table, so indexing one would copy it per element and per tape. + _ => None, + } + } + + /// Rebuild an elementwise node on indexed operands, broadcasting the + /// width-1 ones unchanged. + fn push_elementwise(&mut self, node: &Node, index: usize) -> Option { + let mut children = Vec::new(); + node.for_each_child(|child| children.push(child)); + let mut pushed = Vec::with_capacity(children.len()); + for &child in &children { + let at = if self.width(child) == 1 { 0 } else { index }; + pushed.push(self.push(child, at)?); + } + // `map_children` visits children in `for_each_child` order, so the + // mapped ids line up one for one. + let mut mapped = pushed.into_iter(); + let indexed = node.map_children(|_| mapped.next().expect("one push per child")); + Some(self.alloc(indexed)) + } + + /// Row `row` of a constant matrix contracted against `vector`. + /// + /// Structural zeros are skipped, and so are stored zeros of a dense + /// operand: this is how a 256-wide block stops being one term per column. + fn push_matmul_row(&mut self, matrix: NodeId, vector: NodeId, row: usize) -> Option { + let terms: Vec<(f64, usize)> = match self.source.get(matrix) { + Node::SparseMatrix(csr) => { + let (start, end) = (csr.indptr()[row], csr.indptr()[row + 1]); + csr.data()[start..end] + .iter() + .copied() + .zip(csr.indices()[start..end].iter().copied()) + .collect() + }, + Node::Array(array) => { + let cols = array.shape().cols; + array.data()[row * cols..(row + 1) * cols] + .iter() + .copied() + .enumerate() + .filter(|&(_, value)| !is_zero(value)) + .map(|(col, value)| (value, col)) + .collect() + }, + // Only a constant left operand has rows to index; anything else is + // rejected before lowering anyway. + _ => return None, + }; + + let mut sum: Option = None; + for (coefficient, col) in terms { + let element = self.push(vector, col)?; + let term = self.scaled(coefficient, element); + // Left-to-right, matching the accumulation order the vector matmul + // evaluates in. + sum = Some(sum.map_or(term, |accumulated| self.alloc(Node::Add(accumulated, term)))); + } + Some(sum.unwrap_or_else(|| self.scalar(0.0))) + } + + /// `coefficient * element`, folding the unit coefficients a discretisation + /// matrix is mostly made of. Both foldings are exact. + fn scaled(&mut self, coefficient: f64, element: NodeId) -> NodeId { + if coefficient.to_bits() == 1.0_f64.to_bits() { + return element; + } + if coefficient.to_bits() == (-1.0_f64).to_bits() { + return self.alloc(Node::Neg(element)); + } + let scale = self.scalar(coefficient); + self.alloc(Node::Mul(scale, element)) + } +} + +/// Whether indexing this node is just indexing each of its children. +/// +/// `Conditional` qualifies because its selector is width 1 and so indexes to +/// itself, leaving only the branches to narrow. +/// +/// Exhaustive on purpose: a new [`Node`] variant must be classified here, or it +/// would silently stop extraction and decline the split with nothing to show. +const fn is_elementwise(node: &Node) -> bool { + match node { + Node::Add(..) + | Node::Sub(..) + | Node::Mul(..) + | Node::Div(..) + | Node::Pow(..) + | Node::Minimum(..) + | Node::Maximum(..) + | Node::Modulo(..) + | Node::Hypot(..) + | Node::EqualHeaviside(..) + | Node::NotEqualHeaviside(..) + | Node::Equality(..) + | Node::Neg(_) + | Node::Abs(_) + | Node::Sqrt(_) + | Node::Exp(_) + | Node::Log(_) + | Node::Sin(_) + | Node::Cos(_) + | Node::Tanh(_) + | Node::Sinh(_) + | Node::Cosh(_) + | Node::Arcsinh(_) + | Node::Arctan(_) + | Node::Erf(_) + | Node::Sign(_) + | Node::Floor(_) + | Node::Ceiling(_) + | Node::Conditional { .. } => true, + // Handled by `push_uncached` ahead of this check, or genuinely not + // indexable: reductions collapse width, interpolants would copy their + // table per element, and a matmul needs its row expanded. + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::InputParameter { .. } + | Node::Time + | Node::MatMul(..) + | Node::Index { .. } + | Node::Concat(_) + | Node::Interpolant1DLinear { .. } + | Node::Interpolant1DCubic { .. } + | Node::InterpolantNd { .. } + | Node::InterpolantNdPartial { .. } + | Node::MaxReduce(_) + | Node::MinReduce(_) + | Node::ReduceArgSelect { .. } + | Node::TangentStateVector { .. } + | Node::TangentParameter { .. } + | Node::Interpolant1DLinearDeriv { .. } + | Node::Interpolant1DCubicDeriv { .. } => false, + } +} + +/// Both signed zeros. Spelled through the bit pattern because a dropped term +/// must be one the matmul could not have contributed to. +const fn is_zero(value: f64) -> bool { + value.to_bits() == 0.0_f64.to_bits() || value.to_bits() == (-0.0_f64).to_bits() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::CompiledExpr; + use crate::node::{ArrayData, CsrData, Shape}; + + /// Element `row` of `root`, evaluated through an extracted scalar tape. + fn extracted_value(arena: &Arena, root: NodeId, row: usize, y: &[f64], p: &[f64]) -> f64 { + let block = extract_scalar_rows(arena, root, &[row]) + .unwrap_or_else(|| panic!("row {row} must extract")); + let expr = CompiledExpr::new(&block.arena, block.root); + let mut scratch = vec![0.0; expr.scratch_len()]; + expr.eval(&mut scratch, 0.5, y, &[], p)[0] + } + + /// Element `row` of `root` evaluated the ordinary way, as the reference. + fn reference_value(arena: &Arena, root: NodeId, row: usize, y: &[f64], p: &[f64]) -> f64 { + let expr = CompiledExpr::new(arena, root); + let mut scratch = vec![0.0; expr.scratch_len()]; + expr.eval(&mut scratch, 0.5, y, &[], p)[row] + } + + /// Every element of `root` extracts to the value the whole expression puts + /// at that position. + fn assert_extracts_elementwise(arena: &Arena, root: NodeId, width: usize, y: &[f64]) { + for row in 0..width { + let extracted = extracted_value(arena, root, row, y, &[]); + let reference = reference_value(arena, root, row, y, &[]); + assert_eq!( + extracted.to_bits(), + reference.to_bits(), + "row {row}: extracted {extracted} != reference {reference}" + ); + } + } + + fn states(arena: &mut Arena, n: usize) -> NodeId { + arena.alloc(Node::StateVector { start: 0, end: n }) + } + + fn sample(n: usize) -> Vec { + (0..n) + .map(|i| (i as f64).mul_add(0.37, 1.1).sin() + 1.5) + .collect() + } + + #[test] + fn state_vector_rows_narrow_to_one_entry() { + let mut arena = Arena::new(); + let y = states(&mut arena, 5); + assert_extracts_elementwise(&arena, y, 5, &sample(5)); + } + + #[test] + fn elementwise_rows_index_both_operands() { + let mut arena = Arena::new(); + let y = states(&mut arena, 6); + let scalar = arena.alloc(Node::Scalar(2.5)); + // A vector op, a broadcast scalar operand, and a unary on top. + let product = arena.alloc(Node::Mul(y, y)); + let shifted = arena.alloc(Node::Add(product, scalar)); + let root = arena.alloc(Node::Tanh(shifted)); + assert_extracts_elementwise(&arena, root, 6, &sample(6)); + } + + #[test] + fn broadcast_operand_is_not_indexed() { + // `Div` with the vector on the right: the width-1 side must stay whole + // rather than be indexed alongside it. + let mut arena = Arena::new(); + let y = states(&mut arena, 4); + let numerator = arena.alloc(Node::Scalar(3.0)); + let root = arena.alloc(Node::Div(numerator, y)); + assert_extracts_elementwise(&arena, root, 4, &sample(4)); + } + + #[test] + fn index_and_concat_rows_resolve_through_structure() { + let mut arena = Arena::new(); + let y = states(&mut arena, 8); + let tail = arena.alloc(Node::Index { + child: y, + start: 5, + end: 8, + }); + let head = arena.alloc(Node::Sin(y)); + let root = arena.alloc(Node::Concat(vec![head, tail])); + assert_extracts_elementwise(&arena, root, 11, &sample(8)); + } + + #[test] + fn literal_rows_extract_their_own_element() { + let mut arena = Arena::new(); + let y = states(&mut arena, 3); + let array = arena.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![1.5, -2.5, 4.0], Shape::vector(3)).expect("valid array"), + ))); + let zeros = arena.alloc(Node::ZeroVector { len: 3 }); + let sum = arena.alloc(Node::Add(array, zeros)); + let root = arena.alloc(Node::Mul(sum, y)); + assert_extracts_elementwise(&arena, root, 3, &sample(3)); + } + + #[test] + fn conditional_rows_index_branches_and_keep_the_selector() { + let mut arena = Arena::new(); + let y = states(&mut arena, 4); + let selector = arena.alloc(Node::Scalar(2.0)); + let first = arena.alloc(Node::Sin(y)); + let second = arena.alloc(Node::Cos(y)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![first, second], + }); + assert_extracts_elementwise(&arena, root, 4, &sample(4)); + } + + #[test] + fn sparse_matmul_rows_expand_to_their_stored_entries() { + let mut arena = Arena::new(); + let n = 6; + let y = states(&mut arena, n); + let squared = arena.alloc(Node::Mul(y, y)); + // Rows of 2, 1 and 0 entries, so the empty-row and single-term paths + // are both exercised alongside the ordinary sum. + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 3, 3], + vec![0, 3, 5], + vec![2.0, -1.0, 1.0], + Shape::matrix(3, n), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, squared)); + assert_extracts_elementwise(&arena, root, 3, &sample(n)); + } + + #[test] + fn dense_matmul_rows_expand_row_major_and_skip_zeros() { + let mut arena = Arena::new(); + let n = 4; + let y = states(&mut arena, n); + let matrix = arena.alloc(Node::Array(Box::new( + ArrayData::try_new( + vec![ + 1.0, 0.0, 0.5, 0.0, // + 0.0, 0.0, 0.0, 0.0, // + -1.0, 2.0, 0.0, 3.0, + ], + Shape::matrix(3, n), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, y)); + assert_extracts_elementwise(&arena, root, 3, &sample(n)); + } + + #[test] + fn vector_input_parameters_index_by_element() { + let mut arena = Arena::new(); + let y = states(&mut arena, 3); + let parameter = arena.alloc(Node::InputParameter { + name: "k".into(), + index: 0, + offset: 0, + width: 3, + }); + let root = arena.alloc(Node::Mul(parameter, y)); + let y_values = sample(3); + let p = [0.25, -1.5, 2.0]; + for row in 0..3 { + assert_eq!( + extracted_value(&arena, root, row, &y_values, &p).to_bits(), + reference_value(&arena, root, row, &y_values, &p).to_bits() + ); + } + } + + #[test] + fn unsupported_nodes_stop_the_walk() { + // An interpolant is deliberately not indexable: its table would be + // copied per element. + let mut arena = Arena::new(); + let y = states(&mut arena, 4); + let interpolated = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + crate::node::InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 1.0, 4.0]) + .expect("valid table"), + ), + child: y, + }); + assert!(extract_scalar_rows(&arena, interpolated, &[1]).is_none()); + } + + #[test] + fn out_of_range_rows_are_rejected() { + let mut arena = Arena::new(); + let y = states(&mut arena, 3); + assert!(extract_scalar_rows(&arena, y, &[3]).is_none()); + } + + /// A row that is already its own width-1 block is returned as it stands, + /// so the pre-existing scalar path costs no synthesis. + #[test] + fn scalar_blocks_extract_without_synthesis() { + let mut arena = Arena::new(); + let n = 8; + let y = states(&mut arena, n); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, y)); + // Matrix, state read, matmul, and the one-element Concat wrapping them: + // the row as it already stood, nothing synthesised. + let block = extract_scalar_rows(&arena, root, &[0]).expect("no synthesis needed"); + assert_eq!(block.arena.len(), 4); + assert_extracts_elementwise(&arena, root, 1, &sample(n)); + } + + /// Rows of one block read the same upstream, and the shared tape must hold + /// it once. Two rows over the same `n` squares cost one set of squares. + #[test] + fn rows_share_the_upstream_they_have_in_common() { + let mut arena = Arena::new(); + let n = 64; + let y = states(&mut arena, n); + let squared = arena.alloc(Node::Mul(y, y)); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n, 2 * n], + (0..n).chain(0..n).collect(), + vec![1.0; 2 * n], + Shape::matrix(2, n), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, squared)); + let one = extract_scalar_rows(&arena, root, &[0]).expect("one row"); + let both = extract_scalar_rows(&arena, root, &[0, 1]).expect("both rows"); + // Only the second row's own sum tree is added, not another copy of the + // n squares and n state reads. + assert!( + both.arena.len() < 2 * one.arena.len(), + "two rows took {} nodes against {} for one", + both.arena.len(), + one.arena.len() + ); + } + + /// Every block element is the row its `rows` entry names, whatever order + /// the rows were requested in. + #[test] + fn block_elements_name_the_rows_they_hold() { + let mut arena = Arena::new(); + let n = 6; + let y = states(&mut arena, n); + let root = arena.alloc(Node::Sin(y)); + let wanted = [4usize, 1, 3]; + let block = extract_scalar_rows(&arena, root, &wanted).expect("rows"); + let y_values = sample(n); + assert_eq!( + block.rows, wanted, + "the block must keep the requested order" + ); + let expr = CompiledExpr::new(&block.arena, block.root); + let mut scratch = vec![0.0; expr.scratch_len()]; + let got = expr.eval(&mut scratch, 0.5, &y_values, &[], &[]).to_vec(); + for (&value, &row) in got.iter().zip(&block.rows) { + assert_eq!( + value.to_bits(), + reference_value(&arena, root, row, &y_values, &[]).to_bits() + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/simplify.rs b/packages/pybamm-rust/pybamm-core/src/simplify.rs new file mode 100644 index 0000000000..f9aa159fae --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/simplify.rs @@ -0,0 +1,2108 @@ +//! Expression simplification pass for symbolic differentiation. +//! +//! Without simplification, derivative expressions explode. For example, +//! `d(a*b*c)/dx` produces terms like `Mul(Scalar(0), b)` that should fold to zero. +//! +//! This module provides two modes: +//! - **Conservative** (default): Skips rules that could change NaN/infinity behavior +//! - **Aggressive**: Applies all algebraic simplifications +//! +//! # Examples +//! +//! ``` +//! use pybamm_core::{Arena, Node, simplify}; +//! +//! let mut arena = Arena::new(); +//! let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); +//! let zero = arena.alloc(Node::Scalar(0.0)); +//! let expr = arena.alloc(Node::Add(x, zero)); +//! +//! let simplified = simplify(&mut arena, expr); +//! assert_eq!(simplified, x); +//! ``` + +use crate::arena::{Arena, NodeId, NodeMap}; +use crate::eval::{erf_approx, sign}; +use crate::node::Node; + +/// Controls which simplification rules are applied. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum SimplifyMode { + /// Value-preserving rewrites only. Rules that would change a NaN result are + /// skipped: `0 * x -> 0` (x=Inf), `0 / x -> 0` and `x / x -> 1` (x=0), and + /// `x - x -> 0` (x=Inf). So are the two sub rules that would change a zero's + /// sign: `x - 0 -> x` from `-0.0` and `0 - x -> -x` from `+0.0`. + /// + /// Values are exact but the sign of a zero result is not: `-0.0 + 0.0` + /// is `+0.0` where folding the `+ 0` away yields `-0.0`. No `PyBaMM` path + /// distinguishes them, so the fold stays rather than bloating every tape + /// `zero_propagate` feeds. + #[default] + Conservative, + + /// Apply all algebraic simplifications. + /// Use when you know inputs are well-behaved (no NaN/Inf). + Aggressive, +} + +/// Constant-fold and apply algebraic identities in `Conservative` mode. +#[must_use] +pub fn simplify(arena: &mut Arena, root: NodeId) -> NodeId { + simplify_with_mode(arena, root, SimplifyMode::Conservative) +} + +/// As [`simplify`], with the rule set chosen by `mode`. +#[must_use] +pub fn simplify_with_mode(arena: &mut Arena, root: NodeId, mode: SimplifyMode) -> NodeId { + let mut memo: NodeMap = NodeMap::new(arena.len()); + let mut lens: NodeMap = NodeMap::new(arena.len()); + simplify_node(arena, root, mode, &mut memo, &mut lens) +} + +/// Check if a node is a scalar with a specific value. +/// +/// Exact literal match: tiny nonzero coefficients (e.g. 1e-17) must not be +/// treated as 0/1. Mirrors `zero_propagate`'s exact `== 0.0` policy. +#[allow(clippy::float_cmp)] +fn is_scalar(arena: &Arena, id: NodeId, value: f64) -> bool { + match arena.get(id) { + Node::Scalar(v) => *v == value, + _ => false, + } +} + +/// Check if a node is the scalar 0.0. +fn is_zero(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, 0.0) +} + +/// Check if a node is the scalar `+0.0`, distinguished from `-0.0`. +/// +/// `is_zero` cannot tell the two apart (`-0.0 == 0.0`), but the additive +/// identities are only bit-exact for one sign of zero each. +fn is_positive_zero(arena: &Arena, id: NodeId) -> bool { + matches!(arena.get(id), Node::Scalar(v) if v.to_bits() == 0.0_f64.to_bits()) +} + +/// Check if a node is the scalar `-0.0`. See [`is_positive_zero`]. +fn is_negative_zero(arena: &Arena, id: NodeId) -> bool { + matches!(arena.get(id), Node::Scalar(v) if v.to_bits() == (-0.0_f64).to_bits()) +} + +/// Check if a node is the scalar 1.0. +fn is_one(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, 1.0) +} + +/// Check if a node is the scalar -1.0. +fn is_neg_one(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, -1.0) +} + +/// Get the scalar value if the node is a scalar. +fn get_scalar(arena: &Arena, id: NodeId) -> Option { + match arena.get(id) { + Node::Scalar(v) => Some(*v), + _ => None, + } +} + +/// Broadcast length of a node, mirroring `ir::infer_sizes` semantics. +/// +/// Memoized via `lens`. Used by folds that replace an expression with a +/// constant, so the replacement keeps the expression's length; also reused by +/// `jacobian::mask_scalar_rows` and `row_extract` to walk `Concat` offsets. +pub(crate) fn node_len(arena: &Arena, id: NodeId, lens: &mut NodeMap) -> usize { + if let Some(&len) = lens.get(id) { + return len; + } + let len = match arena.get(id) { + Node::Scalar(_) + | Node::Time + | Node::TangentParameter { .. } + | Node::MaxReduce(_) + | Node::MinReduce(_) + | Node::ReduceArgSelect { .. } => 1, + // Width, not 1: a vector input broadcasts like any other vector, which + // is how `infer_sizes` sizes it. + Node::InputParameter { width, .. } => *width, + Node::Array(arr) => arr.data.len(), + Node::ZeroVector { len } => *len, + Node::SparseMatrix(_) => 0, + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } => end - start, + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => { + let (a, b) = (*a, *b); + node_len(arena, a, lens).max(node_len(arena, b, lens)) + }, + Node::MatMul(a, b) => match arena.get(*a) { + Node::SparseMatrix(csr) => csr.shape.rows, + Node::Array(arr) => arr.shape.rows, + _ => node_len(arena, *b, lens), + }, + Node::Concat(children) => { + let mut total = 0; + for &c in children { + total += node_len(arena, c, lens); + } + total + }, + Node::Neg(a) + | Node::Abs(a) + | Node::Sqrt(a) + | Node::Exp(a) + | Node::Log(a) + | Node::Sin(a) + | Node::Cos(a) + | Node::Tanh(a) + | Node::Sinh(a) + | Node::Cosh(a) + | Node::Arcsinh(a) + | Node::Arctan(a) + | Node::Erf(a) + | Node::Sign(a) + | Node::Floor(a) + | Node::Ceiling(a) => node_len(arena, *a, lens), + Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => node_len(arena, *child, lens), + Node::InterpolantNd { children, .. } | Node::InterpolantNdPartial { children, .. } => { + let mut max = 1; + for &c in children { + max = max.max(node_len(arena, c, lens)); + } + max + }, + Node::Conditional { branches, .. } => { + let mut max = 1; + for &b in branches { + max = max.max(node_len(arena, b, lens)); + } + max + }, + }; + lens.insert(id, len); + len +} + +/// Allocate a zero with the given broadcast length. +fn zero_of_len(arena: &mut Arena, len: usize) -> NodeId { + if len == 1 { + arena.alloc(Node::Scalar(0.0)) + } else { + arena.alloc(Node::ZeroVector { len }) + } +} + +/// Recursively simplify a node, using memoization to avoid redundant work. +fn simplify_node( + arena: &mut Arena, + id: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + // Check memo first + if let Some(&cached) = memo.get(id) { + return cached; + } + + let result = match arena.get(id).clone() { + // Leaves - no simplification needed + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::InputParameter { .. } + | Node::Time + | Node::TangentStateVector { .. } + | Node::TangentParameter { .. } => id, + + // Binary operations + Node::Add(lhs, rhs) => simplify_add(arena, lhs, rhs, mode, memo, lens), + Node::Sub(lhs, rhs) => simplify_sub(arena, lhs, rhs, mode, memo, lens), + Node::Mul(lhs, rhs) => simplify_mul(arena, lhs, rhs, mode, memo, lens), + Node::Div(lhs, rhs) => simplify_div(arena, lhs, rhs, mode, memo, lens), + Node::Pow(base, exp) => simplify_pow(arena, base, exp, mode, memo, lens), + + // Other binary ops - just simplify children, no algebraic rules + Node::MatMul(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::MatMul(lhs_s, rhs_s)) + } + }, + Node::Minimum(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a.min(b))); + } + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::Minimum(lhs_s, rhs_s)) + } + }, + Node::Maximum(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a.max(b))); + } + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::Maximum(lhs_s, rhs_s)) + } + }, + Node::Modulo(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a % b)); + } + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::Modulo(lhs_s, rhs_s)) + } + }, + Node::Hypot(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a.hypot(b))); + } + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::Hypot(lhs_s, rhs_s)) + } + }, + Node::EqualHeaviside(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::EqualHeaviside(lhs_s, rhs_s)) + } + }, + Node::NotEqualHeaviside(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::NotEqualHeaviside(lhs_s, rhs_s)) + } + }, + Node::Equality(lhs, rhs) => { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + if lhs_s == lhs && rhs_s == rhs { + id + } else { + arena.alloc(Node::Equality(lhs_s, rhs_s)) + } + }, + + // Unary operations + Node::Neg(child) => simplify_neg(arena, child, mode, memo, lens), + Node::Abs(child) => simplify_unary(arena, child, mode, memo, lens, f64::abs, Node::Abs), + Node::Sqrt(child) => simplify_unary(arena, child, mode, memo, lens, f64::sqrt, Node::Sqrt), + Node::Exp(child) => simplify_unary(arena, child, mode, memo, lens, f64::exp, Node::Exp), + Node::Log(child) => simplify_unary(arena, child, mode, memo, lens, f64::ln, Node::Log), + Node::Sin(child) => simplify_unary(arena, child, mode, memo, lens, f64::sin, Node::Sin), + Node::Cos(child) => simplify_unary(arena, child, mode, memo, lens, f64::cos, Node::Cos), + Node::Tanh(child) => simplify_unary(arena, child, mode, memo, lens, f64::tanh, Node::Tanh), + Node::Sinh(child) => simplify_unary(arena, child, mode, memo, lens, f64::sinh, Node::Sinh), + Node::Cosh(child) => simplify_unary(arena, child, mode, memo, lens, f64::cosh, Node::Cosh), + Node::Arcsinh(child) => { + simplify_unary(arena, child, mode, memo, lens, f64::asinh, Node::Arcsinh) + }, + Node::Arctan(child) => { + simplify_unary(arena, child, mode, memo, lens, f64::atan, Node::Arctan) + }, + Node::Erf(child) => simplify_unary(arena, child, mode, memo, lens, erf_approx, Node::Erf), + Node::Sign(child) => simplify_unary(arena, child, mode, memo, lens, sign, Node::Sign), + Node::Floor(child) => { + simplify_unary(arena, child, mode, memo, lens, f64::floor, Node::Floor) + }, + Node::Ceiling(child) => { + simplify_unary(arena, child, mode, memo, lens, f64::ceil, Node::Ceiling) + }, + Node::MaxReduce(child) => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::MaxReduce(child_s)) + } + }, + Node::MinReduce(child) => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::MinReduce(child_s)) + } + }, + Node::ReduceArgSelect { + basis, + picker, + is_max, + } => { + let basis_s = simplify_node(arena, basis, mode, memo, lens); + // basis all-zero -> selecting any element yields the scalar 0 + if is_zero(arena, basis_s) || matches!(arena.get(basis_s), Node::ZeroVector { .. }) { + arena.alloc(Node::Scalar(0.0)) + } else { + let picker_s = simplify_node(arena, picker, mode, memo, lens); + if basis_s == basis && picker_s == picker { + id + } else { + arena.alloc(Node::ReduceArgSelect { + basis: basis_s, + picker: picker_s, + is_max, + }) + } + } + }, + + // Structural operations + Node::Index { child, start, end } => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::Index { + child: child_s, + start, + end, + }) + } + }, + Node::Concat(children) => { + let children_s: Vec = children + .iter() + .map(|&c| simplify_node(arena, c, mode, memo, lens)) + .collect(); + if children_s == children { + id + } else { + arena.alloc(Node::Concat(children_s)) + } + }, + + Node::Interpolant1DLinear { data, child } => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::Interpolant1DLinear { + data, + child: child_s, + }) + } + }, + Node::Interpolant1DLinearDeriv { + slopes, + x_data, + child, + } => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::Interpolant1DLinearDeriv { + slopes, + x_data, + child: child_s, + }) + } + }, + Node::Interpolant1DCubic { data, child } => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::Interpolant1DCubic { + data, + child: child_s, + }) + } + }, + Node::Interpolant1DCubicDeriv { data, child } => { + let child_s = simplify_node(arena, child, mode, memo, lens); + if child_s == child { + id + } else { + arena.alloc(Node::Interpolant1DCubicDeriv { + data, + child: child_s, + }) + } + }, + Node::InterpolantNd { data, children } => { + let children_s: Vec = children + .iter() + .map(|&c| simplify_node(arena, c, mode, memo, lens)) + .collect(); + if children_s == children { + id + } else { + arena.alloc(Node::InterpolantNd { + data, + children: children_s, + }) + } + }, + Node::InterpolantNdPartial { + data, + children, + axis, + } => { + let children_s: Vec = children + .iter() + .map(|&c| simplify_node(arena, c, mode, memo, lens)) + .collect(); + if children_s == children { + id + } else { + arena.alloc(Node::InterpolantNdPartial { + data, + children: children_s, + axis, + }) + } + }, + + Node::Conditional { selector, branches } => { + let selector_s = simplify_node(arena, selector, mode, memo, lens); + let branches_s: Vec = branches + .iter() + .map(|&b| simplify_node(arena, b, mode, memo, lens)) + .collect(); + if selector_s == selector && branches_s == branches { + id + } else { + arena.alloc(Node::Conditional { + selector: selector_s, + branches: branches_s, + }) + } + }, + }; + + memo.insert(id, result); + result +} + +/// `x + 0 -> x`, `0 + x -> x`, `const + const -> const`. +fn simplify_add( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a + b)); + } + + // 0 + x -> x + if is_zero(arena, lhs_s) { + return rhs_s; + } + + // x + 0 -> x + if is_zero(arena, rhs_s) { + return lhs_s; + } + + arena.alloc(Node::Add(lhs_s, rhs_s)) +} + +/// `x - 0 -> x`, `0 - x -> -x`, `x - x -> 0`, `const - const -> const`. +fn simplify_sub( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a - b)); + } + + // x - 0 -> x, exact only from +0.0 (x = -0.0 would flip positive). + // Aggressive already normalises zero's sign, as `0 - x` below does. + if is_positive_zero(arena, rhs_s) || (mode == SimplifyMode::Aggressive && is_zero(arena, rhs_s)) + { + return lhs_s; + } + + // 0 - x -> -x, exact only from -0.0: `(+0.0) - (+0.0)` is `+0.0` where + // `-(+0.0)` is `-0.0`, so from +0.0 it is aggressive-only like `x - x`. + if is_negative_zero(arena, lhs_s) || (mode == SimplifyMode::Aggressive && is_zero(arena, lhs_s)) + { + return arena.alloc(Node::Neg(rhs_s)); + } + + // x - x -> 0 (aggressive only: Inf - Inf = NaN); zero keeps x's length + if mode == SimplifyMode::Aggressive && lhs_s == rhs_s { + let len = node_len(arena, lhs_s, lens); + return zero_of_len(arena, len); + } + + arena.alloc(Node::Sub(lhs_s, rhs_s)) +} + +/// `1 * x -> x`, `x * 1 -> x`, `-1 * x -> -x`, `const * const -> const`, and +/// `0 * x -> 0` in aggressive mode. +fn simplify_mul( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a * b)); + } + + // 1 * x -> x + if is_one(arena, lhs_s) { + return rhs_s; + } + + // x * 1 -> x + if is_one(arena, rhs_s) { + return lhs_s; + } + + // -1 * x -> -x + if is_neg_one(arena, lhs_s) { + return arena.alloc(Node::Neg(rhs_s)); + } + + // x * -1 -> -x + if is_neg_one(arena, rhs_s) { + return arena.alloc(Node::Neg(lhs_s)); + } + + // Aggressive mode only: 0 * x -> 0, x * 0 -> 0 + // The zero keeps the broadcast length of the product. + if mode == SimplifyMode::Aggressive && (is_zero(arena, lhs_s) || is_zero(arena, rhs_s)) { + let len = node_len(arena, lhs_s, lens).max(node_len(arena, rhs_s, lens)); + return zero_of_len(arena, len); + } + + arena.alloc(Node::Mul(lhs_s, rhs_s)) +} + +/// `x / 1 -> x`, `const / const -> const`, and `0 / x -> 0`, `x / x -> 1` in +/// aggressive mode. +fn simplify_div( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let lhs_s = simplify_node(arena, lhs, mode, memo, lens); + let rhs_s = simplify_node(arena, rhs, mode, memo, lens); + + // Constant folding (uses IEEE semantics, so 1/0 = inf, 0/0 = nan) + if let (Some(a), Some(b)) = (get_scalar(arena, lhs_s), get_scalar(arena, rhs_s)) { + return arena.alloc(Node::Scalar(a / b)); + } + + // x / 1 -> x + if is_one(arena, rhs_s) { + return lhs_s; + } + + // Aggressive mode only + if mode == SimplifyMode::Aggressive { + // 0 / x -> 0 (the zero keeps the broadcast length of the quotient) + if is_zero(arena, lhs_s) { + let len = node_len(arena, rhs_s, lens); + return zero_of_len(arena, len); + } + + // x / x -> 1 (scalar x only: a vector x / x is a ones vector) + if lhs_s == rhs_s && node_len(arena, lhs_s, lens) == 1 { + return arena.alloc(Node::Scalar(1.0)); + } + } + + arena.alloc(Node::Div(lhs_s, rhs_s)) +} + +/// `x^0 -> 1`, `x^1 -> x`, `1^x -> 1`, `const^const -> const`. +fn simplify_pow( + arena: &mut Arena, + base: NodeId, + exp: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let base_s = simplify_node(arena, base, mode, memo, lens); + let exp_s = simplify_node(arena, exp, mode, memo, lens); + + // Constant folding + if let (Some(a), Some(b)) = (get_scalar(arena, base_s), get_scalar(arena, exp_s)) { + return arena.alloc(Node::Scalar(a.powf(b))); + } + + // x^0 -> 1 (scalar base only: a vector x^0 is a ones vector) + if is_zero(arena, exp_s) && node_len(arena, base_s, lens) == 1 { + return arena.alloc(Node::Scalar(1.0)); + } + + // x^1 -> x + if is_one(arena, exp_s) { + return base_s; + } + + // 1^x -> 1 (scalar exponent only: result broadcasts to the exponent's length) + if is_one(arena, base_s) && node_len(arena, exp_s, lens) == 1 { + return base_s; + } + + // powf is ~30x an elementwise multiply, and a multiply/divide chain keeps + // the IEEE special cases (sign of zero, infinities, NaN) exact. + if let Some(e) = get_scalar(arena, exp_s) + && let Some(id) = lower_int_pow(arena, base_s, e) + { + return id; + } + + arena.alloc(Node::Pow(base_s, exp_s)) +} + +/// Lower `base^e` for e in {2, 3, 4, -1, -2} to a Mul/Div chain, else `None`. +#[allow(clippy::float_cmp)] // exact literal exponents only +fn lower_int_pow(arena: &mut Arena, base: NodeId, e: f64) -> Option { + if e == 2.0 { + return Some(arena.alloc(Node::Mul(base, base))); + } + if e == 3.0 { + let sq = arena.alloc(Node::Mul(base, base)); + return Some(arena.alloc(Node::Mul(sq, base))); + } + if e == 4.0 { + let sq = arena.alloc(Node::Mul(base, base)); + return Some(arena.alloc(Node::Mul(sq, sq))); + } + if e == -1.0 { + let one = arena.alloc(Node::Scalar(1.0)); + return Some(arena.alloc(Node::Div(one, base))); + } + if e == -2.0 { + // (1/x)*(1/x), not 1/(x*x): squaring first overflows to 0 for large + // finite x (e.g. 1e160), whereas this form stays representable like powf. + let one = arena.alloc(Node::Scalar(1.0)); + let inv = arena.alloc(Node::Div(one, base)); + return Some(arena.alloc(Node::Mul(inv, inv))); + } + None +} + +/// `-const -> const`, `--x -> x`. +fn simplify_neg( + arena: &mut Arena, + child: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let child_s = simplify_node(arena, child, mode, memo, lens); + + // Constant folding + if let Some(v) = get_scalar(arena, child_s) { + return arena.alloc(Node::Scalar(-v)); + } + + // --x -> x (double negation) + if let Node::Neg(inner) = arena.get(child_s) { + return *inner; + } + + arena.alloc(Node::Neg(child_s)) +} + +/// Helper for unary operations with constant folding +fn simplify_unary( + arena: &mut Arena, + child: NodeId, + mode: SimplifyMode, + memo: &mut NodeMap, + lens: &mut NodeMap, + fold_fn: F, + constructor: G, +) -> NodeId +where + F: Fn(f64) -> f64, + G: Fn(NodeId) -> Node, +{ + let child_s = simplify_node(arena, child, mode, memo, lens); + + // Constant folding + if let Some(v) = get_scalar(arena, child_s) { + return arena.alloc(Node::Scalar(fold_fn(v))); + } + + arena.alloc(constructor(child_s)) +} + +/// Common Subexpression Elimination +/// +/// Identifies structurally equivalent subtrees and reuses them. +/// Returns a new arena containing only unique nodes and the remapped root. +/// +/// This is important for derivative expressions which duplicate primal computations. +#[must_use] +pub fn cse(arena: &Arena, root: NodeId) -> (Arena, NodeId) { + use rustc_hash::FxHashMap; + use rustc_hash::FxHasher; + use std::hash::Hasher; + + use crate::node::structural_hash; + + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + let mut hash_to_candidates: FxHashMap> = FxHashMap::default(); + + let order = arena.topological_order(root); + + for old_id in order { + let node = arena.get(old_id); + let remapped = node.map_children(|c| { + old_to_new + .get(c) + .copied() + .expect("child precedes parent in topo order") + }); + let mut hasher = FxHasher::default(); + structural_hash(&remapped, &mut hasher, |c| c); + let hash = hasher.finish(); + + let canonical = hash_to_candidates.get(&hash).and_then(|cands| { + cands + .iter() + .copied() + .find(|&existing| new_arena.get(existing) == &remapped) + }); + + if let Some(id) = canonical { + old_to_new.insert(old_id, id); + } else { + let new_id = new_arena.alloc(remapped); + old_to_new.insert(old_id, new_id); + hash_to_candidates.entry(hash).or_default().push(new_id); + } + } + + let new_root = old_to_new + .get(root) + .copied() + .expect("root must be processed"); + (new_arena, new_root) +} + +/// The full compile-prep simplification pipeline: +/// aggressive simplify -> zero propagation -> CSE -> DCE. +/// +/// The single home for the pass order shared by primal compiles, tangent +/// tapes, jacobian prep, sensitivities and algebraic blocks. Takes the +/// arena by value, callers clone (or own) the arena they want rewritten. +#[must_use] +pub fn simplify_pipeline(mut arena: Arena, root: NodeId) -> (Arena, NodeId) { + use crate::zero_propagate::zero_propagate; + + let root = simplify_with_mode(&mut arena, root, SimplifyMode::Aggressive); + let (arena, root) = zero_propagate(&arena, root); + let (arena, root) = cse(&arena, root); + dce(&arena, root) +} + +/// Dead Code Elimination +/// +/// Removes nodes not reachable from the root. +/// Returns a new arena containing only reachable nodes and the remapped root. +#[must_use] +pub fn dce(arena: &Arena, root: NodeId) -> (Arena, NodeId) { + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + + // Get topological order (only visits reachable nodes) + let order = arena.topological_order(root); + + for old_id in order { + let node = arena.get(old_id); + let remapped_node = node.map_children(|c| { + old_to_new + .get(c) + .copied() + .expect("child precedes parent in topo order") + }); + let new_id = new_arena.alloc(remapped_node); + old_to_new.insert(old_id, new_id); + } + + let new_root = old_to_new + .get(root) + .copied() + .expect("root must be processed"); + (new_arena, new_root) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + #[test] + fn simplify_pipeline_matches_manual_pass_order() { + use crate::zero_propagate::zero_propagate; + + // decorated: (sin(y0) * 1.0) + 0.0, must reduce like the manual + // simplify -> zero_propagate -> cse -> dce chain + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s = arena.alloc(Node::Sin(y0)); + let one = arena.alloc(Node::Scalar(1.0)); + let zero = arena.alloc(Node::Scalar(0.0)); + let m = arena.alloc(Node::Mul(s, one)); + let decorated = arena.alloc(Node::Add(m, zero)); + + let (pipe_arena, pipe_root) = simplify_pipeline(arena.clone(), decorated); + + let mut manual = arena; + let r = simplify_with_mode(&mut manual, decorated, SimplifyMode::Aggressive); + let (za, r) = zero_propagate(&manual, r); + let (ca, r) = cse(&za, r); + let (ma, r) = dce(&ca, r); + + assert_eq!(pipe_arena.len(), ma.len()); + assert_eq!(pipe_root, r); + assert_eq!(pipe_arena.get(pipe_root), ma.get(r)); + } + + #[test] + fn test_fold_multiply_by_one() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Mul(x, one)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_fold_one_times_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Mul(one, x)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_fold_add_zero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Add(x, zero)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_fold_zero_plus_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Add(zero, x)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins value-exactness + fn conservative_simplify_may_normalise_the_sign_of_zero() { + // `-0.0 + 0.0` is `+0.0`; folding the `+ 0` away yields `-0.0`. Both + // are zero, so the fold is permitted. + use crate::eval::CompiledExpr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_y = arena.alloc(Node::Neg(y)); + let zero = arena.alloc(Node::Scalar(0.0)); + let root = arena.alloc(Node::Add(neg_y, zero)); + + let mut folded = arena.clone(); + let folded_root = simplify_with_mode(&mut folded, root, SimplifyMode::Conservative); + + let before = CompiledExpr::new(&arena, root); + let after = CompiledExpr::new(&folded, folded_root); + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.0], &[], &[])[0]; + let b = after.eval(&mut s2, 0.0, &[0.0], &[], &[])[0]; + + assert_eq!(a, b, "values must be equal"); + assert!(a == 0.0 && b == 0.0, "both must be zero"); + // The sign is explicitly NOT guaranteed; assert only that we know which + // way it went, so a future change to the contract fails loudly here. + assert!(a.is_sign_positive(), "unfolded -0.0 + 0.0 is +0.0"); + assert!(b.is_sign_negative(), "folded form keeps -0.0"); + } + + #[test] + fn test_fold_sub_zero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Sub(x, zero)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + /// `(+0.0) - x -> -x` is sign-of-zero lossy, so conservative mode keeps the + /// `Sub`; aggressive mode still folds it. + #[test] + fn test_fold_zero_minus_x() { + let build = || { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Sub(zero, x)); + (arena, x, expr) + }; + + let (mut arena, _x, expr) = build(); + let conservative = simplify(&mut arena, expr); + assert!( + matches!(arena.get(conservative), Node::Sub(_, _)), + "conservative mode must not turn (+0) - x into -x: for x = +0 that \ + yields -0 where the subtraction yields +0" + ); + + let (mut arena, x, expr) = build(); + let aggressive = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + match arena.get(aggressive) { + Node::Neg(inner) => assert_eq!(*inner, x), + other => panic!("aggressive mode should fold to Neg, got {other:?}"), + } + } + + /// `-0.0 - x` really is `-x` for every `x`, so that one folds even in + /// conservative mode. + #[test] + fn test_fold_negative_zero_minus_x_is_exact() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_zero = arena.alloc(Node::Scalar(-0.0)); + let expr = arena.alloc(Node::Sub(neg_zero, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Neg(inner) => assert_eq!(*inner, x), + other => panic!("expected Neg node, got {other:?}"), + } + } + + /// `x - (-0.0)` is `x + 0.0`, which loses a negative-zero `x`, so only the + /// `+0.0` form may be folded away. + #[test] + fn test_x_minus_negative_zero_is_not_folded() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_zero = arena.alloc(Node::Scalar(-0.0)); + let expr = arena.alloc(Node::Sub(x, neg_zero)); + + let simplified = simplify(&mut arena, expr); + + assert!( + matches!(arena.get(simplified), Node::Sub(_, _)), + "x - (-0.0) must keep its Sub: for x = -0 it evaluates to +0" + ); + } + + /// Aggressive mode ignores the sign of zero already (`x - x -> 0` below + /// treats `Inf - Inf` as an acceptable NaN loss), so it should fold a + /// vector `x - (-0.0)` too, matching `0 - x`'s existing aggressive-only + /// arm. This also covers the broadcast shape (`x` a vector, the zero a + /// scalar) that `simplify` folds without a shape guard. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins value-exactness + fn test_aggressive_fold_vector_minus_negative_zero() { + use crate::eval::CompiledExpr; + + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let neg_zero = arena.alloc(Node::Scalar(-0.0)); + let expr = arena.alloc(Node::Sub(x, neg_zero)); + + let mut folded_arena = arena.clone(); + let simplified = simplify_with_mode(&mut folded_arena, expr, SimplifyMode::Aggressive); + assert_eq!(simplified, x); + + // Tape shape alone doesn't prove the fold is value-preserving; check + // the trade is confined to a zero element's sign, others untouched. + let before = CompiledExpr::new(&arena, expr); + let after = CompiledExpr::new(&folded_arena, simplified); + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let y = [-0.0, 3.5]; + let unfolded = before.eval(&mut s1, 0.0, &y, &[], &[]).to_vec(); + let folded = after.eval(&mut s2, 0.0, &y, &[], &[]).to_vec(); + + assert_eq!(unfolded, folded, "values must be equal element-wise"); + assert!( + unfolded[0].is_sign_positive(), + "x - (-0.0) turns -0.0 positive" + ); + assert!(folded[0].is_sign_negative(), "the fold keeps -0.0's sign"); + assert_eq!( + unfolded[1], 3.5, + "a nonzero element passes through untouched" + ); + } + + /// Regression for the shrunk `simplify_conservative_preserves_eval` + /// counterexample: `sin((2 - 2) - (y - y))` is `sin(+0.0)`, and conservative + /// simplification must not turn it into `sin(-(y - y))` = `sin(-0.0)`. + #[test] + fn test_conservative_simplify_preserves_sign_of_zero() { + use crate::eval::CompiledExpr; + + let mut arena = Arena::new(); + let two_a = arena.alloc(Node::Scalar(2.0)); + let two_b = arena.alloc(Node::Scalar(2.0)); + let const_zero = arena.alloc(Node::Sub(two_a, two_b)); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y_again = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let state_zero = arena.alloc(Node::Sub(y, y_again)); + let difference = arena.alloc(Node::Sub(const_zero, state_zero)); + let root = arena.alloc(Node::Sin(difference)); + + let before = CompiledExpr::new(&arena, root); + let mut arena_copy = arena.clone(); + let simplified = simplify(&mut arena_copy, root); + let after = CompiledExpr::new(&arena_copy, simplified); + + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.1], &[], &[])[0]; + let b = after.eval(&mut s2, 0.0, &[0.1], &[], &[])[0]; + assert_eq!( + a.to_bits(), + b.to_bits(), + "conservative simplify changed the sign of zero: {a} vs {b}" + ); + assert_eq!(a.to_bits(), 0.0_f64.to_bits(), "expected +0.0"); + } + + #[test] + fn test_conservative_does_not_fold_x_minus_x() { + // Conservative mode should NOT fold x - x -> 0 because Inf - Inf = NaN + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Sub(x, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Sub(_, _) => {}, + _ => panic!("Conservative mode should not simplify x - x"), + } + } + + #[test] + fn test_aggressive_folds_x_minus_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Sub(x, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v).abs() < f64::EPSILON), + _ => panic!("Aggressive mode should simplify x - x to 0"), + } + } + + #[test] + fn test_fold_div_by_one() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Div(x, one)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_fold_neg_one_times_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_one = arena.alloc(Node::Scalar(-1.0)); + let expr = arena.alloc(Node::Mul(neg_one, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Neg(inner) => assert_eq!(*inner, x), + _ => panic!("Expected Neg node"), + } + } + + #[test] + fn test_fold_pow_zero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Pow(x, zero)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 1.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(1.0)"), + } + } + + #[test] + fn test_fold_pow_one() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Pow(x, one)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_fold_one_pow_x() { + let mut arena = Arena::new(); + let one = arena.alloc(Node::Scalar(1.0)); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Pow(one, x)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, one); + } + + #[test] + fn test_double_negation() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_x = arena.alloc(Node::Neg(x)); + let neg_neg_x = arena.alloc(Node::Neg(neg_x)); + + let simplified = simplify(&mut arena, neg_neg_x); + + assert_eq!(simplified, x); + } + + /// Evaluate the simplified form of `y0 ^ e` at `y`. + fn eval_pow_of_state(e: f64, y: f64) -> f64 { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let exp = arena.alloc(Node::Scalar(e)); + let expr = arena.alloc(Node::Pow(x, exp)); + let simplified = simplify(&mut arena, expr); + let compiled = crate::eval::CompiledExpr::new(&arena, simplified); + let mut scratch = vec![0.0; compiled.scratch_len()]; + compiled.eval(&mut scratch, 0.0, &[y], &[], &[])[0] + } + + #[test] + fn test_int_pow_lowers_to_mul_chain() { + // x^2 becomes Mul(x, x), not a runtime powf. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Pow(x, two)); + let simplified = simplify(&mut arena, expr); + match arena.get(simplified) { + Node::Mul(a, b) => { + assert_eq!(*a, x); + assert_eq!(*b, x); + }, + other => panic!("expected Mul(x, x), got {other:?}"), + } + } + + #[test] + fn test_int_pow_chain_values() { + for e in [2.0, 3.0, 4.0, -1.0, -2.0] { + for y in [0.7, -1.3, 2.5] { + let got = eval_pow_of_state(e, y); + let want = y.powf(e); + assert!( + (got - want).abs() <= 2.0 * f64::EPSILON * want.abs(), + "y={y}, e={e}: chain {got} != powf {want}" + ); + } + } + } + + #[test] + fn test_int_pow_chain_preserves_special_cases() { + // The chains must keep IEEE special cases (the reason they are safe + // in Conservative mode): signed zero, infinities, NaN. + let cases = [ + (2.0, f64::NEG_INFINITY), + (3.0, f64::NEG_INFINITY), + (3.0, -0.0), + (-1.0, 0.0), + (-1.0, -0.0), + (-1.0, f64::INFINITY), + (-2.0, 0.0), + ]; + for (e, y) in cases { + let got = eval_pow_of_state(e, y); + let want = y.powf(e); + assert_eq!( + got.to_bits(), + want.to_bits(), + "y={y}, e={e}: chain {got} != powf {want}" + ); + } + // NaN propagates (payload bits are not portable, so assert NaN-ness). + for e in [2.0, -2.0] { + assert!(eval_pow_of_state(e, f64::NAN).is_nan()); + } + } + + #[test] + fn test_non_integer_pow_stays_pow() { + // No chain form: 0.5 and 1.3 must remain runtime Pow (sqrt rewrites + // would change pow(-inf, 0.5) semantics, so simplify never emits them). + for e in [0.5, 1.3, 5.0, -3.0] { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let exp = arena.alloc(Node::Scalar(e)); + let expr = arena.alloc(Node::Pow(x, exp)); + let simplified = simplify(&mut arena, expr); + match arena.get(simplified) { + Node::Pow(_, _) => {}, + other => panic!("x^{e} must stay Pow, got {other:?}"), + } + } + } + + #[test] + fn test_constant_fold_add() { + let mut arena = Arena::new(); + let two = arena.alloc(Node::Scalar(2.0)); + let three = arena.alloc(Node::Scalar(3.0)); + let expr = arena.alloc(Node::Add(two, three)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 5.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(5.0)"), + } + } + + #[test] + fn test_constant_fold_sub() { + let mut arena = Arena::new(); + let five = arena.alloc(Node::Scalar(5.0)); + let three = arena.alloc(Node::Scalar(3.0)); + let expr = arena.alloc(Node::Sub(five, three)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 2.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(2.0)"), + } + } + + #[test] + fn test_constant_fold_mul() { + let mut arena = Arena::new(); + let three = arena.alloc(Node::Scalar(3.0)); + let four = arena.alloc(Node::Scalar(4.0)); + let expr = arena.alloc(Node::Mul(three, four)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 12.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(12.0)"), + } + } + + #[test] + fn test_constant_fold_div() { + let mut arena = Arena::new(); + let twelve = arena.alloc(Node::Scalar(12.0)); + let four = arena.alloc(Node::Scalar(4.0)); + let expr = arena.alloc(Node::Div(twelve, four)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 3.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(3.0)"), + } + } + + #[test] + fn test_constant_fold_pow() { + let mut arena = Arena::new(); + let two = arena.alloc(Node::Scalar(2.0)); + let three = arena.alloc(Node::Scalar(3.0)); + let expr = arena.alloc(Node::Pow(two, three)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 8.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(8.0)"), + } + } + + #[test] + fn test_constant_fold_neg() { + let mut arena = Arena::new(); + let five = arena.alloc(Node::Scalar(5.0)); + let expr = arena.alloc(Node::Neg(five)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v + 5.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(-5.0)"), + } + } + + #[test] + fn test_constant_fold_unary_ops() { + let mut arena = Arena::new(); + + // Test sin(0) = 0 + let zero = arena.alloc(Node::Scalar(0.0)); + let sin_zero = arena.alloc(Node::Sin(zero)); + let simplified = simplify(&mut arena, sin_zero); + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + _ => panic!("Expected Scalar(0.0)"), + } + + // Test cos(0) = 1 + let cos_zero = arena.alloc(Node::Cos(zero)); + let simplified = simplify(&mut arena, cos_zero); + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 1.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(1.0)"), + } + + // Test exp(0) = 1 + let exp_zero = arena.alloc(Node::Exp(zero)); + let simplified = simplify(&mut arena, exp_zero); + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 1.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(1.0)"), + } + + // Test sqrt(4) = 2 + let four = arena.alloc(Node::Scalar(4.0)); + let sqrt_four = arena.alloc(Node::Sqrt(four)); + let simplified = simplify(&mut arena, sqrt_four); + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 2.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(2.0)"), + } + } + + #[test] + fn test_nested_simplification() { + // (x + 0) * 1 -> x + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Mul(x_plus_zero, one)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_nested_simplification_complex() { + // ((x * 1) + 0) - (y * 0) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let one = arena.alloc(Node::Scalar(1.0)); + + let x_times_one = arena.alloc(Node::Mul(x, one)); + let lhs = arena.alloc(Node::Add(x_times_one, zero)); + let y_times_zero = arena.alloc(Node::Mul(y, zero)); + let expr = arena.alloc(Node::Sub(lhs, y_times_zero)); + + // Conservative keeps `y * 0` as a Mul, so the Sub survives. + let simplified_conservative = simplify(&mut arena, expr); + match arena.get(simplified_conservative) { + Node::Sub(l, _r) => assert_eq!(*l, x), + _ => panic!("Expected Sub node in conservative mode"), + } + + let simplified_aggressive = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + assert_eq!(simplified_aggressive, x); + } + + #[test] + fn test_conservative_does_not_fold_zero_times_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Mul(zero, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Mul(_, _) => {}, + _ => panic!("Conservative mode should not simplify 0 * x"), + } + } + + #[test] + fn test_aggressive_folds_zero_times_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Mul(zero, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + _ => panic!("Aggressive mode should simplify 0 * x to 0"), + } + } + + #[test] + fn test_aggressive_folds_x_times_zero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Mul(x, zero)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + _ => panic!("Aggressive mode should simplify x * 0 to 0"), + } + } + + #[test] + fn test_aggressive_does_not_fold_tiny_scalar_coefficient() { + // 1e-17 is below f64::EPSILON but is a genuine nonzero coefficient; + // it must NOT be folded away as if it were exactly zero. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let tiny = arena.alloc(Node::Scalar(1e-17)); + let expr = arena.alloc(Node::Mul(tiny, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Mul(_, _) => {}, // expected: coefficient survives + other => panic!("tiny nonzero coefficient must survive; got {other:?}"), + } + } + + #[test] + fn test_conservative_does_not_fold_zero_div_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Div(zero, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Div(_, _) => {}, + _ => panic!("Conservative mode should not simplify 0 / x"), + } + } + + #[test] + fn test_aggressive_folds_zero_div_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Div(zero, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + _ => panic!("Aggressive mode should simplify 0 / x to 0"), + } + } + + #[test] + fn test_conservative_does_not_fold_x_div_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Div(x, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Div(_, _) => {}, + _ => panic!("Conservative mode should not simplify x / x"), + } + } + + #[test] + fn test_aggressive_folds_x_div_x() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Div(x, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Scalar(v) => assert!((*v - 1.0).abs() < f64::EPSILON), + _ => panic!("Aggressive mode should simplify x / x to 1"), + } + } + + #[test] + fn test_aggressive_zero_times_vector_preserves_shape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Mul(zero, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + assert_eq!(arena.get(simplified), &Node::ZeroVector { len: 3 }); + } + + #[test] + fn test_aggressive_vector_times_zero_preserves_shape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Mul(x, zero)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + assert_eq!(arena.get(simplified), &Node::ZeroVector { len: 3 }); + } + + #[test] + fn test_aggressive_vector_sub_self_preserves_shape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::Sub(x, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + assert_eq!(arena.get(simplified), &Node::ZeroVector { len: 3 }); + } + + #[test] + fn test_aggressive_zero_div_vector_preserves_shape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Div(zero, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + assert_eq!(arena.get(simplified), &Node::ZeroVector { len: 3 }); + } + + #[test] + fn test_aggressive_vector_div_self_not_folded_to_scalar() { + // x / x on a vector is a ones *vector*; without a shaped ones node + // the fold must not fire. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::Div(x, x)); + + let simplified = simplify_with_mode(&mut arena, expr, SimplifyMode::Aggressive); + + match arena.get(simplified) { + Node::Div(_, _) => {}, + other => panic!("vector x / x must stay Div, got {other:?}"), + } + } + + #[test] + fn test_vector_pow_zero_not_folded_to_scalar() { + // x^0 on a vector is a ones *vector*; the fold must not fire. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let expr = arena.alloc(Node::Pow(x, zero)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Pow(_, _) => {}, + other => panic!("vector x^0 must stay Pow, got {other:?}"), + } + } + + #[test] + fn test_one_pow_vector_not_folded_to_scalar() { + // 1^x broadcasts to len(x); folding to scalar 1 narrows the output. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Pow(one, x)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Pow(_, _) => {}, + other => panic!("1^x with vector x must stay Pow, got {other:?}"), + } + } + + #[test] + #[allow(clippy::float_cmp)] // exact fold/runtime agreement is the property under test + fn test_erf_fold_zero_is_exact() { + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let erf = arena.alloc(Node::Erf(zero)); + + let simplified = simplify(&mut arena, erf); + + match arena.get(simplified) { + Node::Scalar(v) => assert_eq!(*v, 0.0, "erf(0) must fold to exactly 0"), + other => panic!("Expected folded Scalar, got {other:?}"), + } + } + + #[test] + fn test_erf_fold_bitwise_matches_runtime() { + let mut arena = Arena::new(); + let c = arena.alloc(Node::Scalar(0.5)); + let erf = arena.alloc(Node::Erf(c)); + + let compiled = crate::eval::CompiledExpr::new(&arena, erf); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let runtime = compiled.eval(&mut scratch, 0.0, &[], &[], &[])[0]; + + let simplified = simplify(&mut arena, erf); + + match arena.get(simplified) { + Node::Scalar(v) => assert_eq!(v.to_bits(), runtime.to_bits()), + other => panic!("Expected folded Scalar, got {other:?}"), + } + } + + #[test] + #[allow(clippy::float_cmp)] // exact fold/runtime agreement is the property under test + fn test_sign_fold_zero_matches_runtime() { + // Runtime sign(0) = 0; f64::signum(0.0) is 1, so the fold must not use it. + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let sign = arena.alloc(Node::Sign(zero)); + + let simplified = simplify(&mut arena, sign); + + match arena.get(simplified) { + Node::Scalar(v) => assert_eq!(*v, 0.0, "sign(0) must fold to 0"), + other => panic!("Expected folded Scalar, got {other:?}"), + } + } + + #[test] + fn test_simplify_index() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 10 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Index { + child: x_plus_zero, + start: 0, + end: 5, + }); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Index { child, start, end } => { + assert_eq!(*child, x); + assert_eq!(*start, 0); + assert_eq!(*end, 5); + }, + _ => panic!("Expected Index node"), + } + } + + #[test] + fn test_simplify_concat() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Concat(vec![x_plus_zero, y])); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Concat(children) => { + assert_eq!(children.len(), 2); + assert_eq!(children[0], x); + assert_eq!(children[1], y); + }, + _ => panic!("Expected Concat node"), + } + } + + #[test] + fn test_memoization() { + // Ensure same subexpression is not processed twice + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + // Use x_plus_zero twice + let expr = arena.alloc(Node::Add(x_plus_zero, x_plus_zero)); + + let simplified = simplify(&mut arena, expr); + + // Should simplify to x + x + match arena.get(simplified) { + Node::Add(l, r) => { + assert_eq!(*l, x); + assert_eq!(*r, x); + }, + _ => panic!("Expected Add node"), + } + } + + #[test] + fn test_leaf_nodes_unchanged() { + let mut arena = Arena::new(); + + let scalar = arena.alloc(Node::Scalar(42.0)); + assert_eq!(simplify(&mut arena, scalar), scalar); + + let sv = arena.alloc(Node::StateVector { start: 0, end: 10 }); + assert_eq!(simplify(&mut arena, sv), sv); + + let time = arena.alloc(Node::Time); + assert_eq!(simplify(&mut arena, time), time); + + let param = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + assert_eq!(simplify(&mut arena, param), param); + } + + #[test] + fn test_cse_deduplicates_identical_subtrees() { + // (x + y) * (x + y) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum1 = arena.alloc(Node::Add(x, y)); + let sum2 = arena.alloc(Node::Add(x, y)); + let product = arena.alloc(Node::Mul(sum1, sum2)); + + let (new_arena, new_root) = cse(&arena, product); + + // 5 nodes in, 4 out: the two identical `x + y` subtrees collapse. + assert_eq!(new_arena.len(), 4); + + match new_arena.get(new_root) { + Node::Mul(l, r) => { + assert_eq!(l, r, "Both operands should be the same deduplicated node"); + }, + _ => panic!("Expected Mul node"), + } + } + + #[test] + fn test_cse_preserves_different_subtrees() { + // Build: (x + y) * (x - y) + // These are different, should not be merged + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum = arena.alloc(Node::Add(x, y)); + let diff = arena.alloc(Node::Sub(x, y)); + let product = arena.alloc(Node::Mul(sum, diff)); + + let (new_arena, new_root) = cse(&arena, product); + + // All 5 nodes should be preserved (x, y shared, but sum and diff are different) + assert_eq!(new_arena.len(), 5); + + match new_arena.get(new_root) { + Node::Mul(l, r) => { + assert_ne!(l, r, "Different operations should not be merged"); + }, + _ => panic!("Expected Mul node"), + } + } + + #[test] + fn test_cse_handles_scalars() { + // Build: 2.0 + 2.0 + // Both scalars are identical and should be deduplicated + let mut arena = Arena::new(); + let two1 = arena.alloc(Node::Scalar(2.0)); + let two2 = arena.alloc(Node::Scalar(2.0)); + let sum = arena.alloc(Node::Add(two1, two2)); + + let (new_arena, new_root) = cse(&arena, sum); + + // Should have: one scalar(2.0), one Add = 2 nodes + assert_eq!(new_arena.len(), 2); + + match new_arena.get(new_root) { + Node::Add(l, r) => { + assert_eq!(l, r, "Identical scalars should be deduplicated"); + }, + _ => panic!("Expected Add node"), + } + } + + #[test] + fn test_cse_nested_common_subexpressions() { + // Build: sin(x + y) + sin(x + y) + // Both sin(x + y) should be deduplicated + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum1 = arena.alloc(Node::Add(x, y)); + let sum2 = arena.alloc(Node::Add(x, y)); + let sin1 = arena.alloc(Node::Sin(sum1)); + let sin2 = arena.alloc(Node::Sin(sum2)); + let result = arena.alloc(Node::Add(sin1, sin2)); + + let (new_arena, _new_root) = cse(&arena, result); + + // Should have: x, y, Add(x,y), Sin(Add), Add(Sin,Sin) = 5 nodes + // (both sum1/sum2 collapse to one, both sin1/sin2 collapse to one) + assert_eq!(new_arena.len(), 5); + } + + #[test] + fn test_dce_removes_unreachable_nodes() { + // Build arena with some unreachable nodes + let mut arena = Arena::new(); + let _unreachable1 = arena.alloc(Node::Scalar(999.0)); // Not referenced + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let _unreachable2 = arena.alloc(Node::Scalar(888.0)); // Not referenced + let one = arena.alloc(Node::Scalar(1.0)); + let root = arena.alloc(Node::Add(x, one)); + + // Original arena has 5 nodes + assert_eq!(arena.len(), 5); + + let (new_arena, new_root) = dce(&arena, root); + + // DCE should remove unreachable nodes, keeping only x, one, and Add + assert_eq!(new_arena.len(), 3); + + // Verify structure is preserved + match new_arena.get(new_root) { + Node::Add(l, r) => { + match new_arena.get(*l) { + Node::StateVector { start: 0, end: 1 } => {}, + _ => panic!("Expected StateVector"), + } + match new_arena.get(*r) { + Node::Scalar(v) => assert!((*v - 1.0).abs() < f64::EPSILON), + _ => panic!("Expected Scalar(1.0)"), + } + }, + _ => panic!("Expected Add node"), + } + } + + #[test] + fn test_dce_preserves_all_reachable_nodes() { + // Build a tree where all nodes are reachable + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum = arena.alloc(Node::Add(x, y)); + let neg = arena.alloc(Node::Neg(sum)); + let root = arena.alloc(Node::Abs(neg)); + + assert_eq!(arena.len(), 5); + + let (new_arena, _new_root) = dce(&arena, root); + + // All nodes are reachable, so none should be removed + assert_eq!(new_arena.len(), 5); + } + + #[test] + fn test_dce_handles_diamond_pattern() { + // Build diamond: root = (x + y) * (x + y) where x and y are shared + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum = arena.alloc(Node::Add(x, y)); + let root = arena.alloc(Node::Mul(sum, sum)); // Same node referenced twice + + assert_eq!(arena.len(), 4); + + let (new_arena, new_root) = dce(&arena, root); + + // All 4 nodes are reachable + assert_eq!(new_arena.len(), 4); + + // Structure should be preserved with shared reference + match new_arena.get(new_root) { + Node::Mul(l, r) => { + assert_eq!(l, r, "Both operands should reference the same Add node"); + }, + _ => panic!("Expected Mul node"), + } + } + + #[test] + fn test_cse_then_dce_pipeline() { + let mut arena = Arena::new(); + let _garbage = arena.alloc(Node::Scalar(123.0)); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum1 = arena.alloc(Node::Add(x, y)); + let sum2 = arena.alloc(Node::Add(x, y)); + let _more_garbage = arena.alloc(Node::Scalar(456.0)); + let product = arena.alloc(Node::Mul(sum1, sum2)); + + // Original: 7 nodes (2 garbage + 2 duplicate sums + x + y + mul) + assert_eq!(arena.len(), 7); + + // CSE first (works on reachable nodes from root) + let (cse_arena, cse_root) = cse(&arena, product); + + // After CSE: x, y, sum (deduplicated), mul = 4 nodes + // Note: CSE only processes reachable nodes, so garbage is already gone + assert_eq!(cse_arena.len(), 4); + + // DCE (should be a no-op after CSE since CSE only copies reachable) + let (final_arena, _final_root) = dce(&cse_arena, cse_root); + assert_eq!(final_arena.len(), 4); + } + + #[test] + fn test_simplify_reduce_arg_select_zero_basis_folds() { + // basis all-zero -> selecting any element is 0 -> Scalar(0.0) + let mut arena = Arena::new(); + let picker = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let basis = arena.alloc(Node::ZeroVector { len: 3 }); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let simplified = simplify(&mut arena, node); + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + other => panic!("expected Scalar(0.0), got {other:?}"), + } + } + + #[test] + fn test_simplify_reduce_arg_select_scalar_zero_basis_folds() { + // width-1 case: basis is Scalar(0.0) -> fold to Scalar(0.0) + let mut arena = Arena::new(); + let picker = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let basis = arena.alloc(Node::Scalar(0.0)); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let simplified = simplify(&mut arena, node); + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + other => panic!("expected Scalar(0.0), got {other:?}"), + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/snapshot.rs b/packages/pybamm-rust/pybamm-core/src/snapshot.rs new file mode 100644 index 0000000000..9f49e41a1c --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/snapshot.rs @@ -0,0 +1,25 @@ +#![cfg(feature = "serialize")] + +use crate::arena::{Arena, NodeId}; +use crate::node::CsrData; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct DagSnapshot { + pub arena: Arena, + pub root: NodeId, + pub n_states: usize, + pub n_params: usize, + pub mass_matrix: Option, + pub model_name: String, +} + +impl DagSnapshot { + pub fn to_bytes(&self) -> Vec { + bincode::serialize(self).expect("DagSnapshot serialization failed") + } + + pub fn from_bytes(data: &[u8]) -> Self { + bincode::deserialize(data).expect("DagSnapshot deserialization failed") + } +} 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..33df3d9608 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/batch.rs @@ -0,0 +1,285 @@ +//! Solving many input sets concurrently. +//! +//! The one method here is a fan-out over [`PreparedSolver::solve`], never a +//! second implementation of a solve. That is what makes a batch bit-identical to +//! the serial loop by construction: one immutable tape shared through the `Arc`, +//! one fresh `Workspace` minted inside each `solve()`, and rayon only choosing +//! the order the independent calls run in. +//! +//! The [`SolveRequest`] is shared across the batch rather than per set, which +//! matches the callers: `DiffsolSolver` computes one output grid for every set, +//! `BaseSolver` refuses to run sets whose discontinuities differ, and the +//! payload flags come from the model. Only [`InputSet`] varies per set. +//! +//! Scheduling is the caller's: these run on the ambient rayon pool, so a caller +//! that wants a specific width wraps the call in `ThreadPool::install`. + +use rayon::prelude::*; + +use super::solve::{InputSet, PreparedSolver, SolveOutcome, SolveRequest}; +use crate::error::CoreError; + +/// Check that every per-set argument agrees on the batch width. +/// +/// A `&[InputSet]` cannot itself disagree, so this is for callers assembling the +/// sets from separate untyped columns — the FFI boundary, where a `y0` array and +/// an `inputs` array arrive with independent row counts. `y0_sens` is `None` when +/// no seeds were supplied. +pub const fn check_batch_widths( + y0: usize, + inputs: usize, + y0_sens: Option, +) -> Result<(), CoreError> { + if y0 != inputs { + return Err(CoreError::BatchWidths { y0, inputs }); + } + if let Some(got) = y0_sens + && got != y0 + { + return Err(CoreError::BatchSensWidth { got, expected: y0 }); + } + Ok(()) +} + +impl PreparedSolver { + /// Solve one trajectory per input set, concurrently. + /// + /// Every set answers the same `request`, so the batch axis costs one method + /// rather than one per payload combination. A set that fails to integrate + /// keeps its own `Err` at its own index, so a caller can report which sets + /// failed rather than collapsing the batch into one error. + pub fn solve_batch( + &self, + request: SolveRequest<'_>, + sets: &[InputSet<'_>], + ) -> Vec> { + sets.par_iter() + .map(|set| self.solve(request, *set)) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::Arena; + use crate::model::{CompiledModelOptions, ModelEvaluator}; + use crate::node::{CsrData, Node, Shape}; + + /// `dy/dt = -p * y` with an event at `y - 0.4`, so a set's trajectory ends + /// either at the root or at the final time depending on its own `p`. The + /// event makes the batch exercise the root-finding and wind-back path, the + /// part of a solve most likely to leak state between sets if any did. + /// Sensitivities are requested for `p`, so one fixture serves both the plain + /// and the sensitivity request. + fn build_decay_with_event() -> ModelEvaluator { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let rate = arena.alloc(Node::InputParameter { + name: "rate".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let neg_rate = arena.alloc(Node::Mul(neg, rate)); + let rhs = arena.alloc(Node::Mul(neg_rate, sv)); + + let threshold = arena.alloc(Node::Scalar(0.4)); + let event = arena.alloc(Node::Sub(sv, threshold)); + + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape { rows: 1, cols: 1 }, + }; + let mut model = ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 1, + 1, + CompiledModelOptions::new().with_sensitivities(&[0]), + ); + model.add_event(&arena, event); + model + } + + fn prepared() -> PreparedSolver { + PreparedSolver::new(build_decay_with_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, so it reaches the + /// event last and takes the most steps. + fn rates(n: usize) -> Vec> { + (0..n).map(|i| vec![0.4 + 0.35 * (i as f64)]).collect() + } + + /// One set per rate, all starting from the same state. + 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) + } + + /// Every field a batch is expected to reproduce exactly from the serial + /// solve of the same set, payload included. + // Bit-identical is the property under test, so the comparisons are exact by + // design rather than by omission. + #[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.flag, want.flag, "{context}: flags differ"); + assert_eq!(got.t_event, want.t_event, "{context}: root times differ"); + assert_eq!(got.y_event, want.y_event, "{context}: root states differ"); + assert_eq!( + got.statistics.number_of_steps, want.statistics.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)")); + } + } + } + + /// The batch axis is one method for every payload combination, so the + /// sensitivity blocks have to survive the fan-out as exactly as the states + /// do — the property the four separate batch methods each needed their own + /// test for. + #[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] + // Result ordering is the property under test; the values are compared + // exactly for the same reason as above. + #[allow(clippy::float_cmp)] + fn results_follow_input_order_under_heterogeneous_cost() { + let solver = prepared(); + let times = t_eval(); + // Widely spread rates, so a batch that returned completion order rather + // than input order would reverse them. + let rates: Vec> = vec![vec![0.2], vec![1.0], vec![3.0], vec![9.0]]; + let y0 = [1.0]; + let sets = sets(&y0, &rates); + let request = SolveRequest::new(×); + + let batched = run_in_pool(4, || solver.solve_batch(request, &sets)); + + for (i, (got, set)) in batched.into_iter().zip(&sets).enumerate() { + let got = got.expect("set failed"); + let want = solver.solve(request, *set).expect("reference solve failed"); + assert_eq!(got.t, want.t, "set {i} landed out of order"); + assert_eq!(got.y, want.y, "set {i} landed out of order"); + } + } + + #[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 is handed the wrong input width, the one per-set failure a test + // can provoke without depending on how the integrator diverges. + 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"); + } + } + } + + #[test] + fn mismatched_batch_widths_are_rejected_before_any_solve() { + let err = check_batch_widths(2, 1, None).expect_err("mismatched widths accepted"); + assert!(matches!(err, CoreError::BatchWidths { y0: 2, inputs: 1 })); + + let err = check_batch_widths(2, 2, Some(1)).expect_err("mismatched seed count accepted"); + assert!(matches!( + err, + CoreError::BatchSensWidth { + got: 1, + expected: 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..7382213694 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/equations.rs @@ -0,0 +1,193 @@ +//! The equation bundle diffsol solves. +//! +//! [`Equations`] owns what a solve shares: the compiled model, workspace, this solve's +//! parameter values and its sensitivity-parameter indices. It mints a fresh +//! operator view whenever diffsol asks for one, because the operator traits +//! borrow for the length of a call rather than for the solve. + +use std::cell::RefCell; +use std::rc::Rc; +use std::sync::Arc; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + FaerContext, OdeEquations, OdeEquationsImplicitSens, OdeEquationsRef, Op, VectorHost, +}; + +use super::FaerSparsity; +use super::init::InitOp; +use super::mass::MassOp; +use super::observable::ObservableOp; +use super::reset::ResetOp; +use super::rhs::RhsOp; +use crate::model::{CompiledModel, Workspace}; +use crate::observable::ObservableKind; + +/// Local ODE-equations container implementing diffsol's public `OdeEquations` +/// trait. +/// +/// diffsol mints an operator view per callback invocation — `eqn.rhs()` on every +/// residual call, and again for every `nstates()` — so the views borrow from +/// here instead of owning: a mint is a few pointer copies, and the compiler, +/// not a refcount, keeps this alive for as long as a view exists. +/// +/// `Arc` marks the state shared with the `PreparedSolver` across solves, which +/// is `Send + Sync`; everything built per solve is owned outright. +pub struct Equations { + pub(crate) compiled: Arc, + /// Shared with the caller, which evaluates output tapes against the same + /// scratch between steps on the output-variable path. + pub(crate) ws: Rc>, + pub(crate) params: Vec, + /// Global parameter index of each sensitivity column; the identity on the + /// ordinary no-sensitivity path. + pub(crate) sens_params: Arc<[usize]>, + pub(crate) y0: Vec, + /// `dy0/dp`, column-major `n_states x sens_params.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, shared with every `MassOp` view. + pub(crate) mass_csc_values: Arc<[f64]>, + pub(crate) context: FaerContext, + pub(crate) n_states: usize, + pub(crate) n_event_outputs: usize, + pub(crate) n_outputs: usize, + pub(crate) with_output: bool, +} + +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.n_states) + // sens_params is the identity 0..n_params on the plain path, so this + // is nparams(), not a sensitivity-subset count, until sens is active. + .field("nparams", &self.sens_params.len()) + .field("n_event_outputs", &self.n_event_outputs) + .field("n_outputs", &self.n_outputs) + .field("with_output", &self.with_output) + .finish_non_exhaustive() + } +} + +impl Equations { + /// Mint an operator view onto one observable family. + fn observable_op(&self, kind: ObservableKind) -> ObservableOp<'_> { + ObservableOp { + compiled: &self.compiled, + ws: &self.ws, + inputs: &self.params, + sens_params: &self.sens_params, + n_states: self.n_states, + kind, + n_out: match kind { + ObservableKind::Outputs => self.n_outputs, + ObservableKind::Events => self.n_event_outputs, + }, + context: self.context, + } + } +} + +impl Op for Equations { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.n_states + } + fn nout(&self) -> usize { + if self.with_output && self.n_outputs > 0 { + self.n_outputs + } else { + self.n_states + } + } + fn nparams(&self) -> usize { + self.sens_params.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>; + type Out = ObservableOp<'a>; + /// Never constructed: `reset()` always returns `None` for `PyBaMM` + /// (events terminate, never reset). `ResetOp` satisfies the bound. + type Reset = ResetOp; +} + +impl OdeEquations for Equations { + fn rhs(&self) -> RhsOp<'_> { + RhsOp { + compiled: &self.compiled, + ws: &self.ws, + inputs: &self.params, + sens_params: &self.sens_params, + jac_sparsity: &self.jac_sparsity, + n_states: self.n_states, + context: self.context, + } + } + + fn mass(&self) -> Option> { + Some(MassOp { + compiled: &self.compiled, + ws: &self.ws, + sparsity: &self.mass_sparsity, + csc_values: &self.mass_csc_values, + n_states: self.n_states, + context: self.context, + }) + } + + fn root(&self) -> Option> { + (self.n_event_outputs > 0).then(|| self.observable_op(ObservableKind::Events)) + } + + fn out(&self) -> Option> { + (self.with_output && self.n_outputs > 0) + .then(|| self.observable_op(ObservableKind::Outputs)) + } + + fn init(&self) -> InitOp<'_> { + InitOp { + y0: &self.y0, + y0_sens: &self.y0_sens, + n_states: self.n_states, + n_sens_params: self.sens_params.len(), + context: self.context, + } + } + + /// Splice the `k` sensitivity values into the full input vector at their + /// global indices, leaving the carried-but-not-differentiated inputs alone. + fn set_params(&mut self, p: &FaerVec) { + for (&global, &v) in self.sens_params.iter().zip(p.as_slice()) { + self.params[global] = v; + } + } + + /// Gather the `k` sensitivity entries out of the full input vector. diffsol + /// sizes `p` by `nparams()`, so this cannot be a wholesale copy. + fn get_params(&self, p: &mut FaerVec) { + for (dst, &global) in p.as_mut_slice().iter_mut().zip(self.sens_params.iter()) { + *dst = self.params[global]; + } + } +} + +// Pins that `Equations` satisfies the forward-sensitivity bounds, which +// diffsol's blanket `impl OdeEquationsImplicitSens` needs. +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..56409a104d --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/init.rs @@ -0,0 +1,143 @@ +//! Initial-condition operator. +//! +//! `y0` is supplied per solve, and so is its parameter derivative `dy0/dp`: a +//! parameter that feeds an initial condition seeds the sensitivity system with +//! a non-zero column, exactly as IDAS does. An empty seed means `dy0/dp = 0`. + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ConstantOp, ConstantOpSens, FaerContext, Op, VectorHost}; + +/// Initial-condition operator: returns the stored y0 at solve time, borrowed +/// from the [`Equations`](super::equations::Equations) that owns the solve. +#[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_states: usize, + 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.n_states + } + 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) { + y.as_mut_slice().copy_from_slice(self.y0); + } +} + +impl ConstantOpSens for InitOp<'_> { + /// `dy0/dp · v`. diffsol drives this with unit vectors to read off one + /// column at a time, then its consistent-IC augmentation overwrites the + /// algebraic rows; the differential rows are taken as seeded here. + fn sens_mul_inplace(&self, _t: f64, v: &FaerVec, y: &mut FaerVec) { + let out = y.as_mut_slice(); + out.fill(0.0); + if self.y0_sens.is_empty() { + return; + } + for (col, &v_col) in v.as_slice().iter().enumerate() { + if v_col == 0.0 { + continue; + } + let offset = col * self.n_states; + for (out_i, &s_i) in out + .iter_mut() + .zip(&self.y0_sens[offset..offset + self.n_states]) + { + *out_i += v_col * s_i; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use diffsol::{Context, Vector}; + + fn init_op<'a>(y0: &'a [f64], y0_sens: &'a [f64], n_sens_params: usize) -> InitOp<'a> { + InitOp { + y0, + y0_sens, + n_states: y0.len(), + 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.n_states); + op.sens_mul_inplace(0.0, &v, &mut y); + y.as_slice().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 unit_directions_read_off_columns() { + // Column-major 3x2: column 0 = [1,2,3], column 1 = [4,5,6]. + 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); + assert_eq!(sens_column(&op, 0), vec![1.0, 2.0, 3.0]); + assert_eq!(sens_column(&op, 1), vec![4.0, 5.0, 6.0]); + } + + #[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!(y.as_slice(), &[2.0 - 4.0, 4.0 - 5.0, 6.0 - 6.0]); + } + + #[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!(y.as_slice(), &[1.0, 2.0, 3.0]); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/linear.rs b/packages/pybamm-rust/pybamm-core/src/solver/linear.rs new file mode 100644 index 0000000000..e9c1d68819 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/linear.rs @@ -0,0 +1,380 @@ +//! Allocation-free faer sparse-LU linear solver. +//! +//! diffsol's stock `FaerSparseLU` allocates on every linearisation refresh (a +//! clone of the symbolic structure, a fresh `NumericLu`, and its factorisation +//! workspace) and on every triangular solve (the solve scratch, hundreds of +//! times per solve inside Newton). This solver keeps the symbolic +//! factorisation, the numeric storage, and both workspaces alive for its whole +//! lifetime, so a refresh runs only the numeric kernel and a Newton iteration +//! runs only the triangular kernels. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, LaError, Matrix, MatrixCommon, Vector}; +use diffsol_la::error::LinearSolverError; +use diffsol_la::{LinearOp, LinearSolver}; +use dyn_stack::{MemBuffer, MemStack}; +use faer::reborrow::Reborrow; +use faer::sparse::linalg::lu::{ + LuRef, LuSymbolicParams, NumericLu, SymbolicLu, factorize_symbolic_lu, +}; +use faer::{Par, Spec}; + +/// Problem sizes here are a few thousand states at most, where rayon dispatch +/// costs more than it saves. +const PAR: Par = Par::Seq; + +/// A [`LinearSolver`] over faer's sparse LU that reuses every buffer. +pub struct ReusedFaerLu { + symbolic: Option>, + numeric: NumericLu, + matrix: Option>, + factorize_buf: Option, + /// `solve_in_place` takes `&self`, so its scratch hides behind a `RefCell`. + solve_buf: RefCell>, + factorized: bool, +} + +impl Default for ReusedFaerLu { + fn default() -> Self { + Self { + symbolic: None, + numeric: NumericLu::new(), + matrix: None, + factorize_buf: None, + solve_buf: RefCell::new(None), + factorized: false, + } + } +} + +impl std::fmt::Debug for ReusedFaerLu { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ReusedFaerLu") + .field("factorized", &self.factorized) + .finish_non_exhaustive() + } +} + +impl LinearSolver> for ReusedFaerLu { + fn set_sparsity< + C: LinearOp, M = FaerSparseMat, C = FaerContext>, + >( + &mut self, + op: &C, + ) { + let matrix = FaerSparseMat::::new_from_sparsity( + op.nrows(), + op.ncols(), + op.sparsity(), + *op.context(), + ); + let symbolic = + factorize_symbolic_lu(matrix.inner().symbolic(), LuSymbolicParams::default()) + .expect("Failed to create symbolic LU"); + self.factorize_buf = Some(MemBuffer::new( + symbolic.factorize_numeric_lu_scratch::(PAR, Spec::default()), + )); + self.solve_buf = RefCell::new(Some(MemBuffer::new( + symbolic.solve_in_place_scratch::(1, PAR), + ))); + self.symbolic = Some(symbolic); + self.matrix = Some(matrix); + self.factorized = false; + } + + fn set_linearisation< + C: LinearOp, M = FaerSparseMat, C = FaerContext>, + >( + &mut self, + op: &C, + ) { + let matrix = self.matrix.as_mut().expect("Matrix not set"); + op.matrix_inplace(matrix); + let symbolic = self.symbolic.as_ref().expect("Sparsity not set"); + let stack = MemStack::new(self.factorize_buf.as_mut().expect("Sparsity not set")); + symbolic + .factorize_numeric_lu( + &mut self.numeric, + matrix.inner().rb(), + PAR, + stack, + Spec::default(), + ) + .expect("Failed to factorise matrix"); + self.factorized = true; + } + + fn solve_in_place(&self, x: &mut FaerVec) -> Result<(), LaError> { + if !self.factorized { + return Err(LinearSolverError::LuNotInitialized.into()); + } + let symbolic = self.symbolic.as_ref().expect("factorized implies symbolic"); + let lu = LuRef::new_unchecked(symbolic, &self.numeric); + let mut buf = self.solve_buf.borrow_mut(); + let stack = MemStack::new(buf.as_mut().expect("factorized implies solve scratch")); + lu.solve_in_place_with_conj(faer::Conj::No, x.inner_mut().as_mat_mut(), PAR, stack); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use diffsol::{Context, VectorHost}; + + /// The pattern is asymmetric so the solve exercises both permutations. + struct TestOp { + matrix: FaerSparseMat, + context: FaerContext, + } + + impl TestOp { + fn new() -> Self { + let context = FaerContext::default(); + let indices = vec![(0, 0), (0, 2), (1, 1), (2, 1), (2, 2)]; + let values = vec![2.0, 1.0, 3.0, 4.0, 5.0]; + let matrix = FaerSparseMat::try_from_triplets(3, 3, indices, values, context) + .expect("bad triplets"); + Self { matrix, context } + } + } + + impl LinearOp for TestOp { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nrows(&self) -> usize { + 3 + } + fn ncols(&self) -> usize { + 3 + } + fn context(&self) -> &FaerContext { + &self.context + } + fn matrix_inplace(&self, y: &mut FaerSparseMat) { + y.copy_from(&self.matrix); + } + fn sparsity(&self) -> Option< as Matrix>::Sparsity> { + self.matrix.sparsity().map(|s| s.to_owned().unwrap()) + } + } + + #[test] + fn matches_stock_faer_sparse_lu() { + let op = TestOp::new(); + + let mut reused = ReusedFaerLu::default(); + reused.set_sparsity(&op); + reused.set_linearisation(&op); + + let mut stock = diffsol::FaerSparseLU::::default(); + stock.set_sparsity(&op); + stock.set_linearisation(&op); + + let ctx = FaerContext::default(); + let mut x_reused = ctx.vector_zeros::>(3); + x_reused.as_mut_slice().copy_from_slice(&[5.0, 6.0, 23.0]); + let mut x_stock = x_reused.clone(); + + LinearSolver::solve_in_place(&reused, &mut x_reused).unwrap(); + LinearSolver::solve_in_place(&stock, &mut x_stock).unwrap(); + + // Same factorisation algorithm, same arithmetic: bitwise equality. + assert_eq!(x_reused.as_slice(), x_stock.as_slice()); + // A x = [5, 6, 23] with the matrix above has x = [1, 2, 3]. + for (got, want) in x_reused.as_slice().iter().zip(&[1.0, 2.0, 3.0]) { + assert!((got - want).abs() < 1e-12, "got {got}, want {want}"); + } + } + + #[test] + fn refactorisation_reuses_buffers_and_stays_correct() { + let op = TestOp::new(); + let mut solver = ReusedFaerLu::default(); + solver.set_sparsity(&op); + + let ctx = FaerContext::default(); + for _ in 0..3 { + solver.set_linearisation(&op); + let mut x = ctx.vector_zeros::>(3); + x.as_mut_slice().copy_from_slice(&[5.0, 6.0, 23.0]); + LinearSolver::solve_in_place(&solver, &mut x).unwrap(); + for (got, want) in x.as_slice().iter().zip(&[1.0, 2.0, 3.0]) { + assert!((got - want).abs() < 1e-12, "got {got}, want {want}"); + } + } + } + + #[test] + fn solve_before_linearisation_is_an_error() { + let op = TestOp::new(); + let mut solver = ReusedFaerLu::default(); + solver.set_sparsity(&op); + let ctx = FaerContext::default(); + let mut x = ctx.vector_zeros::>(3); + assert!(LinearSolver::solve_in_place(&solver, &mut x).is_err()); + } + + /// An op over an arbitrary square CSC matrix built from triplets. + struct TripletOp { + matrix: FaerSparseMat, + context: FaerContext, + n: usize, + } + + impl TripletOp { + fn new(n: usize, triplets: &[(usize, usize, f64)]) -> Self { + let context = FaerContext::default(); + let indices: Vec<(usize, usize)> = triplets.iter().map(|&(r, c, _)| (r, c)).collect(); + let values: Vec = triplets.iter().map(|&(.., v)| v).collect(); + let matrix = FaerSparseMat::try_from_triplets(n, n, indices, values, context) + .expect("bad triplets"); + Self { matrix, context, n } + } + } + + impl LinearOp for TripletOp { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nrows(&self) -> usize { + self.n + } + fn ncols(&self) -> usize { + self.n + } + fn context(&self) -> &FaerContext { + &self.context + } + fn matrix_inplace(&self, y: &mut FaerSparseMat) { + y.copy_from(&self.matrix); + } + fn sparsity(&self) -> Option< as Matrix>::Sparsity> { + self.matrix.sparsity().map(|s| s.to_owned().unwrap()) + } + } + + #[test] + fn set_sparsity_again_rebuilds_for_a_new_pattern_and_size() { + // One solver instance must survive a sparsity change: the symbolic + // factorisation, matrix, and both scratch buffers are all size-bound. + let mut solver = ReusedFaerLu::default(); + let ctx = FaerContext::default(); + + let op3 = TestOp::new(); + solver.set_sparsity(&op3); + solver.set_linearisation(&op3); + let mut x = ctx.vector_zeros::>(3); + x.as_mut_slice().copy_from_slice(&[5.0, 6.0, 23.0]); + LinearSolver::solve_in_place(&solver, &mut x).unwrap(); + + // 4x4 diagonal: a different pattern, dimension, and pivot structure. + let op4 = TripletOp::new(4, &[(0, 0, 2.0), (1, 1, 4.0), (2, 2, 8.0), (3, 3, 16.0)]); + solver.set_sparsity(&op4); + // The old factorisation must not survive the sparsity change. + let mut stale = ctx.vector_zeros::>(4); + assert!(LinearSolver::solve_in_place(&solver, &mut stale).is_err()); + + solver.set_linearisation(&op4); + let mut x = ctx.vector_zeros::>(4); + x.as_mut_slice().copy_from_slice(&[2.0, 4.0, 8.0, 16.0]); + LinearSolver::solve_in_place(&solver, &mut x).unwrap(); + for (got, want) in x.as_slice().iter().zip(&[1.0, 1.0, 1.0, 1.0]) { + assert!((got - want).abs() < 1e-12, "got {got}, want {want}"); + } + } + + #[test] + fn dae_structured_matrix_matches_stock() { + // A DAE-shaped iteration matrix: tridiagonal differential block plus + // two non-diagonally-dominant algebraic rows, so pivoting must permute. + let triplets = vec![ + (0, 0, 4.0), + (0, 1, -1.0), + (1, 0, -1.0), + (1, 1, 4.0), + (1, 2, -1.0), + (2, 1, -1.0), + (2, 2, 4.0), + (2, 3, -1.0), + (3, 3, 1e-3), + (3, 4, 1.0), + (4, 0, 1.0), + (4, 4, -1.0), + ]; + let op = TripletOp::new(5, &triplets); + + let mut reused = ReusedFaerLu::default(); + reused.set_sparsity(&op); + let mut stock = diffsol::FaerSparseLU::::default(); + stock.set_sparsity(&op); + + let ctx = FaerContext::default(); + // Refactorise repeatedly, as Newton does, and compare every solve. + for k in 0..3 { + reused.set_linearisation(&op); + stock.set_linearisation(&op); + let mut b = ctx.vector_zeros::>(5); + let rhs: Vec = (0..5).map(|i| f64::from(k * 3 + i + 1)).collect(); + b.as_mut_slice().copy_from_slice(&rhs); + let mut b_stock = b.clone(); + LinearSolver::solve_in_place(&reused, &mut b).unwrap(); + LinearSolver::solve_in_place(&stock, &mut b_stock).unwrap(); + assert_eq!(b.as_slice(), b_stock.as_slice()); + } + } + + #[test] + fn singular_matrix_behaviour_matches_stock() { + // Newton can hit a singular iteration matrix mid-solve; panic or + // poison, the reused solver must match stock so diffsol recovers alike. + let singular = || TripletOp::new(2, &[(0, 0, 1.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 1.0)]); + + let reused_outcome = std::panic::catch_unwind(|| { + let op = singular(); + let mut solver = ReusedFaerLu::default(); + solver.set_sparsity(&op); + solver.set_linearisation(&op); + let ctx = FaerContext::default(); + let mut x = ctx.vector_zeros::>(2); + x.as_mut_slice().copy_from_slice(&[1.0, 2.0]); + LinearSolver::solve_in_place(&solver, &mut x).unwrap(); + x.as_slice() + .iter() + .map(|v| v.is_finite()) + .collect::>() + }); + let stock_outcome = std::panic::catch_unwind(|| { + let op = singular(); + let mut solver = diffsol::FaerSparseLU::::default(); + solver.set_sparsity(&op); + solver.set_linearisation(&op); + let ctx = FaerContext::default(); + let mut x = ctx.vector_zeros::>(2); + x.as_mut_slice().copy_from_slice(&[1.0, 2.0]); + LinearSolver::solve_in_place(&solver, &mut x).unwrap(); + x.as_slice() + .iter() + .map(|v| v.is_finite()) + .collect::>() + }); + + match (reused_outcome, stock_outcome) { + (Ok(reused), Ok(stock)) => assert_eq!(reused, stock), + (Err(_), Err(_)) => {}, + (reused, stock) => panic!( + "singular-matrix behaviour diverged from stock: reused ok={}, stock ok={}", + reused.is_ok(), + stock.is_ok() + ), + } + } +} 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..84179cb923 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/mass.rs @@ -0,0 +1,89 @@ +//! Mass matrix operator. +//! +//! `M` is constant and arrives from Python in CSR, so the operator only applies +//! it. A zero row is how an algebraic equation reaches the solver: it makes the +//! system a DAE in that row rather than an ODE. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, LinearOp, Matrix, Op, VectorHost}; + +use super::FaerSparsity; +use crate::model::{CompiledModel, Workspace}; + +/// Mass operator M (constant in t and p), borrowed from the +/// [`Equations`](super::equations::Equations) that owns the solve. +/// +/// diffsol mints one of these per callback invocation, so every field is a +/// reference into the equations: a mint copies pointers, never data. +pub struct MassOp<'a> { + pub compiled: &'a CompiledModel, + pub ws: &'a RefCell, + pub sparsity: &'a FaerSparsity, + /// M's values in the CSC order of `sparsity` (see `csr_mass_to_faer_csc`), + /// so `matrix_inplace` is a copy rather than diffsol's column probing. + pub csc_values: &'a [f64], + pub n_states: usize, + 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.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.n_states + } + fn nout(&self) -> usize { + self.n_states + } + fn nparams(&self) -> usize { + 0 + } + fn context(&self) -> &Self::C { + &self.context + } +} + +impl LinearOp for MassOp<'_> { + /// y = M @ x + beta * y. `mass_action` computes M @ x with no beta, so we + /// fold beta ourselves via the workspace `mv_buffer` when beta != 0. + fn gemv_inplace(&self, x: &FaerVec, _t: f64, beta: f64, y: &mut FaerVec) { + if beta == 0.0 { + self.compiled.mass_action(x.as_slice(), y.as_mut_slice()); + } else { + let mut ws = self.ws.borrow_mut(); + self.compiled.mass_action(x.as_slice(), &mut ws.mv_buffer); + let ys = y.as_mut_slice(); + for (yi, &mvi) in ys.iter_mut().zip(ws.mv_buffer.iter()).take(self.n_states) { + *yi = beta.mul_add(*yi, mvi); + } + } + } + + /// M is constant: copy the precomputed CSC values into `y`, which diffsol + /// allocated from this operator's `sparsity`. Overrides the trait default, + /// which probes M one unit-vector gemv per column (O(n²) per Jacobian) + fn matrix_inplace(&self, _t: f64, y: &mut FaerSparseMat) { + y.inner_mut().val_mut().copy_from_slice(self.csc_values); + } + + /// Deep-copies the borrowed pattern, which is what the signature requires. + /// diffsol asks once per problem build, so the copy is paid there rather + /// than on every operator mint. + 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..c878ca9806 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/mod.rs @@ -0,0 +1,913 @@ +//! In-process DAE solving through diffsol. +//! +//! diffsol drives a model through one operator trait per callback it needs, and +//! this module supplies them from a [`ModelEvaluator`](crate::ModelEvaluator): +//! `rhs` for `f(t, y; p)` and its Jacobian, `mass` for `M`, `init` for `y0`, +//! `root` for events, `output` for observed variables, and `reset` as a required +//! no-op. `equations` binds those into the `OdeEquations` diffsol consumes, +//! `solve` owns problem setup and the solve loop, and `batch` fans that loop out +//! over many input sets on rayon. +//! +//! The operators share a field vocabulary: `compiled` is the shared immutable +//! [`CompiledModel`](crate::model::CompiledModel), `ws` the solve-local +//! scratch, `inputs` this solve's parameter values, and `context` faer's +//! allocator handle. Because `ws` is per solve, one solve's scratch is never +//! visible to another. +//! +//! The functions here translate our sparsity patterns into faer's, which is the +//! matrix backend diffsol is instantiated with throughout. + +pub mod batch; +pub mod equations; +pub mod init; +pub mod linear; +pub mod mass; +pub mod observable; +pub mod options; +pub mod reset; +pub mod rhs; +pub mod solve; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use faer::sparse::SymbolicSparseColMat; + +use crate::jacobian::CscPattern; +use crate::node::CsrData; + +pub use self::equations::Equations; +pub use self::options::SolverOptions; + +/// Sparsity type of the faer sparse matrix diffsol is instantiated with. +pub type FaerSparsity = as diffsol::Matrix>::Sparsity; + +/// Reinterpret a CSC Jacobian pattern as faer's symbolic sparsity. +/// +/// Both sides are column-major, so the index arrays transfer verbatim; faer's +/// checked constructor is what validates them. +pub fn csc_to_faer_sparsity(csc: &CscPattern) -> FaerSparsity { + let col_ptrs: Vec = csc.colptr.clone(); + let row_indices: Vec = csc.rowind.clone(); + SymbolicSparseColMat::new_checked(csc.nrows, csc.ncols, col_ptrs, None, row_indices) +} + +/// Build a dense CSC sparsity for an `nrows × ncols` matrix. +/// +/// Used for the `df/dp` sensitivity matrix where every state may depend on +/// every parameter. +pub 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 column-major sparsity plus the +/// value array in that CSC order. +/// +/// The mass matrix arrives from Python in CSR, and its pattern is not assumed +/// symmetric, so the entries are redistributed by counting sort rather than +/// reinterpreted in place. +pub fn csr_mass_to_faer_csc(mass: &CsrData) -> (FaerSparsity, Vec) { + let nrows = mass.shape.rows; + let ncols = mass.shape.cols; + + let mut col_counts = vec![0usize; ncols + 1]; + for row in 0..nrows { + for idx in mass.indptr[row]..mass.indptr[row + 1] { + col_counts[mass.indices[idx] + 1] += 1; + } + } + for c in 1..=ncols { + col_counts[c] += col_counts[c - 1]; + } + + let nnz = mass.indptr[nrows]; + let mut row_indices = vec![0usize; nnz]; + let mut values = vec![0.0f64; nnz]; + let mut current = col_counts.clone(); + for row in 0..nrows { + for idx in mass.indptr[row]..mass.indptr[row + 1] { + let col = mass.indices[idx]; + row_indices[current[col]] = row; + values[current[col]] = mass.data[idx]; + current[col] += 1; + } + } + + let sparsity = SymbolicSparseColMat::new_checked(nrows, ncols, col_counts, None, row_indices); + (sparsity, values) +} + +/// Transpose a CSR mass matrix's pattern into faer's column-major sparsity. +pub fn csr_mass_to_faer_sparsity(mass: &CsrData) -> FaerSparsity { + csr_mass_to_faer_csc(mass).0 +} + +#[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, SolveRequest}; + use super::*; + use crate::arena::Arena; + use crate::model::{CompiledModel, ModelEvaluator, Workspace}; + use crate::node::{CsrData, Node, Shape}; + use diffsol::vector::faer_serial::FaerVec; + use diffsol::{ + ConstantOp, Context, FaerContext, LinearOp, Matrix, NonLinearOp, NonLinearOpJacobian, Op, + VectorHost, + }; + + /// Owns what a borrowed operator view points at, as `Equations` does in a + /// real solve, so one op can be exercised on its own. + struct OpFixture { + compiled: Arc, + ws: RefCell, + jac_sparsity: FaerSparsity, + mass_sparsity: FaerSparsity, + mass_csc_values: Vec, + inputs: Vec, + sens_params: Vec, + n_states: usize, + context: FaerContext, + } + + impl OpFixture { + fn new(model: ModelEvaluator) -> Self { + let compiled = model.into_compiled(); + 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 n_params = compiled.n_params(); + Self { + n_states: compiled.n_states(), + compiled, + ws, + jac_sparsity, + mass_sparsity, + mass_csc_values, + inputs: vec![0.0; n_params], + sens_params: (0..n_params).collect(), + context: FaerContext::default(), + } + } + + fn rhs(&self) -> RhsOp<'_> { + RhsOp { + compiled: &self.compiled, + ws: &self.ws, + inputs: &self.inputs, + sens_params: &self.sens_params, + jac_sparsity: &self.jac_sparsity, + n_states: self.n_states, + context: self.context, + } + } + + fn mass(&self) -> MassOp<'_> { + MassOp { + compiled: &self.compiled, + ws: &self.ws, + sparsity: &self.mass_sparsity, + csc_values: &self.mass_csc_values, + n_states: self.n_states, + context: self.context, + } + } + } + + fn build_2state_model() -> ModelEvaluator { + // dy0/dt = -y0, dy1/dt = -2*y1 + // Jacobian = [[-1, 0], [0, -2]] + 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])); + + let mass = CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![1.0, 1.0], + shape: Shape { rows: 2, cols: 2 }, + }; + + ModelEvaluator::new(&arena, rhs, mass, 2, 0) + } + + fn build_1state_model() -> ModelEvaluator { + // dy/dt = -y + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs_expr = arena.alloc(Node::Mul(neg, sv)); + + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape { rows: 1, cols: 1 }, + }; + + ModelEvaluator::new(&arena, rhs_expr, mass, 1, 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.nrows(), csc.nrows); + assert_eq!(sparsity.ncols(), csc.ncols); + assert_eq!(sparsity.row_idx().len(), csc.rowind.len()); + + // Verify col_ptrs match exactly + let faer_col_ptrs = sparsity.col_ptr(); + assert_eq!( + faer_col_ptrs.len(), + csc.colptr.len(), + "col_ptrs length mismatch" + ); + for (i, (&faer_val, &csc_val)) in faer_col_ptrs.iter().zip(&csc.colptr).enumerate() { + assert_eq!(faer_val, csc_val, "col_ptrs[{i}] mismatch"); + } + + // Verify row_indices match exactly + let faer_row_idx = sparsity.row_idx(); + assert_eq!( + faer_row_idx.len(), + csc.rowind.len(), + "row_idx length mismatch" + ); + for (i, (&faer_val, &csc_val)) in faer_row_idx.iter().zip(&csc.rowind).enumerate() { + assert_eq!(faer_val, csc_val, "row_idx[{i}] mismatch"); + } + } + + #[test] + fn csr_mass_to_faer_preserves_entries() { + let model = build_2state_model(); + let mass = model.mass_matrix(); + let sparsity = csr_mass_to_faer_sparsity(mass); + + assert_eq!(sparsity.nrows(), mass.shape.rows); + assert_eq!(sparsity.ncols(), mass.shape.cols); + + let nnz_csr = mass.indptr[mass.shape.rows]; + assert_eq!( + sparsity.row_idx().len(), + nnz_csr, + "nnz mismatch between CSR and faer CSC" + ); + + // Collect (row, col) pairs from CSR + let mut csr_entries: Vec<(usize, usize)> = Vec::new(); + for row in 0..mass.shape.rows { + for idx in mass.indptr[row]..mass.indptr[row + 1] { + csr_entries.push((row, mass.indices[idx])); + } + } + csr_entries.sort_unstable(); + + // Collect (row, col) pairs from faer CSC sparsity + let col_ptrs = sparsity.col_ptr(); + let row_idx = sparsity.row_idx(); + let mut faer_entries: Vec<(usize, usize)> = Vec::new(); + for col in 0..sparsity.ncols() { + for &row in &row_idx[col_ptrs[col]..col_ptrs[col + 1]] { + faer_entries.push((row, col)); + } + } + faer_entries.sort_unstable(); + + assert_eq!( + csr_entries, faer_entries, + "structural entries differ between CSR source and faer CSC" + ); + } + + #[test] + fn solve_1state_exponential_decay() { + // dy/dt = -y, y(0) = 1. Exact: y(t) = exp(-t) + let model = build_1state_model(); + + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let result = + solve::solve(model, &t_eval, &[], &[1.0], &[], 1e-8, &[1e-10]).expect("solve failed"); + + assert_eq!(result.flag, 0, "should complete without events"); + assert!(result.t_event.is_none(), "no events expected"); + assert!(result.y_event.is_none(), "no event state expected"); + assert_eq!(result.n_rows, 1); + assert!( + result.n_times >= 10, + "should have at least 10 output points, got {}", + result.n_times + ); + + // Verify monotonically increasing time + for w in result.t.windows(2) { + assert!(w[1] > w[0], "time should be monotonically increasing"); + } + + // Verify accuracy against exact solution + for (j, &t) in result.t.iter().enumerate() { + let exact = (-t).exp(); + let y_val = result.y[j * result.n_rows]; + let error = (y_val - exact).abs(); + assert!( + error < 1e-6, + "t={t:.1}: y={y_val:.8}, exact={exact:.8}, error={error:.2e}", + ); + } + } + + #[test] + fn solve_2state_independent_decay() { + // dy0/dt = -y0, dy1/dt = -2*y1 from y = [1, 1], so the exact solution is + // y0(t) = exp(-t), y1(t) = exp(-2t). + let model = build_2state_model(); + + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.05).collect(); + let result = solve::solve(model, &t_eval, &[], &[1.0, 1.0], &[], 1e-8, &[1e-10, 1e-10]) + .expect("solve failed"); + + assert_eq!(result.flag, 0); + + for (j, &t) in result.t.iter().enumerate() { + let y0 = result.y[j * result.n_rows]; + let y1 = result.y[j * result.n_rows + 1]; + + let exact0 = (-t).exp(); + let exact1 = (-2.0 * t).exp(); + let err0 = (y0 - exact0).abs(); + let err1 = (y1 - exact1).abs(); + + assert!( + err0 < 1e-6, + "t={t:.2}: y0={y0:.8}, exact={exact0:.8}, err={err0:.2e}", + ); + assert!( + err1 < 1e-6, + "t={t:.2}: y1={y1:.8}, exact={exact1:.8}, err={err1:.2e}", + ); + } + } + + #[test] + fn prepared_problem_repeated_solves_match() { + let model = build_2state_model(); + let atol = vec![1e-10, 1e-10]; + let t_eval: Vec = (0..=20).map(|i| f64::from(i) * 0.05).collect(); + + // Reference: single solve through the free function + let ref_result = solve::solve( + build_2state_model(), + &t_eval, + &[], + &[1.0, 1.0], + &[], + 1e-8, + &atol, + ) + .expect("reference solve failed"); + + // PreparedSolver: multiple solves on the same problem + let prepared = + solve::PreparedSolver::new(model, 1e-8, &atol).expect("PreparedSolver creation failed"); + + let result1 = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0, 1.0], &[])) + .expect("first prepared solve failed"); + let result2 = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0, 1.0], &[])) + .expect("second prepared solve failed"); + + // Both prepared solves should match the reference + assert_eq!(result1.n_times, ref_result.n_times); + assert_eq!(result2.n_times, ref_result.n_times); + assert_eq!(result1.y.len(), ref_result.y.len(), "result lengths differ"); + for i in 0..result1.y.len() { + assert!( + (result1.y[i] - ref_result.y[i]).abs() < 1e-10, + "result1.y[{i}] mismatch: {} vs {}", + result1.y[i], + ref_result.y[i] + ); + assert!( + (result2.y[i] - ref_result.y[i]).abs() < 1e-10, + "result2.y[{i}] mismatch: {} vs {}", + result2.y[i], + ref_result.y[i] + ); + } + } + + fn build_1state_model_with_output() -> ModelEvaluator { + // dy/dt = -y, output = 2*y + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs_expr = arena.alloc(Node::Mul(neg, sv)); + + let two = arena.alloc(Node::Scalar(2.0)); + let output_expr = arena.alloc(Node::Mul(two, sv)); + + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape { rows: 1, cols: 1 }, + }; + + let mut model = ModelEvaluator::new(&arena, rhs_expr, mass, 1, 0); + model.add_output(&arena, output_expr); + model + } + + #[test] + fn prepared_problem_outputs_request_matches_full() { + // dy/dt = -y, y(0) = 1. Output = 2*y. + // Verify outputs = 2 * states from full solve. + let model_full = build_1state_model_with_output(); + let model_out = build_1state_model_with_output(); + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + + // Full solve + let prepared_full = + solve::PreparedSolver::new(model_full, 1e-8, &atol).expect("full setup failed"); + let full_result = prepared_full + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .expect("full solve failed"); + + // Output-only solve + let prepared_out = + solve::PreparedSolver::new(model_out, 1e-8, &atol).expect("output setup failed"); + let out_result = prepared_out + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[]), + ) + .expect("output solve failed"); + + // Output = 2*y, so outputs should be 2x the state values + assert_eq!(out_result.n_rows, 1); + assert_eq!(out_result.n_times, full_result.n_times); + for j in 0..out_result.n_times { + let state = full_result.y[j * full_result.n_rows]; + let output = out_result.y[j * out_result.n_rows]; + let expected = 2.0 * state; + assert!( + (output - expected).abs() < 1e-8, + "t={}: output={output}, expected={expected}", + out_result.t[j] + ); + } + } + + #[test] + fn prepared_problem_different_y0() { + let model = build_1state_model(); + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + + let prepared = + solve::PreparedSolver::new(model, 1e-8, &atol).expect("PreparedSolver creation failed"); + + let result_a = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .expect("solve with y0=1 failed"); + let result_b = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[2.0], &[])) + .expect("solve with y0=2 failed"); + + // y(t) = y0 * exp(-t), so result_b should be 2x result_a + for (j, &t) in result_a.t.iter().enumerate() { + let ya = result_a.y[j]; + let yb = result_b.y[j]; + let exact_a = (-t).exp(); + let exact_b = 2.0 * (-t).exp(); + assert!( + (ya - exact_a).abs() < 1e-6, + "y0=1: t={t}, y={ya}, expected {exact_a}" + ); + assert!( + (yb - exact_b).abs() < 1e-6, + "y0=2: t={t}, y={yb}, expected {exact_b}" + ); + } + } + + #[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); + x.as_mut_slice().copy_from_slice(&[3.0, 5.0]); + let mut y = ctx.vector_zeros::>(2); + rhs.call_inplace(&x, 0.0, &mut y); + + // f = [-y0, -2*y1] => [-3, -10] + assert!((y.as_slice()[0] - (-3.0)).abs() < 1e-12); + assert!((y.as_slice()[1] - (-10.0)).abs() < 1e-12); + } + + #[test] + fn rhs_jac_mul_2state() { + let fixture = OpFixture::new(build_2state_model()); + let rhs = fixture.rhs(); + let ctx = FaerContext::default(); + + let mut x = ctx.vector_zeros::>(2); + x.as_mut_slice().copy_from_slice(&[1.0, 1.0]); + let mut v = ctx.vector_zeros::>(2); + v.as_mut_slice().copy_from_slice(&[2.0, 3.0]); + let mut y = ctx.vector_zeros::>(2); + + rhs.jac_mul_inplace(&x, 0.0, &v, &mut y); + + // df/dy = diag(-1, -2), so (df/dy) @ v = [-v0, -2*v1] = [-2, -6] + assert!((y.as_slice()[0] - (-2.0)).abs() < 1e-12); + assert!((y.as_slice()[1] - (-6.0)).abs() < 1e-12); + } + + #[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); + x.as_mut_slice().copy_from_slice(&[1.0, 1.0]); + + // Allocate the sparse matrix from the op's own sparsity and assemble via + // jacobian_inplace (the diffsol `jacobian` default does exactly this). + let jac = rhs.jacobian(&x, 0.0); + + // df/dy = diag(-1, -2). Collect the assembled triplets and check. + 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_identity_beta_zero() { + let fixture = OpFixture::new(build_2state_model()); + let mass = fixture.mass(); + let ctx = FaerContext::default(); + + let mut x = ctx.vector_zeros::>(2); + x.as_mut_slice().copy_from_slice(&[7.0, -4.0]); + let mut y = ctx.vector_zeros::>(2); + + // beta == 0: y = M @ x. Mass is identity => y == x. + mass.gemv_inplace(&x, 0.0, 0.0, &mut y); + assert!((y.as_slice()[0] - 7.0).abs() < 1e-12); + assert!((y.as_slice()[1] - (-4.0)).abs() < 1e-12); + } + + #[test] + fn mass_gemv_identity_beta_nonzero() { + let fixture = OpFixture::new(build_2state_model()); + let mass = fixture.mass(); + let ctx = FaerContext::default(); + + let mut x = ctx.vector_zeros::>(2); + x.as_mut_slice().copy_from_slice(&[1.0, 2.0]); + let mut y = ctx.vector_zeros::>(2); + y.as_mut_slice().copy_from_slice(&[10.0, 20.0]); + + // beta == 3: y = M @ x + beta * y_old = x + 3*y_old + // = [1 + 30, 2 + 60] = [31, 62] + mass.gemv_inplace(&x, 0.0, 3.0, &mut y); + assert!( + (y.as_slice()[0] - 31.0).abs() < 1e-12, + "y0={}", + y.as_slice()[0] + ); + assert!( + (y.as_slice()[1] - 62.0).abs() < 1e-12, + "y1={}", + y.as_slice()[1] + ); + } + + #[test] + fn mass_has_sparsity() { + let fixture = OpFixture::new(build_2state_model()); + assert!(fixture.mass().sparsity().is_some()); + } + + #[test] + fn mass_matrix_inplace_matches_probing_default() { + // The copy-based override must reproduce diffsol's column-probing + // default entry for entry, on identity and DAE (zero-row) masses. + 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_states: 3, + n_sens_params: 0, + context: ctx, + }; + let mut y = ctx.vector_zeros::>(3); + init.call_inplace(0.0, &mut y); + assert_eq!(y.as_slice(), &[1.5, -2.5, 3.0]); + } + + #[test] + fn sequential_solves_are_isolated() { + // Same PreparedSolver, two solves with different y0. The fresh + // Workspace-per-solve must keep them fully independent. + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = solve::PreparedSolver::new(build_1state_model(), 1e-8, &atol).unwrap(); + + let a = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .unwrap(); + let b = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[5.0], &[])) + .unwrap(); + + for (j, &t) in a.t.iter().enumerate() { + let exact = (-t).exp(); + assert!( + (a.y[j * a.n_rows] - exact).abs() < 1e-6, + "a t={t}: {} vs {exact}", + a.y[j * a.n_rows] + ); + assert!( + 5.0f64.mul_add(-exact, b.y[j * b.n_rows]).abs() < 1e-6, + "b t={t}: {} vs {}", + b.y[j * b.n_rows], + 5.0 * exact + ); + } + } + + #[test] + fn solve_takes_shared_ref() { + // `prepared` is intentionally NOT `mut`; `solve` takes `&self`. + let atol = vec![1e-10]; + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.1).collect(); + let prepared = solve::PreparedSolver::new(build_1state_model(), 1e-8, &atol).unwrap(); + let r1 = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .unwrap(); + let r2 = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .unwrap(); + assert_eq!(r1.n_times, r2.n_times); + } + + #[test] + fn prepared_problem_is_send_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + } + + #[test] + fn an_outputs_request_matches_analytic() { + // dy/dt = -y, y(0) = 1, output = 2*y => 2*exp(-t). + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = + solve::PreparedSolver::new(build_1state_model_with_output(), 1e-8, &atol).unwrap(); + let r = prepared + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[]), + ) + .unwrap(); + + assert_eq!(r.n_rows, 1); + for (j, &t) in r.t.iter().enumerate() { + let expected = 2.0 * (-t).exp(); + let out = r.y[j * r.n_rows]; + assert!( + (out - expected).abs() < 1e-6, + "t={t}: output={out}, expected={expected}" + ); + } + } + + #[test] + fn an_outputs_request_at_final_time_carries_the_terminal_full_state() { + // Without a state trajectory, y_event is the only state a caller can + // restart from, so it must be present on final-time termination too. + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared_full = + solve::PreparedSolver::new(build_1state_model_with_output(), 1e-8, &atol).unwrap(); + let full = prepared_full + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .unwrap(); + let prepared_out = + solve::PreparedSolver::new(build_1state_model_with_output(), 1e-8, &atol).unwrap(); + let r = prepared_out + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[]), + ) + .unwrap(); + + assert_eq!(r.flag, 0); + assert!(r.t_event.is_none()); + let y_event = r.y_event.expect("terminal state missing"); + assert_eq!(y_event.len(), 1); + let full_last = full.y[(full.n_times - 1) * full.n_rows]; + assert!( + (y_event[0] - full_last).abs() < 1e-8, + "terminal state {} diverges from full solve {full_last}", + y_event[0] + ); + } + + fn build_decay_model_with_output_and_event() -> ModelEvaluator { + // dy/dt = -y, y(0) = 1, output = 2*y, event = y - 0.5 (root at t = ln 2). + // Output and state differ there, so y_event tells them apart. + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg = arena.alloc(Node::Scalar(-1.0)); + let rhs_expr = arena.alloc(Node::Mul(neg, sv)); + + let two = arena.alloc(Node::Scalar(2.0)); + let output_expr = arena.alloc(Node::Mul(two, sv)); + + let half = arena.alloc(Node::Scalar(0.5)); + let event_expr = arena.alloc(Node::Sub(sv, half)); + + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape { rows: 1, cols: 1 }, + }; + + let mut model = ModelEvaluator::new(&arena, rhs_expr, mass, 1, 0); + model.add_output(&arena, output_expr); + model.add_event(&arena, event_expr); + model + } + + #[test] + // The trajectory's final time is the event time stored verbatim, so the + // assertion checks for a bit-identical value, not an approximate one. + #[allow(clippy::float_cmp)] + fn solve_event_ends_trajectory_at_root_with_full_state() { + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = + solve::PreparedSolver::new(build_decay_model_with_output_and_event(), 1e-8, &atol) + .unwrap(); + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[])) + .unwrap(); + + let ln2 = std::f64::consts::LN_2; + assert_eq!(r.flag, 1, "expected event termination"); + let t_event = r.t_event.expect("t_event missing"); + assert!( + (t_event - ln2).abs() < 1e-6, + "t_event={t_event}, expected {ln2}" + ); + // The trajectory includes the root time and state as its final column. + assert_eq!(r.n_times, r.t.len()); + assert_eq!(*r.t.last().unwrap(), t_event); + let y_last = r.y[(r.n_times - 1) * r.n_rows]; + assert!( + (y_last - 0.5).abs() < 1e-6, + "y at event={y_last}, expected 0.5" + ); + // y_event is the full state at the root. + let y_event = r.y_event.expect("y_event missing"); + assert_eq!(y_event.len(), 1); + assert!( + (y_event[0] - 0.5).abs() < 1e-6, + "y_event={}, expected state 0.5", + y_event[0] + ); + } + + #[test] + // The trajectory's final time is the event time stored verbatim, so the + // assertion checks for a bit-identical value, not an approximate one. + #[allow(clippy::float_cmp)] + fn an_outputs_request_returns_a_full_state_y_event_on_an_event() { + let atol = vec![1e-10]; + let t_eval: Vec = (0..=10).map(|i| f64::from(i) * 0.1).collect(); + let prepared = + solve::PreparedSolver::new(build_decay_model_with_output_and_event(), 1e-8, &atol) + .unwrap(); + let r = prepared + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[]), + ) + .unwrap(); + + let ln2 = std::f64::consts::LN_2; + assert_eq!(r.flag, 1, "expected event termination"); + let t_event = r.t_event.expect("t_event missing"); + assert!( + (t_event - ln2).abs() < 1e-6, + "t_event={t_event}, expected {ln2}" + ); + // The trajectory includes the root time and output value as its final column. + assert_eq!(r.n_times, r.t.len()); + assert_eq!(*r.t.last().unwrap(), t_event); + let out_last = r.y[(r.n_times - 1) * r.n_rows]; + assert!( + (out_last - 1.0).abs() < 1e-6, + "output at event={out_last}, expected 1.0" + ); + // y_event is the full state at the root, not the outputs row. + let y_event = r.y_event.expect("y_event missing"); + assert_eq!(y_event.len(), 1); + assert!( + (y_event[0] - 0.5).abs() < 1e-6, + "y_event={}, expected state 0.5, not output 1.0", + y_event[0] + ); + } + + fn build_dae_model() -> ModelEvaluator { + // Differential: y0' = -y0. Algebraic: 0 = 2*y0 - y1 (=> y1 = 2*y0). + // Mass = 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])); + let mass = CsrData { + indptr: vec![0, 1, 1], // row0: (0,0)=1; row1: empty => 0 + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(2, 2), + }; + ModelEvaluator::new(&arena, rhs, 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 = solve::PreparedSolver::new(build_dae_model(), 1e-8, &atol).unwrap(); + // y0=[1.0, 0.0]: algebraic y1 START is deliberately INCONSISTENT (should be 2.0). + let r = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0, 0.0], &[])) + .unwrap(); + assert_eq!(r.flag, 0); + 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..20261155cf --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/observable.rs @@ -0,0 +1,108 @@ +//! Observable operator: one family of observables as a diffsol operator. +//! +//! diffsol wants both of `PyBaMM`'s observable families as a vector-valued +//! function of `(t, y; p)` with a jacobian and a sens action: output variables +//! as its `out`, events as its `root`. That is the same operator over a +//! different [`ObservableKind`], so it is written once and mounted twice. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op, VectorHost}; + +use crate::model::{CompiledModel, Workspace}; +use crate::observable::ObservableKind; + +/// One observable family as `H(t, y; p)`. +/// +/// Acting on an event root is the caller's job, not the operator's, so events +/// need nothing here that outputs do not. +pub struct ObservableOp<'a> { + pub compiled: &'a CompiledModel, + pub ws: &'a RefCell, + pub inputs: &'a [f64], + /// Global parameter index of each sensitivity column. + pub sens_params: &'a [usize], + pub n_states: usize, + pub kind: ObservableKind, + /// The family's concatenated length, which diffsol asks for per callback. + pub n_out: usize, + pub context: FaerContext, +} + +impl std::fmt::Debug for ObservableOp<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ObservableOp") + .field("kind", &self.kind) + .field("n_out", &self.n_out) + .finish_non_exhaustive() + } +} + +impl Op for ObservableOp<'_> { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn nstates(&self) -> usize { + self.n_states + } + fn nout(&self) -> usize { + self.n_out + } + fn nparams(&self) -> usize { + self.sens_params.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, + self.kind, + t, + x.as_slice(), + self.inputs, + y.as_mut_slice(), + ); + } +} + +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, + self.kind, + t, + x.as_slice(), + self.inputs, + v.as_slice(), + y.as_mut_slice(), + ); + } +} + +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, + self.kind, + t, + x.as_slice(), + self.inputs, + self.sens_params, + v.as_slice(), + y.as_mut_slice(), + ); + } +} 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..375d6f1be7 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/options.rs @@ -0,0 +1,280 @@ +//! Integrator tuning, one field per diffsol `OdeSolverOptions` knob. +//! +//! Every default is read from diffsol rather than copied, so a caller +//! overriding one knob does not silently inherit a different value for the +//! rest and an upstream change cannot be pinned here unnoticed. The single +//! deliberate departure is [`SolverOptions::max_nonlinear_solver_failures`]. + +use diffsol::{OdeEquations, OdeSolverOptions, OdeSolverProblem}; + +use crate::error::CoreError; + +/// Our replacement for diffsol's whole-solve nonlinear-failure budget, on the +/// scale of IDAKLU's `max_num_steps`. +/// +/// This counter is cumulative over a solve, where diffsol's error-test counter +/// and IDA's `IDASetMaxConvFails` are consecutive per step, so 50 bounds solve +/// *length*, not divergence: a DFN pulse train recovers from ~86 of them per +/// 1800 s. `min_timestep` is what actually catches divergence. +const MAX_NONLINEAR_SOLVER_FAILURES: usize = 100_000; + +/// Tuning for the diffsol BDF integrator. +/// +/// Field names match diffsol's `OdeSolverOptions` so a knob can be traced from +/// `PyBaMM`'s `options` dict to the integrator without a translation table. +#[derive(Debug, Clone, PartialEq)] +pub struct SolverOptions { + /// Newton iterations allowed per nonlinear solve. + pub max_nonlinear_solver_iterations: usize, + /// Consecutive step rejections allowed before the solve aborts. + pub max_error_test_failures: usize, + /// Nonlinear-solver convergence failures allowed over the whole solve. + pub max_nonlinear_solver_failures: usize, + /// Newton convergence-test scaling factor. + pub nonlinear_solver_tolerance: f64, + /// Smallest step the controller may take before the solve aborts. + pub min_timestep: f64, + /// Upper bound on step-size growth; `None` keeps diffsol's own. + pub max_timestep_growth: Option, + /// Lower bound of the step-growth dead zone; `None` keeps diffsol's own. + pub min_timestep_growth: Option, + /// Upper bound of the step-shrink dead zone; `None` keeps diffsol's own. + pub max_timestep_shrink: Option, + /// Absolute lower bound on step-size reduction; `None` keeps diffsol's own. + pub min_timestep_shrink: Option, + /// Steps between linear-solver setups. + pub update_jacobian_after_steps: usize, + /// Steps between full RHS Jacobian re-evaluations. + pub update_rhs_jacobian_after_steps: usize, + /// Relative step-size change that forces a Jacobian update. + pub threshold_to_update_jacobian: f64, + /// Relative step-size change that forces an RHS Jacobian update. + pub threshold_to_update_rhs_jacobian: f64, + /// PI step controller proportional gain. + pub pi_control_proportional: f64, + /// PI step controller integral gain. + pub pi_control_integral: f64, +} + +impl Default for SolverOptions { + fn default() -> Self { + Self { + max_nonlinear_solver_failures: MAX_NONLINEAR_SOLVER_FAILURES, + ..Self::diffsol_defaults() + } + } +} + +impl SolverOptions { + /// diffsol's own defaults, including the failure budget we otherwise raise. + /// + /// Read straight off `OdeSolverOptions` so a bump changes this with it. + #[must_use] + pub fn diffsol_defaults() -> Self { + let defaults = OdeSolverOptions::::default(); + Self { + max_nonlinear_solver_iterations: defaults.max_nonlinear_solver_iterations, + max_error_test_failures: defaults.max_error_test_failures, + max_nonlinear_solver_failures: defaults.max_nonlinear_solver_failures, + nonlinear_solver_tolerance: defaults.nonlinear_solver_tolerance, + min_timestep: defaults.min_timestep, + max_timestep_growth: defaults.max_timestep_growth, + min_timestep_growth: defaults.min_timestep_growth, + max_timestep_shrink: defaults.max_timestep_shrink, + min_timestep_shrink: defaults.min_timestep_shrink, + update_jacobian_after_steps: defaults.update_jacobian_after_steps, + update_rhs_jacobian_after_steps: defaults.update_rhs_jacobian_after_steps, + threshold_to_update_jacobian: defaults.threshold_to_update_jacobian, + threshold_to_update_rhs_jacobian: defaults.threshold_to_update_rhs_jacobian, + pi_control_proportional: defaults.pi_control_proportional, + pi_control_integral: defaults.pi_control_integral, + } + } + + /// Reject values diffsol would take at face value and then misbehave on. + /// + /// Counts are unsigned, so only the floating-point knobs need checking: all + /// must be finite, and those acting as a scale or a bound must be positive. + /// + /// # Errors + /// + /// [`CoreError::SolverOption`] naming the first offending field. + pub fn validate(&self) -> Result<(), CoreError> { + let positive = [ + ( + "nonlinear_solver_tolerance", + self.nonlinear_solver_tolerance, + ), + ("min_timestep", self.min_timestep), + ( + "threshold_to_update_jacobian", + self.threshold_to_update_jacobian, + ), + ( + "threshold_to_update_rhs_jacobian", + self.threshold_to_update_rhs_jacobian, + ), + ]; + for (name, value) in positive { + if !value.is_finite() || value <= 0.0 { + return Err(CoreError::SolverOption { + name: name.to_string(), + got: value, + }); + } + } + + let optional = [ + ("max_timestep_growth", self.max_timestep_growth), + ("min_timestep_growth", self.min_timestep_growth), + ("max_timestep_shrink", self.max_timestep_shrink), + ("min_timestep_shrink", self.min_timestep_shrink), + ]; + for (name, value) in optional { + if let Some(value) = value + && (!value.is_finite() || value <= 0.0) + { + return Err(CoreError::SolverOption { + name: name.to_string(), + got: value, + }); + } + } + + // A PI gain of zero disables that term, so only finiteness applies. + let finite = [ + ("pi_control_proportional", self.pi_control_proportional), + ("pi_control_integral", self.pi_control_integral), + ]; + for (name, value) in finite { + if !value.is_finite() { + return Err(CoreError::SolverOption { + name: name.to_string(), + got: value, + }); + } + } + + Ok(()) + } + + /// Stamp these options onto a freshly built problem. + /// + /// `OdeBuilder` exposes no setters for them, so the problem's public + /// `ode_options` is the wiring point; it is read when the solver is created. + #[must_use] + pub const fn apply>( + &self, + mut problem: OdeSolverProblem, + ) -> OdeSolverProblem { + let options = &mut problem.ode_options; + options.max_nonlinear_solver_iterations = self.max_nonlinear_solver_iterations; + options.max_error_test_failures = self.max_error_test_failures; + options.max_nonlinear_solver_failures = self.max_nonlinear_solver_failures; + options.nonlinear_solver_tolerance = self.nonlinear_solver_tolerance; + options.min_timestep = self.min_timestep; + options.max_timestep_growth = self.max_timestep_growth; + options.min_timestep_growth = self.min_timestep_growth; + options.max_timestep_shrink = self.max_timestep_shrink; + options.min_timestep_shrink = self.min_timestep_shrink; + options.update_jacobian_after_steps = self.update_jacobian_after_steps; + options.update_rhs_jacobian_after_steps = self.update_rhs_jacobian_after_steps; + options.threshold_to_update_jacobian = self.threshold_to_update_jacobian; + options.threshold_to_update_rhs_jacobian = self.threshold_to_update_rhs_jacobian; + options.pi_control_proportional = self.pi_control_proportional; + options.pi_control_integral = self.pi_control_integral; + problem + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_the_failure_budget_departs_from_diffsol() { + let ours = SolverOptions::default(); + 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 the_raised_budget_clears_the_measured_worst_case() { + // DFN pulse_train spends ~141 with sensitivities on; a default that + // merely doubled diffsol's 50 would still fail it. + assert!(SolverOptions::default().max_nonlinear_solver_failures > 10_000); + } + + #[test] + fn defaults_validate() { + SolverOptions::default() + .validate() + .expect("defaults invalid"); + } + + #[test] + fn non_positive_and_non_finite_scales_are_rejected() { + for bad in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let options = SolverOptions { + nonlinear_solver_tolerance: bad, + ..Default::default() + }; + let Err(err) = options.validate() else { + panic!("{bad} should have been rejected"); + }; + assert!(matches!(err, CoreError::SolverOption { ref name, .. } + if name == "nonlinear_solver_tolerance")); + } + } + + #[test] + fn an_absent_optional_bound_is_not_validated() { + let options = SolverOptions { + max_timestep_growth: None, + ..Default::default() + }; + options.validate().expect("None must be accepted"); + } + + #[test] + fn a_present_optional_bound_is_validated() { + let options = SolverOptions { + max_timestep_growth: Some(-2.0), + ..Default::default() + }; + let err = options.validate().expect_err("negative growth accepted"); + assert!(matches!(err, CoreError::SolverOption { ref name, .. } + if name == "max_timestep_growth")); + } + + #[test] + fn a_zero_pi_gain_is_accepted_because_it_disables_the_term() { + let options = SolverOptions { + pi_control_proportional: 0.0, + ..Default::default() + }; + options.validate().expect("zero gain must be accepted"); + } + + #[test] + fn a_non_finite_pi_gain_is_rejected() { + let options = SolverOptions { + pi_control_integral: f64::NAN, + ..Default::default() + }; + let err = options.validate().expect_err("NaN gain accepted"); + assert!(matches!(err, CoreError::SolverOption { ref name, .. } + if name == "pi_control_integral")); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/solver/reset.rs b/packages/pybamm-rust/pybamm-core/src/solver/reset.rs new file mode 100644 index 0000000000..ebddd0fc8f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/reset.rs @@ -0,0 +1,79 @@ +//! Placeholder reset operator. +//! +//! `PyBaMM` events terminate the solve rather than triggering a state reset, so +//! `Equations::reset()` always returns `None` and `ResetOp` is never +//! constructed. It exists solely to satisfy the `Reset` associated-type bounds +//! of `OdeEquationsImplicitSens`. + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{FaerContext, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op, VectorHost}; + +/// Never-constructed reset operator; satisfies the diffsol `Reset` bounds. +pub struct ResetOp { + /// Number of state variables. + pub n_states: usize, + /// Number of sensitivity columns. + pub n_sens_params: usize, + /// Faer execution context. + pub context: FaerContext, +} + +impl std::fmt::Debug for ResetOp { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResetOp") + .field("n_states", &self.n_states) + .field("n_sens_params", &self.n_sens_params) + .finish_non_exhaustive() + } +} + +impl Op for ResetOp { + type T = f64; + type V = FaerVec; + type M = FaerSparseMat; + type C = FaerContext; + + fn context(&self) -> &Self::C { + &self.context + } + + fn nstates(&self) -> usize { + self.n_states + } + + fn nout(&self) -> usize { + self.n_states + } + + fn nparams(&self) -> usize { + self.n_sens_params + } +} + +impl NonLinearOp for ResetOp { + /// Identity reset: new state equals old state. + fn call_inplace(&self, x: &FaerVec, _t: f64, y: &mut FaerVec) { + y.as_mut_slice().copy_from_slice(x.as_slice()); + } +} + +impl NonLinearOpJacobian for ResetOp { + /// Identity Jacobian-vector product. + fn jac_mul_inplace(&self, _x: &FaerVec, _t: f64, v: &FaerVec, y: &mut FaerVec) { + y.as_mut_slice().copy_from_slice(v.as_slice()); + } +} + +impl NonLinearOpSens for ResetOp { + /// Zero parameter sensitivity. + fn sens_mul_inplace( + &self, + _x: &FaerVec, + _t: f64, + _v: &FaerVec, + y: &mut FaerVec, + ) { + y.as_mut_slice().fill(0.0); + } +} 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..29c6a81462 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/rhs.rs @@ -0,0 +1,142 @@ +//! Right-hand-side operator. +//! +//! Supplies `f(t, y; p)`, its sparse `df/dy` assembled through the compiled +//! coloring, and the `df/dp` columns forward sensitivities need. Those columns +//! are numbered by position within the solve's parameter subset; `sens_params` +//! maps each back to its global parameter index. + +use std::cell::RefCell; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + FaerContext, Matrix, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, Op, VectorHost, +}; + +use super::{FaerSparsity, dense_faer_sparsity}; +use crate::model::{CompiledModel, Workspace}; + +/// RHS operator f(t, y; p), borrowed from the [`Equations`](super::equations::Equations) +/// that owns the solve. +/// +/// diffsol mints one of these per callback invocation, so it holds nothing of +/// its own: every field is a reference into the equations, making a mint a +/// handful of pointer copies. +pub struct RhsOp<'a> { + pub compiled: &'a CompiledModel, + pub ws: &'a RefCell, + pub inputs: &'a [f64], + /// Global parameter index of each sensitivity column. + pub sens_params: &'a [usize], + pub jac_sparsity: &'a FaerSparsity, + pub n_states: usize, + 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.n_states) + .field("n_sens_params", &self.sens_params.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.n_states + } + fn nout(&self) -> usize { + self.n_states + } + fn nparams(&self) -> usize { + self.sens_params.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, x.as_slice(), self.inputs, y.as_mut_slice()); + } +} + +impl NonLinearOpJacobian for RhsOp<'_> { + fn jac_mul_inplace(&self, x: &FaerVec, t: f64, v: &FaerVec, y: &mut FaerVec) { + let mut ws = self.ws.borrow_mut(); + self.compiled.jac_action( + &mut ws, + t, + x.as_slice(), + self.inputs, + v.as_slice(), + y.as_mut_slice(), + ); + } + + fn jacobian_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerSparseMat) { + let mut ws = self.ws.borrow_mut(); + let values = y.inner_mut().val_mut(); + self.compiled + .assemble_jacobian_csc_no_mass(&mut ws, t, x.as_slice(), self.inputs, values); + } + + /// Deep-copies the borrowed pattern: diffsol asks once per problem build, + /// so the owned copy its signature requires is paid there rather than on + /// every operator mint. + 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, + x.as_slice(), + self.inputs, + self.sens_params, + v.as_slice(), + y.as_mut_slice(), + ); + } + + /// Batched df/dp: one primal pass per (t, y), then a tangent-only sweep per + /// column, instead of diffsol's default per-column primal recompute. + /// + /// Writes each column straight into `y`'s value array. diffsol builds `y` + /// from [`Self::sens_sparsity`], whose dense CSC pattern puts column `dst` + /// at `dst * n_states`, so no intermediate column vector is needed. + fn sens_inplace(&self, x: &FaerVec, t: f64, y: &mut FaerSparseMat) { + let mut ws = self.ws.borrow_mut(); + self.compiled + .sens_primal_pass(&mut ws, t, x.as_slice(), self.inputs); + let values = y.inner_mut().val_mut(); + debug_assert_eq!( + values.len(), + self.n_states * self.sens_params.len(), + "df/dp matrix was not built from sens_sparsity", + ); + for (dst, ¶m) in self.sens_params.iter().enumerate() { + let column = &mut values[dst * self.n_states..(dst + 1) * self.n_states]; + self.compiled.sens_tangent_column(&mut ws, param, column); + } + } + + /// Dense (all-entries) pattern for df/dp over the requested subset: every + /// state may depend on every parameter. + fn sens_sparsity(&self) -> Option { + Some(dense_faer_sparsity(self.n_states, self.sens_params.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..84607d9825 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/solver/solve.rs @@ -0,0 +1,3072 @@ +//! Problem setup and the solve loop. +//! +//! [`PreparedSolver`] holds everything reusable about a model, meaning the +//! immutable compiled model, the converted sparsity patterns and the sizes, so +//! repeated solves pay setup once. Each solve then builds its own workspace, its +//! own diffsol `OdeSolverProblem` and its own diffsol instance, which is what +//! keeps sequential solves independent and the prepared handle `Send + Sync`. +//! +//! What a solve carries back — states or output variables, with or without +//! sensitivities — is asked for on [`SolveRequest`] and answered in one +//! [`SolveOutcome`], so the payload combinations share one trajectory layout and +//! one set of termination fields rather than a result type each. + +use std::cell::RefCell; +use std::ops::Range; +use std::rc::Rc; +use std::sync::Arc; + +use diffsol::matrix::sparse_faer::FaerSparseMat; +use diffsol::ode_solver::OdeSolverStatistics; +use diffsol::vector::faer_serial::FaerVec; +use diffsol::{ + AugmentedOdeEquationsImplicit, AugmentedOdeSolverMethod, Context, FaerContext, + NewtonNonlinearSolver, NoLineSearch, NonLinearOp, NonLinearOpJacobian, NonLinearOpSens, + OdeBuilder, OdeEquations, OdeEquationsImplicit, OdeSolverMethod, OdeSolverProblem, + OdeSolverStopReason, StateRefMut, Vector, VectorHost, +}; + +use super::equations::Equations; +use super::linear::ReusedFaerLu; +use super::options::SolverOptions; +use super::{FaerSparsity, csc_to_faer_sparsity, csr_mass_to_faer_csc}; +use crate::error::CoreError; +use crate::model::{CompiledModel, Workspace}; +use crate::node::CsrData; + +/// Default multiplier applied to `atol` on differential rows to form the +/// forward-sensitivity absolute tolerance floor. +pub const DEFAULT_SENS_ATOL_FACTOR: f64 = 1e-3; + +/// Solver statistics from a BDF solve. +#[derive(Debug, Clone)] +pub struct SolverStatistics { + /// Accepted BDF steps. + pub number_of_steps: usize, + /// Jacobian/LU setups in total; the `_from_*` fields below break down why. + pub number_of_linear_solver_setups: usize, + /// Newton iterations across all steps. + pub number_of_nonlinear_solver_iterations: usize, + /// Newton solves that failed to converge. + pub number_of_nonlinear_solver_fails: usize, + /// Steps rejected by the local error test. + pub number_of_error_test_failures: usize, + /// Jacobian/LU setups triggered by checkpoint or reinitialisation. + pub number_of_linear_solver_setups_from_checkpoint: usize, + /// Jacobian/LU setups triggered by a first nonlinear convergence failure. + pub number_of_linear_solver_setups_from_first_convergence_fail: usize, + /// Jacobian/LU setups triggered by a second nonlinear convergence failure. + pub number_of_linear_solver_setups_from_second_convergence_fail: usize, + /// Jacobian/LU setups triggered by a local error test failure. + pub number_of_linear_solver_setups_from_error_test_fail: usize, + /// Jacobian/LU setups triggered by the normal step-success heuristic. + pub number_of_linear_solver_setups_from_step_success: usize, + /// Time spent computing consistent initial conditions (seconds). + pub ic_time_secs: f64, + /// Time spent creating the BDF solver instance (seconds). + pub solver_setup_time_secs: f64, + /// Wall-clock time for the whole solve (seconds), covering the two phases + /// above and the integration itself but no FFI marshalling. Measured here + /// rather than by the caller because a batched solve has no caller-side + /// moment that corresponds to one set's integration. + pub integration_time_secs: f64, + /// True when the sensitivity solve failed under error control and was + /// retried with sensitivities excluded from it. + pub sens_error_control_relaxed: bool, +} + +impl From<&OdeSolverStatistics> for SolverStatistics { + fn from(stats: &OdeSolverStatistics) -> Self { + Self { + number_of_steps: stats.number_of_steps, + number_of_linear_solver_setups: stats.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations: stats.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails: stats.number_of_nonlinear_solver_fails, + number_of_error_test_failures: stats.number_of_error_test_failures, + number_of_linear_solver_setups_from_checkpoint: stats + .number_of_linear_solver_setups_from_checkpoint, + number_of_linear_solver_setups_from_first_convergence_fail: stats + .number_of_linear_solver_setups_from_first_convergence_fail, + number_of_linear_solver_setups_from_second_convergence_fail: stats + .number_of_linear_solver_setups_from_second_convergence_fail, + number_of_linear_solver_setups_from_error_test_fail: stats + .number_of_linear_solver_setups_from_error_test_fail, + number_of_linear_solver_setups_from_step_success: stats + .number_of_linear_solver_setups_from_step_success, + ic_time_secs: 0.0, + solver_setup_time_secs: 0.0, + integration_time_secs: 0.0, + sens_error_control_relaxed: false, + } + } +} + +/// What one solve returns, whichever payloads were asked for. +/// +/// The termination fields live here once rather than once per payload +/// combination. `flag` is 0 when the requested time span completed and 1 when an +/// event root stopped the solve; on a root, `t` ends at the root time rather +/// than at the last requested point and `t_event` repeats it. +#[derive(Debug)] +pub struct SolveOutcome { + /// Output times, one per trajectory column. + pub t: Vec, + /// Flat trajectory in column-major order: `y[i + j * n_rows]` = row i at + /// time j. Rows are states, or the model's output variables when + /// [`SolveRequest::outputs`] asked for them. + pub y: Vec, + /// Rows in `y`, and in each sensitivity block. + pub n_rows: usize, + /// Columns in `y`, equal to `t.len()`. + pub n_times: usize, + /// Row time derivatives sharing `y`'s layout. Present only on a state + /// trajectory, when [`PreparedSolver::with_store_yp`] is set and the grid is + /// under [`MAX_HERMITE_COLUMNS`]: each column is the derivative of the BDF + /// interpolating polynomial at `t[j]`, the knot slope cubic-Hermite output + /// interpolation needs. + pub yp: Option>, + /// One flat block per requested sensitivity parameter, sharing `y`'s layout, + /// or `None` when the request asked for none. Blocks follow the order the + /// sensitivity parameters were requested in, not global parameter order. On + /// an outputs request each block already carries the full derivative + /// `dg/dp + dg/dy · y_s`, not the `dg/dy` term alone. + pub sensitivities: Option>>, + /// Root time, `None` unless an event stopped the solve. + pub t_event: Option, + /// The full state where the trajectory ends, never an outputs row: on a root + /// the state at the root time, and on an outputs request the terminal state + /// as well, the only one such a caller can restart from. + pub y_event: Option>, + /// 0 for a completed span, 1 for an event root. + pub flag: i32, + /// Step and setup counters from the diffsol run, including whether + /// sensitivity error control was relaxed on a retry. + pub statistics: SolverStatistics, +} + +impl SolveOutcome { + /// Assemble the outcome from what an engine produced. + /// + /// The one place a solve's payloads are laid out, so the four payload + /// combinations cannot drift apart in the fields they carry. + fn from_parts( + trajectory: DenseTrajectory, + y_event: Option>, + sensitivities: Option>>, + statistics: SolverStatistics, + ) -> Self { + Self { + n_rows: trajectory.n_rows, + n_times: trajectory.n_cols(), + t: trajectory.t, + y: trajectory.y, + yp: trajectory.yp, + sensitivities, + t_event: trajectory.t_event, + y_event, + flag: trajectory.flag, + statistics, + } + } +} + +/// What to integrate over and which payloads to report, shared by every input +/// set of a batch. +/// +/// The two payload axes are fields rather than one entry point each: `outputs` +/// swaps the trajectory's rows from states to the model's output variables, and +/// `sensitivities` adds the forward-sensitivity blocks. They compose, so a third +/// axis costs a field rather than doubling the entry points. +#[derive(Clone, Copy, Debug)] +#[must_use] +pub struct SolveRequest<'a> { + /// Times the solution is reported at. + pub t_eval: &'a [f64], + /// Discontinuity times the integrator lands on exactly and restarts from. + /// Every entry must also appear in `t_eval`; those that do not 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, +} + +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, + } + } + + /// 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 + } +} + +/// 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> { + /// Initial state, one entry per state. + pub y0: &'a [f64], + /// Flat input-parameter vector. + pub inputs: &'a [f64], + /// `dy0/dp` in column-major `n_states x k` order over the requested + /// sensitivity subset; empty is the all-zero seed. Read only when the + /// request asks for sensitivities. + 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 the shared immutable `CompiledModel`, the tolerances and integrator +/// options, pre-converted sparsity patterns, and the system sizes. It is *not* +/// a fully specified problem: `y0`, `t_eval` and the parameter vector all arrive +/// per call, and each call to [`solve`](Self::solve) builds its own +/// `Workspace`, its own diffsol `OdeSolverProblem` and its own diffsol solver +/// from them. Carrying no +/// per-solve mutable state is what makes this `Send + Sync` and safe to share +/// across threads. +pub struct PreparedSolver { + compiled: Arc, + rtol: f64, + atol: Vec, + /// Multiplier applied to `atol` on differential rows only. + sens_atol_factor: f64, + /// Integrator tuning, stamped onto every problem this handle builds. + options: SolverOptions, + /// Store `yp` alongside `y` on the state-trajectory paths. + store_yp: bool, + /// True where the mass-matrix row is empty or all zeros. + algebraic_rows: Vec, + /// M's diagonal, or `None` if M has an off-diagonal entry. + mass_diagonal: Option>, + /// Shared with every solve's `Equations`, whose operator views borrow them; + /// `Arc` because this handle is `Send + Sync` and outlives every solve. + jac_sparsity: Arc, + mass_sparsity: Arc, + /// M's values in `mass_sparsity` order, precomputed once per problem. + mass_csc_values: Arc<[f64]>, + context: FaerContext, + n_states: usize, + n_params: usize, + n_events: usize, + n_event_outputs: usize, + n_outputs: usize, + /// `0..n_params`, the identity subset the plain solve path uses. + all_param_indices: Arc<[usize]>, + /// The configured sensitivity subset, precomputed so no solve allocates it. + sens_param_indices: Arc<[usize]>, +} + +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.n_states) + .field("n_params", &self.n_params) + .field("n_events", &self.n_events) + .field("n_outputs", &self.n_outputs) + .finish_non_exhaustive() + } +} + +impl PreparedSolver { + /// Build a prepared problem from a compiled model and tolerances. + /// + /// Performs all expensive one-time setup: sparsity conversion and size + /// extraction. Takes the shared artifact, which is what it keeps -- each + /// solve mints its own [`Workspace`] -- though an evaluator converts in. + pub fn new( + model: impl Into>, + rtol: f64, + atol: &[f64], + ) -> Result { + let model = model.into(); + let n_states = model.n_states(); + if atol.len() != n_states { + return Err(CoreError::AtolLength { + got: atol.len(), + expected: n_states, + }); + } + let n_params = model.n_params(); + let n_events = model.n_events(); + let n_event_outputs = model.total_event_len(); + let n_outputs = model.total_output_len(); + + let jac_sparsity = Arc::new(csc_to_faer_sparsity(model.csc_sparsity())); + let (mass_sparsity, mass_csc_values) = csr_mass_to_faer_csc(model.mass_matrix()); + let mass_sparsity = Arc::new(mass_sparsity); + let mass_csc_values: Arc<[f64]> = mass_csc_values.into(); + + let mass = model.mass_matrix(); + if mass.indptr().len() != n_states + 1 { + return Err(CoreError::Csr(format!( + "mass matrix has {} rows but the model has {n_states} states", + mass.indptr().len().saturating_sub(1) + ))); + } + // An all-zero (or empty) mass row is an algebraic state. `all` on an + // empty slice is true, so both cases fall out of the same test. + let algebraic_rows: Vec = (0..n_states) + .map(|i| { + mass.data()[mass.indptr()[i]..mass.indptr()[i + 1]] + .iter() + .all(|v| *v == 0.0) + }) + .collect(); + let mass_diagonal = mass_diagonal(mass); + + let compiled = model; + let all_param_indices: Arc<[usize]> = (0..n_params).collect(); + let sens_param_indices: Arc<[usize]> = compiled.sens_param_indices().into(); + let context = FaerContext::default(); + + Ok(Self { + compiled, + rtol, + atol: atol.to_vec(), + sens_atol_factor: DEFAULT_SENS_ATOL_FACTOR, + options: SolverOptions::default(), + store_yp: false, + algebraic_rows, + mass_diagonal, + jac_sparsity, + mass_sparsity, + mass_csc_values, + context, + n_states, + n_params, + n_events, + n_event_outputs, + n_outputs, + all_param_indices, + sens_param_indices, + }) + } + + /// Build a fresh local `Equations` for one solve. + /// + /// The equations share the immutable compiled model and the supplied solve-local + /// workspace, and carry a clone of the solve's input parameters. + fn build_eqn( + &self, + y0: &[f64], + y0_sens: &[f64], + ws: &Rc>, + with_output: bool, + inputs: &[f64], + sens_params: &Arc<[usize]>, + ) -> Equations { + Equations { + compiled: Arc::clone(&self.compiled), + ws: Rc::clone(ws), + params: inputs.to_vec(), + sens_params: Arc::clone(sens_params), + y0: y0.to_vec(), + y0_sens: 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: self.context, + n_states: self.n_states, + n_event_outputs: self.n_event_outputs, + n_outputs: self.n_outputs, + with_output, + } + } + + /// Build the diffsol problem for one solve. + /// + /// The single place tolerances, `t0`, the parameter vector and + /// [`SolverOptions`] are stamped onto a problem, so a new knob reaches + /// every path at once. + /// + /// `sens_scales` carries the per-parameter scales when forward + /// sensitivities are under error control, and is `None` on the plain paths + /// and on the relaxed retry. + /// + /// # Errors + /// Propagates whichever `build_from_eqn` rejects the equations for. + fn build_problem( + &self, + eqn: Equations, + t0: f64, + p: Vec, + sens_scales: Option>, + ) -> Result, CoreError> { + let mut builder = OdeBuilder::>::new() + .rtol(self.rtol) + .atol(self.atol.clone()); + if let Some(scales) = sens_scales { + builder = builder + .sens_rtol(self.rtol) + .sens_atol(self.sens_atol()) + .param_scales(scales); + } + Ok(self.options.apply(builder.t0(t0).p(p).build_from_eqn(eqn)?)) + } + + /// Set the multiplier applied to `atol` on differential rows when forming + /// the forward-sensitivity tolerance floor. + /// + /// # Errors + /// Returns [`CoreError::SensAtolFactor`] if `factor` is not finite and > 0. + 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. + /// + /// # Errors + /// Returns [`CoreError::SolverOption`] if any option is out of range. + pub fn with_options(mut self, options: SolverOptions) -> Result { + options.validate()?; + self.options = options; + Ok(self) + } + + /// Store the state time derivatives (`yp`) alongside `y` on the + /// state-trajectory paths, giving downstream cubic-Hermite interpolation + /// its knot slopes. Doubles trajectory memory, so it is opt-in; the + /// output-variable paths never store it (they carry no state trajectory), + /// and nor do grids past [`MAX_HERMITE_COLUMNS`]. + #[must_use] + pub const fn with_store_yp(mut self, store_yp: bool) -> Self { + self.store_yp = store_yp; + self + } + + /// Whether a solve reporting `n_columns` columns stores `yp` alongside `y`. + /// + /// Opting in asks for the slopes where they buy accuracy, not for them + /// unconditionally; see [`MAX_HERMITE_COLUMNS`]. + const fn store_yp(&self, n_columns: usize) -> bool { + self.store_yp && n_columns <= MAX_HERMITE_COLUMNS + } + + /// Mask of algebraic states, one entry per state. + fn algebraic_rows(&self) -> &[bool] { + &self.algebraic_rows + } + + /// Per-state sensitivity absolute tolerance. + /// + /// Differential rows carry the tightened floor; algebraic rows keep the + /// state `atol`, because tightening them fails DAE sensitivity solves. + fn sens_atol(&self) -> Vec { + self.atol + .iter() + .zip(self.algebraic_rows()) + .map(|(a, is_algebraic)| { + if *is_algebraic { + *a + } else { + a * self.sens_atol_factor + } + }) + .collect() + } + + /// Per-parameter scales for the sensitivity absolute tolerances, `atol / |scale|`. + /// + /// Uses each parameter's own magnitude, which is the scale `dy/dp_j` is expressed in. + /// Zero and non-finite inputs fall back to 1.0: diffsol rejects them, and a parameter + /// that is currently zero carries no magnitude information to scale by. + /// + /// `sens_inputs` is subset-space (length `k`, already narrowed to the requested + /// sensitivity columns), not the global `n_params`-length input vector. + 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() + } + + /// BDF dense-solve kernel behind [`solve`](Self::solve). + /// + /// Builds a fresh `Workspace`, runs the timed IC / solver-setup sequence, + /// then steps each `t_stop` segment in turn, interpolating every output + /// column straight into the trajectory. diffsol's `solve_dense` would + /// return the same columns via an intermediate zero-initialised dense + /// matrix; at DFN size that allocation and its copy-out are a measurable + /// slice of the whole solve, so the loop writes the final layout directly. + fn run_dense( + &self, + times: TimePlan<'_>, + y0: &[f64], + inputs: &[f64], + ) -> Result<(DenseTrajectory, Option>, SolverStatistics), CoreError> { + let started = std::time::Instant::now(); + let ws = Rc::new(RefCell::new(self.compiled.create_workspace())); + let eqn = self.build_eqn(y0, &[], &ws, false, inputs, &self.all_param_indices); + + let problem = self.build_problem(eqn, times.eval[0], inputs.to_vec(), None)?; + + let ic_start = std::time::Instant::now(); + let state = problem.bdf_state::()?; + let ic_time_secs = ic_start.elapsed().as_secs_f64(); + let setup_start = std::time::Instant::now(); + let mut solver = problem.bdf_solver::(state)?; + let solver_setup_time_secs = setup_start.elapsed().as_secs_f64(); + + let store_yp = self.store_yp(times.eval.len()); + let mut trajectory = + DenseTrajectory::with_capacity(self.n_states, times.eval.len(), store_yp); + let mut interp = self.context.vector_zeros::>(self.n_states); + let mut interp_dy = + store_yp.then(|| self.context.vector_zeros::>(self.n_states)); + let mut restarter = BreakpointRestarter::new(self); + let root = drive_segments( + &mut solver, + times, + |solver, t_next| restarter.restart(solver, t_next), + |solver, column| { + match column { + ColumnSource::Interpolated(t) => { + solver.interpolate_inplace(t, &mut interp)?; + if let Some(dy) = interp_dy.as_mut() { + solver.interpolate_dy_inplace(t, dy)?; + } + trajectory.push_column( + t, + interp.as_slice(), + interp_dy.as_ref().map(VectorHost::as_slice), + ); + }, + ColumnSource::CurrentState(t) => { + // interpolate_dy_inplace rejects t >= state.t; state.dy + // is already wound back to an event root alongside y. + let state = solver.state(); + trajectory.push_column( + t, + state.y.as_slice(), + store_yp.then(|| state.dy.as_slice()), + ); + }, + } + Ok(()) + }, + )?; + + let y_root = root.map(|t_root| { + trajectory.t_event = Some(t_root); + trajectory.flag = 1; + solver.state().y.as_slice().to_vec() + }); + + let mut statistics = SolverStatistics::from(solver.get_statistics()); + statistics.ic_time_secs = ic_time_secs; + statistics.solver_setup_time_secs = solver_setup_time_secs; + statistics.integration_time_secs = started.elapsed().as_secs_f64(); + + Ok((trajectory, y_root, statistics)) + } + + /// Validate caller-supplied solve arguments against the model dimensions. + /// + /// Returns a [`CoreError`] (surfaced as `ValueError` at the Python boundary) + /// for an empty or non-increasing `t_eval`, a mismatched initial state, or a + /// mismatched packed input array, so malformed arguments produce ordinary + /// errors instead of panicking deep inside integration. + /// + /// The ordering check is the one diffsol's `solve_dense` used to run for us; + /// [`drive_segments`] drains the grid in order and would otherwise report a + /// truncated trajectory as a successful solve. + const fn validate_args( + &self, + t_eval: &[f64], + y0: &[f64], + inputs: &[f64], + ) -> Result<(), CoreError> { + if t_eval.is_empty() { + return Err(CoreError::EmptyTimePoints); + } + // Indexed rather than `windows(2)` to stay const; equal times are the + // discontinuity brackets, so only a decrease (or a NaN) is an error. + let mut i = 1; + while i < t_eval.len() { + let (previous, got) = (t_eval[i - 1], t_eval[i]); + if got < previous || got.is_nan() || previous.is_nan() { + return Err(CoreError::UnsortedTimePoints { + index: i, + got, + previous, + }); + } + i += 1; + } + if y0.len() != self.n_states { + return Err(CoreError::Y0Length { + got: y0.len(), + expected: self.n_states, + }); + } + if inputs.len() != self.n_params { + return Err(CoreError::InputsLength { + got: inputs.len(), + expected: self.n_params, + }); + } + Ok(()) + } + + /// Reject an output-variable solve on a model that registered none. + /// + /// The trajectory width comes from `n_outputs`, while the operator the + /// kernels branch on is only minted when there is an output to evaluate; a + /// zero-output model would otherwise report state values under a zero row + /// count. + const fn validate_has_outputs(&self) -> Result<(), CoreError> { + if self.n_outputs == 0 { + return Err(CoreError::NoOutputVariables); + } + Ok(()) + } + + /// Integrate the model, returning whichever payloads `request` asked for. + /// + /// The one solve entry point: the request's payload flags pick the engine, + /// and every combination of them lands in one [`SolveOutcome`]. Each call + /// constructs its own `Workspace` and its own BDF solver, which is what + /// makes concurrent calls through a shared `&self` sound. + pub fn solve( + &self, + request: SolveRequest<'_>, + set: InputSet<'_>, + ) -> Result { + let InputSet { + y0, + inputs, + y0_sens, + } = set; + self.validate_args(request.t_eval, y0, inputs)?; + if request.sensitivities { + self.validate_y0_sens(y0_sens)?; + } + if request.outputs { + self.validate_has_outputs()?; + } + let times = TimePlan::new(request.t_eval, request.t_stop); + + if request.sensitivities { + let (trajectory, sensitivities, y_event, statistics) = + self.run_dense_sensitivities(times, y0, y0_sens, inputs, request.outputs)?; + self.debug_assert_rows(&trajectory, request.outputs); + return Ok(SolveOutcome::from_parts( + trajectory, + y_event, + Some(sensitivities), + statistics, + )); + } + + let (trajectory, y_event, statistics) = if request.outputs { + self.run_dense_outputs(times, y0, inputs)? + } else { + self.run_dense(times, y0, inputs)? + }; + self.debug_assert_rows(&trajectory, request.outputs); + Ok(SolveOutcome::from_parts( + trajectory, y_event, None, statistics, + )) + } + + /// Check a trajectory's rows against what the request asked for. + /// + /// One layout contract covering every payload combination, where the + /// output-variable path used to assert its own and the others none. + fn debug_assert_rows(&self, trajectory: &DenseTrajectory, outputs: bool) { + let expected = if outputs { + self.n_outputs + } else { + self.n_states + }; + debug_assert_eq!( + trajectory.n_rows, expected, + "trajectory row count mismatch: expected {expected}, got {}", + trajectory.n_rows, + ); + } + + /// Run the dense output-variable solve. + /// + /// Integrates plain states and batch-evaluates the output tapes over staged + /// windows of [`OUTPUT_BATCH_LANES`] points, amortising interpreter + /// dispatch; values match the per-point path bitwise. The returned state is + /// always the full state where the trajectory ends (root time on an event, + /// stop time otherwise) — the only full state an outputs-only caller can + /// restart from. + fn run_dense_outputs( + &self, + times: TimePlan<'_>, + y0: &[f64], + inputs: &[f64], + ) -> Result<(DenseTrajectory, Option>, SolverStatistics), CoreError> { + let started = std::time::Instant::now(); + let ws = Rc::new(RefCell::new(self.compiled.create_workspace())); + // with_output stays false: diffsol integrates states, outputs are + // evaluated in windows below. + let eqn = self.build_eqn(y0, &[], &ws, false, inputs, &self.all_param_indices); + + let problem = self.build_problem(eqn, times.eval[0], inputs.to_vec(), None)?; + + let ic_start = std::time::Instant::now(); + let state = problem.bdf_state::()?; + let ic_time_secs = ic_start.elapsed().as_secs_f64(); + let setup_start = std::time::Instant::now(); + let mut solver = problem.bdf_solver::(state)?; + let solver_setup_time_secs = setup_start.elapsed().as_secs_f64(); + + let n_states = self.n_states; + let n_out_total = self.n_outputs; + let mut window = OutputBatchWindow::new(n_states, n_out_total, times.eval.len()); + let mut interp = self.context.vector_zeros::>(n_states); + let mut restarter = BreakpointRestarter::new(self); + let root = drive_segments( + &mut solver, + times, + |solver, t_next| restarter.restart(solver, t_next), + |solver, column| { + match column { + ColumnSource::Interpolated(t) => { + solver.interpolate_inplace(t, &mut interp)?; + window.stage(t, interp.as_slice()); + }, + ColumnSource::CurrentState(t) => { + window.stage(t, solver.state().y.as_slice()); + }, + } + window.flush_if_full(&self.compiled, &mut ws.borrow_mut(), inputs); + Ok(()) + }, + )?; + window.flush(&self.compiled, &mut ws.borrow_mut(), inputs); + + let (t_event, flag) = root.map_or((None, 0), |t_root| (Some(t_root), 1)); + // Always the full state, the only one an outputs-only caller can + // restart from; on a root the state is already wound back to it. + let y_event = Some(solver.state().y.as_slice().to_vec()); + + let mut statistics = SolverStatistics::from(solver.get_statistics()); + statistics.ic_time_secs = ic_time_secs; + statistics.solver_setup_time_secs = solver_setup_time_secs; + statistics.integration_time_secs = started.elapsed().as_secs_f64(); + + let (t, y) = window.into_trajectory(); + Ok(( + DenseTrajectory { + t, + y, + yp: None, + n_rows: n_out_total, + t_event, + flag, + }, + y_event, + statistics, + )) + } + + /// Run the dense sensitivity solve, retrying once with sensitivities + /// excluded from error control if the controlled solve fails. + /// + /// Stiff DAEs can fail under a tightened sensitivity floor; the retry keeps + /// this path no worse than excluding sensitivities from error control + /// entirely, and flags the downgrade in the returned statistics. + #[allow(clippy::type_complexity)] + fn run_dense_sensitivities( + &self, + times: TimePlan<'_>, + y0: &[f64], + y0_sens: &[f64], + inputs: &[f64], + with_output: bool, + ) -> Result< + ( + DenseTrajectory, + Vec>, + Option>, + SolverStatistics, + ), + CoreError, + > { + // Spans both attempts, as a caller-side timer around this call would. + let started = std::time::Instant::now(); + let mut result = + match self.run_dense_sensitivities_inner(times, y0, y0_sens, inputs, with_output, true) + { + Ok(result) => result, + // Only an integration failure is worth retrying; a config error + // fails the relaxed attempt identically and must surface as-is. + Err(controlled) if !matches!(controlled, CoreError::Diffsol(_)) => { + return Err(controlled); + }, + Err(controlled) => { + match self.run_dense_sensitivities_inner( + times, + y0, + y0_sens, + inputs, + with_output, + false, + ) { + Ok((trajectory, sens_flat, y_root, mut statistics)) => { + statistics.sens_error_control_relaxed = true; + (trajectory, sens_flat, y_root, statistics) + }, + // Both attempts failed: the controlled cause is the useful one. + Err(relaxed) => { + return Err(CoreError::SensRetryFailed { + controlled: controlled.to_string(), + relaxed: relaxed.to_string(), + }); + }, + } + }, + }; + result.3.integration_time_secs = started.elapsed().as_secs_f64(); + Ok(result) + } + + /// BDF dense sensitivity-solve kernel behind + /// [`SolveRequest::sensitivities`], reporting output-variable rows rather + /// than states when `with_output` is set. + /// + /// Mirrors [`run_dense`](Self::run_dense) on the diffsol forward-sensitivity + /// chain (`bdf_state_sens` / `bdf_solver_sens`), interpolating each output + /// column and its `k` sensitivity columns straight into the flat trajectory. + /// Returns the trajectory, one flat block per parameter, the full state at + /// the root time (if an event fired), and timing-enriched statistics. + #[allow(clippy::type_complexity)] + fn run_dense_sensitivities_inner( + &self, + times: TimePlan<'_>, + y0: &[f64], + y0_sens: &[f64], + inputs: &[f64], + with_output: bool, + sens_error_control: bool, + ) -> Result< + ( + DenseTrajectory, + Vec>, + Option>, + SolverStatistics, + ), + CoreError, + > { + let sens_params = &self.sens_param_indices; + if sens_params.is_empty() { + return Err(CoreError::NoSensitivityParams); + } + let ws = Rc::new(RefCell::new(self.compiled.create_workspace())); + let eqn = self.build_eqn(y0, y0_sens, &ws, with_output, inputs, sens_params); + // The ops report the subset width, so p and the scales narrow with them. + let sens_inputs: Vec = sens_params.iter().map(|&i| inputs[i]).collect(); + + // param_scales is IDAS's pbar; without it Chen2020 collapses to ~97k steps. + // sens_atol tightens differential rows only, per the tolerance-structure spec. + let scales = sens_error_control.then(|| Self::param_scales(&sens_inputs)); + let problem = self.build_problem(eqn, times.eval[0], sens_inputs, scales)?; + + let ic_start = std::time::Instant::now(); + let state = problem.bdf_state_sens::()?; + let ic_time_secs = ic_start.elapsed().as_secs_f64(); + let setup_start = std::time::Instant::now(); + let mut solver = problem.bdf_solver_sens::(state)?; + let solver_setup_time_secs = setup_start.elapsed().as_secs_f64(); + + let n_columns = sens_params.len(); + debug_assert_eq!( + n_columns, + self.compiled.n_sens_params(), + "sens subset width diverged from the compiled model", + ); + let n_rows = if with_output { + self.n_outputs + } else { + self.n_states + }; + let n_points = times.eval.len(); + // Output mode never stores yp: its trajectory holds outputs, not states. + let store_yp = self.store_yp(n_points) && !with_output; + let mut trajectory = DenseTrajectory::with_capacity(n_rows, n_points, store_yp); + let mut sens_flat: Vec> = vec![Vec::with_capacity(n_rows * n_points); n_columns]; + let mut y_column = self.context.vector_zeros::>(self.n_states); + let mut dy_column = + store_yp.then(|| self.context.vector_zeros::>(self.n_states)); + let mut sens_columns: Vec> = (0..n_columns) + .map(|_| self.context.vector_zeros::>(self.n_states)) + .collect(); + let mut out_values = self.context.vector_zeros::>(self.n_outputs); + let mut out_chain = self.context.vector_zeros::>(self.n_outputs); + let mut out_direct = self.context.vector_zeros::>(self.n_outputs); + let mut unit_param = self.context.vector_zeros::>(n_columns); + + let mut restarter = BreakpointRestarter::new(self); + let root = drive_segments( + &mut solver, + times, + |solver, t_next| restarter.restart_sens(solver, t_next), + |solver, column| { + let t = column.time(); + match column { + ColumnSource::Interpolated(_) => { + solver.interpolate_inplace(t, &mut y_column)?; + solver.interpolate_sens_inplace(t, &mut sens_columns)?; + if let Some(dy) = dy_column.as_mut() { + solver.interpolate_dy_inplace(t, dy)?; + } + }, + ColumnSource::CurrentState(_) => { + let state = solver.state(); + debug_assert_eq!( + state.s.len(), + sens_columns.len(), + "augmented state carries a different sens width" + ); + y_column.copy_from(state.y); + if let Some(dy) = dy_column.as_mut() { + dy.copy_from(state.dy); + } + for (dst, src) in sens_columns.iter_mut().zip(state.s) { + dst.copy_from(src); + } + }, + } + // Output mode reports g and dg/dp + dg/dy · y_s, the chain rule + // diffsol's own dense sensitivity write-out applies. + if let Some(out) = solver.problem().eqn.out() { + out.call_inplace(&y_column, t, &mut out_values); + trajectory.push_column(t, out_values.as_slice(), None); + for (j, (flat, s_j)) in sens_flat.iter_mut().zip(&sens_columns).enumerate() { + out.jac_mul_inplace(&y_column, t, s_j, &mut out_chain); + unit_param.set_index(j, 1.0); + out.sens_mul_inplace(&y_column, t, &unit_param, &mut out_direct); + unit_param.set_index(j, 0.0); + flat.extend( + out_chain + .as_slice() + .iter() + .zip(out_direct.as_slice()) + .map(|(chain, direct)| chain + direct), + ); + } + } else { + trajectory.push_column( + t, + y_column.as_slice(), + dy_column.as_ref().map(VectorHost::as_slice), + ); + for (flat, s_j) in sens_flat.iter_mut().zip(&sens_columns) { + flat.extend_from_slice(s_j.as_slice()); + } + } + Ok(()) + }, + )?; + + let mut y_root = root.map(|t_root| { + trajectory.t_event = Some(t_root); + trajectory.flag = 1; + solver.state().y.as_slice().to_vec() + }); + // In output mode the trajectory never holds states, so the terminal + // state is the caller's only one. + if with_output && y_root.is_none() { + y_root = Some(solver.state().y.as_slice().to_vec()); + } + let mut statistics = SolverStatistics::from(solver.get_statistics()); + statistics.ic_time_secs = ic_time_secs; + statistics.solver_setup_time_secs = solver_setup_time_secs; + + Ok((trajectory, sens_flat, y_root, statistics)) + } + + /// Validate the caller-supplied `dy0/dp` seed. + /// + /// Empty is the "no parameter reaches an initial condition" case; anything + /// else must be a column-major `n_states x k` block over the requested + /// sensitivity subset. + /// + /// A model compiled without sensitivity parameters is reported as such + /// first: its expected width is 0, so the length complaint would otherwise + /// mask the real configuration mistake behind "must be empty or 0". + fn validate_y0_sens(&self, y0_sens: &[f64]) -> Result<(), CoreError> { + if self.sens_param_indices.is_empty() { + return Err(CoreError::NoSensitivityParams); + } + let expected = self.n_states * self.sens_param_indices.len(); + if y0_sens.is_empty() || y0_sens.len() == expected { + return Ok(()); + } + Err(CoreError::Y0SensLength { + got: y0_sens.len(), + expected, + }) + } + + /// Flatten a dense matrix `[nrows × ncols]` in column-major order: + /// `out[col * nrows + row]` (time-outer/row-inner). + /// + /// Only the solve paths that still call diffsol's `solve_dense` need this, + /// which is now just the LU-equivalence test. + #[cfg(test)] + fn dense_to_column_major( + m: & as diffsol::vector::DefaultDenseMatrix>::M, + ) -> Vec { + use diffsol::MatrixCommon; + + let mut flat = Vec::with_capacity(m.nrows() * m.ncols()); + for j in 0..m.ncols() { + flat.extend_from_slice(m.inner().col_as_slice(j)); + } + flat + } +} + +/// How many interpolation points are staged before one batched output +/// evaluation. Bounds the extra state storage at `n_states * OUTPUT_BATCH_LANES` +/// while still amortising per-instruction dispatch across the window. +const OUTPUT_BATCH_LANES: usize = 128; + +/// Staging buffer for the batched output-variable path. +/// +/// Collects interpolated `(t, y)` columns until [`OUTPUT_BATCH_LANES`] are +/// pending, then evaluates every output tape across the window in one +/// `eval_batch` pass and appends the results to the growing trajectory. +struct OutputBatchWindow { + n_states: usize, + n_out_total: usize, + staged_t: Vec, + /// `(n_states, k)` F-contiguous staged states. + staged_y: Vec, + /// `(n_out_total, k)` F-contiguous per-window results. + window_out: Vec, + t_out: Vec, + outputs: Vec, +} + +impl OutputBatchWindow { + fn new(n_states: usize, n_out_total: usize, n_points_hint: usize) -> Self { + Self { + n_states, + n_out_total, + staged_t: Vec::with_capacity(OUTPUT_BATCH_LANES), + staged_y: Vec::with_capacity(n_states * OUTPUT_BATCH_LANES), + window_out: vec![0.0; n_out_total * OUTPUT_BATCH_LANES], + t_out: Vec::with_capacity(n_points_hint), + outputs: Vec::with_capacity(n_out_total * n_points_hint), + } + } + + fn stage(&mut self, t: f64, y: &[f64]) { + debug_assert!(self.staged_t.len() < OUTPUT_BATCH_LANES); + self.staged_t.push(t); + self.staged_y.extend_from_slice(&y[..self.n_states]); + } + + fn flush_if_full(&mut self, compiled: &CompiledModel, ws: &mut Workspace, inputs: &[f64]) { + if self.staged_t.len() == OUTPUT_BATCH_LANES { + self.flush(compiled, ws, inputs); + } + } + + fn flush(&mut self, compiled: &CompiledModel, ws: &mut Workspace, inputs: &[f64]) { + let k = self.staged_t.len(); + if k == 0 { + return; + } + compiled.eval_outputs_batch( + ws, + k, + &self.staged_t, + &self.staged_y, + inputs, + &mut self.window_out[..self.n_out_total * k], + ); + self.outputs + .extend_from_slice(&self.window_out[..self.n_out_total * k]); + self.t_out.extend_from_slice(&self.staged_t); + self.staged_t.clear(); + self.staged_y.clear(); + } + + /// The accumulated `(t, outputs)` trajectory; every window must be flushed. + fn into_trajectory(self) -> (Vec, Vec) { + debug_assert!(self.staged_t.is_empty(), "unflushed output window"); + (self.t_out, self.outputs) + } +} + +/// Prepare and run a one-shot solve. +/// +/// Setup is discarded afterwards, so callers that solve the same model more than +/// once should keep a [`PreparedSolver`] instead. `atol` is per state. +pub fn solve( + model: impl Into>, + t_eval: &[f64], + t_stop: &[f64], + y0: &[f64], + inputs: &[f64], + rtol: f64, + atol: &[f64], +) -> Result { + let prepared = PreparedSolver::new(model, rtol, atol)?; + prepared.solve( + SolveRequest::new(t_eval).with_stop_times(t_stop), + InputSet::new(y0, inputs), + ) +} + +/// M's diagonal, or `None` if any stored nonzero sits off it. +/// +/// Only a diagonal M makes `M dy = f` an element-wise divide in +/// [`BreakpointRestarter`]; `PyBaMM`'s discretisation always produces one. +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 reusable half of a breakpoint restart. +/// +/// One of these lives for a whole solve. Its Newton solvers each keep a +/// [`ReusedFaerLu`], whose entire purpose is to hold a symbolic factorisation +/// across refreshes; building them per segment threw that away, and `PyBaMM` +/// hands down one stop time per output point, so a dense grid restarts +/// hundreds of times. The two solvers stay separate because they factorise +/// different sparsity patterns and sharing one would thrash both. +struct BreakpointRestarter { + root_solver: NewtonNonlinearSolver, ReusedFaerLu, NoLineSearch>, + sens_solver: NewtonNonlinearSolver, ReusedFaerLu, NoLineSearch>, + /// M's diagonal, only when the model has no algebraic state and M is + /// diagonal; `None` leaves `dy` to `set_consistent`. + ode_mass_diagonal: Option>, + /// Holds `f(t, y)` while `dy` is rebuilt from it. + scratch: FaerVec, +} + +impl BreakpointRestarter { + fn new(problem: &PreparedSolver) -> Self { + let has_algebraic = problem + .algebraic_rows + .iter() + .any(|is_algebraic| *is_algebraic); + Self { + root_solver: NewtonNonlinearSolver::new(ReusedFaerLu::default(), NoLineSearch), + sens_solver: NewtonNonlinearSolver::new(ReusedFaerLu::default(), NoLineSearch), + ode_mass_diagonal: (!has_algebraic) + .then(|| problem.mass_diagonal.clone()) + .flatten(), + scratch: problem + .context + .vector_zeros::>(problem.n_states), + } + } + + /// Rebuild `dy` from `M dy = f(t, y)` when nothing else will. + /// + /// diffsol's `set_consistent` returns early once it finds no zero diagonal + /// in M, leaving `dy` holding the slope from *before* the discontinuity. + /// The next step seeds the BDF difference array straight from it + /// (`initialise_diff_to_first_order`), so without this the restart + /// re-derives the old branch — the exact failure the restart exists to + /// prevent, fixed for DAEs and silently broken for ODEs. + /// + /// `t_next` is the incoming segment's first output time. `PyBaMM` brackets + /// a discontinuity with a pair of times one ULP apart and the restart sits + /// on the earlier one, so `f(state.t)` is still the old branch; when + /// `t_next` is within round-off it is the far side of the corner, the + /// branch the segment actually integrates. + fn refresh_ode_dy( + &mut self, + problem: &OdeSolverProblem, + state: &mut StateRefMut<'_, FaerVec>, + t_next: f64, + ) where + Eqn: OdeEquationsImplicit, M = FaerSparseMat, C = FaerContext>, + { + let Some(diagonal) = self.ode_mass_diagonal.as_deref() else { + return; + }; + let t = if stop_already_reached(*state.t, t_next) { + t_next + } else { + *state.t + }; + problem + .eqn + .rhs() + .call_inplace(state.y, t, &mut self.scratch); + for (i, m) in diagonal.iter().enumerate() { + state.dy[i] = self.scratch[i] / m; + } + } + + /// Restart the integrator at a breakpoint, as IDAKLU's `HandleBreakpoint` does. + /// + /// Recomputing the consistent state (the diffsol equivalent of `IDACalcIC` + /// with `IDA_YA_YDP_INIT`) and a fresh step size is what makes the restart + /// land on the new branch of the solution. + fn restart<'a, Eqn, S>(&mut self, solver: &mut S, t_next: f64) -> Result<(), CoreError> + where + Eqn: OdeEquationsImplicit, M = FaerSparseMat, C = FaerContext> + + 'a, + S: OdeSolverMethod<'a, Eqn>, + { + // Taken before the mutable borrow; the problem outlives the solver. + let problem = solver.problem(); + let mut state = solver.state_mut(); + state.set_consistent(problem, &mut self.root_solver)?; + self.refresh_ode_dy(problem, &mut state, t_next); + state.set_step_size(problem.h0, &problem.atol, problem.rtol, &problem.eqn, 1); + Ok(()) + } + + /// [`Self::restart`] for a solve carrying forward sensitivities. + /// + /// The sensitivity difference arrays are seeded from `ds`, which + /// `set_consistent_augmented` rebuilds from `y` and `dy`; running it after + /// the `dy` refresh is what keeps the two consistent at the corner. + fn restart_sens<'a, Eqn, AugEqn, S>( + &mut self, + solver: &mut S, + t_next: f64, + ) -> Result<(), CoreError> + where + Eqn: OdeEquationsImplicit, M = FaerSparseMat, C = FaerContext> + + 'a, + AugEqn: AugmentedOdeEquationsImplicit + std::fmt::Debug, + S: AugmentedOdeSolverMethod<'a, Eqn, AugEqn>, + { + let problem = solver.problem(); + { + let mut state = solver.state_mut(); + state.set_consistent(problem, &mut self.root_solver)?; + self.refresh_ode_dy(problem, &mut state, t_next); + } + + // Reads the y and dy just made consistent, so it has to follow them. + if let Some((mut state, augmented_eqn)) = solver.state_and_augmented_eqn_mut() { + state.set_consistent_augmented(problem, augmented_eqn, &mut self.sens_solver)?; + } + + solver + .state_mut() + .set_step_size(problem.h0, &problem.atol, problem.rtol, &problem.eqn, 1); + Ok(()) + } +} + +/// The times a solve reports at, together with the stop times inside them. +/// +/// `stop` is a subset of `eval`: a stop time absent from the output grid cannot +/// end a segment, because the segment's last output time is what the integrator +/// is told to stop on. +#[derive(Clone, Copy, Debug)] +struct TimePlan<'a> { + eval: &'a [f64], + stop: &'a [f64], +} + +impl<'a> TimePlan<'a> { + const fn new(eval: &'a [f64], stop: &'a [f64]) -> Self { + Self { eval, stop } + } + + fn segments(&self) -> Vec> { + segment_ranges(self.eval, self.stop) + } +} + +/// Where one output column's values come from. +/// +/// Columns inside a step are interpolated. The two that land where the +/// integrator already sits are read off the state instead: the final column of +/// a segment, which diffsol may report as reached from either side of its +/// round-off window, and an event's column after the state is wound back to it. +#[derive(Clone, Copy, Debug)] +enum ColumnSource { + Interpolated(f64), + CurrentState(f64), +} + +impl ColumnSource { + const fn time(self) -> f64 { + match self { + Self::Interpolated(t) | Self::CurrentState(t) => t, + } + } +} + +/// Step one `t_stop` segment at a time, handing every output column to `emit`. +/// +/// Sole home of the stop-time round-off rule and the event semantics, so all +/// three solve lanes report the same times. Returns the root time if an event +/// ended the solve, with the solver wound back to it. +/// +/// diffsol's `solve_dense` would route the same columns through a +/// zero-initialised `n_states x n_times` matrix and copy it out again. +fn drive_segments<'a, Eqn, S>( + solver: &mut S, + times: TimePlan<'_>, + mut restart: impl FnMut(&mut S, f64) -> Result<(), CoreError>, + mut emit: impl FnMut(&S, ColumnSource) -> Result<(), CoreError>, +) -> Result, CoreError> +where + Eqn: OdeEquationsImplicit, M = FaerSparseMat, C = FaerContext> + + 'a, + S: OdeSolverMethod<'a, Eqn>, +{ + debug_assert!( + solver.problem().eqn.reset().is_none(), + "a reset operator would need the apply_reset branch solve_dense has", + ); + 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 incoming segment's first output time: the far side of a + // ULP-bracketed corner, where the new branch is evaluable. + restart(solver, t_eval[range.start])?; + } + 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 { + emit(solver, ColumnSource::Interpolated(t_eval[col]))?; + col += 1; + } + match reason { + OdeSolverStopReason::InternalTimestep => {}, + OdeSolverStopReason::TstopReached => { + // diffsol leaves `state.t` a round-off short of tstop, so + // the next segment would extrapolate this column. + if col == last { + emit(solver, ColumnSource::CurrentState(t_eval[last]))?; + col += 1; + } + debug_assert_eq!(col, range.end, "segment ended with undrained output points"); + break; + }, + OdeSolverStopReason::RootFound(t_root, _) => { + solver.state_mut_back(t_root)?; + // A root column only when output points remain, so a root + // at the segment's last time adds none. + if col <= last { + emit(solver, ColumnSource::CurrentState(t_root))?; + } + return Ok(Some(t_root)); + }, + } + } + } + Ok(None) +} + +/// Whether diffsol would already consider a stop at `to` reached from `from`. +/// +/// Mirrors the round-off window in its `handle_tstop`, where asking to stop is +/// an error rather than a no-op. `PyBaMM` brackets each constant-time +/// discontinuity with a pair of times one ULP apart, so the pair has to collapse +/// to the earlier one: that is the side the pre-discontinuity branch holds on. +fn stop_already_reached(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 only ends a segment when it is itself an output time. The Python +/// layer guarantees that by unioning `t_eval` into the output grid, so a stop +/// time that is missing from it means the two disagree, and integrating straight +/// through is safer than shifting an output column off the time it belongs to. +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` is skipped rather than + // reversing a range into a slice panic. + let previous = cuts.last().copied().unwrap_or(0); + if idx <= previous { + continue; + } + if stop_already_reached(t_eval[previous], t_eval[idx]) + || stop_already_reached(t_eval[idx], t_eval[last]) + { + 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`. +/// +/// The chord's error falls as the column count squared, so past a few thousand +/// columns it reaches the integration error floor Hermite already sits on while +/// the store's cost keeps growing linearly. +const MAX_HERMITE_COLUMNS: usize = 4096; + +/// A column-major trajectory accumulated one output column at a time. +/// +/// Segments are solved one at a time so the BDF can restart at each stop time; +/// their columns land here in order, together with the root-time bookkeeping +/// when an event ends the solve. +struct DenseTrajectory { + t: Vec, + y: Vec, + /// `dy/dt` columns sharing `y`'s layout; `Some` only when the solve + /// stores them for downstream Hermite interpolation. + yp: Option>, + /// Trajectory width: states, or outputs on the output paths. + n_rows: usize, + t_event: Option, + flag: i32, +} + +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, + t_event: None, + flag: 0, + } + } + + 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); + } + } + + const fn n_cols(&self) -> usize { + self.t.len() + } +} + +#[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() { + // PyBaMM brackets `t < 2` with 2.0 and its neighbour one ULP below; + // diffsol rejects the second as a stop time it has 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_stop_time_a_hair_before_the_end_leaves_no_degenerate_segment() { + 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..4]); + } +} + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + use std::rc::Rc; + + use diffsol::ode_solver::sensitivities::SensitivitiesOdeSolverMethod; + use diffsol::{Matrix, NonLinearOpSens, OdeBuilder, OdeEquations, Op, Vector}; + + use super::*; + use crate::arena::Arena; + use crate::model::{CompiledModelOptions, ModelEvaluator}; + use crate::node::{CsrData, Node, Shape}; + + /// Build a small DAE model with one parameter: + /// `dy/dt = -a * y` (`a = inputs[0]`) + /// `0 = y - z` (algebraic: z follows y) + /// Mass = diag(1, 0), `n_states` = 2, `n_params` = 1, sens wrt param 0. + #[cfg(test)] + pub fn build_small_dae_with_param() -> ModelEvaluator { + 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 = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 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)) + }; + // Algebraic residual: y - z = 0 + let algebraic_residual = arena.alloc(Node::Sub(sv0, sv1)); + let rhs = arena.alloc(Node::Concat(vec![neg_a_y, algebraic_residual])); + + let mass = CsrData { + indptr: vec![0, 1, 1], // row 0: (0,0)=1; row 1: empty + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(2, 2), + }; + + ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 2, + 1, + CompiledModelOptions::new().with_sensitivities(&[0]), + ) + } + + /// Like `build_small_dae_with_param` but with a second input `b` in the ODE + /// (`dy/dt = -a*y + b`), with sensitivities requested for `sens`. + fn build_small_dae_two_inputs(sens: &[usize]) -> ModelEvaluator { + 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 = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b = arena.alloc(Node::InputParameter { + name: "b".to_string(), + index: 1, + offset: 1, + width: 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])); + + let mass = CsrData { + indptr: vec![0, 1, 1], // row 0: (0,0)=1; row 1: empty + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(2, 2), + }; + + ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 2, + 2, + CompiledModelOptions::new().with_sensitivities(sens), + ) + } + + /// The two-input fixture with sensitivities requested for `a` (index 0) only. + fn build_small_dae_two_inputs_one_sens() -> ModelEvaluator { + build_small_dae_two_inputs(&[0]) + } + + /// Return a `PreparedSolver` ready for a sensitivity solve on the small DAE. + #[cfg(test)] + pub 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") + } + + /// The sensitivity blocks of an outcome whose request asked for them. + #[cfg(test)] + fn blocks(outcome: &SolveOutcome) -> &[Vec] { + outcome + .sensitivities + .as_deref() + .expect("the request asked for sensitivities") + } + + /// Return `(y0, inputs, t_eval)` for the small-DAE fixture. + #[cfg(test)] + pub fn small_dae_setup() -> (Vec, Vec, Vec) { + let y0 = vec![1.0, 0.0]; // algebraic IC deliberately inconsistent + let inputs = vec![1.0]; // a = 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`, identity mass. No algebraic row, which is + /// the branch `set_consistent` returns early from. + fn build_small_ode_with_param() -> ModelEvaluator { + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + 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, sv0)); + + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(1, 1), + }; + + ModelEvaluator::new_with_options(&arena, rhs, mass, 1, 1, CompiledModelOptions::new()) + } + + #[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 ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + let eqn = prepared.build_eqn(&y0, &[], &ws, false, &inputs, &prepared.all_param_indices); + let problem = prepared.options.apply( + OdeBuilder::>::new() + .rtol(prepared.rtol) + .atol(prepared.atol.clone()) + .t0(0.0) + .p(inputs.to_vec()) + .build_from_eqn(eqn) + .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 difference-array slope rather + // than the seed it was initialised with. + solver.solve_dense(&[0.0, 0.5, 1.0]).expect("solve"); + + let stale = solver.state().dy[0]; + BreakpointRestarter::new(&prepared) + .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] + fn solve_empty_t_eval_returns_error_not_panic() { + let model = build_small_dae_with_param(); + let (y0, inputs, _t) = small_dae_setup(); + let err = solve(model, &[], &[], &y0, &inputs, 1e-8, &[1e-8, 1e-8]) + .expect_err("empty t_eval must be a returned error, not a panic"); + assert!( + matches!(err, CoreError::EmptyTimePoints), + "expected EmptyTimePoints, got: {err:?}" + ); + } + + #[test] + fn solve_decreasing_t_eval_returns_error_not_a_truncated_solve() { + // The segment loop drains the grid in order, so without this check a + // decrease silently reports the points it did reach as a whole solve. + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _t) = small_dae_setup(); + for t_eval in [ + vec![0.0, 2.0, 1.0], + vec![0.0, 0.5, 0.2], + vec![0.0, f64::NAN, 1.0], + ] { + let err = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &inputs)) + .expect_err("a decreasing t_eval must be an error, not a short trajectory"); + assert!( + matches!(err, CoreError::UnsortedTimePoints { .. }), + "expected UnsortedTimePoints for {t_eval:?}, got: {err:?}" + ); + } + } + + #[test] + fn output_solves_reject_a_model_without_output_variables() { + // Without outputs the kernels fall through to the state branch, which + // would report state values under a zero row count. + let prepared = build_small_dae_prepared_with_sens(); // no add_output + 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() { + // Equal consecutive times bracket a discontinuity; only a decrease is + // rejected. + 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.n_times, t_eval.len()); + } + + #[test] + fn a_sensitivity_request_rejects_a_decreasing_t_eval() { + // Every public entry point validates through the same helper. + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, _t) = small_dae_setup(); + let err = prepared + .solve( + SolveRequest::new(&[0.0, 2.0, 1.0]).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect_err("a decreasing t_eval must be an error on the sensitivity path too"); + assert!( + matches!(err, CoreError::UnsortedTimePoints { .. }), + "expected UnsortedTimePoints, got: {err:?}" + ); + } + + #[test] + fn solve_y0_length_mismatch_returns_error_not_panic() { + let model = build_small_dae_with_param(); // n_states = 2 + let (_y0, inputs, t_eval) = small_dae_setup(); + let err = solve(model, &t_eval, &[], &[1.0], &inputs, 1e-8, &[1e-8, 1e-8]) + .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 model = build_small_dae_with_param(); // needs 1 packed input + let (y0, _inputs, t_eval) = small_dae_setup(); + let err = solve(model, &t_eval, &[], &y0, &[], 1e-8, &[1e-8, 1e-8]) + .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 equations_implements_implicit_sens_and_solves() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, t_eval) = small_dae_setup(); + + let ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + let eqn = prepared.build_eqn(&y0, &[], &ws, false, &inputs, &prepared.all_param_indices); + + let problem = OdeBuilder::>::new() + .rtol(1e-8) + .atol(vec![1e-8, 1e-8]) + .t0(t_eval[0]) + .p(inputs.clone()) + .build_from_eqn(eqn) + .expect("build_from_eqn failed"); + + let state = problem + .bdf_state_sens::() + .expect("bdf_state_sens failed"); + let mut solver = problem + .bdf_solver_sens::(state) + .expect("bdf_solver_sens failed"); + + let (_y, ys, _stop) = solver + .solve_dense_sensitivities(&t_eval) + .expect("solve_dense_sensitivities failed"); + + // One sensitivity matrix per parameter + assert_eq!(ys.len(), inputs.len()); + } + + #[test] + fn reused_lu_end_to_end_matches_stock_lu() { + // Same factorisation algorithm, same arithmetic: swapping the stock + // solver for the buffer-reusing one must not move a single bit. + let run = |stock: bool| -> Vec { + let model = build_small_dae_with_param(); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let (y0, inputs, t_eval) = small_dae_setup(); + let ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + let eqn = + prepared.build_eqn(&y0, &[], &ws, false, &inputs, &prepared.all_param_indices); + let problem = OdeBuilder::>::new() + .rtol(1e-8) + .atol(vec![1e-8, 1e-8]) + .t0(t_eval[0]) + .p(inputs) + .build_from_eqn(eqn) + .expect("build_from_eqn failed"); + let y_mat = if stock { + type Stock = diffsol::FaerSparseLU; + let state = problem.bdf_state::().expect("bdf_state failed"); + let mut solver = problem + .bdf_solver::(state) + .expect("bdf_solver failed"); + solver.solve_dense(&t_eval).expect("solve_dense failed").0 + } else { + let state = problem + .bdf_state::() + .expect("bdf_state failed"); + let mut solver = problem + .bdf_solver::(state) + .expect("bdf_solver failed"); + solver.solve_dense(&t_eval).expect("solve_dense failed").0 + }; + PreparedSolver::dense_to_column_major(&y_mat) + }; + assert_eq!(run(false), run(true)); + } + + #[test] + fn sens_solve_integrates_only_the_requested_column() { + // Two inputs, sensitivities requested for `a` alone: one integrated column, and + // it has to be d/da rather than the d/db a mis-mapped subset would hand back. + let model = build_small_dae_two_inputs_one_sens(); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let y0 = vec![1.0, 0.0]; + let inputs = vec![1.0, 0.5]; + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + + let res = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("sensitivity solve failed"); + assert_eq!(blocks(&res).len(), 1, "2 inputs, 1 requested"); + + let h = 1e-6; + let mut ip = inputs.clone(); + ip[0] += h; + let rp = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &ip)) + .unwrap(); + let mut im = inputs; + im[0] -= h; + let rm = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &im)) + .unwrap(); + let n = res.n_rows; + let last = (res.n_times - 1) * n; + for j in 0..n { + let fd = (rp.y[last + j] - rm.y[last + j]) / (2.0 * h); + let got = blocks(&res)[0][last + j]; + assert!((got - fd).abs() < 1e-4, "d/da[{j}]: got {got} fd {fd}"); + } + } + + #[test] + fn sens_solve_integrates_a_non_prefix_subset_column() { + // A prefix subset would still pass under a truncating seed; starting past + // position 0 is what makes the end-to-end solve name d/db rather than d/da. + let model = build_small_dae_two_inputs(&[1]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let y0 = vec![1.0, 0.0]; + let inputs = vec![1.0, 0.5]; + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + + let res = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("sensitivity solve failed"); + assert_eq!(blocks(&res).len(), 1, "2 inputs, 1 requested"); + + let h = 1e-6; + let mut ip = inputs.clone(); + ip[1] += h; + let rp = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &ip)) + .unwrap(); + let mut im = inputs; + im[1] -= h; + let rm = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &im)) + .unwrap(); + let n = res.n_rows; + let last = (res.n_times - 1) * n; + for j in 0..n { + let fd = (rp.y[last + j] - rm.y[last + j]) / (2.0 * h); + let got = blocks(&res)[0][last + j]; + assert!((got - fd).abs() < 1e-4, "d/db[{j}]: got {got} fd {fd}"); + } + } + + #[test] + fn sens_solve_without_requested_params_is_an_error() { + // Nothing to differentiate is a caller mistake, not a zero-column solve: the + // wrapper has no columns to present and the augmented state has no shape. + let prepared = build_decay_without_sens_params(); + let err = prepared + .solve( + SolveRequest::new(&[0.0, 0.1]).with_sensitivities(), + InputSet::new(&[1.0], &[]), + ) + .expect_err("a solve with no requested sensitivities must be an error"); + // The config error must surface as-is: routing it through the relaxed + // retry would double-report it as a misleading SensRetryFailed. + assert!( + matches!(err, CoreError::NoSensitivityParams), + "expected NoSensitivityParams, got: {err:?}" + ); + } + + /// `dy/dt = -y`, one state, compiled with no sensitivity parameters. + fn build_decay_without_sens_params() -> PreparedSolver { + 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 mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(1, 1), + }; + let model = ModelEvaluator::new(&arena, rhs, mass, 1, 0); + PreparedSolver::new(model, 1e-8, &[1e-8]).expect("PreparedSolver failed") + } + + #[test] + fn a_seed_without_requested_params_reports_the_missing_params_not_its_width() { + // Expected width is 0 here, so the length complaint ("must be empty or + // 0") would describe the seed instead of the real mistake. + let prepared = build_decay_without_sens_params(); + let err = prepared + .solve( + SolveRequest::new(&[0.0, 0.1]).with_sensitivities(), + InputSet::new(&[1.0], &[]).with_sens_seed(&[0.5]), + ) + .expect_err("a non-empty seed with no requested sensitivities must be an error"); + assert!( + matches!(err, CoreError::NoSensitivityParams), + "expected NoSensitivityParams, got: {err:?}" + ); + } + + #[test] + fn a_doubly_misconfigured_request_reports_the_seed_before_the_row_space() { + // One entry point validates for every payload combination, so the order + // it checks in is a contract: the missing sensitivity parameters are the + // reason this call cannot run at all, while the absent output variables + // are only reachable once it can. + let prepared = build_decay_without_sens_params(); // no sens params, no outputs + let err = prepared + .solve( + SolveRequest::new(&[0.0, 0.1]) + .with_outputs() + .with_sensitivities(), + InputSet::new(&[1.0], &[]), + ) + .expect_err("a request for payloads the model cannot supply must be an error"); + assert!( + matches!(err, CoreError::NoSensitivityParams), + "expected NoSensitivityParams, got: {err:?}" + ); + } + + /// `dy/dt = -a*y` (`a = inputs[0]`), output `2*y`, event `y - 0.4`, + /// sensitivities wrt `a`. Analytic: `y = exp(-a*t)`, `dy/da = -t*exp(-a*t)`. + fn build_decay_param_output_event() -> PreparedSolver { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 1 }); + 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, sv)); + let two = arena.alloc(Node::Scalar(2.0)); + let output = arena.alloc(Node::Mul(two, sv)); + let threshold = arena.alloc(Node::Scalar(0.4)); + let event = arena.alloc(Node::Sub(sv, threshold)); + let mass = CsrData { + indptr: vec![0, 1], + indices: vec![0], + data: vec![1.0], + shape: Shape::matrix(1, 1), + }; + let mut model = ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 1, + 1, + CompiledModelOptions::new().with_sensitivities(&[0]), + ); + model.add_output(&arena, output); + model.add_event(&arena, event); + PreparedSolver::new(model, 1e-10, &[1e-10]).expect("PreparedSolver failed") + } + + /// Two states, three outputs and one sensitivity parameter, with an event: + /// the row spaces differ, so a payload reported under the wrong row count is + /// visible. `store_yp` is on so the `yp` rule has something to report. + 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 = 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 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 = CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![1.0, 1.0], + shape: Shape::matrix(2, 2), + }; + let mut model = ModelEvaluator::new_with_options( + &arena, + rhs, + mass, + 2, + 1, + CompiledModelOptions::new().with_sensitivities(&[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.add_output(&arena, out0); + model.add_output(&arena, out1); + model.add_output(&arena, out2); + let threshold = arena.alloc(Node::Scalar(0.4)); + let event = arena.alloc(Node::Sub(y0, threshold)); + model.add_event(&arena, event); + + PreparedSolver::new(model, 1e-10, &[1e-10, 1e-10]) + .expect("PreparedSolver failed") + .with_store_yp(true) + } + + /// One outcome type means one layout contract, so it can be asserted once + /// over every payload combination instead of restated per result type — the + /// prose invariant that the four types stopped honouring when `yp` landed on + /// two of them. + #[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 and has a `t_event`/`y_event` to report. + 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.n_times, outcome.t.len(), "{label}: n_times"); + assert_eq!( + outcome.y.len(), + outcome.n_rows * outcome.n_times, + "{label}: trajectory size" + ); + + // yp is the state trajectory's slopes, so it is present exactly + // when the rows are states and the solver stores them. + 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"); + } + } + + assert_eq!(outcome.flag, 1, "{label}: the event should have fired"); + let t_event = outcome.t_event.expect("t_event missing"); + assert!( + (t_event - outcome.t[outcome.n_times - 1]).abs() < 1e-12, + "{label}: the trajectory should end at the root" + ); + // Always a full state, never an outputs row: the only thing a + // caller can restart from. + 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"); + } + } + } + + /// The two requests that differ only in their row space share the + /// integration, so their termination fields agree bit for bit rather than + /// merely in shape. + #[test] + // Bit-identical is the property under test, so the comparison is exact. + #[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.n_times, outputs.n_times); + assert_eq!(states.flag, outputs.flag); + assert_eq!(states.t_event, outputs.t_event); + assert_eq!(states.y_event, outputs.y_event); + assert_eq!( + states.statistics.number_of_steps, + outputs.statistics.number_of_steps + ); + } + + #[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.flag, 0); + assert_eq!(r.n_rows, 1); + assert_eq!(blocks(&r).len(), 1); + assert_eq!(blocks(&r)[0].len(), r.n_times); + 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}" + ); + } + // Final-time outputs solves must still carry a restartable full state. + assert!(r.t_event.is_none()); + 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() { + // Every column before the root keeps its requested time and the root adds + // exactly one; nothing is relabelled to the root time or dropped. + 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"); + + assert_eq!(r.flag, 1, "expected event termination"); + let t_event = r.t_event.expect("t_event missing"); + let drained = t_eval.iter().filter(|&&t| t <= t_event).count(); + assert_eq!( + r.n_times, + drained + 1, + "one root column past the grid points" + ); + assert_eq!(r.t.len(), r.n_times); + 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 store_yp_fills_the_polynomial_derivative_on_the_state_path() { + // dy/dt = -a*y with z = y, so every knot's slope is -a*y(t) on both + // rows; the algebraic row differentiates the same polynomial in z. + 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 result = prepared + .solve(SolveRequest::new(&t_eval), 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.n_times { + 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}", + ); + } + } + } + + #[test] + fn every_solve_path_reports_its_own_integration_time() { + // All four payload combinations, so a kernel that stopped stamping would + // report a zero the caller could not tell from a very fast solve. + let prepared = build_decay_param_output_event(); + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.1).collect(); + let statistics = [ + ( + "states", + prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&[1.0], &[1.0])) + .expect("solve") + .statistics, + ), + ( + "outputs", + prepared + .solve( + SolveRequest::new(&t_eval).with_outputs(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("output solve") + .statistics, + ), + ( + "sensitivities", + prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("sensitivity solve") + .statistics, + ), + ( + "outputs + sensitivities", + prepared + .solve( + SolveRequest::new(&t_eval) + .with_outputs() + .with_sensitivities(), + InputSet::new(&[1.0], &[1.0]), + ) + .expect("output sensitivity solve") + .statistics, + ), + ]; + + for (name, statistics) in statistics { + assert!( + statistics.integration_time_secs > 0.0, + "{name} reported no integration time", + ); + assert!( + statistics.integration_time_secs + >= statistics.ic_time_secs + statistics.solver_setup_time_secs, + "{name} reported less time than the phases it contains", + ); + } + } + + /// A uniform grid of `n` columns over `[0, 1]`, the small-DAE span. + 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 a_grid_at_the_hermite_column_limit_still_stores_yp() { + let prepared = build_small_dae_prepared_with_sens().with_store_yp(true); + let (y0, inputs, _) = small_dae_setup(); + let result = prepared + .solve( + SolveRequest::new(&grid(MAX_HERMITE_COLUMNS)), + InputSet::new(&y0, &inputs), + ) + .expect("solve"); + assert_eq!(result.n_times, MAX_HERMITE_COLUMNS); + assert!(result.yp.is_some(), "the limit itself must keep yp"); + } + + #[test] + fn a_grid_past_the_hermite_column_limit_gives_up_yp() { + let prepared = build_small_dae_prepared_with_sens().with_store_yp(true); + let (y0, inputs, _) = small_dae_setup(); + let columns = MAX_HERMITE_COLUMNS + 1; + let result = prepared + .solve( + SolveRequest::new(&grid(columns)), + InputSet::new(&y0, &inputs), + ) + .expect("solve"); + assert_eq!(result.n_times, columns, "every column is still reported"); + assert!( + result.yp.is_none(), + "{columns} columns should have dropped yp" + ); + } + + #[test] + fn the_hermite_column_limit_applies_to_the_sensitivity_path() { + let prepared = build_small_dae_prepared_with_sens().with_store_yp(true); + let (y0, inputs, _) = small_dae_setup(); + let result = prepared + .solve( + SolveRequest::new(&grid(MAX_HERMITE_COLUMNS + 1)).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("solve"); + assert!( + result.yp.is_none(), + "the sens path drops yp on the same rule" + ); + } + + #[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_sensitivity_state_path_stores_yp_too() { + let prepared = build_small_dae_prepared_with_sens().with_store_yp(true); + let (y0, inputs, t_eval) = small_dae_setup(); + let result = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("solve"); + let yp = result.yp.expect("yp requested"); + assert_eq!(yp.len(), result.y.len()); + let last = (result.n_times - 1) * result.n_rows; + let expected = -inputs[0] * result.y[last]; + assert!( + (yp[last] - expected).abs() <= 1e-6 * expected.abs(), + "terminal yp = {} but -a*y = {expected}", + yp[last], + ); + } + + #[test] + fn the_root_column_carries_the_wound_back_derivative() { + // The root column is read off the state after state_mut_back, so its + // yp must be the root-time slope, 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_eq!(r.flag, 1, "expected event termination"); + let yp = r.yp.expect("yp requested"); + let root = r.n_times - 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): the root must end the + // trajectory and the sensitivity columns must span the same times. + 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"); + + assert_eq!(r.flag, 1, "expected event termination"); + let t_event = r.t_event.expect("t_event missing"); + let ln2p5 = 2.5f64.ln(); + assert!( + (t_event - ln2p5).abs() < 1e-6, + "t_event={t_event}, want {ln2p5}" + ); + assert_eq!(r.n_times, r.t.len()); + assert_eq!(blocks(&r)[0].len(), r.n_times); + // y_event is the full state at the root, not the outputs row. + 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.n_times - 1]; + assert!((out_last - 0.8).abs() < 1e-6, "out at root={out_last}"); + let sens_last = blocks(&r)[0][r.n_times - 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}" + ); + } + + /// `(ws, eqn)` over the two-input fixture with the given sensitivity subset, so a + /// parameter-indexed `f_p` is distinguishable from a subset-indexed one. + fn two_input_rhs_at_y0(sens: &[usize]) -> (Rc>, 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()); + let ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + // a = 1.0, b = 0.5, y0 = [1, 0] + let eqn = prepared.build_eqn( + &[1.0, 0.0], + &[], + &ws, + false, + &[1.0, 0.5], + &prepared.sens_param_indices, + ); + (ws, eqn) + } + + /// Collect the dense `n x ncols` matrix `m` in column-major order. + 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, so a mis-mapped scatter shows up as swapped columns. + let (_ws, eqn) = two_input_rhs_at_y0(&[1, 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_narrows_to_the_requested_column() { + // Subset [b] alone: one column, and it must be d/db. A prefix-truncating + // seed would hand back d/da instead. + let (_ws, eqn) = two_input_rhs_at_y0(&[1]); + 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, 1, "2 inputs, 1 requested"); + + 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), vec![1.0, 0.0]); + } + + #[test] + fn rhs_sens_inplace_matches_per_column_sens_mul_on_a_subset() { + // Sharing one primal pass must not move the answer, including for a + // selection that is reordered relative to the input vector. + let (_ws, eqn) = two_input_rhs_at_y0(&[1, 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(col.as_slice()); + 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() { + // The batching mechanism, not a timing: one shared primal pass feeds every + // requested column, at both k = 1 and k = 2. + for sens in [&[0][..], &[1, 0][..]] { + let (ws, 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); + + ws.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!(ws.borrow().sens_primal_passes, 1, "k = {np}"); + } + } + + #[test] + fn set_params_splices_only_the_sensitivity_slots() { + // k values in, spliced at their global indices: the surplus input keeps + // the value the solve was handed. + let model = build_small_dae_two_inputs(&[1]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + let mut eqn = prepared.build_eqn( + &[1.0, 0.0], + &[], + &ws, + false, + &[1.0, 0.5], + &prepared.sens_param_indices, + ); + assert_eq!(eqn.nparams(), 1); + + let ctx = *eqn.context(); + let mut p = FaerVec::::zeros(1, ctx); + p.set_index(0, 7.0); + eqn.set_params(&p); + + assert_eq!(eqn.params.as_slice(), &[1.0, 7.0]); + } + + #[test] + fn get_params_round_trips_the_subset() { + let model = build_small_dae_two_inputs(&[1]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let ws = Rc::new(RefCell::new(prepared.compiled.create_workspace())); + let eqn = prepared.build_eqn( + &[1.0, 0.0], + &[], + &ws, + false, + &[1.0, 0.5], + &prepared.sens_param_indices, + ); + + let ctx = *eqn.context(); + let mut p = FaerVec::::zeros(eqn.nparams(), ctx); + eqn.get_params(&mut p); + + assert_eq!(p.as_slice(), &[0.5]); + } + + #[test] + fn plain_solve_of_a_sens_model_takes_the_full_input_vector() { + // The identity subset on the plain path: a model compiled with a + // one-parameter subset still solves against both inputs. + let model = build_small_dae_two_inputs(&[1]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + + let res = prepared + .solve( + SolveRequest::new(&t_eval), + InputSet::new(&[1.0, 0.0], &[1.0, 0.5]), + ) + .expect("plain solve of a sens-compiled model must accept the full input vector"); + + // dy/dt = -a*y + b with a = 1, b = 0.5 relaxes towards b/a = 0.5. + let last = (res.n_times - 1) * res.n_rows; + assert!( + res.y[last] > 0.5 && res.y[last] < 1.0, + "y = {}", + res.y[last] + ); + } + + #[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. + // Zeroing the seed (the old behaviour) misses the first contribution. + let prepared = build_small_dae_prepared_with_sens(); + let (_, inputs, t_eval) = small_dae_setup(); + let a = inputs[0]; + let y0 = vec![a, a]; + let seed = [1.0, 1.0]; + + let res = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs).with_sens_seed(&seed), + ) + .expect("seeded sensitivity solve failed"); + + let h = 1e-6; + let rp = prepared + .solve( + SolveRequest::new(&t_eval), + InputSet::new(&[a + h, a + h], &[a + h]), + ) + .expect("perturbed solve failed"); + let rm = prepared + .solve( + SolveRequest::new(&t_eval), + InputSet::new(&[a - h, a - h], &[a - h]), + ) + .expect("perturbed solve failed"); + let n = res.n_rows; + let last = (res.n_times - 1) * n; + for j in 0..n { + let fd = (rp.y[last + j] - rm.y[last + j]) / (2.0 * h); + let got = blocks(&res)[0][last + j]; + assert!((got - fd).abs() < 1e-4, "d/da[{j}]: got {got} fd {fd}"); + } + + let unseeded = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .expect("unseeded sensitivity solve failed"); + assert!( + (blocks(&unseeded)[0][last] - blocks(&res)[0][last]).abs() > 1e-3, + "the seed has to move the answer, else the test proves nothing", + ); + } + + #[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 a_sensitivity_request_matches_finite_difference_on_inputs() { + let prepared = build_small_dae_prepared_with_sens(); + let (y0, inputs, t_eval) = small_dae_setup(); + let res = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&y0, &inputs), + ) + .unwrap(); + assert_eq!(blocks(&res).len(), inputs.len()); + // Finite-difference dy/dp0 at the last time point vs the sens-block tail. + let h = 1e-6; + let mut ip = inputs.clone(); + ip[0] += h; + let rp = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &ip)) + .unwrap(); + let mut im = inputs; // last user of inputs; move instead of clone + im[0] -= h; + let rm = prepared + .solve(SolveRequest::new(&t_eval), InputSet::new(&y0, &im)) + .unwrap(); + let n = res.n_rows; + let last = (res.n_times - 1) * n; + for j in 0..n { + let fd = (rp.y[last + j] - rm.y[last + j]) / (2.0 * h); + let got = blocks(&res)[0][last + j]; + assert!( + (got - fd).abs() < 1e-4, + "sensitivity block 0, entry {j}: got {got} fd {fd}" + ); + } + } + + #[test] + fn algebraic_rows_are_detected_from_the_mass_matrix() { + // build_small_dae_two_inputs has mass row 0 = [1.0] and row 1 empty. + let model = build_small_dae_two_inputs(&[0]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + assert_eq!(prepared.algebraic_rows(), &[false, true]); + } + + #[test] + fn sens_atol_tightens_only_differential_rows() { + let model = build_small_dae_two_inputs(&[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 default_sens_atol_factor_tightens_differential_rows() { + let model = build_small_dae_two_inputs(&[0]); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let sens_atol = prepared.sens_atol(); + assert!( + (sens_atol[0] / (1e-8 * DEFAULT_SENS_ATOL_FACTOR) - 1.0).abs() < 1e-12, + "default factor not applied: {}", + sens_atol[0] + ); + } + + #[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(&[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 successful_sens_solve_does_not_report_a_relaxed_error_control() { + let model = build_small_dae_two_inputs_one_sens(); + let prepared = + PreparedSolver::new(model, 1e-8, &[1e-8, 1e-8]).expect("PreparedSolver failed"); + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + let result = prepared + .solve( + SolveRequest::new(&t_eval).with_sensitivities(), + InputSet::new(&[1.0, 0.0], &[1.0, 0.5]), + ) + .expect("sens solve failed"); + assert!(!result.statistics.sens_error_control_relaxed); + } + + #[test] + fn relaxed_error_control_still_integrates_the_sensitivity() { + // The retry path must produce the same answer, just under looser control. + let prepared = || { + PreparedSolver::new( + build_small_dae_two_inputs_one_sens(), + 1e-10, + &[1e-10, 1e-10], + ) + .expect("PreparedSolver failed") + }; + let t_eval: Vec = (0..=5).map(|i| f64::from(i) * 0.2).collect(); + let (y0, inputs) = ([1.0, 0.0], [1.0, 0.5]); + + let controlled = prepared() + .run_dense_sensitivities_inner( + TimePlan::new(&t_eval, &[]), + &y0, + &[], + &inputs, + false, + true, + ) + .expect("controlled solve failed"); + let relaxed = prepared() + .run_dense_sensitivities_inner( + TimePlan::new(&t_eval, &[]), + &y0, + &[], + &inputs, + false, + false, + ) + .expect("relaxed solve failed"); + + let last = t_eval.len() - 1; + let controlled_s = controlled.1[0][last * controlled.0.n_rows]; + let relaxed_s = relaxed.1[0][last * relaxed.0.n_rows]; + assert!( + (controlled_s - relaxed_s).abs() < 1e-5 * controlled_s.abs().max(1e-8), + "relaxed sens {relaxed_s} diverged from controlled {controlled_s}" + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/sparsity.rs b/packages/pybamm-rust/pybamm-core/src/sparsity.rs new file mode 100644 index 0000000000..3bb9970397 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/sparsity.rs @@ -0,0 +1,1119 @@ +//! Per-output sparsity of `d(root)/dy`, read off the DAG without evaluating it. +//! +//! Each node is annotated with the set of state indices its value can depend on, +//! propagated bottom-up as bitsets; the sets reaching each element of the root +//! become one row of a CSR [`SparsityPattern`]. Tracking outputs individually +//! rather than unioning them is what makes coloring pay: a tridiagonal Jacobian +//! keeps its three colors instead of collapsing to a dense union. +//! +//! The result is structural and conservative: an entry may be present and +//! evaluate to zero, but a missing entry is zero for every `y`, which is the +//! direction coloring and assembly depend on. + +use crate::arena::{Arena, NodeId}; +use crate::node::Node; + +/// Packed bitset over n state indices. Backed by `Vec`. +/// +/// `iter()` yields set indices in ascending order, which matches the +/// CSR column-sorted requirement of `SparsityPattern`, callers do not +/// need to sort after collecting from a `BitSet`. +#[derive(Clone, Debug, PartialEq, Eq)] +struct BitSet { + bits: Vec, + n: usize, +} + +impl BitSet { + fn zeros(n: usize) -> Self { + Self { + bits: vec![0u64; n.div_ceil(64)], + n, + } + } + + fn insert(&mut self, i: usize) { + debug_assert!(i < self.n, "BitSet::insert({i}) out of range n={}", self.n); + self.bits[i / 64] |= 1u64 << (i % 64); + } + + fn union_with(&mut self, other: &Self) { + debug_assert_eq!(self.n, other.n, "BitSet::union_with shape mismatch"); + for (a, &b) in self.bits.iter_mut().zip(other.bits.iter()) { + *a |= b; + } + } + + fn iter(&self) -> impl Iterator + '_ { + self.bits + .iter() + .enumerate() + .flat_map(|(word_idx, &word)| BitsIter { + word, + base: word_idx * 64, + }) + } +} + +struct BitsIter { + word: u64, + base: usize, +} + +impl Iterator for BitsIter { + type Item = usize; + fn next(&mut self) -> Option { + if self.word == 0 { + return None; + } + let tz = self.word.trailing_zeros() as usize; + self.word &= self.word - 1; // clear lowest set bit + Some(self.base + tz) + } +} + +/// CSR structure of a Jacobian: which entries can be non-zero, without values. +/// +/// Column indices are ascending within each row, which [`merge_with`] relies on: +/// it is a two-way sorted merge, so unsorted input silently yields a wrong union +/// rather than an error. +/// +/// [`merge_with`]: Self::merge_with +#[derive(Debug, Clone)] +pub struct SparsityPattern { + /// Rows, one per output element. + pub nrows: usize, + /// Columns, one per state (or per parameter for a `df/dp` pattern). + pub ncols: usize, + /// Row start offsets into `indices`, length `nrows + 1`. + pub indptr: Vec, + /// Column index of every structural entry, row by row. + pub indices: Vec, +} + +impl SparsityPattern { + /// An all-zero pattern of the given shape. + pub fn new(nrows: usize, ncols: usize) -> Self { + Self { + nrows, + ncols, + indptr: vec![0; nrows + 1], + indices: Vec::new(), + } + } + + /// Structural non-zeros, which is the value-buffer length assembly needs. + pub const fn nnz(&self) -> usize { + self.indices.len() + } + + /// Row of every CSR entry, i.e. the inverse of `indptr`. + pub fn entry_rows(&self) -> Vec { + let mut rows = vec![0usize; self.nnz()]; + for row in 0..self.nrows { + rows[self.indptr[row]..self.indptr[row + 1]].fill(row); + } + rows + } + + /// Entries of each row, as the widths `detect_dense_rows` measures. + pub fn row_widths(&self) -> Vec { + (0..self.nrows) + .map(|row| self.indptr[row + 1] - self.indptr[row]) + .collect() + } + + /// Merge another sparsity pattern into this one (union of nonzero positions). + /// + /// Combines df/dy sparsity with mass-matrix sparsity for the full + /// `J = df/dy - cj*M` pattern KLU needs. + pub fn merge_with(&mut self, other: &Self) { + assert_eq!(self.nrows, other.nrows, "Row count mismatch"); + assert_eq!(self.ncols, other.ncols, "Column count mismatch"); + + let mut new_indices = Vec::with_capacity(self.indices.len() + other.indices.len()); + let mut new_indptr = vec![0usize; self.nrows + 1]; + + for row in 0..self.nrows { + let self_start = self.indptr[row]; + let self_end = self.indptr[row + 1]; + let other_start = other.indptr[row]; + let other_end = other.indptr[row + 1]; + + // Merge sorted column indices + let mut i = self_start; + let mut j = other_start; + while i < self_end && j < other_end { + match self.indices[i].cmp(&other.indices[j]) { + std::cmp::Ordering::Less => { + new_indices.push(self.indices[i]); + i += 1; + }, + std::cmp::Ordering::Greater => { + new_indices.push(other.indices[j]); + j += 1; + }, + std::cmp::Ordering::Equal => { + new_indices.push(self.indices[i]); + i += 1; + j += 1; + }, + } + } + while i < self_end { + new_indices.push(self.indices[i]); + i += 1; + } + while j < other_end { + new_indices.push(other.indices[j]); + j += 1; + } + new_indptr[row + 1] = new_indices.len(); + } + + self.indices = new_indices; + self.indptr = new_indptr; + } + + /// Fully dense pattern: every (row, col) entry present. + /// + /// Used for df/dp jacobians where parameter sparsity is not detected; + /// coloring a dense pattern yields one color per column (= one JVP + /// sweep per parameter, matching the per-parameter cost of unit seeds). + pub fn dense(nrows: usize, ncols: usize) -> Self { + let indptr = (0..=nrows).map(|r| r * ncols).collect(); + let indices = (0..nrows).flat_map(|_| 0..ncols).collect(); + Self { + nrows, + ncols, + indptr, + indices, + } + } + + /// Create sparsity pattern from CSR data. + pub fn from_csr_data(csr: &crate::node::CsrData) -> Self { + Self { + nrows: csr.shape.rows, + ncols: csr.shape.cols, + indptr: csr.indptr.clone(), + indices: csr.indices.clone(), + } + } +} + +/// Per-element dependency information +/// Each element in the output may depend on different state variables +#[derive(Debug, Clone)] +enum ElementDeps { + /// Scalar with specific state dependencies + Scalar(BitSet), + /// Vector where each element has its own dependencies + Vector(Vec), +} + +impl ElementDeps { + /// Create scalar dependencies with no state variables + fn scalar_empty(n_states: usize) -> Self { + Self::Scalar(BitSet::zeros(n_states)) + } + + /// Get the number of elements + const fn len(&self) -> usize { + match self { + Self::Scalar(_) => 1, + Self::Vector(v) => v.len(), + } + } + + /// Get dependencies for a specific element, broadcasting scalars + fn get(&self, idx: usize) -> &BitSet { + match self { + Self::Scalar(deps) => deps, + Self::Vector(v) => &v[idx], + } + } + + /// Get union of all dependencies + fn union_all(&self, n_states: usize) -> BitSet { + match self { + // Scalar already carries its own size; n_states only sizes the + // accumulator in the Vector arm. + Self::Scalar(deps) => deps.clone(), + Self::Vector(v) => { + let mut result = BitSet::zeros(n_states); + for deps in v { + result.union_with(deps); + } + result + }, + } + } + + /// Convert to vector form with given length (broadcasts scalar) + fn to_vector(&self, len: usize, n_states: usize) -> Vec { + match self { + Self::Scalar(deps) => vec![deps.clone(); len], + Self::Vector(v) => { + if v.len() == len { + v.clone() + } else if v.len() == 1 { + vec![v[0].clone(); len] + } else { + // Unexpected shape mismatch - return union for safety + let union = self.union_all(n_states); + vec![union; len] + } + }, + } + } +} + +/// Bottom-up memoized analysis. For each reachable node, compute its +/// `ElementDeps` exactly once. Children are guaranteed in `deps` before +/// the parent is processed because `topological_order` visits them +/// first. +fn analyze_output_dependencies(arena: &Arena, root: NodeId, n_states: usize) -> ElementDeps { + let order = arena.topological_order(root); + let mut deps: Vec> = (0..arena.len()).map(|_| None).collect(); + for &id in &order { + let computed = compute_node_deps(arena, id, n_states, &deps); + deps[id.index()] = Some(computed); + } + deps[root.index()] + .take() + .expect("root must be in topological order") +} + +fn compute_node_deps( + arena: &Arena, + id: NodeId, + n_states: usize, + deps: &[Option], +) -> ElementDeps { + match arena.get(id) { + Node::StateVector { start, end } => { + let len = end - start; + if len == 1 { + let mut bs = BitSet::zeros(n_states); + bs.insert(*start); + ElementDeps::Scalar(bs) + } else { + let v: Vec = (*start..*end) + .map(|i| { + let mut bs = BitSet::zeros(n_states); + bs.insert(i); + bs + }) + .collect(); + ElementDeps::Vector(v) + } + }, + Node::StateVectorDot { start, end } | Node::TangentStateVector { start, end } => { + let len = end - start; + if len == 1 { + ElementDeps::scalar_empty(n_states) + } else { + ElementDeps::Vector(vec![BitSet::zeros(n_states); len]) + } + }, + Node::TangentParameter { .. } + | Node::Scalar(_) + | Node::Time + | Node::InputParameter { .. } => ElementDeps::scalar_empty(n_states), + Node::ZeroVector { len } => { + if *len == 1 { + ElementDeps::scalar_empty(n_states) + } else { + ElementDeps::Vector(vec![BitSet::zeros(n_states); *len]) + } + }, + Node::Array(arr) => { + if arr.data.len() == 1 { + ElementDeps::scalar_empty(n_states) + } else { + ElementDeps::Vector(vec![BitSet::zeros(n_states); arr.data.len()]) + } + }, + Node::SparseMatrix(csr) => { + let total = csr.shape.rows * csr.shape.cols; + if total == 1 { + ElementDeps::scalar_empty(n_states) + } else { + ElementDeps::Vector(vec![BitSet::zeros(n_states); csr.shape.rows]) + } + }, + Node::Neg(a) + | Node::Abs(a) + | Node::Sqrt(a) + | Node::Exp(a) + | Node::Log(a) + | Node::Sin(a) + | Node::Cos(a) + | Node::Tanh(a) + | Node::Sinh(a) + | Node::Cosh(a) + | Node::Arcsinh(a) + | Node::Arctan(a) + | Node::Erf(a) + | Node::Sign(a) + | Node::Floor(a) + | Node::Ceiling(a) => child_deps(deps, *a).clone(), + Node::MaxReduce(a) | Node::MinReduce(a) => { + ElementDeps::Scalar(child_deps(deps, *a).union_all(n_states)) + }, + Node::ReduceArgSelect { basis, picker, .. } => { + // Argmax/argmin is runtime-dependent, so the static pattern must + // union deps over the whole reduced vector (deps(basis) subset of deps(picker)). + let mut combined = child_deps(deps, *basis).union_all(n_states); + combined.union_with(&child_deps(deps, *picker).union_all(n_states)); + ElementDeps::Scalar(combined) + }, + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => { + combine_binary_deps(child_deps(deps, *a), child_deps(deps, *b), n_states) + }, + Node::MatMul(mat_id, vec_id) => { + let vec_d = child_deps(deps, *vec_id); + match arena.get(*mat_id) { + Node::SparseMatrix(csr) => { + let n_rows = csr.shape.rows; + let mut result = Vec::with_capacity(n_rows); + for row in 0..n_rows { + let row_start = csr.indptr[row]; + let row_end = csr.indptr[row + 1]; + let mut row_deps = BitSet::zeros(n_states); + for &col in &csr.indices[row_start..row_end] { + row_deps.union_with(vec_d.get(col)); + } + result.push(row_deps); + } + if result.len() == 1 { + ElementDeps::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDeps::Vector(result) + } + }, + Node::Array(arr) => { + let u = vec_d.union_all(n_states); + let n_rows = arr.shape.rows; + if n_rows == 1 { + ElementDeps::Scalar(u) + } else { + ElementDeps::Vector(vec![u; n_rows]) + } + }, + _ => { + let mat_d = child_deps(deps, *mat_id); + let mut combined = mat_d.union_all(n_states); + combined.union_with(&vec_d.union_all(n_states)); + ElementDeps::Scalar(combined) + }, + } + }, + Node::Index { child, start, end } => { + let cd = child_deps(deps, *child); + let len = end - start; + match cd { + ElementDeps::Scalar(b) => { + if len == 1 { + ElementDeps::Scalar(b.clone()) + } else { + ElementDeps::Vector(vec![b.clone(); len]) + } + }, + ElementDeps::Vector(v) => { + let subset: Vec<_> = v[*start..*end].to_vec(); + if subset.len() == 1 { + ElementDeps::Scalar(subset.into_iter().next().unwrap()) + } else { + ElementDeps::Vector(subset) + } + }, + } + }, + Node::Concat(children) => { + let mut result = Vec::new(); + for c in children { + match child_deps(deps, *c) { + ElementDeps::Scalar(b) => result.push(b.clone()), + ElementDeps::Vector(v) => result.extend(v.iter().cloned()), + } + } + if result.len() == 1 { + ElementDeps::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDeps::Vector(result) + } + }, + Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => child_deps(deps, *child).clone(), + Node::InterpolantNd { children, .. } | Node::InterpolantNdPartial { children, .. } => { + let mut acc = child_deps(deps, children[0]).clone(); + for &c in &children[1..] { + acc = combine_binary_deps(&acc, child_deps(deps, c), n_states); + } + acc + }, + Node::Conditional { selector, branches } => { + let sel = child_deps(deps, *selector); + let bd: Vec<&ElementDeps> = branches.iter().map(|b| child_deps(deps, *b)).collect(); + let output_len = bd.iter().map(|d| d.len()).max().unwrap_or(1); + let sel_union = sel.union_all(n_states); + let mut result = Vec::with_capacity(output_len); + for i in 0..output_len { + let mut e = sel_union.clone(); + for bdj in &bd { + if let Some(last_idx) = bdj.len().checked_sub(1) { + e.union_with(bdj.get(i.min(last_idx))); + } + } + result.push(e); + } + if result.len() == 1 { + ElementDeps::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDeps::Vector(result) + } + }, + } +} + +fn child_deps(deps: &[Option], id: NodeId) -> &ElementDeps { + deps[id.index()] + .as_ref() + .expect("child must be processed before parent in topological order") +} + +/// Combine dependencies for binary operations with broadcast semantics +fn combine_binary_deps(a: &ElementDeps, b: &ElementDeps, n_states: usize) -> ElementDeps { + match (a, b) { + // Both scalars: union + (ElementDeps::Scalar(da), ElementDeps::Scalar(db)) => { + let mut c = da.clone(); + c.union_with(db); + ElementDeps::Scalar(c) + }, + // Scalar + Vector: broadcast scalar to each element + (ElementDeps::Scalar(scalar_deps), ElementDeps::Vector(vec_deps)) => { + let result: Vec<_> = vec_deps + .iter() + .map(|vd| { + let mut c = scalar_deps.clone(); + c.union_with(vd); + c + }) + .collect(); + ElementDeps::Vector(result) + }, + // Vector + Scalar: broadcast scalar to each element + (ElementDeps::Vector(vec_deps), ElementDeps::Scalar(scalar_deps)) => { + let result: Vec<_> = vec_deps + .iter() + .map(|vd| { + let mut c = vd.clone(); + c.union_with(scalar_deps); + c + }) + .collect(); + ElementDeps::Vector(result) + }, + // Vector + Vector: element-wise union + (ElementDeps::Vector(va), ElementDeps::Vector(vb)) => { + let len = va.len().max(vb.len()); + let va_expanded = a.to_vector(len, n_states); + let vb_expanded = b.to_vector(len, n_states); + + let result: Vec<_> = va_expanded + .iter() + .zip(vb_expanded.iter()) + .map(|(da, db)| { + let mut c = da.clone(); + c.union_with(db); + c + }) + .collect(); + ElementDeps::Vector(result) + }, + } +} + +/// Detect sparsity pattern with per-output dependency tracking +/// +/// Tracks which state variables each individual output element depends on, +/// rather than a conservative union across all outputs. This enables efficient +/// coloring for banded Jacobians (e.g., tridiagonal patterns get O(3) colors +/// instead of O(n)). +pub fn detect_sparsity_per_output( + arena: &Arena, + root: NodeId, + n_outputs: usize, + n_states: usize, +) -> SparsityPattern { + let output_deps = analyze_output_dependencies(arena, root, n_states); + + // Build CSR pattern + let mut pattern = SparsityPattern::new(n_outputs, n_states); + + // Convert to vector form to iterate over rows + let deps_vec = output_deps.to_vector(n_outputs, n_states); + + for (row, bs) in deps_vec.iter().enumerate() { + pattern.indptr[row] = pattern.indices.len(); + pattern.indices.extend(bs.iter()); + } + pattern.indptr[n_outputs] = pattern.indices.len(); + + pattern +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::coloring::color_columns; + use crate::node::{CsrData, Shape}; + + /// Collect every state index that any reachable `StateVector` node uses. + /// One topological walk, `O(arena.len())` + total `StateVector` slot count. + fn collect_state_deps(arena: &Arena, root: NodeId, n_states: usize) -> BitSet { + let mut acc = BitSet::zeros(n_states); + for id in arena.topological_order(root) { + if let Node::StateVector { start, end } = arena.get(id) { + for i in *start..*end { + acc.insert(i); + } + } + } + acc + } + + /// Conservative sparsity oracle: every output row gets the union of all + /// state dependencies. A superset of `detect_sparsity_per_output`, which it + /// cross-checks in tests. + fn detect_sparsity_simple( + arena: &Arena, + root: NodeId, + n_outputs: usize, + n_states: usize, + ) -> SparsityPattern { + let deps = collect_state_deps(arena, root, n_states); + let mut pattern = SparsityPattern::new(n_outputs, n_states); + let sorted_deps: Vec = deps.iter().collect(); + for row in 0..n_outputs { + pattern.indptr[row] = pattern.indices.len(); + pattern.indices.extend_from_slice(&sorted_deps); + } + pattern.indptr[n_outputs] = pattern.indices.len(); + pattern + } + + #[test] + fn test_sparsity_pattern_new() { + let pattern = SparsityPattern::new(3, 4); + assert_eq!(pattern.nrows, 3); + assert_eq!(pattern.ncols, 4); + assert_eq!(pattern.indptr.len(), 4); // nrows + 1 + assert_eq!(pattern.nnz(), 0); + } + + #[test] + fn test_detect_sparsity_scalar() { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, sv)); + + let pattern = detect_sparsity_simple(&arena, expr, 2, 4); + + // Output depends on states 0 and 1 + assert_eq!(pattern.nrows, 2); + assert_eq!(pattern.ncols, 4); + // Each row should have indices [0, 1] + assert_eq!( + &pattern.indices[pattern.indptr[0]..pattern.indptr[1]], + &[0, 1] + ); + } + + #[test] + fn test_detect_sparsity_disjoint() { + let mut arena = Arena::new(); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sv1 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let expr = arena.alloc(Node::Add(sv0, sv1)); + + let pattern = detect_sparsity_simple(&arena, expr, 1, 4); + + // Output depends on states 0 and 2 + let row_indices = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert!(row_indices.contains(&0)); + assert!(row_indices.contains(&2)); + assert!(!row_indices.contains(&1)); + } + + #[test] + fn test_per_output_sparsity_diagonal() { + // f(y) = [y0, y1, y2] - each output depends only on its corresponding input + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let expr = arena.alloc(Node::Concat(vec![y0, y1, y2])); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + // Row 0 should only depend on column 0 + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[0], "Row 0 should only depend on state 0"); + + // Row 1 should only depend on column 1 + let row1 = &pattern.indices[pattern.indptr[1]..pattern.indptr[2]]; + assert_eq!(row1, &[1], "Row 1 should only depend on state 1"); + + // Row 2 should only depend on column 2 + let row2 = &pattern.indices[pattern.indptr[2]..pattern.indptr[3]]; + assert_eq!(row2, &[2], "Row 2 should only depend on state 2"); + + // Diagonal matrix should need only 1 color + let coloring = color_columns(&pattern); + assert_eq!(coloring.n_colors, 1, "Diagonal should need only 1 color"); + } + + #[test] + fn test_per_output_sparsity_identity_passthrough() { + // f(y) = y (identity function on vector) + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + + let pattern = detect_sparsity_per_output(&arena, y, 4, 4); + + // Each row i should depend only on column i + for i in 0..4 { + let row = &pattern.indices[pattern.indptr[i]..pattern.indptr[i + 1]]; + assert_eq!(row, &[i], "Row {i} should only depend on state {i}"); + } + + // Diagonal needs 1 color + let coloring = color_columns(&pattern); + assert_eq!(coloring.n_colors, 1); + } + + #[test] + fn test_sparsity_tridiagonal() { + // f(y) = [y0+y1, y0+y1+y2, y1+y2+y3, y2+y3] + // Should have tridiagonal-like sparsity + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let y3 = arena.alloc(Node::StateVector { start: 3, end: 4 }); + + let f0 = arena.alloc(Node::Add(y0, y1)); + let f1_partial = arena.alloc(Node::Add(y0, y1)); + let f1 = arena.alloc(Node::Add(f1_partial, y2)); + let f2_partial = arena.alloc(Node::Add(y1, y2)); + let f2 = arena.alloc(Node::Add(f2_partial, y3)); + let f3 = arena.alloc(Node::Add(y2, y3)); + + let expr = arena.alloc(Node::Concat(vec![f0, f1, f2, f3])); + + let pattern = detect_sparsity_per_output(&arena, expr, 4, 4); + + // Check row dependencies + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[0, 1], "Row 0 should depend on states 0, 1"); + + let row1 = &pattern.indices[pattern.indptr[1]..pattern.indptr[2]]; + assert_eq!(row1, &[0, 1, 2], "Row 1 should depend on states 0, 1, 2"); + + let row2 = &pattern.indices[pattern.indptr[2]..pattern.indptr[3]]; + assert_eq!(row2, &[1, 2, 3], "Row 2 should depend on states 1, 2, 3"); + + let row3 = &pattern.indices[pattern.indptr[3]..pattern.indptr[4]]; + assert_eq!(row3, &[2, 3], "Row 3 should depend on states 2, 3"); + + // Check that coloring can exploit sparsity + let coloring = color_columns(&pattern); + // Tridiagonal needs at most 3 colors + assert!( + coloring.n_colors <= 3, + "Expected <= 3 colors, got {}", + coloring.n_colors + ); + } + + #[test] + fn test_per_output_sparsity_scalar_broadcast() { + // f(y) = [y0 * c, y1 * c, y2 * c] where c is a scalar + // Each output should depend only on its corresponding input + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let c = arena.alloc(Node::Scalar(2.0)); + + let f0 = arena.alloc(Node::Mul(y0, c)); + let f1 = arena.alloc(Node::Mul(y1, c)); + let f2 = arena.alloc(Node::Mul(y2, c)); + let expr = arena.alloc(Node::Concat(vec![f0, f1, f2])); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + // Each row should depend only on its corresponding column + for i in 0..3 { + let row = &pattern.indices[pattern.indptr[i]..pattern.indptr[i + 1]]; + assert_eq!(row, &[i], "Row {i} should only depend on state {i}"); + } + } + + #[test] + fn test_per_output_sparsity_unary_preserves() { + // f(y) = [exp(y0), sin(y1), log(y2)] + // Unary ops should preserve per-element structure + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + + let f0 = arena.alloc(Node::Exp(y0)); + let f1 = arena.alloc(Node::Sin(y1)); + let f2 = arena.alloc(Node::Log(y2)); + let expr = arena.alloc(Node::Concat(vec![f0, f1, f2])); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + // Each row should depend only on its corresponding column + for i in 0..3 { + let row = &pattern.indices[pattern.indptr[i]..pattern.indptr[i + 1]]; + assert_eq!(row, &[i], "Row {i} should only depend on state {i}"); + } + } + + #[test] + fn test_per_output_sparsity_index() { + // f(y) = y[1:3] where y has 4 elements + // Output[0] should depend on state 1, output[1] on state 2 + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let expr = arena.alloc(Node::Index { + child: y, + start: 1, + end: 3, + }); + + let pattern = detect_sparsity_per_output(&arena, expr, 2, 4); + + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[1], "Row 0 should only depend on state 1"); + + let row1 = &pattern.indices[pattern.indptr[1]..pattern.indptr[2]]; + assert_eq!(row1, &[2], "Row 1 should only depend on state 2"); + } + + #[test] + fn test_per_output_sparsity_reduction() { + // f(y) = max(y) where y = [y0, y1, y2] + // Scalar output should depend on all inputs + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::MaxReduce(y)); + + let pattern = detect_sparsity_per_output(&arena, expr, 1, 3); + + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[0, 1, 2], "Reduction should depend on all inputs"); + } + + #[test] + fn test_per_output_sparsity_sparse_matmul() { + // Sparse matrix M @ y where M is tridiagonal + // Each output row i depends only on states i-1, i, i+1 + let mut arena = Arena::new(); + + // Tridiagonal 3x3: [1 1 0; 1 1 1; 0 1 1] + let csr = CsrData { + indptr: vec![0, 2, 5, 7], + indices: vec![0, 1, 0, 1, 2, 1, 2], + data: vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + shape: Shape::matrix(3, 3), + }; + let mat = arena.alloc(Node::SparseMatrix(Box::new(csr))); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::MatMul(mat, y)); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[0, 1], "Row 0 should depend on states 0, 1"); + + let row1 = &pattern.indices[pattern.indptr[1]..pattern.indptr[2]]; + assert_eq!(row1, &[0, 1, 2], "Row 1 should depend on states 0, 1, 2"); + + let row2 = &pattern.indices[pattern.indptr[2]..pattern.indptr[3]]; + assert_eq!(row2, &[1, 2], "Row 2 should depend on states 1, 2"); + + // Should need only 3 colors + let coloring = color_columns(&pattern); + assert!(coloring.n_colors <= 3); + } + + #[test] + fn test_per_output_vs_simple_diagonal() { + // Compare per-output vs simple detection on diagonal case + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let expr = arena.alloc(Node::Concat(vec![y0, y1, y2])); + + let pattern_simple = detect_sparsity_simple(&arena, expr, 3, 3); + let pattern_per_output = detect_sparsity_per_output(&arena, expr, 3, 3); + + // Simple: all rows have same dependencies [0, 1, 2] + let coloring_simple = color_columns(&pattern_simple); + assert_eq!(coloring_simple.n_colors, 3, "Simple should need 3 colors"); + + // Per-output: diagonal, only needs 1 color + let coloring_per_output = color_columns(&pattern_per_output); + assert_eq!( + coloring_per_output.n_colors, 1, + "Per-output should need only 1 color" + ); + } + + #[test] + fn test_per_output_sparsity_vector_add() { + // f(y) = y + y (element-wise add) + // Each output i should depend only on input i + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::Add(y, y)); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + for i in 0..3 { + let row = &pattern.indices[pattern.indptr[i]..pattern.indptr[i + 1]]; + assert_eq!(row, &[i], "Row {i} should only depend on state {i}"); + } + } + + #[test] + fn test_per_output_sparsity_cross_dependency() { + // f(y) = [y0 + y1, y1 + y2] + // Creates off-diagonal dependencies + 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 y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + + let f0 = arena.alloc(Node::Add(y0, y1)); + let f1 = arena.alloc(Node::Add(y1, y2)); + let expr = arena.alloc(Node::Concat(vec![f0, f1])); + + let pattern = detect_sparsity_per_output(&arena, expr, 2, 3); + + let row0 = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row0, &[0, 1], "Row 0 should depend on states 0, 1"); + + let row1 = &pattern.indices[pattern.indptr[1]..pattern.indptr[2]]; + assert_eq!(row1, &[1, 2], "Row 1 should depend on states 1, 2"); + + // Should need 2 colors (columns 0 and 2 can share, column 1 needs its own) + let coloring = color_columns(&pattern); + assert!(coloring.n_colors <= 2); + } + + #[test] + fn test_per_output_empty_deps() { + // f(y) = [1.0, 2.0, 3.0] - constants with no state deps + let mut arena = Arena::new(); + let c1 = arena.alloc(Node::Scalar(1.0)); + let c2 = arena.alloc(Node::Scalar(2.0)); + let c3 = arena.alloc(Node::Scalar(3.0)); + let expr = arena.alloc(Node::Concat(vec![c1, c2, c3])); + + let pattern = detect_sparsity_per_output(&arena, expr, 3, 3); + + // All rows should have no dependencies + for i in 0..3 { + let row = &pattern.indices[pattern.indptr[i]..pattern.indptr[i + 1]]; + assert!(row.is_empty(), "Row {i} should have no dependencies"); + } + + // Zero nonzeros + assert_eq!(pattern.nnz(), 0); + } + + #[test] + fn test_large_tridiagonal_coloring_efficiency() { + // Large tridiagonal system to verify O(3) coloring vs O(n) + let n = 100; + let mut arena = Arena::new(); + + // Create y[i] for i in 0..n + let y_nodes: Vec = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + + // f[i] = y[i-1] + y[i] + y[i+1] (with boundary handling) + let f_nodes: Vec = (0..n) + .map(|i| { + if i == 0 { + arena.alloc(Node::Add(y_nodes[0], y_nodes[1])) + } else if i == n - 1 { + arena.alloc(Node::Add(y_nodes[n - 2], y_nodes[n - 1])) + } else { + let partial = arena.alloc(Node::Add(y_nodes[i - 1], y_nodes[i])); + arena.alloc(Node::Add(partial, y_nodes[i + 1])) + } + }) + .collect(); + + let expr = arena.alloc(Node::Concat(f_nodes)); + + // Compare simple vs per-output + let pattern_simple = detect_sparsity_simple(&arena, expr, n, n); + let pattern_per_output = detect_sparsity_per_output(&arena, expr, n, n); + + let coloring_simple = color_columns(&pattern_simple); + let coloring_per_output = color_columns(&pattern_per_output); + + // Simple should need n colors (dense) + assert_eq!( + coloring_simple.n_colors, n, + "Simple detection should need {n} colors" + ); + + // Per-output should need only 3 colors (tridiagonal) + assert!( + coloring_per_output.n_colors <= 3, + "Per-output should need <= 3 colors, got {}", + coloring_per_output.n_colors + ); + } + + #[test] + fn test_bitset_zeros_empty() { + let bs = BitSet::zeros(100); + assert_eq!(bs.iter().count(), 0); + } + + #[test] + fn test_bitset_insert_and_iter_ascending() { + let mut bs = BitSet::zeros(200); + bs.insert(150); + bs.insert(3); + bs.insert(64); + bs.insert(199); + let v: Vec = bs.iter().collect(); + assert_eq!(v, vec![3, 64, 150, 199]); + } + + #[test] + fn test_bitset_insert_idempotent() { + let mut bs = BitSet::zeros(10); + bs.insert(5); + bs.insert(5); + let v: Vec = bs.iter().collect(); + assert_eq!(v, vec![5]); + } + + #[test] + fn test_bitset_union_with() { + let mut a = BitSet::zeros(128); + a.insert(1); + a.insert(70); + let mut b = BitSet::zeros(128); + b.insert(70); + b.insert(127); + a.union_with(&b); + let v: Vec = a.iter().collect(); + assert_eq!(v, vec![1, 70, 127]); + } + + #[test] + fn test_bitset_clone_eq() { + let mut bs = BitSet::zeros(64); + bs.insert(0); + bs.insert(63); + let clone = bs.clone(); + assert_eq!(bs, clone); + } + + #[test] + fn test_bitset_word_boundaries() { + // Indices at u64-word boundaries + let mut bs = BitSet::zeros(200); + for i in [0, 63, 64, 127, 128, 191, 192, 199] { + bs.insert(i); + } + let v: Vec = bs.iter().collect(); + assert_eq!(v, vec![0, 63, 64, 127, 128, 191, 192, 199]); + } + + #[test] + fn test_deep_chain_does_not_stack_overflow() { + // Locks in stack safety for deep chains. The iterative analyzer + // over Arena::topological_order makes this trivially bounded by + // heap, not stack. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let mut current = y; + for _ in 0..5000 { + current = arena.alloc(Node::Sin(current)); + current = arena.alloc(Node::Neg(current)); + } + let pattern = detect_sparsity_per_output(&arena, current, 1, 1); + let row = &pattern.indices[pattern.indptr[0]..pattern.indptr[1]]; + assert_eq!(row, &[0], "Deep chain output depends only on state 0"); + } + + #[test] + fn dense_pattern_has_all_entries() { + let p = SparsityPattern::dense(2, 3); + assert_eq!(p.nrows, 2); + assert_eq!(p.ncols, 3); + assert_eq!(p.indptr, vec![0, 3, 6]); + assert_eq!(p.indices, vec![0, 1, 2, 0, 1, 2]); + assert_eq!(p.nnz(), 6); + } + + #[test] + fn dense_pattern_zero_cols_is_empty() { + let p = SparsityPattern::dense(2, 0); + assert_eq!(p.indptr, vec![0, 0, 0]); + assert_eq!(p.nnz(), 0); + } + + /// 3x3 CSR fixture: row 0 has 1 entry, row 1 has 3 entries, row 2 has 1 entry. + fn make_row_filter_fixture() -> SparsityPattern { + let mut pattern = SparsityPattern::new(3, 3); + pattern.indptr = vec![0, 1, 4, 5]; + pattern.indices = vec![0, 0, 1, 2, 2]; + pattern + } + + #[test] + fn entry_rows_and_row_widths_invert_indptr() { + let pattern = make_row_filter_fixture(); + assert_eq!(pattern.entry_rows(), vec![0, 1, 1, 1, 2]); + assert_eq!(pattern.row_widths(), vec![1, 3, 1]); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/tangent.rs b/packages/pybamm-rust/pybamm-core/src/tangent.rs new file mode 100644 index 0000000000..85b9a66fe2 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/tangent.rs @@ -0,0 +1,2073 @@ +//! Symbolic differentiation for forward-mode automatic differentiation. +//! +//! Rewrites a primal expression into one that evaluates a Jacobian-vector +//! product: `TangentStateVector` and `TangentParameter` nodes stand in for the +//! differentiated variable and read their seed at evaluation time, so one +//! compiled tape serves every direction. + +use std::collections::HashSet; +use std::collections::hash_map::RandomState; + +use crate::arena::{Arena, NodeId, NodeMap}; +use crate::ir::infer_sizes; +use crate::node::{InterpolantData, Node}; + +/// Allocate a zero matching a primal subtree's structural width. +/// +/// A zero derivative must keep the width of the node it replaces. A bare +/// `Scalar(0.0)` collapses a vector derivative to length 1, shifting every +/// downstream `Concat` offset and truncating the output; `len == 0` (a +/// pure-algebraic model's empty `concatenated_rhs`) would widen to 1 and shift +/// the same offsets the other way, so it stays a `ZeroVector`. +fn zero_of_width(arena: &mut Arena, len: usize) -> NodeId { + if len == 1 { + arena.alloc(Node::Scalar(0.0)) + } else { + arena.alloc(Node::ZeroVector { len }) + } +} + +/// Output width of every node in `arena`, indexed by `NodeId::index()`. +/// +/// Computed before differentiation mutates the arena. The recursion only ever +/// queries primal nodes, so appending tangent nodes cannot invalidate it. +fn primal_widths(arena: &Arena, root: NodeId) -> Vec { + let order = arena.topological_order(root); + infer_sizes(arena, &order) +} + +/// Which variables carry a tangent; every other leaf differentiates to zero. +/// +/// Also names the axis a Jacobian differentiates along, so `jacobian` re-exports +/// it rather than defining a second enum over the same two cases. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DiffTarget { + /// `d/dy`: seed `StateVector` nodes from `TangentInputs::dy`; Jacobian + /// columns are states. + States, + /// `d/dp`: seed `InputParameter` nodes from `TangentInputs::dp`; Jacobian + /// columns are parameters. + Params, +} + +/// Build the tangent expression for `(d expr/dy) @ dy`. +/// +/// `StateVector` nodes become `TangentStateVector` nodes that read `dy` at +/// evaluation time; parameters and time are constants. +#[must_use] +pub fn tangent_wrt_states(arena: &mut Arena, expr: NodeId) -> NodeId { + let mut memo: NodeMap = NodeMap::new(arena.len()); + let no_filter: Option<&HashSet> = None; + let widths = primal_widths(arena, expr); + differentiate( + arena, + expr, + DiffTarget::States, + &mut memo, + no_filter, + &widths, + ) +} + +/// Build the tangent expression for `(d expr/dp) @ dp`. +/// +/// `InputParameter` nodes become `TangentParameter` nodes that read `dp` at +/// evaluation time; states and time are constants. +#[must_use] +pub fn tangent_wrt_params(arena: &mut Arena, expr: NodeId) -> NodeId { + let mut memo: NodeMap = NodeMap::new(arena.len()); + let no_filter: Option<&HashSet> = None; + let widths = primal_widths(arena, expr); + differentiate( + arena, + expr, + DiffTarget::Params, + &mut memo, + no_filter, + &widths, + ) +} + +/// Build the tangent expression for `(d expr/dy) @ dy`, restricted to +/// `StateVector` nodes whose index range overlaps `active_indices`. +/// +/// State vectors outside the active set are constants. Overlap is +/// all-or-nothing: a node straddling the boundary is seeded over its whole +/// range. Used for the algebraic Jacobian sub-block. +#[must_use] +pub fn tangent_wrt_subset( + arena: &mut Arena, + expr: NodeId, + active_indices: &HashSet, +) -> NodeId { + let mut memo: NodeMap = NodeMap::new(arena.len()); + let widths = primal_widths(arena, expr); + differentiate( + arena, + expr, + DiffTarget::States, + &mut memo, + Some(active_indices), + &widths, + ) +} + +/// The literal value of `id`, or `None` for any other node — nothing is folded, +/// so a constant-valued subtree still reads as non-scalar here. +fn get_scalar(arena: &Arena, id: NodeId) -> Option { + match arena.get(id) { + Node::Scalar(v) => Some(*v), + _ => None, + } +} + +/// Whether `v` is a whole number, to an absolute tolerance of `f64::EPSILON`. +fn is_integer(v: f64) -> bool { + (v.round() - v).abs() < f64::EPSILON +} + +/// Differentiate `id`, memoising one derivative per primal node so a shared +/// subtree is differentiated once. +/// +/// `widths` carries the primal output widths that zero derivatives must match. +#[allow(clippy::match_same_arms, clippy::branches_sharing_code)] +fn differentiate( + arena: &mut Arena, + id: NodeId, + mode: DiffTarget, + memo: &mut NodeMap, + state_filter: Option<&HashSet>, + widths: &[usize], +) -> NodeId { + if let Some(&cached) = memo.get(id) { + return cached; + } + + let width = widths.get(id.index()).copied().unwrap_or(1); + + let result = match arena.get(id).clone() { + // Constants - derivative is 0 + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::Time => zero_of_width(arena, width), + + // d(y')/dy is the mass matrix, which the solver supplies separately. + Node::StateVectorDot { .. } => zero_of_width(arena, width), + + // Only first order: a tangent node has no tangent of its own. + Node::TangentStateVector { .. } | Node::TangentParameter { .. } => { + zero_of_width(arena, width) + }, + + // Active variables + Node::StateVector { start, end } => { + if mode == DiffTarget::States { + let active = + state_filter.is_none_or(|filter| (start..end).any(|i| filter.contains(&i))); + if active { + arena.alloc(Node::TangentStateVector { start, end }) + } else { + zero_of_width(arena, width) + } + } else { + // dy/dp = 0 + zero_of_width(arena, width) + } + }, + + Node::InputParameter { + index, + width: param_width, + .. + } => { + if mode == DiffTarget::Params { + // dp/dp = tangent_p, replicated to the packed width so a width-1 + // tangent under an Index slice can't read outside its buffer slot. + let tangent = arena.alloc(Node::TangentParameter { index }); + if param_width > 1 { + arena.alloc(Node::Concat(vec![tangent; param_width])) + } else { + tangent + } + } else { + // dp/dy = 0 + zero_of_width(arena, width) + } + }, + + // Binary operations + Node::Add(a, b) => { + // d(a + b) = da + db + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + arena.alloc(Node::Add(da, db)) + }, + + Node::Sub(a, b) => { + // d(a - b) = da - db + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + arena.alloc(Node::Sub(da, db)) + }, + + Node::Mul(a, b) => { + // d(a * b) = a * db + da * b (product rule) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + let a_db = arena.alloc(Node::Mul(a, db)); + let da_b = arena.alloc(Node::Mul(da, b)); + arena.alloc(Node::Add(a_db, da_b)) + }, + + Node::Div(a, b) => { + // d(a / b) = (da * b - a * db) / b^2 (quotient rule) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + let da_b = arena.alloc(Node::Mul(da, b)); + let a_db = arena.alloc(Node::Mul(a, db)); + let numer = arena.alloc(Node::Sub(da_b, a_db)); + let two = arena.alloc(Node::Scalar(2.0)); + let b_sq = arena.alloc(Node::Pow(b, two)); + arena.alloc(Node::Div(numer, b_sq)) + }, + + Node::Pow(base, exp) => { + // A constant exponent takes the power rule, which holds for a + // negative base; a varying one needs log(base). + if let Some(n) = get_scalar(arena, exp) { + // The integer test only changes which node carries the exponent. + if is_integer(n) { + // Integer power: d(a^n) = n * a^(n-1) * da + let da = differentiate(arena, base, mode, memo, state_filter, widths); + let n_scalar = arena.alloc(Node::Scalar(n)); + let n_minus_1 = arena.alloc(Node::Scalar(n - 1.0)); + let base_pow_nm1 = arena.alloc(Node::Pow(base, n_minus_1)); + let n_times_pow = arena.alloc(Node::Mul(n_scalar, base_pow_nm1)); + arena.alloc(Node::Mul(n_times_pow, da)) + } else { + // Non-integer constant exponent: d(a^b) = b * a^(b-1) * da + let da = differentiate(arena, base, mode, memo, state_filter, widths); + let b_minus_1 = arena.alloc(Node::Scalar(n - 1.0)); + let base_pow_bm1 = arena.alloc(Node::Pow(base, b_minus_1)); + let b_times_pow = arena.alloc(Node::Mul(exp, base_pow_bm1)); + arena.alloc(Node::Mul(b_times_pow, da)) + } + } else { + // General case: d(a^b) = a^b * (b * da/a + log(a) * db) + let da = differentiate(arena, base, mode, memo, state_filter, widths); + let db = differentiate(arena, exp, mode, memo, state_filter, widths); + let a_pow_b = arena.alloc(Node::Pow(base, exp)); + let log_a = arena.alloc(Node::Log(base)); + let da_over_a = arena.alloc(Node::Div(da, base)); + let b_da_over_a = arena.alloc(Node::Mul(exp, da_over_a)); + let log_a_db = arena.alloc(Node::Mul(log_a, db)); + let inner = arena.alloc(Node::Add(b_da_over_a, log_a_db)); + arena.alloc(Node::Mul(a_pow_b, inner)) + } + }, + + // Unary operations + Node::Neg(a) => { + let da = differentiate(arena, a, mode, memo, state_filter, widths); + arena.alloc(Node::Neg(da)) + }, + + Node::Abs(a) => { + // d(|a|) = sign(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let sign_a = arena.alloc(Node::Sign(a)); + arena.alloc(Node::Mul(sign_a, da)) + }, + + Node::Sqrt(a) => { + // d(sqrt(a)) = da / (2 * sqrt(a)) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let two = arena.alloc(Node::Scalar(2.0)); + let sqrt_a = arena.alloc(Node::Sqrt(a)); + let two_sqrt_a = arena.alloc(Node::Mul(two, sqrt_a)); + arena.alloc(Node::Div(da, two_sqrt_a)) + }, + + Node::Exp(a) => { + // d(exp(a)) = exp(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let exp_a = arena.alloc(Node::Exp(a)); + arena.alloc(Node::Mul(exp_a, da)) + }, + + Node::Log(a) => { + // d(log(a)) = da / a + let da = differentiate(arena, a, mode, memo, state_filter, widths); + arena.alloc(Node::Div(da, a)) + }, + + Node::Sin(a) => { + // d(sin(a)) = cos(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let cos_a = arena.alloc(Node::Cos(a)); + arena.alloc(Node::Mul(cos_a, da)) + }, + + Node::Cos(a) => { + // d(cos(a)) = -sin(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let sin_a = arena.alloc(Node::Sin(a)); + let neg_sin_a = arena.alloc(Node::Neg(sin_a)); + arena.alloc(Node::Mul(neg_sin_a, da)) + }, + + Node::Tanh(a) => { + // d(tanh(a)) = (1 - tanh(a)^2) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let tanh_a = arena.alloc(Node::Tanh(a)); + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let tanh_sq = arena.alloc(Node::Pow(tanh_a, two)); + let one_minus_tanh_sq = arena.alloc(Node::Sub(one, tanh_sq)); + arena.alloc(Node::Mul(one_minus_tanh_sq, da)) + }, + + Node::Sinh(a) => { + // d(sinh(a)) = cosh(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let cosh_a = arena.alloc(Node::Cosh(a)); + arena.alloc(Node::Mul(cosh_a, da)) + }, + + Node::Cosh(a) => { + // d(cosh(a)) = sinh(a) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let sinh_a = arena.alloc(Node::Sinh(a)); + arena.alloc(Node::Mul(sinh_a, da)) + }, + + Node::Arcsinh(a) => { + // d(arcsinh(a)) = da / sqrt(1 + a^2) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let one_plus_a_sq = arena.alloc(Node::Add(one, a_sq)); + let sqrt_denom = arena.alloc(Node::Sqrt(one_plus_a_sq)); + arena.alloc(Node::Div(da, sqrt_denom)) + }, + + Node::Arctan(a) => { + // d(arctan(a)) = da / (1 + a^2) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let one_plus_a_sq = arena.alloc(Node::Add(one, a_sq)); + arena.alloc(Node::Div(da, one_plus_a_sq)) + }, + + Node::Erf(a) => { + // d(erf(a)) = 2/sqrt(pi) * exp(-a^2) * da + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let two_over_sqrt_pi = arena.alloc(Node::Scalar(2.0 / std::f64::consts::PI.sqrt())); + let neg_one = arena.alloc(Node::Scalar(-1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let neg_a_sq = arena.alloc(Node::Mul(neg_one, a_sq)); + let exp_neg_a_sq = arena.alloc(Node::Exp(neg_a_sq)); + let coeff = arena.alloc(Node::Mul(two_over_sqrt_pi, exp_neg_a_sq)); + arena.alloc(Node::Mul(coeff, da)) + }, + + // Step functions: zero away from the jumps, which are not modelled. + Node::Sign(_) | Node::Floor(_) | Node::Ceiling(_) => zero_of_width(arena, width), + Node::EqualHeaviside(_, _) | Node::NotEqualHeaviside(_, _) | Node::Equality(_, _) => { + zero_of_width(arena, width) + }, + + Node::Minimum(a, b) => { + // Subgradient: the (a <= b) indicator picks da, its complement db. + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + let selector = arena.alloc(Node::EqualHeaviside(a, b)); + let one = arena.alloc(Node::Scalar(1.0)); + let one_minus_sel = arena.alloc(Node::Sub(one, selector)); + let sel_da = arena.alloc(Node::Mul(selector, da)); + let not_sel_db = arena.alloc(Node::Mul(one_minus_sel, db)); + arena.alloc(Node::Add(sel_da, not_sel_db)) + }, + + Node::Maximum(a, b) => { + // Same, with the operands swapped: EqualHeaviside(b, a) is (a >= b). + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + let selector = arena.alloc(Node::EqualHeaviside(b, a)); + let one = arena.alloc(Node::Scalar(1.0)); + let one_minus_sel = arena.alloc(Node::Sub(one, selector)); + let sel_da = arena.alloc(Node::Mul(selector, da)); + let not_sel_db = arena.alloc(Node::Mul(one_minus_sel, db)); + arena.alloc(Node::Add(sel_da, not_sel_db)) + }, + + Node::Modulo(a, _b) => { + // d(a % b) = da (ignoring discontinuities at integer multiples of b) + differentiate(arena, a, mode, memo, state_filter, widths) + }, + + Node::Hypot(a, b) => { + // d(hypot(a, b)) = (a*da + b*db) / hypot(a, b) + let da = differentiate(arena, a, mode, memo, state_filter, widths); + let db = differentiate(arena, b, mode, memo, state_filter, widths); + let a_da = arena.alloc(Node::Mul(a, da)); + let b_db = arena.alloc(Node::Mul(b, db)); + let numer = arena.alloc(Node::Add(a_da, b_db)); + let hypot_ab = arena.alloc(Node::Hypot(a, b)); + arena.alloc(Node::Div(numer, hypot_ab)) + }, + + Node::MatMul(a, b) => { + // The matrix is always constant, so it contributes no term of its own. + let db = differentiate(arena, b, mode, memo, state_filter, widths); + arena.alloc(Node::MatMul(a, db)) + }, + + Node::Index { child, start, end } => { + // d(v[start:end]) = dv[start:end] + let d_child = differentiate(arena, child, mode, memo, state_filter, widths); + arena.alloc(Node::Index { + child: d_child, + start, + end, + }) + }, + + Node::Concat(children) => { + // d(concat(a, b, ...)) = concat(da, db, ...) + let d_children: Vec = children + .iter() + .map(|&c| differentiate(arena, c, mode, memo, state_filter, widths)) + .collect(); + arena.alloc(Node::Concat(d_children)) + }, + + Node::Interpolant1DLinear { data, child } => { + let d_child = differentiate(arena, child, mode, memo, state_filter, widths); + + // Slopes are baked into the node here, so evaluation only has to + // pick a segment. + let slopes = compute_interpolant_slopes(&data); + + let deriv_interp = arena.alloc(Node::Interpolant1DLinearDeriv { + slopes: slopes.into_boxed_slice(), + x_data: data.x_data.clone().into_boxed_slice(), + child, + }); + + arena.alloc(Node::Mul(deriv_interp, d_child)) + }, + + // Only first order: the solver Jacobian never needs an interpolant's + // second derivative. + Node::Interpolant1DLinearDeriv { .. } => zero_of_width(arena, width), + + // Cubic/pchip interpolation: d(interp(x)) = interp'(x) * dx. + Node::Interpolant1DCubic { data, child } => { + let d_child = differentiate(arena, child, mode, memo, state_filter, widths); + let deriv_interp = arena.alloc(Node::Interpolant1DCubicDeriv { data, child }); + arena.alloc(Node::Mul(deriv_interp, d_child)) + }, + // Only first order, as for the linear interpolant above. + Node::Interpolant1DCubicDeriv { .. } => zero_of_width(arena, width), + + // Multivariate chain rule: sum over axes of ∂interp/∂x_a · dg_a. + Node::InterpolantNd { data, children } => { + let mut sum: Option = None; + for (axis, &child) in children.iter().enumerate() { + let d_child = differentiate(arena, child, mode, memo, state_filter, widths); + let partial = arena.alloc(Node::InterpolantNdPartial { + data: data.clone(), + children: children.clone(), + axis: u32::try_from(axis).expect("axis index fits in u32"), + }); + let term = arena.alloc(Node::Mul(partial, d_child)); + sum = Some(sum.map_or(term, |acc| arena.alloc(Node::Add(acc, term)))); + } + sum.expect("InterpolantNd has at least one child") + }, + // Only first order, matching the 1D interpolant-derivative treatment. + Node::InterpolantNdPartial { .. } => zero_of_width(arena, width), + + Node::Conditional { selector, branches } => { + // The selector only switches, so it carries no derivative. + let d_branches: Vec = branches + .iter() + .map(|&b| differentiate(arena, b, mode, memo, state_filter, widths)) + .collect(); + arena.alloc(Node::Conditional { + selector, + branches: d_branches, + }) + }, + + // Picks the tangent component at the argmax, ties going to the first + // occurrence, as `ReduceArgSelect` evaluation does for the primal. + Node::MaxReduce(a) => { + let da = differentiate(arena, a, mode, memo, state_filter, widths); + arena.alloc(Node::ReduceArgSelect { + basis: da, + picker: a, + is_max: true, + }) + }, + Node::MinReduce(a) => { + let da = differentiate(arena, a, mode, memo, state_filter, widths); + arena.alloc(Node::ReduceArgSelect { + basis: da, + picker: a, + is_max: false, + }) + }, + + // Only first order: this node exists only inside a derivative tape, + // which is never differentiated again. + Node::ReduceArgSelect { .. } => zero_of_width(arena, width), + }; + + memo.insert(id, result); + result +} + +/// Segment slopes of a piecewise linear interpolant: `n - 1` values for `n` +/// breakpoints, and zero across a segment whose breakpoints coincide. +pub(crate) fn compute_interpolant_slopes(data: &InterpolantData) -> Vec { + let n = data.x_data.len(); + if n < 2 { + return vec![]; + } + + let mut slopes = Vec::with_capacity(n - 1); + for i in 0..n - 1 { + let dx = data.x_data[i + 1] - data.x_data[i]; + let dy = data.y_data[i + 1] - data.y_data[i]; + slopes.push(if dx.abs() > f64::EPSILON { + dy / dx + } else { + 0.0 + }); + } + slopes +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::{CompiledExpr, TangentInputs}; + use crate::ir::TypedIr; + use crate::simplify::{SimplifyMode, cse, dce, simplify, simplify_with_mode}; + use crate::zero_propagate::zero_propagate; + + /// Run the binding/jacobian tangent pipeline end to end and return the + /// compiled tape. Mirrors `function.rs::tangent_expr` / + /// `jacobian.rs::finish` exactly so width-collapse regressions surface here. + fn compile_tangent_pipeline(arena: &Arena, root: NodeId, wrt_params: bool) -> CompiledExpr { + let mut a = arena.clone(); + let root = if wrt_params { + tangent_wrt_params(&mut a, root) + } else { + tangent_wrt_states(&mut a, root) + }; + let root = simplify_with_mode(&mut a, root, SimplifyMode::Aggressive); + let (za, root) = zero_propagate(&a, root); + let (ca, root) = cse(&za, root); + let (da, root) = dce(&ca, root); + CompiledExpr::from_ir(TypedIr::from_arena(&da, root)) + } + + #[test] + fn test_jvp_width_vector_plus_scalar_param() { + // f = y[0:3] + p0, so df/dp @ [1] must be [1, 1, 1]: the tangent_p tape + // must not collapse to length 1 and truncate the JVP to [1, 0, 0]. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let p = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Add(y, p)); + + let tp = compile_tangent_pipeline(&arena, root, true); + assert_eq!( + tp.output_len(), + 3, + "tangent_p tape collapsed below primal width" + ); + let mut s = vec![0.0; tp.scratch_len()]; + let tangent = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let r = tp.eval_with_tangent(&mut s, 0.0, &[0.0, 0.0, 0.0], &[], &[5.0], &tangent); + assert_eq!(r, &[1.0, 1.0, 1.0]); + } + + #[test] + fn test_jvp_width_wide_param_index_slice() { + // f = b[1]*y0, b width-2: a width-1 tangent under Index{1..2} would read + // outside its buffer slot. Seeded direction gives y0; zero direction gives 0. + let mut arena = Arena::new(); + let b = arena.alloc(Node::InputParameter { + name: "b".into(), + index: 0, + offset: 0, + width: 2, + }); + let b1 = arena.alloc(Node::Index { + child: b, + start: 1, + end: 2, + }); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Mul(b1, y0)); + + let tp = compile_tangent_pipeline(&arena, root, true); + let mut s = vec![0.0; tp.scratch_len()]; + let seeded = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let r = tp.eval_with_tangent(&mut s, 0.0, &[3.0], &[], &[10.0, 20.0], &seeded); + assert_eq!(r, &[3.0]); + + let mut s2 = vec![0.0; tp.scratch_len()]; + let unseeded = TangentInputs { + dy: None, + dp: Some(&[0.0]), + }; + let r0 = tp.eval_with_tangent(&mut s2, 0.0, &[3.0], &[], &[10.0, 20.0], &unseeded); + assert_eq!(r0, &[0.0]); + } + + #[test] + fn test_jvp_width_zero_dfdy_keeps_length() { + // f = const_vector([1,2,3]) + p0; df/dy ≡ 0. Repro 2. + // The tangent_y tape must keep length 3 (all zeros), not collapse to [0]. + use crate::node::ArrayData; + let mut arena = Arena::new(); + let cv = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: crate::node::Shape::vector(3), + }))); + let p = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Add(cv, p)); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!( + ty.output_len(), + 3, + "tangent_y tape collapsed below primal width" + ); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[0.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[], &[], &[5.0], &tangent); + assert_eq!(r, &[0.0, 0.0, 0.0]); + } + + #[test] + fn test_jvp_width_concat_collapsed_child_offset() { + // f = concat(const_vec[1,2], y0); df/dy must place y0's derivative at + // row 2, not shift it to row 0 when the const-vector child collapses. + use crate::node::ArrayData; + let mut arena = Arena::new(); + let cv = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0], + shape: crate::node::Shape::vector(2), + }))); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![cv, y0])); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!(ty.output_len(), 3, "concat tangent width must match primal"); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[7.0], &[], &[], &tangent); + // rows 0,1 come from the constant vector (df/dy = 0); row 2 = d(y0)/dy = 1. + assert_eq!(r, &[0.0, 0.0, 1.0]); + } + + #[test] + fn test_jvp_width_concat_empty_child_offset() { + // f = concat(empty_vec[], y0 + 1); mirrors a pure-algebraic PyBaMM model + // (concatenated_rhs is a length-0 Vector). df/dy0 must land at row 0, + // not be swallowed by a phantom length-1 zero for the length-0 child. + use crate::node::ArrayData; + let mut arena = Arena::new(); + let empty = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![], + shape: crate::node::Shape::vector(0), + }))); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Add(y0, one)); + let root = arena.alloc(Node::Concat(vec![empty, expr])); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!( + ty.output_len(), + 1, + "concat tangent width must match primal (empty child contributes 0)" + ); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[7.0], &[], &[], &tangent); + assert_eq!( + r, + &[1.0], + "d(y0 + 1)/dy0 must be 1, not swallowed by the empty child" + ); + } + + #[test] + fn test_product_rule() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let product = arena.alloc(Node::Mul(x, y)); + + let jac = tangent_wrt_states(&mut arena, product); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // d(x*y) @ [1, 0] = y = 4 + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[2.0, 4.0], &[], &[], &tangent); + assert!( + (result[0] - 4.0).abs() < 1e-14, + "Expected 4.0, got {}", + result[0] + ); + + // d(x*y) @ [0, 1] = x = 2 + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[2.0, 4.0], &[], &[], &tangent); + assert!( + (result[0] - 2.0).abs() < 1e-14, + "Expected 2.0, got {}", + result[0] + ); + } + + #[test] + fn test_diff_max_reduce_selects_argmax_tangent() { + // d(max(y)) @ seed = seed[argmax(y)]; y=[0.3,0.9,0.1] -> k=1 -> seed[1]=20 + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::MaxReduce(y)); + let jac = tangent_wrt_states(&mut arena, expr); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s = vec![0.0; compiled.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[10.0, 20.0, 30.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s, 0.0, &[0.3, 0.9, 0.1], &[], &[], &tangent); + assert!( + (result[0] - 20.0).abs() < 1e-14, + "expected 20.0, got {}", + result[0] + ); + } + + #[test] + fn test_diff_min_reduce_selects_argmin_tangent() { + // d(min(y)) @ seed = seed[argmin(y)]; y=[0.3,0.9,0.1] -> k=2 -> seed[2]=30 + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::MinReduce(y)); + let jac = tangent_wrt_states(&mut arena, expr); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s = vec![0.0; compiled.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[10.0, 20.0, 30.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s, 0.0, &[0.3, 0.9, 0.1], &[], &[], &tangent); + assert!( + (result[0] - 30.0).abs() < 1e-14, + "expected 30.0, got {}", + result[0] + ); + } + + #[test] + fn test_chain_rule() { + // d(exp(2*x)) / dx = 2 * exp(2*x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let two_x = arena.alloc(Node::Mul(two, x)); + let expr = arena.alloc(Node::Exp(two_x)); + + let jac = tangent_wrt_states(&mut arena, expr); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.0; + let expected = 2.0 * (2.0 * x_val).exp(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_quotient_rule() { + // d(x / y) = (y - x) / y^2 when dx=1, dy=1 + // At x=1, y=2: (2 - 1) / 4 = 0.25 + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let quotient = arena.alloc(Node::Div(x, y)); + + let jac = tangent_wrt_states(&mut arena, quotient); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // d(x/y) @ [1, 0] = 1/y = 0.5 + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[1.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] - 0.5).abs() < 1e-14, + "Expected 0.5, got {}", + result[0] + ); + + // d(x/y) @ [0, 1] = -x/y^2 = -1/4 = -0.25 + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[1.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] + 0.25).abs() < 1e-14, + "Expected -0.25, got {}", + result[0] + ); + } + + #[test] + fn test_power_rule_integer() { + // d(x^3) / dx = 3x^2 + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let three = arena.alloc(Node::Scalar(3.0)); + let x_cubed = arena.alloc(Node::Pow(x, three)); + + let jac = tangent_wrt_states(&mut arena, x_cubed); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val = 2.0; + let expected = 3.0 * x_val * x_val; // 12.0 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_power_rule_negative_base() { + // d(x^2) / dx = 2x (works for negative x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let x_sq = arena.alloc(Node::Pow(x, two)); + + let jac = tangent_wrt_states(&mut arena, x_sq); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val = -3.0; + let expected = 2.0 * x_val; // -6.0 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_sqrt_derivative() { + // d(sqrt(x)) / dx = 1 / (2*sqrt(x)) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sqrt_x = arena.alloc(Node::Sqrt(x)); + + let jac = tangent_wrt_states(&mut arena, sqrt_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 4.0; + let expected = 0.5 / x_val.sqrt(); // 0.25 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_log_derivative() { + // d(log(x)) / dx = 1/x + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let log_x = arena.alloc(Node::Log(x)); + + let jac = tangent_wrt_states(&mut arena, log_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val = 2.0; + let expected = 1.0 / x_val; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_sin_derivative() { + // d(sin(x)) / dx = cos(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sin_x = arena.alloc(Node::Sin(x)); + + let jac = tangent_wrt_states(&mut arena, sin_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val = std::f64::consts::PI / 4.0; + let expected = x_val.cos(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_cos_derivative() { + // d(cos(x)) / dx = -sin(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let cos_x = arena.alloc(Node::Cos(x)); + + let jac = tangent_wrt_states(&mut arena, cos_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val = std::f64::consts::PI / 4.0; + let expected = -x_val.sin(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_tanh_derivative() { + // d(tanh(x)) / dx = 1 - tanh(x)^2 + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let tanh_x = arena.alloc(Node::Tanh(x)); + + let jac = tangent_wrt_states(&mut arena, tanh_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.0; + let tanh_val = x_val.tanh(); + let expected = tanh_val.mul_add(-tanh_val, 1.0); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_sinh_derivative() { + // d(sinh(x)) / dx = cosh(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sinh_x = arena.alloc(Node::Sinh(x)); + + let jac = tangent_wrt_states(&mut arena, sinh_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.5; + let expected = x_val.cosh(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_cosh_derivative() { + // d(cosh(x)) / dx = sinh(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let cosh_x = arena.alloc(Node::Cosh(x)); + + let jac = tangent_wrt_states(&mut arena, cosh_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.5; + let expected = x_val.sinh(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_arcsinh_derivative() { + // d(arcsinh(x)) / dx = 1 / sqrt(1 + x^2) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let arcsinh_x = arena.alloc(Node::Arcsinh(x)); + + let jac = tangent_wrt_states(&mut arena, arcsinh_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 2.0; + let expected = 1.0 / x_val.mul_add(x_val, 1.0).sqrt(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_arctan_derivative() { + // d(arctan(x)) / dx = 1 / (1 + x^2) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let arctan_x = arena.alloc(Node::Arctan(x)); + + let jac = tangent_wrt_states(&mut arena, arctan_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.0; + let expected = 1.0 / x_val.mul_add(x_val, 1.0); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_neg_derivative() { + // d(-x) / dx = -1 + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_x = arena.alloc(Node::Neg(x)); + + let jac = tangent_wrt_states(&mut arena, neg_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0], &[], &[], &tangent); + assert!( + (result[0] + 1.0).abs() < 1e-14, + "Expected -1.0, got {}", + result[0] + ); + } + + #[test] + fn test_abs_derivative() { + // d(|x|) / dx = sign(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let abs_x = arena.alloc(Node::Abs(x)); + + let jac = tangent_wrt_states(&mut arena, abs_x); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0], &[], &[], &tangent); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[-5.0], &[], &[], &tangent); + assert!( + (result[0] + 1.0).abs() < 1e-14, + "Expected -1.0, got {}", + result[0] + ); + } + + #[test] + fn test_index_derivative() { + // d(x[1:3]) / dx = dx[1:3] + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let slice = arena.alloc(Node::Index { + child: x, + start: 1, + end: 3, + }); + + let jac = tangent_wrt_states(&mut arena, slice); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0, 2.0, 0.0, 0.0]), + dp: None, + }; + let result = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[0.0, 1.0, 2.0, 3.0, 4.0], + &[], + &[], + &tangent, + ); + assert_eq!(result.len(), 2); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + assert!( + (result[1] - 2.0).abs() < 1e-14, + "Expected 2.0, got {}", + result[1] + ); + } + + #[test] + fn test_concat_derivative() { + // d(concat(x, y)) = concat(dx, dy) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y = arena.alloc(Node::StateVector { start: 2, end: 4 }); + let concat = arena.alloc(Node::Concat(vec![x, y])); + + let jac = tangent_wrt_states(&mut arena, concat); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[1.0, 2.0, 3.0, 4.0]), + dp: None, + }; + let result = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[0.0, 0.0, 0.0, 0.0], + &[], + &[], + &tangent, + ); + assert_eq!(result.len(), 4); + assert!((result[0] - 1.0).abs() < 1e-14); + assert!((result[1] - 2.0).abs() < 1e-14); + assert!((result[2] - 3.0).abs() < 1e-14); + assert!((result[3] - 4.0).abs() < 1e-14); + } + + #[test] + fn test_parameter_differentiation() { + // d(p * x) / dp = x + let mut arena = Arena::new(); + let p = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Mul(p, x)); + + let jac = tangent_wrt_params(&mut arena, expr); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let x_val = 3.0; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[2.0], &tangent); + assert!( + (result[0] - x_val).abs() < 1e-14, + "Expected {}, got {}", + x_val, + result[0] + ); + } + + #[test] + fn test_constant_derivative() { + let mut arena = Arena::new(); + let c = arena.alloc(Node::Scalar(42.0)); + + let jac = tangent_wrt_states(&mut arena, c); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[], &[], &[], &tangent); + assert!((result[0]).abs() < 1e-14, "Expected 0.0, got {}", result[0]); + } + + #[test] + fn test_time_derivative() { + // d(t) / dy = 0 + let mut arena = Arena::new(); + let t = arena.alloc(Node::Time); + + let jac = tangent_wrt_states(&mut arena, t); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 5.0, &[], &[], &[], &tangent); + assert!((result[0]).abs() < 1e-14, "Expected 0.0, got {}", result[0]); + } + + #[test] + fn test_minimum_derivative() { + // d(min(x, y)) = dx if x <= y, else dy + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let min_xy = arena.alloc(Node::Minimum(x, y)); + + let jac = tangent_wrt_states(&mut arena, min_xy); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // x < y, so derivative is dx + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[1.0, 3.0], &[], &[], &tangent); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + + // x > y, so derivative is dy + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + } + + #[test] + fn test_maximum_derivative() { + // d(max(x, y)) = dx if x >= y, else dy + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let max_xy = arena.alloc(Node::Maximum(x, y)); + + let jac = tangent_wrt_states(&mut arena, max_xy); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // x > y, so derivative is dx + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + + // x < y, so derivative is dy + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[1.0, 3.0], &[], &[], &tangent); + assert!( + (result[0] - 1.0).abs() < 1e-14, + "Expected 1.0, got {}", + result[0] + ); + } + + #[test] + fn test_hypot_derivative() { + // d(hypot(x, y)) = (x*dx + y*dy) / hypot(x, y) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let hypot_xy = arena.alloc(Node::Hypot(x, y)); + + let jac = tangent_wrt_states(&mut arena, hypot_xy); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // At (3, 4), hypot = 5 + // d(hypot)/dx = 3/5 = 0.6 + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[3.0, 4.0], &[], &[], &tangent); + assert!( + (result[0] - 0.6).abs() < 1e-12, + "Expected 0.6, got {}", + result[0] + ); + + // d(hypot)/dy = 4/5 = 0.8 + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[3.0, 4.0], &[], &[], &tangent); + assert!( + (result[0] - 0.8).abs() < 1e-12, + "Expected 0.8, got {}", + result[0] + ); + } + + #[test] + fn test_interpolation_derivative() { + // d(interp(x)) / dx = slope at x + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(InterpolantData { + x_data: vec![0.0, 1.0, 2.0], + y_data: vec![0.0, 10.0, 30.0], // slopes: 10, 20 + }), + child: x, + }); + + let jac = tangent_wrt_states(&mut arena, interp); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // x=0.5 lands in segment [0, 1], slope 10 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[0.5], &[], &[], &tangent); + assert!( + (result[0] - 10.0).abs() < 1e-12, + "Expected 10.0, got {}", + result[0] + ); + + // x=1.5 lands in segment [1, 2], slope 20 + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[1.5], &[], &[], &tangent); + assert!( + (result[0] - 20.0).abs() < 1e-12, + "Expected 20.0, got {}", + result[0] + ); + + // Out-of-domain: extend with boundary-segment slope (not 0) + // Below x=0.0: first segment slope = 10 + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[-0.5], &[], &[], &tangent); + assert!( + (result[0] - 10.0).abs() < 1e-12, + "Expected 10.0 (first-segment slope) below domain, got {}", + result[0] + ); + + // Above x=2.0: last segment slope = 20 + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[2.5], &[], &[], &tangent); + assert!( + (result[0] - 20.0).abs() < 1e-12, + "Expected 20.0 (last-segment slope) above domain, got {}", + result[0] + ); + } + + #[test] + fn test_cubic_interpolation_derivative() { + use crate::node::CubicInterpolantData; + // p(dx) = 1 + 2*dx + 3*dx^2 + 4*dx^3 on [0, 5]; p'(dx) = 2 + 6*dx + 12*dx^2. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 5.0], + coeffs: vec![[1.0, 2.0, 3.0, 4.0]], + }), + child: x, + }); + let jac = tangent_wrt_states(&mut arena, interp); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s = vec![0.0; compiled.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + + // At x=1 (dx=1): p' = 2 + 6 + 12 = 20 + let r = compiled.eval_with_tangent(&mut s, 0.0, &[1.0], &[], &[], &tangent); + assert!((r[0] - 20.0).abs() < 1e-12, "expected 20, got {}", r[0]); + // At x=0 (dx=0): p' = 2 + let r = compiled.eval_with_tangent(&mut s, 0.0, &[0.0], &[], &[], &tangent); + assert!((r[0] - 2.0).abs() < 1e-12, "expected 2, got {}", r[0]); + // Below domain x=-1 (clamps to interval 0, dx=-1): p' = 2 - 6 + 12 = 8 + let r = compiled.eval_with_tangent(&mut s, 0.0, &[-1.0], &[], &[], &tangent); + assert!((r[0] - 8.0).abs() < 1e-12, "expected 8, got {}", r[0]); + // Above domain x=6 (clamps to interval 0, dx=6): p' = 2 + 36 + 432 = 470 + let r = compiled.eval_with_tangent(&mut s, 0.0, &[6.0], &[], &[], &tangent); + assert!((r[0] - 470.0).abs() < 1e-12, "expected 470, got {}", r[0]); + } + + #[test] + fn test_nd_interpolation_partial_derivatives() { + use crate::node::NdInterpolantData; + // p = 3*dx0^2 + 5*dx1^3 + dx0*dx1 on one cell [0,4]x[0,4]. + // ∂p/∂x0 = 6*dx0 + dx1; ∂p/∂x1 = 15*dx1^2 + dx0. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let mut coeffs = vec![0.0; 16]; + coeffs[2 * 4] = 3.0; // dx0^2 + coeffs[3] = 5.0; // dx1^3 + coeffs[4 + 1] = 1.0; // dx0*dx1 + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 4.0], vec![0.0, 4.0]], + coeffs, + order: 4, + }), + children: vec![x0, x1], + }); + let jac = tangent_wrt_states(&mut arena, interp); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s = vec![0.0; compiled.scratch_len()]; + + // At (2,1): ∂p/∂x0 = 12 + 1 = 13; ∂p/∂x1 = 15 + 2 = 17 (seeded + // independently so each partial is verified on its own). + let t0 = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let r = compiled.eval_with_tangent(&mut s, 0.0, &[2.0, 1.0], &[], &[], &t0); + assert!((r[0] - 13.0).abs() < 1e-12, "expected 13, got {}", r[0]); + let t1 = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let r = compiled.eval_with_tangent(&mut s, 0.0, &[2.0, 1.0], &[], &[], &t1); + assert!((r[0] - 17.0).abs() < 1e-12, "expected 17, got {}", r[0]); + // Directional derivative sums both partials: 13 + 17 = 30. + let tb = TangentInputs { + dy: Some(&[1.0, 1.0]), + dp: None, + }; + let r = compiled.eval_with_tangent(&mut s, 0.0, &[2.0, 1.0], &[], &[], &tb); + assert!((r[0] - 30.0).abs() < 1e-12, "expected 30, got {}", r[0]); + } + + #[test] + fn test_nd_interpolation_bilinear_partials() { + use crate::node::NdInterpolantData; + // p = 2 + 3*dx0 + 4*dx1 + 5*dx0*dx1 on one cell [0,2]x[0,2]. + // ∂p/∂x0 = 3 + 5*dx1; ∂p/∂x1 = 4 + 5*dx0. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 2.0], vec![0.0, 2.0]], + coeffs: vec![2.0, 4.0, 3.0, 5.0], + order: 2, + }), + children: vec![x0, x1], + }); + let jac = tangent_wrt_states(&mut arena, interp); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s = vec![0.0; compiled.scratch_len()]; + + // At (1,2): ∂p/∂x0 = 3 + 10 = 13; ∂p/∂x1 = 4 + 5 = 9. + let t0 = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let r = compiled.eval_with_tangent(&mut s, 0.0, &[1.0, 2.0], &[], &[], &t0); + assert!((r[0] - 13.0).abs() < 1e-12, "expected 13, got {}", r[0]); + let t1 = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let r = compiled.eval_with_tangent(&mut s, 0.0, &[1.0, 2.0], &[], &[], &t1); + assert!((r[0] - 9.0).abs() < 1e-12, "expected 9, got {}", r[0]); + } + + #[test] + fn test_memoization() { + // x * x reaches the same child twice, so the memo is what keeps one + // tangent of x rather than two. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x_sq = arena.alloc(Node::Mul(x, x)); + + let jac = tangent_wrt_states(&mut arena, x_sq); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // d(x^2) @ [1] = 2x at x=3 => 6 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[3.0], &[], &[], &tangent); + assert!( + (result[0] - 6.0).abs() < 1e-12, + "Expected 6.0, got {}", + result[0] + ); + } + + #[test] + fn test_conditional_derivative() { + // d(cond(sel, [b1, b2])) = cond(sel, [db1, db2]) + let mut arena = Arena::new(); + let selector = arena.alloc(Node::Scalar(1.0)); // Select first branch + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let branch1 = arena.alloc(Node::Mul(two, x)); // 2x + let three = arena.alloc(Node::Scalar(3.0)); + let branch2 = arena.alloc(Node::Mul(three, x)); // 3x + let cond = arena.alloc(Node::Conditional { + selector, + branches: vec![branch1, branch2], + }); + + let jac = tangent_wrt_states(&mut arena, cond); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // With selector=1, should select d(2x) = 2 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0], &[], &[], &tangent); + assert!( + (result[0] - 2.0).abs() < 1e-12, + "Expected 2.0, got {}", + result[0] + ); + } + + #[test] + fn test_matmul_derivative() { + // d(A @ v) = A @ dv + use crate::node::{CsrData, Shape}; + + let mut arena = Arena::new(); + // 2x3, one entry per row: A @ dv = [2*dv[0], 3*dv[1]] + let sparse = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], // Diagonal with 2 and 3 + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let matmul = arena.alloc(Node::MatMul(sparse, v)); + + let jac = tangent_wrt_states(&mut arena, matmul); + let jac = simplify(&mut arena, jac); + + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // A @ [1, 0, 0] = [2, 0] + let tangent = TangentInputs { + dy: Some(&[1.0, 0.0, 0.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[0.0, 0.0, 0.0], &[], &[], &tangent); + assert_eq!(result.len(), 2); + assert!( + (result[0] - 2.0).abs() < 1e-14, + "Expected 2.0, got {}", + result[0] + ); + assert!((result[1]).abs() < 1e-14, "Expected 0.0, got {}", result[1]); + } + + #[test] + fn test_erf_derivative() { + // d(erf(x))/dx = 2/sqrt(pi) * exp(-x^2) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let erf_x = arena.alloc(Node::Erf(x)); + let jac = tangent_wrt_states(&mut arena, erf_x); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 1.0; + let two_over_sqrt_pi = 2.0 / std::f64::consts::PI.sqrt(); + let expected = two_over_sqrt_pi * (-x_val * x_val).exp(); + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + } + + #[test] + fn test_power_rule_fractional() { + // d(x^1.5) / dx = 1.5 * x^0.5 + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let exp = arena.alloc(Node::Scalar(1.5)); + let x_pow = arena.alloc(Node::Pow(x, exp)); + + let jac = tangent_wrt_states(&mut arena, x_pow); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 4.0; + let expected = 1.5 * x_val.sqrt(); // 1.5 * 2.0 = 3.0 + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut s_compiled, 0.0, &[x_val], &[], &[], &tangent); + assert!( + (result[0] - expected).abs() < 1e-12, + "Expected {}, got {}", + expected, + result[0] + ); + + let mut arena2 = Arena::new(); + let x2 = arena2.alloc(Node::StateVector { start: 0, end: 1 }); + let half = arena2.alloc(Node::Scalar(0.5)); + let x_sqrt = arena2.alloc(Node::Pow(x2, half)); + let jac2 = tangent_wrt_states(&mut arena2, x_sqrt); + let jac2 = simplify(&mut arena2, jac2); + let compiled2 = CompiledExpr::new(&arena2, jac2); + let mut s_compiled2 = vec![0.0; compiled2.scratch_len()]; + + // d(x^0.5) / dx = 0.5 * x^(-0.5) = 0.5 / sqrt(x) + let x_val2: f64 = 9.0; + let expected2 = 0.5 / x_val2.sqrt(); // 0.5 / 3.0 ≈ 0.1667 + let result2 = + compiled2.eval_with_tangent(&mut s_compiled2, 0.0, &[x_val2], &[], &[], &tangent); + assert!( + (result2[0] - expected2).abs() < 1e-12, + "Expected {}, got {}", + expected2, + result2[0] + ); + } + + #[test] + fn test_power_rule_variable_exponent() { + // d(x^y)/dx @ [1, 0] = y * x^(y-1), d(x^y)/dy @ [0, 1] = x^y * log(x) + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let x_pow_y = arena.alloc(Node::Pow(x, y)); + + let jac = tangent_wrt_states(&mut arena, x_pow_y); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(&arena, jac); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let x_val: f64 = 2.0; + let y_val: f64 = 3.0; + + // d(x^y)/dx: tangent [1, 0] -> y * x^(y-1) = 3 * 2^2 = 12 + let tangent_dx = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result_dx = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[x_val, y_val], + &[], + &[], + &tangent_dx, + ); + let expected_dx = y_val * x_val.powf(y_val - 1.0); // 3 * 4 = 12 + assert!( + (result_dx[0] - expected_dx).abs() < 1e-10, + "d(x^y)/dx: Expected {}, got {}", + expected_dx, + result_dx[0] + ); + + // d(x^y)/dy: tangent [0, 1] -> x^y * log(x) = 8 * ln(2) ≈ 5.545 + let tangent_dy = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result_dy = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[x_val, y_val], + &[], + &[], + &tangent_dy, + ); + let expected_dy = x_val.powf(y_val) * x_val.ln(); + assert!( + (result_dy[0] - expected_dy).abs() < 1e-10, + "d(x^y)/dy: Expected {}, got {}", + expected_dy, + result_dy[0] + ); + } + + #[test] + fn test_tangent_wrt_subset() { + // f = x0 * x1, differentiate only w.r.t. x1 + // df/dx1 = x0 (x0 treated as constant) + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let expr = arena.alloc(Node::Mul(x0, x1)); + + let subset = HashSet::from([1usize]); + let deriv = tangent_wrt_subset(&mut arena, expr, &subset); + let deriv = simplify(&mut arena, deriv); + + let topo = arena.topological_order(deriv); + for &nid in &topo { + if let Node::TangentStateVector { start, .. } = *arena.get(nid) { + assert_ne!(start, 0, "Should not differentiate w.r.t. x0"); + } + } + assert!( + topo.iter().any(|&nid| matches!( + arena.get(nid), + Node::TangentStateVector { start: 1, end: 2 } + )), + "Should contain TangentStateVector for x1" + ); + } + + #[test] + fn test_tangent_wrt_subset_numerical() { + // f = x0^2 + 3*x1, differentiate only w.r.t. x1 + // df/dx1 = 3 (x0 treated as constant) + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let x0_sq = arena.alloc(Node::Mul(x0, x0)); + let three = arena.alloc(Node::Scalar(3.0)); + let three_x1 = arena.alloc(Node::Mul(three, x1)); + let expr = arena.alloc(Node::Add(x0_sq, three_x1)); + + let subset = HashSet::from([1usize]); + let deriv = tangent_wrt_subset(&mut arena, expr, &subset); + let deriv = simplify(&mut arena, deriv); + let compiled = CompiledExpr::new(&arena, deriv); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + // With tangent [0, 1]: df/dx1 = 3 + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] - 3.0).abs() < 1e-12, + "Expected 3.0, got {}", + result[0] + ); + + // With tangent [1, 0]: should be 0 (x0 not in active set) + let tangent_x0 = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result_x0 = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent_x0); + assert!( + result_x0[0].abs() < 1e-12, + "Expected 0.0, got {}", + result_x0[0] + ); + } + + #[test] + fn test_tangent_wrt_subset_all_active_matches_full() { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum = arena.alloc(Node::Add(x0, x1)); + let expr = arena.alloc(Node::Mul(sum, x0)); // f = (x0 + x1) * x0 + + let all_active = HashSet::from([0usize, 1]); + + let mut arena_full = arena.clone(); + let full = tangent_wrt_states(&mut arena_full, expr); + let full = simplify(&mut arena_full, full); + let compiled_full = CompiledExpr::new(&arena_full, full); + let mut s_compiled_full = vec![0.0; compiled_full.scratch_len()]; + + let subset = tangent_wrt_subset(&mut arena, expr, &all_active); + let subset = simplify(&mut arena, subset); + let compiled_subset = CompiledExpr::new(&arena, subset); + let mut s_compiled_subset = vec![0.0; compiled_subset.scratch_len()]; + + for dy in &[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] { + let tangent = TangentInputs { + dy: Some(dy), + dp: None, + }; + let r_full = compiled_full.eval_with_tangent( + &mut s_compiled_full, + 0.0, + &[3.0, 7.0], + &[], + &[], + &tangent, + ); + let r_sub = compiled_subset.eval_with_tangent( + &mut s_compiled_subset, + 0.0, + &[3.0, 7.0], + &[], + &[], + &tangent, + ); + assert!( + (r_full[0] - r_sub[0]).abs() < 1e-12, + "Mismatch for dy={dy:?}: full={}, subset={}", + r_full[0], + r_sub[0] + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs b/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs new file mode 100644 index 0000000000..ec7a621ce9 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs @@ -0,0 +1,445 @@ +//! Batched tangent sweeps: one walk of the tangent tape for `K` colour seeds. +//! +//! Coloured Jacobian assembly runs the tangent section once per colour, and on +//! larger models nearly all of that time sits in the CSR gather inside +//! [`Instruction::MatMul`] — a scattered read whose cache line serves a single +//! lane. Widening the tangent region to `K` lanes per element turns that gather +//! into one contiguous `K`-wide load reused by every lane, so the operator is +//! streamed once per block instead of once per colour. +//! +//! The split-eval layout makes the split cheap to exploit: primal slots occupy +//! `[0, primal_buffer_size)` and are shared across lanes unchanged, while +//! tangent slots live above it and carry one value per lane. Each lane +//! accumulates in the same order as the scalar sweep, so results are bitwise +//! identical to [`CompiledExpr::run_tangent_section`]. +//! +//! [`CompiledExpr::run_tangent_section`]: crate::eval::CompiledExpr + +use crate::eval::{EQUALITY_EPS, erf_approx, sign}; +use crate::ir::{BinaryOp, BroadcastKind, ConstPool, Instruction, TypedIr, UnaryOp}; + +/// Lane counts the batched sweep is instantiated for, widest first. +pub const SUPPORTED_LANES: [usize; 2] = [8, 4]; + +/// An operand's home: primal slots are shared across lanes, tangent slots are +/// lane-minor. +#[derive(Clone, Copy)] +enum Operand { + /// Base index into the primal region. + Primal(usize), + /// Base index into the lane-minor tangent region, already scaled by `K`. + Tangent(usize), +} + +/// Classify a slot offset against the primal/tangent boundary. +#[inline] +const fn operand(offset: usize, primal_len: usize) -> Operand { + if offset < primal_len { + Operand::Primal(offset) + } else { + Operand::Tangent((offset - primal_len) * K) + } +} + +/// Whether every instruction in the tangent section has a batched form. +/// +/// Tapes carrying anything else (reductions, interpolants, branch dispatch) +/// fall back to the per-colour scalar sweep. +pub fn is_batchable(ir: &TypedIr) -> bool { + let Some(split) = ir.split_eval_info() else { + return false; + }; + ir.instructions()[split.primal_end..].iter().all(|instr| { + matches!( + instr, + Instruction::LoadTangentState { .. } + | Instruction::LoadScalar { .. } + | Instruction::LoadArray { .. } + | Instruction::FillZero { .. } + | Instruction::Binary { .. } + | Instruction::Unary { .. } + | Instruction::Index { .. } + | Instruction::Concat { .. } + | Instruction::MatMul { .. } + ) + }) +} + +/// Whether the root slot sits in the primal pool, which happens when the +/// tangent folds to a constant and so depends on no seed. +const fn root_is_primal(ir: &TypedIr, primal_len: usize) -> bool { + ir.root_slot().offset_usize() < primal_len +} + +/// Scratch length the tangent region needs for `lanes` lanes. +/// +/// A primal root has no lane-minor home of its own, so the buffer carries a +/// tail the sweep broadcasts it into, keeping one return shape for callers. +pub fn tangent_scratch_len(ir: &TypedIr, lanes: usize) -> usize { + ir.split_eval_info().map_or(0, |s| { + let tangent = (ir.buffer_size() - s.primal_buffer_size) * lanes; + let spill = if root_is_primal(ir, s.primal_buffer_size) { + ir.root_slot().len_usize() * lanes + } else { + 0 + }; + tangent + spill + }) +} + +/// Run the tangent section for `K` seed vectors at once. +/// +/// `primal` is the primal region a prior primal sweep filled, `tan` the +/// lane-minor tangent region, and `seeds` the `(n_states, K)` lane-minor seed +/// matrix. Returns the root slot as `(out_len, K)` lane-minor. +/// +/// # Panics +/// +/// Panics if the tape is not split-eval or carries an instruction outside the +/// batchable set; call [`is_batchable`] first. +pub fn run_tangent_batch<'t, const K: usize>( + ir: &TypedIr, + primal: &[f64], + tan: &'t mut [f64], + seeds: &[f64], +) -> &'t [f64] { + let split = ir + .split_eval_info() + .expect("run_tangent_batch requires a split-eval IR"); + let primal_len = split.primal_buffer_size; + let consts = ir.consts(); + + for instr in &ir.instructions()[split.primal_end..] { + exec::(*instr, primal, tan, seeds, primal_len, consts); + } + + let root = ir.root_slot(); + if root_is_primal(ir, primal_len) { + let spill = (ir.buffer_size() - primal_len) * K; + for e in 0..root.len_usize() { + tan[spill + e * K..spill + (e + 1) * K].fill(primal[root.offset_usize() + e]); + } + return &tan[spill..spill + root.len_usize() * K]; + } + let base = (root.offset_usize() - primal_len) * K; + &tan[base..base + root.len_usize() * K] +} + +#[allow(clippy::too_many_lines)] +fn exec( + instr: Instruction, + primal: &[f64], + tan: &mut [f64], + seeds: &[f64], + primal_len: usize, + consts: &ConstPool, +) { + match instr { + Instruction::LoadTangentState { start, end, dst } => { + let dst = tangent_base::(dst as usize, primal_len); + let len = (end - start) as usize * K; + tan[dst..dst + len].copy_from_slice(&seeds[start as usize * K..][..len]); + }, + + Instruction::LoadScalar { value, dst } => { + let dst = tangent_base::(dst as usize, primal_len); + tan[dst..dst + K].fill(value); + }, + + Instruction::FillZero { dst, len } => { + let dst = tangent_base::(dst as usize, primal_len); + tan[dst..dst + len as usize * K].fill(0.0); + }, + + Instruction::LoadArray { data_idx, len, dst } => { + let src = consts.get_array(data_idx, len); + let dst = tangent_base::(dst as usize, primal_len); + for (e, &value) in src.iter().enumerate() { + tan[dst + e * K..dst + (e + 1) * K].fill(value); + } + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + let dst = tangent_base::(dst as usize, primal_len); + let len = len as usize; + match operand::(src as usize + start as usize, primal_len) { + Operand::Tangent(src) => tan.copy_within(src..src + len * K, dst), + Operand::Primal(src) => { + for e in 0..len { + tan[dst + e * K..dst + (e + 1) * K].fill(primal[src + e]); + } + }, + } + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write = tangent_base::(dst as usize, primal_len); + for i in 0..sources_len as usize { + let (offset, len) = consts.concat_sources[sources_idx as usize + i]; + let len = len as usize; + match operand::(offset as usize, primal_len) { + Operand::Tangent(src) => tan.copy_within(src..src + len * K, write), + Operand::Primal(src) => { + for e in 0..len { + tan[write + e * K..write + (e + 1) * K].fill(primal[src + e]); + } + }, + } + write += len * K; + } + }, + + Instruction::Unary { op, src, dst, len } => unary::( + op, + operand::(src as usize, primal_len), + tangent_base::(dst as usize, primal_len), + len as usize, + primal, + tan, + ), + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => binary_op::( + op, + operand::(a as usize, primal_len), + operand::(b as usize, primal_len), + tangent_base::(dst as usize, primal_len), + len as usize, + kind, + primal, + tan, + ), + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx as usize]; + let dst = tangent_base::(dst as usize, primal_len); + match operand::(vec_src as usize, primal_len) { + Operand::Tangent(vec) => { + for row in 0..csr.shape.rows { + let span = csr.indptr[row]..csr.indptr[row + 1]; + let mut acc = [0.0_f64; K]; + for (&col, &value) in csr.indices[span.clone()].iter().zip(&csr.data[span]) + { + let lanes = &tan[vec + col * K..][..K]; + for (a, &v) in acc.iter_mut().zip(lanes) { + *a += value * v; + } + } + tan[dst + row * K..][..K].copy_from_slice(&acc); + } + }, + Operand::Primal(vec) => { + for row in 0..csr.shape.rows { + let span = csr.indptr[row]..csr.indptr[row + 1]; + let mut sum = 0.0; + for (&col, &value) in csr.indices[span.clone()].iter().zip(&csr.data[span]) + { + sum += value * primal[vec + col]; + } + tan[dst + row * K..][..K].fill(sum); + } + }, + } + }, + + other => unreachable!("instruction {other:?} is not batchable; check is_batchable first"), + } +} + +/// Tangent-region base index for a slot that must live above the primal pool. +#[inline] +fn tangent_base(offset: usize, primal_len: usize) -> usize { + debug_assert!( + offset >= primal_len, + "tangent-section writes target the tangent pool" + ); + (offset - primal_len) * K +} + +/// Dispatch a binary op to a monomorphised [`binary`] loop, mirroring +/// `eval::eval_binary_op` so both paths share one definition per op. +#[allow(clippy::too_many_arguments)] +fn binary_op( + op: BinaryOp, + a: Operand, + b: Operand, + dst: usize, + len: usize, + kind: BroadcastKind, + primal: &[f64], + tan: &mut [f64], +) { + macro_rules! apply { + ($f:expr) => { + binary::($f, a, b, dst, len, kind, primal, tan) + }; + } + match op { + BinaryOp::Add => apply!(|x, y| x + y), + BinaryOp::Sub => apply!(|x, y| x - y), + BinaryOp::Mul => apply!(|x, y| x * y), + BinaryOp::Div => apply!(|x, y| x / y), + BinaryOp::Pow => apply!(f64::powf), + BinaryOp::Minimum => apply!(f64::min), + BinaryOp::Maximum => apply!(f64::max), + BinaryOp::Modulo => apply!(|x, y| x % y), + BinaryOp::Hypot => apply!(f64::hypot), + BinaryOp::EqualHeaviside => apply!(|x, y| if x <= y { 1.0 } else { 0.0 }), + BinaryOp::NotEqualHeaviside => apply!(|x, y| if x < y { 1.0 } else { 0.0 }), + BinaryOp::Equality => { + apply!(|x: f64, y: f64| if (x - y).abs() < EQUALITY_EPS { + 1.0 + } else { + 0.0 + }); + }, + } +} + +/// Dispatch a unary op to a monomorphised [`unary_apply_lanes`] loop, mirroring +/// `eval::eval_unary_op`. +fn unary( + op: UnaryOp, + src: Operand, + dst: usize, + len: usize, + primal: &[f64], + tan: &mut [f64], +) { + macro_rules! apply { + ($f:expr) => { + unary_apply_lanes::($f, src, dst, len, primal, tan) + }; + } + match op { + UnaryOp::Neg => apply!(|x: f64| -x), + UnaryOp::Abs => apply!(f64::abs), + UnaryOp::Sqrt => apply!(f64::sqrt), + UnaryOp::Exp => apply!(f64::exp), + UnaryOp::Log => apply!(f64::ln), + UnaryOp::Sin => apply!(f64::sin), + UnaryOp::Cos => apply!(f64::cos), + UnaryOp::Tanh => apply!(f64::tanh), + UnaryOp::Sinh => apply!(f64::sinh), + UnaryOp::Cosh => apply!(f64::cosh), + UnaryOp::Arcsinh => apply!(f64::asinh), + UnaryOp::Arctan => apply!(f64::atan), + UnaryOp::Erf => apply!(erf_approx), + UnaryOp::Sign => apply!(sign), + UnaryOp::Floor => apply!(f64::floor), + UnaryOp::Ceiling => apply!(f64::ceil), + } +} + +fn unary_apply_lanes f64>( + f: F, + src: Operand, + dst: usize, + len: usize, + primal: &[f64], + tan: &mut [f64], +) { + match src { + Operand::Tangent(src) => { + for e in 0..len { + let mut out = [0.0_f64; K]; + for (o, &x) in out.iter_mut().zip(&tan[src + e * K..][..K]) { + *o = f(x); + } + tan[dst + e * K..][..K].copy_from_slice(&out); + } + }, + Operand::Primal(src) => { + for e in 0..len { + tan[dst + e * K..][..K].fill(f(primal[src + e])); + } + }, + } +} + +/// Element strides `(a, b)` for a broadcast kind: a broadcast operand holds +/// one element and never advances. +#[inline] +const fn broadcast_strides(kind: BroadcastKind) -> (usize, usize) { + match kind { + BroadcastKind::ScalarScalar => (0, 0), + BroadcastKind::ScalarVector => (0, 1), + BroadcastKind::VectorScalar => (1, 0), + BroadcastKind::VectorVector => (1, 1), + } +} + +/// Every lane loop lands each element in a stack array before storing it. +/// Operand and destination slots are disjoint, but they share one `&mut [f64]`, +/// so reading and writing it in the same expression leaves the compiler unable +/// to rule out aliasing and it emits a scalar, ordered loop. +#[allow(clippy::too_many_arguments)] +fn binary f64>( + f: F, + a: Operand, + b: Operand, + dst: usize, + len: usize, + kind: BroadcastKind, + primal: &[f64], + tan: &mut [f64], +) { + let (sa, sb) = broadcast_strides(kind); + match (a, b) { + (Operand::Tangent(a), Operand::Tangent(b)) => { + for e in 0..len { + let (x, y) = (&tan[a + e * sa * K..][..K], &tan[b + e * sb * K..][..K]); + let mut out = [0.0_f64; K]; + for (o, (&x, &y)) in out.iter_mut().zip(x.iter().zip(y)) { + *o = f(x, y); + } + tan[dst + e * K..][..K].copy_from_slice(&out); + } + }, + (Operand::Primal(a), Operand::Tangent(b)) => { + for e in 0..len { + let x = primal[a + e * sa]; + let mut out = [0.0_f64; K]; + for (o, &y) in out.iter_mut().zip(&tan[b + e * sb * K..][..K]) { + *o = f(x, y); + } + tan[dst + e * K..][..K].copy_from_slice(&out); + } + }, + (Operand::Tangent(a), Operand::Primal(b)) => { + for e in 0..len { + let y = primal[b + e * sb]; + let mut out = [0.0_f64; K]; + for (o, &x) in out.iter_mut().zip(&tan[a + e * sa * K..][..K]) { + *o = f(x, y); + } + tan[dst + e * K..][..K].copy_from_slice(&out); + } + }, + (Operand::Primal(a), Operand::Primal(b)) => { + for e in 0..len { + let value = f(primal[a + e * sa], primal[b + e * sb]); + tan[dst + e * K..][..K].fill(value); + } + }, + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs b/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs new file mode 100644 index 0000000000..ca505344bb --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs @@ -0,0 +1,1299 @@ +//! Shape-aware zero propagation for dead branch elimination. +//! +//! Uses a three-valued lattice to prevent unsafe folds: +//! - `0 / Unknown` must NOT fold to 0 (divisor could be 0 → NaN) +//! - `log(AllZero)` must NOT fold to 0 (log(0) = -∞) + +use crate::arena::NodeMap; +use crate::arena::{Arena, NodeId}; +use crate::node::{ArrayData, Node}; + +/// Zero-status lattice for safe constant folding. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ZeroStatus { + /// Provably all zeros (from `Scalar(0.0)`, `ZeroVector`, or `Array{[0,0,...]}`) + AllZero, + /// Provably no zeros (from non-zero constants like `Scalar(5.0)`) + DefinitelyNonZero, + /// Could be anything at runtime (state vectors, parameters, etc.) + Unknown, +} + +/// Shape and zero-status for a node. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShapeInfo { + /// Length: 1 for scalar, n for vector + pub len: usize, + /// Zero status from the lattice + pub zero_status: ZeroStatus, +} + +impl ShapeInfo { + /// A node of `len` elements with an explicit lattice status. + pub const fn new(len: usize, zero_status: ZeroStatus) -> Self { + Self { len, zero_status } + } + + /// Every element is provably zero, which is what licenses a fold. + pub const fn all_zero(len: usize) -> Self { + Self { + len, + zero_status: ZeroStatus::AllZero, + } + } + + /// No element can be zero, which is what makes a divisor safe to fold under. + pub const fn definitely_nonzero(len: usize) -> Self { + Self { + len, + zero_status: ZeroStatus::DefinitelyNonZero, + } + } + + /// Nothing is known, the conservative default that blocks folding. + pub const fn unknown(len: usize) -> Self { + Self { + len, + zero_status: ZeroStatus::Unknown, + } + } + + /// Whether the value is provably all zeros. + pub const fn is_all_zero(&self) -> bool { + matches!(self.zero_status, ZeroStatus::AllZero) + } + + /// Whether the value is provably free of zeros. + pub const fn is_definitely_nonzero(&self) -> bool { + matches!(self.zero_status, ZeroStatus::DefinitelyNonZero) + } +} + +/// Infer shape for a single node given shapes of its children. +fn infer_shape(arena: &Arena, id: NodeId, shapes: &[ShapeInfo]) -> ShapeInfo { + match arena.get(id) { + // Scalar literals + Node::Scalar(v) => { + if *v == 0.0 { + ShapeInfo::all_zero(1) + } else { + ShapeInfo::definitely_nonzero(1) + } + }, + + // Zero vector (first-class) + Node::ZeroVector { len } => ShapeInfo::all_zero(*len), + + // Arrays - check if all zeros + Node::Array(arr) => { + let len = arr.shape.rows * arr.shape.cols; + if arr.data.iter().all(|&v| v == 0.0) { + ShapeInfo::all_zero(len) + } else if arr.data.iter().all(|&v| v != 0.0) { + ShapeInfo::definitely_nonzero(len) + } else { + ShapeInfo::unknown(len) + } + }, + + // State vectors - Unknown (could be anything at runtime) + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } => ShapeInfo::unknown(end - start), + + // Runtime values (time, parameters); InputParameter keeps its packed width + Node::InputParameter { width, .. } => ShapeInfo::unknown(*width), + Node::TangentParameter { .. } | Node::Time => ShapeInfo::unknown(1), + + // Sparse matrix (matrix, Unknown) + Node::SparseMatrix(csr) => { + let len = csr.shape.rows * csr.shape.cols; + ShapeInfo::unknown(len) + }, + + // Binary operations + Node::Add(a, b) | Node::Sub(a, b) => { + let sa = &shapes[a.index()]; + let sb = &shapes[b.index()]; + let len = sa.len.max(sb.len); + let status = if sa.is_all_zero() && sb.is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(len, status) + }, + + Node::Mul(a, b) => { + let sa = &shapes[a.index()]; + let sb = &shapes[b.index()]; + let len = sa.len.max(sb.len); + let status = if sa.is_all_zero() || sb.is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(len, status) + }, + + Node::Div(a, b) => { + let sa = &shapes[a.index()]; + let sb = &shapes[b.index()]; + let len = sa.len.max(sb.len); + // CRITICAL: Only fold 0/x when x is DefinitelyNonZero + let status = if sa.is_all_zero() && sb.is_definitely_nonzero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(len, status) + }, + + Node::Pow(base, exp) => { + let sb = &shapes[base.index()]; + let se = &shapes[exp.index()]; + ShapeInfo::unknown(sb.len.max(se.len)) + }, + + Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Hypot(a, b) + | Node::Modulo(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => { + let sa = &shapes[a.index()]; + let sb = &shapes[b.index()]; + let len = sa.len.max(sb.len); + ShapeInfo::unknown(len) + }, + + Node::MatMul(a, b) => { + let sb = &shapes[b.index()]; + let result_len = match arena.get(*a) { + Node::SparseMatrix(csr) => csr.shape.rows, + Node::Array(arr) => arr.shape.rows, + _ => sb.len, + }; + let status = if sb.is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(result_len, status) + }, + + // Unary operations - preserve ZeroStatus for Neg/Abs + Node::Neg(c) | Node::Abs(c) => { + let sc = &shapes[c.index()]; + ShapeInfo::new(sc.len, sc.zero_status) + }, + + // Nonzero by range: exp(x) > 0 and cosh(x) >= 1 for every x. + Node::Exp(c) | Node::Cosh(c) => { + let sc = &shapes[c.index()]; + ShapeInfo::definitely_nonzero(sc.len) + }, + + // log(0) = -∞, NOT 0 + Node::Log(c) => { + let sc = &shapes[c.index()]; + ShapeInfo::unknown(sc.len) + }, + + // f(0) = 0 for these functions + Node::Sqrt(c) + | Node::Sin(c) + | Node::Tanh(c) + | Node::Sinh(c) + | Node::Arcsinh(c) + | Node::Arctan(c) + | Node::Erf(c) + | Node::Sign(c) + | Node::Floor(c) + | Node::Ceiling(c) => { + let sc = &shapes[c.index()]; + let status = if sc.is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(sc.len, status) + }, + + // cos(0) = 1, so NOT zero + Node::Cos(c) => { + let sc = &shapes[c.index()]; + if sc.is_all_zero() { + ShapeInfo::definitely_nonzero(sc.len) + } else { + ShapeInfo::unknown(sc.len) + } + }, + + Node::MaxReduce(c) | Node::MinReduce(c) => { + let sc = &shapes[c.index()]; + let status = if sc.is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(1, status) + }, + + Node::ReduceArgSelect { basis, .. } => { + let status = if shapes[basis.index()].is_all_zero() { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(1, status) + }, + + // Structural nodes + Node::Index { child, start, end } => { + let sc = &shapes[child.index()]; + ShapeInfo::new(end - start, sc.zero_status) + }, + + Node::Concat(children) => { + let total_len: usize = children.iter().map(|c| shapes[c.index()].len).sum(); + let all_zero = children.iter().all(|c| shapes[c.index()].is_all_zero()); + let status = if all_zero { + ZeroStatus::AllZero + } else { + ZeroStatus::Unknown + }; + ShapeInfo::new(total_len, status) + }, + + // Interpolation - Unknown + Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => { + let sc = &shapes[child.index()]; + ShapeInfo::unknown(sc.len) + }, + + // N-D interpolation - Unknown, element-wise over children + Node::InterpolantNd { children, .. } | Node::InterpolantNdPartial { children, .. } => { + let len = children + .iter() + .map(|c| shapes[c.index()].len) + .max() + .unwrap_or(1); + ShapeInfo::unknown(len) + }, + + // Conditional - conservatively Unknown; length mirrors `ir::infer_sizes` + Node::Conditional { branches, .. } => { + let len = branches + .iter() + .map(|b| shapes[b.index()].len) + .max() + .unwrap_or(1); + ShapeInfo::unknown(len) + }, + } +} + +/// Analyze shapes for all nodes reachable from root. +pub fn analyze_shapes(arena: &Arena, root: NodeId) -> Vec { + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let order = arena.topological_order(root); + + for id in order { + let info = infer_shape(arena, id, &shapes); + shapes[id.index()] = info; + } + + shapes +} + +/// Create a zero node of the appropriate size. +fn make_zero(arena: &mut Arena, len: usize) -> NodeId { + if len == 1 { + arena.alloc(Node::Scalar(0.0)) + } else { + arena.alloc(Node::ZeroVector { len }) + } +} + +/// Propagate zero information and eliminate dead branches. +/// +/// Value-preserving, not sign-of-zero-preserving, matching the relaxed +/// guarantee documented on +/// [`SimplifyMode::Conservative`](crate::simplify::SimplifyMode::Conservative). +/// The folds `0 + b -> b`, `a + 0 -> a`, `a - 0 -> a` and `0 - b -> -b` are +/// guarded on an operand this pass *proved* all-zero, not on a literal `+0.0`, +/// so they reach computed zeros as well and a folded zero result may carry the +/// opposite sign from an unfolded one. +/// +/// `a - 0 -> a` and `0 - b -> -b` therefore go further than `Conservative` +/// `simplify`, which restricts them to the zero-sign that is exact. This pass has +/// no mode switch and always runs in `simplify_pipeline`, so the relaxed +/// guarantee above is what licenses the difference. +pub fn zero_propagate(arena: &Arena, root: NodeId) -> (Arena, NodeId) { + let shapes = analyze_shapes(arena, root); + + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + + let order = arena.topological_order(root); + + for old_id in order { + let shape = &shapes[old_id.index()]; + + // If this node is provably AllZero, replace with zero literal + if shape.is_all_zero() { + let new_id = make_zero(&mut new_arena, shape.len); + old_to_new.insert(old_id, new_id); + continue; + } + + // Otherwise, rebuild the node with potential simplifications + let new_id = match arena.get(old_id) { + // Leaf nodes - copy directly + Node::Scalar(v) => new_arena.alloc(Node::Scalar(*v)), + Node::ZeroVector { len } => new_arena.alloc(Node::ZeroVector { len: *len }), + Node::Array(arr) => new_arena.alloc(Node::Array(Box::new(ArrayData { + data: arr.data.clone(), + shape: arr.shape, + }))), + Node::SparseMatrix(csr) => new_arena.alloc(Node::SparseMatrix(csr.clone())), + Node::StateVector { start, end } => new_arena.alloc(Node::StateVector { + start: *start, + end: *end, + }), + Node::StateVectorDot { start, end } => new_arena.alloc(Node::StateVectorDot { + start: *start, + end: *end, + }), + Node::TangentStateVector { start, end } => new_arena.alloc(Node::TangentStateVector { + start: *start, + end: *end, + }), + Node::InputParameter { + name, + index, + offset, + width, + } => new_arena.alloc(Node::InputParameter { + name: name.clone(), + index: *index, + offset: *offset, + width: *width, + }), + Node::TangentParameter { index } => { + new_arena.alloc(Node::TangentParameter { index: *index }) + }, + Node::Time => new_arena.alloc(Node::Time), + + // Binary ops with zero elimination + Node::Add(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + let same_shape = shapes[a.index()].len == shapes[b.index()].len; + if same_shape && shapes[a.index()].is_all_zero() { + nb // 0 + b = b + } else if same_shape && shapes[b.index()].is_all_zero() { + na // a + 0 = a + } else { + new_arena.alloc(Node::Add(na, nb)) + } + }, + + Node::Sub(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + let same_shape = shapes[a.index()].len == shapes[b.index()].len; + if same_shape && shapes[b.index()].is_all_zero() { + na // a - 0 = a + } else if same_shape && shapes[a.index()].is_all_zero() { + new_arena.alloc(Node::Neg(nb)) // 0 - b = -b + } else { + new_arena.alloc(Node::Sub(na, nb)) + } + }, + + Node::Mul(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Mul(na, nb)) + }, + + Node::Div(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Div(na, nb)) + }, + + Node::Pow(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Pow(na, nb)) + }, + + Node::MatMul(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::MatMul(na, nb)) + }, + + Node::Minimum(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Minimum(na, nb)) + }, + + Node::Maximum(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Maximum(na, nb)) + }, + + Node::Modulo(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Modulo(na, nb)) + }, + + Node::Hypot(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Hypot(na, nb)) + }, + + Node::EqualHeaviside(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::EqualHeaviside(na, nb)) + }, + + Node::NotEqualHeaviside(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::NotEqualHeaviside(na, nb)) + }, + + Node::Equality(a, b) => { + let na = old_to_new + .get(*a) + .copied() + .expect("child must be processed"); + let nb = old_to_new + .get(*b) + .copied() + .expect("child must be processed"); + new_arena.alloc(Node::Equality(na, nb)) + }, + + // Unary ops + Node::Neg(c) => new_arena.alloc(Node::Neg( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Abs(c) => new_arena.alloc(Node::Abs( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Sqrt(c) => new_arena.alloc(Node::Sqrt( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Exp(c) => new_arena.alloc(Node::Exp( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Log(c) => new_arena.alloc(Node::Log( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Sin(c) => new_arena.alloc(Node::Sin( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Cos(c) => new_arena.alloc(Node::Cos( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Tanh(c) => new_arena.alloc(Node::Tanh( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Sinh(c) => new_arena.alloc(Node::Sinh( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Cosh(c) => new_arena.alloc(Node::Cosh( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Arcsinh(c) => new_arena.alloc(Node::Arcsinh( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Arctan(c) => new_arena.alloc(Node::Arctan( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Erf(c) => new_arena.alloc(Node::Erf( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Sign(c) => new_arena.alloc(Node::Sign( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Floor(c) => new_arena.alloc(Node::Floor( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::Ceiling(c) => new_arena.alloc(Node::Ceiling( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::MaxReduce(c) => new_arena.alloc(Node::MaxReduce( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::MinReduce(c) => new_arena.alloc(Node::MinReduce( + old_to_new + .get(*c) + .copied() + .expect("child must be processed"), + )), + Node::ReduceArgSelect { + basis, + picker, + is_max, + } => new_arena.alloc(Node::ReduceArgSelect { + basis: old_to_new + .get(*basis) + .copied() + .expect("child must be processed"), + picker: old_to_new + .get(*picker) + .copied() + .expect("child must be processed"), + is_max: *is_max, + }), + + // Structural nodes + Node::Index { child, start, end } => new_arena.alloc(Node::Index { + child: old_to_new + .get(*child) + .copied() + .expect("child must be processed"), + start: *start, + end: *end, + }), + + Node::Concat(children) => { + let new_children: Vec = children + .iter() + .map(|c| { + old_to_new + .get(*c) + .copied() + .expect("child must be processed") + }) + .collect(); + new_arena.alloc(Node::Concat(new_children)) + }, + + // Interpolation nodes + Node::Interpolant1DLinear { data, child } => { + new_arena.alloc(Node::Interpolant1DLinear { + data: data.clone(), + child: old_to_new + .get(*child) + .copied() + .expect("child must be processed"), + }) + }, + + Node::Interpolant1DLinearDeriv { + slopes, + x_data, + child, + } => new_arena.alloc(Node::Interpolant1DLinearDeriv { + slopes: slopes.clone(), + x_data: x_data.clone(), + child: old_to_new + .get(*child) + .copied() + .expect("child must be processed"), + }), + Node::Interpolant1DCubic { data, child } => new_arena.alloc(Node::Interpolant1DCubic { + data: data.clone(), + child: old_to_new + .get(*child) + .copied() + .expect("child must be processed"), + }), + Node::Interpolant1DCubicDeriv { data, child } => { + new_arena.alloc(Node::Interpolant1DCubicDeriv { + data: data.clone(), + child: old_to_new + .get(*child) + .copied() + .expect("child must be processed"), + }) + }, + Node::InterpolantNd { data, children } => { + let new_children: Vec = children + .iter() + .map(|c| { + old_to_new + .get(*c) + .copied() + .expect("child must be processed") + }) + .collect(); + new_arena.alloc(Node::InterpolantNd { + data: data.clone(), + children: new_children, + }) + }, + Node::InterpolantNdPartial { + data, + children, + axis, + } => { + let new_children: Vec = children + .iter() + .map(|c| { + old_to_new + .get(*c) + .copied() + .expect("child must be processed") + }) + .collect(); + new_arena.alloc(Node::InterpolantNdPartial { + data: data.clone(), + children: new_children, + axis: *axis, + }) + }, + + Node::Conditional { selector, branches } => { + let new_branches: Vec = branches + .iter() + .map(|b| { + old_to_new + .get(*b) + .copied() + .expect("child must be processed") + }) + .collect(); + new_arena.alloc(Node::Conditional { + selector: old_to_new + .get(*selector) + .copied() + .expect("child must be processed"), + branches: new_branches, + }) + }, + }; + + old_to_new.insert(old_id, new_id); + } + + let new_root = old_to_new + .get(root) + .copied() + .expect("root must be processed"); + (new_arena, new_root) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Shape; + + #[test] + fn test_zero_status_lattice() { + let zero = ShapeInfo::all_zero(10); + assert!(zero.is_all_zero()); + assert!(!zero.is_definitely_nonzero()); + + let nonzero = ShapeInfo::definitely_nonzero(1); + assert!(!nonzero.is_all_zero()); + assert!(nonzero.is_definitely_nonzero()); + + let unknown = ShapeInfo::unknown(5); + assert!(!unknown.is_all_zero()); + assert!(!unknown.is_definitely_nonzero()); + } + + #[test] + fn test_infer_shape_scalar_zero() { + let mut arena = Arena::new(); + let scalar = arena.alloc(Node::Scalar(0.0)); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, scalar, &shapes); + assert_eq!(info.len, 1); + assert!(info.is_all_zero()); + } + + #[test] + fn test_infer_shape_scalar_nonzero() { + let mut arena = Arena::new(); + let scalar = arena.alloc(Node::Scalar(5.0)); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, scalar, &shapes); + assert_eq!(info.len, 1); + assert!(info.is_definitely_nonzero()); + } + + #[test] + fn test_infer_shape_state_vector() { + let mut arena = Arena::new(); + let sv = arena.alloc(Node::StateVector { start: 0, end: 10 }); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, sv, &shapes); + assert_eq!(info.len, 10); + assert!(!info.is_all_zero()); + assert!(!info.is_definitely_nonzero()); + } + + #[test] + fn test_infer_shape_zero_vector() { + let mut arena = Arena::new(); + let zv = arena.alloc(Node::ZeroVector { len: 100 }); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, zv, &shapes); + assert_eq!(info.len, 100); + assert!(info.is_all_zero()); + } + + #[test] + fn test_infer_shape_array_all_zero() { + let mut arena = Arena::new(); + let arr = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![0.0, 0.0, 0.0], + shape: Shape::vector(3), + }))); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, arr, &shapes); + assert_eq!(info.len, 3); + assert!(info.is_all_zero()); + } + + #[test] + fn test_infer_shape_array_tiny_nonzero_not_zero() { + let mut arena = Arena::new(); + let arr = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0e-17, -2.0e-17], + shape: Shape::vector(2), + }))); + let shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + let info = infer_shape(&arena, arr, &shapes); + assert_eq!(info.len, 2); + assert!( + !info.is_all_zero(), + "tiny non-zero coefficients must not be erased as zeros" + ); + } + + #[test] + fn test_infer_shape_mul_zero_times_unknown() { + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let mul = arena.alloc(Node::Mul(zero, x)); + + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + shapes[zero.index()] = ShapeInfo::all_zero(1); + shapes[x.index()] = ShapeInfo::unknown(5); + + let info = infer_shape(&arena, mul, &shapes); + assert_eq!(info.len, 5); + assert!(info.is_all_zero(), "AllZero * Unknown should be AllZero"); + } + + #[test] + fn test_infer_shape_div_zero_by_nonzero() { + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let five = arena.alloc(Node::Scalar(5.0)); + let div = arena.alloc(Node::Div(zero, five)); + + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + shapes[zero.index()] = ShapeInfo::all_zero(1); + shapes[five.index()] = ShapeInfo::definitely_nonzero(1); + + let info = infer_shape(&arena, div, &shapes); + assert!( + info.is_all_zero(), + "AllZero / DefinitelyNonZero should be AllZero" + ); + } + + #[test] + fn test_infer_shape_div_zero_by_unknown() { + // CRITICAL: 0 / Unknown must NOT be AllZero (divisor could be 0 → NaN) + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let div = arena.alloc(Node::Div(zero, x)); + + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + shapes[zero.index()] = ShapeInfo::all_zero(1); + shapes[x.index()] = ShapeInfo::unknown(1); + + let info = infer_shape(&arena, div, &shapes); + assert!(!info.is_all_zero(), "AllZero / Unknown must NOT be AllZero"); + } + + #[test] + fn test_infer_shape_exp() { + // exp(x) > 0 for all x, so always DefinitelyNonZero + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let exp = arena.alloc(Node::Exp(zero)); + + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + shapes[zero.index()] = ShapeInfo::all_zero(1); + + let info = infer_shape(&arena, exp, &shapes); + assert!( + info.is_definitely_nonzero(), + "exp(0) = 1, which is DefinitelyNonZero" + ); + } + + #[test] + fn test_infer_shape_log_zero() { + // log(0) = -∞, NOT 0, so must be Unknown + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let log = arena.alloc(Node::Log(zero)); + + let mut shapes = vec![ + ShapeInfo { + len: 0, + zero_status: ZeroStatus::Unknown + }; + arena.len() + ]; + shapes[zero.index()] = ShapeInfo::all_zero(1); + + let info = infer_shape(&arena, log, &shapes); + assert!( + !info.is_all_zero(), + "log(AllZero) must NOT be AllZero (log(0) = -∞)" + ); + } + + #[test] + fn test_analyze_shapes_simple() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let mul = arena.alloc(Node::Mul(zero, x)); + + let shapes = analyze_shapes(&arena, mul); + + assert!(shapes[zero.index()].is_all_zero()); + assert!(!shapes[x.index()].is_all_zero()); + assert!( + shapes[mul.index()].is_all_zero(), + "AllZero * Unknown should be AllZero" + ); + } + + #[test] + fn test_analyze_shapes_nested() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let mul = arena.alloc(Node::Mul(zero, x)); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let add = arena.alloc(Node::Add(mul, y)); + + let shapes = analyze_shapes(&arena, add); + + assert!(shapes[mul.index()].is_all_zero()); + assert!(!shapes[y.index()].is_all_zero()); + assert!( + !shapes[add.index()].is_all_zero(), + "AllZero + Unknown is Unknown" + ); + } + + #[test] + fn test_zero_propagate_eliminates_mul_zero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let mul = arena.alloc(Node::Mul(zero, x)); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let add = arena.alloc(Node::Add(mul, y)); + + let (new_arena, new_root) = zero_propagate(&arena, add); + + // Result should be equivalent to just y (Add(0, y) => y) + match new_arena.get(new_root) { + Node::StateVector { start: 5, end: 10 } => {}, + Node::Add(_, _) => { + panic!("Add should have been simplified to just y"); + }, + other => panic!("Unexpected result: {other:?}"), + } + } + + #[test] + fn test_zero_propagate_preserves_nonzero() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let add = arena.alloc(Node::Add(x, y)); + + let (new_arena, new_root) = zero_propagate(&arena, add); + + // Should remain an Add + match new_arena.get(new_root) { + Node::Add(_, _) => {}, + _ => panic!("Expected Add to be preserved"), + } + } + + #[test] + fn test_zero_propagate_preserves_tiny_nonzero_derivative_coefficients() { + use crate::eval::{CompiledExpr, TangentInputs}; + use crate::tangent::tangent_wrt_states; + + let mut arena = Arena::new(); + let coeff = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0e-17, -2.0e-17], + shape: Shape::vector(2), + }))); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let expr = arena.alloc(Node::Mul(coeff, y)); + + let tangent = tangent_wrt_states(&mut arena, expr); + let (new_arena, new_root) = zero_propagate(&arena, tangent); + let compiled = CompiledExpr::new(&new_arena, new_root); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + let tangent_inputs = TangentInputs { + dy: Some(&[1.0, 1.0]), + dp: None, + }; + + let result = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[3.0, 4.0], + &[], + &[], + &tangent_inputs, + ); + assert_eq!(result.len(), 2); + assert!((result[0] - 1.0e-17).abs() < f64::EPSILON); + assert!((result[1] - (-2.0e-17)).abs() < f64::EPSILON); + } + + #[test] + fn test_zero_propagate_matmul_zero() { + use crate::node::CsrData; + + let mut arena = Arena::new(); + let matrix = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, 2, 4], + indices: vec![0, 1, 0, 1], + data: vec![1.0, 2.0, 3.0, 4.0], + shape: Shape::matrix(2, 2), + }))); + let zero_vec = arena.alloc(Node::ZeroVector { len: 2 }); + let matmul = arena.alloc(Node::MatMul(matrix, zero_vec)); + + let (new_arena, new_root) = zero_propagate(&arena, matmul); + + // Result should be a ZeroVector (or Scalar for len=1) + match new_arena.get(new_root) { + Node::ZeroVector { len } => { + assert_eq!(*len, 2); + }, + Node::Scalar(v) if v.abs() < f64::EPSILON => { + // Also acceptable for scalar case + }, + other => panic!("Expected ZeroVector or zero Scalar, got {other:?}"), + } + } + + #[test] + fn test_analyze_shapes_dense_returns_per_node_info() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(0.0)); + let b = arena.alloc(Node::Scalar(3.0)); + let c = arena.alloc(Node::Add(a, b)); + + let shapes = analyze_shapes(&arena, c); + assert!(shapes[a.index()].is_all_zero()); + assert!(shapes[b.index()].is_definitely_nonzero()); + } + + #[test] + fn test_pow_shape_uses_broadcast_len() { + // Pow broadcasts like the other binaries (`ir::infer_sizes` uses + // max(base, exp)); a scalar^vector result has the exponent's length. + let mut arena = Arena::new(); + let base = arena.alloc(Node::Scalar(2.0)); + let exp = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let pow = arena.alloc(Node::Pow(base, exp)); + + let shapes = analyze_shapes(&arena, pow); + assert_eq!(shapes[pow.index()].len, 3); + } + + #[test] + fn test_zero_times_pow_vector_exponent_preserves_shape() { + // 0 * 2^y with y of length 3 is a zero *vector*; the materialized + // zero must keep the broadcast length. + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let base = arena.alloc(Node::Scalar(2.0)); + let exp = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let pow = arena.alloc(Node::Pow(base, exp)); + let expr = arena.alloc(Node::Mul(zero, pow)); + + let (new_arena, new_root) = zero_propagate(&arena, expr); + + assert_eq!(new_arena.get(new_root), &Node::ZeroVector { len: 3 }); + } + + /// Documents `zero_propagate`'s contract (see its doc comment): value-exact, + /// but a zero result's sign may be normalised. `-0.0 + 0.0` is `+0.0`; + /// folding the `+ 0` away yields `-0.0`. Both are zero, so this is permitted. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins value-exactness + fn zero_propagate_may_normalise_the_sign_of_zero() { + use crate::eval::CompiledExpr; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let neg_y = arena.alloc(Node::Neg(y)); + let zero = arena.alloc(Node::Scalar(0.0)); + let root = arena.alloc(Node::Add(neg_y, zero)); + + let before = CompiledExpr::new(&arena, root); + let (folded_arena, folded_root) = zero_propagate(&arena, root); + let after = CompiledExpr::new(&folded_arena, folded_root); + + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.0], &[], &[])[0]; + let b = after.eval(&mut s2, 0.0, &[0.0], &[], &[])[0]; + + assert_eq!(a, b, "values must be equal"); + assert!(a == 0.0 && b == 0.0, "both must be zero"); + // The sign is explicitly NOT guaranteed; assert only that we know + // which way it went, so a contract change fails loudly here. + assert!(a.is_sign_positive(), "unfolded -0.0 + 0.0 is +0.0"); + assert!(b.is_sign_negative(), "folded form keeps -0.0"); + } + + /// Sibling of the Add case above, for the Sub-family fold named in + /// `zero_propagate`'s doc comment: `0 - b -> -b` is sign-exact only from + /// `-0.0`, but this pass folds it for any provably-all-zero left operand. + /// `0.0 - 0.0` is `+0.0`; folding to `Neg(y)` yields `-0.0`. + #[test] + #[allow(clippy::float_cmp)] // exact equality is the point: pins value-exactness + fn zero_propagate_sub_fold_may_normalise_the_sign_of_zero() { + use crate::eval::CompiledExpr; + + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Sub(zero, y)); + + let before = CompiledExpr::new(&arena, root); + let (folded_arena, folded_root) = zero_propagate(&arena, root); + let after = CompiledExpr::new(&folded_arena, folded_root); + + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.0], &[], &[])[0]; + let b = after.eval(&mut s2, 0.0, &[0.0], &[], &[])[0]; + + assert_eq!(a, b, "values must be equal"); + assert!(a == 0.0 && b == 0.0, "both must be zero"); + assert!(a.is_sign_positive(), "0.0 - 0.0 is +0.0"); + assert!(b.is_sign_negative(), "the fold to Neg(y) keeps -0.0"); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/common/cases.rs b/packages/pybamm-rust/pybamm-core/tests/common/cases.rs new file mode 100644 index 0000000000..a8c99b3cf0 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/cases.rs @@ -0,0 +1,880 @@ +use proptest::prelude::*; +use pybamm_core::{Arena, CompiledExpr, Node, NodeId, TypedIr}; + +// Case structs + +#[derive(Clone, Debug)] +pub struct DagCase { + pub arena: Arena, + pub root: NodeId, + pub y: Vec, + pub n_states: usize, +} + +#[derive(Clone, Debug)] +// shared scaffolding: not every integration-test binary reads every field +#[allow(dead_code)] +pub struct TangentCase { + pub arena: Arena, + pub root: NodeId, + pub y: Vec, + pub seeds: Vec>, + pub n_states: usize, +} + +#[allow(dead_code)] +pub fn eval_dag(arena: &Arena, root: NodeId, t: f64, y: &[f64], inputs: &[f64]) -> Vec { + let ir = TypedIr::from_arena(arena, root); + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + expr.eval(&mut s, t, y, &[], inputs).to_vec() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +// shared scaffolding: not every integration-test binary uses every class +#[allow(dead_code)] +enum ValueClass { + AnyFinite, + PositiveFinite, + NonZeroFinite, + BoundedSmall, + SmoothSafe, + SelectorScalar, +} + +#[derive(Clone, Copy, Debug)] +enum UnarySpec { + Neg, + Abs, + Sqrt, + Exp, + Sin, + Cos, + Tanh, + Sinh, + Cosh, + Arcsinh, + Arctan, + Erf, +} + +#[derive(Clone, Copy, Debug)] +enum BinarySpec { + Add, + Sub, + Mul, + Div, + EqualHeaviside, +} + +#[derive(Clone, Debug)] +// shared scaffolding: not every integration-test binary builds every variant +#[allow(dead_code)] +enum ExprSpec { + Scalar(f64, ValueClass), + StateSlice { + start: usize, + end: usize, + class: ValueClass, + }, + Unary { + op: UnarySpec, + child: Box, + len: usize, + class: ValueClass, + }, + Binary { + op: BinarySpec, + lhs: Box, + rhs: Box, + len: usize, + class: ValueClass, + }, + Concat(Vec, usize), + Index { + child: Box, + start: usize, + end: usize, + class: ValueClass, + }, + Conditional { + selector: Box, + branches: Vec, + len: usize, + class: ValueClass, + }, + Interpolant1DLinear { + x_data: Vec, + y_data: Vec, + child: Box, + len: usize, + class: ValueClass, + }, + SparseMatMul { + indptr: Vec, + indices: Vec, + data: Vec, + nrows: usize, + ncols: usize, + rhs: Box, + class: ValueClass, + }, +} + +#[derive(Clone, Debug)] +struct TypedEntry { + spec: ExprSpec, + len: usize, + class: ValueClass, +} + +const fn can_broadcast(a_len: usize, b_len: usize) -> bool { + a_len == b_len || a_len == 1 || b_len == 1 +} + +fn lower_expr(arena: &mut Arena, spec: &ExprSpec) -> NodeId { + match spec { + ExprSpec::Scalar(v, _) => arena.alloc(Node::Scalar(*v)), + ExprSpec::StateSlice { start, end, .. } => arena.alloc(Node::StateVector { + start: *start, + end: *end, + }), + ExprSpec::Unary { op, child, .. } => { + let child = lower_expr(arena, child); + match op { + UnarySpec::Neg => arena.alloc(Node::Neg(child)), + UnarySpec::Abs => arena.alloc(Node::Abs(child)), + UnarySpec::Sqrt => arena.alloc(Node::Sqrt(child)), + UnarySpec::Exp => arena.alloc(Node::Exp(child)), + UnarySpec::Sin => arena.alloc(Node::Sin(child)), + UnarySpec::Cos => arena.alloc(Node::Cos(child)), + UnarySpec::Tanh => arena.alloc(Node::Tanh(child)), + UnarySpec::Sinh => arena.alloc(Node::Sinh(child)), + UnarySpec::Cosh => arena.alloc(Node::Cosh(child)), + UnarySpec::Arcsinh => arena.alloc(Node::Arcsinh(child)), + UnarySpec::Arctan => arena.alloc(Node::Arctan(child)), + UnarySpec::Erf => arena.alloc(Node::Erf(child)), + } + }, + ExprSpec::Binary { op, lhs, rhs, .. } => { + let lhs = lower_expr(arena, lhs); + let rhs = lower_expr(arena, rhs); + match op { + BinarySpec::Add => arena.alloc(Node::Add(lhs, rhs)), + BinarySpec::Sub => arena.alloc(Node::Sub(lhs, rhs)), + BinarySpec::Mul => arena.alloc(Node::Mul(lhs, rhs)), + BinarySpec::Div => arena.alloc(Node::Div(lhs, rhs)), + BinarySpec::EqualHeaviside => arena.alloc(Node::EqualHeaviside(lhs, rhs)), + } + }, + ExprSpec::Concat(children, _) => { + let children: Vec = children.iter().map(|c| lower_expr(arena, c)).collect(); + arena.alloc(Node::Concat(children)) + }, + ExprSpec::Index { + child, start, end, .. + } => { + let child = lower_expr(arena, child); + arena.alloc(Node::Index { + child, + start: *start, + end: *end, + }) + }, + ExprSpec::Conditional { + selector, branches, .. + } => { + let selector = lower_expr(arena, selector); + let branches: Vec = branches.iter().map(|b| lower_expr(arena, b)).collect(); + arena.alloc(Node::Conditional { selector, branches }) + }, + ExprSpec::Interpolant1DLinear { + x_data, + y_data, + child, + .. + } => { + let child = lower_expr(arena, child); + arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + pybamm_core::InterpolantData::try_new(x_data.clone(), y_data.clone()) + .expect("valid interpolant"), + ), + child, + }) + }, + ExprSpec::SparseMatMul { + indptr, + indices, + data, + nrows, + ncols, + rhs, + .. + } => { + let mat = arena.alloc(Node::SparseMatrix(Box::new( + pybamm_core::CsrData::try_new( + indptr.clone(), + indices.clone(), + data.clone(), + pybamm_core::Shape::matrix(*nrows, *ncols), + ) + .expect("valid test matrix"), + ))); + let rhs = lower_expr(arena, rhs); + arena.alloc(Node::MatMul(mat, rhs)) + }, + } +} + +#[allow(clippy::cast_possible_truncation)] +fn seed_set_from_entropy(n_states: usize, entropy: &[u8]) -> Vec> { + let mut seeds: Vec> = (0..n_states) + .map(|j| { + let mut v = vec![0.0; n_states]; + v[j] = 1.0; + v + }) + .collect(); + + seeds.push(vec![1.0; n_states]); + seeds.push( + (0..n_states) + .map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }) + .collect(), + ); + + let mut dense = vec![0.0; n_states]; + for (i, slot) in dense.iter_mut().enumerate() { + let b = entropy + .get(i) + .copied() + .unwrap_or_else(|| (i as u8).wrapping_mul(17)); + *slot = (f64::from(b) / 127.5) - 1.0; + } + let max_abs = dense + .iter() + .fold(0.0_f64, |acc, v| acc.max(v.abs())) + .max(1e-12); + for slot in &mut dense { + *slot /= max_abs; + } + seeds.push(dense); + seeds +} + +const fn unary_input_class(op: UnarySpec, mode: GenMode) -> ValueClass { + match op { + UnarySpec::Exp => ValueClass::BoundedSmall, + // Sinh/Cosh grow like Exp, so unbounded inputs push derivatives past the + // FD oracle's validity: sinh(sinh(x)) sweeps the probe radians of a cos. + UnarySpec::Sinh | UnarySpec::Cosh => match mode { + GenMode::Smooth => ValueClass::BoundedSmall, + GenMode::Full | GenMode::TangentSafe => ValueClass::AnyFinite, + }, + UnarySpec::Sqrt => ValueClass::PositiveFinite, + UnarySpec::Arcsinh + | UnarySpec::Arctan + | UnarySpec::Erf + | UnarySpec::Sin + | UnarySpec::Cos + | UnarySpec::Tanh + | UnarySpec::Neg + | UnarySpec::Abs => ValueClass::AnyFinite, + } +} + +const fn unary_output_class(op: UnarySpec, child_class: ValueClass) -> ValueClass { + match op { + UnarySpec::Exp | UnarySpec::Sqrt | UnarySpec::Abs => ValueClass::PositiveFinite, + UnarySpec::Sin | UnarySpec::Cos | UnarySpec::Tanh | UnarySpec::Erf | UnarySpec::Arctan => { + ValueClass::BoundedSmall + }, + UnarySpec::Neg => match child_class { + ValueClass::BoundedSmall => ValueClass::BoundedSmall, + ValueClass::NonZeroFinite => ValueClass::NonZeroFinite, + _ => ValueClass::AnyFinite, + }, + UnarySpec::Sinh | UnarySpec::Cosh | UnarySpec::Arcsinh => ValueClass::AnyFinite, + } +} + +const fn binary_output_class(op: BinarySpec, a: ValueClass, b: ValueClass) -> ValueClass { + match op { + BinarySpec::Add | BinarySpec::Sub => { + if matches!(a, ValueClass::BoundedSmall) && matches!(b, ValueClass::BoundedSmall) { + ValueClass::BoundedSmall + } else { + ValueClass::AnyFinite + } + }, + BinarySpec::Div => { + if matches!(a, ValueClass::PositiveFinite) && matches!(b, ValueClass::PositiveFinite) { + ValueClass::PositiveFinite + } else { + ValueClass::AnyFinite + } + }, + BinarySpec::Mul | BinarySpec::EqualHeaviside => ValueClass::AnyFinite, + } +} + +const fn class_satisfies(actual: ValueClass, required: ValueClass) -> bool { + use ValueClass::*; + match required { + AnyFinite => true, + BoundedSmall => matches!(actual, BoundedSmall | SelectorScalar), + PositiveFinite => matches!(actual, PositiveFinite), + NonZeroFinite => matches!(actual, NonZeroFinite | PositiveFinite), + SmoothSafe => matches!(actual, SmoothSafe | BoundedSmall | PositiveFinite), + SelectorScalar => matches!(actual, SelectorScalar), + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum GenMode { + Full, + TangentSafe, + Smooth, +} + +#[allow(clippy::cast_possible_truncation)] +fn build_random_spec(entropy: &[u8], n_states: usize, mode: GenMode) -> ExprSpec { + let mut pool: Vec = Vec::new(); + let state_class = if mode == GenMode::Smooth { + ValueClass::BoundedSmall + } else { + ValueClass::PositiveFinite + }; + + pool.push(TypedEntry { + spec: ExprSpec::StateSlice { + start: 0, + end: n_states, + class: state_class, + }, + len: n_states, + class: state_class, + }); + + for i in 0..n_states.min(4) { + pool.push(TypedEntry { + spec: ExprSpec::StateSlice { + start: i, + end: i + 1, + class: state_class, + }, + len: 1, + class: state_class, + }); + } + + let scalar_vals = [0.5, 1.0, 2.0, -1.0]; + for &v in &scalar_vals { + let class = if v > 0.0 { + ValueClass::PositiveFinite + } else { + ValueClass::NonZeroFinite + }; + pool.push(TypedEntry { + spec: ExprSpec::Scalar(v, class), + len: 1, + class, + }); + } + + let smooth_unary = [ + UnarySpec::Neg, + UnarySpec::Sin, + UnarySpec::Cos, + UnarySpec::Tanh, + UnarySpec::Sinh, + UnarySpec::Cosh, + UnarySpec::Arcsinh, + UnarySpec::Arctan, + UnarySpec::Erf, + ]; + let extra_unary = [UnarySpec::Abs, UnarySpec::Sqrt]; + let exp_unary = [UnarySpec::Exp]; + + let smooth_binary = [BinarySpec::Add, BinarySpec::Sub, BinarySpec::Mul]; + + let target_successful_ops = if mode == GenMode::Smooth { 12 } else { 16 }; + let max_attempts = target_successful_ops * 3; + let mut successful_ops = 0; + let mut attempts = 0; + let mut last_good_idx = 0; + let mut byte_cursor = 0; + + let next_byte = |cursor: &mut usize| -> u8 { + let b = entropy + .get(*cursor) + .copied() + .unwrap_or_else(|| (*cursor as u8).wrapping_mul(31)); + *cursor += 1; + b + }; + + while successful_ops < target_successful_ops && attempts < max_attempts { + attempts += 1; + let op_kind = next_byte(&mut byte_cursor) % 5; + + match op_kind { + 0 | 1 => { + let all_unary: Vec = if mode == GenMode::Smooth { + smooth_unary + .iter() + .chain(exp_unary.iter()) + .copied() + .collect() + } else { + smooth_unary + .iter() + .chain(exp_unary.iter()) + .chain(extra_unary.iter()) + .copied() + .collect() + }; + let op = all_unary[next_byte(&mut byte_cursor) as usize % all_unary.len()]; + let required = unary_input_class(op, mode); + + let child_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + if class_satisfies(pool[child_idx].class, required) { + let child = &pool[child_idx]; + let out_class = unary_output_class(op, child.class); + let out_len = child.len; + pool.push(TypedEntry { + spec: ExprSpec::Unary { + op, + child: Box::new(child.spec.clone()), + len: out_len, + class: out_class, + }, + len: out_len, + class: out_class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + 2 | 3 => { + let all_binary: Vec = match mode { + GenMode::Smooth => smooth_binary.to_vec(), + GenMode::TangentSafe => { + let mut v = smooth_binary.to_vec(); + v.push(BinarySpec::Div); + v + }, + GenMode::Full => { + let mut v = smooth_binary.to_vec(); + v.push(BinarySpec::Div); + v.push(BinarySpec::EqualHeaviside); + v + }, + }; + let op = all_binary[next_byte(&mut byte_cursor) as usize % all_binary.len()]; + + let a_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + let b_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + let a = &pool[a_idx]; + let b = &pool[b_idx]; + + if !can_broadcast(a.len, b.len) { + continue; + } + + let ok = match op { + BinarySpec::Div => class_satisfies(b.class, ValueClass::NonZeroFinite), + _ => true, + }; + + if ok { + let out_len = a.len.max(b.len); + let out_class = binary_output_class(op, a.class, b.class); + pool.push(TypedEntry { + spec: ExprSpec::Binary { + op, + lhs: Box::new(a.spec.clone()), + rhs: Box::new(b.spec.clone()), + len: out_len, + class: out_class, + }, + len: out_len, + class: out_class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + 4 => { + let do_index = next_byte(&mut byte_cursor) % 2 == 0; + if do_index { + let child_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + let child = &pool[child_idx]; + if child.len > 1 { + let start = next_byte(&mut byte_cursor) as usize % child.len; + let max_end = child.len; + let end = + start + 1 + (next_byte(&mut byte_cursor) as usize % (max_end - start)); + let end = end.min(max_end); + pool.push(TypedEntry { + spec: ExprSpec::Index { + child: Box::new(child.spec.clone()), + start, + end, + class: child.class, + }, + len: end - start, + class: child.class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + } else { + let n_parts = 2 + (next_byte(&mut byte_cursor) as usize % 3); + let parts: Vec = (0..n_parts) + .map(|_| { + let idx = next_byte(&mut byte_cursor) as usize % pool.len(); + pool[idx].clone() + }) + .collect(); + let total_len: usize = parts.iter().map(|p| p.len).sum(); + let specs: Vec = parts.iter().map(|p| p.spec.clone()).collect(); + pool.push(TypedEntry { + spec: ExprSpec::Concat(specs, total_len), + len: total_len, + class: ValueClass::AnyFinite, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + _ => {}, + } + } + + pool[last_good_idx].spec.clone() +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_eval_case(entropy: Vec, n_states: usize, y: Vec) -> DagCase { + let spec = build_random_spec(&entropy, n_states, GenMode::Full); + let mut arena = Arena::new(); + let root = lower_expr(&mut arena, &spec); + DagCase { + arena, + root, + y, + n_states, + } +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_smooth_tangent_case(entropy: Vec, n_states: usize, y: Vec) -> TangentCase { + let spec = build_random_spec(&entropy, n_states, GenMode::Smooth); + let mut arena = Arena::new(); + let root = lower_expr(&mut arena, &spec); + let seeds = seed_set_from_entropy(n_states, &entropy); + TangentCase { + arena, + root, + y, + seeds, + n_states, + } +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_split_eval_case(entropy: Vec, n_states: usize, y: Vec) -> TangentCase { + let spec = build_random_spec(&entropy, n_states, GenMode::TangentSafe); + let mut arena = Arena::new(); + let root = lower_expr(&mut arena, &spec); + let seeds = seed_set_from_entropy(n_states, &entropy); + TangentCase { + arena, + root, + y, + seeds, + n_states, + } +} + +// Targeted shape builders + +fn tangent_case_from_dag(case: DagCase, entropy: &[u8]) -> TangentCase { + TangentCase { + arena: case.arena, + root: case.root, + y: case.y, + seeds: seed_set_from_entropy(case.n_states, entropy), + n_states: case.n_states, + } +} + +/// Deep linear chain: y[0..n] + 1 + 1 + ... (depth additions). +#[allow(dead_code)] +pub fn deep_chain_case(n_states: usize, depth: usize) -> DagCase { + let mut arena = Arena::new(); + let mut cur = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + for _ in 0..depth { + let one = arena.alloc(Node::Scalar(1.0)); + cur = arena.alloc(Node::Add(cur, one)); + } + let y = vec![1.0; n_states]; + DagCase { + arena, + root: cur, + y, + n_states, + } +} + +/// Wide fan-out: Concat(2*y[0], 2*y[1], ..., 2*y[fanout-1]). +#[allow(dead_code)] +pub fn wide_concat_case(fanout: usize) -> DagCase { + let mut arena = Arena::new(); + let mut terms = Vec::with_capacity(fanout); + for i in 0..fanout { + let yi = arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + terms.push(arena.alloc(Node::Mul(yi, two))); + } + let root = arena.alloc(Node::Concat(terms)); + let y: Vec = (0..fanout) + .map(|i| (i as f64).mul_add(0.03, 0.05)) + .collect(); + DagCase { + arena, + root, + y, + n_states: fanout, + } +} + +#[allow(dead_code)] +pub fn broadcast_mix_case(n_states: usize) -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + let one = arena.alloc(Node::Scalar(1.0)); + let shifted = arena.alloc(Node::Add(y, one)); + let scale = arena.alloc(Node::Scalar(0.5)); + let root = arena.alloc(Node::Mul(shifted, scale)); + DagCase { + arena, + root, + y: vec![1.0; n_states], + n_states, + } +} + +#[allow(dead_code)] +pub fn index_slice_case() -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let root = arena.alloc(Node::Index { + child: y, + start: 1, + end: 3, + }); + DagCase { + arena, + root, + y: vec![0.5, 1.0, 1.5, 2.0], + n_states: 4, + } +} + +#[allow(dead_code)] +pub fn conditional_case() -> DagCase { + let mut arena = Arena::new(); + let selector = arena.alloc(Node::Scalar(1.0)); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let s2 = arena.alloc(Node::Scalar(2.0)); + let b1 = arena.alloc(Node::Add(y0, s2)); + let s3 = arena.alloc(Node::Scalar(3.0)); + let b2 = arena.alloc(Node::Mul(y1, s3)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + DagCase { + arena, + root, + y: vec![1.0, 2.0], + n_states: 2, + } +} + +#[allow(dead_code)] +pub fn interpolant_case() -> DagCase { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + pybamm_core::InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 20.0]) + .expect("valid interpolant"), + ), + child: x, + }); + DagCase { + arena, + root, + y: vec![1.5], + n_states: 1, + } +} + +#[allow(dead_code)] +pub fn sparse_matmul_case() -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + pybamm_core::CsrData::try_new( + vec![0, 2, 4], + vec![0, 1, 2, 3], + vec![1.0, -1.0, 2.0, 3.0], + pybamm_core::Shape::matrix(2, 4), + ) + .expect("valid test matrix"), + ))); + let root = arena.alloc(Node::MatMul(mat, y)); + DagCase { + arena, + root, + y: vec![0.5, 1.0, 1.5, 2.0], + n_states: 4, + } +} + +/// Intentional structural duplicate: (3*y[0]) + (3*y[0]) built twice. +#[allow(dead_code)] +pub fn duplicate_subexpr_case() -> DagCase { + let mut arena = Arena::new(); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s1 = arena.alloc(Node::Scalar(3.0)); + let m1 = arena.alloc(Node::Mul(y1, s1)); + let y2 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s2 = arena.alloc(Node::Scalar(3.0)); + let m2 = arena.alloc(Node::Mul(y2, s2)); + let root = arena.alloc(Node::Add(m1, m2)); + DagCase { + arena, + root, + y: vec![1.25], + n_states: 1, + } +} + +/// Deep smooth composition: repeatedly apply tanh(0.5*x + 0.25). +#[allow(dead_code)] +pub fn smooth_composition_case(n_states: usize, depth: usize) -> TangentCase { + let mut arena = Arena::new(); + let mut cur = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + for _ in 0..depth { + let half = arena.alloc(Node::Scalar(0.5)); + let scaled = arena.alloc(Node::Mul(cur, half)); + let bias = arena.alloc(Node::Scalar(0.25)); + let shifted = arena.alloc(Node::Add(scaled, bias)); + cur = arena.alloc(Node::Tanh(shifted)); + } + TangentCase { + arena, + root: cur, + y: vec![0.2; n_states], + seeds: seed_set_from_entropy(n_states, &[31, 32, 33, 34]), + n_states, + } +} + +#[allow(dead_code)] +pub fn targeted_eval_cases() -> Vec { + vec![ + deep_chain_case(1, 64), + deep_chain_case(8, 100), + wide_concat_case(32), + wide_concat_case(64), + broadcast_mix_case(8), + index_slice_case(), + conditional_case(), + interpolant_case(), + sparse_matmul_case(), + duplicate_subexpr_case(), + ] +} + +#[allow(dead_code)] +pub fn targeted_smooth_tangent_cases() -> Vec { + vec![ + smooth_composition_case(1, 64), + smooth_composition_case(4, 32), + tangent_case_from_dag(wide_concat_case(8), &[5, 6, 7, 8]), + tangent_case_from_dag(broadcast_mix_case(6), &[9, 10, 11, 12]), + tangent_case_from_dag(index_slice_case(), &[13, 14, 15, 16]), + ] +} + +#[allow(dead_code)] +pub fn targeted_split_eval_cases() -> Vec { + vec![ + tangent_case_from_dag(wide_concat_case(8), &[17, 18, 19, 20]), + tangent_case_from_dag(index_slice_case(), &[21, 22, 23, 24]), + tangent_case_from_dag(conditional_case(), &[25, 26, 27, 28]), + tangent_case_from_dag(sparse_matmul_case(), &[29, 30, 31, 32]), + ] +} + +// Proptest strategies + +fn arb_random_eval_case() -> impl Strategy { + (2_usize..=8).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..48), + prop::collection::vec(0.1_f64..10.0, n_states), + ) + .prop_map(move |(entropy, y)| build_random_eval_case(entropy, n_states, y)) + }) +} + +#[allow(dead_code)] +pub fn arb_eval_case() -> impl Strategy { + arb_random_eval_case() +} + +fn arb_random_smooth_tangent_case() -> impl Strategy { + (2_usize..=6).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(-1.5_f64..1.5, n_states), + ) + .prop_map(move |(entropy, y)| build_random_smooth_tangent_case(entropy, n_states, y)) + }) +} + +#[allow(dead_code)] +pub fn arb_smooth_tangent_case() -> impl Strategy { + arb_random_smooth_tangent_case() +} + +#[allow(dead_code)] +pub fn arb_split_eval_case() -> impl Strategy { + (2_usize..=6).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(0.1_f64..10.0, n_states), + ) + .prop_map(move |(entropy, y)| build_random_split_eval_case(entropy, n_states, y)) + }) +} diff --git a/packages/pybamm-rust/pybamm-core/tests/common/mod.rs b/packages/pybamm-rust/pybamm-core/tests/common/mod.rs new file mode 100644 index 0000000000..5ba0a9b804 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/mod.rs @@ -0,0 +1,2 @@ +pub mod cases; +pub mod numeric_eq; diff --git a/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs b/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs new file mode 100644 index 0000000000..37561858a6 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs @@ -0,0 +1,76 @@ +#[allow(dead_code)] +#[track_caller] +pub fn assert_bitwise_eq(lhs: &[f64], rhs: &[f64]) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + assert_eq!( + l.to_bits(), + r.to_bits(), + "bit mismatch at index {i}: {l} vs {r}" + ); + } +} + +/// Conservative-mode simplification (and the `zero_propagate` pass it feeds) +/// is value-exact but may normalise the sign of a zero result — see the doc +/// comments on `SimplifyMode::Conservative` and `zero_propagate`. Compare +/// under that contract: a real value divergence still fails, a ±0 +/// difference does not. Every other proptest must stay strictly bitwise. +#[allow(dead_code)] +#[track_caller] +pub fn assert_conservative_eq(lhs: &[f64], rhs: &[f64]) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + if *l == 0.0 && *r == 0.0 { + continue; // ±0 both permitted under the sign-of-zero carve-out + } + assert_eq!( + l.to_bits(), + r.to_bits(), + "mismatch at index {i} (beyond the sign-of-zero carve-out): {l} vs {r}" + ); + } +} + +#[allow(dead_code)] +#[track_caller] +pub fn assert_close(lhs: &[f64], rhs: &[f64], rtol: f64, atol: f64) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + assert!( + !l.is_nan() && !r.is_nan(), + "NaN mismatch at index {i}: {l} vs {r}" + ); + if !l.is_finite() || !r.is_finite() { + assert!( + l.to_bits() == r.to_bits(), + "non-finite mismatch at index {i}: {l} vs {r}" + ); + continue; + } + let tol = rtol.mul_add(l.abs().max(r.abs()), atol); + assert!( + (l - r).abs() <= tol, + "mismatch at index {i}: {l} vs {r} (diff={}, tol={tol})", + (l - r).abs() + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_dfn.bin b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_dfn.bin new file mode 100644 index 0000000000..272ad312b7 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_dfn.bin differ diff --git a/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spm.bin b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spm.bin new file mode 100644 index 0000000000..b88eafeaa5 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spm.bin differ diff --git a/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spme.bin b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spme.bin new file mode 100644 index 0000000000..60b0a68660 Binary files /dev/null and b/packages/pybamm-rust/pybamm-core/tests/fixtures/sparsity_spme.bin differ diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_ad.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_ad.proptest-regressions new file mode 100644 index 0000000000..13e471f95f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_ad.proptest-regressions @@ -0,0 +1,11 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc aaa2952dbc6d8b2241160652b63cea8cf0ac4a87e5c2fe83d339b5837d56cc63 # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 1 }, Scalar(0.5), Neg(NodeId(1)), Add(NodeId(0), NodeId(2)), StateVector { start: 0, end: 6 }, Neg(NodeId(4)), Erf(NodeId(5)), Sub(NodeId(3), NodeId(6))] }, root: NodeId(7), y: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], seeds: [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0, 1.0, -1.0], [-0.8823529411764706, -1.0, -1.0, -0.9843137254901961, -1.0, -1.0]], n_states: 6 } +cc 0ed3d137f33f0e6ae944417543b9a826be5fa07071a55f8a9edd2478a3a40816 # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 1 }, Scalar(0.5), Add(NodeId(0), NodeId(1)), Scalar(0.5), Scalar(0.5), Scalar(0.5), Scalar(0.5), Concat([NodeId(3), NodeId(4), NodeId(5), NodeId(6)]), Sin(NodeId(7)), Sub(NodeId(2), NodeId(8))] }, root: NodeId(9), y: [0.0, 0.0, 0.0], seeds: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 1.0, 1.0], [1.0, -1.0, 1.0], [0.6912751677852348, -0.731543624161074, -1.0]], n_states: 3 } +cc af379777e6ccc6173fc7b66c56bc8167b983ec5d19ec35f473ab924521bee40b # shrinks to case = TangentCase { arena: Arena { nodes: [Scalar(0.5), Scalar(0.5), Sub(NodeId(0), NodeId(1)), StateVector { start: 0, end: 2 }, Mul(NodeId(2), NodeId(3)), Sin(NodeId(4)), StateVector { start: 0, end: 2 }, StateVector { start: 0, end: 2 }, Add(NodeId(6), NodeId(7)), Mul(NodeId(5), NodeId(8))] }, root: NodeId(9), y: [0.0, 0.0], seeds: [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0], [0.22790697674418606, 1.0]], n_states: 2 } +cc 398ffdb6597bc55fc2d1048079e5f9ec60437257d9cf95f76299e86b8a0ea32b # shrinks to case = TangentCase { arena: Arena { nodes: [Scalar(1.0), Scalar(1.0), Sub(NodeId(0), NodeId(1)), Erf(NodeId(2)), StateVector { start: 0, end: 5 }, Index { child: NodeId(4), start: 0, end: 2 }, Add(NodeId(3), NodeId(5)), Arcsinh(NodeId(6)), Arctan(NodeId(7))] }, root: NodeId(8), y: [0.0, 0.6674426596538593, 0.0, 0.0, 0.0], seeds: [[1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0, 1.0], [-0.7448559670781892, -0.49794238683127573, -1.0, -1.0, -0.5967078189300412]], n_states: 5 } +cc 9c7c5d8e3054a4eebf3301a8bae95cf0acec450699c16180e820aff3007a12ef # shrinks to case = TangentCase { arena: Arena { nodes: [Scalar(2.0), StateVector { start: 0, end: 2 }, Sub(NodeId(0), NodeId(1)), Sinh(NodeId(2)), Sinh(NodeId(3)), Cos(NodeId(4)), StateVector { start: 0, end: 2 }, Neg(NodeId(6)), Sub(NodeId(5), NodeId(7)), StateVector { start: 0, end: 2 }, Neg(NodeId(9)), Sub(NodeId(8), NodeId(10))] }, root: NodeId(11), y: [0.0, -1.0741483803953313], seeds: [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0], [0.21568627450980382, -1.0]], n_states: 2 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_ad.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_ad.rs new file mode 100644 index 0000000000..81f37fc4cf --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_ad.rs @@ -0,0 +1,677 @@ +mod common; + +use common::cases::{ + TangentCase, arb_smooth_tangent_case, eval_dag, targeted_smooth_tangent_cases, +}; +use common::numeric_eq::assert_close; +use proptest::prelude::*; +use pybamm_core::{ + Arena, ArrayData, CompiledExpr, CsrData, CubicInterpolantData, InterpolantData, JacobianData, + JacobianScratch, NdInterpolantData, Node, NodeId, Shape, SimplifyMode, TangentInputs, TypedIr, + adjoint::AdjointTape, cse, dce, simplify_with_mode, tangent_wrt_states, zero_propagate, +}; + +const FD_STEP: f64 = 1e-5; +/// Central FD with step h=1e-5 has O(h^2) truncation ≈ 1e-10. +/// Deep smooth compositions (12 layers) amplify higher-order terms +/// to ~1e-4 relative error on moderate-valued derivatives. +const RTOL: f64 = 1e-4; +/// For small derivatives (~1e-3), truncation error ~1e-10 is amplified +/// by chain-rule factors to ~1.5e-6 absolute. Observed worst case across +/// 10k+ proptest trials is 1.52e-6, so 2e-6 gives margin without +/// hiding real errors. +const ATOL: f64 = 2e-6; + +// Reusable check functions (called from both proptests and targeted tests) + +/// Check that forward-mode AD agrees with central finite differences +/// for every seed direction in the case. +/// +/// Panics on mismatch (suitable for both proptest and `#[test]` contexts). +fn check_ad_matches_fd(case: &TangentCase) { + // Smooth generator contract: primal evaluation must be finite. + let primal = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + assert!( + primal.iter().all(|v| v.is_finite()), + "smooth generator produced non-finite primal output: {primal:?}" + ); + + // This check is the AD oracle, so it must not depend on the rewrite pipeline + // beyond `tangent_wrt_states` itself. + let mut ad_arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut ad_arena, case.root); + + // Test every seed direction in the case: basis + dense directions. + for seed in &case.seeds { + // AD evaluation via split-eval + let ir = TypedIr::from_arena_split_eval(&ad_arena, tangent_root); + let compiled = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; compiled.scratch_len()]; + let mut cache = compiled.eval_primal(&mut s, 0.0, &case.y, &[], &[]); + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + let ad_result = cache.eval_tangent(&tangent_inputs).to_vec(); + + // Central finite difference: (f(y + h*seed) - f(y - h*seed)) / (2h) + let y_plus: Vec = case + .y + .iter() + .zip(seed.iter()) + .map(|(yi, si)| yi + FD_STEP * si) + .collect(); + let y_minus: Vec = case + .y + .iter() + .zip(seed.iter()) + .map(|(yi, si)| yi - FD_STEP * si) + .collect(); + + let f_plus = eval_dag(&case.arena, case.root, 0.0, &y_plus, &[]); + let f_minus = eval_dag(&case.arena, case.root, 0.0, &y_minus, &[]); + + assert!( + f_plus.iter().chain(f_minus.iter()).all(|v| v.is_finite()), + "smooth generator produced non-finite FD probe values:\n\ + seed={seed:?}\n\ + f_plus={f_plus:?}\n\ + f_minus={f_minus:?}" + ); + + let fd_result: Vec = f_plus + .iter() + .zip(f_minus.iter()) + .map(|(fp, fm)| (fp - fm) / (2.0 * FD_STEP)) + .collect(); + + assert_close(&ad_result, &fd_result, RTOL, ATOL); + } +} + +/// Reverse gradient of a scalar row vs (a) the forward-JVP row assembled through +/// the crate's trusted forward-mode path, tight (1e-12/1e-14), catching any +/// AD-path divergence, and (b) central FD: loose, independent. `ctx` labels +/// the failing case. Row of `J @ e_j` is `df/dy_j`. +fn assert_reverse_matches_forward_and_fd(ctx: &str, arena: &Arena, root: NodeId, y: &[f64]) { + let n = y.len(); + + // Reverse row. + let tape = AdjointTape::new(arena, root, n); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; n]; + tape.assemble(&mut scratch, &mut bar, &mut grad, 0.5, y, &[], &[]); + + // Forward-JVP row via the trusted forward-mode path, and central FD. + let mut ad_arena = arena.clone(); + let tangent_root = tangent_wrt_states(&mut ad_arena, root); + let compiled = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&ad_arena, tangent_root)); + let mut s = vec![0.0; compiled.scratch_len()]; + let mut seed = vec![0.0; n]; + let primal = CompiledExpr::new(arena, root); + let mut ps = vec![0.0; primal.scratch_len()]; + let eps = 1e-6; + + for j in 0..n { + // Tight forward-JVP oracle: row of J @ e_j. + seed.fill(0.0); + seed[j] = 1.0; + let mut cache = compiled.eval_primal(&mut s, 0.5, y, &[], &[]); + let jvp = cache.eval_tangent(&TangentInputs { + dy: Some(&seed), + dp: None, + })[0]; + assert!( + (grad[j] - jvp).abs() <= 1e-12f64.mul_add(jvp.abs(), 1e-14), + "{ctx} col {j}: reverse {} vs forward-JVP {jvp}", + grad[j] + ); + + // Independent central-FD oracle. + let mut yp = y.to_vec(); + let mut ym = y.to_vec(); + yp[j] += eps; + ym[j] -= eps; + let fp = primal.eval(&mut ps, 0.5, &yp, &[], &[])[0]; + let fm = primal.eval(&mut ps, 0.5, &ym, &[], &[])[0]; + let fd = (fp - fm) / (2.0 * eps); + assert!( + (grad[j] - fd).abs() <= 1e-4 * (1.0 + fd.abs()), + "{ctx} col {j}: reverse {} vs fd {fd}", + grad[j] + ); + } +} + +/// The crate intentionally treats first-derivative interpolant nodes as +/// first-order terminals. Verify reverse and forward AD both return zero. +fn assert_reverse_and_forward_are_zero(ctx: &str, arena: &Arena, root: NodeId, y: &[f64]) { + let n = y.len(); + let tape = AdjointTape::new(arena, root, n); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; n]; + tape.assemble(&mut scratch, &mut bar, &mut grad, 0.5, y, &[], &[]); + + let mut ad_arena = arena.clone(); + let tangent_root = tangent_wrt_states(&mut ad_arena, root); + let compiled = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&ad_arena, tangent_root)); + let mut ad_scratch = vec![0.0; compiled.scratch_len()]; + let mut seed = vec![0.0; n]; + for j in 0..n { + seed.fill(0.0); + seed[j] = 1.0; + let result = compiled.eval_with_tangent( + &mut ad_scratch, + 0.5, + y, + &[], + &[], + &TangentInputs { + dy: Some(&seed), + dp: None, + }, + ); + assert!( + grad[j].abs() <= f64::EPSILON, + "{ctx} reverse col {j}: {}", + grad[j] + ); + assert!( + result[0].abs() <= f64::EPSILON, + "{ctx} forward col {j}: {}", + result[0] + ); + } +} + +/// Check that the optimized derivative DAG (after the production pipeline) +/// agrees with the unoptimized derivative DAG for every seed direction. +/// +/// Production pipeline: `tangent_wrt_states` -> simplify(Aggressive) +/// -> `zero_propagate` -> cse -> dce +fn check_optimized_matches_unoptimized(case: &TangentCase) { + // Build unoptimized tangent + let mut unopt_arena = case.arena.clone(); + let unopt_root = tangent_wrt_states(&mut unopt_arena, case.root); + + // Build optimized tangent (mirrors model.rs production pipeline) + let mut opt_arena = case.arena.clone(); + let opt_root = tangent_wrt_states(&mut opt_arena, case.root); + let opt_root = simplify_with_mode(&mut opt_arena, opt_root, SimplifyMode::Aggressive); + let (opt_arena, opt_root) = zero_propagate(&opt_arena, opt_root); + let (opt_arena, opt_root) = cse(&opt_arena, opt_root); + let (opt_arena, opt_root) = dce(&opt_arena, opt_root); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + // Unoptimized eval + let ir_unopt = TypedIr::from_arena_split_eval(&unopt_arena, unopt_root); + let compiled_unopt = CompiledExpr::from_ir(ir_unopt); + let mut s_unopt = vec![0.0; compiled_unopt.scratch_len()]; + let mut cache_unopt = compiled_unopt.eval_primal(&mut s_unopt, 0.0, &case.y, &[], &[]); + let unopt_result = cache_unopt.eval_tangent(&tangent_inputs).to_vec(); + + // Optimized eval + let ir_opt = TypedIr::from_arena_split_eval(&opt_arena, opt_root); + let compiled_opt = CompiledExpr::from_ir(ir_opt); + let mut s_opt = vec![0.0; compiled_opt.scratch_len()]; + let mut cache_opt = compiled_opt.eval_primal(&mut s_opt, 0.0, &case.y, &[], &[]); + let opt_result = cache_opt.eval_tangent(&tangent_inputs).to_vec(); + + assert!( + unopt_result.iter().all(|v| v.is_finite()) && opt_result.iter().all(|v| v.is_finite()), + "smooth derivative generator produced non-finite tangent output:\n\ + seed={seed:?}\n\ + unoptimized={unopt_result:?}\n\ + optimized={opt_result:?}" + ); + + assert_eq!( + unopt_result.len(), + opt_result.len(), + "optimization pipeline changed output dimensionality:\n\ + seed={seed:?}\n\ + unoptimized ({} outputs)={unopt_result:?}\n\ + optimized ({} outputs)={opt_result:?}", + unopt_result.len(), + opt_result.len(), + ); + assert_close(&unopt_result, &opt_result, 1e-12, 1e-14); + } +} + +// Proptest properties + +proptest! { + #![proptest_config(ProptestConfig::with_cases(200))] + + /// Forward-mode AD must agree with central finite differences + /// on smooth, finite-valued expressions. + #[test] + fn ad_matches_finite_differences(case in arb_smooth_tangent_case()) { + check_ad_matches_fd(&case); + } + + /// The optimized derivative DAG (after the production pipeline) + /// must agree with the unoptimized derivative DAG. + /// + /// Production pipeline: tangent_wrt_states -> simplify(Aggressive) + /// -> zero_propagate -> cse -> dce + #[test] + fn optimized_derivative_matches_unoptimized(case in arb_smooth_tangent_case()) { + check_optimized_matches_unoptimized(&case); + } + + /// Assembled df/dy matches central finite differences column-by-column + /// on smooth random trees. + #[test] + fn jacobian_data_wrt_states_matches_fd(case in arb_smooth_tangent_case()) { + let TangentCase { arena, root, y, n_states, .. } = case; + prop_assume!(n_states > 0); + let primal = CompiledExpr::new(&arena, root); + let n_rows = primal.output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n_states); + + // The production driver, at the lane width this tape would really run. + let layout = jac.layout(); + let mut scratch = JacobianScratch::new(&jac); + let mut data = vec![0.0; layout.n_slots()]; + jac.assemble_into(&mut scratch, layout, 0.5, &y, &[], &[], &mut data); + + let mut s = vec![0.0; primal.scratch_len()]; + let eps = 1e-6; + for col in 0..jac.n_cols() { + let mut yp = y.clone(); yp[col] += eps; + let mut ym = y.clone(); ym[col] -= eps; + let fp = primal.eval(&mut s, 0.5, &yp, &[], &[]).to_vec(); + let fm = primal.eval(&mut s, 0.5, &ym, &[], &[]).to_vec(); + let (lo, hi) = (jac.csc().colptr[col], jac.csc().colptr[col + 1]); + for (&dk, &row) in data[lo..hi].iter().zip(&jac.csc().rowind[lo..hi]) { + let fd = (fp[row] - fm[row]) / (2.0 * eps); + prop_assert!((dk - fd).abs() <= 1e-4 * (1.0 + fd.abs()), + "entry ({},{}): assembled {} vs fd {}", row, col, dk, fd); + } + } + } + + /// Reverse gradient of each output row matches the forward-JVP row (tight) and + /// central FD (loose) across generated smooth expressions. + #[test] + fn reverse_row_gradient_matches_forward_and_fd(case in arb_smooth_tangent_case()) { + let TangentCase { mut arena, root, y, n_states, .. } = case; + prop_assume!(n_states > 0); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + for r in 0..n_rows { + let row = arena.alloc(Node::Index { child: root, start: r, end: r + 1 }); + assert_reverse_matches_forward_and_fd(&format!("row {r}"), &arena, row, &y); + } + } + + /// df/dp on random trees via parameter grafting: for g = f * p0 + p1, + /// dg/dp0 == f(y) and dg/dp1 == 1, exactly, no FD tolerance needed. + #[test] + fn jacobian_data_wrt_params_matches_grafted_analytic(case in arb_smooth_tangent_case()) { + let TangentCase { mut arena, root, y, n_states, .. } = case; + prop_assume!(n_states > 0); + let p0 = arena.alloc(Node::InputParameter { + name: "p0".into(), + index: 0, + offset: 0, + width: 1, + }); + let p1 = arena.alloc(Node::InputParameter { + name: "p1".into(), + index: 1, + offset: 1, + width: 1, + }); + let scaled = arena.alloc(Node::Mul(root, p0)); + let grafted = arena.alloc(Node::Add(scaled, p1)); + + let primal = CompiledExpr::new(&arena, grafted); + let n_rows = primal.output_len(); + let jac = JacobianData::new_wrt_params(&arena, grafted, n_rows, 2); + + let p = [1.7, -0.3]; + let layout = jac.layout(); + let mut scratch = JacobianScratch::new(&jac); + let mut data = vec![0.0; layout.n_slots()]; + jac.assemble_into(&mut scratch, layout, 0.5, &y, &[], &p, &mut data); + + // Reference: f(y) on the un-grafted tree. + let base = CompiledExpr::new(&arena, root); + let mut s = vec![0.0; base.scratch_len()]; + let f_val = base.eval(&mut s, 0.5, &y, &[], &[]).to_vec(); + + for col in 0..2 { + let (lo, hi) = (jac.csc().colptr[col], jac.csc().colptr[col + 1]); + for (&dk, &row) in data[lo..hi].iter().zip(&jac.csc().rowind[lo..hi]) { + let expected = if col == 0 { f_val[row] } else { 1.0 }; + prop_assert!((dk - expected).abs() <= 1e-10 * (1.0 + expected.abs()), + "dg/dp{} row {}: assembled {} vs analytic {}", col, row, dk, expected); + } + } + } +} + +// Targeted test: named deep smooth-composition and shape-stress cases + +#[test] +fn targeted_smooth_tangent_cases_pass() { + for case in targeted_smooth_tangent_cases() { + check_ad_matches_fd(&case); + check_optimized_matches_unoptimized(&case); + } +} + +/// AD must stay exact on a stiff double-exponential composition where the FD +/// oracle is invalid: f(y) = cos(sinh(sinh(2 - y))) + 2y has a diagonal Jacobian +/// entry of ~2.6e5, so an h=1e-5 central-difference probe sweeps the cosine +/// argument by ~2.6 radians. Checked against the analytic derivative instead. +#[test] +fn ad_exact_on_stiff_double_exponential() { + use pybamm_core::{Arena, Node}; + + let mut arena = Arena::new(); + let two = arena.alloc(Node::Scalar(2.0)); + let sv0 = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sub = arena.alloc(Node::Sub(two, sv0)); + let sinh1 = arena.alloc(Node::Sinh(sub)); + let sinh2 = arena.alloc(Node::Sinh(sinh1)); + let cos = arena.alloc(Node::Cos(sinh2)); + let sv1 = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let neg1 = arena.alloc(Node::Neg(sv1)); + let sub1 = arena.alloc(Node::Sub(cos, neg1)); + let sv2 = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let neg2 = arena.alloc(Node::Neg(sv2)); + let root = arena.alloc(Node::Sub(sub1, neg2)); + + let y = [0.0, -1.074_148_380_395_331_3]; + let tangent_root = tangent_wrt_states(&mut arena, root); + let ir = TypedIr::from_arena_split_eval(&arena, tangent_root); + let compiled = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; compiled.scratch_len()]; + let mut cache = compiled.eval_primal(&mut s, 0.0, &y, &[], &[]); + + // Analytic: df_i/dy_i = sin(sinh(sinh(u))) * cosh(sinh(u)) * cosh(u) + 2 + // with u = 2 - y_i; off-diagonal entries are exactly zero. + let analytic: Vec = y + .iter() + .map(|yi| { + let u = 2.0 - yi; + let w = u.sinh(); + (w.sinh().sin() * w.cosh()).mul_add(u.cosh(), 2.0) + }) + .collect(); + + for (i, &di) in analytic.iter().enumerate() { + let mut seed = vec![0.0; y.len()]; + seed[i] = 1.0; + let tangent_inputs = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let ad = cache.eval_tangent(&tangent_inputs).to_vec(); + + let mut expected = vec![0.0; y.len()]; + expected[i] = di; + assert_close(&ad, &expected, 1e-12, 1e-12); + } +} + +#[test] +fn reverse_targeted_instruction_families() { + fn idx(a: &mut Arena, y: NodeId, i: usize) -> NodeId { + a.alloc(Node::Index { + child: y, + start: i, + end: i + 1, + }) + } + fn vec3(a: &mut Arena, y: NodeId) -> NodeId { + a.alloc(Node::Index { + child: y, + start: 1, + end: 4, + }) + } + // Sum a width-`w` vector to a scalar via a constant 1×w ones row (MatMul). + fn sum_row(a: &mut Arena, v: NodeId, w: usize) -> NodeId { + let ones = a.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, w], + (0..w).collect(), + vec![1.0; w], + Shape::matrix(1, w), + ) + .unwrap(), + ))); + a.alloc(Node::MatMul(ones, v)) + } + + // y chosen so every domain is valid (log/sqrt need >0; min/max branches differ). + type CaseBuilder = fn(&mut Arena, NodeId) -> NodeId; + let y = [2.25_f64, 3.25, 4.25, 5.25]; + let cases: Vec<(&str, CaseBuilder)> = vec![ + // Binary scalar-scalar partials the smooth generator never emits. + ("sub", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Sub(x0, x1)) + }), + ("div", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Div(x0, x1)) + }), + ("pow", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Pow(x0, x1)) + }), + ("min", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Minimum(x0, x1)) + }), + ("max", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Maximum(x0, x1)) + }), + ("modulo", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Modulo(x0, x1)) + }), + ("hypot", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Hypot(x0, x1)) + }), + ("equal_heaviside", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::EqualHeaviside(x0, x1)) + }), + ("not_equal_heaviside", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::NotEqualHeaviside(x0, x1)) + }), + ("equality", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Equality(x0, x1)) + }), + // All four broadcast kinds (scalar*scalar via `mul_ss`, then s*v, v*s, v*v summed). + ("mul_ss", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + a.alloc(Node::Mul(x0, x1)) + }), + ("mul_sv", |a, y| { + let x = idx(a, y, 0); + let v = vec3(a, y); + let p = a.alloc(Node::Mul(x, v)); + sum_row(a, p, 3) + }), + ("mul_vs", |a, y| { + let v = vec3(a, y); + let x = idx(a, y, 0); + let p = a.alloc(Node::Mul(v, x)); + sum_row(a, p, 3) + }), + ("mul_vv", |a, y| { + let v = vec3(a, y); + let p = a.alloc(Node::Mul(v, v)); + sum_row(a, p, 3) + }), + // Unary derivatives not exercised elsewhere (valid at y0=2). + ("neg", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Neg(x)) + }), + ("abs", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Abs(x)) + }), + ("sqrt", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Sqrt(x)) + }), + ("log", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Log(x)) + }), + ("sign", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Sign(x)) + }), + ("floor", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Floor(x)) + }), + ("ceiling", |a, y| { + let x = idx(a, y, 0); + a.alloc(Node::Ceiling(x)) + }), + // Reductions route the scalar bar to the argmax/argmin element. + ("max_reduce", |a, y| { + let v = vec3(a, y); + a.alloc(Node::MaxReduce(v)) + }), + ("min_reduce", |a, y| { + let v = vec3(a, y); + a.alloc(Node::MinReduce(v)) + }), + // Concat then reduce (bar split across source ranges). + ("concat", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + let c = a.alloc(Node::Concat(vec![x0, x1])); + sum_row(a, c, 2) + }), + // Dense matmul: constant 1×3 row @ vec3 (transpose-scatter adjoint). + ("dense_matmul", |a, y| { + let v = vec3(a, y); + let mat = a.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![2.0, -1.0, 0.5], Shape::matrix(1, 3)).unwrap(), + ))); + a.alloc(Node::MatMul(mat, v)) + }), + // 1-D linear interpolant of a scalar (interp'(x) · bar adjoint). + ("interp_1d_linear", |a, y| { + let x = idx(a, y, 0); + let data = InterpolantData::try_new(vec![0.0, 2.5, 5.0], vec![1.0, 4.0, 9.0]).unwrap(); + a.alloc(Node::Interpolant1DLinear { + data: Box::new(data), + child: x, + }) + }), + ("interp_1d_cubic", |a, y| { + let x = idx(a, y, 0); + let data = + CubicInterpolantData::try_new(vec![0.0, 5.0], vec![[1.0, 2.0, 3.0, 4.0]]).unwrap(); + a.alloc(Node::Interpolant1DCubic { + data: Box::new(data), + child: x, + }) + }), + ("interp_nd", |a, y| { + let x0 = idx(a, y, 0); + let x1 = idx(a, y, 1); + let data = NdInterpolantData::try_new( + vec![vec![0.0, 5.0], vec![0.0, 5.0]], + vec![1.0, 4.0, 3.0, 2.0], + 2, + ) + .unwrap(); + a.alloc(Node::InterpolantNd { + data: Box::new(data), + children: vec![x0, x1], + }) + }), + // Conditional routes bar to the selected branch (selector 1.0 → branch 0). + ("conditional", |a, y| { + let selector = a.alloc(Node::Scalar(1.0)); + let b0 = idx(a, y, 0); + let b1 = idx(a, y, 1); + a.alloc(Node::Conditional { + selector, + branches: vec![b0, b1], + }) + }), + ]; + + for (label, build) in cases { + let mut arena = Arena::new(); + let yv = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let root = build(&mut arena, yv); + assert_reverse_matches_forward_and_fd(label, &arena, root, &y); + } +} + +#[test] +fn reverse_derivative_interpolants_follow_first_order_policy() { + let mut linear_arena = Arena::new(); + let linear_y = linear_arena.alloc(Node::StateVector { start: 0, end: 1 }); + let linear = linear_arena.alloc(Node::Interpolant1DLinearDeriv { + slopes: vec![2.0, 3.0].into_boxed_slice(), + x_data: vec![0.0, 2.0, 5.0].into_boxed_slice(), + child: linear_y, + }); + assert_reverse_and_forward_are_zero("linear derivative", &linear_arena, linear, &[1.0]); + + let mut cubic_arena = Arena::new(); + let cubic_y = cubic_arena.alloc(Node::StateVector { start: 0, end: 1 }); + let cubic_data = + CubicInterpolantData::try_new(vec![0.0, 5.0], vec![[1.0, 2.0, 3.0, 4.0]]).unwrap(); + let cubic = cubic_arena.alloc(Node::Interpolant1DCubicDeriv { + data: Box::new(cubic_data), + child: cubic_y, + }); + assert_reverse_and_forward_are_zero("cubic derivative", &cubic_arena, cubic, &[1.0]); + + let mut nd_arena = Arena::new(); + let nd_y0 = nd_arena.alloc(Node::StateVector { start: 0, end: 1 }); + let nd_y1 = nd_arena.alloc(Node::StateVector { start: 1, end: 2 }); + let nd_data = NdInterpolantData::try_new( + vec![vec![0.0, 5.0], vec![0.0, 5.0]], + vec![1.0, 4.0, 3.0, 2.0], + 2, + ) + .unwrap(); + let nd = nd_arena.alloc(Node::InterpolantNdPartial { + data: Box::new(nd_data), + children: vec![nd_y0, nd_y1], + axis: 0, + }); + assert_reverse_and_forward_are_zero("ND partial", &nd_arena, nd, &[1.0, 2.0]); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs new file mode 100644 index 0000000000..2d59d1ace5 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs @@ -0,0 +1,102 @@ +//! Differential test for the lane-batched observe evaluator. +//! +//! `eval_batch` reorders the interpreter's loops (K lanes per tape pass) but +//! performs the identical per-element float operations in the identical order, +//! so its output must be **bitwise identical** to `k` independent `eval` calls. +//! This is the primary correctness gate: random DAGs from the shared generator +//! plus the targeted cases that cover instructions the random generator does +//! not emit (conditional, interpolant, sparse matmul). + +mod common; + +use common::cases::{DagCase, arb_eval_case, targeted_eval_cases}; +use proptest::prelude::*; +use pybamm_core::{Arena, CompiledExpr, NodeId}; + +/// Assert `eval_batch` over `k` lanes equals `k` scalar `eval`s, bit for bit. +#[track_caller] +fn assert_batch_matches_scalar( + arena: &Arena, + root: NodeId, + n_states: usize, + k: usize, + ts: &[f64], + y_cols: &[f64], +) { + let expr = CompiledExpr::new(arena, root); + let out_len = expr.output_len(); + + // Scalar reference: one eval per lane, column-major into `scalar`. + let mut scalar = vec![0.0_f64; out_len * k]; + let mut s = vec![0.0_f64; expr.scratch_len()]; + for l in 0..k { + let y = &y_cols[l * n_states..(l + 1) * n_states]; + let res = expr.eval(&mut s, ts[l], y, &[], &[]); + scalar[l * out_len..(l + 1) * out_len].copy_from_slice(res); + } + + // Batched: one pass over all lanes. + let mut batch_scratch = vec![0.0_f64; expr.scratch_len() * k]; + let root_slice = expr + .eval_batch(&mut batch_scratch, k, ts, y_cols, &[]) + .expect("primal tape must batch-evaluate"); + + for l in 0..k { + for e in 0..out_len { + let got = root_slice[e * k + l]; + let want = scalar[l * out_len + e]; + assert_eq!( + got.to_bits(), + want.to_bits(), + "lane {l}, elem {e}: batch {got} != scalar {want} (k={k})" + ); + } + } +} + +/// `(case, k, ts[k], y_cols[n_states*k])` with `k` sampled from {1, 2, 7, 32}. +fn arb_batch_case() -> impl Strategy, Vec)> { + arb_eval_case() + .prop_flat_map(|case| (Just(case), prop::sample::select(vec![1_usize, 2, 7, 32]))) + .prop_flat_map(|(case, k)| { + let n_states = case.n_states; + ( + Just(case), + Just(k), + prop::collection::vec(-2.0_f64..5.0, k), + prop::collection::vec(0.1_f64..10.0, n_states * k), + ) + }) +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(400))] + + #[test] + fn batch_matches_scalar_random_dag((case, k, ts, y_cols) in arb_batch_case()) { + assert_batch_matches_scalar(&case.arena, case.root, case.n_states, k, &ts, &y_cols); + } +} + +/// Targeted DAGs cover instructions the random generator never emits +/// (conditional, interpolant, sparse matmul), swept across lane counts +/// including the ragged tail (`k = 1`). +#[test] +fn targeted_cases_batch_match_scalar() { + for case in targeted_eval_cases() { + let n = case.n_states; + for &k in &[1_usize, 2, 7, 32] { + // Deterministic per-lane state, varied by lane but kept near the + // case's own domain so guarded ops (sqrt/div) stay in range. + let mut y_cols = vec![0.0_f64; n * k]; + for l in 0..k { + for i in 0..n { + y_cols[l * n + i] = + case.y[i].mul_add(0.03_f64.mul_add(l as f64, 1.0), 0.01 * i as f64); + } + } + let ts: Vec = (0..k).map(|l| l as f64 * 0.1).collect(); + assert_batch_matches_scalar(&case.arena, case.root, n, k, &ts, &y_cols); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_const_entries.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_const_entries.rs new file mode 100644 index 0000000000..ca30dd494a --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_const_entries.rs @@ -0,0 +1,173 @@ +//! Soundness of the constant-entry classifier, and of assembling on it. +//! +//! Two properties, both over random graphs at random states. First: an entry +//! the classifier calls constant must equal what the tangent tape produces for +//! that column, at every state — a wrong constant does not crash, it quietly +//! degrades Newton convergence, so this is the property the whole scheme rests +//! on. Second: assembling with the split on must reproduce the unsplit +//! assembly entry for entry. +//! +//! Non-finite tape values are exempt from both: where a tape overflows, a term +//! the fold drops as exactly zero evaluates to `inf * 0.0` and poisons the +//! sweep, so the two legitimately disagree (see `const_entries`). + +mod common; + +use common::cases::{TangentCase, arb_split_eval_case, targeted_split_eval_cases}; +use proptest::prelude::*; +use pybamm_core::const_entries::classify_constant_entries; +use pybamm_core::jacobian::{JacobianData, JacobianScratch}; +use pybamm_core::{ + CompiledExpr, TangentInputs, TypedIr, detect_sparsity_per_output, simplify_pipeline, + tangent_wrt_states, +}; + +/// States to probe each case at, spread so a coefficient that merely looks +/// constant at one point does not survive. Contractive only: the generators +/// promise a finite tape at their own `y`, and scaling up walks into the +/// overflow regime where the two paths are allowed to differ. +fn probe_states(y: &[f64]) -> Vec> { + let mut states = vec![y.to_vec()]; + for factor in [0.83_f64, 0.61, 0.42, 0.25] { + states.push( + y.iter() + .enumerate() + .map(|(i, value)| value * factor.powi(1 + i32::try_from(i % 3).expect("small"))) + .collect(), + ); + } + states +} + +/// Every column of `d(root)/dy` the tape produces, one unit seed at a time. +/// +/// Takes an already-compiled tape: the compile is state-independent, and the +/// callers probe several states per case. +fn sweep_columns( + expr: &CompiledExpr, + scratch: &mut [f64], + n_states: usize, + t: f64, + y: &[f64], +) -> Vec> { + let mut cache = expr.eval_primal(scratch, t, y, &[], &[]); + (0..n_states) + .map(|col| { + let mut seed = vec![0.0; n_states]; + seed[col] = 1.0; + cache + .eval_tangent(&TangentInputs { + dy: Some(&seed), + dp: None, + }) + .to_vec() + }) + .collect() +} + +/// Returns how many entries the classifier proved, so a caller can check the +/// property is not passing vacuously. +fn check_constants_match_the_tape(case: &TangentCase) -> usize { + let n_states = case.n_states; + let n_rows = CompiledExpr::new(&case.arena, case.root).output_len(); + + let mut diff_arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut diff_arena, case.root); + let (diff_arena, tangent_root) = simplify_pipeline(diff_arena, tangent_root); + let pattern = detect_sparsity_per_output(&case.arena, case.root, n_rows, n_states); + let (varying, entries) = classify_constant_entries(&diff_arena, tangent_root, &pattern); + + assert_eq!( + varying.iter().filter(|&&v| v).count() + entries.len(), + pattern.nnz(), + "every pattern entry is either swept or known" + ); + + let rows = pattern.entry_rows(); + // The tape the fold is checked against is the one just built, compiled once. + let expr = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&diff_arena, tangent_root)); + let mut scratch = vec![0.0; expr.scratch_len()]; + + for y in probe_states(&case.y) { + let columns = sweep_columns(&expr, &mut scratch, n_states, 0.5, &y); + for &(csr_idx, value) in &entries { + let (row, col) = (rows[csr_idx], pattern.indices[csr_idx]); + let swept = columns[col][row]; + if !swept.is_finite() { + continue; + } + assert!( + value.to_bits() == swept.to_bits() || (value == 0.0 && swept == 0.0), + "entry ({row}, {col}) folded to {value}, tape gave {swept} at y={y:?}" + ); + } + } + entries.len() +} + +#[allow(clippy::float_cmp)] // exact equality is the point: pins the two paths +fn check_split_assembly_matches_unsplit(case: &TangentCase) { + let n_states = case.n_states; + let n_rows = CompiledExpr::new(&case.arena, case.root).output_len(); + let split = JacobianData::new_wrt_states(&case.arena, case.root, n_rows, n_states); + let reference = JacobianData::new_wrt_states_unsplit(&case.arena, case.root, n_rows, n_states); + + // Scratch and layout depend only on the artifact, so they are minted once + // here rather than per probe state. + let buffers = |jac: &JacobianData| { + ( + JacobianScratch::new(jac), + vec![f64::NAN; jac.layout().n_slots()], + ) + }; + let (mut split_bufs, mut reference_bufs) = (buffers(&split), buffers(&reference)); + let assemble = |jac: &JacobianData, bufs: &mut (JacobianScratch, Vec), y: &[f64]| { + let (scratch, data) = bufs; + jac.assemble_into(scratch, jac.layout(), 0.5, y, &[], &[], data); + data.clone() + }; + + for y in probe_states(&case.y) { + let (actual, expected) = ( + assemble(&split, &mut split_bufs, &y), + assemble(&reference, &mut reference_bufs, &y), + ); + for (csc_idx, (&got, &want)) in actual.iter().zip(&expected).enumerate() { + if !want.is_finite() { + continue; + } + let (row, col) = split.csc().csc_to_csr_map[csc_idx]; + assert!( + got == want, + "entry ({row}, {col}): split gave {got}, unsplit {want} at y={y:?}" + ); + } + } +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(64))] + + #[test] + fn constants_match_the_tape(case in arb_split_eval_case()) { + check_constants_match_the_tape(&case); + } + + #[test] + fn split_assembly_matches_unsplit(case in arb_split_eval_case()) { + check_split_assembly_matches_unsplit(&case); + } +} + +#[test] +fn targeted_cases_classify_soundly() { + let mut proved = 0; + for case in targeted_split_eval_cases() { + proved += check_constants_match_the_tape(&case); + check_split_assembly_matches_unsplit(&case); + } + assert!( + proved > 0, + "the targeted cases must exercise the fold at all" + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs new file mode 100644 index 0000000000..45e249874f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs @@ -0,0 +1,57 @@ +mod common; + +use common::cases::{arb_eval_case, eval_dag, targeted_eval_cases}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::{CompiledExpr, TypedIr}; + +proptest! { + #![proptest_config(ProptestConfig::with_cases(250))] + + #[test] + fn eval_is_deterministic(case in arb_eval_case()) { + let result1 = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + let result2 = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + assert_bitwise_eq(&result1, &result2); + } + + #[test] + fn independent_compilations_agree(case in arb_eval_case()) { + let ir1 = TypedIr::from_arena(&case.arena, case.root); + let ir2 = TypedIr::from_arena(&case.arena, case.root); + + let expr1 = CompiledExpr::from_ir(ir1); + let expr2 = CompiledExpr::from_ir(ir2); + let mut s1 = vec![0.0; expr1.scratch_len()]; + let mut s2 = vec![0.0; expr2.scratch_len()]; + + let result1 = expr1.eval(&mut s1, 0.0, &case.y, &[], &[]).to_vec(); + let result2 = expr2.eval(&mut s2, 0.0, &case.y, &[], &[]).to_vec(); + assert_bitwise_eq(&result1, &result2); + } + + #[test] + fn repeated_eval_on_same_compiled_expr(case in arb_eval_case()) { + let ir = TypedIr::from_arena(&case.arena, case.root); + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + + // Warm up with one eval + let _ = expr.eval(&mut s, 0.0, &case.y, &[], &[]).to_vec(); + + // Eval with shifted input, then repeat — catches buffer state leaks + let y2: Vec = case.y.iter().map(|v| v + 0.1).collect(); + let first = expr.eval(&mut s, 0.0, &y2, &[], &[]).to_vec(); + let second = expr.eval(&mut s, 0.0, &y2, &[], &[]).to_vec(); + assert_bitwise_eq(&first, &second); + } +} + +#[test] +fn targeted_eval_cases_are_stable() { + for case in targeted_eval_cases() { + let r1 = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + let r2 = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + assert_bitwise_eq(&r1, &r2); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions new file mode 100644 index 0000000000..4070f0c176 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 9f844515f49dea055f1a6f428f32fa1800f5339048d39418b6e41e46f1093b6e # shrinks to case = DagCase { arena: Arena { nodes: [Scalar(2.0), Scalar(2.0), Sub(NodeId(0), NodeId(1)), Arcsinh(NodeId(2)), Scalar(2.0), Scalar(2.0), Sub(NodeId(4), NodeId(5)), StateVector { start: 0, end: 1 }, StateVector { start: 0, end: 1 }, Sub(NodeId(7), NodeId(8)), Sub(NodeId(6), NodeId(9)), Sin(NodeId(10)), StateVector { start: 0, end: 2 }, StateVector { start: 0, end: 2 }, EqualHeaviside(NodeId(12), NodeId(13)), StateVector { start: 0, end: 2 }, Concat([NodeId(3), NodeId(11), NodeId(14), NodeId(15)])] }, root: NodeId(16), y: [0.1, 0.1], n_states: 2 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs new file mode 100644 index 0000000000..cde7e836c1 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs @@ -0,0 +1,67 @@ +mod common; + +use common::cases::{arb_eval_case, duplicate_subexpr_case, eval_dag, targeted_eval_cases}; +use common::numeric_eq::{assert_bitwise_eq, assert_conservative_eq}; +use proptest::prelude::*; +use pybamm_core::{cse, simplify, zero_propagate}; + +proptest! { + #![proptest_config(ProptestConfig::with_cases(250))] + + #[test] + fn cse_preserves_eval(case in arb_eval_case()) { + let original = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + let (cse_arena, cse_root) = cse(&case.arena, case.root); + let after = eval_dag(&cse_arena, cse_root, 0.0, &case.y, &[]); + assert_bitwise_eq(&original, &after); + } + + #[test] + fn simplify_conservative_preserves_eval(case in arb_eval_case()) { + let original = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + let mut arena_copy = case.arena.clone(); + let simplified_root = simplify(&mut arena_copy, case.root); + let after = eval_dag(&arena_copy, simplified_root, 0.0, &case.y, &[]); + assert_conservative_eq(&original, &after); + } + + #[test] + fn zero_propagate_preserves_eval(case in arb_eval_case()) { + let original = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + let (zp_arena, zp_root) = zero_propagate(&case.arena, case.root); + let after = eval_dag(&zp_arena, zp_root, 0.0, &case.y, &[]); + assert_conservative_eq(&original, &after); + } +} + +#[test] +fn rewrite_passes_preserve_eval_on_targeted_cases() { + for case in targeted_eval_cases() { + let original = eval_dag(&case.arena, case.root, 0.0, &case.y, &[]); + + let (cse_arena, cse_root) = cse(&case.arena, case.root); + let cse_after = eval_dag(&cse_arena, cse_root, 0.0, &case.y, &[]); + assert_bitwise_eq(&original, &cse_after); + + let mut arena_copy = case.arena.clone(); + let simplified_root = simplify(&mut arena_copy, case.root); + let simplified_after = eval_dag(&arena_copy, simplified_root, 0.0, &case.y, &[]); + assert_bitwise_eq(&original, &simplified_after); + + let (zp_arena, zp_root) = zero_propagate(&case.arena, case.root); + let zp_after = eval_dag(&zp_arena, zp_root, 0.0, &case.y, &[]); + assert_bitwise_eq(&original, &zp_after); + } +} + +#[test] +fn cse_reduces_known_duplicate_graph() { + let case = duplicate_subexpr_case(); + let (cse_arena, _) = cse(&case.arena, case.root); + assert!( + cse_arena.len() < case.arena.len(), + "CSE failed to reduce known duplicates: {} -> {}", + case.arena.len(), + cse_arena.len(), + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.proptest-regressions new file mode 100644 index 0000000000..338e8f11bc --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.proptest-regressions @@ -0,0 +1,9 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc aea3e5527a53cec73352753f48e8ea76ef8b6510288d81ec896eb05dd7afe112 # shrinks to case = DagCase { arena: Arena { nodes: [Scalar(-1.0), Scalar(-1.0), EqualHeaviside(NodeId(0), NodeId(1)), Scalar(-1.0), Scalar(-1.0), EqualHeaviside(NodeId(3), NodeId(4)), StateVector { start: 0, end: 2 }, StateVector { start: 0, end: 2 }, Add(NodeId(6), NodeId(7)), Concat([NodeId(5), NodeId(8)]), Sin(NodeId(9)), Scalar(-1.0), Erf(NodeId(11)), StateVector { start: 0, end: 2 }, Concat([NodeId(2), NodeId(10), NodeId(12), NodeId(13)])] }, root: NodeId(14), y: [0.1, 0.1], n_states: 2 } +cc af2d7a51de00f8639b4effade1b9009ad171c91a3fce13a8498948f7bd0bc242 # shrinks to case = DagCase { arena: Arena { nodes: [StateVector { start: 0, end: 3 }, Abs(NodeId(0)), Arcsinh(NodeId(1)), Sinh(NodeId(2)), StateVector { start: 0, end: 1 }, EqualHeaviside(NodeId(3), NodeId(4))] }, root: NodeId(5), y: [5.071448568001731, 0.1, 0.1], n_states: 3 } +cc 8ccf7eb764a00b6573137a1b029b7941d6f4979a36222872a7cf7f9d547618b1 # shrinks to case = DagCase { arena: Arena { nodes: [StateVector { start: 0, end: 6 }, StateVector { start: 0, end: 6 }, Add(NodeId(0), NodeId(1)), Sinh(NodeId(2)), Sin(NodeId(3))] }, root: NodeId(4), y: [0.1, 0.1, 0.1, 0.1, 0.1, 8.629260979374383], n_states: 6 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.rs new file mode 100644 index 0000000000..10a57efc64 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_row_blocks.rs @@ -0,0 +1,187 @@ +//! Grouping split Jacobian rows onto one shared adjoint tape must not change +//! any row's gradient. +//! +//! Two oracles: +//! +//! 1. A grouped tape against a tape holding that row alone. Seeding element `r` +//! of a shared root must recover exactly what a tape built for row `r` does, +//! which is the whole contract `ROWS_PER_TAPE` relies on. +//! 2. Both against central finite differences of the parent expression, so a +//! mistake shared by the two adjoint paths still fails. + +mod common; + +use common::cases::{DagCase, arb_eval_case, targeted_eval_cases}; +use proptest::prelude::*; +use pybamm_core::{Arena, CompiledExpr, NodeId, adjoint::AdjointTape, extract_scalar_rows}; + +/// Central-difference step, and the tolerance the FD oracle is checked at. +const FD_STEP: f64 = 1e-6; +const FD_RTOL: f64 = 1e-4; +const FD_ATOL: f64 = 1e-6; + +/// Assemble `row` of `tape` into a fresh gradient, from its own forward pass. +fn row_gradient(tape: &AdjointTape, row: usize, y: &[f64]) -> Vec { + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; tape.n_states().max(1)]; + tape.eval_forward(&mut scratch, 0.5, y, &[], &[]); + tape.assemble_row(&scratch, &mut bar, &mut grad, row); + grad +} + +/// Central finite difference of element `row` of `root` w.r.t. every state. +/// +/// `None` unless the difference is finite and has converged -- it is retaken at +/// half the step and kept only if the two agree. `sin(sinh(2y))` at `y = 8.6` +/// moves 31 radians across the stencil, so its quotient is noise of plausible +/// magnitude, which would otherwise read as a broken adjoint. +fn fd_row(arena: &Arena, root: NodeId, row: usize, y: &[f64]) -> Option> { + let expr = CompiledExpr::new(arena, root); + let mut scratch = vec![0.0; expr.scratch_len()]; + let mut slope_at = |col: usize, step: f64| { + let mut plus = y.to_vec(); + let mut minus = y.to_vec(); + plus[col] += step; + minus[col] -= step; + let f_plus = expr.eval(&mut scratch, 0.5, &plus, &[], &[])[row]; + let f_minus = expr.eval(&mut scratch, 0.5, &minus, &[], &[])[row]; + (f_plus - f_minus) / (2.0 * step) + }; + (0..y.len()) + .map(|col| { + let coarse = slope_at(col, FD_STEP); + // Halving the step cuts an O(h^2) truncation error fourfold, so the + // two agree only where the difference means something. + let fine = slope_at(col, FD_STEP / 2.0); + let converged = fine.is_finite() + && coarse.is_finite() + && (coarse - fine).abs() <= FD_ATOL + FD_RTOL * fine.abs(); + converged.then_some(fine) + }) + .collect() +} + +/// Every row of `case`, grouped as the builder likes, against all three oracles. +/// +/// Panics on mismatch, so it serves proptest and the targeted tests alike. +/// Returns how many rows the finite-difference oracle ran on, so a caller that +/// means to exercise it can assert that it did. +fn check_row_blocks(case: &DagCase, check_fd: bool) -> usize { + let mut fd_rows = 0; + let width = CompiledExpr::new(&case.arena, case.root).output_len(); + if width < 2 { + return fd_rows; // sharing a tape is only meaningful with rows to share it + } + let rows: Vec = (0..width).collect(); + let Some(block) = extract_scalar_rows(&case.arena, case.root, &rows) else { + return fd_rows; // an unindexable node declined the split, which is allowed + }; + + { + let tape = AdjointTape::new(&block.arena, block.root, case.n_states); + assert_eq!( + tape.n_rows(), + block.rows.len(), + "the tape must hold one element per row it was built for" + ); + + for (element, &row) in block.rows.iter().enumerate() { + let grouped = row_gradient(&tape, element, &case.y); + + // Oracle 1: a tape holding this row alone must agree with the group. + let Some(solo_block) = extract_scalar_rows(&case.arena, case.root, &[row]) else { + continue; + }; + let solo_tape = AdjointTape::new(&solo_block.arena, solo_block.root, case.n_states); + let solo = row_gradient(&solo_tape, 0, &case.y); + assert_close_rows(&grouped, &solo, row, "grouped vs solo"); + + // Oracle 2: an independent check that both are the real derivative. + if check_fd && let Some(fd) = fd_row(&case.arena, case.root, row, &case.y) { + assert_close_rows(&grouped, &fd, row, "adjoint vs finite difference"); + fd_rows += 1; + } + } + } + fd_rows +} + +/// Compare two gradient rows at the FD tolerance, naming the row on failure. +/// +/// Exact equality is checked first so an expression that legitimately overflows +/// to the same infinity on both sides agrees, where `(inf - inf)` would be NaN. +#[track_caller] +fn assert_close_rows(got: &[f64], want: &[f64], row: usize, what: &str) { + assert_eq!(got.len(), want.len(), "row {row}: {what} length mismatch"); + for (col, (&g, &w)) in got.iter().zip(want).enumerate() { + // Bit equality first, so a shared overflow to the same infinity agrees + // where `(inf - inf)` would be NaN. `float_cmp` wants a margin; exactness + // is the point here. + if g.to_bits() == w.to_bits() || (g.is_nan() && w.is_nan()) { + continue; + } + assert!( + (g - w).abs() <= FD_ATOL + FD_RTOL * w.abs(), + "row {row}, col {col}: {what}: {g} vs {w}" + ); + } +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(256))] + + /// The whole invariant over arbitrary DAGs. Finite differences are checked + /// too, since a shared error in both adjoint paths would satisfy oracles 1 + /// and 2 alone. + #[test] + fn grouped_rows_match_solo_rows_and_finite_differences(case in arb_eval_case()) { + check_row_blocks(&case, true); + } +} + +/// The shapes the generator reaches only by luck: a wide `Concat`, a shared +/// sub-expression across rows, and a sparse matmul block. +#[test] +fn targeted_cases_hold_the_row_block_invariants() { + for case in targeted_eval_cases() { + check_row_blocks(&case, false); + } +} + +/// A dense block over shared upstream: every row reads every lane of the same +/// expression, which is the shape the split exists for. +#[test] +fn a_dense_block_over_shared_upstream_holds() { + use pybamm_core::{CsrData, Node, Shape}; + + let (n, width) = (24, 6); + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let shared = arena.alloc(Node::Tanh(y)); + let indptr: Vec = (0..=width).map(|row| row * n).collect(); + let indices: Vec = (0..width).flat_map(|_| 0..n).collect(); + let data: Vec = (0..width * n) + .map(|k| (k as f64).mul_add(0.017, 0.3)) + .collect(); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(indptr, indices, data, Shape::matrix(width, n)).expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(matrix, shared)); + + let case = DagCase { + arena, + root, + y: (0..n) + .map(|i| (i as f64).mul_add(0.13, 0.2).sin()) + .collect(), + n_states: n, + }; + // Well-conditioned by construction, so this is what stops `fd_row`'s + // convergence gate quietly disabling oracle 2 everywhere. + assert_eq!( + check_row_blocks(&case, true), + width, + "the finite-difference oracle must run on every row of a well-conditioned block" + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions new file mode 100644 index 0000000000..6feff31ab2 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions @@ -0,0 +1,9 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 3cbce9a9d1b06125e5f5ffec4e785214e6278679d98aab41ffba0925fad13d52 # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 5 }, StateVector { start: 0, end: 5 }, Neg(NodeId(1)), Tanh(NodeId(2)), EqualHeaviside(NodeId(0), NodeId(3)), Index { child: NodeId(4), start: 2, end: 4 }] }, root: NodeId(5), y: [0.1, 0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0, 1.0], [-1.0, -1.0, -1.0, -1.0, -1.0]], n_states: 5 } +cc f4b85e8b1c0824b4a35cc97b1c38db3de0b8eaec65bbf677203beddc9aaad31b # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 4 }, Abs(NodeId(0)), StateVector { start: 0, end: 4 }, EqualHeaviside(NodeId(1), NodeId(2)), Index { child: NodeId(3), start: 2, end: 4 }] }, root: NodeId(4), y: [0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0], [-0.4893617021276595, -1.0, -0.39574468085106385, 0.5063829787234041]], n_states: 4 } +cc dc5ce2ed538db747c102ea1eed62013153aeb0a83f15896620df7515e8690fdf # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 4 }, StateVector { start: 0, end: 4 }, Add(NodeId(0), NodeId(1)), StateVector { start: 0, end: 4 }, StateVector { start: 0, end: 1 }, EqualHeaviside(NodeId(3), NodeId(4)), StateVector { start: 0, end: 1 }, Concat([NodeId(2), NodeId(5), NodeId(6)]), Sinh(NodeId(7))] }, root: NodeId(8), y: [0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0], [-0.24110671936758896, -0.9762845849802372, -0.6521739130434782, -1.0]], n_states: 4 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs new file mode 100644 index 0000000000..f087e99187 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs @@ -0,0 +1,133 @@ +mod common; + +use common::cases::{TangentCase, arb_split_eval_case, targeted_split_eval_cases}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::{ + CompiledExpr, SimplifyMode, TangentInputs, TypedIr, cse, dce, simplify_with_mode, + tangent_wrt_states, zero_propagate, +}; + +// Reusable check functions (called from both proptests and targeted tests) + +fn check_split_eval_matches_monolithic(case: &TangentCase) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + // Monolithic eval + let ir_mono = TypedIr::from_arena(&arena, tangent_root); + let mono = CompiledExpr::from_ir(ir_mono); + let mut s_mono = vec![0.0; mono.scratch_len()]; + let mono_result = mono + .eval_with_tangent(&mut s_mono, 0.0, &case.y, &[], &[], &tangent_inputs) + .to_vec(); + + // Partitioned tape + let ir_split = TypedIr::from_arena_split_eval(&arena, tangent_root); + let split = CompiledExpr::from_ir(ir_split); + let mut s_split = vec![0.0; split.scratch_len()]; + let mut cache = split.eval_primal(&mut s_split, 0.0, &case.y, &[], &[]); + let split_result = cache.eval_tangent(&tangent_inputs).to_vec(); + + assert_bitwise_eq(&mono_result, &split_result); + } +} + +fn check_repeated_tangent_stable(case: &TangentCase) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + + let ir = TypedIr::from_arena_split_eval(&arena, tangent_root); + let compiled = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; compiled.scratch_len()]; + let mut cache = compiled.eval_primal(&mut s, 0.0, &case.y, &[], &[]); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + let first = cache.eval_tangent(&tangent_inputs).to_vec(); + let second = cache.eval_tangent(&tangent_inputs).to_vec(); + assert_bitwise_eq(&first, &second); + } +} + +fn check_split_eval_after_pipeline(case: &TangentCase) { + let mut arena = case.arena.clone(); + let root = tangent_wrt_states(&mut arena, case.root); + let root = simplify_with_mode(&mut arena, root, SimplifyMode::Aggressive); + let (arena, root) = zero_propagate(&arena, root); + let (arena, root) = cse(&arena, root); + let (arena, root) = dce(&arena, root); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + // Monolithic tape + let ir_mono = TypedIr::from_arena(&arena, root); + let mono = CompiledExpr::from_ir(ir_mono); + let mut s_mono = vec![0.0; mono.scratch_len()]; + let mono_result = mono + .eval_with_tangent(&mut s_mono, 0.0, &case.y, &[], &[], &tangent_inputs) + .to_vec(); + + // Partitioned tape + let ir_split = TypedIr::from_arena_split_eval(&arena, root); + let split = CompiledExpr::from_ir(ir_split); + let mut s_split = vec![0.0; split.scratch_len()]; + let mut cache = split.eval_primal(&mut s_split, 0.0, &case.y, &[], &[]); + let split_result = cache.eval_tangent(&tangent_inputs).to_vec(); + + assert_bitwise_eq(&mono_result, &split_result); + } +} + +// Proptest properties + +proptest! { + #![proptest_config(ProptestConfig::with_cases(200))] + + /// `eval_primal` then `PrimalCache::eval_tangent` must produce bitwise + /// identical tangent output to monolithic `eval_with_tangent`. + #[test] + fn split_eval_matches_monolithic(case in arb_split_eval_case()) { + check_split_eval_matches_monolithic(&case); + } + + /// Repeated tangent-only evaluation after a single primal must be stable. + #[test] + fn split_eval_repeated_tangent_is_stable(case in arb_split_eval_case()) { + check_repeated_tangent_stable(&case); + } + + /// The split tape must still agree with monolithic after the production + /// derivative pipeline `jacobian.rs` runs: tangent_wrt_states -> + /// simplify(Aggressive) -> zero_propagate -> cse -> dce -> + /// from_arena_split_eval. + #[test] + fn split_eval_after_production_pipeline(case in arb_split_eval_case()) { + check_split_eval_after_pipeline(&case); + } +} + +// Targeted test: named Conditional, Index, wide fan-out, sparse matmul cases + +#[test] +fn targeted_split_eval_cases_pass() { + for case in targeted_split_eval_cases() { + check_split_eval_matches_monolithic(&case); + check_repeated_tangent_stable(&case); + check_split_eval_after_pipeline(&case); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions new file mode 100644 index 0000000000..fb82b7cc26 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc fc6afd0373771469220a9e47b281f78c5e79208224feb66daf251380e3107990 # shrinks to case = TangentCase { arena: Arena { nodes: [Scalar(0.5), Neg(NodeId(0)), Arctan(NodeId(1))] }, root: NodeId(2), y: [0.1, 0.1], seeds: [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0], [-0.2784313725490196, -1.0]], n_states: 2 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs new file mode 100644 index 0000000000..24965adf77 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs @@ -0,0 +1,80 @@ +mod common; + +use common::cases::{TangentCase, arb_split_eval_case, targeted_split_eval_cases}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::tangent_batch::{is_batchable, run_tangent_batch, tangent_scratch_len}; +use pybamm_core::{CompiledExpr, TangentInputs, TypedIr, tangent_wrt_states}; + +const LANES: usize = 4; + +/// A batched sweep must reproduce `LANES` independent scalar tangent sweeps bit +/// for bit: each lane accumulates in the same order, so nothing may drift. +fn check_batched_matches_scalar(case: &TangentCase) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + let ir = TypedIr::from_arena_split_eval(&arena, tangent_root); + if !is_batchable(&ir) { + return; + } + let split = ir.split_eval_info().expect("split-eval tape"); + let primal_len = split.primal_buffer_size; + let n_states = case.y.len(); + + // Pad to a full block so short cases still exercise every lane. + let seeds: Vec> = (0..LANES) + .map(|lane| { + case.seeds + .get(lane) + .cloned() + .unwrap_or_else(|| vec![0.0; n_states]) + }) + .collect(); + + let compiled = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; compiled.scratch_len()]; + + let mut cache = compiled.eval_primal(&mut scratch, 0.0, &case.y, &[], &[]); + let scalar: Vec> = seeds + .iter() + .map(|seed| { + cache + .eval_tangent(&TangentInputs { + dy: Some(seed), + dp: None, + }) + .to_vec() + }) + .collect(); + + let mut lane_seeds = vec![0.0; n_states * LANES]; + for (lane, seed) in seeds.iter().enumerate() { + for (state, &value) in seed.iter().enumerate() { + lane_seeds[state * LANES + lane] = value; + } + } + + let ir = compiled.ir(); + let mut tan = vec![0.0; tangent_scratch_len(ir, LANES)]; + compiled.run_primal_section(&mut scratch, 0.0, &case.y, &[], &[]); + let batched = run_tangent_batch::(ir, &scratch[..primal_len], &mut tan, &lane_seeds); + + for (lane, want) in scalar.iter().enumerate() { + let got: Vec = (0..want.len()).map(|e| batched[e * LANES + lane]).collect(); + assert_bitwise_eq(want, &got); + } +} + +#[test] +fn targeted_batched_tangent_matches_scalar() { + for case in targeted_split_eval_cases() { + check_batched_matches_scalar(&case); + } +} + +proptest! { + #[test] + fn batched_tangent_matches_scalar(case in arb_split_eval_case()) { + check_batched_matches_scalar(&case); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs b/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs new file mode 100644 index 0000000000..b16ff88230 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs @@ -0,0 +1,1000 @@ +//! Only-active-branch execution: the invariant is asserted from the executed +//! instruction count, not from a reported tape length, because a count metric +//! alone cannot tell "the work is skipped" from "the work is not reported". + +use pybamm_core::adjoint::AdjointTape; +use pybamm_core::arena::{Arena, NodeId}; +use pybamm_core::eval::{CompiledExpr, TangentInputs}; +use pybamm_core::ir::{Instruction, TypedIr}; +use pybamm_core::node::{CsrData, Node, Shape}; +use pybamm_core::{simplify, tangent_wrt_states}; + +/// `cond(sel, [chain(y, 1), chain(y, 8), chain(y, 16)])` — three exclusive +/// branches of very different sizes over an `InputParameter` selector. +fn uneven_branches() -> (Arena, NodeId, Vec) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let mut branches = Vec::new(); + let mut lens = Vec::new(); + for depth in [1_usize, 8, 16] { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + branches.push(node); + lens.push(depth); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + (arena, cond, lens) +} + +#[test] +fn only_the_active_branch_executes() { + let (arena, root, lens) = uneven_branches(); + let expr = CompiledExpr::from_ir(TypedIr::from_arena(&arena, root)); + let mut scratch = vec![0.0; expr.scratch_len()]; + + // common = y load + selector load + Conditional + Dispatch = 4. + let common = 4; + for (i, &len) in lens.iter().enumerate() { + let sel = (i + 1) as f64; + let (executed, _) = expr.eval_counted(&mut scratch, 0.0, &[0.3], &[], &[sel]); + assert_eq!( + executed, + common + len, + "selector {sel} executed {executed} instructions, expected {}", + common + len + ); + } + + // No match: the dispatch runs, no block does. + let (executed, out) = expr.eval_counted(&mut scratch, 0.0, &[0.3], &[], &[0.0]); + assert_eq!(executed, common); + assert_eq!(out, &[0.0]); +} + +#[test] +fn executed_count_is_independent_of_the_other_branches() { + // Growing branch 3 must not change what branch 1 costs. This is the + // invariant `test_unified_active_branch_independent_of_other_modes` checks. + let cost_of_branch_one = |depth_of_last: usize| { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let mut b2 = y; + for _ in 0..depth_of_last { + b2 = arena.alloc(Node::Cos(b2)); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let expr = CompiledExpr::from_ir(TypedIr::from_arena(&arena, cond)); + let mut scratch = vec![0.0; expr.scratch_len()]; + expr.eval_counted(&mut scratch, 0.0, &[0.3], &[], &[1.0]).0 + }; + assert_eq!(cost_of_branch_one(2), cost_of_branch_one(64)); +} + +#[test] +fn short_circuited_results_are_bitwise_identical() { + let (arena, root, _) = uneven_branches(); + let short = CompiledExpr::from_ir(TypedIr::from_arena(&arena, root)); + // Checked against `expected_conditional`, an oracle written out from the + // semantics contract rather than derived from the evaluator. + let mut s = vec![0.0; short.scratch_len()]; + for sel in [ + 0.0_f64, + 0.49, + 0.5, + 0.51, + 1.0, + 1.49, + 1.5, + 2.0, + 2.5, + 3.0, + 3.49, + 3.5, + 4.0, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = short.eval(&mut s, 0.0, &[0.3], &[], &[sel])[0]; + let expected = expected_conditional(0.3, sel); + assert_eq!( + got.to_bits(), + expected.to_bits(), + "selector {sel}: got {got}, expected {expected}" + ); + } +} + +/// The semantics contract, written out independently of the evaluator. +fn expected_conditional(y: f64, sel: f64) -> f64 { + for (i, depth) in [1_usize, 8, 16].iter().enumerate() { + let idx = (i + 1) as f64; + if sel > idx - 0.5 && sel < idx + 0.5 { + let mut v = y; + for _ in 0..*depth { + v = v.sin(); + } + return v; + } + } + 0.0 +} + +#[test] +fn nested_conditionals_degrade_without_miscompiling() { + // Inner conditional inside outer branch 1. The inner cone is forced common, + // so only the outer conditional short-circuits — and the values are right. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let outer_sel = arena.alloc(Node::InputParameter { + name: "o".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let inner_sel = arena.alloc(Node::InputParameter { + name: "i".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let ia = arena.alloc(Node::Sin(y)); + let ib = arena.alloc(Node::Cos(y)); + let inner = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![ia, ib], + }); + let ob = arena.alloc(Node::Neg(y)); + let outer = arena.alloc(Node::Conditional { + selector: outer_sel, + branches: vec![inner, ob], + }); + + let ir = TypedIr::from_arena(&arena, outer); + // Values alone would pass even if the pass emitted no blocks, so pin the + // outcome: only the outer conditional blocks, one instruction per branch. + assert_eq!( + ir.dispatch_count(), + 1, + "only the outer conditional is blockable" + ); + assert_eq!(ir.branch_block_lens(), vec![1, 1]); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + let y0 = 0.3_f64; + for (o, i, expected) in [ + (1.0, 1.0, y0.sin()), + (1.0, 2.0, y0.cos()), + (1.0, 0.0, 0.0), + (2.0, 1.0, -y0), + (0.0, 1.0, 0.0), + ] { + let got = expr.eval(&mut s, 0.0, &[y0], &[], &[o, i])[0]; + assert_eq!(got.to_bits(), expected.to_bits(), "outer {o} inner {i}"); + } +} + +#[test] +fn a_branch_owning_no_nodes_gets_an_empty_block() { + // Branch 1 is a bare shared state load (common), branch 2 owns one node. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b2 = arena.alloc(Node::Sin(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![y, b2], + }); + let ir = TypedIr::from_arena(&arena, cond); + assert_eq!(ir.branch_block_lens(), vec![0, 1]); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + assert_eq!(expr.eval(&mut s, 0.0, &[0.3], &[], &[1.0]), &[0.3]); + assert_eq!( + expr.eval(&mut s, 0.0, &[0.3], &[], &[2.0]), + &[0.3_f64.sin()] + ); +} + +/// Differential test between the two block-forming layouts, reuse slots +/// (`from_arena`) and pinned SSA slots (`from_arena_pinned`), plus a closed-form +/// oracle, on a graph mixing every hazard: a cone shared by a strict subset of +/// branches, a `Common` node also read outside the conditional, a nested +/// conditional, and a `SparseMatrix` inside a block. +/// +/// The worst failure mode is a node wrongly placed in a block whose value is read +/// from outside. That gives a *different* wrong value per layout, so the two tapes +/// disagree, and the oracle catches a mistake they share. +#[test] +fn scheduled_tape_matches_the_pinned_tape_and_a_closed_form_oracle() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let sel = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let shared_inner = arena.alloc(Node::Exp(y)); + let shared = arena.alloc(Node::Sqrt(shared_inner)); + + // Branch 0 and branch 1 share `shared` (a strict subset of three branches). + let b0 = arena.alloc(Node::Sin(shared)); + let csr = arena.alloc(Node::SparseMatrix( + CsrData::try_new( + vec![0, 1, 2, 3], + vec![1, 2, 0], + vec![2.0, -1.5, 0.75], + Shape { rows: 3, cols: 3 }, + ) + .expect("valid csr") + .into(), + )); + let mm = arena.alloc(Node::MatMul(csr, shared)); + let b1 = arena.alloc(Node::Cos(mm)); + + // Branch 2 is itself a conditional: its cone degrades to `Common`. + let inner_sel = arena.alloc(Node::Time); + let ia = arena.alloc(Node::Tanh(y)); + let ib = arena.alloc(Node::Abs(y)); + let b2 = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![ia, ib], + }); + + let outer = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b0, b1, b2], + }); + // `shared_inner` is also read outside the conditional, so it must stay + // `Common` and be emitted outside every block. + let root = arena.alloc(Node::Add(outer, shared_inner)); + + let scheduled = CompiledExpr::from_ir(TypedIr::from_arena(&arena, root)); + let pinned = CompiledExpr::from_ir(TypedIr::from_arena_pinned(&arena, root)); + // Block lengths count instructions, not nodes: branch 1's `SparseMatrix` emits + // none, and slot allocation cannot change the count, so both layouts must agree. + assert_eq!(scheduled.ir().dispatch_count(), 1); + assert_eq!(scheduled.ir().branch_block_lens(), vec![2, 3, 1]); + assert_eq!(pinned.ir().dispatch_count(), 1); + assert_eq!(pinned.ir().branch_block_lens(), vec![2, 3, 1]); + + let mut sched_scratch = vec![0.0; scheduled.scratch_len()]; + let mut pinned_scratch = vec![0.0; pinned.scratch_len()]; + for sel_val in [0.0_f64, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, f64::NAN] { + for t in [1.0_f64, 2.0, 0.0] { + let state = [0.35, -0.2, 0.9, sel_val]; + let got = scheduled.eval(&mut sched_scratch, t, &state, &[], &[]); + let want = pinned.eval(&mut pinned_scratch, t, &state, &[], &[]); + let oracle = expected_hazard_graph([state[0], state[1], state[2]], sel_val, t); + assert_eq!(got.len(), want.len()); + assert_eq!(got.len(), oracle.len()); + for (i, ((g, w), o)) in got.iter().zip(want).zip(&oracle).enumerate() { + assert_eq!( + g.to_bits(), + w.to_bits(), + "selector {sel_val}, t {t}, element {i}: scheduled {g}, pinned {w}" + ); + assert!( + (g - o).abs() <= 1e-15 * o.abs(), + "selector {sel_val}, t {t}, element {i}: got {g}, oracle {o}" + ); + } + } + } +} + +/// The hazard graph's semantics, written out independently of the evaluator: +/// `cond(sel, [sin(s), cos(A @ s), cond(t, [tanh(y), abs(y)])]) + exp(y)` where +/// `s = sqrt(exp(y))` and `A` is the test's 3x3 CSR matrix. +fn expected_hazard_graph(y: [f64; 3], sel: f64, t: f64) -> Vec { + let exp_y = y.map(f64::exp); + let shared = exp_y.map(f64::sqrt); + let outer = match branch_window(sel, 3) { + Some(0) => shared.map(f64::sin), + Some(1) => [2.0 * shared[1], -1.5 * shared[2], 0.75 * shared[0]].map(f64::cos), + Some(2) => match branch_window(t, 2) { + Some(0) => y.map(f64::tanh), + Some(1) => y.map(f64::abs), + _ => [0.0; 3], + }, + _ => [0.0; 3], + }; + (0..3).map(|i| outer[i] + exp_y[i]).collect() +} + +/// The 1-based round-to-nearest branch window of the semantics contract. +fn branch_window(selector: f64, n_branches: usize) -> Option { + (0..n_branches).find(|&i| { + let idx = (i + 1) as f64; + selector > idx - 0.5 && selector < idx + 0.5 + }) +} + +/// `eval_batch` dispatches the union of branches any lane selects, then lets +/// `Conditional` pick per lane within that union. Lanes select different +/// branches here, so every block stays live — the scenario a batch evaluator +/// that skipped based on a single lane's choice would get wrong. +#[test] +fn batch_eval_over_dispatched_blocks_matches_per_lane_eval() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + // A state-derived selector: `InputParameter` broadcasts to every lane. + let sel = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let mut branches = Vec::new(); + for depth in [1_usize, 4, 9] { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + branches.push(node); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + + let expr = CompiledExpr::from_ir(TypedIr::from_arena(&arena, cond)); + assert_eq!( + expr.ir().branch_block_lens(), + vec![1, 4, 9], + "the tape must actually carry blocks for this to test anything" + ); + + // Four lanes: branch 1, branch 3, branch 2, and no match. + let k = 4; + let ts = vec![0.0; k]; + let y_cols = [0.3, 1.0, 0.4, 3.0, 0.5, 2.0, 0.6, 0.0]; + let n_states = 2; + + let mut per_lane = vec![0.0_f64; k]; + let mut s = vec![0.0; expr.scratch_len()]; + for l in 0..k { + let y_l = &y_cols[l * n_states..(l + 1) * n_states]; + per_lane[l] = expr.eval(&mut s, ts[l], y_l, &[], &[])[0]; + } + + let mut batch_scratch = vec![0.0; expr.scratch_len() * k]; + let batched = expr + .eval_batch(&mut batch_scratch, k, &ts, &y_cols, &[]) + .expect("primal tape batches"); + for l in 0..k { + assert_eq!( + batched[l].to_bits(), + per_lane[l].to_bits(), + "lane {l}: batch {} vs scalar {}", + batched[l], + per_lane[l] + ); + } +} + +/// Two sibling top-level conditionals — the shape `PyBaMM`'s unified experiment +/// model actually builds (a control residual plus a fused termination event). +/// This is the only shape that exercises per-`Dispatch` `blocks_idx` bookkeeping, +/// the `group_at` anchor table, and summing `branch_block_lens` over more than +/// one dispatch. +#[test] +fn two_independent_conditionals_get_one_dispatch_each() { + const DEPTHS_A: [usize; 2] = [1, 3]; + const DEPTHS_B: [usize; 2] = [2, 5]; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel_a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let sel_b = arena.alloc(Node::InputParameter { + name: "b".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let mut a_branches = Vec::new(); + for depth in DEPTHS_A { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + a_branches.push(node); + } + let mut b_branches = Vec::new(); + for depth in DEPTHS_B { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Cos(node)); + } + b_branches.push(node); + } + let cond_a = arena.alloc(Node::Conditional { + selector: sel_a, + branches: a_branches, + }); + let cond_b = arena.alloc(Node::Conditional { + selector: sel_b, + branches: b_branches, + }); + let root = arena.alloc(Node::Add(cond_a, cond_b)); + + let ir = TypedIr::from_arena(&arena, root); + assert_eq!(ir.dispatch_count(), 2); + // `b`'s dispatch is emitted first: the topological walk reaches `cond_b`'s + // cone before `cond_a`'s. Both dispatches' blocks are listed in tape order. + assert_eq!(ir.branch_block_lens(), vec![2, 5, 1, 3]); + // y + two selectors + two Conditionals + Add + two Dispatches. + let common = 8; + assert_eq!(ir.common_instruction_count(), common); + assert_eq!(ir.instructions().len(), 19); + + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.3_f64; + // A selector of 0 or 3 matches no branch of a two-branch conditional, so the + // matrix covers a live block, the other live block, and no-match on each side. + for (a_index, sel_a_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0), (None, 3.0)] { + for (b_index, sel_b_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0), (None, 3.0)] { + let a_cost = a_index.map_or(0, |i| DEPTHS_A[i]); + let b_cost = b_index.map_or(0, |i| DEPTHS_B[i]); + let (executed, out) = + expr.eval_counted(&mut scratch, 0.0, &[y0], &[], &[sel_a_val, sel_b_val]); + assert_eq!( + executed, + common + a_cost + b_cost, + "selectors ({sel_a_val}, {sel_b_val}) executed {executed}" + ); + + let chain = |depth: usize, f: fn(f64) -> f64| { + let mut v = y0; + for _ in 0..depth { + v = f(v); + } + v + }; + let expected = a_index.map_or(0.0, |i| chain(DEPTHS_A[i], f64::sin)) + + b_index.map_or(0.0, |i| chain(DEPTHS_B[i], f64::cos)); + assert_eq!( + out[0].to_bits(), + expected.to_bits(), + "selectors ({sel_a_val}, {sel_b_val}): got {}, expected {expected}", + out[0] + ); + } + } +} + +/// Forward-mode JVP over a dispatched tape. `eval_with_tangent` on a non-split +/// `from_arena` tape is production-live; the split-eval tape reaches the same +/// derivative graph through a different layout — one `Dispatch` per half instead +/// of one for the whole tape — so the two are independent of each other. The +/// closed-form check at the end stops them passing by a shared mistake. +#[test] +fn jvp_over_dispatched_blocks_matches_the_split_eval_tape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + // Nonlinear branches, so each adjoint genuinely needs its own primal value. + let b0 = arena.alloc(Node::Sin(x)); + let b1 = arena.alloc(Node::Exp(x)); + let b2 = arena.alloc(Node::Tanh(x)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b0, b1, b2], + }); + let jac = tangent_wrt_states(&mut arena, cond); + let jac = simplify(&mut arena, jac); + + let scheduled = CompiledExpr::from_ir(TypedIr::from_arena(&arena, jac)); + let split = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&arena, jac)); + assert_eq!(scheduled.ir().dispatch_count(), 1); + assert_eq!(scheduled.ir().branch_block_lens(), vec![2, 2, 5]); + assert_eq!( + split.ir().dispatch_count(), + 2, + "split-eval splits this cone into one block per partition half" + ); + + let mut sched_scratch = vec![0.0; scheduled.scratch_len()]; + let mut split_scratch = vec![0.0; split.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let x0 = 0.4_f64; + for sel_val in [ + 0.0_f64, + 0.5, + 1.0, + 1.5, + 2.0, + 2.5, + 3.0, + 3.5, + 4.0, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = + scheduled.eval_with_tangent(&mut sched_scratch, 0.0, &[x0], &[], &[sel_val], &tangent) + [0]; + let want = + split.eval_with_tangent(&mut split_scratch, 0.0, &[x0], &[], &[sel_val], &tangent)[0]; + assert_eq!( + got.to_bits(), + want.to_bits(), + "selector {sel_val}: scheduled {got}, split {want}" + ); + } + + // And the active branch's derivative is the mathematically right one, so the + // two tapes are not merely agreeing on a shared mistake. + for (sel_val, expected) in [ + (1.0_f64, x0.cos()), + (2.0, x0.exp()), + (3.0, x0.tanh().mul_add(-x0.tanh(), 1.0)), + (0.0, 0.0), + ] { + let got = + scheduled.eval_with_tangent(&mut sched_scratch, 0.0, &[x0], &[], &[sel_val], &tangent) + [0]; + assert!( + (got - expected).abs() < 1e-12, + "selector {sel_val}: got {got}, expected {expected}" + ); + } +} + +#[test] +fn common_instruction_count_excludes_branch_blocks() { + let (arena, root, lens) = uneven_branches(); + let ir = TypedIr::from_arena(&arena, root); + let total: usize = lens.iter().sum(); + assert_eq!( + ir.instructions().len(), + ir.common_instruction_count() + total + ); + // common = y + selector + Dispatch + Conditional + assert_eq!(ir.common_instruction_count(), 4); +} + +/// A conditional whose cone spans the primal/tangent partition still +/// short-circuits in both halves, and `primal_end` remains a valid split point. +#[test] +fn split_eval_blocks_respect_the_primal_tangent_partition() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + // Two branches whose primal work is deliberately unequal. + let mut b1 = y; + for _ in 0..3 { + b1 = arena.alloc(Node::Sin(b1)); + } + let mut b2 = y; + for _ in 0..12 { + b2 = arena.alloc(Node::Exp(b2)); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let tangent_root = tangent_wrt_states(&mut arena, cond); + + let ir = TypedIr::from_arena_split_eval(&arena, tangent_root); + let split = ir.split_eval_info().expect("split-eval info"); + // One block per half, so the cone really is split rather than degraded. + assert_eq!(ir.dispatch_count(), 2); + + // primal_end must not land inside a block span. + for (i, instr) in ir.instructions().iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len as usize { + let (rel, len) = ir.consts().branch_blocks[blocks_idx as usize + b]; + let (start, end) = (i + rel as usize, i + rel as usize + len as usize); + assert!( + split.primal_end <= start || split.primal_end >= end, + "primal_end {} splits block [{start}, {end})", + split.primal_end + ); + } + } + } + + // Both halves short-circuit: the executed count differs by branch. + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + let seed = [1.0, 0.0]; + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let mut counts = Vec::new(); + for sel_val in [1.0_f64, 2.0] { + let (executed, _) = + expr.eval_counted_with_tangent(&mut s, 0.0, &[0.3, 0.4], &[], &[sel_val], &tangent); + counts.push(executed); + } + assert!( + counts[0] < counts[1], + "branch 1 ({}) should cost less than branch 2 ({})", + counts[0], + counts[1] + ); +} + +/// The split tape must match the monolithic tape bitwise for every selector, +/// including no-match. +#[test] +fn split_eval_matches_monolithic_for_every_selector() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Exp(y)); + let b3 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + let tangent_root = tangent_wrt_states(&mut arena, cond); + + let mono = CompiledExpr::from_ir(TypedIr::from_arena(&arena, tangent_root)); + let split = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&arena, tangent_root)); + let seed = [1.0, 1.0]; + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let mut sm = vec![0.0; mono.scratch_len()]; + let mut ss = vec![0.0; split.scratch_len()]; + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 1.5, f64::NAN] { + let a = mono + .eval_with_tangent(&mut sm, 0.0, &[0.3, 0.4], &[], &[sel_val], &tangent) + .to_vec(); + let mut cache = split.eval_primal(&mut ss, 0.0, &[0.3, 0.4], &[], &[sel_val]); + let b = cache.eval_tangent(&tangent).to_vec(); + assert_eq!(a.len(), b.len()); + for (x, z) in a.iter().zip(&b) { + assert_eq!(x.to_bits(), z.to_bits(), "selector {sel_val}"); + } + } +} + +/// A `Dispatch` reads its selector's slot, so a group anchored at the end of one +/// partition half is only safe when that half also computes the selector. A +/// tangent-tainted selector puts it in the *tangent* pool while the branch-owned +/// primal nodes stay in the primal half — the primal `Dispatch` would then read +/// an unwritten slot and skip a block whose value the tangent half still reads. +/// `assert_block_slots_private` cannot see this (there is no writer to blame), +/// so the scheduler must degrade instead. +#[test] +fn a_tangent_selector_degrades_rather_than_dispatching_on_an_unwritten_slot() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let dy = arena.alloc(Node::TangentStateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + // Tangent-tainted, so the selector lands in the tangent slot pool. + let sel = arena.alloc(Node::Add(dy, one)); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Exp(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let mono = CompiledExpr::from_ir(TypedIr::from_arena(&arena, cond)); + let split = CompiledExpr::from_ir(TypedIr::from_arena_split_eval(&arena, cond)); + let mut sm = vec![0.0; mono.scratch_len()]; + let mut ss = vec![0.0; split.scratch_len()]; + let y0 = 0.4_f64; + // `dy + 1` sweeps no-match, both branches and a half-integer boundary. + for seed in [-1.0_f64, 0.0, 0.5, 1.0, 2.0, f64::NAN] { + let seeds = [seed]; + let tangent = TangentInputs { + dy: Some(&seeds), + dp: None, + }; + let want = mono.eval_with_tangent(&mut sm, 0.0, &[y0], &[], &[], &tangent)[0]; + let mut cache = split.eval_primal(&mut ss, 0.0, &[y0], &[], &[]); + let got = cache.eval_tangent(&tangent)[0]; + let expected = match seed + 1.0 { + s if s > 0.5 && s < 1.5 => y0.sin(), + s if s > 1.5 && s < 2.5 => y0.exp(), + _ => 0.0, + }; + assert_eq!(got.to_bits(), want.to_bits(), "seed {seed}"); + assert_eq!(got.to_bits(), expected.to_bits(), "seed {seed}"); + } +} + +#[test] +fn pinned_layout_short_circuits_and_stays_ssa() { + let (arena, root, lens) = uneven_branches(); + let ir = TypedIr::from_arena_pinned(&arena, root); + assert_eq!(ir.branch_block_lens(), vec![1, 8, 16]); + // Blocking must not cost the pinned layout its SSA property, which reverse AD + // needs: with no reuse, `y`, the selector and the `Conditional` account for the 3. + assert_eq!(ir.buffer_size(), 3 + lens.iter().sum::()); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + for (i, &len) in lens.iter().enumerate() { + let sel = (i + 1) as f64; + let (executed, out) = expr.eval_counted(&mut s, 0.0, &[0.3], &[], &[sel]); + assert_eq!(executed, 4 + len); + assert_eq!(out[0].to_bits(), expected_conditional(0.3, sel).to_bits()); + } +} + +/// The backward pass must not walk an inactive branch's block. Counted, not +/// inferred: `assemble` reports what the adjoint replay actually touched. +#[test] +fn reverse_pass_skips_inactive_branch_blocks() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + // `Tanh`, not `Exp`: iterated `exp` overflows to inf by depth 5, which would + // make the gradient check vacuous (inf == inf). + let mut b2 = y; + for _ in 0..20 { + b2 = arena.alloc(Node::Tanh(b2)); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let tape = AdjointTape::new(&arena, cond, 1); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; 1]; + + let walked_1 = tape.assemble(&mut scratch, &mut bar, &mut grad, 0.0, &[0.3], &[], &[1.0]); + let g1 = grad[0]; + let walked_2 = tape.assemble(&mut scratch, &mut bar, &mut grad, 0.0, &[0.3], &[], &[2.0]); + let g2 = grad[0]; + + assert!( + walked_1 < walked_2, + "branch 1 backward walk ({walked_1}) should be shorter than branch 2 ({walked_2})" + ); + // Values still correct: d/dy sin(y) and d/dy tanh^20(y). + assert!((g1 - 0.3_f64.cos()).abs() < 1e-12); + let mut expected = 1.0_f64; + let mut v = 0.3_f64; + for _ in 0..20 { + v = v.tanh(); + expected *= v.mul_add(-v, 1.0); + } + assert!( + (g2 - expected).abs() / expected < 1e-9, + "got {g2}, want {expected}" + ); +} + +/// Two sibling conditionals in one adjoint tape: the backward walk crosses two +/// span ends, so it must resolve each independently rather than treating the +/// first one it meets as the only one. The reverse-side counterpart of +/// `two_independent_conditionals_get_one_dispatch_each`. +#[test] +fn reverse_pass_skips_blocks_of_two_sibling_conditionals() { + const DEPTHS_A: [usize; 2] = [1, 3]; + const DEPTHS_B: [usize; 2] = [2, 5]; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel_a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let sel_b = arena.alloc(Node::InputParameter { + name: "b".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let chain = |arena: &mut Arena, depth: usize, tanh: bool| { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(if tanh { + Node::Tanh(node) + } else { + Node::Sin(node) + }); + } + node + }; + let a_branches = DEPTHS_A.map(|d| chain(&mut arena, d, false)).to_vec(); + let b_branches = DEPTHS_B.map(|d| chain(&mut arena, d, true)).to_vec(); + let cond_a = arena.alloc(Node::Conditional { + selector: sel_a, + branches: a_branches, + }); + let cond_b = arena.alloc(Node::Conditional { + selector: sel_b, + branches: b_branches, + }); + let root = arena.alloc(Node::Add(cond_a, cond_b)); + + let tape = AdjointTape::new(&arena, root, 1); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![0.0; 1]; + // y + two selectors + two Conditionals + Add + two Dispatches. + let common = 8; + let y0 = 0.3_f64; + + // d/dy of an iterated sin / tanh chain, by the chain rule. + let d_chain = |depth: usize, tanh: bool| { + let mut derivative = 1.0_f64; + let mut v = y0; + for _ in 0..depth { + if tanh { + v = v.tanh(); + derivative *= v.mul_add(-v, 1.0); + } else { + derivative *= v.cos(); + v = v.sin(); + } + } + derivative + }; + + for (a_index, sel_a_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0)] { + for (b_index, sel_b_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0)] { + let walked = tape.assemble( + &mut scratch, + &mut bar, + &mut grad, + 0.0, + &[y0], + &[], + &[sel_a_val, sel_b_val], + ); + let a_cost = a_index.map_or(0, |i| DEPTHS_A[i]); + let b_cost = b_index.map_or(0, |i| DEPTHS_B[i]); + assert_eq!( + walked, + common + a_cost + b_cost, + "selectors ({sel_a_val}, {sel_b_val}) walked {walked}" + ); + + let expected = a_index.map_or(0.0, |i| d_chain(DEPTHS_A[i], false)) + + b_index.map_or(0.0, |i| d_chain(DEPTHS_B[i], true)); + assert!( + (grad[0] - expected).abs() < 1e-12, + "selectors ({sel_a_val}, {sel_b_val}): got {}, want {expected}", + grad[0] + ); + } + } +} + +/// No match: no block is walked and the gradient is all zeros. The count is what +/// distinguishes skipping from not skipping — an all-zero gradient alone was +/// already true before blocks existed, because the `Conditional` adjoint seeds no +/// branch when nothing matches. +#[test] +fn reverse_pass_on_no_match_walks_no_block() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let tape = AdjointTape::new(&arena, cond, 1); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![9.0; 1]; + let walked = tape.assemble(&mut scratch, &mut bar, &mut grad, 0.0, &[0.3], &[], &[0.0]); + assert_eq!(grad, vec![0.0]); + // y + selector + Conditional + Dispatch, and neither one-instruction block: + // entering both would be 6, and a blockless tape would be 5. + assert_eq!(walked, 4, "no-match walk touched a branch block"); + + // The contrast: matching a branch walks exactly one more instruction. + let matched = tape.assemble(&mut scratch, &mut bar, &mut grad, 0.0, &[0.3], &[], &[1.0]); + assert_eq!(matched, 5); + assert!((grad[0] - 0.3_f64.cos()).abs() < 1e-12); +} + +/// The semantics contract's edge cases through the backward walk: a half-integer +/// window boundary, NaN, a negative and an infinite selector all match no branch, +/// so no block is walked and the gradient is zero. +#[test] +fn reverse_pass_on_boundary_and_nan_selectors_walks_no_block() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let tape = AdjointTape::new(&arena, cond, 1); + let mut scratch = vec![0.0; tape.scratch_len()]; + let mut bar = vec![0.0; tape.scratch_len()]; + let mut grad = vec![9.0; 1]; + + for sel_val in [0.5_f64, 1.5, 2.5, -1.0, f64::NAN, f64::INFINITY] { + let walked = tape.assemble( + &mut scratch, + &mut bar, + &mut grad, + 0.0, + &[0.3], + &[], + &[sel_val], + ); + // y + selector + Conditional + Dispatch, and neither one-instruction block. + assert_eq!(walked, 4, "selector {sel_val} walked a branch block"); + assert_eq!(grad, vec![0.0], "selector {sel_val}"); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_coloring_correctness.rs b/packages/pybamm-rust/pybamm-core/tests/test_coloring_correctness.rs new file mode 100644 index 0000000000..2327662540 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_coloring_correctness.rs @@ -0,0 +1,753 @@ +//! Correctness tests: verify coloring-based Jacobian assembly against +//! central finite differences. +//! +//! This module provides an independent ground truth for Jacobian correctness +//! that does not depend on any symbolic AD infrastructure. + +use pybamm_core::NodeId; +use pybamm_core::arena::Arena; +use pybamm_core::eval::CompiledExpr; +use pybamm_core::jacobian::{JacobianData, JacobianScratch}; +use pybamm_core::model::ModelEvaluator; +use pybamm_core::node::{CsrData, InterpolantData, Node, Shape}; + +/// Assemble coloring Jacobian and compare each entry against central +/// finite differences. +/// +/// Strategy: +/// 1. Assemble coloring Jacobian → dense n×n matrix (scattered from CSC) +/// 2. For each column j, compute FD column via 2 RHS evaluations +/// 3. Compare entry-by-entry +fn compare_coloring_vs_fd( + model: &mut ModelEvaluator, + y: &[f64], + inputs: &[f64], + h: f64, + rtol: f64, +) { + let n = model.n_states(); + let nnz = model.nnz(); + + // Assemble coloring Jacobian into CSC buffer + let mut jac_csc = vec![0.0_f64; nnz]; + model.assemble_jacobian_csc_into_coloring(0.0, y, inputs, &mut jac_csc); + + // Scatter CSC into dense matrix + let csc_sparsity = model.csc_sparsity(); + let mut dense = vec![vec![0.0_f64; n]; n]; + for (col, colptr) in csc_sparsity.colptr.windows(2).enumerate() { + let range_start = colptr[0]; + let range_end = colptr[1]; + for (local, &val) in jac_csc[range_start..range_end].iter().enumerate() { + let row = csc_sparsity.rowind[range_start + local]; + dense[row][col] = val; + } + } + + // Compute FD columns and compare + // NOTE: eval_rhs writes to output buffer, does NOT return a value + let mut f_plus = vec![0.0_f64; n]; + let mut f_minus = vec![0.0_f64; n]; + + for col in 0..n { + let mut y_plus = y.to_vec(); + let mut y_minus = y.to_vec(); + let scale = h.max(1e-10); + y_plus[col] += scale; + y_minus[col] -= scale; + + model.eval_rhs(0.0, &y_plus, inputs, &mut f_plus); + model.eval_rhs(0.0, &y_minus, inputs, &mut f_minus); + + for row in 0..n { + let fd_val = (f_plus[row] - f_minus[row]) / (2.0 * scale); + let coloring_val = dense[row][col]; + let err = (fd_val - coloring_val).abs(); + let scale_val = fd_val.abs().max(coloring_val.abs()).max(1e-15); + assert!( + err / scale_val < rtol, + "Jacobian({row},{col}): coloring={coloring_val}, fd={fd_val}, err={err}, rel={}", + err / scale_val + ); + } + } +} + +/// Build a banded (tridiagonal) model: `f_i = x_{i-1} + 2*x_i + x_{i+1}` +fn build_banded_model(n: usize) -> ModelEvaluator { + let mut arena = Arena::new(); + let svecs: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + + let two = arena.alloc(Node::Scalar(2.0)); + let rows: Vec<_> = (0..n) + .map(|i| { + let left = svecs[i.saturating_sub(1)]; + let mid = svecs[i]; + let right = svecs[(i + 1).min(n - 1)]; + let two_mid = arena.alloc(Node::Mul(two, mid)); + let sum_lr = arena.alloc(Node::Add(left, right)); + arena.alloc(Node::Add(two_mid, sum_lr)) + }) + .collect(); + + let rhs = arena.alloc(Node::Concat(rows)); + let mass = CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("valid identity mass matrix"); + ModelEvaluator::new(&arena, rhs, mass, n, 0) +} + +/// Build a nonlinear model with cross-coupling: +/// `f_i = sin(y_i) + y_{i-1} * y_{i+1}` (with boundary clamping) +fn build_nonlinear_model(n: usize) -> ModelEvaluator { + let mut arena = Arena::new(); + let svecs: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + + let rows: Vec<_> = (0..n) + .map(|i| { + let self_term = arena.alloc(Node::Sin(svecs[i])); + let left = svecs[i.saturating_sub(1)]; + let right = svecs[(i + 1).min(n - 1)]; + let coupling = arena.alloc(Node::Mul(left, right)); + arena.alloc(Node::Add(self_term, coupling)) + }) + .collect(); + + let rhs = arena.alloc(Node::Concat(rows)); + let mass = CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("valid identity mass matrix"); + ModelEvaluator::new(&arena, rhs, mass, n, 0) +} + +/// Build a model with sparse matrix multiplication: +/// `f = A @ y` where `A` is a constant sparse matrix +fn build_sparse_matmul_model(n: usize) -> ModelEvaluator { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + + // Build a banded sparse matrix A (tridiagonal with 1, 2, 1) + let mut indptr = vec![0usize; n + 1]; + let mut indices = Vec::with_capacity(n * 3); + let mut data = Vec::with_capacity(n * 3); + + for (row, ptr) in indptr.iter_mut().enumerate().take(n) { + *ptr = data.len(); + if row > 0 { + indices.push(row - 1); + data.push(1.0); + } + indices.push(row); + data.push(2.0); + if row + 1 < n { + indices.push(row + 1); + data.push(1.0); + } + } + indptr[n] = data.len(); + + let sparse = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(indptr, indices, data, Shape::matrix(n, n)).expect("valid test matrix"), + ))); + let rhs = arena.alloc(Node::MatMul(sparse, y)); + + let mass = CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("valid identity mass matrix"); + ModelEvaluator::new(&arena, rhs, mass, n, 0) +} + +#[test] +fn test_coloring_vs_fd_banded_5() { + let n = 5; + let mut model = build_banded_model(n); + model.set_cj(0.0); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.5, 0.1).sin()).collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); +} + +#[test] +fn test_coloring_vs_fd_banded_50() { + let n = 50; + let mut model = build_banded_model(n); + model.set_cj(0.0); + + let y: Vec = (0..n).map(|i| ((i as f64 + 1.0) * 0.1).sin()).collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); +} + +#[test] +fn test_coloring_vs_fd_nonlinear_10() { + let n = 10; + let mut model = build_nonlinear_model(n); + model.set_cj(0.0); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.3, 0.5).sin()).collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); +} + +#[test] +fn test_coloring_vs_fd_sparse_matmul_8() { + let n = 8; + let mut model = build_sparse_matmul_model(n); + model.set_cj(0.0); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.4, 0.2).sin()).collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); +} + +#[test] +fn test_coloring_vs_fd_multiple_points() { + let n = 15; + let mut model = build_nonlinear_model(n); + model.set_cj(0.0); + + for seed in [42u64, 137, 2718, 31415, 99991] { + let y: Vec = (0..n) + .map(|i| ((seed as f64 + i as f64) * 0.1).sin()) + .collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); + } +} + +#[test] +fn test_coloring_vs_fd_with_mass_cj() { + // With cj != 0 the assembled Jacobian is df/dy - cj*M, while FD gives df/dy, + // so cj*M is subtracted here. + let n = 10; + let mut model = build_banded_model(n); + let cj = 0.5; + model.set_cj(cj); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.3, 0.5).sin()).collect(); + + let nnz = model.nnz(); + let mut jac_csc = vec![0.0_f64; nnz]; + model.assemble_jacobian_csc_into_coloring(0.0, &y, &[], &mut jac_csc); + + // Compute FD (using output buffer pattern) + let h = 1e-7; + let mut dense_fd = vec![vec![0.0_f64; n]; n]; + let mut f_plus = vec![0.0_f64; n]; + let mut f_minus = vec![0.0_f64; n]; + + for col in 0..n { + let mut y_plus = y.clone(); + let mut y_minus = y.clone(); + y_plus[col] += h; + y_minus[col] -= h; + model.eval_rhs(0.0, &y_plus, &[], &mut f_plus); + model.eval_rhs(0.0, &y_minus, &[], &mut f_minus); + for row in 0..n { + dense_fd[row][col] = (f_plus[row] - f_minus[row]) / (2.0 * h); + } + } + + // Subtract cj * identity from FD (mass matrix is identity) + for (i, row) in dense_fd.iter_mut().enumerate() { + row[i] -= cj; + } + + // Scatter CSC into dense + let csc_sparsity = model.csc_sparsity(); + let mut dense_csc = vec![vec![0.0_f64; n]; n]; + for (col, colptr) in csc_sparsity.colptr.windows(2).enumerate() { + let range_start = colptr[0]; + let range_end = colptr[1]; + for (local, &val) in jac_csc[range_start..range_end].iter().enumerate() { + let row = csc_sparsity.rowind[range_start + local]; + dense_csc[row][col] = val; + } + } + + for row in 0..n { + for col in 0..n { + let err = (dense_csc[row][col] - dense_fd[row][col]).abs(); + let scale = dense_csc[row][col] + .abs() + .max(dense_fd[row][col].abs()) + .max(1e-15); + assert!( + err / scale < 1e-5, + "Jacobian({row},{col}) with cj={cj}: assembled={}, fd={}", + dense_csc[row][col], + dense_fd[row][col] + ); + } + } +} + +#[test] +fn test_coloring_vs_fd_identity_jacobian() { + // f_i = y_i → Jacobian = I + let n = 6; + let mut arena = Arena::new(); + let svecs: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + let rhs = arena.alloc(Node::Concat(svecs)); + + let mass = CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("valid identity mass matrix"); + let mut model = ModelEvaluator::new(&arena, rhs, mass, n, 0); + model.set_cj(0.0); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.5, 0.1).sin()).collect(); + + compare_coloring_vs_fd(&mut model, &y, &[], 1e-7, 1e-5); +} + +// Dense-row split (loop A). These tests drive `JacobianData` directly: build via +// `new_wrt_states` and finite-difference its `assemble_csc_into` output. + +/// Dense-row fixture: `sin(y)` elementwise (a diagonal block) concatenated with +/// one scalar row `ones_row(1 x n) @ (y*y)` that depends on every state — the +/// SPMe/vaas shape of a dense algebraic row atop a sparse structure. +fn build_dense_row_expr(n: usize) -> (Arena, NodeId, usize) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let elt = arena.alloc(Node::Sin(y)); // len n, diagonal + let gy = arena.alloc(Node::Mul(y, y)); // len n, diagonal, nonlinear + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .expect("valid test matrix"), + ))); + let dense = arena.alloc(Node::MatMul(ones, gy)); // len 1, depends on all n + let root = arena.alloc(Node::Concat(vec![elt, dense])); + (arena, root, n) +} + +/// Same shape but the dense row lives inside a length-2 vector block (a `2 x n` +/// matmul), so no scalar sub-node holds it and the split has to synthesise one. +fn build_dense_row_vector_block_expr(n: usize) -> (Arena, NodeId, usize) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let elt = arena.alloc(Node::Sin(y)); + let gy = arena.alloc(Node::Mul(y, y)); + // 2 x n: row 0 dense (all cols), row 1 sparse (col 0 only). + let mut indices: Vec = (0..n).collect(); + indices.push(0); + let a = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n, n + 1], + indices, + vec![1.0; n + 1], + Shape::matrix(2, n), + ) + .expect("valid test matrix"), + ))); + let block = arena.alloc(Node::MatMul(a, gy)); // len 2, row 0 dense + let root = arena.alloc(Node::Concat(vec![elt, block])); + (arena, root, n) +} + +/// Assemble `jac`'s CSC values at `t = 0.5` through the production driver, at the +/// lane width this artifact's tape would actually run. +fn assemble_csc(jac: &JacobianData, y: &[f64]) -> Vec { + let layout = jac.layout(); + let mut scratch = JacobianScratch::new(jac); + let mut data = vec![0.0; layout.n_slots()]; + jac.assemble_into(&mut scratch, layout, 0.5, y, &[], &[], &mut data); + data +} + +/// Assemble the `JacobianData` via loop A and compare every CSC entry against +/// central finite differences of the primal expression. +fn fd_check_jacobian_data(arena: &Arena, root: NodeId, jac: &JacobianData, y: &[f64]) { + let primal = CompiledExpr::new(arena, root); + let data = assemble_csc(jac, y); + + let mut s = vec![0.0; primal.scratch_len()]; + let eps = 1e-6; + for col in 0..jac.n_cols() { + let mut yp = y.to_vec(); + yp[col] += eps; + let mut ym = y.to_vec(); + ym[col] -= eps; + let fp = primal.eval(&mut s, 0.5, &yp, &[], &[]).to_vec(); + let fm = primal.eval(&mut s, 0.5, &ym, &[], &[]).to_vec(); + let (lo, hi) = (jac.csc().colptr[col], jac.csc().colptr[col + 1]); + for (&dk, &row) in data[lo..hi].iter().zip(&jac.csc().rowind[lo..hi]) { + let fd = (fp[row] - fm[row]) / (2.0 * eps); + assert!( + (dk - fd).abs() <= 1e-5 * (1.0 + fd.abs()), + "entry ({row},{col}): assembled {dk} vs fd {fd}" + ); + } + } +} + +/// A block of `width` fully dense rows over `n` states, plus a diagonal +/// remainder: the shape a 2-D current collector produces, where one dense row +/// per collector node sits inside a single vector-valued block. +/// +/// `coupled` builds the block's interior, which is what decides whether the +/// rows can be extracted at all. +fn build_dense_block_expr( + n: usize, + width: usize, + coupled: impl FnOnce(&mut Arena, NodeId) -> NodeId, +) -> (Arena, NodeId, usize) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let diagonal = arena.alloc(Node::Sin(y)); + let coupled = coupled(&mut arena, y); + // width x n, every entry stored: each block row reaches every state. + let indptr: Vec = (0..=width).map(|row| row * n).collect(); + let indices: Vec = (0..width).flat_map(|_| 0..n).collect(); + let data: Vec = (0..width * n) + .map(|k| (k as f64).mul_add(0.03, 0.5)) + .collect(); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(indptr, indices, data, Shape::matrix(width, n)) + .expect("valid test matrix"), + ))); + let block = arena.alloc(Node::MatMul(matrix, coupled)); + let root = arena.alloc(Node::Concat(vec![diagonal, block])); + (arena, root, n) +} + +/// Assemble `jac` into a dense matrix, for comparing two builds of the same +/// derivative entry by entry. +fn assemble_dense(jac: &JacobianData, y: &[f64]) -> Vec> { + let data = assemble_csc(jac, y); + let mut dense = vec![vec![0.0; jac.n_cols()]; jac.n_rows()]; + for (col, span) in jac.csc().colptr.windows(2).enumerate() { + for k in span[0]..span[1] { + dense[jac.csc().rowind[k]][col] = data[k]; + } + } + dense +} + +#[test] +fn test_dense_block_splits_every_row_and_matches_the_unsplit_build() { + // A minority of the states, as a 2-D collector's block is: the regime where + // deleting one colour per state pays for the reverse passes. + let (n, width) = (64, 16); + let (arena, root, n_states) = + build_dense_block_expr(n, width, |arena, y| arena.alloc(Node::Mul(y, y))); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n_states); + + assert_eq!(jac.n_dense_rows(), width, "the whole plateau must split"); + assert_eq!(jac.n_candidate_rows(), width); + assert!( + jac.coloring().n_colors <= width + 1, + "the split must leave the sparse remainder cheap to colour, got {}", + jac.coloring().n_colors + ); + + // A full-column subset builds the same derivative with splitting off, so + // it is the reference the split must reproduce. + let all: Vec = (0..n_states).collect(); + let reference = JacobianData::new_wrt_state_subset(&arena, root, n_rows, n_states, &all); + assert!(reference.dense_rows().is_empty()); + assert!(reference.coloring().n_colors > jac.coloring().n_colors); + + let y: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.13, 0.7).sin()) + .collect(); + let split = assemble_dense(&jac, &y); + let unsplit = assemble_dense(&reference, &y); + for (row, (split_row, unsplit_row)) in split.iter().zip(&unsplit).enumerate() { + for (col, (&a, &b)) in split_row.iter().zip(unsplit_row).enumerate() { + assert!( + (a - b).abs() <= 1e-12 * b.abs().mul_add(1.0, 1.0), + "entry ({row},{col}): split {a} vs unsplit {b}" + ); + } + } + fd_check_jacobian_data(&arena, root, &jac, &y); +} + +/// An interpolant has no cheap indexed form, so a block containing one keeps +/// the wide colouring. The decline has to stay visible in the telemetry. +#[test] +fn test_unextractable_dense_block_declines_visibly() { + let (n, width) = (64, 16); + let (arena, root, n_states) = build_dense_block_expr(n, width, |arena, y| { + arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![-2.0, 0.0, 2.0], vec![-1.0, 0.5, 3.0]) + .expect("valid table"), + ), + child: y, + }) + }); + + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n_states); + assert!( + jac.dense_rows().is_empty(), + "an interpolant must stop the walk" + ); + assert_eq!( + jac.n_candidate_rows(), + width, + "a declined split must still report what it wanted" + ); + + let y_values: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.11, 0.2).sin()) + .collect(); + fd_check_jacobian_data(&arena, root, &jac, &y_values); +} + +#[test] +fn test_dense_row_split_matches_fd() { + let n = 30; + let (arena, root, n_states) = build_dense_row_expr(n); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n_states); + + // 30-nnz row >= DENSE_ROW_MIN_NNZ, extracts to a scalar block, and the + // reduced coloring (1) strictly beats the full coloring (30) -> split. + assert_eq!(jac.n_dense_rows(), 1, "dense row must be split out"); + assert!( + jac.coloring().n_colors <= 4, + "residual coloring must reflect the sparse structure, got {}", + jac.coloring().n_colors + ); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.1, 0.3).sin()).collect(); + fd_check_jacobian_data(&arena, root, &jac, &y); +} + +#[test] +fn test_dense_conditional_ignores_invalid_inactive_branch() { + let mut arena = Arena::new(); + let n = 20; + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .unwrap(), + ))); + let active = arena.alloc(Node::MatMul(ones, y)); + let y0 = arena.alloc(Node::Index { + child: y, + start: 0, + end: 1, + }); + let inactive = arena.alloc(Node::Sqrt(y0)); + let selector = arena.alloc(Node::Scalar(1.0)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![active, inactive], + }); + let jac = JacobianData::new_wrt_states(&arena, root, 1, n); + assert_eq!(jac.n_dense_rows(), 1, "dense row must use reverse AD"); + + fd_check_jacobian_data(&arena, root, &jac, &vec![-1.0; n]); +} + +// Dense-row split, model-level loops C/D. A green `compare_coloring_vs_fd` proves +// the model consumes `dense_rows` rather than aliasing the reduced coloring. + +/// Square DAE: `f_i = sin(y_i)` for the first `n-1` (differential) rows and +/// `f_{n-1} = sum_j y_j^2` (algebraic dense row over all `n` states). +fn build_dense_row_model(n: usize) -> ModelEvaluator { + let mut arena = Arena::new(); + let y_full = arena.alloc(Node::StateVector { start: 0, end: n }); + let y_head = arena.alloc(Node::StateVector { + start: 0, + end: n - 1, + }); + let diff = arena.alloc(Node::Sin(y_head)); + let gy = arena.alloc(Node::Mul(y_full, y_full)); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .expect("valid test matrix"), + ))); + let dense = arena.alloc(Node::MatMul(ones, gy)); + let rhs = arena.alloc(Node::Concat(vec![diff, dense])); + let mass = CsrData::try_new( + (0..n).chain(std::iter::once(n - 1)).collect(), + (0..n - 1).collect(), + vec![1.0; n - 1], + Shape::matrix(n, n), + ) + .expect("valid test mass matrix"); + ModelEvaluator::new(&arena, rhs, mass, n, 0) +} + +#[test] +fn test_coloring_vs_fd_dense_row_20() { + let n = 20; + let mut model = build_dense_row_model(n); + // The split must be active: reduced coloring is far below the dense row's nnz. + assert!( + model.coloring().n_colors <= 4, + "dense-row split must reduce the coloring, got {}", + model.coloring().n_colors + ); + model.set_cj(0.0); + + let y: Vec = (0..n) + .map(|i| (i as f64).mul_add(0.1, 0.3).sin() + 0.7) + .collect(); + compare_coloring_vs_fd(&mut model, &y, &[], 1e-6, 1e-5); +} + +#[test] +fn test_dense_row_stats_report_the_split_and_its_tape() { + let model = build_dense_row_model(20); + let stats = model.jacobian_stats(); + assert_eq!(stats.n_dense_rows, 1); + // The split's compiled-memory cost, which nothing else asserts. + assert!(stats.dense_row_tape_instructions > 0); +} + +#[test] +fn test_dense_row_inside_a_vector_block_still_splits() { + let n = 30; + let (arena, root, n_states) = build_dense_row_vector_block_expr(n); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n_states); + + // The dense row is element 0 of a len-2 matmul block: no scalar sub-node + // exists for it, so it is synthesised by pushing the index into the matmul. + assert_eq!(jac.n_dense_rows(), 1, "vector-block dense row must split"); + assert_eq!(jac.dense_rows()[0].rows, vec![n]); + assert_eq!(jac.n_candidate_rows(), 1); + assert!( + jac.coloring().n_colors < n, + "the split must lower the colour count, got {}", + jac.coloring().n_colors + ); + + // Masking the row out of the parent tape needs a width-1 tangent node, which + // a row inside a vector block has none of, so pruning is lost but not sense. + assert!( + std::ptr::eq( + std::sync::Arc::as_ptr(jac.action_tape()), + std::sync::Arc::as_ptr(jac.assembly_tape()) + ), + "an unmaskable split must fall back to the unpruned tape" + ); + // No colour may scatter into the split row -- `layout_in` asserts that for + // every layout it builds, and building one here is what exercises it. + let _ = jac.layout(); + + let y: Vec = (0..n).map(|i| (i as f64).mul_add(0.1, 0.3).sin()).collect(); + fd_check_jacobian_data(&arena, root, &jac, &y); +} + +#[test] +fn test_two_dense_rows_match_fd() { + // Two wide rows over a sparse diagonal remainder; only the nonlinear one + // is worth a reverse pass, since the linear one is known at compile time. + let mut arena = Arena::new(); + let n = 20; + let y = arena.alloc(Node::StateVector { start: 0, end: n }); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + vec![1.0; n], + Shape::matrix(1, n), + ) + .unwrap(), + ))); + let weights = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, n], + (0..n).collect(), + (0..n).map(|i| 1.0 + i as f64).collect(), + Shape::matrix(1, n), + ) + .unwrap(), + ))); + let sq = arena.alloc(Node::Mul(y, y)); + let dense0 = arena.alloc(Node::MatMul(ones, sq)); // sum y_i^2 + let dense1 = arena.alloc(Node::MatMul(weights, y)); // sum (1+i) y_i + let mut rows = vec![dense0, dense1]; + for i in 2..n { + let yi = arena.alloc(Node::Index { + child: y, + start: i, + end: i + 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + rows.push(arena.alloc(Node::Mul(two, yi))); // sparse diagonal remainder + } + let root = arena.alloc(Node::Concat(rows)); + let n_rows = CompiledExpr::new(&arena, root).output_len(); + let jac = JacobianData::new_wrt_states(&arena, root, n_rows, n); + assert_eq!(jac.n_dense_rows(), 1, "only the nonlinear row is split out"); + assert_eq!(jac.dense_rows()[0].rows[0], 0); + let row1 = jac.sparsity().indptr[1]..jac.sparsity().indptr[2]; + assert_eq!( + jac.constant_csr_entries() + .iter() + .filter(|&&(csr_idx, _)| row1.contains(&csr_idx)) + .count(), + n, + "the linear row comes wholly from the constant table" + ); + + let yv: Vec = (0..n).map(|i| (i as f64).mul_add(0.1, 0.3)).collect(); + fd_check_jacobian_data(&arena, root, &jac, &yv); + + // With the split off, both wide rows fall back to reverse mode. + let reference = JacobianData::new_wrt_states_unsplit(&arena, root, n_rows, n); + assert_eq!(reference.n_dense_rows(), 2); + fd_check_jacobian_data(&arena, root, &reference, &yv); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_cse.rs b/packages/pybamm-rust/pybamm-core/tests/test_cse.rs new file mode 100644 index 0000000000..f3121a2d4f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_cse.rs @@ -0,0 +1,36 @@ +use pybamm_core::{Arena, Node, NodeId, cse}; + +#[test] +fn cse_canonicalizes_structurally_equal_nodes() { + let mut arena = Arena::new(); + let s1 = arena.alloc(Node::Scalar(2.0)); + let s2 = arena.alloc(Node::Scalar(2.0)); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y2 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let m1 = arena.alloc(Node::Mul(s1, y1)); + let m2 = arena.alloc(Node::Mul(s2, y2)); + let root = arena.alloc(Node::Add(m1, m2)); + + let (cse_arena, _root) = cse(&arena, root); + let n_muls = (0..cse_arena.len()) + .filter(|&i| matches!(cse_arena.get(NodeId::from(i)), Node::Mul(_, _))) + .count(); + assert_eq!(n_muls, 1, "structurally identical Muls must canonicalize"); +} + +#[test] +fn cse_distinguishes_state_vectors_with_different_ranges() { + let mut arena = Arena::new(); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y2 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = arena.alloc(Node::Add(y1, y2)); + + let (cse_arena, _root) = cse(&arena, root); + let n_states = (0..cse_arena.len()) + .filter(|&i| matches!(cse_arena.get(NodeId::from(i)), Node::StateVector { .. })) + .count(); + assert_eq!( + n_states, 2, + "StateVectors with different ranges must remain distinct" + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_ffi_abi_contract.rs b/packages/pybamm-rust/pybamm-core/tests/test_ffi_abi_contract.rs new file mode 100644 index 0000000000..d954fe98f9 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_ffi_abi_contract.rs @@ -0,0 +1,354 @@ +//! ABI contract drift test (T4.1). +//! +//! The C++ IDAKLU consumer resolves every Rust FFI entry point by name via +//! `dlsym`, which matches on name only — a signature change in `ffi.rs` would be +//! called through a wrong-typed function pointer (silent UB). This test parses +//! the Rust `extern "C"` exports and the C++ consumer's `RustFfi` typedefs and +//! asserts every consumed symbol matches (name + normalized arg/return types), +//! and that the two ABI version numbers are equal. + +use std::collections::HashMap; +use std::fs; + +const FFI_RS: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/src/ffi.rs"); +const CONSUMER_HEADER: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h" +); +const CONSUMER_IMPL: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp" +); + +/// Map a Rust or C type spelling to a shared canonical token. Panics on an +/// unrecognized spelling so a new FFI type must be taught to both sides here. +fn canon(raw: &str) -> String { + let t = raw.split_whitespace().collect::>().join(" "); + let canonical = match t.as_str() { + "f64" | "double" => "f64", + "u32" | "uint32_t" => "u32", + "c_int" | "int" => "i32", + "*const f64" | "const double*" | "const double *" => "*const f64", + "*mut f64" | "double*" | "double *" => "*mut f64", + "*mut c_int" | "int*" | "int *" => "*mut i32", + "*mut i64" | "int64_t*" | "int64_t *" => "*mut i64", + "*const c_void" | "const void*" | "const void *" => "*const void", + "*mut c_void" | "void*" | "void *" => "*mut void", + other => panic!( + "ABI drift test: unrecognized type spelling {other:?}. \ + Add it to `canon()` on both language sides." + ), + }; + canonical.to_string() +} + +/// Parse the Rust `extern "C"` exports: symbol -> (arg canon types, return canon). +fn parse_rust_exports(src: &str) -> HashMap, String)> { + let mut out = HashMap::new(); + let bytes = src.as_bytes(); + let needle = "extern \"C\" fn "; + let mut cursor = 0; + while let Some(rel) = src[cursor..].find(needle) { + let name_start = cursor + rel + needle.len(); + let paren = src[name_start..].find('(').expect("fn without '('"); + let name = src[name_start..name_start + paren].trim().to_string(); + + let args_start = name_start + paren + 1; + let mut depth = 1; + let mut i = args_start; + while depth > 0 { + match bytes[i] { + b'(' => depth += 1, + b')' => depth -= 1, + _ => {}, + } + i += 1; + } + let args_str = &src[args_start..i - 1]; + + let brace = src[i..].find('{').expect("fn without body"); + let ret_str = src[i..i + brace].trim(); + let ret = ret_str + .strip_prefix("->") + .map_or_else(|| "void".to_string(), |r| canon(r.trim())); + + out.insert(name, (parse_rust_args(args_str), ret)); + cursor = i; + } + out +} + +fn parse_rust_args(s: &str) -> Vec { + s.split(',') + .filter(|p| !p.trim().is_empty()) + .map(|p| { + let ty = p + .split_once(':') + .map(|(_, ty)| ty) + .expect("arg must be `name: type`"); + canon(ty.trim()) + }) + .collect() +} + +/// Parse `pub const RUST_ABI_VERSION: u32 = N;` from source text. +fn parse_rust_abi_version(src: &str) -> u32 { + src.lines() + .find_map(|l| { + l.trim() + .strip_prefix("pub const RUST_ABI_VERSION: u32 =") + .map(|rest| { + rest.trim() + .trim_end_matches(';') + .trim() + .parse::() + .expect("RUST_ABI_VERSION must be a u32 literal") + }) + }) + .expect("ffi.rs must declare `pub const RUST_ABI_VERSION: u32 = N;`") +} + +/// FNV-1a 64-bit — a tiny, portable, deterministic hash. Chosen over +/// `std::hash::DefaultHasher`, whose output is not guaranteed stable across +/// Rust versions or platforms; this value is committed as a golden constant. +fn stable_hash(bytes: &[u8]) -> u64 { + const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const FNV_PRIME: u64 = 0x0000_0100_0000_01b3; + let mut hash = FNV_OFFSET; + for &byte in bytes { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(FNV_PRIME); + } + hash +} + +/// Canonical, order-independent fingerprint of the parsed export set. Sorted by +/// symbol name because the dlsym contract is by-name — reordering exports is a +/// no-op for the ABI, so ordering must not affect the hash. +fn compute_abi_hash(exports: &HashMap, String)>) -> u64 { + let mut entries: Vec = exports + .iter() + .map(|(name, (args, ret))| format!("{name}({})->{ret}", args.join(","))) + .collect(); + entries.sort(); + stable_hash(entries.join("\n").as_bytes()) +} + +/// Parse `pub const EXPECTED_ABI_HASH: u64 = 0x...;` from source text (hex, +/// `_` separators allowed), mirroring `parse_rust_abi_version`. +fn parse_expected_abi_hash(src: &str) -> u64 { + src.lines() + .find_map(|l| { + l.trim() + .strip_prefix("pub const EXPECTED_ABI_HASH: u64 =") + .map(|rest| { + let raw = rest.trim().trim_end_matches(';').trim(); + let raw = raw.strip_prefix("0x").unwrap_or(raw); + let cleaned: String = raw.chars().filter(|c| *c != '_').collect(); + u64::from_str_radix(&cleaned, 16) + .expect("EXPECTED_ABI_HASH must be a u64 hex literal") + }) + }) + .expect("ffi.rs must declare `pub const EXPECTED_ABI_HASH: u64 = 0x...;`") +} + +struct Consumer { + define_version: u32, + typedefs: HashMap, String)>, + bindings: Vec<(String, String)>, +} + +fn parse_consumer(src: &str) -> Consumer { + let define_version = src + .lines() + .find_map(|l| { + l.trim() + .strip_prefix("#define PYBAMM_RUST_ABI_VERSION") + .map(|rest| { + rest.trim() + .parse::() + .expect("PYBAMM_RUST_ABI_VERSION must be a u32 literal") + }) + }) + .expect("consumer header must `#define PYBAMM_RUST_ABI_VERSION N`"); + + // `using NAME = RET (*)(ARGS);` — may span multiple physical lines. + let mut typedefs = HashMap::new(); + let mut rest = src; + while let Some(p) = rest.find("using ") { + let after = &rest[p + "using ".len()..]; + let semi = after.find(';').expect("`using` without ';'"); + let stmt = &after[..semi]; + rest = &after[semi + 1..]; + if let Some((name, def)) = stmt.split_once('=') { + let def = def.trim(); + if let Some(star) = def.find("(*)") { + let ret = canon(def[..star].trim()); + let tail = &def[star + 3..]; + let open = tail.find('(').expect("typedef without '('"); + let close = tail.rfind(')').expect("typedef without ')'"); + let args = parse_c_args(&tail[open + 1..close]); + typedefs.insert(name.trim().to_string(), (args, ret)); + } + } + } + + // `load_symbol("symbol")` — the authoritative typedef->symbol map. + let mut bindings = Vec::new(); + let mut rest = src; + while let Some(p) = rest.find("load_symbol<") { + let after = &rest[p + "load_symbol<".len()..]; + let gt = after.find('>').expect("load_symbol<...> without '>'"); + let typedef = after[..gt].trim().to_string(); + let tail = &after[gt + 1..]; + let q1 = tail.find('"').expect("load_symbol without symbol string"); + let q2 = tail[q1 + 1..] + .find('"') + .expect("unterminated symbol string"); + let symbol = tail[q1 + 1..q1 + 1 + q2].to_string(); + bindings.push((typedef, symbol)); + rest = &tail[q1 + 1 + q2..]; + } + + Consumer { + define_version, + typedefs, + bindings, + } +} + +fn parse_c_args(s: &str) -> Vec { + let s = s.trim(); + if s.is_empty() || s == "void" { + return Vec::new(); + } + s.split(',') + .filter(|p| !p.trim().is_empty()) + .map(|p| canon(p.trim())) + .collect() +} + +#[test] +fn c_consumer_matches_rust_exports() { + let rust_src = + fs::read_to_string(FFI_RS).unwrap_or_else(|e| panic!("cannot read {FFI_RS}: {e}")); + let c_src = fs::read_to_string(CONSUMER_HEADER) + .unwrap_or_else(|e| panic!("cannot read consumer header {CONSUMER_HEADER}: {e}")); + + let exports = parse_rust_exports(&rust_src); + let consumer = parse_consumer(&c_src); + + // Version numbers must be in lockstep. + let rust_version = parse_rust_abi_version(&rust_src); + assert_eq!( + rust_version, consumer.define_version, + "RUST_ABI_VERSION ({rust_version}) != PYBAMM_RUST_ABI_VERSION ({})", + consumer.define_version + ); + + // Every consumed symbol must match its Rust export exactly. + for (typedef, symbol) in &consumer.bindings { + let (c_args, c_ret) = consumer + .typedefs + .get(typedef) + .unwrap_or_else(|| panic!("no typedef `{typedef}` for symbol `{symbol}`")); + let (r_args, r_ret) = exports.get(symbol).unwrap_or_else(|| { + panic!("consumer binds `{symbol}` but ffi.rs exports no such function") + }); + assert_eq!(c_ret, r_ret, "return type mismatch for `{symbol}`"); + assert_eq!(c_args, r_args, "argument type mismatch for `{symbol}`"); + } +} + +#[test] +fn abi_export_set_hash_is_pinned() { + let rust_src = + fs::read_to_string(FFI_RS).unwrap_or_else(|e| panic!("cannot read {FFI_RS}: {e}")); + let exports = parse_rust_exports(&rust_src); + let actual = compute_abi_hash(&exports); + let expected = parse_expected_abi_hash(&rust_src); + assert_eq!( + actual, expected, + "ABI export set changed: computed hash {actual:#018x} != EXPECTED_ABI_HASH \ + {expected:#018x}. The Rust extern \"C\" surface changed. Update \ + EXPECTED_ABI_HASH in ffi.rs to {actual:#018x} AND bump both RUST_ABI_VERSION \ + and PYBAMM_RUST_ABI_VERSION in lockstep." + ); +} + +/// Every Rust FFI entry point returns a status (`SUCCESS` or a negative +/// `ERROR_*`), and a panic caught at the boundary returns `ERROR_PANIC` without +/// having written the output buffer. A call that drops the status therefore lets +/// the previous step's stale data flow on into SUNDIALS as a valid evaluation. +/// +/// Enforced structurally: the consumer must route every call through +/// `PYBAMM_RUST_CALL` (status-returning) or `PYBAMM_RUST_VALUE` (count-returning), +/// both of which raise on a failure code. No raw `rust_ffi().x(...)` may remain. +#[test] +fn consumer_checks_every_ffi_return_status() { + let src = fs::read_to_string(CONSUMER_IMPL) + .unwrap_or_else(|e| panic!("cannot read {CONSUMER_IMPL}: {e}")); + + let unchecked: Vec<&str> = src + .lines() + .map(str::trim) + .filter(|line| !line.starts_with("//") && !line.starts_with('*')) + .filter(|line| line.contains("rust_ffi().")) + .collect(); + + assert!( + unchecked.is_empty(), + "these calls drop the Rust FFI status code; wrap each in PYBAMM_RUST_CALL \ + or PYBAMM_RUST_VALUE:\n {}", + unchecked.join("\n ") + ); +} + +/// `rust_ffi()` resolves its whole table eagerly and throws on the first missing +/// symbol, so a bound-but-uncalled entry point couples IDAKLU to a Rust export it +/// does not need: renaming that export breaks the Rust path at first use for no +/// reason. `RustFfi` must therefore stay a subset of `ffi.rs` covering exactly +/// what the consumer calls, not a mirror of the full export surface. +#[test] +fn every_bound_ffi_entry_point_is_called() { + // Resolved as part of the version handshake in `rust_ffi()` itself rather + // than through the call macros. + const RESOLVER_INTERNAL: [&str; 1] = ["abi_version"]; + + let header = fs::read_to_string(CONSUMER_HEADER) + .unwrap_or_else(|e| panic!("cannot read consumer header {CONSUMER_HEADER}: {e}")); + let impl_src = fs::read_to_string(CONSUMER_IMPL) + .unwrap_or_else(|e| panic!("cannot read {CONSUMER_IMPL}: {e}")); + + let open = header + .find("struct RustFfi {") + .expect("consumer header must declare `struct RustFfi {`"); + let body_start = open + "struct RustFfi {".len(); + let body_len = header[body_start..] + .find('}') + .expect("`struct RustFfi` without closing '}'"); + // Each field is `rust__t ;`. + let fields: Vec<&str> = header[body_start..body_start + body_len] + .lines() + .map(str::trim) + .filter(|l| l.ends_with(';')) + .filter_map(|l| l.trim_end_matches(';').split_whitespace().nth(1)) + .collect(); + assert!(!fields.is_empty(), "parsed no fields from `struct RustFfi`"); + + let uncalled: Vec<&str> = fields + .into_iter() + .filter(|f| !RESOLVER_INTERNAL.contains(f)) + .filter(|f| { + !impl_src.contains(&format!("PYBAMM_RUST_CALL({f},")) + && !impl_src.contains(&format!("PYBAMM_RUST_VALUE({f},")) + }) + .collect(); + + assert!( + uncalled.is_empty(), + "these `RustFfi` entry points are resolved but never called; drop them \ + from the table (and re-add when a call site lands):\n {}", + uncalled.join("\n ") + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_residual_bit_exactness.rs b/packages/pybamm-rust/pybamm-core/tests/test_residual_bit_exactness.rs new file mode 100644 index 0000000000..2deb02ab74 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_residual_bit_exactness.rs @@ -0,0 +1,119 @@ +//! The residual tape `CompiledModel::new` compiles must be bit-identical to a +//! direct evaluation of the DAG Python handed over. +//! +//! Run with: cargo test --features serialize --test `test_residual_bit_exactness` +//! +//! IDA's Newton is sensitive to the residual at ULP scale, so the compile path +//! may only apply passes that move no bits. This is an in-process A/B rather +//! than a stored fingerprint on purpose: `exp`/`pow`/`tanh` results differ +//! between libm implementations, so a golden float file would fail on a +//! different platform for reasons that have nothing to do with the tape. +//! +//! A bit-exact pass (CSE, DCE, a no-op-`Index` elision) may be added to +//! `CompiledModel::new` freely and this test keeps passing. A fold that shifts +//! ULPs — `simplify`'s int-pow lowering is worth 4096 ULP on the DFN residual — +//! fails it immediately, which is the point. + +#![cfg(feature = "serialize")] + +use pybamm_core::{CompiledExpr, CompiledModel, CsrData, DagSnapshot, Shape}; + +const SPM: &[u8] = include_bytes!("../benches/fixtures/spm.bin"); +const SPME: &[u8] = include_bytes!("../benches/fixtures/spme.bin"); +const DFN: &[u8] = include_bytes!("../benches/fixtures/dfn.bin"); + +const TRIALS: usize = 400; + +/// xorshift64*, so the sweep is reproducible without a rand dependency. +struct XorShift(u64); + +impl XorShift { + const fn next_u64(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + fn unit(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 + } + + /// A positive state spanning 1e-6..1e4, the range battery concentrations + /// and potentials cover, plus an occasional exact zero. + fn state(&mut self) -> f64 { + let draw = self.next_u64(); + if draw.is_multiple_of(17) { + return 0.0; + } + let exponent = (draw % 11) as i32 - 6; + self.unit().mul_add(0.9, 0.1) * 10f64.powi(exponent) + } + + fn states(&mut self, n: usize) -> Vec { + (0..n).map(|_| self.state()).collect() + } +} + +fn identity_mass(n: usize) -> CsrData { + CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("identity mass matrix") +} + +fn check(name: &str, bytes: &[u8]) { + let snap = DagSnapshot::from_bytes(bytes); + let mass = snap + .mass_matrix + .clone() + .unwrap_or_else(|| identity_mass(snap.n_states)); + let model = CompiledModel::new(&snap.arena, snap.root, mass, snap.n_states, snap.n_params); + let mut ws = model.create_workspace(); + + // The reference: the DAG exactly as handed over, no compile-path passes. + let reference = CompiledExpr::new(&snap.arena, snap.root); + let mut scratch = vec![0.0; reference.scratch_len()]; + + let inputs = vec![0.0; snap.n_params]; + // f, not the assembled M*y' - f: comparing through the mass subtraction + // would test `yp - (yp - f) == f`, which is not a floating-point identity. + let mut f = vec![0.0; model.output_len()]; + let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); + + for trial in 0..TRIALS { + let y = rng.states(snap.n_states); + let t = rng.unit() * 3600.0; + + let want = reference.eval(&mut scratch, t, &y, &[], &inputs).to_vec(); + model.eval_rhs(&mut ws, t, &y, &inputs, &mut f); + + for (i, (&got, &want_i)) in f.iter().zip(&want).enumerate() { + assert_eq!( + got.to_bits(), + want_i.to_bits(), + "{name} trial {trial}: residual bit mismatch at output {i}: \ + compiled {got} vs raw-DAG {want_i}. The rhs compile path has \ + gained a pass that moves bits; see the guard in CompiledModel::new." + ); + } + } +} + +#[test] +fn residual_is_compiled_bit_exactly_spm() { + check("spm", SPM); +} + +#[test] +fn residual_is_compiled_bit_exactly_spme() { + check("spme", SPME); +} + +#[test] +fn residual_is_compiled_bit_exactly_dfn() { + check("dfn", DFN); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_sparsity_fixtures.rs b/packages/pybamm-rust/pybamm-core/tests/test_sparsity_fixtures.rs new file mode 100644 index 0000000000..666402e222 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_sparsity_fixtures.rs @@ -0,0 +1,109 @@ +//! Snapshot regression tests for sparsity patterns on real models. +//! +//! Run with: cargo test --features serialize --test `test_sparsity_fixtures` +//! +//! To regenerate the .bin fixtures (e.g., after a deliberate, audited +//! semantic change), run the #[ignore]'d test: +//! cargo test --features serialize --test `test_sparsity_fixtures` \ +//! -- --ignored `regenerate_golden_fixtures` --nocapture + +#![cfg(feature = "serialize")] + +use pybamm_core::{DagSnapshot, JacobianData, TypedIr, detect_sparsity_per_output}; +use std::fs; +use std::path::Path; + +const SPM: &[u8] = include_bytes!("../benches/fixtures/spm.bin"); +const SPME: &[u8] = include_bytes!("../benches/fixtures/spme.bin"); +const DFN: &[u8] = include_bytes!("../benches/fixtures/dfn.bin"); + +/// `(constant entries, colours)` of each fixture's compiled state Jacobian. +/// +/// Inline rather than in the `.bin` fixtures: these are two integers a human +/// has to weigh on every classifier change, so a change should read as a diff +/// here instead of as three rewritten binary blobs. +const JACOBIAN_COUNTS: [(&str, usize, usize); 3] = + [("spm", 116, 0), ("spme", 116, 3), ("dfn", 2316, 9)]; + +fn compute_pattern(bytes: &[u8]) -> (usize, usize, Vec, Vec) { + let snap = DagSnapshot::from_bytes(bytes); + let ir = TypedIr::from_arena(&snap.arena, snap.root); + let n_outputs = ir.output_len(); + let p = detect_sparsity_per_output(&snap.arena, snap.root, n_outputs, snap.n_states); + (p.nrows, p.ncols, p.indptr.clone(), p.indices) +} + +#[test] +#[ignore = "regeneration only — produces .bin fixture files"] +fn regenerate_golden_fixtures() { + let dir = Path::new("tests/fixtures"); + fs::create_dir_all(dir).expect("create fixtures dir"); + for (name, bytes) in [("spm", SPM), ("spme", SPME), ("dfn", DFN)] { + let payload = compute_pattern(bytes); + let encoded = bincode::serialize(&payload).expect("serialize"); + let path = dir.join(format!("sparsity_{name}.bin")); + fs::write(&path, &encoded).expect("write fixture"); + println!( + "wrote {} ({} bytes, nrows={}, ncols={}, nnz={})", + path.display(), + encoded.len(), + payload.0, + payload.1, + payload.3.len() + ); + } +} + +fn load_golden(name: &str) -> (usize, usize, Vec, Vec) { + let path = format!("tests/fixtures/sparsity_{name}.bin"); + let bytes = fs::read(&path).unwrap_or_else(|e| { + panic!("failed to read {path}: {e}. Run --ignored regenerate_golden_fixtures.") + }); + bincode::deserialize(&bytes).expect("deserialize golden") +} + +fn check_against_golden(name: &str, bytes: &[u8]) { + let actual = compute_pattern(bytes); + let golden = load_golden(name); + assert_eq!(actual.0, golden.0, "{name}: nrows"); + assert_eq!(actual.1, golden.1, "{name}: ncols"); + assert_eq!(actual.2, golden.2, "{name}: indptr"); + assert_eq!(actual.3, golden.3, "{name}: indices"); +} + +/// A classifier that silently proved nothing would cost only sweeps, so the +/// pattern fixtures above would still pass. These pin what it resolved. +#[test] +fn jacobian_counts_match_the_pinned_table() { + for (name, constants, colors) in JACOBIAN_COUNTS { + let bytes = match name { + "spm" => SPM, + "spme" => SPME, + _ => DFN, + }; + let snap = DagSnapshot::from_bytes(bytes); + let n_outputs = TypedIr::from_arena(&snap.arena, snap.root).output_len(); + let jac = JacobianData::new_wrt_states(&snap.arena, snap.root, n_outputs, snap.n_states); + assert_eq!( + jac.constant_csr_entries().len(), + constants, + "{name}: constant entries" + ); + assert_eq!(jac.coloring().n_colors, colors, "{name}: colours"); + } +} + +#[test] +fn sparsity_spm_matches_golden() { + check_against_golden("spm", SPM); +} + +#[test] +fn sparsity_spme_matches_golden() { + check_against_golden("spme", SPME); +} + +#[test] +fn sparsity_dfn_matches_golden() { + check_against_golden("dfn", DFN); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_sparsity_oracle.rs b/packages/pybamm-rust/pybamm-core/tests/test_sparsity_oracle.rs new file mode 100644 index 0000000000..16f777b90f --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_sparsity_oracle.rs @@ -0,0 +1,477 @@ +//! Random-DAG property test: assert `detect_sparsity_per_output` produces +//! output bit-identical to a vendored copy of the pre-refactor recursive +//! implementation. +//! +//! The oracle code below is a verbatim copy of `sparsity.rs` at the +//! commit that starts this refactor. It is frozen — do NOT update it +//! when `sparsity.rs` changes. That is the whole point of the oracle. + +use pybamm_core::arena::{Arena, NodeId}; +use pybamm_core::node::Node; +use pybamm_core::{SparsityPattern, detect_sparsity_per_output}; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::collections::HashSet; + +// Frozen copy of `sparsity.rs::analyze_output_dependencies` and its helpers, +// renamed with an `Oracle` suffix, kept as an independent implementation. + +#[derive(Clone, Debug)] +enum ElementDepsOracle { + Scalar(HashSet), + Vector(Vec>), +} + +impl ElementDepsOracle { + fn scalar_empty() -> Self { + Self::Scalar(HashSet::new()) + } + const fn len(&self) -> usize { + match self { + Self::Scalar(_) => 1, + Self::Vector(v) => v.len(), + } + } + fn get(&self, idx: usize) -> &HashSet { + match self { + Self::Scalar(deps) => deps, + Self::Vector(v) => &v[idx], + } + } + fn union_all(&self) -> HashSet { + match self { + Self::Scalar(deps) => deps.clone(), + Self::Vector(v) => { + let mut result = HashSet::new(); + for deps in v { + result.extend(deps); + } + result + }, + } + } + fn to_vector(&self, len: usize) -> Vec> { + match self { + Self::Scalar(deps) => vec![deps.clone(); len], + Self::Vector(v) => { + if v.len() == len { + v.clone() + } else if v.len() == 1 { + vec![v[0].clone(); len] + } else { + let union = self.union_all(); + vec![union; len] + } + }, + } + } +} + +fn analyze_oracle(arena: &Arena, node_id: NodeId) -> ElementDepsOracle { + match arena.get(node_id) { + // StateVector: each element depends only on its corresponding state index + Node::StateVector { start, end } => { + let deps: Vec> = (*start..*end) + .map(|i| { + let mut set = HashSet::new(); + set.insert(i); + set + }) + .collect(); + if deps.len() == 1 { + ElementDepsOracle::Scalar(deps.into_iter().next().unwrap()) + } else { + ElementDepsOracle::Vector(deps) + } + }, + + // No state dependencies for these types + Node::StateVectorDot { start, end } | Node::TangentStateVector { start, end } => { + let len = end - start; + if len == 1 { + ElementDepsOracle::scalar_empty() + } else { + ElementDepsOracle::Vector(vec![HashSet::new(); len]) + } + }, + Node::TangentParameter { .. } + | Node::Scalar(_) + | Node::Time + | Node::InputParameter { .. } => ElementDepsOracle::scalar_empty(), + Node::ZeroVector { len } => { + if *len == 1 { + ElementDepsOracle::scalar_empty() + } else { + ElementDepsOracle::Vector(vec![HashSet::new(); *len]) + } + }, + Node::Array(arr) => { + if arr.data().len() == 1 { + ElementDepsOracle::scalar_empty() + } else { + ElementDepsOracle::Vector(vec![HashSet::new(); arr.data().len()]) + } + }, + Node::SparseMatrix(csr) => { + // Sparse matrix is constant, no state dependencies + // Output is rows x cols, but typically used as matrix not vector + let total = csr.shape().rows * csr.shape().cols; + if total == 1 { + ElementDepsOracle::scalar_empty() + } else { + ElementDepsOracle::Vector(vec![HashSet::new(); csr.shape().rows]) + } + }, + + // Unary operations: preserve per-element structure + Node::Neg(a) + | Node::Abs(a) + | Node::Sqrt(a) + | Node::Exp(a) + | Node::Log(a) + | Node::Sin(a) + | Node::Cos(a) + | Node::Tanh(a) + | Node::Sinh(a) + | Node::Cosh(a) + | Node::Arcsinh(a) + | Node::Arctan(a) + | Node::Erf(a) + | Node::Sign(a) + | Node::Floor(a) + | Node::Ceiling(a) => analyze_oracle(arena, *a), + + // Reductions: output is scalar with union of all input dependencies + Node::MaxReduce(a) | Node::MinReduce(a) => { + let child_deps = analyze_oracle(arena, *a); + ElementDepsOracle::Scalar(child_deps.union_all()) + }, + + // Internal-only node, never produced by `random_dag` below; no + // pre-refactor oracle behavior exists to freeze for it. + Node::ReduceArgSelect { .. } => unreachable!("not generated by random_dag"), + + // Binary operations: combine with broadcast semantics + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => { + let deps_a = analyze_oracle(arena, *a); + let deps_b = analyze_oracle(arena, *b); + combine_binary_oracle(&deps_a, &deps_b) + }, + + // MatMul: row i depends on columns with non-zeros + // For dense case: row i of output depends on all elements of vector + Node::MatMul(mat_id, vec_id) => { + let vec_deps = analyze_oracle(arena, *vec_id); + + // Check if matrix is sparse + match arena.get(*mat_id) { + Node::SparseMatrix(csr) => { + // For sparse matrix: row i depends on columns where mat[i,:] is non-zero + let n_rows = csr.shape().rows; + let mut result = Vec::with_capacity(n_rows); + + for row in 0..n_rows { + let row_start = csr.indptr()[row]; + let row_end = csr.indptr()[row + 1]; + let mut row_deps = HashSet::new(); + + // Row i depends on vec[col] for each non-zero mat[i, col] + for &col in &csr.indices()[row_start..row_end] { + row_deps.extend(vec_deps.get(col)); + } + result.push(row_deps); + } + + if result.len() == 1 { + ElementDepsOracle::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDepsOracle::Vector(result) + } + }, + Node::Array(arr) => { + // Dense matrix: each output row depends on all vector elements + let union = vec_deps.union_all(); + let n_rows = arr.shape().rows; + if n_rows == 1 { + ElementDepsOracle::Scalar(union) + } else { + ElementDepsOracle::Vector(vec![union; n_rows]) + } + }, + _ => { + // For computed matrices, assume dense (conservative) + let mat_deps = analyze_oracle(arena, *mat_id); + let union_mat = mat_deps.union_all(); + let union_vec = vec_deps.union_all(); + let mut combined = union_mat; + combined.extend(union_vec); + ElementDepsOracle::Scalar(combined) + }, + } + }, + + // Index: select subset of child dependencies + Node::Index { child, start, end } => { + let child_deps = analyze_oracle(arena, *child); + let len = end - start; + + match &child_deps { + ElementDepsOracle::Scalar(deps) => { + // Indexing a scalar (should be 0..1) + if len == 1 { + ElementDepsOracle::Scalar(deps.clone()) + } else { + ElementDepsOracle::Vector(vec![deps.clone(); len]) + } + }, + ElementDepsOracle::Vector(v) => { + let subset: Vec<_> = v[*start..*end].to_vec(); + if subset.len() == 1 { + ElementDepsOracle::Scalar(subset.into_iter().next().unwrap()) + } else { + ElementDepsOracle::Vector(subset) + } + }, + } + }, + + // Concat: concatenate child dependencies + Node::Concat(children) => { + let mut result = Vec::new(); + for child in children { + let child_deps = analyze_oracle(arena, *child); + match child_deps { + ElementDepsOracle::Scalar(deps) => result.push(deps), + ElementDepsOracle::Vector(v) => result.extend(v), + } + } + if result.len() == 1 { + ElementDepsOracle::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDepsOracle::Vector(result) + } + }, + + // Interpolant: output depends on child (the x value being interpolated) + Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => analyze_oracle(arena, *child), + + // N-D interpolant: element-wise, union of all children's deps + Node::InterpolantNd { children, .. } | Node::InterpolantNdPartial { children, .. } => { + let mut acc = analyze_oracle(arena, children[0]); + for &c in &children[1..] { + let next = analyze_oracle(arena, c); + acc = combine_binary_oracle(&acc, &next); + } + acc + }, + + // Conditional: union of selector and all branches + Node::Conditional { selector, branches } => { + let selector_deps = analyze_oracle(arena, *selector); + + // Get dependencies from all branches + let branch_deps: Vec<_> = branches.iter().map(|b| analyze_oracle(arena, *b)).collect(); + + // Determine output length from branches + let output_len = branch_deps + .iter() + .map(ElementDepsOracle::len) + .max() + .unwrap_or(1); + + // Combine: each output element depends on selector + corresponding branch elements + let selector_union = selector_deps.union_all(); + let mut result = Vec::with_capacity(output_len); + + for i in 0..output_len { + let mut elem_deps = selector_union.clone(); + for bd in &branch_deps { + if let Some(last_idx) = bd.len().checked_sub(1) { + elem_deps.extend(bd.get(i.min(last_idx))); + } + } + result.push(elem_deps); + } + + if result.len() == 1 { + ElementDepsOracle::Scalar(result.into_iter().next().unwrap()) + } else { + ElementDepsOracle::Vector(result) + } + }, + } +} + +fn combine_binary_oracle(a: &ElementDepsOracle, b: &ElementDepsOracle) -> ElementDepsOracle { + match (a, b) { + // Both scalars: union + (ElementDepsOracle::Scalar(da), ElementDepsOracle::Scalar(db)) => { + let mut combined = da.clone(); + combined.extend(db); + ElementDepsOracle::Scalar(combined) + }, + // Scalar + Vector: broadcast scalar to each element + (ElementDepsOracle::Scalar(scalar_deps), ElementDepsOracle::Vector(vec_deps)) => { + let result: Vec<_> = vec_deps + .iter() + .map(|vd| { + let mut combined = scalar_deps.clone(); + combined.extend(vd); + combined + }) + .collect(); + ElementDepsOracle::Vector(result) + }, + // Vector + Scalar: broadcast scalar to each element + (ElementDepsOracle::Vector(vec_deps), ElementDepsOracle::Scalar(scalar_deps)) => { + let result: Vec<_> = vec_deps + .iter() + .map(|vd| { + let mut combined = vd.clone(); + combined.extend(scalar_deps); + combined + }) + .collect(); + ElementDepsOracle::Vector(result) + }, + // Vector + Vector: element-wise union + (ElementDepsOracle::Vector(va), ElementDepsOracle::Vector(vb)) => { + let len = va.len().max(vb.len()); + let va_expanded = a.to_vector(len); + let vb_expanded = b.to_vector(len); + + let result: Vec<_> = va_expanded + .iter() + .zip(vb_expanded.iter()) + .map(|(da, db)| { + let mut combined = da.clone(); + combined.extend(db); + combined + }) + .collect(); + ElementDepsOracle::Vector(result) + }, + } +} + +fn detect_oracle( + arena: &Arena, + root: NodeId, + n_outputs: usize, + n_states: usize, +) -> SparsityPattern { + let output_deps = analyze_oracle(arena, root); + let mut pattern = SparsityPattern::new(n_outputs, n_states); + let deps_vec = output_deps.to_vector(n_outputs); + for (row, deps) in deps_vec.iter().enumerate() { + pattern.indptr[row] = pattern.indices.len(); + let mut sorted: Vec = deps.iter().copied().collect(); + sorted.sort_unstable(); + pattern.indices.extend(sorted); + } + pattern.indptr[n_outputs] = pattern.indices.len(); + pattern +} + +// Random DAG generator + +fn random_dag(rng: &mut StdRng, n_states: usize, target_nodes: usize) -> (Arena, NodeId, usize) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + let two = arena.alloc(Node::Scalar(2.0)); + let mut pool: Vec = vec![y, two]; + + // Add some single-element state vectors + for i in 0..n_states { + pool.push(arena.alloc(Node::StateVector { + start: i, + end: i + 1, + })); + } + + while arena.len() < target_nodes { + let op = rng.random_range(0..7); + let id = match op { + 0 => { + let a = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Sin(a)) + }, + 1 => { + let a = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Exp(a)) + }, + 2 => { + let a = pool[rng.random_range(0..pool.len())]; + let b = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Add(a, b)) + }, + 3 => { + let a = pool[rng.random_range(0..pool.len())]; + let b = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Mul(a, b)) + }, + 4 => { + let a = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Neg(a)) + }, + 5 => { + // Concat of 2–4 pool elements + let k = rng.random_range(2..=4); + let children: Vec = (0..k) + .map(|_| pool[rng.random_range(0..pool.len())]) + .collect(); + arena.alloc(Node::Concat(children)) + }, + _ => { + let a = pool[rng.random_range(0..pool.len())]; + arena.alloc(Node::Sub(a, a)) + }, + }; + pool.push(id); + } + + // Wrap with a final Concat to fix the output shape + let n_out = n_states; + let final_children: Vec = (0..n_out) + .map(|_| pool[rng.random_range(0..pool.len())]) + .collect(); + let root = arena.alloc(Node::Concat(final_children)); + (arena, root, n_out) +} + +// The property test + +#[test] +fn property_sparsity_matches_oracle() { + let mut rng = StdRng::seed_from_u64(0x00C0_FFEE); + for trial in 0..200 { + let n_states = rng.random_range(2..30); + let target = rng.random_range(20..80); + let (arena, root, n_outputs) = random_dag(&mut rng, n_states, target); + + let actual = detect_sparsity_per_output(&arena, root, n_outputs, n_states); + let oracle = detect_oracle(&arena, root, n_outputs, n_states); + + assert_eq!( + (actual.nrows, actual.ncols, &actual.indptr, &actual.indices), + (oracle.nrows, oracle.ncols, &oracle.indptr, &oracle.indices), + "trial {trial}: mismatch on n_states={n_states} target_nodes={target}", + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_split_eval.rs b/packages/pybamm-rust/pybamm-core/tests/test_split_eval.rs new file mode 100644 index 0000000000..33805da4c4 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_split_eval.rs @@ -0,0 +1,282 @@ +//! Bitwise-equivalence tests for split primal/tangent evaluation. +//! +//! Verifies that `eval_primal()` + `PrimalCache::eval_tangent()` produces identical +//! results to `eval_with_tangent()` for various expressions and seed vectors. +use pybamm_core::arena::Arena; +use pybamm_core::eval::{CompiledExpr, TangentInputs}; +use pybamm_core::ir::TypedIr; +use pybamm_core::node::{CsrData, Node, Shape}; +use pybamm_core::tangent::tangent_wrt_states; + +/// Build a tangent expression for `expr` w.r.t. all states, then compile it +/// both with standard IR and split-eval IR, and verify bitwise equality +/// across multiple seed vectors. +fn check_split_eval_equivalence( + arena: &Arena, + expr: pybamm_core::arena::NodeId, + _n_states: usize, + y: &[f64], + seeds: &[Vec], +) { + let mut diff_arena = arena.clone(); + let jac_y = tangent_wrt_states(&mut diff_arena, expr); + + let standard_ir = TypedIr::from_arena(&diff_arena, jac_y); + let split_ir = TypedIr::from_arena_split_eval(&diff_arena, jac_y); + + assert!( + split_ir.split_eval_info().is_some(), + "split IR should have SplitEvalInfo" + ); + + let standard_expr = CompiledExpr::from_ir(standard_ir); + let split_expr = CompiledExpr::from_ir(split_ir); + let mut s_standard = vec![0.0; standard_expr.scratch_len()]; + let mut s_split = vec![0.0; split_expr.scratch_len()]; + + // Evaluate primal once + let mut cache = split_expr.eval_primal(&mut s_split, 0.0, y, &[], &[]); + + for (seed_idx, seed) in seeds.iter().enumerate() { + let tangent = TangentInputs { + dy: Some(seed), + dp: None, + }; + + let standard_result = + standard_expr.eval_with_tangent(&mut s_standard, 0.0, y, &[], &[], &tangent); + let split_result = cache.eval_tangent(&tangent); + + assert_eq!( + standard_result.len(), + split_result.len(), + "Output length mismatch at seed {seed_idx}" + ); + + for (i, (s, sp)) in standard_result.iter().zip(split_result.iter()).enumerate() { + assert_eq!( + s.to_bits(), + sp.to_bits(), + "Bitwise mismatch at output[{i}], seed {seed_idx}: standard={s}, split={sp}" + ); + } + } +} + +#[test] +fn test_split_eval_y_squared() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let expr = arena.alloc(Node::Mul(y, y)); + + let seeds = vec![ + vec![1.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0], + vec![0.0, 0.0, 1.0], + vec![1.0, 1.0, 1.0], + vec![0.3, -0.7, 1.2], + ]; + + let y_vals = [2.0, 3.0, 4.0]; + check_split_eval_equivalence(&arena, expr, 3, &y_vals, &seeds); +} + +#[test] +fn test_split_eval_exp_sin() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let exp_y = arena.alloc(Node::Exp(y)); + let sin_y = arena.alloc(Node::Sin(y)); + let expr = arena.alloc(Node::Mul(exp_y, sin_y)); + + let seeds = vec![ + vec![1.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0], + vec![0.0, 0.0, 1.0], + vec![0.5, -0.3, 0.8], + ]; + + let y_vals = [0.1, 0.5, 1.0]; + check_split_eval_equivalence(&arena, expr, 3, &y_vals, &seeds); +} + +#[test] +fn test_split_eval_tridiagonal() { + let n = 10usize; + let mut arena = Arena::new(); + + let svecs: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + + let two = arena.alloc(Node::Scalar(2.0)); + + let rows: Vec<_> = (0..n) + .map(|i| { + let left = svecs[i.saturating_sub(1)]; + let mid = svecs[i]; + let right = svecs[(i + 1).min(n - 1)]; + + let two_mid = arena.alloc(Node::Mul(two, mid)); + let sum_lr = arena.alloc(Node::Add(left, right)); + arena.alloc(Node::Add(two_mid, sum_lr)) + }) + .collect(); + + let rhs = arena.alloc(Node::Concat(rows)); + + let seeds: Vec> = (0..n) + .map(|i| { + let mut s = vec![0.0; n]; + s[i] = 1.0; + s + }) + .collect(); + + let y: Vec = (0..n).map(|i| (i as f64 * 0.3).sin()).collect(); + check_split_eval_equivalence(&arena, rhs, n, &y, &seeds); +} + +#[test] +fn test_split_eval_sparse_matmul() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 5, 8, 10], + vec![0, 1, 0, 1, 2, 1, 2, 3, 2, 3], + vec![1.0; 10], + Shape::matrix(4, 4), + ) + .expect("valid test matrix"), + ))); + let expr = arena.alloc(Node::MatMul(mat, y)); + + let seeds = vec![ + vec![1.0, 0.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + vec![0.0, 0.0, 0.0, 1.0], + vec![1.0, 0.0, 1.0, 0.0], + vec![0.2, -0.4, 0.6, -0.8], + ]; + + let y_vals = [0.1, 0.2, 0.3, 0.4]; + check_split_eval_equivalence(&arena, expr, 4, &y_vals, &seeds); +} + +#[test] +fn test_split_eval_multi_seed_stability() { + // Verify calling eval_tangent many times after a single eval_primal + // produces correct results every time (no buffer corruption). + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let exp_y = arena.alloc(Node::Exp(y)); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, exp_y)); + + let mut diff_arena = arena.clone(); + let jac_y = tangent_wrt_states(&mut diff_arena, expr); + + let standard_ir = TypedIr::from_arena(&diff_arena, jac_y); + let split_ir = TypedIr::from_arena_split_eval(&diff_arena, jac_y); + + let y_vals = [0.1, 0.2, 0.3, 0.4, 0.5]; + + let standard_expr = CompiledExpr::from_ir(standard_ir); + let split_expr = CompiledExpr::from_ir(split_ir); + let mut s_standard = vec![0.0; standard_expr.scratch_len()]; + let mut s_split = vec![0.0; split_expr.scratch_len()]; + + let mut cache = split_expr.eval_primal(&mut s_split, 0.0, &y_vals, &[], &[]); + + // Run 20 different seeds after a single primal eval + for seed_idx in 0..20u64 { + let seed: Vec = (0_u32..5) + .map(|i| ((seed_idx as f64 + f64::from(i)) * 0.7).sin()) + .collect(); + + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + + let standard_result = + standard_expr.eval_with_tangent(&mut s_standard, 0.0, &y_vals, &[], &[], &tangent); + let split_result = cache.eval_tangent(&tangent); + + for (i, (s, sp)) in standard_result.iter().zip(split_result.iter()).enumerate() { + assert_eq!( + s.to_bits(), + sp.to_bits(), + "Bitwise mismatch at output[{i}], seed {seed_idx}: standard={s}, split={sp}" + ); + } + } +} + +#[test] +fn test_split_eval_jacobian_assembly() { + // Verify the full Jacobian assembly path with split eval works correctly. + // Numerical correctness is verified by test_coloring_correctness (finite differences). + use pybamm_core::model::ModelEvaluator; + use pybamm_core::node::{CsrData, Shape}; + + let n = 20usize; + let mut arena = Arena::new(); + + let svecs: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + + let two = arena.alloc(Node::Scalar(2.0)); + + let rows: Vec<_> = (0..n) + .map(|i| { + let left = svecs[i.saturating_sub(1)]; + let mid = svecs[i]; + let right = svecs[(i + 1).min(n - 1)]; + + let two_mid = arena.alloc(Node::Mul(two, mid)); + let sum_lr = arena.alloc(Node::Add(left, right)); + arena.alloc(Node::Add(two_mid, sum_lr)) + }) + .collect(); + + let rhs = arena.alloc(Node::Concat(rows)); + + let mass = CsrData::try_new( + (0..=n).collect(), + (0..n).collect(), + vec![1.0; n], + Shape::matrix(n, n), + ) + .expect("valid identity mass matrix"); + + let mut model = ModelEvaluator::new(&arena, rhs, mass, n, 0); + model.set_cj(0.0); + + // Just verify assembly completes without panicking. + // Numerical correctness is tested in test_coloring_correctness. + let nnz = model.nnz(); + let mut jac = vec![0.0_f64; nnz]; + let y = vec![0.1_f64; n]; + model.assemble_jacobian_csc_into_coloring(0.0, &y, &[], &mut jac); + + // Verify non-trivial result (not all zeros) + let nonzero_count = jac.iter().filter(|&&v| v.abs() > 1e-15).count(); + assert!( + nonzero_count > 0, + "Jacobian should have nonzero entries, got {nonzero_count}" + ); +} diff --git a/packages/pybamm-rust/pybamm-python/Cargo.toml b/packages/pybamm-rust/pybamm-python/Cargo.toml new file mode 100644 index 0000000000..42b841e911 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "pybamm-python" +version = "0.1.0" +edition = "2024" +rust-version.workspace = true + +[lib] +name = "_core" +crate-type = ["cdylib"] + +[features] +default = ["serialize", "diffsol"] +serialize = ["pybamm-core/serialize", "dep:serde", "dep:bincode"] +diffsol = ["pybamm-core/diffsol", "dep:rayon"] + +[dependencies] +pybamm-core = { path = "../pybamm-core" } +pyo3 = { version = "0.28" } +numpy = "0.28" +serde = { version = "1", features = ["derive"], optional = true } +bincode = { version = "1", optional = true } +# Only the feature-gated `pool` module uses it. +rayon = { version = "1", optional = true } + +[lints] +workspace = true diff --git a/packages/pybamm-rust/pybamm-python/src/errors.rs b/packages/pybamm-rust/pybamm-python/src/errors.rs new file mode 100644 index 0000000000..ef3d699cb4 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/errors.rs @@ -0,0 +1,19 @@ +//! Conversion from `pybamm_core` errors to Python exceptions. + +use pybamm_core::CoreError; +use pyo3::PyErr; +use pyo3::exceptions::PyValueError; + +/// Map a core error to a Python exception. Caller-argument and +/// invalid-data problems (empty `t_eval`, mismatched `y0`/inputs/atol, or a +/// malformed matrix/interpolant crossing the boundary) become `ValueError`; +/// an integration failure inside diffsol becomes `RuntimeError`. +pub fn core_err_to_py(err: CoreError) -> PyErr { + match err { + #[cfg(feature = "diffsol")] + CoreError::Diffsol(source) => { + pyo3::exceptions::PyRuntimeError::new_err(format!("diffsol error: {source}")) + }, + other => PyValueError::new_err(other.to_string()), + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/evaluator_pool.rs b/packages/pybamm-rust/pybamm-python/src/evaluator_pool.rs new file mode 100644 index 0000000000..053b355872 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/evaluator_pool.rs @@ -0,0 +1,102 @@ +//! N independent [`ModelEvaluator`]s over one shared compiled tape. +//! +//! An `IDAKLUSolverGroup` runs one solver per OpenMP thread, and each solver +//! evaluates through the FFI with `&mut ModelEvaluator`. The tape is not what +//! stops that being concurrent — the scratch is, so the pool mints one +//! `Workspace` per solver against the one immutable `CompiledModel` behind the +//! `Arc`. N independent `CompiledModel`s would mean N lowerings and N tape +//! copies, which is the cost the `Arc` exists to avoid. + +use std::cell::UnsafeCell; +use std::sync::atomic::{AtomicBool, Ordering}; + +use pyo3::exceptions::{PyIndexError, PyRuntimeError}; +use pyo3::prelude::*; + +use std::sync::Arc; + +use pybamm_core::{CompiledModel, ModelEvaluator}; + +// `as_ptr` moves exclusive access to one evaluator onto a C++ solver thread, +// so a thread-bound field added to `Workspace` must fail here, not data-race. +const _: () = { + const fn assert_send() {} + assert_send::(); +}; + +/// N [`ModelEvaluator`]s over one shared compiled tape. +/// +/// `UnsafeCell` because C++ writes through each address for as long as the +/// solver group lives, so a `*const` derived from a shared reference and cast +/// to `*mut` would be the wrong provenance under Stacked/Tree Borrows. The +/// `Vec` is filled once in [`from_compiled`](Self::from_compiled) and never +/// resized, so the addresses it hands out stay valid for the pool's life. +#[pyclass(module = "pybamm.rust")] +pub struct EvaluatorPool { + evaluators: Vec>, + /// Take-once flag per evaluator, set when `as_ptr` hands its address out. + taken: Vec, +} + +// SAFETY: the pool never reads or writes through a cell itself, and `as_ptr` +// hands each evaluator's address out at most once (`taken`), so every address +// has exactly one writer; the `CompiledModel` behind the `Arc` is immutable. +#[allow(unsafe_code)] +unsafe impl Sync for EvaluatorPool {} + +impl std::fmt::Debug for EvaluatorPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EvaluatorPool") + .field("len", &self.evaluators.len()) + .finish_non_exhaustive() + } +} + +impl EvaluatorPool { + /// Build a pool of `n` evaluators over one shared compiled model, which + /// allocates `n` scratches and no tapes. + pub(crate) fn from_compiled(compiled: &Arc, n: usize) -> Self { + Self { + evaluators: (0..n) + .map(|_| UnsafeCell::new(ModelEvaluator::from_compiled(Arc::clone(compiled)))) + .collect(), + taken: (0..n).map(|_| AtomicBool::new(false)).collect(), + } + } +} + +#[pymethods] +impl EvaluatorPool { + /// Address of evaluator `index`, for the C++ solver that will drive it. + /// + /// Each address is handed out at most once — a second take raises + /// `RuntimeError` — so no two solvers can be given the same evaluator; + /// build a new pool per solver group rather than re-taking from one. + /// The caller must keep this pool alive for as long as the address is used; + /// `IDAKLUSolverGroup` does that by holding the pool itself. + fn as_ptr(&self, index: usize) -> PyResult { + let cell = self.evaluators.get(index).ok_or_else(|| { + PyIndexError::new_err(format!( + "evaluator index {index} is out of range for a pool of {}", + self.evaluators.len() + )) + })?; + if self.taken[index].swap(true, Ordering::Relaxed) { + return Err(PyRuntimeError::new_err(format!( + "evaluator {index} was already handed to a solver; each address \ + is given out once — build a new pool for a new solver group" + ))); + } + // expose_provenance, not `as usize`: the integer crosses to C++ and is + // cast straight back to a pointer that is then written through. + Ok(cell.get().expose_provenance()) + } + + const fn __len__(&self) -> usize { + self.evaluators.len() + } + + fn __repr__(&self) -> String { + format!("EvaluatorPool(len={})", self.evaluators.len()) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/expr.rs b/packages/pybamm-rust/pybamm-python/src/expr.rs new file mode 100644 index 0000000000..9ea4926af2 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/expr.rs @@ -0,0 +1,962 @@ +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::collections::HashMap; + +use numpy::{PyArray1, PyReadonlyArray1}; +use pyo3::prelude::*; + +use pybamm_core::{ + Arena, ArrayData, CompiledExpr, CubicInterpolantData, InterpolantData, NdInterpolantData, Node, + NodeId, Shape, +}; + +use crate::errors::core_err_to_py; + +/// Validate a half-open `[start, end)` extent, returning a `ValueError` when +/// inverted so the `end - start` size computation cannot wrap in release. +fn check_range(start: usize, end: usize, what: &str) -> PyResult<()> { + if start > end { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "{what}: start ({start}) must not exceed end ({end})" + ))); + } + Ok(()) +} + +// `module` is required so pickle can locate the class as `pybamm.rust.ExprGraph` +// instead of the pyo3 default `builtins.ExprGraph`, which pickle cannot import. +#[pyclass(module = "pybamm.rust")] +#[derive(Debug)] +pub struct ExprGraph { + arena: Arena, + input_map: HashMap, + /// Packed width per registered input, indexed by registration order + /// (parallel to `input_map`'s values). + input_widths: Vec, +} + +/// Serialized form of every `ExprGraph` field, used by the pickle protocol. +#[cfg(feature = "serialize")] +#[derive(serde::Serialize, serde::Deserialize)] +struct ExprGraphState { + arena: Arena, + input_map: HashMap, + input_widths: Vec, +} + +#[pyclass(frozen, module = "pybamm.rust")] +pub struct Expr { + node_id: NodeId, + graph: Py, +} + +impl std::fmt::Debug for Expr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Expr") + .field("node_id", &self.node_id) + .finish_non_exhaustive() + } +} + +impl Expr { + /// Node id, validated to belong to `graph`. Guards against using an `Expr` + /// from one `ExprGraph` in another, which would index the wrong arena. + pub(crate) fn node_id_in(&self, graph: &Py) -> PyResult { + if self.graph.as_ptr() == graph.as_ptr() { + Ok(self.node_id) + } else { + Err(pyo3::exceptions::PyValueError::new_err( + "Expr belongs to a different ExprGraph than the graph being built into", + )) + } + } +} + +impl ExprGraph { + /// Get a reference to the arena (crate-internal). + pub(crate) const fn arena(&self) -> &Arena { + &self.arena + } + + /// Mutable arena reference for build-time node allocation (crate-internal). + pub(crate) const fn arena_mut(&mut self) -> &mut Arena { + &mut self.arena + } + + /// Input names ordered by registration index (crate-internal). + pub(crate) fn input_names(&self) -> Vec { + let mut names = vec![String::new(); self.input_map.len()]; + for (name, &idx) in &self.input_map { + names[idx].clone_from(name); + } + names + } + + /// Packed width per input, ordered by registration index (crate-internal). + pub(crate) fn input_widths(&self) -> Vec { + self.input_widths.clone() + } +} + +/// Validate that the graph rooted at `root` can be lowered safely. +/// +/// Called at every binding entry point so unsupported nodes and invalid shape +/// relationships surface as catchable Python errors instead of FFI panics. +pub fn check_supported(arena: &Arena, root: NodeId) -> PyResult<()> { + if let Some(msg) = pybamm_core::first_unsupported(arena, root) { + return Err(pyo3::exceptions::PyNotImplementedError::new_err(format!( + "Rust conversion does not support this expression: {msg}. Use a \ + CasADi-backed model (set `model.convert_to_format = 'casadi'`)." + ))); + } + if let Some(msg) = pybamm_core::first_invalid(arena, root) { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "Invalid Rust expression graph: {msg}" + ))); + } + Ok(()) +} + +#[pymethods] +impl ExprGraph { + #[new] + fn new() -> Self { + Self { + arena: Arena::new(), + input_map: HashMap::new(), + input_widths: Vec::new(), + } + } + + fn scalar(slf: Py, py: Python<'_>, value: f64) -> Expr { + let id = slf.borrow_mut(py).arena.alloc(Node::Scalar(value)); + Expr { + node_id: id, + graph: slf, + } + } + + fn time(slf: Py, py: Python<'_>) -> Expr { + let id = slf.borrow_mut(py).arena.alloc(Node::Time); + Expr { + node_id: id, + graph: slf, + } + } + + fn state_vector(slf: Py, py: Python<'_>, start: usize, end: usize) -> PyResult { + check_range(start, end, "state_vector")?; + let id = slf + .borrow_mut(py) + .arena + .alloc(Node::StateVector { start, end }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn state_vector_dot(slf: Py, py: Python<'_>, start: usize, end: usize) -> PyResult { + check_range(start, end, "state_vector_dot")?; + let id = slf + .borrow_mut(py) + .arena + .alloc(Node::StateVectorDot { start, end }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + /// Total packed width of every input parameter registered in this graph + /// (sum of `input_widths`, not the count of distinct names), this is + /// the length the FFI/solver boundary expects the stacked `p` array to + /// have, and what `n_inputs=` on `CompiledModel.from_expr` must receive. + fn n_inputs(&self) -> usize { + self.input_widths.iter().sum() + } + + /// Number of nodes in the expression arena (introspection). + #[getter] + const fn n_nodes(&self) -> usize { + self.arena.len() + } + + /// Register (or re-look-up) a named input parameter. + /// + /// `width` is the number of packed values the parameter occupies (>1 for + /// vector-valued inputs). Re-registering an existing name must repeat the + /// same width, mismatches would silently corrupt every other parameter's + /// offset into the packed `p` array. + #[pyo3(signature = (name, width = 1))] + fn input_parameter(slf: Py, py: Python<'_>, name: &str, width: usize) -> PyResult { + let mut graph = slf.borrow_mut(py); + let index = if let Some(&idx) = graph.input_map.get(name) { + let existing_width = graph.input_widths[idx]; + if existing_width != width { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "input parameter '{name}' re-registered with width {width}, \ + previously registered with width {existing_width}" + ))); + } + idx + } else { + let idx = graph.input_map.len(); + graph.input_map.insert(name.to_string(), idx); + graph.input_widths.push(width); + idx + }; + let offset: usize = graph.input_widths[..index].iter().sum(); + let id = graph.arena.alloc(Node::InputParameter { + name: name.to_string(), + index, + offset, + width, + }); + drop(graph); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn array(slf: Py, py: Python<'_>, data: PyReadonlyArray1<'_, f64>) -> PyResult { + let vec = data.as_slice()?.to_vec(); + let len = vec.len(); + let array = ArrayData::try_new(vec, Shape::vector(len)).map_err(core_err_to_py)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Array(Box::new(array))); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + /// Dense matrix constant from row-major data (rows × cols). + fn dense_matrix( + slf: Py, + py: Python<'_>, + data: PyReadonlyArray1<'_, f64>, + rows: usize, + cols: usize, + ) -> PyResult { + let vec = data.as_slice()?.to_vec(); + let array = ArrayData::try_new(vec, Shape::matrix(rows, cols)).map_err(core_err_to_py)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Array(Box::new(array))); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn add(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Add(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sub(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Sub(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn mul(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Mul(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn div(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Div(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn neg(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Neg(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn abs(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Abs(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn pow(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Pow(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sqrt(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Sqrt(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn exp(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Exp(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn log(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Log(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sin(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Sin(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn cos(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Cos(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn tanh(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Tanh(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sinh(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Sinh(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn cosh(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Cosh(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn arcsinh(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Arcsinh(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn arctan(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Arctan(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn erf(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Erf(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sign(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Sign(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn floor(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Floor(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn ceiling(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Ceiling(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn max_reduce(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::MaxReduce(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn min_reduce(slf: Py, py: Python<'_>, a: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::MinReduce(a_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn interpolant_1d_linear( + slf: Py, + py: Python<'_>, + x_data: Vec, + y_data: Vec, + child: &Expr, + ) -> PyResult { + let child_id = child.node_id_in(&slf)?; + let data = InterpolantData::try_new(x_data, y_data).map_err(core_err_to_py)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Interpolant1DLinear { + data: Box::new(data), + child: child_id, + }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + /// Build a 1D cubic/pchip interpolant. `coeffs` is flat row-major + /// `[c0,c1,c2,c3, c0,c1,c2,c3, ...]`, length `4 * (breakpoints.len() - 1)`. + /// `breakpoints` must have at least 2 entries (one segment). + fn interpolant_1d_cubic( + slf: Py, + py: Python<'_>, + breakpoints: Vec, + coeffs: Vec, + child: &Expr, + ) -> PyResult { + // Flat coeffs are per-segment power-basis groups of 4; guard the grouping + // before chunking (which would silently drop a non-multiple-of-4 tail). + if !coeffs.len().is_multiple_of(4) { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "interpolant_1d_cubic: coeffs length must be a multiple of 4, got {}", + coeffs.len() + ))); + } + let child_id = child.node_id_in(&slf)?; + let coeffs: Vec<[f64; 4]> = coeffs + .chunks_exact(4) + .map(|c| [c[0], c[1], c[2], c[3]]) + .collect(); + let data = CubicInterpolantData::try_new(breakpoints, coeffs).map_err(core_err_to_py)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Interpolant1DCubic { + data: Box::new(data), + child: child_id, + }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + /// Build an N-D (2 or 3 axis) tensor-product interpolant. `breakpoints` + /// holds the per-axis knot vectors. `coeffs` is flat: cell-major (axis-0 + /// segment slowest), then `order^ndim` power coeffs per cell (axis-0 + /// power slowest, ascending powers). `order` is 2 (multilinear) or 4 + /// (tensor cubic). One child per axis, evaluated element-wise. + fn interpolant_nd( + slf: Py, + py: Python<'_>, + breakpoints: Vec>, + coeffs: Vec, + order: usize, + children: Vec>, + ) -> PyResult { + // Child count is a graph-construction concern, one child per axis; the + // table's own invariants are checked by `NdInterpolantData::try_new`. + let ndim = breakpoints.len(); + if children.len() != ndim { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "interpolant_nd: expected {ndim} children (one per axis), got {}", + children.len() + ))); + } + let order_u32 = u32::try_from(order).map_err(|_| { + pyo3::exceptions::PyValueError::new_err(format!( + "interpolant_nd: order {order} is too large" + )) + })?; + let data = + NdInterpolantData::try_new(breakpoints, coeffs, order_u32).map_err(core_err_to_py)?; + let child_ids: Vec = children + .iter() + .map(|c| c.borrow().node_id_in(&slf)) + .collect::>()?; + let id = slf.borrow_mut(py).arena.alloc(Node::InterpolantNd { + data: Box::new(data), + children: child_ids, + }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn index( + slf: Py, + py: Python<'_>, + child: &Expr, + start: usize, + end: usize, + ) -> PyResult { + check_range(start, end, "index")?; + let child_id = child.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Index { + child: child_id, + start, + end, + }); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn concat(slf: Py, py: Python<'_>, children: Vec>) -> PyResult { + let child_ids: Vec = children + .iter() + .map(|c| c.borrow().node_id_in(&slf)) + .collect::>()?; + let id = slf.borrow_mut(py).arena.alloc(Node::Concat(child_ids)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn matmul(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::MatMul(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn minimum(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Minimum(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn maximum(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Maximum(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn modulo(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Modulo(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn hypot(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Hypot(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn equal_heaviside(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf + .borrow_mut(py) + .arena + .alloc(Node::EqualHeaviside(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn not_equal_heaviside(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf + .borrow_mut(py) + .arena + .alloc(Node::NotEqualHeaviside(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn equality(slf: Py, py: Python<'_>, a: &Expr, b: &Expr) -> PyResult { + let a_id = a.node_id_in(&slf)?; + let b_id = b.node_id_in(&slf)?; + let id = slf.borrow_mut(py).arena.alloc(Node::Equality(a_id, b_id)); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + /// Create a conditional expression with selector and branches. + /// Branch i is active when i - 0.5 < selector < i + 0.5 (1-based indexing). + fn conditional( + slf: Py, + py: Python<'_>, + selector: &Expr, + branches: Vec>, + ) -> PyResult { + if branches.is_empty() { + return Err(pyo3::exceptions::PyValueError::new_err( + "Conditional requires at least one branch", + )); + } + + let selector_id = selector.node_id_in(&slf)?; + let branch_ids: Vec = branches + .iter() + .map(|b| b.borrow().node_id_in(&slf)) + .collect::>()?; + let id = slf.borrow_mut(py).arena.alloc(Node::Conditional { + selector: selector_id, + branches: branch_ids, + }); + + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn sparse_matrix( + slf: Py, + py: Python<'_>, + indptr: Vec, + indices: Vec, + data: PyReadonlyArray1<'_, f64>, + rows: usize, + cols: usize, + ) -> PyResult { + let data = data.as_slice()?.to_vec(); + let csr = pybamm_core::CsrData::try_new(indptr, indices, data, Shape::matrix(rows, cols)) + .map_err(core_err_to_py)?; + let id = slf + .borrow_mut(py) + .arena + .alloc(Node::SparseMatrix(Box::new(csr))); + Ok(Expr { + node_id: id, + graph: slf, + }) + } + + fn eval_to_float( + &self, + expr: &Expr, + t: f64, + y: Vec, + y_dot: Vec, + inputs: Vec, + ) -> PyResult { + check_supported(&self.arena, expr.node_id)?; + let compiled = CompiledExpr::new(&self.arena, expr.node_id); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut scratch, t, &y, &y_dot, &inputs); + result.first().copied().ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "eval_to_float: expression evaluates to an empty array", + ) + }) + } + + fn eval_to_array<'py>( + &self, + py: Python<'py>, + expr: &Expr, + t: f64, + y: PyReadonlyArray1<'_, f64>, + y_dot: PyReadonlyArray1<'_, f64>, + inputs: Vec, + ) -> PyResult>> { + check_supported(&self.arena, expr.node_id)?; + let y_slice = y.as_slice()?; + let y_dot_slice = y_dot.as_slice()?; + let compiled = CompiledExpr::new(&self.arena, expr.node_id); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let result = compiled.eval(&mut scratch, t, y_slice, y_dot_slice, &inputs); + Ok(PyArray1::from_slice(py, result)) + } + + /// Compile an expression into an immutable, shareable `CompiledFunction`. + /// + /// `n_states` overrides the scanned state extent so partial-group + /// expressions can carry the full system width. + #[pyo3(signature = (expr, name = None, n_states = None))] + fn compile( + slf: Py, + py: Python<'_>, + expr: &Expr, + name: Option, + n_states: Option, + ) -> PyResult { + let node_id = expr.node_id_in(&slf)?; + check_supported(slf.borrow(py).arena(), node_id)?; + crate::function::CompiledFunction::build(py, slf, node_id, name, n_states) + } + + /// Compile a named set of outputs into ONE tape with cross-output + /// sharing: synthetic concat root + recorded slice offsets. + #[pyo3(signature = (outputs, name = None, n_states = None))] + fn compile_group( + slf: Py, + py: Python<'_>, + outputs: &Bound<'_, pyo3::types::PyDict>, + name: Option, + n_states: Option, + ) -> PyResult { + let mut names = Vec::with_capacity(outputs.len()); + let mut ids = Vec::with_capacity(outputs.len()); + for (key, value) in outputs.iter() { + names.push(key.extract::()?); + let expr = value.extract::>()?; + ids.push(expr.node_id_in(&slf)?); + } + if ids.is_empty() { + return Err(pyo3::exceptions::PyValueError::new_err( + "compile_group requires at least one output", + )); + } + { + let g = slf.borrow(py); + for &id in &ids { + check_supported(g.arena(), id)?; + } + } + crate::group::CompiledFunctionGroup::build(py, slf, names, ids, name, n_states) + } + + /// Get the number of colors needed for Jacobian computation + #[cfg(feature = "serialize")] + fn dump_dag( + &self, + expr: &Expr, + path: &str, + model_name: &str, + n_states: usize, + n_params: usize, + ) -> PyResult<()> { + let snapshot = pybamm_core::DagSnapshot { + arena: self.arena.clone(), + root: expr.node_id, + n_states, + n_params, + mass_matrix: None, + model_name: model_name.to_string(), + }; + let bytes = snapshot.to_bytes(); + std::fs::write(path, bytes) + .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?; + Ok(()) + } + + /// Serialize the arena and registered inputs to bytes (pickle protocol). + #[cfg(feature = "serialize")] + fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { + let state = ExprGraphState { + arena: self.arena.clone(), + input_map: self.input_map.clone(), + input_widths: self.input_widths.clone(), + }; + let bytes = bincode::serialize(&state).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!("ExprGraph serialize failed: {e}")) + })?; + Ok(pyo3::types::PyBytes::new(py, &bytes)) + } + + /// Restore the arena and registered inputs from bytes (pickle protocol). + #[cfg(feature = "serialize")] + fn __setstate__(&mut self, state: &[u8]) -> PyResult<()> { + let state: ExprGraphState = bincode::deserialize(state).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!("ExprGraph deserialize failed: {e}")) + })?; + self.arena = state.arena; + self.input_map = state.input_map; + self.input_widths = state.input_widths; + Ok(()) + } + + /// Empty args tuple for pickle (zero-arg `#[new]`); explicit `PyTuple` + /// because pyo3 maps a `()` return to Python `None`, which pickle rejects. + #[cfg(feature = "serialize")] + #[allow(clippy::unused_self)] + fn __getnewargs__<'py>(&self, py: Python<'py>) -> Bound<'py, pyo3::types::PyTuple> { + pyo3::types::PyTuple::empty(py) + } +} + +#[pymethods] +impl Expr { + #[getter] + const fn id(&self) -> u32 { + self.node_id.raw() + } + + fn __add__(&self, py: Python<'_>, other: &Self) -> PyResult { + let other_id = other.node_id_in(&self.graph)?; + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Add(self.node_id, other_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } + + fn __sub__(&self, py: Python<'_>, other: &Self) -> PyResult { + let other_id = other.node_id_in(&self.graph)?; + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Sub(self.node_id, other_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } + + fn __mul__(&self, py: Python<'_>, other: &Self) -> PyResult { + let other_id = other.node_id_in(&self.graph)?; + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Mul(self.node_id, other_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } + + fn __truediv__(&self, py: Python<'_>, other: &Self) -> PyResult { + let other_id = other.node_id_in(&self.graph)?; + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Div(self.node_id, other_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } + + fn __pow__( + &self, + py: Python<'_>, + other: &Self, + _modulo: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let other_id = other.node_id_in(&self.graph)?; + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Pow(self.node_id, other_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } + + fn __neg__(&self, py: Python<'_>) -> PyResult { + let id = self + .graph + .borrow_mut(py) + .arena + .alloc(Node::Neg(self.node_id)); + Ok(Self { + node_id: id, + graph: self.graph.clone_ref(py), + }) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/function.rs b/packages/pybamm-rust/pybamm-python/src/function.rs new file mode 100644 index 0000000000..80d6fe048a --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/function.rs @@ -0,0 +1,903 @@ +//! `CompiledFunction`: the central prepared artifact. + +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::{Arc, OnceLock}; + +use numpy::ndarray::ShapeBuilder; +use numpy::{ + AllowTypeChange, PyArray1, PyArray2, PyArrayLike1, PyReadonlyArray1, PyReadonlyArray2, + PyReadwriteArray1, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyDict; + +use pybamm_core::{ + Arena, CompiledExpr, NodeId, TangentInputs, TypedIr, scan_state_usage, simplify_pipeline, + tangent_wrt_params, tangent_wrt_states, +}; + +use crate::expr::ExprGraph; +use crate::jacobian::CompiledJacobian; +use crate::scratch::{Buffer, ScratchPool}; +use crate::signature::FunctionSignature; + +/// Target working-set size for a lane-batched trajectory tile (~1 MiB). The +/// tile lane count is `clamp(TARGET_BYTES / (8 * scratch_len), 8, 64)`. +const TARGET_BYTES: usize = 1 << 20; + +/// Memoised tangent tape plus its scratch pool (one per seed class). +struct TangentEntry { + expr: CompiledExpr, + pool: ScratchPool, +} + +/// `(graph, root, name, n_states)`, the `__reduce__`/`_rebuild` argument +/// tuple for `CompiledFunction`'s pickle protocol. +type RebuildArgs = (Py, u32, Option, Option); + +// `module` is required so pickle can locate the class as `pybamm.rust.CompiledFunction` +// instead of the pyo3 default `builtins.CompiledFunction`, which pickle cannot import. +#[pyclass(frozen, module = "pybamm.rust")] +pub struct CompiledFunction { + pub(crate) expr: Arc, + pub(crate) sig: FunctionSignature, + pub(crate) pool: ScratchPool, + /// Retained for lazy derivation (jacobian/jvp). Always present: + /// bundle views retain the bundle's graph too. + pub(crate) graph: Py, + pub(crate) root: NodeId, + /// Memoised JVP tape, wrt y (lazily derived on first jvp call). + tangent_y: OnceLock, + /// Memoised JVP tape, wrt p (lazily derived on first jvp call with vp). + tangent_p: OnceLock, + /// Memoised prepared jacobian, wrt y (lazily derived on first call). + pub(crate) jac_y: PyOnceLock>, + /// Memoised prepared jacobian, wrt p (lazily derived on first call). + pub(crate) jac_p: PyOnceLock>, +} + +impl std::fmt::Debug for CompiledFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledFunction") + .field("sig", &self.sig) + .finish_non_exhaustive() + } +} + +impl CompiledFunction { + pub(crate) fn build( + py: Python<'_>, + graph: Py, + root: NodeId, + name: Option, + n_states_override: Option, + ) -> PyResult { + let g = graph.borrow(py); + // Scan the ORIGINAL root: simplification can shrink the reachable state + // extent, and a post-simplify scan would reject valid-length y. + let usage = scan_state_usage(g.arena(), root); + let n_states = resolve_n_states(usage.n_states, n_states_override, name.as_deref())?; + let expr = Arc::new(compile_simplified(g.arena(), root)); + let sig = FunctionSignature { + input_names: g.input_names(), + input_widths: g.input_widths(), + n_states, + uses_y_dot: usage.uses_y_dot, + output_len: expr.output_len(), + name, + }; + drop(g); + let pool = ScratchPool::new(Buffer(expr.scratch_len())); + Ok(Self { + expr, + sig, + pool, + graph, + root, + tangent_y: OnceLock::new(), + tangent_p: OnceLock::new(), + jac_y: PyOnceLock::new(), + jac_p: PyOnceLock::new(), + }) + } + + /// Bundle-view constructor: shares an existing tape (no recompilation). + /// + /// Unlike [`build`](Self::build), the primal expression is already + /// compiled (the bundle's shared `Arc`); we only scan the + /// retained graph for the signature's `uses_y_dot` flag and reachable + /// width. Tangent/jacobian caches start empty and derive lazily on demand. + pub(crate) fn from_shared( + py: Python<'_>, + expr: Arc, + graph: Py, + root: NodeId, + n_states: usize, + name: Option, + ) -> PyResult { + let g = graph.try_borrow(py).map_err(|_| { + PyValueError::new_err("from_shared: graph borrow conflict while preparing bundle view") + })?; + let usage = scan_state_usage(g.arena(), root); + let sig = FunctionSignature { + input_names: g.input_names(), + input_widths: g.input_widths(), + n_states, // bundle views always carry the full system width + uses_y_dot: usage.uses_y_dot, + output_len: expr.output_len(), + name, + }; + drop(g); + let pool = ScratchPool::new(Buffer(expr.scratch_len())); + Ok(Self { + expr, + sig, + pool, + graph, + root, + tangent_y: OnceLock::new(), + tangent_p: OnceLock::new(), + jac_y: PyOnceLock::new(), + jac_p: PyOnceLock::new(), + }) + } + + fn eval_inner( + &self, + t: f64, + y: &[f64], + p: &[f64], + y_dot: Option<&[f64]>, + out: &mut [f64], + ) -> PyResult<()> { + self.sig.check_y(y.len())?; + if self.sig.uses_y_dot { + match y_dot { + Some(yd) => self.sig.check_y_dot(yd.len())?, + None => { + return Err(PyValueError::new_err(format!( + "{}: expression uses y_dot; pass y_dot=", + self.sig.display_name() + ))); + }, + } + } + let mut scratch = self.pool.acquire(); + self.expr + .eval_into(&mut scratch, t, y, y_dot.unwrap_or(&[]), p, out); + self.pool.release(scratch); + Ok(()) + } + + /// Lazily derive and cache the tangent tape + scratch pool for one + /// seed class. + fn tangent_entry(&self, py: Python<'_>, wrt_params: bool) -> PyResult<&TangentEntry> { + let cell = if wrt_params { + &self.tangent_p + } else { + &self.tangent_y + }; + if let Some(cached) = cell.get() { + return Ok(cached); + } + let g = self.graph.try_borrow(py).map_err(|_| { + PyValueError::new_err(format!( + "{}: graph borrow conflict during tangent derivation", + self.sig.display_name() + )) + })?; + let mut diff_arena = g.arena().clone(); + drop(g); // release the graph borrow before the compile pipeline + let root = if wrt_params { + tangent_wrt_params(&mut diff_arena, self.root) + } else { + tangent_wrt_states(&mut diff_arena, self.root) + }; + let (da, root) = simplify_pipeline(diff_arena, root); + let expr = CompiledExpr::from_ir(TypedIr::from_arena(&da, root)); + let pool = ScratchPool::new(Buffer(expr.scratch_len())); + Ok(cell.get_or_init(|| TangentEntry { expr, pool })) + } +} + +#[pymethods] +impl CompiledFunction { + #[pyo3(signature = (t, y, p, y_dot = None))] + fn __call__<'py>( + &self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + y_dot: Option>, + ) -> PyResult>> { + let packed = self.sig.extract_p(p)?; + let mut out = vec![0.0; self.sig.output_len]; + let yd = y_dot.as_ref().map(|a| a.as_slice()).transpose()?; + self.eval_inner(t, y.as_slice()?, &packed, yd, &mut out)?; + Ok(PyArray1::from_vec(py, out)) + } + + /// Alias for `__call__`. + #[pyo3(signature = (t, y, p, y_dot = None))] + fn eval<'py>( + &self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + y_dot: Option>, + ) -> PyResult>> { + self.__call__(py, t, y, p, y_dot) + } + + /// Evaluate into a pre-allocated output array (no intermediate allocation). + #[pyo3(signature = (t, y, p, out, y_dot = None))] + fn eval_into( + &self, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + mut out: PyReadwriteArray1<'_, f64>, + y_dot: Option>, + ) -> PyResult<()> { + let packed = self.sig.extract_p(p)?; + let out_slice = out.as_slice_mut()?; + if out_slice.len() != self.sig.output_len { + return Err(PyValueError::new_err(format!( + "{}: expected out of length {}, got {}", + self.sig.display_name(), + self.sig.output_len, + out_slice.len() + ))); + } + let yd = y_dot.as_ref().map(|a| a.as_slice()).transpose()?; + self.eval_inner(t, y.as_slice()?, &packed, yd, out_slice) + } + + /// Pack a {name: value} mapping into the stacked input layout. + fn pack<'py>( + &self, + py: Python<'py>, + mapping: &Bound<'_, PyDict>, + ) -> PyResult>> { + Ok(PyArray1::from_vec(py, self.sig.pack(mapping)?)) + } + + #[getter] + fn input_names<'py>(&self, py: Python<'py>) -> PyResult> { + pyo3::types::PyTuple::new(py, &self.sig.input_names) + } + /// Registered-name count (the `vp`/parameter-tangent seed length), NOT + /// the packed input width, unlike `ExprGraph::n_inputs`. + #[getter] + const fn n_inputs(&self) -> usize { + self.sig.input_names.len() + } + #[getter] + const fn n_states(&self) -> usize { + self.sig.n_states + } + #[getter] + const fn output_len(&self) -> usize { + self.sig.output_len + } + #[getter] + const fn uses_y_dot(&self) -> bool { + self.sig.uses_y_dot + } + #[getter] + fn name(&self) -> Option { + self.sig.name.clone() + } + /// Instruction count excluding conditional branch blocks: the common tape + /// plus one dispatch per conditional. Makes cross-output CSE observable and + /// is directly comparable to `casadi.Function.n_instructions()`. + #[getter] + fn n_instructions(&self) -> usize { + self.expr.ir().common_instruction_count() + } + /// Raw tape length, branch blocks included. + #[getter] + fn n_instructions_total(&self) -> usize { + self.expr.ir().instructions().len() + } + /// How many dispatches `n_instructions` includes, one per short-circuited + /// conditional. The part of the reported count that is control flow rather + /// than always-run work, which `branch_block_lens` cannot recover. + #[getter] + fn n_dispatches(&self) -> usize { + self.expr.ir().dispatch_count() + } + /// Per-branch block lengths, in tape order. + #[getter] + fn branch_block_lens<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + pyo3::types::PyTuple::new(py, self.expr.ir().branch_block_lens()) + } + + /// Forward-mode JVP: df/dy @ vy (+ df/dp @ vp when given). + #[pyo3(signature = (t, y, p, vy, vp = None))] + fn jvp<'py>( + &self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + vy: PyReadonlyArray1<'_, f64>, + vp: Option>, + ) -> PyResult>> { + // The tangent tape would slice an empty y_dot and panic; + // d/d(y_dot) seeding is a solver concern (cj). + self.sig.reject_y_dot("jvp")?; + let packed = self.sig.extract_p(p)?; + let y_slice = y.as_slice()?; + self.sig.check_y(y_slice.len())?; + let vy_slice = vy.as_slice()?; + if vy_slice.len() != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: expected vy of length {}, got {}", + self.sig.display_name(), + self.sig.n_states, + vy_slice.len() + ))); + } + + let ty = self.tangent_entry(py, false)?; + let mut scratch = ty.pool.acquire(); + let tangent = TangentInputs { + dy: Some(vy_slice), + dp: None, + }; + let mut out: Vec = ty + .expr + .eval_with_tangent(&mut scratch, t, y_slice, &[], &packed, &tangent) + .to_vec(); + ty.pool.release(scratch); + + if let Some(vp) = vp { + let vp_slice = vp.as_slice()?; + self.sig.check_vp(vp_slice.len())?; + let tp = self.tangent_entry(py, true)?; + let mut scratch_p = tp.pool.acquire(); + let tangent_p = TangentInputs { + dy: None, + dp: Some(vp_slice), + }; + let contrib = + tp.expr + .eval_with_tangent(&mut scratch_p, t, y_slice, &[], &packed, &tangent_p); + for (o, c) in out.iter_mut().zip(contrib) { + *o += *c; + } + tp.pool.release(scratch_p); + } + Ok(PyArray1::from_vec(py, out)) + } + + /// Lazy, cached-per-wrt prepared jacobian. + #[pyo3(signature = (wrt = "y"))] + fn jacobian(&self, py: Python<'_>, wrt: &str) -> PyResult> { + // Guard BEFORE prep: assembly evaluates with an empty y_dot, which the + // tape would slice and panic on. cj-weighted systems stay solver-side. + self.sig.reject_y_dot("jacobian")?; + let cell = match wrt { + "y" => &self.jac_y, + "p" => &self.jac_p, + other => { + return Err(PyValueError::new_err(format!( + "{}: wrt must be 'y' or 'p', got {:?}", + self.sig.display_name(), + other + ))); + }, + }; + cell.get_or_try_init(py, || { + let g = self.graph.try_borrow(py).map_err(|_| { + PyValueError::new_err(format!( + "{}: graph borrow conflict during jacobian derivation", + self.sig.display_name() + )) + })?; + let data = match wrt { + "y" => pybamm_core::JacobianData::new_wrt_states( + g.arena(), + self.root, + self.sig.output_len, + self.sig.n_states, + ), + _ => pybamm_core::JacobianData::new_wrt_params( + g.arena(), + self.root, + self.sig.output_len, + self.sig.input_names.len(), + ), + }; + drop(g); + Py::new( + py, + CompiledJacobian::build(py, Arc::new(data), self.sig.clone())?, + ) + }) + .map(|j| j.clone_ref(py)) + } + + /// Evaluate along a trajectory: one `extract_p` per sweep, GIL released + /// for the inner loop, scratch reused across columns. + fn eval_trajectory<'py>( + &self, + py: Python<'py>, + ts: PyArrayLike1<'_, f64, AllowTypeChange>, + y_traj: PyReadonlyArray2<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult>> { + let packed = self.sig.extract_p(p)?; + let ts_view = ts.as_array(); + let ts_slice = contiguous_slice_1d(&ts_view); + let view = y_traj.as_array(); + let (n_rows, n_t) = (view.shape()[0], view.shape()[1]); + if n_rows != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: Y.shape[0] must equal n_states ({}), got {}", + self.sig.display_name(), + self.sig.n_states, + n_rows + ))); + } + if ts_slice.len() != n_t { + return Err(PyValueError::new_err(format!( + "{}: len(ts) ({}) must equal Y.shape[1] ({})", + self.sig.display_name(), + ts_slice.len(), + n_t + ))); + } + self.sig.reject_y_dot("eval_trajectory")?; + + // Zero-copy when Y is already F-contiguous; one gathering copy otherwise. + let y_cols = columns_slice(&view); + + let out_len = self.sig.output_len; + let mut out = vec![0.0_f64; out_len * n_t]; + let expr = Arc::clone(&self.expr); + let scratch_len = expr.scratch_len(); + // Lane-batched tiling: ~1 MiB working set, 8..=64 lanes, capped at n_t. + let k_max = (TARGET_BYTES / (8 * scratch_len.max(1))).clamp(8, 64); + let k_alloc = k_max.min(n_t.max(1)); + let mut scratch = vec![0.0_f64; scratch_len * k_alloc]; + let batch_result = py.detach(|| -> Result<(), pybamm_core::BatchEvalError> { + let mut j0 = 0; + while j0 < n_t { + let k = (n_t - j0).min(k_max); + let ts_tile = &ts_slice[j0..j0 + k]; + let y_tile = &y_cols[j0 * n_rows..(j0 + k) * n_rows]; + let root = expr.eval_batch(&mut scratch, k, ts_tile, y_tile, &packed)?; + // root is (out_len, k) lane-minor: element e lane l at root[e*k + l]. + for l in 0..k { + let dst = &mut out[(j0 + l) * out_len..(j0 + l + 1) * out_len]; + for (e, o) in dst.iter_mut().enumerate() { + *o = root[e * k + l]; + } + } + j0 += k; + } + Ok(()) + }); + batch_result.map_err(|e| PyValueError::new_err(e.to_string()))?; + + // out is laid out column-major: (out_len, n_t) in F order. + // from_shape_vec is infallible here, out has exactly out_len * n_t elements. + let arr = numpy::ndarray::Array2::from_shape_vec((out_len, n_t).f(), out) + .map_err(|e| PyValueError::new_err(format!("eval_trajectory shape error: {e}")))?; + Ok(PyArray2::from_owned_array(py, arr)) + } + + /// Forward-mode JVP swept along a trajectory: for each time column `j`, + /// `df/dy(t_j, y_j) @ vy_j (+ df/dp(t_j, y_j) @ vp when given)`. + /// + /// `vy_traj` is `(n_states, n_t)`, one `yS` parameter-column over time; + /// `vp` is the constant parameter direction `e_k`, one entry per + /// registered parameter name (`n_inputs`, not the packed width). + /// Returns `(output_len, n_t)` F-contiguous, mirroring `eval_trajectory`: + /// one tangent-tape eval per column, scratch reused, GIL released. The + /// tangent tapes are the same lazily-derived, cached tapes `jvp` uses. + #[pyo3(signature = (ts, y_traj, p, vy_traj, vp = None))] + fn jvp_trajectory<'py>( + &self, + py: Python<'py>, + ts: PyArrayLike1<'_, f64, AllowTypeChange>, + y_traj: PyReadonlyArray2<'_, f64>, + p: &Bound<'_, PyAny>, + vy_traj: PyReadonlyArray2<'_, f64>, + vp: Option>, + ) -> PyResult>> { + // The tangent tape slices an empty y_dot and would panic; d/d(y_dot) + // seeding is a solver concern (cj), mirroring jvp/eval_trajectory. + self.sig.reject_y_dot("jvp_trajectory")?; + let packed = self.sig.extract_p(p)?; + let ts_view = ts.as_array(); + let ts_slice = contiguous_slice_1d(&ts_view); + let view = y_traj.as_array(); + let (n_rows, n_t) = (view.shape()[0], view.shape()[1]); + if n_rows != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: Y.shape[0] must equal n_states ({}), got {}", + self.sig.display_name(), + self.sig.n_states, + n_rows + ))); + } + if ts_slice.len() != n_t { + return Err(PyValueError::new_err(format!( + "{}: len(ts) ({}) must equal Y.shape[1] ({})", + self.sig.display_name(), + ts_slice.len(), + n_t + ))); + } + let vy_view = vy_traj.as_array(); + if vy_view.shape()[0] != self.sig.n_states || vy_view.shape()[1] != n_t { + return Err(PyValueError::new_err(format!( + "{}: vy_traj shape must be (n_states, n_t) = ({}, {}), got ({}, {})", + self.sig.display_name(), + self.sig.n_states, + n_t, + vy_view.shape()[0], + vy_view.shape()[1] + ))); + } + + // Zero-copy when F-contiguous; one gathering copy otherwise. + let y_cols = columns_slice(&view); + let vy_cols = columns_slice(&vy_view); + + // y-tangent tape (always); p-tangent tape only when vp is supplied. + let ty = self.tangent_entry(py, false)?; + let vp_slice: Option<&[f64]> = match &vp { + Some(arr) => { + let s = arr.as_slice()?; + self.sig.check_vp(s.len())?; + Some(s) + }, + None => None, + }; + let tp = match vp_slice { + Some(_) => Some(self.tangent_entry(py, true)?), + None => None, + }; + + let out_len = self.sig.output_len; + let mut out = vec![0.0_f64; out_len * n_t]; + let mut scratch_y = ty.pool.acquire(); + let mut scratch_p = tp.map(|t| t.pool.acquire()); + py.detach(|| { + for j in 0..n_t { + let y = &y_cols[j * n_rows..(j + 1) * n_rows]; + let vy = &vy_cols[j * n_rows..(j + 1) * n_rows]; + let dst = &mut out[j * out_len..(j + 1) * out_len]; + + let tangent_y = TangentInputs { + dy: Some(vy), + dp: None, + }; + let res_y = ty.expr.eval_with_tangent( + &mut scratch_y, + ts_slice[j], + y, + &[], + &packed, + &tangent_y, + ); + dst.copy_from_slice(res_y); + + if let (Some(tp), Some(scratch_p), Some(vp)) = (tp, scratch_p.as_mut(), vp_slice) { + let tangent_p = TangentInputs { + dy: None, + dp: Some(vp), + }; + let res_p = tp.expr.eval_with_tangent( + scratch_p, + ts_slice[j], + y, + &[], + &packed, + &tangent_p, + ); + for (o, c) in dst.iter_mut().zip(res_p) { + *o += *c; + } + } + } + }); + ty.pool.release(scratch_y); + if let (Some(tp), Some(scratch_p)) = (tp, scratch_p) { + tp.pool.release(scratch_p); + } + + // Column-major (out_len, n_t) in F order; from_shape_vec is infallible + // here, out has exactly out_len * n_t elements. + let arr = numpy::ndarray::Array2::from_shape_vec((out_len, n_t).f(), out) + .map_err(|e| PyValueError::new_err(format!("jvp_trajectory shape error: {e}")))?; + Ok(PyArray2::from_owned_array(py, arr)) + } + + /// Cubic-Hermite reconstruct the state at each `t_query` from the solver + /// knots (`ts`, `ys`, `yps`), then evaluate the compiled graph. Mirrors the + /// C++ `observe.cpp::observe_hermite_interp` math on the Rust evaluator; + /// returns `(output_len, n_query)` F-contiguous like `eval_trajectory`. + #[pyo3(signature = (t_query, ts, ys, yps, p))] + #[allow(clippy::suboptimal_flops)] // mixed-sign cubic-Hermite basis; mul_add would obscure it, matches eval.rs precedent + fn eval_trajectory_hermite<'py>( + &self, + py: Python<'py>, + t_query: PyArrayLike1<'_, f64, AllowTypeChange>, + ts: PyArrayLike1<'_, f64, AllowTypeChange>, + ys: PyReadonlyArray2<'_, f64>, + yps: PyReadonlyArray2<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult>> { + self.sig.reject_y_dot("eval_trajectory_hermite")?; + let packed = self.sig.extract_p(p)?; + let tq_view = t_query.as_array(); + let tq = contiguous_slice_1d(&tq_view); + let ts_view = ts.as_array(); + let ts_slice = contiguous_slice_1d(&ts_view); + let y_view = ys.as_array(); + let yp_view = yps.as_array(); + let (n_rows, n_knots) = (y_view.shape()[0], y_view.shape()[1]); + if n_rows != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: Y.shape[0] must equal n_states ({}), got {}", + self.sig.display_name(), + self.sig.n_states, + n_rows + ))); + } + if ts_slice.len() != n_knots { + return Err(PyValueError::new_err(format!( + "{}: len(ts) ({}) must equal Y.shape[1] ({})", + self.sig.display_name(), + ts_slice.len(), + n_knots + ))); + } + if yp_view.shape()[0] != n_rows || yp_view.shape()[1] != n_knots { + return Err(PyValueError::new_err(format!( + "{}: yps shape must equal ys shape ({}, {}), got ({}, {})", + self.sig.display_name(), + n_rows, + n_knots, + yp_view.shape()[0], + yp_view.shape()[1] + ))); + } + if n_knots < 2 { + return Err(PyValueError::new_err(format!( + "{}: need >= 2 knots for Hermite interpolation, got {}", + self.sig.display_name(), + n_knots + ))); + } + + // Zero-copy when F-contiguous; one gathering copy otherwise. + let y_cols = columns_slice(&y_view); + let yp_cols = columns_slice(&yp_view); + + let out_len = self.sig.output_len; + let n_query = tq.len(); + let mut out = vec![0.0_f64; out_len * n_query]; + let expr = Arc::clone(&self.expr); + let scratch_len = expr.scratch_len(); + // Lane-batched tiling: ~1 MiB working set, 8..=64 lanes, capped at n_query. + let k_max = (TARGET_BYTES / (8 * scratch_len.max(1))).clamp(8, 64); + let k_alloc = k_max.min(n_query.max(1)); + let mut scratch = vec![0.0_f64; scratch_len * k_alloc]; + // Reconstructed state columns for one tile, (n_states, k) F-contiguous. + let mut y_tile = vec![0.0_f64; n_rows * k_alloc]; + let batch_result = py.detach(|| -> Result<(), pybamm_core::BatchEvalError> { + let mut q0 = 0; + while q0 < n_query { + let k = (n_query - q0).min(k_max); + for l in 0..k { + let q = q0 + l; + let i = locate_interval(&ts_slice, tq[q]); + let h = ts_slice[i + 1] - ts_slice[i]; + let s = if h > 0.0 { + (tq[q] - ts_slice[i]) / h + } else { + 0.0 + }; + let (s2, s3) = (s * s, s * s * s); + // cubic-Hermite basis; derivative terms scaled by the step h + let h00 = 2.0 * s3 - 3.0 * s2 + 1.0; + let h10 = s3 - 2.0 * s2 + s; + let h01 = -2.0 * s3 + 3.0 * s2; + let h11 = s3 - s2; + let yi = &y_cols[i * n_rows..(i + 1) * n_rows]; + let yi1 = &y_cols[(i + 1) * n_rows..(i + 2) * n_rows]; + let ypi = &yp_cols[i * n_rows..(i + 1) * n_rows]; + let ypi1 = &yp_cols[(i + 1) * n_rows..(i + 2) * n_rows]; + let col = &mut y_tile[l * n_rows..(l + 1) * n_rows]; + for m in 0..n_rows { + col[m] = h00 * yi[m] + h10 * h * ypi[m] + h01 * yi1[m] + h11 * h * ypi1[m]; + } + } + let root = expr.eval_batch( + &mut scratch, + k, + &tq[q0..q0 + k], + &y_tile[..k * n_rows], + &packed, + )?; + // root is (out_len, k) lane-minor: element e lane l at root[e*k + l]. + for l in 0..k { + let dst = &mut out[(q0 + l) * out_len..(q0 + l + 1) * out_len]; + for (e, o) in dst.iter_mut().enumerate() { + *o = root[e * k + l]; + } + } + q0 += k; + } + Ok(()) + }); + batch_result.map_err(|e| PyValueError::new_err(e.to_string()))?; + + let arr = + numpy::ndarray::Array2::from_shape_vec((out_len, n_query).f(), out).map_err(|e| { + PyValueError::new_err(format!("eval_trajectory_hermite shape error: {e}")) + })?; + Ok(PyArray2::from_owned_array(py, arr)) + } + + fn __repr__(&self) -> String { + format!( + "CompiledFunction(name={:?}, inputs={:?}, n_states={}, output_len={}, uses_y_dot={})", + self.sig.display_name(), + self.sig.input_names, + self.sig.n_states, + self.sig.output_len, + self.sig.uses_y_dot + ) + } + + /// Rebuild from the retained `(graph, root)` derivation source (pickle + /// protocol). Recompiles the tape and resets the lazily-derived + /// tangent/jacobian caches; `n_states` is re-pinned exactly so a + /// user-widened system round-trips at the same width. + #[staticmethod] + fn _rebuild( + py: Python<'_>, + graph: Py, + root: u32, + name: Option, + n_states: Option, + ) -> PyResult { + let root = NodeId::from(root); + { + // Reachable from Python with an arbitrary root: validate before + // lowering, mirroring the `graph.compile` entry point. + let g = graph.borrow(py); + if root.index() >= g.arena().len() { + return Err(PyValueError::new_err(format!( + "_rebuild: root {} is out of range for a graph of {} nodes", + root.raw(), + g.arena().len() + ))); + } + crate::expr::check_supported(g.arena(), root)?; + } + Self::build(py, graph, root, name, n_states) + } + + /// `(callable, args)` pair for the pickle protocol: rebuild from the + /// retained graph and root rather than serializing derived state. + fn __reduce__<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyAny>, RebuildArgs)> { + let this = slf.get(); + let rebuild = slf.get_type().getattr("_rebuild")?; + Ok(( + rebuild, + ( + this.graph.clone_ref(py), + this.root.raw(), + this.name(), + Some(this.n_states()), + ), + )) + } +} + +/// Column-major slice of Y: zero-copy when Y is F-contiguous (strides[0] == 1), +/// otherwise at most one gathering copy. +pub fn columns_slice<'a>( + view: &'a numpy::ndarray::ArrayView2<'a, f64>, +) -> std::borrow::Cow<'a, [f64]> { + match view.as_slice_memory_order() { + // strides[0] == 1 and all strides non-negative ⇒ true F-contiguous column-major + Some(s) if view.strides()[0] == 1 && view.strides().iter().all(|&st| st >= 0) => { + std::borrow::Cow::Borrowed(s) + }, + _ => { + let (n_rows, n_t) = (view.shape()[0], view.shape()[1]); + let mut owned = Vec::with_capacity(n_rows * n_t); + for j in 0..n_t { + owned.extend(view.column(j).iter()); + } + std::borrow::Cow::Owned(owned) + }, + } +} + +/// Contiguous view of a 1-D time array: zero-copy when already contiguous, +/// otherwise one gathering copy. `PyArrayLike1` only +/// materialises a fresh copy on dtype mismatch or non-ndarray input; an +/// already-f64 but strided ndarray passes through unmaterialized, so +/// `.as_slice()` alone would still reject it. +pub fn contiguous_slice_1d<'a>( + view: &'a numpy::ndarray::ArrayView1<'a, f64>, +) -> std::borrow::Cow<'a, [f64]> { + view.as_slice().map_or_else( + || std::borrow::Cow::Owned(view.to_vec()), + std::borrow::Cow::Borrowed, + ) +} + +/// Bracketing interval index for `x` in ascending `ts`, clamped to `[0, n-2]` +/// (extends the boundary segment). `ts.len() >= 2` is enforced by the caller. +pub fn locate_interval(ts: &[f64], x: f64) -> usize { + let nseg = ts.len() - 1; + if x <= ts[0] { + return 0; + } + if x >= ts[nseg] { + return nseg - 1; + } + let (mut lo, mut hi) = (0usize, nseg); + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if ts[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + lo +} + +/// Compile prep = simplify + instruction tape: same pass order as +/// `CompiledModel::new`, run on a clone so the retained graph stays +/// the unsimplified derivation source for jvp / jacobian (which re-run +/// the pipeline on their own tangent arenas). Without this, user-compiled +/// functions get systematically worse tapes than the bundle's, and group +/// CSE only catches shared `NodeIds`, not structurally identical nodes +/// from separate builds. +pub fn compile_simplified(arena: &Arena, root: NodeId) -> CompiledExpr { + let (da, root) = simplify_pipeline(arena.clone(), root); + CompiledExpr::from_ir(TypedIr::from_arena(&da, root)) +} + +/// Resolve the user-supplied `n_states` override against the scanned +/// extent: an override below it would admit a too-short `y` and panic +/// inside the tape. +pub fn resolve_n_states( + scanned: usize, + requested: Option, + name: Option<&str>, +) -> PyResult { + match requested { + Some(n) if n < scanned => Err(PyValueError::new_err(format!( + "{}: n_states={n} is below the expression's state extent {scanned}", + name.unwrap_or(""), + ))), + Some(n) => Ok(n), + None => Ok(scanned), + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/group.rs b/packages/pybamm-rust/pybamm-python/src/group.rs new file mode 100644 index 0000000000..37cd344ab4 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/group.rs @@ -0,0 +1,404 @@ +//! `CompiledFunctionGroup`: shared-tape multi-output artifact. + +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::Arc; + +use numpy::ndarray::ShapeBuilder; +use numpy::{ + AllowTypeChange, PyArray1, PyArray2, PyArrayLike1, PyReadonlyArray1, PyReadonlyArray2, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +use pybamm_core::{CompiledExpr, Node, NodeId, TypedIr, scan_state_usage, simplify_pipeline}; + +use crate::expr::ExprGraph; +use crate::scratch::{Buffer, ScratchPool}; +use crate::signature::FunctionSignature; + +/// A named set of outputs compiled into ONE shared tape. Cross-output +/// common subexpressions are evaluated once; per-output results are recovered via +/// recorded `(offset, len)` slices into the synthetic concat root. +#[pyclass(frozen, module = "pybamm.rust")] +pub struct CompiledFunctionGroup { + expr: Arc, + sig: FunctionSignature, + pool: ScratchPool, + out_names: Vec, + /// Per-output (offset, len) into the concat result. + slices: Vec<(usize, usize)>, +} + +impl std::fmt::Debug for CompiledFunctionGroup { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledFunctionGroup") + .field("sig", &self.sig) + .field("out_names", &self.out_names) + .finish_non_exhaustive() + } +} + +impl CompiledFunctionGroup { + pub(crate) fn build( + py: Python<'_>, + graph: Py, + names: Vec, + ids: Vec, + name: Option, + n_states_override: Option, + ) -> PyResult { + let g = graph + .try_borrow(py) + .map_err(|_| PyValueError::new_err("compile_group: graph borrow conflict"))?; + // Per-output lengths from per-child IRs (build-time only). + let lens: Vec = ids + .iter() + .map(|&id| TypedIr::from_arena(g.arena(), id).output_len()) + .collect(); + // The synthetic concat root lives on a clone: compiling a group + // must not grow the caller's graph. + let mut work = g.arena().clone(); + let input_names = g.input_names(); + let input_widths = g.input_widths(); + drop(g); + let root = work.alloc(Node::Concat(ids)); + // Scan the ORIGINAL root (pre-simplify): the signature is the + // user-declared contract. + let usage = scan_state_usage(&work, root); + // Running the pipeline over the combined reachable set gives cross-output + // CSE for structurally-identical nodes, not just shared NodeIds. + let (work, root) = simplify_pipeline(work, root); + let expr = Arc::new(CompiledExpr::from_ir(TypedIr::from_arena(&work, root))); + + let mut slices = Vec::with_capacity(lens.len()); + let mut offset = 0; + for &len in &lens { + slices.push((offset, len)); + offset += len; + } + if offset != expr.output_len() { + return Err(PyValueError::new_err(format!( + "compile_group: internal error: output slices tile {} values \ + but the compiled tape produces {}", + offset, + expr.output_len() + ))); + } + let sig = FunctionSignature { + input_names, + input_widths, + n_states: crate::function::resolve_n_states( + usage.n_states, + n_states_override, + name.as_deref(), + )?, + uses_y_dot: usage.uses_y_dot, + output_len: expr.output_len(), + name, + }; + let pool = ScratchPool::new(Buffer(expr.scratch_len())); + Ok(Self { + expr, + sig, + pool, + out_names: names, + slices, + }) + } + + fn eval_columns( + &self, + py: Python<'_>, + ts: &[f64], + y_cols: &[f64], + n_rows: usize, + packed: &[f64], + ) -> Vec { + let n_t = ts.len(); + let total = self.sig.output_len; + let mut out = vec![0.0; total * n_t]; + let expr = &self.expr; + let mut scratch = self.pool.acquire(); + py.detach(|| { + for j in 0..n_t { + let y = &y_cols[j * n_rows..(j + 1) * n_rows]; + let dst = &mut out[j * total..(j + 1) * total]; + expr.eval_into(&mut scratch, ts[j], y, &[], packed, dst); + } + }); + self.pool.release(scratch); + out + } +} + +#[pymethods] +impl CompiledFunctionGroup { + /// Evaluate all outputs once over the shared tape; returns a list of + /// arrays in declared order. + fn __call__<'py>( + &self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult>>> { + // The eval paths slice an empty y_dot, which the tape would index and panic + // on. Build stays allowed so signature introspection still works. + self.sig.reject_y_dot("group eval")?; + let packed = self.sig.extract_p(p)?; + let y_slice = y.as_slice()?; + self.sig.check_y(y_slice.len())?; + + let mut full = vec![0.0; self.sig.output_len]; + let mut scratch = self.pool.acquire(); + self.expr + .eval_into(&mut scratch, t, y_slice, &[], &packed, &mut full); + self.pool.release(scratch); + + Ok(self + .slices + .iter() + .map(|&(off, len)| PyArray1::from_slice(py, &full[off..off + len])) + .collect()) + } + + /// Trajectory sweep over the shared tape: one crossing, one tape eval + /// per column regardless of output count. + fn eval_trajectory<'py>( + &self, + py: Python<'py>, + ts: PyArrayLike1<'_, f64, AllowTypeChange>, + y_traj: PyReadonlyArray2<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult>>> { + self.sig.reject_y_dot("group eval")?; + let packed = self.sig.extract_p(p)?; + let ts_view = ts.as_array(); + let ts_slice = crate::function::contiguous_slice_1d(&ts_view); + let view = y_traj.as_array(); + let (n_rows, n_t) = (view.shape()[0], view.shape()[1]); + if n_rows != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: Y.shape[0] must equal n_states ({}), got {}", + self.sig.display_name(), + self.sig.n_states, + n_rows + ))); + } + if ts_slice.len() != n_t { + return Err(PyValueError::new_err(format!( + "{}: len(ts) ({}) must equal Y.shape[1] ({})", + self.sig.display_name(), + ts_slice.len(), + n_t + ))); + } + let y_cols = crate::function::columns_slice(&view); + + let flat = self.eval_columns(py, &ts_slice, &y_cols, n_rows, &packed); + let total = self.sig.output_len; + self.slices + .iter() + .map(|&(off, len)| { + let mut data = Vec::with_capacity(len * n_t); + for j in 0..n_t { + data.extend_from_slice(&flat[j * total + off..j * total + off + len]); + } + let arr = + numpy::ndarray::Array2::from_shape_vec((len, n_t).f(), data).map_err(|e| { + PyValueError::new_err(format!("eval_trajectory shape error: {e}")) + })?; + Ok(PyArray2::from_owned_array(py, arr)) + }) + .collect() + } + + /// Cubic-Hermite reconstruct the state at each `t_query` from the solver + /// knots, then evaluate the shared tape once per query and slice per + /// output. API parity with `CompiledFunction::eval_trajectory_hermite`; + /// full-state observation routes through per-variable `CompiledFunction`, + /// not this group method. + #[pyo3(signature = (t_query, ts, ys, yps, p))] + #[allow(clippy::suboptimal_flops)] // mixed-sign cubic-Hermite basis, mirrors function.rs + fn eval_trajectory_hermite<'py>( + &self, + py: Python<'py>, + t_query: PyArrayLike1<'_, f64, AllowTypeChange>, + ts: PyArrayLike1<'_, f64, AllowTypeChange>, + ys: PyReadonlyArray2<'_, f64>, + yps: PyReadonlyArray2<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult>>> { + self.sig.reject_y_dot("group eval_trajectory_hermite")?; + let packed = self.sig.extract_p(p)?; + let tq_view = t_query.as_array(); + let tq = crate::function::contiguous_slice_1d(&tq_view); + let ts_view = ts.as_array(); + let ts_slice = crate::function::contiguous_slice_1d(&ts_view); + let y_view = ys.as_array(); + let yp_view = yps.as_array(); + let (n_rows, n_knots) = (y_view.shape()[0], y_view.shape()[1]); + if n_rows != self.sig.n_states { + return Err(PyValueError::new_err(format!( + "{}: Y.shape[0] must equal n_states ({}), got {}", + self.sig.display_name(), + self.sig.n_states, + n_rows + ))); + } + if ts_slice.len() != n_knots { + return Err(PyValueError::new_err(format!( + "{}: len(ts) ({}) must equal Y.shape[1] ({})", + self.sig.display_name(), + ts_slice.len(), + n_knots + ))); + } + if yp_view.shape()[0] != n_rows || yp_view.shape()[1] != n_knots { + return Err(PyValueError::new_err(format!( + "{}: yps shape must equal ys shape ({}, {}), got ({}, {})", + self.sig.display_name(), + n_rows, + n_knots, + yp_view.shape()[0], + yp_view.shape()[1] + ))); + } + if n_knots < 2 { + return Err(PyValueError::new_err(format!( + "{}: need >= 2 knots for Hermite interpolation, got {}", + self.sig.display_name(), + n_knots + ))); + } + + // Zero-copy when Y/YP are already F-contiguous; one gathering copy otherwise. + let y_cols = crate::function::columns_slice(&y_view); + let yp_cols = crate::function::columns_slice(&yp_view); + + let total = self.sig.output_len; + let n_query = tq.len(); + let mut flat = vec![0.0_f64; total * n_query]; + let expr = &self.expr; + let mut scratch = self.pool.acquire(); + let mut y_interp = vec![0.0_f64; n_rows]; + py.detach(|| { + for q in 0..n_query { + let i = crate::function::locate_interval(&ts_slice, tq[q]); + let h = ts_slice[i + 1] - ts_slice[i]; + let s = if h > 0.0 { + (tq[q] - ts_slice[i]) / h + } else { + 0.0 + }; + let (s2, s3) = (s * s, s * s * s); + // cubic-Hermite basis; derivative terms scaled by the step h + let h00 = 2.0 * s3 - 3.0 * s2 + 1.0; + let h10 = s3 - 2.0 * s2 + s; + let h01 = -2.0 * s3 + 3.0 * s2; + let h11 = s3 - s2; + let yi = &y_cols[i * n_rows..(i + 1) * n_rows]; + let yi1 = &y_cols[(i + 1) * n_rows..(i + 2) * n_rows]; + let ypi = &yp_cols[i * n_rows..(i + 1) * n_rows]; + let ypi1 = &yp_cols[(i + 1) * n_rows..(i + 2) * n_rows]; + for m in 0..n_rows { + y_interp[m] = h00 * yi[m] + h10 * h * ypi[m] + h01 * yi1[m] + h11 * h * ypi1[m]; + } + let dst = &mut flat[q * total..(q + 1) * total]; + expr.eval_into(&mut scratch, tq[q], &y_interp, &[], &packed, dst); + } + }); + self.pool.release(scratch); + + self.slices + .iter() + .map(|&(off, len)| { + let mut data = Vec::with_capacity(len * n_query); + for q in 0..n_query { + data.extend_from_slice(&flat[q * total + off..q * total + off + len]); + } + let arr = numpy::ndarray::Array2::from_shape_vec((len, n_query).f(), data) + .map_err(|e| { + PyValueError::new_err(format!("eval_trajectory_hermite shape error: {e}")) + })?; + Ok(PyArray2::from_owned_array(py, arr)) + }) + .collect() + } + + #[getter] + fn names<'py>(&self, py: Python<'py>) -> PyResult> { + pyo3::types::PyTuple::new(py, &self.out_names) + } + #[getter] + fn output_lens(&self) -> Vec { + self.slices.iter().map(|&(_, len)| len).collect() + } + /// Pack a {name: value} mapping into the stacked input layout + /// (same signature surface as `CompiledFunction`). + fn pack<'py>( + &self, + py: Python<'py>, + mapping: &Bound<'_, pyo3::types::PyDict>, + ) -> PyResult>> { + Ok(PyArray1::from_vec(py, self.sig.pack(mapping)?)) + } + + #[getter] + fn input_names<'py>(&self, py: Python<'py>) -> PyResult> { + pyo3::types::PyTuple::new(py, &self.sig.input_names) + } + /// Registered-name count (the `vp`/parameter-tangent seed length), NOT + /// the packed input width, unlike `ExprGraph::n_inputs`. + #[getter] + const fn n_inputs(&self) -> usize { + self.sig.input_names.len() + } + #[getter] + const fn n_states(&self) -> usize { + self.sig.n_states + } + #[getter] + const fn output_len(&self) -> usize { + self.sig.output_len + } + #[getter] + const fn uses_y_dot(&self) -> bool { + self.sig.uses_y_dot + } + /// Instruction count excluding conditional branch blocks: the common tape + /// plus one dispatch per conditional. Makes cross-output CSE observable and + /// is directly comparable to `casadi.Function.n_instructions()`. + #[getter] + fn n_instructions(&self) -> usize { + self.expr.ir().common_instruction_count() + } + /// Raw tape length, branch blocks included. + #[getter] + fn n_instructions_total(&self) -> usize { + self.expr.ir().instructions().len() + } + /// Per-branch block lengths, in tape order. + #[getter] + fn branch_block_lens<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + pyo3::types::PyTuple::new(py, self.expr.ir().branch_block_lens()) + } + #[getter] + fn name(&self) -> Option { + self.sig.name.clone() + } + + fn __repr__(&self) -> String { + format!( + "CompiledFunctionGroup(name={:?}, outputs={:?}, n_states={})", + self.sig.display_name(), + self.out_names, + self.sig.n_states + ) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/jacobian.rs b/packages/pybamm-rust/pybamm-python/src/jacobian.rs new file mode 100644 index 0000000000..ca0f0f7857 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/jacobian.rs @@ -0,0 +1,159 @@ +//! `CompiledJacobian`: the prepared derivative artifact. + +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::Arc; + +use numpy::{PyArray1, PyReadonlyArray1}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::IntoPyDict; + +use pybamm_core::JacobianData; + +use crate::scratch::ScratchPool; +use crate::signature::FunctionSignature; + +#[pyclass(frozen, module = "pybamm.rust")] +pub struct CompiledJacobian { + pub(crate) data: Arc, + pub(crate) sig: FunctionSignature, + scratch_pool: ScratchPool>, + /// CSC index arrays, built once: **int32** so scipy adopts + /// them without a cast-copy, and **read-only** so in-place + /// canonicalisation of one returned matrix (`sort_indices`, + /// `sum_duplicates`) raises instead of corrupting the pattern + /// shared by every other matrix. Each call allocates only the + /// data array. + indptr: Py>, + indices: Py>, +} + +impl std::fmt::Debug for CompiledJacobian { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledJacobian") + .field("sig", &self.sig) + .field("wrt", &self.data.wrt()) + .field("shape", &(self.data.n_rows(), self.data.n_cols())) + .field("nnz", &self.data.nnz()) + .finish_non_exhaustive() + } +} + +fn to_readonly_i32(py: Python<'_>, values: &[usize]) -> PyResult>> { + let v: Vec = values + .iter() + .map(|&x| { + i32::try_from(x).map_err(|_| { + PyValueError::new_err("jacobian pattern exceeds the int32 index range") + }) + }) + .collect::>()?; + let arr = PyArray1::from_vec(py, v); + arr.call_method("setflags", (), Some(&[("write", false)].into_py_dict(py)?))?; + Ok(arr.unbind()) +} + +impl CompiledJacobian { + pub(crate) fn build( + py: Python<'_>, + data: Arc, + sig: FunctionSignature, + ) -> PyResult { + Ok(Self { + scratch_pool: ScratchPool::new(Arc::clone(&data)), + indptr: to_readonly_i32(py, &data.csc().colptr)?, + indices: to_readonly_i32(py, &data.csc().rowind)?, + data, + sig, + }) + } +} + +#[pymethods] +impl CompiledJacobian { + /// Assemble and return a `scipy.sparse.csc_matrix`. Per-call work is + /// `n_colors` JVP sweeps + linear scatter; only the data array allocates. + fn __call__<'py>( + &self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + p: &Bound<'_, PyAny>, + ) -> PyResult> { + let packed = self.sig.extract_p(p)?; + let y_slice = y.as_slice()?; + self.sig.check_y(y_slice.len())?; + + let mut values = vec![0.0; self.data.nnz()]; + let mut scratch = self.scratch_pool.acquire(); + let data = &self.data; + py.detach(|| { + data.assemble_into( + &mut scratch, + data.layout(), + t, + y_slice, + &[], + &packed, + &mut values, + ); + }); + self.scratch_pool.release(scratch); + + let data_arr = PyArray1::from_vec(py, values); + let csc = py.import("scipy.sparse")?.getattr("csc_matrix")?.call1(( + (data_arr, self.indices.bind(py), self.indptr.bind(py)), + (self.data.n_rows(), self.data.n_cols()), + ))?; + Ok(csc) + } + + /// CSC pattern as (indptr, indices), the cached read-only arrays. + fn sparsity<'py>( + &self, + py: Python<'py>, + ) -> (Bound<'py, PyArray1>, Bound<'py, PyArray1>) { + (self.indptr.bind(py).clone(), self.indices.bind(py).clone()) + } + + #[getter] + fn nnz(&self) -> usize { + self.data.nnz() + } + #[getter] + fn shape(&self) -> (usize, usize) { + (self.data.n_rows(), self.data.n_cols()) + } + #[getter] + fn n_colors(&self) -> usize { + self.data.n_colors() + } + /// Number of dense rows split out of the column coloring; they share one + /// reverse-mode tape, seeded once per row. + #[getter] + fn n_dense_rows(&self) -> usize { + self.data.n_dense_rows() + } + #[getter] + fn wrt(&self) -> &'static str { + match self.data.wrt() { + pybamm_core::DiffTarget::States => "y", + pybamm_core::DiffTarget::Params => "p", + } + } + + fn __repr__(&self) -> String { + format!( + "CompiledJacobian(of={:?}, wrt={:?}, shape=({}, {}), nnz={}, n_colors={}, n_dense_rows={})", + self.sig.display_name(), + self.wrt(), + self.data.n_rows(), + self.data.n_cols(), + self.data.nnz(), + self.data.n_colors(), + self.data.n_dense_rows() + ) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/lib.rs b/packages/pybamm-rust/pybamm-python/src/lib.rs new file mode 100644 index 0000000000..8757ae93a8 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/lib.rs @@ -0,0 +1,53 @@ +//! `PyO3` bindings that expose `pybamm-core` to Python as `pybamm.rust._core`. +//! +//! Python builds a model by allocating nodes into an `ExprGraph`, then asks for +//! the artifact it needs: a `CompiledFunction` or `CompiledFunctionGroup` for +//! plain evaluation, a `CompiledJacobian` for derivatives, a `CompiledModel` +//! for the DAE the IDAKLU bridge drives, or a `PreparedSolver` to integrate in +//! Rust. Each names the core type it is the Python face of; core's per-solve +//! `ModelEvaluator` is reached through `CompiledModel.evaluator_pool`. +//! Each artifact is prepared once and owns pooled scratch, so a call reuses +//! that scratch rather than re-deriving it; the packed parameter vector and the +//! returned array are still allocated per call. +//! +//! Bindings own the boundary checks: array lengths and index ranges are validated +//! here so core code can assume them, and a core `CoreError` is mapped to +//! `ValueError` or `RuntimeError` rather than escaping as a panic. + +use pyo3::prelude::*; + +mod errors; +mod evaluator_pool; +mod expr; +mod function; +mod group; +mod jacobian; +mod model; +#[cfg(feature = "diffsol")] +mod pool; +mod scratch; +mod signature; +#[cfg(feature = "diffsol")] +mod solver; + +#[pymodule] +fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + #[cfg(feature = "diffsol")] + { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(solver::default_solver_options, m)?)?; + m.add_function(wrap_pyfunction!(pool::_pool_ids, m)?)?; + } + + Ok(()) +} diff --git a/packages/pybamm-rust/pybamm-python/src/model.rs b/packages/pybamm-rust/pybamm-python/src/model.rs new file mode 100644 index 0000000000..798a627db9 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/model.rs @@ -0,0 +1,890 @@ +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::Arc; + +use numpy::{PyArray1, PyReadonlyArray1, PyReadwriteArray1}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyDict; + +use pybamm_core::model as core_model; +use pybamm_core::{CsrData, NodeId, ObservableKind, Shape}; + +use crate::errors::core_err_to_py; +use crate::evaluator_pool::EvaluatorPool; +use crate::expr::{Expr, ExprGraph}; +use crate::function::CompiledFunction; +use crate::jacobian::CompiledJacobian; +use crate::signature::FunctionSignature; + +/// `(graph, rhs_root, output_roots, event_roots, algebraic_root, algebraic_variable_indices, +/// mass_data, mass_indptr, mass_indices, n_inputs, sens_param_indices)`, the +/// `__reduce__`/`_rebuild` argument tuple for `CompiledModel`'s pickle protocol. +type RebuildArgs = ( + Py, + u32, + Vec, + Vec, + Option, + Vec, + Vec, + Vec, + Vec, + usize, + Vec, +); + +/// The Python face of [`core_model::CompiledModel`]: the same immutable +/// artifact behind the same `Arc`, plus the retained derivation graph and the +/// bundle accessors (`rhs`, `jacobian`, `outputs`, `events`, +/// `algebraic_residual`, `algebraic_jacobian`) composed over the same tapes. +/// +/// Holding the artifact rather than an evaluator is what keeps it shareable; +/// solvers take their per-solve state from `evaluator_pool`. +// `module` is required so pickle can locate the class as `pybamm.rust.CompiledModel` +// instead of the pyo3 default `builtins.CompiledModel`, which pickle cannot import. +#[pyclass(module = "pybamm.rust")] +pub struct CompiledModel { + pub(crate) compiled: Arc, + /// Scratch for the direct `eval_*`/`assemble_*` helpers, bound on first + /// use, so a model that is only lowered and handed on never allocates one. + scratch: Option, + /// Retained derivation source for the bundle views (jacobian / jvp / the + /// algebraic-subset jacobian all re-run the AD pipeline on this arena). + graph: Py, + rhs_root: NodeId, + output_roots: Vec, + event_roots: Vec, + algebraic_root: Option, + /// Strictly-ascending global state indices of the algebraic block. + algebraic_variable_indices: Vec, + // Cached views: every accessor hands back the SAME prepared artifact, so pools, + // lazy tangent tapes and sparsity/coloring prep amortise across accesses. + rhs_view: PyOnceLock>, + jac_view: PyOnceLock>, + outputs_view: PyOnceLock>>, + events_view: PyOnceLock>>, + algebraic_residual_view: PyOnceLock>>, + algebraic_jacobian_view: PyOnceLock>>, +} + +impl std::fmt::Debug for CompiledModel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledModel") + .field("n_states", &self.compiled.n_states()) + .field("n_inputs", &self.compiled.n_params()) + .field("output_len", &self.compiled.output_len()) + .field("n_colors", &self.compiled.coloring().n_colors) + .finish_non_exhaustive() + } +} + +impl CompiledModel { + /// Run `body` against the artifact and a [`core_model::Workspace`] bound on + /// first use. Only the direct `eval_*`/`assemble_*` helpers need one. + fn with_scratch( + &mut self, + body: impl FnOnce(&core_model::CompiledModel, &mut core_model::Workspace) -> R, + ) -> R { + let Self { + compiled, scratch, .. + } = self; + body( + compiled, + scratch.get_or_insert_with(|| compiled.create_workspace()), + ) + } + + /// One `CompiledFunction` view per observable of `kind`, named `label[i]`. + /// + /// `roots` are this family's retained arena roots, in the order they were + /// compiled, so a view carries the node it came from as well as its tape. + fn observable_views( + &self, + py: Python<'_>, + kind: ObservableKind, + roots: &[NodeId], + label: &str, + ) -> PyResult>> { + let set = self.compiled.observables(kind); + (0..set.count()) + .map(|i| { + Py::new( + py, + CompiledFunction::from_shared( + py, + set.expr_arc(i), + self.graph.clone_ref(py), + roots[i], + self.compiled.n_states(), + Some(format!("{label}[{i}]")), + )?, + ) + }) + .collect() + } + + /// Build a [`FunctionSignature`] for a bundle view over the retained graph. + /// + /// Bundle views always carry the full system width and never use `y_dot` + /// (the residual's `cj`/mass coupling stays solver-side). + fn view_signature( + &self, + py: Python<'_>, + output_len: usize, + name: Option, + ) -> PyResult { + let g = self.graph.try_borrow(py).map_err(|_| { + PyValueError::new_err("CompiledModel: graph borrow conflict building view signature") + })?; + Ok(FunctionSignature { + input_names: g.input_names(), + input_widths: g.input_widths(), + n_states: self.compiled.n_states(), + uses_y_dot: false, + output_len, + name, + }) + } + + /// Shared builder over already-resolved roots, the raw CSR mass matrix and + /// the optional artifacts. + /// + /// Both `from_expr` (roots from validated `Expr`s) and `_rebuild` (roots + /// bounds-checked against the arena) funnel through here, so the + /// `check_supported` gate below runs exactly once per public entry point. + #[allow(clippy::too_many_arguments)] + fn build_from_parts( + py: Python<'_>, + graph: Py, + rhs_root: NodeId, + output_roots: Vec, + event_roots: Vec, + algebraic_root: Option, + algebraic_variable_indices: Vec, + mass_data: Vec, + mass_indptr: Vec, + mass_indices: Vec, + n_inputs: usize, + sens_param_indices: Vec, + ) -> PyResult { + // Infer n_states from indptr length + if mass_indptr.is_empty() { + return Err(PyValueError::new_err( + "mass_indptr must have at least 1 element", + )); + } + let n = mass_indptr.len() - 1; + + // scipy CSR arrays arrive as i64; convert with a bounds check so a + // negative entry becomes a clear error rather than a wrapped `usize`. + let indptr: Vec = mass_indptr + .iter() + .map(|&x| usize::try_from(x)) + .collect::>() + .map_err(|_| PyValueError::new_err("mass_indptr entries must be non-negative"))?; + let indices: Vec = mass_indices + .iter() + .map(|&x| usize::try_from(x)) + .collect::>() + .map_err(|_| PyValueError::new_err("mass_indices entries must be non-negative"))?; + + let mass = CsrData::try_new(indptr, indices, mass_data, Shape::matrix(n, n)) + .map_err(core_err_to_py)?; + + // `new_wrt_state_subset` asserts a strictly-ascending subset; callers pass + // a contiguous range, but normalise so the invariant holds regardless. + let mut algebraic_variable_indices = algebraic_variable_indices; + algebraic_variable_indices.sort_unstable(); + algebraic_variable_indices.dedup(); + + let g = graph.try_borrow(py).map_err(|_| { + PyValueError::new_err("CompiledModel.build_from_parts: graph borrow conflict") + })?; + let arena = g.arena(); + + // Validate every root before lowering so unsupported nodes or invalid + // shape relationships surface as catchable Python errors. + crate::expr::check_supported(arena, rhs_root)?; + for &root in &output_roots { + crate::expr::check_supported(arena, root)?; + } + for &root in &event_roots { + crate::expr::check_supported(arena, root)?; + } + if let Some(root) = algebraic_root { + crate::expr::check_supported(arena, root)?; + } + + let options = algebraic_root.map_or_else( + || core_model::CompiledModelOptions::new().with_sensitivities(&sens_param_indices), + |algebraic| { + core_model::CompiledModelOptions::new() + .with_sensitivities(&sens_param_indices) + .with_algebraic(core_model::CompiledModelAlgebraicBlock::new( + algebraic, + &algebraic_variable_indices, + )) + }, + ); + + // Composed before the `Arc`, so appending an output or event is a plain + // `&mut self` call rather than a copy-on-write plus workspace rebuild. + let mut compiled = core_model::CompiledModel::new_with_options( + arena, rhs_root, mass, n, n_inputs, options, + ); + + // Output-variable nodes must come from the same arena as the rhs, which + // holds because `PyBaMM` builds both through one `ExprGraph`. + for &root in &output_roots { + compiled.add_output(arena, root); + } + + // Compile each event expression for root-finding during integration. + for &root in &event_roots { + compiled.add_event(arena, root); + } + drop(g); + + // Fuse the events into one tape so both hot loops evaluate shared event + // subgraphs once. Needs a mutable arena borrow to alloc the `Concat` root. + if event_roots.len() >= 2 { + let mut g = graph.try_borrow_mut(py).map_err(|_| { + PyValueError::new_err("CompiledModel.build_from_parts: graph borrow conflict") + })?; + compiled.fuse_events(g.arena_mut(), &event_roots); + } + + Ok(Self { + compiled: Arc::new(compiled), + scratch: None, + graph, + rhs_root, + output_roots, + event_roots, + algebraic_root, + algebraic_variable_indices, + rhs_view: PyOnceLock::new(), + jac_view: PyOnceLock::new(), + outputs_view: PyOnceLock::new(), + events_view: PyOnceLock::new(), + algebraic_residual_view: PyOnceLock::new(), + algebraic_jacobian_view: PyOnceLock::new(), + }) + } +} + +#[pymethods] +impl CompiledModel { + /// Create a compiled model from an expression graph. + /// + /// # Arguments + /// + /// * `graph` - The expression graph containing the RHS expression + /// * `expr` - The root expression node (f(t, y)) + /// * `mass_data` - Non-zero values of mass matrix (CSR data array) + /// * `mass_indptr` - CSR indptr array (length `n_states` + 1) + /// * `mass_indices` - CSR column indices array + /// * `n_inputs` - Number of input parameters (default 0) + /// + /// # Returns + /// + /// A new `CompiledModel` ready for evaluation + #[staticmethod] + #[pyo3(signature = ( + graph, + expr, + mass_data, + mass_indptr, + mass_indices, + n_inputs = 0, + sens_param_indices = vec![], + output_exprs = vec![], + algebraic_expr = None, + algebraic_variable_indices = vec![], + event_exprs = vec![], + ))] + #[allow(clippy::too_many_arguments)] // PyO3 keyword args, all distinct concerns + fn from_expr( + py: Python<'_>, + graph: Py, + expr: &Expr, + mass_data: PyReadonlyArray1<'_, f64>, + mass_indptr: PyReadonlyArray1<'_, i64>, + mass_indices: PyReadonlyArray1<'_, i64>, + n_inputs: usize, + sens_param_indices: Vec, + output_exprs: Vec>, + algebraic_expr: Option<&Expr>, + algebraic_variable_indices: Vec, + event_exprs: Vec>, + ) -> PyResult { + // numpy -> Vec conversions; `build_from_parts` owns the CSR assembly + // so it is shared with the `_rebuild` pickle path. + let mass_data: Vec = mass_data.as_slice()?.to_vec(); + let mass_indptr: Vec = mass_indptr.as_slice()?.to_vec(); + let mass_indices: Vec = mass_indices.as_slice()?.to_vec(); + + // Expr -> NodeId conversions: the bundle retains the roots so the view + // accessors can compose prepared artifacts over the shared tapes. + let rhs_root = expr.node_id_in(&graph)?; + let output_roots: Vec = output_exprs + .iter() + .map(|e| e.node_id_in(&graph)) + .collect::>()?; + let event_roots: Vec = event_exprs + .iter() + .map(|e| e.node_id_in(&graph)) + .collect::>()?; + let algebraic_root = algebraic_expr + .as_ref() + .map(|e| e.node_id_in(&graph)) + .transpose()?; + + Self::build_from_parts( + py, + graph, + rhs_root, + output_roots, + event_roots, + algebraic_root, + algebraic_variable_indices, + mass_data, + mass_indptr, + mass_indices, + n_inputs, + sens_param_indices, + ) + } + + /// Rebuild from the retained `(graph, roots, mass CSR, options)` + /// derivation source (pickle protocol). Recompiles every tape and resets + /// the lazily-derived bundle-view caches. + #[staticmethod] + #[allow(clippy::too_many_arguments)] + fn _rebuild( + py: Python<'_>, + graph: Py, + rhs_root: u32, + output_roots: Vec, + event_roots: Vec, + algebraic_root: Option, + algebraic_variable_indices: Vec, + mass_data: Vec, + mass_indptr: Vec, + mass_indices: Vec, + n_inputs: usize, + sens_param_indices: Vec, + ) -> PyResult { + let rhs_root = NodeId::from(rhs_root); + let output_roots: Vec = output_roots.into_iter().map(NodeId::from).collect(); + let event_roots: Vec = event_roots.into_iter().map(NodeId::from).collect(); + let algebraic_root = algebraic_root.map(NodeId::from); + + { + // Reachable from Python with arbitrary ids: bounds-check every + // root before `build_from_parts` indexes the arena with it. + let g = graph.borrow(py); + let n_nodes = g.arena().len(); + let check_bounds = |root: NodeId| -> PyResult<()> { + if root.index() >= n_nodes { + return Err(PyValueError::new_err(format!( + "_rebuild: root {} is out of range for a graph of {n_nodes} nodes", + root.raw(), + ))); + } + Ok(()) + }; + check_bounds(rhs_root)?; + for &root in &output_roots { + check_bounds(root)?; + } + for &root in &event_roots { + check_bounds(root)?; + } + if let Some(root) = algebraic_root { + check_bounds(root)?; + } + } + + Self::build_from_parts( + py, + graph, + rhs_root, + output_roots, + event_roots, + algebraic_root, + algebraic_variable_indices, + mass_data, + mass_indptr, + mass_indices, + n_inputs, + sens_param_indices, + ) + } + + /// `(callable, args)` pair for the pickle protocol: rebuild from the + /// retained graph, roots, and mass CSR rather than serializing derived + /// evaluator/coloring state. + fn __reduce__<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyAny>, RebuildArgs)> { + // `CompiledModel` is not frozen (it has `&mut self` eval methods), + // so `Bound::get` is unavailable; borrow the wrapped value instead. + let this = slf.borrow(); + let rebuild = slf.get_type().getattr("_rebuild")?; + let mass = this.compiled.mass_matrix(); + // Mirrors `mass_indptr`/`mass_indices` back to the i64 width they + // arrived in from scipy via `from_expr`. + #[allow(clippy::cast_possible_wrap)] + let mass_indptr: Vec = mass.indptr().iter().map(|&x| x as i64).collect(); + #[allow(clippy::cast_possible_wrap)] + let mass_indices: Vec = mass.indices().iter().map(|&x| x as i64).collect(); + Ok(( + rebuild, + ( + this.graph.clone_ref(py), + this.rhs_root.raw(), + this.output_roots.iter().map(|r| r.raw()).collect(), + this.event_roots.iter().map(|r| r.raw()).collect(), + this.algebraic_root.map(NodeId::raw), + this.algebraic_variable_indices.clone(), + mass.data().to_vec(), + mass_indptr, + mass_indices, + this.compiled.n_params(), + this.compiled.sens_param_indices().to_vec(), + ), + )) + } + + // Each view is built ONCE from the bundle's shared `Arc`s and cached, so + // repeated access is a `clone_ref` and the prepared state amortises. + + /// Primal f(t, y, p) as a shareable `CompiledFunction` view. + #[getter] + fn rhs(&self, py: Python<'_>) -> PyResult> { + self.rhs_view + .get_or_try_init(py, || { + Py::new( + py, + CompiledFunction::from_shared( + py, + self.compiled.primal_expr_arc(), + self.graph.clone_ref(py), + self.rhs_root, + self.compiled.n_states(), + Some("rhs".to_string()), + )?, + ) + }) + .map(|f| f.clone_ref(py)) + } + + /// The retained derivation arena. Observation lowers a new root into it + /// (`symbol.to_rust(graph)` + `graph.compile`) so the observed expression + /// shares the solve's input-parameter and state indices by construction. + #[getter] + fn graph(&self, py: Python<'_>) -> Py { + self.graph.clone_ref(py) + } + + /// Pure df/dy (cj = 0, no mass): the bundle's composed `JacobianData`. + /// + /// Note: `model.rhs.jacobian()` re-derives its own artifact from the + /// graph; prefer this composed one (zero extra prep). + #[getter] + fn jacobian(&self, py: Python<'_>) -> PyResult> { + self.jac_view + .get_or_try_init(py, || { + let sig = + self.view_signature(py, self.compiled.output_len(), Some("rhs".to_string()))?; + Py::new( + py, + CompiledJacobian::build(py, self.compiled.jacobian_data(), sig)?, + ) + }) + .map(|j| j.clone_ref(py)) + } + + /// Output-variable expressions as shareable `CompiledFunction` views. + #[getter] + fn outputs(&self, py: Python<'_>) -> PyResult>> { + let views = self.outputs_view.get_or_try_init(py, || { + self.observable_views(py, ObservableKind::Outputs, &self.output_roots, "output") + })?; + Ok(views.iter().map(|f| f.clone_ref(py)).collect()) + } + + /// Event expressions as shareable `CompiledFunction` views. + #[getter] + fn events(&self, py: Python<'_>) -> PyResult>> { + let views = self.events_view.get_or_try_init(py, || { + self.observable_views(py, ObservableKind::Events, &self.event_roots, "event") + })?; + Ok(views.iter().map(|f| f.clone_ref(py)).collect()) + } + + /// Algebraic residual g(t, y, p) as a shareable view, or `None` for ODEs. + #[getter] + fn algebraic_residual(&self, py: Python<'_>) -> PyResult>> { + let cached = self.algebraic_residual_view.get_or_try_init(py, || { + let Some(expr) = self.compiled.algebraic_expr_arc() else { + return PyResult::Ok(None); + }; + let root = self.algebraic_root.ok_or_else(|| { + PyValueError::new_err("algebraic_expr implies algebraic_root: internal invariant") + })?; + Ok(Some(Py::new( + py, + CompiledFunction::from_shared( + py, + expr, + self.graph.clone_ref(py), + root, + self.compiled.n_states(), + Some("algebraic_residual".to_string()), + )?, + )?)) + })?; + Ok(cached.as_ref().map(|f| f.clone_ref(py))) + } + + /// `dg/dy_alg` as a standalone prepared jacobian (`n_algebraic` x `n_algebraic`), or + /// `None` for ODEs. + /// + /// A view onto the artifact the model already compiled, as `algebraic_residual` + /// is of the residual, so it costs a shared handle rather than a second + /// tangent transform, sparsity detection and colouring of the same expression. + #[getter] + fn algebraic_jacobian(&self, py: Python<'_>) -> PyResult>> { + let cached = self.algebraic_jacobian_view.get_or_try_init(py, || { + let Some(data) = self.compiled.algebraic_jacobian_data() else { + return PyResult::Ok(None); + }; + let sig = + self.view_signature(py, data.n_rows(), Some("algebraic_residual".to_string()))?; + Ok(Some(Py::new(py, CompiledJacobian::build(py, data, sig)?)?)) + })?; + Ok(cached.as_ref().map(|j| j.clone_ref(py))) + } + + /// Compute residual r = M*y' - f(t,y) and return as a new array. + /// + /// This is the DAE residual function used by IDAKLU and other DAE solvers. + /// For performance-critical code, use `residual_into` to avoid allocation. + /// + /// # Arguments + /// + /// * `t` - Time value + /// * `y` - State vector + /// * `yp` - Time derivative of state vector y' + /// * `inputs` - Input parameters (can be empty array) + /// + /// # Returns + /// + /// The residual as a numpy array + fn eval_residual<'py>( + &mut self, + py: Python<'py>, + t: f64, + y: PyReadonlyArray1<'_, f64>, + yp: PyReadonlyArray1<'_, f64>, + inputs: PyReadonlyArray1<'_, f64>, + ) -> PyResult>> { + let y_slice = y.as_slice()?; + let yp_slice = yp.as_slice()?; + let inputs_slice = inputs.as_slice()?; + let mut output = vec![0.0; self.compiled.output_len()]; + + self.with_scratch(|compiled, ws| { + compiled.eval_residual(ws, t, y_slice, yp_slice, inputs_slice, &mut output); + }); + + Ok(PyArray1::from_vec(py, output)) + } + + /// Compute residual r = M*y' - f(t,y) into a pre-allocated output array. + /// + /// This avoids allocation overhead and is preferred for solver hot paths. + /// + /// # Arguments + /// + /// * `t` - Time value + /// * `y` - State vector + /// * `yp` - Time derivative of state vector y' + /// * `inputs` - Input parameters (can be empty array) + /// * `output` - Pre-allocated output array (length `n_states`) + fn eval_residual_into( + &mut self, + t: f64, + y: PyReadonlyArray1<'_, f64>, + yp: PyReadonlyArray1<'_, f64>, + inputs: PyReadonlyArray1<'_, f64>, + mut output: PyReadwriteArray1<'_, f64>, + ) -> PyResult<()> { + let y_slice = y.as_slice()?; + let yp_slice = yp.as_slice()?; + let inputs_slice = inputs.as_slice()?; + let output_slice = output.as_slice_mut()?; + + if output_slice.len() < self.compiled.n_states() { + return Err(PyValueError::new_err(format!( + "output array too small: need {} elements, got {}", + self.compiled.n_states(), + output_slice.len() + ))); + } + + self.with_scratch(|compiled, ws| { + compiled.eval_residual(ws, t, y_slice, yp_slice, inputs_slice, output_slice); + }); + Ok(()) + } + + /// Get the number of states in the model. + #[getter] + fn n_states(&self) -> usize { + self.compiled.n_states() + } + + /// Get the number of input parameters in the model. + /// + /// Core calls these `n_params`; we expose them as `n_inputs` at the Python + /// boundary to match the C ABI naming. + #[getter] + fn n_inputs(&self) -> usize { + self.compiled.n_params() + } + + /// Get algebraic-state IDs as a `(n_states,)` numpy array using IDA's + /// convention: `1.0` for differential states, `0.0` for algebraic states. + fn algebraic_ids<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + let n = self.compiled.n_states(); + let mut buf = vec![0.0; n]; + self.compiled.algebraic_ids_f64(&mut buf); + PyArray1::from_vec(py, buf) + } + + /// Whether the model has algebraic sub-block expressions. + #[getter] + fn has_algebraic(&self) -> bool { + self.compiled.has_algebraic() + } + + /// Number of algebraic states in the compiled sub-block. + #[getter] + fn n_algebraic(&self) -> usize { + self.compiled.n_algebraic() + } + + /// Number of non-zeros in the assembled algebraic Jacobian. + #[getter] + fn algebraic_jacobian_nnz(&self) -> usize { + self.compiled.algebraic_jacobian_nnz() + } + + /// Get the algebraic Jacobian sparsity pattern as COO `(rows, cols)`. + #[allow(clippy::type_complexity)] + fn algebraic_jacobian_sparsity_pattern<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + if !self.compiled.has_algebraic() { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "Model has no algebraic block", + )); + } + Ok(( + PyArray1::from_vec(py, self.compiled.algebraic_jacobian_row_indices().to_vec()), + PyArray1::from_vec(py, self.compiled.algebraic_jacobian_col_indices().to_vec()), + )) + } + + /// Number of forward-sensitivity parameters configured. + #[getter] + fn n_sens_params(&self) -> usize { + self.compiled.n_sens_params() + } + + /// Number of compiled output-variable expressions. + #[getter] + fn n_outputs(&self) -> usize { + self.compiled.n_outputs() + } + + /// Number of compiled event expressions. + #[getter] + fn n_events(&self) -> usize { + self.compiled.n_events() + } + + /// Get the output length of f(t, y). + #[getter] + fn output_len(&self) -> usize { + self.compiled.output_len() + } + + /// Get the sparsity pattern of df/dy as (indptr, indices). + /// + /// Returns the sparsity pattern in CSR format for use with + /// sparse matrix construction. + #[allow(clippy::type_complexity)] + fn sparsity_pattern<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + let pattern = self.compiled.sparsity(); + let indptr = PyArray1::from_vec(py, pattern.indptr.clone()); + let indices = PyArray1::from_vec(py, pattern.indices.clone()); + Ok((indptr, indices)) + } + + /// Number of colors in the Jacobian sparsity pattern's graph coloring, each + /// filling one or more columns in a single sweep. This is the reduced coloring + /// when a dense-row split was adopted, so it need not cover every row. + #[getter] + fn n_colors(&self) -> usize { + self.compiled.coloring().n_colors + } + + /// Assemble the Jacobian into a pre-allocated CSC data buffer. + /// + /// This is the zero-allocation version for performance-critical code. + /// The caller must pre-allocate `jac_data` with length `nnz`. + /// + /// # Arguments + /// + /// * `t` - Time value + /// * `y` - State vector + /// * `cj` - Jacobian coefficient from solver (for J = df/dy - cj*M) + /// * `inputs` - Input parameters (can be empty array) + /// * `jac_data` - Pre-allocated output buffer in CSC order (length `nnz`) + fn assemble_jacobian_csc_into( + &mut self, + t: f64, + y: PyReadonlyArray1<'_, f64>, + cj: f64, + inputs: PyReadonlyArray1<'_, f64>, + mut jac_data: PyReadwriteArray1<'_, f64>, + ) -> PyResult<()> { + let y_slice = y.as_slice()?; + let inputs_slice = inputs.as_slice()?; + let jac_slice = jac_data.as_slice_mut()?; + + let nnz = self.compiled.nnz(); + if jac_slice.len() < nnz { + return Err(PyValueError::new_err(format!( + "jac_data buffer too small: need {nnz} elements, got {}", + jac_slice.len() + ))); + } + + self.with_scratch(|compiled, ws| { + ws.set_cj(cj); + compiled.assemble_jacobian_csc_into(ws, t, y_slice, inputs_slice, jac_slice); + }); + Ok(()) + } + + /// Get the number of non-zeros in the Jacobian. + #[getter] + #[allow(clippy::missing_const_for_fn)] // PyO3 getters can't be const + fn nnz(&self) -> usize { + self.compiled.nnz() + } + + /// Jacobian assembly strategy name this model compiled to. + #[getter] + fn jacobian_strategy(&self) -> &'static str { + self.compiled.jacobian_strategy().as_str() + } + + /// Jacobian assembly stats for benchmarking and debug attribution. + fn jacobian_stats<'py>(&self, py: Python<'py>) -> PyResult> { + let stats = self.compiled.jacobian_stats(); + let dict = PyDict::new(py); + dict.set_item("strategy", stats.strategy.as_str())?; + dict.set_item("n_colors", stats.n_colors)?; + dict.set_item("nnz", stats.nnz)?; + dict.set_item("n_dense_rows", stats.n_dense_rows)?; + dict.set_item("n_dense_row_candidates", stats.n_dense_row_candidates)?; + dict.set_item("n_constant_entries", stats.n_constant_entries)?; + dict.set_item("n_swept_columns", stats.n_swept_columns)?; + dict.set_item("jac_lane_width", stats.jac_lane_width)?; + dict.set_item("dense_row_entries", stats.dense_row_entries)?; + dict.set_item( + "dense_row_tape_instructions", + stats.dense_row_tape_instructions, + )?; + // primal_end is an index into the raw tape, not an instruction count, + // so it is left as a raw quantity here. + dict.set_item( + "split_eval_primal_instructions", + stats.split_eval_primal_instructions, + )?; + dict.set_item( + "split_eval_total_instructions", + stats.split_eval_total_instructions, + )?; + dict.set_item( + "split_eval_raw_instructions", + stats.split_eval_raw_instructions, + )?; + dict.set_item("split_eval_dispatch_count", stats.split_eval_dispatch_count)?; + dict.set_item( + "branch_block_lens", + pyo3::types::PyTuple::new(py, &stats.branch_block_lens)?, + )?; + Ok(dict) + } + + /// `(csc_idx, value)` for every entry proved constant at compile time. + /// + /// Indices address the same buffer `assemble_jacobian_csc_into` fills, so + /// these are the slots no colour sweep writes. + #[allow(clippy::type_complexity)] + fn constant_jacobian_entries<'py>( + &self, + py: Python<'py>, + ) -> (Bound<'py, PyArray1>, Bound<'py, PyArray1>) { + let entries = self.compiled.constant_jacobian_entries(); + let (indices, values): (Vec, Vec) = entries.iter().copied().unzip(); + ( + PyArray1::from_vec(py, indices), + PyArray1::from_vec(py, values), + ) + } + + /// Get the CSC sparsity pattern for KLU compatibility. + /// + /// Returns `(colptr, rowind)` where: + /// - `colptr` has length `n_states + 1` + /// - `rowind` has length `nnz` + #[allow(clippy::type_complexity)] + fn csc_sparsity_pattern<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + let csc = self.compiled.csc_sparsity(); + let colptr = PyArray1::from_vec(py, csc.colptr.clone()); + let rowind = PyArray1::from_vec(py, csc.rowind.clone()); + Ok((colptr, rowind)) + } + + /// `n` independent evaluators over this model's tape, one per parallel solver. + /// + /// Rejects `n == 0` rather than handing back an empty pool: the caller would + /// then build a solver group with no solvers and divide by its size. + fn evaluator_pool(&self, n: usize) -> PyResult { + if n == 0 { + return Err(PyValueError::new_err( + "evaluator_pool needs at least one evaluator, got 0", + )); + } + Ok(EvaluatorPool::from_compiled(&self.compiled, n)) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/pool.rs b/packages/pybamm-rust/pybamm-python/src/pool.rs new file mode 100644 index 0000000000..bda58db187 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/pool.rs @@ -0,0 +1,79 @@ +//! Process-wide rayon pools, one per distinct thread count. +//! +//! Pool ownership sits in the binding rather than in `pybamm-core`: a numerics +//! library should not own process-wide OS threads, while a binding is the layer +//! that knows "one process, many solver objects". Non-Python consumers of the +//! core install their own pool and lose nothing. +//! +//! One solver per experiment step is the common case, so pools are keyed by +//! count rather than owned per solver: ten `DiffsolSolver`s asking for 8 threads +//! share one pool of 8, not eighty threads. Idle rayon workers park on a condvar +//! after a short spin, so a cached pool costs stack address space and no CPU. +//! +//! `ThreadPool::install` re-entered from inside the same pool runs inline, so +//! nested batch calls cannot deadlock. + +use std::collections::HashMap; +use std::sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError}; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use rayon::{ThreadPool, ThreadPoolBuilder}; + +/// Populated lazily, so a process that never asks for more than one thread +/// never constructs rayon at all. +static POOLS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Lock the pool cache, recovering from a poisoned lock. +/// +/// A panic inside a locked section can only have left the `HashMap` mid-insert, +/// which no reader can observe as corrupt, so poisoning is not a reason to fail +/// every later solve. +fn lock_cache() -> MutexGuard<'static, HashMap>> { + POOLS.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// The shared pool of `threads` workers, building it on first request. +/// +/// Workers are named `pybamm-solve-` so `top -H`, `perf` and py-spy identify +/// them. The build runs under the lock, so two solvers constructed concurrently +/// share one pool rather than racing to create two. +pub fn pool_for(threads: usize) -> PyResult> { + let mut pools = lock_cache(); + if let Some(pool) = pools.get(&threads) { + return Ok(Arc::clone(pool)); + } + let pool = Arc::new( + ThreadPoolBuilder::new() + .num_threads(threads) + .thread_name(|i| format!("pybamm-solve-{i}")) + .build() + .map_err(|e| { + PyRuntimeError::new_err(format!( + "could not build a rayon pool of {threads} thread(s): {e}" + )) + })?, + ); + pools.insert(threads, Arc::clone(&pool)); + // Explicit, so the guard's scope stays tight enough for clippy. + drop(pools); + Ok(pool) +} + +/// The cached pools as `{thread_count: pool identity}`. +/// +/// Introspection for the tests that pin the caching rules — that solvers asking +/// for the same width share one pool, and that a default-configured process +/// builds none. +#[pyfunction] +pub fn _pool_ids(py: Python<'_>) -> PyResult> { + let dict = PyDict::new(py); + // Integer keys run no arbitrary Python, so set_item cannot re-enter + // pool_for and the cache stays locked for the whole build. + for (threads, pool) in lock_cache().iter() { + dict.set_item(threads, Arc::as_ptr(pool) as usize)?; + } + Ok(dict.unbind()) +} diff --git a/packages/pybamm-rust/pybamm-python/src/scratch.rs b/packages/pybamm-rust/pybamm-python/src/scratch.rs new file mode 100644 index 0000000000..80d379d075 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/scratch.rs @@ -0,0 +1,115 @@ +use std::sync::Mutex; + +use pybamm_core::{JacobianData, JacobianScratch}; + +/// How a [`ScratchPool`] makes a buffer when it has none to hand out. +/// +/// The recipe rather than a spare buffer: a pool that minted by cloning a +/// prototype would both retain one extra set of buffers for its lifetime and +/// memcpy them on every miss, where the shape alone is what a caller needs. +pub trait Mint { + /// What this recipe produces. + type Item; + + /// A fresh item, with every slot zeroed. + fn mint(&self) -> Self::Item; +} + +/// Scratch for one `CompiledExpr`, named by the buffer length it needs. +#[derive(Debug)] +pub struct Buffer(pub usize); + +impl Mint for Buffer { + type Item = Box<[f64]>; + + fn mint(&self) -> Self::Item { + vec![0.0; self.0].into_boxed_slice() + } +} + +impl Mint for std::sync::Arc { + type Item = JacobianScratch; + + fn mint(&self) -> Self::Item { + JacobianScratch::new(self) + } +} + +/// Reuse pool for the per-call scratch an evaluation needs. +/// +/// `PyO3` methods take `&self`, so scratch cannot live in the object; pooling it +/// keeps a steady-state call from allocating. Concurrent calls that find the pool +/// locked mint their own rather than block. +#[derive(Debug)] +pub struct ScratchPool { + buffers: Mutex>, + mint: M, +} + +impl ScratchPool { + /// Retained-buffer cap: bounds steady-state memory under thread churn. + pub const MAX_POOLED: usize = 8; + + pub const fn new(mint: M) -> Self { + Self { + buffers: Mutex::new(Vec::new()), + mint, + } + } + + /// Pooled buffers may carry stale data from a prior evaluation; + /// callers must overwrite every slot before reading (the tape in + /// `CompiledExpr::eval` satisfies this — every slot is written + /// before any downstream read). + pub fn acquire(&self) -> M::Item { + if let Ok(mut pool) = self.buffers.try_lock() + && let Some(buf) = pool.pop() + { + return buf; + } + self.mint.mint() + } + + pub fn release(&self, buf: M::Item) { + if let Ok(mut pool) = self.buffers.try_lock() + && pool.len() < Self::MAX_POOLED + { + pool.push(buf); + } + } + + #[cfg(test)] + pub fn pooled_len(&self) -> usize { + self.buffers.lock().map_or(0, |p| p.len()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn acquire_release_reuses_buffer() { + let pool = ScratchPool::new(Buffer(16)); + let a = pool.acquire(); + assert_eq!(a.len(), 16); + let ptr = a.as_ptr(); + pool.release(a); + let b = pool.acquire(); + // ownership moves into the pool and back out — same allocation, + // so pointer identity is guaranteed (not allocator-dependent) + assert_eq!(b.as_ptr(), ptr, "steady state must reuse the buffer"); + } + + #[test] + fn pool_caps_retained_buffers() { + let pool = ScratchPool::new(Buffer(4)); + let bufs: Vec<_> = (0..(ScratchPool::::MAX_POOLED + 4)) + .map(|_| pool.acquire()) + .collect(); + for b in bufs { + pool.release(b); + } + assert!(pool.pooled_len() <= ScratchPool::::MAX_POOLED); + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/signature.rs b/packages/pybamm-rust/pybamm-python/src/signature.rs new file mode 100644 index 0000000000..1c98328b91 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/signature.rs @@ -0,0 +1,212 @@ +use numpy::PyReadonlyArray1; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +#[derive(Debug, Clone)] +pub struct FunctionSignature { + /// Input names ordered by registration index (== stacked layout). + pub input_names: Vec, + /// Packed width per input (>1 for vector-valued inputs), parallel to `input_names`. + pub input_widths: Vec, + pub n_states: usize, + pub uses_y_dot: bool, + pub output_len: usize, + pub name: Option, +} + +impl FunctionSignature { + pub fn display_name(&self) -> &str { + self.name.as_deref().unwrap_or("") + } + + /// Total packed width across all inputs (sum of `input_widths`). + fn total_width(&self) -> usize { + self.input_widths.iter().sum() + } + + /// Coerce a single input value to a scalar `f64`. Accepts a real number + /// or a length-1 numeric array. + fn extract_scalar(&self, v: &Bound<'_, PyAny>, name: &str) -> PyResult { + if let Ok(x) = v.extract::() { + return Ok(x); + } + let bad = || { + PyValueError::new_err(format!( + "{}: input '{}' must be a real scalar or length-1 numeric array", + self.display_name(), + name + )) + }; + let size: usize = v + .getattr("size") + .and_then(|s| s.extract::()) + .map_err(|_| bad())?; + if size != 1 { + return Err(PyValueError::new_err(format!( + "{}: input '{}' must be a scalar or length-1 array, got size {}", + self.display_name(), + name, + size + ))); + } + v.call_method0("item")?.extract::().map_err(|_| bad()) + } + + /// Coerce a single input value to a flat `width`-length `Vec`. + /// Accepts any array shape whose flattened size matches `width` + /// (e.g. the column-vector convention `np.array([[..]])` used elsewhere). + fn extract_vector(&self, v: &Bound<'_, PyAny>, name: &str, width: usize) -> PyResult> { + let bad_shape = |size: usize| { + PyValueError::new_err(format!( + "{}: input '{}' must have {} value(s), got {}", + self.display_name(), + name, + width, + size + )) + }; + let bad_type = || { + PyValueError::new_err(format!( + "{}: input '{}' must be a numeric array of length {}", + self.display_name(), + name, + width + )) + }; + let raveled = v.call_method0("ravel").map_err(|_| bad_type())?; + let arr: PyReadonlyArray1<'_, f64> = raveled.extract().map_err(|_| bad_type())?; + let slice = arr.as_slice()?; + if slice.len() != width { + return Err(bad_shape(slice.len())); + } + Ok(slice.to_vec()) + } + + /// Pack a {name: value} mapping into the stacked layout. Validates + /// keys both ways, names offenders, and checks each value's length + /// against the parameter's declared width. + pub fn pack(&self, dict: &Bound<'_, PyDict>) -> PyResult> { + let mut packed = vec![0.0; self.total_width()]; + let mut offset = 0usize; + for (i, name) in self.input_names.iter().enumerate() { + let width = self.input_widths[i]; + match dict.get_item(name)? { + Some(v) => { + if width == 1 { + packed[offset] = self.extract_scalar(&v, name)?; + } else { + let values = self.extract_vector(&v, name, width)?; + packed[offset..offset + width].copy_from_slice(&values); + } + }, + None => { + return Err(PyValueError::new_err(format!( + "{}: missing input '{}'; expected inputs {:?}", + self.display_name(), + name, + self.input_names + ))); + }, + } + offset += width; + } + if dict.len() != self.input_names.len() { + for key in dict.keys() { + let k: String = key.extract()?; + if !self.input_names.iter().any(|n| n == &k) { + return Err(PyValueError::new_err(format!( + "{}: unknown input '{}'; expected inputs {:?}", + self.display_name(), + k, + self.input_names + ))); + } + } + } + Ok(packed) + } + + /// `p` as a stacked slice copy, packing a {name: value} mapping when given + /// one. The single home for the dict-vs-array branch on every call path. + pub fn extract_p(&self, p: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(dict) = p.cast::() { + return self.pack(dict); + } + let arr: PyReadonlyArray1<'_, f64> = p.extract().map_err(|_| { + PyValueError::new_err(format!( + "{}: p must be a 1-D float64 array or a {{name: value}} mapping", + self.display_name() + )) + })?; + let slice = arr.as_slice()?; + self.check_p(slice.len())?; + Ok(slice.to_vec()) + } + + pub fn check_y(&self, len: usize) -> PyResult<()> { + if len != self.n_states { + return Err(PyValueError::new_err(format!( + "{}: expected y of length {}, got {}", + self.display_name(), + self.n_states, + len + ))); + } + Ok(()) + } + + pub fn check_p(&self, len: usize) -> PyResult<()> { + let total_width = self.total_width(); + if len != total_width { + return Err(PyValueError::new_err(format!( + "{}: expected {} input values ({} parameters), got {}", + self.display_name(), + total_width, + self.input_names.len(), + len + ))); + } + Ok(()) + } + + /// Validate a parameter-direction vector (`vp`/`dp` seed): one scalar + /// direction per registered parameter name (registration-index space, + /// what `TangentParameter` reads), NOT the packed total width. + pub fn check_vp(&self, len: usize) -> PyResult<()> { + if len != self.input_names.len() { + return Err(PyValueError::new_err(format!( + "{}: expected vp of length {} (one direction per parameter), got {}", + self.display_name(), + self.input_names.len(), + len + ))); + } + Ok(()) + } + + pub fn check_y_dot(&self, len: usize) -> PyResult<()> { + if len != self.n_states { + return Err(PyValueError::new_err(format!( + "{}: expected y_dot of length {}, got {}", + self.display_name(), + self.n_states, + len + ))); + } + Ok(()) + } + + /// Reject `y_dot`-using expressions for operations that evaluate with + /// an empty one, the tape would slice it and panic across `PyO3`. + pub fn reject_y_dot(&self, op: &str) -> PyResult<()> { + if self.uses_y_dot { + return Err(PyValueError::new_err(format!( + "{}: {} is not defined for expressions that use y_dot", + self.display_name(), + op + ))); + } + Ok(()) + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/solver.rs b/packages/pybamm-rust/pybamm-python/src/solver.rs new file mode 100644 index 0000000000..445c1b5cb8 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/solver.rs @@ -0,0 +1,603 @@ +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::Arc; + +use numpy::ndarray::{Array2, ShapeBuilder}; +use numpy::{ + IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods, +}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use pybamm_core::CoreError; +use pybamm_core::solver::SolverOptions; +use pybamm_core::solver::batch::check_batch_widths; +// Core and binding share one name per concept, so the boundary crossing is +// spelled out at every use site rather than hidden in a mutated noun. +use pybamm_core::solver::solve as core_solve; + +use crate::errors::core_err_to_py; +use crate::model::CompiledModel; +use crate::pool::pool_for; + +/// Integrator tuning as `PyBaMM`'s `options` dict, extracted item by item. +/// +/// Every field is required: the Python solver owns the defaults and always +/// sends a complete dict, so a key missing here is a wiring bug rather than +/// something to paper over with a fallback. +#[derive(Debug, FromPyObject)] +#[pyo3(from_item_all)] +pub struct PySolverOptions { + max_nonlinear_solver_iterations: usize, + max_error_test_failures: usize, + max_nonlinear_solver_failures: usize, + nonlinear_solver_tolerance: f64, + min_timestep: f64, + max_timestep_growth: Option, + min_timestep_growth: Option, + max_timestep_shrink: Option, + min_timestep_shrink: Option, + update_jacobian_after_steps: usize, + update_rhs_jacobian_after_steps: usize, + threshold_to_update_jacobian: f64, + threshold_to_update_rhs_jacobian: f64, + pi_control_proportional: f64, + pi_control_integral: f64, +} + +/// The integrator defaults, as the `options` dict `PyBaMM` overlays onto. +/// +/// Exposed so `pybamm.DiffsolSolver.DEFAULT_OPTIONS` can be pinned against +/// diffsol's own defaults rather than trusted to stay a faithful hand-copy. +#[pyfunction] +pub fn default_solver_options(py: Python<'_>) -> PyResult> { + let defaults = SolverOptions::default(); + let dict = PyDict::new(py); + dict.set_item( + "max_nonlinear_solver_iterations", + defaults.max_nonlinear_solver_iterations, + )?; + dict.set_item("max_error_test_failures", defaults.max_error_test_failures)?; + dict.set_item( + "max_nonlinear_solver_failures", + defaults.max_nonlinear_solver_failures, + )?; + dict.set_item( + "nonlinear_solver_tolerance", + defaults.nonlinear_solver_tolerance, + )?; + dict.set_item("min_timestep", defaults.min_timestep)?; + dict.set_item("max_timestep_growth", defaults.max_timestep_growth)?; + dict.set_item("min_timestep_growth", defaults.min_timestep_growth)?; + dict.set_item("max_timestep_shrink", defaults.max_timestep_shrink)?; + dict.set_item("min_timestep_shrink", defaults.min_timestep_shrink)?; + dict.set_item( + "update_jacobian_after_steps", + defaults.update_jacobian_after_steps, + )?; + dict.set_item( + "update_rhs_jacobian_after_steps", + defaults.update_rhs_jacobian_after_steps, + )?; + dict.set_item( + "threshold_to_update_jacobian", + defaults.threshold_to_update_jacobian, + )?; + dict.set_item( + "threshold_to_update_rhs_jacobian", + defaults.threshold_to_update_rhs_jacobian, + )?; + dict.set_item("pi_control_proportional", defaults.pi_control_proportional)?; + dict.set_item("pi_control_integral", defaults.pi_control_integral)?; + Ok(dict.unbind()) +} + +impl From for SolverOptions { + fn from(options: PySolverOptions) -> Self { + Self { + max_nonlinear_solver_iterations: options.max_nonlinear_solver_iterations, + max_error_test_failures: options.max_error_test_failures, + max_nonlinear_solver_failures: options.max_nonlinear_solver_failures, + nonlinear_solver_tolerance: options.nonlinear_solver_tolerance, + min_timestep: options.min_timestep, + max_timestep_growth: options.max_timestep_growth, + min_timestep_growth: options.min_timestep_growth, + max_timestep_shrink: options.max_timestep_shrink, + min_timestep_shrink: options.min_timestep_shrink, + update_jacobian_after_steps: options.update_jacobian_after_steps, + update_rhs_jacobian_after_steps: options.update_rhs_jacobian_after_steps, + threshold_to_update_jacobian: options.threshold_to_update_jacobian, + threshold_to_update_rhs_jacobian: options.threshold_to_update_rhs_jacobian, + pi_control_proportional: options.pi_control_proportional, + pi_control_integral: options.pi_control_integral, + } + } +} + +/// BDF solver statistics exposed to Python. +#[derive(Debug, Clone)] +// Output-only type: opt out of the (now deprecated) automatic FromPyObject derive. +#[pyclass(skip_from_py_object, module = "pybamm.rust")] +pub struct SolverStatistics { + #[pyo3(get)] + number_of_steps: usize, + #[pyo3(get)] + number_of_linear_solver_setups: usize, + #[pyo3(get)] + number_of_nonlinear_solver_iterations: usize, + #[pyo3(get)] + number_of_nonlinear_solver_fails: usize, + #[pyo3(get)] + number_of_error_test_failures: usize, + #[pyo3(get)] + number_of_linear_solver_setups_from_checkpoint: usize, + #[pyo3(get)] + number_of_linear_solver_setups_from_first_convergence_fail: usize, + #[pyo3(get)] + number_of_linear_solver_setups_from_second_convergence_fail: usize, + #[pyo3(get)] + number_of_linear_solver_setups_from_error_test_fail: usize, + #[pyo3(get)] + number_of_linear_solver_setups_from_step_success: usize, + #[pyo3(get)] + ic_time_secs: f64, + #[pyo3(get)] + solver_setup_time_secs: f64, + #[pyo3(get)] + integration_time_secs: f64, + #[pyo3(get)] + sens_error_control_relaxed: bool, +} + +#[pymethods] +impl SolverStatistics { + fn __repr__(&self) -> String { + format!( + "SolverStatistics(steps={}, linear_setups={} \ + [checkpoint={}, 1st_conv_fail={}, 2nd_conv_fail={}, err_fail={}, heuristic={}], \ + nl_iters={}, nl_fails={}, err_fails={}, \ + ic_time={:.3}ms, solver_setup={:.3}ms, integration={:.3}ms)", + self.number_of_steps, + self.number_of_linear_solver_setups, + self.number_of_linear_solver_setups_from_checkpoint, + self.number_of_linear_solver_setups_from_first_convergence_fail, + self.number_of_linear_solver_setups_from_second_convergence_fail, + self.number_of_linear_solver_setups_from_error_test_fail, + self.number_of_linear_solver_setups_from_step_success, + self.number_of_nonlinear_solver_iterations, + self.number_of_nonlinear_solver_fails, + self.number_of_error_test_failures, + self.ic_time_secs * 1000.0, + self.solver_setup_time_secs * 1000.0, + self.integration_time_secs * 1000.0, + ) + } +} + +impl From for SolverStatistics { + fn from(s: core_solve::SolverStatistics) -> Self { + Self { + number_of_steps: s.number_of_steps, + number_of_linear_solver_setups: s.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations: s.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails: s.number_of_nonlinear_solver_fails, + number_of_error_test_failures: s.number_of_error_test_failures, + number_of_linear_solver_setups_from_checkpoint: s + .number_of_linear_solver_setups_from_checkpoint, + number_of_linear_solver_setups_from_first_convergence_fail: s + .number_of_linear_solver_setups_from_first_convergence_fail, + number_of_linear_solver_setups_from_second_convergence_fail: s + .number_of_linear_solver_setups_from_second_convergence_fail, + number_of_linear_solver_setups_from_error_test_fail: s + .number_of_linear_solver_setups_from_error_test_fail, + number_of_linear_solver_setups_from_step_success: s + .number_of_linear_solver_setups_from_step_success, + ic_time_secs: s.ic_time_secs, + solver_setup_time_secs: s.solver_setup_time_secs, + integration_time_secs: s.integration_time_secs, + sens_error_control_relaxed: s.sens_error_control_relaxed, + } + } +} + +/// Result of a diffsol solve, exposed to Python. +/// +/// Wraps core's `SolveOutcome` and converts data to numpy arrays on access. One +/// type for every payload combination: `y` holds whichever rows the solve was +/// asked for, and the payloads it was not asked for read as `None`. +#[derive(Debug)] +#[pyclass(module = "pybamm.rust")] +pub struct SolveOutcome { + /// Solver flag: 0 = success, 1 = root found + #[pyo3(get)] + flag: i32, + /// Time at which an event was triggered, if any + #[pyo3(get)] + t_event: Option, + /// BDF solver statistics + #[pyo3(get)] + statistics: SolverStatistics, + /// Time points returned by the solver + t_vec: Vec, + /// Flat trajectory in column-major order: states, or output variables when + /// the solve was asked for them. + y_flat: Vec, + /// Flat row-derivative trajectory matching `y_flat`, when stored. + yp_flat: Option>, + /// Per-parameter sensitivity blocks, `None` when none were requested. + sens_blocks: Option>>, + n_rows: usize, + n_times: usize, + /// Full state where the trajectory ends, if the solve reported one + y_event_vec: Option>, + /// Cached y matrix to avoid repeated allocation/copy + y_cache: Option>>, + /// Cached yp matrix to avoid repeated allocation/copy + yp_cache: Option>>, +} + +#[pymethods] +impl SolveOutcome { + /// Time points as a numpy array. + #[getter] + fn t<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + PyArray1::from_slice(py, &self.t_vec) + } + + /// Trajectory matrix with shape `(n_rows, n_times)`. + /// + /// Rows are states, or the model's output variables when the solve was asked + /// for `outputs`; row index varies along rows, time along columns. + #[getter] + fn y<'py>(&mut self, py: Python<'py>) -> PyResult>> { + take_cached_matrix( + py, + &mut self.y_cache, + &mut self.y_flat, + self.n_rows, + self.n_times, + ) + } + + /// Row-derivative matrix with shape `(n_rows, n_times)`, or `None` when the + /// solver was built with `store_yp=False` or the solve returned outputs. + #[getter] + fn yp<'py>(&mut self, py: Python<'py>) -> PyResult>>> { + take_cached_matrix_opt( + py, + &mut self.yp_cache, + &mut self.yp_flat, + self.n_rows, + self.n_times, + ) + } + + /// Per-parameter sensitivities, each matching the flat layout of `y`, or + /// `None` when the solve was not asked for them. + #[getter] + #[pyo3(name = "yS")] + fn param_sensitivities<'py>(&self, py: Python<'py>) -> Option>>> { + self.sens_blocks + .as_ref() + .map(|blocks| sens_blocks(blocks, py)) + } + + /// Full state where the trajectory ends, never an outputs row. + #[getter] + fn y_event<'py>(&self, py: Python<'py>) -> Option>> { + self.y_event_vec + .as_ref() + .map(|v| PyArray1::from_slice(py, v)) + } +} + +impl From for SolveOutcome { + fn from(outcome: core_solve::SolveOutcome) -> Self { + Self { + flag: outcome.flag, + t_event: outcome.t_event, + statistics: SolverStatistics::from(outcome.statistics), + t_vec: outcome.t, + y_flat: outcome.y, + yp_flat: outcome.yp, + sens_blocks: outcome.sensitivities, + n_rows: outcome.n_rows, + n_times: outcome.n_times, + y_event_vec: outcome.y_event, + y_cache: None, + yp_cache: None, + } + } +} + +/// Consume a column-major flat buffer into a cached `(nrows, ncols)` F-order +/// array: zero-copy on first access, cached reference thereafter. +fn take_cached_matrix<'py>( + py: Python<'py>, + cache: &mut Option>>, + flat: &mut Vec, + nrows: usize, + ncols: usize, +) -> PyResult>> { + if let Some(cached) = cache { + return Ok(cached.bind(py).clone()); + } + let data = std::mem::take(flat); + // Not a double-access guard: the cache above makes that unreachable. This + // fires only when the result was built with empty data for a non-empty + // shape, which is a producer bug rather than anything the caller did. + if data.is_empty() && nrows > 0 && ncols > 0 { + return Err(PyRuntimeError::new_err(format!( + "result constructed with empty data for a {nrows}x{ncols} matrix" + ))); + } + let arr = Array2::from_shape_vec((nrows, ncols).f(), data) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + let py_arr = arr.into_pyarray(py); + *cache = Some(py_arr.clone().unbind()); + Ok(py_arr) +} + +/// [`take_cached_matrix`] for an optional flat buffer: `None` stays `None`. +fn take_cached_matrix_opt<'py>( + py: Python<'py>, + cache: &mut Option>>, + flat: &mut Option>, + nrows: usize, + ncols: usize, +) -> PyResult>>> { + flat.as_mut().map_or_else( + || Ok(None), + |flat| take_cached_matrix(py, cache, flat, nrows, ncols).map(Some), + ) +} + +/// The `dy0/dp` seed as a flat slice; a missing seed is the empty (all-zero) one. +/// +/// A seed handed to a solve that was not asked for sensitivities is rejected +/// rather than ignored: the two arguments are one intent, and silently dropping +/// the seed would return zero sensitivities that look computed. +fn y0_sens_slice<'a>( + y0_sens: Option<&'a PyReadonlyArray1<'_, f64>>, + sensitivities: bool, +) -> PyResult<&'a [f64]> { + match y0_sens { + Some(_) if !sensitivities => Err(seed_without_sensitivities()), + Some(arr) => Ok(arr.as_slice()?), + None => Ok(&[]), + } +} + +/// The rejection both seed helpers raise. +fn seed_without_sensitivities() -> PyErr { + PyValueError::new_err("y0_sens was given but sensitivities=False") +} + +/// Per-parameter flat sensitivity blocks as a list of 1-D numpy arrays. +fn sens_blocks<'py>(blocks: &[Vec], py: Python<'py>) -> Vec>> { + blocks + .iter() + .map(|blk| PyArray1::from_slice(py, blk)) + .collect() +} + +/// Absolute tolerance as `PyBaMM` sends it: one value shared by every state, or +/// one value per state. +/// +/// The per-state arm is tried first because a length-1 array also satisfies the +/// uniform arm, and would then silently broadcast its single entry. +#[derive(Debug, FromPyObject)] +pub enum PyAtol<'py> { + PerState(PyReadonlyArray1<'py, f64>), + Uniform(f64), +} + +impl PyAtol<'_> { + /// Widen to the per-state vector the core takes. + /// + /// A per-state array passes through; a length mismatch is + /// `PreparedSolver::new`'s to reject against the model's own state count. + fn into_vec(self, n_states: usize) -> Vec { + match self { + Self::PerState(atol) => atol.as_array().to_vec(), + Self::Uniform(atol) => vec![atol; n_states], + } + } +} + +/// Prepare-once/execute-many solver for repeated integrations of one model. +/// +/// Built once during setup, then driven by many `solve*()` calls with different +/// initial conditions and inputs. `y0`, `t_eval` and the inputs are per call; +/// only the model, tolerances and integrator options are retained. +#[derive(Debug)] +#[pyclass(module = "pybamm.rust")] +pub struct PreparedSolver { + inner: core_solve::PreparedSolver, +} + +#[pymethods] +impl PreparedSolver { + /// Prepare a solver for repeated use with the given model and tolerances. + /// + /// `atol` is either one tolerance for every state or one per state; + /// `options` is `PyBaMM`'s integrator tuning dict, and omitting it keeps the + /// defaults in `SolverOptions`. `store_yp` additionally stores the state + /// time derivatives on the state-trajectory paths (`result.yp`), the knot + /// slopes cubic-Hermite output interpolation needs. + #[new] + #[pyo3(signature = (model, rtol=1e-6, atol=PyAtol::Uniform(1e-6), sens_atol_factor=1e-3, options=None, store_yp=false))] + fn new( + model: &CompiledModel, + rtol: f64, + atol: PyAtol<'_>, + sens_atol_factor: f64, + options: Option, + store_yp: bool, + ) -> PyResult { + let compiled = Arc::clone(&model.compiled); + let atol_vec = atol.into_vec(compiled.n_states()); + let options = options.map_or_else(SolverOptions::default, SolverOptions::from); + + let prepared = core_solve::PreparedSolver::new(compiled, rtol, &atol_vec) + .and_then(|p| p.with_sens_atol_factor(sens_atol_factor)) + .and_then(|p| p.with_options(options)) + .map_err(core_err_to_py)? + .with_store_yp(store_yp); + + Ok(Self { inner: prepared }) + } + + /// Solve the model over the given time span. + /// + /// `t_stop` holds the discontinuity times the integrator must land on + /// exactly, restarting there; every one of them must also appear in + /// `t_eval`, which is where the solution is reported. + /// + /// `outputs` reports the model's registered output variables instead of the + /// full state, which cuts the FFI transfer when only a few variables are + /// wanted, and `sensitivities` adds the forward-sensitivity blocks, seeded by + /// `y0_sens` (flattened `dy0/dp`, column-major over the requested + /// parameters). The two compose, so the four payload combinations are these + /// two flags rather than four entry points. + #[pyo3(signature = (t_eval, t_stop, y0, inputs, *, outputs=false, sensitivities=false, y0_sens=None))] + // Over clippy's threshold because pyo3 signatures are flat: the payload flags + // that replaced four entry points have to be spelled out as arguments. + #[allow(clippy::too_many_arguments)] + fn solve( + &self, + py: Python<'_>, + t_eval: PyReadonlyArray1<'_, f64>, + t_stop: PyReadonlyArray1<'_, f64>, + y0: PyReadonlyArray1<'_, f64>, + inputs: PyReadonlyArray1<'_, f64>, + outputs: bool, + sensitivities: bool, + y0_sens: Option>, + ) -> PyResult { + let request = solve_request(&t_eval, &t_stop, outputs, sensitivities)?; + let set = core_solve::InputSet::new(y0.as_slice()?, inputs.as_slice()?) + .with_sens_seed(y0_sens_slice(y0_sens.as_ref(), sensitivities)?); + + let inner = &self.inner; + let outcome = py + .detach(|| inner.solve(request, set)) + .map_err(core_err_to_py)?; + + Ok(SolveOutcome::from(outcome)) + } + + /// Solve every input set in `y0`/`inputs`, `num_threads` at a time. + /// + /// `y0` and `inputs` are C-contiguous 2-D arrays with one row per input set, + /// and `y0_sens` one seed row per set; `t_eval`, `t_stop` and the payload + /// flags are shared, as they already are on the callers. The returned list + /// has one entry per row, in row order: a result object, or — for a set that + /// failed — the exception *instance*, unraised, which is what keeps the + /// failing set's identity that one collapsed error would lose. + /// + /// One `py.detach()` covers the batch, so Ctrl-C lands when it returns. + #[pyo3(signature = (t_eval, t_stop, y0, inputs, num_threads, *, outputs=false, sensitivities=false, y0_sens=None))] + // Over the threshold for the same reason as `solve`, plus the pool width. + #[allow(clippy::too_many_arguments)] + fn solve_batch( + &self, + py: Python<'_>, + t_eval: PyReadonlyArray1<'_, f64>, + t_stop: PyReadonlyArray1<'_, f64>, + y0: PyReadonlyArray2<'_, f64>, + inputs: PyReadonlyArray2<'_, f64>, + num_threads: usize, + outputs: bool, + sensitivities: bool, + y0_sens: Option>, + ) -> PyResult>> { + let request = solve_request(&t_eval, &t_stop, outputs, sensitivities)?; + let y0_rows = batch_rows(&y0)?; + let inputs_rows = batch_rows(&inputs)?; + let seed_rows = batch_sens_rows(y0_sens.as_ref(), sensitivities, y0_rows.len())?; + // The widths a `&[InputSet]` cannot express: three arrays arrive here + // with independent row counts, and core's rule names the mismatch. + let seeds_given = y0_sens.is_some().then_some(seed_rows.len()); + check_batch_widths(y0_rows.len(), inputs_rows.len(), seeds_given) + .map_err(core_err_to_py)?; + let sets: Vec> = y0_rows + .iter() + .zip(&inputs_rows) + .zip(&seed_rows) + .map(|((y0, inputs), seed)| core_solve::InputSet::new(y0, inputs).with_sens_seed(seed)) + .collect(); + + let inner = &self.inner; + let outcomes = py.detach(|| { + let pool = pool_for(num_threads)?; + Ok::<_, PyErr>(pool.install(|| inner.solve_batch(request, &sets))) + })?; + batch_entries(py, outcomes) + } +} + +/// Build the shared half of a solve from the arguments Python sent. +fn solve_request<'a>( + t_eval: &'a PyReadonlyArray1<'_, f64>, + t_stop: &'a PyReadonlyArray1<'_, f64>, + outputs: bool, + sensitivities: bool, +) -> PyResult> { + Ok(core_solve::SolveRequest { + t_eval: t_eval.as_slice()?, + t_stop: t_stop.as_slice()?, + outputs, + sensitivities, + }) +} + +/// One borrowed row per input set of a C-contiguous `(n_sets, width)` array. +/// +/// A zero-width array (a model with no input parameters) yields `n_sets` empty +/// rows rather than no rows, which is what keeps the batch width equal to the +/// number of sets. +fn batch_rows<'a>(array: &'a PyReadonlyArray2<'_, f64>) -> PyResult> { + let shape = array.shape(); + let (n_sets, width) = (shape[0], shape[1]); + // as_slice enforces C-contiguity, which is what makes the row split sound. + let flat = array.as_slice()?; + Ok((0..n_sets) + .map(|i| &flat[i * width..(i + 1) * width]) + .collect()) +} + +/// The per-set `dy0/dp` seeds, or `n_sets` empty (all-zero) seeds when omitted. +/// +/// Rejects a seed array without `sensitivities` for the same reason +/// [`y0_sens_slice`] does. +fn batch_sens_rows<'a>( + y0_sens: Option<&'a PyReadonlyArray2<'_, f64>>, + sensitivities: bool, + n_sets: usize, +) -> PyResult> { + match y0_sens { + Some(_) if !sensitivities => Err(seed_without_sensitivities()), + Some(array) => batch_rows(array), + None => Ok(vec![&[][..]; n_sets]), + } +} + +/// Convert per-set core outcomes into the Python list `solve_batch` returns. +/// +/// A failed set contributes its exception instance, built but never raised, so +/// the caller keeps the index alongside the cause. +fn batch_entries( + py: Python<'_>, + outcomes: Vec>, +) -> PyResult>> { + outcomes + .into_iter() + .map(|outcome| match outcome { + Ok(value) => Ok(Py::new(py, SolveOutcome::from(value))?.into_any()), + Err(error) => Ok(core_err_to_py(error).into_value(py).into_any()), + }) + .collect() +} diff --git a/packages/pybamm-rust/rustfmt.toml b/packages/pybamm-rust/rustfmt.toml new file mode 100644 index 0000000000..988bbfa968 --- /dev/null +++ b/packages/pybamm-rust/rustfmt.toml @@ -0,0 +1,17 @@ +edition = "2024" + +max_width = 100 +use_small_heuristics = "Default" +tab_spaces = 4 +newline_style = "Auto" + +# Import organization (stable) +reorder_imports = true +reorder_modules = true + +# Expression formatting (stable) +use_field_init_shorthand = true +use_try_shorthand = true + +# Control flow formatting (stable) +match_block_trailing_comma = true diff --git a/packages/pybamm/hatch_build.py b/packages/pybamm/hatch_build.py new file mode 100644 index 0000000000..88d29f1f63 --- /dev/null +++ b/packages/pybamm/hatch_build.py @@ -0,0 +1,173 @@ +"""Hatchling build hook that compiles the Rust extension into the pybamm wheel. + +maturin builds the ``pybamm-rust`` crate into its own wheel; this hook runs that +build, lifts the ``_core`` extension out of the result into ``pybamm/rust/``, and +re-tags the pybamm wheel as platform-specific so the extension is not shipped as +pure Python. +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +import tempfile +import zipfile +from pathlib import Path + +from hatchling.builders.hooks.plugin.interface import BuildHookInterface + +# Buys one wheel per platform for CPython 3.10-3.14, at +20 ns per FFI call. +FEATURES = "pyo3/abi3-py310,diffsol" + +CRATE_MANIFEST = Path("pybamm-python") / "Cargo.toml" +DESTINATION = Path("src") / "pybamm" / "rust" +EXTENSION_SUFFIXES = (".so", ".pyd") + +CARGO_MISSING = """\ +Building PyBaMM from source requires a Rust toolchain (cargo >= 1.89), which was +not found on PATH. Install it from https://rustup.rs, or install a prebuilt wheel +with `pip install --only-binary pybamm pybamm`. +""" + + +def wheel_tag(wheel_name: str) -> str: + """Return the ``--`` tag from a wheel filename. + + Parameters + ---------- + wheel_name : str + Basename of a wheel, e.g. ``pybamm-1.0-cp310-abi3-win_amd64.whl``. + + Returns + ------- + str + The trailing three tag components. + + Raises + ------ + RuntimeError + If the name has too few components to carry a tag. + """ + parts = wheel_name.removesuffix(".whl").split("-") + if len(parts) < 5: + raise RuntimeError(f"Cannot parse a wheel tag from {wheel_name!r}") + return "-".join(parts[-3:]) + + +def extract_extension(wheel: Path, destination: Path) -> Path: + """Copy the ``_core`` extension out of ``wheel`` into ``destination``. + + Parameters + ---------- + wheel : Path + Wheel produced by maturin. + destination : Path + Directory to write the extension into. + + Returns + ------- + Path + The written extension. + + Raises + ------ + RuntimeError + If the wheel does not hold exactly one ``_core`` extension. + """ + with zipfile.ZipFile(wheel) as archive: + members = [ + name + for name in archive.namelist() + if Path(name).name.startswith("_core") and name.endswith(EXTENSION_SUFFIXES) + ] + if len(members) != 1: + raise RuntimeError( + f"Expected exactly one _core extension in {wheel.name}, found {members}" + ) + + destination.mkdir(parents=True, exist_ok=True) + # A stale version-specific .so shadows a new abi3 one (importlib checks it + # first); only binaries are stale — the checked-in _core.pyi must survive. + for stale in destination.glob("_core*"): + if stale.name.endswith(EXTENSION_SUFFIXES): + stale.unlink() + + target = destination / Path(members[0]).name + staged = target.with_name(f"{target.name}.tmp") + with archive.open(members[0]) as source, staged.open("wb") as sink: + shutil.copyfileobj(source, sink) + staged.replace(target) + + return target + + +class RustBuildHook(BuildHookInterface): + """Build the Rust extension into ``pybamm/rust/`` and tag the wheel for the platform.""" + + # Local hook files load through hatchling's `custom` plugin, which overwrites + # PLUGIN_NAME on the instance anyway. Declaring it matching avoids confusion. + PLUGIN_NAME = "custom" + + def initialize(self, version: str, build_data: dict) -> None: + """Build the extension and register it with the wheel being assembled. + + Parameters + ---------- + version : str + Hatchling's build version; ``"editable"`` keeps hatchling's own tag, + since an editable install is not redistributed. + build_data : dict + Mutated in place to mark the wheel platform-specific, list the + extension as an artifact, and set the wheel tag. + """ + artifact, tag = self._build_extension(self._crate_root()) + + # Without this the wheel is declared pure-Python, and .gitignore's `*.so` + # would hide the artifact from hatchling's file selection. + build_data["pure_python"] = False + build_data["artifacts"].append(f"/{DESTINATION.as_posix()}/{artifact.name}") + + if version != "editable": + build_data["tag"] = tag + + def _crate_root(self) -> Path: + """Locate the Cargo workspace: ``../pybamm-rust`` in the monorepo, ``./pybamm-rust`` in an sdist.""" + root = Path(self.root) + for candidate in (root.parent / "pybamm-rust", root / "pybamm-rust"): + if (candidate / CRATE_MANIFEST).is_file(): + return candidate + raise RuntimeError( + f"Could not find {CRATE_MANIFEST} under {root.parent / 'pybamm-rust'} " + f"or {root / 'pybamm-rust'}" + ) + + def _build_extension(self, crate_root: Path) -> tuple[Path, str]: + if shutil.which("cargo") is None: + raise RuntimeError(CARGO_MISSING) + + command = [ + sys.executable, + "-m", + "maturin", + "build", + "--release", + "--locked", + "--manifest-path", + str(CRATE_MANIFEST), + "--features", + FEATURES, + ] + if sys.platform.startswith("linux"): + # Let auditwheel own manylinux compliance and retagging downstream. + command += ["--compatibility", "linux", "--auditwheel", "skip"] + + with tempfile.TemporaryDirectory() as staging: + # cwd must be the crate root: cargo discovers .cargo/config.toml from + # its cwd, not from --manifest-path. + subprocess.run([*command, "--out", staging], cwd=crate_root, check=True) + wheel = next(Path(staging).glob("*.whl")) + return ( + extract_extension(wheel, Path(self.root) / DESTINATION), + wheel_tag(wheel.name), + ) diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index 945f58bc6b..1a6f7733b8 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["hatchling>=1.31.0", "hatch-vcs>=0.5.0"] +requires = ["hatchling>=1.31.0", "hatch-vcs>=0.5.0", "maturin>=1.14.1,<2.0"] build-backend = "hatchling.build" [project] @@ -113,6 +113,13 @@ dev = [ "importlib-metadata; python_version < '3.10'", # For property based testing "hypothesis", + # Imported by packages/pybamm/hatch_build.py, which has unit tests. + "hatchling>=1.31.0", + # Only for `mypy.stubtest`, which pins pybamm/rust/_core.pyi against the + # built extension. Also in the root group, which nox does not install here. + "mypy>=1.14", + # The stub imports scipy.sparse; without these stubtest refuses to run. + "scipy-stubs", ] [project.entry-points."pybamm_parameter_sets"] @@ -147,9 +154,41 @@ Yang2017 = "pybamm.models.full_battery_models.lithium_ion.Yang2017:Yang2017" [tool.hatch] version.source = "vcs" version.fallback-version = "0.0.0" -build.targets.sdist.include = ["src/pybamm", "CITATION.cff"] +build.targets.sdist.include = ["src/pybamm", "CITATION.cff", "hatch_build.py"] build.hooks.vcs.version-file = "src/pybamm/_version.py" +# Only the files a cargo build needs. A force-included DIRECTORY is walked +# recursively and ignores `exclude`, so never point an entry at ../pybamm-rust. +[tool.hatch.build.targets.sdist.force-include] +"../pybamm-rust/Cargo.toml" = "pybamm-rust/Cargo.toml" +"../pybamm-rust/Cargo.lock" = "pybamm-rust/Cargo.lock" +"../pybamm-rust/pybamm-core/Cargo.toml" = "pybamm-rust/pybamm-core/Cargo.toml" +"../pybamm-rust/pybamm-core/src" = "pybamm-rust/pybamm-core/src" +# cargo validates every declared [[bench]] target exists, even for a lib-only +# build. tests/ is auto-discovered, so it stays out. +"../pybamm-rust/pybamm-core/benches" = "pybamm-rust/pybamm-core/benches" +"../pybamm-rust/pybamm-python/Cargo.toml" = "pybamm-rust/pybamm-python/Cargo.toml" +"../pybamm-rust/pybamm-python/src" = "pybamm-rust/pybamm-python/src" + +# Local hook files load through hatchling's `custom` plugin. `editable` is a +# version of the wheel target, not a target, so one registration covers both. +[tool.hatch.build.targets.wheel.hooks.custom] +path = "hatch_build.py" + +# Replaces uv's default of just pyproject.toml, so every build input must appear. +# Scoped to src/ trees: a `**/*.rs` glob would walk the 21 GB Cargo target/ dir. +[tool.uv] +cache-keys = [ + { file = "pyproject.toml" }, + { file = "hatch_build.py" }, + { file = "../pybamm-rust/Cargo.toml" }, + { file = "../pybamm-rust/Cargo.lock" }, + { file = "../pybamm-rust/pybamm-core/Cargo.toml" }, + { file = "../pybamm-rust/pybamm-python/Cargo.toml" }, + { file = "../pybamm-rust/pybamm-core/src/**/*.rs" }, + { file = "../pybamm-rust/pybamm-python/src/**/*.rs" }, +] + # Monorepo tag scoping: PyBaMM releases are tagged `pybamm-v` (e.g. # pybamm-v27.1.0.0). Scope git-describe to that namespace so pybammsolvers tags # (`pybammsolvers-v*`) and legacy bare solver tags (`v0.x`) never drive PyBaMM's @@ -185,6 +224,9 @@ minversion = "9.0" required_plugins = ["pytest-xdist", "pytest-mock"] addopts = ["-nauto", "-vra", "--strict-config", "--strict-markers"] testpaths = ["tests"] +# Repo root, so benchmark-backed tests can `import benchmarks...` without each +# one re-deriving the path from its own depth. +pythonpath = ["../.."] console_output_style = "progress" xfail_strict = true filterwarnings = [ diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index 8370ce6f58..96cd21580a 100644 --- a/packages/pybamm/src/pybamm/__init__.py +++ b/packages/pybamm/src/pybamm/__init__.py @@ -195,6 +195,7 @@ # Solver classes from .solvers.solution import ( + SolverStatistics, SolutionBase, Solution, EISSolution, @@ -219,6 +220,7 @@ from .solvers.idaklu_jax import IDAKLUJax from .solvers.idaklu_solver import IDAKLUSolver +from .solvers.diffsol_solver import DiffsolSolver from .solvers.nonlinear_solver import NonlinearSolver # Experiments diff --git a/packages/pybamm/src/pybamm/codegen/compilation.py b/packages/pybamm/src/pybamm/codegen/compilation.py index a0b756ca0f..95489ca16c 100644 --- a/packages/pybamm/src/pybamm/codegen/compilation.py +++ b/packages/pybamm/src/pybamm/codegen/compilation.py @@ -6,6 +6,10 @@ import tempfile import time import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass import casadi @@ -40,6 +44,44 @@ _ALLOWED_COMPILERS = frozenset({"gcc", "clang", "cc", "g++", "clang++"}) +@dataclass(frozen=True) +class _AotCompileEvent: + """One AOT compilation request and its cache/phase telemetry.""" + + function_names: tuple[str, ...] + cache_status: str + cache_key: str | None + codegen_ms: float + compiler_ms: float + load_ms: float + total_ms: float + library_path: str | None + library_size_bytes: int | None + error: str | None = None + + +_CAPTURED_EVENTS: ContextVar[list[_AotCompileEvent] | None] = ContextVar( + "pybamm_aot_compile_events", default=None +) + + +@contextmanager +def _capture_aot_compile_events() -> Iterator[list[_AotCompileEvent]]: + """Capture AOT cache and phase telemetry within the current context.""" + events: list[_AotCompileEvent] = [] + token = _CAPTURED_EVENTS.set(events) + try: + yield events + finally: + _CAPTURED_EVENTS.reset(token) + + +def _record_compile_event(event: _AotCompileEvent) -> None: + events = _CAPTURED_EVENTS.get() + if events is not None: + events.append(event) + + def _default_cache_dir() -> str: d = os.environ.get("PYBAMM_CASADI_AOT_CACHE") if d: @@ -88,12 +130,26 @@ def aot_compile(fn_or_fns, **kwargs): """ is_single = isinstance(fn_or_fns, casadi.Function) fns = [fn_or_fns] if is_single else list(fn_or_fns) + start = time.perf_counter() try: - out = _aot_compile(fns, **kwargs) + out, event = _aot_compile(fns, **kwargs) except (OSError, ValueError, RuntimeError, subprocess.SubprocessError) as e: names = ", ".join(fn.name() for fn in fns) logger.warning(f"Failed to compile [{names}] with error: {e}") out = list(fns) + event = _AotCompileEvent( + function_names=tuple(fn.name() for fn in fns), + cache_status="fallback", + cache_key=None, + codegen_ms=0.0, + compiler_ms=0.0, + load_ms=0.0, + total_ms=(time.perf_counter() - start) * 1000.0, + library_path=None, + library_size_bytes=None, + error=f"{type(e).__name__}: {e}", + ) + _record_compile_event(event) return out[0] if is_single else out @@ -103,14 +159,25 @@ def _aot_compile( cache_dir: str | None = None, compiler: str | None = None, flags: tuple[str, ...] | None = None, -) -> list[casadi.Function]: +) -> tuple[list[casadi.Function], _AotCompileEvent]: + start = time.perf_counter() # Pass-through Externals; compile the rest together in one TU. result: list[casadi.Function] = list(fns) indices_to_compile = [ i for i, fn in enumerate(fns) if fn.class_name() != "External" ] if not indices_to_compile: - return result + return result, _AotCompileEvent( + function_names=tuple(fn.name() for fn in fns), + cache_status="external", + cache_key=None, + codegen_ms=0.0, + compiler_ms=0.0, + load_ms=0.0, + total_ms=(time.perf_counter() - start) * 1000.0, + library_path=None, + library_size_bytes=None, + ) # Cache key: ordered hash of each fn's name + serialized form. hasher = hashlib.sha1(usedforsecurity=False) @@ -126,7 +193,17 @@ def _aot_compile( if cached is not None: for idx, ext_fn in zip(indices_to_compile, cached, strict=True): result[idx] = ext_fn - return result + return result, _AotCompileEvent( + function_names=tuple(fns[idx].name() for idx in indices_to_compile), + cache_status="memory", + cache_key=key, + codegen_ms=0.0, + compiler_ms=0.0, + load_ms=0.0, + total_ms=(time.perf_counter() - start) * 1000.0, + library_path=None, + library_size_bytes=None, + ) if compiler is None: compiler = "gcc" @@ -148,12 +225,17 @@ def _aot_compile( stem = f"{_TMP_FILE_PREFIX}{label}_{key}" ext = _shared_ext() sofile = os.path.join(cdir, stem + ext) + cache_status = "disk" if os.path.exists(sofile) else "miss" + codegen_ms = 0.0 + compiler_ms = 0.0 - if not os.path.exists(sofile): + if cache_status == "miss": + codegen_start = time.perf_counter() gen = casadi.CodeGenerator(stem, {"with_header": False}) for fn in fns_to_compile: gen.add(fn) c_source = gen.dump() + codegen_ms = (time.perf_counter() - codegen_start) * 1000.0 bundled = {fn.name() for fn in fns_to_compile} externs = set(_EXTERN_DECL.findall(c_source)) - bundled @@ -173,10 +255,12 @@ def _aot_compile( try: with open(tmp_cfile, "w") as f: f.write(c_source) + compiler_start = time.perf_counter() subprocess.run( # nosec B603 B607 - compiler validated against allowlist [compiler, *flags, "-shared", tmp_cfile, "-o", tmp_sofile], check=True, ) + compiler_ms = (time.perf_counter() - compiler_start) * 1000.0 os.replace(tmp_sofile, sofile) if os.environ.get("PYBAMM_CASADI_AOT_KEEP_C"): os.replace(tmp_cfile, os.path.join(cdir, stem + ".c")) @@ -187,13 +271,29 @@ def _aot_compile( except OSError: pass + load_start = time.perf_counter() ext_fns: list[casadi.Function] = [] for idx, fn in zip(indices_to_compile, fns_to_compile, strict=True): ext_fn = casadi.external(fn.name(), sofile) result[idx] = ext_fn ext_fns.append(ext_fn) + load_ms = (time.perf_counter() - load_start) * 1000.0 _CACHE[key] = ext_fns - return result + try: + library_size_bytes = os.path.getsize(sofile) + except OSError: + library_size_bytes = None + return result, _AotCompileEvent( + function_names=tuple(fn.name() for fn in fns_to_compile), + cache_status=cache_status, + cache_key=key, + codegen_ms=codegen_ms, + compiler_ms=compiler_ms, + load_ms=load_ms, + total_ms=(time.perf_counter() - start) * 1000.0, + library_path=sofile, + library_size_bytes=library_size_bytes, + ) def _maybe_sweep_stale(cdir: str) -> None: diff --git a/packages/pybamm/src/pybamm/expression_tree/array.py b/packages/pybamm/src/pybamm/expression_tree/array.py index b509c07d6f..b381339880 100644 --- a/packages/pybamm/src/pybamm/expression_tree/array.py +++ b/packages/pybamm/src/pybamm/expression_tree/array.py @@ -132,6 +132,26 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): """See :meth:`pybamm.Symbol._to_casadi()`.""" return casadi.MX(self.evaluate(t, y, y_dot, inputs)) + def _to_rust(self, graph, rust_symbols): + entries = self.evaluate() + if issparse(entries): + csr = csr_matrix(entries) + return graph.sparse_matrix( + csr.indptr.tolist(), + csr.indices.tolist(), + csr.data.astype(float), + csr.shape[0], + csr.shape[1], + ) + entries = np.asarray(entries, dtype=float) + if entries.ndim == 2 and entries.shape[1] > 1: + return graph.dense_matrix( + np.ascontiguousarray(entries).ravel(), + entries.shape[0], + entries.shape[1], + ) + return graph.array(entries.flatten()) + def _jac(self, variable) -> pybamm.Matrix: """See :meth:`pybamm.Symbol._jac()`.""" # Return zeros of correct size diff --git a/packages/pybamm/src/pybamm/expression_tree/binary_operators.py b/packages/pybamm/src/pybamm/expression_tree/binary_operators.py index 1cc41f2e89..164a71289e 100644 --- a/packages/pybamm/src/pybamm/expression_tree/binary_operators.py +++ b/packages/pybamm/src/pybamm/expression_tree/binary_operators.py @@ -197,6 +197,41 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ) return self._casadi_evaluate(converted_left, converted_right) + def _to_rust(self, graph, rust_symbols): + left, right = self._children_to_rust(graph, rust_symbols) + if isinstance(self, Addition): + return graph.add(left, right) + elif isinstance(self, Subtraction): + return graph.sub(left, right) + elif isinstance(self, Multiplication): + return graph.mul(left, right) + elif isinstance(self, Division): + return graph.div(left, right) + elif isinstance(self, Power): + return graph.pow(left, right) + elif isinstance(self, MatrixMultiplication): + return graph.matmul(left, right) + elif isinstance(self, Inner): + return graph.mul(left, right) + elif isinstance(self, Minimum): + return graph.minimum(left, right) + elif isinstance(self, Maximum): + return graph.maximum(left, right) + elif isinstance(self, Modulo): + return graph.modulo(left, right) + elif isinstance(self, Hypot): + return graph.hypot(left, right) + elif isinstance(self, EqualHeaviside): + return graph.equal_heaviside(left, right) + elif isinstance(self, NotEqualHeaviside): + return graph.not_equal_heaviside(left, right) + elif isinstance(self, Equality): + return graph.equality(left, right) + else: + raise TypeError( + f"Cannot convert binary operator of type '{type(self)}' to Rust" + ) + def _evaluates_on_edges(self, dimension: str) -> bool: """See :meth:`pybamm.Symbol._evaluates_on_edges()`.""" return self.left.evaluates_on_edges(dimension) or self.right.evaluates_on_edges( diff --git a/packages/pybamm/src/pybamm/expression_tree/concatenations.py b/packages/pybamm/src/pybamm/expression_tree/concatenations.py index b1773edfbb..1cf0c788f1 100644 --- a/packages/pybamm/src/pybamm/expression_tree/concatenations.py +++ b/packages/pybamm/src/pybamm/expression_tree/concatenations.py @@ -246,6 +246,10 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ) return casadi.vertcat(*converted_children) + def _to_rust(self, graph, rust_symbols): + converted_children = self._children_to_rust(graph, rust_symbols) + return graph.concat(converted_children) + def _concatenation_jac(self, children_jacs): """See :meth:`pybamm.Concatenation.concatenation_jac()`.""" children = self.children @@ -428,6 +432,27 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ) return casadi.vertcat(*all_child_vectors) + def _to_rust(self, graph, rust_symbols): + converted_children = self._children_to_rust(graph, rust_symbols) + slice_starts = [] + all_child_exprs = [] + for i in range(self.secondary_dimensions_npts): + child_exprs = [] + for child_var, slices in zip( + converted_children, self._children_slices, strict=True + ): + for child_dom, child_slice in slices.items(): + slice_starts.append(self._slices[child_dom][i].start) + child_exprs.append( + graph.index( + child_var, child_slice[i].start, child_slice[i].stop + ) + ) + all_child_exprs.extend( + [v for _, v in sorted(zip(slice_starts, child_exprs, strict=False))] + ) + return graph.concat(all_child_exprs) + def _concatenation_jac(self, children_jacs): """See :meth:`pybamm.Concatenation.concatenation_jac()`.""" # note that this assumes that the children are in the right order and only have @@ -523,6 +548,10 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ) return casadi.vertcat(*converted_children) + def _to_rust(self, graph, rust_symbols): + converted_children = self._children_to_rust(graph, rust_symbols) + return graph.concat(converted_children) + def _concatenation_new_copy(self, children, perform_simplifications=True): """See :meth:`pybamm.Concatenation._concatenation_new_copy()`.""" return SparseStack(*children) diff --git a/packages/pybamm/src/pybamm/expression_tree/conditional.py b/packages/pybamm/src/pybamm/expression_tree/conditional.py index 9154701a5d..024a38f1f7 100644 --- a/packages/pybamm/src/pybamm/expression_tree/conditional.py +++ b/packages/pybamm/src/pybamm/expression_tree/conditional.py @@ -191,6 +191,17 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): index = casadi.if_else(on_boundary, -1, floored - 1) return switch(index, *shared) + def _to_rust(self, graph, rust_symbols): + """Convert to Rust expression graph. + + See :meth:`pybamm.Symbol._to_rust()`. + """ + converted_selector = self.selector._to_rust_inner(graph, rust_symbols) + converted_branches = [ + branch._to_rust_inner(graph, rust_symbols) for branch in self.branches + ] + return graph.conditional(converted_selector, converted_branches) + def to_equation(self): if self.print_name is not None: return sympy.Symbol(self.print_name) diff --git a/packages/pybamm/src/pybamm/expression_tree/discrete_time_sum.py b/packages/pybamm/src/pybamm/expression_tree/discrete_time_sum.py index afee5b43a1..f35d4b53c6 100644 --- a/packages/pybamm/src/pybamm/expression_tree/discrete_time_sum.py +++ b/packages/pybamm/src/pybamm/expression_tree/discrete_time_sum.py @@ -114,3 +114,8 @@ def sum_times(self): def _unary_evaluate(self, child): # return result of evaluating the child, we'll only implement the sum once the model is solved (in pybamm.ProcessedVariable) return child + + def _to_rust(self, graph, rust_symbols): + # Pass-through: actual summation happens in ProcessedVariable + (child,) = self._children_to_rust(graph, rust_symbols) + return child diff --git a/packages/pybamm/src/pybamm/expression_tree/functions.py b/packages/pybamm/src/pybamm/expression_tree/functions.py index dca18046be..e39f35e9d6 100644 --- a/packages/pybamm/src/pybamm/expression_tree/functions.py +++ b/packages/pybamm/src/pybamm/expression_tree/functions.py @@ -178,6 +178,21 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ) return self._casadi_evaluate(*converted_children) + def _to_rust(self, graph, rust_symbols): + """See :meth:`pybamm.Symbol._to_rust()`.""" + converted_children = self._children_to_rust(graph, rust_symbols) + return self._rust_evaluate(graph, *converted_children) + + def _rust_evaluate(self, graph, *converted_children): + raise TypeError( + f"Cannot convert function '{self.name}' of type " + f"'{type(self).__name__}' to the Rust backend: generic Python " + "callables cannot be lowered. Use a CasADi-backed model (set " + "`model.convert_to_format = 'casadi'`) for models using " + "this function, or override _rust_evaluate in the subclass to add " + "native support." + ) + def create_copy( self, new_children: list[pybamm.Symbol] | None = None, @@ -301,6 +316,16 @@ def _casadi_evaluate(self, child): f"{self.__class__} does not implement _casadi_evaluate." ) + def _rust_evaluate(self, graph, child): + raise TypeError( + f"Cannot convert function '{self.name}' of type " + f"'{type(self).__name__}' to the Rust backend: no native Rust " + "implementation exists for this function. Use a CasADi-backed model " + "(set `model.convert_to_format = 'casadi'`) for models " + "using this function, or override _rust_evaluate in the subclass to " + "add native support." + ) + def to_json(self): """ Method to serialise a SpecificFunction object into JSON. @@ -355,6 +380,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.arcsinh(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.arcsinh(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Symbol._function_diff()`.""" return 1 / sqrt(children[0] ** 2 + 1) @@ -473,6 +502,19 @@ def _casadi_evaluate(self, a, b): b_eff = sign_b * casadi.hypot(b, self.eps) return casadi.arcsinh(a / b_eff) + def _rust_evaluate(self, graph, a, b): + """See :meth:`pybamm.Function._rust_evaluate()`.""" + zero = graph.scalar(0.0) + one = graph.scalar(1.0) + two = graph.scalar(2.0) + eps = graph.scalar(self.eps) + # sign_b = 2 * (b >= 0) - 1, treating sign(0) as 1 + # equal_heaviside(a, b) returns 1 if a <= b, so use (zero, b) for b >= 0 + b_ge_zero = graph.equal_heaviside(zero, b) + sign_b = graph.sub(graph.mul(two, b_ge_zero), one) + b_eff = graph.mul(sign_b, graph.hypot(b, eps)) + return graph.arcsinh(graph.div(a, b_eff)) + def _sympy_operator(self, a, b): """Convert to SymPy expression.""" # sign(b) but treat sign(0) as non-zero @@ -551,6 +593,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.arctan(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.arctan(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return 1 / (children[0] ** 2 + 1) @@ -582,6 +628,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.cos(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.cos(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Symbol._function_diff()`.""" return -sin(children[0]) @@ -609,6 +659,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.cosh(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.cosh(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return sinh(children[0]) @@ -636,6 +690,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.erf(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.erf(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return 2 / np.sqrt(np.pi) * exp(-(children[0] ** 2)) @@ -668,6 +726,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.exp(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.exp(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return exp(children[0]) @@ -695,6 +757,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.log(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.log(child) + def _function_evaluate(self, evaluated_children): # don't raise RuntimeWarning for NaNs with np.errstate(invalid="ignore"): @@ -736,6 +802,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.mmax(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.max_reduce(child) + def _evaluate_for_shape(self): """See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()`""" # Max will always return a scalar @@ -767,6 +837,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.mmin(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.min_reduce(child) + def _evaluate_for_shape(self): """See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()`""" # Min will always return a scalar @@ -803,6 +877,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.sin(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.sin(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return cos(children[0]) @@ -830,6 +908,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.sinh(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.sinh(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return cosh(children[0]) @@ -857,6 +939,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.sqrt(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.sqrt(child) + def _function_evaluate(self, evaluated_children): # don't raise RuntimeWarning for NaNs with np.errstate(invalid="ignore"): @@ -872,6 +958,40 @@ def sqrt(child: pybamm.Symbol): return simplified_function(Sqrt, child) +def _positive_base_pow_chain(graph, base, exponent): + """Lower ``base ** exponent`` to a sqrt/div chain for select exponents. + + Valid only for strictly positive ``base`` (sqrt of a negative base would + change NaN semantics). Returns the graph expression, or ``None`` when the + exponent has no chain form and the caller must fall back to ``pow``. + + Parameters + ---------- + graph : ExprGraph + The Rust expression graph being built. + base : Expr + Graph expression for the (positive) base. + exponent : float + Constant exponent value. + + Returns + ------- + Expr or None + Chain expression for exponents in {1, 0.5, -0.5, 0.25, -0.25}. + """ + if exponent == 1.0: + return base + if exponent == 0.5: + return graph.sqrt(base) + if exponent == -0.5: + return graph.div(graph.scalar(1.0), graph.sqrt(base)) + if exponent == 0.25: + return graph.sqrt(graph.sqrt(base)) + if exponent == -0.25: + return graph.div(graph.scalar(1.0), graph.sqrt(graph.sqrt(base))) + return None + + class RegPower(Function): """ Regularised power: |x|^a * sign(x) with finite derivative at x=0. @@ -939,6 +1059,36 @@ def _function_evaluate(self, evaluated_children): """See :meth:`pybamm.Function._function_evaluate()`.""" return self._reg_power_evaluate(*evaluated_children) + def _rust_evaluate(self, graph, base, exponent, scale): + """See :meth:`pybamm.Function._rust_evaluate()`. + + Mirrors :meth:`_reg_power_evaluate`: + y = (base/scale) * ((base/scale)^2 + delta^2) ^ ((exponent-1)/2) + * scale^exponent + + For constant exponents, the inner power lowers to a sqrt/div chain: + its base is >= delta^2 > 0 by construction, so the chain is exact- + domain, and runtime powf costs ~15x a hardware sqrt per element. + """ + x = graph.div(base, scale) + x_sq = graph.mul(x, x) + delta_sq = graph.scalar(self.delta * self.delta) + x_sq_plus_d_sq = graph.add(x_sq, delta_sq) + pow_x = None + exp_child = self.children[1] + if isinstance(exp_child, pybamm.Scalar): + half_a_minus_1_value = (float(exp_child.value) - 1.0) / 2.0 + pow_x = _positive_base_pow_chain( + graph, x_sq_plus_d_sq, half_a_minus_1_value + ) + if pow_x is None: + one = graph.scalar(1.0) + half = graph.scalar(0.5) + half_a_minus_1 = graph.mul(graph.sub(exponent, one), half) + pow_x = graph.pow(x_sq_plus_d_sq, half_a_minus_1) + pow_scale = graph.pow(scale, exponent) + return graph.mul(graph.mul(x, pow_x), pow_scale) + def _function_diff(self, children, idx): """ Derivative with respect to child number 'idx'. @@ -1143,6 +1293,10 @@ def _casadi_evaluate(self, child): """See :meth:`pybamm.SpecificFunction._casadi_evaluate()`.""" return casadi.tanh(child) + def _rust_evaluate(self, graph, child): + """See :meth:`pybamm.SpecificFunction._rust_evaluate()`.""" + return graph.tanh(child) + def _function_diff(self, children, idx): """See :meth:`pybamm.Function._function_diff()`.""" return sech(children[0]) ** 2 diff --git a/packages/pybamm/src/pybamm/expression_tree/independent_variable.py b/packages/pybamm/src/pybamm/expression_tree/independent_variable.py index a71ba5d83a..c5c9ccbaae 100644 --- a/packages/pybamm/src/pybamm/expression_tree/independent_variable.py +++ b/packages/pybamm/src/pybamm/expression_tree/independent_variable.py @@ -117,6 +117,9 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): """See :meth:`pybamm.Symbol._to_casadi()`.""" return casadi.MX(self.evaluate(t, y, y_dot, inputs)) + def _to_rust(self, graph, rust_symbols): + return graph.time() + def to_equation(self): """Convert the node and its subtree into a SymPy equation.""" return sympy.Symbol("t") diff --git a/packages/pybamm/src/pybamm/expression_tree/input_parameter.py b/packages/pybamm/src/pybamm/expression_tree/input_parameter.py index f71b78861f..c1d5eb7119 100644 --- a/packages/pybamm/src/pybamm/expression_tree/input_parameter.py +++ b/packages/pybamm/src/pybamm/expression_tree/input_parameter.py @@ -89,6 +89,9 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): """See :meth:`pybamm.Symbol._to_casadi()`.""" return casadi.MX(self.evaluate(t, y, y_dot, inputs)) + def _to_rust(self, graph, rust_symbols): + return graph.input_parameter(self.name, self._expected_size or 1) + def _jac(self, variable: pybamm.StateVector) -> pybamm.Matrix: """See :meth:`pybamm.Symbol._jac()`.""" n_variable = variable.evaluation_array.count(True) diff --git a/packages/pybamm/src/pybamm/expression_tree/interpolant.py b/packages/pybamm/src/pybamm/expression_tree/interpolant.py index a8cbca5560..70f48dae0c 100644 --- a/packages/pybamm/src/pybamm/expression_tree/interpolant.py +++ b/packages/pybamm/src/pybamm/expression_tree/interpolant.py @@ -3,6 +3,8 @@ # from __future__ import annotations +import itertools +import math import numbers from collections.abc import Sequence from typing import Any @@ -399,6 +401,128 @@ def _pchip_to_casadi(self, converted_children): result = result.T return result + def _to_rust(self, graph, rust_symbols): + converted_children = self._children_to_rust(graph, rust_symbols) + if self.dimension == 1 and self.y.ndim > 2: + raise NotImplementedError( + f"Rust conversion of 1D {self.interpolator} interpolation does " + f"not support y data with more than two dimensions " + f"(y.ndim={self.y.ndim}). Use a CasADi-backed model (set " + "`model.convert_to_format = 'casadi'`) for this model." + ) + # Vector-valued y (n, m): the constructor guarantees a size-1 child, so + # stack one interpolant per column in evaluate()'s column order. + if self.dimension == 1 and self.interpolator == "linear": + x_data = self.x[0].tolist() + child = converted_children[0] + if self.y.ndim == 2: + columns = [ + graph.interpolant_1d_linear(x_data, self.y[:, j].tolist(), child) + for j in range(self.y.shape[1]) + ] + return graph.concat(columns) + return graph.interpolant_1d_linear(x_data, self.y.tolist(), child) + elif self.dimension == 1 and self.interpolator in ("cubic", "pchip"): + # Coefficients from the same PPoly that evaluate() uses (exact oracle + # match); reshape adds the column axis vector-valued y already has. + ppoly = self.function + breakpoints = ppoly.x.tolist() + c = ppoly.c.reshape(ppoly.c.shape[0], ppoly.c.shape[1], -1) + nseg = c.shape[1] + deg = c.shape[0] - 1 + if deg > 3: + raise ValueError( + f"Unexpected PPoly degree {deg} for {self.interpolator} " + "interpolation (expected <= 3)" + ) + child = converted_children[0] + columns = [] + for col in range(c.shape[2]): + coeffs: list[float] = [] + for i in range(nseg): + # Horner power basis [c0, c1, c2, c3]; pad degrees < 3 with zeros. + power = [0.0, 0.0, 0.0, 0.0] + for j in range(deg + 1): + power[deg - j] = float(c[j, i, col]) + coeffs.extend(power) + columns.append(graph.interpolant_1d_cubic(breakpoints, coeffs, child)) + if self.y.ndim == 2: + return graph.concat(columns) + return columns[0] + elif self.dimension in (2, 3): + # The constructor rejects ND pchip, so this is linear or cubic. + breakpoints = [np.asarray(xi, dtype=float).tolist() for xi in self.x] + if self.interpolator == "linear": + order, coeffs = self._nd_linear_coeffs() + else: + order, coeffs = self._nd_cubic_coeffs() + return graph.interpolant_nd( + breakpoints, coeffs.tolist(), order, converted_children + ) + else: + raise NotImplementedError( + f"Rust conversion is not implemented for {self.dimension}D " + f"{self.interpolator} interpolation. Use a CasADi-backed model " + "(set `model.convert_to_format = 'casadi'`) for models " + "using this interpolant." + ) + + def _nd_tensor_layout(self, coeff): + """Flatten per-axis (order, nseg_a) coefficient pairs to the Rust + layout: cell-major (axis-0 segment slowest), then the order**ndim + within-cell power tensor (axis-0 power slowest, ascending powers).""" + ndim = self.dimension + seg_axes = tuple(2 * a + 1 for a in range(ndim)) + pow_axes = tuple(2 * a for a in range(ndim)) + return np.ascontiguousarray(np.transpose(coeff, seg_axes + pow_axes)).ravel() + + def _nd_linear_coeffs(self): + """Per-cell multilinear tensor coefficients via per-axis finite + differences. Matches RegularGridInterpolator(method="linear", + fill_value=None) exactly, including multilinear extension outside + the domain (pinned <=1e-15).""" + coeff = np.asarray(self.y, dtype=float) + for a in range(self.dimension): + pos = 2 * a # raw data axis a position after processing axes < a + coeff = np.moveaxis(coeff, pos, 0) + h = np.diff(self.x[a]).reshape((-1,) + (1,) * (coeff.ndim - 1)) + coeff = np.stack([coeff[:-1], np.diff(coeff, axis=0) / h]) + coeff = np.moveaxis(coeff, (0, 1), (pos, pos + 1)) + return 2, self._nd_tensor_layout(coeff) + + def _nd_cubic_coeffs(self): + """Per-cell tensor power coefficients for ND cubic, extracted from + self.function (the RegularGridInterpolator evaluate() uses). + + scipy >= 1.13 pre-fits a tensor-product NdBSpline with an iterative + solver, so RGI cubic is only approximately interpolating; the + coefficients must be Taylor-extracted from that fitted spline via the + public nu= derivative API, not refit (refit differs ~1e-5; extraction + pinned <=5e-15). On scipy < 1.13 (no nu=), RGI cubic is + the recursive not-a-knot tensor spline, reproduced exactly by a + sequential CubicSpline fit (pinned <=5e-14).""" + ndim = self.dimension + mesh = np.meshgrid(*(np.asarray(xi)[:-1] for xi in self.x), indexing="ij") + corners = np.column_stack([m.ravel() for m in mesh]) + try: + self.function(corners[:1], nu=(0,) * ndim) + has_nu = True + except TypeError: # scipy < 1.13 + has_nu = False + if has_nu: + coeffs = np.empty((corners.shape[0],) + (4,) * ndim) + for nu in itertools.product(range(4), repeat=ndim): + scale = math.prod(math.factorial(v) for v in nu) + coeffs[(slice(None), *nu)] = self.function(corners, nu=nu) / scale + return 4, coeffs.ravel() + coeff = np.asarray(self.y, dtype=float) + for a in range(ndim): + pos = 2 * a + coeff = np.moveaxis(coeff, pos, 0) + coeff = interpolate.CubicSpline(self.x[a], coeff, axis=0).c[::-1] + coeff = np.moveaxis(coeff, (0, 1), (pos, pos + 1)) + return 4, self._nd_tensor_layout(coeff) + def _function_diff(self, children: Sequence[pybamm.Symbol], idx: float): """ Derivative with respect to child number 'idx'. diff --git a/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py b/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py index 5c25087b1f..adb5ffa4df 100644 --- a/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py +++ b/packages/pybamm/src/pybamm/expression_tree/operations/serialise.py @@ -185,6 +185,7 @@ def serialise_model( "pybamm_version": pybamm.__version__, "name": model.name, "options": model.options, + "convert_to_format": model.convert_to_format, "bounds": [bound.tolist() for bound in model.bounds], # type: ignore[attr-defined] "concatenated_rhs": encode(model._concatenated_rhs), "concatenated_algebraic": encode(model._concatenated_algebraic), @@ -306,6 +307,7 @@ def load_model( recon_model_dict = { "name": model_data["name"], "options": self._convert_options(model_data["options"]), + "convert_to_format": model_data.get("convert_to_format", "casadi"), "bounds": tuple(np.array(bound) for bound in model_data["bounds"]), "concatenated_rhs": self._decode_model_node(model_data["concatenated_rhs"]), "concatenated_algebraic": self._decode_model_node( @@ -480,6 +482,7 @@ def serialise_custom_model(model: pybamm.BaseModel, compress: bool = False) -> d "base_class": base_cls_str, "base_class_mro": base_class_mro, "options": getattr(model, "options", {}), + "convert_to_format": getattr(model, "convert_to_format", "casadi"), "rhs": [ ( convert_symbol_to_json(variable), @@ -1560,6 +1563,9 @@ def load_custom_model(filename: str | dict) -> pybamm.BaseModel: f"Failed to convert variable '{variable_name}': {e!s}" ) from e + # Restore convert_to_format + model.convert_to_format = model_data.get("convert_to_format", "casadi") + # Restore observable state model._solution_observable = False diff --git a/packages/pybamm/src/pybamm/expression_tree/scalar.py b/packages/pybamm/src/pybamm/expression_tree/scalar.py index dac857b7e9..c6ee2ceb70 100644 --- a/packages/pybamm/src/pybamm/expression_tree/scalar.py +++ b/packages/pybamm/src/pybamm/expression_tree/scalar.py @@ -82,6 +82,9 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): """See :meth:`pybamm.Symbol._to_casadi()`.""" return casadi.MX(self.evaluate(t, y, y_dot, inputs)) + def _to_rust(self, graph, rust_symbols): + return graph.scalar(float(self.value)) + def _jac(self, variable: pybamm.Variable) -> pybamm.Scalar: """See :meth:`pybamm.Symbol._jac()`.""" return pybamm.Scalar(0) diff --git a/packages/pybamm/src/pybamm/expression_tree/state_vector.py b/packages/pybamm/src/pybamm/expression_tree/state_vector.py index c91265494b..95dc363b9c 100644 --- a/packages/pybamm/src/pybamm/expression_tree/state_vector.py +++ b/packages/pybamm/src/pybamm/expression_tree/state_vector.py @@ -319,6 +319,12 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): raise ValueError("Must provide a 'y' for converting state vectors") return casadi.vertcat(*[y[y_slice] for y_slice in self.y_slices]) + def _to_rust(self, graph, rust_symbols): + exprs = [graph.state_vector(s.start, s.stop) for s in self._y_slices] + if len(exprs) == 1: + return exprs[0] + return graph.concat(exprs) + def _jac(self, variable: pybamm.StateVector | pybamm.StateVectorDot): if isinstance(variable, pybamm.StateVector): return self._jac_same_vector(variable) @@ -406,6 +412,12 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): raise ValueError("Must provide a 'y_dot' for converting state vectors") return casadi.vertcat(*[y_dot[y_slice] for y_slice in self.y_slices]) + def _to_rust(self, graph, rust_symbols): + exprs = [graph.state_vector_dot(s.start, s.stop) for s in self._y_slices] + if len(exprs) == 1: + return exprs[0] + return graph.concat(exprs) + def _jac(self, variable: pybamm.StateVector | pybamm.StateVectorDot): if isinstance(variable, pybamm.StateVectorDot): return self._jac_same_vector(variable) diff --git a/packages/pybamm/src/pybamm/expression_tree/symbol.py b/packages/pybamm/src/pybamm/expression_tree/symbol.py index f2cff8a1c4..9e94a08666 100644 --- a/packages/pybamm/src/pybamm/expression_tree/symbol.py +++ b/packages/pybamm/src/pybamm/expression_tree/symbol.py @@ -1075,6 +1075,48 @@ def _children_to_casadi( for child in self.children ] + def to_rust(self, graph, rust_symbols=None): + """ + Allocate this expression into a Rust expression graph. + + Parameters + ---------- + graph : :class:`pybamm.rust.ExprGraph` + The graph to allocate nodes into. The returned handles are only valid + against this graph. + rust_symbols : dict, optional + A shared cache mapping symbol ids to already-allocated handles. Pass a + single ``{}`` when converting several expressions over one graph so a + shared subgraph is allocated once and stays shared in Rust. + + Returns + ------- + :class:`pybamm.rust.Expr` + Handle to the root node of the converted expression. + """ + if rust_symbols is None: + rust_symbols = {} + return self._to_rust_inner(graph, rust_symbols) + + def _to_rust_inner(self, graph, rust_symbols): + cached = rust_symbols.get(self.id, None) + if cached is not None: + return cached + result = self._to_rust(graph, rust_symbols) + rust_symbols[self.id] = result + return result + + def _to_rust(self, graph, rust_symbols): + raise TypeError( + f"Cannot convert symbol of type '{type(self).__name__}' to the Rust " + "backend. Use a CasADi-backed model (set " + "`model.convert_to_format = 'casadi'`) for models using " + "this symbol." + ) + + def _children_to_rust(self, graph, rust_symbols): + return [child._to_rust_inner(graph, rust_symbols) for child in self.children] + def _children_for_copying(self, children: list[Symbol] | None = None) -> Symbol: """ Gets existing children for a symbol being copied if they aren't provided. diff --git a/packages/pybamm/src/pybamm/expression_tree/unary_operators.py b/packages/pybamm/src/pybamm/expression_tree/unary_operators.py index b6a50821ba..9a8299fe57 100644 --- a/packages/pybamm/src/pybamm/expression_tree/unary_operators.py +++ b/packages/pybamm/src/pybamm/expression_tree/unary_operators.py @@ -105,6 +105,44 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): ] return self._casadi_evaluate(converted_child) + def _to_rust(self, graph, rust_symbols): + (child,) = self._children_to_rust(graph, rust_symbols) + if isinstance(self, Negate): + return graph.neg(child) + elif isinstance(self, AbsoluteValue): + return graph.abs(child) + elif isinstance(self, Sign): + return graph.sign(child) + elif isinstance(self, Floor): + return graph.floor(child) + elif isinstance(self, Ceiling): + return graph.ceiling(child) + elif isinstance(self, Index): + s = self.slice + if s.step is not None and s.step != 1: + raise NotImplementedError( + "Cannot convert strided Index (step != 1) to Rust" + ) + size = self.children[0].size + start = 0 if s.start is None else s.start + stop = size if s.stop is None else s.stop + if start < 0: + start += size + if stop < 0: + stop += size + # Clamp to [0, size] to match numpy/CasADi slice semantics and + # guarantee non-negative bounds for the usize Rust builder. + start = max(0, min(start, size)) + stop = max(0, min(stop, size)) + return graph.index(child, start, stop) + elif isinstance(self, (ExplicitTimeIntegral, NotConstant)): + # Pass-through operators: actual computation happens in ProcessedVariable + return child + else: + raise TypeError( + f"Cannot convert unary operator of type '{type(self)}' to Rust" + ) + def evaluate( self, t: float | None = None, diff --git a/packages/pybamm/src/pybamm/expression_tree/vector_field.py b/packages/pybamm/src/pybamm/expression_tree/vector_field.py index 2268a56434..5dbc2dbc48 100644 --- a/packages/pybamm/src/pybamm/expression_tree/vector_field.py +++ b/packages/pybamm/src/pybamm/expression_tree/vector_field.py @@ -82,6 +82,10 @@ def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols): *self._children_to_casadi(t, y, y_dot, inputs, casadi_symbols) ) + def _to_rust(self, graph, rust_symbols): + """See :meth:`pybamm.Symbol._to_rust()`.""" + return graph.concat(self._children_to_rust(graph, rust_symbols)) + def evaluates_on_edges(self, dimension: str) -> bool: statuses = [c.evaluates_on_edges(dimension) for c in self.components] if all(statuses): diff --git a/packages/pybamm/src/pybamm/models/base_model.py b/packages/pybamm/src/pybamm/models/base_model.py index 21ff2f0f75..2296128546 100644 --- a/packages/pybamm/src/pybamm/models/base_model.py +++ b/packages/pybamm/src/pybamm/models/base_model.py @@ -93,6 +93,8 @@ class BaseModel: - "python": convert to Python code for evaluating `evaluate(t, y)` on expressions. - "casadi": convert to CasADi expression tree for Jacobian calculation. - "jax": convert to JAX expression tree. + - "rust": convert to a compiled Rust expression graph for + solve-time evaluation (IDAKLU, diffsol, scipy, algebraic solvers). Default is "casadi". is_discretised: bool @@ -102,6 +104,9 @@ class BaseModel: different submodels in the full concatenated solution vector. """ + _DEFAULT_CONVERT_TO_FORMAT = "rust" + _VALID_CONVERT_TO_FORMATS = (None, "python", "casadi", "jax", "rust") + def __init__(self, name="Unnamed model"): self.name = name self._options = {} @@ -141,7 +146,7 @@ def __init__(self, name="Unnamed model"): # Default behaviour is to use the jacobian self.use_jacobian = True - self.convert_to_format = "casadi" + self.convert_to_format = self._DEFAULT_CONVERT_TO_FORMAT # Model is not initially discretised or parameterised self.is_discretised = False @@ -177,6 +182,7 @@ def deserialise(cls, properties: dict): @classmethod def generic_deserialise(cls, instance, properties): # Initialise model with stored variables that have already been discretised + instance.convert_to_format = properties.get("convert_to_format", "casadi") instance._concatenated_rhs = properties["concatenated_rhs"] instance._concatenated_algebraic = properties["concatenated_algebraic"] instance._concatenated_initial_conditions = properties[ @@ -237,6 +243,23 @@ def name(self): def name(self, value): self._name = value + @property + def convert_to_format(self): + return self._convert_to_format + + @convert_to_format.setter + def convert_to_format(self, value): + if value not in self._VALID_CONVERT_TO_FORMATS: + valid = ", ".join(repr(v) for v in self._VALID_CONVERT_TO_FORMATS) + raise ValueError(f"convert_to_format must be one of {valid}, got {value!r}") + self._convert_to_format = value + + @property + def uses_stacked_inputs(self): + """Whether solver evaluators expect inputs stacked into one flat + vector (casadi/rust convention) rather than a dict (python/jax).""" + return self._convert_to_format in ("casadi", "rust") + @property def rhs(self): """Returns a dictionary mapping expressions (variables) to expressions that represent diff --git a/packages/pybamm/src/pybamm/rust/__init__.py b/packages/pybamm/src/pybamm/rust/__init__.py new file mode 100644 index 0000000000..78b27ad506 --- /dev/null +++ b/packages/pybamm/src/pybamm/rust/__init__.py @@ -0,0 +1,64 @@ +"""Access point for the compiled Rust extension. + +``pybamm.rust._core`` is the built extension module; import its classes from +``pybamm.rust`` rather than reaching into ``_core``, so the private module stays +free to move. Importing this package fails if the extension was not built, which +is why callers that must tolerate its absence import it inside a function. +""" + +import contextlib +import os +import sys + + +@contextlib.contextmanager +def _global_symbol_visibility(): + """Open extension modules into the process-global symbol scope on Linux. + + pybammsolvers cannot link against the Rust core (the dependency only flows + the other way), so it resolves the FFI entry points with + ``dlsym(RTLD_DEFAULT)``. That only finds symbols from libraries loaded with + ``RTLD_GLOBAL``, and CPython defaults to ``RTLD_LOCAL`` on Linux. macOS + exports them globally already, so this is a no-op there. + """ + if not sys.platform.startswith("linux"): + yield + return + + previous_flags = sys.getdlopenflags() + sys.setdlopenflags(previous_flags | os.RTLD_GLOBAL) + try: + yield + finally: + sys.setdlopenflags(previous_flags) + + +with _global_symbol_visibility(): + from pybamm.rust._core import ( + CompiledFunction, + CompiledFunctionGroup, + CompiledJacobian, + CompiledModel, + EvaluatorPool, + Expr, + ExprGraph, + PreparedSolver, + SolveOutcome, + SolverStatistics, + default_solver_options, + ) + +# Without this, ``contextlib``/``os``/``sys`` are also public names here. +__all__ = [ + "CompiledFunction", + "CompiledFunctionGroup", + "CompiledJacobian", + "CompiledModel", + "EvaluatorPool", + "Expr", + "ExprGraph", + "PreparedSolver", + "SolveOutcome", + "SolverStatistics", + "default_solver_options", +] diff --git a/packages/pybamm/src/pybamm/rust/_core.pyi b/packages/pybamm/src/pybamm/rust/_core.pyi new file mode 100644 index 0000000000..cfa40c075b --- /dev/null +++ b/packages/pybamm/src/pybamm/rust/_core.pyi @@ -0,0 +1,674 @@ +"""Hand-maintained stubs for the compiled extension ``pybamm.rust._core``. + +Any change to the Python-visible API in ``packages/pybamm-rust/pybamm-python/src`` +must update this file in the same commit; ``mypy.stubtest`` pins names, arities +and defaults against the built extension (``tests/unit/test_rust_stubs.py``), +while the types themselves are review-enforced. +""" + +from collections.abc import Callable, Sequence +from typing import Literal, TypeAlias, TypedDict, final + +import numpy as np +import numpy.typing as npt +from scipy.sparse import csc_matrix + +__all__ = [ + "CompiledFunction", + "CompiledFunctionGroup", + "CompiledJacobian", + "CompiledModel", + "EvaluatorPool", + "Expr", + "ExprGraph", + "PreparedSolver", + "SolveOutcome", + "SolverStatistics", + "_pool_ids", + "default_solver_options", +] + +_FloatArray: TypeAlias = npt.NDArray[np.float64] +# The packed parameter vector, or a {name: value} mapping packed on entry. +_Params: TypeAlias = dict[str, float | _FloatArray] | _FloatArray +# Arguments extracted into Rust vectors: any sequence of floats converts. +_FloatSequence: TypeAlias = Sequence[float] | _FloatArray +_IntSequence: TypeAlias = Sequence[int] | npt.NDArray[np.integer] +# Time grids convert from any 1-D numeric array-like (dtype changes allowed). +_TimeGrid: TypeAlias = ( + Sequence[float] | npt.NDArray[np.floating] | npt.NDArray[np.integer] +) + +class _SolverOptions(TypedDict): + """Integrator tuning dict; every key is required (defaults live Python-side).""" + + max_nonlinear_solver_iterations: int + max_error_test_failures: int + max_nonlinear_solver_failures: int + nonlinear_solver_tolerance: float + min_timestep: float + max_timestep_growth: float | None + min_timestep_growth: float | None + max_timestep_shrink: float | None + min_timestep_shrink: float | None + update_jacobian_after_steps: int + update_rhs_jacobian_after_steps: int + threshold_to_update_jacobian: float + threshold_to_update_rhs_jacobian: float + pi_control_proportional: float + pi_control_integral: float + +class _JacobianStats(TypedDict): + """Assembly stats returned by :meth:`CompiledModel.jacobian_stats`.""" + + strategy: str + n_colors: int + nnz: int + n_dense_rows: int + n_dense_row_candidates: int + n_constant_entries: int + n_swept_columns: int + jac_lane_width: int + dense_row_entries: int + dense_row_tape_instructions: int + split_eval_primal_instructions: int + split_eval_total_instructions: int + split_eval_raw_instructions: int + split_eval_dispatch_count: int + branch_block_lens: tuple[int, ...] + +@final +class ExprGraph: + """Arena of expression nodes; the build surface for every compiled artifact.""" + + def __new__(cls) -> ExprGraph: ... + @property + def n_nodes(self) -> int: + """Number of nodes in the expression arena.""" + + def n_inputs(self) -> int: + """Total packed width of every registered input (sum of widths, not names).""" + + def scalar(self, value: float) -> Expr: ... + def time(self) -> Expr: ... + def state_vector(self, start: int, end: int) -> Expr: ... + def state_vector_dot(self, start: int, end: int) -> Expr: ... + def input_parameter(self, name: str, width: int = 1) -> Expr: + """Register (or re-look-up) a named input; re-registering must repeat the width.""" + + def array(self, data: _FloatArray) -> Expr: ... + def dense_matrix(self, data: _FloatArray, rows: int, cols: int) -> Expr: + """Dense matrix constant from flat row-major ``data``.""" + + def sparse_matrix( + self, + indptr: _IntSequence, + indices: _IntSequence, + data: _FloatArray, + rows: int, + cols: int, + ) -> Expr: + """CSR matrix constant.""" + + def add(self, a: Expr, b: Expr) -> Expr: ... + def sub(self, a: Expr, b: Expr) -> Expr: ... + def mul(self, a: Expr, b: Expr) -> Expr: ... + def div(self, a: Expr, b: Expr) -> Expr: ... + def neg(self, a: Expr) -> Expr: ... + def abs(self, a: Expr) -> Expr: ... + def pow(self, a: Expr, b: Expr) -> Expr: ... + def sqrt(self, a: Expr) -> Expr: ... + def exp(self, a: Expr) -> Expr: ... + def log(self, a: Expr) -> Expr: ... + def sin(self, a: Expr) -> Expr: ... + def cos(self, a: Expr) -> Expr: ... + def tanh(self, a: Expr) -> Expr: ... + def sinh(self, a: Expr) -> Expr: ... + def cosh(self, a: Expr) -> Expr: ... + def arcsinh(self, a: Expr) -> Expr: ... + def arctan(self, a: Expr) -> Expr: ... + def erf(self, a: Expr) -> Expr: ... + def sign(self, a: Expr) -> Expr: ... + def floor(self, a: Expr) -> Expr: ... + def ceiling(self, a: Expr) -> Expr: ... + def max_reduce(self, a: Expr) -> Expr: ... + def min_reduce(self, a: Expr) -> Expr: ... + def matmul(self, a: Expr, b: Expr) -> Expr: ... + def minimum(self, a: Expr, b: Expr) -> Expr: ... + def maximum(self, a: Expr, b: Expr) -> Expr: ... + def modulo(self, a: Expr, b: Expr) -> Expr: ... + def hypot(self, a: Expr, b: Expr) -> Expr: ... + def equal_heaviside(self, a: Expr, b: Expr) -> Expr: ... + def not_equal_heaviside(self, a: Expr, b: Expr) -> Expr: ... + def equality(self, a: Expr, b: Expr) -> Expr: ... + def index(self, child: Expr, start: int, end: int) -> Expr: ... + def concat(self, children: Sequence[Expr]) -> Expr: ... + def conditional(self, selector: Expr, branches: Sequence[Expr]) -> Expr: + """Branch ``i`` is active when ``i - 0.5 < selector < i + 0.5`` (1-based).""" + + def interpolant_1d_linear( + self, x_data: _FloatSequence, y_data: _FloatSequence, child: Expr + ) -> Expr: ... + def interpolant_1d_cubic( + self, breakpoints: _FloatSequence, coeffs: _FloatSequence, child: Expr + ) -> Expr: + """``coeffs`` is flat row-major ``[c0..c3]`` groups, one per segment.""" + + def interpolant_nd( + self, + breakpoints: Sequence[_FloatSequence], + coeffs: _FloatSequence, + order: int, + children: Sequence[Expr], + ) -> Expr: + """Tensor-product interpolant over 2 or 3 axes; one child per axis.""" + + def eval_to_float( + self, + expr: Expr, + t: float, + y: _FloatSequence, + y_dot: _FloatSequence, + inputs: _FloatSequence, + ) -> float: + """Test/debug helper; inputs are not length-validated.""" + + def eval_to_array( + self, + expr: Expr, + t: float, + y: _FloatArray, + y_dot: _FloatArray, + inputs: _FloatSequence, + ) -> _FloatArray: + """Test/debug helper; inputs are not length-validated.""" + + def compile( + self, expr: Expr, name: str | None = None, n_states: int | None = None + ) -> CompiledFunction: + """Compile ``expr`` into an immutable, shareable :class:`CompiledFunction`.""" + + def compile_group( + self, + outputs: dict[str, Expr], + name: str | None = None, + n_states: int | None = None, + ) -> CompiledFunctionGroup: + """Compile named outputs into ONE shared tape with cross-output CSE.""" + + def dump_dag( + self, expr: Expr, path: str, model_name: str, n_states: int, n_params: int + ) -> None: + """Serialize the graph rooted at ``expr`` to ``path`` (debug snapshot).""" + + def __getstate__(self) -> bytes: ... + def __setstate__(self, state: bytes) -> None: ... + def __getnewargs__(self) -> tuple[()]: ... + +@final +class Expr: + """Handle to one node in an :class:`ExprGraph`; only combines within its graph.""" + + @property + def id(self) -> int: + """Raw node id within the owning graph.""" + + def __add__(self, other: Expr, /) -> Expr: ... + def __sub__(self, other: Expr, /) -> Expr: ... + def __mul__(self, other: Expr, /) -> Expr: ... + def __truediv__(self, other: Expr, /) -> Expr: ... + def __pow__(self, other: Expr, modulo: object = None, /) -> Expr: ... + def __neg__(self) -> Expr: ... + +@final +class CompiledFunction: + """Prepared evaluation artifact for one expression: eval, JVP and jacobians.""" + + def __call__( + self, + t: float, + y: _FloatArray, + p: _Params, + y_dot: _FloatArray | None = None, + ) -> _FloatArray: ... + def eval( + self, + t: float, + y: _FloatArray, + p: _Params, + y_dot: _FloatArray | None = None, + ) -> _FloatArray: + """Alias for :meth:`__call__`.""" + + def eval_into( + self, + t: float, + y: _FloatArray, + p: _Params, + out: _FloatArray, + y_dot: _FloatArray | None = None, + ) -> None: + """Evaluate into pre-allocated ``out`` (length ``output_len``).""" + + def pack(self, mapping: dict[str, float | _FloatArray]) -> _FloatArray: + """Pack a ``{name: value}`` mapping into the stacked input layout.""" + + def jvp( + self, + t: float, + y: _FloatArray, + p: _Params, + vy: _FloatArray, + vp: _FloatArray | None = None, + ) -> _FloatArray: + """Forward-mode JVP: ``df/dy @ vy`` (+ ``df/dp @ vp`` when given).""" + + def jacobian(self, wrt: Literal["y", "p"] = "y") -> CompiledJacobian: + """Lazy, cached-per-``wrt`` prepared jacobian.""" + + def eval_trajectory( + self, ts: _TimeGrid, y_traj: _FloatArray, p: _Params + ) -> _FloatArray: + """Evaluate per time column; returns ``(output_len, n_t)`` F-contiguous.""" + + def jvp_trajectory( + self, + ts: _TimeGrid, + y_traj: _FloatArray, + p: _Params, + vy_traj: _FloatArray, + vp: _FloatArray | None = None, + ) -> _FloatArray: + """Per-column JVP along a trajectory; returns ``(output_len, n_t)``.""" + + def eval_trajectory_hermite( + self, + t_query: _TimeGrid, + ts: _TimeGrid, + ys: _FloatArray, + yps: _FloatArray, + p: _Params, + ) -> _FloatArray: + """Cubic-Hermite reconstruct the state at ``t_query``, then evaluate.""" + + @property + def input_names(self) -> tuple[str, ...]: ... + @property + def n_inputs(self) -> int: + """Registered-name count (the ``vp`` seed length), NOT the packed width.""" + + @property + def n_states(self) -> int: ... + @property + def output_len(self) -> int: ... + @property + def uses_y_dot(self) -> bool: ... + @property + def name(self) -> str | None: ... + @property + def n_instructions(self) -> int: + """Tape length excluding conditional branch blocks (one dispatch each).""" + + @property + def n_instructions_total(self) -> int: + """Raw tape length, branch blocks included.""" + + @property + def n_dispatches(self) -> int: ... + @property + def branch_block_lens(self) -> tuple[int, ...]: ... + @staticmethod + def _rebuild( + graph: ExprGraph, root: int, name: str | None, n_states: int | None + ) -> CompiledFunction: ... + def __reduce__( + self, + ) -> tuple[ + Callable[..., CompiledFunction], + tuple[ExprGraph, int, str | None, int | None], + ]: ... + +@final +class CompiledFunctionGroup: + """Named outputs compiled into ONE shared tape; results sliced per output.""" + + def __call__(self, t: float, y: _FloatArray, p: _Params) -> list[_FloatArray]: ... + def eval_trajectory( + self, ts: _TimeGrid, y_traj: _FloatArray, p: _Params + ) -> list[_FloatArray]: + """One ``(output_len_i, n_t)`` F-contiguous array per output, in order.""" + + def eval_trajectory_hermite( + self, + t_query: _TimeGrid, + ts: _TimeGrid, + ys: _FloatArray, + yps: _FloatArray, + p: _Params, + ) -> list[_FloatArray]: + """Cubic-Hermite reconstruct at ``t_query``, then evaluate and slice.""" + + def pack(self, mapping: dict[str, float | _FloatArray]) -> _FloatArray: + """Pack a ``{name: value}`` mapping into the stacked input layout.""" + + @property + def names(self) -> tuple[str, ...]: + """Output names in declared order.""" + + @property + def output_lens(self) -> list[int]: ... + @property + def input_names(self) -> tuple[str, ...]: ... + @property + def n_inputs(self) -> int: + """Registered-name count (the ``vp`` seed length), NOT the packed width.""" + + @property + def n_states(self) -> int: ... + @property + def output_len(self) -> int: + """Total length across all outputs.""" + + @property + def uses_y_dot(self) -> bool: ... + @property + def n_instructions(self) -> int: ... + @property + def n_instructions_total(self) -> int: ... + @property + def branch_block_lens(self) -> tuple[int, ...]: ... + @property + def name(self) -> str | None: ... + +@final +class CompiledJacobian: + """Prepared sparse jacobian: colored JVP sweeps assembled into scipy CSC.""" + + def __call__(self, t: float, y: _FloatArray, p: _Params) -> csc_matrix: + """Assemble and return a ``scipy.sparse.csc_matrix``.""" + + def sparsity(self) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]: + """CSC pattern as ``(indptr, indices)``, cached read-only int32 arrays.""" + + @property + def nnz(self) -> int: ... + @property + def shape(self) -> tuple[int, int]: ... + @property + def n_colors(self) -> int: ... + @property + def n_dense_rows(self) -> int: ... + @property + def wrt(self) -> Literal["y", "p"]: ... + +@final +class CompiledModel: + """Compiled DAE bundle: residual/jacobian FFI surface plus shareable views.""" + + @staticmethod + def from_expr( + graph: ExprGraph, + expr: Expr, + mass_data: _FloatArray, + mass_indptr: npt.NDArray[np.int64], + mass_indices: npt.NDArray[np.int64], + n_inputs: int = 0, + sens_param_indices: Sequence[int] = ..., + output_exprs: Sequence[Expr] = ..., + algebraic_expr: Expr | None = None, + algebraic_variable_indices: Sequence[int] = ..., + event_exprs: Sequence[Expr] = ..., + ) -> CompiledModel: + """Build from an rhs expression and a CSR mass matrix.""" + + @staticmethod + def _rebuild( + graph: ExprGraph, + rhs_root: int, + output_roots: Sequence[int], + event_roots: Sequence[int], + algebraic_root: int | None, + algebraic_variable_indices: Sequence[int], + mass_data: Sequence[float], + mass_indptr: Sequence[int], + mass_indices: Sequence[int], + n_inputs: int, + sens_param_indices: Sequence[int], + ) -> CompiledModel: ... + def __reduce__( + self, + ) -> tuple[ + Callable[..., CompiledModel], + tuple[ + ExprGraph, + int, + list[int], + list[int], + int | None, + list[int], + list[float], + list[int], + list[int], + int, + list[int], + ], + ]: ... + @property + def rhs(self) -> CompiledFunction: + """Primal ``f(t, y, p)`` as a shareable view.""" + + @property + def graph(self) -> ExprGraph: + """The retained derivation arena (observation lowers new roots into it).""" + + @property + def jacobian(self) -> CompiledJacobian: + """Pure ``df/dy`` (``cj = 0``, no mass), composed over the bundle's tapes.""" + + @property + def outputs(self) -> list[CompiledFunction]: ... + @property + def events(self) -> list[CompiledFunction]: ... + @property + def algebraic_residual(self) -> CompiledFunction | None: + """``g(t, y, p)`` as a shareable view, or ``None`` for ODEs.""" + + @property + def algebraic_jacobian(self) -> CompiledJacobian | None: + """``dg/dy_alg`` (``n_algebraic`` square), or ``None`` for ODEs.""" + + def eval_residual( + self, + t: float, + y: _FloatArray, + yp: _FloatArray, + inputs: _FloatArray, + ) -> _FloatArray: + """DAE residual ``r = M @ yp - f(t, y)`` as a new array (packed inputs only).""" + + def eval_residual_into( + self, + t: float, + y: _FloatArray, + yp: _FloatArray, + inputs: _FloatArray, + output: _FloatArray, + ) -> None: + """DAE residual into pre-allocated ``output`` (length ``n_states``).""" + + def assemble_jacobian_csc_into( + self, + t: float, + y: _FloatArray, + cj: float, + inputs: _FloatArray, + jac_data: _FloatArray, + ) -> None: + """Assemble ``df/dy - cj * M`` into a pre-allocated CSC data buffer.""" + + def algebraic_ids(self) -> _FloatArray: + """IDA-convention ids: ``1.0`` differential, ``0.0`` algebraic, per state.""" + + def sparsity_pattern( + self, + ) -> tuple[npt.NDArray[np.uintp], npt.NDArray[np.uintp]]: + """CSR ``(indptr, indices)`` of ``df/dy``.""" + + def csc_sparsity_pattern( + self, + ) -> tuple[npt.NDArray[np.uintp], npt.NDArray[np.uintp]]: + """CSC ``(colptr, rowind)`` for KLU compatibility.""" + + def algebraic_jacobian_sparsity_pattern( + self, + ) -> tuple[npt.NDArray[np.uintp], npt.NDArray[np.uintp]]: + """COO ``(rows, cols)`` of the algebraic jacobian; raises for ODEs.""" + + def constant_jacobian_entries( + self, + ) -> tuple[npt.NDArray[np.uintp], _FloatArray]: + """``(csc_idx, value)`` for entries proved constant at compile time.""" + + def jacobian_stats(self) -> _JacobianStats: ... + def evaluator_pool(self, n: int) -> EvaluatorPool: + """``n`` independent evaluators over this model's tape (``n >= 1``).""" + + @property + def n_states(self) -> int: ... + @property + def n_inputs(self) -> int: + """Packed input width (core's ``n_params``), matching the C ABI naming.""" + + @property + def has_algebraic(self) -> bool: ... + @property + def n_algebraic(self) -> int: ... + @property + def algebraic_jacobian_nnz(self) -> int: ... + @property + def n_sens_params(self) -> int: ... + @property + def n_outputs(self) -> int: ... + @property + def n_events(self) -> int: ... + @property + def output_len(self) -> int: ... + @property + def n_colors(self) -> int: ... + @property + def nnz(self) -> int: ... + @property + def jacobian_strategy(self) -> str: ... + +@final +class EvaluatorPool: + """N independent evaluators over one shared tape, one per parallel solver.""" + + def as_ptr(self, index: int) -> int: + """Address of evaluator ``index``; the pool must outlive every use of it. + + Each address is handed out at most once; a second take raises + ``RuntimeError``, so two solvers can never share one evaluator. + """ + + def __len__(self) -> int: ... + +@final +class SolverStatistics: + """BDF solver statistics for one solve.""" + + number_of_steps: int + number_of_linear_solver_setups: int + number_of_nonlinear_solver_iterations: int + number_of_nonlinear_solver_fails: int + number_of_error_test_failures: int + number_of_linear_solver_setups_from_checkpoint: int + number_of_linear_solver_setups_from_first_convergence_fail: int + number_of_linear_solver_setups_from_second_convergence_fail: int + number_of_linear_solver_setups_from_error_test_fail: int + number_of_linear_solver_setups_from_step_success: int + ic_time_secs: float + solver_setup_time_secs: float + integration_time_secs: float + sens_error_control_relaxed: bool + +@final +class SolveOutcome: + """What one diffsol solve returns, whichever payloads it was asked for.""" + + flag: int + """0 = success, 1 = root found.""" + t_event: float | None + statistics: SolverStatistics + @property + def t(self) -> _FloatArray: ... + @property + def y(self) -> _FloatArray: + """Trajectory, shape ``(n_rows, n_times)``; rows are states, or the + model's output variables when the solve asked for ``outputs``.""" + + @property + def yp(self) -> _FloatArray | None: + """Row derivatives ``(n_rows, n_times)``, or ``None`` without + ``store_yp`` or on an ``outputs`` solve.""" + + @property + def yS(self) -> list[_FloatArray] | None: + """Per-parameter sensitivities, each matching the flat layout of ``y``, + or ``None`` when the solve asked for none.""" + + @property + def y_event(self) -> _FloatArray | None: + """The full state where the trajectory ends, never an outputs row.""" + +@final +class PreparedSolver: + """Prepare-once/execute-many diffsol solver for repeated integrations.""" + + def __new__( + cls, + model: CompiledModel, + rtol: float = 1e-6, + atol: float | _FloatArray = ..., + sens_atol_factor: float = 1e-3, + options: _SolverOptions | None = None, + store_yp: bool = False, + ) -> PreparedSolver: ... + def solve( + self, + t_eval: _FloatArray, + t_stop: _FloatArray, + y0: _FloatArray, + inputs: _FloatArray, + *, + outputs: bool = False, + sensitivities: bool = False, + y0_sens: _FloatArray | None = None, + ) -> SolveOutcome: + """Integrate over ``t_eval``, landing exactly on each ``t_stop``. + + ``outputs`` reports the model's output variables rather than the full + state; ``sensitivities`` adds the forward-sensitivity blocks, seeded by + flat ``dy0/dp`` in ``y0_sens``. + """ + + def solve_batch( + self, + t_eval: _FloatArray, + t_stop: _FloatArray, + y0: _FloatArray, + inputs: _FloatArray, + num_threads: int, + *, + outputs: bool = False, + sensitivities: bool = False, + y0_sens: _FloatArray | None = None, + ) -> list[SolveOutcome | Exception]: + """One row of ``y0``/``inputs``/``y0_sens`` per set, answering one + shared request; a failed set yields its exception instance, unraised.""" + +def default_solver_options() -> _SolverOptions: + """Diffsol's own integrator defaults, the dict PyBaMM overlays onto.""" + +def _pool_ids() -> dict[int, int]: + """Cached rayon pools as ``{thread_count: pool identity}`` (test introspection).""" diff --git a/packages/pybamm/src/pybamm/simulation/eis_simulation.py b/packages/pybamm/src/pybamm/simulation/eis_simulation.py index 0c7ebdfdef..ab9504ae30 100644 --- a/packages/pybamm/src/pybamm/simulation/eis_simulation.py +++ b/packages/pybamm/src/pybamm/simulation/eis_simulation.py @@ -183,13 +183,17 @@ def _build_matrix_problem(self, inputs_dict=None): model = self._built_model inputs_dict = inputs_dict or {} - # Convert inputs to casadi format for Jacobian evaluation + # Inputs in the format the compiled Jacobian expects for each backend if model.convert_to_format == "casadi": from casadi import vertcat - casadi_inputs = vertcat(*inputs_dict.values()) if inputs_dict else [] + eval_inputs = vertcat(*inputs_dict.values()) if inputs_dict else [] + elif model.convert_to_format == "rust": + from pybamm.solvers.base_solver import stack_inputs + + eval_inputs = stack_inputs(inputs_dict, "rust") else: - casadi_inputs = inputs_dict + eval_inputs = inputs_dict # Only compile Jacobian/model functions on first call; the compiled # functions persist on the model and work with any inputs/y0 values. @@ -198,8 +202,10 @@ def _build_matrix_problem(self, inputs_dict=None): solver.set_up(model, inputs=inputs_dict) y0 = model.concatenated_initial_conditions.evaluate(0, inputs=inputs_dict) - J_sparse = model.jac_rhs_algebraic_eval(0, y0, casadi_inputs).sparse() - neg_J = -J_sparse if isinstance(J_sparse, csc_matrix) else -csc_matrix(J_sparse) + J = model.jac_rhs_algebraic_eval(0, y0, eval_inputs) + if model.convert_to_format == "casadi": + J = J.sparse() + neg_J = -J if isinstance(J, csc_matrix) else -csc_matrix(J) # M and b are independent of operating point and cached after first call, # but we have a defensive guard to invalidate it if the state vector size changes diff --git a/packages/pybamm/src/pybamm/solvers/__init__.py b/packages/pybamm/src/pybamm/solvers/__init__.py index cfb1727300..05d7beae66 100644 --- a/packages/pybamm/src/pybamm/solvers/__init__.py +++ b/packages/pybamm/src/pybamm/solvers/__init__.py @@ -5,6 +5,7 @@ "casadi_algebraic_solver", "casadi_solver", "composite_solver", + "diffsol_solver", "dummy_solver", "idaklu_jax", "idaklu_solver", diff --git a/packages/pybamm/src/pybamm/solvers/algebraic_solver.py b/packages/pybamm/src/pybamm/solvers/algebraic_solver.py index c84bffd182..6764a779b8 100644 --- a/packages/pybamm/src/pybamm/solvers/algebraic_solver.py +++ b/packages/pybamm/src/pybamm/solvers/algebraic_solver.py @@ -7,6 +7,7 @@ from scipy.sparse import issparse import pybamm +from pybamm.solvers.base_solver import stack_inputs class AlgebraicSolver(pybamm.BaseSolver): @@ -105,8 +106,8 @@ def _integrate_single(self, model, t_eval, inputs_dict, y0): as well as various diagnostic messages. """ inputs_dict = inputs_dict or {} - if model.convert_to_format == "casadi": - inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) + if model.uses_stacked_inputs: + inputs = stack_inputs(inputs_dict, model.convert_to_format) else: inputs = inputs_dict @@ -220,7 +221,9 @@ def root_norm(y): else: def jac_norm(y, jac_fn=jac_fn): - return np.sum(2 * root_fun(y) * jac_fn(y), 0) + # Gradient of sum(f**2) is 2*J^T*f; keep f as a column + # so numpy broadcasts rows like casadi DM does. + return np.sum(2 * root_fun(y)[:, np.newaxis] * jac_fn(y), 0) if self.method == "minimize": method = None diff --git a/packages/pybamm/src/pybamm/solvers/base_solver.py b/packages/pybamm/src/pybamm/solvers/base_solver.py index b3c7ab5445..6414df4312 100644 --- a/packages/pybamm/src/pybamm/solvers/base_solver.py +++ b/packages/pybamm/src/pybamm/solvers/base_solver.py @@ -8,6 +8,7 @@ import casadi import numpy as np +import numpy.typing as npt import pybamm from pybamm import ParameterValues @@ -25,8 +26,10 @@ class BaseSolver: The method to use for integration, specific to each solver rtol : float, optional The relative tolerance for the solver (default is 1e-6). - atol : float, optional - The absolute tolerance for the solver (default is 1e-6). + atol : float or :class:`numpy.ndarray`, optional + The absolute tolerance for the solver, either shared by every state or + one entry per state (default is 1e-6). Per-state tolerances are honoured + by :class:`pybamm.IDAKLUSolver` and :class:`pybamm.DiffsolSolver`. root_method : str or pybamm algebraic solver class, optional The method to use to find initial conditions (for DAE solvers). Default is "nonlinear_solver", which uses a custom Newton solver for consistent @@ -60,6 +63,14 @@ class BaseSolver: Default is False. """ + #: Subclasses that integrate via a whole-model CompiledModel artifact + #: (IDAKLUSolver, DiffsolSolver) set this True to skip per-group lowering. + _integrates_via_compiled_model = False + + # True once this solver's native (Rust) observation backend reaches parity + # with CasADi; combined with `model.convert_to_format == "rust"` to activate. + _observes_via_compiled_model = False + def __init__( self, method=None, @@ -181,8 +192,9 @@ def root_method(self): @root_method.setter def root_method(self, method): if method == "nonlinear_solver": - # use the tighter of the two tolerances - atol = min(self.root_tol, self.atol) + # use the tighter of the two tolerances; a per-state atol contributes + # its tightest entry, since the root solver takes a single tolerance + atol = min(self.root_tol, float(np.min(self.atol))) method = pybamm.NonlinearSolver( atol=atol, rtol=self.rtol, on_failure=self.on_failure ) @@ -217,6 +229,66 @@ def copy(self): new_solver._model_set_up = {} return new_solver + def _check_atol_type(self, atol, model): + """Widen an absolute tolerance to one entry per state. + + Parameters + ---------- + atol : float or array-like + One tolerance shared by every state, or one per state. + model : :class:`pybamm.BaseModel` + The discretised model, whose state count a per-state ``atol`` is + checked against. + + Returns + ------- + :class:`numpy.ndarray` + Float64 tolerances with shape ``(model.len_rhs_and_alg,)``. + + Raises + ------ + :class:`pybamm.SolverError` + If ``atol`` is neither a number nor a sequence of numbers of the + model's width. + """ + if isinstance(atol, numbers.Real): + return np.full(model.len_rhs_and_alg, float(atol)) + # A list arrives from a to_config round trip, where JSON has no arrays. + elif isinstance(atol, (np.ndarray, list, tuple)): + try: + atol = np.asarray(atol, dtype=np.float64) + except (TypeError, ValueError) as error: + raise pybamm.SolverError( + "Absolute tolerances must all be numbers" + ) from error + if atol.shape != (model.len_rhs_and_alg,): + raise pybamm.SolverError( + f"Absolute tolerances have shape {atol.shape} but " + f"({model.len_rhs_and_alg},) was expected (one per state)" + ) + return atol + else: + raise pybamm.SolverError( + "Absolute tolerances must be a float or one value per state" + ) + + def _lowers_via_compiled_model(self, model) -> bool: + """Whether one shared Rust lowering serves every evaluator of this solve. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model being set up. + + Returns + ------- + bool + True when the solver integrates a ``CompiledModel`` lowered from a + Rust model, so lowering an expression here would build a second graph + beside the solver's own. + """ + return self._integrates_via_compiled_model and model.convert_to_format == "rust" + def set_up( self, model: pybamm.BaseModel, @@ -272,32 +344,39 @@ def set_up( pybamm.logger.info("Finish solver set-up") return - # Process rhs, algebraic, residual and event expressions - # and wrap in callables - rhs, jac_rhs, jacp_rhs, jac_rhs_action = process( - model.concatenated_rhs, "RHS", vars_for_processing - ) - - algebraic, jac_algebraic, jacp_algebraic, jac_algebraic_action = process( - model.concatenated_algebraic, "algebraic", vars_for_processing - ) - - # combine rhs and algebraic functions - if len(model.rhs) == 0: - rhs_algebraic = model.concatenated_algebraic - elif len(model.algebraic) == 0: - rhs_algebraic = model.concatenated_rhs + # Each group here would be a second graph beside the solver's own. + skip_per_group = self._lowers_via_compiled_model(model) + if skip_per_group: + rhs = jac_rhs = jacp_rhs = jac_rhs_action = None + algebraic = jac_algebraic = jacp_algebraic = jac_algebraic_action = None + rhs_algebraic = jac_rhs_algebraic = jacp_rhs_algebraic = ( + jac_rhs_algebraic_action + ) = None else: - rhs_algebraic = pybamm.NumpyConcatenation( - model.concatenated_rhs, model.concatenated_algebraic + rhs, jac_rhs, jacp_rhs, jac_rhs_action = process( + model.concatenated_rhs, "RHS", vars_for_processing ) - ( - rhs_algebraic, - jac_rhs_algebraic, - jacp_rhs_algebraic, - jac_rhs_algebraic_action, - ) = process(rhs_algebraic, "rhs_algebraic", vars_for_processing) + algebraic, jac_algebraic, jacp_algebraic, jac_algebraic_action = process( + model.concatenated_algebraic, "algebraic", vars_for_processing + ) + + # combine rhs and algebraic functions + if len(model.rhs) == 0: + rhs_algebraic = model.concatenated_algebraic + elif len(model.algebraic) == 0: + rhs_algebraic = model.concatenated_rhs + else: + rhs_algebraic = pybamm.NumpyConcatenation( + model.concatenated_rhs, model.concatenated_algebraic + ) + + ( + rhs_algebraic, + jac_rhs_algebraic, + jacp_rhs_algebraic, + jac_rhs_algebraic_action, + ) = process(rhs_algebraic, "rhs_algebraic", vars_for_processing) ( casadi_switch_events, @@ -365,19 +444,24 @@ def set_up( self.computed_dvar_dp_fcns = {} self._time_integral_vars = {} for key in self.output_variables: - # Check for any ExplicitTimeIntegral or DiscreteTimeSum variables + # Check for any ExplicitTimeIntegral or DiscreteTimeSum variables. + # We will evaluate the sum node in the solver and sum it afterwards processed_time_integral = ( pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( model.get_processed_variable_or_event(key), model.len_rhs_and_alg, ) ) - # We will evaluate the sum node in the solver and sum it afterwards + if processed_time_integral is not None: + self._time_integral_vars[key] = processed_time_integral + if model.convert_to_format == "rust": + # Output values flow natively via CompiledModel.outputs, so the + # casadi computed_var_fcns path is unused on rust. + continue if processed_time_integral is None: var = model.get_processed_variable_or_event(key) else: var = processed_time_integral.sum_node - self._time_integral_vars[key] = processed_time_integral # Generate Casadi function to calculate variable and derivates # to enable sensitivites to be computed within the solver @@ -401,13 +485,12 @@ def _set_initial_conditions(self, model, time, inputs: list[dict]): len_tot = model.len_rhs_and_alg y_zero = np.zeros((len_tot, 1)) - casadi_format = model.convert_to_format == "casadi" + stacked = model.uses_stacked_inputs model.y0_list = [] model.y0S_list = [] if model.jacp_initial_conditions_eval is not None else None for ipts in inputs: - if casadi_format: - # stack inputs - inputs_y0_ics = casadi.vertcat(*[x for x in ipts.values()]) + if stacked: + inputs_y0_ics = stack_inputs(ipts, model.convert_to_format) else: inputs_y0_ics = ipts @@ -416,7 +499,7 @@ def _set_initial_conditions(self, model, time, inputs: list[dict]): ) if model.jacp_initial_conditions_eval is not None: - if casadi_format: + if stacked: inputs_jacp_ics = inputs_y0_ics else: # we are calculating the derivative wrt the inputs @@ -488,21 +571,34 @@ def _check_and_prepare_model_inplace(self, model): f"model should be discretised before solving ({e})" ) from e - if ( - isinstance(self, pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver) - ) and model.convert_to_format != "casadi": - pybamm.logger.warning( - f"Converting {model.name} to CasADi for solving with CasADi solver" - ) - model.convert_to_format = "casadi" - if ( - isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) - and model.convert_to_format != "casadi" - ): - pybamm.logger.warning( - f"Converting {model.name} to CasADi for calculating ICs with CasADi" - ) - model.convert_to_format = "casadi" + if model.convert_to_format == "rust": + if isinstance(self, pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver): + raise pybamm.SolverError( + f"{self.__class__.__name__} does not support convert_to_format='rust'. " + "Use IDAKLUSolver or DiffsolSolver (or ScipySolver for ODE-only models, " + "NonlinearSolver for algebraic-only models)." + ) + if isinstance(self.root_method, pybamm.CasadiAlgebraicSolver): + pybamm.logger.warning( + f"Switching root method to NonlinearSolver for rust model {model.name}" + ) + self.root_method = "nonlinear_solver" + else: + if ( + isinstance(self, pybamm.CasadiSolver | pybamm.CasadiAlgebraicSolver) + ) and model.convert_to_format != "casadi": + pybamm.logger.warning( + f"Converting {model.name} to CasADi for solving with CasADi solver" + ) + model.convert_to_format = "casadi" + if ( + isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) + and model.convert_to_format != "casadi" + ): + pybamm.logger.warning( + f"Converting {model.name} to CasADi for calculating ICs with CasADi" + ) + model.convert_to_format = "casadi" @staticmethod def _get_vars_for_processing(model, inputs: dict): @@ -510,7 +606,11 @@ def _get_vars_for_processing(model, inputs: dict): "model": model, } - if model.convert_to_format != "casadi": + if model.convert_to_format == "rust": + vars_for_processing.update({"input_names": list(inputs.keys())}) + return vars_for_processing + + elif model.convert_to_format != "casadi": # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector(slice(0, model.len_rhs_and_alg)) # set up Jacobian object, for re-use of dict @@ -619,6 +719,8 @@ def supports_t_eval_discontinuities(expr): terminate_events = [] interpolant_extrapolation_events = [] discontinuity_events = [] + # Discovery above is backend-agnostic; only the lowering below moves. + skip_event_lowering = self._lowers_via_compiled_model(model) for n, event in enumerate(model.events): if event.event_type == pybamm.EventType.DISCONTINUITY: # discontinuity events are evaluated before the solver is called, @@ -658,6 +760,10 @@ def supports_t_eval_discontinuities(expr): )[0] # use the actual casadi object as this will go into the rhs casadi_switch_events.append(event_casadi) + elif skip_event_lowering: + # Termination events are served from the solver's shared lowering; + # nothing reads the extrapolation slot. + continue else: # use the function call event_call = process( @@ -805,15 +911,13 @@ def _integrate( inputs_list = inputs_list or [{}] - ninputs = len(inputs_list) - if ninputs == 1: - new_solution = self._integrate_single( - model, - t_eval, - inputs_list[0], - model.y0_list[0], - ) - new_solutions = [new_solution] + if len(inputs_list) == 1 or nproc == 1: + # A one-process pool cannot beat a loop: same solves in the same + # order, plus a spawn and a pickle of the model per input set. + new_solutions = [ + self._integrate_single(model, t_eval, inputs, y0) + for inputs, y0 in zip(inputs_list, model.y0_list, strict=True) + ] else: with mp.get_context(self._mp_context).Pool(processes=nproc) as p: model_list = [model] * len(inputs_list) @@ -1217,17 +1321,18 @@ def _check_event_violation(self, t_eval, model, y0, inputs_dict): if num_terminate_events == 0: return - if model.convert_to_format == "casadi": - inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) + inputs = inputs_dict + if model.uses_stacked_inputs: + inputs = stack_inputs(inputs_dict, model.convert_to_format) events_eval = np.empty(num_terminate_events) for idx, event in enumerate(model.terminate_events_eval): - if model.convert_to_format == "casadi": + if model.uses_stacked_inputs: event_eval = event(t_eval[0], y0, inputs) - elif model.convert_to_format in ["python", "jax"]: + else: event_eval = event(t=t_eval[0], y=y0, inputs=inputs_dict) - if not isinstance(event_eval, float): - event_eval = event_eval.item() + if not isinstance(event_eval, float): + event_eval = np.asarray(event_eval).item() events_eval[idx] = event_eval if events_eval.min() <= 0: @@ -1241,6 +1346,69 @@ def _check_event_violation(self, t_eval, model, y0, inputs_dict): f"Events {event_names} are non-positive at initial conditions with inputs {inputs_dict}" ) + @staticmethod + def _overlay_options( + defaults: dict, user_options: dict | None, *, solver_name: str + ) -> dict: + """Overlay ``user_options`` on ``defaults``, rejecting unknown keys. + + Parameters + ---------- + defaults : dict + One entry per known option. + user_options : dict or None + Caller overrides. + solver_name : str + Named in the error message, so it points at the right option list. + + Returns + ------- + dict + ``defaults`` with ``user_options`` applied. + + Raises + ------ + :class:`pybamm.SolverError` + If a key is not a known option, which would otherwise be dropped + silently and leave the caller with the default. + """ + user_options = user_options or {} + unknown = sorted(set(user_options) - set(defaults)) + if unknown: + raise pybamm.SolverError( + f"Unknown {solver_name} solver option(s): {', '.join(unknown)}. " + f"Known options are: {', '.join(sorted(defaults))}." + ) + return defaults | user_options + + @staticmethod + def _check_restart_sensitivities(old_solution: pybamm.Solution) -> None: + """Reject a sensitivity restart from a solution that carries none. + + Parameters + ---------- + old_solution : :class:`pybamm.Solution` + The solution the next step would restart from. + + Raises + ------ + :class:`pybamm.SolverError` + If ``old_solution`` returned output variables only and so holds no + state sensitivities to seed ``dy0/dp`` from. + """ + if isinstance(old_solution, pybamm.EmptySolution): + return + # Not gated on _all_sensitivities: IDAKLU populates it at output width, + # which cannot seed a state-width dy0/dp. + if old_solution.variables_returned: + raise pybamm.SolverError( + "Cannot continue a sensitivity solve from a solution that " + "returned output variables only: the step boundary has no " + "state sensitivities to seed dy0/dp from. Drop " + "'output_variables' or 'calculate_sensitivities' for " + "multi-step solves." + ) + def _set_sens_initial_conditions_from( self, solution: pybamm.Solution, model: pybamm.BaseModel ) -> tuple: @@ -1489,6 +1657,8 @@ def step( pybamm.logger.verbose(f"Start stepping {model.name} with {self.name}") using_sensitivities = len(model.calculate_sensitivities) > 0 + if using_sensitivities: + self._check_restart_sensitivities(old_solution) if isinstance(old_solution, pybamm.EmptySolution): if not first_step_this_model: @@ -1827,6 +1997,81 @@ def _set_up_model_inputs(model: pybamm.BaseModel, inputs: dict): return ordered_inputs +def rust_input_parameter_widths(model: pybamm.BaseModel) -> dict[str, int]: + """Map each input parameter name to its packed width for the Rust backend. + + Every `ExprGraph.input_parameter` pre-registration site needs this so a + vector-valued (`expected_size > 1`) input reserves the right number of + slots in the packed `p` array before conversion re-registers it. + """ + return {ip.name: ip._expected_size or 1 for ip in model.input_parameters} + + +def validate_rust_sensitivity_widths( + model: pybamm.BaseModel, calculate_sensitivities: list[str] +) -> None: + """Reject sensitivity requests for vector-width inputs on the Rust backend. + + The Rust tangent/JVP machinery seeds one scalar direction per named + parameter (``sens_param_indices`` is registration-index based), so a + sensitivity direction for a ``width > 1`` parameter cannot represent + per-element derivatives and would silently compute the wrong values. + + Raises + ------ + pybamm.SolverError + If ``calculate_sensitivities`` includes a ``width > 1`` parameter. + """ + widths = rust_input_parameter_widths(model) + wide_params = sorted( + name for name in calculate_sensitivities if widths.get(name, 1) > 1 + ) + if wide_params: + raise pybamm.SolverError( + "convert_to_format='rust' sensitivities are not supported " + f"for vector-width input parameters: {wide_params}. The Rust " + "backend's tangent/JVP machinery assumes one scalar direction " + "per named parameter, so seeding a sensitivity direction for " + "a vector input would silently compute the wrong derivative. " + "Use convert_to_format='casadi' for sensitivities with " + "respect to these inputs." + ) + + +def flatten_inputs(inputs_dict: dict) -> npt.NDArray[np.float64]: + """Flatten ``{name: value}`` into the packed vector the Rust backend takes. + + Values are concatenated in dict-key order, which is the order every Rust + lowering assigns parameter slots in, and a vector-valued input contributes + its whole block. + + Parameters + ---------- + inputs_dict : dict + Input values for one input set. + + Returns + ------- + :class:`numpy.ndarray` + Flat ``float64`` vector, empty for an empty dict. + """ + if not inputs_dict: + return np.zeros(0, dtype=np.float64) + return np.concatenate( + [ + np.atleast_1d(np.asarray(v, dtype=np.float64)).ravel() + for v in inputs_dict.values() + ] + ) + + +def stack_inputs(inputs_dict: dict, convert_to_format: str): + """Stack input values into the flat vector stacked-input evaluators expect.""" + if convert_to_format == "rust": + return flatten_inputs(inputs_dict) + return casadi.vertcat(*[x for x in inputs_dict.values()]) + + def process( symbol, name, vars_for_processing, use_jacobian=None, return_jacp_stacked=None ): @@ -1848,25 +2093,29 @@ def process( ------- func: :class:`pybamm.EvaluatorPython` or :class:`pybamm.EvaluatorJax` or - :class:`casadi.Function` + :class:`casadi.Function` or + :class:`pybamm.solvers.rust_evaluator.RustEvaluator` evaluator for the function $f(y, t, p)$ given by `symbol` jac: :class:`pybamm.EvaluatorPython` or :class:`pybamm.EvaluatorJaxJacobian` or - :class:`casadi.Function` + :class:`casadi.Function` or + :class:`pybamm.solvers.rust_evaluator.RustEvaluator` evaluator for the Jacobian $\\frac{\\partial f}{\\partial y}$ of the function given by `symbol` jacp: :class:`pybamm.EvaluatorPython` or :class:`pybamm.EvaluatorJaxSensitivities` or - :class:`casadi.Function` + :class:`casadi.Function` or + :class:`pybamm.solvers.rust_evaluator.RustEvaluator` evaluator for the parameter sensitivities $\frac{\\partial f}{\\partial p}$ of the function given by `symbol` jac_action: :class:`pybamm.EvaluatorPython` or :class:`pybamm.EvaluatorJax` or - :class:`casadi.Function` + :class:`casadi.Function` or + :class:`pybamm.solvers.rust_evaluator.RustEvaluator` evaluator for product of the Jacobian with a vector $v$, i.e. $\\frac{\\partial f}{\\partial y} * v$ """ @@ -1899,6 +2148,49 @@ def report(string): jac = None jac_action = None + elif model.convert_to_format == "rust": + from pybamm.solvers.rust_evaluator import RustEvaluator + from pybamm.solvers.rust_lowering import rust_graph_with_inputs + + report(f"Converting {name} to Rust") + input_names = vars_for_processing["input_names"] + + graph = rust_graph_with_inputs(model, input_names) + rust_expr = symbol.to_rust(graph, {}) + + # n_states carries the full system width so rectangular sub-groups + # still differentiate wrt all states (design invariant 2) + cf = graph.compile(rust_expr, name=name, n_states=model.len_rhs_and_alg) + if return_jacp_stacked: + raise pybamm.SolverError( + "return_jacp_stacked is not supported for convert_to_format='rust'. " + "Use the per-parameter tuple returned by jacp instead." + ) + + func = RustEvaluator(cf, "func") + + if model.calculate_sensitivities: + report( + f"Calculating sensitivities for {name} with respect " + f"to parameters {model.calculate_sensitivities} using Rust" + ) + sens_indices = [ + input_names.index(p) + for p in model.calculate_sensitivities + if p in input_names + ] + jacp = RustEvaluator(cf, "jacp", sens_indices=sens_indices) + else: + jacp = None + + if use_jacobian: + report(f"Calculating jacobian for {name} using Rust") + jac = RustEvaluator(cf, "jac") + jac_action = RustEvaluator(cf, "jac_action") + else: + jac = None + jac_action = None + elif model.convert_to_format != "casadi": y = vars_for_processing["y"] jacobian = vars_for_processing["jacobian"] diff --git a/packages/pybamm/src/pybamm/solvers/diffsol_solver.py b/packages/pybamm/src/pybamm/solvers/diffsol_solver.py new file mode 100644 index 0000000000..1f60fb4860 --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/diffsol_solver.py @@ -0,0 +1,780 @@ +"""Solver class wrapping the Rust diffsol BDF integrator.""" + +import numbers + +import numpy as np +import numpy.typing as npt + +import pybamm +from pybamm.solvers.base_solver import ( + flatten_inputs, + validate_rust_sensitivity_widths, +) +from pybamm.solvers.observation import ( + NativeComputedObservation, + NativeInterpolatingObservation, + OutputAssembly, +) +from pybamm.solvers.rust_lowering import RustModelLowering + + +def _as_flat_float64(y0) -> npt.NDArray[np.float64]: + """Flatten an initial state to a float64 vector (handles ``casadi.DM``).""" + if hasattr(y0, "full"): + y0 = y0.full() + return np.asarray(y0, dtype=np.float64).flatten() + + +def _flatten_y0_sens(y0S, n_states: int, n_params: int) -> npt.NDArray[np.float64]: + """Flatten ``dy0/dp`` to the column-major block the Rust solver expects. + + Every producer of ``model.y0S_list`` hands over one column per sensitivity + parameter: ``(n_states, 1)`` from ``jacp``, bare ``(n_states,)`` from a step + restart. A single ``(n_states, n_params)`` matrix is accepted too. All + normalise to ``n_states * n_params`` values, parameter-outer/state-inner. + + Parameters + ---------- + y0S : array-like or sequence of array-like + Initial-condition sensitivities for one input set. + n_states : int + Number of states in the model. + n_params : int + Number of requested sensitivity parameters. + + Returns + ------- + :class:`numpy.ndarray` + Flat ``(n_states * n_params,)`` seed, or empty for an all-zero one. + + Raises + ------ + :class:`pybamm.SolverError` + If the flattened seed does not have ``n_states * n_params`` entries. + """ + blocks = [y0S] if hasattr(y0S, "full") or isinstance(y0S, np.ndarray) else list(y0S) + if not blocks: + # Empty is the all-zero seed. A shape complaint here would mask the + # solver's clearer error for a parameter it cannot differentiate. + return np.zeros(0, dtype=np.float64) + # column_stack lifts a bare (n_states,) column and leaves (n_states, 1) alone. + matrix = np.column_stack( + [ + np.asarray( + block.full() if hasattr(block, "full") else block, dtype=np.float64 + ) + for block in blocks + ] + ) + if matrix.shape != (n_states, n_params): + raise pybamm.SolverError( + f"Initial-condition sensitivities have shape {matrix.shape} but " + f"({n_states}, {n_params}) was expected (states x sensitivity " + "parameters)." + ) + return matrix.ravel(order="F") + + +class DiffsolSolver(pybamm.BaseSolver): + """Solve a discretised model using the Rust diffsol BDF solver. + + This solver delegates time integration to a compiled Rust BDF (backward + differentiation formula) implementation backed by the ``diffsol`` crate. + The model's expression graph is lowered to Rust via ``to_rust`` and + compiled into a ``CompiledModel`` that the Rust solver evaluates + directly, avoiding per-step Python callbacks. + + The full state trajectory (and, by default, its time derivatives) is + stored at every output time; the solvers differ only in which times those + are. :class:`pybamm.IDAKLUSolver` uses its internal integrator steps as + the output grid when ``t_interp`` is omitted, whereas diffsol evaluates + its error-controlled dense output at the requested times alone. A bare + ``t_eval=[t0, tf]`` span is answered on a uniform 100-point grid; any + explicit output grid is honoured exactly. + + Parameters + ---------- + rtol : float, optional + The relative tolerance for the solver (default is 1e-6). + atol : float or :class:`numpy.ndarray`, optional + The absolute tolerance for the solver, either shared by every state or + one entry per state (default is 1e-6). A per-state array is the way to + tolerance states of different magnitudes, since ``rtol`` already scales + with each state's own value. + root_method : str or pybamm algebraic solver class, optional + The method to use to find initial conditions (for DAE solvers). + root_tol : float, optional + The tolerance for the initial-condition solver. Default is 1e-6. + extrap_tol : float, optional + The tolerance to assert whether extrapolation occurs or not. + on_extrapolation : str, optional + What to do if the solver is extrapolating. Options are "warn", + "error", or "ignore". Default is "warn". + on_failure : str, optional + What to do if a solver error flag occurs. Options are "warn", + "error", or "ignore". Default is "error". + output_variables : list[str], optional + List of variables to calculate and return. If none are specified + then the complete state vector is returned (default is []). + calc_ic : bool, optional + If True, use native diffsol initial condition calculation instead of + Python-side root-finding. Default is False. + hermite_interpolation : bool, optional + If True (default), also store the state time derivatives so off-grid + ``sol[...](t)`` reads interpolate with cubic Hermite between output + points, as :class:`pybamm.IDAKLUSolver` does. Disabling it halves + trajectory memory; off-grid reads then interpolate linearly. Has no + effect with ``output_variables``, whose solves store no state + trajectory to interpolate. A solve asked for more than 4096 output + points also drops the derivatives, the chord already sitting on the + integration error floor by that density; + ``Solution.hermite_interpolation`` reports which way a solve went. + sens_atol_factor : float, optional + Multiplier applied to ``atol`` on differential states to form the + forward-sensitivity absolute tolerance floor (default 1e-3). Algebraic + states keep ``atol``. Raise it towards 1.0 if a stiff DAE fails its + sensitivity solve; lower it for tighter sensitivities at more steps. + options : dict, optional + Integrator tuning, one key per diffsol ``OdeSolverOptions`` knob, plus + ``num_threads``: how many input sets solve concurrently. Unset keys take + the defaults in :attr:`DEFAULT_OPTIONS`; an unknown key raises rather + than being silently ignored. Leave ``num_threads`` at 1 under an + outer-parallel caller (PyBOP, joblib, a ``ThreadPoolExecutor``), whose + thread count would otherwise multiply with this one. + """ + + _integrates_via_compiled_model = True + # Native (Rust) observation is the default for diffsol: the attach on the + # full-state solve fires, and an un-lowerable variable hard-fails by design. + _observes_via_compiled_model = True + + #: diffsol's own defaults, except the cumulative-per-solve + #: ``max_nonlinear_solver_failures``: at diffsol's 50 it caps solve length + #: rather than divergence, which ``min_timestep`` catches instead. Kept + #: literal so importing ``pybamm`` never needs the extension; pinned + #: against it by ``test_defaults_match_the_rust_defaults``. + _INTEGRATOR_DEFAULTS = { + "max_nonlinear_solver_iterations": 10, + "max_error_test_failures": 40, + "max_nonlinear_solver_failures": 100000, + "nonlinear_solver_tolerance": 0.2, + "min_timestep": 1e-13, + "max_timestep_growth": None, + "min_timestep_growth": None, + "max_timestep_shrink": None, + "min_timestep_shrink": None, + "update_jacobian_after_steps": 20, + "update_rhs_jacobian_after_steps": 50, + "threshold_to_update_jacobian": 0.3, + "threshold_to_update_rhs_jacobian": 0.2, + "pi_control_proportional": 0.0, + "pi_control_integral": 0.5, + } + + #: The integrator knobs plus ``num_threads``, which says how solves are + #: executed rather than how one integrates and so never reaches diffsol. + DEFAULT_OPTIONS = _INTEGRATOR_DEFAULTS | {"num_threads": 1} + + def __init__( + self, + rtol=1e-6, + atol=1e-6, + root_method=None, + root_tol=1e-6, + extrap_tol=None, + on_extrapolation=None, + on_failure=None, + output_variables=None, + calc_ic=False, + hermite_interpolation=True, + sens_atol_factor=1e-3, + options=None, + ): + super().__init__( + "problem dependent", + rtol, + atol, + root_method, + root_tol, + extrap_tol, + on_extrapolation, + on_failure, + output_variables, + ) + self.name = "diffsol solver (bdf)" + self._calc_ic = calc_ic + self._hermite_interpolation = bool(hermite_interpolation) + self._supports_interp = True + self._supports_t_eval_discontinuities = True + self._options = self._combine_options(options) + self._options["num_threads"] = self._checked_num_threads( + self._options["num_threads"] + ) + + try: + factor = float(sens_atol_factor) + except (TypeError, ValueError) as exc: + raise pybamm.SolverError( + f"sens_atol_factor must be a finite number > 0, got {sens_atol_factor!r}" + ) from exc + if not np.isfinite(factor) or factor <= 0: + raise pybamm.SolverError( + f"sens_atol_factor must be a finite number > 0, got {sens_atol_factor!r}" + ) + self._sens_atol_factor = factor + + if root_method is None and not calc_ic: + self._use_default_root_method = True + else: + self._use_default_root_method = False + + @classmethod + def _combine_options(cls, user_options: dict | None) -> dict: + """Overlay ``user_options`` on :attr:`DEFAULT_OPTIONS`. + + Parameters + ---------- + user_options : dict or None + Overrides, keyed as diffsol's option names or ``num_threads``. + + Returns + ------- + dict + One entry per known option. + + Raises + ------ + :class:`pybamm.SolverError` + If a key is not a known option. The Rust side requires every + integrator key, so a typo would otherwise be dropped silently. + """ + return cls._overlay_options( + cls.DEFAULT_OPTIONS, user_options, solver_name="diffsol" + ) + + @staticmethod + def _checked_num_threads(num_threads) -> int: + """Validate ``num_threads`` as a count of concurrent input sets. + + Parameters + ---------- + num_threads : object + The ``num_threads`` option as the caller supplied it. + + Returns + ------- + int + The validated count. + + Raises + ------ + :class:`pybamm.SolverError` + If it is not an integer of at least 1. Rust takes it as a ``usize``, + which would reject a negative with a bare ``OverflowError``. + """ + if ( + isinstance(num_threads, bool) + or not isinstance(num_threads, numbers.Integral) + or num_threads < 1 + ): + raise pybamm.SolverError( + f"num_threads must be an integer >= 1, got {num_threads!r}" + ) + return int(num_threads) + + def _integrator_options(self) -> dict: + """The subset of ``self._options`` diffsol's ``OdeSolverOptions`` takes.""" + return {key: self._options[key] for key in self._INTEGRATOR_DEFAULTS} + + @property + def _internal_initialisation(self) -> bool: + """Return True if using native diffsol IC calculation.""" + return self._calc_ic + + def set_up(self, model, inputs=None, t_eval=None, ics_only=False): + """Set up the solver, building the Rust compiled model. + + Delegates to ``BaseSolver.set_up`` for initial-condition processing, + then lowers the discretised model to a ``CompiledModel`` for the + Rust diffsol backend. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model whose solution to calculate. + inputs : dict or list of dict, optional + Any input parameters to pass to the model when solving. + t_eval : numeric type, optional + The times at which to stop the integration due to a + discontinuity in time. + ics_only : bool, optional + If True, only process initial conditions (skip full setup). + """ + if model.convert_to_format != "rust": + pybamm.logger.info( + f"Converting {model.name} to Rust for solving with DiffsolSolver" + ) + model.convert_to_format = "rust" + + # Auto-select nonlinear_solver for DAE models if user didn't specify root_method + if self._use_default_root_method and model.len_alg > 0: + self.root_method = "nonlinear_solver" + + base_set_up_return = super().set_up(model, inputs, t_eval, ics_only) + + if ics_only: + return base_set_up_return + + if isinstance(inputs, list): + inputs_dict = inputs[0] + else: + inputs_dict = inputs or {} + + self._build_rust_model(model, inputs_dict) + + return base_set_up_return + + def _build_rust_model(self, model, inputs_dict): + """Lower the discretised PyBaMM model to a Rust ``CompiledModel``. + + Builds the expression graph, mass matrix, output variable + expressions, and termination event expressions required by the + Rust diffsol solver. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The discretised model. + inputs_dict : dict + Input parameter values. + """ + lowering = RustModelLowering(model, inputs_dict) + lowering.state_residual() + + if model.calculate_sensitivities: + validate_rust_sensitivity_widths(model, model.calculate_sensitivities) + lowering.sensitivity_indices(model.calculate_sensitivities) + + # A time-integral output lowers to its integrand trajectory; the postfix + # sum runs post-solve, in the assembly. + _, output_lens = lowering.outputs( + self.output_variables, time_integral_vars=self._time_integral_vars + ) + self._output_assembly = OutputAssembly( + self.output_variables, + output_lens, + time_integrals=self._time_integral_vars, + ) + + lowering.termination_events() + + self._rust_model = lowering.compile() + lowering.bind_generic_evaluators(self._rust_model) + + # Observation tapes cached 1:1 with the rust model so repeated solves + # reuse them instead of recompiling against the retained graph. + self._rust_observation_cache: dict = {} + + from pybamm.rust import PreparedSolver + + self._prepared_solver = PreparedSolver( + self._rust_model, + float(self.rtol), + self._check_atol_type(self.atol, model), + self._sens_atol_factor, + self._integrator_options(), + # Outputs-only solves store no state trajectory to Hermite between. + self._hermite_interpolation and not self.output_variables, + ) + + def _set_consistent_initialization(self, model, time, inputs_list): + super()._set_consistent_initialization(model, time, inputs_list) + # first_state on an outputs-only solution rebuilds from y0full. + model.y0full = [_as_flat_float64(y0) for y0 in model.y0_list] + + def _integrate( + self, + model, + t_eval, + inputs_list=None, + t_interp=None, + nproc=1, + ): + """Integrate the model using the diffsol BDF solver. + + Overrides the base class for two reasons. The diffsol backend uses dense + output, so `t_eval` and `t_interp` merge into a single sorted array of + output times, with `t_eval` also passed down as the stop times the + integrator must land on and restart from. And concurrency over input + sets is the solver's own, through the `num_threads` option and a rayon + pool inside Rust: `nproc` is accepted for signature compatibility and + ignored, because `PreparedSolver` is a PyO3 object that cannot be + pickled to worker processes. + """ + if not hasattr(self, "_prepared_solver"): + raise RuntimeError("DiffsolSolver requires set_up() before solve()") + + inputs_list = inputs_list or [{}] + # Shared by every set, so converted once rather than per solve. + t_solve = np.asarray(self._output_times(t_eval, t_interp), dtype=np.float64) + t_stop = np.asarray(t_eval, dtype=np.float64) + # dy0/dp is per input set, like y0; absent when nothing is differentiated. + y0S_list = getattr(model, "y0S_list", None) or [None] * len(model.y0_list) + + n_sets = len(inputs_list) + if self._options["num_threads"] > 1 and n_sets > 1: + results = self._solve_batch(model, t_solve, t_stop, inputs_list, y0S_list) + else: + results = [ + self._solve_one(model, t_solve, t_stop, i, inputs_dict, y0, y0S) + for i, (inputs_dict, y0, y0S) in enumerate( + zip(inputs_list, model.y0_list, y0S_list, strict=True) + ) + ] + + return [ + self._build_solution(model, inputs_dict, result) + for inputs_dict, result in zip(inputs_list, results, strict=True) + ] + + def _payload_flags(self, model) -> dict[str, bool]: + """The payload the configured solve mode asks the Rust solver for. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model being solved, for whether sensitivities were requested. + + Returns + ------- + dict + ``outputs`` and ``sensitivities`` keyword arguments for + :meth:`pybamm.rust.PreparedSolver.solve`. + """ + return { + "outputs": bool(self.output_variables), + "sensitivities": bool(model.calculate_sensitivities), + } + + @staticmethod + def _output_times(t_eval, t_interp) -> npt.NDArray[np.float64]: + """The grid the solver reports its dense output on. + + Parameters + ---------- + t_eval : array-like + Requested times, which are also the integrator's stop times. + t_interp : array-like or None + Extra interpolation times. + + Returns + ------- + :class:`numpy.ndarray` + Sorted output times. + """ + if t_interp is not None and len(t_interp) > 0: + return np.union1d(t_eval, t_interp) + if len(t_eval) == 2: + # A bare [t0, tf] span leaves the output times to the solver (IDAKLU + # returns its steps); two points would interpolate as one chord. + return np.linspace(t_eval[0], t_eval[-1], 100) + return t_eval + + def _solve_one(self, model, t_eval, t_stop, index, inputs_dict, y0, y0S=None): + """Integrate one input set, returning the raw Rust result. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model whose solution to calculate. + t_eval : :class:`numpy.array`, size (k,) + The times at which to compute the solution. + t_stop : :class:`numpy.array` + Discontinuity times the integrator must land on exactly and restart + from. Every entry must also appear in ``t_eval``; those that do not + are integrated through. + index : int + Position of this set in the sweep, named if it fails. + inputs_dict : dict, optional + Any input parameters to pass to the model when solving. + y0 : array-like + The initial conditions for the model. + y0S : array-like or sequence of array-like, optional + ``dy0/dp`` for this input set, one column per requested sensitivity + parameter. ``None`` seeds the sensitivity system with zeros. + + Returns + ------- + :class:`pybamm.rust.SolveOutcome` + The solve's payloads, whichever the configured mode asked for. + + Raises + ------ + :class:`pybamm.SolverError` + If the integration fails. + """ + y0_np = _as_flat_float64(y0) + y0_sens = None + if model.calculate_sensitivities and y0S is not None: + y0_sens = _flatten_y0_sens( + y0S, y0_np.size, len(model.calculate_sensitivities) + ) + + try: + return self._prepared_solver.solve( + t_eval, + t_stop, + y0_np, + flatten_inputs(inputs_dict), + y0_sens=y0_sens, + **self._payload_flags(model), + ) + except RuntimeError as error: + # Integration failures cross the FFI boundary as RuntimeError. + self._raise_for_set(index, len(model.y0_list), inputs_dict, error) + + def _solve_batch(self, model, t_eval, t_stop, inputs_list, y0S_list): + """Integrate every input set concurrently, returning the raw Rust results. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model whose solution to calculate. + t_eval : :class:`numpy.array` + Output times, shared by every set, as + :meth:`_output_times` computed them. + t_stop : :class:`numpy.array` + Discontinuity times, shared by every set. + inputs_list : list of dict + One input dict per set, in the order the results are returned in. + y0S_list : list + One ``dy0/dp`` seed per set, entries ``None`` where absent. + + Returns + ------- + list + One :class:`pybamm.rust.SolveOutcome` per set, in input order. + + Raises + ------ + :class:`pybamm.SolverError` + If any set fails, naming which one and why. + """ + y0 = np.vstack([_as_flat_float64(y0) for y0 in model.y0_list]) + inputs = np.vstack([flatten_inputs(d) for d in inputs_list]) + y0_sens = None + if model.calculate_sensitivities: + y0_sens = self._stack_y0_sens(model, y0S_list, y0.shape[1]) + + results = self._prepared_solver.solve_batch( + t_eval, + t_stop, + y0, + inputs, + # The configured width, not the sweep's: keying the process-wide pool + # cache on the workload would build a fresh pool per distinct sweep size. + self._options["num_threads"], + y0_sens=y0_sens, + **self._payload_flags(model), + ) + + # A failed set carries its exception rather than raising it, so the index + # survives the crossing and the message can name the set. + for i, result in enumerate(results): + if isinstance(result, BaseException): + self._raise_for_set(i, len(results), inputs_list[i], result) + return results + + @staticmethod + def _raise_for_set(index, n_sets, inputs_dict, error): + """Re-raise ``error`` as a :class:`pybamm.SolverError`. + + Names the input set when there is more than one, so a sweep's failure + is attributable whether it ran batched or serially. + + Parameters + ---------- + index : int + Position of the failing set. + n_sets : int + Sets in the sweep. + inputs_dict : dict + The failing set's inputs. + error : BaseException + The underlying failure. + + Raises + ------ + :class:`pybamm.SolverError` + Always. + """ + if n_sets > 1: + raise pybamm.SolverError( + f"input set {index} of {n_sets} ({inputs_dict}) failed: {error}" + ) from error + raise pybamm.SolverError(str(error)) from error + + @staticmethod + def _stack_y0_sens(model, y0S_list, n_states): + """The batch's ``dy0/dp`` seeds as one row per set, or ``None``. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + Solved model, for the requested sensitivity parameters. + y0S_list : list + One seed per set, entries ``None`` where absent. + n_states : int + States in the model. + + Returns + ------- + :class:`numpy.ndarray` or None + ``(n_sets, n_states * n_params)`` seeds, or ``None`` when every set's + seed is the all-zero one and Rust can default it. + """ + if not model.calculate_sensitivities: + return None + n_params = len(model.calculate_sensitivities) + + # An empty seed and an all-zero one mean the same thing, so both + # normalise to None here and the rectangular array is filled below. + def seed(y0S): + if y0S is None: + return None + flat = _flatten_y0_sens(y0S, n_states, n_params) + return flat if flat.size else None + + rows = [seed(y0S) for y0S in y0S_list] + if all(row is None for row in rows): + return None + zeros = np.zeros(n_states * n_params, dtype=np.float64) + return np.vstack([zeros if row is None else row for row in rows]) + + def _build_solution(self, model, inputs_dict, result): + """Turn one raw Rust result into a :class:`pybamm.Solution`. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The solved model. + inputs_dict : dict + The input set this result came from. + result : :class:`pybamm.rust.SolveOutcome` + An outcome from :meth:`_solve_one` or :meth:`_solve_batch`. + + Returns + ------- + :class:`pybamm.Solution` + Solution object with times, states, and event data. + """ + t = result.t + termination = {0: "final time", 1: "event"}.get(result.flag, "failure") + t_event = None + if result.t_event is not None: + t_event = np.array([result.t_event]) + + if self.output_variables: + # Always present on this path: the terminal (or root) full state, + # the only state an outputs-only caller can restart from. + y_event = np.asarray(result.y_event).reshape(-1, 1) + + sol = pybamm.Solution( + t, + # Zero-row, not None: experiment-step stitching slices all_ys + # unconditionally. + np.zeros((0, t.size)), + model, + inputs_dict, + t_event, + y_event, + termination, + variables_returned=True, + ) + sensitivity_names = model.calculate_sensitivities or [] + sensitivities = None + if sensitivity_names: + # Bind yS once — the FFI getter rebuilds the list on every access. + sensitivities = self._output_assembly.stack_parameter_blocks( + result.yS, t.size, sensitivity_names + ) + # `result.y` holds output rows, not states, because the request asked + # for them; Rust lays them out output-major, the assembly wants time. + self._output_assembly.attach( + sol, + np.asarray(result.y).T, + sensitivities=sensitivities, + sensitivity_names=sensitivity_names, + ) + else: + y = result.y # shape (n_states, n_times) from Rust + + y_event = None + if t_event is not None and result.y_event is not None: + y_event = np.asarray(result.y_event).reshape(-1, 1) + + yS_out = {} + if model.calculate_sensitivities: + sensitivity_names = model.calculate_sensitivities + # Bind yS once — the getter rebuilds the list on every access. + yS_list = result.yS + yS_out = { + name: np.asarray(yS_list[i]).reshape(-1, 1) + for i, name in enumerate(sensitivity_names) + } + yS_out["all"] = np.hstack([yS_out[name] for name in sensitivity_names]) + + sol = pybamm.Solution( + t, + y, + model, + inputs_dict, + t_event, + y_event, + termination, + all_sensitivities=yS_out, + all_yps=result.yp, + ) + + # Outputs-only too: first_state/last_state carry observable full states. + if self._observes_via_compiled_model and model.convert_to_format == "rust": + # Stored yps reroute observation to IDAKLU's yp-consuming path; + # without them the solution is only read on its own grid. + backend = ( + NativeInterpolatingObservation + if sol.hermite_interpolation + else NativeComputedObservation + ) + sol.observation = backend.uniform( + self._rust_model, + len(sol.all_ys), + cache=self._rust_observation_cache, + ) + + # Bind statistics once — the getter clones the struct on every access. + statistics = result.statistics + # Measured in Rust, not around this call: under a batch the wall clocks + # overlap, so every set would be stamped with the batch duration. + sol.integration_time = statistics.integration_time_secs + sol.solver_statistics = pybamm.SolverStatistics( + number_of_steps=statistics.number_of_steps, + number_of_linear_solver_setups=statistics.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations=statistics.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails=statistics.number_of_nonlinear_solver_fails, + number_of_error_test_failures=statistics.number_of_error_test_failures, + number_of_linear_solver_setups_from_checkpoint=statistics.number_of_linear_solver_setups_from_checkpoint, + number_of_linear_solver_setups_from_first_convergence_fail=statistics.number_of_linear_solver_setups_from_first_convergence_fail, + number_of_linear_solver_setups_from_second_convergence_fail=statistics.number_of_linear_solver_setups_from_second_convergence_fail, + number_of_linear_solver_setups_from_error_test_fail=statistics.number_of_linear_solver_setups_from_error_test_fail, + number_of_linear_solver_setups_from_step_success=statistics.number_of_linear_solver_setups_from_step_success, + ic_time_secs=statistics.ic_time_secs, + solver_setup_time_secs=statistics.solver_setup_time_secs, + sens_error_control_relaxed=statistics.sens_error_control_relaxed, + ) + if sol.solver_statistics.sens_error_control_relaxed: + pybamm.logger.warning( + "The diffsol sensitivity solve failed under error control and was " + "retried with sensitivities excluded from it. Sensitivity accuracy " + "may be reduced; try increasing `sens_atol_factor`." + ) + return sol diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index c6e5596c60..2369fc9212 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -1,6 +1,5 @@ # mypy: ignore-errors import logging -import math import numbers import warnings from enum import IntEnum @@ -12,15 +11,51 @@ import pybamm from pybamm.codegen.compilation import aot_compile +from pybamm.solvers.base_solver import ( + flatten_inputs, + stack_inputs, + validate_rust_sensitivity_widths, +) +from pybamm.solvers.observation import ( + NativeInterpolatingObservation, + OutputAssembly, +) +from pybamm.solvers.rust_lowering import RustModelLowering _UNSET = object() -def _flatten_inputs(inputs_dict): - """Flatten ``{name: value}`` into a 1-D float array in dict-key order.""" - if not inputs_dict: - return np.zeros(0) - return np.concatenate([np.asarray(v).reshape(-1) for v in inputs_dict.values()]) +def _sensitivity_scales(inputs_dict, sensitivity_names): + """IDAS ``pbar``: the magnitude of each differentiated parameter. + + IDAS weights the *scaled* sensitivity ``pbar_i * dy/dp_i`` like a state, so + ``pbar_i = |p_i|`` makes a column of ``dy/dp`` comparable to ``y`` however + small the parameter is. Left at the default 1.0, a diffusivity around + 1e-15 produces a sensitivity column around 1e14 that no absolute tolerance + can accommodate, and the corrector fails to converge. + + Magnitudes are handed over raw; the solver clamps the zero and non-finite + cases IDAS rejects. A vector-valued input takes one scale for the block. + + Parameters + ---------- + inputs_dict : dict + Input values for one input set. + sensitivity_names : list of str + Differentiated parameter names, in the solver's column order. + + Returns + ------- + :class:`numpy.ndarray` + One scale per differentiated parameter, ``(len(sensitivity_names),)``. + """ + return np.asarray( + [ + np.abs(np.asarray(inputs_dict[name], dtype=np.float64)).max() + for name in sensitivity_names + ], + dtype=np.float64, + ) # Mirrors SUNDIALS ``IDA_ROOT_RETURN`` in ``sundials/include/ida/ida.h``. @@ -67,8 +102,9 @@ class IDAKLUSolver(pybamm.BaseSolver): ---------- rtol : float, optional The relative tolerance for the solver (default is 1e-4). - atol : float, optional - The absolute tolerance for the solver (default is 1e-6). + atol : float or :class:`numpy.ndarray`, optional + The absolute tolerance for the solver, either shared by every state or + one entry per state (default is 1e-6). root_method : str or pybamm algebraic solver class, optional The method to use to find initial conditions (for DAE solvers). Default is None, which uses a custom Newton solver for consistent @@ -228,11 +264,17 @@ class IDAKLUSolver(pybamm.BaseSolver): "t_no_progress": 0.0, } - Note: These options only have an effect if model.convert_to_format == 'casadi' + Note: Options apply to both 'casadi' and 'rust' formats except where noted. """ + _integrates_via_compiled_model = True + + # idaklu observation is native (Rust): full-state values via native + # cubic-Hermite, sensitivities via the native chain rule. + _observes_via_compiled_model = True + class StateID(IntEnum): ALGEBRAIC = 0 DIFFERENTIAL = 1 @@ -279,6 +321,37 @@ def __init__( pybamm.citations.register("Hindmarsh2000") pybamm.citations.register("Hindmarsh2005") + def _validate_rust_compatibility(self, model, inputs): + """Runtime guards for unsupported Rust backend features.""" + if not model.use_jacobian: + raise pybamm.SolverError("KLU requires the Jacobian") + # Equality makes SetupOptions derive one thread per solver, so nobody + # gets OpenMP N_Vectors: a pessimisation, and not run-to-run reproducible. + num_threads = self._options["num_threads"] + if self._options.get("num_solvers", num_threads) != num_threads: + pybamm.logger.warning( + "The Rust backend runs one solver per thread; num_solvers is " + f"set to num_threads ({num_threads})." + ) + self._options["num_solvers"] = num_threads + if model.calculate_sensitivities: + validate_rust_sensitivity_widths(model, model.calculate_sensitivities) + + def _is_identity_matrix(self, matrix): + """Check if sparse matrix is identity.""" + import scipy.sparse as sp + + if not sp.issparse(matrix): + return False + n = matrix.shape[0] + if matrix.shape[0] != matrix.shape[1]: + return False + if matrix.nnz != n: + return False + # Check diagonal is all ones + diag = matrix.diagonal() + return np.allclose(diag, 1.0) + def _combine_options(self, user_options: dict | None) -> dict: user_options = user_options or {} num_solvers = user_options.get("num_threads", 1) @@ -323,13 +396,10 @@ def _combine_options(self, user_options: dict | None) -> dict: "num_steps_no_progress": 0, "t_no_progress": 0.0, } - if not user_options: - return default_options - - options = default_options | user_options - + options = self._overlay_options( + default_options, user_options, solver_name="IDAKLU" + ) self._check_options(options) - return options def _check_options(self, options: dict): @@ -353,18 +423,118 @@ def _check_options(self, options: dict): if not isinstance(options["compile"], bool): raise pybamm.SolverError("compile must be a bool") - def _check_atol_type(self, atol, model): - if isinstance(atol, float): - return np.full(model.len_rhs_and_alg, atol) - elif isinstance(atol, np.ndarray): - return atol - else: - raise pybamm.SolverError( - "Absolute tolerances must be a numpy array or float" + def _set_up_rust(self, model, inputs_dict, y0): + """Set up solver with Rust evaluation backend.""" + lowering = RustModelLowering(model, inputs_dict) + lowering.state_residual() + + _, sensitivity_names = lowering.sensitivity_indices( + model.calculate_sensitivities + ) + + if self._options["hermite_reduction_factor"] > 1.0 and sensitivity_names: + warnings.warn( + "Setting hermite_reduction_factor > 1.0 is not currently supported " + "with sensitivities. The hermite_reduction_factor option will be " + "ignored.", + pybamm.SolverWarning, + stacklevel=2, + ) + + # A time-integral output lowers to its integrand trajectory; the postfix + # sum runs post-solve. + _, output_lens = lowering.outputs( + self.output_variables, time_integral_vars=self._time_integral_vars + ) + + lowering.algebraic_block() + lowering.termination_events() + + n_inputs = lowering.n_inputs + rust_model = lowering.compile() + lowering.bind_generic_evaluators(rust_model) + + # Get sparsity pattern (CSC format) + colptrs, rowinds = rust_model.csc_sparsity_pattern() + nnz = rust_model.nnz + # Mirror the existing CasADi path: only expose algebraic-subblock + # callbacks when Newton IC mode is "auto". + enable_alg_subblock = ( + self._options.get("newton_mode", "auto") == "auto" + and rust_model.has_algebraic + ) + if enable_alg_subblock: + alg_jac_rows, alg_jac_cols = ( + rust_model.algebraic_jacobian_sparsity_pattern() ) + alg_jac_nnz = rust_model.algebraic_jacobian_nnz + n_alg = rust_model.n_algebraic + else: + alg_jac_rows = np.array([], dtype=np.int64) + alg_jac_cols = np.array([], dtype=np.int64) + alg_jac_nnz = 0 + n_alg = 0 + + # IDA `id` vector: 1.0 for differential, 0.0 for algebraic. + # Inferred from the mass matrix diagonal by the Rust core. + ids = np.asarray(rust_model.algebraic_ids(), dtype=np.float64) + + atol = self._check_atol_type(getattr(model, "atol", self.atol), model) + + # Store for solver creation + self._setup = { + "rust_model": rust_model, + # One evaluator per parallel solver, all sharing rust_model's tape. + "rust_evaluators": rust_model.evaluator_pool(self._options["num_solvers"]), + "number_of_states": len(y0), + "number_of_inputs": n_inputs, + "number_of_sensitivity_parameters": rust_model.n_sens_params, + "sensitivity_names": sensitivity_names, + "output_lens": output_lens, + "output_assembly": OutputAssembly( + self.output_variables, + output_lens, + time_integrals=self._time_integral_vars, + ), + "jac_colptrs": np.asarray(colptrs, dtype=np.int64), + "jac_rowvals": np.asarray(rowinds, dtype=np.int64), + "jac_nnz": nnz, + "n_alg": n_alg, + "alg_jac_rowvals": np.asarray(alg_jac_rows, dtype=np.int64), + "alg_jac_colvals": np.asarray(alg_jac_cols, dtype=np.int64), + "alg_jac_nnz": alg_jac_nnz, + "ids": ids, + "atol": atol, + # `num_of_events` mirrors the CasADi setup key so the shared + # `_post_process_solution` can read the event count uniformly. + "num_of_events": rust_model.n_events, + # Observation tapes cached 1:1 with the rust model; recreated on every + # `set_up`, so it can never serve tapes from a stale graph. + "rust_observation_cache": {}, + } + + self._setup["solver"] = idaklu.create_rust_solver_group( + rust_evaluators=self._setup["rust_evaluators"], + number_of_states=self._setup["number_of_states"], + number_of_inputs=self._setup["number_of_inputs"], + number_of_sens_params=self._setup["number_of_sensitivity_parameters"], + number_of_events=self._setup["num_of_events"], + output_lens=self._setup["output_lens"], + jac_colptrs=self._setup["jac_colptrs"], + jac_rowvals=self._setup["jac_rowvals"], + jac_nnz=self._setup["jac_nnz"], + n_alg=self._setup["n_alg"], + alg_jac_rowvals=self._setup["alg_jac_rowvals"], + alg_jac_colvals=self._setup["alg_jac_colvals"], + alg_jac_nnz=self._setup["alg_jac_nnz"], + rhs_alg_id=self._setup["ids"], + atol=self._setup["atol"], + rtol=self.rtol, + options=self._options, + ) def set_up(self, model, inputs=None, t_eval=None, ics_only=False): - if model.convert_to_format != "casadi": + if model.convert_to_format not in ("casadi", "rust"): pybamm.logger.warning( f"Converting {model.name} to CasADi for solving with IDAKLUSolver" ) @@ -392,6 +562,11 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): if ics_only: return base_set_up_return + if model.convert_to_format == "rust": + self._validate_rust_compatibility(model, inputs_dict) + self._set_up_rust(model, inputs_dict, y0) + return base_set_up_return + if model.convert_to_format != "casadi": msg = "The python-idaklu solver has been deprecated." warnings.warn(msg, DeprecationWarning, stacklevel=2) @@ -577,6 +752,11 @@ def to_idaklu(fn): "number_of_sensitivity_parameters": number_of_sensitivity_parameters, "standard_form_dae": model.is_standard_form_dae, "output_variables": self.output_variables, + "output_assembly": OutputAssembly.from_casadi( + self.output_variables, + self.computed_var_fcns, + time_integrals=self._time_integral_vars, + ), "var_fcns": self.computed_var_fcns, "var_idaklu_fcns": [], "dvar_dy_idaklu_fcns": [], @@ -673,17 +853,25 @@ def _integrate( """ Overloads the _integrate method from BaseSolver to use the IDAKLU solver """ - if model.convert_to_format != "casadi": # pragma: no cover + if model.convert_to_format not in ("casadi", "rust"): # pragma: no cover raise pybamm.SolverError("Unsupported IDAKLU solver configuration.") inputs_list = inputs_list or [{}] # stack inputs so that they are a 2D array of shape (number_of_inputs, number_of_parameters) if inputs_list and inputs_list[0]: - inputs = np.vstack([_flatten_inputs(d) for d in inputs_list]) + inputs = np.vstack([flatten_inputs(d) for d in inputs_list]) else: inputs = np.array([[]] * len(inputs_list)) + sensitivity_names = self._setup["sensitivity_names"] + if sensitivity_names and inputs_list and inputs_list[0]: + pbar = np.vstack( + [_sensitivity_scales(d, sensitivity_names) for d in inputs_list] + ) + else: + pbar = np.empty((0, 0)) + # y0full is now a list with length = number of input sets y0full = np.vstack(model.y0full) ydot0full = np.vstack(model.ydot0full) @@ -703,6 +891,7 @@ def _integrate( y0full, ydot0full, inputs, + pbar, logger=logger, ) except ValueError as e: @@ -821,92 +1010,62 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev options=solution_options, ) - # Set closest_event_idx so BaseSolver.get_termination_reason doesn't - # re-walk every event's symbolic expression on the Python side. - if sol.flag == _IDA_ROOT_RETURN and self._setup["num_of_events"] > 0: - event_values = np.asarray( - self._setup["rootfn_casadi"]( - float(sol.t[-1]), - np.asarray(y_event).reshape(-1), - _flatten_inputs(inputs_dict), + # Fast path via the compiled events (`rootfn_casadi`, or the rust model's + # own event tapes); else BaseSolver re-walks events symbolically. + if sol.flag == _IDA_ROOT_RETURN and self._setup.get("num_of_events", 0) > 0: + t_event = float(sol.t[-1]) + y_event_flat = np.asarray(y_event).reshape(-1) + if "rootfn_casadi" in self._setup: + event_values = np.asarray( + self._setup["rootfn_casadi"]( + t_event, + y_event_flat, + flatten_inputs(inputs_dict), + ) + ).reshape(-1) + newsol.closest_event_idx = int(np.nanargmin(np.abs(event_values))) + elif model.convert_to_format == "rust": + stacked = stack_inputs(inputs_dict, "rust") + event_values = np.array( + [ + float(np.asarray(fn(t_event, y_event_flat, stacked)).ravel()[0]) + for fn in model.terminate_events_eval + ] ) - ).reshape(-1) - newsol.closest_event_idx = int(np.nanargmin(np.abs(event_values))) + newsol.closest_event_idx = int(np.nanargmin(np.abs(event_values))) newsol.integration_time = integration_time + newsol.solver_statistics = pybamm.SolverStatistics( + number_of_steps=sol.stats.number_of_steps, + number_of_linear_solver_setups=sol.stats.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations=sol.stats.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails=sol.stats.number_of_nonlinear_solver_fails, + number_of_error_test_failures=sol.stats.number_of_error_test_failures, + ) + if self._observes_via_compiled_model and model.convert_to_format == "rust": + # Attached on the outputs-only path too: first_state/last_state carry + # full states, and their observation must match the full-state path. + # IDAKLU always stores yps, so every solution reads off-grid. + newsol.observation = NativeInterpolatingObservation.uniform( + self._setup["rust_model"], + len(newsol.all_ys), + cache=self._setup["rust_observation_cache"], + ) if not save_outputs_only: return newsol - # Populate variables and sensitivities dictionaries directly - number_of_samples = sol.y.shape[0] // number_of_timesteps - sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) - sensitivity_params = ( - model.calculate_sensitivities if model.calculate_sensitivities else [] + # On this path `sol.y` carries the concatenated outputs rather than the + # states, and `sol.yS` their sensitivities as (n_t, n_rows, n_p). + self._setup["output_assembly"].attach( + newsol, + np.asarray(sol.y).reshape(number_of_timesteps, -1), + sensitivities=( + np.asarray(sol.yS) if number_of_sensitivity_parameters else None + ), + sensitivity_names=sensitivity_names, ) - - start_idx = 0 - for var in self.output_variables: - var_nnz, var_shape, base_variables = self._get_variable_info(model, var) - end_idx = start_idx + var_nnz - data = sol.y[:, start_idx:end_idx] - time_indep = False - - # handle any time integral variables - if var in self._time_integral_vars: - # time integral variables should all be 1D - tiv = self._time_integral_vars[var] - data = tiv.postfix(data.reshape(-1), sol.t, inputs_dict) - time_indep = True - - newsol._variables[var] = pybamm.ProcessedVariableComputed( - [model.get_processed_variable_or_event(var)], - base_variables, - [data], - newsol, - time_indep=time_indep, - ) - - # Add sensitivities - newsol[var]._sensitivities = {} - if sensitivity_params: - if var_nnz != math.prod(var_shape): - raise pybamm.SolverError( - f"Sensitivity of sparse variables not supported. {var} is a sparse variable with number of non-zeros {var_nnz} and shape {var_shape}" - ) - sens_data = sol.yS[:, start_idx:end_idx, :] - sens_data = sens_data.reshape( - number_of_timesteps * (end_idx - start_idx), - number_of_sensitivity_parameters, - ) - if var in self._time_integral_vars: - tiv = self._time_integral_vars[var] - sens_data = tiv.postfix_sensitivities( - var, data, sol.t, inputs_dict, sens_data - ) - newsol[var]._sensitivities["all"] = sens_data - - # Add the individual sensitivity - for i, name in enumerate(inputs_dict.keys()): - sens = newsol[var]._sensitivities["all"][:, i : i + 1].reshape(-1) - newsol[var]._sensitivities[name] = sens - - start_idx += var_nnz return newsol - def _get_variable_info(self, model, var) -> tuple: - """Get variable length and base variables based on model format.""" - if model.convert_to_format == "casadi": - base_var = self._setup["var_fcns"][var] - var_eval = base_var(0.0, 0.0, 0.0) - var_nnz = var_eval.sparsity().nnz() - var_shape = var_eval.shape - return var_nnz, var_shape, [base_var] - else: # pragma: no cover - raise pybamm.SolverError( - f"Unsupported evaluation engine for convert_to_format=" - f"{model.convert_to_format}" - ) - def _set_consistent_initialization(self, model, time, inputs_list): """ Initialize y0 and ydot0 for the solver. In addition to calculating @@ -925,8 +1084,6 @@ def _set_consistent_initialization(self, model, time, inputs_list): # set model.y0_list super()._set_consistent_initialization(model, time, inputs_list) - casadi_format = model.convert_to_format == "casadi" - def handle_y0(y0): if isinstance(y0, casadi.DM): y0 = y0.full() @@ -944,7 +1101,7 @@ def handle_y0(y0): else: ydot0_list = [np.zeros_like(y0) for y0 in y0_list] - sensitivity = model.y0S_list and casadi_format + sensitivity = model.y0S_list and model.uses_stacked_inputs if sensitivity: y0S_list = model.y0S_list y0full = [] @@ -982,20 +1139,13 @@ def _rhs_dot_consistent_initialization(self, y0, model, time, inputs_dict): Any input parameters to pass to the model when solving. """ - casadi_format = model.convert_to_format == "casadi" - inputs_dict = inputs_dict or {} - # stack inputs - if inputs_dict: - arrays_to_stack = [np.array(x).reshape(-1, 1) for x in inputs_dict.values()] - inputs = np.vstack(arrays_to_stack) + if model.uses_stacked_inputs: + input_eval = stack_inputs(inputs_dict, model.convert_to_format) else: - inputs = np.array([[]]) + input_eval = inputs_dict ydot0 = np.zeros_like(y0) - # calculate the time derivatives of the differential equations - input_eval = inputs if casadi_format else inputs_dict - rhs0 = model.rhs_eval(time, y0, input_eval) if isinstance(rhs0, casadi.DM): rhs0 = rhs0.full() @@ -1031,10 +1181,10 @@ def _sensitivity_consistent_initialization(self, y0, ydot0, y0S, time, inputs_di """ - if isinstance(y0S, casadi.DM): + if isinstance(y0S, casadi.DM | np.ndarray): y0S = (y0S,) - if isinstance(y0S[0], casadi.DM): + if y0S and isinstance(y0S[0], casadi.DM): y0S = (x.full() for x in y0S) y0S = [x.flatten() for x in y0S] @@ -1183,7 +1333,9 @@ def reduce_solution( options=solution.user_options, ) - # Propagate metadata from the original solution + # Propagate metadata from the original solution. Thinning knots leaves + # the segments and their models alone, so the backend carries over as is. + new_sol.observation = solution.observation new_sol._all_inputs_stacked = solution.all_inputs_stacked new_sol._all_inputs_casadi = solution.all_inputs_casadi new_sol.closest_event_idx = solution.closest_event_idx diff --git a/packages/pybamm/src/pybamm/solvers/nonlinear_solver.py b/packages/pybamm/src/pybamm/solvers/nonlinear_solver.py index 4cef3ec3f9..a3207cf87b 100644 --- a/packages/pybamm/src/pybamm/solvers/nonlinear_solver.py +++ b/packages/pybamm/src/pybamm/solvers/nonlinear_solver.py @@ -6,6 +6,8 @@ import pybamm from pybamm.codegen.compilation import aot_compile +from pybamm.solvers.base_solver import flatten_inputs +from pybamm.solvers.rust_lowering import RustModelLowering _DEFAULT_OPTIONS = { "compile": False, @@ -15,19 +17,23 @@ class _NonlinearSolverSetup: """Pickle-safe wrapper around StandaloneNewtonSolver""" - __slots__ = ["_setup"] + __slots__ = ["_keepalive", "_setup"] - def __init__(self, setup: idaklu.StandaloneNewtonSolver): + def __init__(self, setup: idaklu.StandaloneNewtonSolver, keepalive=None): self._setup = setup + self._keepalive = keepalive # EvaluatorPool backing the C++ raw ptr def __bool__(self): + # falsy once the handle is gone -> caller rebuilds, never reuses a + # dangling pointer (e.g. after a pickle round-trip) return self._setup is not None def __getstate__(self): - return {"_setup": None} + return {"_setup": None, "_keepalive": None} def __setstate__(self, state): self._setup = None + self._keepalive = None def solve_batch(self, *args, **kwargs): return self._setup.solve_batch(*args, **kwargs) @@ -91,6 +97,17 @@ def __init__( self._user_options = options or {} self._options = _DEFAULT_OPTIONS | self._user_options + def __getstate__(self): + # _model_set_up holds unpicklable rust artifacts, so clear it: the solver + # stays picklable and the next solve() rebuilds from the model. + state = self.__dict__.copy() + state["_model_set_up"] = {} + return state + + def __setstate__(self, state): + self.__dict__.update(state) + self._model_set_up = {} + @staticmethod def _check_tolerance(value): if value < 0: @@ -132,6 +149,9 @@ def _set_up_root_solver(self, model, inputs_dict, t_eval): experiment (e.g. each unique step model) get independently sized Newton solvers without one stomping on another's cache. """ + if model.convert_to_format == "rust": + return self._set_up_root_solver_rust(model, inputs_dict) + pybamm.logger.info(f"Start building {self.name}") y0 = model.y0_list[0] @@ -175,6 +195,65 @@ def _set_up_root_solver(self, model, inputs_dict, t_eval): return self._build_newton_solver(res_fn, jac_fn, len_alg) + def _set_up_root_solver_rust(self, model, inputs_dict): + """Build the Rust-backed standalone Newton solver. + + Lowers rhs+algebraic into one CompiledModel (same layout as the + IDAKLU rust set-up) and hands its algebraic sub-block to the C++ + Newton driver via the dlsym FFI. + """ + pybamm.logger.info(f"Start building {self.name} (rust)") + if self._options["compile"]: + raise pybamm.SolverError( + "options['compile'] is CasADi-only; not supported with " + "convert_to_format='rust'" + ) + len_rhs = 0 if model.rhs == {} else model.len_rhs + # An oversized consistent y0 means the rhs/alg block was extended with + # sensitivities, which the Rust Newton driver does not support. + y0_list = getattr(model, "y0_list", None) + if y0_list and model.len_rhs_and_alg != np.asarray(y0_list[0]).shape[0]: + raise pybamm.SolverError( + "The Rust Newton root solver does not support " + "sensitivity-extended states; use convert_to_format='casadi' " + "for this configuration" + ) + + lowering = RustModelLowering(model, inputs_dict) + lowering.state_residual(algebraic_only=len_rhs == 0) + lowering.algebraic_block(first_algebraic_index=len_rhs) + rust_model = lowering.compile() + jac_rows, jac_cols = rust_model.algebraic_jacobian_sparsity_pattern() + len_alg = model.len_rhs_and_alg - len_rhs + # algebraic_jacobian_sparsity_pattern returns global state columns, but the C++ Newton + # builds an n_alg x n_alg system, so localise to the algebraic block. + jac_cols = [int(c) - len_rhs for c in np.asarray(jac_cols)] + jac_rows = np.asarray(jac_rows).tolist() + if any(c < 0 or c >= len_alg for c in jac_cols): + raise pybamm.SolverError( + "Rust algebraic jacobian has columns outside the algebraic " + f"block [0, {len_alg}); cannot localise for the Newton solver" + ) + + # The C++ Newton driver mutates the evaluator it drives, so its address + # must come from the pool's exclusive handout, not a rust_model borrow. + pool = rust_model.evaluator_pool(1) + _setup = idaklu.StandaloneNewtonSolver( + rust_model=pool.as_ptr(0), + n_rhs=len_rhs, + n_alg=len_alg, + jac_rows=jac_rows, + jac_cols=jac_cols, + atol=np.full(len_alg, float(self.atol)).tolist(), + rtol=float(self.rtol), + step_tol=float(self.step_tol), + max_iter=int(self.max_iter), + max_backtracks=int(self.max_backtracks), + eps_newt=float(self.eps_newt), + use_sparse=bool(self.use_sparse), + ) + return _NonlinearSolverSetup(_setup, keepalive=pool) + def _build_newton_solver(self, res_fn, jac_fn, len_alg): _setup = idaklu.StandaloneNewtonSolver( residual=idaklu.generate_function(res_fn.serialize()), @@ -195,11 +274,7 @@ def _integrate_single(self, model, t_eval, inputs_dict, y0): root_solver = self.get_root_solver(model, inputs_dict, t_eval) len_rhs = model.len_rhs - inputs_flat = ( - np.concatenate([np.atleast_1d(v).ravel() for v in inputs_dict.values()]) - if inputs_dict - else np.empty(0) - ) + inputs_flat = flatten_inputs(inputs_dict) y0_np = np.asarray(y0).ravel() y0_diff = y0_np[:len_rhs] diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py new file mode 100644 index 0000000000..b53909e81c --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -0,0 +1,754 @@ +"""Observation backends: how a :class:`pybamm.Solution` reads its variables. + +A ``Solution`` holds exactly one :class:`ObservationBackend`, chosen when the +solve finishes. The backend owns everything needed to turn a variable name into +a ready-to-evaluate processed variable — CasADi conversion for the default, or +the retained Rust graph, its compiled tapes and their cache for the native one. +Derived solutions (``first_state``, ``last_state``, ``copy``, ``__add__``, +``from_sub_solutions``) carry that one field, so nothing re-decides which +backend is in play. + +An ``output_variables`` solve reads its variables the other way round: the +solver computed them already and hands back a concatenated payload instead of a +state trajectory. :class:`OutputAssembly` is that path's counterpart to a +backend — it owns the payload's row layout and populates the Solution eagerly, +for whichever solver produced it. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from itertools import accumulate, pairwise + +import numpy as np + +import pybamm +from pybamm.solvers.variable_observer import ( + CasadiObserver, + NativeObserver, + check_variable_in_solve, + native_sensitivities, +) + + +class ObservationBackend(ABC): + """How a Solution lowers a variable name into evaluable leaves. + + Backends are immutable values covering an ordered run of a Solution's + sub-solutions. ``backend[key]`` restricts them to a slice of that run and + :func:`join_observations` concatenates runs, so a derived Solution copies + one field instead of a bundle of parallel ones. + """ + + @abstractmethod + def __getitem__(self, key: slice) -> ObservationBackend: + """This backend restricted to a slice of the Solution's segments.""" + + @abstractmethod + def build_variable(self, solution, name): + """The processed variable for ``name``, ready to evaluate. + + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution being observed; supplies the trajectories, models and + inputs. Its segments are 1:1 with this backend's. + name : str + Variable name, as registered on the models. + + Returns + ------- + :class:`pybamm.solvers.base_processed_variable.BaseProcessedVariable` + """ + + +class CasadiObservation(ObservationBackend): + """The default backend: variables are converted to CasADi and read by IDAKLU. + + Stateless -- every per-segment artifact it needs is reached through the + Solution it is handed -- so :data:`CASADI_OBSERVATION` is shared by every + Solution that has not been given a native backend. + """ + + def __getitem__(self, key): + return self + + def build_variable(self, solution, name): + time_integral = None + pybamm.logger.debug(f"Post-processing {name}") + + # Iterate through all models, some may be in the list several times and + # therefore only get set up once + vars_pybamm = [ + model.get_processed_variable_or_event(name) for model in solution.all_models + ] + vars_casadi = [None] * len(solution.all_models) + for i, (model, ys, inputs) in enumerate( + zip(solution.all_models, solution.all_ys, solution.all_inputs, strict=True) + ): + _var_pybamm = vars_pybamm[i] + check_variable_in_solve(solution, name, _var_pybamm) + if isinstance(_var_pybamm, pybamm.VectorField): + comp_casadi = [] + for k, comp in enumerate(_var_pybamm.components): + cc, _, _ = self._model_leaf( + solution, + model, + comp, + inputs=inputs, + ys_shape=ys.shape, + time_integral=None, + cache_key=f"{name}[{k}]", + ) + comp_casadi.append(cc) + vars_casadi[i] = comp_casadi + else: + var_casadi, var_pybamm, time_integral = self._model_leaf( + solution, + model, + _var_pybamm, + inputs=inputs, + ys_shape=ys.shape, + time_integral=time_integral, + cache_key=name, + ) + vars_pybamm[i] = var_pybamm + vars_casadi[i] = var_casadi + return pybamm.process_variable( + name, + vars_pybamm, + CasadiObserver(vars_casadi), + solution, + time_integral=time_integral, + ) + + @staticmethod + def _model_leaf( + solution, + model, + var_pybamm, + time_integral, + inputs, + ys_shape, + cache_key, + ): + """One model's CasADi leaf, memoised on the model unless time-integrated.""" + _var_casadi = model._variables_casadi.get(cache_key) + if _var_casadi is not None: + return _var_casadi, var_pybamm, time_integral + + var_casadi, var_pybamm, time_integral = solution._convert_to_casadi( + var_pybamm, inputs, ys_shape + ) + + # Only cache if it's not a time integral + if time_integral is None: + model._variables_casadi[cache_key] = var_casadi + return var_casadi, var_pybamm, time_integral + + +class NativeObservation(ObservationBackend): + """Observation through a compiled Rust model's retained expression graph. + + Observed variables are lowered as new roots into the graph the solve was + built from and compiled to tapes, so no CasADi conversion happens. The + graph never leaves this class. + + Abstract: a solve that stores state derivatives can be read at arbitrary + times and wants :class:`NativeInterpolatingObservation`, while one read + only on its own grid wants :class:`NativeComputedObservation`. Which of + the two a solver constructs is what picks the kind of processed variable + its solutions hand back. + + Parameters + ---------- + models : list of :class:`pybamm.rust.CompiledModel` + One model per sub-solution, in solve order. + cache : dict, optional + Cache of compiled tapes and time-integral analyses, 1:1 with the + models. Sharing the solver's dict lets repeated solves reuse tapes + instead of growing the retained graph; omitting it starts a fresh one. + """ + + def __init__(self, models, *, cache=None): + self._models = list(models) + self._cache = {} if cache is None else cache + + @classmethod + def uniform(cls, model, n_segments, **kwargs) -> NativeObservation: + """One model observing every segment, which is what a single solve produces.""" + return cls([model] * n_segments, **kwargs) + + @classmethod + def _adopting(cls, models, *, cache) -> NativeObservation: + """Take ownership of an already-fresh list instead of copying it again.""" + backend = cls.__new__(cls) + backend._models = models + backend._cache = cache + return backend + + @property + def n_segments(self): + return len(self._models) + + @property + def compile_cache(self): + """The shared cache of compiled tapes and time-integral analyses.""" + return self._cache + + @property + def segment_models(self): + """The per-segment compiled models, in solve order.""" + return self._models + + @property + def primary_model(self): + """The model owning the graph new observation roots are lowered into.""" + return self._models[0] + + def __getitem__(self, key): + return self._adopting(self._models[key], cache=self._cache) + + def _segment_leaf(self, solution, name, model, rust_model, nstates): + """One segment's ``(pybamm variable, time integral, compiled leaf)``. + + A time-integrated variable contributes its integrand, so the leaf is + what the postfix sum consumes rather than the variable itself. + """ + var_pybamm = model.get_processed_variable_or_event(name) + check_variable_in_solve(solution, name, var_pybamm) + time_integral = self._time_integral(name, model, var_pybamm, nstates) + integrand = ( + time_integral.sum_node.child if time_integral is not None else var_pybamm + ) + return var_pybamm, time_integral, self._leaf(name, integrand, rust_model) + + def _time_integral(self, name, model, var_pybamm, nstates): + """Time-integral classification of `name`, memoised across solves. + + ``from_pybamm_var`` walks the variable's full expression tree, which + costs more than a small model's solve, so both build paths share this + memo. ``None`` results are stored too, hence key membership decides a hit. + """ + ti_key = ("__time_integral__", name, id(model), nstates) + if ti_key in self._cache: + return self._cache[ti_key] + time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( + var_pybamm, nstates + ) + self._cache[ti_key] = time_integral + return time_integral + + def _leaf(self, name, integrand, rust_model): + """Compile `integrand` once into the retained graph, cached by name+model.""" + # model id keys the cache so distinct models don't alias; models are + # retained for the backend's lifetime, so the id cannot be reused. + cache_key = (name, id(rust_model)) + fn = self._cache.get(cache_key) + if fn is None: + graph = rust_model.graph + rust_expr = integrand.to_rust(graph) + fn = graph.compile(rust_expr, name=name, n_states=rust_model.n_states) + self._cache[cache_key] = fn + return fn + + def _post_sum_leaf(self, name, time_integral, n_inner): + """Compile a time-integral's ``post_sum_node`` against the retained graph. + + The post-sum node's synthetic StateVector input is the integrated inner + value, so the compiled "state" is that ``n_inner`` vector (the eval + point is the postfix VALUE, not the full state trajectory). + """ + # Key on the variable name, not id(post_sum_node): the time_integral is a + # transient local whose id can be reused and alias the wrong compiled fn. + cache_key = ("__post_sum__", name, n_inner) + fn = self._cache.get(cache_key) + if fn is None: + graph = self.primary_model.graph + rust_expr = time_integral.post_sum_node.to_rust(graph) + fn = graph.compile(rust_expr, name="post_sum", n_states=n_inner) + self._cache[cache_key] = fn + return fn + + def postfix_sensitivities( + self, + times, + name, + time_integral, + postfix_value, + inner_sens, + inputs, + sens_names, + ): + """``dvar/dp`` for a time-integral variable. + + Integrates the inner-variable sensitivities over time, then applies the + post-sum chain rule ``dpost_dy @ s_integral + dpost_dp`` using a + Rust-lowered jacobian of ``post_sum_node`` evaluated at the postfix value. + + Parameters + ---------- + times : numpy.ndarray + The solution's time points, shape ``(n_t,)``. + name : str + Variable name; keys the compiled post-sum cache. + time_integral : pybamm.ProcessedVariableTimeIntegral + Time-integral descriptor for the variable. + postfix_value : numpy.ndarray + Postfix value, shape ``(n_inner,)``; the eval point for the post-sum + jacobians. + inner_sens : numpy.ndarray + Inner-variable sensitivities, shape ``(n_t * n_inner, n_p)``, + time-outer/inner-inner. + inputs : dict + Input-parameter values for the solve. + sens_names : list[str] + Sensitivity-parameter names; column order of ``inner_sens`` and the + returned array. + + Returns + ------- + numpy.ndarray + Variable sensitivities ``dvar/dp``, shape ``(n_inner, n_p)``. + + Raises + ------ + pybamm.SolverError + If a discrete-time sum's times do not match the solution's. + """ + time_integral.check_discrete_times(name, times) + + # Integrate the inner sensitivities over time, shape (n_inner, n_p). + s_integral = time_integral.postfix_sum(inner_sens, times) + if time_integral.post_sum_node is None: + return s_integral + + n_inner = int(np.asarray(time_integral.sum_node.evaluate_for_shape()).shape[0]) + post_fn = self._post_sum_leaf(name, time_integral, n_inner) + + # Evaluate dpost_dy / dpost_dp ONCE at the postfix VALUE (the CasADi + # `entries` eval point), NOT the raw trajectory. + entries = np.ascontiguousarray( + np.asarray(postfix_value, dtype=np.float64).ravel() + ) + p_stacked = post_fn.pack(inputs) + dpost_dy = post_fn.jacobian("y")(0.0, entries, p_stacked).toarray() + dpost_dp_full = post_fn.jacobian("p")(0.0, entries, p_stacked).toarray() + # Slice/reorder dpost_dp columns to the sensitivity-parameter order. + input_names = list(post_fn.input_names) + sens_idx = [input_names.index(param) for param in sens_names] + dpost_dp = dpost_dp_full[:, sens_idx] + + # dpost_dy @ s_integral + dpost_dp, mirroring postfix_sensitivities exactly. + return dpost_dy @ s_integral + dpost_dp + + +class NativeInterpolatingObservation(NativeObservation): + """Native observation of a solve that stored state derivatives. + + Its leaves can be evaluated at arbitrary times, so variables are built as + lazily-evaluated :class:`pybamm.ProcessedVariable`s that cubic-Hermite + reconstruct the state off the solution's own grid. Hermite is a no-op when + a particular solution carries no ``yps``, which matches the CasADi path. + """ + + def build_variable(self, solution, name): + """A lazily-evaluated ProcessedVariable whose leaves are compiled tapes.""" + vars_pybamm = [] + leaves = [] + time_integral = None + nstates = solution.all_ys[0].shape[0] + for model, rust_model in zip(solution.all_models, self._models, strict=True): + var_pybamm, ti, leaf = self._segment_leaf( + solution, name, model, rust_model, nstates + ) + leaves.append(leaf) + vars_pybamm.append(var_pybamm) + time_integral = ti if ti is not None else time_integral + placeholder_states = ( + [model.len_rhs_and_alg for model in solution.all_models] + if solution.variables_returned + else None + ) + return pybamm.process_variable( + name, + vars_pybamm, + NativeObserver(leaves, self, placeholder_states), + solution, + time_integral=time_integral, + ) + + +class NativeComputedObservation(NativeObservation): + """Native observation of a solve read only on its own time points. + + Variables are evaluated eagerly into grid-aligned + :class:`pybamm.ProcessedVariableComputed`s, which interpolate in time + themselves rather than re-entering the compiled leaves. + """ + + def build_variable(self, solution, name): + """An eager ProcessedVariableComputed evaluated on the solution's own grid.""" + base_variables = [] + base_variables_data = [] + # time integrals: accumulate integrands, postfix once after the loop + integrand_segments = [] + first_time_integral = None + first_var_pybamm = None + # per-segment (compiled fn, ts, ys, inputs) for the sensitivity chain rule + sens_segments = [] + for model, ts, ys, inputs, rust_model in zip( + solution.all_models, + solution.all_ts, + solution.all_ys, + solution.all_inputs, + self._models, + strict=True, + ): + if solution.variables_returned: + # No states stored; the variable is state-free, so evaluate on a + # shaped placeholder trajectory. + ys = np.zeros((model.len_rhs_and_alg, ts.size)) + var_pybamm, time_integral, leaf = self._segment_leaf( + solution, name, model, rust_model, ys.shape[0] + ) + sens_segments.append((leaf, ts, ys, inputs)) + # returns (output, n_times) F-contiguous; .T is a zero-copy time-major view + data = np.asarray(leaf.eval_trajectory(ts, ys, inputs)).T + if time_integral is not None: + integrand_segments.append(data.reshape(-1)) + if first_time_integral is None: + first_time_integral = time_integral + first_var_pybamm = var_pybamm + else: + base_variables.append(var_pybamm) + base_variables_data.append(data) + + the_integral = None + if first_time_integral is not None: + # integrate once over the full trajectory (integrand lines up 1:1 with t) + full_integrand = np.concatenate(integrand_segments) + the_integral = first_time_integral.postfix( + full_integrand, solution.t, solution.all_inputs[0] + ) + base_variables = [first_var_pybamm] + base_variables_data = [the_integral] + + var = pybamm.ProcessedVariableComputed( + base_variables, + [None] * len(base_variables), + base_variables_data, + solution, + time_indep=first_time_integral is not None, + ) + var._sensitivities = self._computed_sensitivities( + solution, name, sens_segments, first_time_integral, the_integral + ) + return var + + def _computed_sensitivities( + self, solution, name, sens_segments, time_integral, postfix_value + ): + """Forward sensitivities for the eager path, or ``{}`` when there are none.""" + if not solution.has_sensitivities(): + return {} + return native_sensitivities( + sens_segments, + solution._all_sensitivities["all"], + solution.sensitivity_names, + time_integral=time_integral, + postfix=lambda inner, sens_names: self.postfix_sensitivities( + solution.t, + name, + time_integral, + postfix_value, + inner, + solution.all_inputs[0], + sens_names, + ), + ) + + +class OutputAssembly: + """Attaching a solver's concatenated outputs-only payload to a Solution. + + A solve run with ``output_variables`` propagates no state trajectory: the + solver evaluates the requested variables itself and returns them as one row + per *flattened output component*, in variable order, so a vector variable + spans ``lens[i]`` consecutive rows. Slicing that payload by variable ordinal + instead would drop a vector variable's tail and shift every variable after + it. Every solver that can produce such a payload assembles it through this + one object, so the layout — and the time-integral postfix riding on it — + lives in one place. + + Parameters + ---------- + names : list of str + The output variables, in row order. + lens : list of int + Flattened component count per variable, so ``names[i]`` owns rows + ``sum(lens[:i])`` up to ``sum(lens[: i + 1])``. + time_integrals : dict, optional + Map of name to :class:`pybamm.ProcessedVariableTimeIntegral` for the + outputs whose rows carry an integrand rather than the variable itself. + Their postfix sum runs here, once the trajectory is in hand. + casadi_fns : dict, optional + Map of name to the CasADi function the rows were evaluated by, on the + CasADi path only. It carries the sparsity a sparse variable is unrolled + through; native rows are dense and pass nothing. + + Raises + ------ + :class:`pybamm.SolverError` + If ``names`` and ``lens`` are not 1:1. + """ + + def __init__(self, names, lens, *, time_integrals=None, casadi_fns=None): + self._names = tuple(names) + self._lens = tuple(int(length) for length in lens) + if len(self._names) != len(self._lens): + raise pybamm.SolverError( + f"Output layout mismatch: {len(self._names)} output variables but " + f"{len(self._lens)} row lengths." + ) + self._time_integrals = dict(time_integrals or {}) + self._casadi_fns = dict(casadi_fns or {}) + offsets = list(accumulate(self._lens, initial=0)) + self._n_rows = offsets[-1] + # The layout itself: one slice of payload rows per output variable. + self._rows = tuple(slice(start, end) for start, end in pairwise(offsets)) + + @classmethod + def from_casadi(cls, names, casadi_fns, *, time_integrals=None): + """The layout of a CasADi-evaluated payload, whose rows are its non-zeros. + + Parameters + ---------- + names : list of str + The output variables, in row order. + casadi_fns : dict + Map of name to the CasADi function evaluating it. + time_integrals : dict, optional + As for the constructor. + + Returns + ------- + OutputAssembly + """ + lens = [casadi_fns[name](0.0, 0.0, 0.0).sparsity().nnz() for name in names] + return cls(names, lens, time_integrals=time_integrals, casadi_fns=casadi_fns) + + @property + def names(self) -> tuple[str, ...]: + """The output variables, in row order.""" + return self._names + + @property + def lens(self) -> tuple[int, ...]: + """Flattened component count per output variable.""" + return self._lens + + @property + def n_rows(self) -> int: + """Rows in one time point of the payload.""" + return self._n_rows + + def attach(self, solution, data, *, sensitivities=None, sensitivity_names=()): + """Populate ``solution``'s variables from one outputs-only payload. + + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution to populate, built with ``variables_returned=True``. + data : array-like + Output trajectory of shape ``(n_t, n_rows)``, time-outer. + sensitivities : array-like, optional + Output sensitivities of shape ``(n_t, n_rows, n_p)``. Omit when the + solve carried none, which leaves every variable's sensitivities empty + rather than lazily recomputed — an outputs-only solve keeps no state + to recompute them from. + sensitivity_names : list of str, optional + Sensitivity-parameter names, in ``sensitivities``' column order. + + Raises + ------ + :class:`pybamm.SolverError` + If the payload does not match this layout, or if sensitivities were + requested for a variable whose CasADi rows are sparse. + """ + data = self._checked_rows(data) + if sensitivities is not None: + sensitivities = self._checked_sensitivities( + sensitivities, data.shape[0], sensitivity_names + ) + model = solution.all_models[0] + for name, rows in zip(self._names, self._rows, strict=True): + time_integral = self._time_integrals.get(name) + values = np.ascontiguousarray(data[:, rows]) + if time_integral is not None: + # These rows are the integrand's trajectory, not the variable's. + values = time_integral.postfix( + values.reshape(-1), solution.t, solution.all_inputs[0] + ) + variable = pybamm.ProcessedVariableComputed( + [model.get_processed_variable_or_event(name)], + [self._casadi_fns.get(name)], + [values], + solution, + time_indep=time_integral is not None, + ) + variable._sensitivities = ( + {} + if sensitivities is None + else self._variable_sensitivities( + name, values, sensitivities[:, rows, :], solution, sensitivity_names + ) + ) + solution._variables[name] = variable + + def stack_parameter_blocks(self, blocks, n_timesteps, sensitivity_names): + """``(n_t, n_rows, n_p)`` sensitivities from one flat block per parameter. + + Parameters + ---------- + blocks : sequence of array-like + One block per sensitivity parameter, each ``n_t * n_rows`` values in + time-outer/output-inner order. + n_timesteps : int + Number of solution time points. + sensitivity_names : list of str + Sensitivity-parameter names, in ``blocks`` order. + + Returns + ------- + :class:`numpy.ndarray` + Sensitivities laid out for :meth:`attach`. + + Raises + ------ + :class:`pybamm.SolverError` + If there is not one block per named parameter. + """ + if len(blocks) != len(sensitivity_names): + raise pybamm.SolverError( + f"Sensitivity block count mismatch: expected " + f"{len(sensitivity_names)} parameter blocks (from " + f"model.calculate_sensitivities) but the solver returned " + f"{len(blocks)}." + ) + return np.stack( + [np.asarray(block).reshape(n_timesteps, self.n_rows) for block in blocks], + axis=-1, + ) + + def _checked_rows(self, data): + """``data`` as a ``(n_t, n_rows)`` array, or a complaint about its width.""" + array = np.asarray(data) + if array.ndim != 2 or array.shape[1] != self.n_rows: + raise pybamm.SolverError( + f"Output row count mismatch: expected {self.n_rows} rows (the total " + f"flattened length of {len(self._names)} output variables) but the " + f"solver returned an array of shape {array.shape}." + ) + return array + + def _checked_sensitivities(self, sensitivities, n_timesteps, sensitivity_names): + """``sensitivities`` as a ``(n_t, n_rows, n_p)`` array, or a complaint.""" + array = np.asarray(sensitivities) + expected = (n_timesteps, self.n_rows, len(sensitivity_names)) + if array.shape != expected: + raise pybamm.SolverError( + f"Output sensitivity shape mismatch: expected {expected} (times, " + f"flattened outputs, parameters) but the solver returned " + f"{array.shape}." + ) + return array + + def _variable_sensitivities(self, name, values, block, solution, sensitivity_names): + """One variable's ``"all"`` block plus a flat vector per parameter.""" + self._reject_sparse(name) + n_timesteps, var_len, n_params = block.shape + all_sens = block.reshape(n_timesteps * var_len, n_params) + time_integral = self._time_integrals.get(name) + if time_integral is not None: + all_sens = time_integral.postfix_sensitivities( + name, values, solution.t, solution.all_inputs[0], all_sens + ) + sensitivities = {"all": all_sens} + for i, param in enumerate(sensitivity_names): + sensitivities[param] = all_sens[:, i : i + 1].reshape(-1) + return sensitivities + + def _reject_sparse(self, name): + """A CasADi-sparse variable's rows unroll, so its sensitivities cannot.""" + casadi_fn = self._casadi_fns.get(name) + if casadi_fn is None: + return + evaluated = casadi_fn(0.0, 0.0, 0.0) + sparsity = evaluated.sparsity() + if sparsity.nnz() == sparsity.numel(): + return + raise pybamm.SolverError( + f"Sensitivity of sparse variables not supported. {name} is a sparse " + f"variable with number of non-zeros {sparsity.nnz()} and shape " + f"{evaluated.shape}" + ) + + +def join_observations(runs) -> ObservationBackend: + """One backend covering ``runs``' segments, concatenated in order. + + Parameters + ---------- + runs : list[tuple[ObservationBackend, int]] + One ``(backend, n_segments)`` pair per Solution being joined, in order. + The count comes from the caller because only a native backend holds + per-segment state; the default CasADi one covers any run. + + Returns + ------- + ObservationBackend + :data:`CASADI_OBSERVATION` when no run is native, else a native backend + spanning every segment: interpolating if any native run was, so a join + never narrows how a solution can be read. A native run wins over a + CasADi one, whose segments are observed by the first native model -- + the behaviour from before this seam existed, and sound only because the + runs an experiment stitches share one discretised model. Nothing checks + that, so do not rely on it for a genuine mix. + """ + natives = [backend for backend, _ in runs if isinstance(backend, NativeObservation)] + if not natives: + return CASADI_OBSERVATION + stand_in = natives[0].primary_model + models = [] + for backend, n_segments in runs: + if isinstance(backend, NativeObservation): + models.extend(backend.segment_models) + else: + models.extend([stand_in] * n_segments) + joined = ( + NativeInterpolatingObservation + if any(isinstance(b, NativeInterpolatingObservation) for b in natives) + else NativeComputedObservation + ) + return joined._adopting( + models, cache=_join_caches([backend.compile_cache for backend in natives]) + ) + + +def _join_caches(caches): + """The compile cache for a joined backend, earlier runs winning on collision. + + Identity is preserved when every run already shares one dict, so a merged + solution keeps compiling into the solver's cache instead of into a copy. + """ + if all(cache is caches[0] for cache in caches): + return caches[0] + merged = {} + for cache in reversed(caches): + merged.update(cache) + return merged + + +#: Shared by every Solution that has not been given a native backend. +CASADI_OBSERVATION = CasadiObservation() diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable.py b/packages/pybamm/src/pybamm/solvers/processed_variable.py index 263becd5ed..b6bb598806 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable.py @@ -1,13 +1,10 @@ -import bisect - -import casadi import numpy as np import xarray as xr -from pybammsolvers import idaklu import pybamm from .base_processed_variable import BaseProcessedVariable +from .variable_observer import as_observer class ProcessedVariable(BaseProcessedVariable): @@ -27,9 +24,10 @@ class ProcessedVariable(BaseProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves. A bare list of + :class:`casadi.Function` is accepted and wrapped in a + :class:`CasadiObserver`, one function per sub-solution. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables time_integral : :class:`pybamm.ProcessedVariableTimeIntegral`, optional @@ -40,22 +38,20 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): self._name = name self.base_variables = base_variables - self.base_variables_casadi = base_variables_casadi + self._observer = as_observer(observer) self.all_ts = solution.all_ts self.all_ys = solution.all_ys self.all_yps = solution.all_yps self.all_inputs = solution.all_inputs self.all_inputs_stacked = solution.all_inputs_stacked - self.sensitivity_names = [ - name for name in solution._all_sensitivities if name != "all" - ] + self.sensitivity_names = solution.sensitivity_names self.mesh = base_variables[0].mesh self.domain = base_variables[0].domain @@ -124,67 +120,11 @@ def observe_raw(self): t = self.t_pts return self._observe_postfix(self._observe_raw(), t) - def _setup_inputs(self, t, full_range): - pybamm.logger.debug("Setting up C++ interpolation inputs") - - ts = self.all_ts - ys = self.all_ys - yps = self.all_yps - inputs = self.all_inputs_stacked - - # Remove all empty ts - idxs = np.where([ti.size > 0 for ti in ts])[0] - - # Find the indices of the time points to observe - if not full_range: - ts_nonempty = [ts[idx] for idx in idxs] - idxs_subset = _find_ts_indices(ts_nonempty, t) - idxs = idxs[idxs_subset] - - # Extract the time points and inputs - ts = [ts[idx] for idx in idxs] - ys = [ys[idx] for idx in idxs] - if self.hermite_interpolation: - yps = [yps[idx] for idx in idxs] - inputs = [inputs[idx] for idx in idxs] - - is_f_contiguous = _is_f_contiguous(ys) - - ts = idaklu.VectorRealtypeNdArray(ts) - ys = idaklu.VectorRealtypeNdArray(ys) - if self.hermite_interpolation: - yps = idaklu.VectorRealtypeNdArray(yps) - else: - yps = None - inputs = idaklu.VectorRealtypeNdArray(inputs) - - # Generate the serialized C++ functions only once - funcs_unique = {} - funcs = [None] * len(idxs) - for i in range(len(idxs)): - vars = self.base_variables_casadi[idxs[i]] - if vars not in funcs_unique: - funcs_unique[vars] = vars.serialize() - funcs[i] = funcs_unique[vars] - - return ts, ys, yps, funcs, inputs, is_f_contiguous - - def _observe_hermite(self, t): - pybamm.logger.debug("Observing and Hermite interpolating the variable") - - ts, ys, yps, funcs, inputs, _ = self._setup_inputs(t, full_range=False) - shapes = self._shape(t) - return idaklu.observe_hermite_interp(t, ts, ys, yps, inputs, funcs, shapes) - def _observe_raw(self): - pybamm.logger.debug("Observing the variable raw data") - t = self.t_pts - ts, ys, _, funcs, inputs, is_f_contiguous = self._setup_inputs( - t, full_range=True - ) - shapes = self._shape(self.t_pts) + return self._observer.observe_raw(self) - return idaklu.observe(ts, ys, inputs, funcs, is_f_contiguous, shapes) + def _observe_hermite(self, t): + return self._observer.observe_hermite(self, t) def _observe_postfix(self, entries, t): return entries @@ -284,7 +224,9 @@ def __call__( else: processed_entries = entries - if not is_sorted: + # Only the hermite path consumed the sorted times; the xarray path was + # handed the original t, so its output is already in query order. + if not is_sorted and hermite_time_interp: idxs_unsort = np.empty_like(idxs_sort) idxs_unsort[idxs_sort] = np.arange(len(t_observe)) @@ -371,6 +313,11 @@ def _check_observe_raw(self, t): return t_observe, observe_raw + @property + def name(self) -> str: + """The variable's name, as registered on the model.""" + return self._name + @property def entries(self): """ @@ -420,83 +367,7 @@ def sensitivities(self): def initialise_sensitivity_explicit_forward(self): "Set up the sensitivity dictionary" - - all_S_var = [] - for ts, ys, inputs, base_variable, dy_dp in zip( - self.all_ts, - self.all_ys, - self.all_inputs, - self.base_variables, - self.all_solution_sensitivities["all"], - strict=True, - ): - sensitivity_inputs = { - name: inputs[name] for name in self.sensitivity_names if name in inputs - } - sensitivity_inputs_stacked = casadi.vertcat( - *[sensitivity_inputs[name] for name in self.sensitivity_names] - ) - - # Set up symbolic variables - t_casadi = casadi.MX.sym("t") - y_casadi = casadi.MX.sym("y", ys.shape[0]) - p_casadi = { - name: casadi.MX.sym(name, value.shape[0]) - for name, value in sensitivity_inputs.items() - } - - p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) - - # Symbolic for sensitivity targets, concrete for the rest. Non-target - # inputs may still appear in the expression tree (e.g. from - # experiment steps) so they must be present for casadi conversion. - inputs_for_casadi = {**inputs, **p_casadi} - - var_casadi = base_variable.to_casadi( - t_casadi, y_casadi, inputs=inputs_for_casadi - ) - dvar_dy = casadi.jacobian(var_casadi, y_casadi) - dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) - - # Convert to functions and evaluate index-by-index - dvar_dy_func = casadi.Function( - "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] - ) - dvar_dp_func = casadi.Function( - "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp] - ) - dvar_dy_eval = casadi.diagcat( - *[ - dvar_dy_func(t, ys[:, idx], sensitivity_inputs_stacked) - for idx, t in enumerate(ts) - ] - ) - dvar_dp_eval = casadi.vertcat( - *[ - dvar_dp_func(t, ys[:, idx], sensitivity_inputs_stacked) - for idx, t in enumerate(ts) - ] - ) - - # Compute sensitivity - S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval - - if self.time_integral is not None: - S_var = self.time_integral.postfix_sensitivities( - self._name, self.data, ts, inputs, S_var - ) - - all_S_var.append(S_var) - - S_var = np.vstack(all_S_var) - sensitivities = {"all": S_var} - - # Add the individual sensitivity - for i, name in enumerate(self.sensitivity_names): - sensitivities[name] = S_var[:, i : i + 1].reshape(-1) - - # Save attribute - self._sensitivities = sensitivities + self._sensitivities = self._observer.sensitivities(self) def _is_discrete_time_method(self): """Check if using discrete time integral method""" @@ -553,7 +424,7 @@ def __init__( cpv = pybamm.ProcessedVariableComputed( self.base_variables, - self.base_variables_casadi, + self._observer.leaves, base_data, _stub_solution(self), ) @@ -570,7 +441,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -578,7 +449,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -618,9 +489,9 @@ class ProcessedVariable1D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ @@ -629,7 +500,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -637,7 +508,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -703,9 +574,9 @@ class ProcessedVariable2D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ @@ -714,7 +585,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -722,7 +593,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -849,9 +720,9 @@ class ProcessedVariable2DSciKitFEM(ProcessedVariable2D): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ @@ -860,7 +731,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -868,7 +739,7 @@ def __init__( super(ProcessedVariable2D, self).__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -906,7 +777,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -914,7 +785,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -994,9 +865,9 @@ class ProcessedVariable3D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ @@ -1005,7 +876,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1013,7 +884,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1189,9 +1060,9 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ @@ -1200,7 +1071,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1208,7 +1079,7 @@ def __init__( super(ProcessedVariable3D, self).__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1280,9 +1151,9 @@ class ProcessedVariableUnstructured(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables time_integral : pybamm.ProcessedVariableTimeIntegral, optional @@ -1294,7 +1165,7 @@ def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1302,7 +1173,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1434,20 +1305,6 @@ def process_variable(name: str, base_variables, *args, **kwargs): raise NotImplementedError(f"Shape not recognized for {base_variables[0]}") -def _is_f_contiguous(all_ys): - """ - Check if all the ys are f-contiguous in memory - - Args: - all_ys (list of np.ndarray): list of all ys - - Returns: - bool: True if all ys are f-contiguous - """ - - return all(isinstance(y, np.ndarray) and y.data.f_contiguous for y in all_ys) - - def _is_sorted(t): """ Check if an array is sorted @@ -1459,39 +1316,3 @@ def _is_sorted(t): bool: True if array is sorted """ return np.all(t[:-1] <= t[1:]) - - -def _find_ts_indices(ts, t): - """ - Parameters: - - ts: A list of numpy arrays (each sorted) whose values are successively increasing. - - t: A sorted list or array of values to find within ts. - - Returns: - - indices: A list of indices from `ts` such that at least one value of `t` falls within ts[idx]. - """ - - indices = [] - - # Get the minimum and maximum values of the target values `t` - t_min, t_max = t[0], t[-1] - - # Step 1: Use binary search to find the range of `ts` arrays where t_min and t_max could lie - low_idx = bisect.bisect_left([ts_arr[-1] for ts_arr in ts], t_min) - high_idx = bisect.bisect_right([ts_arr[0] for ts_arr in ts], t_max) - - # Step 2: Iterate over the identified range - for idx in range(low_idx, high_idx): - ts_min, ts_max = ts[idx][0], ts[idx][-1] - - # Binary search within `t` to check if any value falls within [ts_min, ts_max] - i = bisect.bisect_left(t, ts_min) - if i < len(t) and t[i] <= ts_max: - # At least one value of t is within ts[idx] - indices.append(idx) - - # extrapolating - if (t[-1] > ts[-1][-1]) and (len(indices) == 0 or indices[-1] != len(ts) - 1): - indices.append(len(ts) - 1) - - return indices diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable_computed.py b/packages/pybamm/src/pybamm/solvers/processed_variable_computed.py index 6045869e24..74234a638d 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable_computed.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable_computed.py @@ -85,6 +85,10 @@ def __init__( # initialise_* runs lazily on first read of `entries` / `_xr_data_array`. self._initialised = False + # Building the xr.DataArray costs more than evaluating a small + # variable, so `.data` reads must not pay for it. + self._xr_interp_args = None + self._xr_data_array_cache = None self._initialise_method, self.dimensions = self._resolve_initialise_method() def _resolve_initialise_method(self): @@ -149,6 +153,10 @@ def entries(self): @property def _xr_data_array(self): self._materialise() + if self._xr_data_array_cache is None and self._xr_interp_args is not None: + data, coords = self._xr_interp_args + self._xr_data_array_cache = xr.DataArray(data, coords=coords) + self._xr_interp_args = None return self._xr_data_array_cache def as_computed(self) -> ProcessedVariableComputed: @@ -164,11 +172,13 @@ def _unroll_nnz(self, realdata=None): # unroll in nnz != numel, otherwise copy if realdata is None: realdata = self.base_variables_data - if isinstance(self.base_variables_casadi[0], casadi.Function): # casadi fcn - sp = self.base_variables_casadi[0](0, 0, 0).sparsity() - nnz = sp.nnz() - numel = sp.numel() - row = sp.row() + # Native observation carries no CasADi function: data is already dense. + if not isinstance(self.base_variables_casadi[0], casadi.Function): + return realdata + sp = self.base_variables_casadi[0](0, 0, 0).sparsity() + nnz = sp.nnz() + numel = sp.numel() + row = sp.row() if nnz != numel: data = [None] * len(realdata) for datak in range(len(realdata)): @@ -232,12 +242,9 @@ def unroll_3D( n_dim2 = self.unroll_params["n_dim2"] n_dim3 = self.unroll_params["n_dim3"] axis_swaps = self.unroll_params["axis_swaps"] - entries = ( - np.concatenate(self._unroll_nnz(realdata), axis=0) - .transpose() - .reshape( - (len(self.t_pts), n_dim1, n_dim2, n_dim3), - ) + # time-major (n_t, output), like unroll_1D/2D + entries = np.concatenate(self._unroll_nnz(realdata), axis=0).reshape( + (len(self.t_pts), n_dim1, n_dim2, n_dim3) ) for a, b in axis_swaps: entries = np.moveaxis(entries, a, b) @@ -257,7 +264,7 @@ def unroll(self, realdata=None): def initialise_time_independent(self): self._entries = self.unroll_0D() - self._xr_data_array_cache = None + self._xr_interp_args = None def initialise_0D(self): entries = self.unroll_0D() @@ -268,7 +275,7 @@ def initialise_0D(self): ) # set up interpolation - self._xr_data_array_cache = xr.DataArray(entries, coords=[("t", self.t_pts)]) + self._xr_interp_args = (entries, [("t", self.t_pts)]) self._entries = entries @@ -323,9 +330,9 @@ def initialise_1D(self): self.first_dim_pts = edges # set up interpolation - self._xr_data_array_cache = xr.DataArray( + self._xr_interp_args = ( entries_for_interp, - coords=[(self.first_dimension, pts_for_interp), ("t", self.t_pts)], + [(self.first_dimension, pts_for_interp), ("t", self.t_pts)], ) def initialise_2D(self): @@ -452,9 +459,9 @@ def initialise_2D(self): self.second_dim_pts = second_dim_edges # set up interpolation - self._xr_data_array_cache = xr.DataArray( + self._xr_interp_args = ( entries_for_interp, - coords={ + { self.first_dimension: first_dim_pts_for_interp, self.second_dimension: second_dim_pts_for_interp, "t": self.t_pts, @@ -483,9 +490,9 @@ def initialise_2D_scikit_fem(self): self.second_dim_pts = z_sol # set up interpolation - self._xr_data_array_cache = xr.DataArray( + self._xr_interp_args = ( entries, - coords={"y": y_sol, "z": z_sol, "t": self.t_pts}, + {"y": y_sol, "z": z_sol, "t": self.t_pts}, ) def initialise_3D(self): @@ -628,9 +635,9 @@ def initialise_3D(self): self.third_dim_pts = third_dim_edges # set up interpolation - self._xr_data_array_cache = xr.DataArray( + self._xr_interp_args = ( entries_for_interp, - coords={ + { self.first_dimension: first_dim_pts_for_interp, self.second_dimension: second_dim_pts_for_interp, self.third_dimension: third_dim_pts_for_interp, @@ -672,9 +679,9 @@ def initialise_3D_scikit_fem(self): self.third_dim_pts = z_sol # set up interpolation - self._xr_data_array_cache = xr.DataArray( + self._xr_interp_args = ( entries, - coords={"x": x_sol, "y": y_sol, "z": z_sol, "t": self.t_pts}, + {"x": x_sol, "y": y_sol, "z": z_sol, "t": self.t_pts}, ) def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None): @@ -684,6 +691,14 @@ def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None): """ if self.time_indep: return self.entries + if ( + t is not None + and self.dimensions == 0 + and all(arg is None for arg in (x, r, y, z, R)) + ): + # np.interp is far cheaper than xr.interp for time-only reads; + # NaN outside the range matches xarray's fill behaviour. + return np.interp(t, self.t_pts, self.entries, left=np.nan, right=np.nan) kwargs = {"t": t, "x": x, "r": r, "y": y, "z": z, "R": R} # Remove any None arguments kwargs = {key: value for key, value in kwargs.items() if value is not None} diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py b/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py index 0fd846dea3..c9d38240fe 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py @@ -45,6 +45,29 @@ def postfix(self, entries, t_pts, inputs) -> np.ndarray: ret = self.post_sum(0.0, the_integral, inputs).full().reshape(-1) return ret + def check_discrete_times(self, var_name, t_pts) -> None: + """Reject a discrete-time sum whose times are not the solution's. + + The postfix sum only lands on the requested times when the two grids + agree, so both the CasADi and the native sensitivity paths gate on this. + + Raises + ------ + pybamm.SolverError + If this is a discrete-time sum and ``t_pts`` are not its times. + """ + if self.method != "discrete": + return + if len(t_pts) == len(self.discrete_times) and np.allclose( + t_pts, self.discrete_times, atol=1e-10 + ): + return + raise pybamm.SolverError( + f'Processing discrete-time-sum variable "{var_name}": solution times ' + "and discrete times of the time integral are not equal. Set 't_interp=discrete_sum_times' to " + f"ensure the correct times are used.\nSolution times: {t_pts}\nDiscrete Sum times: {self.discrete_times}" + ) + def postfix_sensitivities( self, var_name, @@ -53,18 +76,7 @@ def postfix_sensitivities( inputs, sensitivities, ) -> np.ndarray: - # post fix for discrete time integral won't give correct result - # if ts are not equal to the discrete times. Raise error - # in this case - if self.method == "discrete" and not ( - len(t_pts) == len(self.discrete_times) - and np.allclose(t_pts, self.discrete_times, atol=1e-10) - ): - raise pybamm.SolverError( - f'Processing discrete-time-sum variable "{var_name}": solution times ' - "and discrete times of the time integral are not equal. Set 't_interp=discrete_sum_times' to " - f"ensure the correct times are used.\nSolution times: {t_pts}\nDiscrete Sum times: {self.discrete_times}" - ) + self.check_discrete_times(var_name, t_pts) the_integral = self.postfix_sum(sensitivities, t_pts) if self.post_sum_node is None: diff --git a/packages/pybamm/src/pybamm/solvers/rust_evaluator.py b/packages/pybamm/src/pybamm/solvers/rust_evaluator.py new file mode 100644 index 0000000000..a78f80f7cf --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/rust_evaluator.py @@ -0,0 +1,67 @@ +"""Callable wrappers around Rust prep-artifacts for ``BaseSolver.process``. + +Rust evaluators follow the casadi calling convention (positional +``(t, y, stacked_inputs)``, all inputs stacked into one flat vector) but take and +return plain numpy / scipy-sparse objects, never ``casadi.DM``. Every role wraps one +``CompiledFunction``; the ``CompiledJacobian`` backing the ``jac``/``jacp`` roles is +derived lazily so groups whose jacobian is never evaluated (idaklu/diffsol, which +integrate via ``CompiledModel``) pay nothing for it. See the bundle-accessor API in +``packages/pybamm-rust/pybamm-python/src/{function,jacobian}.rs``. +""" + +import numpy as np + + +def _as_1d(arr): + return np.ascontiguousarray(np.asarray(arr, dtype=np.float64)).ravel() + + +class RustEvaluator: + """Role-specific view over a per-group ``CompiledFunction``. + + Roles: + - ``func`` -> column vector ``f(t, y, p)`` + - ``jac`` -> ``df/dy`` as scipy CSC (lazy ``cf.jacobian("y")``) + - ``jac_action`` -> ``(df/dy) @ v`` column (``cf.jvp``) + - ``jacp`` -> tuple of ``df/dp_i`` columns, one per sensitivity + parameter (lazy ``cf.jacobian("p")`` sliced to ``sens_indices``), matching + casadi's multi-output ``jacp`` convention + + Parameters + ---------- + cf : :class:`pybamm.rust.CompiledFunction` + The compiled function this evaluator views. + role : str + One of the roles above; it fixes what calling the evaluator returns. + sens_indices : list[int], optional + Positions in ``cf``'s parameter vector to keep for the ``jacp`` role, + in the order the solver expects its sensitivity columns. + """ + + def __init__(self, cf, role, sens_indices=None): + self._cf = cf + self._role = role + self._sens_indices = sens_indices + self._jac = None # CompiledJacobian, derived on first jac/jacp call + + def __getstate__(self): + state = self.__dict__.copy() + state["_jac"] = None # derived cache; rebuilt lazily after unpickle + return state + + def _jacobian(self, wrt): + if self._jac is None: + self._jac = self._cf.jacobian(wrt) + return self._jac + + def __call__(self, t, y, inputs, v=None): + t, y, p = float(t), _as_1d(y), _as_1d(inputs) + if self._role == "func": + return np.asarray(self._cf(t, y, p)).reshape(-1, 1) + if self._role == "jac": + return self._jacobian("y")(t, y, p) # scipy CSC + if self._role == "jac_action": + return np.asarray(self._cf.jvp(t, y, p, _as_1d(v))).reshape(-1, 1) + # jacp: slice df/dp columns to the requested sensitivity params + dense = self._jacobian("p")(t, y, p).toarray() + return tuple(dense[:, i].reshape(-1, 1) for i in self._sens_indices) diff --git a/packages/pybamm/src/pybamm/solvers/rust_lowering.py b/packages/pybamm/src/pybamm/solvers/rust_lowering.py new file mode 100644 index 0000000000..e6d4a3a844 --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/rust_lowering.py @@ -0,0 +1,303 @@ +"""Shared lowering of a discretised model into a Rust ``CompiledModel``.""" + +from __future__ import annotations + +import numpy as np +import numpy.typing as npt + +import pybamm + + +def rust_graph_with_inputs(model: pybamm.BaseModel, input_names): + """Build an ``ExprGraph`` with ``input_names`` registered in order. + + Graph input indices are positional, and both the stacked-input convention + shared with CasADi and the sensitivity index mapping rely on that position + matching the caller's ordering, so registration happens before any lowering. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + Discretised model, used to size vector-valued input parameters. + input_names : iterable of str + Input parameter names, in the order the solver will stack them. + + Returns + ------- + ExprGraph + Graph with every input registered. + """ + from pybamm.rust import ExprGraph + from pybamm.solvers.base_solver import rust_input_parameter_widths + + graph = ExprGraph() + widths = rust_input_parameter_widths(model) + for name in input_names: + graph.input_parameter(name, widths.get(name, 1)) + return graph + + +class RustModelLowering: + """Lower one discretised model into a Rust ``CompiledModel``. + + Holds the single ``ExprGraph`` that the state residual, output variables, + events and algebraic block are all lowered into, so common subexpressions are + shared and every part sees the same input registration. Callers compose only + the parts their solver needs, then call :meth:`compile`. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The discretised model to lower. + inputs_dict : dict + Input parameter values; only the key order is used. + """ + + def __init__(self, model: pybamm.BaseModel, inputs_dict: dict): + self.model = model + self.input_name_order = list(inputs_dict) + self.graph = rust_graph_with_inputs(model, self.input_name_order) + self._symbols: dict = {} + self._state_residual = None + self._output_exprs: list = [] + self._output_lens: list[int] = [] + self._event_exprs: list = [] + self._algebraic_expr = None + self._algebraic_var_indices: list[int] = [] + self._sens_param_indices: list[int] = [] + + def lower(self, symbol: pybamm.Symbol): + """Lower one symbol into this graph, reusing already-lowered subtrees.""" + return symbol.to_rust(self.graph, self._symbols) + + def state_residual(self, algebraic_only: bool = False): + """Lower the right-hand side, concatenated with the algebraic block. + + A DAE's residual must produce ``len_rhs_and_alg`` outputs to line up with + the mass matrix, whose algebraic rows are empty. + + Parameters + ---------- + algebraic_only : bool, optional + Lower only the algebraic block, for a model with no ``rhs`` (default + False). + + Returns + ------- + Expr + The lowered residual. + """ + model = self.model + if algebraic_only: + residual = model.concatenated_algebraic + elif model.len_alg > 0: + residual = pybamm.numpy_concatenation( + model.concatenated_rhs, model.concatenated_algebraic + ) + else: + residual = model.concatenated_rhs + self._state_residual = self.lower(residual) + return self._state_residual + + def algebraic_block(self, first_algebraic_index: int | None = None): + """Lower the algebraic sub-block used by the Newton initialisation. + + Parameters + ---------- + first_algebraic_index : int, optional + Global state index the algebraic block starts at. Defaults to + ``model.len_rhs``. + + Returns + ------- + tuple + ``(expr, var_indices)``, both empty-ish for a pure ODE. + """ + model = self.model + if model.len_alg == 0: + return None, [] + start = ( + model.len_rhs if first_algebraic_index is None else first_algebraic_index + ) + self._algebraic_expr = self.lower(model.concatenated_algebraic) + self._algebraic_var_indices = list(range(start, model.len_rhs_and_alg)) + return self._algebraic_expr, self._algebraic_var_indices + + def outputs(self, output_variables, time_integral_vars: dict | None = None): + """Lower the requested output variables and record their lengths. + + Lengths are flattened component counts, which is how the Rust core lays + out the output rows; slicing by variable ordinal instead would drop a + vector variable's tail and shift every variable after it. + + Parameters + ---------- + output_variables : list of str + Variable names to lower. + time_integral_vars : dict, optional + Map of name to :class:`pybamm.ProcessedVariableTimeIntegral`. For a + listed name the integrand's ``sum_node`` is lowered, leaving the + postfix summation to run after the solve. Omit to reject such + variables instead. + + Returns + ------- + tuple + ``(exprs, lengths)`` in ``output_variables`` order. + + Raises + ------ + :class:`pybamm.SolverError` + If a time-integral variable is requested but unsupported, or if a + variable is a tensor field. + """ + time_integral_vars = time_integral_vars or {} + for name in output_variables: + time_integral = time_integral_vars.get(name) + if time_integral is not None: + symbol = time_integral.sum_node + else: + symbol = self.model.get_processed_variable_or_event(name) + if isinstance(symbol, pybamm.TensorField): + raise pybamm.SolverError( + f"Output variable '{name}' is a tensor field, which cannot be " + "read back from a Solution. Request a scalar component with " + "pybamm.Component instead." + ) + self._output_exprs.append(self.lower(symbol)) + shape = getattr(symbol, "shape", ()) + self._output_lens.append(int(np.prod(shape)) if shape else 1) + return self._output_exprs, self._output_lens + + def termination_events(self): + """Lower every termination event expression, for root-finding.""" + self._event_exprs = [ + self.lower(event.expression) + for event in self.model.events + if event.event_type == pybamm.EventType.TERMINATION + ] + return self._event_exprs + + def rhs_evaluator(self): + """Build an evaluator for the differential block alone, over this graph. + + The consistent initialisation of ``ydot0`` needs the ``len_rhs`` rows that + the concatenated residual does not expose on a DAE, so the right-hand side + gets its own root here rather than its own graph -- every node the residual + already allocated is reused. + + Returns + ------- + :class:`pybamm.solvers.rust_evaluator.RustEvaluator` + Callable ``(t, y, p)`` returning the right-hand side as a column. + """ + from pybamm.solvers.rust_evaluator import RustEvaluator + + compiled = self.graph.compile( + self.lower(self.model.concatenated_rhs), + name="RHS", + n_states=self.model.len_rhs_and_alg, + ) + return RustEvaluator(compiled, "func") + + def bind_generic_evaluators(self, rust_model) -> None: + """Serve the model's backend-agnostic evaluator slots from this lowering. + + ``BaseSolver.set_up`` leaves ``rhs_eval`` and ``terminate_events_eval`` + unset on the native path, so the shared helpers that read them + (consistent initialisation, the event-violation check, event attribution) + are served from here instead of from a graph per expression. The event + evaluators are views onto ``rust_model``'s own event tapes -- the roots + its fused root-finding tape is built from -- so a Python-side event value + cannot drift from the one the integrator roots on. + + Parameters + ---------- + rust_model : CompiledModel + The model compiled from this lowering. Its ``events`` follow + :meth:`termination_events` order, which is ``model.events`` order. + """ + from pybamm.solvers.rust_evaluator import RustEvaluator + + self.model.rhs_eval = self.rhs_evaluator() + self.model.terminate_events_eval = [ + RustEvaluator(compiled, "func") for compiled in rust_model.events + ] + + def sensitivity_indices(self, sensitivity_names): + """Map sensitivity parameter names to graph input indices. + + The Rust core differentiates against positions in the global input array, + so names absent from the solve's inputs are dropped rather than indexed. + + Parameters + ---------- + sensitivity_names : list of str + Requested parameter names, typically ``model.calculate_sensitivities``. + + Returns + ------- + tuple + ``(indices, names)``, filtered to the inputs actually supplied and + kept in the requested order. + """ + indices: list[int] = [] + names: list[str] = [] + for name in sensitivity_names or []: + if name in self.input_name_order: + indices.append(self.input_name_order.index(name)) + names.append(name) + self._sens_param_indices = indices + return indices, names + + @property + def n_inputs(self) -> int: + """Packed width of every input registered, including any added while + lowering, which may exceed the inputs the solver passes.""" + return self.graph.n_inputs() + + def mass_matrix_csr( + self, + ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.int64], npt.NDArray[np.int64]]: + """Return the mass matrix as the ``(data, indptr, indices)`` CSR triple.""" + mass = self.model.mass_matrix.entries + return ( + mass.data.astype(np.float64), + mass.indptr.astype(np.int64), + mass.indices.astype(np.int64), + ) + + def compile(self): + """Compile everything lowered so far into a ``CompiledModel``. + + Returns + ------- + CompiledModel + The compiled model. + + Raises + ------ + :class:`pybamm.SolverError` + If :meth:`state_residual` has not been called. + """ + from pybamm.rust import CompiledModel + + if self._state_residual is None: + raise pybamm.SolverError( + "Cannot compile a Rust model before lowering its state residual; " + "call state_residual() first." + ) + data, indptr, indices = self.mass_matrix_csr() + return CompiledModel.from_expr( + self.graph, + self._state_residual, + data, + indptr, + indices, + n_inputs=self.n_inputs, + sens_param_indices=self._sens_param_indices, + output_exprs=self._output_exprs, + algebraic_expr=self._algebraic_expr, + algebraic_variable_indices=self._algebraic_var_indices, + event_exprs=self._event_exprs, + ) diff --git a/packages/pybamm/src/pybamm/solvers/scipy_solver.py b/packages/pybamm/src/pybamm/solvers/scipy_solver.py index 35d126dfbc..0e4a3299b3 100644 --- a/packages/pybamm/src/pybamm/solvers/scipy_solver.py +++ b/packages/pybamm/src/pybamm/solvers/scipy_solver.py @@ -8,6 +8,7 @@ import scipy.integrate as it import pybamm +from pybamm.solvers.base_solver import stack_inputs class ScipySolver(pybamm.BaseSolver): @@ -94,10 +95,10 @@ def _integrate_single(self, model, t_eval, inputs_dict, y0): "Sensitivity analysis is not implemented for the Scipy solver." ) - # Save inputs dictionary, and if necessary convert inputs to a casadi vector + # Save inputs dictionary, and if necessary convert inputs to a stacked vector inputs_dict = inputs_dict or {} - if model.convert_to_format == "casadi": - inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) + if model.uses_stacked_inputs: + inputs = stack_inputs(inputs_dict, model.convert_to_format) else: inputs = inputs_dict @@ -133,7 +134,7 @@ def rhs(t, y): def event_wrapper(event): def event_fn(t, y): - return event(t, y, inputs) + return np.asarray(event(t, y, inputs)).item() event_fn.terminal = True return event_fn diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index b67ce33c7e..46cd7618da 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -6,6 +6,7 @@ import json import numbers import pickle +from dataclasses import dataclass from functools import cached_property from itertools import chain @@ -16,6 +17,11 @@ import pybamm from pybamm.codegen.compilation import aot_compile +from pybamm.solvers.observation import ( + CASADI_OBSERVATION, + ObservationBackend, + join_observations, +) class NumpyEncoder(json.JSONEncoder): @@ -32,6 +38,29 @@ def default(self, obj): return json.JSONEncoder.default(self, obj) # pragma: no cover +@dataclass +class SolverStatistics: + """Solver statistics common across all ODE/DAE solver backends. + + Populated by solvers that track internal step/iteration counts. + Fields mirror the SUNDIALS IDA / diffsol BDF statistics. + """ + + number_of_steps: int + number_of_linear_solver_setups: int + number_of_nonlinear_solver_iterations: int + number_of_nonlinear_solver_fails: int + number_of_error_test_failures: int + number_of_linear_solver_setups_from_checkpoint: int = 0 + number_of_linear_solver_setups_from_first_convergence_fail: int = 0 + number_of_linear_solver_setups_from_second_convergence_fail: int = 0 + number_of_linear_solver_setups_from_error_test_fail: int = 0 + number_of_linear_solver_setups_from_step_success: int = 0 + ic_time_secs: float = 0.0 + solver_setup_time_secs: float = 0.0 + sens_error_control_relaxed: bool = False + + class SolutionBase: """Base class for all PyBaMM solution types (time-series and EIS). @@ -262,7 +291,9 @@ def __init__( all_ts = [all_ts] if _validate_time_structure: self._ensure_sorted_t(all_ts, "all_ts") - if not isinstance(all_ys, list): + if all_ys is None: + all_ys = [None] + elif not isinstance(all_ys, list): all_ys = [all_ys] if not isinstance(all_models, list): all_models = [all_models] @@ -289,6 +320,10 @@ def __init__( # computed lazily on first access; see the `observable` property self._observable = None + # How variables are read back; a solver replaces it wholesale and + # every derived Solution carries it. + self._observation = CASADI_OBSERVATION + # Set up inputs if not isinstance(all_inputs, list): all_inputs_copy = dict(all_inputs) @@ -320,6 +355,7 @@ def __init__( super().__init__() self.integration_time = None + self.solver_statistics = None self._all_inputs_stacked = None self._all_inputs_casadi = None @@ -357,6 +393,11 @@ def __init__( def has_sensitivities(self) -> bool: return len(self._all_sensitivities) > 0 + @property + def sensitivity_names(self) -> list[str]: + """Sensitivity-parameter names, the column order of every ``"all"`` block.""" + return [key for key in self._all_sensitivities if key != "all"] + @staticmethod def _ensure_t_evals(all_ts, all_t_evals): # all_ts is already checked upstream @@ -471,6 +512,9 @@ def check_ys_are_not_too_large(self): # We only care about the cases where y is growing too large without any # restraint, so if y gets large in the middle then comes back down that is ok t, y, model = self.all_ts[-1], self.all_ys[-1], self.all_models[-1] + if y is None: + # Outputs-only solve - no full state to check + return t = t[-1] y = y[:, -1] @@ -509,8 +553,18 @@ def all_models(self): @property def all_inputs_stacked(self) -> list[np.ndarray]: if self._all_inputs_stacked is None: + # Flatten each value first: a mix of scalar and vector-width inputs is + # a ragged sequence numpy refuses to coerce into one array. self._all_inputs_stacked = [ - np.asarray(list(inp.values())).reshape(-1) for inp in self.all_inputs + np.concatenate( + [ + np.atleast_1d(np.asarray(v, dtype=float)).ravel() + for v in inp.values() + ] + ) + if inp + else np.array([]) + for inp in self.all_inputs ] return self._all_inputs_stacked @@ -617,6 +671,7 @@ def first_state(self): ) # stacked/casadi stay lazy; built from all_inputs[:1] on first access new_sol._sub_solutions = self.sub_solutions[:1] + new_sol._observation = self._observation[:1] new_sol.solve_time = 0 new_sol.integration_time = 0 @@ -661,6 +716,7 @@ def last_state(self): ) # stacked/casadi stay lazy; built from all_inputs[-1:] on first access new_sol._sub_solutions = self.sub_solutions[-1:] + new_sol._observation = self._observation[-1:] new_sol.solve_time = 0 new_sol.integration_time = 0 new_sol.set_up_time = 0 @@ -694,6 +750,20 @@ def update_summary_variables(self, all_summary_variables): self, cycle_summary_variables=all_summary_variables ) + @property + def observation(self) -> ObservationBackend: + """How this Solution reads its variables back. + + Defaults to CasADi conversion; a solver assigns a native backend when + the model was lowered to Rust. The backend must cover this Solution's + sub-solutions, so it has one segment per entry of ``all_ys``. + """ + return self._observation + + @observation.setter + def observation(self, backend: ObservationBackend) -> None: + self._observation = backend + def update(self, variables: str | list[str]): """Add ProcessedVariables to the dictionary of variables in the solution""" # Single variable @@ -704,79 +774,8 @@ def update(self, variables: str | list[str]): for variable in variables: self._update_variable(variable) - def _update_model_variable( - self, - model: pybamm.BaseModel, - var_pybamm: pybamm.Symbol, - time_integral: pybamm.ProcessedVariableTimeIntegral | None, - inputs: dict, - ys_shape: tuple, - cache_key, - ): - _var_casadi = model._variables_casadi.get(cache_key) - if _var_casadi is not None: - return _var_casadi, var_pybamm, time_integral - - var_casadi, var_pybamm, time_integral = self._convert_to_casadi( - var_pybamm, inputs, ys_shape - ) - - # Only cache if it's not a time integral - if time_integral is None: - model._variables_casadi[cache_key] = var_casadi - return var_casadi, var_pybamm, time_integral - def _update_variable(self, name: str): - time_integral = None - pybamm.logger.debug(f"Post-processing {name}") - - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_pybamm = [ - model.get_processed_variable_or_event(name) for model in self.all_models - ] - vars_casadi = [None] * len(self.all_models) - for i, (model, ys, inputs) in enumerate( - zip(self.all_models, self.all_ys, self.all_inputs, strict=True) - ): - _var_pybamm = vars_pybamm[i] - if self.variables_returned and _var_pybamm.has_symbol_of_classes( - pybamm.expression_tree.state_vector.StateVector - ): - raise KeyError( - f"Cannot process variable '{name}' as it was not part of the " - "solve. Please re-run the solve with `output_variables` set to " - "include this variable." - ) - if isinstance(_var_pybamm, pybamm.VectorField): - comp_casadi = [] - for k, comp in enumerate(_var_pybamm.components): - cc, _, _ = self._update_model_variable( - model, - comp, - inputs=inputs, - ys_shape=ys.shape, - time_integral=None, - cache_key=f"{name}[{k}]", - ) - comp_casadi.append(cc) - vars_casadi[i] = comp_casadi - else: - var_casadi, var_pybamm, time_integral = self._update_model_variable( - model, - _var_pybamm, - inputs=inputs, - ys_shape=ys.shape, - time_integral=time_integral, - cache_key=name, - ) - vars_pybamm[i] = var_pybamm - vars_casadi[i] = var_casadi - var = pybamm.process_variable( - name, vars_pybamm, vars_casadi, self, time_integral=time_integral - ) - - self._variables[name] = var + self._variables[name] = self._observation.build_variable(self, name) def observe(self, symbol: pybamm.Symbol) -> pybamm.ProcessedVariable: """ @@ -1180,6 +1179,14 @@ def __add__(self, other): # Set sub_solutions new_sol._sub_solutions = self.sub_solutions + other.sub_solutions + # Segment count is unchanged by the leading-sample drop that + # _segment_series applies to the ys data, so the runs just append. + new_sol._observation = join_observations( + [ + (self._observation, len(self.all_ys)), + (other._observation, len(other.all_ys)), + ] + ) # update variables which were derived at the solver stage if any([self.variables_returned, other.variables_returned]): @@ -1286,6 +1293,9 @@ def from_sub_solutions(cls, sub_solutions): new_sol.closest_event_idx = segments[-1].closest_event_idx # leave stacked/casadi unset; built lazily from all_inputs (casadi is costly) new_sol._sub_solutions = sub_sols + new_sol._observation = join_observations( + [(s._observation, len(s.all_ys)) for s in segments] + ) for attr in ["solve_time", "integration_time", "set_up_time"]: vals = [getattr(s, attr, None) for s in segments] @@ -1330,6 +1340,7 @@ def copy(self): new_sol.solve_time = self.solve_time new_sol.integration_time = self.integration_time new_sol.set_up_time = self.set_up_time + new_sol._observation = self._observation # copy over variables which were derived at the solver stage if self._variables and all( diff --git a/packages/pybamm/src/pybamm/solvers/variable_observer.py b/packages/pybamm/src/pybamm/solvers/variable_observer.py new file mode 100644 index 0000000000..5a3c4ebf8c --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/variable_observer.py @@ -0,0 +1,531 @@ +"""Per-variable observation strategies for :class:`pybamm.ProcessedVariable`. + +A :class:`VariableObserver` owns one variable's per-sub-solution leaves and +knows how to evaluate them on the solution grid, off-grid via cubic Hermite, +and through the forward chain rule for sensitivities. ``ProcessedVariable`` +holds exactly one, chosen when it is built, so nothing downstream re-decides +which backend is in play. +""" + +from __future__ import annotations + +import bisect +from abc import ABC, abstractmethod + +import casadi +import numpy as np +from pybammsolvers import idaklu + +import pybamm + + +class SegmentSelector: + """The single rule for which sub-solutions cover a set of query times. + + A variable's ``all_ts`` is frozen for its lifetime, so the non-empty + segments and their end times are found once here rather than on every + observe call. + + Parameters + ---------- + all_ts : list[numpy.ndarray] + Per-sub-solution time arrays, successively increasing. + """ + + def __init__(self, all_ts): + self.indices = np.where([ti.size > 0 for ti in all_ts])[0] + self._starts = [all_ts[idx][0] for idx in self.indices] + self._ends = [all_ts[idx][-1] for idx in self.indices] + + def select(self, t, full_range): + """Indices into ``all_ts`` covering ``t``, ascending. + + Empty segments are always dropped; when ``full_range`` is False only + the segments whose span contains at least one of ``t`` are kept. + """ + if full_range: + return self.indices + return self.indices[_find_ts_indices(self._starts, self._ends, t)] + + +class VariableObserver(ABC): + """How one variable's leaves are evaluated over a solution's segments. + + Implementations read only these attributes of the ``variable`` handed to + them: ``all_ts``, ``all_ys``, ``all_yps``, ``all_inputs``, + ``all_inputs_stacked``, ``t_pts``, ``hermite_interpolation``, + ``time_integral``, ``base_variables``, ``sensitivity_names``, + ``all_solution_sensitivities`` and ``_shape``. + """ + + #: Derived caches: built on demand, never pickled (see __getstate__). + _selector = None + _serialised = None + + def __getstate__(self): + """Pickle without the derived caches, which a Solution should not carry.""" + state = self.__dict__.copy() + for key in ("_selector", "_serialised"): + state.pop(key, None) + return state + + @property + @abstractmethod + def leaves(self) -> list: + """The variable's per-sub-solution evaluable leaves, in solve order.""" + + def segments(self, variable, t, full_range): + """Indices of ``variable``'s sub-solutions covering ``t``.""" + if self._selector is None: + self._selector = SegmentSelector(variable.all_ts) + return self._selector.select(t, full_range) + + @abstractmethod + def observe_raw(self, variable): + """Evaluate on the solution's own time points, shaped by ``_shape``.""" + + @abstractmethod + def observe_hermite(self, variable, t): + """Evaluate at arbitrary sorted times ``t``, cubic-Hermite in state.""" + + @abstractmethod + def sensitivities(self, variable) -> dict: + """``{"all": (N, n_p), param: (N,)}`` forward sensitivities of the variable.""" + + +class CasadiObserver(VariableObserver): + """Observation through serialised CasADi functions and the IDAKLU kernels. + + Parameters + ---------- + leaves : list of :class:`casadi.Function` + One function per sub-solution, evaluating ``(t, y, p_stacked)``. + """ + + def __init__(self, leaves): + self._leaves = leaves + + @property + def leaves(self): + return self._leaves + + def _serialise(self, idxs): + """Serialised leaves for `idxs`, memoised by leaf identity. + + CasADi functions are immutable and serialising one is about half the + cost of an observe call, so the bytes are built once per leaf. + """ + if self._serialised is None: + self._serialised = {} + serialised = self._serialised + funcs = [None] * len(idxs) + for i, idx in enumerate(idxs): + leaf = self._leaves[idx] + key = id(leaf) + if key not in serialised: + serialised[key] = leaf.serialize() + funcs[i] = serialised[key] + return funcs + + def _setup(self, variable, t, full_range): + """Per-segment IDAKLU inputs: ``(ts, ys, yps, funcs, inputs, is_f_contiguous)``.""" + pybamm.logger.debug("Setting up C++ interpolation inputs") + idxs = self.segments(variable, t, full_range) + hermite = variable.hermite_interpolation + all_ts, all_ys = variable.all_ts, variable.all_ys + + ts = [all_ts[idx] for idx in idxs] + ys = [all_ys[idx] for idx in idxs] + yps = [variable.all_yps[idx] for idx in idxs] if hermite else None + inputs = [variable.all_inputs_stacked[idx] for idx in idxs] + + is_f_contiguous = _is_f_contiguous(ys) + + ts = idaklu.VectorRealtypeNdArray(ts) + ys = idaklu.VectorRealtypeNdArray(ys) + yps = idaklu.VectorRealtypeNdArray(yps) if hermite else None + inputs = idaklu.VectorRealtypeNdArray(inputs) + + return ts, ys, yps, self._serialise(idxs), inputs, is_f_contiguous + + def observe_raw(self, variable): + pybamm.logger.debug("Observing the variable raw data") + t = variable.t_pts + ts, ys, _, funcs, inputs, is_f_contiguous = self._setup( + variable, t, full_range=True + ) + return idaklu.observe( + ts, ys, inputs, funcs, is_f_contiguous, variable._shape(t) + ) + + def observe_hermite(self, variable, t): + pybamm.logger.debug("Observing and Hermite interpolating the variable") + ts, ys, yps, funcs, inputs, _ = self._setup(variable, t, full_range=False) + return idaklu.observe_hermite_interp( + t, ts, ys, yps, inputs, funcs, variable._shape(t) + ) + + def sensitivities(self, variable): + sensitivity_names = variable.sensitivity_names + all_S_var = [] + for ts, ys, inputs, base_variable, dy_dp in zip( + variable.all_ts, + variable.all_ys, + variable.all_inputs, + variable.base_variables, + variable.all_solution_sensitivities["all"], + strict=True, + ): + sensitivity_inputs = { + name: inputs[name] for name in sensitivity_names if name in inputs + } + sensitivity_inputs_stacked = casadi.vertcat( + *[sensitivity_inputs[name] for name in sensitivity_names] + ) + + # Set up symbolic variables + t_casadi = casadi.MX.sym("t") + y_casadi = casadi.MX.sym("y", ys.shape[0]) + p_casadi = { + name: casadi.MX.sym(name, value.shape[0]) + for name, value in sensitivity_inputs.items() + } + + p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) + + # Symbolic for sensitivity targets, concrete for the rest. Non-target + # inputs may still appear in the expression tree (e.g. from + # experiment steps) so they must be present for casadi conversion. + inputs_for_casadi = {**inputs, **p_casadi} + + var_casadi = base_variable.to_casadi( + t_casadi, y_casadi, inputs=inputs_for_casadi + ) + dvar_dy = casadi.jacobian(var_casadi, y_casadi) + dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) + + # Convert to functions and evaluate index-by-index + dvar_dy_func = casadi.Function( + "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] + ) + dvar_dp_func = casadi.Function( + "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp] + ) + dvar_dy_eval = casadi.diagcat( + *[ + dvar_dy_func(t, ys[:, idx], sensitivity_inputs_stacked) + for idx, t in enumerate(ts) + ] + ) + dvar_dp_eval = casadi.vertcat( + *[ + dvar_dp_func(t, ys[:, idx], sensitivity_inputs_stacked) + for idx, t in enumerate(ts) + ] + ) + + # Compute sensitivity + S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval + + if variable.time_integral is not None: + S_var = variable.time_integral.postfix_sensitivities( + variable.name, variable.data, ts, inputs, S_var + ) + + all_S_var.append(S_var) + + return pack_sensitivity_dict(np.vstack(all_S_var), sensitivity_names) + + +class NativeObserver(VariableObserver): + """Observation through compiled Rust tapes lowered into the retained graph. + + Parameters + ---------- + leaves : list of :class:`pybamm.rust.CompiledFunction` + One compiled tape per sub-solution, evaluating ``(t, y, inputs)``. + backend : :class:`pybamm.solvers.observation.NativeObservation` + Owner of the retained graph and the compile cache; consulted for the + time-integral post-sum tape. + placeholder_states : list[int] or None + Per-sub-solution state count for outputs-only solves, which store no + states and need a shaped zero trajectory. ``None`` when states are real. + """ + + def __init__(self, leaves, backend, placeholder_states=None): + self._leaves = leaves + self._backend = backend + self._placeholder_states = placeholder_states + + @property + def leaves(self): + return self._leaves + + def _setup(self, variable, t, full_range): + """Per-segment ``(ts, ys, yps, inputs, leaves)`` as plain numpy and dicts. + + Inputs stay dict-shaped: the compiled tape packs by name, so there are + no stacking-order concerns. + """ + idxs = self.segments(variable, t, full_range) + ts = [variable.all_ts[idx] for idx in idxs] + if self._placeholder_states is None: + ys = [variable.all_ys[idx] for idx in idxs] + else: + ys = [ + np.zeros((self._placeholder_states[idx], ti.size)) + for idx, ti in zip(idxs, ts, strict=True) + ] + yps = ( + [variable.all_yps[idx] for idx in idxs] + if variable.hermite_interpolation + else None + ) + inputs = [variable.all_inputs[idx] for idx in idxs] + leaves = [self._leaves[idx] for idx in idxs] + return ts, ys, yps, inputs, leaves + + def observe_raw(self, variable): + # eval_trajectory returns (output_len, n_t) F-contiguous, so the flat + # spatial index varies fastest and reshape must use order="F". + t = variable.t_pts + ts, ys, _, inputs, leaves = self._setup(variable, t, full_range=True) + cols = [ + np.asarray(leaf.eval_trajectory(t_i, y_i, inp_i)) + for leaf, t_i, y_i, inp_i in zip(leaves, ts, ys, inputs, strict=True) + ] + return np.concatenate(cols, axis=1).reshape(variable._shape(t), order="F") + + def observe_hermite(self, variable, t): + # Route each query time to its segment (mirrors observe.cpp's sequential + # <=/> knot-window scan), Hermite-reconstruct, then concatenate. + ts, ys, yps, inputs, leaves = self._setup(variable, t, full_range=False) + n_segments = len(ts) + cols = [] + i = 0 + n = len(t) + for seg_idx, (leaf, t_i, y_i, yp_i, inp_i) in enumerate( + zip(leaves, ts, ys, yps, inputs, strict=True) + ): + if i >= n: + break + is_last_segment = seg_idx == n_segments - 1 + if is_last_segment: + j = n + else: + j = i + np.searchsorted(t[i:], t_i[-1], side="right") + if j <= i: + continue + query = t[i:j] + if t_i.size < 2: + # No interval to Hermite-interpolate within: fall back to a + # direct eval, holding the single known state constant. + y_query = np.repeat(y_i, len(query), axis=1) + cols.append(np.asarray(leaf.eval_trajectory(query, y_query, inp_i))) + else: + cols.append( + np.asarray( + leaf.eval_trajectory_hermite(query, t_i, y_i, yp_i, inp_i) + ) + ) + i = j + return np.concatenate(cols, axis=1).reshape(variable._shape(t), order="F") + + def sensitivities(self, variable): + segments = list( + zip( + self._leaves, + variable.all_ts, + variable.all_ys, + variable.all_inputs, + strict=True, + ) + ) + return native_sensitivities( + segments, + variable.all_solution_sensitivities["all"], + variable.sensitivity_names, + time_integral=variable.time_integral, + postfix=lambda inner, sens_names: self._backend.postfix_sensitivities( + variable.t_pts, + variable.name, + variable.time_integral, + variable.entries, + inner, + variable.all_inputs[0], + sens_names, + ), + ) + + +def as_observer(leaves) -> VariableObserver: + """Coerce a :class:`ProcessedVariable` leaf argument to an observer. + + A bare list of :class:`casadi.Function` — the historical form, and what the + default backend passes — becomes a :class:`CasadiObserver`. + """ + if isinstance(leaves, VariableObserver): + return leaves + return CasadiObserver(leaves) + + +def check_variable_in_solve(solution, name, var_pybamm) -> None: + """Reject a state-dependent variable an outputs-only solve did not store. + + Raises + ------ + KeyError + If the solve returned variables only and ``var_pybamm`` reads states. + """ + if not solution.variables_returned: + return + if var_pybamm.has_symbol_of_classes( + pybamm.expression_tree.state_vector.StateVector + ): + raise KeyError( + f"Cannot process variable '{name}' as it was not part of the " + "solve. Please re-run the solve with `output_variables` set to " + "include this variable." + ) + + +def native_sensitivities(segments, dy_dp_segments, sens_names, time_integral, postfix): + """Forward sensitivities of a natively-observed variable, packed for reading. + + Parameters + ---------- + segments : list[tuple] + One ``(compiled leaf, ts, ys, inputs)`` per sub-solution, in solve order. + dy_dp_segments : list[numpy.ndarray] + The matching state sensitivities, one block per sub-solution. + sens_names : list[str] + Sensitivity-parameter names; the column order throughout. + time_integral : pybamm.ProcessedVariableTimeIntegral or None + Set when the variable is time-integrated, so the chain rule runs on the + integrand and ``postfix`` finishes it. + postfix : callable + ``(inner_sensitivities, sens_names) -> dvar/dp``, applying the post-sum + chain rule at the postfix value. Only called for a time integral. + + Returns + ------- + dict + ``{"all": (N, n_p), param: (N,)}``, as :meth:`ProcessedVariable.sensitivities`. + """ + inner = np.vstack( + [ + chain_rule_sensitivities(leaf, ts, ys, inputs, dy_dp, sens_names) + for (leaf, ts, ys, inputs), dy_dp in zip( + segments, dy_dp_segments, strict=True + ) + ] + ) + if time_integral is None: + return pack_sensitivity_dict(inner, sens_names) + return pack_sensitivity_dict(postfix(inner, sens_names), sens_names) + + +def chain_rule_sensitivities(cf, ts, ys, inputs, dy_dp, sens_names): + """Variable sensitivities ``dvar/dp`` for one sub-solution via jvp_trajectory. + + Computes ``S_var[:, k] = dvar_dy(t)·yS_k(t) + dvar_dp(t)·e_k`` for every + sensitivity parameter, one Rust forward sweep per parameter. + + Parameters + ---------- + cf : pybamm.rust.CompiledFunction + Compiled observed-variable function. + ts : numpy.ndarray + Sub-solution times, shape ``(n_t,)``. + ys : numpy.ndarray + State trajectory, shape ``(n_states, n_t)``. + inputs : dict or numpy.ndarray + Input-parameter values for this sub-solution. + dy_dp : numpy.ndarray + State sensitivities, shape ``(n_t * n_states, n_p)``, time-outer / + state-inner, columns ordered as ``sens_names``. + sens_names : list[str] + Sensitivity-parameter names; column order of ``dy_dp``. + + Returns + ------- + numpy.ndarray + Variable sensitivities, shape ``(n_t * output_len, n_p)``, time-outer / + output-inner — the layout the CasADi forward path produces. + """ + n_states, n_t = ys.shape + out_len = cf.output_len + input_names = cf.input_names + sensitivities = np.empty((n_t * out_len, len(sens_names))) + for k, name in enumerate(sens_names): + # yS for this parameter, reshaped (n_t, n_states) then to (n_states, n_t) + vy_k = np.ascontiguousarray(dy_dp[:, k].reshape(n_t, n_states).T) + if name in input_names: + vp = np.zeros(cf.n_inputs) + vp[input_names.index(name)] = 1.0 + col = cf.jvp_trajectory(ts, ys, inputs, vy_k, vp=vp) + else: + # variable has no direct dependence on this parameter; dvar_dp == 0 + col = cf.jvp_trajectory(ts, ys, inputs, vy_k) + # (out_len, n_t) -> time-outer / output-inner flat column + sensitivities[:, k] = np.asarray(col).T.reshape(-1) + return sensitivities + + +def pack_sensitivity_dict(S_var, sens_names): + """Pack a sensitivity matrix ``(N, n_p)`` into the ``{"all", per-param}`` dict. + + Matches :meth:`ProcessedVariable.sensitivities`: the full block under + ``"all"`` plus one flat ``(N,)`` vector per parameter in ``sens_names`` order. + """ + sensitivities = {"all": S_var} + for i, name in enumerate(sens_names): + sensitivities[name] = S_var[:, i : i + 1].reshape(-1) + return sensitivities + + +def _is_f_contiguous(all_ys): + """ + Check if all the ys are f-contiguous in memory + + Args: + all_ys (list of np.ndarray): list of all ys + + Returns: + bool: True if all ys are f-contiguous + """ + + return all(isinstance(y, np.ndarray) and y.data.f_contiguous for y in all_ys) + + +def _find_ts_indices(starts, ends, t): + """ + Parameters: + - starts, ends: First and last time of each segment, successively increasing. + - t: A sorted list or array of values to find within the segments. + + Returns: + - indices: Positions in `starts`/`ends` whose segment contains a value of `t`. + """ + + indices = [] + + # Get the minimum and maximum values of the target values `t` + t_min, t_max = t[0], t[-1] + + # Step 1: Use binary search to find the range of segments where t_min and t_max could lie + low_idx = bisect.bisect_left(ends, t_min) + high_idx = bisect.bisect_right(starts, t_max) + + # Step 2: Iterate over the identified range + for idx in range(low_idx, high_idx): + # Binary search within `t` to check if any value falls within the segment + i = bisect.bisect_left(t, starts[idx]) + if i < len(t) and t[i] <= ends[idx]: + indices.append(idx) + + # extrapolating + if (t_max > ends[-1]) and (len(indices) == 0 or indices[-1] != len(ends) - 1): + indices.append(len(ends) - 1) + + return indices diff --git a/packages/pybamm/src/pybamm/spatial_methods/finite_volume.py b/packages/pybamm/src/pybamm/spatial_methods/finite_volume.py index 902d199fdf..58ff7f5e97 100644 --- a/packages/pybamm/src/pybamm/spatial_methods/finite_volume.py +++ b/packages/pybamm/src/pybamm/spatial_methods/finite_volume.py @@ -63,9 +63,7 @@ def spatial_variable(self, symbol): if symbol.evaluates_on_edges("primary"): if hasattr(symbol_mesh, "length"): edges = self._get_edges_symbolic_mesh(symbol.domains["primary"]) - entries = pybamm.kronecker_product( - pybamm.Matrix(np.ones(repeats)), edges - ) + entries = self._repeat_vector(edges, repeats) entries.domains = symbol.domains else: entries = pybamm.Vector( @@ -74,9 +72,7 @@ def spatial_variable(self, symbol): else: if hasattr(symbol_mesh, "length"): nodes = self._get_nodes_symbolic_mesh(symbol.domains["primary"]) - entries = pybamm.kronecker_product( - pybamm.Matrix(np.ones(repeats)), nodes - ) + entries = self._repeat_vector(nodes, repeats) entries.domains = symbol.domains else: entries = pybamm.Vector( @@ -85,6 +81,29 @@ def spatial_variable(self, symbol): return entries + def _repeat_vector(self, vector, repeats): + """Stack ``repeats`` copies of ``vector``, as ``np.tile`` does for a fixed mesh. + + Parameters + ---------- + vector : :class:`pybamm.Symbol` + The column vector to repeat. + repeats : int + Number of copies to stack. + + Returns + ------- + :class:`pybamm.Symbol` + The stacked vector. + """ + # A Kronecker product against a ones-column says the same thing, but stays + # symbolic here and no evaluator backend can lower that. + if not isinstance(vector, pybamm.Symbol): + return pybamm.Vector(np.tile(vector, repeats)) + if repeats == 1: + return vector + return pybamm.numpy_concatenation(*([vector] * repeats)) + def gradient(self, symbol, discretised_symbol, boundary_conditions): """Matrix-vector multiplication to implement the gradient operator. See :meth:`pybamm.SpatialMethod.gradient` @@ -102,10 +121,12 @@ def gradient(self, symbol, discretised_symbol, boundary_conditions): ) # note in 1D cartesian, cylindrical and spherical grad are the same - gradient_matrix = self.gradient_matrix(domain, symbol.domains) + gradient_matrix, row_scale = self._gradient_parts(domain, symbol.domains) # Multiply by gradient matrix out = gradient_matrix @ discretised_symbol + if row_scale is not None: + out = row_scale * out # Add Neumann boundary conditions, if defined if symbol in boundary_conditions: @@ -212,11 +233,30 @@ def gradient_matrix(self, domain, domains): :class:`pybamm.Matrix` The (sparse) finite volume gradient matrix for the domain """ + matrix, row_scale = self._gradient_parts(domain, domains) + return matrix if row_scale is None else matrix * row_scale + + def _gradient_parts(self, domain, domains): + """The gradient matrix, split into a constant stencil and any row scaling. + + Parameters + ---------- + domain : str + The domain in which to compute the gradient matrix. + domains : dict + The domain and auxiliary domains of the symbol being differentiated. + + Returns + ------- + tuple[:class:`pybamm.Matrix`, :class:`pybamm.Symbol` or None] + The stencil, and the row scaling to apply after it or ``None`` when it + already folds into the stencil. + """ # Create appropriate submesh by combining submeshes in primary domain submesh = self.mesh[domain] - if hasattr(submesh, "length"): - d_nodes = self._get_d_nodes_symbolic_mesh(domain) - e = 1 / d_nodes + symbolic = hasattr(submesh, "length") + if symbolic: + e = 1 / self._get_d_nodes_symbolic_mesh(domain) else: e = 1 / submesh.d_nodes @@ -226,7 +266,11 @@ def gradient_matrix(self, domain, domains): diags([-1.0], [0], shape=(n - 1, n), dtype=None) ) sub_matrix_plus = pybamm.Matrix(diags([1.0], [1], shape=(n - 1, n), dtype=None)) - sub_matrix = (sub_matrix_minus + sub_matrix_plus) * e + sub_matrix = sub_matrix_minus + sub_matrix_plus + # On a symbolic mesh the scaling is left for the caller to apply to the + # result: folded in here it would make the Kronecker product symbolic. + if not symbolic: + sub_matrix = sub_matrix * e # number of repeats second_dim_repeats = self._get_auxiliary_domain_repeats(domains) @@ -239,7 +283,8 @@ def gradient_matrix(self, domain, domains): matrix = pybamm.kronecker_product( pybamm.Matrix(eye(second_dim_repeats, dtype=np.float64)), sub_matrix ) - return matrix + row_scale = self._repeat_vector(e, second_dim_repeats) if symbolic else None + return matrix, row_scale def divergence(self, symbol, discretised_symbol, boundary_conditions): """Matrix-vector multiplication to implement the divergence operator. @@ -247,7 +292,7 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): """ submesh = self.mesh[symbol.domain] - divergence_matrix = self.divergence_matrix(symbol.domains) + divergence_matrix, row_scale = self._divergence_parts(symbol.domains) # check coordinate system if submesh.coord_sys in ["cylindrical polar", "spherical polar"]: @@ -258,9 +303,7 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): else: edges = submesh.edges - r_edges = pybamm.kronecker_product( - pybamm.Matrix(np.ones(second_dim_repeats)), edges - ) + r_edges = self._repeat_vector(edges, second_dim_repeats) if submesh.coord_sys == "spherical polar": out = divergence_matrix @ ((r_edges**2) * discretised_symbol) elif submesh.coord_sys == "cylindrical polar": @@ -268,6 +311,9 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): else: out = divergence_matrix @ discretised_symbol + if row_scale is not None: + out = row_scale * out + return out def divergence_matrix(self, domains): @@ -285,8 +331,26 @@ def divergence_matrix(self, domains): :class:`pybamm.Matrix` The (sparse) finite volume divergence matrix for the domain """ + matrix, row_scale = self._divergence_parts(domains) + return matrix if row_scale is None else matrix * row_scale + + def _divergence_parts(self, domains): + """The divergence matrix, split into a constant stencil and any row scaling. + + Parameters + ---------- + domains : dict + The domain and auxiliary domain in which to compute the matrix. + + Returns + ------- + tuple[:class:`pybamm.Matrix`, :class:`pybamm.Symbol` or None] + The stencil, and the row scaling to apply after it or ``None`` when it + already folds into the stencil. + """ # Create appropriate submesh by combining submeshes in domain submesh = self.mesh[domains["primary"]] + symbolic = hasattr(submesh, "length") if hasattr(submesh, "length"): d_edges = self._get_d_edges_symbolic_mesh(domains["primary"]) else: @@ -313,7 +377,11 @@ def divergence_matrix(self, domains): diags([-1.0], [0], shape=(n - 1, n), dtype=None) ) sub_matrix_plus = pybamm.Matrix(diags([1.0], [1], shape=(n - 1, n), dtype=None)) - sub_matrix = (sub_matrix_minus + sub_matrix_plus) * e + sub_matrix = sub_matrix_minus + sub_matrix_plus + # On a symbolic mesh the scaling is left for the caller to apply to the + # result: folded in here it would make the Kronecker product symbolic. + if not symbolic: + sub_matrix = sub_matrix * e # repeat matrix for each node in secondary dimensions second_dim_repeats = self._get_auxiliary_domain_repeats(domains) @@ -321,7 +389,8 @@ def divergence_matrix(self, domains): matrix = pybamm.kronecker_product( pybamm.Matrix(eye(second_dim_repeats, dtype=np.float64)), sub_matrix ) - return matrix + row_scale = self._repeat_vector(e, second_dim_repeats) if symbolic else None + return matrix, row_scale def laplacian(self, symbol, discretised_symbol, boundary_conditions): """ @@ -1347,9 +1416,9 @@ def boundary_value_or_flux(self, symbol, discretised_child, bcs=None): # issue matrix = csr_matrix(kron(eye(repeats, dtype=np.float64), sub_matrix)) - # Return boundary value with domain given by symbol - matrix = pybamm.Matrix(matrix) * multiplicative - boundary_value = matrix @ discretised_child + # Return boundary value with domain given by symbol. The scalar scaling + # goes on the result: in the matrix it would make it symbolic on a symbolic mesh. + boundary_value = multiplicative * (pybamm.Matrix(matrix) @ discretised_child) boundary_value.copy_domains(symbol) additive.copy_domains(symbol) diff --git a/packages/pybamm/tests/__init__.py b/packages/pybamm/tests/__init__.py index 56f64e381e..f99e12bbd0 100644 --- a/packages/pybamm/tests/__init__.py +++ b/packages/pybamm/tests/__init__.py @@ -47,6 +47,7 @@ get_unit_2p1D_mesh_for_testing, get_cylindrical_discretisation_for_testing, get_base_model_with_battery_geometry, + get_broken_input_model, get_required_distribution_deps, get_optional_distribution_deps, get_present_optional_import_deps, diff --git a/packages/pybamm/tests/conftest_rust.py b/packages/pybamm/tests/conftest_rust.py new file mode 100644 index 0000000000..cc42e82630 --- /dev/null +++ b/packages/pybamm/tests/conftest_rust.py @@ -0,0 +1,170 @@ +# tests/conftest_rust.py +"""Pytest fixtures for Rust expression backend testing.""" + +import numpy as np +import pytest + +import pybamm + + +@pytest.fixture +def rust_backend(monkeypatch): + """Make newly constructed models default to the Rust backend.""" + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "rust") + yield + + +def evaluate_with_rust(expr, t=0.0, y=None, y_dot=None, inputs=None): + """Evaluate a PyBaMM expression using the Rust backend. + + Parameters + ---------- + expr : pybamm.Symbol + The expression to evaluate + t : float + Time value + y : array-like, optional + State vector + y_dot : array-like, optional + State vector derivative + inputs : dict, optional + Input parameters as name->value dict + + Returns + ------- + numpy.ndarray or float + The evaluation result + """ + from pybamm.rust import ExprGraph + + graph = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(graph, rust_symbols) + + # Convert inputs dict to ordered list - order matches registration during to_rust() + # We need to extract input parameter names from the expression tree + if inputs is None: + inputs_list = [] + else: + # Get input parameter names in the order they appear in the expression + input_params = [ + node.name + for node in expr.pre_order() + if isinstance(node, pybamm.InputParameter) + ] + # Deduplicate, preserving order. + seen = set() + unique_params = [] + for name in input_params: + if name not in seen: + seen.add(name) + unique_params.append(name) + inputs_list = [inputs.get(name, 0.0) for name in unique_params] + + y_np = ( + np.asarray(y, dtype=np.float64).ravel() + if y is not None + else np.array([], dtype=np.float64) + ) + y_dot_np = ( + np.asarray(y_dot, dtype=np.float64).ravel() + if y_dot is not None + else np.array([], dtype=np.float64) + ) + + # Determine scalar vs array by checking PyBaMM's evaluation result shape + pybamm_result = expr.evaluate(t=t, y=y, y_dot=y_dot, inputs=inputs) + is_scalar = np.isscalar(pybamm_result) or ( + hasattr(pybamm_result, "shape") and pybamm_result.size == 1 + ) + + if is_scalar: + return graph.eval_to_float( + rust_expr, t, y_np.tolist(), y_dot_np.tolist(), inputs_list + ) + else: + return np.array(graph.eval_to_array(rust_expr, t, y_np, y_dot_np, inputs_list)) + + +def evaluate_with_casadi(expr, t=0.0, y=None, y_dot=None, inputs=None): + """Evaluate a PyBaMM expression using the CasADi backend. + + Parameters + ---------- + expr : pybamm.Symbol + The expression to evaluate + t : float + Time value + y : array-like, optional + State vector + y_dot : array-like, optional + State vector derivative + inputs : dict, optional + Input parameters as name->value dict + + Returns + ------- + numpy.ndarray or float + The evaluation result + """ + import casadi + + # Determine sizes + y_size = 0 if y is None else np.asarray(y).size + y_dot_size = 0 if y_dot is None else np.asarray(y_dot).size + + # Create CasADi symbols + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", y_size) if y_size > 0 else casadi.MX.sym("y", 0) + y_dot_sym = ( + casadi.MX.sym("y_dot", y_dot_size) + if y_dot_size > 0 + else casadi.MX.sym("y_dot", 0) + ) + + # Build inputs symbols + if inputs is None: + inputs = {} + inputs_sym = {name: casadi.MX.sym(name) for name in inputs} + + casadi_symbols = {"t": t_sym, "y": y_sym, "y_dot": y_dot_sym, "inputs": inputs_sym} + + # Convert expression + casadi_expr = expr.to_casadi(t_sym, y_sym, y_dot_sym, inputs_sym, casadi_symbols) + + # Build function inputs list + func_inputs = [t_sym, y_sym, y_dot_sym, *inputs_sym.values()] + func = casadi.Function("f", func_inputs, [casadi_expr]) + + y_val = np.array([]) if y is None else np.asarray(y).flatten() + y_dot_val = np.array([]) if y_dot is None else np.asarray(y_dot).flatten() + input_vals = [inputs[name] for name in inputs_sym] + + result = func(t, y_val, y_dot_val, *input_vals) + return np.asarray(result).flatten() + + +@pytest.fixture +def dual_backend_compare(): + """Fixture that compares Rust and CasADi evaluation results. + + Usage: + def test_something(dual_backend_compare): + expr = pybamm.Scalar(1.0) + pybamm.Scalar(2.0) + dual_backend_compare(expr) # asserts results match + """ + + def _compare(expr, t=0.0, y=None, y_dot=None, inputs=None, rtol=1e-10, atol=1e-14): + rust_result = evaluate_with_rust(expr, t, y, y_dot, inputs) + casadi_result = evaluate_with_casadi(expr, t, y, y_dot, inputs) + + np.testing.assert_allclose( + rust_result, + casadi_result, + rtol=rtol, + atol=atol, + err_msg=f"Rust vs CasADi mismatch for {type(expr).__name__}", + ) + return rust_result + + return _compare diff --git a/packages/pybamm/tests/integration/test_convert_to_format_rust.py b/packages/pybamm/tests/integration/test_convert_to_format_rust.py new file mode 100644 index 0000000000..b168c98574 --- /dev/null +++ b/packages/pybamm/tests/integration/test_convert_to_format_rust.py @@ -0,0 +1,188 @@ +import numpy as np +import pytest + +import pybamm + + +def _solve_spm(fmt, solver, t_eval, **solve_kwargs): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = fmt + sim = pybamm.Simulation(model, solver=solver) + return sim.solve(t_eval, **solve_kwargs) + + +class TestIdakluSpmRustParity: + def test_spm_idaklu_rust_vs_casadi(self): + # SPM is a DAE under the voltage-as-a-state default, so it needs IDAKLU; + # t_interp forces both paths to output exactly on t_eval. + t_eval = np.linspace(0, 3600, 50) + sols = { + fmt: _solve_spm( + fmt, + pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-10), + t_eval, + t_interp=t_eval, + ) + for fmt in ("casadi", "rust") + } + np.testing.assert_allclose( + sols["rust"]["Voltage [V]"].entries, + sols["casadi"]["Voltage [V]"].entries, + rtol=1e-6, + ) + + +class TestAlgebraicRustParity: + def test_pure_algebraic_rust_vs_casadi(self): + results = {} + for fmt in ("casadi", "rust"): + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.algebraic = {v: v**2 - 3 * v + 2} + model.initial_conditions = {v: 3.0} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = fmt + sol = pybamm.AlgebraicSolver().solve(model, [0]) + results[fmt] = sol.y[0][0] + assert results["rust"] == pytest.approx(results["casadi"], rel=1e-8) + assert results["rust"] == pytest.approx(2.0, rel=1e-6) + + +class TestIdakluRustParity: + def test_dfn_chen2020_rust_vs_casadi(self): + # t_interp forces IDAKLU to output exactly on t_eval for both paths. + t_eval = np.linspace(0, 3500, 50) + sols = {} + for fmt in ("casadi", "rust"): + model = pybamm.lithium_ion.DFN() + model.events = [] + param = pybamm.ParameterValues("Chen2020") + model.convert_to_format = fmt + sim = pybamm.Simulation( + model, + parameter_values=param, + solver=pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-10), + ) + sols[fmt] = sim.solve(t_eval, t_interp=t_eval) + np.testing.assert_allclose( + sols["rust"]["Voltage [V]"].entries, + sols["casadi"]["Voltage [V]"].entries, + rtol=1e-5, + ) + + +class TestRustNewton: + @staticmethod + def _pure_algebraic_rust(): + model = pybamm.BaseModel() + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + model.algebraic = {v: v**2 - a * v + 2} + model.initial_conditions = {v: 3.0} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = "rust" + return model + + def test_pure_algebraic_nonlinear_solver_parity(self): + results = {} + for fmt in ("casadi", "rust"): + model = pybamm.BaseModel() + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + model.algebraic = {v: v**2 - a * v + 2} + model.initial_conditions = {v: 3.0} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = fmt + sol = pybamm.NonlinearSolver().solve(model, [0], inputs={"a": 3.0}) + results[fmt] = sol.y[0][0] + assert results["rust"] == pytest.approx(results["casadi"], rel=1e-10) + + def test_rust_newton_setup_pickles_to_inert(self): + # The C++ StandaloneNewtonSolver holds a RAW pointer into the CompiledModel, + # so pickling must drop the handle and keepalive and come back falsy. + import pickle + + from pybamm.solvers.nonlinear_solver import _NonlinearSolverSetup + + solver = pybamm.NonlinearSolver() + model = self._pure_algebraic_rust() + setup = solver._set_up_root_solver_rust(model, {"a": 3.0}) + assert bool(setup) is True and setup._keepalive is not None + restored = pickle.loads(pickle.dumps(setup)) + assert isinstance(restored, _NonlinearSolverSetup) + assert restored._setup is None and restored._keepalive is None + assert bool(restored) is False # falsy -> triggers rebuild, no UAF + + def test_rust_newton_solver_resolves_after_pickle(self): + # A pickled-then-unpickled rust NonlinearSolver must rebuild its Newton + # setup and match a fresh solve, proving the raw pointer is never reused. + import pickle + + fresh = pybamm.NonlinearSolver().solve( + self._pure_algebraic_rust(), [0], inputs={"a": 3.0} + ) + solver = pybamm.NonlinearSolver() + solver.solve(self._pure_algebraic_rust(), [0], inputs={"a": 3.0}) + revived = pickle.loads(pickle.dumps(solver)) + sol = revived.solve(self._pure_algebraic_rust(), [0], inputs={"a": 3.0}) + assert sol.y[0][0] == pytest.approx(fresh.y[0][0], rel=1e-10) + + def test_dae_rust_newton_parity_vs_casadi(self): + # len_rhs>0 DAE: exercises the global->local algebraic jacobian column + # remap; t_interp pins both paths to t_eval as their step counts differ. + t_eval = np.linspace(0, 1, 20) + sols = {} + for fmt in ("casadi", "rust"): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + w = pybamm.Variable("w") + a = pybamm.InputParameter("a") + model.rhs = {u: -a * u} + model.algebraic = {w: w**2 - a * w + 2 * u} + model.initial_conditions = {u: 1.0, w: 3.0} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = fmt + solver = pybamm.IDAKLUSolver( + rtol=1e-8, atol=1e-10, root_method="nonlinear_solver" + ) + sols[fmt] = solver.solve(model, t_eval, inputs={"a": 3.0}, t_interp=t_eval) + np.testing.assert_allclose(sols["rust"].y, sols["casadi"].y, rtol=1e-6) + + +class TestRustRootResolution: + def test_casadi_root_method_switches_to_rust_newton(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + solver = pybamm.IDAKLUSolver(root_method="casadi", options={"calc_ic": False}) + solver._check_and_prepare_model_inplace(model) + assert isinstance(solver.root_method, pybamm.NonlinearSolver) + + def test_diffsol_normalises_model_to_rust(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("casadi") + solver = pybamm.DiffsolSolver() + solver.solve(model, np.linspace(0, 1, 10), inputs={"a": 0.5}) + assert model.convert_to_format == "rust" + + def test_dae_consistent_ic_parity_rust_newton_vs_casadi(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + y0 = {} + for fmt in ("casadi", "rust"): + model = _toy_dae(fmt) + model.initial_conditions = { + k: v for k, v in model.initial_conditions.items() + } + solver = pybamm.IDAKLUSolver( + root_method="nonlinear_solver", options={"calc_ic": False} + ) + solver.set_up(model, inputs=[{"a": 0.5}]) + solver._set_consistent_initialization(model, 0.0, [{"a": 0.5}]) + y0[fmt] = np.asarray(model.y0_list[0]).ravel() + np.testing.assert_allclose(y0["rust"], y0["casadi"], rtol=1e-8) diff --git a/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py b/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py new file mode 100644 index 0000000000..f6f39e5840 --- /dev/null +++ b/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py @@ -0,0 +1,170 @@ +"""diffsol under a current profile with many discontinuities. + +A ramped pulse train puts a corner in the current roughly every 60 s. PyBaMM +hands every ``t_eval`` entry to the integrator as a stop time to land on and +restart from, which is how IDAKLU meets those corners; before diffsol did the +same it stepped an extrapolating high-order BDF straight across each one and +paid for it in rejected steps. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import pybamm + +PULSE_AMPLITUDE_A = 5.0 +PULSE_ON_S = 60.0 +PULSE_REST_S = 120.0 +PULSE_RAMP_S = 2.0 +DURATION_S = 1800.0 +SOLVER_TOL = 1e-6 + + +def _pulse_train_breakpoints() -> tuple[np.ndarray, np.ndarray]: + """Corners of a ramped pulse/rest train; finite ramps keep it Lipschitz.""" + times, values = [0.0], [0.0] + start = 0.0 + period = PULSE_ON_S + PULSE_REST_S + while start < DURATION_S: + times.extend( + [ + start + PULSE_RAMP_S, + start + PULSE_ON_S, + start + PULSE_ON_S + PULSE_RAMP_S, + start + period, + ] + ) + values.extend([PULSE_AMPLITUDE_A, PULSE_AMPLITUDE_A, 0.0, 0.0]) + start += period + times_arr = np.asarray(times, dtype=np.float64) + keep = times_arr <= DURATION_S + return times_arr[keep], np.asarray(values, dtype=np.float64)[keep] + + +def _pulse_train_parameter_values(inputs: dict[str, str] | None = None): + parameter_values = pybamm.ParameterValues("Chen2020") + times, values = _pulse_train_breakpoints() + parameter_values["Current function [A]"] = pybamm.Interpolant( + times, values, pybamm.t, interpolator="linear" + ) + for name, input_name in (inputs or {}).items(): + parameter_values[name] = pybamm.InputParameter(input_name) + return parameter_values + + +INFERENCE_INPUTS = { + "Negative particle diffusivity [m2.s-1]": "D_n", + "Positive particle diffusivity [m2.s-1]": "D_p", + "Negative electrode active material volume fraction": "eps_n", + "Positive electrode active material volume fraction": "eps_p", +} +INPUT_VALUES = {"D_n": 3.3e-14, "D_p": 4.0e-15, "eps_n": 0.75, "eps_p": 0.665} + + +def _solve( + solver, + calculate_sensitivities=False, + inputs=None, + options=None, + t_eval=None, + model_factory=pybamm.lithium_ion.DFN, +): + model = model_factory() + model.convert_to_format = "casadi" if solver == "casadi_idaklu" else "rust" + parameter_values = _pulse_train_parameter_values( + INFERENCE_INPUTS if inputs else None + ) + if solver == "diffsol": + instance = pybamm.DiffsolSolver( + rtol=SOLVER_TOL, atol=SOLVER_TOL, options=options + ) + else: + instance = pybamm.IDAKLUSolver(rtol=SOLVER_TOL, atol=SOLVER_TOL) + breakpoints, _ = _pulse_train_breakpoints() + simulation = pybamm.Simulation( + model, parameter_values=parameter_values, solver=instance + ) + return simulation.solve( + breakpoints if t_eval is None else t_eval, + t_interp=np.linspace(0.0, DURATION_S, 100), + initial_soc=0.5, + inputs=inputs, + calculate_sensitivities=calculate_sensitivities, + ) + + +class TestDiffsolPulseTrain: + def test_the_workload_still_costs_nonlinear_solver_failures(self): + # Guards the guard below: an option capping failures can only be shown + # to bite on a workload that spends some. + solution = _solve("diffsol") + assert solution.solver_statistics.number_of_nonlinear_solver_fails > 5 + + def test_a_long_pulse_train_completes(self): + solution = _solve("diffsol") + assert solution.t[-1] == pytest.approx(DURATION_S) + assert np.all(np.isfinite(np.asarray(solution["Voltage [V]"](solution.t)))) + + def test_the_failure_budget_option_reaches_diffsol(self): + # A budget below what this workload spends must stop the solve, which is + # what shows `options` is plumbed through rather than silently dropped. + with pytest.raises(pybamm.SolverError, match=r"nonlinear solver failures"): + _solve("diffsol", options={"max_nonlinear_solver_failures": 5}) + + def test_values_match_idaklu(self): + native = _solve("diffsol") + reference = _solve("casadi_idaklu") + t = np.linspace(0.0, DURATION_S, 100) + np.testing.assert_allclose( + np.asarray(native["Voltage [V]"](t)), + np.asarray(reference["Voltage [V]"](t)), + rtol=1e-4, + atol=1e-5, + ) + + def test_t_eval_breakpoints_are_stop_times_not_just_output_times(self): + # Equal cost either way would mean t_eval never reached the step control. + given = _solve("diffsol") + hidden = _solve("diffsol", t_eval=np.array([0.0, DURATION_S])) + assert ( + given.solver_statistics.number_of_nonlinear_solver_fails + < hidden.solver_statistics.number_of_nonlinear_solver_fails + ) + assert ( + given.solver_statistics.number_of_error_test_failures + < hidden.solver_statistics.number_of_error_test_failures + ) + + def test_gradients_hold_error_control_across_the_corners(self): + # Falling back re-solves the whole trajectory, so the corners have to be + # met with consistent sensitivity derivatives, not just consistent states. + solution = _solve( + "diffsol", calculate_sensitivities=sorted(INPUT_VALUES), inputs=INPUT_VALUES + ) + assert not solution.solver_statistics.sens_error_control_relaxed + + def test_a_pure_ode_model_meets_the_corners_too(self): + # SPM has no algebraic state, and diffsol's `set_consistent` returns + # early there without refreshing `dy`; every other test here is a DAE. + native = _solve("diffsol", model_factory=pybamm.lithium_ion.SPM) + reference = _solve("casadi_idaklu", model_factory=pybamm.lithium_ion.SPM) + t = np.linspace(0.0, DURATION_S, 100) + np.testing.assert_allclose( + np.asarray(native["Voltage [V]"](t)), + np.asarray(reference["Voltage [V]"](t)), + rtol=1e-4, + atol=1e-5, + ) + + def test_a_long_pulse_train_completes_with_sensitivities(self): + inputs = INPUT_VALUES + names = sorted(inputs) + solution = _solve("diffsol", calculate_sensitivities=names, inputs=inputs) + + assert solution.t[-1] == pytest.approx(DURATION_S) + for name in names: + gradient = np.asarray(solution["Voltage [V]"].sensitivities[name]).ravel() + assert np.all(np.isfinite(gradient)) + assert np.abs(gradient).max() > 0.0 diff --git a/packages/pybamm/tests/integration/test_diffsol_spm.py b/packages/pybamm/tests/integration/test_diffsol_spm.py new file mode 100644 index 0000000000..41d0c8cb87 --- /dev/null +++ b/packages/pybamm/tests/integration/test_diffsol_spm.py @@ -0,0 +1,312 @@ +"""Integration test: SPM solved via Rust diffsol BDF solver. + +End-to-end test for the DiffsolSolver. Builds an SPM model, +solves it via the Rust diffsol BDF backend, and compares +against the CasADi solver for parity. +""" + +import numpy as np +import pytest + +import pybamm + + +class TestDiffsolSPM: + """Integration tests for DiffsolSolver with SPM model.""" + + def test_spm_basic_solve(self): + """Verify DiffsolSolver can solve a basic SPM discharge.""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 3600, 100) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + assert len(solution.t) > 0 + + voltage = solution["Voltage [V]"] + assert voltage is not None + # SPM voltage should be in a reasonable range + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_spm_voltage_trajectory(self): + """Verify voltage trajectory is physically reasonable.""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 1800, 100) + sol = sim.solve(t_eval) + + voltage = sol["Voltage [V]"].entries + assert voltage[0] > 3.5, "Initial voltage should be > 3.5V" + assert voltage[-1] > 3.0, "Final voltage at 1800s should be > 3V" + assert voltage[-1] < voltage[0], "Voltage should decrease during discharge" + + def test_diffsol_matches_casadi(self): + """Verify DiffsolSolver produces results close to CasadiSolver. + + SPM carries an algebraic ``Voltage [V]`` equation (the voltage-as-a-state + default), so this is a DAE. The diffsol and CasADi BDF implementations + agree to ~1e-7 on the DAE path, so the parity tolerance is looser than + the solver tolerance (1e-8) while still well below the 1e-4 used by the + other cross-implementation parity tests in this module. + """ + model_ds = pybamm.lithium_ion.SPM() + model_cs = pybamm.lithium_ion.SPM() + model_cs.convert_to_format = "casadi" + + t_eval = np.linspace(0, 3600, 100) + + sol_ds = pybamm.Simulation( + model_ds, solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + ).solve(t_eval) + sol_cs = pybamm.Simulation( + model_cs, solver=pybamm.CasadiSolver(rtol=1e-8, atol=1e-8) + ).solve(t_eval) + + v_ds = sol_ds["Voltage [V]"].entries + v_cs = sol_cs["Voltage [V]"].entries + + np.testing.assert_allclose( + v_ds, + v_cs, + rtol=1e-6, + atol=1e-7, + err_msg="DiffsolSolver voltage differs from CasadiSolver", + ) + + def test_diffsol_explicit_nonlinear_solver_no_crash(self): + """Verify DiffsolSolver(root_method='nonlinear_solver') doesn't crash on construction.""" + solver = pybamm.DiffsolSolver(root_method="nonlinear_solver") + assert solver.root_method is not None + assert solver.root_tol == 1e-6 + + def test_diffsol_calc_ic_parameter_exists(self): + """Verify calc_ic parameter is accepted and creates _internal_initialisation property.""" + solver = pybamm.DiffsolSolver(calc_ic=True) + assert solver._internal_initialisation is True + + solver_default = pybamm.DiffsolSolver() + assert solver_default._internal_initialisation is False + + def test_diffsol_dfn_default_solves(self): + """Verify DiffsolSolver() solves DFN (DAE) model without explicit root_method.""" + model = pybamm.lithium_ion.DFN() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + voltage = solution["Voltage [V]"] + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_diffsol_explicit_casadi_dfn(self): + """Verify DiffsolSolver(root_method='casadi') works for DFN.""" + model = pybamm.lithium_ion.DFN() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6, root_method="casadi") + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + def test_diffsol_calc_ic_native_spm(self): + """Verify DiffsolSolver(calc_ic=True) works on SPM (DAE via voltage state).""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6, calc_ic=True) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + def test_diffsol_spm_unchanged_after_dae_fix(self): + """Regression test: SPM should still work after DAE IC fix.""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver() + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 3600, 100) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + voltage = solution["Voltage [V]"] + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_diffsol_output_variables(self): + """DiffsolSolver with output_variables should request output rows.""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(output_variables=["Voltage [V]"]) + sim = pybamm.Simulation(model, solver=solver) + + sol = sim.solve([0, 3600]) + + # Should be able to access the requested variable + voltage = sol["Voltage [V]"].data + assert len(voltage) > 0 + assert voltage[0] > 3.0 # Reasonable voltage range + assert voltage[0] < 5.0 + + def test_diffsol_output_variables_with_event(self): + """Event detection should work with output-only solve path.""" + model = pybamm.lithium_ion.SPM() + # Add a voltage cutoff event + model.events.append( + pybamm.Event( + "Voltage cutoff", + model.variables["Voltage [V]"] - 3.2, + pybamm.EventType.TERMINATION, + ) + ) + + solver = pybamm.DiffsolSolver(output_variables=["Voltage [V]"]) + sim = pybamm.Simulation(model, solver=solver) + + sol = sim.solve([0, 7200]) # Long enough to hit cutoff + + # Should terminate early due to event + assert sol.termination.startswith("event") + assert sol.t[-1] < 7200 + # Voltage at end should be near cutoff + voltage_end = sol["Voltage [V]"].data[-1] + assert abs(voltage_end - 3.2) < 0.01 + + +class TestDiffsolCubicOCP: + """SPM with cubic data-interpolant OCPs solves on the diffsol path.""" + + @staticmethod + def _cubic_ocp_parameter_values(): + # Sample Chen2020's closed-form OCPs into data and rebuild them as CUBIC + # interpolants, so the model carries 1D cubic interpolants. + pv = pybamm.ParameterValues("Chen2020") + sto = np.linspace(0.0, 1.0, 200) + un_cb = pv["Negative electrode OCP [V]"] + up_cb = pv["Positive electrode OCP [V]"] + un_data = np.array([float(un_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + up_data = np.array([float(up_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + + def make_cubic_ocp(x_data, y_data, name): + def ocp(stoich): + return pybamm.Interpolant( + x_data, y_data, stoich, name=name, interpolator="cubic" + ) + + return ocp + + pv["Negative electrode OCP [V]"] = make_cubic_ocp(sto, un_data, "Un") + pv["Positive electrode OCP [V]"] = make_cubic_ocp(sto, up_data, "Up") + return pv + + def test_cubic_ocp_diffsol_vs_casadi(self): + model = pybamm.lithium_ion.SPM() + model.events = [] + pv = self._cubic_ocp_parameter_values() + t_eval = np.linspace(0, 1800, 50) + + # Pass t_interp so IDAKLU outputs exactly on t_eval (not internal steps). + ref_model = model.new_copy() + ref_model.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + ref_model, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval, t_interp=t_eval) + + sol_diffsol = pybamm.Simulation( + model.new_copy(), + parameter_values=pv, + solver=pybamm.DiffsolSolver(), + ).solve(t_eval) + + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-4, + atol=1e-6, + err_msg="Diffsol cubic-OCP voltage differs from CasADi", + ) + + +class TestDiffsolECMNDInterpolants: + """Thevenin ECM with 3D/2D data tables solves on the diffsol path.""" + + def test_ecm_diffsol_vs_casadi(self): + model = pybamm.equivalent_circuit.Thevenin() + model.events = [] + pv = pybamm.ParameterValues("ECM_Example") + t_eval = np.linspace(0, 600, 50) + + # Pass t_interp so IDAKLU outputs exactly on t_eval (not internal steps). + ref_model = model.new_copy() + ref_model.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + ref_model, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval, t_interp=t_eval) + + sol_diffsol = pybamm.Simulation( + model.new_copy(), + parameter_values=pv, + solver=pybamm.DiffsolSolver(), + ).solve(t_eval) + + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-4, + atol=1e-6, + err_msg="Diffsol ECM ND-interpolant voltage differs from CasADi", + ) + + +class TestDiffsolExperiment: + """Experiments re-build the rust model per step and stitch segment + solutions; diffsol had no coverage of either.""" + + def test_two_step_experiment_matches_idaklu(self): + experiment = pybamm.Experiment( + [ + "Discharge at 1C for 5 minutes", + "Rest for 5 minutes", + ] + ) + + def run(solver): + sim = pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment, solver=solver + ) + return sim.solve() + + sol_diffsol = run(pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8)) + sol_idaklu = run(pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8)) + + # A flat instruction list makes one cycle per instruction. + assert len(sol_diffsol.cycles) == 2 + assert sol_diffsol.t[-1] == pytest.approx(600.0) + # Compare on diffsol's own grid: its observation is grid-aligned, so + # off-grid points would measure interpolation error, not solver error. + t_common = sol_diffsol.t + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"](t_common), + sol_idaklu["Voltage [V]"](t_common), + rtol=1e-5, + atol=1e-6, + ) diff --git a/packages/pybamm/tests/integration/test_models/standard_output_tests.py b/packages/pybamm/tests/integration/test_models/standard_output_tests.py index 5eb4680b7f..cb48aca88c 100644 --- a/packages/pybamm/tests/integration/test_models/standard_output_tests.py +++ b/packages/pybamm/tests/integration/test_models/standard_output_tests.py @@ -186,8 +186,10 @@ def test_overpotentials(self): elif self.operating_condition == "off": np.testing.assert_allclose(self.eta_r_av(self.t), 0, rtol=1e-7, atol=1e-6) np.testing.assert_allclose(self.eta_e_av(self.t), 0, rtol=1e-12, atol=1e-11) + # Analytically zero; measures round-off of ~4 V potentials, whose + # noise floor varies by backend and platform (observed max ~2.2e-14). np.testing.assert_allclose( - self.delta_phi_s_av(self.t), 0, atol=2e-14, rtol=1e-16 + self.delta_phi_s_av(self.t), 0, atol=5e-14, rtol=1e-16 ) def test_ocps(self): diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py index 1d90f93792..4639bebb6e 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_asymptotics_convergence.py @@ -15,6 +15,8 @@ def test_leading_order_convergence(self): # Create models leading_order_model = pybamm.lead_acid.LOQS() full_model = pybamm.lead_acid.Full() + for model in (leading_order_model, full_model): + model.convert_to_format = "casadi" # Same parameters, same geometry parameter_values = full_model.default_parameter_values diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py index 21858a3d4f..3cea18cd03 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_basic_models.py @@ -10,6 +10,8 @@ class TestCompareBasicModels: def test_compare_full(self): basic_full = pybamm.lead_acid.BasicFull() full = pybamm.lead_acid.Full() + for model in (basic_full, full): + model.convert_to_format = "casadi" parameter_values = pybamm.ParameterValues("Sulzer2019") parameter_values["Current function [A]"] = 10 diff --git a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py index 82c43b8c60..9d6655ed14 100644 --- a/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py +++ b/packages/pybamm/tests/integration/test_models/test_full_battery_models/test_lead_acid/test_compare_outputs.py @@ -19,6 +19,7 @@ def test_compare_averages_asymptotics(self): param = models[0].default_parameter_values param.update({"Current function [A]": 1}) for model in models: + model.convert_to_format = "casadi" param.process_model(model) # set mesh @@ -61,6 +62,7 @@ def test_compare_outputs_surface_form(self): param = models[0].default_parameter_values param.update({"Current function [A]": 1}) for model in models: + model.convert_to_format = "casadi" param.process_model(model) # set mesh diff --git a/packages/pybamm/tests/integration/test_rust_adjoint_dense_rows.py b/packages/pybamm/tests/integration/test_rust_adjoint_dense_rows.py new file mode 100644 index 0000000000..4c656199e6 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_adjoint_dense_rows.py @@ -0,0 +1,56 @@ +"""End-to-end coverage for Rust adjoint dense-row assembly.""" + +import numpy as np + +import pybamm + + +class TestRustAdjointDenseRows: + def test_spme_uses_one_adjoint_sweep_for_voltage_row(self): + # voltage-as-a-state gives SPMe the dense voltage row this test targets; + # the option's default reverted to "false" on main (#5670). + model = pybamm.lithium_ion.SPMe(options={"voltage as a state": "true"}) + model.convert_to_format = "rust" + simulation = pybamm.Simulation(model) + solution = simulation.solve(np.linspace(0.0, 600.0, 20)) + + stats = simulation.solver._setup["rust_model"].jacobian_stats() + assert stats["n_dense_rows"] == 1 + assert stats["dense_row_tape_instructions"] > 0 + assert stats["dense_row_entries"] > stats["n_colors"] + + t = np.linspace(0.0, 600.0, 25) + np.testing.assert_allclose( + solution["Voltage [V]"](t), + solution["Voltage expression [V]"](t), + rtol=1e-5, + atol=1e-5, + ) + + def test_pouch_cv_hold_splits_its_outlier_row(self): + # Many current-collector rows sit just over the dense-row threshold, so + # only the far wider CV constraint row may drive the split. + model = pybamm.lithium_ion.DFN( + {"current collector": "potential pair", "dimensionality": 2} + ) + model.convert_to_format = "rust" + simulation = pybamm.Simulation( + model, + var_pts={"x_n": 4, "x_s": 4, "x_p": 4, "r_n": 4, "r_p": 4, "y": 8, "z": 8}, + experiment=pybamm.Experiment(["Hold at 4.1 V until C/20"]), + ) + simulation.build_for_experiment() + built = next( + m + for key, m in simulation.steps_to_built_models.items() + if key.startswith("Voltage") + ) + solver = pybamm.IDAKLUSolver() + solver.set_up(built, inputs={}, t_eval=np.array([0.0, 1.0])) + + stats = solver._setup["rust_model"].jacobian_stats() + assert stats["n_dense_rows"] == 1 + assert stats["dense_row_tape_instructions"] > 0 + # unsplit, the constraint row alone would force a colour per column + assert stats["n_colors"] < stats["dense_row_entries"] // 10 + assert stats["jac_lane_width"] > 1 diff --git a/packages/pybamm/tests/integration/test_rust_function_api_parity.py b/packages/pybamm/tests/integration/test_rust_function_api_parity.py new file mode 100644 index 0000000000..2d9f8f8f27 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_function_api_parity.py @@ -0,0 +1,178 @@ +"""CasADi parity for the prep-artifact API. + +eval / jacobian(wrt="y"/"p") / jvp / eval_trajectory vs casadi.Function +equivalents, on a toy expression and a discretised SPM rhs. +""" + +import numpy as np +import pytest + +casadi = pytest.importorskip("casadi") + +import pybamm +from pybamm.rust import ExprGraph + + +def _toy_expr(): + """f(t, y, p): 2 states, 2 inputs, smooth, t-dependent.""" + y0 = pybamm.StateVector(slice(0, 1)) + y1 = pybamm.StateVector(slice(1, 2)) + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.NumpyConcatenation( + a * y0 * y1 + pybamm.t, pybamm.sin(y0) * b + pybamm.exp(-y1) + ) + return expr, 2, ["a", "b"] + + +def _casadi_artifacts(expr, n_states, input_names): + """casadi Function quartet: f, df/dy, df/dp_stacked, jvp_y.""" + t = casadi.MX.sym("t") + y = casadi.MX.sym("y", n_states) + y_dot = casadi.MX.sym("y_dot", n_states) + p_syms = {name: casadi.MX.sym(name) for name in input_names} + casadi_symbols = {"t": t, "y": y, "y_dot": y_dot, "inputs": p_syms} + cexpr = expr.to_casadi(t, y, y_dot, p_syms, casadi_symbols) + p_stacked = casadi.vertcat(*p_syms.values()) + v = casadi.MX.sym("v", n_states) + return ( + casadi.Function("f", [t, y, p_stacked], [cexpr]), + casadi.Function("jy", [t, y, p_stacked], [casadi.jacobian(cexpr, y)]), + casadi.Function("jp", [t, y, p_stacked], [casadi.jacobian(cexpr, p_stacked)]), + casadi.Function("jvp", [t, y, p_stacked, v], [casadi.jtimes(cexpr, y, v)]), + ) + + +class TestToyParity: + def setup_method(self): + expr, self.n, names = _toy_expr() + self.cf, self.cjy, self.cjp, self.cjvp = _casadi_artifacts(expr, self.n, names) + g = ExprGraph() + self.f = g.compile(expr.to_rust(g, {}), name="toy", n_states=self.n) + # registration order == pre-order appearance == p_stacked order + assert self.f.input_names == tuple(names) + self.t = 0.7 + self.y = np.array([0.3, 1.2]) + self.p = np.array([2.5, -0.8]) + + def test_eval(self): + np.testing.assert_allclose( + self.f(self.t, self.y, self.p), + np.asarray(self.cf(self.t, self.y, self.p)).ravel(), + rtol=1e-12, + atol=1e-14, + ) + + def test_jacobian_wrt_y(self): + np.testing.assert_allclose( + self.f.jacobian()(self.t, self.y, self.p).toarray(), + np.asarray(self.cjy(self.t, self.y, self.p)), + rtol=1e-12, + atol=1e-14, + ) + + def test_jacobian_wrt_p(self): + # the spec-named check: jacobian(wrt="p") vs casadi.jacobian(expr, p_stacked) + np.testing.assert_allclose( + self.f.jacobian(wrt="p")(self.t, self.y, self.p).toarray(), + np.asarray(self.cjp(self.t, self.y, self.p)), + rtol=1e-12, + atol=1e-14, + ) + + def test_jvp(self): + v = np.array([0.6, -1.1]) + np.testing.assert_allclose( + self.f.jvp(self.t, self.y, self.p, v), + np.asarray(self.cjvp(self.t, self.y, self.p, v)).ravel(), + rtol=1e-12, + atol=1e-14, + ) + + def test_eval_trajectory(self): + n_t = 40 + ts = np.linspace(0.0, 2.0, n_t) + Y = np.vstack([np.linspace(0.1, 1.0, n_t), np.linspace(-0.5, 1.5, n_t)]) + out = self.f.eval_trajectory(ts, Y, self.p) + ref = np.column_stack( + [ + np.asarray(self.cf(tj, Y[:, j], self.p)).ravel() + for j, tj in enumerate(ts) + ] + ) + np.testing.assert_allclose(out, ref, rtol=1e-12, atol=1e-14) + + +class TestSPMParity: + @pytest.fixture(scope="class") + def spm(self): + sim = pybamm.Simulation(pybamm.lithium_ion.SPM()) + sim.build() + return sim.built_model + + @staticmethod + def _y0(spm): + # concatenated_initial_conditions is populated by build(); y0_list only + # exists after a solver set_up runs. + return np.asarray( + spm.concatenated_initial_conditions.evaluate(), dtype=np.float64 + ).ravel() + + @staticmethod + def _casadi_fn(expr, n): + t = casadi.MX.sym("t") + y = casadi.MX.sym("y", n) + y_dot = casadi.MX.sym("y_dot", n) + symbols = {"t": t, "y": y, "y_dot": y_dot, "inputs": {}} + cexpr = expr.to_casadi(t, y, y_dot, {}, symbols) + return ( + casadi.Function("f", [t, y], [cexpr]), + casadi.Function("j", [t, y], [casadi.jacobian(cexpr, y)]), + ) + + def test_rhs_eval_and_jacobians(self, spm): + rhs = spm.concatenated_rhs + y0 = self._y0(spm) + # rhs.size counts only the differential equations, but the solver evaluates + # rhs against the full state vector, so size the input to y0. + n = y0.shape[0] + g = ExprGraph() + f = g.compile(rhs.to_rust(g, {}), name="SPM_rhs", n_states=n) + cf, cj = self._casadi_fn(rhs, n) + p = np.array([]) + np.testing.assert_allclose( + f(0.0, y0, p), np.asarray(cf(0.0, y0)).ravel(), rtol=1e-9, atol=1e-12 + ) + np.testing.assert_allclose( + f.jacobian()(0.0, y0, p).toarray(), + np.asarray(cj(0.0, y0)), + rtol=1e-9, + atol=1e-12, + ) + + def test_output_group_trajectory(self, spm): + rust_symbols = {} + g = ExprGraph() + names = ["Voltage [V]", "Current [A]"] + # get_processed_variable_or_event returns the discretised expression the + # solver compiles; raw variables_and_events hold unlowerable nodes. + exprs = {name: spm.get_processed_variable_or_event(name) for name in names} + y0 = self._y0(spm) + # outputs reference the algebraic voltage state, so size to the full DAE + # state vector rather than the rhs (differential-only) size. + n = y0.shape[0] + group = g.compile_group( + {name: e.to_rust(g, rust_symbols) for name, e in exprs.items()}, + n_states=n, + ) + n_t = 25 + ts = np.linspace(0.0, 3600.0, n_t) + Y = np.tile(y0[:, None], (1, n_t)) * np.linspace(1.0, 1.05, n_t) + p = np.array([]) + results = group.eval_trajectory(ts, Y, p) + for name, out in zip(names, results, strict=True): + cf, _ = self._casadi_fn(exprs[name], n) + ref = np.column_stack( + [np.asarray(cf(tj, Y[:, j])).ravel() for j, tj in enumerate(ts)] + ) + np.testing.assert_allclose(out, ref, rtol=1e-9, atol=1e-12) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu.py b/packages/pybamm/tests/integration/test_rust_idaklu.py new file mode 100644 index 0000000000..731fa26a22 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu.py @@ -0,0 +1,425 @@ +# tests/integration/test_rust_idaklu.py +"""Integration tests verifying Rust evaluator produces IDAKLU-compatible Jacobians. + +These tests validate that the Rust `CompiledModel` produces correct Jacobian +values that match CasADi. This is the foundation for integrating Rust evaluation +into the IDAKLU solver pipeline. + +The tests focus on: +1. RHS evaluation (f(t, y)) matching between Rust and CasADi +2. Jacobian-vector product (df/dy @ v - cj*M @ v) matching +3. Assembled Jacobian correctness via consistency with JVP +""" + +import numpy as np +import pytest +from scipy.sparse import diags, eye + +import pybamm + +casadi = pytest.importorskip("casadi") + +# Try to import Rust extension +try: + from pybamm.rust import CompiledModel, ExprGraph +except ImportError: + pytest.skip( + "Rust extension not available. Build with: uv sync", + allow_module_level=True, + ) + + +class TestRustIDAKLUJacobianParity: + """Tests verifying Rust Jacobian evaluation matches CasADi for IDAKLU.""" + + @staticmethod + def build_diffusion_reaction_expr(n_states: int): + """Build a realistic diffusion-reaction expression like in battery models. + + f(y) = D * L @ y + k * exp(-y/T) * (1 - y) + + where: + - L is a tridiagonal Laplacian matrix (diffusion) + - k * exp(-y/T) * (1 - y) is a Butler-Volmer-like reaction term + + This mimics the structure of SPM/DFN equations. + """ + sv = pybamm.StateVector(slice(0, n_states)) + t = pybamm.Time() + D = pybamm.InputParameter("diffusivity") # Diffusion coefficient + k = pybamm.InputParameter("rate_constant") # Reaction rate + T = pybamm.InputParameter("temperature") + + # Tridiagonal Laplacian (FD discretization of d^2/dx^2) + diag_data = [ + np.ones(n_states) * -2, + np.ones(n_states - 1), + np.ones(n_states - 1), + ] + laplacian = diags(diag_data, [0, -1, 1], format="csr") + L = pybamm.Matrix(laplacian) + + # Diffusion term: D * L @ y + diffusion = D * (L @ sv) + + # Reaction term: k * exp(-y/T) * (1 - y) * (1 + 0.01*t) + reaction = ( + k + * pybamm.exp(-sv / T) + * (pybamm.Scalar(1.0) - sv) + * (pybamm.Scalar(1.0) + pybamm.Scalar(0.01) * t) + ) + + return diffusion + reaction + + @staticmethod + def build_stiff_ode_expr(n_states: int): + """Build a stiff ODE system like in electrochemistry. + + f(y) = -A @ y + nonlinear_source(y) + + where A has eigenvalues spanning several orders of magnitude. + """ + sv = pybamm.StateVector(slice(0, n_states)) + t = pybamm.Time() + alpha = pybamm.InputParameter("alpha") + beta = pybamm.InputParameter("beta") + + # Stiff matrix: band structure with varying eigenvalues + diag_main = np.linspace(1.0, 1000.0, n_states) # Eigenvalues 1 to 1000 + diag_off = np.ones(n_states - 1) * 0.1 + A_sparse = diags([diag_main, -diag_off, -diag_off], [0, -1, 1], format="csr") + A = pybamm.Matrix(A_sparse) + + # Linear decay + linear = -alpha * (A @ sv) + + # Nonlinear source: beta * tanh(y) * exp(-0.01 * y^2) + nonlinear = ( + beta + * pybamm.tanh(sv) + * pybamm.exp(-pybamm.Scalar(0.01) * sv * sv) + * (pybamm.Scalar(1.0) + pybamm.Scalar(0.001) * t) + ) + + return linear + nonlinear + + @staticmethod + def _build_casadi_functions(expr, n_states: int, inputs: dict): + """Build CasADi functions for RHS and Jacobian-vector product.""" + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", n_states) + y_dot_sym = casadi.MX.sym("y_dot", 0) + inputs_sym = {name: casadi.MX.sym(name) for name in inputs} + + casadi_symbols = { + "t": t_sym, + "y": y_sym, + "y_dot": y_dot_sym, + "inputs": inputs_sym, + } + f_casadi = expr.to_casadi(t_sym, y_sym, y_dot_sym, inputs_sym, casadi_symbols) + + # RHS function + func_inputs = [t_sym, y_sym, *inputs_sym.values()] + f_fn = casadi.Function("f", func_inputs, [f_casadi]) + + # Full Jacobian function + jac_expr = casadi.jacobian(f_casadi, y_sym) + jac_fn = casadi.Function("jacobian_fn", func_inputs, [jac_expr]) + + # Jacobian-vector product: J @ v - cj * M @ v (M = I for this test) + v_sym = casadi.MX.sym("v", n_states) + cj_sym = casadi.MX.sym("cj") + jac_action = casadi.jtimes(f_casadi, y_sym, v_sym) - cj_sym * v_sym + jac_inputs = [t_sym, y_sym, *inputs_sym.values(), cj_sym, v_sym] + jac_action_fn = casadi.Function("jac_action", jac_inputs, [jac_action]) + + return f_fn, jac_fn, jac_action_fn + + @staticmethod + def _build_rust_model(expr, n_states: int): + """Build Rust CompiledModel from PyBaMM expression.""" + graph = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(graph, rust_symbols) + + # Mass matrix (identity for ODE) + mass = eye(n_states, format="csr") + model = CompiledModel.from_expr( + graph, + rust_expr, + mass.data.astype(np.float64), + mass.indptr.astype(np.int64), + mass.indices.astype(np.int64), + ) + return model + + @pytest.mark.parametrize("n_states", [10, 50, 100]) + def test_rhs_parity_diffusion_reaction(self, n_states): + """Test that Rust RHS evaluation matches CasADi for diffusion-reaction.""" + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + f_fn, _, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + t = 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + casadi_result = np.array(f_fn(t, y, *inputs.values())).flatten() + rust_result = np.array(rust_model.rhs(t, y, inputs_arr)) + + np.testing.assert_allclose(rust_result, casadi_result, rtol=1e-10, atol=1e-14) + + @pytest.mark.parametrize("n_states", [10, 50, 100]) + def test_rhs_parity_stiff_ode(self, n_states): + """Test that Rust RHS evaluation matches CasADi for stiff ODE.""" + inputs = {"alpha": 1.0, "beta": 0.5} + expr = self.build_stiff_ode_expr(n_states) + + # Build models + f_fn, _, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + t = 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + casadi_result = np.array(f_fn(t, y, *inputs.values())).flatten() + rust_result = np.array(rust_model.rhs(t, y, inputs_arr)) + + np.testing.assert_allclose(rust_result, casadi_result, rtol=1e-10, atol=1e-14) + + @pytest.mark.parametrize("n_states", [10, 50, 100]) + def test_jac_mul_parity(self, n_states): + """Test that Rust Jacobian-vector product matches CasADi.""" + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + _, _, jac_action_fn = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + v = np.random.randn(n_states) + t = 0.5 + cj = 1.0 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Evaluate: J @ v - cj * M @ v; M = I for ODE tests + casadi_result = np.array(jac_action_fn(t, y, *inputs.values(), cj, v)).flatten() + J = rust_model.jacobian(t, y, inputs_arr) + M = eye(n_states, format="csr") + rust_result = np.array(J @ v - cj * (M @ v)) + + np.testing.assert_allclose(rust_result, casadi_result, rtol=1e-9, atol=1e-12) + + @pytest.mark.parametrize("n_states", [10, 50]) + def test_assembled_jacobian_parity(self, n_states): + """Test that Rust assembled Jacobian matches CasADi full Jacobian.""" + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + _, jac_fn, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + t = 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # CasADi full Jacobian + casadi_jac = np.array(jac_fn(t, y, *inputs.values())) + + # Rust assembled Jacobian (CSC via bundle accessor); model.jacobian = pure df/dy + J = rust_model.jacobian(t, y, inputs_arr) + rust_jac = J.toarray() + + np.testing.assert_allclose(rust_jac, casadi_jac, rtol=1e-9, atol=1e-12) + + @pytest.mark.parametrize("n_states", [10, 50]) + def test_assembled_jacobian_with_cj(self, n_states): + """Test that assembled Jacobian with cj correctly computes df/dy - cj*M.""" + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + _, jac_fn, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + t = 0.5 + cj = 1.5 # Non-trivial cj value + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Expected: df/dy - cj*M = df/dy - cj*I + casadi_jac = np.array(jac_fn(t, y, *inputs.values())) + expected_jac = casadi_jac - cj * np.eye(n_states) + + # Rust Jacobian (CSC) minus cj*M; M = I for ODE tests + J = rust_model.jacobian(t, y, inputs_arr) + M = eye(n_states, format="csr") + rust_jac = (J - cj * M).toarray() + + np.testing.assert_allclose(rust_jac, expected_jac, rtol=1e-9, atol=1e-12) + + @pytest.mark.parametrize("n_states", [10, 50]) + def test_jacobian_jvp_consistency(self, n_states): + """Cross-check two independent Rust kernels on J @ v. + + The matrix-free tangent JVP (rhs view, one tangent sweep with the + caller's seed) against the colored-assembly Jacobian (jacobian + view, n_colors sweeps + scatter). A coloring/scatter bug that + affects only one path fails this; J itself is pinned against + casadi in test_assembled_jacobian_parity. + """ + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + v = np.random.randn(n_states) + t = 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Kernel 1: matrix-free tangent JVP via the rhs view + jvp_result = rust_model.rhs.jvp(t, y, inputs_arr, v) + + # Kernel 2: colored assembly via the jacobian view, then matmul + J = rust_model.jacobian(t, y, inputs_arr) + assembled_result = np.asarray(J @ v) + + np.testing.assert_allclose(jvp_result, assembled_result, rtol=1e-9, atol=1e-12) + + def test_sparsity_pattern_structure(self): + """Test that Rust model correctly reports sparsity pattern.""" + n_states = 20 + expr = self.build_diffusion_reaction_expr(n_states) + rust_model = self._build_rust_model(expr, n_states) + + # Get sparsity pattern + indptr, indices = rust_model.sparsity_pattern() + indptr = np.array(indptr) + indices = np.array(indices) + + # Basic structural checks + assert len(indptr) == n_states + 1, "indptr should have n_states + 1 elements" + assert indptr[0] == 0, "indptr should start at 0" + assert indptr[-1] == len(indices), "indptr[-1] should equal nnz" + + # All indices should be valid column indices + assert np.all(indices >= 0), "all indices should be non-negative" + assert np.all(indices < n_states), "all indices should be < n_states" + + def test_residual_evaluation(self): + """Test DAE residual evaluation r = M*y' - f(t,y).""" + n_states = 20 + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + f_fn, _, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + # Test points + y = np.random.randn(n_states) * 0.1 + 0.5 + yp = np.random.randn(n_states) * 0.01 # y' (time derivative) + t = 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Expected residual: M*y' - f = I*y' - f = y' - f (for M = I) + f_val = np.array(f_fn(t, y, *inputs.values())).flatten() + expected_residual = yp - f_val + + rust_residual = np.array(rust_model.eval_residual(t, y, yp, inputs_arr)) + + np.testing.assert_allclose( + rust_residual, expected_residual, rtol=1e-10, atol=1e-14 + ) + + @pytest.mark.parametrize("n_states", [10, 50, 100]) + def test_into_methods_match_allocating(self, n_states): + """Test _into methods match allocating versions (rhs, residual). + + The jacobian block is arithmetic-only: the new bundle jacobian has + no _into variant (assemble_jacobian_csc_into covers the solver path). + """ + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + rust_model = self._build_rust_model(expr, n_states) + + y = np.random.randn(n_states) * 0.1 + 0.5 + v = np.random.randn(n_states) + yp = np.random.randn(n_states) * 0.01 + t = 0.5 + cj = 1.0 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Pre-allocate output buffers + rhs_output = np.zeros(rust_model.output_len) + residual_output = np.zeros(rust_model.n_states) + + # Compare rhs allocating vs eval_into + rhs_alloc = np.array(rust_model.rhs(t, y, inputs_arr)) + rust_model.rhs.eval_into(t, y, inputs_arr, rhs_output) + np.testing.assert_array_equal(rhs_output, rhs_alloc) + + # Jacobian consistency: two equivalent forms of (J - cj*M) @ v must agree + # to within floating-point rounding (different summation order → ULP diffs). + M = eye(n_states, format="csr") + J = rust_model.jacobian(t, y, inputs_arr) + jvp_via_matmul = np.array(J @ v - cj * (M @ v)) + jvp_via_sparse = np.array((J - cj * M) @ v) + np.testing.assert_allclose( + jvp_via_matmul, jvp_via_sparse, rtol=1e-14, atol=1e-14 + ) + + res_alloc = np.array(rust_model.eval_residual(t, y, yp, inputs_arr)) + rust_model.eval_residual_into(t, y, yp, inputs_arr, residual_output) + np.testing.assert_array_equal(residual_output, res_alloc) + + def test_varying_time_values(self): + """Test that time-dependent expressions evaluate correctly at multiple times.""" + n_states = 20 + inputs = {"diffusivity": 1e-5, "rate_constant": 1.0, "temperature": 298.15} + expr = self.build_diffusion_reaction_expr(n_states) + + # Build models + f_fn, _, _ = self._build_casadi_functions(expr, n_states, inputs) + rust_model = self._build_rust_model(expr, n_states) + + y = np.random.randn(n_states) * 0.1 + 0.5 + inputs_arr = np.array([inputs[name] for name in inputs]) + + # Test at multiple time values + for t in [0.0, 0.1, 0.5, 1.0, 10.0, 100.0]: + casadi_result = np.array(f_fn(t, y, *inputs.values())).flatten() + rust_result = np.array(rust_model.rhs(t, y, inputs_arr)) + np.testing.assert_allclose( + rust_result, + casadi_result, + rtol=1e-10, + atol=1e-14, + err_msg=f"Mismatch at t={t}", + ) + + def test_model_properties(self): + """Test that model properties are correctly reported.""" + n_states = 25 + expr = self.build_diffusion_reaction_expr(n_states) + rust_model = self._build_rust_model(expr, n_states) + + assert rust_model.n_states == n_states + assert rust_model.output_len == n_states + # n_colors should be positive and <= n_states + assert rust_model.n_colors > 0 + assert rust_model.n_colors <= n_states diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_all_features.py b/packages/pybamm/tests/integration/test_rust_idaklu_all_features.py new file mode 100644 index 0000000000..acf64b0074 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_all_features.py @@ -0,0 +1,86 @@ +import numpy as np +import pytest + +import pybamm + +pytest.importorskip("casadi") + + +def _build_spm_with_input_current(): + """SPM with `Current function [A]` parameterized as InputParameter `I`. + + Removes events so the ODE-only Rust path applies; `Discharge capacity [A.h]` + is used as the output variable because it lowers cleanly to Rust (unlike + `Voltage [V]`, which routes through `RegPower` not yet supported by the + Rust converter). + """ + model = pybamm.lithium_ion.SPM() + model.events = [] + return model + + +def _chen_with_input_current(): + param = pybamm.ParameterValues("Chen2020") + param["Current function [A]"] = pybamm.InputParameter("I") + return param + + +T_EVAL = np.linspace(0, 100, 15) +INPUTS = {"I": 0.5} + + +def test_spm_inputs_plus_outputs_parity(): + """SPM + InputParameter + output_variables: trajectory and output match CasADi.""" + output_vars = ["Discharge capacity [A.h]", "Current [A]"] + + model_casadi = _build_spm_with_input_current() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=_chen_with_input_current(), + solver=pybamm.IDAKLUSolver(output_variables=output_vars), + ).solve(T_EVAL, inputs=INPUTS) + + model_rust = _build_spm_with_input_current() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=_chen_with_input_current(), + solver=pybamm.IDAKLUSolver(output_variables=output_vars), + ).solve(T_EVAL, inputs=INPUTS) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) + for var in output_vars: + np.testing.assert_allclose( + sol_rust[var].entries, + sol_casadi[var].entries, + rtol=1e-5, + atol=1e-7, + ) + + +def test_spm_inputs_plus_sensitivities_parity(): + """SPM + InputParameter + calculate_sensitivities: sensitivity matches CasADi.""" + model_casadi = _build_spm_with_input_current() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=_chen_with_input_current(), + solver=pybamm.IDAKLUSolver(), + ).solve(T_EVAL, inputs=INPUTS, calculate_sensitivities=["I"]) + + model_rust = _build_spm_with_input_current() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=_chen_with_input_current(), + solver=pybamm.IDAKLUSolver(), + ).solve(T_EVAL, inputs=INPUTS, calculate_sensitivities=["I"]) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) + np.testing.assert_allclose( + sol_rust.sensitivities["I"], + sol_casadi.sensitivities["I"], + rtol=1e-4, + atol=1e-7, + ) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_dae.py b/packages/pybamm/tests/integration/test_rust_idaklu_dae.py new file mode 100644 index 0000000000..0409486321 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_dae.py @@ -0,0 +1,117 @@ +"""DFN-driven Rust IDAKLU DAE parity tests. + +DFN is a real DAE — 962 differential states + 100 algebraic constraints at +default `var_pts` — so each parity assertion exercises the full +ABI/converter/algebraic-IC path on production-shape math. +""" + +import os +import sys + +import numpy as np +import pytest + +pytest.importorskip("casadi") + + +# Reuse the benchmarks/ DFN A/B harness — same builder used by the timing +# benchmarks, so test parity and benchmark fairness stay aligned. +sys.path.insert( + 0, + os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "benchmarks") + ), +) +from dfn_ab_harness import ( + build_dfn_ab, + casadi_jacobian_dense, + sample_state, +) + + +@pytest.fixture(scope="module") +def ab_small(): + """Small DFN (npts=5 -> 92 states) — fast, still mixed RHS/algebraic.""" + return build_dfn_ab(npts=5) + + +def _assemble_rust_jacobian_dense(ab, t, y, cj=0.0): + """Run Rust assembled Jacobian and reconstruct dense (n,n) numpy matrix.""" + jac_data = np.zeros(ab.nnz) + ab.rust_model.assemble_jacobian_csc_into(t, y, cj, ab.inputs_array, jac_data) + Jr = np.zeros((ab.n_states, ab.n_states)) + for col in range(ab.n_states): + for k in range(ab.csc_colptrs[col], ab.csc_colptrs[col + 1]): + Jr[ab.csc_rowinds[k], col] = jac_data[k] + return Jr + + +def test_dfn_set_up_succeeds(ab_small): + """Rust setup must produce a model with consistent metadata for DFN.""" + assert ab_small.n_states > 50, "DFN should yield > 50 states even at npts=5" + ids = np.asarray(ab_small.rust_model.algebraic_ids()) + jac_stats = ab_small.rust_model.jacobian_stats() + assert ids.shape == (ab_small.n_states,) + # Mixed DAE: must have both rhs (1.0) and algebraic (0.0) ids. + assert np.any(ids == 1.0), "Expected at least one differential row" + assert np.any(ids == 0.0), "Expected at least one algebraic row" + assert jac_stats["n_colors"] > 0 + assert jac_stats["nnz"] == ab_small.nnz + assert jac_stats["strategy"] == "coloring" + + +def test_dfn_residual_parity(ab_small): + """Rust rhs and CasADi rhs_algebraic_eval agree on f(t, y, p). + + Both backends consume the *same* discretised model (built once by the + harness), so residual parity verifies the Rust expression converter + produced a graph numerically equivalent to CasADi's symbolic form. + """ + rng = np.random.default_rng(7) + for _ in range(3): + y = sample_state(ab_small, perturb=1e-3, seed=rng.integers(0, 1 << 32)) + rc = np.asarray( + ab_small.casadi_residual(0.0, y, ab_small.inputs_array) + ).reshape(-1) + rr = ab_small.rust_model.rhs(0.0, y, ab_small.inputs_array) + np.testing.assert_allclose(rr, rc, rtol=1e-10, atol=1e-12) + + +def test_dfn_jacobian_alg_rows_parity(ab_small): + """Rust assembled Jacobian matches CasADi exactly on algebraic rows.""" + y = sample_state(ab_small) + Jc = casadi_jacobian_dense(ab_small, 0.0, y) + Jr = _assemble_rust_jacobian_dense(ab_small, 0.0, y) + ids = np.asarray(ab_small.rust_model.algebraic_ids()) + alg_mask = ids == 0.0 + np.testing.assert_allclose(Jr[alg_mask], Jc[alg_mask], rtol=1e-10, atol=1e-12) + + +def test_dfn_jacobian_rhs_rows_parity(ab_small): + """Rust assembled Jacobian matches CasADi exactly on RHS rows.""" + y = sample_state(ab_small) + Jc = casadi_jacobian_dense(ab_small, 0.0, y) + Jr = _assemble_rust_jacobian_dense(ab_small, 0.0, y) + ids = np.asarray(ab_small.rust_model.algebraic_ids()) + rhs_mask = ids == 1.0 + np.testing.assert_allclose(Jr[rhs_mask], Jc[rhs_mask], rtol=1e-10, atol=1e-12) + + +def test_dfn_jacobian_nonzero_cj_merges_mass(ab_small): + """Rust-side cj*mass merge with a non-identity (DAE) mass matrix. + + DFN mass is diag(algebraic_ids) — ones on differential rows, zeros on + algebraic rows — so nonzero cj must shift exactly the differential + diagonal: J(cj) == J(0) - cj * M. This exercises the Rust merged-mass + scatter (mass_to_csc_map) itself, not a scipy reconstruction. + """ + y = sample_state(ab_small) + cj = 1.5 + J0 = _assemble_rust_jacobian_dense(ab_small, 0.0, y) + Jcj = _assemble_rust_jacobian_dense(ab_small, 0.0, y, cj=cj) + + ids = np.asarray(ab_small.rust_model.algebraic_ids()) + # mixed DAE: the merge must bite on differential rows and skip algebraic + assert np.any(ids == 1.0) and np.any(ids == 0.0) + M = np.diag(ids) + np.testing.assert_allclose(Jcj, J0 - cj * M, rtol=1e-12, atol=1e-12) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_dfn.py b/packages/pybamm/tests/integration/test_rust_idaklu_dfn.py new file mode 100644 index 0000000000..750adde490 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_dfn.py @@ -0,0 +1,76 @@ +"""Integration test: DFN solved via Rust-backed IDAKLU. + +DFN exercises the `RegPower` Rust converter (used in OCP / regularised +intercalation flux); without it the model build raises `TypeError` at +`_to_rust`. This test pins parity against the CasADi backend. +""" + +import numpy as np +import pytest + +import pybamm + + +class TestRustIDAKLUDFN: + """Full DFN simulation via Rust IDAKLU backend; parity vs CasADi.""" + + @pytest.fixture + def dfn_model_no_events(self): + model = pybamm.lithium_ion.DFN() + model.events = [] + return model + + @pytest.fixture + def parameter_values(self): + return pybamm.ParameterValues("Chen2020") + + def test_dfn_uses_reg_power(self, dfn_model_no_events, parameter_values): + """Sanity-check: DFN's processed equations contain RegPower. + + If a refactor removes RegPower from DFN, the Rust converter + coverage for it is no longer load-bearing for this model — but + this test fails so we notice and re-evaluate the integration. + """ + param = parameter_values.copy() + model = dfn_model_no_events.new_copy() + param.process_model(model) + + def _has_reg_power(symbol): + if isinstance(symbol, pybamm.expression_tree.functions.RegPower): + return True + return any(_has_reg_power(c) for c in symbol.children) + + equations = list(model.rhs.values()) + list(model.algebraic.values()) + assert any(_has_reg_power(eq) for eq in equations), ( + "DFN no longer contains RegPower; this gate test is stale." + ) + + def test_dfn_rust_vs_casadi_voltage(self, dfn_model_no_events, parameter_values): + """Voltage trajectory matches CasADi within IDA tolerance.""" + t_eval = np.linspace(0, 100, 25) + + model_casadi = dfn_model_no_events.new_copy() + model_casadi.convert_to_format = "casadi" + sim_casadi = pybamm.Simulation( + model_casadi, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ) + sol_casadi = sim_casadi.solve(t_eval) + + model_rust = dfn_model_no_events.new_copy() + model_rust.convert_to_format = "rust" + sim_rust = pybamm.Simulation( + model_rust, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ) + sol_rust = sim_rust.solve(t_eval) + + np.testing.assert_allclose( + sol_rust["Voltage [V]"](t_eval), + sol_casadi["Voltage [V]"](t_eval), + rtol=1e-4, + atol=1e-5, + err_msg="DFN voltage trajectory differs between Rust and CasADi", + ) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_inputs.py b/packages/pybamm/tests/integration/test_rust_idaklu_inputs.py new file mode 100644 index 0000000000..fd29ea9f07 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_inputs.py @@ -0,0 +1,35 @@ +"""Parity test: Rust IDAKLU with InputParameter models vs CasADi.""" + +import numpy as np +import pytest + +import pybamm + +pytest.importorskip("casadi") + + +def _build_input_param_model(): + """ODE: dy/dt = -k * y, with k as InputParameter. Pure ODE, identity mass.""" + model = pybamm.BaseModel() + y = pybamm.Variable("y") + k = pybamm.InputParameter("k") + model.rhs = {y: -k * y} + model.initial_conditions = {y: 1.0} + return model + + +@pytest.mark.parametrize("k_val", [0.1, 1.0, 5.0]) +def test_input_parameter_solve_parity(k_val): + """Solve dy/dt = -k*y with both backends, compare trajectories.""" + t_eval = np.linspace(0, 5, 50) + inputs = {"k": k_val} + + model_casadi = _build_input_param_model() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.IDAKLUSolver().solve(model_casadi, t_eval, inputs=inputs) + + model_rust = _build_input_param_model() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.IDAKLUSolver().solve(model_rust, t_eval, inputs=inputs) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_outputs.py b/packages/pybamm/tests/integration/test_rust_idaklu_outputs.py new file mode 100644 index 0000000000..8bbe175782 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_outputs.py @@ -0,0 +1,83 @@ +"""Parity test: Rust IDAKLU output_variables vs CasADi. + +The Rust expression converter does not yet support every node used in real +battery models, so we exercise output-variable plumbing on a synthetic +ODE that uses only supported operations. The test confirms: + - the Python-side `output_exprs` arg flows through to CompiledModel, + - `model.outputs[i](t, y, p)` produces the correct values, + - configured output expressions don't disturb the rhs solve. +""" + +import numpy as np +import pytest + +import pybamm + +pytest.importorskip("casadi") + +try: + from pybamm.rust import CompiledModel, ExprGraph +except ImportError: + pytest.skip( + "Rust extension not available. Build with: uv sync", + allow_module_level=True, + ) + + +def _build_decay_model_with_outputs(): + """ODE: dy/dt = -k*y, output: 2*y. Uses StateVector to bypass discretization.""" + graph = ExprGraph() + rust_symbols: dict = {} + + # StateVector -> Rust StateVector node directly (no discretization needed). + y = pybamm.StateVector(slice(0, 1)) + k = pybamm.InputParameter("k") + rhs_sym = -k * y + output_sym = 2 * y + + rhs_expr = rhs_sym.to_rust(graph, rust_symbols) + output_expr = output_sym.to_rust(graph, rust_symbols) + + return graph, rhs_expr, output_expr + + +def test_pycompiled_model_eval_output_matches_analytical(): + """Output 2*y at y=3 with k=1 should equal 6, regardless of k.""" + graph, rhs_expr, output_expr = _build_decay_model_with_outputs() + + # 1-state ODE, identity mass, 1 input, 1 output + mass_data = np.ones(1) + mass_indptr = np.array([0, 1], dtype=np.int64) + mass_indices = np.array([0], dtype=np.int64) + + model = CompiledModel.from_expr( + graph, + rhs_expr, + mass_data, + mass_indptr, + mass_indices, + n_inputs=1, + output_exprs=[output_expr], + ) + + assert model.n_outputs == 1 + assert [f.output_len for f in model.outputs] == [1] + + rhs = model.rhs(0.0, np.asarray([3.0]), np.asarray([1.0])) + np.testing.assert_allclose(rhs, [-3.0], atol=1e-12) + + # output 2*y at y=3: 6 + out = model.outputs[0](0.0, np.asarray([3.0]), np.asarray([1.0])) + np.testing.assert_allclose(out, [6.0], atol=1e-12) + + +def test_pycompiled_model_no_outputs_default(): + """`output_exprs` defaults to empty -> n_outputs == 0, output_lens empty.""" + graph = ExprGraph() + y = graph.state_vector(0, 1) + mass_data = np.ones(1) + mass_indptr = np.array([0, 1], dtype=np.int64) + mass_indices = np.array([0], dtype=np.int64) + model = CompiledModel.from_expr(graph, y, mass_data, mass_indptr, mass_indices) + assert model.n_outputs == 0 + assert [f.output_len for f in model.outputs] == [] diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_parity.py b/packages/pybamm/tests/integration/test_rust_idaklu_parity.py new file mode 100644 index 0000000000..cc02edb600 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_parity.py @@ -0,0 +1,151 @@ +"""E2E parity tests: Rust vs CasADi IDAKLU backends. + +These tests verify that the Rust evaluator produces identical results +to CasADi by running both backends at test time and comparing outputs. +""" + +import numpy as np +import pytest +from scipy.sparse import diags, eye + +import pybamm + +casadi = pytest.importorskip("casadi") + +try: + from pybamm.rust import CompiledModel, ExprGraph +except ImportError: + pytest.skip( + "Rust extension not available. Build with: uv sync", + allow_module_level=True, + ) + + +def build_diffusion_reaction_expr(n_states: int): + """Build a diffusion-reaction expression for testing. + + f(y) = D * L @ y + k * exp(-y/T) * (1 - y) + """ + sv = pybamm.StateVector(slice(0, n_states)) + D = pybamm.Scalar(1e-5) + k = pybamm.Scalar(1.0) + T = pybamm.Scalar(298.15) + + diag_data = [ + np.ones(n_states) * -2, + np.ones(n_states - 1), + np.ones(n_states - 1), + ] + laplacian = diags(diag_data, [0, -1, 1], format="csr") + L = pybamm.Matrix(laplacian) + + diffusion = D * (L @ sv) + reaction = k * pybamm.exp(-sv / T) * (pybamm.Scalar(1.0) - sv) + + return diffusion + reaction + + +def evaluate_expr_casadi(expr, n_states, t, y): + """Evaluate expression using CasADi.""" + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", n_states) + y_dot_sym = casadi.MX.sym("y_dot", 0) + + casadi_symbols = {"t": t_sym, "y": y_sym, "y_dot": y_dot_sym, "inputs": {}} + f_casadi = expr.to_casadi(t_sym, y_sym, y_dot_sym, {}, casadi_symbols) + f_fn = casadi.Function("f", [t_sym, y_sym], [f_casadi]) + + return np.array(f_fn(t, y)).flatten() + + +def evaluate_expr_rust(expr, n_states, t, y): + """Evaluate expression using Rust.""" + graph = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(graph, rust_symbols) + + mass = eye(n_states, format="csr") + model = CompiledModel.from_expr( + graph, + rust_expr, + mass.data.astype(np.float64), + mass.indptr.astype(np.int64), + mass.indices.astype(np.int64), + ) + + return np.array(model.rhs(t, y, np.array([]))) + + +class TestRustCasADiRHSParity: + """Test RHS evaluation parity between Rust and CasADi.""" + + @pytest.mark.parametrize("n_states", [10, 50, 100]) + def test_rhs_parity_diffusion_reaction(self, n_states): + """RHS evaluation matches for diffusion-reaction expression.""" + expr = build_diffusion_reaction_expr(n_states) + y = np.random.randn(n_states) * 0.1 + 0.5 + t = 0.0 + + casadi_result = evaluate_expr_casadi(expr, n_states, t, y) + rust_result = evaluate_expr_rust(expr, n_states, t, y) + + np.testing.assert_allclose(rust_result, casadi_result, rtol=1e-10, atol=1e-14) + + +class TestRustCasADiSolveParity: + """Test full solve parity between Rust and CasADi IDAKLU backends.""" + + @pytest.fixture + def spm_model_no_events(self): + """SPM model with events removed.""" + model = pybamm.lithium_ion.SPM() + model.events = [] + return model + + @pytest.fixture + def parameter_values(self): + return pybamm.ParameterValues("Chen2020") + + def test_short_solve_parity(self, spm_model_no_events, parameter_values): + """Short solve: Rust trajectory matches CasADi.""" + t_eval = np.linspace(0, 100, 10) + + model_casadi = spm_model_no_events + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + model_rust = spm_model_no_events.new_copy() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) + + def test_full_discharge_parity(self, spm_model_no_events, parameter_values): + """Full discharge: Rust trajectory matches CasADi.""" + t_eval = np.linspace(0, 3600, 100) + + model_casadi = spm_model_no_events + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + model_rust = spm_model_no_events.new_copy() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_sens.py b/packages/pybamm/tests/integration/test_rust_idaklu_sens.py new file mode 100644 index 0000000000..62d5d341b0 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_sens.py @@ -0,0 +1,174 @@ +"""Parity test: Rust IDAKLU forward sensitivities vs CasADi.""" + +import numpy as np +import pytest + +import pybamm + +pytest.importorskip("casadi") + + +def _build_decay_model(): + """ODE: dy/dt = -k * y, with k as InputParameter (sensitivity target).""" + model = pybamm.BaseModel() + y = pybamm.Variable("y") + k = pybamm.InputParameter("k") + model.rhs = {y: -k * y} + model.initial_conditions = {y: 1.0} + return model + + +def test_sensitivity_parity_decay(): + """Both backends agree on dy/dk for dy/dt = -k*y, y(0)=1. + + Analytical: y(t) = exp(-k*t), so dy/dk = -t * exp(-k*t). + """ + t_eval = np.linspace(0, 5, 30) + inputs = {"k": 0.5} + + model_casadi = _build_decay_model() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.IDAKLUSolver().solve( + model_casadi, + t_eval, + inputs=inputs, + calculate_sensitivities=["k"], + ) + model_rust = _build_decay_model() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.IDAKLUSolver().solve( + model_rust, + t_eval, + inputs=inputs, + calculate_sensitivities=["k"], + ) + + # Trajectory parity + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-5, atol=1e-8) + + # Sensitivity parity + casadi_sens = sol_casadi.sensitivities["k"] + rust_sens = sol_rust.sensitivities["k"] + np.testing.assert_allclose(rust_sens, casadi_sens, rtol=1e-4, atol=1e-7) + + +DECOUPLED_RATES = {"p0": 0.3, "p1": 1.7, "p2": 0.9} + + +def _build_decoupled_decay_model(names): + """One decay per parameter: dy_i/dt = -p_i * y_i, y_i(0) = 1. + + States are decoupled, so dy_i/dp_j is -t*exp(-p_i*t) when i == j and + exactly zero otherwise. + """ + model = pybamm.BaseModel() + states = [pybamm.Variable(f"y{i}") for i in range(len(names))] + model.rhs = { + y: -pybamm.InputParameter(name) * y + for y, name in zip(states, names, strict=True) + } + model.initial_conditions = dict.fromkeys(states, pybamm.Scalar(1.0)) + return model + + +def _build_coupled_model(): + """Every parameter drives every state, so no sensitivity column is sparse.""" + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + c = pybamm.InputParameter("c") + model.rhs = {u: -(a + b + c) * u + b * v, v: -(a + 2 * b + 3 * c) * v + c * u} + model.initial_conditions = {u: 1.0, v: 2.0} + return model + + +def test_multi_param_sensitivity_columns_land_in_their_own_parameter(): + """Each parameter's df/dp column must reach that parameter's own slot. + + The rust core evaluates every parameter column in one call and the C++ + consumer scatters them into SUNDIALS' per-parameter buffers, so a stride + or ordering slip swaps columns between parameters. Decoupled states make + that visible: p_i drives state i alone, so a mis-scatter shows up as + sensitivity mass on an off-diagonal state. + """ + names = list(DECOUPLED_RATES) + n_states = len(names) + t_interp = np.linspace(0, 2, 9) + + model = _build_decoupled_decay_model(names) + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + sol = pybamm.IDAKLUSolver(atol=1e-10, rtol=1e-10).solve( + model, + [0, 2], + inputs=DECOUPLED_RATES, + calculate_sensitivities=names, + t_interp=t_interp, + ) + + # Pin the state order the sensitivity assertions below index against. + for i, name in enumerate(names): + np.testing.assert_allclose( + sol.y[i, :], + np.exp(-DECOUPLED_RATES[name] * sol.t), + rtol=1e-6, + atol=1e-8, + err_msg=f"state {i} is not the decay driven by '{name}'", + ) + + for i, name in enumerate(names): + expected = np.zeros((len(sol.t), n_states)) + expected[:, i] = -sol.t * np.exp(-DECOUPLED_RATES[name] * sol.t) + np.testing.assert_allclose( + np.asarray(sol.sensitivities[name]).reshape(len(sol.t), n_states), + expected, + rtol=1e-5, + atol=1e-8, + err_msg=f"'{name}' sensitivity landed on the wrong state", + ) + + +def test_multi_param_dense_sensitivities_match_casadi(): + """Three-parameter parity where every column is dense. + + Complements the decoupled test: with no structural zeros, a column read at + the wrong stride produces plausible-looking numbers that only a value + comparison catches. + """ + inputs = {"a": 0.4, "b": 0.25, "c": 0.15} + t_interp = np.linspace(0, 3, 13) + + sols = {} + for backend in ("casadi", "rust"): + model = _build_coupled_model() + model.convert_to_format = backend + pybamm.Discretisation().process_model(model) + sols[backend] = pybamm.IDAKLUSolver(atol=1e-10, rtol=1e-10).solve( + model, + [0, 3], + inputs=inputs, + calculate_sensitivities=list(inputs), + t_interp=t_interp, + ) + + np.testing.assert_allclose(sols["rust"].y, sols["casadi"].y, rtol=1e-6, atol=1e-9) + + per_param = {} + for name in inputs: + rust_sens = np.asarray(sols["rust"].sensitivities[name]) + per_param[name] = rust_sens + np.testing.assert_allclose( + rust_sens, + np.asarray(sols["casadi"].sensitivities[name]), + rtol=1e-5, + atol=1e-8, + err_msg=f"sensitivity mismatch for parameter '{name}'", + ) + + # Parity only discriminates if the columns actually differ from each other. + for lhs, rhs in (("a", "b"), ("a", "c"), ("b", "c")): + assert not np.allclose(per_param[lhs], per_param[rhs], rtol=1e-3, atol=1e-6), ( + f"columns '{lhs}' and '{rhs}' are too alike to detect a swap" + ) diff --git a/packages/pybamm/tests/integration/test_rust_idaklu_spm.py b/packages/pybamm/tests/integration/test_rust_idaklu_spm.py new file mode 100644 index 0000000000..35e6be51a3 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_idaklu_spm.py @@ -0,0 +1,205 @@ +"""Integration test: SPM solved via Rust-backed IDAKLU. + +Gate test for Rust core IDAKLU integration. Builds SPM model, +converts to Rust, solves via IDAKLU, compares against CasADi. +""" + +import numpy as np +import pytest + +import pybamm + + +class TestRustIDAKLUSPM: + """Gate test: Full SPM simulation via Rust IDAKLU backend.""" + + @pytest.fixture + def spm_model_no_events(self): + """SPM model with events removed for fixed-time solve.""" + model = pybamm.lithium_ion.SPM() + model.events = [] + return model + + @pytest.fixture + def parameter_values(self): + return pybamm.ParameterValues("Chen2020") + + def test_spm_rust_vs_casadi_solve(self, spm_model_no_events, parameter_values): + """Compare full SPM solve: Rust backend vs CasADi backend.""" + t_eval = np.linspace(0, 3600, 100) + + # CasADi reference + model_casadi = spm_model_no_events + model_casadi.convert_to_format = "casadi" + sim_casadi = pybamm.Simulation( + model_casadi, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ) + sol_casadi = sim_casadi.solve(t_eval) + + # Rust backend + model_rust = model_casadi.new_copy() + model_rust.convert_to_format = "rust" + sim_rust = pybamm.Simulation( + model_rust, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(), + ) + sol_rust = sim_rust.solve(t_eval) + + np.testing.assert_allclose( + sol_rust.y, + sol_casadi.y, + rtol=1e-5, + atol=1e-8, + err_msg="Solution trajectories differ", + ) + + def test_spm_rust_voltage_trajectory(self, spm_model_no_events, parameter_values): + """Verify voltage trajectory is physically reasonable.""" + model = spm_model_no_events + model.convert_to_format = "rust" + # Use 1800s (half discharge) so the battery stays above 2.5V cutoff + t_eval = np.linspace(0, 1800, 100) + + sim = pybamm.Simulation( + model, parameter_values=parameter_values, solver=pybamm.IDAKLUSolver() + ) + sol = sim.solve(t_eval) + voltage = sol["Voltage [V]"].data + + assert voltage[0] > 4.0, "Initial voltage should be > 4V" + assert voltage[-1] > 3.0, "Final voltage at 1800s should be > 3V" + assert voltage[-1] < voltage[0], "Voltage should decrease" + + +class TestRustIDAKLUEvents: + """Verify termination events fire correctly via the Rust IDAKLU backend.""" + + def test_event_termination_matches_casadi(self): + """A voltage cutoff event must terminate the Rust solve at the same + time, on the same event, as the CasADi backend.""" + # High discharge current so the "Minimum voltage [V]" event fires well + # before the final time. + parameter_values = pybamm.ParameterValues("Marquis2019") + parameter_values["Current function [A]"] = 5.0 + t_eval = [0, 3600] + t_interp = np.linspace(0, 3600, 1000) + + def solve(fmt): + model = pybamm.lithium_ion.SPM() # has events by default + model.convert_to_format = fmt + solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8) + sim = pybamm.Simulation( + model, parameter_values=parameter_values, solver=solver + ) + return sim.solve(t_eval=t_eval, t_interp=t_interp) + + sol_casadi = solve("casadi") + sol_rust = solve("rust") + + # The event actually fired (did not run to the final time). + assert sol_rust.t[-1] < 3600.0 + assert str(sol_rust.termination).startswith("event") + # Same event identified, and same termination time as CasADi. + assert sol_rust.termination == sol_casadi.termination + np.testing.assert_allclose(sol_rust.t[-1], sol_casadi.t[-1], rtol=1e-6) + np.testing.assert_allclose( + sol_rust["Voltage [V]"].entries[-1], + sol_casadi["Voltage [V]"].entries[-1], + rtol=1e-6, + atol=1e-8, + ) + + +class TestRustCubicOCP: + """SPM with cubic data-interpolant OCPs solves on the Rust paths.""" + + @staticmethod + def _cubic_ocp_parameter_values(): + # Sample Chen2020's closed-form OCPs into data, rebuild as CUBIC + # interpolants so the model carries 1D cubic interpolants. + pv = pybamm.ParameterValues("Chen2020") + sto = np.linspace(0.0, 1.0, 200) + un_cb = pv["Negative electrode OCP [V]"] + up_cb = pv["Positive electrode OCP [V]"] + un_data = np.array([float(un_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + up_data = np.array([float(up_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + + def make_cubic_ocp(x_data, y_data, name): + def ocp(stoich): + return pybamm.Interpolant( + x_data, y_data, stoich, name=name, interpolator="cubic" + ) + + return ocp + + pv["Negative electrode OCP [V]"] = make_cubic_ocp(sto, un_data, "Un") + pv["Positive electrode OCP [V]"] = make_cubic_ocp(sto, up_data, "Up") + return pv + + def test_cubic_ocp_idaklu_rust_vs_casadi(self): + model = pybamm.lithium_ion.SPM() + model.events = [] + pv = self._cubic_ocp_parameter_values() + t_eval = np.linspace(0, 1800, 50) + + model_casadi = model.new_copy() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + model_rust = model.new_copy() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + np.testing.assert_allclose( + sol_rust["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-5, + atol=1e-6, + err_msg="Rust cubic-OCP voltage differs from CasADi", + ) + + +class TestRustECMNDInterpolants: + """Thevenin ECM with 3D (r0/r1/c1) and 2D (dUdT) data tables solves on + the Rust evaluator path and matches CasADi.""" + + def test_ecm_idaklu_rust_vs_casadi(self): + model = pybamm.equivalent_circuit.Thevenin() + model.events = [] + pv = pybamm.ParameterValues("ECM_Example") + t_eval = np.linspace(0, 600, 50) + + model_casadi = model.new_copy() + model_casadi.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + model_casadi, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + model_rust = model.new_copy() + model_rust.convert_to_format = "rust" + sol_rust = pybamm.Simulation( + model_rust, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval) + + np.testing.assert_allclose( + sol_rust["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-5, + atol=1e-6, + err_msg="Rust ECM ND-interpolant voltage differs from CasADi", + ) diff --git a/packages/pybamm/tests/integration/test_rust_observability_inference_parameters.py b/packages/pybamm/tests/integration/test_rust_observability_inference_parameters.py new file mode 100644 index 0000000000..ac00b76739 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_observability_inference_parameters.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from benchmarks.rust_observability.registry import ( + INFERENCE_INPUTS, + get_inference_scenarios, + get_protocol_names, + inference_nominal_values, +) +from benchmarks.rust_observability.runners import ( + _build_and_time, + _build_simulation, + _make_solver, + _solve_kwargs, +) + +_BACKEND = "casadi_idaklu" + + +def _scenarios(): + """Every model x protocol the inference lane actually runs.""" + return get_inference_scenarios( + ["SPM", "SPMe", "DFN"], get_protocol_names(), output_points=50 + ) + + +class TestInferenceParametersAreLive: + @pytest.mark.parametrize( + "scenario", _scenarios(), ids=lambda s: f"{s.name}-{s.protocol}" + ) + def test_every_fitted_parameter_moves_the_voltage(self, scenario): + """A dead parameter would inflate the count without perturbing the graph.""" + # INFERENCE_COMPLEMENTS makes dV/d(eps) a different quantity, so the + # lane's own builders have to be what is measured. + nominal = inference_nominal_values() + solver = _make_solver(_BACKEND, atol=scenario.atol, rtol=scenario.rtol) + simulation = _build_simulation(scenario, solver, _BACKEND) + _build_and_time(simulation, scenario, inputs=nominal) + names = sorted(INFERENCE_INPUTS.values()) + solution = simulation.solve( + **_solve_kwargs( + scenario, + {"inputs": nominal, "calculate_sensitivities": names}, + ) + ) + + sensitivities = solution["Voltage [V]"].sensitivities + for input_name in names: + magnitude = float(np.max(np.abs(sensitivities[input_name]))) + assert magnitude > 0.0, ( + f"{input_name} has zero sensitivity on {scenario.name} under " + f"{scenario.protocol}; it is not used by this model and would " + "be a dead fitted parameter" + ) diff --git a/packages/pybamm/tests/integration/test_rust_observability_lane_smoke.py b/packages/pybamm/tests/integration/test_rust_observability_lane_smoke.py new file mode 100644 index 0000000000..f727a39f23 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_observability_lane_smoke.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import pytest + +import pybamm +from benchmarks.run_rust_observability import build_parser, repeats_for +from benchmarks.rust_observability.registry import ( + INFERENCE_INPUTS, + get_inference_scenarios, + get_protocol_names, + get_solver_scenarios, + inference_nominal_values, +) +from benchmarks.rust_observability.report import ( + _fits, + render_inference_table, + render_sensitivity_table, + render_solver_table, + suite_to_jsonable, +) +from benchmarks.rust_observability.runners import ( + run_inference_lane, + run_sensitivity_lane, + run_solver_lane, + sample_input_vectors, +) + +# Enough points for the observation grid to interpolate, small enough to run +# every protocol on every lane. +SMOKE_POINTS = 20 +PROTOCOLS = get_protocol_names() + + +def assert_renders(lane: str, results, renderer): + """Every row must render and serialise, whatever its status.""" + assert results + for width in (250, 120, 80): + table = renderer(results, width=width) + assert _fits(table, width) + payload = suite_to_jsonable(lane, results) + assert len(payload["results"]) == len(results) + return {result.backend: result.status for result in results} + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +class TestLanesRunForEveryProtocol: + """Wiring a protocol into the registry does not prove a lane can run it. + + Each lane is exercised end to end so a protocol that builds but cannot be + solved, observed, differentiated or rendered fails here rather than in a + benchmarking session. + """ + + def test_solver_lane(self, protocol): + results = run_solver_lane( + get_solver_scenarios(["SPM"], [protocol], output_points=SMOKE_POINTS), + repeats=1, + warmup=0, + include_aot=False, + ) + statuses = assert_renders("solver", results, render_solver_table) + + # Gated like any other row now that it is scored against the reference. + assert statuses["casadi_idaklu"] in {"pass", "warn"} + assert all(result.reference_tolerance for result in results if result.supported) + assert any(backend.startswith("rust_") for backend in statuses) + for result in results: + if result.supported and result.trajectory_comparison is not None: + assert result.trajectory_comparison.coverage > 0.0 + + def test_sensitivity_lane(self, protocol): + results = run_sensitivity_lane( + get_solver_scenarios(["SPM"], [protocol], output_points=SMOKE_POINTS), + repeats=1, + warmup=0, + include_aot=False, + ) + statuses = assert_renders("sensitivity", results, render_sensitivity_table) + + assert statuses["casadi_idaklu"] in {"pass", "warn"} + # An Interpolant or Experiment current cannot be a sensitivity input. + expected = ( + {"eps_p"} + if protocol in {"drive_cycle", "pulse_train", "experiment"} + else {"I", "eps_p"} + ) + for result in results: + if result.supported: + assert set(result.sensitivity_parameters) == expected + + def test_inference_lane(self, protocol): + results = run_inference_lane( + get_inference_scenarios(["SPM"], [protocol], output_points=SMOKE_POINTS), + repeats=2, + warmup=0, + include_aot=False, + ) + statuses = assert_renders("inference", results, render_inference_table) + + assert statuses["casadi_idaklu"] in {"pass", "warn"} + for result in results: + if result.supported: + assert result.output_comparison is not None + assert result.trajectory_comparison is not None + assert result.sensitivity_comparison is None + # The baseline has nothing to differ from; every other row does. + is_baseline = result.backend == "casadi_idaklu" + assert (result.baseline_delta is None) is is_baseline + if result.supported: + assert result.cold_observe_ms > 0.0 + assert len(result.eval_samples_ms) == 2 + + +class TestDefaultSamplingDepthIsFeasible: + """The lane draws `warmup + repeats` vectors, so a hazard can hide past a short draw. + + A fitted maximum concentration made the eleventh draw start the cell above its + voltage cutoff, which every smoke test at two or four vectors missed. The depth + is read from the lane's own default rather than pinned here. + """ + + def test_every_default_depth_draw_solves(self): + defaults = build_parser().parse_args([]) + results = run_inference_lane( + get_inference_scenarios( + ["SPM"], ["cc_discharge"], output_points=SMOKE_POINTS + ), + repeats=repeats_for("inference", defaults.repeats), + warmup=defaults.warmup, + seed=defaults.inference_seed, + include_aot=False, + ) + + unsupported = {r.backend: r.reason for r in results if not r.supported} + assert not unsupported, f"default-depth draws are infeasible: {unsupported}" + + def test_no_fitted_parameter_defines_the_initial_state(self): + # The base set fixes the initial concentration absolutely, so fitting a + # maximum concentration moves the stoichiometry rather than the capacity. + assert not [ + name + for name in INFERENCE_INPUTS + if name.startswith("Maximum concentration") + ] + + def test_sampled_draws_keep_the_initial_stoichiometry_physical(self): + defaults = build_parser().parse_args([]) + scenario = get_inference_scenarios(["SPM"], ["cc_discharge"])[0] + vectors = sample_input_vectors( + inference_nominal_values(), + defaults.warmup + repeats_for("inference", defaults.repeats), + seed=defaults.inference_seed, + ) + built = scenario.parameter_values_builder() + + for electrode in ("negative", "positive"): + initial = built[f"Initial concentration in {electrode} electrode [mol.m-3]"] + maximum = built[f"Maximum concentration in {electrode} electrode [mol.m-3]"] + for vector in vectors: + stoichiometry = pybamm.Scalar(1) * initial / maximum + value = float(stoichiometry.evaluate(inputs=vector)) + assert 0.0 < value < 1.0, ( + f"{electrode} stoichiometry {value:.4f} is outside [0, 1] " + f"for inputs {vector}" + ) + + +class TestInferenceGradientsAreValidated: + def test_requesting_sensitivities_compares_them(self): + results = run_inference_lane( + get_inference_scenarios( + ["SPM"], ["cc_discharge"], output_points=SMOKE_POINTS + ), + repeats=2, + warmup=0, + sensitivities=True, + include_aot=False, + ) + + candidates = [ + r for r in results if r.backend != "casadi_idaklu" and r.supported + ] + assert candidates + for result in candidates: + assert result.sensitivity_comparison is not None + + def test_gradients_cost_more_than_values_alone(self): + common = { + "scenarios": get_inference_scenarios( + ["SPM"], ["cc_discharge"], output_points=SMOKE_POINTS + ), + "repeats": 2, + "warmup": 1, + "include_aot": False, + } + values_only = run_inference_lane(**common, sensitivities=False) + with_gradients = run_inference_lane(**common, sensitivities=True) + + def baseline(results): + return next(r for r in results if r.backend == "casadi_idaklu") + + # Materialising the chain rule is real work; if it were free the lane + # would not be observing the gradient at all. + assert ( + baseline(with_gradients).observe_median_ms + > baseline(values_only).observe_median_ms + ) diff --git a/packages/pybamm/tests/integration/test_rust_observation.py b/packages/pybamm/tests/integration/test_rust_observation.py new file mode 100644 index 0000000000..f1494ee048 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_observation.py @@ -0,0 +1,423 @@ +"""Parity: native (Rust) observation vs CasADi, for both rust-core solvers.""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import NativeObservation + + +def _solve(solver_cls, native, model_cls=pybamm.lithium_ion.SPM, **opts): + model = model_cls() + solver = solver_cls(**opts) + if native: + model.convert_to_format = "rust" + # Force the per-solver capability flag on; the instance attribute survives + # the shallow solver.copy() that Simulation performs. + solver._observes_via_compiled_model = True + sim = pybamm.Simulation(model, solver=solver) + return sim.solve([0, 3600]) + + +def _native_idaklu(): + """An IDAKLUSolver with native observation forced on (foundation tests only). + + idaklu's `_observes_via_compiled_model` stays False until prereqs A + B; the + instance attribute set here survives the shallow solver.copy() Simulation does. + Pair with `model.convert_to_format = "rust"`. + """ + solver = pybamm.IDAKLUSolver() + solver._observes_via_compiled_model = True + return solver + + +SOLVERS = [ + ("idaklu", pybamm.IDAKLUSolver, {}), + ("diffsol", pybamm.DiffsolSolver, {}), +] +NATIVE_IDS = [s[0] for s in SOLVERS] + +MODELS = [ + ("spm", pybamm.lithium_ion.SPM), + ("spme", pybamm.lithium_ion.SPMe), +] +MODEL_IDS = [m[0] for m in MODELS] + +# 0D scalars plus 1D particle (r-axis) and electrolyte (x-axis) fields, so the +# unroll/transpose contract is exercised on two distinct 1D axes. +DIRECT_VARS = [ + "Terminal voltage [V]", # 0D + "Current [A]", # 0D + "Negative particle concentration [mol.m-3]", # 1D (r-axis) + "Electrolyte concentration [mol.m-3]", # 1D (x-axis) +] + + +@pytest.mark.parametrize("solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS) +@pytest.mark.parametrize("model_name,model_cls", MODELS, ids=MODEL_IDS) +@pytest.mark.parametrize("var", DIRECT_VARS) +def test_direct_variable_parity( + solver_name, solver_cls, native_opts, model_name, model_cls, var +): + native_sol = _solve(solver_cls, native=True, model_cls=model_cls, **native_opts) + casadi_sol = _solve(solver_cls, native=False, model_cls=model_cls) + np.testing.assert_allclose( + native_sol[var].entries, casadi_sol[var].entries, rtol=1e-9, atol=1e-11 + ) + + +class TestZeroCopyStateTrajectory: + """Pin the F-contiguity precondition for zero-copy `eval_trajectory`. + + `all_ys[i]` is exactly what the native path hands to + `CompiledFunction.eval_trajectory`. It MUST be F-contiguous `(n_states, + n_times)`; otherwise every observe gather-copies the whole state matrix + (silent perf regression, no error). + """ + + @pytest.mark.parametrize( + "solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS + ) + def test_full_state_trajectory_is_f_contiguous( + self, solver_name, solver_cls, native_opts + ): + sol = _solve(solver_cls, native=True, **native_opts) + assert isinstance(sol.observation, NativeObservation) # native path active + for ys in sol.all_ys: + arr = np.asarray(ys) + assert arr.ndim == 2 + # F-contiguous ⇒ columns_slice borrows (zero-copy). + assert arr.flags["F_CONTIGUOUS"], ( + f"{solver_name}: all_ys segment is not F-contiguous " + f"(shape={arr.shape}, strides={arr.strides}) — eval_trajectory " + "would gather-copy the whole state matrix on every observe." + ) + + def test_composed_trajectory_segments_are_f_contiguous(self): + # Stepping accumulates per-segment arrays into one Solution, and each segment + # is consumed by its own eval_trajectory, so each must stay F-contiguous. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + for _ in range(3): + sim.step(360) + for ys in sim.solution.all_ys: + arr = np.asarray(ys) + assert arr.flags["F_CONTIGUOUS"], ( + f"composed all_ys segment is not F-contiguous " + f"(shape={arr.shape}, strides={arr.strides})" + ) + + +@pytest.mark.parametrize("solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS) +def test_time_integral_parity(solver_name, solver_cls, native_opts): + native_sol = _solve(solver_cls, native=True, **native_opts) + casadi_sol = _solve(solver_cls, native=False) + np.testing.assert_allclose( + native_sol["Discharge capacity [A.h]"].entries, + casadi_sol["Discharge capacity [A.h]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + +@pytest.mark.parametrize("solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS) +def test_explicit_time_integral_parity(solver_name, solver_cls, native_opts): + # A genuine ExplicitTimeIntegral output, unlike "Discharge capacity [A.h]", so + # this drives the native backend's time-integral branch, not the direct one. + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( + model.variables["Current [A]"], pybamm.Scalar(0.0) + ) + solver = solver_cls(**native_opts) + if native: + model.convert_to_format = "rust" + solver._observes_via_compiled_model = True + return pybamm.Simulation(model, solver=solver).solve([0, 3600]) + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert isinstance(native_sol.observation, NativeObservation) # native path active + native_entries = native_sol["Integrated current [A.s]"].entries + # time-integral output is time-independent → a single scalar + assert native_entries.shape == (1,), ( + f"Expected scalar shape (1,) for time-integral output, got {native_entries.shape}" + ) + np.testing.assert_allclose( + native_entries, + casadi_sol["Integrated current [A.s]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + +def test_stepping_addition_parity(): + # __add__: Simulation.step accumulates sub-solutions into one combined + # Solution; observe on it. `sim.solution` is the running combined result. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + for _ in range(5): + sim.step(360) + native_v = sim.solution["Terminal voltage [V]"].entries + + cmodel = pybamm.lithium_ion.SPM() + csim = pybamm.Simulation(cmodel, solver=pybamm.IDAKLUSolver()) + for _ in range(5): + csim.step(360) + casadi_v = csim.solution["Terminal voltage [V]"].entries + np.testing.assert_allclose(native_v, casadi_v, rtol=1e-8, atol=1e-10) + + +def test_first_last_state_parity(): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + sol = sim.solve([0, 3600]) + # Observation must still work through first/last_state, which carry the context. + assert sol.first_state["Terminal voltage [V]"].entries.size >= 1 + assert sol.last_state["Terminal voltage [V]"].entries.size >= 1 + + +def test_experiment_parity(): + experiment = pybamm.Experiment( + ["Discharge at 1C until 3.0 V", "Charge at 1C until 4.0 V"] + ) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation( + model, + experiment=experiment, + solver=_native_idaklu(), + ) + sol = sim.solve() + # Verify the experiment result genuinely used the native path. + assert isinstance(sol.observation, NativeObservation), ( + "Experiment solution lost its native backend — native path not active" + ) + native_v = sol["Terminal voltage [V]"].entries + + cmodel = pybamm.lithium_ion.SPM() + csim = pybamm.Simulation( + cmodel, experiment=experiment, solver=pybamm.IDAKLUSolver() + ) + csol = csim.solve() + np.testing.assert_allclose( + native_v, csol["Terminal voltage [V]"].entries, rtol=1e-7, atol=1e-9 + ) + + +def test_diffsol_native_by_default(): + # After the flip, a plain DiffsolSolver observes natively without any opt-in. + t_eval = np.linspace(0, 3600, 51) + + model = pybamm.lithium_ion.SPM() + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve(t_eval=t_eval) + assert isinstance(sol.observation, NativeObservation) # native path active + + # Native diffsol vs CasADi-backed IDAKLU at their default tolerances + # (1e-6 and 1e-4), so rtol=1e-4 is the right cross-solver bound. + cmodel = pybamm.lithium_ion.SPM() + csim = pybamm.Simulation(cmodel, solver=pybamm.IDAKLUSolver()) + csol = csim.solve(t_eval=t_eval, t_interp=t_eval) + np.testing.assert_allclose( + sol["Terminal voltage [V]"].entries, + csol["Terminal voltage [V]"].entries, + rtol=1e-4, + atol=1e-6, + ) + + +def test_diffsol_does_not_build_casadi_observation(monkeypatch): + import pybamm.solvers.solution as solution_mod + + model = pybamm.lithium_ion.SPM() + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve([0, 600]) + + called = {"casadi": False} + orig = solution_mod.Solution._convert_to_casadi + + def spy(self, *a, **k): + called["casadi"] = True + return orig(self, *a, **k) + + monkeypatch.setattr(solution_mod.Solution, "_convert_to_casadi", spy) + _ = sol["Terminal voltage [V]"] + assert called["casadi"] is False + + +def test_multistep_time_integral_parity(): + # Phase-2 x Phase-3: a time integral observed over a MULTI-SEGMENT solution + # must integrate ONCE over the full trajectory (shape (1,)), not per segment. + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( + model.variables["Current [A]"], pybamm.Scalar(0.0) + ) + solver = pybamm.IDAKLUSolver() + if native: + model.convert_to_format = "rust" + solver._observes_via_compiled_model = True + sim = pybamm.Simulation(model, solver=solver) + for _ in range(3): + sim.step(360) + return sim.solution + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert isinstance(native_sol.observation, NativeObservation) + nv = native_sol["Integrated current [A.s]"].entries + cv = casadi_sol["Integrated current [A.s]"].entries + # Single integral over the whole 3-step trajectory, NOT one value per step. + assert nv.shape == (1,), f"expected single scalar, got shape {nv.shape}" + np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) + + +def test_discrete_time_sum_parity(): + # Both paths sum the integrand over solution.t, so t_interp == discrete times + # makes the sum run over exactly those times. + discrete_times = np.linspace(0.0, 3600.0, 11) + data = pybamm.DiscreteTimeData(discrete_times, np.zeros(11), "td") + + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["DT sum"] = pybamm.DiscreteTimeSum( + model.variables["Voltage [V]"] - data + ) + solver = pybamm.IDAKLUSolver() + if native: + model.convert_to_format = "rust" + solver._observes_via_compiled_model = True + return pybamm.Simulation(model, solver=solver).solve( + t_eval=[0, 3600], t_interp=discrete_times + ) + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert isinstance(native_sol.observation, NativeObservation) + nv = native_sol["DT sum"].entries + cv = casadi_sol["DT sum"].entries + assert nv.shape == (1,), f"expected a single scalar, got shape {nv.shape}" + np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) + + +def test_native_observed_variable_sensitivities_require_calculate_sensitivities(): + # Without calculate_sensitivities, a native solve with an input parameter must + # raise the same informative error as CasADi, not return None or {}. + def _build(calculate_sensitivities): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + param = model.default_parameter_values + param["Current function [A]"] = pybamm.InputParameter("I") + sim = pybamm.Simulation(model, parameter_values=param, solver=_native_idaklu()) + return sim.solve( + [0, 3600], + inputs={"I": 0.5}, + calculate_sensitivities=calculate_sensitivities, + ) + + sol = _build(calculate_sensitivities=False) + assert isinstance(sol.observation, NativeObservation) + v = sol["Terminal voltage [V]"] + assert len(sol.all_inputs[0]) > 0 # inputs present, so the no-input shortcut + with pytest.raises(ValueError, match=r"calculate_sensitivities"): + _ = v.sensitivities + + # Requesting sensitivities: the native path returns real values, not {}. + sens_sol = _build(calculate_sensitivities=True) + assert sens_sol["Terminal voltage [V]"].sensitivities != {} + + +@pytest.mark.parametrize("input_val", [1, np.int64(1), np.float32(1.0)]) +def test_native_observation_nonfloat_input_parity(input_val): + # non-float inputs (int64/float32) must match CasADi, not crash the input pack. + def build(native): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + param = model.default_parameter_values + param["Current function [A]"] = pybamm.InputParameter("I") + solver = pybamm.IDAKLUSolver() + solver._observes_via_compiled_model = native + sim = pybamm.Simulation(model, parameter_values=param, solver=solver) + return sim.solve([0, 3600], inputs={"I": input_val}) + + native_sol = build(native=True) + casadi_sol = build(native=False) + assert isinstance(native_sol.observation, NativeObservation) + np.testing.assert_allclose( + native_sol["Terminal voltage [V]"].entries, + casadi_sol["Terminal voltage [V]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + +@pytest.mark.parametrize("solver_cls", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver]) +def test_dense_trajectory_parity(solver_cls): + # parity must hold over a dense t_interp (large n_t), not only adaptive steps. + t_interp = np.linspace(0, 3600, 500) + + def build(native): + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + solver = solver_cls() + solver._observes_via_compiled_model = native + return pybamm.Simulation(model, solver=solver).solve( + t_eval=[0, 3600], t_interp=t_interp + ) + + native_sol = build(native=True) + casadi_sol = build(native=False) + assert isinstance(native_sol.observation, NativeObservation) + assert native_sol["Terminal voltage [V]"].entries.shape[-1] == len(t_interp) + for var in ("Terminal voltage [V]", "Electrolyte concentration [mol.m-3]"): + np.testing.assert_allclose( + native_sol[var].entries, casadi_sol[var].entries, rtol=1e-8, atol=1e-10 + ) + + +def test_2d_variable_native_parity(): + # A 2D field (resolved across both r and x in DFN) drives initialise_2D and + # locks the native column-major -> .T -> reshape contract against CasADi. + var = "Negative particle concentration [mol.m-3]" + native_sol = _solve( + pybamm.IDAKLUSolver, native=True, model_cls=pybamm.lithium_ion.DFN + ) + casadi_sol = _solve( + pybamm.IDAKLUSolver, native=False, model_cls=pybamm.lithium_ion.DFN + ) + native_entries = native_sol[var].entries + casadi_entries = casadi_sol[var].entries + # (r, x, t): genuinely 2D in space, so this is the 2D unroll path. + assert native_entries.ndim == 3, ( + f"expected a 2D field (r, x, t), got ndim={native_entries.ndim}" + ) + assert native_entries.shape == casadi_entries.shape + # This test guards the 2D layout contract. Native and CasADi use different eval + # cores, whose tiny RHS differences DFN amplifies to ~1e-6, over the trajectory. + np.testing.assert_allclose(native_entries, casadi_entries, rtol=1e-4, atol=1e-3) + + +def test_3d_variable_native_parity(): + # 3D field (MPM, r x R x x): drives initialise_3D / unroll_3D. Order-pinning + # test for the unroll_3D fix; a scrambled field is non-uniform at t=0. + var = "Negative particle concentration distribution [mol.m-3]" + native_sol = _solve( + pybamm.IDAKLUSolver, native=True, model_cls=pybamm.lithium_ion.MPM + ) + casadi_sol = _solve( + pybamm.IDAKLUSolver, native=False, model_cls=pybamm.lithium_ion.MPM + ) + native_entries = native_sol[var].entries + casadi_entries = casadi_sol[var].entries + # (r, R, x, t): genuinely 3D in space, so this is the 3D unroll path. + assert native_entries.ndim == 4, ( + f"expected a 3D field (r, R, x, t), got ndim={native_entries.ndim}" + ) + assert native_entries.shape == casadi_entries.shape + # Layout contract, not numerical equality (cross eval-core drift); a wrong + # 3D reshape/axis-swap would mismatch by orders of magnitude. + np.testing.assert_allclose(native_entries, casadi_entries, rtol=1e-4, atol=1e-3) diff --git a/packages/pybamm/tests/integration/test_rust_parity.py b/packages/pybamm/tests/integration/test_rust_parity.py new file mode 100644 index 0000000000..ff79038d44 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_parity.py @@ -0,0 +1,119 @@ +# tests/integration/test_rust_parity.py +"""Dual-backend parity tests: verify Rust evaluation matches CasADi.""" + +import numpy as np + +import pybamm + +# Import fixtures from conftest_rust +pytest_plugins = ["tests.conftest_rust"] + + +class TestRustCasadiParity: + """Tests that verify Rust backend produces same results as CasADi.""" + + def test_scalar_arithmetic(self, dual_backend_compare): + """Basic scalar arithmetic.""" + a = pybamm.Scalar(2.0) + b = pybamm.Scalar(3.0) + + dual_backend_compare(a + b) + dual_backend_compare(a - b) + dual_backend_compare(a * b) + dual_backend_compare(a / b) + dual_backend_compare(a**b) + + def test_unary_functions(self, dual_backend_compare): + """Unary math functions.""" + x = pybamm.InputParameter("x") + + dual_backend_compare(pybamm.sqrt(x), inputs={"x": 4.0}) + dual_backend_compare(pybamm.exp(x), inputs={"x": 1.0}) + dual_backend_compare(pybamm.log(x), inputs={"x": np.e}) + dual_backend_compare(pybamm.sin(x), inputs={"x": np.pi / 6}) + dual_backend_compare(pybamm.cos(x), inputs={"x": np.pi / 3}) + dual_backend_compare(pybamm.tanh(x), inputs={"x": 1.0}) + + def test_state_vector_operations(self, dual_backend_compare): + """StateVector arithmetic.""" + sv = pybamm.StateVector(slice(0, 3)) + expr = pybamm.Scalar(2.0) * sv + pybamm.Scalar(1.0) + + y = np.array([1.0, 2.0, 3.0]) + dual_backend_compare(expr, y=y) + + def test_interpolation_1d(self, dual_backend_compare): + """1D linear interpolation.""" + x_data = np.linspace(0, 1, 50) + y_data = 2 * x_data + sv = pybamm.StateVector(slice(0, 2)) + interp = pybamm.Interpolant(x_data, y_data, sv) + + y = np.array([0.25, 0.75]) + dual_backend_compare(interp, y=y) + + def test_interpolation_1d_vector_valued_y(self, dual_backend_compare): + """1D interpolation with vector-valued y (one output per column).""" + x_data = np.linspace(0, 1, 50) + y_data = np.column_stack([2 * x_data, np.sin(x_data)]) + sv = pybamm.StateVector(slice(0, 1)) + + y = np.array([0.25]) + for interpolator in ["linear", "cubic", "pchip"]: + interp = pybamm.Interpolant(x_data, y_data, sv, interpolator=interpolator) + dual_backend_compare(interp, y=y) + + def test_concatenation(self, dual_backend_compare): + """Vector concatenation.""" + v1 = pybamm.Vector(np.array([1.0, 2.0])) + v2 = pybamm.Vector(np.array([3.0, 4.0, 5.0])) + expr = pybamm.NumpyConcatenation(v1, v2) + + dual_backend_compare(expr) + + def test_sparse_matmul(self, dual_backend_compare): + """Sparse matrix-vector multiplication.""" + from scipy.sparse import csr_matrix + + data = np.array([1.0, 2.0, 3.0]) + row = np.array([0, 1, 2]) + col = np.array([0, 1, 2]) + sparse_mat = csr_matrix((data, (row, col)), shape=(3, 3)) + + mat = pybamm.Matrix(sparse_mat) + sv = pybamm.StateVector(slice(0, 3)) + expr = mat @ sv + + y = np.array([1.0, 2.0, 3.0]) + dual_backend_compare(expr, y=y) + + def test_composed_expression(self, dual_backend_compare): + """Complex composed expression.""" + sv = pybamm.StateVector(slice(0, 1)) + t = pybamm.Time() + expr = pybamm.sqrt(sv**2 + pybamm.Scalar(1.0)) * pybamm.exp(-t) + + y = np.array([3.0]) + dual_backend_compare(expr, t=2.0, y=y, rtol=1e-8) + + def test_conditional_parity(self, dual_backend_compare): + """Conditional branch selection matches CasADi.""" + selector = pybamm.InputParameter("s") + branch1 = pybamm.InputParameter("a") * pybamm.Scalar(2.0) + branch2 = pybamm.InputParameter("b") + pybamm.Scalar(5.0) + branch3 = pybamm.Scalar(42.0) + expr = pybamm.Conditional(selector, branch1, branch2, branch3) + + # Test each branch + dual_backend_compare( + expr, inputs={"s": 1.0, "a": 10.0, "b": 20.0} + ) # branch1: 20 + dual_backend_compare( + expr, inputs={"s": 2.0, "a": 10.0, "b": 20.0} + ) # branch2: 25 + dual_backend_compare( + expr, inputs={"s": 3.0, "a": 10.0, "b": 20.0} + ) # branch3: 42 + dual_backend_compare( + expr, inputs={"s": 0.0, "a": 10.0, "b": 20.0} + ) # no branch: 0 diff --git a/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py b/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py index 5e24c85860..3e2031179c 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py +++ b/packages/pybamm/tests/integration/test_solvers/test_casadi_mode.py @@ -25,6 +25,7 @@ def test_casadi_solver_mode(self): for solver in solvers: # define model model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" # solve simulation sim = pybamm.Simulation( @@ -57,6 +58,7 @@ def test_casadi_fast_matches_safe_default_ode(self): solutions = [] for mode in ["safe", "fast"]: model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(mode=mode, atol=1e-6, rtol=1e-6) sim = pybamm.Simulation(model, solver=solver) solutions.append(sim.solve(t_eval)) diff --git a/packages/pybamm/tests/integration/test_solvers/test_idaklu.py b/packages/pybamm/tests/integration/test_solvers/test_idaklu.py index 69999bb527..b57d9ab44d 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_idaklu.py +++ b/packages/pybamm/tests/integration/test_solvers/test_idaklu.py @@ -323,6 +323,7 @@ class TestIDAKLUSolverAOTCompilation: @staticmethod def _build_model(): model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "casadi" geometry = model.default_geometry param = model.default_parameter_values param.process_model(model) diff --git a/packages/pybamm/tests/integration/test_voltage_as_state.py b/packages/pybamm/tests/integration/test_voltage_as_state.py index f188ce48cc..289e568c37 100644 --- a/packages/pybamm/tests/integration/test_voltage_as_state.py +++ b/packages/pybamm/tests/integration/test_voltage_as_state.py @@ -64,6 +64,7 @@ def test_default_voltage_is_expression(self, model_cls): @pytest.mark.parametrize("model_cls", REDUCED_MODELS) def test_opt_in_solvable_by_casadi_safe(self, model_cls): model = model_cls(options=OPT_IN) + model.convert_to_format = "casadi" sim = pybamm.Simulation(model, solver=pybamm.CasadiSolver(mode="safe")) sol = sim.solve([0, 3600]) v = sol["Voltage [V]"].entries diff --git a/packages/pybamm/tests/shared.py b/packages/pybamm/tests/shared.py index bd31c3f8b9..f0423af5c4 100644 --- a/packages/pybamm/tests/shared.py +++ b/packages/pybamm/tests/shared.py @@ -433,6 +433,22 @@ def get_base_model_with_battery_geometry(**kwargs): return model +def get_broken_input_model(convert_to_format="rust"): + """``dv/dt = -sqrt(k)*v``, whose residual is NaN for a negative ``k``. + + The one per-input-set failure a sweep can provoke without depending on how + the integrator diverges, so only the sets given a negative ``k`` fail. + """ + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: -pybamm.sqrt(pybamm.InputParameter("k")) * v} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + model.convert_to_format = convert_to_format + pybamm.Discretisation().process_model(model) + return model + + def get_required_distribution_deps(package_name): pattern = re.compile(r"(?!.*extra\b)^([^<>=;\[]+)\b.*$") if json_deps := importlib_metadata.metadata(package_name).json.get("requires_dist"): @@ -529,3 +545,49 @@ def get_cylindrical_mesh_for_testing_symbolic(): var_pts = {cylindrical_r: 15} mesh = pybamm.Mesh(geometry, submesh_types, var_pts) return mesh + + +# 2+1D pouch DFN, the model whose Jacobian colouring the constant split changes. +POUCH_OPTIONS = {"current collector": "potential pair", "dimensionality": 2} +POUCH_PTS = {"x_n": 4, "x_s": 4, "x_p": 4, "r_n": 4, "r_p": 4, "y": 8, "z": 8} + + +def build_rust_model(model, var_pts=None): + """Set up `model` on the Rust backend, returning its built model and core.""" + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, var_pts=var_pts) + sim.build() + sim.solver.set_up(sim.built_model, inputs=[{}]) + return sim.built_model, sim.solver._setup["rust_model"] + + +def dense_rust_jacobian(built_model, rust_model): + """Assemble the Rust merged CSC at `y0` and scatter it dense.""" + import numpy as np + from scipy import sparse + + n = built_model.len_rhs_and_alg + y = np.ascontiguousarray(np.asarray(built_model.y0_list[0]).flatten().astype(float)) + jac_data = np.zeros(rust_model.nnz) + rust_model.assemble_jacobian_csc_into(0.0, y, 0.0, np.array([]), jac_data) + colptr, rowind = rust_model.csc_sparsity_pattern() + dense = sparse.csc_matrix( + (jac_data, np.asarray(rowind), np.asarray(colptr)), shape=(n, n) + ).toarray() + return y, dense + + +def dense_casadi_jacobian(model, y, var_pts=None): + """The same Jacobian off the CasADi backend, for a value comparison.""" + import casadi + import numpy as np + + model.convert_to_format = "casadi" + sim = pybamm.Simulation(model, var_pts=var_pts) + sim.build() + sim.solver.set_up(sim.built_model, inputs=[{}]) + return np.array( + sim.built_model.jac_rhs_algebraic_eval( + casadi.DM(0.0), casadi.DM(y), casadi.DM(np.zeros((0, 1))) + ) + ) diff --git a/packages/pybamm/tests/strategies/__init__.py b/packages/pybamm/tests/strategies/__init__.py index b0ddbb5fda..54eb0b6421 100644 --- a/packages/pybamm/tests/strategies/__init__.py +++ b/packages/pybamm/tests/strategies/__init__.py @@ -1,4 +1,4 @@ -"""Hypothesis strategies and shared settings for the serialisation property tests. +"""Hypothesis strategies and shared settings for the property tests. ``serialisation_settings`` is applied as a decorator to the individual serialisation property/smoke tests rather than registered as a suite-wide @@ -11,7 +11,7 @@ import os -from hypothesis import settings +from hypothesis import HealthCheck, settings #: Shared Hypothesis settings for the serialisation property tests # Increase `max_examples` for a larger amount of randomised draws @@ -19,3 +19,13 @@ max_examples=500, deadline=None, ) + +#: Shared Hypothesis settings for properties whose example runs a solve. +# Far fewer examples than the serialisation properties: each integrates a sweep twice. +solve_settings = settings( + max_examples=20, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) + +__all__ = ["serialisation_settings", "solve_settings"] diff --git a/packages/pybamm/tests/strategies/input_sweeps.py b/packages/pybamm/tests/strategies/input_sweeps.py new file mode 100644 index 0000000000..ee53b67f1a --- /dev/null +++ b/packages/pybamm/tests/strategies/input_sweeps.py @@ -0,0 +1,38 @@ +"""Hypothesis strategies for sweeps over solver input sets. + +A sweep is what ``num_threads`` parallelises over, so the property these feed is +that solving one concurrently is indistinguishable from solving it in a loop. +""" + +from __future__ import annotations + +import hypothesis.strategies as st + + +def decay_rate_sweeps( + min_sets: int = 2, max_sets: int = 8 +) -> st.SearchStrategy[list[float]]: + """Sweeps of distinct positive decay rates. + + Rates span two decades, so a drawn sweep is heterogeneous in solve cost and + in whether each set reaches an event, which is where a batch that shared + scratch or returned completion order would show it. + + Parameters + ---------- + min_sets : int, optional + Fewest input sets to draw (default 2, the smallest batch). + max_sets : int, optional + Most input sets to draw (default 8). + + Returns + ------- + :class:`hypothesis.strategies.SearchStrategy` + Lists of rates, each list free of duplicates. + """ + return st.lists( + st.floats(min_value=0.05, max_value=5.0, allow_nan=False, allow_infinity=False), + min_size=min_sets, + max_size=max_sets, + unique=True, + ) diff --git a/packages/pybamm/tests/strategies/serialise_values.py b/packages/pybamm/tests/strategies/serialise_values.py index fece2bb239..80b1db557e 100644 --- a/packages/pybamm/tests/strategies/serialise_values.py +++ b/packages/pybamm/tests/strategies/serialise_values.py @@ -31,6 +31,7 @@ pybamm.CasadiAlgebraicSolver, pybamm.NonlinearSolver, pybamm.CompositeSolver, + pybamm.DiffsolSolver, ) @@ -87,6 +88,11 @@ def _strategy(draw): elif cls is pybamm.NonlinearSolver: kwargs["rtol"] = draw(_POSITIVE_TOL) kwargs["atol"] = draw(_POSITIVE_TOL) + elif cls is pybamm.DiffsolSolver: + kwargs["rtol"] = draw(_POSITIVE_TOL) + kwargs["atol"] = draw(_POSITIVE_TOL) + kwargs["root_tol"] = draw(_POSITIVE_TOL) + kwargs["calc_ic"] = draw(st.booleans()) elif cls is pybamm.CompositeSolver: # Draw two simple solver classes (possibly the same) for sub_solvers. sub_cls_a, sub_cls_b = draw( diff --git a/packages/pybamm/tests/unit/snapshots/test_rust_observability_benchmark/test_solver_table_snapshot_and_json_samples/rust_observability_solver_table.snapshot b/packages/pybamm/tests/unit/snapshots/test_rust_observability_benchmark/test_solver_table_snapshot_and_json_samples/rust_observability_solver_table.snapshot new file mode 100644 index 0000000000..92df37d638 --- /dev/null +++ b/packages/pybamm/tests/unit/snapshots/test_rust_observability_benchmark/test_solver_table_snapshot_and_json_samples/rust_observability_solver_table.snapshot @@ -0,0 +1,15 @@ +Δ is the raw difference from casadi_idaklu at the scenario tolerance; no converged reference was run. + +Timings (ms) +Scenario Backend Protocol Pts Build Prep Cold Warm Solve Wall Int Obs E2E +------------------------------------------------------------------------------------------------------------------- +SPM rust_idaklu cc_discharge 100 10.00 5.20 16.50 0.10 1.20 1.30 1.10 0.20 1.50 +SPMe rust_idaklu cc_discharge 100 10.00 5.20 16.50 0.10 1.20 1.30 1.10 0.20 1.50 +DFN rust_idaklu cc_discharge 100 10.00 5.20 16.50 0.10 1.20 1.30 1.10 0.20 1.50 + +Validation +Scenario Backend Protocol Cover Δt Clr Rows Entry Tape State Δ Output Δ Base Δ Status Reason +----------------------------------------------------------------------------------------------------------------------- +SPM rust_idaklu cc_discharge 1.000 0.00e+00 5 0 0 0 1.00e-08 1.00e-08 - pass +SPMe rust_idaklu cc_discharge 1.000 0.00e+00 3 1 65 4096 1.00e-08 1.00e-08 - pass +DFN rust_idaklu cc_discharge 1.000 0.00e+00 9 0 0 0 1.00e-08 1.00e-08 - pass diff --git a/packages/pybamm/tests/unit/test_codegen/test_compilation.py b/packages/pybamm/tests/unit/test_codegen/test_compilation.py index c720cde89e..0372c7553b 100644 --- a/packages/pybamm/tests/unit/test_codegen/test_compilation.py +++ b/packages/pybamm/tests/unit/test_codegen/test_compilation.py @@ -15,6 +15,7 @@ _CACHE, _PER_ATTEMPT_TOKEN, _STALE_TMP_AGE_S, + _capture_aot_compile_events, _default_cache_dir, _maybe_sweep_stale, _shared_ext, @@ -90,13 +91,19 @@ def test_outputs_match_original(self, cache_dir): def test_in_memory_cache_hit(self, cache_dir): f = _make_simple_fn("test_aot_inmem_cache") - g1 = aot_compile(f, cache_dir=cache_dir) - g2 = aot_compile(f, cache_dir=cache_dir) + with _capture_aot_compile_events() as events: + g1 = aot_compile(f, cache_dir=cache_dir) + g2 = aot_compile(f, cache_dir=cache_dir) assert g1 is g2 + assert [event.cache_status for event in events] == ["miss", "memory"] + assert events[0].cache_key == events[1].cache_key + assert events[0].compiler_ms > 0 + assert events[0].library_size_bytes > 0 def test_on_disk_cache_skips_recompile(self, cache_dir): f = _make_simple_fn("test_aot_disk_cache") - _ = aot_compile(f, cache_dir=cache_dir) + with _capture_aot_compile_events() as miss_events: + _ = aot_compile(f, cache_dir=cache_dir) sofile = next( os.path.join(cache_dir, p) for p in os.listdir(cache_dir) @@ -115,11 +122,19 @@ def fake_run(*args, **kwargs): # pragma: no cover - fail path compilation_module.subprocess.run = fake_run try: - g2 = aot_compile(f, cache_dir=cache_dir) + with _capture_aot_compile_events() as events: + g2 = aot_compile(f, cache_dir=cache_dir) finally: compilation_module.subprocess.run = original_run assert called == [], "gcc was invoked despite an existing on-disk .dylib" + assert len(events) == 1 + assert events[0].cache_status == "disk" + assert events[0].cache_key == miss_events[0].cache_key + assert events[0].library_path == sofile + assert events[0].library_size_bytes == os.path.getsize(sofile) + assert events[0].compiler_ms == 0 + assert events[0].load_ms > 0 assert g2.class_name() == "External" assert os.path.getmtime(sofile) == mtime_before np.testing.assert_allclose( @@ -145,12 +160,15 @@ def fake_run(*args, **kwargs): # pragma: no cover - fail path compilation_module.subprocess.run = fake_run try: - g2 = aot_compile(g, cache_dir=cache_dir) + with _capture_aot_compile_events() as events: + g2 = aot_compile(g, cache_dir=cache_dir) finally: compilation_module.subprocess.run = original_run assert g2 is g assert called == [] + assert len(events) == 1 + assert events[0].cache_status == "external" n_dylibs_after = sum( 1 for p in os.listdir(cache_dir) if p.endswith(_shared_ext()) ) @@ -158,13 +176,40 @@ def fake_run(*args, **kwargs): # pragma: no cover - fail path def test_compiler_failure_returns_original(self, cache_dir, caplog): f = _make_simple_fn("test_aot_compiler_failure") - with caplog.at_level(logging.WARNING, logger="pybamm.logger"): + with ( + caplog.at_level(logging.WARNING, logger="pybamm.logger"), + _capture_aot_compile_events() as events, + ): g = aot_compile(f, cache_dir=cache_dir, compiler="nonexistent_compiler_x") assert g is f + assert len(events) == 1 + assert events[0].cache_status == "fallback" + assert "ValueError" in events[0].error assert not any(p.endswith(_shared_ext()) for p in os.listdir(cache_dir)) assert len(_CACHE) == 0 assert any("Failed to compile" in r.getMessage() for r in caplog.records) + def test_library_size_telemetry_is_best_effort( + self, cache_dir, monkeypatch, caplog + ): + f = _make_simple_fn("test_aot_library_size_best_effort") + + def fail_getsize(_path): + raise OSError("metadata unavailable") + + monkeypatch.setattr(compilation_module.os.path, "getsize", fail_getsize) + with ( + caplog.at_level(logging.WARNING, logger="pybamm.logger"), + _capture_aot_compile_events() as events, + ): + g = aot_compile(f, cache_dir=cache_dir) + + assert g.class_name() == "External" + assert len(events) == 1 + assert events[0].cache_status == "miss" + assert events[0].library_size_bytes is None + assert not caplog.records + def test_atomic_install_no_partial_dylib_on_failure(self, cache_dir): f = _make_simple_fn("test_aot_atomic_install") diff --git a/packages/pybamm/tests/unit/test_eis_simulation.py b/packages/pybamm/tests/unit/test_eis_simulation.py index 9d30746f94..ecf5b81867 100644 --- a/packages/pybamm/tests/unit/test_eis_simulation.py +++ b/packages/pybamm/tests/unit/test_eis_simulation.py @@ -236,6 +236,18 @@ def test_initial_soc_forwards_inputs(self): z_v = eis_sim.solve(frequencies, inputs={"vf_solid": 0.94}, initial_soc="0.5 V") assert z_v.impedance.shape == (5,) + def test_eis_rust_backend_matches_casadi(self): + def impedance(fmt): + model = pybamm.lithium_ion.SPM( + options={"surface form": "differential"}, name="SPM" + ) + model.convert_to_format = fmt + eis_sim = pybamm.EISSimulation(model) + frequencies = np.logspace(-2, 2, 10) + return eis_sim.solve(frequencies).impedance + + np.testing.assert_allclose(impedance("rust"), impedance("casadi"), rtol=1e-6) + def test_high_freq_intercept_matches_contact_resistance(self): model = pybamm.lithium_ion.SPM( options={"surface form": "differential", "contact resistance": "true"} diff --git a/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py b/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py index 13f6c719e8..b51a4a3ad6 100644 --- a/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py +++ b/packages/pybamm/tests/unit/test_experiments/test_simulation_with_experiment.py @@ -5,6 +5,7 @@ import re from datetime import datetime from types import SimpleNamespace +from typing import NamedTuple import casadi import numpy as np @@ -19,25 +20,78 @@ def _default_timespan(self, value): return 1 -def _unified_rhs_jac_n_instructions(unique_step_strings): - """rhs_algebraic and jacobian top-function instruction counts for a unified-mode - SPM experiment with the given distinct steps.""" +class _TopFunctionSize(NamedTuple): + """A unified-model top function's reported instruction count, split into + always-run work and the control-flow nodes that pick a branch.""" + + n_instructions: int + n_control_flow: int + + @property + def always_run_work(self) -> int: + """Reported count less its control-flow nodes: one Rust ``Dispatch`` per + short-circuited conditional (one per half in a split-eval tape), or one + `CasADi` switch call.""" + return self.n_instructions - self.n_control_flow + + +def _casadi_switch_calls(fn): + """How many of ``fn``'s instructions call another function; its Switch is one.""" + return sum( + 1 for k in range(fn.n_instructions()) if fn.instruction_id(k) == casadi.OP_CALL + ) + + +def _unified_rhs_jac_n_instructions(unique_step_strings, convert_to_format): + """Unified-model residual and Jacobian top-function sizes for one backend.""" experiment = pybamm.Experiment([tuple(unique_step_strings)]) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format sim = pybamm.Simulation( - pybamm.lithium_ion.SPM(), + model, experiment=experiment, solver=pybamm.IDAKLUSolver(), experiment_model_mode="unified", ) sim.solve() assert sim._experiment_uses_unified_model - model = sim._built_experiment_model + if convert_to_format == "casadi": + model = sim._built_experiment_model + return tuple( + _TopFunctionSize(fn.n_instructions(), _casadi_switch_calls(fn)) + for fn in (model.rhs_algebraic_eval, model.jac_rhs_algebraic_eval) + ) + + model = sim._built_experiment_solver._setup["rust_model"] + stats = model.jacobian_stats() return ( - model.rhs_algebraic_eval.n_instructions(), - model.jac_rhs_algebraic_eval.n_instructions(), + _TopFunctionSize(model.rhs.n_instructions, model.rhs.n_dispatches), + _TopFunctionSize( + stats["split_eval_total_instructions"], stats["split_eval_dispatch_count"] + ), ) +def _assert_no_always_run_growth(base, extended, what): + """Assert adding control modes did not grow the always-run tape. + + The margin is derived, not fixed: a short-circuited conditional adds one + control-flow instruction to the reported count, which genuinely executes + every call and so stays in the metric. Short-circuiting one more conditional + is therefore allowed, while any per-mode growth of the always-run work still + fails. + """ + assert extended.always_run_work <= base.always_run_work, ( + f"{what} grew the always-run tape: {base} -> {extended}" + ) + + +def _native_jacobian_instruction_count(sim): + """Instruction count for IDAKLU's Rust-native Jacobian artifact.""" + model = sim._built_experiment_solver._setup["rust_model"] + return model.jacobian_stats()["split_eval_total_instructions"] + + def _largest_generated_fn_lines(fn): """Body-line count of the largest ``casadi_fN`` in ``fn``'s generated C.""" gen = casadi.CodeGenerator("probe", {"with_header": False}) @@ -641,37 +695,59 @@ def max_conditional_branches(model): assert set(indices) == set(range(1, n_unique + 1)) assert indices == [(i % n_unique) + 1 for i in range(len(indices))] - def test_unified_switch_top_functions_flat_in_unique_steps(self): - # Per-branch dispatch keeps the top rhs/jac flat in unique step count. + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_unified_switch_top_functions_flat_in_unique_steps(self, convert_to_format): + # Per-branch dispatch keeps the compiled residual and Jacobian flat in step count. rhs_1, jac_1 = _unified_rhs_jac_n_instructions( - [f"Discharge at {0.4 + 0.1 * i:.2f}C for 10 s" for i in range(1)] + [f"Discharge at {0.4 + 0.1 * i:.2f}C for 10 s" for i in range(1)], + convert_to_format, ) rhs_8, jac_8 = _unified_rhs_jac_n_instructions( - [f"Discharge at {0.4 + 0.1 * i:.2f}C for 10 s" for i in range(8)] + [f"Discharge at {0.4 + 0.1 * i:.2f}C for 10 s" for i in range(8)], + convert_to_format, ) - assert rhs_8 <= rhs_1 + 2, ( - f"unified rhs top function grows with unique steps: {rhs_1} -> {rhs_8}" + _assert_no_always_run_growth(rhs_1, rhs_8, "8 unique steps in the rhs") + _assert_no_always_run_growth(jac_1, jac_8, "8 unique steps in the jac") + + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_unified_active_branch_independent_of_other_modes(self, convert_to_format): + # Extra control modes must not inflate the active residual or Jacobian work. + rhs_cc, jac_cc = _unified_rhs_jac_n_instructions( + ["Discharge at 1C for 10 s"], convert_to_format ) - assert jac_8 <= jac_1 + 2, ( - f"unified jac top function grows with unique steps: {jac_1} -> {jac_8}" + rhs_multi, jac_multi = _unified_rhs_jac_n_instructions( + [ + "Discharge at 1C for 10 s", + "Charge at 0.5C until 4.2 V", + "Hold at 4.2 V until C/50", + ], + convert_to_format, ) + _assert_no_always_run_growth(rhs_cc, rhs_multi, "adding modes to the rhs") + _assert_no_always_run_growth(jac_cc, jac_multi, "adding modes to the jac") - def test_unified_active_branch_independent_of_other_modes(self): - # Each mode is a separate branch function, so adding modes doesn't grow the top - # rhs/jac and inactive modes aren't evaluated during an active step. - rhs_cc, jac_cc = _unified_rhs_jac_n_instructions(["Discharge at 1C for 10 s"]) + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_unified_voltage_and_power_modes_share_no_active_work( + self, convert_to_format + ): + # Voltage and power both need the voltage expression, so CSE shares that + # cone; without privatisation it is unowned and stays in the always-run tape. + rhs_cc, jac_cc = _unified_rhs_jac_n_instructions( + ["Discharge at 1C for 10 s"], convert_to_format + ) rhs_multi, jac_multi = _unified_rhs_jac_n_instructions( [ "Discharge at 1C for 10 s", - "Charge at 0.5C until 4.2 V", "Hold at 4.2 V until C/50", - ] + "Discharge at 5 W for 10 s", + ], + convert_to_format, ) - assert rhs_multi <= rhs_cc + 2, ( - f"adding modes grew the rhs top function: {rhs_cc} -> {rhs_multi}" + _assert_no_always_run_growth( + rhs_cc, rhs_multi, "voltage+power modes in the rhs" ) - assert jac_multi <= jac_cc + 2, ( - f"adding modes grew the jac top function: {jac_cc} -> {jac_multi}" + _assert_no_always_run_growth( + jac_cc, jac_multi, "voltage+power modes in the jac" ) def test_unified_switch_matches_legacy_voltage(self): @@ -698,15 +774,15 @@ def voltage(mode): assert voltage("unified") == pytest.approx(voltage("legacy"), abs=1e-6) - def test_unified_control_row_jacobian_is_sparse(self): - # CasADi declares a Switch's jacobian structurally dense, so the control - # equation would otherwise be a full-bandwidth row and balloon KLU work. The - # solver projects it back onto the true (union-of-branches) sparsity; assert no - # jacobian row spans the full bandwidth. + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_unified_control_row_jacobian_is_sparse(self, convert_to_format): + # The conditional Jacobian must use the union of branch sparsities, not a full row. from collections import Counter + model = pybamm.lithium_ion.DFN() + model.convert_to_format = convert_to_format sim = pybamm.Simulation( - pybamm.lithium_ion.DFN(), + model, experiment=pybamm.Experiment( [("Discharge at 1C until 3.3 V", "Hold at 3.3 V until C/20")] ), @@ -714,9 +790,16 @@ def test_unified_control_row_jacobian_is_sparse(self): experiment_model_mode="unified", ) sim.solve() - sparsity = sim._built_experiment_model.jac_rhs_algebraic_eval.sparsity_out(0) - n = sparsity.size2() - densest_row_nnz = max(Counter(sparsity.get_triplet()[0]).values()) + if convert_to_format == "casadi": + jacobian = sim._built_experiment_model.jac_rhs_algebraic_eval + sparsity = jacobian.sparsity_out(0) + rows = sparsity.get_triplet()[0] + n = sparsity.size2() + else: + model = sim._built_experiment_solver._setup["rust_model"] + _, rows = model.csc_sparsity_pattern() + n = model.rhs.n_states + densest_row_nnz = max(Counter(rows).values()) assert densest_row_nnz <= n // 2, ( f"a jacobian row has {densest_row_nnz}/{n} nonzeros (near full bandwidth); " "the control-row union-sparsity projection did not take effect" @@ -768,27 +851,31 @@ def test_non_unified_model_jacobian_is_plain_casadi(self): model.build_casadi_jacobian(expr, x), casadi.jacobian(expr, x), 20 ) - def test_unified_aot_compile_bounded_in_unique_steps(self): - # -O3 is superlinear in single-function size, so per-branch dispatch must keep - # the largest generated jac function flat as unique steps grow. Deterministic. - def largest_jac_fn_lines(n): + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_unified_aot_compile_bounded_in_unique_steps(self, convert_to_format): + # Compilation must stay flat as equivalent step values are added. + def jacobian_size(n): ops = [f"Discharge at {0.4 + 0.1 * i:.2f}C for 10 s" for i in range(n)] + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = convert_to_format sim = pybamm.Simulation( - pybamm.lithium_ion.SPMe(), + model, experiment=pybamm.Experiment([tuple(ops)]), solver=pybamm.IDAKLUSolver(), experiment_model_mode="unified", ) sim.solve() - return _largest_generated_fn_lines( - sim._built_experiment_model.jac_rhs_algebraic_eval - ) + if convert_to_format == "casadi": + return _largest_generated_fn_lines( + sim._built_experiment_model.jac_rhs_algebraic_eval + ) + return _native_jacobian_instruction_count(sim) - big_2 = largest_jac_fn_lines(2) - big_8 = largest_jac_fn_lines(8) + big_2 = jacobian_size(2) + big_8 = jacobian_size(8) assert big_8 <= big_2 + 50, ( - "largest unified jac function grew with unique steps: " - f"{big_2} -> {big_8} lines (per-branch dispatch should keep it flat)" + "unified Jacobian compilation grew with equivalent steps: " + f"{big_2} -> {big_8}" ) def test_experiment_state_mapper_has_full_state_size_for_2d_current_collector(self): @@ -1216,8 +1303,15 @@ def test_skip_ok(self): sol = sim.solve() assert sol.termination == "Event exceeded in initial conditions" - def test_skip_ok_with_multiple_infeasible_terminations_in_unified_model(self): + @pytest.mark.parametrize( + ("convert_to_format", "solver_class"), + [("casadi", pybamm.CasadiSolver), ("rust", pybamm.IDAKLUSolver)], + ) + def test_skip_ok_with_multiple_infeasible_terminations_in_unified_model( + self, convert_to_format, solver_class + ): model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format experiment = pybamm.Experiment( [ pybamm.step.Current( @@ -1231,7 +1325,7 @@ def test_skip_ok_with_multiple_infeasible_terminations_in_unified_model(self): sim = pybamm.Simulation( model, experiment=experiment, - solver=pybamm.CasadiSolver(), + solver=solver_class(), experiment_model_mode="unified", ) @@ -1304,6 +1398,8 @@ def test_run_experiment_cccv_solvers(self): solutions = {} for name, solver in solvers.items(): model = pybamm.lithium_ion.SPM() + if isinstance(solver, pybamm.CasadiSolver): + model.convert_to_format = "casadi" sim = pybamm.Simulation(model, experiment=experiment_2step, solver=solver) solution = sim.solve() assert solution.t[-1] == pytest.approx(3600 * len(experiment_2step.steps)) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_rust.py b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_rust.py new file mode 100644 index 0000000000..8a3c9a0bff --- /dev/null +++ b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_convert_to_rust.py @@ -0,0 +1,1640 @@ +import itertools + +import numpy as np +import pytest +from numpy.testing import assert_array_almost_equal + + +class TestRustExprGraph: + """Tests for the Rust expression graph via PyO3 bindings.""" + + def test_import(self): + from pybamm.rust import Expr, ExprGraph # noqa: F401 + + def test_scalar(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + s = g.scalar(3.14) + result = g.eval_to_float(s, 0.0, [], [], []) + assert result == pytest.approx(3.14) + + def test_time(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + t = g.time() + result = g.eval_to_float(t, 2.5, [], [], []) + assert result == pytest.approx(2.5) + + def test_state_vector(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + sv = g.state_vector(1, 3) + y = np.array([10.0, 20.0, 30.0, 40.0]) + result = g.eval_to_array(sv, 0.0, y, np.array([]), []) + assert_array_almost_equal(result, [20.0, 30.0]) + + def test_state_vector_dot(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + sv_dot = g.state_vector_dot(0, 2) + y_dot = np.array([100.0, 200.0, 300.0]) + result = g.eval_to_array(sv_dot, 0.0, np.array([]), y_dot, []) + assert_array_almost_equal(result, [100.0, 200.0]) + + def test_input_parameter(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + p = g.input_parameter("C_rate") + result = g.eval_to_float(p, 0.0, [], [], [1.5]) + assert result == pytest.approx(1.5) + + def test_array(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.array(np.array([1.0, 2.0, 3.0])) + result = g.eval_to_array(a, 0.0, np.array([]), np.array([]), []) + assert_array_almost_equal(result, [1.0, 2.0, 3.0]) + + def test_add_scalars(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(2.0) + b = g.scalar(3.0) + c = g.add(a, b) + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(5.0) + + def test_add_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(2.0) + b = g.scalar(3.0) + c = a + b + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(5.0) + + def test_mul_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(4.0) + b = g.scalar(5.0) + c = a * b + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(20.0) + + def test_sub_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(10.0) + b = g.scalar(3.0) + c = a - b + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(7.0) + + def test_neg_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(5.0) + b = -a + result = g.eval_to_float(b, 0.0, [], [], []) + assert result == pytest.approx(-5.0) + + def test_truediv_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(10.0) + b = g.scalar(4.0) + c = a / b + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(2.5) + + def test_pow_dunder(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.scalar(3.0) + b = g.scalar(2.0) + c = a**b + result = g.eval_to_float(c, 0.0, [], [], []) + assert result == pytest.approx(9.0) + + def test_nested_expression(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + two = g.scalar(2.0) + three = g.scalar(3.0) + four = g.scalar(4.0) + one = g.scalar(1.0) + result_expr = (two + three) * four - one + result = g.eval_to_float(result_expr, 0.0, [], [], []) + assert result == pytest.approx(19.0) + + def test_vector_add(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + a = g.array(np.array([1.0, 2.0, 3.0])) + b = g.array(np.array([10.0, 20.0, 30.0])) + c = a + b + result = g.eval_to_array(c, 0.0, np.array([]), np.array([]), []) + assert_array_almost_equal(result, [11.0, 22.0, 33.0]) + + def test_scalar_times_vector(self): + from pybamm.rust import ExprGraph + + g = ExprGraph() + s = g.scalar(2.0) + v = g.array(np.array([1.0, 2.0, 3.0])) + c = s * v + result = g.eval_to_array(c, 0.0, np.array([]), np.array([]), []) + assert_array_almost_equal(result, [2.0, 4.0, 6.0]) + + +class TestSymbolToRust: + """Tests that _to_rust() conversion matches Symbol.evaluate().""" + + def test_scalar_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + s = pybamm.Scalar(3.14) + g = ExprGraph() + rust_symbols = {} + expr = s.to_rust(g, rust_symbols) + result = g.eval_to_float(expr, 0.0, [], [], []) + assert result == pytest.approx(s.evaluate()) + + def test_array_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + a = pybamm.Vector(np.array([1.0, 2.0, 3.0])) + g = ExprGraph() + rust_symbols = {} + expr = a.to_rust(g, rust_symbols) + result = g.eval_to_array(expr, 0.0, np.array([]), np.array([]), []) + expected = a.evaluate() + assert_array_almost_equal(result, expected.flatten()) + + def test_addition_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + a = pybamm.Scalar(2.0) + b = pybamm.Scalar(3.0) + expr_pybamm = a + b + g = ExprGraph() + rust_symbols = {} + expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(expr, 0.0, [], [], []) + assert result == pytest.approx(expr_pybamm.evaluate()) + + def test_nested_arithmetic_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + a = pybamm.Scalar(2.0) + b = pybamm.Scalar(3.0) + c = pybamm.Scalar(4.0) + expr_pybamm = (a + b) * c - pybamm.Scalar(1.0) + g = ExprGraph() + rust_symbols = {} + expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(expr, 0.0, [], [], []) + assert result == pytest.approx(expr_pybamm.evaluate()) + + def test_state_vector_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + sv = pybamm.StateVector(slice(1, 4)) + y = np.array([0.0, 10.0, 20.0, 30.0, 40.0]) + g = ExprGraph() + rust_symbols = {} + expr = sv.to_rust(g, rust_symbols) + result = g.eval_to_array(expr, 0.0, y, np.array([]), []) + expected = sv.evaluate(y=y) + assert_array_almost_equal(result, expected.flatten()) + + def test_caching(self): + import pybamm + from pybamm.rust import ExprGraph + + s = pybamm.Scalar(3.14) + g = ExprGraph() + rust_symbols = {} + expr1 = s.to_rust(g, rust_symbols) + expr2 = s.to_rust(g, rust_symbols) + assert expr1.id == expr2.id + + def test_negation_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + a = pybamm.Scalar(5.0) + expr_pybamm = -a + g = ExprGraph() + rust_symbols = {} + expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(expr, 0.0, [], [], []) + assert result == pytest.approx(expr_pybamm.evaluate()) + + def test_abs_to_rust(self): + import pybamm + from pybamm.rust import ExprGraph + + a = pybamm.Scalar(-5.0) + expr_pybamm = pybamm.AbsoluteValue(a) + g = ExprGraph() + rust_symbols = {} + expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(expr, 0.0, [], [], []) + assert result == pytest.approx(expr_pybamm.evaluate()) + + +class TestSpecificFunctionToRust: + """Tests that SpecificFunction._to_rust() conversions match Symbol.evaluate().""" + + def _eval_scalar(self, pybamm_expr, g=None): + """Helper: convert pybamm_expr to Rust and eval as float.""" + from pybamm.rust import ExprGraph + + if g is None: + g = ExprGraph() + rust_symbols = {} + expr = pybamm_expr.to_rust(g, rust_symbols) + return g.eval_to_float(expr, 0.0, [], [], []) + + def test_sqrt_scalar(self): + import pybamm + + expr = pybamm.sqrt(pybamm.Scalar(9.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_sqrt_vector(self): + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([1.0, 4.0, 9.0, 16.0])) + expr_pybamm = pybamm.sqrt(v) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + expected = expr_pybamm.evaluate().flatten() + assert_array_almost_equal(result, expected) + + def test_exp_scalar(self): + import pybamm + + expr = pybamm.exp(pybamm.Scalar(2.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_log_scalar(self): + import pybamm + + expr = pybamm.log(pybamm.Scalar(np.e)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_sin_scalar(self): + import pybamm + + expr = pybamm.sin(pybamm.Scalar(np.pi / 6)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_cos_scalar(self): + import pybamm + + expr = pybamm.cos(pybamm.Scalar(np.pi / 3)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_tanh_scalar(self): + import pybamm + + expr = pybamm.tanh(pybamm.Scalar(1.5)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_sinh_scalar(self): + import pybamm + + expr = pybamm.sinh(pybamm.Scalar(1.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_cosh_scalar(self): + import pybamm + + expr = pybamm.cosh(pybamm.Scalar(1.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_arcsinh_scalar(self): + import pybamm + + expr = pybamm.arcsinh(pybamm.Scalar(2.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_arctan_scalar(self): + import pybamm + from pybamm.rust import ExprGraph + + # Use InputParameter to prevent constant folding + p = pybamm.InputParameter("x") + expr = pybamm.arctan(p) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + result = g.eval_to_float(rust_expr, 0.0, [], [], [1.0]) + expected = float(expr.evaluate(inputs={"x": 1.0})) + assert result == pytest.approx(expected) + + def test_erf_scalar(self): + import pybamm + from pybamm.rust import ExprGraph + + # Use InputParameter to prevent constant folding + p = pybamm.InputParameter("x") + expr = pybamm.erf(p) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + result = g.eval_to_float(rust_expr, 0.0, [], [], [1.0]) + expected = float(expr.evaluate(inputs={"x": 1.0})) + # erf approximation has ~1.5e-7 max error + assert result == pytest.approx(expected, rel=1e-5) + + def test_sign_positive(self): + import pybamm + + expr = pybamm.sign(pybamm.Scalar(3.5)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_sign_negative(self): + import pybamm + + expr = pybamm.sign(pybamm.Scalar(-7.0)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_floor_scalar(self): + import pybamm + + expr = pybamm.Floor(pybamm.Scalar(3.7)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_ceiling_scalar(self): + import pybamm + + expr = pybamm.Ceiling(pybamm.Scalar(3.2)) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + def test_nested_function(self): + """Test nested: exp(sqrt(x)) for x=4.0 should give e^2.""" + import pybamm + + expr = pybamm.exp(pybamm.sqrt(pybamm.Scalar(4.0))) + result = self._eval_scalar(expr) + assert result == pytest.approx(float(expr.evaluate())) + + @pytest.mark.parametrize( + "x_val,exponent", + [ + (4.0, 0.5), + (-4.0, 0.5), + (1.5, 1.5), + (1e-12, 0.5), # near zero, regularisation regime dominates + (0.0, 0.5), # exactly zero + # One case per chain in `_positive_base_pow_chain`, which keys off the + # inner exponent (a-1)/2 in {-0.5, -0.25, 0.25, 0.5, 1}, plus a + # non-chain exponent that falls back to runtime pow. + (4.0, 0.0), + (-2.0, 0.0), + (4.0, 1.5), + (4.0, 2.0), + (-3.0, 2.0), + (4.0, 3.0), + (4.0, 0.7), + (-4.0, 0.7), + ], + ) + def test_reg_power_scalar(self, x_val, exponent): + """RegPower(x, a) on the Rust path matches the numpy evaluator.""" + import pybamm + from pybamm.rust import ExprGraph + + p = pybamm.InputParameter("x") + expr = pybamm.reg_power(p, exponent) + g = ExprGraph() + rust_expr = expr.to_rust(g, {}) + result = g.eval_to_float(rust_expr, 0.0, [], [], [x_val]) + expected = float(expr.evaluate(inputs={"x": x_val})) + assert result == pytest.approx(expected, rel=1e-12, abs=1e-15) + + def test_reg_power_with_scale(self): + """RegPower with explicit non-unit scale matches the numpy evaluator.""" + import pybamm + from pybamm.rust import ExprGraph + + p = pybamm.InputParameter("x") + expr = pybamm.reg_power(p, 0.7, scale=2.5) + g = ExprGraph() + rust_expr = expr.to_rust(g, {}) + for x_val in (-3.0, -0.1, 1e-9, 0.0, 0.4, 5.0): + result = g.eval_to_float(rust_expr, 0.0, [], [], [x_val]) + expected = float(expr.evaluate(inputs={"x": x_val})) + assert result == pytest.approx(expected, rel=1e-12, abs=1e-14) + + +class TestIndexToRust: + """Tests for Index (slicing) node conversion to Rust.""" + + def test_index_vector_slice(self): + """Index a Vector with slice(1, 3) → elements at positions 1 and 2.""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([10.0, 20.0, 30.0, 40.0, 50.0])) + expr_pybamm = pybamm.Index(v, slice(1, 3)) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + expected = expr_pybamm.evaluate().flatten() + assert_array_almost_equal(result, expected) + + def test_index_state_vector_slice(self): + """Index a StateVector with slice(2, 4) using specific y values.""" + import pybamm + from pybamm.rust import ExprGraph + + sv = pybamm.StateVector(slice(0, 6)) + expr_pybamm = pybamm.Index(sv, slice(2, 4)) + y = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, y, np.array([]), []) + expected = expr_pybamm.evaluate(y=y.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected) + + def test_index_negative_last_element(self): + """Index(vector, -1) → last element, via slice(-1, None).""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([10.0, 20.0, 30.0, 40.0, 50.0])) + expr_pybamm = pybamm.Index(v, -1) + g = ExprGraph() + rust_expr = expr_pybamm.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + expected = expr_pybamm.evaluate().flatten() # [50.0] + assert_array_almost_equal(result, expected) + + def test_index_step_raises(self): + """A strided Index (step != 1) must raise, not silently return a + contiguous slice.""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([10.0, 20.0, 30.0, 40.0, 50.0])) + expr_pybamm = pybamm.Index(v, slice(0, 4, 2)) + g = ExprGraph() + with pytest.raises(NotImplementedError, match="step"): + expr_pybamm.to_rust(g, {}) + + def test_index_over_negative_start_clamps(self): + """A start more negative than -size clamps to 0 (numpy semantics).""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([10.0, 20.0, 30.0, 40.0, 50.0])) + expr_pybamm = pybamm.Index(v, slice(-7, 5)) + g = ExprGraph() + rust_expr = expr_pybamm.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + expected = expr_pybamm.evaluate().flatten() # full [10,20,30,40,50] + assert_array_almost_equal(result, expected) + + +class TestNewBinaryOperatorsToRust: + """Tests for new binary operators: Minimum, Maximum, Modulo, Hypot, + EqualHeaviside, NotEqualHeaviside, Equality, and sparse MatMul.""" + + def _eval_scalar(self, pybamm_expr, inputs=None): + from pybamm.rust import ExprGraph + + g = ExprGraph() + rust_symbols = {} + expr = pybamm_expr.to_rust(g, rust_symbols) + inputs_list = list(inputs.values()) if inputs else [] + return g.eval_to_float(expr, 0.0, [], [], inputs_list) + + def _eval_array(self, pybamm_expr, y=None, inputs=None): + from pybamm.rust import ExprGraph + + g = ExprGraph() + rust_symbols = {} + expr = pybamm_expr.to_rust(g, rust_symbols) + y_arr = np.array(y) if y is not None else np.array([]) + inputs_list = list(inputs.values()) if inputs else [] + return g.eval_to_array(expr, 0.0, y_arr, np.array([]), inputs_list) + + def test_minimum_scalars(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.minimum(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 5.0}) + expected = float(expr.evaluate(inputs={"a": 3.0, "b": 5.0})) + assert result == pytest.approx(expected) + + def test_minimum_reversed(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.minimum(a, b) + result = self._eval_scalar(expr, inputs={"a": 7.0, "b": 2.0}) + expected = float(expr.evaluate(inputs={"a": 7.0, "b": 2.0})) + assert result == pytest.approx(expected) + + def test_maximum_scalars(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.maximum(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 5.0}) + expected = float(expr.evaluate(inputs={"a": 3.0, "b": 5.0})) + assert result == pytest.approx(expected) + + def test_maximum_reversed(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.maximum(a, b) + result = self._eval_scalar(expr, inputs={"a": 7.0, "b": 2.0}) + expected = float(expr.evaluate(inputs={"a": 7.0, "b": 2.0})) + assert result == pytest.approx(expected) + + def test_modulo_scalar(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.Modulo(a, b) + result = self._eval_scalar(expr, inputs={"a": 10.0, "b": 3.0}) + expected = float(expr.evaluate(inputs={"a": 10.0, "b": 3.0})) + assert result == pytest.approx(expected) + + def test_hypot_scalar(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.hypot(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 4.0}) + expected = float(expr.evaluate(inputs={"a": 3.0, "b": 4.0})) + assert result == pytest.approx(expected) + + def test_equal_heaviside_true(self): + """3.0 <= 5.0 should return 1.""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.EqualHeaviside(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 5.0}) + assert result == pytest.approx(1.0) + + def test_equal_heaviside_equal(self): + """5.0 <= 5.0 should return 1 (equal case).""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.EqualHeaviside(a, b) + result = self._eval_scalar(expr, inputs={"a": 5.0, "b": 5.0}) + assert result == pytest.approx(1.0) + + def test_equal_heaviside_false(self): + """7.0 <= 5.0 should return 0.""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.EqualHeaviside(a, b) + result = self._eval_scalar(expr, inputs={"a": 7.0, "b": 5.0}) + assert result == pytest.approx(0.0) + + def test_not_equal_heaviside_true(self): + """3.0 < 5.0 should return 1.""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.NotEqualHeaviside(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 5.0}) + assert result == pytest.approx(1.0) + + def test_not_equal_heaviside_equal(self): + """5.0 < 5.0 should return 0 (strict inequality).""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.NotEqualHeaviside(a, b) + result = self._eval_scalar(expr, inputs={"a": 5.0, "b": 5.0}) + assert result == pytest.approx(0.0) + + def test_equality_equal(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.Equality(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 3.0}) + assert result == pytest.approx(1.0) + + def test_equality_not_equal(self): + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.Equality(a, b) + result = self._eval_scalar(expr, inputs={"a": 3.0, "b": 4.0}) + assert result == pytest.approx(0.0) + + def test_sparse_matrix_matmul(self): + """Test sparse Matrix @ StateVector conversion to Rust.""" + from scipy.sparse import csr_matrix as scipy_csr + + import pybamm + from pybamm.rust import ExprGraph + + # 2x3 sparse matrix: [1 0 0; 0 0 2] + data = np.array([1.0, 2.0]) + row = np.array([0, 1]) + col = np.array([0, 2]) + sparse_mat = scipy_csr((data, (row, col)), shape=(2, 3)) + + mat = pybamm.Matrix(sparse_mat) + sv = pybamm.StateVector(slice(0, 3)) + expr = mat @ sv + + y = np.array([10.0, 20.0, 30.0]) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, y, np.array([]), []) + expected = expr.evaluate(y=y.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected) + + +class TestInterpolantToRust: + """Tests for Interpolant._to_rust() — 1D linear, cubic, and pchip interpolation.""" + + def test_1d_linear(self): + """1D linear interpolant: y=2x on [0,1], evaluate at x=0.4 → ~0.8.""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.linspace(0, 1, 50) + y = 2 * x + sv = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant(x, y, sv) + + y_test = np.array([0.4]) + g = ExprGraph() + rust_symbols = {} + rust_expr = interp.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, y_test, np.array([]), []) + expected = interp.evaluate(y=y_test.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=5) + + def test_1d_linear_multiple_values(self): + """1D linear interpolant: y=2x on [0,1], evaluate at [0.4, 0.6] → [0.8, 1.2].""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.linspace(0, 1, 50) + y = 2 * x + sv = pybamm.StateVector(slice(0, 2)) + interp = pybamm.Interpolant(x, y, sv) + + y_test = np.array([0.4, 0.6]) + g = ExprGraph() + rust_symbols = {} + rust_expr = interp.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, y_test, np.array([]), []) + expected = interp.evaluate(y=y_test.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=5) + + def test_1d_linear_extrapolation_matches_evaluate(self): + """Out-of-domain linear interp extends, matching Symbol.evaluate().""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.linspace(0.0, 1.0, 11) + y = 3.0 * x + 1.0 # slope 3 + sv = pybamm.StateVector(slice(0, 2)) + interp = pybamm.Interpolant(x, y, sv) # default linear, extrapolate=True + + y_test = np.array([-0.5, 1.5]) # both out of [0, 1] + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, y_test, np.array([]), []) + expected = interp.evaluate(y=y_test.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + def _check_interp_parity(self, interpolator): + import pybamm + from pybamm.rust import ExprGraph + + # Non-uniform grid so coefficients are non-trivial. + x = np.array([0.0, 0.5, 1.0, 2.0, 3.5, 5.0]) + y = np.array([0.0, 0.4, 1.1, 3.9, 12.0, 24.0]) + sv = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant(x, y, sv, interpolator=interpolator) + + g = ExprGraph() + rust_symbols = {} + rust_expr = interp.to_rust(g, rust_symbols) + # In-domain, at breakpoints, and out-of-domain (extend) on both sides. + for q in [-1.0, 0.0, 0.5, 1.3, 2.0, 3.5, 5.0, 7.5]: + yq = np.array([q]) + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = interp.evaluate(y=yq.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + def test_1d_cubic_matches_evaluate(self): + self._check_interp_parity("cubic") + + def test_1d_pchip_matches_evaluate(self): + self._check_interp_parity("pchip") + + def test_1d_cubic_vector_child(self): + """Element-wise over a vector child matches evaluate().""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.linspace(0.0, 4.0, 20) + y = np.sin(x) + sv = pybamm.StateVector(slice(0, 3)) + interp = pybamm.Interpolant(x, y, sv, interpolator="cubic") + + y_test = np.array([0.7, 2.1, 5.0]) # last is out-of-domain (extend) + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, y_test, np.array([]), []) + expected = interp.evaluate(y=y_test.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic", "pchip"]) + def test_1d_extrapolate_false_extends(self, interpolator): + """extrapolate=False is intentionally ignored on the Rust path (parity + spec Decision 3): in-domain matches evaluate(), out-of-domain extends + the boundary polynomial instead of returning NaN (a NaN would poison + the solver residual; the domain is guarded by extrapolation events).""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.array([0.0, 0.5, 1.0, 2.0, 3.5, 5.0]) + y = np.array([0.0, 0.4, 1.1, 3.9, 12.0, 24.0]) + sv = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant( + x, y, sv, interpolator=interpolator, extrapolate=False + ) + extend = pybamm.Interpolant(x, y, sv, interpolator=interpolator) + + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + # In-domain: parity with evaluate(). + yq = np.array([1.3]) + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = interp.evaluate(y=yq.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + # Out-of-domain: evaluate() gives NaN, Rust extends (extrapolate=True). + for q in [-1.0, 7.5]: + yq = np.array([q]) + assert np.isnan(interp.evaluate(y=yq.reshape(-1, 1))).all() + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = extend.evaluate(y=yq.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic", "pchip"]) + def test_1d_vector_valued_y_matches_evaluate(self, interpolator): + """Vector-valued y (y.ndim == 2) stacks one interpolant per column in + evaluate()'s column order.""" + import pybamm + from pybamm.rust import ExprGraph + + # Non-uniform grid and distinct nonlinear columns so a column swap or + # wrong-length table cannot cancel out. + x = np.array([0.0, 0.5, 1.0, 2.0, 3.5, 5.0]) + y = np.column_stack([np.sin(x), np.cos(x), x**2]) + sv = pybamm.StateVector(slice(0, 1)) # scalar child (constructor requires) + interp = pybamm.Interpolant(x, y, sv, interpolator=interpolator) + + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + # In-domain, at breakpoints, and out-of-domain (extend) on both sides. + for q in [-1.0, 0.0, 0.5, 1.3, 2.0, 3.5, 5.0, 7.5]: + yq = np.array([q]) + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = interp.evaluate(y=yq.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic", "pchip"]) + def test_1d_vector_valued_y_jacobian_matches_casadi(self, interpolator): + """d/dy of the stacked columns matches the CasADi conversion.""" + import casadi + + import pybamm + from pybamm.rust import ExprGraph + + x = np.array([0.0, 0.5, 1.0, 2.0, 3.5, 5.0]) + y = np.column_stack([np.sin(x), np.cos(x), x**2]) + sv = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant(x, y, sv, interpolator=interpolator) + + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + f = g.compile(rust_expr, n_states=1) + yq = np.array([1.3]) # strictly inside a segment (linear has knot kinks) + rust_jac = f.jacobian("y")(0.0, yq, np.array([])).toarray() + + ys = casadi.MX.sym("y", 1) + casadi_jac = casadi.Function( + "J", [ys], [casadi.jacobian(interp.to_casadi(y=ys), ys)] + ) + np.testing.assert_allclose( + rust_jac, np.array(casadi_jac(yq)), rtol=1e-8, atol=1e-10 + ) + + def test_1d_y_with_more_than_two_dims_raises(self): + """y.ndim > 2 has no defined column order, so refuse loudly.""" + import pybamm + from pybamm.rust import ExprGraph + + x = np.linspace(0.0, 1.0, 5) + y = np.ones((5, 2, 2)) + sv = pybamm.StateVector(slice(0, 1)) + interp = pybamm.Interpolant(x, y, sv) + with pytest.raises(NotImplementedError, match=r"more than two dimensions"): + interp.to_rust(ExprGraph(), {}) + + +class TestInterpolantNDToRust: + """2D/3D interpolants lower to the Rust ND tensor-product node and match + Symbol.evaluate() (RegularGridInterpolator) in-domain, at every grid + point, and out-of-domain (extend), including all-axes-out corners.""" + + # Non-uniform grids so coefficients and cell strides are non-trivial. + x0 = np.array([0.0, 0.5, 1.3, 2.0, 3.1, 4.0]) + x1 = np.array([-1.0, -0.3, 0.4, 1.0, 2.2]) + x2 = np.array([0.0, 1.0, 2.5, 3.5, 5.0]) + + @staticmethod + def _pts(axes): + rng = np.random.default_rng(7) + ins = np.column_stack([rng.uniform(a[0], a[-1], 12) for a in axes]) + grid = np.array(list(itertools.product(*axes))) + mid = [0.5 * (a[0] + a[-1]) for a in axes] + outs = [] + for i, a in enumerate(axes): + lo = list(mid) + lo[i] = a[0] - 0.7 + outs.append(lo) + hi = list(mid) + hi[i] = a[-1] + 0.9 + outs.append(hi) + outs.append([a[0] - 0.8 for a in axes]) # all-axes-out corner + outs.append([a[-1] + 1.1 for a in axes]) # all-axes-out corner + return np.vstack([ins, grid, np.asarray(outs)]) + + def _check_parity(self, x, y, interpolator): + import pybamm + from pybamm.rust import ExprGraph + + ndim = len(x) + svs = tuple(pybamm.StateVector(slice(i, i + 1)) for i in range(ndim)) + interp = pybamm.Interpolant(x, y, svs, interpolator=interpolator) + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + for p in self._pts(list(x)): + yq = np.asarray(p, dtype=float) + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = np.asarray(interp.evaluate(y=yq.reshape(-1, 1))).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic"]) + def test_2d_matches_evaluate(self, interpolator): + X0, X1 = np.meshgrid(self.x0, self.x1, indexing="ij") + y = np.sin(X0) * np.exp(-0.3 * X1) + 0.1 * X0 * X1 + self._check_parity((self.x0, self.x1), y, interpolator) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic"]) + def test_3d_matches_evaluate(self, interpolator): + X0, X1, X2 = np.meshgrid(self.x0, self.x1, self.x2, indexing="ij") + y = np.cos(X0) * X1 + 0.5 * np.sqrt(1 + X2) + 0.05 * X0 * X1 * X2 + self._check_parity((self.x0, self.x1, self.x2), y, interpolator) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic"]) + def test_2d_vector_children(self, interpolator): + """Element-wise over equal-length vector children matches evaluate().""" + import pybamm + from pybamm.rust import ExprGraph + + X0, X1 = np.meshgrid(self.x0, self.x1, indexing="ij") + y = X0**2 + X1 + sv0 = pybamm.StateVector(slice(0, 3)) + sv1 = pybamm.StateVector(slice(3, 6)) + interp = pybamm.Interpolant( + (self.x0, self.x1), y, (sv0, sv1), interpolator=interpolator + ) + # Last pair is out-of-domain on both axes (extend). + y_test = np.array([0.7, 2.1, 5.0, 0.0, 1.5, -1.4]) + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, y_test, np.array([]), []) + expected = np.asarray(interp.evaluate(y=y_test.reshape(-1, 1))).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + @pytest.mark.parametrize("interpolator", ["linear", "cubic"]) + def test_2d_extrapolate_false_extends(self, interpolator): + """extrapolate=False is intentionally ignored on the Rust path (parity + spec Decision 3): in-domain matches evaluate(), out-of-domain extends + instead of returning NaN.""" + import pybamm + from pybamm.rust import ExprGraph + + X0, X1 = np.meshgrid(self.x0, self.x1, indexing="ij") + y = np.sin(X0) * np.exp(-0.3 * X1) + 0.1 * X0 * X1 + svs = (pybamm.StateVector(slice(0, 1)), pybamm.StateVector(slice(1, 2))) + interp = pybamm.Interpolant( + (self.x0, self.x1), y, svs, interpolator=interpolator, extrapolate=False + ) + extend = pybamm.Interpolant( + (self.x0, self.x1), y, svs, interpolator=interpolator + ) + + g = ExprGraph() + rust_expr = interp.to_rust(g, {}) + # In-domain: parity with evaluate(). + yq = np.array([1.3, 0.5]) + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = np.asarray(interp.evaluate(y=yq.reshape(-1, 1))).flatten() + assert_array_almost_equal(result, expected, decimal=10) + # Out-of-domain: evaluate() gives NaN, Rust extends (extrapolate=True). + for p in [[-0.7, 0.5], [4.9, 3.1]]: + yq = np.asarray(p) + assert np.isnan(interp.evaluate(y=yq.reshape(-1, 1))).all() + result = g.eval_to_array(rust_expr, 0.0, yq, np.array([]), []) + expected = np.asarray(extend.evaluate(y=yq.reshape(-1, 1))).flatten() + assert_array_almost_equal(result, expected, decimal=10) + + +class TestReductionToRust: + """Tests for MaxReduce and MinReduce node conversion to Rust.""" + + def test_max_reduce(self): + """pybamm.max(Vector([1, 5, 3])) should reduce to 5.0.""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([1.0, 5.0, 3.0])) + expr_pybamm = pybamm.max(v) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(rust_expr, 0.0, [], [], []) + assert result == pytest.approx(5.0) + + def test_min_reduce(self): + """pybamm.min(Vector([1, 5, 3])) should reduce to 1.0.""" + import pybamm + from pybamm.rust import ExprGraph + + v = pybamm.Vector(np.array([1.0, 5.0, 3.0])) + expr_pybamm = pybamm.min(v) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_float(rust_expr, 0.0, [], [], []) + assert result == pytest.approx(1.0) + + +class TestConcatenationToRust: + """Tests for concatenation node conversion to Rust.""" + + def test_numpy_concatenation(self): + """NumpyConcatenation of Vector([1,2]) and Vector([3,4,5]) → [1,2,3,4,5].""" + import pybamm + from pybamm.rust import ExprGraph + + v1 = pybamm.Vector(np.array([1.0, 2.0])) + v2 = pybamm.Vector(np.array([3.0, 4.0, 5.0])) + expr_pybamm = pybamm.NumpyConcatenation(v1, v2) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + expected = expr_pybamm.evaluate().flatten() + assert_array_almost_equal(result, expected) + + def test_state_vector_concatenation(self): + """NumpyConcatenation of StateVector(0:3) and StateVector(3:5), y=[1,2,3,4,5].""" + import pybamm + from pybamm.rust import ExprGraph + + sv1 = pybamm.StateVector(slice(0, 3)) + sv2 = pybamm.StateVector(slice(3, 5)) + expr_pybamm = pybamm.NumpyConcatenation(sv1, sv2) + y = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + g = ExprGraph() + rust_symbols = {} + rust_expr = expr_pybamm.to_rust(g, rust_symbols) + result = g.eval_to_array(rust_expr, 0.0, y, np.array([]), []) + expected = expr_pybamm.evaluate(y=y.reshape(-1, 1)).flatten() + assert_array_almost_equal(result, expected) + + +class TestVectorFieldToRust: + """Tests for vector field conversion to Rust, which stacks its components.""" + + def test_vector_field(self): + """VectorField of Vector([1,2]) and Vector([3,4]) → [1,2,3,4].""" + import pybamm + from pybamm.rust import ExprGraph + + v1 = pybamm.Vector(np.array([1.0, 2.0])) + v2 = pybamm.Vector(np.array([3.0, 4.0])) + expr_pybamm = pybamm.VectorField(v1, v2) + g = ExprGraph() + rust_expr = expr_pybamm.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, np.array([]), np.array([]), []) + assert_array_almost_equal(result, [1.0, 2.0, 3.0, 4.0]) + + def test_vector_field_matches_casadi(self): + """Three-component state-dependent field matches the CasADi conversion.""" + import casadi + + import pybamm + from pybamm.rust import ExprGraph + + components = [ + pybamm.StateVector(slice(0, 2)), + 2 * pybamm.StateVector(slice(2, 4)), + pybamm.StateVector(slice(0, 2)) + pybamm.StateVector(slice(2, 4)), + ] + expr_pybamm = pybamm.VectorField(*components) + y = np.array([1.0, 2.0, 3.0, 4.0]) + + g = ExprGraph() + rust_expr = expr_pybamm.to_rust(g, {}) + result = g.eval_to_array(rust_expr, 0.0, y, np.array([]), []) + + t_sym = casadi.MX.sym("t") + y_sym = casadi.MX.sym("y", 4) + y_dot_sym = casadi.MX.sym("y_dot", 4) + casadi_func = casadi.Function( + "f", [t_sym, y_sym], [expr_pybamm.to_casadi(t_sym, y_sym, y_dot_sym, {})] + ) + expected = np.array(casadi_func(0.0, y)).flatten() + assert_array_almost_equal(result, expected) + + +class TestRustMatchesPyBaMM: + """Cross-validation: _to_rust().eval() matches symbol.evaluate() for all supported node types.""" + + def _assert_rust_eval_matches(self, expr, t=0.0, y=None, y_dot=None, inputs=None): + from pybamm.rust import ExprGraph + + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + + # Always flatten y/y_dot to 1D float64 for the Rust bindings + y_np = ( + np.asarray(y, dtype=np.float64).ravel() + if y is not None + else np.array([], dtype=np.float64) + ) + y_dot_np = ( + np.asarray(y_dot, dtype=np.float64).ravel() + if y_dot is not None + else np.array([], dtype=np.float64) + ) + inputs_list = inputs if inputs is not None else [] + pybamm_result = expr.evaluate(t=t, y=y, y_dot=y_dot) + + if np.isscalar(pybamm_result) or ( + hasattr(pybamm_result, "shape") and pybamm_result.size == 1 + ): + result = g.eval_to_float( + rust_expr, + t, + y_np.tolist(), + y_dot_np.tolist(), + inputs_list, + ) + assert result == pytest.approx(float(pybamm_result), rel=1e-5, abs=1e-10) + else: + result = g.eval_to_array(rust_expr, t, y_np, y_dot_np, inputs_list) + np.testing.assert_allclose( + result, np.asarray(pybamm_result).flatten(), rtol=1e-5, atol=1e-10 + ) + + def test_scalar_arithmetic(self): + """Test +, -, *, /, **, neg, abs on scalars.""" + import pybamm + + a = pybamm.Scalar(6.0) + b = pybamm.Scalar(4.0) + + self._assert_rust_eval_matches(a + b) + self._assert_rust_eval_matches(a - b) + self._assert_rust_eval_matches(a * b) + self._assert_rust_eval_matches(a / b) + self._assert_rust_eval_matches(a**b) + self._assert_rust_eval_matches(-a) + self._assert_rust_eval_matches(pybamm.AbsoluteValue(pybamm.Scalar(-7.5))) + + def test_special_functions(self): + """Test sqrt, exp, log, sin, cos, tanh, sinh, cosh, arcsinh, erf.""" + import pybamm + + # Use InputParameter to prevent constant folding where needed + p = pybamm.InputParameter("x") + + self._assert_rust_eval_matches(pybamm.sqrt(pybamm.Scalar(9.0))) + self._assert_rust_eval_matches(pybamm.exp(pybamm.Scalar(2.0))) + self._assert_rust_eval_matches(pybamm.log(pybamm.Scalar(np.e))) + self._assert_rust_eval_matches(pybamm.sin(pybamm.Scalar(np.pi / 6))) + self._assert_rust_eval_matches(pybamm.cos(pybamm.Scalar(np.pi / 3))) + self._assert_rust_eval_matches(pybamm.tanh(pybamm.Scalar(1.5))) + self._assert_rust_eval_matches(pybamm.sinh(pybamm.Scalar(1.0))) + self._assert_rust_eval_matches(pybamm.cosh(pybamm.Scalar(1.0))) + self._assert_rust_eval_matches(pybamm.arcsinh(pybamm.Scalar(2.0))) + + # erf via InputParameter (approximation has ~1.5e-7 max error, rel=1e-5 handles it) + from pybamm.rust import ExprGraph + + expr_erf = pybamm.erf(p) + g = ExprGraph() + rust_symbols = {} + re = expr_erf.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [1.0]) + pybamm_result = float(expr_erf.evaluate(inputs={"x": 1.0})) + assert rust_result == pytest.approx(pybamm_result, rel=1e-5) + + def test_floor_ceiling(self): + """Test Floor and Ceiling on 3.3.""" + import pybamm + + # Use InputParameter to avoid constant folding + p = pybamm.InputParameter("x") + expr_floor = pybamm.Floor(p) + expr_ceil = pybamm.Ceiling(p) + + from pybamm.rust import ExprGraph + + for expr in (expr_floor, expr_ceil): + g = ExprGraph() + rust_symbols = {} + re = expr.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [3.3]) + pybamm_result = float(expr.evaluate(inputs={"x": 3.3})) + assert rust_result == pytest.approx(pybamm_result) + + def test_sign(self): + """Test sign on positive and negative values.""" + import pybamm + + # Use InputParameter to prevent constant folding + p = pybamm.InputParameter("x") + expr = pybamm.sign(p) + + from pybamm.rust import ExprGraph + + # Positive value: sign(5.0) == 1 + g = ExprGraph() + rust_symbols = {} + re = expr.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [5.0]) + pybamm_result = float(expr.evaluate(inputs={"x": 5.0})) + assert rust_result == pytest.approx(pybamm_result) + + # Negative value: sign(-3.0) == -1 + g2 = ExprGraph() + rust_symbols2 = {} + re2 = expr.to_rust(g2, rust_symbols2) + rust_result2 = g2.eval_to_float(re2, 0.0, [], [], [-3.0]) + pybamm_result2 = float(expr.evaluate(inputs={"x": -3.0})) + assert rust_result2 == pytest.approx(pybamm_result2) + + # Zero: sign(0.0) == 0.0 (matches numpy convention) + g3 = ExprGraph() + rust_symbols3 = {} + re3 = expr.to_rust(g3, rust_symbols3) + rust_zero = g3.eval_to_float(re3, 0.0, [], [], [0.0]) + pybamm_zero = float(expr.evaluate(inputs={"x": 0.0})) + assert rust_zero == pytest.approx(pybamm_zero) + + def test_modulo_minimum_maximum(self): + """Test Modulo(7,3), Minimum(7,3), Maximum(7,3).""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + + from pybamm.rust import ExprGraph + + for expr_factory, expected in [ + (pybamm.Modulo, 1.0), + (pybamm.minimum, 3.0), + (pybamm.maximum, 7.0), + ]: + expr = expr_factory(a, b) + g = ExprGraph() + rust_symbols = {} + re = expr.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [7.0, 3.0]) + pybamm_result = float(expr.evaluate(inputs={"a": 7.0, "b": 3.0})) + assert rust_result == pytest.approx(pybamm_result) + assert rust_result == pytest.approx(expected) + + def test_hypot(self): + """Test Hypot(3,4) == 5.""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + expr = pybamm.hypot(a, b) + + from pybamm.rust import ExprGraph + + g = ExprGraph() + rust_symbols = {} + re = expr.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [3.0, 4.0]) + pybamm_result = float(expr.evaluate(inputs={"a": 3.0, "b": 4.0})) + assert rust_result == pytest.approx(pybamm_result) + assert rust_result == pytest.approx(5.0) + + def test_state_vector_arithmetic(self): + """2*sv + 1 evaluated at y=[1,2,3].""" + import pybamm + + sv = pybamm.StateVector(slice(0, 3)) + expr = pybamm.Scalar(2.0) * sv + pybamm.Scalar(1.0) + y = np.array([1.0, 2.0, 3.0]) + self._assert_rust_eval_matches(expr, y=y.reshape(-1, 1)) + + def test_multi_slice_state_vector_and_dot(self): + """Multi-slice StateVector/StateVectorDot lower as concat of slices.""" + import pybamm + + sv = pybamm.StateVector(slice(0, 2), slice(4, 6)) + y = np.arange(6.0) + self._assert_rust_eval_matches(sv, y=y.reshape(-1, 1)) + + sv_dot = pybamm.StateVectorDot(slice(0, 2), slice(4, 6)) + y_dot = np.arange(10.0, 16.0) + self._assert_rust_eval_matches(sv_dot, y_dot=y_dot.reshape(-1, 1)) + + def test_heaviside(self): + """Test EqualHeaviside(2,3) and NotEqualHeaviside(3,3).""" + import pybamm + + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + + from pybamm.rust import ExprGraph + + # EqualHeaviside(2, 3): 2 <= 3 → 1 (True) + expr_eh = pybamm.EqualHeaviside(a, b) + g = ExprGraph() + rust_symbols = {} + re = expr_eh.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(re, 0.0, [], [], [2.0, 3.0]) + pybamm_result = float(expr_eh.evaluate(inputs={"a": 2.0, "b": 3.0})) + assert rust_result == pytest.approx(pybamm_result) + assert rust_result == pytest.approx(1.0) + + # NotEqualHeaviside(3, 3): 3 < 3 → 0 (False, strict inequality) + expr_neh = pybamm.NotEqualHeaviside(a, b) + g2 = ExprGraph() + rust_symbols2 = {} + re2 = expr_neh.to_rust(g2, rust_symbols2) + rust_result2 = g2.eval_to_float(re2, 0.0, [], [], [3.0, 3.0]) + pybamm_result2 = float(expr_neh.evaluate(inputs={"a": 3.0, "b": 3.0})) + assert rust_result2 == pytest.approx(pybamm_result2) + assert rust_result2 == pytest.approx(0.0) + + def test_interpolation_1d_linear(self): + """Interpolant with y=2x, evaluate at [0.4, 0.6].""" + import pybamm + + x = np.linspace(0, 1, 50) + y_data = 2 * x + sv = pybamm.StateVector(slice(0, 2)) + interp = pybamm.Interpolant(x, y_data, sv) + + y_test = np.array([0.4, 0.6]) + self._assert_rust_eval_matches(interp, y=y_test.reshape(-1, 1)) + + def test_max_min_reduction(self): + """pybamm.max and pybamm.min on arrays.""" + import pybamm + + v = pybamm.Vector(np.array([1.0, 5.0, 3.0])) + self._assert_rust_eval_matches(pybamm.max(v)) + self._assert_rust_eval_matches(pybamm.min(v)) + + def test_index_slice(self): + """Index(Vector, slice(1,3)) → elements at positions 1 and 2.""" + import pybamm + + v = pybamm.Vector(np.array([10.0, 20.0, 30.0, 40.0, 50.0])) + expr = pybamm.Index(v, slice(1, 3)) + self._assert_rust_eval_matches(expr) + + def test_concatenation(self): + """NumpyConcatenation of two vectors.""" + import pybamm + + v1 = pybamm.Vector(np.array([1.0, 2.0])) + v2 = pybamm.Vector(np.array([3.0, 4.0, 5.0])) + expr = pybamm.NumpyConcatenation(v1, v2) + self._assert_rust_eval_matches(expr) + + def test_sparse_matmul(self): + """Sparse Matrix @ StateVector.""" + from scipy.sparse import csr_matrix as scipy_csr + + import pybamm + + data = np.array([1.0, 2.0]) + row = np.array([0, 1]) + col = np.array([0, 2]) + sparse_mat = scipy_csr((data, (row, col)), shape=(2, 3)) + mat = pybamm.Matrix(sparse_mat) + sv = pybamm.StateVector(slice(0, 3)) + expr = mat @ sv + + y = np.array([10.0, 20.0, 30.0]) + self._assert_rust_eval_matches(expr, y=y.reshape(-1, 1)) + + def test_composed_expression(self): + """Realistic composed expression: sqrt(sv**2 + 1) * exp(-t) with StateVector and Time.""" + import pybamm + + sv = pybamm.StateVector(slice(0, 1)) + t_sym = pybamm.t + expr = pybamm.sqrt(sv**2 + pybamm.Scalar(1.0)) * pybamm.exp(-t_sym) + + y = np.array([3.0]) + # sqrt(9 + 1) * exp(-2) ≈ 0.4280 + from pybamm.rust import ExprGraph + + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + rust_result = g.eval_to_float(rust_expr, 2.0, y.tolist(), [], []) + pybamm_result = float( + np.asarray(expr.evaluate(t=2.0, y=y.reshape(-1, 1))).flat[0] + ) + assert rust_result == pytest.approx(pybamm_result, rel=1e-5) + + def test_conditional_branch_selection(self): + """Test Conditional with different selector values selecting different branches.""" + import pybamm + from pybamm.rust import ExprGraph + + selector = pybamm.InputParameter("s") + branch1 = pybamm.Scalar(10.0) + branch2 = pybamm.Scalar(20.0) + branch3 = pybamm.Scalar(30.0) + expr = pybamm.Conditional(selector, branch1, branch2, branch3) + + # Selector = 1.0 → branch 1 (10.0) + g1 = ExprGraph() + rust_symbols1 = {} + re1 = expr.to_rust(g1, rust_symbols1) + result1 = g1.eval_to_float(re1, 0.0, [], [], [1.0]) + assert result1 == pytest.approx(10.0) + + # Selector = 2.0 → branch 2 (20.0) + g2 = ExprGraph() + rust_symbols2 = {} + re2 = expr.to_rust(g2, rust_symbols2) + result2 = g2.eval_to_float(re2, 0.0, [], [], [2.0]) + assert result2 == pytest.approx(20.0) + + # Selector = 3.0 → branch 3 (30.0) + g3 = ExprGraph() + rust_symbols3 = {} + re3 = expr.to_rust(g3, rust_symbols3) + result3 = g3.eval_to_float(re3, 0.0, [], [], [3.0]) + assert result3 == pytest.approx(30.0) + + # Selector = 0.0 → no branch active, returns 0 + g4 = ExprGraph() + rust_symbols4 = {} + re4 = expr.to_rust(g4, rust_symbols4) + result4 = g4.eval_to_float(re4, 0.0, [], [], [0.0]) + assert result4 == pytest.approx(0.0) + + def test_conditional_vs_python(self): + """Test Conditional Rust evaluation matches Python evaluation.""" + import pybamm + from pybamm.rust import ExprGraph + + selector = pybamm.InputParameter("s") + branch1 = pybamm.InputParameter("a") + branch2 = pybamm.InputParameter("b") + expr = pybamm.Conditional(selector, branch1, branch2) + + for s_val in [0.5, 1.0, 1.5, 2.0, 2.5]: + inputs = {"s": s_val, "a": 100.0, "b": 200.0} + + # Rust evaluation - inputs order matches expression traversal: s, a, b + g = ExprGraph() + rust_symbols = {} + rust_expr = expr.to_rust(g, rust_symbols) + inputs_list = [inputs["s"], inputs["a"], inputs["b"]] + rust_result = g.eval_to_float(rust_expr, 0.0, [], [], inputs_list) + + # Python evaluation + pybamm_result = float(expr.evaluate(inputs=inputs)) + + assert rust_result == pytest.approx(pybamm_result, rel=1e-5) + + +class TestRustUnsupportedNodeErrors: + """Unsupported nodes raise explicit, actionable errors naming CasADi.""" + + def test_unsupported_4d_interpolant_message(self): + import pybamm + from pybamm.rust import ExprGraph + + # 2D/3D now lower natively; 4D+ keeps the actionable error. + x = tuple(np.linspace(0, 1, 5) for _ in range(4)) + z = np.zeros((5, 5, 5, 5)) + svs = tuple(pybamm.StateVector(slice(i, i + 1)) for i in range(4)) + interp = pybamm.Interpolant(x, z, svs, interpolator="linear") + with pytest.raises(NotImplementedError, match=r"convert_to_format"): + interp.to_rust(ExprGraph(), {}) + + def test_unsupported_base_symbol_message(self): + import pybamm + from pybamm.rust import ExprGraph + + # SpatialVariable has no _to_rust -> base Symbol._to_rust raises. + sym = pybamm.SpatialVariable("x", domain=["negative electrode"]) + with pytest.raises(TypeError, match=r"convert_to_format") as exc: + sym.to_rust(ExprGraph(), {}) + # Must point at the live API, not the removed IDAKLUSolver(evaluator=...). + assert "evaluator=" not in str(exc.value) + + def test_unsupported_generic_function_message(self): + import pybamm + from pybamm.rust import ExprGraph + + # A generic Function wrapping a Python callable hits Function._rust_evaluate + # rather than Symbol._to_rust, so it needs its own actionable message. + sv = pybamm.StateVector(slice(0, 1)) + fun = pybamm.Function(np.sin, sv) + with pytest.raises(TypeError, match=r"convert_to_format"): + fun.to_rust(ExprGraph(), {}) + + +class TestCrossGraphGuard: + """An Expr from one ExprGraph must not be usable in another.""" + + def test_builder_rejects_foreign_expr(self): + import pytest + + from pybamm.rust import ExprGraph + + g_a = ExprGraph() + g_b = ExprGraph() + a = g_a.scalar(2.0) + b = g_b.scalar(3.0) + with pytest.raises(ValueError, match="different ExprGraph"): + g_a.add(a, b) # b belongs to g_b + + def test_dunder_rejects_foreign_expr(self): + import pytest + + from pybamm.rust import ExprGraph + + g_a = ExprGraph() + g_b = ExprGraph() + a = g_a.scalar(2.0) + b = g_b.scalar(3.0) + with pytest.raises(ValueError, match="different ExprGraph"): + _ = a + b # __add__ on a, other from g_b + + def test_compile_rejects_foreign_expr(self): + import pytest + + from pybamm.rust import ExprGraph + + g_a = ExprGraph() + g_b = ExprGraph() + a = g_a.scalar(2.0) + with pytest.raises(ValueError, match="different ExprGraph"): + g_b.compile(a) # a belongs to g_a + + def test_compile_group_rejects_foreign_expr(self): + import pytest + + from pybamm.rust import ExprGraph + + g_a = ExprGraph() + g_b = ExprGraph() + a = g_a.scalar(2.0) + with pytest.raises(ValueError, match="different ExprGraph"): + g_b.compile_group({"x": a}) # a belongs to g_a + + +class TestRustReduceSubgradientParity: + """pybamm.max/min Rust Jacobian matches casadi.mmax/mmin (argmax indicator).""" + + def _rust_jac_row(self, expr_pybamm, n, y): + import numpy as np + + from pybamm.rust import ExprGraph + + g = ExprGraph() + expr = expr_pybamm.to_rust(g, {}) + f = g.compile(expr, n_states=n) + return f.jacobian("y")(0.0, y, np.array([])).toarray() + + def test_max_jacobian_matches_casadi(self): + import casadi + import numpy as np + + import pybamm + + n = 4 + y = np.array([0.3, 0.9, 0.1, 0.5]) # unique argmax at index 1 + sv = pybamm.StateVector(slice(0, n)) + rust_row = self._rust_jac_row(pybamm.max(sv), n, y) + + ys = casadi.MX.sym("y", n) + jac = casadi.Function("J", [ys], [casadi.jacobian(casadi.mmax(ys), ys)]) + casadi_row = np.array(jac(y)).reshape(1, n) + + expected = np.zeros((1, n)) + expected[0, 1] = 1.0 + np.testing.assert_allclose(rust_row, casadi_row, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(rust_row, expected, rtol=1e-12, atol=1e-12) + + def test_min_jacobian_matches_casadi(self): + import casadi + import numpy as np + + import pybamm + + n = 4 + y = np.array([0.3, 0.9, 0.1, 0.5]) # unique argmin at index 2 + sv = pybamm.StateVector(slice(0, n)) + rust_row = self._rust_jac_row(pybamm.min(sv), n, y) + + ys = casadi.MX.sym("y", n) + jac = casadi.Function("J", [ys], [casadi.jacobian(casadi.mmin(ys), ys)]) + casadi_row = np.array(jac(y)).reshape(1, n) + + expected = np.zeros((1, n)) + expected[0, 2] = 1.0 + np.testing.assert_allclose(rust_row, casadi_row, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(rust_row, expected, rtol=1e-12, atol=1e-12) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_state_vector.py b/packages/pybamm/tests/unit/test_expression_tree/test_state_vector.py index a534511bbc..34fa1556f6 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_state_vector.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_state_vector.py @@ -99,6 +99,18 @@ def test_to_from_json(self, mocker): # Turn debug mode back to what is was before pybamm.settings.debug_mode = original_debug_mode + def test_multi_slice_to_rust(self): + from pybamm.rust import ExprGraph + + sv = pybamm.StateVector(slice(0, 2), slice(4, 6)) + graph = ExprGraph() + cf = graph.compile(sv.to_rust(graph, {}), name="sv", n_states=6) + y = np.arange(6.0) + np.testing.assert_array_equal( + np.asarray(cf(0.0, y, np.array([]))).ravel(), + sv.evaluate(y=y.reshape(-1, 1)).ravel(), + ) + class TestStateVectorDot: def test_evaluate(self): diff --git a/packages/pybamm/tests/unit/test_hatch_build.py b/packages/pybamm/tests/unit/test_hatch_build.py new file mode 100644 index 0000000000..872c9c9840 --- /dev/null +++ b/packages/pybamm/tests/unit/test_hatch_build.py @@ -0,0 +1,216 @@ +"""Unit tests for the maturin build hook (``packages/pybamm/hatch_build.py``). + +Synthetic wheels stand in for maturin's output, so none of this needs a Rust +toolchain. +""" + +from __future__ import annotations + +import importlib.util +import sys +import zipfile +from pathlib import Path + +import pytest + +HOOK_PATH = Path(__file__).resolve().parents[2] / "hatch_build.py" + + +def load_hook_module(): + spec = importlib.util.spec_from_file_location("pybamm_hatch_build", HOOK_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def make_hook(root): + """Build a hook instance. ``root`` is a read-only property, so it must go + through ``__init__``; the other arguments are unused by the paths tested.""" + module = load_hook_module() + hook = module.RustBuildHook( + root=str(root), + config={}, + build_config=None, + metadata=None, + directory="", + target_name="wheel", + ) + return module, hook + + +def make_wheel(path, members): + with zipfile.ZipFile(path, "w") as archive: + for name in members: + archive.writestr(name, b"\x7fELF fake") + return path + + +class TestWheelTag: + @pytest.mark.parametrize( + ("name", "expected"), + [ + ( + "pybamm_rust-0.1.0-cp310-abi3-macosx_11_0_arm64.whl", + "cp310-abi3-macosx_11_0_arm64", + ), + ( + "pybamm_rust-0.1.0-cp314-cp314-linux_x86_64.whl", + "cp314-cp314-linux_x86_64", + ), + ("pybamm-26.7.1.1-cp310-abi3-win_amd64.whl", "cp310-abi3-win_amd64"), + ], + ) + def test_extracts_the_trailing_three_components(self, name, expected): + module = load_hook_module() + assert module.wheel_tag(name) == expected + + def test_rejects_a_malformed_name(self): + module = load_hook_module() + with pytest.raises(RuntimeError, match=r"Cannot parse a wheel tag"): + module.wheel_tag("not-a-wheel.whl") + + +class TestExtractExtension: + def test_copies_the_single_extension(self, tmp_path): + module = load_hook_module() + wheel = make_wheel( + tmp_path / "w.whl", ["_core/_core.abi3.so", "_core/__init__.py"] + ) + target = module.extract_extension(wheel, tmp_path / "dest") + assert target.name == "_core.abi3.so" + assert target.read_bytes() == b"\x7fELF fake" + + def test_removes_a_stale_version_specific_artifact(self, tmp_path): + module = load_hook_module() + destination = tmp_path / "dest" + destination.mkdir() + stale = destination / "_core.cpython-314-darwin.so" + stale.write_bytes(b"old") + (destination / "__init__.py").write_text("keep me\n") + + wheel = make_wheel(tmp_path / "w.whl", ["_core/_core.abi3.so"]) + module.extract_extension(wheel, destination) + + assert not stale.exists() + assert sorted(p.name for p in destination.glob("_core*")) == ["_core.abi3.so"] + assert (destination / "__init__.py").read_text() == "keep me\n" + + def test_keeps_the_checked_in_stub_file(self, tmp_path): + module = load_hook_module() + destination = tmp_path / "dest" + destination.mkdir() + stub = destination / "_core.pyi" + stub.write_text("keep me\n") + + wheel = make_wheel(tmp_path / "w.whl", ["_core/_core.abi3.so"]) + module.extract_extension(wheel, destination) + + assert stub.read_text() == "keep me\n" + + def test_leaves_no_staging_file(self, tmp_path): + module = load_hook_module() + wheel = make_wheel(tmp_path / "w.whl", ["_core/_core.abi3.so"]) + destination = tmp_path / "dest" + module.extract_extension(wheel, destination) + assert list(destination.glob("*.tmp")) == [] + + def test_rejects_a_wheel_with_no_extension(self, tmp_path): + module = load_hook_module() + wheel = make_wheel(tmp_path / "w.whl", ["_core/__init__.py"]) + with pytest.raises(RuntimeError, match=r"Expected exactly one _core extension"): + module.extract_extension(wheel, tmp_path / "dest") + + def test_rejects_a_wheel_with_two_extensions(self, tmp_path): + module = load_hook_module() + wheel = make_wheel( + tmp_path / "w.whl", + ["_core/_core.abi3.so", "_core/_core.cpython-314-darwin.so"], + ) + with pytest.raises(RuntimeError, match=r"Expected exactly one _core extension"): + module.extract_extension(wheel, tmp_path / "dest") + + +class TestCrateRootResolution: + def test_finds_crate_in_monorepo_layout(self, tmp_path): + (tmp_path / "packages" / "pybamm-rust" / "pybamm-python").mkdir(parents=True) + (tmp_path / "packages" / "pybamm-rust" / "pybamm-python" / "Cargo.toml").touch() + package = tmp_path / "packages" / "pybamm" + package.mkdir() + _, hook = make_hook(package) + assert hook._crate_root() == tmp_path / "packages" / "pybamm-rust" + + def test_finds_crate_in_sdist_layout(self, tmp_path): + (tmp_path / "pybamm-rust" / "pybamm-python").mkdir(parents=True) + (tmp_path / "pybamm-rust" / "pybamm-python" / "Cargo.toml").touch() + _, hook = make_hook(tmp_path) + assert hook._crate_root() == tmp_path / "pybamm-rust" + + def test_raises_when_crate_is_absent(self, tmp_path): + _, hook = make_hook(tmp_path) + with pytest.raises(RuntimeError, match=r"Could not find"): + hook._crate_root() + + +class TestBuildInvocation: + def test_missing_cargo_raises_actionable_error(self, tmp_path, monkeypatch): + module, hook = make_hook(tmp_path) + monkeypatch.setattr(module.shutil, "which", lambda _: None) + with pytest.raises(RuntimeError, match=r"requires a Rust toolchain"): + hook._build_extension(tmp_path) + + def test_maturin_is_invoked_through_the_interpreter_with_locked( + self, tmp_path, monkeypatch + ): + module, hook = make_hook(tmp_path) + monkeypatch.setattr(module.shutil, "which", lambda _: "/usr/bin/cargo") + recorded = {} + + def fake_run(command, cwd, check): + recorded["command"] = command + recorded["cwd"] = cwd + out = Path(command[command.index("--out") + 1]) + make_wheel( + out / "pybamm_rust-0.1.0-cp310-abi3-linux_x86_64.whl", + ["_core/_core.abi3.so"], + ) + + monkeypatch.setattr(module.subprocess, "run", fake_run) + artifact, tag = hook._build_extension(tmp_path) + + assert recorded["command"][:4] == [sys.executable, "-m", "maturin", "build"] + assert "--locked" in recorded["command"] + assert "--release" in recorded["command"] + assert recorded["cwd"] == tmp_path + assert tag == "cp310-abi3-linux_x86_64" + assert artifact.name == "_core.abi3.so" + + +class TestBuildDataWiring: + def _initialize(self, tmp_path, monkeypatch, version): + _, hook = make_hook(tmp_path) + monkeypatch.setattr(hook, "_crate_root", lambda: tmp_path) + monkeypatch.setattr( + hook, + "_build_extension", + lambda _: (Path("_core.abi3.so"), "cp310-abi3-linux_x86_64"), + ) + build_data = { + "artifacts": [], + "force_include": {}, + "infer_tag": False, + "pure_python": True, + } + hook.initialize(version, build_data) + return build_data + + def test_standard_build_sets_the_tag_and_artifact(self, tmp_path, monkeypatch): + build_data = self._initialize(tmp_path, monkeypatch, "standard") + assert build_data["tag"] == "cp310-abi3-linux_x86_64" + assert build_data["pure_python"] is False + assert build_data["artifacts"] == ["/src/pybamm/rust/_core.abi3.so"] + + def test_editable_build_does_not_set_a_tag(self, tmp_path, monkeypatch): + build_data = self._initialize(tmp_path, monkeypatch, "editable") + assert "tag" not in build_data + assert build_data["artifacts"] == ["/src/pybamm/rust/_core.abi3.so"] diff --git a/packages/pybamm/tests/unit/test_models/test_base_model.py b/packages/pybamm/tests/unit/test_models/test_base_model.py index dbd423e803..58c76fa750 100644 --- a/packages/pybamm/tests/unit/test_models/test_base_model.py +++ b/packages/pybamm/tests/unit/test_models/test_base_model.py @@ -15,6 +15,8 @@ import pybamm +pytest_plugins = ["tests.conftest_rust"] + class TestBaseModel: def test_rhs_set_get(self): @@ -1752,3 +1754,42 @@ def test_save_load_model(self): def test_y0_property(self): model = pybamm.BaseModel() assert model.y0 is None + + def test_convert_to_format_validated(self): + # Which backend ships as the default is asserted by + # test_convert_to_format_default_constant; this covers the setter only. + model = pybamm.BaseModel() + for valid in (None, "python", "casadi", "jax", "rust"): + model.convert_to_format = valid + assert model.convert_to_format == valid + with pytest.raises(ValueError, match="convert_to_format must be one of"): + model.convert_to_format = "fortran" + + def test_convert_to_format_default_constant(self, monkeypatch): + # Drive the constant to each backend in turn: an __init__ that hardcoded + # any one format would satisfy at most one of these. Never reads the + # shipped default, so it holds whichever backend ships as the default. + for patched in ("casadi", "rust", "python"): + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", patched) + assert pybamm.BaseModel().convert_to_format == patched + # explicit assignment still wins over the class default + model = pybamm.BaseModel() + for explicit in ("casadi", "rust"): + model.convert_to_format = explicit + assert model.convert_to_format == explicit + + def test_convert_to_format_serialisation_roundtrip(self, tmp_path): + model = pybamm.BaseModel(name="rt") + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1.0} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = "rust" + path = str(tmp_path / "rt_model") + model.save_model(filename=path) + loaded = pybamm.load_model(path + ".json") + assert loaded.convert_to_format == "rust" + + def test_rust_backend_fixture_drives_default(self, rust_backend): + assert pybamm.BaseModel().convert_to_format == "rust" diff --git a/packages/pybamm/tests/unit/test_parameters/test_parameter_values_serialisation.py b/packages/pybamm/tests/unit/test_parameters/test_parameter_values_serialisation.py index b749d42f8d..b782243494 100644 --- a/packages/pybamm/tests/unit/test_parameters/test_parameter_values_serialisation.py +++ b/packages/pybamm/tests/unit/test_parameters/test_parameter_values_serialisation.py @@ -528,8 +528,8 @@ def test_full_parameter_set_roundtrip_ai2020(self): # 2-arg diffusivities (sto, T) y_sto_T = np.vstack([sto_vals, T_vals]) for diff_name in [ - "Negative electrode diffusivity [m2.s-1]", - "Positive electrode diffusivity [m2.s-1]", + "Negative particle diffusivity [m2.s-1]", + "Positive particle diffusivity [m2.s-1]", ]: _assert_evaluate_equal_array( pv, diff --git a/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py b/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py index 45956fc020..91666e7505 100644 --- a/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py +++ b/packages/pybamm/tests/unit/test_plotting/test_quick_plot.py @@ -12,6 +12,8 @@ class TestQuickPlot: @pytest.mark.parametrize("solver", _solver_args) def test_simple_ode_model(self, solver): model = pybamm.lithium_ion.BaseModel(name="Simple ODE Model") + if isinstance(solver, pybamm.CasadiSolver): + model.convert_to_format = "casadi" whole_cell = ["negative electrode", "separator", "positive electrode"] # Create variables: domain is explicitly empty since these variables are only diff --git a/packages/pybamm/tests/unit/test_rust_boundary_validation.py b/packages/pybamm/tests/unit/test_rust_boundary_validation.py new file mode 100644 index 0000000000..d046146b55 --- /dev/null +++ b/packages/pybamm/tests/unit/test_rust_boundary_validation.py @@ -0,0 +1,100 @@ +"""Boundary validation for the Rust ExprGraph constructors. + +Malformed ranges, matrices, and interpolation tables supplied at the Python +boundary must raise ordinary ``ValueError``s rather than panicking, wrapping in +release arithmetic, or reading out of bounds inside evaluation. +""" + +import numpy as np +import pytest + +from pybamm.rust import ExprGraph + + +class TestStateVectorValidation: + def test_state_vector_inverted_range_raises(self): + g = ExprGraph() + with pytest.raises(ValueError, match=r"start"): + g.state_vector(2, 1) + + def test_state_vector_dot_inverted_range_raises(self): + g = ExprGraph() + with pytest.raises(ValueError, match=r"start"): + g.state_vector_dot(2, 1) + + def test_index_inverted_range_raises(self): + g = ExprGraph() + y = g.state_vector(0, 4) + with pytest.raises(ValueError, match=r"start"): + g.index(y, 3, 1) + + def test_valid_ranges_still_work(self): + g = ExprGraph() + y = g.state_vector(0, 4) + # end == start (empty) and normal ranges must not raise + g.state_vector(1, 1) + g.index(y, 1, 3) + + +class TestSparseMatrixValidation: + def test_non_monotonic_indptr_raises(self): + g = ExprGraph() + with pytest.raises(ValueError, match=r"CSR|indptr"): + g.sparse_matrix([0, 2, 1], [0, 1], np.array([1.0, 2.0]), 2, 2) + + def test_out_of_range_column_raises(self): + g = ExprGraph() + with pytest.raises(ValueError, match=r"CSR|column"): + g.sparse_matrix([0, 1, 2], [0, 5], np.array([1.0, 2.0]), 2, 2) + + def test_indptr_wrong_length_raises(self): + g = ExprGraph() + with pytest.raises(ValueError, match=r"CSR|indptr"): + g.sparse_matrix([0, 2], [0, 1], np.array([1.0, 2.0]), 2, 2) + + def test_valid_sparse_matrix_still_works(self): + g = ExprGraph() + g.sparse_matrix([0, 1, 2], [0, 1], np.array([1.0, 2.0]), 2, 2) + + +class TestInterpolantValidation: + def test_empty_grid_raises(self): + g = ExprGraph() + y = g.state_vector(0, 1) + with pytest.raises(ValueError, match=r"interpolant|non-empty"): + g.interpolant_1d_linear([], [], y) + + def test_length_mismatch_raises(self): + g = ExprGraph() + y = g.state_vector(0, 1) + with pytest.raises(ValueError, match=r"interpolant|length"): + g.interpolant_1d_linear([0.0, 1.0], [1.0], y) + + def test_non_increasing_grid_raises(self): + g = ExprGraph() + y = g.state_vector(0, 1) + with pytest.raises(ValueError, match=r"interpolant|increasing"): + g.interpolant_1d_linear([0.0, 2.0, 1.0], [1.0, 2.0, 3.0], y) + + def test_valid_interpolant_still_works(self): + g = ExprGraph() + y = g.state_vector(0, 1) + g.interpolant_1d_linear([0.0, 1.0, 2.0], [10.0, 20.0, 30.0], y) + + +class TestEvalHelperValidation: + def test_eval_to_array_non_contiguous_y_raises(self): + # TypeError (not PanicException): the layout error every other eval + # path raises for a strided array. + g = ExprGraph() + expr = g.state_vector(0, 3) + strided = np.arange(6.0)[::2] + with pytest.raises(TypeError, match=r"contiguous"): + g.eval_to_array(expr, 0.0, strided, np.array([]), []) + + def test_eval_to_float_empty_result_raises(self): + g = ExprGraph() + y = g.state_vector(0, 4) + empty = g.index(y, 1, 1) + with pytest.raises(ValueError): + g.eval_to_float(empty, 0.0, [0.0, 0.0, 0.0, 0.0], [], []) diff --git a/packages/pybamm/tests/unit/test_rust_extension_loading.py b/packages/pybamm/tests/unit/test_rust_extension_loading.py new file mode 100644 index 0000000000..a5f6992686 --- /dev/null +++ b/packages/pybamm/tests/unit/test_rust_extension_loading.py @@ -0,0 +1,86 @@ +"""How ``pybamm.rust`` loads the compiled extension. + +pybammsolvers resolves the Rust FFI entry points with ``dlsym(RTLD_DEFAULT)``, +which only finds them when the extension is in the process-global symbol scope. +CPython opens extension modules ``RTLD_LOCAL`` on Linux, so the facade has to ask +for global visibility there; macOS already exports them globally. +""" + +import os +import sys + +import pytest + +import pybamm.rust + +# The Linux cases stub sys.getdlopenflags/setdlopenflags and compare against +# os.RTLD_*, none of which exist off POSIX. +requires_posix_dlopen = pytest.mark.skipif( + not hasattr(sys, "getdlopenflags"), reason="POSIX dlopen flags only" +) + + +class TestGlobalSymbolVisibility: + @requires_posix_dlopen + def test_sets_rtld_global_on_linux(self, monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "getdlopenflags", lambda: os.RTLD_NOW) + applied = [] + monkeypatch.setattr(sys, "setdlopenflags", applied.append) + + with pybamm.rust._global_symbol_visibility(): + pass + + assert applied, "no dlopen flags were set on Linux" + assert applied[0] & os.RTLD_GLOBAL + + @requires_posix_dlopen + def test_restores_the_previous_flags_on_linux(self, monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "getdlopenflags", lambda: os.RTLD_NOW) + applied = [] + monkeypatch.setattr(sys, "setdlopenflags", applied.append) + + with pybamm.rust._global_symbol_visibility(): + pass + + assert applied[-1] == os.RTLD_NOW + + @requires_posix_dlopen + def test_restores_the_previous_flags_when_the_import_raises(self, monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "getdlopenflags", lambda: os.RTLD_NOW) + applied = [] + monkeypatch.setattr(sys, "setdlopenflags", applied.append) + + with ( + pytest.raises(ImportError), + pybamm.rust._global_symbol_visibility(), + ): + raise ImportError("extension missing") + + assert applied[-1] == os.RTLD_NOW + + def test_leaves_flags_untouched_off_linux(self, monkeypatch): + monkeypatch.setattr(sys, "platform", "darwin") + # raising=False: Windows has no setdlopenflags, and is itself off-Linux. + monkeypatch.setattr( + sys, + "setdlopenflags", + lambda _: pytest.fail("must not touch dlopen flags"), + raising=False, + ) + + with pybamm.rust._global_symbol_visibility(): + pass + + @requires_posix_dlopen + def test_importing_the_facade_does_not_leak_flags(self): + """The real import already ran at collection; flags must be back to normal.""" + assert not sys.getdlopenflags() & os.RTLD_GLOBAL + + +class TestFacadeExports: + def test_core_is_importable_through_the_facade(self): + assert pybamm.rust.CompiledModel is not None + assert pybamm.rust.ExprGraph is not None diff --git a/packages/pybamm/tests/unit/test_rust_function_api.py b/packages/pybamm/tests/unit/test_rust_function_api.py new file mode 100644 index 0000000000..236569f470 --- /dev/null +++ b/packages/pybamm/tests/unit/test_rust_function_api.py @@ -0,0 +1,946 @@ +"""Unit tests for the prep-artifact binding API (CompiledFunction).""" + +import importlib +from pathlib import Path + +import numpy as np +import pytest +import scipy.sparse + +import pybamm +from pybamm.rust import ExprGraph + + +def _make_fn(name=None): + """f(t, y, p) = [p_a * y0 * y1, sin(y0) * p_b]; 2 states, 2 inputs.""" + g = ExprGraph() + pa = g.input_parameter("a") + pb = g.input_parameter("b") + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + r0 = g.mul(g.mul(pa, y0), y1) + r1 = g.mul(g.sin(y0), pb) + expr = g.concat([r0, r1]) + return g, g.compile(expr, name=name) + + +def _full_residual_symbol(built_model): + """rhs concatenated with algebraic, matching the solver's residual. + + Inlined from `benchmarks/rust_observability/runners.py` — `benchmarks/` + is not importable under this test tree's pytest rootdir. + """ + if built_model.len_alg > 0: + return pybamm.numpy_concatenation( + built_model.concatenated_rhs, + built_model.concatenated_algebraic, + ) + return built_model.concatenated_rhs + + +class TestSignature: + def test_properties(self): + _, f = _make_fn(name="RHS") + assert f.input_names == ("a", "b") + assert f.n_inputs == 2 + assert f.n_states == 2 + assert f.output_len == 2 + assert f.uses_y_dot is False + assert f.name == "RHS" + assert "RHS" in repr(f) + + def test_explicit_n_states_override(self): + g = ExprGraph() + y0 = g.state_vector(0, 1) # touches only state 0 of a 3-state system + f = g.compile(g.mul(y0, y0), n_states=3) + assert f.n_states == 3 + + def test_n_states_override_below_extent_rejected(self): + # an override below the scanned extent would admit a too-short y + # and panic inside the tape + g = ExprGraph() + y = g.state_vector(0, 5) + with pytest.raises(ValueError, match=r"n_states=2.*state extent 5"): + g.compile(g.mul(y, y), n_states=2) + + +class TestCompilePrep: + def test_compile_runs_simplify_pipeline(self): + # compile prep includes simplification — decorating an + # expression with *1 and +0 must not lengthen the tape + g = ExprGraph() + y0 = g.state_vector(0, 1) + base = g.mul(g.sin(y0), g.input_parameter("a")) + decorated = g.add(g.mul(base, g.scalar(1.0)), g.scalar(0.0)) + assert g.compile(decorated).n_instructions == g.compile(base).n_instructions + + def test_signature_scans_pre_simplify_extent(self): + # y1 - y1 cancels, but the user-declared contract still spans both + # states: a length-2 y must stay valid at the call site + g = ExprGraph() + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + f = g.compile(g.add(y0, g.sub(y1, y1))) + assert f.n_states == 2 + f(0.0, np.array([1.0, 5.0]), np.array([])) # must not raise + + +class TestCall: + def test_call_stacked(self): + _, f = _make_fn() + y = np.array([0.5, 2.0]) + p = np.array([3.0, 4.0]) + out = f(0.0, y, p) + np.testing.assert_allclose(out, [3.0 * 0.5 * 2.0, np.sin(0.5) * 4.0]) + + def test_call_dict(self): + _, f = _make_fn() + y = np.array([0.5, 2.0]) + out = f(0.0, y, {"a": 3.0, "b": 4.0}) + np.testing.assert_allclose(out, [3.0, np.sin(0.5) * 4.0], rtol=1e-12) + + def test_call_dict_with_1element_array_values(self): + # PyBaMM stores solve inputs as 1-element arrays, e.g. {'a': array([3.0])}. + # Native observation passes solution.all_inputs[i] straight to the binding, + # so pack must accept length-1 arrays as scalars, not only Python floats. + _, f = _make_fn() + y = np.array([0.5, 2.0]) + out = f(0.0, y, {"a": np.array([3.0]), "b": np.array([4.0])}) + np.testing.assert_allclose(out, [3.0, np.sin(0.5) * 4.0], rtol=1e-12) + + def test_call_dict_with_nonfloat64_values(self): + # int64/float32 arrays and numpy scalars must coerce to f64 like CasADi. + _, f = _make_fn() + y = np.array([0.5, 2.0]) + expected = [3.0, np.sin(0.5) * 4.0] + cases = [ + {"a": np.array([3], dtype=np.int64), "b": np.array([4], dtype=np.int64)}, + {"a": np.array([3.0], dtype=np.float32), "b": np.float32(4.0)}, + {"a": np.int64(3), "b": 4}, + ] + for inputs in cases: + np.testing.assert_allclose(f(0.0, y, inputs), expected, rtol=1e-6) + + def test_dict_and_stacked_bitwise_equal(self): + _, f = _make_fn() + y = np.array([0.5, 2.0]) + a = f(0.0, y, np.array([3.0, 4.0])) + b = f(0.0, y, {"a": 3.0, "b": 4.0}) + assert (a == b).all() + + def test_pack(self): + _, f = _make_fn() + np.testing.assert_array_equal(f.pack({"a": 3.0, "b": 4.0}), [3.0, 4.0]) + + def test_eval_into(self): + _, f = _make_fn() + out = np.zeros(2) + f.eval_into(0.0, np.array([0.5, 2.0]), np.array([3.0, 4.0]), out) + np.testing.assert_allclose(out, [3.0, np.sin(0.5) * 4.0]) + + def test_shared_function_is_reusable(self): + # &self eval: two interleaved calls on one object, no state bleed + _, f = _make_fn() + y1, y2 = np.array([0.5, 2.0]), np.array([1.0, 1.0]) + p = np.array([1.0, 1.0]) + a1 = f(0.0, y1, p) + b1 = f(0.0, y2, p) + a2 = f(0.0, y1, p) + assert (a1 == a2).all() and not (a1 == b1).all() + + +class TestValidation: + def test_wrong_y_length(self): + _, f = _make_fn(name="RHS") + with pytest.raises(ValueError, match=r"RHS.*expected y of length 2.*got 3"): + f(0.0, np.zeros(3), np.zeros(2)) + + def test_wrong_p_length(self): + _, f = _make_fn(name="RHS") + with pytest.raises( + ValueError, match=r"expected 2 input values.*2 parameters.*got 1" + ): + f(0.0, np.zeros(2), np.zeros(1)) + + def test_wrong_y_dot_length(self): + # a short y_dot must raise, not panic inside the tape + g = ExprGraph() + f = g.compile(g.state_vector_dot(0, 2), name="resid") + with pytest.raises( + ValueError, match=r"resid.*expected y_dot of length 2.*got 1" + ): + f(0.0, np.zeros(2), np.array([]), y_dot=np.zeros(1)) + + def test_dict_missing_key(self): + _, f = _make_fn() + with pytest.raises(ValueError, match=r"missing input 'b'"): + f(0.0, np.zeros(2), {"a": 1.0}) + + def test_dict_unknown_key(self): + _, f = _make_fn() + with pytest.raises(ValueError, match=r"unknown input 'c'"): + f(0.0, np.zeros(2), {"a": 1.0, "b": 2.0, "c": 3.0}) + + def test_dict_rejects_multielement_array(self): + # a length>1 array is an error, not a silent take-first. + _, f = _make_fn() + with pytest.raises(ValueError, match=r"input 'a'.*scalar or length-1"): + f(0.0, np.zeros(2), {"a": np.array([1.0, 2.0]), "b": 3.0}) + + def test_eval_into_rejects_aliased_out(self): + # out aliasing y would be UB without the borrow guard + _, f = _make_fn() + arr = np.array([0.5, 2.0]) + with pytest.raises(BaseException, match=r"AlreadyBorrowed"): + f.eval_into(0.0, arr, np.array([3.0, 4.0]), arr) + + def test_non_contiguous_y_is_a_clear_type_error(self): + # A strided view must fail at conversion with a TypeError, never + # reach the slice unwrap inside the FFI. + _, f = _make_fn() + y_strided = np.arange(4.0).reshape(2, 2)[:, 0] + assert not y_strided.flags["C_CONTIGUOUS"] + with pytest.raises(TypeError, match=r"contiguous"): + f(0.0, y_strided, np.zeros(2)) + + +class TestJvp: + def test_jvp_wrt_y_matches_fd(self): + _, f = _make_fn() + y = np.array([0.5, 2.0]) + p = np.array([3.0, 4.0]) + vy = np.array([1.0, -0.5]) + eps = 1e-7 + fd = (f(0.0, y + eps * vy, p) - f(0.0, y - eps * vy, p)) / (2 * eps) + np.testing.assert_allclose(f.jvp(0.0, y, p, vy), fd, rtol=1e-6) + + def test_jvp_with_vp_sums_contributions(self): + _, f = _make_fn() + y = np.array([0.5, 2.0]) + p = np.array([3.0, 4.0]) + vy = np.array([1.0, -0.5]) + vp = np.array([0.25, -1.0]) + eps = 1e-7 + fd = ( + f(0.0, y + eps * vy, p + eps * vp) - f(0.0, y - eps * vy, p - eps * vp) + ) / (2 * eps) + np.testing.assert_allclose(f.jvp(0.0, y, p, vy, vp=vp), fd, rtol=1e-6) + + def test_jvp_tangent_tape_is_cached(self): + import time + + _, f = _make_fn() + y, p, vy = np.array([0.5, 2.0]), np.array([3.0, 4.0]), np.array([1.0, 0.0]) + t0 = time.perf_counter() + f.jvp(0.0, y, p, vy) + first = time.perf_counter() - t0 + t0 = time.perf_counter() + for _ in range(100): + f.jvp(0.0, y, p, vy) + per_call = (time.perf_counter() - t0) / 100 + assert per_call < first, "derivation must be one-time prep, not per-call" + + def test_jvp_rejects_y_dot_expressions(self): + # LoadStateVectorDot slices y_dot unconditionally, so without this guard + # the call would panic across PyO3. + g = ExprGraph() + f = g.compile(g.state_vector_dot(0, 1), name="resid") + with pytest.raises(ValueError, match=r"resid.*y_dot"): + f.jvp(0.0, np.zeros(1), np.zeros(0), np.zeros(1)) + + def test_jvp_vp_broadcast_param_keeps_full_width(self): + # f = y[0:3] + p0, so df/dp @ [1] must contribute [1, 1, 1]: the tangent_p + # tape must not collapse to length 1 and truncate it to [1, 0, 0]. + g = ExprGraph() + sv = g.state_vector(0, 3) + f = g.compile(g.add(sv, g.input_parameter("a")), name="vecp") + assert f.output_len == 3 + y = np.array([10.0, 20.0, 30.0]) + out = f.jvp(0.0, y, np.array([5.0]), np.zeros(3), vp=np.array([1.0])) + np.testing.assert_array_equal(out, [1.0, 1.0, 1.0]) + + def test_jvp_zero_dfdy_returns_full_width_zeros(self): + # Regression: f = const_vector([1,2,3]) + p0; df/dy ≡ 0. + # jvp wrt y must return length-3 zeros, not a collapsed length-1 [0.]. + g = ExprGraph() + cv = g.array(np.array([1.0, 2.0, 3.0])) + f = g.compile(g.add(cv, g.input_parameter("a")), name="zeroy", n_states=2) + assert f.output_len == 3 + out = f.jvp(0.0, np.zeros(2), np.array([5.0]), np.zeros(2)) + np.testing.assert_array_equal(out, [0.0, 0.0, 0.0]) + + def test_jvp_rejects_wrong_vy_length(self): + _, f = _make_fn(name="rhs") + y, p = np.array([0.5, 2.0]), np.array([3.0, 4.0]) + with pytest.raises(ValueError, match=r"rhs.*vy.*length"): + f.jvp(0.0, y, p, np.zeros(3)) + + def test_jvp_rejects_wrong_vp_length(self): + _, f = _make_fn(name="rhs") + y, p, vy = np.array([0.5, 2.0]), np.array([3.0, 4.0]), np.zeros(2) + with pytest.raises(ValueError, match=r"rhs"): + f.jvp(0.0, y, p, vy, vp=np.zeros(3)) + + +class TestJacobian: + def test_wrt_y_matches_fd(self): + _, f = _make_fn() + J = f.jacobian() + y = np.array([0.5, 2.0]) + p = np.array([3.0, 4.0]) + mat = J(0.0, y, p) + assert isinstance(mat, scipy.sparse.csc_matrix) + assert mat.shape == (2, 2) + eps = 1e-7 + for j in range(2): + e = np.zeros(2) + e[j] = 1.0 + fd = (f(0.0, y + eps * e, p) - f(0.0, y - eps * e, p)) / (2 * eps) + np.testing.assert_allclose(mat.toarray()[:, j], fd, rtol=1e-6, atol=1e-10) + + def test_wrt_p_matches_fd(self): + _, f = _make_fn() + Jp = f.jacobian(wrt="p") + y = np.array([0.5, 2.0]) + p = np.array([3.0, 4.0]) + mat = Jp(0.0, y, p) + assert mat.shape == (2, 2) + eps = 1e-7 + for j in range(2): + e = np.zeros(2) + e[j] = 1.0 + fd = (f(0.0, y, p + eps * e) - f(0.0, y, p - eps * e)) / (2 * eps) + np.testing.assert_allclose(mat.toarray()[:, j], fd, rtol=1e-6, atol=1e-10) + + def test_cached_per_wrt(self): + _, f = _make_fn() + assert f.jacobian() is f.jacobian() + assert f.jacobian(wrt="p") is f.jacobian(wrt="p") + assert f.jacobian() is not f.jacobian(wrt="p") + + def test_introspection(self): + _, f = _make_fn() + J = f.jacobian() + assert J.shape == (2, 2) + assert J.wrt == "y" + assert J.nnz >= 3 # (0,0),(0,1),(1,0) at minimum + assert J.n_colors >= 1 + indptr, _indices = J.sparsity() + assert len(indptr) == 3 # CSC: n_cols + 1 + + def test_rectangular_partial_group(self): + g = ExprGraph() + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + f = g.compile(g.mul(y0, y1), n_states=2) # 1 output row, 2 states + assert f.jacobian().shape == (1, 2) + + def test_invalid_wrt(self): + _, f = _make_fn() + with pytest.raises(ValueError, match="wrt"): + f.jacobian(wrt="t") + + def test_rejects_y_dot_expressions(self): + # see TestJvp.test_jvp_rejects_y_dot_expressions — same panic guard + g = ExprGraph() + f = g.compile(g.state_vector_dot(0, 1), name="resid") + with pytest.raises(ValueError, match=r"resid.*y_dot"): + f.jacobian() + + def test_csc_buffers_shared_and_readonly(self): + # Only the data array is allocated per call: the cached int32 index arrays + # are shared with scipy without a cast-copy, and are read-only. + _, f = _make_fn() + J = f.jacobian() + y, p = np.array([0.5, 2.0]), np.array([3.0, 4.0]) + m1, m2 = J(0.0, y, p), J(0.0, y, p) + assert m1.indices.dtype == np.int32 # scipy's native index dtype + assert np.shares_memory(m1.indices, m2.indices) + assert np.shares_memory(m1.indptr, m2.indptr) + assert not m1.indices.flags.writeable + assert not m1.indptr.flags.writeable + with pytest.raises(ValueError): + m1.indices[0] = m1.indices[0] + assert m1.data.flags.writeable # only the pattern is frozen + + def test_spme_vaas_dense_row_coloring(self): + # SPMe with voltage-as-a-state has one ~65-nnz voltage row (default + # grid); without the dense-row split, column coloring needs 65 colors. + model = pybamm.lithium_ion.SPMe(options={"voltage as a state": "true"}) + model.events = [] + sim = pybamm.Simulation(model) + sim.build() + built = sim.built_model + g = ExprGraph() + fn = g.compile( + _full_residual_symbol(built).to_rust(g, {}), + name="spme", + n_states=built.len_rhs_and_alg, + ) + jac = fn.jacobian() + assert jac.n_dense_rows == 1 + assert jac.n_colors <= 8 # 65 before the split; sparse remainder ~5 + + # numeric parity: split assembly vs finite differences on a few columns + y0 = built.concatenated_initial_conditions.evaluate().flatten() + J = jac(0.0, y0, np.array([])).toarray() + h = 1e-7 + for col in [0, built.len_rhs_and_alg // 2, built.len_rhs_and_alg - 1]: + yp, ym = y0.copy(), y0.copy() + yp[col] += h + ym[col] -= h + fd = (fn(0.0, yp, np.array([])) - fn(0.0, ym, np.array([]))) / (2 * h) + np.testing.assert_allclose( + J[:, col], np.asarray(fd).flatten(), rtol=1e-4, atol=1e-6 + ) + + +class TestEvalTrajectory: + def _setup(self): + _, f = _make_fn() + n_t = 50 + ts = np.linspace(0.0, 1.0, n_t) + Y = np.vstack([np.linspace(0.1, 1.0, n_t), np.linspace(2.0, 3.0, n_t)]) + p = np.array([3.0, 4.0]) + return f, ts, Y, p + + def test_matches_per_column_loop(self): + f, ts, Y, p = self._setup() + out = f.eval_trajectory(ts, Y, p) + assert out.shape == (2, len(ts)) + for j, t in enumerate(ts): + np.testing.assert_array_equal(out[:, j], f(t, Y[:, j].copy(), p)) + + def test_accepts_c_and_f_order(self): + f, ts, Y, p = self._setup() + a = f.eval_trajectory(ts, np.ascontiguousarray(Y), p) + b = f.eval_trajectory(ts, np.asfortranarray(Y), p) + np.testing.assert_array_equal(a, b) + + def test_f_order_matches_per_column_loop(self): + # pins the zero-copy borrow path against ground truth + f, ts, Y, p = self._setup() + Yf = np.asfortranarray(Y) + out = f.eval_trajectory(ts, Yf, p) + for j, t in enumerate(ts): + np.testing.assert_array_equal(out[:, j], f(t, Yf[:, j].copy(), p)) + + def test_reversed_view_matches_per_column_loop(self): + # negative-stride views must gather, not borrow the raw buffer + f, ts, Y, p = self._setup() + Yr = np.asfortranarray(Y)[:, ::-1] + out = f.eval_trajectory(ts, Yr, p) + for j, t in enumerate(ts): + np.testing.assert_array_equal(out[:, j], f(t, Yr[:, j].copy(), p)) + + def test_shape_validation(self): + f, ts, Y, p = self._setup() + with pytest.raises(ValueError, match=r"Y.shape\[0\]"): + f.eval_trajectory(ts, Y[:1, :], p) + with pytest.raises(ValueError, match=r"len\(ts\)"): + f.eval_trajectory(ts[:-1], Y, p) + + +class TestGroup: + def _setup(self): + """Two outputs sharing the subexpression sin(y0)*a.""" + g = ExprGraph() + a = g.input_parameter("a") + y0 = g.state_vector(0, 1) + shared = g.mul(g.sin(y0), a) # shared NodeId + out1 = g.add(shared, g.scalar(1.0)) + out2 = g.mul(shared, g.scalar(2.0)) + group = g.compile_group({"plus_one": out1, "doubled": out2}) + f1 = g.compile(out1) + f2 = g.compile(out2) + return g, group, f1, f2 + + def test_outputs_match_siso_compiles(self): + _, group, f1, f2 = self._setup() + y = np.array([0.7]) + p = np.array([2.0]) + r1, r2 = group(0.0, y, p) + assert (r1 == f1(0.0, y, p)).all() + assert (r2 == f2(0.0, y, p)).all() + + def test_names_and_lens(self): + _, group, _, _ = self._setup() + assert group.names == ("plus_one", "doubled") + assert group.output_lens == [1, 1] + + def test_signature_surface(self): + # groups expose the same signature surface as CompiledFunction + _, group, _, _ = self._setup() + assert group.input_names == ("a",) + assert group.n_inputs == 1 + assert group.output_len == 2 + assert group.uses_y_dot is False + np.testing.assert_array_equal(group.pack({"a": 2.0}), [2.0]) + + def test_cse_observable_in_tape_length(self): + _, group, f1, f2 = self._setup() + # shared work appears once in the group tape + assert group.n_instructions < f1.n_instructions + f2.n_instructions + + def test_cse_dedupes_structurally_identical_outputs(self): + # The same subexpression built twice (distinct NodeIds, identical structure) + # collapses via cse() at compile, matching the shared-NodeId group's tape. + g = ExprGraph() + a = g.input_parameter("a") + s1 = g.mul(g.sin(g.state_vector(0, 1)), a) + s2 = g.mul(g.sin(g.state_vector(0, 1)), a) # structural twin of s1 + twin = g.compile_group( + {"x": g.add(s1, g.scalar(1.0)), "y": g.mul(s2, g.scalar(2.0))} + ) + shared = g.compile_group( + {"x": g.add(s1, g.scalar(1.0)), "y": g.mul(s1, g.scalar(2.0))} + ) + assert twin.n_instructions == shared.n_instructions + + def test_eval_trajectory(self): + _, group, f1, f2 = self._setup() + n_t = 20 + ts = np.linspace(0.0, 1.0, n_t) + Y = np.linspace(0.1, 1.0, n_t).reshape(1, n_t) + p = np.array([2.0]) + r1, r2 = group.eval_trajectory(ts, Y, p) + assert r1.shape == (1, n_t) and r2.shape == (1, n_t) + for j, t in enumerate(ts): + np.testing.assert_array_equal(r1[:, j], f1(t, Y[:, j].copy(), p)) + np.testing.assert_array_equal(r2[:, j], f2(t, Y[:, j].copy(), p)) + + def test_group_rejects_y_dot_at_call(self): + g = ExprGraph() + grp = g.compile_group({"r": g.state_vector_dot(0, 1)}, name="resid") + with pytest.raises(ValueError, match=r"resid.*y_dot"): + grp(0.0, np.zeros(1), np.array([])) + + def test_group_n_states_override_below_extent_rejected(self): + g = ExprGraph() + e = g.mul(g.state_vector(0, 3), g.state_vector(0, 3)) + with pytest.raises(ValueError, match=r"n_states=1.*state extent 3"): + g.compile_group({"out": e}, n_states=1) + + def test_compile_group_leaves_graph_unchanged(self): + # compiling a group must not mutate the shared graph arena + g = ExprGraph() + e = g.mul(g.sin(g.state_vector(0, 1)), g.input_parameter("a")) + before = g.n_nodes + g.compile_group({"out": e}) + assert g.n_nodes == before + + def test_n_instructions_excludes_conditional_branch_blocks(self): + # `n_instructions` reports the common tape plus one dispatch, matching + # what casadi's Switch-lowered outer function reports. + g = ExprGraph() + y0 = g.state_vector(0, 1) + selector = g.input_parameter("step") + small = g.sin(y0) + big = y0 + for _ in range(12): + big = g.exp(big) + f = g.compile(g.conditional(selector, [small, big])) + + assert f.branch_block_lens == (1, 12) + assert f.n_instructions_total == f.n_instructions + 13 + # common = y0 + selector + dispatch + conditional + assert f.n_instructions == 4 + # And the reported count says how much of itself is control flow. + assert f.n_dispatches == 1 + + def test_n_instructions_flat_in_branch_count(self): + # An extra branch must not grow the reported count. + def reported(n_branches): + g = ExprGraph() + y0 = g.state_vector(0, 1) + selector = g.input_parameter("step") + branches = [] + for i in range(n_branches): + node = y0 + for _ in range(3 + i): + node = g.exp(node) + branches.append(node) + return g.compile(g.conditional(selector, branches)).n_instructions + + assert reported(2) == reported(6) + + +class TestCompiledModelBundle: + def _setup(self): + from pybamm.rust import CompiledModel + + g = ExprGraph() + a = g.input_parameter("a") + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + rhs = g.concat([g.mul(a, y1), g.neg(y0)]) + out = g.mul(y0, y1) + event = g.sub(y0, g.scalar(10.0)) + mass = np.array([1.0, 1.0]) + indptr = np.array([0, 1, 2], dtype=np.int64) + indices = np.array([0, 1], dtype=np.int64) + model = CompiledModel.from_expr( + g, + rhs, + mass, + indptr, + indices, + n_inputs=1, + output_exprs=[out], + event_exprs=[event], + ) + return g, model + + def test_rhs_view(self): + _, model = self._setup() + f = model.rhs + y = np.array([1.0, 2.0]) + p = np.array([3.0]) + np.testing.assert_array_equal(f(0.0, y, p), [6.0, -1.0]) + assert f.n_states == 2 + + def test_jacobian_view_is_pure_df_dy(self): + _, model = self._setup() + J = model.jacobian + mat = J(0.0, np.array([1.0, 2.0]), np.array([3.0])) + # df/dy = [[0, a], [-1, 0]] — no mass, no cj + np.testing.assert_allclose(mat.toarray(), [[0.0, 3.0], [-1.0, 0.0]]) + + def test_outputs_and_events_views(self): + _, model = self._setup() + (out_fn,) = model.outputs + (event_fn,) = model.events + y = np.array([1.0, 2.0]) + p = np.array([3.0]) + np.testing.assert_array_equal(out_fn(0.0, y, p), [2.0]) + np.testing.assert_array_equal(event_fn(0.0, y, p), [-9.0]) + + def test_algebraic_views_none_for_ode(self): + _, model = self._setup() + assert model.algebraic_residual is None + assert model.algebraic_jacobian is None + + def test_algebraic_jacobian_view_is_the_compiled_artifact(self): + # The standalone dg/dy_alg view must lend the artifact the model already + # compiled, not re-derive one: a second tangent transform could disagree + # with the block the solver drives and nothing else would notice. + model = pybamm.lithium_ion.DFN() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model) + sim.build() + sim.solver.set_up(sim.built_model, inputs=[{}]) + rust_model = sim.solver._setup["rust_model"] + assert rust_model.has_algebraic + + view = rust_model.algebraic_jacobian + n_algebraic = rust_model.n_algebraic + assert view.shape == (n_algebraic, n_algebraic) + assert view.nnz == rust_model.algebraic_jacobian_nnz + rows, _ = rust_model.algebraic_jacobian_sparsity_pattern() + assert sorted(np.asarray(rows).tolist()) == sorted( + np.asarray(view.sparsity()[1]).tolist() + ) + + def test_views_are_cached_objects(self): + # accessors hand back the SAME prepared artifacts on every access — + # identity, not timing, is the contract + _, model = self._setup() + assert model.rhs is model.rhs + assert model.jacobian is model.jacobian + assert model.outputs[0] is model.outputs[0] + assert model.events[0] is model.events[0] + + def test_removed_methods_are_gone(self): + _, model = self._setup() + for gone in [ + # gone because it derived the C++ Newton driver's write-through + # address from a shared borrow; the FFI goes via evaluator_pool now + "as_ptr", + "eval_rhs", + "eval_rhs_into", + "jac_mul", + "jac_mul_into", + "assemble_jacobian", + "eval_sens_all", + "eval_sens", + "eval_output", + "output_len_at", + "output_lens", + "eval_events", + "eval_events_into", + "total_event_len", + ]: + assert not hasattr(model, gone), gone + + def test_kept_solver_surface(self): + _, model = self._setup() + assert model.n_states == 2 + assert model.nnz >= 2 + assert model.n_colors >= 1 + assert isinstance(model.evaluator_pool(1).as_ptr(0), int) + model.csc_sparsity_pattern() + stats = model.jacobian_stats() + assert stats["n_dense_rows"] == 0 + assert stats["dense_row_entries"] == 0 + assert stats["dense_row_tape_instructions"] == 0 + + +class TestJvpTrajectory: + def _setup(self): + _, f = _make_fn() + n_t = 20 + ts = np.linspace(0.0, 1.0, n_t) + Y = np.vstack([np.linspace(0.1, 1.0, n_t), np.linspace(2.0, 3.0, n_t)]) + VY = np.vstack([np.linspace(0.5, 1.5, n_t), np.linspace(-1.0, 1.0, n_t)]) + p = np.array([3.0, 4.0]) + return f, ts, Y, VY, p + + def test_matches_per_column_jvp(self): + # jvp_trajectory wrt y must equal the per-column jvp wrt y, bitwise: + # same tangent tape, same inputs. + f, ts, Y, VY, p = self._setup() + out = f.jvp_trajectory(ts, Y, p, VY) + assert out.shape == (2, len(ts)) + for j, t in enumerate(ts): + expected = f.jvp(t, Y[:, j].copy(), p, VY[:, j].copy()) + np.testing.assert_array_equal(out[:, j], expected) + + def test_with_vp_matches_per_column_jvp(self): + # With a parameter direction, each column equals jvp(.., vp=vp): + # df/dy @ vy_j + df/dp @ vp. + f, ts, Y, VY, p = self._setup() + vp = np.array([0.25, -1.0]) + out = f.jvp_trajectory(ts, Y, p, VY, vp=vp) + assert out.shape == (2, len(ts)) + for j, t in enumerate(ts): + expected = f.jvp(t, Y[:, j].copy(), p, VY[:, j].copy(), vp=vp) + np.testing.assert_array_equal(out[:, j], expected) + + def test_matches_numpy_jacobian_oracle(self): + # Independent oracle: build the full dvar_dy / dvar_dp jacobians per column + # and matmul dvar_dy @ vy + dvar_dp @ vp. + f, ts, Y, VY, p = self._setup() + vp = np.array([0.25, -1.0]) + Jy = f.jacobian("y") + Jp = f.jacobian("p") + out = f.jvp_trajectory(ts, Y, p, VY, vp=vp) + for j, t in enumerate(ts): + dvar_dy = Jy(t, Y[:, j].copy(), p).toarray() + dvar_dp = Jp(t, Y[:, j].copy(), p).toarray() + oracle = dvar_dy @ VY[:, j] + dvar_dp @ vp + np.testing.assert_allclose(out[:, j], oracle, rtol=1e-12, atol=1e-12) + + def test_preserves_zero_derivative_width(self): + # f = const_vector([1,2,3]) + p0, so df/dy is identically zero: the + # trajectory jvp wrt y must stay full width, not collapse to one row. + g = ExprGraph() + cv = g.array(np.array([1.0, 2.0, 3.0])) + f = g.compile(g.add(cv, g.input_parameter("a")), name="zeroy", n_states=2) + assert f.output_len == 3 + n_t = 5 + ts = np.linspace(0.0, 1.0, n_t) + Y = np.zeros((2, n_t)) + VY = np.ones((2, n_t)) + out = f.jvp_trajectory(ts, Y, np.array([5.0]), VY) + assert out.shape == (3, n_t) + np.testing.assert_array_equal(out, np.zeros((3, n_t))) + + def test_shape_validation(self): + f, ts, Y, VY, p = self._setup() + with pytest.raises(ValueError, match=r"Y.shape\[0\]"): + f.jvp_trajectory(ts, Y[:1, :], p, VY) + with pytest.raises(ValueError, match=r"len\(ts\)"): + f.jvp_trajectory(ts[:-1], Y, p, VY) + with pytest.raises(ValueError, match=r"vy_traj"): + f.jvp_trajectory(ts, Y, p, VY[:1, :]) + # the same guard's time-dimension arm: vy_traj.shape[1] != n_t + with pytest.raises(ValueError, match=r"vy_traj"): + f.jvp_trajectory(ts, Y, p, VY[:, :-1]) + + def test_accepts_c_and_f_order(self): + # zero-copy F path and gather C path must agree, for both y and vy. + f, ts, Y, VY, p = self._setup() + vp = np.array([0.25, -1.0]) + a = f.jvp_trajectory( + ts, np.ascontiguousarray(Y), p, np.ascontiguousarray(VY), vp=vp + ) + b = f.jvp_trajectory(ts, np.asfortranarray(Y), p, np.asfortranarray(VY), vp=vp) + np.testing.assert_array_equal(a, b) + + def test_reversed_view_matches_per_column_jvp(self): + # negative-stride views must gather, not borrow the raw buffer + f, ts, Y, VY, p = self._setup() + Yr = np.asfortranarray(Y)[:, ::-1] + VYr = np.asfortranarray(VY)[:, ::-1] + out = f.jvp_trajectory(ts, Yr, p, VYr) + for j, t in enumerate(ts): + expected = f.jvp(t, Yr[:, j].copy(), p, VYr[:, j].copy()) + np.testing.assert_array_equal(out[:, j], expected) + + def test_rejects_y_dot_expressions(self): + # same panic guard as TestJvp.test_jvp_rejects_y_dot_expressions: + # the tangent tape slices an empty y_dot, so y_dot funcs are rejected + g = ExprGraph() + f = g.compile(g.state_vector_dot(0, 1), name="resid") + with pytest.raises(ValueError, match=r"jvp_trajectory.*y_dot"): + f.jvp_trajectory( + np.zeros(1), np.zeros((1, 1)), np.zeros(0), np.zeros((1, 1)) + ) + + +class TestEvalTrajectoryHermite: + def _setup(self): + _, f = _make_fn() + n_knots = 12 + ts = np.linspace(0.0, 1.0, n_knots) + Y = np.vstack([np.sin(ts), np.cos(ts)]) + YP = np.vstack([np.cos(ts), -np.sin(ts)]) + p = np.array([3.0, 4.0]) + return f, ts, Y, YP, p + + def test_reduces_to_eval_on_knots(self): + # querying exactly at the knots reproduces eval_trajectory + f, ts, Y, YP, p = self._setup() + out = f.eval_trajectory_hermite(ts, ts, Y, YP, p) + ref = f.eval_trajectory(ts, Y, p) + assert out.shape == (2, len(ts)) + np.testing.assert_allclose(out, ref, rtol=1e-12, atol=1e-12) + + def test_linear_state_exact(self): + # Hermite reproduces a linear-in-t state exactly at off-grid points + _, f = _make_fn() + ts = np.linspace(0.0, 1.0, 5) + Y = np.vstack([2.0 * ts + 0.5, -3.0 * ts + 1.0]) + YP = np.vstack([np.full_like(ts, 2.0), np.full_like(ts, -3.0)]) + p = np.array([3.0, 4.0]) + tq = np.linspace(0.0, 1.0, 37) + out = f.eval_trajectory_hermite(tq, ts, Y, YP, p) + Yq = np.ascontiguousarray(np.vstack([2.0 * tq + 0.5, -3.0 * tq + 1.0])) + ref = f.eval_trajectory(tq, Yq, p) + np.testing.assert_allclose(out, ref, rtol=1e-12, atol=1e-12) + + def test_matches_scipy_cubic_hermite_oracle(self): + from scipy.interpolate import CubicHermiteSpline + + f, ts, Y, YP, p = self._setup() + tq = np.linspace(0.0, 1.0, 41) + spline = CubicHermiteSpline(ts, Y.T, YP.T) # (n_query, n_states) + Yq = np.ascontiguousarray(spline(tq).T) + ref = f.eval_trajectory(tq, Yq, p) + out = f.eval_trajectory_hermite(tq, ts, Y, YP, p) + np.testing.assert_allclose(out, ref, rtol=1e-10, atol=1e-10) + + def test_shape_validation(self): + f, ts, Y, YP, p = self._setup() + with pytest.raises(ValueError, match=r"yps shape"): + f.eval_trajectory_hermite(ts, ts, Y, YP[:1, :], p) + with pytest.raises(ValueError, match=r"len\(ts\)"): + f.eval_trajectory_hermite(ts, ts[:-1], Y, YP, p) + + +class TestGroupEvalTrajectoryHermite: + def test_matches_per_output_hermite(self): + # group hermite == per-CompiledFunction hermite, per output + _, group, f1, f2 = TestGroup()._setup() + n_knots = 10 + ts = np.linspace(0.0, 1.0, n_knots) + Y = np.sin(ts).reshape(1, n_knots) + YP = np.cos(ts).reshape(1, n_knots) + p = np.array([2.0]) + tq = np.linspace(0.0, 1.0, 31) + r1, r2 = group.eval_trajectory_hermite(tq, ts, Y, YP, p) + np.testing.assert_allclose( + r1, f1.eval_trajectory_hermite(tq, ts, Y, YP, p), rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + r2, f2.eval_trajectory_hermite(tq, ts, Y, YP, p), rtol=1e-12, atol=1e-12 + ) + + +class TestConcurrency: + def test_shared_function_and_jacobian_under_threads(self): + import concurrent.futures + + _, f = _make_fn() + J = f.jacobian() + rng = np.random.default_rng(42) + cases = [ + (float(t), rng.uniform(0.1, 2.0, 2), rng.uniform(0.5, 5.0, 2)) + for t in range(64) + ] + serial = [(f(t, y, p), J(t, y, p).toarray()) for t, y, p in cases] + + def work(case): + t, y, p = case + return f(t, y, p), J(t, y, p).toarray() + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + threaded = list(pool.map(work, cases * 4)) + + for i, (fv, jv) in enumerate(threaded): + ref_f, ref_j = serial[i % len(cases)] + assert (fv == ref_f).all(), "threaded eval must be bitwise-identical" + assert (jv == ref_j).all(), "threaded jacobian must be bitwise-identical" + + +class TestExtensionModuleLocation: + def test_extension_ships_inside_pybamm_package(self): + from pybamm.rust import _core + + assert Path(_core.__file__).parent == Path(pybamm.__file__).parent / "rust" + + def test_public_names_resolve_through_pybamm_rust(self): + import pybamm.rust + + for name in ( + "CompiledFunction", + "CompiledModel", + "ExprGraph", + "PreparedSolver", + ): + assert hasattr(pybamm.rust, name), name + + def test_every_core_class_is_re_exported(self): + import pybamm.rust + from pybamm.rust import _core + + registered = { + name + for name, obj in vars(_core).items() + if isinstance(obj, type) and not name.startswith("_") + } + assert registered <= set(pybamm.rust.__all__), sorted( + registered - set(pybamm.rust.__all__) + ) + + def test_pyclass_module_is_the_public_path(self): + import pybamm.rust + + for name in pybamm.rust.__all__: + obj = getattr(pybamm.rust, name) + if isinstance(obj, type): + assert obj.__module__ == "pybamm.rust", name + + def test_no_public_name_carries_the_py_prefix(self): + # `Py*` is PyO3's own namespace for its smart pointers and native type + # bindings; a leaked Rust-side prefix reads as a stutter from Python. + import pybamm.rust + + leaked = [name for name in pybamm.rust.__all__ if name.startswith("Py")] + assert not leaked, leaked + + def test_only_one_core_binary_is_installed(self): + # A stale version-specific .so would shadow the abi3 one at import time. + # The checked-in _core.pyi stub sits alongside and is not a binary. + import pybamm.rust + + artifacts = sorted( + p.name + for p in Path(pybamm.rust.__file__).parent.glob("_core*") + if p.suffix in (".so", ".pyd") + ) + assert len(artifacts) == 1, artifacts + + def test_pybamm_rust_distribution_is_gone(self): + with pytest.raises(ImportError): + importlib.import_module("pybamm_rust") diff --git a/packages/pybamm/tests/unit/test_rust_observability_benchmark.py b/packages/pybamm/tests/unit/test_rust_observability_benchmark.py new file mode 100644 index 0000000000..1acdd15bfd --- /dev/null +++ b/packages/pybamm/tests/unit/test_rust_observability_benchmark.py @@ -0,0 +1,1626 @@ +from __future__ import annotations + +import json +import pickle +from types import SimpleNamespace + +import numpy as np +import pytest + +import benchmarks.run_rust_observability as run_module +import pybamm +from benchmarks.run_rust_observability import build_parser, include_aot_for +from benchmarks.rust_observability import runners +from benchmarks.rust_observability.registry import ( + BASE_PARAMETER_SET, + CHARGE_C_RATE, + DEFAULT_OUTPUT_POINTS, + INFERENCE_COMPLEMENTS, + INFERENCE_INPUTS, + INFERENCE_SPREADS, + TRIANGLE_AMPLITUDE_A, + TRIANGLE_PERIOD_S, + get_inference_scenarios, + get_protocol_names, + get_solver_scenarios, + inference_nominal_values, +) +from benchmarks.rust_observability.report import ( + _fits, + render_inference_table, + render_sensitivity_table, + render_solver_table, + suite_to_jsonable, +) +from benchmarks.rust_observability.runners import ( + _BASELINE_CASE, + DEFAULT_REFERENCE_TOLERANCE, + AotProfile, + ComparisonSummary, + InferenceResult, + JacobianTelemetry, + PhaseTiming, + RepeatObservation, + SensitivityResult, + SolverResult, + TimingSamples, + TrajectorySummary, + _aot_worker_payload, + _build_and_time, + _build_sensitivity_parameters, + _comparable_length, + _get_jacobian_telemetry, + _observation_grid, + _observed_sensitivities, + _observed_values, + _reference_ladder, + _reference_tolerances, + _resolve_reference, + _sensitivity_tolerances, + _shuffled_backend_cases, + _solve_kwargs, + _summarize_cache_statuses, + _summarize_timing_samples, + _worst_repeat_comparison, + _worst_repeat_trajectory, + backend_cases, + resolve_reference_tolerance, + sample_input_vectors, + summarize_diff, + summarize_trajectory, +) + + +def make_repeat( + values, *, times=None, sensitivities=None, final_time=None, termination="final time" +): + """A ``RepeatObservation`` from just the values, for comparison tests.""" + values = np.asarray(values, dtype=np.float64) + times = np.arange(values.shape[0], dtype=np.float64) if times is None else times + return RepeatObservation( + values=values, + times=np.asarray(times, dtype=np.float64), + sensitivities=sensitivities, + sensitivity_times=None if sensitivities is None else times, + final_time=float(times[-1]) if final_time is None else final_time, + termination=termination, + ) + + +class TestProtocolRegistry: + def test_default_protocol_preserves_scenario_identity(self): + scenarios = get_solver_scenarios(["SPM", "SPMe", "DFN"]) + + assert [s.name for s in scenarios] == ["SPM", "SPMe", "DFN"] + assert {s.protocol for s in scenarios} == {"cc_discharge"} + assert all(s.initial_soc is None for s in scenarios) + assert all(s.plan.experiment is None for s in scenarios) + + def test_cross_product_is_model_major(self): + scenarios = get_solver_scenarios( + ["SPM", "DFN"], ["cc_discharge", "drive_cycle"] + ) + + assert [(s.name, s.protocol) for s in scenarios] == [ + ("SPM", "cc_discharge"), + ("SPM", "drive_cycle"), + ("DFN", "cc_discharge"), + ("DFN", "drive_cycle"), + ] + + def test_unknown_protocol_is_rejected(self): + with pytest.raises(ValueError, match="Unknown names requested: bogus"): + get_solver_scenarios(["SPM"], ["bogus"]) + + def test_interpolant_protocols_pass_breakpoints_as_t_eval(self): + for protocol in ("drive_cycle", "pulse_train"): + scenario = get_solver_scenarios(["SPM"], [protocol])[0] + parameter_values = scenario.parameter_values_builder() + current = parameter_values["Current function [A]"] + + assert isinstance(current, pybamm.Interpolant) + assert scenario.initial_soc == 0.5 + # Every breakpoint must be in t_eval or PyBaMM warns about resolution. + breakpoints = np.asarray(current.x[0], dtype=np.float64) + assert np.isin( + np.round(breakpoints, 12), np.round(scenario.plan.t_eval, 12) + ).all() + + def test_triangle_wave_is_exact_through_its_vertices(self): + scenario = get_solver_scenarios(["SPM"], ["drive_cycle"])[0] + current = scenario.parameter_values_builder()["Current function [A]"] + # Interpolant.x is a list of arrays; Interpolant.y is a plain ndarray. + vertices = np.asarray(current.x[0], dtype=np.float64) + values = np.asarray(current.y, dtype=np.float64).reshape(-1) + + dense_t = np.linspace(vertices[0], vertices[-1], 2001) + expected = ( + TRIANGLE_AMPLITUDE_A + * (2.0 / np.pi) + * np.arcsin(np.sin(2.0 * np.pi * dense_t / TRIANGLE_PERIOD_S)) + ) + np.testing.assert_allclose( + np.interp(dense_t, vertices, values), expected, atol=1e-12 + ) + + def test_charge_and_experiment_plans(self): + charge = get_solver_scenarios(["SPM"], ["cc_charge"])[0] + assert charge.initial_soc == 0.0 + parameter_values = charge.parameter_values_builder() + current = parameter_values["Current function [A]"] + nominal_capacity = float(parameter_values["Nominal cell capacity [A.h]"]) + assert float(current) == pytest.approx(-CHARGE_C_RATE * nominal_capacity) + + experiment = get_solver_scenarios(["SPM"], ["experiment"])[0] + assert experiment.plan.experiment is not None + assert experiment.plan.t_interp is None + # Period-driven grid, so the output-points knob does not apply. + assert experiment.plan.requested_points == 0 + + def test_output_points_only_resizes_grid_protocols(self): + grid = get_solver_scenarios(["SPM"], ["cc_discharge"], output_points=250)[0] + assert grid.plan.t_interp.size == 250 + assert grid.plan.requested_points == 250 + + experiment = get_solver_scenarios(["SPM"], ["experiment"], output_points=250)[0] + assert experiment.plan.t_interp is None + + +class TestRustObservabilityBenchmark: + def test_output_points_are_explicit(self): + default = get_solver_scenarios(["SPM"])[0] + dense = get_solver_scenarios(["SPM"], output_points=1000)[0] + + assert default.plan.t_interp.size == DEFAULT_OUTPUT_POINTS + assert dense.plan.t_interp.size == 1000 + np.testing.assert_array_equal( + default.plan.t_interp[[0, -1]], dense.plan.t_interp[[0, -1]] + ) + + with pytest.raises(ValueError, match="at least 2"): + get_solver_scenarios(["SPM"], output_points=1) + + def test_e2e_uses_paired_wall_samples(self): + samples = TimingSamples( + warm_set_up_ms=(0.1, 0.1, 0.1), + solve_ms=(1.0, 2.0, 100.0), + wall_solve_ms=(1.1, 2.1, 100.1), + integration_ms=(0.9, 1.9, 99.9), + observe_ms=(100.0, 2.0, 1.0), + e2e_ms=(101.1, 4.1, 101.1), + ) + + timing = _summarize_timing_samples( + samples, + build_ms=3.0, + cold_set_up_ms=4.0, + cold_observe_ms=5.0, + cold_startup_ms=12.0, + ) + + assert timing.prepare_ms == pytest.approx(9.0) + assert timing.cold_startup_ms == pytest.approx(12.0) + assert timing.cold_startup_ms >= timing.build_ms + timing.prepare_ms + assert timing.solve_ms == pytest.approx(2.0) + assert timing.observe_ms == pytest.approx(2.0) + assert timing.e2e_ms == pytest.approx(101.1) + assert timing.e2e_ms != pytest.approx(timing.solve_ms + timing.observe_ms) + + def test_trajectory_rejects_early_termination(self): + baseline = SimpleNamespace( + t=np.array([0.0, 1.0, 2.0]), termination="final time" + ) + candidate = SimpleNamespace(t=np.array([0.0, 1.0]), termination="event") + + summary = summarize_trajectory(baseline, candidate, atol=1e-9, rtol=1e-9) + + assert summary.status == "warn" + assert summary.common_points == 2 + assert summary.coverage == pytest.approx(2 / 3) + + def test_trajectory_allows_terminal_roundoff(self): + baseline = SimpleNamespace( + t=np.array([0.0, 1.0, 2.0]), termination="final time" + ) + candidate = SimpleNamespace( + t=np.array([0.0, 1.0, 2.0 + 1e-6]), termination="final time" + ) + + summary = summarize_trajectory(baseline, candidate, atol=1e-9, rtol=1e-6) + + assert summary.status == "pass" + assert summary.coverage == 1.0 + + def test_shape_mismatch_warns(self): + summary = summarize_diff( + np.zeros((2, 3)), + np.zeros((1, 3)), + atol=1e-9, + rtol=1e-9, + ) + + assert summary.status == "warn" + assert np.isinf(summary.max_abs_diff) + + def test_report_order_compares_backends_within_output_mode(self): + shuffled_backends = [ + "rust_diffsol_out", + "rust_idaklu", + "casadi_idaklu_aot_out", + "casadi_idaklu", + "rust_diffsol", + "casadi_idaklu_out", + "rust_idaklu_out", + "casadi_idaklu_aot", + ] + results = [ + SolverResult( + scenario="SPM", + backend=backend, + timings=PhaseTiming(), + requested_output_points=100, + ) + for backend in shuffled_backends + ] + expected = [ + "casadi_idaklu", + "casadi_idaklu_aot", + "rust_idaklu", + "rust_diffsol", + "casadi_idaklu_out", + "casadi_idaklu_aot_out", + "rust_idaklu_out", + "rust_diffsol_out", + ] + + # Wide enough that the layout is the one-table one, whatever columns it + # carries, so this asserts on ordering rather than on the layout choice. + report = render_solver_table(results, width=250) + rendered = [ + line.split()[1] for line in report.splitlines() if line.startswith("SPM ") + ] + assert rendered == expected + + payload = suite_to_jsonable("solver", results) + assert [result["backend"] for result in payload["results"]] == expected + + def test_aot_profile_is_rendered_and_serialized(self): + profile = AotProfile( + fresh_cache_statuses=("miss", "miss"), + disk_cache_statuses=("disk", "disk"), + codegen_ms=100.0, + compiler_ms=17000.0, + fresh_load_ms=1.0, + disk_load_ms=2.0, + fresh_total_ms=17101.0, + disk_total_ms=3.0, + disk_prepare_ms=45.0, + disk_cold_startup_ms=110.0, + library_size_bytes=2 * 1024**2, + verified=True, + ) + result = SolverResult( + scenario="DFN", + backend="casadi_idaklu_aot", + timings=PhaseTiming(prepare_ms=17150.0, cold_startup_ms=17220.0), + requested_output_points=100, + aot_profile=profile, + ) + + report = render_solver_table([result], width=120) + assert "AOT profile (isolated cache" in report + assert "missx2" in report + assert "diskx2" in report + assert "17000.00" in report + + payload = suite_to_jsonable("solver", [result]) + encoded = json.loads(json.dumps(payload)) + assert encoded["results"][0]["aot_profile"]["verified"] is True + assert encoded["results"][0]["aot_profile"]["compiler_ms"] == 17000.0 + + def test_sensitivity_table_matches_solver_width_layout(self): + trajectory = TrajectorySummary( + baseline_points=100, + candidate_points=100, + common_points=100, + coverage=1.0, + max_time_diff=0.0, + final_time_diff=0.0, + baseline_termination="final time", + candidate_termination="final time", + status="pass", + ) + comparison = ComparisonSummary(0.0, 0.0, 0.0) + timing = PhaseTiming( + build_ms=22.81, + prepare_ms=29.49, + cold_startup_ms=57.86, + warm_set_up_ms=0.08, + solve_ms=4.43, + wall_solve_ms=4.55, + integration_ms=4.32, + observe_ms=16.39, + e2e_ms=20.97, + ) + results = [ + SensitivityResult( + scenario="SPM", + backend="casadi_idaklu", + timings=timing, + requested_output_points=1000, + ), + SensitivityResult( + scenario="SPM", + backend="casadi_idaklu_aot", + timings=timing, + requested_output_points=1000, + state_sens_comparison=comparison, + output_sens_comparison=comparison, + trajectory_comparison=trajectory, + ), + ] + + compact = render_sensitivity_table(results, width=120) + assert "Timings (ms)" in compact + assert "Validation" in compact + assert _fits(compact, 120) + + stacked = render_sensitivity_table(results, width=80) + assert "One block per backend" in stacked + assert _fits(stacked, 80) + + def test_solver_table_snapshot_and_json_samples(self, monkeypatch, snapshot): + trajectory = TrajectorySummary( + baseline_points=100, + candidate_points=100, + common_points=100, + coverage=1.0, + max_time_diff=0.0, + final_time_diff=0.0, + baseline_termination="event", + candidate_termination="event", + status="pass", + ) + comparison = ComparisonSummary(1e-8, 1e-7, 0.01) + samples = TimingSamples( + warm_set_up_ms=(0.1,), + solve_ms=(1.2,), + wall_solve_ms=(1.3,), + integration_ms=(1.1,), + observe_ms=(0.2,), + e2e_ms=(1.5,), + ) + timing = PhaseTiming( + build_ms=10.0, + prepare_ms=5.2, + cold_startup_ms=16.5, + set_up_ms=5.0, + warm_set_up_ms=0.1, + solve_ms=1.2, + wall_solve_ms=1.3, + integration_ms=1.1, + observe_ms=0.2, + e2e_ms=1.5, + ) + telemetry = { + "SPM": JacobianTelemetry( + strategy="coloring", + n_colors=5, + nnz=123, + n_dense_rows=0, + dense_row_entries=0, + dense_row_tape_instructions=0, + split_eval_primal_instructions=10, + split_eval_total_instructions=20, + split_eval_raw_instructions=20, + branch_block_lens=(), + ), + "SPMe": JacobianTelemetry( + strategy="coloring", + n_colors=3, + nnz=361, + n_dense_rows=1, + dense_row_entries=65, + dense_row_tape_instructions=4096, + split_eval_primal_instructions=10, + split_eval_total_instructions=20, + split_eval_raw_instructions=20, + branch_block_lens=(), + ), + "DFN": JacobianTelemetry( + strategy="coloring", + n_colors=9, + nnz=3673, + n_dense_rows=0, + dense_row_entries=0, + dense_row_tape_instructions=0, + split_eval_primal_instructions=10, + split_eval_total_instructions=20, + split_eval_raw_instructions=20, + branch_block_lens=(), + ), + } + results = [ + SolverResult( + scenario=name, + protocol="cc_discharge", + backend="rust_idaklu", + timings=timing, + requested_output_points=100, + timing_samples=samples, + state_comparison=comparison, + output_comparison=comparison, + trajectory_comparison=trajectory, + jacobian_telemetry=stats, + ) + for name, stats in telemetry.items() + ] + + compact = render_solver_table(results, width=120) + snapshot.assert_match( + compact + "\n", + "rust_observability_solver_table.snapshot", + ) + assert _fits(compact, 120) + monkeypatch.setenv("COLUMNS", "120") + assert render_solver_table(results) == compact + + wide = render_solver_table(results, width=240) + assert "Prep" in wide + assert "Cold" in wide + assert "Validation\n" not in wide + + narrow = render_solver_table(results, width=80) + assert _fits(narrow, 80) + assert "One block per backend" in narrow + + payload = suite_to_jsonable("solver", results) + encoded = json.loads(json.dumps(payload)) + assert encoded["results"][0]["timings"]["prepare_ms"] == 5.2 + assert encoded["results"][0]["timings"]["cold_startup_ms"] == 16.5 + assert encoded["results"][0]["timing_samples"]["e2e_ms"] == [1.5] + + +class TestProtocolSolveWiring: + def test_grid_protocol_solve_kwargs(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + + kwargs = _solve_kwargs(scenario) + + assert kwargs["t_eval"] == [0.0, 3600.0] + np.testing.assert_array_equal(kwargs["t_interp"], scenario.plan.t_interp) + assert "experiment" not in kwargs + + def test_interpolant_protocol_passes_breakpoints(self): + scenario = get_solver_scenarios(["SPM"], ["pulse_train"])[0] + + kwargs = _solve_kwargs(scenario) + + np.testing.assert_array_equal(kwargs["t_eval"], scenario.plan.t_eval) + + def test_experiment_protocol_omits_time_arguments(self): + scenario = get_solver_scenarios(["SPM"], ["experiment"])[0] + + kwargs = _solve_kwargs(scenario) + + assert kwargs == {} + + def test_extra_kwargs_are_merged(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + + kwargs = _solve_kwargs(scenario, {"inputs": {"D_n": 1.0}}) + + assert kwargs["inputs"] == {"D_n": 1.0} + assert kwargs["t_eval"] == [0.0, 3600.0] + + +class _BuildSpy: + """Stands in for a ``pybamm.Simulation``, recording ``build`` calls.""" + + def __init__(self): + self.build_calls: list[tuple[float | None, dict | None]] = [] + + def build(self, initial_soc=None, inputs=None): + self.build_calls.append((initial_soc, inputs)) + + +class TestExperimentBuildInteraction: + """Regression tests for two bugs surfaced by wiring the experiment protocol. + + Pre-building an experiment-attached ``Simulation`` reparameterises the same + model that ``Simulation.solve`` then parameterises again per step, tripping + PyBaMM's reparameterised-model guard and making every variable + unprocessable. Separately, "Current function [A]" cannot be a sensitivity + input once an ``Experiment`` supplies its own control law. + """ + + def test_build_and_time_skips_build_for_experiment_plan(self): + scenario = get_solver_scenarios(["SPM"], ["experiment"])[0] + simulation = _BuildSpy() + + build_ms = _build_and_time(simulation, scenario) + + # No `.build()` call means this is pure guard-check overhead, not a real build. + assert build_ms < 1.0 + assert simulation.build_calls == [] + + def test_build_and_time_builds_grid_plan_with_initial_soc(self): + scenario = get_solver_scenarios(["SPM"], ["cc_charge"])[0] + simulation = _BuildSpy() + + _build_and_time(simulation, scenario) + + assert simulation.build_calls == [(scenario.initial_soc, None)] + + def test_build_and_time_forwards_inputs_for_the_initial_state(self): + scenario = get_solver_scenarios(["SPM"], ["cc_charge"])[0] + simulation = _BuildSpy() + inputs = {"eps_p": 0.66} + + _build_and_time(simulation, scenario, inputs=inputs) + + # Mapping initial_soc to concentrations runs an ElectrodeSOH solve, which + # cannot evaluate a symbolic parameter without its value. + assert simulation.build_calls == [(0.0, inputs)] + + def test_sensitivity_parameters_drop_dead_current_input_under_experiment(self): + scenario = get_solver_scenarios(["SPM"], ["experiment"])[0] + parameter_values, inputs = _build_sensitivity_parameters(scenario) + + assert "I" not in inputs + assert "eps_p" in inputs + assert not isinstance( + parameter_values["Current function [A]"], pybamm.InputParameter + ) + + def test_sensitivity_parameters_include_current_input_for_grid_protocol(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + parameter_values, inputs = _build_sensitivity_parameters(scenario) + + assert "I" in inputs + assert isinstance( + parameter_values["Current function [A]"], pybamm.InputParameter + ) + + def test_jacobian_telemetry_handles_never_set_up_solver(self): + # `_setup` is only assigned inside `set_up()`, never in `__init__`; the + # experiment path solves via a per-step copy, leaving this one bare. + solver = pybamm.IDAKLUSolver() + + assert _get_jacobian_telemetry(solver, "rust_idaklu") is None + + +class TestAotGating: + def test_backend_cases_drop_aot_when_disabled(self): + with_aot = backend_cases(include_aot=True) + without_aot = backend_cases(include_aot=False) + + assert any(backend == "casadi_idaklu_aot" for backend, _ in with_aot) + assert not any(backend == "casadi_idaklu_aot" for backend, _ in without_aot) + assert len(with_aot) - len(without_aot) == 2 # full-state and output-only rows + + def test_disk_worker_payload_is_names_not_objects(self): + scenario = get_solver_scenarios(["SPM"], ["drive_cycle"])[0] + + payload = _aot_worker_payload( + "solver", scenario, output_only=False, cache_dir="/tmp/x", output_points=100 + ) + + # Rebuilt by name in the worker, so no PyBaMM object crosses the boundary. + assert payload == ("solver", "SPM", "drive_cycle", 100, False, "/tmp/x") + pickle.loads(pickle.dumps(payload)) + + +class TestInferenceLane: + def test_input_vectors_are_seed_reproducible_and_varying(self): + nominal = {"D_n": 3.3e-14, "eps_p": 0.665} + + first = sample_input_vectors(nominal, 5, seed=0) + again = sample_input_vectors(nominal, 5, seed=0) + other = sample_input_vectors(nominal, 5, seed=1) + + assert first == again # identical across backends for a given seed + assert first != other + assert len(first) == 5 + assert all(set(vector) == set(nominal) for vector in first) + # Every repeat differs, or the lane is not measuring changing inputs. + assert len({vector["D_n"] for vector in first}) == 5 + + def test_input_vectors_stay_within_spread(self): + nominal = {"D_n": 3.3e-14} + + vectors = sample_input_vectors(nominal, 200, seed=3, spread=0.2) + + values = np.array([vector["D_n"] for vector in vectors]) + assert values.min() >= 3.3e-14 * 0.8 + assert values.max() <= 3.3e-14 * 1.2 + + def test_spread_can_differ_per_parameter(self): + """One width cannot suit parameters of different natures. + + 20% is routine for a diffusivity but takes a volume fraction to a porosity + DFN cannot converge at and SPMe saturates on charge. + """ + nominal = {"D_n": 3.3e-14, "eps_n": 0.75} + + vectors = sample_input_vectors( + nominal, 200, seed=3, spread={"D_n": 0.2, "eps_n": 0.05} + ) + + diffusivity = np.array([vector["D_n"] for vector in vectors]) + fraction = np.array([vector["eps_n"] for vector in vectors]) + assert diffusivity.min() < 3.3e-14 * 0.85 + assert fraction.min() >= 0.75 * 0.95 + assert fraction.max() <= 0.75 * 1.05 + + def test_a_missing_per_parameter_spread_is_rejected(self): + # A new fitted parameter must not silently inherit another's width. + with pytest.raises(KeyError): + sample_input_vectors( + {"D_n": 3.3e-14, "eps_n": 0.75}, 2, seed=0, spread={"D_n": 0.2} + ) + + def test_every_fitted_parameter_has_a_spread(self): + assert set(INFERENCE_SPREADS) == set(INFERENCE_INPUTS.values()) + + def test_sampled_porosity_stays_physical(self): + """Porosity is the complement of a fitted fraction, so its width is derived. + + At 20% it ranged 0.165-0.39 about a nominal 0.25, well outside anything the + base set describes. + """ + vectors = sample_input_vectors( + inference_nominal_values(), 200, seed=0, spread=INFERENCE_SPREADS + ) + + for input_name in INFERENCE_COMPLEMENTS.values(): + porosity = np.array([1.0 - vector[input_name] for vector in vectors]) + assert porosity.min() > 0.2 + assert porosity.max() < 0.4 + + def test_inference_scenario_makes_every_parameter_an_input(self): + scenario = get_inference_scenarios(["SPM"])[0] + parameter_values = scenario.parameter_values_builder() + + for pybamm_name, input_name in INFERENCE_INPUTS.items(): + assert isinstance(parameter_values[pybamm_name], pybamm.InputParameter) + assert parameter_values[pybamm_name].name == input_name + + def test_inference_scenarios_keep_protocol_identity(self): + scenarios = get_inference_scenarios(["SPM", "DFN"], ["cc_discharge"]) + + assert [(s.name, s.protocol) for s in scenarios] == [ + ("SPM", "cc_discharge"), + ("DFN", "cc_discharge"), + ] + + @pytest.mark.parametrize( + "protocol", ["cc_discharge", "cc_charge", "drive_cycle", "pulse_train"] + ) + def test_inference_preserves_the_protocol_control_law(self, protocol): + """The inference lane layers inputs on the protocol, it does not replace it. + + Replacing the builder silently reverted every protocol to a plain + constant-current discharge while the table still named the protocol. + """ + solver = get_solver_scenarios(["SPM"], [protocol])[0] + inference = get_inference_scenarios(["SPM"], [protocol])[0] + key = "Current function [A]" + + expected = solver.parameter_values_builder()[key] + actual = inference.parameter_values_builder()[key] + + assert type(actual) is type(expected) + if isinstance(expected, pybamm.Interpolant): + np.testing.assert_array_equal(actual.x[0], expected.x[0]) + np.testing.assert_array_equal(actual.y, expected.y) + else: + assert float(actual) == pytest.approx(float(expected)) + + def test_inference_keeps_electrode_volume_fractions_feasible(self): + scenario = get_inference_scenarios(["SPM"], ["cc_discharge"])[0] + parameter_values = scenario.parameter_values_builder() + nominal = inference_nominal_values() + + for porosity_name, input_name in INFERENCE_COMPLEMENTS.items(): + porosity = parameter_values[porosity_name] + # Solid and pore volume must still sum to 1 at any sampled value. + assert porosity.evaluate(inputs={input_name: 0.9}) == pytest.approx(0.1) + assert porosity.evaluate( + inputs={input_name: nominal[input_name]} + ) == pytest.approx(1.0 - nominal[input_name]) + + def test_every_protocol_shares_one_base_parameter_set(self): + base = pybamm.ParameterValues(BASE_PARAMETER_SET) + + for protocol in get_protocol_names(): + scenario = get_solver_scenarios(["SPM"], [protocol])[0] + parameter_values = scenario.parameter_values_builder() + # Only the current law may differ, or rows are not comparable. + assert float(parameter_values["Nominal cell capacity [A.h]"]) == float( + base["Nominal cell capacity [A.h]"] + ) + + def test_nominal_values_cover_every_fitted_parameter(self): + nominal = inference_nominal_values() + + assert set(nominal) == set(INFERENCE_INPUTS.values()) + assert all(np.isfinite(value) and value > 0.0 for value in nominal.values()) + + def test_observation_grid_interpolates_between_solver_nodes(self): + scenario = get_inference_scenarios(["SPM"], ["cc_discharge"], output_points=5)[ + 0 + ] + + nodes = scenario.plan.t_interp + grid = _observation_grid(scenario) + + # Midpoints, so no observation time coincides with a solver output node. + assert not np.isin(grid, nodes).any() + # One per interval, less the first: it straddles the initial transient, + # which no two-point interpolant can reach from its endpoints. + assert grid.size == nodes.size - 2 + assert grid[0] > nodes[1] + np.testing.assert_allclose(grid, 0.5 * (nodes[1:-1] + nodes[2:])) + + def test_observation_grid_is_empty_without_a_declared_grid(self): + scenario = get_inference_scenarios(["SPM"], ["experiment"])[0] + + assert _observation_grid(scenario).size == 0 + + def test_inference_result_reports_spread_and_one_time_costs(self): + result = InferenceResult( + scenario="SPM", + protocol="cc_discharge", + backend="rust_idaklu", + build_ms=10.0, + setup_ms=5.0, + aot_cache_status="disk", + eval_samples_ms=(1.0, 2.0, 3.0, 4.0, 5.0), + solve_samples_ms=(0.8, 1.8, 2.8, 3.8, 4.8), + observe_samples_ms=(0.2, 0.2, 0.2, 0.2, 0.2), + requested_output_points=100, + ) + + assert result.eval_median_ms == pytest.approx(3.0) + assert result.eval_p10_ms == pytest.approx(1.4) + assert result.eval_p90_ms == pytest.approx(4.6) + assert result.solve_median_ms == pytest.approx(2.8) + assert result.observe_median_ms == pytest.approx(0.2) + assert result.status == "baseline" + + def test_unsupported_inference_result_reports_reason(self): + result = InferenceResult( + scenario="SPM", + protocol="experiment", + backend="rust_diffsol", + build_ms=0.0, + setup_ms=0.0, + aot_cache_status="-", + eval_samples_ms=(), + solve_samples_ms=(), + observe_samples_ms=(), + requested_output_points=0, + supported=False, + reason="SolverError: nope", + ) + + assert result.status == "unsupported" + + def test_cache_status_reports_what_the_compiler_did(self): + assert _summarize_cache_statuses(None) == "-" + assert _summarize_cache_statuses([]) == "-" + assert ( + _summarize_cache_statuses([SimpleNamespace(cache_status="miss")] * 3) + == "miss" + ) + # A warm in-process reuse must not be reported as a fresh compile. + assert ( + _summarize_cache_statuses( + [ + SimpleNamespace(cache_status="memory"), + SimpleNamespace(cache_status="disk"), + ] + ) + == "disk+memory" + ) + + def _compare(self, baseline, candidate, scenario): + return _worst_repeat_comparison( + baseline, + candidate, + select=_observed_values, + atol=scenario.atol, + rtol=scenario.rtol, + ) + + def test_ragged_repeats_compare_on_the_aligned_prefix(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + # Both grids are prefixes of the same observation times, so trimming to + # the shorter one compares like for like rather than shifting by one. + trace = np.linspace(4.0, 3.0, 12) + times = np.arange(12.0) + summary = self._compare( + [make_repeat(trace, times=times)], + [make_repeat(trace[:9], times=times[:9], final_time=11.0)], + scenario, + ) + + assert summary.max_abs_diff == pytest.approx(0.0) + assert summary.status == "pass" + + def test_worst_repeat_drives_the_comparison(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + off = np.full(10, 4.0) + off[0] = 3.6 + # Second repeat is far off; the worst case must win, not the last. + summary = self._compare( + [make_repeat(np.full(10, 4.0)), make_repeat(np.full(10, 4.0))], + [make_repeat(np.full(10, 4.0)), make_repeat(off)], + scenario, + ) + + assert summary.max_abs_diff == pytest.approx(0.4) + assert summary.status == "warn" + + def test_a_failing_repeat_is_never_masked_by_a_larger_permitted_one(self): + """Tolerance scales with the baseline magnitude, so a repeat sitting at a + lower voltage can breach it on a *smaller* absolute difference than one + sitting higher up is allowed. + """ + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + tight = np.full(20, 2.5) + tight_off = tight.copy() + tight_off[5] += 4.0e-6 # breaches atol + rtol * 2.5 + loose = np.full(20, 4.2) + loose_off = loose.copy() + loose_off[5] += 4.5e-6 # larger, but inside atol + rtol * 4.2 + + assert self._compare([make_repeat(loose)], [make_repeat(loose_off)], scenario) + assert ( + self._compare([make_repeat(tight)], [make_repeat(tight_off)], scenario) + ).status == "warn" + summary = self._compare( + [make_repeat(tight), make_repeat(loose)], + [make_repeat(tight_off), make_repeat(loose_off)], + scenario, + ) + + assert summary.status == "warn" + assert summary.max_abs_diff == pytest.approx(4.0e-6) + + def test_endpoint_window_is_measured_not_a_fixed_point_count(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + times = np.arange(10.0) + # Terminations 2 s apart, so the last two points measure the endpoint gap. + baseline = make_repeat(np.full(10, 4.0), times=times, final_time=9.0) + candidate = make_repeat(np.full(10, 4.0), times=times, final_time=7.0) + candidate.values[-2:] = 3.0 + + assert _comparable_length(times, times, endpoint_gap=0.0) == 10 + assert _comparable_length(times, times, endpoint_gap=2.0) == 8 + assert self._compare([baseline], [candidate], scenario).status == "pass" + + def test_matching_terminations_drop_nothing(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + times = np.arange(10.0) + candidate = make_repeat(np.full(10, 4.0), times=times) + candidate.values[-1] = 3.0 + baseline = make_repeat(np.full(10, 4.0), times=times) + + # Terminations coincide, so a final-point disagreement is real, not tail + # noise, and must not be trimmed away. + assert self._compare([baseline], [candidate], scenario).status == "warn" + + def test_early_termination_cannot_pass_on_a_matching_prefix(self): + """A candidate stopping a tenth of the way in agrees perfectly wherever + both were observed, so coverage and termination have to carry the verdict. + """ + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + times = np.linspace(0.0, 3600.0, 1000) + trace = np.linspace(4.2, 2.5, 1000) + baseline = make_repeat(trace, times=times, termination="final time") + early = make_repeat( + trace[:100], times=times[:100], termination="event: Minimum voltage [V]" + ) + + # Nothing outlives an endpoint gap this wide, so the values are not + # comparable at all rather than agreeing on a truncated prefix. + values = self._compare([baseline], [early], scenario) + assert values.status == "warn" + assert values.max_abs_diff == float("inf") + + trajectory = _worst_repeat_trajectory([baseline], [early], scenario) + assert trajectory.status == "warn" + assert trajectory.coverage == pytest.approx(0.1) + assert trajectory.candidate_termination.startswith("event") + assert trajectory.final_time_diff == pytest.approx(times[-1] - times[99]) + + def test_gradient_comparison_is_skipped_when_not_requested(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + atol, rtol = _sensitivity_tolerances(scenario) + repeats = [make_repeat(np.full(4, 4.0))] + + assert ( + _worst_repeat_comparison( + repeats, repeats, select=_observed_sensitivities, atol=atol, rtol=rtol + ) + is None + ) + + def test_gradient_tolerance_is_looser_than_the_state_tolerance(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + atol, rtol = _sensitivity_tolerances(scenario) + + # Forward sensitivities are not error-controlled to the state tolerance, + # and the gate sits a decade above that noise floor rather than on it. + assert atol == pytest.approx(10.0 * scenario.atol**0.5) + assert rtol == pytest.approx(10.0 * scenario.rtol**0.5) + assert atol > scenario.atol and rtol > scenario.rtol + + def test_gradient_comparison_catches_a_broken_chain_rule(self): + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + atol, rtol = _sensitivity_tolerances(scenario) + gradient = np.tile(np.array([[0.34, 18.0]]), (6, 1)) + baseline = [make_repeat(np.full(6, 4.0), sensitivities=gradient)] + noisy = [ + make_repeat(np.full(6, 4.0), sensitivities=gradient + 2.0e-6) + ] # integrator noise + broken = [make_repeat(np.full(6, 4.0), sensitivities=gradient * 0.2)] + + def compare(repeats): + return _worst_repeat_comparison( + baseline, repeats, select=_observed_sensitivities, atol=atol, rtol=rtol + ) + + assert compare(noisy).status == "pass" + assert compare(broken).status == "warn" + + +class TestProtocolReporting: + def test_protocol_column_is_rendered_and_serialized(self): + results = [ + SolverResult( + scenario="SPM", + protocol=protocol, + backend="rust_idaklu", + timings=PhaseTiming(build_ms=1.0), + requested_output_points=100, + ) + for protocol in ("cc_discharge", "drive_cycle") + ] + + table = render_solver_table(results, width=240) + + assert "Protocol" in table + assert "drive_cycle" in table + payload = suite_to_jsonable("solver", results) + assert [r["protocol"] for r in payload["results"]] == [ + "cc_discharge", + "drive_cycle", + ] + + def test_period_driven_grid_reports_no_point_count(self): + results = [ + SolverResult( + scenario="SPM", + protocol="experiment", + backend="rust_idaklu", + timings=PhaseTiming(build_ms=1.0), + requested_output_points=0, + ) + ] + + # 0 means "the protocol owns its grid", which must not print as "0". + assert " 0 " not in render_solver_table(results, width=240) + assert "-" in render_solver_table(results, width=240) + + def test_inference_table_shows_one_time_and_per_eval_costs(self): + results = [ + InferenceResult( + scenario="SPM", + protocol="cc_discharge", + backend="casadi_idaklu", + build_ms=10.0, + setup_ms=5.0, + aot_cache_status="-", + eval_samples_ms=(1.0, 2.0, 3.0), + solve_samples_ms=(0.8, 1.8, 2.8), + observe_samples_ms=(0.2, 0.2, 0.2), + requested_output_points=99, + ) + ] + + table = render_inference_table(results, width=200) + + assert "Build" in table and "Setup" in table + assert "Eval p50" in table + assert "p10-p90" in table + assert "baseline" in table + payload = suite_to_jsonable("inference", results) + encoded = json.loads(json.dumps(payload)) + assert encoded["results"][0]["eval_samples_ms"] == [1.0, 2.0, 3.0] + # Medians are properties, so JSON carries the raw samples by design. + assert "eval_median_ms" not in encoded["results"][0] + + def test_inference_table_handles_an_unsupported_row(self): + results = [ + InferenceResult( + scenario="SPM", + protocol="experiment", + backend="rust_diffsol", + build_ms=0.0, + setup_ms=0.0, + aot_cache_status="-", + eval_samples_ms=(), + solve_samples_ms=(), + observe_samples_ms=(), + requested_output_points=0, + supported=False, + reason="SolverError: unsupported", + ) + ] + + table = render_inference_table(results, width=200) + + assert "unsupported" in table + assert "SolverError" in table + + def test_empty_inference_results_render_a_message(self): + assert render_inference_table([]) == "No inference results." + + def test_inference_table_reports_coverage_and_gradient_parity(self): + trajectory = TrajectorySummary( + baseline_points=99, + candidate_points=40, + common_points=40, + coverage=0.404, + max_time_diff=0.0, + final_time_diff=2100.0, + baseline_termination="final time", + candidate_termination="event", + status="warn", + ) + result = InferenceResult( + scenario="SPM", + protocol="cc_discharge", + backend="rust_diffsol", + build_ms=10.0, + setup_ms=5.0, + cold_observe_ms=7.5, + aot_cache_status="-", + eval_samples_ms=(1.0, 2.0, 3.0), + solve_samples_ms=(0.8, 1.8, 2.8), + observe_samples_ms=(0.2, 0.2, 0.2), + requested_output_points=99, + output_comparison=ComparisonSummary(1e-8, 1e-8, 0.01), + sensitivity_comparison=ComparisonSummary(15.6, 28.0, 28.0), + trajectory_comparison=trajectory, + ) + + table = render_inference_table([result], width=250) + + assert "ColdObs" in table and "7.50" in table + assert "Cover" in table and "0.404" in table + assert "Sens Δ" in table and "1.56e+01" in table + # A clean value comparison must not read as an overall pass. + assert result.status == "warn" + + def test_cold_observation_is_reported_separately_from_setup(self): + result = InferenceResult( + scenario="SPM", + protocol="cc_discharge", + backend="casadi_idaklu", + build_ms=10.0, + setup_ms=5.0, + cold_observe_ms=9.15, + aot_cache_status="-", + eval_samples_ms=(1.0,), + solve_samples_ms=(0.8,), + observe_samples_ms=(0.2,), + requested_output_points=99, + ) + + encoded = json.loads(json.dumps(suite_to_jsonable("inference", [result]))) + assert encoded["results"][0]["cold_observe_ms"] == 9.15 + assert encoded["results"][0]["setup_ms"] == 5.0 + + +class TestReportAttribution: + def _solver_rows(self, protocols): + comparison = ComparisonSummary(1e-8, 1e-7, 0.01) + trajectory = TrajectorySummary( + baseline_points=100, + candidate_points=100, + common_points=100, + coverage=1.0, + max_time_diff=0.0, + final_time_diff=0.0, + baseline_termination="final time", + candidate_termination="final time", + status="pass", + ) + return [ + SolverResult( + scenario="SPM", + protocol=protocol, + backend="rust_idaklu", + timings=PhaseTiming(), + requested_output_points=100, + state_comparison=comparison, + output_comparison=comparison, + trajectory_comparison=trajectory, + ) + for protocol in protocols + ] + + def test_compact_validation_table_carries_the_protocol(self): + rows = self._solver_rows(("cc_discharge", "drive_cycle")) + + compact = render_solver_table(rows, width=120) + validation = compact.split("\n\nValidation\n", maxsplit=1)[1] + + # Same scenario and backend on both rows, so without the protocol the + # two validation lines are indistinguishable. + assert "Protocol" in validation.splitlines()[0] + assert "cc_discharge" in validation and "drive_cycle" in validation + assert _fits(compact, 120) + + def test_aot_profile_carries_the_protocol(self): + profile = AotProfile( + fresh_cache_statuses=("miss",), + disk_cache_statuses=("disk",), + codegen_ms=1.0, + compiler_ms=2.0, + fresh_load_ms=1.0, + disk_load_ms=1.0, + fresh_total_ms=4.0, + disk_total_ms=2.0, + disk_prepare_ms=1.0, + disk_cold_startup_ms=1.0, + library_size_bytes=1024, + verified=True, + ) + rows = [ + SolverResult( + scenario="SPM", + protocol=protocol, + backend="casadi_idaklu_aot", + timings=PhaseTiming(), + requested_output_points=100, + aot_profile=profile, + ) + for protocol in ("cc_discharge", "drive_cycle") + ] + + for width in (200, 120, 80): + report = render_solver_table(rows, width=width) + aot = report.split("AOT profile", maxsplit=1)[1] + assert "cc_discharge" in aot and "drive_cycle" in aot + assert _fits(report, width) + + def test_sensitivity_rows_record_the_parameters_actually_differentiated(self): + rows = [ + SensitivityResult( + scenario="SPM", + protocol="cc_discharge", + backend="rust_idaklu", + timings=PhaseTiming(), + requested_output_points=100, + sensitivity_parameters=("I", "eps_p"), + ), + SensitivityResult( + scenario="SPM", + protocol="drive_cycle", + backend="rust_idaklu", + timings=PhaseTiming(), + requested_output_points=100, + # An Interpolant current cannot be an input parameter. + sensitivity_parameters=("eps_p",), + ), + ] + + for width in (250, 120, 80): + table = render_sensitivity_table(rows, width=width) + assert "I,eps_p" in table + assert _fits(table, width) + + encoded = json.loads(json.dumps(suite_to_jsonable("sensitivity", rows))) + assert encoded["results"][1]["sensitivity_parameters"] == ["eps_p"] + + def test_run_metadata_distinguishes_two_dirty_trees(self, monkeypatch): + outputs = {"diff": "--- a\n+++ b\n+one implementation\n"} + + def fake_git(*arguments): + if arguments[0] == "rev-parse": + return "abc123\n" + if arguments[0] == "status": + return " M benchmarks/x.py\n" + return outputs["diff"] + + monkeypatch.setattr(run_module, "_git", fake_git) + first = run_module._git_metadata() + outputs["diff"] = "--- a\n+++ b\n+a different implementation\n" + second = run_module._git_metadata() + + assert first["git_revision"] == second["git_revision"] + assert first["git_dirty"] is True + # The revision alone cannot tell two local implementations apart. + assert first["git_diff_digest"] != second["git_diff_digest"] + + def test_clean_tree_has_no_diff_digest(self, monkeypatch): + def fake_git(*arguments): + return "abc123\n" if arguments[0] == "rev-parse" else "" + + monkeypatch.setattr(run_module, "_git", fake_git) + + metadata = run_module._git_metadata() + + assert metadata["git_dirty"] is False + assert metadata["git_diff_digest"] is None + + +class TestBackendOrdering: + def test_baseline_is_shuffled_in_with_the_candidates(self): + cases = backend_cases(include_aot=True) + + assert _BASELINE_CASE in cases + # Pinning the baseline first measures it on a systematically colder + # machine than everything it is compared against. + orders = { + _shuffled_backend_cases( + 0, f"solver:SPM:{protocol}", include_aot=True + ).index(_BASELINE_CASE) + for protocol in get_protocol_names() + } + assert orders != {0} + + def test_shuffle_key_separates_protocols_of_one_model(self): + orders = { + tuple(_shuffled_backend_cases(0, f"solver:SPM:{p}", include_aot=True)) + for p in get_protocol_names() + } + + # One key per model gave every protocol the same order, so a slow first + # slot always landed on the same backend. + assert len(orders) > 1 + + def test_shuffle_is_reproducible_for_a_fixed_seed(self): + first = _shuffled_backend_cases(3, "solver:DFN:cc_charge", include_aot=False) + again = _shuffled_backend_cases(3, "solver:DFN:cc_charge", include_aot=False) + other = _shuffled_backend_cases(4, "solver:DFN:cc_charge", include_aot=False) + + assert first == again + assert first != other + assert sorted(first) == sorted(backend_cases(include_aot=False)) + + +class TestConvergedReference: + """The oracle every row is judged against, and its fallbacks.""" + + def test_a_reference_no_tighter_than_the_scenario_is_rejected(self): + scenarios = get_solver_scenarios(["SPM"]) + + # 1e-6 scenario, so anything looser than 1e-8 measures two approximations + # against each other -- the artifact the reference exists to remove. + for tolerance in (1e-6, 1e-7): + with pytest.raises(ValueError, match="decades tighter"): + resolve_reference_tolerance(scenarios, tolerance) + assert resolve_reference_tolerance(scenarios, 1e-8) == 1e-8 + assert resolve_reference_tolerance(scenarios, 1e-10) == 1e-10 + + def test_zero_disables_the_reference(self): + scenarios = get_solver_scenarios(["SPM"]) + + assert resolve_reference_tolerance(scenarios, 0.0) is None + assert resolve_reference_tolerance(scenarios, None) is None + + def test_ladder_loosens_a_decade_at_a_time_up_to_the_ceiling(self): + scenario = get_solver_scenarios(["SPM"])[0] + + ladder = _reference_ladder(scenario, 1e-11) + + np.testing.assert_allclose(ladder, [1e-11, 1e-10, 1e-9, 1e-8], rtol=1e-12) + + def test_the_tightest_converging_tolerance_wins(self): + scenario = get_solver_scenarios(["SPM"])[0] + attempted = [] + + def solve(tolerance): + attempted.append(tolerance) + # DFN under a ramping current behaves like this: converged only after + # loosening, so a usable reference exists but not at the tightest rung. + if tolerance < 1e-9: + raise pybamm.SolverError("IDA_BAD_K") + return "converged" + + used, value = _resolve_reference(scenario, 1e-11, solve, what="reference") + + assert (used, value) == (1e-9, "converged") + np.testing.assert_allclose(attempted, [1e-11, 1e-10, 1e-9], rtol=1e-12) + + def test_an_unreachable_reference_degrades_rather_than_raising(self): + scenario = get_solver_scenarios(["SPM"])[0] + + def solve(tolerance): + raise pybamm.SolverError("IDA_BAD_K") + + assert _resolve_reference(scenario, 1e-10, solve, what="reference") == ( + None, + None, + ) + + def test_a_failed_reference_keeps_the_gradient_gate(self, monkeypatch): + # Status has to flip on backend correctness, not on whether the + # reference happened to converge; the state gate is decades tighter. + scenario = get_solver_scenarios(["SPM"], ["cc_discharge"])[0] + _, gradient_rtol = _sensitivity_tolerances(scenario) + captured: list[tuple[float | None, float]] = [] + + def result(backend): + return SensitivityResult( + scenario=scenario.name, + backend=backend, + timings=PhaseTiming(), + requested_output_points=scenario.plan.requested_points, + ) + + monkeypatch.setattr( + runners, + "_execute_backend_cases", + lambda *a, **k: { + runners._BASELINE_CASE: (result("casadi_idaklu"), None, None, None), + ("rust_diffsol", False): (result("rust_diffsol"), None, None, None), + }, + ) + monkeypatch.setattr( + runners, "_sensitivity_reference", lambda *a, **k: (None, None) + ) + + def spy(*args, atol, rtol, **kwargs): + captured.append((atol, rtol)) + return (None, None, None) + + monkeypatch.setattr(runners, "_compare_sensitivity_pair", spy) + + runners.run_sensitivity_lane( + [scenario], repeats=1, warmup=0, reference_tolerance=None + ) + + assert captured == [(None, pytest.approx(gradient_rtol))] + + def test_reference_gate_is_looser_than_one_tolerance_unit(self): + scenario = get_solver_scenarios(["SPM"])[0] + + atol, rtol = _reference_tolerances(scenario) + + # A tolerance bounds one step's local error; the global error accumulates + # it, so a correct solve lands several tolerance units from the answer. + assert atol > scenario.atol + assert rtol > scenario.rtol + + def test_a_doomed_rung_retries_the_failing_draw_before_the_rest(self, monkeypatch): + # Re-solving converged draws for a rung that cannot converge costs + # rungs x repeats converged DFN solves instead of rungs + repeats. + scenario = get_inference_scenarios(["SPM"], ["cc_charge"])[0] + vectors = [{"eps_p": 0.60}, {"eps_p": 0.62}, {"eps_p": 0.64}, {"eps_p": 0.66}] + attempted: list[dict] = [] + + def solve(**kwargs): + inputs = kwargs["inputs"] + attempted.append(inputs) + if inputs == vectors[3]: + raise pybamm.SolverError("IDA_BAD_K") + + monkeypatch.setattr(runners, "_make_solver", lambda *a, **k: None) + monkeypatch.setattr( + runners, "_build_simulation", lambda *a, **k: SimpleNamespace(solve=solve) + ) + monkeypatch.setattr(runners, "_build_and_time", lambda *a, **k: 0.0) + monkeypatch.setattr( + runners, "_observe_inference", lambda solution, inputs, **kwargs: inputs + ) + + used, _ = runners._inference_reference( + scenario, + vectors, + warmup=1, + grid=np.zeros(1), + sensitivities=False, + tolerance=1e-11, + ) + + assert used is None + assert attempted == [vectors[1], vectors[2], vectors[3], *[vectors[3]] * 3] + + def test_inference_reference_fixes_y0_at_the_first_draw(self, monkeypatch): + # The lane holds y0 at vectors[0]; a reference resolved from a later draw + # starts the cell at a different SOC and moves the event by ~80 s. + scenario = get_inference_scenarios(["SPM"], ["cc_charge"])[0] + vectors = [{"eps_p": 0.60}, {"eps_p": 0.62}, {"eps_p": 0.64}, {"eps_p": 0.66}] + built_with = [] + observed = [] + + monkeypatch.setattr(runners, "_make_solver", lambda *a, **k: None) + monkeypatch.setattr( + runners, "_build_simulation", lambda *a, **k: SimpleNamespace(solve=dict) + ) + monkeypatch.setattr( + runners, + "_build_and_time", + lambda simulation, scenario, inputs=None: built_with.append(inputs) or 0.0, + ) + monkeypatch.setattr( + runners, + "_observe_inference", + lambda solution, inputs, **kwargs: observed.append(inputs), + ) + + used, _ = runners._inference_reference( + scenario, + vectors, + warmup=2, + grid=np.zeros(1), + sensitivities=False, + tolerance=1e-10, + ) + + assert used == 1e-10 + assert built_with == [vectors[0]] + assert observed == vectors[2:] + + +class TestReferenceReporting: + def _rows(self, reference_tolerance): + comparison = ComparisonSummary( + max_abs_diff=1e-6, + max_rel_diff=1e-6, + max_normalized_error=0.1, + ) + return [ + SolverResult( + scenario="SPM", + backend=backend, + timings=PhaseTiming(), + requested_output_points=100, + state_comparison=comparison, + output_comparison=comparison, + reference_tolerance=reference_tolerance, + baseline_delta=None if backend == "casadi_idaklu" else comparison, + ) + for backend in ("casadi_idaklu", "rust_idaklu") + ] + + def test_the_baseline_row_is_gated_once_a_reference_exists(self): + gated, ungated = ( + self._rows(1e-10)[0], + SolverResult( + scenario="SPM", + backend="casadi_idaklu", + timings=PhaseTiming(), + requested_output_points=100, + ), + ) + + assert gated.status == "pass" + assert ungated.status == "baseline" + + def test_the_table_names_what_the_deltas_are_measured_against(self): + with_reference = render_solver_table(self._rows(1e-10), width=200) + without = render_solver_table(self._rows(None), width=200) + + assert "converged casadi_idaklu reference at atol=rtol=1e-10" in with_reference + assert "no converged reference was run" in without + + def test_every_comparison_lane_carries_the_reference_columns(self): + comparison = ComparisonSummary( + max_abs_diff=1e-6, + max_rel_diff=1e-6, + max_normalized_error=0.1, + ) + sensitivity = render_sensitivity_table( + [ + SensitivityResult( + scenario="SPM", + backend="rust_diffsol", + timings=PhaseTiming(), + requested_output_points=100, + output_sens_comparison=comparison, + reference_tolerance=1e-10, + baseline_delta=comparison, + ) + ], + width=200, + ) + inference = render_inference_table( + [ + InferenceResult( + scenario="SPM", + protocol="cc_discharge", + backend="rust_diffsol", + build_ms=1.0, + setup_ms=1.0, + aot_cache_status="-", + eval_samples_ms=(1.0,), + solve_samples_ms=(1.0,), + observe_samples_ms=(1.0,), + requested_output_points=100, + output_comparison=comparison, + reference_tolerance=1e-10, + baseline_delta=comparison, + ) + ], + width=200, + ) + + for table in (sensitivity, inference): + assert "converged casadi_idaklu reference at atol=rtol=1e-10" in table + assert "Base Δ" in table + + def test_the_cross_backend_delta_is_reported_beside_the_reference_error(self): + table = render_solver_table(self._rows(1e-10), width=200) + payload = suite_to_jsonable("solver", self._rows(1e-10)) + + assert "Base Δ" in table + assert payload["results"][0]["baseline_delta"] is None + assert payload["results"][1]["baseline_delta"]["max_abs_diff"] == 1e-6 + assert payload["results"][1]["reference_tolerance"] == 1e-10 + + +class TestCli: + def test_reference_tolerance_flag(self): + parser = build_parser() + + assert parser.parse_args([]).reference_tolerance == DEFAULT_REFERENCE_TOLERANCE + assert ( + parser.parse_args(["--reference-tolerance", "0"]).reference_tolerance == 0 + ) + + def test_lane_and_protocol_flags(self): + parser = build_parser() + + args = parser.parse_args( + ["--lane", "inference", "--protocols", "drive_cycle", "pulse_train"] + ) + + assert args.lane == "inference" + assert args.protocols == ["drive_cycle", "pulse_train"] + assert args.aot == "solver" + assert args.inference_sensitivities is False + assert args.inference_seed == 0 + + def test_default_protocols_preserve_the_baseline_run(self): + args = build_parser().parse_args([]) + + assert args.protocols == ["cc_discharge"] + assert args.lane == "all" + + def test_an_empty_protocol_list_means_all_of_them_as_it_does_for_models(self): + # `--protocols` and `--models` are parallel flags, so a bare one has to + # mean the same thing on both. + every_protocol = get_protocol_names() + selected = {scenario.protocol for scenario in get_solver_scenarios(["SPM"], [])} + assert selected == set(every_protocol) + assert [ + scenario.protocol for scenario in get_solver_scenarios(["SPM"], None) + ] == ["cc_discharge"] + + def test_aot_choices(self): + parser = build_parser() + + assert parser.parse_args(["--aot", "none"]).aot == "none" + assert parser.parse_args(["--aot", "all"]).aot == "all" + with pytest.raises(SystemExit): + parser.parse_args(["--aot", "sometimes"]) + + def test_include_aot_resolution(self): + assert include_aot_for("solver", "solver") is True + assert include_aot_for("sensitivity", "solver") is False + assert include_aot_for("sensitivity", "all") is True + assert include_aot_for("solver", "all") is True + assert include_aot_for("solver", "none") is False + assert include_aot_for("sensitivity", "none") is False + + def test_removed_diffsol_flag_is_gone(self): + with pytest.raises(SystemExit): + build_parser().parse_args(["--include-diffsol"]) diff --git a/packages/pybamm/tests/unit/test_rust_stubs.py b/packages/pybamm/tests/unit/test_rust_stubs.py new file mode 100644 index 0000000000..9d6314c0d9 --- /dev/null +++ b/packages/pybamm/tests/unit/test_rust_stubs.py @@ -0,0 +1,45 @@ +"""``pybamm/rust/_core.pyi`` stays in sync with the built extension. + +The stub is hand-maintained; ``mypy.stubtest`` compares it against the runtime +module, so an added, removed or renamed method, argument or default in the +bindings fails here until the stub is updated. Types are not runtime-checkable +on an extension module and stay review-enforced. +""" + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pybamm.rust + + +class TestRustStubs: + def test_stub_matches_the_built_extension(self, tmp_path): + stub = Path(pybamm.rust.__file__).parent / "_core.pyi" + assert stub.is_file(), f"missing stub file {stub}" + + # A hermetic stub-only package tree: empty ancestor stubs keep mypy from + # analysing the full pybamm package just to resolve `pybamm.rust._core`. + stub_root = tmp_path / "stubs" + package = stub_root / "pybamm" / "rust" + package.mkdir(parents=True) + (stub_root / "pybamm" / "__init__.pyi").touch() + (package / "__init__.pyi").touch() + shutil.copyfile(stub, package / "_core.pyi") + + result = subprocess.run( + [sys.executable, "-m", "mypy.stubtest", "pybamm.rust._core"], + env={**os.environ, "MYPYPATH": str(stub_root)}, + # tmp_path, so mypy's cache directory never lands in the repo. + cwd=tmp_path, + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert result.returncode == 0, ( + f"stubtest found mismatches between pybamm/rust/_core.pyi and the " + f"built extension:\n{result.stdout}\n{result.stderr}" + ) diff --git a/packages/pybamm/tests/unit/test_simulation.py b/packages/pybamm/tests/unit/test_simulation.py index 3ab2feba4a..247973aaa3 100644 --- a/packages/pybamm/tests/unit/test_simulation.py +++ b/packages/pybamm/tests/unit/test_simulation.py @@ -903,6 +903,7 @@ def test_drive_cycle_t_eval_warnings_for_missing_points_and_resolution(self): param["Current function [A]"] = pybamm.Interpolant( drive_cycle[:, 0], drive_cycle[:, 1], pybamm.t ) + model.convert_to_format = "casadi" sim = pybamm.Simulation( model, parameter_values=param, solver=pybamm.CasadiSolver() ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_base_solver.py b/packages/pybamm/tests/unit/test_solvers/test_base_solver.py index 72b97563a4..fa9954fe27 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_base_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_base_solver.py @@ -2,6 +2,7 @@ # Tests for the Base Solver class # +import multiprocessing as mp import re import casadi @@ -24,6 +25,14 @@ def test_base_solver_init(self): solver.rtol = 1e-7 assert solver.rtol == 1e-7 + def test_root_method_takes_the_tightest_entry_of_a_per_state_atol(self): + # The root solver takes a single tolerance, so a per-state atol has to + # reduce rather than reach the comparison against root_tol as an array. + solver = pybamm.BaseSolver( + atol=np.array([1e-6, 1e-9]), root_method="nonlinear_solver", root_tol=1e-3 + ) + assert solver.root_method.atol == 1e-9 + def test_root_method_init(self): solver = pybamm.BaseSolver(root_method="nonlinear_solver") assert isinstance(solver.root_method, pybamm.NonlinearSolver) @@ -161,6 +170,7 @@ def __init__(self): "alg", [t, y, p], [self.algebraic_eval(t, y, p)] ) self.convert_to_format = "casadi" + self.uses_stacked_inputs = True self.bounds = (np.array([-np.inf]), np.array([np.inf])) self.len_rhs_and_alg = 1 self.events = [] @@ -200,6 +210,7 @@ def __init__(self): "alg", [t, y, p], [self.algebraic_eval(t, y, p)] ) self.convert_to_format = "casadi" + self.uses_stacked_inputs = True self.bounds = (-np.inf * np.ones(4), np.inf * np.ones(4)) self.len_rhs = 1 self.len_rhs_and_alg = 4 @@ -526,6 +537,26 @@ def test_integrate_single_error(self): ): solver._integrate_single(model, np.array([0, 1]), {}, np.array([1])) + def test_one_process_solves_in_this_process(self, monkeypatch): + # calculate_consistent_state asks for one process, and paid ~6 s per 32 + # DFN sets in spawn and pickle costs for no concurrency at all. + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: -pybamm.InputParameter("k") * v} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + pybamm.Discretisation().process_model(model) + + def no_pools(*args, **kwargs): + raise AssertionError("a solve with nproc=1 built a process pool") + + monkeypatch.setattr(mp, "get_context", no_pools) + + solutions = pybamm.IDAKLUSolver().solve( + model, np.linspace(0, 1, 10), inputs=[{"k": 1.0}, {"k": 2.0}], nproc=1 + ) + assert len(solutions) == 2 + def test_discontinuity_events_different_times_error(self): # Test that an error is raised when discontinuity events occur at different # times for different input parameter sets @@ -550,3 +581,46 @@ def test_discontinuity_events_different_times_error(self): match="Discontinuity events occur at different times between input parameter sets", ): solver.solve(model, t_eval, inputs=inputs_list) + + @pytest.mark.parametrize( + "all_sensitivities", + [ + pytest.param({}, id="no_sensitivities_stored"), + # IDAKLU stores output-width sensitivities here, which cannot seed + # a state-width dy0/dp. + pytest.param( + {"p": np.zeros((2, 1)), "all": np.zeros((2, 1))}, id="output_width" + ), + ], + ) + def test_check_restart_sensitivities_rejects_outputs_only(self, all_sensitivities): + solution = pybamm.Solution( + [np.array([0.0, 1.0])], + [np.zeros((0, 2))], + pybamm.BaseModel(), + [{}], + variables_returned=True, + ) + solution._all_sensitivities = all_sensitivities + + with pytest.raises( + pybamm.SolverError, + match=r"Cannot continue a sensitivity solve from a solution that " + r"returned output variables only", + ): + pybamm.BaseSolver._check_restart_sensitivities(solution) + + def test_check_restart_sensitivities_allows_full_state(self): + solution = pybamm.Solution( + [np.array([0.0, 1.0])], + [np.zeros((3, 2))], + pybamm.BaseModel(), + [{}], + variables_returned=False, + ) + solution._all_sensitivities = {"p": np.zeros((6, 1))} + + pybamm.BaseSolver._check_restart_sensitivities(solution) + + def test_check_restart_sensitivities_allows_empty_solution(self): + pybamm.BaseSolver._check_restart_sensitivities(pybamm.EmptySolution()) diff --git a/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py b/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py index 4aa36d92e3..b1b8c8917c 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_casadi_algebraic_solver.py @@ -6,6 +6,12 @@ from tests import get_discretisation_for_testing +@pytest.fixture(autouse=True) +def _casadi_default(monkeypatch): + # This file tests the CasADi solvers; pin the model default they require. + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi") + + class TestCasadiAlgebraicSolver: def test_algebraic_solver_init(self): solver = pybamm.CasadiAlgebraicSolver(step_tol=1e-6, tol=1e-4) diff --git a/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py b/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py index a941074446..8ca5438211 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_casadi_solver.py @@ -7,6 +7,12 @@ from tests import get_discretisation_for_testing, get_mesh_for_testing +@pytest.fixture(autouse=True) +def _casadi_default(monkeypatch): + # This file tests the CasADi solvers; pin the model default they require. + monkeypatch.setattr(pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi") + + class TestCasadiSolver: def test_no_sensitivities_error(self): model = pybamm.lithium_ion.SPM() diff --git a/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py b/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py index 76538f760f..e9969e7cc7 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_composite_solver.py @@ -59,6 +59,7 @@ def solve(self, *args, **kwargs): failing_solver = FailingSolver() working_solver = pybamm.CasadiAlgebraicSolver() + model.convert_to_format = "casadi" composite = pybamm.CompositeSolver([failing_solver, working_solver]) solution = composite.solve(model) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py new file mode 100644 index 0000000000..4951d106f2 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py @@ -0,0 +1,61 @@ +"""Diffsol with several input sets: one solution per set. + +``BaseSolver`` fans multiple input sets out across a process pool, which diffsol +cannot use — ``PreparedSolver`` is a PyO3 object and will not pickle. The solver +runs its own fan-out instead (see ``TestDiffsolNumThreads`` in +``test_diffsol_solver.py``), defaulting to a serial loop, and either way must +honour the contract of returning one result per input set. +""" + +import numpy as np + +import pybamm + +_SOLVER_TOL = 1e-8 +_CURRENTS = [0.5, 1.5] + + +class TestDiffsolBatchedInputs: + def _build(self): + model = pybamm.lithium_ion.SPM() + params = model.default_parameter_values + params["Current function [A]"] = "[input]" + sim = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ) + sim.build() + return sim + + def test_returns_one_solution_per_input_set(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + assert isinstance(solutions, list) + assert len(solutions) == len(inputs) + + def test_each_solution_carries_its_own_input_set(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + for solution, current in zip(solutions, _CURRENTS, strict=True): + assert solution.all_inputs[0]["Current function [A]"] == np.array([current]) + + def test_distinct_input_sets_give_distinct_trajectories(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + first = solutions[0]["Voltage [V]"](t_eval) + second = solutions[1]["Voltage [V]"](t_eval) + # A 3x larger current must discharge measurably faster. + assert np.all(second < first) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py new file mode 100644 index 0000000000..fa937ac180 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py @@ -0,0 +1,212 @@ +"""Diffsol ``output_variables``: layout parity against the full-state solve. + +Rust returns one row per flattened output component, so a vector output variable +occupies several rows. These tests pin that the Python side slices by component +count rather than by variable ordinal, and that the results expose the same +interpolating interface as every other solver. +""" + +import numpy as np +import pytest + +import pybamm + +# A 20-node spatial variable sandwiched between two scalars: an ordinal-indexed +# reader returns one component for it and shifts everything after it. +_VECTOR_VAR = "Negative particle surface concentration [mol.m-3]" +_OUTPUT_VARIABLES = ["Voltage [V]", _VECTOR_VAR, "Current [A]"] + +_SOLVER_TOL = 1e-8 + + +class TestDiffsolOutputVariables: + @pytest.fixture(scope="class") + def solutions(self): + """Solve SPM twice: once output-only, once for the full state.""" + t_eval = np.linspace(0, 600, 11) + + model_outputs = pybamm.lithium_ion.SPM() + sol_outputs = pybamm.Simulation( + model_outputs, + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, + atol=_SOLVER_TOL, + output_variables=_OUTPUT_VARIABLES, + ), + ).solve(t_eval) + + model_full = pybamm.lithium_ion.SPM() + sol_full = pybamm.Simulation( + model_full, + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ).solve(t_eval) + + return sol_outputs, sol_full, t_eval + + def test_vector_output_keeps_every_component(self, solutions): + sol_outputs, sol_full, _ = solutions + + expected = np.asarray(sol_full[_VECTOR_VAR].entries) + actual = np.asarray(sol_outputs[_VECTOR_VAR].entries) + + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-8) + + def test_scalar_after_a_vector_is_not_shifted(self, solutions): + sol_outputs, sol_full, _ = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs["Current [A]"].entries), + np.asarray(sol_full["Current [A]"].entries), + rtol=1e-6, + atol=1e-8, + ) + + def test_scalar_before_a_vector_is_unaffected(self, solutions): + sol_outputs, sol_full, _ = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs["Voltage [V]"].entries), + np.asarray(sol_full["Voltage [V]"].entries), + rtol=1e-6, + atol=1e-8, + ) + + def test_vector_output_primal_and_sensitivity_layouts_agree(self): + """The sensitivity path always sliced by component count; the primal path + sliced by variable ordinal. For a vector variable the two disagreed.""" + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + + model = pybamm.lithium_ion.SPM() + model.events = [] + sol = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, atol=_SOLVER_TOL, output_variables=[_VECTOR_VAR] + ), + ).solve( + np.linspace(0, 100, 6), inputs={"I": 0.5}, calculate_sensitivities=["I"] + ) + + variable = sol[_VECTOR_VAR] + assert variable.entries.size == variable.sensitivities["I"].shape[0] + + def test_output_rows_stay_aligned_across_the_batch_window(self): + """200 output points cross the 128-lane staging window in Rust; a flush + off-by-one would shift every row after the window boundary.""" + t_eval = np.linspace(0, 600, 200) + + sol_outputs = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, + atol=_SOLVER_TOL, + output_variables=_OUTPUT_VARIABLES, + ), + ).solve(t_eval) + sol_full = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ).solve(t_eval) + + for name in _OUTPUT_VARIABLES: + np.testing.assert_allclose( + np.asarray(sol_outputs[name](t_eval)), + np.asarray(sol_full[name](t_eval)), + rtol=1e-6, + atol=1e-8, + err_msg=name, + ) + + @pytest.mark.parametrize("name", _OUTPUT_VARIABLES) + def test_output_variables_are_callable(self, solutions, name): + """`solution[name](t)` is the documented read interface for every solver. + + Compared against the full-state solve rather than against ``.entries``: + for a spatial variable the call interface interpolates onto the node grid + including its two boundary points, so the two shapes differ by design. + """ + sol_outputs, sol_full, t_eval = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs[name](t_eval)), + np.asarray(sol_full[name](t_eval)), + rtol=1e-6, + atol=1e-8, + ) + + +class TestDiffsolOutputVariablesAcrossExperimentSteps: + """Stitching outputs-only segments together, as an ``Experiment`` does.""" + + @staticmethod + def _solve(solver_cls, convert_to_format, **solver_kwargs): + experiment = pybamm.Experiment( + ["Discharge at 1C for 200 seconds", "Rest for 100 seconds"] + ) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format + return pybamm.Simulation( + model, + parameter_values=pybamm.ParameterValues("Chen2020"), + experiment=experiment, + solver=solver_cls(rtol=_SOLVER_TOL, atol=_SOLVER_TOL, **solver_kwargs), + ).solve() + + def test_experiment_with_output_variables_stitches(self): + # Segment stitching slices all_ys for every sub-solution; a None there + # raised TypeError before diffsol supplied a zero-row array. + sol = self._solve( + pybamm.DiffsolSolver, "rust", output_variables=["Voltage [V]"] + ) + assert len(sol.all_ts) == 2 + assert sol.variables_returned + assert all(y.shape[0] == 0 for y in sol.all_ys) + assert np.all(np.isfinite(np.asarray(sol["Voltage [V]"](sol.t)))) + + def test_experiment_outputs_match_the_full_state_solve(self): + sol_out = self._solve( + pybamm.DiffsolSolver, "rust", output_variables=["Voltage [V]"] + ) + sol_full = self._solve(pybamm.DiffsolSolver, "rust") + t = sol_full.t + np.testing.assert_allclose( + np.asarray(sol_out["Voltage [V]"](t)), + np.asarray(sol_full["Voltage [V]"](t)), + rtol=1e-6, + atol=1e-8, + ) + + def test_restarting_with_sensitivities_is_refused_not_seeded_with_zeros(self): + # An outputs-only step boundary carries no state sensitivities, so the + # next step would silently restart dy0/dp from zero. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Negative particle diffusivity [m2.s-1]"] = ( + pybamm.InputParameter("D_n") + ) + solver = pybamm.DiffsolSolver(output_variables=["Voltage [V]"]) + simulation = pybamm.Simulation( + model, parameter_values=parameter_values, solver=solver + ) + simulation.build() + inputs = {"D_n": 3.3e-14} + + first = solver.step( + pybamm.EmptySolution(), + simulation.built_model, + dt=100.0, + inputs=inputs, + calculate_sensitivities=["D_n"], + ) + with pytest.raises(pybamm.SolverError, match=r"output variables only"): + solver.step( + first, + simulation.built_model, + dt=100.0, + inputs=inputs, + calculate_sensitivities=["D_n"], + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py new file mode 100644 index 0000000000..eaa1f6e5e1 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py @@ -0,0 +1,936 @@ +"""Parity tests: diffsol native state sensitivities vs CasADi-IDAKLU oracle.""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import NativeObservation + +# Parity assertion tolerances (mandated by task brief; do not weaken). +_PARITY_RTOL = 1e-5 +_PARITY_ATOL = 1e-8 + +# 0D outputs: 1e-9 is the loosest tolerance at which BDF/IDA cross-integrator +# trajectory differences stay within the parity bounds above. +_SCALAR_SOLVER_TOL = 1e-9 + +# D_n's |p * dV/dp| ~ 1e-4 makes 1e-9's ~6e-6 oracle error a 1.7x parity margin; +# 1e-11 moves the oracle close enough that the assertion measures native error. +_TWO_PARAM_SOLVER_TOL = 1e-11 + +# Spatial outputs: concentration Jacobian (~246) amplifies cross-integrator +# differences; 1e-12 is required to stay within the parity bounds. +_SPATIAL_SOLVER_TOL = 1e-12 + + +class TestDiffsolSensitivities: + def _solve_both( + self, output_name, calc, extra_inputs=None, solver_tol=_SCALAR_SOLVER_TOL + ): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + if extra_inputs: + for pybamm_name, (input_name, value) in extra_inputs.items(): + params[pybamm_name] = pybamm.InputParameter(input_name) + inputs[input_name] = value + t_eval = np.linspace(0, 100, 15) + + m_native = pybamm.lithium_ion.SPM() + m_native.events = [] + m_native.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m_native, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=solver_tol, atol=solver_tol), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=calc) + + m_casadi = pybamm.lithium_ion.SPM() + m_casadi.events = [] + m_casadi.convert_to_format = "casadi" + # t_interp=t_eval forces IDAKLUSolver to output exactly at t_eval so + # both solutions share the same time grid for comparison. + sol_c = pybamm.Simulation( + m_casadi, + parameter_values=params, + solver=pybamm.IDAKLUSolver(rtol=solver_tol, atol=solver_tol), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=calc, t_interp=t_eval) + return sol_n[output_name], sol_c[output_name] + + def test_sensitivities_match_casadi_0d(self): + var_n, var_c = self._solve_both("Terminal voltage [V]", ["I"]) + sens_n, sens_c = var_n.sensitivities, var_c.sensitivities + assert set(sens_n) == set(sens_c) + # Non-vacuous: a real multi-point, nonzero sensitivity (not collapsed/all-zero). + assert sens_c["I"].size > 1 and np.any(sens_c["I"] != 0) + for key in sens_c: + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_match_casadi_spatial(self): + name = "Negative particle concentration [mol.m-3]" + var_n, var_c = self._solve_both(name, ["I"], solver_tol=_SPATIAL_SOLVER_TOL) + assert var_c.entries.shape[0] > 1 + sens_n, sens_c = var_n.sensitivities, var_c.sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_all_block_two_params(self): + var_n, var_c = self._solve_both( + "Terminal voltage [V]", + ["D_n", "I"], + extra_inputs={"Negative particle diffusivity [m2.s-1]": ("D_n", 3.3e-14)}, + solver_tol=_TWO_PARAM_SOLVER_TOL, + ) + a_n, a_c = var_n.sensitivities["all"], var_c.sensitivities["all"] + assert a_n.shape == a_c.shape and a_n.shape[1] == 2 + assert np.any(var_n.sensitivities["D_n"] != 0) + # "all" column 0 aligns with the named block (assembly-order check)... + np.testing.assert_allclose( + a_n[:, 0], var_n.sensitivities["D_n"], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + # ...and each named block matches the oracle, not just itself. + for k in ("D_n", "I"): + np.testing.assert_allclose( + var_n.sensitivities[k], + var_c.sensitivities[k], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + np.testing.assert_allclose(a_n, a_c, rtol=_PARITY_RTOL, atol=_PARITY_ATOL) + + def test_state_sensitivity_matches_finite_difference(self): + # Finite-difference diffsol's OWN values: no cross-integrator trajectory + # noise, so the sensitivity formula is validated without ultra-tight tols. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + i0, h = 0.5, 1e-4 + + def voltage(current): + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8), + ).solve(t_eval, inputs={"I": current}) + return sol[name].entries.ravel() + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8), + ).solve(t_eval, inputs={"I": i0}, calculate_sensitivities=["I"]) + analytic = sol[name].sensitivities["I"].ravel() + fd = (voltage(i0 + h) - voltage(i0 - h)) / (2 * h) + assert np.any(analytic != 0) + np.testing.assert_allclose(analytic, fd, rtol=2e-3, atol=1e-6) + + def test_output_variable_sensitivities_match_casadi(self): + # output_variables + calculate_sensitivities routes diffsol through the + # native outputs-and-sensitivities request, oracled by CasADi-IDAKLU. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL, output_variables=[name] + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["I"] + ) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_subset_of_inputs_match_casadi(self): + # Two input parameters (I, D_n) but sensitivities requested for one (I): the + # native output path must return just that block, not one per input. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL, output_variables=[name] + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["I"] + ) + + # Oracle is the IDAKLU STATE path: its output path mislabels columns when + # calculate_sensitivities is a strict subset, and dV/dI is the same either way. + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + assert set(sol_n[name].sensitivities) == {"I", "all"} + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_state_sensitivities_subset_of_inputs_match_casadi(self): + # With sensitivities requested for I alone, the state path must label the + # single block "I" and match CasADi regardless of the surplus input. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"]) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + assert set(sol_n[name].sensitivities) == {"I", "all"} + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + @staticmethod + def _build_time_integral_model(fmt): + # A genuine ExplicitTimeIntegral whose post_sum_node is the squaring. ``c`` + # is a state, so the integral discretises to StateVector(0:1). + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + integral = pybamm.ExplicitTimeIntegral(c, 0) ** 2 + model = pybamm.BaseModel(name="time_integral_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["integral"] = integral + model.variables["c"] = c + model.convert_to_format = fmt + pybamm.Discretisation().process_model(model) + return model + + def test_time_integral_sensitivities_match_casadi(self): + # Confirm the variable is a time-integral with a non-None post_sum_node. + model_check = self._build_time_integral_model("casadi") + var = model_check.get_processed_variable_or_event("integral") + time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var(var, 2) + assert time_integral is not None + assert time_integral.post_sum_node is not None + + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + + # The squaring post-sum amplifies BDF-vs-IDA trajectory differences, so pin + # both solvers at the tight spatial tolerance. + m_native = self._build_time_integral_model("rust") + sol_n = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + m_native, + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + assert isinstance(sol_n.observation, NativeObservation) + + m_casadi = self._build_time_integral_model("casadi") + sol_c = pybamm.IDAKLUSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + m_casadi, + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + + sens_n, sens_c = ( + sol_n["integral"].sensitivities, + sol_c["integral"].sensitivities, + ) + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_output_variables_time_integral_matches_the_full_state_solve(self): + # The rows of a time-integral output carry its integrand; the postfix sum + # runs in the shared outputs assembly, so diffsol reads the same value and + # the same sensitivities as its own full-state solve. + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + solve_kwargs = { + "t_eval": [times[0], times[-1]], + "t_interp": times, + "inputs": inputs, + "calculate_sensitivities": ["a", "b"], + } + + sol_out = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, + atol=_SPATIAL_SOLVER_TOL, + output_variables=["integral"], + ).solve(self._build_time_integral_model("rust"), **solve_kwargs) + sol_full = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve(self._build_time_integral_model("rust"), **solve_kwargs) + + np.testing.assert_allclose( + np.asarray(sol_out["integral"].entries), + np.asarray(sol_full["integral"].entries), + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + sens_out, sens_full = ( + sol_out["integral"].sensitivities, + sol_full["integral"].sensitivities, + ) + for name in inputs: + np.testing.assert_allclose( + sens_out[name], + sens_full[name], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + err_msg=f"sensitivity mismatch for param '{name}'", + ) + + @staticmethod + def _build_two_integrals_model(fmt): + # Two time-integral variables with the SAME n_inner (scalar) but DIFFERENT + # post-sum operations (square vs cube), to catch post-sum cache aliasing. + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + model = pybamm.BaseModel(name="two_integrals_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["sq"] = pybamm.ExplicitTimeIntegral(c, 0) ** 2 + model.variables["cube"] = pybamm.ExplicitTimeIntegral(c, 0) ** 3 + model.variables["c"] = c + model.convert_to_format = fmt + pybamm.Discretisation().process_model(model) + return model + + def test_distinct_time_integrals_same_n_inner_dont_alias(self): + # Both post-sum fns share n_inner=1; the cache must key them apart by + # variable name, else "sq" and "cube" swap post-sum jacobians. + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + + sol_n = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + self._build_two_integrals_model("rust"), + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + sol_c = pybamm.IDAKLUSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + self._build_two_integrals_model("casadi"), + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + + for var in ("sq", "cube"): + sens_n, sens_c = sol_n[var].sensitivities, sol_c[var].sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_with_event_termination_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 5.0} # high current trips the lower voltage cutoff early + t_eval = np.linspace(0, 3600, 50) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() # events retained (no m.events = []) + m.convert_to_format = "rust" + # t_interp pins output to the t_eval grid so both solvers compare aligned + # time points up to the event. + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_n.termination.startswith("event") + + m_c = pybamm.lithium_ion.SPM() + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + # t_interp forces IDAKLU to output exactly at t_eval so both solutions + # share the same time grid for comparison. + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_c.termination.startswith("event") + + # t_interp pins both to one grid, so allow at most a single grid-point + # difference in the event window and require it to be non-trivial. + len_n = sol_n[name].sensitivities["I"].shape[0] + len_c = sol_c[name].sensitivities["I"].shape[0] + assert abs(len_n - len_c) <= 1, ( + f"solvers disagree on event window length: diffsol={len_n}, idaklu={len_c}" + ) + n = min(len_n, len_c) + assert n > 1 + np.testing.assert_allclose( + sol_n[name].sensitivities["I"][:n], + sol_c[name].sensitivities["I"][:n], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_with_event_termination_match_casadi(self): + # Exercises an outputs-and-sensitivities request under event truncation. The single + # input is the sole requested sensitivity, so IDAKLU's output path oracles it. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 5.0} + t_eval = np.linspace(0, 3600, 50) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() # events retained + m.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_n.termination.startswith("event") + + m_c = pybamm.lithium_ion.SPM() + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + len_n = sol_n[name].sensitivities["I"].shape[0] + len_c = sol_c[name].sensitivities["I"].shape[0] + assert abs(len_n - len_c) <= 1, ( + f"solvers disagree on event window length: diffsol={len_n}, idaklu={len_c}" + ) + n = min(len_n, len_c) + assert n > 1 + np.testing.assert_allclose( + sol_n[name].sensitivities["I"][:n], + sol_c[name].sensitivities["I"][:n], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_two_params_match_casadi(self): + # Locking test: insertion order (I, D_n) differs from sorted + # calculate_sensitivities order (D_n, I), exposing column mislabelling. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + # Insertion order: I first, D_n second — intentionally != sorted order. + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_TWO_PARAM_SOLVER_TOL, + atol=_TWO_PARAM_SOLVER_TOL, + output_variables=[name], + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["D_n", "I"] + ) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_TWO_PARAM_SOLVER_TOL, + atol=_TWO_PARAM_SOLVER_TOL, + output_variables=[name], + ), + ).solve( + t_eval, inputs=inputs, calculate_sensitivities=["D_n", "I"], t_interp=t_eval + ) + + for param_name in ("I", "D_n"): + np.testing.assert_allclose( + sol_n[name].sensitivities[param_name], + sol_c[name].sensitivities[param_name], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + err_msg=f"sensitivity mislabelled for param '{param_name}'", + ) + + def test_discrete_time_sum_guard(self): + # A DiscreteTimeSum solved on times that do not match the discrete sum + # times must raise when its sensitivities are accessed natively. + data_times = np.linspace(0, 1, 10) + ref = pybamm.BaseModel(name="ref") + c_ref = pybamm.Variable("c") + ref.rhs = {c_ref: -2 * c_ref} + ref.initial_conditions = {c_ref: 1} + ref.variables["c"] = c_ref + pybamm.Discretisation().process_model(ref) + data_values = ( + pybamm.IDAKLUSolver() + .solve(ref, t_eval=[data_times[0], data_times[-1]], t_interp=data_times)[ + "c" + ] + .entries + ) + data = pybamm.DiscreteTimeData(data_times, data_values, "test_data") + + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + model = pybamm.BaseModel(name="dts_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["data_comparison"] = pybamm.DiscreteTimeSum((c - data) ** 2) + model.variables["c"] = c + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + # Solve on a mismatched grid (t_interp != data_times) so the guard fires. + sol = pybamm.DiffsolSolver().solve( + model, + t_eval=[0, 1], + t_interp=np.linspace(0, 1, 5), + inputs={"a": 0.7, "b": 1.0}, + calculate_sensitivities=["a", "b"], + ) + assert isinstance(sol.observation, NativeObservation) + with pytest.raises(pybamm.SolverError, match="discrete times"): + _ = sol["data_comparison"].sensitivities + + def test_calculate_sensitivities_rejects_vector_width_input_parameter(self): + # sens_param_indices seeds one scalar direction per named parameter; a + # width>1 parameter must raise instead of silently under-seeding (idaklu guard). + model = pybamm.BaseModel() + u = pybamm.Variable("u") + b = pybamm.InputParameter("b", expected_size=2) + model.rhs = {u: -(pybamm.Index(b, 0) + pybamm.Index(b, 1)) * u} + model.initial_conditions = {u: 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + with pytest.raises( + pybamm.SolverError, match=r"vector-width input parameters.*'b'" + ): + pybamm.DiffsolSolver().solve( + model, + t_eval=np.linspace(0, 1, 10), + inputs={"b": np.array([0.2, 0.3])}, + calculate_sensitivities=["b"], + ) + + def _chen_spm_params(self, extra): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + for pybamm_name, input_name in extra.items(): + params[pybamm_name] = pybamm.InputParameter(input_name) + return params + + def _solve(self, fmt, tol, params, inputs, t_eval, calc, t_interp=None): + model = pybamm.lithium_ion.SPM() + model.events = [] + model.convert_to_format = fmt + solver = ( + pybamm.DiffsolSolver(rtol=tol, atol=tol) + if fmt == "rust" + else pybamm.IDAKLUSolver(rtol=tol, atol=tol) + ) + kwargs = {"inputs": inputs, "calculate_sensitivities": calc} + # Both backends must land on the same grid or the arrays are not + # comparable; IDAKLU additionally needs t_interp to avoid its own knots. + if t_interp is not None: + kwargs["t_interp"] = t_interp + elif fmt == "casadi": + kwargs["t_interp"] = t_eval + return pybamm.Simulation(model, parameter_values=params, solver=solver).solve( + t_eval, **kwargs + ) + + def test_weakly_influential_parameter_sensitivity_is_accurate(self): + # D_n has |p * dV/dp| ~ 1e-4; a state-sized atol floor swamps its column. + # Assert against a converged reference: a same-tol oracle shares the defect. + params = self._chen_spm_params( + {"Negative particle diffusivity [m2.s-1]": "D_n"} + ) + inputs = {"D_n": 3.3e-14, "I": 0.5} + t_eval = np.linspace(0, 100, 15) + calc = ["D_n", "I"] + + reference = self._solve("casadi", 1e-12, params, inputs, t_eval, calc) + native = self._solve("rust", 1e-9, params, inputs, t_eval, calc) + + ref = np.asarray(reference["Terminal voltage [V]"].sensitivities["D_n"]).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities["D_n"]).ravel() + assert np.any(ref != 0) + # Pre-fix this column lands at 2.1e-05; the converged value is ~1.6e-08. + np.testing.assert_allclose(got, ref, rtol=1e-6, atol=0.0) + + def test_early_transient_sensitivity_is_accurate(self): + # Error peaks near t=0.1s where |s| is a few percent of its peak; a uniform + # grid samples this weakly, so use a log-spaced early grid. + params = self._chen_spm_params( + {"Positive electrode active material volume fraction": "eps_p"} + ) + inputs = {"I": 5.0, "eps_p": 0.665} + grid = np.unique( + np.concatenate([np.array([0.0]), np.logspace(-3, np.log10(30.0), 25)]) + ) + span = np.array([0.0, 30.0]) + calc = ["I", "eps_p"] + + reference = self._solve( + "casadi", 1e-11, params, inputs, span, calc, t_interp=grid + ) + native = self._solve("rust", 1e-6, params, inputs, span, calc, t_interp=grid) + + for key in calc: + ref = np.asarray( + reference["Terminal voltage [V]"].sensitivities[key] + ).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities[key]).ravel() + peak = np.max(np.abs(ref)) + assert peak > 0 + # Pre-fix (this span): I=4.99e-04, eps_p=1.04e-03 peak-normalised. + # Not the design doc's 3000 s-horizon figures; those don't apply here. + assert np.max(np.abs(got - ref)) / peak < 1e-6, ( + f"early-time {key} sensitivity error too large" + ) + + def test_parameter_entering_initial_conditions_is_seeded(self): + # diffsol once held dy0/dp at zero, reading 5x low at t=0 and drifting + # from there; both IDAKLU backends seed it from jacp_initial_conditions. + name = "Maximum concentration in negative electrode [mol.m-3]" + params = self._chen_spm_params({name: "c_n_max"}) + inputs = {"I": 0.5, "c_n_max": 33133.0} + t_eval = np.linspace(0, 100, 15) + calc = ["I", "c_n_max"] + + reference = self._solve("casadi", 1e-9, params, inputs, t_eval, calc) + native = self._solve("rust", 1e-9, params, inputs, t_eval, calc) + + ref = np.asarray( + reference["Terminal voltage [V]"].sensitivities["c_n_max"] + ).ravel() + got = np.asarray( + native["Terminal voltage [V]"].sensitivities["c_n_max"] + ).ravel() + # Non-vacuous: the t=0 value is the one an unseeded dy0/dp gets wrong. + assert abs(ref[0]) > 0 + np.testing.assert_allclose(got, ref, rtol=1e-5, atol=0.0) + + def test_dy0_dp_seed_does_not_smear_across_columns(self): + # I never reaches y0, so its seed column is zero while c_n_max's is not; + # a seed written to the wrong column would move this one too. + name = "Maximum concentration in negative electrode [mol.m-3]" + params = self._chen_spm_params({name: "c_n_max"}) + inputs = {"I": 0.5, "c_n_max": 33133.0} + t_eval = np.linspace(0, 100, 15) + calc = ["I", "c_n_max"] + + reference = self._solve("casadi", 1e-9, params, inputs, t_eval, calc) + native = self._solve("rust", 1e-9, params, inputs, t_eval, calc) + + ref = np.asarray(reference["Terminal voltage [V]"].sensitivities["I"]).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities["I"]).ravel() + assert np.any(ref != 0) + np.testing.assert_allclose(got, ref, rtol=_PARITY_RTOL, atol=_PARITY_ATOL) + + @pytest.mark.parametrize("parameter_set", ["default", "Chen2020"]) + def test_dfn_sensitivities_solve(self, parameter_set): + # DFN is a DAE (algebraic block ~100 states). Sensitivity error control + # broke this entirely once; both parameter sets must stay solvable. + model = pybamm.lithium_ion.DFN() + params = ( + model.default_parameter_values.copy() + if parameter_set == "default" + else pybamm.ParameterValues("Chen2020") + ) + sens = { + "Current function [A]": "I", + "Positive electrode active material volume fraction": "eps_p", + } + inputs = {} + for pybamm_name, input_name in sens.items(): + inputs[input_name] = float(params[pybamm_name]) + params[pybamm_name] = pybamm.InputParameter(input_name) + model.convert_to_format = "rust" + + grid = np.linspace(0.0, 600.0, 50) + solution = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6), + ).solve( + [float(grid[0]), float(grid[-1])], + t_interp=grid, + inputs=inputs, + calculate_sensitivities=sorted(sens.values()), + ) + + assert solution.solver_statistics.sens_error_control_relaxed is False + for key in sorted(sens.values()): + block = np.asarray(solution["Voltage [V]"].sensitivities[key]).ravel() + assert block.size > 1 + assert np.all(np.isfinite(block)) + assert np.any(block != 0) + + def test_state_sensitivities_subset_selects_the_extreme_scale_input(self): + # Guards the subset mapping against applying I's parameter scale to the + # D_n column: the included parameter is the extreme-magnitude one. + params = self._chen_spm_params( + {"Negative particle diffusivity [m2.s-1]": "D_n"} + ) + inputs = {"D_n": 3.3e-14, "I": 0.5} + t_eval = np.linspace(0, 100, 15) + + native = self._solve( + "rust", _TWO_PARAM_SOLVER_TOL, params, inputs, t_eval, ["D_n"] + ) + casadi = self._solve( + "casadi", _TWO_PARAM_SOLVER_TOL, params, inputs, t_eval, ["D_n"] + ) + + sens_n = native["Terminal voltage [V]"].sensitivities + sens_c = casadi["Terminal voltage [V]"].sensitivities + assert "D_n" in sens_n + assert "I" not in sens_n + assert np.any(np.asarray(sens_n["D_n"]) != 0) + np.testing.assert_allclose( + sens_n["D_n"], sens_c["D_n"], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + +class TestDiffsolSensAtolFactor: + def test_default_sens_atol_factor(self): + solver = pybamm.DiffsolSolver() + assert solver._sens_atol_factor == pytest.approx(1e-3) + + def test_sens_atol_factor_is_configurable(self): + solver = pybamm.DiffsolSolver(sens_atol_factor=1e-2) + assert solver._sens_atol_factor == pytest.approx(1e-2) + + @pytest.mark.parametrize( + "bad", [0, -1.0, float("nan"), float("inf"), "not-a-number", None] + ) + def test_sens_atol_factor_rejects_invalid(self, bad): + with pytest.raises(pybamm.SolverError, match=r"sens_atol_factor"): + pybamm.DiffsolSolver(sens_atol_factor=bad) + + def test_sens_atol_factor_reaches_rust(self): + # Pins the factor actually reaching the Rust binding (not just stored + # on self): a tighter floor forces more BDF steps under error control. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + t_eval = np.linspace(0, 100, 15) + + def solve(sens_atol_factor): + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + return pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=1e-9, atol=1e-9, sens_atol_factor=sens_atol_factor + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"]) + + sol_default = solve(1e-3) + sol_relaxed = solve(1.0) + assert ( + sol_default.solver_statistics.number_of_steps + > sol_relaxed.solver_statistics.number_of_steps + ) + + +class TestFlattenY0Sens: + """``dy0/dp`` normalisation, the seed threading that fixed diffsol's + unseeded initial-condition sensitivities. + + Producers of ``model.y0S_list`` disagree on shape: ``jacp`` hands over a + list of ``(n, 1)`` columns, a step restart a tuple of bare ``(n,)`` ones. + Both must reach the Rust solver as one column-major block. + """ + + @staticmethod + def _flatten(*args, **kwargs): + return pybamm.solvers.diffsol_solver._flatten_y0_sens(*args, **kwargs) + + def test_a_list_of_column_vectors_is_column_major(self): + blocks = [np.array([[1.0], [2.0], [3.0]]), np.array([[4.0], [5.0], [6.0]])] + np.testing.assert_array_equal( + self._flatten(blocks, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_tuple_of_bare_columns_is_the_same_block(self): + # The shape a step restart produces, via full_sens[:, i]. + blocks = (np.array([1.0, 2.0, 3.0]), np.array([4.0, 5.0, 6.0])) + np.testing.assert_array_equal( + self._flatten(blocks, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_single_matrix_is_the_same_block(self): + matrix = np.array([[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]]) + np.testing.assert_array_equal( + self._flatten(matrix, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_casadi_column_is_densified(self): + casadi = pybamm.import_optional_dependency("casadi") + blocks = [casadi.DM([1.0, 2.0]), casadi.DM([3.0, 4.0])] + np.testing.assert_array_equal(self._flatten(blocks, 2, 2), [1.0, 2.0, 3.0, 4.0]) + + def test_an_empty_seed_is_the_all_zero_case(self): + assert self._flatten([], 3, 2).size == 0 + + def test_a_wrong_shape_is_a_solver_error(self): + # Silently reshaping would attach one parameter's seed to another. + with pytest.raises(pybamm.SolverError, match=r"\(3, 2\) was expected"): + self._flatten([np.array([1.0, 2.0, 3.0])], 3, 2) + + def test_a_wrong_state_count_is_a_solver_error(self): + with pytest.raises(pybamm.SolverError, match=r"\(3, 1\) was expected"): + self._flatten([np.array([1.0, 2.0])], 3, 1) + + +class TestDiffsolSensitivitiesAcrossExperimentSteps: + """A step restart re-seeds ``dy0/dp`` from the previous segment's terminal + sensitivities, in a different shape from the ``jacp`` path.""" + + @staticmethod + def _solve(convert_to_format): + parameter_values = pybamm.ParameterValues("Chen2020") + name = "Negative particle diffusivity [m2.s-1]" + parameter_values[name] = pybamm.InputParameter("D_n") + model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format + solver = ( + pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + if convert_to_format == "rust" + else pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8) + ) + return pybamm.Simulation( + model, + parameter_values=parameter_values, + experiment=pybamm.Experiment( + ["Discharge at 1C for 200 seconds", "Rest for 100 seconds"], + period="20 seconds", + ), + solver=solver, + ).solve(inputs={"D_n": 3.3e-14}, calculate_sensitivities=["D_n"]) + + def test_stepped_sensitivities_match_casadi(self): + native = self._solve("rust") + reference = self._solve("casadi") + assert len(native.all_ts) == 2 + + got = np.asarray(native["Voltage [V]"].sensitivities["D_n"]).ravel() + ref = np.asarray(reference["Voltage [V]"].sensitivities["D_n"]).ravel() + assert np.any(ref != 0) + assert got.shape == ref.shape + # Scaled by |p|: dV/dD_n is ~1e10, so a bare rtol reads as noise. + scale = np.abs(ref).max() + np.testing.assert_allclose(got / scale, ref / scale, rtol=0.0, atol=1e-4) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py new file mode 100644 index 0000000000..94b5c58ff1 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py @@ -0,0 +1,721 @@ +"""General ``DiffsolSolver`` behaviour on small models: restartable states from +outputs-only solves, ``t_eval``/``t_interp`` merging, failure handling, and +sensitivity configuration errors. +""" + +import numpy as np +import pytest +from hypothesis import given + +import pybamm +from pybamm.solvers.observation import ( + NativeComputedObservation, + NativeObservation, +) +from tests.shared import get_broken_input_model as _broken_input_model +from tests.strategies import solve_settings +from tests.strategies.input_sweeps import decay_rate_sweeps + + +def _decay_model(with_event=False, with_input=False): + """``dv/dt = -v`` (or ``-k*v``), so every value is analytically known.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + rate = -pybamm.InputParameter("k") * v if with_input else -v + model.rhs = {v: rate} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v, "2v": 2 * v} + if with_event: + model.events = [pybamm.Event("v = 0.5", v - 0.5)] + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +class TestDiffsolOutputsOnlyStates: + """Outputs-only solves return no state trajectory, so first_state and + last_state must be rebuilt from the terminal/initial full state and stay + observable, matching the IDAKLU outputs-only behaviour.""" + + def _solve_outputs(self, model, t_eval): + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + return solver.solve(model, t_eval) + + def test_last_state_carries_the_terminal_state(self): + sol = self._solve_outputs(_decay_model(), np.linspace(0, 1, 5)) + assert sol.termination == "final time" + assert isinstance(sol.observation, NativeObservation) + last = sol.last_state + np.testing.assert_allclose(last["v"].data, np.exp(-1.0), rtol=1e-8) + np.testing.assert_allclose(last["2v"].data, 2 * np.exp(-1.0), rtol=1e-8) + + def test_first_state_rebuilds_from_initial_conditions(self): + sol = self._solve_outputs(_decay_model(), np.linspace(0, 1, 5)) + np.testing.assert_allclose(sol.first_state["v"].data, 1.0, rtol=1e-10) + + def test_event_terminated_last_state_is_the_root_state(self): + sol = self._solve_outputs(_decay_model(with_event=True), np.linspace(0, 2, 9)) + assert sol.termination.startswith("event") + np.testing.assert_allclose(sol.t[-1], np.log(2.0), rtol=1e-6) + np.testing.assert_allclose(sol.last_state["v"].data, 0.5, rtol=1e-6) + + def test_last_state_matches_the_full_state_solve(self): + t_eval = np.linspace(0, 1, 5) + sol_out = self._solve_outputs(_decay_model(), t_eval) + sol_full = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval + ) + # Integrator state against dense interpolation: each carries its own + # error of order the solve tolerance, so neither bounds the other at it. + for solution in (sol_out, sol_full): + np.testing.assert_allclose( + solution.last_state["v"].data, np.exp(-1.0), rtol=1e-9 + ) + np.testing.assert_allclose( + sol_out.last_state["v"].data, + sol_full.last_state["v"].data, + rtol=1e-8, + ) + + +class TestDiffsolTInterp: + """Diffsol merges ``t_eval`` and ``t_interp`` into one dense output grid.""" + + def test_t_interp_merges_with_a_multi_point_t_eval(self): + t_eval = np.array([0.0, 0.5, 1.0]) + t_interp = np.linspace(0.0, 1.0, 21) # mostly not in t_eval + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval, t_interp=t_interp + ) + np.testing.assert_array_equal(sol.t, np.union1d(t_eval, t_interp)) + np.testing.assert_allclose(sol["v"](t_interp), np.exp(-t_interp), rtol=1e-8) + + def test_empty_t_interp_keeps_t_eval(self): + t_eval = np.linspace(0, 1, 5) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval, t_interp=np.array([]) + ) + np.testing.assert_array_equal(sol.t, t_eval) + + def test_bare_span_t_eval_densifies_the_output_grid(self): + # The canonical solve([t0, tf]) call: a two-point solution would make + # every later sol[...](t) read a single chord between the endpoints. + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve(_decay_model(), [0, 1]) + np.testing.assert_array_equal(sol.t, np.linspace(0.0, 1.0, 100)) + np.testing.assert_allclose(sol["v"](sol.t), np.exp(-sol.t), rtol=1e-8) + + def test_bare_span_off_grid_reads_are_not_a_chord(self): + # Halfway across [0, 1] a two-point chord is off by ~0.07; the dense + # grid keeps grid interpolation error at the 1e-5 scale. + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve(_decay_model(), [0, 1]) + t_check = np.array([0.31, 0.5, 0.77]) # deliberately off-grid + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=2e-5 + ) + + def test_outputs_only_rows_follow_the_merged_grid(self): + t_eval = np.array([0.0, 1.0]) + t_interp = np.linspace(0.0, 1.0, 33) + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + sol = solver.solve(_decay_model(), t_eval, t_interp=t_interp) + np.testing.assert_allclose( + sol["2v"](t_interp), 2.0 * np.exp(-t_interp), rtol=1e-8 + ) + + +class TestDiffsolHermiteInterpolation: + """Full-state diffsol solves store the state time derivatives and + interpolate off-grid ``sol[...](t)`` reads with cubic Hermite through the + native ProcessedVariable path, matching IDAKLU's read semantics.""" + + def _solve(self, model, t_eval, hermite=True, **solve_kwargs): + solver = pybamm.DiffsolSolver( + rtol=1e-10, atol=1e-10, hermite_interpolation=hermite + ) + return solver.solve(model, t_eval, **solve_kwargs) + + def test_off_grid_reads_beat_the_chord(self): + t_eval = np.linspace(0, 1, 6) + t_check = t_eval[:-1] + 0.1 # knot-interval midpoints + hermite = self._solve(_decay_model(), t_eval) + linear = self._solve(_decay_model(), t_eval, hermite=False) + hermite_error = np.max(np.abs(hermite["v"](t_check) - np.exp(-t_check))) + linear_error = np.max(np.abs(linear["v"](t_check) - np.exp(-t_check))) + # On an h = 0.2 grid the chord midpoint error is ~4e-3; cubic Hermite + # on the same knots sits at the ~4e-6 scale. + assert hermite_error < 1e-5 + assert linear_error > 1e-3 + + def test_yp_matches_the_analytic_derivative(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5)) + assert sol.hermite_interpolation + # The t0 knot's slope comes from the first accepted step's low-order + # polynomial (~1e-6 off), where IDAKLU stores the consistent yp0. + np.testing.assert_allclose(sol.yp[0], -np.exp(-sol.t), rtol=1e-5) + + def test_disabling_hermite_keeps_the_grid_aligned_path(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5), hermite=False) + assert not sol.hermite_interpolation + assert sol.all_yps is None + assert isinstance(sol["v"], pybamm.ProcessedVariableComputed) + + @pytest.mark.parametrize( + ("points", "stored"), [(4096, True), (4097, False)], ids=["limit", "past"] + ) + def test_a_dense_output_grid_gives_up_the_derivatives(self, points, stored): + # Past the limit the chord is already at the integration error floor. + sol = self._solve(_decay_model(), np.linspace(0, 1, points)) + assert sol.hermite_interpolation is stored + assert (sol.all_yps is not None) is stored + np.testing.assert_allclose(sol["v"].data, np.exp(-sol.t), rtol=1e-7) + + def test_a_dense_grid_reads_off_grid_to_the_solver_tolerance(self): + t_check = np.linspace(0.05, 0.95, 41) + sol = self._solve(_decay_model(), np.linspace(0, 1, 8193)) + assert not sol.hermite_interpolation + np.testing.assert_allclose(sol["v"](t_check), np.exp(-t_check), atol=1e-8) + + def test_hermite_observation_is_the_native_processed_variable(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5)) + assert not isinstance(sol["v"], pybamm.ProcessedVariableComputed) + np.testing.assert_allclose(sol["v"].data, np.exp(-sol.t), rtol=1e-8) + + def test_outputs_only_solves_store_no_derivatives(self): + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + sol = solver.solve(_decay_model(), np.linspace(0, 1, 5)) + assert sol.all_yps is None + assert isinstance(sol.observation, NativeComputedObservation) + + def test_the_event_root_column_carries_the_wound_back_slope(self): + # Coarse grid: the root at ln 2 lands mid-interval, so reads between + # the last grid knot and the root exercise the root column's yp. + sol = self._solve(_decay_model(with_event=True), np.linspace(0, 2, 5)) + assert sol.termination.startswith("event") + np.testing.assert_allclose(sol.yp[0, -1], -0.5, rtol=1e-6) + t_check = np.array([0.6, 0.69]) + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=1e-5 + ) + + def test_sensitivities_survive_the_native_hermite_path(self): + # The interpolating backend reroutes sensitivities from the eager + # per-segment chain rule to the native jvp_trajectory path. + sol = self._solve( + _decay_model(with_input=True), + np.linspace(0, 1, 9), + inputs={"k": 0.5}, + calculate_sensitivities=True, + ) + t = sol.t + np.testing.assert_allclose( + sol["v"].sensitivities["k"], + -t * np.exp(-0.5 * t), + rtol=1e-5, + atol=1e-8, + ) + + def test_stepped_solve_keeps_hermite_across_segments(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10) + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=0.5, npts=6) + assert sol.hermite_interpolation + t_check = np.array([0.25, 0.75]) # off both segments' grids + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=1e-6 + ) + + def test_a_time_discontinuity_corner_is_not_smoothed(self): + # The corner's knots are the ULP bracket pair, so each side's Hermite + # arc uses its own branch's slope; a stale dy would bow the hold side. + solver = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + sol = solver.solve(_ramp_then_hold_model(), np.linspace(0, 10, 11)) + t_check = np.array([4.75, 5.25]) + np.testing.assert_allclose( + sol["v"](t_check), np.array([4.75, 5.0]), rtol=0, atol=1e-6 + ) + + +class TestDiffsolStep: + """step() re-enters _integrate per segment and merges the segments' + native-observation models through Solution.__add__.""" + + def test_stepped_solve_stays_observable_across_segments(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10) + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=0.5, npts=6) + + assert len(sol.sub_solutions) == 2 + assert isinstance(sol.observation, NativeObservation) + # One on-grid point per segment, so grid interpolation adds no error. + t_check = np.array([0.2, 0.7]) + np.testing.assert_allclose(sol["v"](t_check), np.exp(-t_check), rtol=1e-6) + np.testing.assert_allclose(sol.last_state["v"].data, np.exp(-1.0), rtol=1e-6) + + +class TestDiffsolFailureHandling: + def test_integration_failure_raises_solver_error(self): + # dv/dt = v^2 with v(0) = 1 blows up at t = 1; integrating past it + # must surface as a SolverError, not a raw FFI RuntimeError. + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: v**2} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + with pytest.raises(pybamm.SolverError, match=r"diffsol error"): + pybamm.DiffsolSolver().solve(model, np.linspace(0, 2, 10)) + + +class TestDiffsolSolverOptions: + """The ``options`` dict and its route to diffsol's ``OdeSolverOptions``.""" + + def test_defaults_are_used_when_no_options_are_given(self): + solver = pybamm.DiffsolSolver() + assert solver._options == pybamm.DiffsolSolver.DEFAULT_OPTIONS + + def test_an_override_leaves_the_other_defaults_alone(self): + solver = pybamm.DiffsolSolver(options={"max_error_test_failures": 7}) + assert solver._options["max_error_test_failures"] == 7 + untouched = set(pybamm.DiffsolSolver.DEFAULT_OPTIONS) - { + "max_error_test_failures" + } + for key in untouched: + assert solver._options[key] == pybamm.DiffsolSolver.DEFAULT_OPTIONS[key] + + def test_defaults_match_the_rust_defaults(self): + # Rust derives its defaults from diffsol's OdeSolverOptions, so this + # turns a silent drift on a diffsol bump into a failing test. + from pybamm.rust import default_solver_options + + assert pybamm.DiffsolSolver._INTEGRATOR_DEFAULTS == default_solver_options() + + def test_an_unknown_option_is_rejected(self): + # The Rust extractor requires every key, so a typo would otherwise be + # dropped and the caller would silently get the default. + with pytest.raises(pybamm.SolverError, match=r"Unknown diffsol solver option"): + pybamm.DiffsolSolver(options={"max_error_test_failure": 7}) + + def test_the_failure_budget_default_is_not_diffsols_own(self): + # diffsol counts these cumulatively per solve, so its own 50 caps solve + # length rather than divergence: a long run spends them recovering. + assert ( + pybamm.DiffsolSolver.DEFAULT_OPTIONS["max_nonlinear_solver_failures"] + > 10000 + ) + + def test_options_reach_the_integrator(self): + # Non-vacuous end to end: a budget of zero must abort a solve that the + # default budget completes, which it can only do if the dict arrived. + model = _decay_model() + starved = pybamm.DiffsolSolver( + options={"max_nonlinear_solver_failures": 0, "min_timestep": 0.5} + ) + with pytest.raises(pybamm.SolverError, match=r"diffsol error"): + starved.solve(model, np.linspace(0, 1, 5)) + + assert pybamm.DiffsolSolver().solve(model, np.linspace(0, 1, 5)) is not None + + def test_an_out_of_range_option_is_rejected_by_the_solver(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(options={"nonlinear_solver_tolerance": -1.0}) + with pytest.raises(ValueError, match=r"nonlinear_solver_tolerance"): + solver.solve(model, np.linspace(0, 1, 5)) + + +def _solve_both_ways(inputs_list, num_threads, t_eval, with_event=True): + """The same sweep solved at ``num_threads`` and serially.""" + parallel = pybamm.DiffsolSolver( + rtol=1e-8, atol=1e-10, options={"num_threads": num_threads} + ).solve( + _decay_model(with_event=with_event, with_input=True), t_eval, inputs=inputs_list + ) + serial = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-10).solve( + _decay_model(with_event=with_event, with_input=True), t_eval, inputs=inputs_list + ) + return parallel, serial + + +def _assert_sweeps_identical(parallel, serial): + """Bit-identical, not close: the batch fans out over the same entry point.""" + for i, (got, want) in enumerate(zip(parallel, serial, strict=True)): + np.testing.assert_array_equal(got.t, want.t, err_msg=f"set {i} times") + np.testing.assert_array_equal(got.y, want.y, err_msg=f"set {i} states") + assert got.termination == want.termination, f"set {i} termination" + assert ( + got.solver_statistics.number_of_steps + == want.solver_statistics.number_of_steps + ), f"set {i} step count" + + +class TestDiffsolNumThreads: + """``num_threads`` means "solve this many input sets at once".""" + + @staticmethod + def _sweep(n): + return [{"k": 0.3 * (i + 1)} for i in range(n)] + + @pytest.mark.parametrize("num_threads", [1, 2, 8]) + @pytest.mark.parametrize("with_event", [False, True]) + def test_a_parallel_sweep_is_identical_to_the_serial_one( + self, num_threads, with_event + ): + parallel, serial = _solve_both_ways( + self._sweep(6), num_threads, np.linspace(0, 3, 40), with_event=with_event + ) + _assert_sweeps_identical(parallel, serial) + + def test_results_follow_input_order_not_completion_order(self): + # Descending solve cost, so returning completion order would reverse it. + t_eval = np.linspace(0, 3, 40) + inputs_list = [{"k": k} for k in (0.1, 0.5, 2.0, 9.0)] + solutions = pybamm.DiffsolSolver( + rtol=1e-10, atol=1e-12, options={"num_threads": 4} + ).solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + + for inputs, solution in zip(inputs_list, solutions, strict=True): + assert solution.all_inputs[0] == inputs + np.testing.assert_allclose( + solution["v"](t_eval), + np.exp(-inputs["k"] * t_eval), + rtol=1e-6, + atol=1e-12, + ) + + @pytest.mark.parametrize("num_threads", [1, 4]) + def test_a_failing_set_is_named_whether_batched_or_serial(self, num_threads): + # The message must not depend on how the sweep was scheduled. + model = _broken_input_model() + inputs_list = [{"k": k} for k in (1.0, 2.0, -1.0, 3.0, 4.0)] + solver = pybamm.DiffsolSolver(options={"num_threads": num_threads}) + with pytest.raises(pybamm.SolverError, match=r"input set 2 of 5") as excinfo: + solver.solve(model, np.linspace(0, 1, 10), inputs=inputs_list) + assert "diffsol error" in str(excinfo.value) + + def test_a_lone_failing_set_is_not_dressed_up_as_a_sweep(self): + with pytest.raises(pybamm.SolverError, match=r"diffsol error") as excinfo: + pybamm.DiffsolSolver().solve( + _broken_input_model(), np.linspace(0, 1, 10), inputs=[{"k": -1.0}] + ) + assert "input set" not in str(excinfo.value) + + def test_one_pool_is_shared_by_every_solver_of_the_same_width(self): + from pybamm.rust._core import _pool_ids + + t_eval = np.linspace(0, 1, 10) + inputs_list = self._sweep(8) + first = pybamm.DiffsolSolver(options={"num_threads": 8}) + first.solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + pool_id = _pool_ids()[8] + + second = pybamm.DiffsolSolver(options={"num_threads": 8}) + second.solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + assert _pool_ids()[8] == pool_id + + def test_the_default_builds_no_pool_at_all(self): + # rayon is never constructed in a default-configured process. + from pybamm.rust._core import _pool_ids + + pybamm.DiffsolSolver().solve( + _decay_model(with_input=True), np.linspace(0, 1, 10), inputs=self._sweep(4) + ) + assert 1 not in _pool_ids() + + def test_pools_are_keyed_on_the_configured_width_not_the_sweep(self): + # Keying on the sweep would build a fresh pool per distinct sweep size, + # so a process solving 2, 3, ... 8 sets would accumulate them all. + from pybamm.rust._core import _pool_ids + + # The cache is process-wide, so only the keys this solver adds are ours. + before = set(_pool_ids()) + model = _decay_model(with_input=True) + solver = pybamm.DiffsolSolver(options={"num_threads": 7}) + for n_sets in (2, 3): + solver.solve(model, np.linspace(0, 1, 10), inputs=self._sweep(n_sets)) + added = set(_pool_ids()) - before + assert 7 in _pool_ids() + assert added.isdisjoint({2, 3}) + + @pytest.mark.parametrize("num_threads", [0, -1, 2.5, "4", True, None]) + def test_a_num_threads_that_is_not_a_count_is_rejected(self, num_threads): + with pytest.raises(pybamm.SolverError, match=r"num_threads must be an integer"): + pybamm.DiffsolSolver(options={"num_threads": num_threads}) + + def test_num_threads_does_not_reach_the_integrator_options(self): + solver = pybamm.DiffsolSolver(options={"num_threads": 4}) + assert "num_threads" not in solver._integrator_options() + assert set(solver._integrator_options()) == set( + pybamm.DiffsolSolver._INTEGRATOR_DEFAULTS + ) + + +@solve_settings +@given(rates=decay_rate_sweeps()) +def test_a_batch_is_indistinguishable_from_the_serial_loop(rates): + """Whatever the sweep, four threads reproduce the serial loop exactly.""" + parallel, serial = _solve_both_ways( + [{"k": rate} for rate in rates], 4, np.linspace(0, 3, 40) + ) + _assert_sweeps_identical(parallel, serial) + + +class TestDiffsolIntegrationTime: + def test_each_set_gets_its_own_integration_time(self): + # Under a batch the wall clocks overlap, so a Python-side timer would + # stamp every set with the batch duration. + inputs_list = [{"k": 0.3 * (i + 1)} for i in range(4)] + solutions = pybamm.DiffsolSolver(options={"num_threads": 4}).solve( + _decay_model(with_input=True), np.linspace(0, 2, 30), inputs=inputs_list + ) + times = [solution.integration_time for solution in solutions] + assert all(time > 0 for time in times) + assert len(set(times)) > 1 + + def test_the_reported_time_covers_setup_and_integration(self): + solution = pybamm.DiffsolSolver().solve(_decay_model(), np.linspace(0, 2, 30)) + statistics = solution.solver_statistics + assert solution.integration_time >= ( + statistics.ic_time_secs + statistics.solver_setup_time_secs + ) + + +class TestDiffsolSensitivityConfigErrors: + def test_unknown_sensitivity_parameter_is_a_single_clear_error(self): + model = _decay_model(with_input=True) + with pytest.raises(ValueError, match=r"no sensitivity parameters") as excinfo: + pybamm.DiffsolSolver().solve( + model, + np.linspace(0, 1, 5), + inputs={"k": 1.0}, + calculate_sensitivities=["not_an_input"], + ) + # A config error must not be routed through the relaxed-retry path, + # which would double-report it as an error-control failure. + assert "retry" not in str(excinfo.value) + + +def _two_scale_model(): + """``du/dt = -u`` and ``dw/dt = -w`` with ``u`` at O(1) and ``w`` at O(1e4). + + Two states four decades apart, which is the case a per-state ``atol`` exists + for. + """ + model = pybamm.BaseModel() + u = pybamm.Variable("u") + w = pybamm.Variable("w") + model.rhs = {u: -u, w: -w} + model.initial_conditions = {u: 1.0, w: 1e4} + model.variables = {"u": u, "w": w} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +@pytest.mark.parametrize("solver_class", [pybamm.DiffsolSolver, pybamm.IDAKLUSolver]) +class TestPerStateAtol: + """``atol`` reaches the integrator as one entry per state, so states of + different magnitudes can be toleranced separately -- which ``rtol``, already + scaled by each state's own value, cannot do.""" + + def _steps(self, solver_class, atol): + solution = solver_class(rtol=1e-6, atol=atol).solve( + _two_scale_model(), np.linspace(0, 1, 5) + ) + return solution.solver_statistics.number_of_steps + + def test_a_uniform_array_matches_the_scalar(self, solver_class): + t_eval = np.linspace(0, 1, 5) + scalar = solver_class(rtol=1e-8, atol=1e-8).solve(_two_scale_model(), t_eval) + array = solver_class(rtol=1e-8, atol=np.full(2, 1e-8)).solve( + _two_scale_model(), t_eval + ) + np.testing.assert_array_equal(array["u"](t_eval), scalar["u"](t_eval)) + np.testing.assert_array_equal(array["w"](t_eval), scalar["w"](t_eval)) + + def test_every_entry_reaches_its_own_state(self, solver_class): + # Broadcasting either entry over both states would leave one of the + # mixed arrays at the all-loose count. + loose = self._steps(solver_class, np.full(2, 1e-1)) + assert self._steps(solver_class, 1e-1) == loose + assert self._steps(solver_class, np.array([1e-12, 1e-1])) > loose + assert self._steps(solver_class, np.array([1e-1, 1e-12])) > loose + + def test_a_per_state_atol_survives_a_dae_solve(self, solver_class): + # The initial-condition root solver takes a single tolerance rather + # than the array; diffsol picks it up by default during set_up. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + z = pybamm.Variable("z") + model.rhs = {u: -u} + model.algebraic = {z: z - 2 * u} + model.initial_conditions = {u: 1.0, z: 2.0} + model.variables = {"u": u, "z": z} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + t_eval = np.linspace(0, 1, 5) + solution = solver_class(rtol=1e-8, atol=np.array([1e-9, 1e-7])).solve( + model, t_eval + ) + np.testing.assert_allclose( + solution["z"](t_eval), 2 * np.exp(-t_eval), rtol=1e-6 + ) + + def test_a_config_round_trip_keeps_the_per_state_atol(self, solver_class): + # to_config puts the array through JSON, which has no arrays to put back. + atol = np.array([1e-1, 1e-12]) + restored = pybamm.BaseSolver.from_config( + solver_class(rtol=1e-6, atol=atol).to_config() + ) + assert self._steps(solver_class, restored.atol) == self._steps( + solver_class, atol + ) + + def test_a_wrong_width_atol_is_rejected(self, solver_class): + solver = solver_class(atol=np.full(3, 1e-6)) + with pytest.raises(pybamm.SolverError, match=r"shape \(3,\) but \(2,\)"): + solver.solve(_two_scale_model(), np.linspace(0, 1, 5)) + + @pytest.mark.parametrize( + ("bad_atol", "message"), + [ + ("tight", r"a float or one value per state"), + (["tight", "loose"], r"must all be numbers"), + (np.array(["tight", "loose"]), r"must all be numbers"), + ], + ) + def test_a_non_numeric_atol_is_rejected(self, solver_class, bad_atol, message): + # NumPy's own ValueError would otherwise escape the documented contract. + solver = solver_class(atol=bad_atol) + with pytest.raises(pybamm.SolverError, match=message): + solver.solve(_two_scale_model(), np.linspace(0, 1, 5)) + + +def _ramp_then_hold_model(): + """``dv/dt = 1`` until t = 5 and 0 after, so ``v(t) = min(t, 5)``.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: pybamm.t < 5} + model.initial_conditions = {v: pybamm.Scalar(0)} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +class TestDiffsolTimeDiscontinuities: + """A constant time discontinuity reaches diffsol as a stop time, the way it + reaches IDAKLU, rather than splitting the run into separate solves.""" + + @pytest.mark.parametrize( + "solver_class", [pybamm.DiffsolSolver, pybamm.IDAKLUSolver] + ) + def test_a_heaviside_in_time_is_solved_in_one_pass(self, solver_class): + solver = solver_class(rtol=1e-8, atol=1e-8) + solution = solver.solve(_ramp_then_hold_model(), np.linspace(0, 10, 21)) + + assert len(solution.sub_solutions) == 1 + np.testing.assert_allclose( + solution["v"](np.array([2.5, 5.0, 7.5])), + np.array([2.5, 5.0, 5.0]), + rtol=1e-6, + atol=1e-6, + ) + + +def _native_prepared_solver(with_sens=False): + """``dy/dt = -k*y`` natively, with two outputs against its single state.""" + from pybamm.rust import CompiledModel, ExprGraph, PreparedSolver + + g = ExprGraph() + y = g.state_vector(0, 1) + rate = g.mul(g.mul(g.scalar(-1.0), g.input_parameter("k")), y) + model = CompiledModel.from_expr( + g, + rate, + np.ones(1), + np.arange(2, dtype=np.int64), + np.zeros(1, dtype=np.int64), + n_inputs=1, + sens_param_indices=[0] if with_sens else [], + output_exprs=[g.mul(g.scalar(2.0), y), g.mul(g.scalar(3.0), y)], + event_exprs=[], + ) + return PreparedSolver(model, 1e-10, 1e-10) + + +class TestNativeSolveRequest: + """The payload flags on the single ``solve`` entry point, at the FFI seam.""" + + _T_EVAL = np.linspace(0.0, 1.0, 11) + _NO_STOPS = np.array([], dtype=np.float64) + + def _solve(self, solver, **flags): + return solver.solve( + self._T_EVAL, self._NO_STOPS, np.array([1.0]), np.array([1.0]), **flags + ) + + def test_the_row_space_follows_the_outputs_flag(self): + solver = _native_prepared_solver() + states = self._solve(solver) + outputs = self._solve(solver, outputs=True) + + # One state against two output expressions, 2y and 3y: the flag swaps the + # rows of the same trajectory field, and the columns are untouched. + assert states.y.shape == (1, self._T_EVAL.size) + assert outputs.y.shape == (2, self._T_EVAL.size) + np.testing.assert_allclose(outputs.y[0], 2.0 * states.y[0], rtol=1e-8) + np.testing.assert_allclose(outputs.y[1], 3.0 * states.y[0], rtol=1e-8) + + def test_payloads_not_asked_for_read_as_none(self): + solver = _native_prepared_solver(with_sens=True) + + assert self._solve(solver).yS is None + assert self._solve(solver, sensitivities=True).yS is not None + # yp is the state trajectory's slopes, and this solver does not store them. + assert self._solve(solver).yp is None + + def test_a_seed_without_sensitivities_is_rejected(self): + solver = _native_prepared_solver(with_sens=True) + + # Silently ignoring the seed would return zero sensitivities that look + # computed, so the inconsistent pair is refused instead. + with pytest.raises(ValueError, match=r"y0_sens was given but"): + self._solve(solver, y0_sens=np.array([0.5])) + + def test_a_batch_answers_one_shared_request(self): + """The payload flags are shared by the batch, so every set comes back + with the same payloads its serial twin would.""" + solver = _native_prepared_solver(with_sens=True) + rates = [0.5, 1.0, 2.0] + + results = solver.solve_batch( + self._T_EVAL, + self._NO_STOPS, + np.ones((len(rates), 1)), + np.array(rates).reshape(-1, 1), + 1, + outputs=True, + sensitivities=True, + ) + + assert len(results) == len(rates) + for i, (result, rate) in enumerate(zip(results, rates, strict=True)): + serial = solver.solve( + self._T_EVAL, + self._NO_STOPS, + np.array([1.0]), + np.array([rate]), + outputs=True, + sensitivities=True, + ) + assert result.yS is not None, f"set {i} dropped its blocks" + np.testing.assert_array_equal(result.y, serial.y, err_msg=f"set {i}") + np.testing.assert_array_equal( + result.yS[0], serial.yS[0], err_msg=f"set {i} sensitivities" + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py index 907d54e2ae..d233cfce21 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -5,6 +5,7 @@ import warnings from contextlib import redirect_stdout +import casadi import numpy as np import pandas as pd import pytest @@ -12,7 +13,28 @@ from scipy.interpolate import CubicHermiteSpline import pybamm -from tests import get_discretisation_for_testing, no_internet_connection +from pybamm.solvers.observation import ( + NativeInterpolatingObservation, + NativeObservation, +) +from pybamm.solvers.variable_observer import NativeObserver +from tests import ( + get_broken_input_model, + get_discretisation_for_testing, + no_internet_connection, +) + + +def _rust_decay_model(with_input=True): + """``dvar/dt = -rate*var`` (or ``-var``), lowered to the Rust backend.""" + model = pybamm.BaseModel() + var = pybamm.Variable("var") + rate = pybamm.InputParameter("rate") if with_input else 1 + model.rhs = {var: -rate * var} + model.initial_conditions = {var: 2 if with_input else 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model def _hermite_wrms(sol_base, sol_reduced, atol, rtol) -> list[tuple[int, float]]: @@ -169,6 +191,74 @@ def test_multiple_inputs(self): rtol=1e-4, ) + def test_multiple_inputs_rust_solves_in_parallel(self): + # One evaluator per solver over a shared tape, so a parallel solve must + # match per-input solves with no cross-input interference. + model = _rust_decay_model() + + t_eval = [0, 1] + t_interp = np.linspace(t_eval[0], t_eval[-1], 10) + inputs_list = [{"rate": 0.01 * (i + 1)} for i in range(4)] + + solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10, options={"num_threads": 4}) + batched_solutions = solver.solve( + model, t_eval, inputs=inputs_list, t_interp=t_interp + ) + + sequential_solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10) + for inputs, batched in zip(inputs_list, batched_solutions, strict=True): + sequential = sequential_solver.solve( + model, t_eval, inputs=inputs, t_interp=t_interp + ) + np.testing.assert_array_equal(batched.t, sequential.t) + np.testing.assert_allclose( + batched.y[0], sequential.y[0], rtol=1e-8, atol=1e-10 + ) + + @pytest.mark.parametrize( + ("num_threads", "num_solvers"), [(4, None), (4, 1), (4, 2), (1, 1)] + ) + def test_rust_runs_one_solver_per_thread(self, num_threads, num_solvers): + # Equality is what makes SetupOptions derive one thread per solver, so no + # solver is given the OpenMP N_Vectors that were 28x slower on DFN x32. + model = _rust_decay_model() + + options = {"num_threads": num_threads} + if num_solvers is not None: + options["num_solvers"] = num_solvers + solver = pybamm.IDAKLUSolver(rtol=1e-5, atol=1e-5, options=options) + solver.solve(model, [0, 1], inputs=[{"rate": 0.01 * (i + 1)} for i in range(4)]) + + assert solver._options["num_threads"] == num_threads + assert solver._options["num_solvers"] == num_threads + assert len(solver._setup["rust_evaluators"]) == num_threads + + def test_rust_evaluator_pool_hands_out_distinct_evaluators(self): + solver = pybamm.IDAKLUSolver(options={"num_threads": 3}) + solver.set_up(_rust_decay_model(with_input=False)) + pool = solver._setup["rust_evaluators"] + + assert len(pool) == 3 + # set_up's solver group consumed every handout: an address is given out + # once, so nothing else can be handed a solver's evaluator. + with pytest.raises(RuntimeError, match=r"already handed to a solver"): + pool.as_ptr(0) + + fresh = solver._setup["rust_model"].evaluator_pool(3) + assert len({fresh.as_ptr(i) for i in range(3)}) == 3 + with pytest.raises(IndexError, match=r"out of range for a pool of 3"): + fresh.as_ptr(3) + with pytest.raises(RuntimeError, match=r"already handed to a solver"): + fresh.as_ptr(1) + + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_a_failing_input_set_is_named(self, convert_to_format): + model = get_broken_input_model(convert_to_format) + solver = pybamm.IDAKLUSolver(options={"num_threads": 4}) + inputs_list = [{"k": k} for k in (1.0, 2.0, -1.0, 3.0, 4.0)] + with pytest.raises(pybamm.SolverError, match=r"input set 2:"): + solver.solve(model, np.linspace(0, 1, 10), inputs=inputs_list) + def test_model_events(self): # Create model model = pybamm.BaseModel() @@ -330,6 +420,72 @@ def test_input_params(self): atol=1e-6, ) + def test_vector_input_parameter_rust(self): + # rust `check_p` must compare packed-input length against total input + # *width*, not input count, or vector inputs are misrejected. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b", expected_size=2) + model.rhs = {u: -a * u * (pybamm.Index(b, 0) + pybamm.Index(b, 1))} + model.initial_conditions = {u: 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver() + inputs = {"a": 0.5, "b": np.array([[0.2], [0.3]])} + sol = solver.solve(model, np.linspace(0, 1, 10), inputs=inputs) + np.testing.assert_allclose( + sol["u"].data, np.exp(-0.5 * 0.5 * sol.t), rtol=1e-3, atol=1e-5 + ) + + def test_calculate_sensitivities_rejects_vector_width_input_parameter(self): + # Rust JVP/tangent seeds one scalar direction per named parameter; a + # width>1 parameter would silently under-seed, so this must raise instead. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + b = pybamm.InputParameter("b", expected_size=2) + model.rhs = {u: -(pybamm.Index(b, 0) + pybamm.Index(b, 1)) * u} + model.initial_conditions = {u: 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver() + with pytest.raises( + pybamm.SolverError, match=r"vector-width input parameters.*'b'" + ): + solver.solve( + model, + np.linspace(0, 1, 10), + inputs={"b": np.array([0.2, 0.3])}, + calculate_sensitivities=["b"], + ) + + def test_scalar_sensitivities_with_wide_input_registered_first(self): + # Guard only rejects sensitivities FOR a wide parameter; a scalar request + # with an earlier-registered width-2 input must still seed correctly. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + a = pybamm.InputParameter("a", expected_size=2) + c = pybamm.InputParameter("c") + model.rhs = {u: -c * (pybamm.Index(a, 0) + pybamm.Index(a, 1)) * u} + model.initial_conditions = {u: 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10) + sol = solver.solve( + model, + np.linspace(0, 1, 11), + inputs={"a": np.array([0.4, 0.6]), "c": 0.5}, + calculate_sensitivities=["c"], + ) + # u = exp(-c*(a0+a1)*t) with a0+a1 = 1: du/dc = -t*exp(-0.5*t) + analytic = -sol.t * np.exp(-0.5 * sol.t) + np.testing.assert_allclose( + np.asarray(sol["u"].sensitivities["c"]).ravel(), + analytic, + rtol=1e-5, + atol=1e-8, + ) + def test_sensitivities_initial_condition(self): for output_variables in [[], ["2v"]]: model = pybamm.BaseModel() @@ -656,6 +812,17 @@ def test_failures(self): with pytest.raises(pybamm.SolverError): solver.solve(model, t_eval) + def test_rust_klu_requires_jacobian(self): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1} + model.use_jacobian = False + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + with pytest.raises(pybamm.SolverError, match=r"KLU requires the Jacobian"): + pybamm.IDAKLUSolver().solve(model, np.array([0.0, 1.0])) + def test_dae_solver_algebraic_model(self): model = pybamm.BaseModel() var = pybamm.Variable("var") @@ -776,6 +943,9 @@ def test_debug_log_flushed_when_solve_raises(self, caplog): model.initial_conditions = {u: 1} model.variables = {"u": u} pybamm.Discretisation().process_model(model) + # The worker-thread flush under test is OpenMP-only; rust clamps to one + # solver, so every trace would come from the calling thread instead. + model.convert_to_format = "casadi" # Two solves per solver, so one solver buffers on a worker thread solver = pybamm.IDAKLUSolver(options={"num_threads": 2}) @@ -786,10 +956,10 @@ def test_debug_log_flushed_when_solve_raises(self, caplog): ): solver.solve(model, np.array([0.0, 5.0]), inputs=inputs) - # Each solver throws on its first solve, and the calling thread streams - # its own, so the second trace can only come from the pre-rethrow flush + # Every set is attempted, and only the calling thread's traces stream + # directly, so the rest can only come from the pre-rethrow flush starts = [m for m in caplog.messages if m.startswith("Integrating from t =")] - assert len(starts) == 2 + assert len(starts) == 4 def test_solve_interrupted_from_debug_logger(self, caplog, monkeypatch): sim = pybamm.Simulation(pybamm.lithium_ion.SPM()) @@ -893,6 +1063,59 @@ def test_setup_options(self): with pytest.raises(ValueError): _ = solver.solve(model, t_eval, t_interp=t_interp) + def test_rust_dense_jacobian_matches_sparse(self): + # rhs depends only on u, leaving a structural zero (row 0, col v) that + # exercises the dense scatter path; nonzero ICs force real Newton iterations. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + model.rhs = {u: -0.1 * u} + model.algebraic = {v: v - u} + model.initial_conditions = {u: 1, v: 1} + model.convert_to_format = "rust" + disc = pybamm.Discretisation() + disc.process_model(model) + + t_eval = np.array([0.0, 1.0]) + t_interp = np.linspace(0, 1, 100) + base = pybamm.IDAKLUSolver(atol=1e-8, rtol=1e-8).solve( + model, t_eval, t_interp=t_interp + ) + for jacobian, linsol in [ + ("dense", "SUNLinSol_Dense"), + ("none", "SUNLinSol_Dense"), + ]: + options = { + "jacobian": jacobian, + "linear_solver": linsol, + "preconditioner": "none", + "max_num_steps": 10_000, + } + solver = pybamm.IDAKLUSolver(atol=1e-8, rtol=1e-8, options=options) + soln = solver.solve(model, t_eval, t_interp=t_interp) + np.testing.assert_allclose(soln.y, base.y, rtol=1e-5, atol=1e-4) + + def test_rust_pure_algebraic_jacobian_diagonal(self): + # Pure-algebraic model (empty rhs): the empty child's zero-derivative must + # not widen to length 1, which would shift the diagonal and zero the Jacobian. + model = pybamm.BaseModel() + var = pybamm.Variable("var") + model.algebraic = {var: var + 1} + model.initial_conditions = {var: 0} + model.convert_to_format = "rust" + + disc = pybamm.Discretisation() + disc.process_model(model) + + solver = pybamm.IDAKLUSolver() + solution = solver.solve(model, [0, 1]) + np.testing.assert_array_equal(solution.y, -1) + + def test_an_unknown_option_is_rejected(self): + # Merging a misspelt key through leaves the caller on the default. + with pytest.raises(pybamm.SolverError, match=r"Unknown IDAKLU solver option"): + pybamm.IDAKLUSolver(options={"num_thread": 4}) + def test_solver_options(self): model = pybamm.BaseModel() u = pybamm.Variable("u") @@ -1085,6 +1308,9 @@ def test_with_sparse_output_variables_and_sensitivities(self): output_variables=["Negative particle flux [mol.m-2.s-1]"], ) model = pybamm.lithium_ion.DFN() + # The nnz-compression this guards against only occurs on the casadi path; + # the rust equivalent is test_sparse_output_variable_sensitivities_rust_matches_casadi. + model.convert_to_format = "casadi" params = model.default_parameter_values params.update({"Current function [A]": "[input]"}) sim = pybamm.Simulation(model, solver=solver, parameter_values=params) @@ -1094,6 +1320,46 @@ def test_with_sparse_output_variables_and_sensitivities(self): ): sim.solve([0, 100], inputs=input_parameters, calculate_sensitivities=True) + def test_sparse_output_variable_sensitivities_rust_matches_casadi(self): + # Rust output lengths are always dense (== prod(shape)), so the casadi + # nnz-compression bug behind the guard above is unreachable on rust. + input_parameters = { + "Current function [A]": 0.222, + "Separator porosity": 0.3, + } + var_name = "Negative particle flux [mol.m-2.s-1]" + + # Reference: full-state casadi solve bypasses the output-variables + # fast path (and its guard) entirely. + model_ref = pybamm.lithium_ion.DFN() + model_ref.convert_to_format = "casadi" + params_ref = model_ref.default_parameter_values + params_ref.update({"Current function [A]": "[input]"}) + sim_ref = pybamm.Simulation( + model_ref, solver=pybamm.IDAKLUSolver(), parameter_values=params_ref + ) + sol_ref = sim_ref.solve( + [0, 100], inputs=input_parameters, calculate_sensitivities=True + ) + + model = pybamm.lithium_ion.DFN() + model.convert_to_format = "rust" + params = model.default_parameter_values + params.update({"Current function [A]": "[input]"}) + solver = pybamm.IDAKLUSolver(output_variables=[var_name]) + sim = pybamm.Simulation(model, solver=solver, parameter_values=params) + sol = sim.solve([0, 100], inputs=input_parameters, calculate_sensitivities=True) + + np.testing.assert_allclose( + sol[var_name].data, sol_ref[var_name].data, rtol=1e-6, atol=1e-10 + ) + np.testing.assert_allclose( + np.asarray(sol[var_name].sensitivities["Current function [A]"]), + np.asarray(sol_ref[var_name].sensitivities["Current function [A]"]), + rtol=1e-5, + atol=1e-9, + ) + def test_with_output_variables_and_sensitivities(self): # Construct a model and solve for all variables, then test # the 'output_variables' option for each variable in turn, confirming @@ -1295,6 +1561,114 @@ def test_pickle_roundtrip_preserves_closest_event_idx(self): "closest_event_idx after a root return" ) + def test_closest_event_idx_set_after_root_return_rust(self): + # Rust-path variant of test_closest_event_idx_set_after_root_return; + # _set_up_rust must bind the compiled model's event tapes for this to pass. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + cycle = ( + "Discharge at 1C until 3.0 V", + "Charge at 1C until 4.2 V", + "Hold at 4.2 V until C/50", + ) + sim = pybamm.Simulation( + model, + experiment=pybamm.Experiment([cycle] * 2, period=300), + solver=pybamm.IDAKLUSolver(output_variables=["Voltage [V]"]), + ) + sim.solve() + + event_steps = [ + step + for cycle_sol in sim.solution.cycles + for step in cycle_sol.steps + if step.termination.startswith("event:") + ] + assert event_steps, "expected at least one event-terminated step" + for step in event_steps: + assert step.all_models[-1].convert_to_format == "rust" + assert step.closest_event_idx is not None, ( + f"event-terminated step {step.termination!r} has " + f"closest_event_idx=None — BaseSolver will fall back to " + f"per-step Python event re-evaluation" + ) + terminate_events = [ + e + for e in step.all_models[-1].events + if e.event_type == pybamm.EventType.TERMINATION + ] + picked = terminate_events[step.closest_event_idx].name + assert step.termination == f"event: {picked}", ( + f"closest_event_idx={step.closest_event_idx} resolves to " + f"{picked!r}, but step.termination is {step.termination!r}" + ) + + def test_rust_fused_events_parity_with_casadi(self): + # SPMe has >= 2 termination events, so rust evaluates them through the fused + # tape; a tightened cut-off fires one, and both paths must agree on which. + def _solve(fmt): + params = pybamm.ParameterValues("Chen2020") + params["Lower voltage cut-off [V]"] = 3.5 + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = fmt + sim = pybamm.Simulation( + model, parameter_values=params, solver=pybamm.IDAKLUSolver() + ) + return sim.solve(np.linspace(0, 3600, 100)) + + sol_rust = _solve("rust") + sol_casadi = _solve("casadi") + + assert sol_rust.termination.startswith("event:"), ( + f"expected an event to fire, got {sol_rust.termination!r}" + ) + assert sol_rust.termination == sol_casadi.termination + assert sol_rust.closest_event_idx == sol_casadi.closest_event_idx + np.testing.assert_allclose( + float(sol_rust.t_event[0]), float(sol_casadi.t_event[0]), rtol=1e-4 + ) + + def test_rust_events_survive_compiled_model_pickle(self): + # Pickling round-trips through `_rebuild`, which re-runs `build_from_parts` + # and thus `fuse_events`, so the rebuilt model must evaluate events the same. + import pickle + + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=pybamm.IDAKLUSolver()) + sol = sim.solve(np.linspace(0, 3600, 10)) + + rust_model = sim._solver._setup["rust_model"] + assert rust_model.n_events >= 2 # fusion is active + + # A representative evaluation point from the solve. + y = np.ascontiguousarray(sol.all_ys[0][:, 0], dtype=np.float64) + p = np.zeros(rust_model.n_inputs, dtype=np.float64) + t = float(sol.all_ts[0][0]) + before = [np.asarray(fn(t, y, p)).ravel() for fn in rust_model.events] + + restored = pickle.loads(pickle.dumps(rust_model)) + assert restored.n_events == rust_model.n_events + after = [np.asarray(fn(t, y, p)).ravel() for fn in restored.events] + for a, b in zip(before, after, strict=True): + np.testing.assert_array_equal(a, b) + + def test_rust_diffsol_event_termination(self): + # Diffsol root-finding uses the same fused event tape (RootOp path). + params = pybamm.ParameterValues("Chen2020") + params["Lower voltage cut-off [V]"] = 3.5 + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + sim = pybamm.Simulation( + model, parameter_values=params, solver=pybamm.DiffsolSolver() + ) + sol = sim.solve(np.linspace(0, 3600, 100)) + assert sol.termination.startswith("event:"), ( + f"expected diffsol root-finding to terminate on an event, " + f"got {sol.termination!r}" + ) + assert float(sol.t[-1]) < 3600.0 + def test_simulation_period(self): model = pybamm.lithium_ion.DFN() parameter_values = pybamm.ParameterValues("Chen2020") @@ -1771,6 +2145,18 @@ def test_hermite_reduction_factor_incompatible(self): with pytest.warns(pybamm.SolverWarning, match="not currently supported"): sim.solve([0, 1], inputs={"I": 1.0}, calculate_sensitivities=True) + def test_hermite_reduction_factor_sensitivities_warning_rust(self): + # Regression: `_set_up_rust` used to return before the casadi path's + # hermite_reduction_factor + sensitivities check, skipping the warning. + model_sens = pybamm.lithium_ion.SPM() + model_sens.convert_to_format = "rust" + param = model_sens.default_parameter_values + param["Current function [A]"] = pybamm.InputParameter("I") + solver = pybamm.IDAKLUSolver(options={"hermite_reduction_factor": 2.0}) + sim = pybamm.Simulation(model_sens, parameter_values=param, solver=solver) + with pytest.warns(pybamm.SolverWarning, match="not currently supported"): + sim.solve([0, 1], inputs={"I": 1.0}, calculate_sensitivities=True) + def test_reduce_solution_basic(self): """Test basic post-hoc reduce_solution: fewer points, finite yps, bounded error.""" model = pybamm.lithium_ion.SPM() @@ -1821,6 +2207,29 @@ def test_reduce_solution_metadata(self): for rte, ste in zip(reduced.all_t_evals, sol.all_t_evals, strict=True): np.testing.assert_array_equal(rte, ste) + def test_reduce_solution_keeps_native_observation(self): + """A reduced solution must not fall back to CasADi observation.""" + model = pybamm.lithium_ion.SPM() + solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-8) + sol = pybamm.Simulation(model, solver=solver).solve([0, 3600]) + assert isinstance(sol.observation, NativeInterpolatingObservation) + + reduced = solver.reduce_solution(sol, hermite_reduction_factor=2.0) + + # Thinning knots leaves the segments and their models alone. + assert isinstance(reduced.observation, NativeInterpolatingObservation) + assert reduced.observation.segment_models == sol.observation.segment_models + assert reduced.observation.compile_cache is sol.observation.compile_cache + assert isinstance(reduced["Terminal voltage [V]"]._observer, NativeObserver) + # ... and the reduced spline still reads within its error budget + t = np.linspace(0, 3600, 51) + np.testing.assert_allclose( + reduced["Terminal voltage [V]"](t), + sol["Terminal voltage [V]"](t), + rtol=0, + atol=1e-4, + ) + def test_reduce_solution_vs_online(self): """Compare post-hoc reduce_solution with online knot reduction on a drive cycle. @@ -1935,3 +2344,697 @@ def test_solution_user_options_survive_pickle(self, tmp_path): np.testing.assert_allclose( loaded["2u"].entries, sol["2u"].entries, rtol=1e-12, atol=1e-12 ) + + def test_idaklu_dispatch_is_flag_driven(self): + with pytest.raises(TypeError): + pybamm.IDAKLUSolver(evaluator="rust") + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + solver = pybamm.IDAKLUSolver() + solver.set_up(model, inputs=[{"a": 0.5}]) + assert "rust_model" in solver._setup + assert model.convert_to_format == "rust" + + def test_idaklu_rust_output_sensitivities_set_up_ok(self): + # output_variables + calculate_sensitivities is supported for + # convert_to_format="rust", so set_up must not raise for the combination. + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + model.calculate_sensitivities = ["a"] + solver = pybamm.IDAKLUSolver(output_variables=["u"]) + solver.set_up(model, inputs=[{"a": 0.5}]) + assert model.convert_to_format == "rust" + + def test_idaklu_rust_output_values_ok_without_sens(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + solver = pybamm.IDAKLUSolver(output_variables=["u"]) + solver.set_up(model, inputs=[{"a": 0.5}]) + assert model.convert_to_format == "rust" + + def test_rhs_dot_consistent_init_rust_inputs(self): + # Regression: _rhs_dot_consistent_initialization must stack inputs for rust, + # not hand a dict to the RustEvaluator. + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + solver = pybamm.IDAKLUSolver() + solver.set_up(model, inputs=[{"a": 0.5}]) + y0 = np.asarray(model.y0_list[0]).ravel() + ydot0 = solver._rhs_dot_consistent_initialization(y0, model, 0.0, {"a": 0.5}) + assert ydot0.shape == y0.shape + + def test_sensitivity_consistent_init_gate_fires_for_rust(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("rust") + model.calculate_sensitivities = ["a"] + solver = pybamm.IDAKLUSolver() + solver.set_up(model, inputs=[{"a": 0.5}]) + solver._set_consistent_initialization(model, 0.0, [{"a": 0.5}]) + # 1 sens param: y0full should be len_rhs_and_alg * 2 + assert model.y0full[0].shape[0] == model.len_rhs_and_alg * 2 + + +class TestIDAKLUNativeObservation: + def _solve(self, convert_to_format, calculate_sensitivities=False, t_interp=None): + model = pybamm.lithium_ion.SPM() + geometry = model.default_geometry + param = model.default_parameter_values + if calculate_sensitivities: + param.update({"Current function [A]": "[input]"}) + param.process_model(model) + param.process_geometry(geometry) + mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver() + t_eval = [0, 3600] + sol = solver.solve( + model, + t_eval, + t_interp=t_interp, + inputs={"Current function [A]": 0.68} if calculate_sensitivities else None, + calculate_sensitivities=calculate_sensitivities, + ) + return sol, solver + + def test_offgrid_values_match_casadi(self): + # An off-grid t_interp exercises native Hermite against CasADi's + # observe_hermite_interp. + t_interp = np.linspace(0, 3600, 97) + cas, _ = self._solve("casadi", t_interp=t_interp) + rust, _ = self._solve("rust", t_interp=t_interp) + for name in [ + "Terminal voltage [V]", + "X-averaged negative particle concentration [mol.m-3]", + ]: + np.testing.assert_allclose( + rust[name](t_interp), cas[name](t_interp), rtol=1e-6, atol=1e-6 + ) + + def test_native_processed_variable_matches_casadi_direct(self): + # Build a native-backed ProcessedVariable by hand (routing is Task A4) + # and assert its native leaves + sensitivities match a CasADi solve. + t_interp = np.linspace(0, 3600, 41) + cas, _ = self._solve("casadi", calculate_sensitivities=True, t_interp=t_interp) + sol, solver = self._solve( + "rust", calculate_sensitivities=True, t_interp=t_interp + ) + + name = "Terminal voltage [V]" + vars_pybamm = [m.get_processed_variable_or_event(name) for m in sol.all_models] + rust_model = solver._setup["rust_model"] + rust_fns = [sol.observation._leaf(name, vp, rust_model) for vp in vars_pybamm] + pv = pybamm.process_variable( + name, + vars_pybamm, + NativeObserver(rust_fns, sol.observation), + sol, + time_integral=None, + ) + + # raw entries (exercises _observe_raw_native) + np.testing.assert_allclose(pv.entries, cas[name].entries, rtol=1e-6, atol=1e-6) + # off-grid query (exercises _observe_hermite_native) + np.testing.assert_allclose( + pv(t_interp), cas[name](t_interp), rtol=1e-6, atol=1e-6 + ) + # sensitivities (exercises _initialise_sensitivity_native) + np.testing.assert_allclose( + pv.sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_flip_compiles_no_casadi_observe(self): + # An idaklu rust-mode full-state solve observes natively: the + # ProcessedVariable is Rust-backed with no CasADi observe function. + rust, _ = self._solve("rust") + v = rust["Terminal voltage [V]"] + _ = v.entries # force observation + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + + def test_offgrid_values_native_backed(self): + # The off-grid parity above only exercises the native path if the observed + # ProcessedVariable is Rust-backed with no CasADi funcs. + rust, _ = self._solve("rust") + for name in [ + "Terminal voltage [V]", + "X-averaged negative particle concentration [mol.m-3]", + ]: + v = rust[name] + _ = v.entries + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + + def test_sensitivities_match_casadi(self): + # 0D non-time-integral variable sensitivities via the native chain rule. + cas, _ = self._solve("casadi", calculate_sensitivities=True) + rust, _ = self._solve("rust", calculate_sensitivities=True) + for name in ["Terminal voltage [V]"]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name].sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def _solve_output_vars( + self, convert_to_format, output_variables, model=None, inputs=None, var_pts=None + ): + if model is None: + model = pybamm.lithium_ion.SPM() + geometry = model.default_geometry + param = model.default_parameter_values + if inputs is None: + inputs = {"Current function [A]": 0.68} + param.update({key: "[input]" for key in inputs}) + param.process_model(model) + param.process_geometry(geometry) + mesh = pybamm.Mesh( + geometry, model.default_submesh_types, var_pts or model.default_var_pts + ) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver(output_variables=output_variables) + # The two lowerings take different adaptive steps, so without t_interp each + # solution lands on its own grid and callers compare mismatched times. + sol = solver.solve( + model, + [0, 3600], + inputs=inputs, + calculate_sensitivities=True, + t_interp=np.linspace(0, 3600, 100), + ) + return sol, solver + + def test_output_variables_sensitivities_match_casadi(self): + # save_outputs_only path: output_variables + calculate_sensitivities + # together exercise the native yS projection consumer (B1+B2+B3). + ov = ["Terminal voltage [V]"] + cas, _ = self._solve_output_vars("casadi", ov) + rust, _ = self._solve_output_vars("rust", ov) + np.testing.assert_allclose( + rust["Terminal voltage [V]"].sensitivities["Current function [A]"], + cas["Terminal voltage [V]"].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_sensitivities_match_casadi_spme(self): + # Repeat the single-param output-variable-sensitivity parity check on + # a second chemistry, guarding against SPM-only coincidences. + ov = ["Terminal voltage [V]"] + cas, _ = self._solve_output_vars("casadi", ov, model=pybamm.lithium_ion.SPMe()) + rust, _ = self._solve_output_vars("rust", ov, model=pybamm.lithium_ion.SPMe()) + np.testing.assert_allclose( + rust["Terminal voltage [V]"].sensitivities["Current function [A]"], + cas["Terminal voltage [V]"].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_sensitivities_match_casadi_dfn(self): + # DFN is a much heavier solve, so a coarse mesh and one output variable bound + # the cost while still exercising the native yS projection. + ov = ["Terminal voltage [V]"] + var_pts = {"x_n": 10, "x_s": 10, "x_p": 10, "r_n": 5, "r_p": 5} + cas, _ = self._solve_output_vars( + "casadi", ov, model=pybamm.lithium_ion.DFN(), var_pts=var_pts + ) + rust, _ = self._solve_output_vars( + "rust", ov, model=pybamm.lithium_ion.DFN(), var_pts=var_pts + ) + np.testing.assert_allclose( + rust["Terminal voltage [V]"].sensitivities["Current function [A]"], + cas["Terminal voltage [V]"].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_multi_param_sensitivities_match_casadi(self): + # The two sensitivity params reach solve() reversed from the sorted order that + # sets yS columns, so keying off insertion order would mislabel one. + ov = ["Terminal voltage [V]"] + inputs = { + "Negative electrode active material volume fraction": 0.6, + "Current function [A]": 0.68, + } + cas, _ = self._solve_output_vars("casadi", ov, inputs=inputs) + rust, _ = self._solve_output_vars("rust", ov, inputs=inputs) + for name in inputs: + np.testing.assert_allclose( + rust["Terminal voltage [V]"].sensitivities[name], + cas["Terminal voltage [V]"].sensitivities[name], + rtol=1e-5, + atol=1e-6, + err_msg=f"sensitivity mismatch for param '{name}'", + ) + + def test_multi_output_multi_param_sensitivities_match_casadi(self): + # Two outputs and two parameters: yS is scattered as (output, param), and + # a transposed write agrees on the diagonal, so it needs both to be > 1. + ov = ["Terminal voltage [V]", "Discharge capacity [A.h]"] + inputs = { + "Negative electrode active material volume fraction": 0.6, + "Current function [A]": 0.68, + } + cas, _ = self._solve_output_vars("casadi", ov, inputs=inputs) + rust, _ = self._solve_output_vars("rust", ov, inputs=inputs) + for var in ov: + for name in inputs: + np.testing.assert_allclose( + rust[var].sensitivities[name], + cas[var].sensitivities[name], + rtol=1e-5, + atol=1e-6, + err_msg=f"sensitivity mismatch for '{var}', parameter '{name}'", + ) + + def test_output_variables_time_integral_sensitivities_match_casadi(self): + # "Discharge capacity [A.h]" is a plain ODE state under the default + # discretisation, so this exercises the sensitivity chain off voltage. + ov = ["Discharge capacity [A.h]"] + cas, _ = self._solve_output_vars("casadi", ov) + rust, _ = self._solve_output_vars("rust", ov) + np.testing.assert_allclose( + rust["Discharge capacity [A.h]"].sensitivities["Current function [A]"], + cas["Discharge capacity [A.h]"].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_sensitivities_no_casadi_compiled(self): + # BaseSolver.set_up skips the computed_var_fcns loop for + # convert_to_format="rust", so no CasADi var/sens keys are ever compiled. + ov = ["Terminal voltage [V]"] + rust, solver = self._solve_output_vars("rust", ov) + assert solver.computed_var_fcns == {} + assert solver.computed_dvar_dy_fcns == {} + assert solver.computed_dvar_dp_fcns == {} + assert "var_fcns" not in solver._setup + assert "dvar_dy_idaklu_fcns" not in solver._setup + assert "dvar_dp_idaklu_fcns" not in solver._setup + assert rust.variables_returned is True + + @staticmethod + def _discretise_dfn(convert_to_format): + model = pybamm.lithium_ion.DFN() + geometry = model.default_geometry + param = model.default_parameter_values + param.process_model(model) + param.process_geometry(geometry) + var_pts = {"x_n": 20, "x_s": 20, "x_p": 20, "r_n": 10, "r_p": 10} + mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + model.convert_to_format = convert_to_format + return model + + def test_spatial_1d_2d_variables_match_casadi(self): + # 1D (x) and 2D (r, x) spatial variables exercise the native + # order="F" reshape/segment layout beyond the 0D cases above. + t_interp = np.linspace(0, 3600, 51) + cas = pybamm.IDAKLUSolver().solve( + self._discretise_dfn("casadi"), [0, 3600], t_interp=t_interp + ) + rust = pybamm.IDAKLUSolver().solve( + self._discretise_dfn("rust"), [0, 3600], t_interp=t_interp + ) + # off-grid query points strictly inside the solved interval + off = t_interp[:-1] + np.diff(t_interp) / 3 + for name in [ + "Electrolyte concentration [mol.m-3]", # 1D in x + "Negative particle concentration [mol.m-3]", # 2D (r, x) + ]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + + @staticmethod + def _discretise_spm(convert_to_format, as_input=False): + model = pybamm.lithium_ion.SPM() + geometry = model.default_geometry + param = model.default_parameter_values + if as_input: + param.update({"Current function [A]": "[input]"}) + param.process_model(model) + param.process_geometry(geometry) + mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + model.convert_to_format = convert_to_format + return model + + def test_multi_segment_experiment_matches_casadi(self): + # A stepped solve builds all_ys with >1 segment, exercising per-segment + # native Hermite routing and flag propagation through Solution.__add__. + def stepped(convert_to_format): + model = self._discretise_spm(convert_to_format) + solver = pybamm.IDAKLUSolver() + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=1800) + return sol + + cas = stepped("casadi") + rust = stepped("rust") + assert len(rust.all_ys) == 2 + assert isinstance(rust.observation, NativeInterpolatingObservation) + off = np.linspace(0, rust.t[-1], 73)[1:-1] + 5.0 + off = off[off < rust.t[-1]] + for name in ["Terminal voltage [V]"]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + + def test_event_termination_matches_casadi(self): + # A voltage-cutoff termination yields a partial final segment; the + # native path must still off-grid-interpolate to match CasADi. + cas = pybamm.IDAKLUSolver().solve(self._discretise_spm("casadi"), [0, 100000]) + rust = pybamm.IDAKLUSolver().solve(self._discretise_spm("rust"), [0, 100000]) + assert rust.termination.startswith("event:") + name = "Terminal voltage [V]" + assert isinstance(rust[name]._observer, NativeObserver) + t_end = min(rust.t[-1], cas.t[-1]) + off = np.linspace(0, t_end, 40)[1:-1] + 1.0 + off = off[off < t_end] + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + + def test_hermite_off_still_native_no_casadi(self): + # hermite_interpolation=False disables yps, but the solve must stay on the + # native-backed ProcessedVariable, not diffsol's ProcessedVariableComputed. + solver_kwargs = {"options": {"hermite_interpolation": False}} + cas = pybamm.IDAKLUSolver(**solver_kwargs).solve( + self._discretise_spm("casadi"), [0, 3600] + ) + rust = pybamm.IDAKLUSolver(**solver_kwargs).solve( + self._discretise_spm("rust"), [0, 3600] + ) + assert not rust.hermite_interpolation + v = rust["Terminal voltage [V]"] + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + np.testing.assert_allclose( + v.entries, cas["Terminal voltage [V]"].entries, rtol=1e-6, atol=1e-6 + ) + + def test_time_integral_sensitivities_match_casadi(self): + # An ExplicitTimeIntegral output variable exercises the native postfix + # (_native_postfix_sensitivities) sensitivity path. + inputs = {"Current function [A]": 0.68} + cas = pybamm.IDAKLUSolver().solve( + self._discretise_spm("casadi", as_input=True), + [0, 3600], + inputs=inputs, + calculate_sensitivities=True, + ) + rust = pybamm.IDAKLUSolver().solve( + self._discretise_spm("rust", as_input=True), + [0, 3600], + inputs=inputs, + calculate_sensitivities=True, + ) + name = "Discharge capacity [A.h]" + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name].sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_first_last_state_observe_natively(self): + # Outputs-only solves must attach the observation context: their first/last + # states carry full state vectors, so summary variables must match exactly. + def solve(output_variables): + model = self._discretise_spm("rust") + solver = pybamm.IDAKLUSolver(output_variables=output_variables) + return solver.solve(model, [0, 3600]) + + full = solve(None) + outputs_only = solve(["Voltage [V]"]) + assert isinstance(outputs_only.observation, NativeObservation) + assert isinstance(outputs_only.last_state.observation, NativeObservation) + name = "Total lithium in electrolyte [mol]" + np.testing.assert_array_equal( + outputs_only.last_state[name].data, full.last_state[name].data + ) + + def test_output_variables_stay_aligned_across_the_batch_window(self): + # 200 points cross the 128-point batched-evaluation window on the rust + # FFI; a flush off-by-one would shift every value after the boundary. + t_interp = np.linspace(0, 3600, 200) + names = [ + "Voltage [V]", + "Negative particle surface concentration [mol.m-3]", + ] + + def solve(output_variables): + model = self._discretise_spm("rust") + solver = pybamm.IDAKLUSolver(output_variables=output_variables) + return solver.solve(model, [0, 3600], t_interp=t_interp) + + full = solve(None) + outputs_only = solve(names) + for name in names: + np.testing.assert_allclose( + np.asarray(outputs_only[name](t_interp)), + np.asarray(full[name](t_interp)), + rtol=1e-7, + atol=1e-9, + err_msg=name, + ) + + def test_output_variables_genuine_time_integral_matches_casadi(self): + # remove_independent_variables_from_rhs=True turns "Discharge capacity [A.h]" + # into a genuine ExplicitTimeIntegral, evaluated natively then postfixed. + def solve(convert_to_format): + model = pybamm.lithium_ion.SPM() + geometry = model.default_geometry + param = model.default_parameter_values + param.update({"Current function [A]": "[input]"}) + param.process_model(model) + param.process_geometry(geometry) + mesh = pybamm.Mesh( + geometry, model.default_submesh_types, model.default_var_pts + ) + disc = pybamm.Discretisation( + mesh, + model.default_spatial_methods, + remove_independent_variables_from_rhs=True, + ) + disc.process_model(model) + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver(output_variables=["Discharge capacity [A.h]"]) + return solver.solve( + model, + [0, 3600], + inputs={"Current function [A]": 0.68}, + calculate_sensitivities=True, + ) + + cas = solve("casadi") + rust = solve("rust") + name = "Discharge capacity [A.h]" + assert rust[name].data.shape == (1,) + np.testing.assert_allclose( + rust[name].data, cas[name].data, rtol=1e-6, atol=1e-8 + ) + np.testing.assert_allclose( + rust[name].sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-6, + atol=1e-8, + ) + + +class TestIDAKLUSensitivityScales: + """IDAS ``pbar``, so a tiny parameter's sensitivity column stays solvable.""" + + @staticmethod + def _tiny_parameter_model(): + """``du/dt = -(a / a0) u`` with ``a0 = 1e-14``, so ``du/da ~ 1e14``. + + Analytically ``u = exp(-t)`` at ``a = a0`` and + ``du/da = -t exp(-t) / a0``, a column no absolute tolerance can hold at + the default ``pbar = 1``. + """ + model = pybamm.BaseModel() + u = pybamm.Variable("u") + a = pybamm.InputParameter("a") + model.rhs = {u: -(a / 1e-14) * u} + model.initial_conditions = {u: 1} + model.variables = {"u": u} + pybamm.Discretisation().process_model(model) + return model + + def test_scales_are_the_parameter_magnitudes(self): + scales = pybamm.solvers.idaklu_solver._sensitivity_scales( + {"a": -3.0, "b": 4e-15, "c": 1.0}, ["b", "a"] + ) + np.testing.assert_allclose(scales, [4e-15, 3.0]) + + def test_magnitudes_are_handed_over_unclamped(self): + # The solver owns the zero/non-finite clamp, so a raw 0.0 reaches it. + scales = pybamm.solvers.idaklu_solver._sensitivity_scales( + {"a": 0.0, "b": 2.0}, ["a", "b"] + ) + np.testing.assert_allclose(scales, [0.0, 2.0]) + + def test_a_zero_parameter_still_solves(self): + # IDAS rejects pbar = 0, so the solver has to clamp it to the unit scale. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + a = pybamm.InputParameter("a") + model.rhs = {u: -u + a} + model.initial_conditions = {u: 1} + model.variables = {"u": u} + pybamm.Discretisation().process_model(model) + + solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8) + sol = solver.solve( + model, + np.linspace(0, 1, 10), + inputs={"a": 0.0}, + calculate_sensitivities=True, + ) + # du/da = 1 - exp(-t) regardless of a, so the clamp must not skew it. + np.testing.assert_allclose( + np.asarray(sol["u"].sensitivities["a"]).ravel(), + 1.0 - np.exp(-sol.t), + rtol=1e-5, + atol=1e-7, + ) + + def test_a_vector_parameter_uses_its_largest_magnitude(self): + scales = pybamm.solvers.idaklu_solver._sensitivity_scales( + {"a": np.array([1e-3, -5e-3, 2e-3])}, ["a"] + ) + np.testing.assert_allclose(scales, [5e-3]) + + def test_scales_reach_the_solver_group(self): + model = self._tiny_parameter_model() + solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-6) + solve_kwargs = {"inputs": {"a": 1e-14}, "calculate_sensitivities": True} + # Solve once so the second solve reuses this group rather than rebuilding + # it, which would discard the spy. + solver.solve(model, [0, 1], **solve_kwargs) + + seen = {} + original = solver._setup["solver"].solve + + def spy(*args, **kwargs): + seen["pbar"] = np.asarray(args[5]) + return original(*args, **kwargs) + + solver._setup["solver"] = type("Spy", (), {"solve": staticmethod(spy)})() + solver.solve(model, [0, 1], **solve_kwargs) + np.testing.assert_allclose(seen["pbar"], [[1e-14]]) + + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_a_tight_dfn_diffusivity_sensitivity_solve_converges( + self, convert_to_format + ): + # D_p is 4e-15, so dy/dD_p reaches ~1e14; unscaled in the corrector's + # weighted norm that is an IDA_CONV_FAIL. + model = pybamm.lithium_ion.DFN() + model.convert_to_format = convert_to_format + name = "Positive particle diffusivity [m2.s-1]" + parameter_values = pybamm.ParameterValues("Chen2020") + nominal = float(parameter_values[name]) + parameter_values[name] = pybamm.InputParameter("D_p") + simulation = pybamm.Simulation( + model, + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8), + ) + sol = simulation.solve( + np.linspace(0, 600, 20), + inputs={"D_p": nominal}, + calculate_sensitivities=["D_p"], + ) + gradient = np.asarray(sol["Voltage [V]"].sensitivities["D_p"]).ravel() + assert np.all(np.isfinite(gradient)) + assert np.abs(gradient).max() > 0.0 + + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_tiny_parameter_sensitivity_matches_the_analytic_column( + self, convert_to_format + ): + # Scaling the weights must not disturb the column itself. + model = self._tiny_parameter_model() + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8) + sol = solver.solve( + model, + np.linspace(0, 3, 20), + inputs={"a": 1e-14}, + calculate_sensitivities=True, + ) + expected = -sol.t * np.exp(-sol.t) / 1e-14 + np.testing.assert_allclose( + np.asarray(sol["u"].sensitivities["a"]).ravel(), + expected, + rtol=1e-5, + atol=1e-5 / 1e-14, + ) + + +class TestIDAKLUSensitivityScaleOrdering: + """One scale per parameter, in the solver's column order rather than the + input dict's insertion order. Both orderings have the right length, so a + mix-up degrades the weighting silently instead of raising.""" + + def test_each_parameter_gets_its_own_magnitude(self): + scales = pybamm.solvers.idaklu_solver._sensitivity_scales( + {"b": 1e2, "a": 1e-14}, ["a", "b"] + ) + np.testing.assert_allclose(scales, [1e-14, 1e2]) + + +class TestIDAKLUModelAtol: + """``model.atol`` overrides the solver's own tolerance, per state as well as + uniformly. Only IDAKLU reads it.""" + + def _two_state_model(self, atol=None): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + w = pybamm.Variable("w") + model.rhs = {u: -u, w: -2 * w} + model.initial_conditions = {u: 1.0, w: 1.0} + model.variables = {"u": u, "w": w} + pybamm.Discretisation().process_model(model) + if atol is not None: + model.atol = atol + return model + + def _steps(self, model, atol): + solution = pybamm.IDAKLUSolver(rtol=1e-6, atol=atol).solve( + model, np.linspace(0, 1, 5) + ) + return solution.solver_statistics.number_of_steps + + def test_a_per_state_model_atol_wins_over_the_solvers(self): + tight = self._steps(self._two_state_model(), 1e-12) + loose = self._steps(self._two_state_model(), 1e-1) + assert self._steps(self._two_state_model(np.full(2, 1e-1)), 1e-12) == loose + assert loose < tight + + def test_a_wrong_width_model_atol_is_rejected(self): + model = self._two_state_model(np.full(3, 1e-6)) + with pytest.raises(pybamm.SolverError, match=r"shape \(3,\) but \(2,\)"): + self._steps(model, 1e-6) diff --git a/packages/pybamm/tests/unit/test_solvers/test_jax_solver.py b/packages/pybamm/tests/unit/test_solvers/test_jax_solver.py index 1291615348..555bf95ee0 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_jax_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_jax_solver.py @@ -134,8 +134,7 @@ def test_solver_only_works_with_jax(self): t_eval = np.linspace(0, 3, 100) - # solver needs a model converted to jax - for convert_to_format in ["casadi", "python", "something_else"]: + for convert_to_format in ["casadi", "python"]: model.convert_to_format = convert_to_format solver = pybamm.JaxSolver() diff --git a/packages/pybamm/tests/unit/test_solvers/test_native_observation.py b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py new file mode 100644 index 0000000000..980353bb99 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py @@ -0,0 +1,408 @@ +"""Unit tests for native (Rust) observation: accessor, producer, compile-once.""" + +import numpy as np + +from pybamm.rust import ExprGraph +from pybamm.solvers.observation import ( + NativeComputedObservation, + NativeObservation, +) +from pybamm.solvers.variable_observer import chain_rule_sensitivities + +# D_n's dust rides IDAKLU's unscaled pbar = 1, splitting the two backends' step +# sequences at default tolerances. 1e-11 converges both; the parity bounds stay. +_TWO_PARAM_SOLVER_TOL = 1e-11 + + +def _make_compiled_model(): + """Minimal 2-state model with one input (mirrors the binding-test fixture). + + rhs = [a * y1, -y0]; identity mass; one input 'a'. + """ + from pybamm.rust import CompiledModel + + g = ExprGraph() + a = g.input_parameter("a") + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + rhs = g.concat([g.mul(a, y1), g.neg(y0)]) + mass = np.array([1.0, 1.0]) + indptr = np.array([0, 1, 2], dtype=np.int64) + indices = np.array([0, 1], dtype=np.int64) + model = CompiledModel.from_expr(g, rhs, mass, indptr, indices, n_inputs=1) + return g, model + + +class TestGraphAccessor: + def test_graph_returns_exprgraph(self): + _, model = _make_compiled_model() + assert isinstance(model.graph, ExprGraph) + + def test_graph_compiles_a_new_observation_root(self): + # Lower a NEW root (2*y0) into the retained arena and compile it. + # The new root has no inputs, so p is the empty stacked array. + _, model = _make_compiled_model() + graph = model.graph + new_root = graph.mul(graph.scalar(2.0), graph.state_vector(0, 1)) + fn = graph.compile(new_root, name="obs", n_states=model.n_states) + out = fn(0.0, np.array([3.0, 5.0]), np.array([1.0])) + np.testing.assert_allclose(out, [6.0]) + + +class TestObservationContext: + def test_set_and_propagate(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" # no native observation to start from + sim = pybamm.Simulation(model) + sol = sim.solve([0, 600]) + assert not isinstance(sol.observation, NativeObservation) + + _, compiled = _make_compiled_model() + sol.observation = NativeComputedObservation.uniform(compiled, len(sol.all_ys)) + assert sol.observation.primary_model is compiled + assert sol.observation.compile_cache == {} + + # Context survives copy and slicing + assert sol.copy().observation.primary_model is compiled + assert sol.first_state.observation.primary_model is compiled + assert sol.last_state.observation.primary_model is compiled + + +def _native_idaklu_sim(model, **solver_kwargs): + """Simulation on idaklu with native observation forced on (test-only). + + The public switch is `convert_to_format == "rust"`; idaklu's + `_observes_via_compiled_model` stays False until prereqs A + B land, so + foundation tests set it on the solver instance. The instance attribute + survives the shallow `solver.copy()` that `Simulation` performs. + """ + import pybamm + + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver(**solver_kwargs) + solver._observes_via_compiled_model = True + return pybamm.Simulation(model, solver=solver) + + +class TestSingleVariableNative: + def test_native_matches_casadi_0d(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 3600]) + v_native = sol["Terminal voltage [V]"].entries + + casadi_model = pybamm.lithium_ion.SPM() + casadi_model.convert_to_format = "casadi" + sim_c = pybamm.Simulation(casadi_model, solver=pybamm.IDAKLUSolver()) + sol_c = sim_c.solve([0, 3600]) + v_casadi = sol_c["Terminal voltage [V]"].entries + + np.testing.assert_allclose(v_native, v_casadi, rtol=1e-10, atol=1e-12) + + def test_hermite_native_matches_casadi_multi_tile(self): + import pybamm + + # A dense 137-point grid off the solver's knots: the lane-batched Hermite + # evaluator spans several tiles plus a ragged tail, and hits s=0/1 and 0 1 exercises the (output_len, n_t) -> time-outer/output-inner + # flatten against the CasADi block-diagonal layout. + name = "Negative particle concentration [mol.m-3]" + var_n, var_c = self._solve_both(name, ["I"]) + # Assert the variable is genuinely spatial (multi-point); if it resolves to + # 0D, the flatten logic under test would not be exercised. + assert var_c.entries.shape[0] > 1, ( + f"{name!r} resolved to 0D (shape {var_c.entries.shape}); " + "test premise violated" + ) + sens_n = var_n.sensitivities + sens_c = var_c.sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose(sens_n[key], sens_c[key], rtol=1e-5, atol=1e-8) + + def test_sensitivities_all_block_two_params(self): + # "I" plus a diffusivity input; sorted column order is ["D_n", "I"]. + var_n, var_c = self._solve_both( + "Terminal voltage [V]", + ["D_n", "I"], + extra_inputs={"Negative particle diffusivity [m2.s-1]": ("D_n", 3.3e-14)}, + solver_tol=_TWO_PARAM_SOLVER_TOL, + ) + a_n = var_n.sensitivities["all"] + a_c = var_c.sensitivities["all"] + assert a_n.shape == a_c.shape # (n_t*output_len, 2) + assert a_n.shape[1] == 2 # block must have exactly 2 parameter columns + assert np.any(var_n.sensitivities["D_n"] != 0) + # Sorted order is ["D_n", "I"]; pin column 0 absolutely to D_n sensitivity. + np.testing.assert_allclose( + a_n[:, 0], var_n.sensitivities["D_n"], rtol=1e-5, atol=1e-8 + ) + np.testing.assert_allclose(a_n, a_c, rtol=1e-5, atol=1e-8) + + def test_no_sensitivities_when_not_requested(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 600]) + assert sol["Terminal voltage [V]"].sensitivities == {} + + +class TestObservationCacheAcrossSolves: + # The compiled-observation cache lives with the solver setup, 1:1 with the rust + # model, so repeated solves reuse tapes instead of growing the retained graph. + + @staticmethod + def _sim(): + # events removed so termination is "final time" and solve() returns the + # post-processed Solution directly (which shares the solver's cache). + import pybamm + + model = pybamm.lithium_ion.SPM() + model.events = [] + return _native_idaklu_sim(model) + + def test_compiled_leaf_reused_across_solves(self): + sim = self._sim() + sol1 = sim.solve([0, 600]) + v1 = sol1["Terminal voltage [V]"] + sol2 = sim.solve([0, 600]) + v2 = sol2["Terminal voltage [V]"] + + # Same solver-owned cache dict, shared by reference across solves. + assert sol2.observation.compile_cache is sol1.observation.compile_cache + assert ( + sol1.observation.compile_cache + is sim._solver._setup["rust_observation_cache"] + ) + # The compiled leaf is reused, not recompiled, on the second solve. + assert v2._observer.leaves[0] is v1._observer.leaves[0] + + def test_retained_graph_stops_growing(self): + sim = self._sim() + sol1 = sim.solve([0, 600]) + _ = sol1["Terminal voltage [V]"] + cache = sim._solver._setup["rust_observation_cache"] + nodes_after_first = sim._solver._setup["rust_model"].graph.n_nodes + cache_len_after_first = len(cache) + + sol2 = sim.solve([0, 600]) + _ = sol2["Terminal voltage [V]"] + # No recompile on the second observe: neither the retained graph's node + # count nor the cache size grows. + assert sim._solver._setup["rust_model"].graph.n_nodes == nodes_after_first + assert len(cache) == cache_len_after_first + + def test_discrete_time_sum_across_two_solves(self): + import pybamm + + # DiscreteTimeSum output var exercises the time-integral memo path. + # The 0D model auto-discretises, so the solver is driven directly. + data_times = np.linspace(0, 1, 10) + data = pybamm.DiscreteTimeData(data_times, np.zeros_like(data_times), "zeros") + + def _build_model(fmt): + m = pybamm.BaseModel(name="dts_model") + c = pybamm.Variable("c") + m.rhs = {c: -2 * c} + m.initial_conditions = {c: 1} + m.variables["c"] = c + # (c - 0)^2 summed over the data times -> sum(exp(-4 t)) + m.variables["dts"] = pybamm.DiscreteTimeSum((c - data) ** 2) + m.convert_to_format = fmt + return m + + model = _build_model("rust") + solver = pybamm.IDAKLUSolver() + solver._observes_via_compiled_model = True + sol1 = solver.solve(model, t_eval=[0, 1], t_interp=data_times) + val1 = sol1["dts"]() + sol2 = solver.solve(model, t_eval=[0, 1], t_interp=data_times) + val2 = sol2["dts"]() + + # Stable across solves (the memoised time-integral analysis is reused). + np.testing.assert_allclose(val2, val1, rtol=1e-12, atol=0) + # Correct against a CasADi reference. + model_c = _build_model("casadi") + sol_c = pybamm.IDAKLUSolver().solve(model_c, t_eval=[0, 1], t_interp=data_times) + np.testing.assert_allclose(val1, sol_c["dts"](), rtol=1e-6, atol=1e-8) + + # The memo stored the (model, name, nstates) analysis, reused on solve 2. + assert sol2.observation.compile_cache is sol1.observation.compile_cache + ti_keys = [ + k + for k in sol1.observation.compile_cache + if isinstance(k, tuple) and k[:2] == ("__time_integral__", "dts") + ] + assert len(ti_keys) == 1 + + def test_pickle_round_trip_after_observe(self): + import pickle + + sim = self._sim() + sol = sim.solve([0, 600]) + expected = sol["Terminal voltage [V]"].entries # populate the cache + + restored = pickle.loads(pickle.dumps(sol)) + np.testing.assert_allclose( + restored["Terminal voltage [V]"].entries, expected, rtol=1e-10, atol=1e-12 + ) + + +class TestDiffsolObservationCache: + def test_compiled_fn_reused_across_solves(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + model.events = [] + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + + sol1 = sim.solve(np.linspace(0, 600, 50)) + _ = sol1["Terminal voltage [V]"] + cache = sim._solver._rust_observation_cache + key = ("Terminal voltage [V]", id(sim._solver._rust_model)) + fn1 = cache[key] + nodes_after_first = sim._solver._rust_model.graph.n_nodes + + sol2 = sim.solve(np.linspace(0, 600, 50)) + _ = sol2["Terminal voltage [V]"] + assert sol2.observation.compile_cache is cache + assert cache[key] is fn1 # reused, not recompiled + assert sim._solver._rust_model.graph.n_nodes == nodes_after_first diff --git a/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py index 0f99c10113..0ae6a401e5 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver.py @@ -196,6 +196,7 @@ def test_matches_casadi_algebraic_solver(self): model2 = _ElectrodeSOH(param=li_param) parameter_values.process_model(model2) + model2.convert_to_format = "casadi" sol_casadi = pybamm.CasadiAlgebraicSolver(tol=1e-10).solve( model2, [0], diff --git a/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver_rust.py b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver_rust.py new file mode 100644 index 0000000000..7c54c81016 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_nonlinear_solver_rust.py @@ -0,0 +1,101 @@ +"""Tests for the Rust-backed ``pybamm.NonlinearSolver`` root solve. + +Exercises the dlsym Rust ``StandaloneNewtonSolver`` end-to-end through +``_set_up_root_solver_rust`` (``model.convert_to_format == "rust"``), the only +path that reaches the Rust Newton C++ constructor. +""" + +import numpy as np +import pytest + +import pybamm + + +class TestNonlinearSolverRust: + def _solve_rust(self, model, t_eval, inputs=None, use_sparse=False): + pybamm.Discretisation().process_model(model) + model.convert_to_format = "rust" + return pybamm.NonlinearSolver(use_sparse=use_sparse).solve( + model, t_eval, inputs=inputs + ) + + def test_simple_root_find_rust(self): + var = pybamm.Variable("var") + model = pybamm.BaseModel() + model.algebraic = {var: var + 2} + model.initial_conditions = {var: 2} + + solution = self._solve_rust(model, np.linspace(0, 1, 10)) + np.testing.assert_allclose(solution.y, -2, atol=1e-8) + + def test_solve_with_input_rust(self): + var = pybamm.Variable("var") + model = pybamm.BaseModel() + model.algebraic = {var: var + pybamm.InputParameter("param")} + model.initial_conditions = {var: 2} + + solution = self._solve_rust(model, np.linspace(0, 1, 10), inputs={"param": 7}) + np.testing.assert_allclose(solution.y, -7, atol=1e-8) + + def test_sparse_matches_dense_rust(self): + def build(): + model = pybamm.BaseModel() + var = pybamm.Variable("var") + model.algebraic = {var: var - 5} + model.initial_conditions = {var: 1} + return model + + t_eval = np.linspace(0, 1, 5) + sol_sparse = self._solve_rust(build(), t_eval, use_sparse=True) + sol_dense = self._solve_rust(build(), t_eval, use_sparse=False) + np.testing.assert_allclose(sol_sparse.y, sol_dense.y, atol=1e-10) + + def test_compile_option_is_rejected(self): + var = pybamm.Variable("var") + model = pybamm.BaseModel() + model.algebraic = {var: var + 2} + model.initial_conditions = {var: 2} + pybamm.Discretisation().process_model(model) + model.convert_to_format = "rust" + + solver = pybamm.NonlinearSolver(options={"compile": True}) + with pytest.raises(pybamm.SolverError, match=r"CasADi-only"): + solver.solve(model, np.linspace(0, 1, 5)) + + def test_sensitivity_extended_states_are_rejected(self): + var = pybamm.Variable("var") + model = pybamm.BaseModel() + model.algebraic = {var: var + pybamm.InputParameter("param")} + model.initial_conditions = {var: 2} + pybamm.Discretisation().process_model(model) + model.convert_to_format = "rust" + # A y0 longer than len_rhs_and_alg is how a sensitivity-extended + # state block reaches the root solver. + model.y0_list = [np.zeros(model.len_rhs_and_alg + 1)] + + solver = pybamm.NonlinearSolver() + with pytest.raises(pybamm.SolverError, match=r"sensitivity-extended"): + solver._set_up_root_solver_rust(model, {"param": 7.0}) + + def test_matches_casadi_newton(self): + """The Rust and CasADi Newton drivers must agree on the same system.""" + + def build(): + model = pybamm.BaseModel() + var1 = pybamm.Variable("var1") + var2 = pybamm.Variable("var2") + model.algebraic = {var1: var1 - 3, var2: 2 * var1 - var2} + model.initial_conditions = { + var1: pybamm.Scalar(1), + var2: pybamm.Scalar(4), + } + return model + + t_eval = np.linspace(0, 1, 5) + sol_rust = self._solve_rust(build(), t_eval) + + model_casadi = build() + pybamm.Discretisation().process_model(model_casadi) + sol_casadi = pybamm.NonlinearSolver().solve(model_casadi, t_eval) + + np.testing.assert_allclose(sol_rust.y, sol_casadi.y, rtol=1e-7, atol=1e-7) diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py new file mode 100644 index 0000000000..b54a259e49 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py @@ -0,0 +1,292 @@ +"""Contract tests for the observation seam: backends and observers standalone.""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import ( + CASADI_OBSERVATION, + CasadiObservation, + NativeComputedObservation, + NativeInterpolatingObservation, + NativeObservation, + join_observations, +) +from pybamm.solvers.variable_observer import ( + CasadiObserver, + NativeObserver, + SegmentSelector, + as_observer, + pack_sensitivity_dict, +) + + +class _Model: + """Stand-in for a CompiledModel: only identity matters to the backend.""" + + def __init__(self, tag): + self.tag = tag + + +class TestSegmentSlicing: + def test_the_casadi_backend_is_a_shared_stateless_value(self): + # It reaches every per-segment artifact through the Solution it is + # handed, so slicing a segment run cannot change it. + assert isinstance(CASADI_OBSERVATION, CasadiObservation) + assert CASADI_OBSERVATION[:1] is CASADI_OBSERVATION + assert CASADI_OBSERVATION[-1:] is CASADI_OBSERVATION + + def test_native_slices_its_per_segment_models(self): + models = [_Model(i) for i in range(3)] + cache = {"tape": object()} + backend = NativeInterpolatingObservation(models, cache=cache) + + assert backend.n_segments == 3 + assert backend.primary_model is models[0] + assert backend[:1].segment_models == models[:1] + assert backend[-1:].segment_models == models[-1:] + # a slice keeps the concrete backend and shares the cache by identity + assert isinstance(backend[1:], NativeInterpolatingObservation) + assert backend[1:].compile_cache is cache + assert backend[-1:].primary_model is models[-1] + + def test_uniform_gives_every_segment_the_same_model(self): + model = _Model("a") + backend = NativeComputedObservation.uniform(model, 3) + assert backend.segment_models == [model] * 3 + assert isinstance(backend, NativeComputedObservation) + + def test_the_native_base_cannot_be_instantiated(self): + # Which concrete backend a solver builds is what picks the kind of + # processed variable its solutions hand back, so there is no default. + with pytest.raises(TypeError, match="abstract"): + NativeObservation([_Model("a")]) + + +class TestJoin: + def test_casadi_only_runs_join_to_the_shared_value(self): + joined = join_observations([(CASADI_OBSERVATION, 2), (CASADI_OBSERVATION, 3)]) + assert joined is CASADI_OBSERVATION + + def test_native_runs_concatenate_their_models(self): + left = NativeComputedObservation([_Model("a")]) + right = NativeComputedObservation([_Model("b"), _Model("c")]) + joined = join_observations([(left, 1), (right, 2)]) + + assert joined.segment_models == left.segment_models + right.segment_models + assert isinstance(joined, NativeComputedObservation) + + def test_joining_never_narrows_how_a_solution_can_be_read(self): + # One interpolating run is enough: the join must not demote the whole + # span to grid-aligned reads. + interpolating = NativeInterpolatingObservation([_Model("a")]) + computed = NativeComputedObservation([_Model("b")]) + + for runs in ( + [(interpolating, 1), (computed, 1)], + [(computed, 1), (interpolating, 1)], + ): + assert isinstance(join_observations(runs), NativeInterpolatingObservation) + + def test_many_runs_join_in_one_pass(self): + runs = [(NativeComputedObservation([_Model(i)]), 1) for i in range(50)] + joined = join_observations(runs) + assert joined.n_segments == 50 + assert [m.tag for m in joined.segment_models] == list(range(50)) + + def test_shared_cache_identity_survives_a_join(self): + cache = {} + runs = [ + (NativeComputedObservation([_Model(i)], cache=cache), 1) for i in range(3) + ] + assert join_observations(runs).compile_cache is cache + + def test_distinct_caches_merge_with_the_earlier_run_winning(self): + left = NativeComputedObservation([_Model("a")], cache={"k": "left"}) + right = NativeComputedObservation( + [_Model("b")], cache={"k": "right", "extra": 1} + ) + merged = join_observations([(left, 1), (right, 1)]).compile_cache + assert merged == {"k": "left", "extra": 1} + + def test_a_casadi_run_is_observed_by_the_first_native_model(self): + native = NativeInterpolatingObservation([_Model("a")]) + joined = join_observations([(CASADI_OBSERVATION, 2), (native, 1)]) + + # Every segment gets a concrete model, so there is no fallback rule + assert joined.segment_models == [native.primary_model] * 3 + assert isinstance(joined, NativeInterpolatingObservation) + + # ... and in the other order + joined = join_observations([(native, 1), (CASADI_OBSERVATION, 1)]) + assert joined.segment_models == [native.primary_model] * 2 + + +class _Variable: + """The slice of ProcessedVariable an observer is allowed to read.""" + + def __init__(self, all_ts, all_ys, all_yps=None): + self.all_ts = all_ts + self.all_ys = all_ys + self.all_yps = all_yps + self.all_inputs = [{} for _ in all_ts] + self.t_pts = np.concatenate(all_ts) + self.time_integral = None + + @property + def hermite_interpolation(self): + return self.all_yps is not None + + def _shape(self, t): + return [len(t)] + + +class _Leaf: + """Stand-in for a compiled tape: records its calls, returns 2*y0.""" + + def __init__(self): + self.trajectory_calls = [] + + def eval_trajectory(self, ts, ys, inputs): + self.trajectory_calls.append((ts, ys)) + return np.asfortranarray(2.0 * np.asarray(ys)[:1, :]) + + +class TestNativeObserverStandalone: + def test_observe_raw_evaluates_every_segment(self): + leaves = [_Leaf(), _Leaf()] + observer = NativeObserver(leaves, backend=None) + variable = _Variable( + [np.array([0.0, 1.0]), np.array([2.0])], + [np.array([[1.0, 2.0]]), np.array([[3.0]])], + ) + + np.testing.assert_allclose(observer.observe_raw(variable), [2.0, 4.0, 6.0]) + assert observer.leaves is leaves + + def test_outputs_only_solves_evaluate_on_shaped_zeros(self): + leaf = _Leaf() + observer = NativeObserver([leaf], backend=None, placeholder_states=[3]) + variable = _Variable([np.array([0.0, 1.0])], [np.zeros((0, 0))]) + + np.testing.assert_allclose(observer.observe_raw(variable), [0.0, 0.0]) + _, ys = leaf.trajectory_calls[0] + assert ys.shape == (3, 2) + + def test_empty_segments_are_skipped(self): + leaves = [_Leaf(), _Leaf()] + observer = NativeObserver(leaves, backend=None) + variable = _Variable( + [np.array([]), np.array([0.0, 1.0])], + [np.zeros((1, 0)), np.array([[1.0, 2.0]])], + ) + + np.testing.assert_allclose(observer.observe_raw(variable), [2.0, 4.0]) + assert leaves[0].trajectory_calls == [] + + def test_segment_selection_is_computed_once_per_observer(self): + observer = NativeObserver([_Leaf()], backend=None) + variable = _Variable([np.array([0.0, 1.0])], [np.array([[1.0, 2.0]])]) + + observer.observe_raw(variable) + selector = observer._selector + observer.observe_raw(variable) + assert observer._selector is selector + + +class TestSegmentSelector: + def test_full_range_keeps_every_nonempty_segment(self): + selector = SegmentSelector( + [np.array([]), np.array([0.0, 1.0]), np.array([2.0])] + ) + np.testing.assert_array_equal( + selector.select(np.array([0.0]), full_range=True), [1, 2] + ) + + def test_restricted_range_keeps_only_covering_segments(self): + selector = SegmentSelector([np.array([0.0, 1.0]), np.array([2.0, 3.0])]) + np.testing.assert_array_equal( + selector.select(np.array([2.5]), full_range=False), [1] + ) + + def test_extrapolating_past_the_end_keeps_the_last_segment(self): + selector = SegmentSelector([np.array([0.0, 1.0]), np.array([2.0, 3.0])]) + np.testing.assert_array_equal( + selector.select(np.array([4.0]), full_range=False), [1] + ) + + +class TestObserverCoercion: + def test_a_bare_casadi_list_becomes_a_casadi_observer(self): + observer = as_observer([None, None]) + assert isinstance(observer, CasadiObserver) + assert observer.leaves == [None, None] + + def test_an_observer_passes_through(self): + observer = NativeObserver([_Leaf()], backend=None) + assert as_observer(observer) is observer + + def test_process_variable_accepts_either_form(self): + # A bare list is only ever CasADi leaves, so drive the CasADi backend. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + solution = pybamm.Simulation(model).solve([0, 600]) + name = "Terminal voltage [V]" + base = [m.get_processed_variable_or_event(name) for m in solution.all_models] + leaves = solution[name]._observer.leaves + + direct = pybamm.process_variable(name, base, leaves, solution) + wrapped = pybamm.process_variable(name, base, as_observer(leaves), solution) + np.testing.assert_allclose(direct.entries, wrapped.entries) + + def test_casadi_leaves_serialise_once_across_calls(self): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + solution = pybamm.Simulation(model).solve([0, 600]) + variable = solution["Terminal voltage [V]"] + + variable.entries + serialised = dict(variable._observer._serialised) + assert serialised + variable(np.linspace(0, 600, 11)) + # the immutable CasADi function is not re-serialised per observe call + assert variable._observer._serialised == serialised + + +class TestPackSensitivityDict: + def test_all_block_plus_one_flat_vector_per_parameter(self): + S_var = np.arange(6.0).reshape(3, 2) + packed = pack_sensitivity_dict(S_var, ["a", "b"]) + + assert set(packed) == {"all", "a", "b"} + np.testing.assert_array_equal(packed["all"], S_var) + np.testing.assert_array_equal(packed["a"], [0.0, 2.0, 4.0]) + np.testing.assert_array_equal(packed["b"], [1.0, 3.0, 5.0]) + + +class TestSolutionCarriesOneField: + @staticmethod + def _solution(t_eval): + return pybamm.Simulation(pybamm.lithium_ion.SPM()).solve(t_eval) + + def test_backend_travels_with_derived_solutions(self): + solution = self._solution([0, 600]) + backend = solution.observation + assert backend.n_segments == len(solution.all_ys) + + for derived in (solution.copy(), solution.first_state, solution.last_state): + assert type(derived.observation) is type(backend) + assert derived.observation.n_segments == len(derived.all_ys) + assert derived.observation.compile_cache is backend.compile_cache + + def test_added_solutions_cover_every_segment(self): + first = self._solution([0, 600]) + second = self._solution([600, 1200]) + joined = first + second + + assert joined.observation.n_segments == len(joined.all_ys) + np.testing.assert_allclose( + joined["Terminal voltage [V]"](300.0), + first["Terminal voltage [V]"](300.0), + rtol=1e-10, + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py new file mode 100644 index 0000000000..c3e9feb2b6 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py @@ -0,0 +1,159 @@ +"""Layout-contract tests for native observation (pure Python, no Rust producer).""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import OutputAssembly + + +def _solve_spm(): + model = pybamm.lithium_ion.SPM() + sim = pybamm.Simulation(model) + return sim.solve([0, 3600]) + + +class TestUnrollNnzDense: + def test_unroll_nnz_returns_dense_when_no_casadi(self): + # A ProcessedVariableComputed built with base_variables_casadi=[None], the + # native contract, must not inspect CasADi sparsity. + sol = _solve_spm() + var = pybamm.StateVector(slice(0, 1)) + n_t = len(sol.t) + data = [np.arange(n_t, dtype=float).reshape(n_t, 1)] + pvc = pybamm.ProcessedVariableComputed([var], [None], data, sol) + + # Directly test _unroll_nnz - it should return data as-is when no CasADi func + result = pvc._unroll_nnz(data) + np.testing.assert_array_equal(result[0], data[0]) + + +class TestLayoutContract: + def test_0d_time_major(self): + # Contract: base_variables_data is one (n_times, output_size) array per + # sub-solution; for 0D output_size == 1 and data[t, 0] is the value at t. + sol = _solve_spm() + base = [pybamm.StateVector(slice(0, 1))] + n_t = len(sol.t) + values = np.linspace(3.0, 4.2, n_t) + data = [values.reshape(n_t, 1)] # time-major, C-contiguous + pvc = pybamm.ProcessedVariableComputed(base, [None], data, sol) + np.testing.assert_allclose(pvc.entries.reshape(-1), values) + + def test_1d_time_major_unrolls_to_space_by_time(self): + # A 1D variable's (n_times, len_space) array must unroll to + # (len_space, n_times) via reshape((n_times, len_space)).transpose(). + sol = _solve_spm() + base_pv = sol["X-averaged negative particle concentration [mol.m-3]"] + var = base_pv.base_variables[0] + len_space = var.shape[0] # 20 radial nodes in the negative particle + n_t = len(sol.t) + rng = np.arange(n_t * len_space, dtype=float).reshape(n_t, len_space) + pvc = pybamm.ProcessedVariableComputed([var], [None], [rng], sol) + # entries is (len_space, n_t): element [k, t] == rng[t, k] + assert pvc.entries.shape[0] == len_space + np.testing.assert_allclose(pvc.entries[:, 0], rng[0, :]) + np.testing.assert_allclose(pvc.entries[:, -1], rng[-1, :]) + + +class TestOutputAssembly: + """The outputs-only payload layout, shared by every solver that produces one. + + A scalar, a 20-component vector and a second scalar: an ordinal-indexed + reader returns one component for the vector and shifts the scalar after it. + """ + + _NAMES = [ + "Voltage [V]", + "X-averaged negative particle concentration [mol.m-3]", + "Current [A]", + ] + + @staticmethod + def _fixture(names=None): + """An SPM solution, an assembly over ``names``, and a row-indexed payload.""" + names = names or TestOutputAssembly._NAMES + solution = _solve_spm() + model = solution.all_models[0] + lens = [ + int(np.prod(model.get_processed_variable_or_event(name).shape)) + for name in names + ] + assembly = OutputAssembly(names, lens) + # data[t, k] == k, so a variable's entries are its own row indices. + data = np.tile(np.arange(assembly.n_rows, dtype=float), (len(solution.t), 1)) + return assembly, solution, data + + def test_rows_are_sliced_by_component_count_not_ordinal(self): + assembly, solution, data = self._fixture() + assembly.attach(solution, data) + + np.testing.assert_allclose(solution["Voltage [V]"].entries, 0.0) + np.testing.assert_allclose( + solution[self._NAMES[1]].entries[:, 0], np.arange(1.0, 21.0) + ) + # 21, not 2: the vector consumed rows 1--20 rather than row 1 alone. + np.testing.assert_allclose(solution["Current [A]"].entries, 21.0) + + def test_a_payload_of_the_wrong_width_is_rejected(self): + assembly, solution, data = self._fixture() + with pytest.raises(pybamm.SolverError, match=r"Output row count mismatch"): + assembly.attach(solution, data[:, :-1]) + + def test_sensitivities_are_named_and_flattened_per_parameter(self): + assembly, solution, data = self._fixture(["Voltage [V]"]) + n_t = len(solution.t) + sensitivities = np.arange(n_t * 2, dtype=float).reshape(n_t, 1, 2) + + assembly.attach( + solution, + data, + sensitivities=sensitivities, + sensitivity_names=["a", "b"], + ) + + # Read the field, not the property: the property short-circuits to {} on + # this input-free SPM solve. + attached = solution["Voltage [V]"]._sensitivities + assert attached["all"].shape == (n_t, 2) + np.testing.assert_allclose(attached["a"], sensitivities[:, 0, 0]) + np.testing.assert_allclose(attached["b"], sensitivities[:, 0, 1]) + + def test_sensitivities_of_the_wrong_shape_are_rejected(self): + assembly, solution, data = self._fixture(["Voltage [V]"]) + n_t = len(solution.t) + with pytest.raises( + pybamm.SolverError, match=r"Output sensitivity shape mismatch" + ): + assembly.attach( + solution, + data, + sensitivities=np.zeros((n_t, 1, 1)), + sensitivity_names=["a", "b"], + ) + + def test_one_block_per_parameter_is_required(self): + assembly, solution, _ = self._fixture(["Voltage [V]"]) + n_t = len(solution.t) + with pytest.raises( + pybamm.SolverError, match=r"Sensitivity block count mismatch" + ): + assembly.stack_parameter_blocks([np.zeros(n_t)], n_t, ["a", "b"]) + + def test_parameter_blocks_stack_into_the_attachable_layout(self): + assembly, solution, _ = self._fixture(["Voltage [V]"]) + n_t = len(solution.t) + blocks = [np.arange(n_t, dtype=float), np.arange(n_t, dtype=float) * -1.0] + + stacked = assembly.stack_parameter_blocks(blocks, n_t, ["a", "b"]) + + assert stacked.shape == (n_t, 1, 2) + np.testing.assert_allclose(stacked[:, 0, 0], blocks[0]) + np.testing.assert_allclose(stacked[:, 0, 1], blocks[1]) + + def test_a_solve_without_sensitivities_leaves_an_empty_mapping(self): + # Not None: an outputs-only solve retains no state to compute them from, + # so the answer is "there are none", not "ask again later". + assembly, solution, data = self._fixture(["Voltage [V]"]) + assembly.attach(solution, data) + assert solution["Voltage [V]"]._sensitivities == {} diff --git a/packages/pybamm/tests/unit/test_solvers/test_process_rust.py b/packages/pybamm/tests/unit/test_solvers/test_process_rust.py new file mode 100644 index 0000000000..58712b1eaf --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_process_rust.py @@ -0,0 +1,777 @@ +import numpy as np +import pytest +import scipy.sparse + +import pybamm +from tests import get_discretisation_for_testing + +casadi = pytest.importorskip("casadi") + + +class TestCompiledFunctionADContract: + def _cf(self, n_states=2): + # f(t, y, a) = [y0*a - y1, y0 + y1*y1] over a 2-state system, 1 input. + from pybamm.rust import ExprGraph + + graph = ExprGraph() + a = graph.input_parameter("a") + y0 = graph.state_vector(0, 1) + y1 = graph.state_vector(1, 2) + expr = graph.concat([y0 * a - y1, y0 + y1 * y1]) + return graph.compile(expr, name="rhs", n_states=n_states) + + def test_primal_eval(self): + cf = self._cf() + out = cf(0.0, np.array([1.5, -2.0]), np.array([3.0])) + np.testing.assert_allclose(np.asarray(out).ravel(), [6.5, 5.5]) + + def test_p_accepts_dict_or_array(self): + cf = self._cf() + y = np.array([1.5, -2.0]) + np.testing.assert_allclose( + np.asarray(cf(0.0, y, np.array([3.0]))).ravel(), + np.asarray(cf(0.0, y, {"a": 3.0})).ravel(), + ) + + def test_jacobian_y_is_scipy_csc(self): + cf = self._cf() + jac = cf.jacobian("y")(0.0, np.array([1.5, -2.0]), np.array([3.0])) + assert scipy.sparse.issparse(jac) + np.testing.assert_allclose(jac.toarray(), [[3.0, -1.0], [1.0, -4.0]]) + + def test_jvp_matches_jac_matvec(self): + cf = self._cf() + y, p, v = np.array([1.5, -2.0]), np.array([3.0]), np.array([0.3, 0.7]) + jv = np.asarray(cf.jvp(0.0, y, p, v)).ravel() + full = cf.jacobian("y")(0.0, y, p).toarray() @ v + np.testing.assert_allclose(jv, full) + + def test_jacobian_p_columns(self): + cf = self._cf() + # df/da = [y0 = 1.5, 0.0] + jp = cf.jacobian("p")(0.0, np.array([1.5, -2.0]), np.array([3.0])) + np.testing.assert_allclose(jp.toarray()[:, 0], [1.5, 0.0]) + + def test_rectangular_group_jacobian(self): + # output_len (1) != n_states (2): the rhs sub-block of a DAE. + from pybamm.rust import ExprGraph + + graph = ExprGraph() + y0 = graph.state_vector(0, 1) + y1 = graph.state_vector(1, 2) + cf = graph.compile(graph.concat([y0 * y1]), name="RHS", n_states=2) + jac = cf.jacobian("y")(0.0, np.array([1.5, -2.0]), np.empty(0)) + assert jac.shape == (1, 2) + np.testing.assert_allclose(jac.toarray(), [[-2.0, 1.5]]) + + +class TestTrajectoryFFILooseDtypes: + """The 1-D time params on the trajectory entry points coerce loose dtypes + (lists, float32, non-contiguous slices) rather than rejecting anything that + is not a strict contiguous float64 ndarray.""" + + def _cf(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + expr = (2 * pybamm.StateVector(slice(0, 1))).to_rust(graph, {}) + return graph.compile(expr, name="f", n_states=1) + + def _group(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + expr = (2 * pybamm.StateVector(slice(0, 1))).to_rust(graph, {}) + return graph.compile_group({"double": expr}) + + def test_eval_trajectory_hermite_accepts_loose_t_query(self): + cf = self._cf() + ts = np.array([0.0, 1.0]) + ys = np.array([[1.0, 2.0]]) + yps = np.array([[1.0, 1.0]]) + query = [0.25, 0.75] + expected = np.asarray( + cf.eval_trajectory_hermite(np.array(query), ts, ys, yps, np.array([])) + ) + + as_list = query + as_float32 = np.array(query, dtype=np.float32) + # genuinely non-contiguous: a length-1 strided slice is trivially + # contiguous, so use two elements to force a real stride gap. + strided = np.array([0.25, -1.0, 0.75, -1.0])[::2] + assert not strided.flags["C_CONTIGUOUS"] + + for t_query in (as_list, as_float32, strided): + out = np.asarray( + cf.eval_trajectory_hermite(t_query, ts, ys, yps, np.array([])) + ) + np.testing.assert_allclose(out, expected) + + def test_eval_trajectory_hermite_accepts_loose_ts(self): + cf = self._cf() + ys = np.array([[1.0, 2.0]]) + yps = np.array([[1.0, 1.0]]) + tq = np.array([0.5]) + expected = np.asarray( + cf.eval_trajectory_hermite(tq, np.array([0.0, 1.0]), ys, yps, np.array([])) + ) + for ts in ([0.0, 1.0], np.array([0.0, 1.0], dtype=np.float32)): + out = np.asarray(cf.eval_trajectory_hermite(tq, ts, ys, yps, np.array([]))) + np.testing.assert_allclose(out, expected) + + def test_eval_trajectory_accepts_loose_ts(self): + cf = self._cf() + y_traj = np.array([[1.0, 2.0, 3.0]]) + p = np.array([]) + expected = np.asarray(cf.eval_trajectory(np.array([0.0, 0.5, 1.0]), y_traj, p)) + for ts in ([0.0, 0.5, 1.0], np.array([0.0, 0.5, 1.0], dtype=np.float32)): + out = np.asarray(cf.eval_trajectory(ts, y_traj, p)) + np.testing.assert_allclose(out, expected) + + def test_jvp_trajectory_accepts_loose_ts(self): + cf = self._cf() + y_traj = np.array([[1.0, 2.0, 3.0]]) + vy_traj = np.array([[0.1, 0.2, 0.3]]) + p = np.array([]) + expected = np.asarray( + cf.jvp_trajectory(np.array([0.0, 0.5, 1.0]), y_traj, p, vy_traj) + ) + for ts in ([0.0, 0.5, 1.0], np.array([0.0, 0.5, 1.0], dtype=np.float32)): + out = np.asarray(cf.jvp_trajectory(ts, y_traj, p, vy_traj)) + np.testing.assert_allclose(out, expected) + + def test_group_eval_trajectory_accepts_loose_ts(self): + group = self._group() + y_traj = np.array([[1.0, 2.0, 3.0]]) + p = np.array([]) + (expected,) = group.eval_trajectory(np.array([0.0, 0.5, 1.0]), y_traj, p) + for ts in ([0.0, 0.5, 1.0], np.array([0.0, 0.5, 1.0], dtype=np.float32)): + (out,) = group.eval_trajectory(ts, y_traj, p) + np.testing.assert_allclose(np.asarray(out), np.asarray(expected)) + + def test_group_eval_trajectory_hermite_accepts_loose_t_query(self): + group = self._group() + ts = np.array([0.0, 1.0]) + ys = np.array([[1.0, 2.0]]) + yps = np.array([[1.0, 1.0]]) + p = np.array([]) + (expected,) = group.eval_trajectory_hermite(np.array([0.5]), ts, ys, yps, p) + for t_query in ([0.5], np.array([0.5], dtype=np.float32)): + (out,) = group.eval_trajectory_hermite(t_query, ts, ys, yps, p) + np.testing.assert_allclose(np.asarray(out), np.asarray(expected)) + + +class TestDenseMatrixToRust: + def test_dense_matrix_matmul_to_rust(self): + from pybamm.rust import ExprGraph + from pybamm.solvers.rust_evaluator import RustEvaluator + + A = np.arange(6.0).reshape(2, 3) + y = pybamm.StateVector(slice(0, 3)) + expr = pybamm.Matrix(A) @ y + graph = ExprGraph() + cf = graph.compile(expr.to_rust(graph, {}), name="f", n_states=3) + yv = np.array([1.0, 2.0, 3.0]) + np.testing.assert_allclose( + np.asarray(cf(0.0, yv, np.array([]))).ravel(), A @ yv + ) + # jacobian path: d(A @ y)/dy == A, exercising sparsity/tangent on a dense LHS + ev = RustEvaluator(cf, "jac") + np.testing.assert_allclose(ev(0.0, yv, np.array([])).toarray(), A) + + +class TestMatMulLhsValidation: + def test_matmul_non_constant_lhs_raises_clean_error(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + lhs = graph.state_vector(0, 2) # computed (non-constant) LHS + rhs = graph.state_vector(0, 2) + mm = graph.matmul(lhs, rhs) + with pytest.raises(NotImplementedError, match=r"MatMul left operand"): + graph.compile(mm, name="f", n_states=2) + + def test_matmul_non_constant_lhs_raises_on_eval_to_array(self): + # eval_to_array bypassed check_supported (direct CompiledExpr::new), + # so a non-constant MatMul LHS panicked instead of raising cleanly. + from pybamm.rust import ExprGraph + + graph = ExprGraph() + sv = graph.state_vector(0, 2) + mm = graph.matmul(sv, sv) + with pytest.raises(NotImplementedError, match=r"MatMul left operand"): + graph.eval_to_array(mm, 0.0, np.array([1.0, 2.0]), np.array([]), []) + + +class TestExpressionShapeValidation: + def test_incompatible_binary_widths_raise_before_evaluation(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + short = graph.state_vector(0, 2) + long = graph.state_vector(2, 5) + invalid = short + long + + with pytest.raises(ValueError, match=r"incompatible operand widths 2 and 3"): + graph.compile(invalid, name="invalid", n_states=5) + + +class TestRustEvaluatorWrappers: + def _cf(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + a = graph.input_parameter("a") + y0 = graph.state_vector(0, 1) + y1 = graph.state_vector(1, 2) + expr = graph.concat([y0 * a - y1, y0 + y1 * y1]) + return graph.compile(expr, name="rhs", n_states=2) + + def test_func_returns_column(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + out = RustEvaluator(cf, "func")(0.0, np.array([[1.5], [-2.0]]), np.array([3.0])) + assert out.shape == (2, 1) + np.testing.assert_allclose(out.ravel(), [6.5, 5.5]) + + def test_jac_returns_scipy_csc(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + jac = RustEvaluator(cf, "jac")(0.0, np.array([1.5, -2.0]), np.array([3.0])) + assert scipy.sparse.issparse(jac) + np.testing.assert_allclose(jac.toarray(), [[3.0, -1.0], [1.0, -4.0]]) + + def test_jac_action_matches_jac_matvec(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + y, p, v = np.array([1.5, -2.0]), np.array([3.0]), np.array([0.3, 0.7]) + jv = RustEvaluator(cf, "jac_action")(0.0, y, p, v) + full = RustEvaluator(cf, "jac")(0.0, y, p).toarray() @ v + np.testing.assert_allclose(jv.ravel(), full) + + def test_jacp_tuple_per_param(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + # sens_indices=[0] selects the "a" column of df/dp + out = RustEvaluator(cf, "jacp", sens_indices=[0])( + 0.0, np.array([1.5, -2.0]), np.array([3.0]) + ) + assert isinstance(out, tuple) and len(out) == 1 + np.testing.assert_allclose(out[0].ravel(), [1.5, 0.0]) + + def test_jacobian_is_derived_lazily(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + ev = RustEvaluator(cf, "jac") + assert ev._jac is None # not derived at construction + ev(0.0, np.array([1.5, -2.0]), np.array([3.0])) + assert ev._jac is not None # derived + cached on first call + + def test_func_accepts_eval_only(self): + from pybamm.solvers.rust_evaluator import RustEvaluator + + cf = self._cf() + out = RustEvaluator(cf, "func")(0.0, np.array([1.5, -2.0]), np.array([3.0])) + assert out.shape == (2, 1) + np.testing.assert_allclose(out.ravel(), [6.5, 5.5]) + + +def _jacp_as_list(result): + """Normalise a jacp call result to a list of 1-D arrays. + + casadi returns a single DM (not a tuple) when there is exactly one output, + while RustEvaluator always returns a tuple. + """ + if isinstance(result, casadi.DM): + result = (result,) + return [np.asarray(j).ravel() for j in result] + + +def _toy_dae(convert_to_format): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + model.rhs = {u: -2 * u + a * v} + model.algebraic = {v: 2 * u - v} + model.initial_conditions = {u: 1.0, v: 2.0} + model.events = [pybamm.Event("u-0.05", u - 0.05)] + model.variables = {"u": u, "v": v} + disc = pybamm.Discretisation() + disc.process_model(model) + model.convert_to_format = convert_to_format + return model + + +class TestProcessRustParity: + @pytest.fixture + def pair(self): + from pybamm.solvers.base_solver import BaseSolver + + inputs = {"a": 0.5} + out = {} + for fmt in ("casadi", "rust"): + model = _toy_dae(fmt) + model.calculate_sensitivities = ["a"] + vars_ = BaseSolver._get_vars_for_processing(model, inputs) + rhs_alg = pybamm.numpy_concatenation( + model.concatenated_rhs, model.concatenated_algebraic + ) + out[fmt] = (model, vars_, rhs_alg) + return out + + def _stack(self, fmt, inputs): + if fmt == "casadi": + return casadi.vertcat(*inputs.values()) + return np.array(list(inputs.values()), dtype=np.float64) + + def test_rhs_algebraic_func_jac_action_jacp(self, pair): + from pybamm.solvers.base_solver import process + + t, y, inputs = 0.3, np.array([1.1, 0.4]), {"a": 0.5} + results = {} + for fmt, (_model, vars_, rhs_alg) in pair.items(): + func, jac, jacp, jac_action = process(rhs_alg, "rhs_algebraic", vars_) + p = self._stack(fmt, inputs) + jac_mat = jac(t, y, p) + results[fmt] = { + "f": np.asarray(func(t, y, p)).ravel(), + "jac": np.asarray( + jac_mat.toarray() if hasattr(jac_mat, "toarray") else jac_mat + ), + "jv": np.asarray(jac_action(t, y, p, np.array([0.2, -0.3]))).ravel(), + "jacp": _jacp_as_list(jacp(t, y, p)), + } + for key in ("f", "jac", "jv"): + np.testing.assert_allclose( + results["rust"][key], results["casadi"][key], rtol=1e-12, atol=1e-14 + ) + for r, c in zip( + results["rust"]["jacp"], results["casadi"]["jacp"], strict=True + ): + np.testing.assert_allclose(r, c, rtol=1e-12, atol=1e-14) + + def test_event_and_ic_eval_only(self, pair): + from pybamm.solvers.base_solver import process + + t, y, inputs = 0.0, np.array([1.0, 2.0]), {"a": 0.5} + vals = {} + for fmt, (model, vars_, _) in pair.items(): + ev = process( + model.events[0].expression, "event_0", vars_, use_jacobian=False + )[0] + ic = process( + model.concatenated_initial_conditions, + "initial_conditions", + vars_, + use_jacobian=False, + )[0] + p = self._stack(fmt, inputs) + vals[fmt] = ( + float(np.asarray(ev(t, y, p)).item()), + np.asarray(ic(t, np.zeros((2, 1)), p)).ravel(), + ) + assert vals["rust"][0] == pytest.approx(vals["casadi"][0], rel=1e-12) + np.testing.assert_allclose(vals["rust"][1], vals["casadi"][1], rtol=1e-12) + + def test_rectangular_group_jacobian_matches_casadi(self, pair): + from pybamm.solvers.base_solver import process + + t, y = 0.3, np.array([1.1, 0.4]) + mats = {} + for fmt, (model, vars_, _) in pair.items(): + _, jac, _, jac_action = process(model.concatenated_rhs, "RHS", vars_) + assert jac is not None and jac_action is not None + p = self._stack(fmt, inputs={"a": 0.5}) + jm = jac(t, y, p) + mats[fmt] = np.asarray(jm.toarray() if hasattr(jm, "toarray") else jm) + np.testing.assert_allclose(mats["rust"], mats["casadi"], rtol=1e-12, atol=1e-14) + + +class TestProcessRustDuplication: + # These exercise the BaseSolver.set_up skip logic in isolation, via a minimal + # whole-model-artifact solver rather than the real IDAKLUSolver rust path. + def test_lazy_jac_not_built_for_uncalled_groups(self): + from pybamm.solvers.base_solver import BaseSolver, process + from pybamm.solvers.rust_evaluator import RustEvaluator + + model = _toy_dae("rust") + model.calculate_sensitivities = [] + vars_ = BaseSolver._get_vars_for_processing(model, {"a": 0.5}) + _, jac, _, jac_action = process(model.concatenated_rhs, "RHS", vars_) + assert ( + isinstance(jac, RustEvaluator) + and jac._jac is None + and jac_action is not None + ) # not yet derived + + def test_whole_model_solver_skips_every_per_group_lowering(self, monkeypatch): + from pybamm.solvers.base_solver import BaseSolver + + # Minimal concrete subclass with the whole-model-artifact flag + class _WholeModelSolver(BaseSolver): + _integrates_via_compiled_model = True + + def _run(self, *a, **kw): # pragma: no cover + raise NotImplementedError + + model = _toy_dae("rust") + solver = _WholeModelSolver() + + # Patch _set_initial_conditions so base set_up can be called without the + # full initial-condition plumbing. + monkeypatch.setattr( + BaseSolver, "_set_initial_conditions", lambda *a, **kw: None + ) + # Patch _set_up_events to return empty structures + monkeypatch.setattr( + BaseSolver, + "_set_up_events", + lambda *a, **kw: ([], [], [], {}, []), + ) + + solver.set_up(model, inputs=[{"a": 0.5}]) + + # Every group is served from the solver's own shared lowering, which fills + # rhs_eval via RustModelLowering.bind_generic_evaluators. + assert model.rhs_eval is None + assert model.jac_rhs_eval is None + assert model.jac_rhs_action_eval is None + assert model.jacp_rhs_eval is None + assert model.algebraic_eval is None + assert model.rhs_algebraic_eval is None + + def test_rust_output_block_skipped(self, monkeypatch): + from pybamm.solvers.base_solver import BaseSolver + + class _WholeModelSolver(BaseSolver): + _integrates_via_compiled_model = True + + def _run(self, *a, **kw): # pragma: no cover + raise NotImplementedError + + model = _toy_dae("rust") + solver = _WholeModelSolver(output_variables=["u"]) + + monkeypatch.setattr( + BaseSolver, "_set_initial_conditions", lambda *a, **kw: None + ) + monkeypatch.setattr( + BaseSolver, + "_set_up_events", + lambda *a, **kw: ([], [], [], {}, []), + ) + + solver.set_up(model, inputs=[{"a": 0.5}]) + + assert solver.computed_var_fcns == {} + + +class TestStackedInputPredicate: + def test_uses_stacked_inputs(self): + model = pybamm.BaseModel() + for fmt, expected in [ + ("casadi", True), + ("rust", True), + ("python", False), + ("jax", False), + (None, False), + ]: + model.convert_to_format = fmt + assert model.uses_stacked_inputs is expected + + def test_stack_inputs_rust_is_ndarray(self): + from pybamm.solvers.base_solver import stack_inputs + + out = stack_inputs({"a": 1.0, "b": np.array([2.0, 3.0])}, "rust") + np.testing.assert_allclose(out, [1.0, 2.0, 3.0]) + assert stack_inputs({}, "rust").size == 0 + + def test_set_initial_conditions_and_event_check_rust(self): + from pybamm.solvers.base_solver import BaseSolver + + class _DaeSolver(BaseSolver): + def _run(self, *a, **kw): # pragma: no cover + raise NotImplementedError + + model = _toy_dae("rust") + solver = _DaeSolver() + # set_up exercises _set_initial_conditions and the rust IC evaluator + solver.set_up(model, inputs=[{"a": 0.5}], ics_only=True) + y0 = np.asarray(model.y0_list[0]).ravel() + np.testing.assert_allclose(y0, [1.0, 2.0]) + # event check must not raise for positive events + model.terminate_events_eval = [] + solver._check_event_violation([0.0], model, y0, {"a": 0.5}) + + +class TestRustSolverGuards: + def test_casadi_solver_rejects_rust_model(self): + model = _toy_dae("rust") + with pytest.raises(pybamm.SolverError, match="convert_to_format='rust'"): + pybamm.CasadiSolver()._check_and_prepare_model_inplace(model) + assert model.convert_to_format == "rust" # never silently forced + + def test_casadi_root_method_switches_to_rust_newton(self): + class _DaeSolver(pybamm.BaseSolver): + def _run(self, *a, **kw): # pragma: no cover + raise NotImplementedError + + model = _toy_dae("rust") + solver = _DaeSolver() + solver.root_method = "casadi" + solver._check_and_prepare_model_inplace(model) + assert isinstance(solver.root_method, pybamm.NonlinearSolver) + assert model.convert_to_format == "rust" # never silently forced + + +class TestRustAlgebraicMinimize: + def test_minimize_converges_on_rust_model(self): + # scipy minimize needs the exact gradient of sum(f**2), and jac_norm's + # broadcasting must hold for numpy jacobians as well as casadi DM. + model = pybamm.BaseModel() + whole_cell = ["negative electrode", "separator", "positive electrode"] + var1 = pybamm.Variable("var1", domain=whole_cell) + var2 = pybamm.Variable("var2", domain=whole_cell) + model.algebraic = {var1: var1 - 3, var2: 2 * var1 - var2} + model.initial_conditions = {var1: pybamm.Scalar(1), var2: pybamm.Scalar(4)} + model.variables = {"var1": var1, "var2": var2} + disc = get_discretisation_for_testing() + disc.process_model(model) + model.convert_to_format = "rust" + + solution = pybamm.AlgebraicSolver("minimize", tol=1e-8).solve(model) + np.testing.assert_allclose( + model.get_processed_variable("var1").evaluate(t=None, y=solution.y), + 3, + rtol=1e-7, + atol=1e-6, + ) + np.testing.assert_allclose( + model.get_processed_variable("var2").evaluate(t=None, y=solution.y), + 6, + rtol=1e-7, + atol=1e-6, + ) + + +class TestVectorWidthInputParameter: + """`ExprGraph.input_parameter(name, width)` — packed-offset support for + vector-valued (`expected_size > 1`) inputs.""" + + def test_default_width_is_scalar(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + graph.input_parameter("a") + assert graph.n_inputs() == 1 + + def test_n_inputs_is_total_packed_width_not_name_count(self): + # n_inputs() sizes the packed `p` buffer at the FFI boundary, so it must be + # the total width, not the count of distinct names (2 names, 3 values). + from pybamm.rust import ExprGraph + + graph = ExprGraph() + graph.input_parameter("a") + graph.input_parameter("b", 2) + assert graph.n_inputs() == 3 + + def test_vector_input_indexes_into_packed_offset(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + a = graph.input_parameter("a") # width 1, offset 0 + b = graph.input_parameter("b", 2) # width 2, offset 1 + y0 = graph.state_vector(0, 1) + b0 = graph.index(b, 0, 1) + b1 = graph.index(b, 1, 2) + expr = graph.concat([(b0 + b1) * a * y0]) + cf = graph.compile(expr, name="f", n_states=1) + out = np.asarray(cf(0.0, np.array([2.0]), np.array([3.0, 0.2, 0.3]))) + np.testing.assert_allclose(out, [3.0 * (0.2 + 0.3) * 2.0]) + + def test_check_p_reports_total_width_not_name_count(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + a = graph.input_parameter("a") + b = graph.input_parameter("b", 2) + expr = graph.concat([a + graph.index(b, 0, 1) + graph.index(b, 1, 2)]) + cf = graph.compile(expr, name="f", n_states=0) + with pytest.raises( + ValueError, + match=r"f: expected 3 input values \(2 parameters\), got 2", + ): + cf(0.0, np.empty(0), np.array([1.0, 2.0])) + + def test_pack_dict_path_validates_vector_length(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + a = graph.input_parameter("a") + b = graph.input_parameter("b", 2) + expr = graph.concat([a + graph.index(b, 0, 1) + graph.index(b, 1, 2)]) + cf = graph.compile(expr, name="f", n_states=0) + # correct-length dict path matches the array path + via_array = np.asarray(cf(0.0, np.empty(0), np.array([1.0, 0.2, 0.3]))) + via_dict = np.asarray( + cf(0.0, np.empty(0), {"a": 1.0, "b": np.array([[0.2], [0.3]])}) + ) + np.testing.assert_allclose(via_array, via_dict) + # wrong-length value for 'b' must be rejected, not silently truncated + with pytest.raises( + ValueError, match=r"input 'b' must have 2 value\(s\), got 3" + ): + cf(0.0, np.empty(0), {"a": 1.0, "b": np.array([0.2, 0.3, 0.4])}) + + def test_reregistering_input_with_different_width_raises(self): + from pybamm.rust import ExprGraph + + graph = ExprGraph() + graph.input_parameter("a") + with pytest.raises(ValueError, match=r"'a'.*width 2.*width 1"): + graph.input_parameter("a", 2) + + +class TestPickle: + def test_expr_graph_pickle_roundtrip(self): + import pickle + + from pybamm.rust import ExprGraph + + graph = ExprGraph() + (pybamm.InputParameter("a") * pybamm.StateVector(slice(0, 2))).to_rust( + graph, {} + ) + g2 = pickle.loads(pickle.dumps(graph)) + # same expression converts and compiles identically on the restored graph + expr2 = (pybamm.InputParameter("a") * pybamm.StateVector(slice(0, 2))).to_rust( + g2, {} + ) + cf = g2.compile(expr2, name="f", n_states=2) + np.testing.assert_allclose( + np.asarray(cf(0.0, np.array([1.0, 2.0]), np.array([3.0]))).ravel(), + [3.0, 6.0], + ) + + def test_pickle_preserves_input_registration_indices(self): + import pickle + + from pybamm.rust import ExprGraph + + graph = ExprGraph() + graph.input_parameter("a") + graph.input_parameter("b") + g2 = pickle.loads(pickle.dumps(graph)) + # "b" must still resolve to index 1 on the restored graph; if the + # input map were lost, re-registration would give "b" index 0 ("a"). + expr_b = (pybamm.InputParameter("b") * pybamm.StateVector(slice(0, 2))).to_rust( + g2, {} + ) + cf = g2.compile(expr_b, name="fb", n_states=2) + np.testing.assert_allclose( + np.asarray(cf(0.0, np.array([1.0, 2.0]), np.array([5.0, 7.0]))).ravel(), + [7.0, 14.0], + ) + + def test_pickle_preserves_input_widths(self): + import pickle + + from pybamm.rust import ExprGraph + + graph = ExprGraph() + graph.input_parameter("a") + graph.input_parameter("b", 2) + g2 = pickle.loads(pickle.dumps(graph)) + # re-registering "b" with a different width on the restored graph + # must still be rejected — the width table round-tripped, not lost. + with pytest.raises(ValueError, match=r"'b'.*width 1.*width 2"): + g2.input_parameter("b", 1) + # "b" must still occupy packed offset 1 (after "a"'s width-1 slot) + # and keep its width-2 slot on the restored graph. + b = pybamm.InputParameter("b", expected_size=2) + expr_b = (pybamm.Index(b, 0) + pybamm.Index(b, 1)).to_rust(g2, {}) + cf = g2.compile(expr_b, name="fb", n_states=0) + np.testing.assert_allclose( + np.asarray(cf(0.0, np.empty(0), np.array([5.0, 7.0, 11.0]))), + [18.0], + ) + + def test_compiled_function_pickle_roundtrip(self): + import pickle + + from pybamm.rust import ExprGraph + + graph = ExprGraph() + expr = (pybamm.InputParameter("a") * pybamm.StateVector(slice(0, 2))).to_rust( + graph, {} + ) + cf = graph.compile(expr, name="f", n_states=2) + cf2 = pickle.loads(pickle.dumps(cf)) + y, p = np.array([1.0, 2.0]), np.array([3.0]) + np.testing.assert_allclose( + np.asarray(cf2(0.0, y, p)), np.asarray(cf(0.0, y, p)) + ) + # jacobian still derivable after the round-trip + np.testing.assert_allclose( + cf2.jacobian("y")(0.0, y, p).toarray(), + cf.jacobian("y")(0.0, y, p).toarray(), + ) + + def test_rebuild_rejects_out_of_range_root(self): + from pybamm.rust import CompiledFunction, ExprGraph + + graph = ExprGraph() + graph.state_vector(0, 2) + with pytest.raises(ValueError, match=r"out of range"): + CompiledFunction._rebuild(graph, 2**31, None, None) + + def test_rebuild_rejects_unsupported_root(self): + from pybamm.rust import CompiledFunction, ExprGraph + + graph = ExprGraph() + # non-constant MatMul LHS is a lowering blocker; _rebuild must apply + # the same check_supported gate as graph.compile + mm = graph.matmul(graph.state_vector(0, 2), graph.state_vector(0, 2)) + with pytest.raises(NotImplementedError, match=r"MatMul left operand"): + CompiledFunction._rebuild(graph, mm.id, None, None) + + def test_compiled_model_pickle_roundtrip(self): + import pickle + + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + model.rhs = {u: 0.1 * v} + model.algebraic = {v: 2 * u - v} + model.initial_conditions = {u: 0, v: 0} + model.convert_to_format = "rust" + disc = pybamm.Discretisation() + disc.process_model(model) + solver = pybamm.IDAKLUSolver() + solver.solve(model, np.array([0.0, 1.0])) + rm = solver._setup["rust_model"] + rm2 = pickle.loads(pickle.dumps(rm)) + assert rm2.nnz == rm.nnz + np.testing.assert_array_equal( + rm2.csc_sparsity_pattern()[0], rm.csc_sparsity_pattern()[0] + ) + + def test_rust_evaluator_pickle_rederives_jacobian(self): + import pickle + + from pybamm.rust import ExprGraph + from pybamm.solvers.rust_evaluator import RustEvaluator + + graph = ExprGraph() + expr = (3 * pybamm.StateVector(slice(0, 2))).to_rust(graph, {}) + ev = RustEvaluator(graph.compile(expr, name="f", n_states=2), "jac") + y, p = np.array([1.0, 2.0]), np.array([]) + expected = ev(0.0, y, p).toarray() # populates the _jac cache + ev2 = pickle.loads(pickle.dumps(ev)) + np.testing.assert_allclose(ev2(0.0, y, p).toarray(), expected) diff --git a/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py b/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py index 82c359cb9d..8153dea72b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py +++ b/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py @@ -1575,6 +1575,31 @@ def spl(t): # Check that the unsorted and sorted arrays are the same assert np.all(y_unsorted == y_sorted[idxs_unsort]) + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_unsorted_t_query_returns_query_order(self, hermite_interp): + # Both interpolation routes must return values in the caller's order: + # only the hermite route consumes the internally sorted times. + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = t * y + model = pybamm.BaseModel() + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) + var_casadi = to_casadi(var, y_sol) + processed_var = pybamm.process_variable( + "test", + [var], + [var_casadi], + self._sol_default(t_sol, y_sol, yp_sol, model), + ) + + t_unsorted = np.array([0.9, 0.3, 0.6]) + # var = t * y = 5 t^2; atol covers linear-interp error on the xr route + np.testing.assert_allclose( + processed_var(t_unsorted), 5 * t_unsorted**2, atol=2e-3 + ) + def test_as_computed_0D(self): # 0D t = pybamm.t diff --git a/packages/pybamm/tests/unit/test_solvers/test_processed_variable_computed.py b/packages/pybamm/tests/unit/test_solvers/test_processed_variable_computed.py index e39bf29be5..dc2f8d7427 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_processed_variable_computed.py +++ b/packages/pybamm/tests/unit/test_solvers/test_processed_variable_computed.py @@ -110,6 +110,52 @@ def test_processed_variable_0D(self): comb_var = processed_var.update(processed_var2, comb_sol) np.testing.assert_array_equal(comb_var.entries, np.append(y_sol, y_sol2)) + def _build_0D_var(self): + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = t * y + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + var_casadi = to_casadi(var, y_sol) + return pybamm.ProcessedVariableComputed( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + ) + + def test_0D_call_matches_the_xarray_route(self): + # __call__ takes an np.interp fast path for time-only 0D queries; it + # must agree with xarray on values, query order, and NaN fill. + processed_var = self._build_0D_var() + t_query = np.array([0.9, 0.3, 0.6]) + + # entries are 5*t, so linear interpolation is exact + np.testing.assert_allclose(processed_var(t_query), 5 * t_query, rtol=1e-12) + np.testing.assert_allclose(processed_var(0.5), 2.5, rtol=1e-12) + np.testing.assert_allclose( + processed_var(t_query), + processed_var._xr_data_array.interp(t=t_query).values, + rtol=1e-14, + ) + # out-of-range queries keep xarray's NaN fill + np.testing.assert_array_equal( + np.isnan(processed_var(np.array([-0.5, 0.5, 2.0]))), + [True, False, True], + ) + + def test_data_read_does_not_build_the_data_array(self): + # .data must stay cheap: the xr.DataArray builds on first + # interpolation read only. + processed_var = self._build_0D_var() + processed_var.data + assert processed_var._xr_data_array_cache is None + assert processed_var._xr_interp_args is not None + + data_array = processed_var._xr_data_array + assert processed_var._xr_data_array_cache is data_array + assert processed_var._xr_interp_args is None + # check empty sensitivity works def test_processed_variable_0D_no_sensitivity(self): # without space @@ -530,10 +576,11 @@ def test_processed_variable_3D_r_R_x(self): var_casadi = to_casadi(var_sol, u_sol) geometry_options = {"options": {"particle size": "distribution"}} model = tests.get_base_model_with_battery_geometry(**geometry_options) + # base_variables_data is time-major (n_t, output); u_sol is (output, n_t) processed_var = pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], - [u_sol], + [u_sol.T], pybamm.Solution(t_sol, u_sol, model, {}), ) @@ -571,10 +618,11 @@ def test_processed_variable_3D_x_y_z(self, edges_eval): u_sol = np.ones(len(x_sol) * len(y_sol) * len(z_sol))[:, np.newaxis] * t_sol var_casadi = to_casadi(var_sol, u_sol) + # base_variables_data is time-major (n_t, output); u_sol is (output, n_t) processed_var = pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], - [u_sol], + [u_sol.T], pybamm.Solution(t_sol, u_sol, pybamm.BaseModel(), {}), ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py b/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py new file mode 100644 index 0000000000..89922df208 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py @@ -0,0 +1,95 @@ +"""The native diffsol solver must release the GIL during integration. + +A pure-Rust call that holds the GIL freezes every other Python thread for its +entire duration (there are no bytecode boundaries, so the periodic GIL switch +never fires). With ``py.detach`` around the integration another thread keeps +getting scheduled throughout the solve. This test guards that release so a +future change that drops the ``detach`` is caught. +""" + +import threading +import time + +import numpy as np + +from pybamm.rust import CompiledModel, ExprGraph, PreparedSolver + + +def _decay_prepared_solver(n_states): + """Native prepared solver for a decoupled linear decay system dy_i/dt=-y_i.""" + g = ExprGraph() + terms = [g.mul(g.scalar(-1.0), g.state_vector(i, i + 1)) for i in range(n_states)] + rhs = g.concat(terms) + model = CompiledModel.from_expr( + g, + rhs, + np.ones(n_states), # identity mass matrix (CSR) + np.arange(n_states + 1, dtype=np.int64), + np.arange(n_states, dtype=np.int64), + n_inputs=0, + sens_param_indices=[], + output_exprs=[], + event_exprs=[], + ) + return PreparedSolver(model, 1e-8, 1e-8) + + +def test_native_solve_releases_gil(): + n_states = 100 + ps = _decay_prepared_solver(n_states) + y0 = np.ones(n_states) + inputs = np.array([], dtype=np.float64) + no_stops = np.array([], dtype=np.float64) + # A fine output grid makes one native solve take hundreds of milliseconds. + t_eval = np.linspace(0.0, 5.0, 400_000) + + ps.solve( + t_eval, no_stops, y0, inputs + ) # warm up (first solve builds internal state) + + stop = threading.Event() + ticks = [] + + def worker(): + # Each tick needs the GIL only briefly, and sleeping keeps this thread off + # the CPU so it never competes with the solve for cores or memory. + while not stop.is_set(): + ticks.append(time.perf_counter()) + time.sleep(0.001) + + t = threading.Thread(target=worker) + t.start() + try: + # The cadence the worker keeps with the main thread idle. Sleep granularity + # varies by platform, so this is the yardstick, not a wall-clock constant. + time.sleep(0.1) + quiet_cadence = np.percentile(np.diff(ticks[:]), 95) + + wall0 = time.perf_counter() + ps.solve(t_eval, no_stops, y0, inputs) + wall1 = time.perf_counter() + finally: + stop.set() + t.join() + + solve_wall = wall1 - wall0 + assert solve_wall > 8 * quiet_cadence, ( + f"native solve took {solve_wall * 1e3:.1f} ms against a worker cadence of " + f"{quiet_cadence * 1e3:.1f} ms, too few ticks to judge the GIL either way; " + "raise the output-point count" + ) + + # A held GIL blocks the worker outright, so it records no tick inside the call. + during = [tick for tick in ticks if wall0 < tick < wall1] + assert during, ( + f"worker ran zero times during a {solve_wall * 1e3:.1f} ms native solve: " + "the GIL was not released" + ) + + # A GIL held for part of the solve leaves one long quiet stretch, which is what + # to bound; the worker's throughput would also track CPU and memory contention. + longest_gap = np.diff([wall0, *during, wall1]).max() + assert longest_gap < 0.5 * solve_wall, ( + f"worker stalled for {longest_gap * 1e3:.1f} ms inside a " + f"{solve_wall * 1e3:.1f} ms native solve: the GIL was held for part of it" + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_assembly_cost.py b/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_assembly_cost.py new file mode 100644 index 0000000000..9d3e026bed --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_assembly_cost.py @@ -0,0 +1,152 @@ +"""Budgets on the tangent-tape walks one Jacobian assembly costs. + +This count is fixed at compile time, so unlike wall times and solver counters it +survives contention and trajectory chaos. Budgets allow a lane-width halving; +raising one means confirming the model grew or that assembly got cheaper. +""" + +import numpy as np +import pytest + +import pybamm +from tests.shared import POUCH_OPTIONS, POUCH_PTS, build_rust_model + + +def sweeps_per_assembly(stats): + """Batched colour sweeps plus one reverse pass per split dense row.""" + batched = -(-stats["n_colors"] // stats["jac_lane_width"]) + return batched + stats["n_dense_rows"] + + +def _stats_for_model(model, var_pts=None): + """`(jacobian_stats, n_states)` for a plainly built model.""" + built, rust_model = build_rust_model(model, var_pts) + return rust_model.jacobian_stats(), built.len_rhs_and_alg + + +def _stats_for_step(model, step, var_pts, key_prefix): + """`(jacobian_stats, n_states)` for one experiment step's own model.""" + model.convert_to_format = "rust" + sim = pybamm.Simulation( + model, var_pts=var_pts, experiment=pybamm.Experiment([step]) + ) + sim.build_for_experiment() + built = next( + m for key, m in sim.steps_to_built_models.items() if key.startswith(key_prefix) + ) + solver = pybamm.IDAKLUSolver() + solver.set_up(built, inputs={}, t_eval=np.array([0.0, 1.0])) + return solver._setup["rust_model"].jacobian_stats(), built.len_rhs_and_alg + + +PLAIN_CASES = [ + ("SPM", lambda: pybamm.lithium_ion.SPM(), None, 1), + ("SPMe", lambda: pybamm.lithium_ion.SPMe(), None, 3), + ( + "SPMe-voltage-as-a-state", + lambda: pybamm.lithium_ion.SPMe({"voltage as a state": "true"}), + None, + 4, + ), + ("DFN", lambda: pybamm.lithium_ion.DFN(), None, 3), + ( + "DFN-fine", + lambda: pybamm.lithium_ion.DFN(), + {"x_n": 30, "x_s": 30, "x_p": 30, "r_n": 30, "r_p": 30}, + 3, + ), +] + +STEP_CASES = [ + ("pouch-2plus1D-CC", POUCH_OPTIONS, "Discharge at 1C until 2.7 V", "CRate", 3), + ("pouch-2plus1D-CV", POUCH_OPTIONS, "Hold at 4.1 V until C/20", "Voltage", 4), + ( + "pouch-1plus1D-CV", + {"current collector": "potential pair", "dimensionality": 1}, + "Hold at 4.1 V until C/20", + "Voltage", + 3, + ), +] + + +class TestAssemblySweepBudget: + @pytest.mark.parametrize( + ("name", "build", "var_pts", "budget"), + PLAIN_CASES, + ids=[case[0] for case in PLAIN_CASES], + ) + def test_plain_model_stays_within_sweep_budget(self, name, build, var_pts, budget): + stats, _ = _stats_for_model(build(), var_pts) + assert sweeps_per_assembly(stats) <= budget, ( + f"{name} assembles in {sweeps_per_assembly(stats)} sweeps " + f"(budget {budget}): {stats}" + ) + + @pytest.mark.parametrize( + ("name", "options", "step", "prefix", "budget"), + STEP_CASES, + ids=[case[0] for case in STEP_CASES], + ) + def test_experiment_step_stays_within_sweep_budget( + self, name, options, step, prefix, budget + ): + var_pts = POUCH_PTS if options["dimensionality"] == 2 else None + if options["dimensionality"] == 1: + var_pts = {"x_n": 4, "x_s": 4, "x_p": 4, "r_n": 4, "r_p": 4, "z": 10} + stats, _ = _stats_for_step( + pybamm.lithium_ion.DFN(options), step, var_pts, prefix + ) + assert sweeps_per_assembly(stats) <= budget, ( + f"{name} assembles in {sweeps_per_assembly(stats)} sweeps " + f"(budget {budget}): {stats}" + ) + + +class TestNoWideRowEscapesTheSplit: + def test_a_constraint_row_never_sets_the_colour_count(self): + """A row far wider than the rest must be split, not coloured against.""" + stats, _ = _stats_for_step( + pybamm.lithium_ion.DFN(POUCH_OPTIONS), + "Hold at 4.1 V until C/20", + POUCH_PTS, + "Voltage", + ) + assert stats["n_dense_rows"] == 1 + assert stats["dense_row_entries"] > 10 * stats["n_colors"] + + def test_batching_engages_on_a_multi_colour_model(self): + """Batching turning itself off multiplies assembly cost by the lane width.""" + for stats, _ in ( + _stats_for_model(pybamm.lithium_ion.DFN()), + _stats_for_step( + pybamm.lithium_ion.DFN(POUCH_OPTIONS), + "Hold at 4.1 V until C/20", + POUCH_PTS, + "Voltage", + ), + ): + assert stats["n_colors"] > 1 + assert stats["jac_lane_width"] > 1 + + +class TestConstantEntriesCarryTheirWeight: + """The classifier silently proving nothing would cost only sweeps, so the + budgets above would still pass. These pin that it fires.""" + + def test_a_dfn_proves_most_of_its_entries_constant(self): + stats, n_states = _stats_for_model(pybamm.lithium_ion.DFN()) + assert stats["n_constant_entries"] > stats["nnz"] // 2 + assert stats["n_swept_columns"] < n_states + + def test_the_pouch_stops_colouring_against_its_constant_rows(self): + stats, n_states = _stats_for_step( + pybamm.lithium_ion.DFN(POUCH_OPTIONS), + "Discharge at 1C until 2.7 V", + POUCH_PTS, + "CRate", + ) + assert stats["n_constant_entries"] > stats["nnz"] // 10 + # The current-collector rows are wholly constant, so their columns are + # never seeded and the coloring does not have to separate them. + assert stats["n_swept_columns"] < n_states diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_batching.py b/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_batching.py new file mode 100644 index 0000000000..68243a7b07 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_jacobian_batching.py @@ -0,0 +1,87 @@ +"""Coloured Jacobian assembly batches its tangent sweeps without changing values. + +The Rust core evaluates the tangent tape once per colour; batching runs several +colours per sweep to amortise the sparse-operator gather. The rust-side property +tests pin bitwise equality of the sweep itself, so these pin the parts visible +from Python: that a real battery model actually engages batching, that a model +too small to benefit does not, and that the assembled matrix still matches +CasADi's Jacobian of the same discretised model. +""" + +import numpy as np +import pytest + +import pybamm +from tests.shared import ( + POUCH_OPTIONS, + POUCH_PTS, + build_rust_model, + dense_casadi_jacobian, + dense_rust_jacobian, +) + +pytest.importorskip("casadi") + + +class TestLaneWidth: + def test_a_dfn_batches_its_colour_sweeps(self): + _, rust_model = build_rust_model(pybamm.lithium_ion.DFN()) + stats = rust_model.jacobian_stats() + assert stats["n_colors"] > 1 + assert stats["jac_lane_width"] > 1 + + def test_a_single_colour_model_stays_scalar(self): + # Nonlinear, so the one entry is genuinely swept rather than folded. + model = pybamm.BaseModel() + model.convert_to_format = "rust" + var = pybamm.Variable("var") + model.rhs = {var: -(var**2)} + model.initial_conditions = {var: 1.0} + solver = pybamm.IDAKLUSolver() + solver.set_up(model, inputs=[{}]) + stats = solver._setup["rust_model"].jacobian_stats() + assert stats["n_colors"] == 1 + assert stats["jac_lane_width"] == 1 + + def test_a_linear_model_needs_no_colour_at_all(self): + # Every entry folds at compile time, so assembly is a table write. + model = pybamm.BaseModel() + model.convert_to_format = "rust" + var = pybamm.Variable("var") + model.rhs = {var: -var} + model.initial_conditions = {var: 1.0} + solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10) + solver.set_up(model, inputs=[{}]) + stats = solver._setup["rust_model"].jacobian_stats() + assert stats["n_colors"] == 0 + assert stats["n_swept_columns"] == 0 + assert stats["n_constant_entries"] == 1 + + t = np.linspace(0, 1, 5) + solution = solver.solve(model, t_eval=[0, 1], t_interp=t) + np.testing.assert_allclose(solution["var"](t), np.exp(-t), rtol=1e-7, atol=1e-7) + + +class TestAssembledValues: + def test_batched_assembly_matches_casadi(self): + """The batched sweep must reproduce CasADi's Jacobian of the same model.""" + built, rust_model = build_rust_model(pybamm.lithium_ion.DFN()) + assert rust_model.jacobian_stats()["jac_lane_width"] > 1 + + y, got = dense_rust_jacobian(built, rust_model) + want = dense_casadi_jacobian(pybamm.lithium_ion.DFN(), y) + np.testing.assert_allclose(got, want, rtol=1e-10, atol=1e-10) + + def test_pouch_assembly_matches_casadi(self): + """The pouch is the model whose colouring the constant split changes.""" + var_pts = POUCH_PTS | {"y": 4, "z": 4} + built, rust_model = build_rust_model( + pybamm.lithium_ion.DFN(POUCH_OPTIONS), var_pts + ) + stats = rust_model.jacobian_stats() + assert stats["n_constant_entries"] > 0 + assert stats["n_swept_columns"] < built.len_rhs_and_alg + + y, got = dense_rust_jacobian(built, rust_model) + want = dense_casadi_jacobian(pybamm.lithium_ion.DFN(POUCH_OPTIONS), y, var_pts) + np.testing.assert_allclose(got, want, rtol=1e-10, atol=1e-10) diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py b/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py new file mode 100644 index 0000000000..e701ee14ea --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py @@ -0,0 +1,248 @@ +"""Invariants every Rust-backed solver's model lowering must share. + +Each solver composes its own lowering, so these pin the parts that must agree +between them: output variables slice by flattened component count, graph input +indices follow the inputs dict, the state residual spans every state, and one +set_up lowers into one graph. +""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.rust_lowering import RustModelLowering + +_VECTOR_VAR = "Negative particle surface concentration [mol.m-3]" +_OUTPUTS = ["Voltage [V]", _VECTOR_VAR, "Current [A]"] + + +def _built_model(with_input=False, options=None): + """Build an SPM pinned to the Rust format, not the ambient default.""" + model = pybamm.lithium_ion.SPM(options) + model.convert_to_format = "rust" + params = model.default_parameter_values + if with_input: + params["Current function [A]"] = "[input]" + sim = pybamm.Simulation(model, parameter_values=params) + sim.build() + built = sim.built_model + built.convert_to_format = "rust" + return built + + +def _termination_events(model): + return [ + event + for event in model.events + if event.event_type == pybamm.EventType.TERMINATION + ] + + +class TestOutputLengths: + def test_lengths_are_the_flattened_component_counts(self): + model = _built_model() + lowering = RustModelLowering(model, {}) + lowering.state_residual() + + _, lengths = lowering.outputs(_OUTPUTS) + + assert lengths[0] == 1 + assert lengths[1] > 1, "a spatial variable must contribute every component" + assert lengths[2] == 1 + + def test_tensor_field_output_is_rejected(self): + """A stacked field would be mis-sliced by its tensor shape, so refuse it.""" + model = _built_model() + lowering = RustModelLowering(model, {}) + lowering.state_residual() + component = model.get_processed_variable_or_event(_VECTOR_VAR) + model.variables_and_events["field"] = pybamm.VectorField(component, component) + + with pytest.raises(pybamm.SolverError, match=r"'field' is a tensor field"): + lowering.outputs(["field"]) + + def test_idaklu_and_diffsol_agree_on_output_lengths(self): + t_eval = np.linspace(0, 100, 5) + + idaklu = pybamm.IDAKLUSolver(output_variables=_OUTPUTS) + idaklu.solve(_built_model(), t_eval) + + diffsol = pybamm.DiffsolSolver(output_variables=_OUTPUTS) + diffsol.solve(_built_model(), t_eval) + + assert idaklu._setup["output_assembly"].lens == diffsol._output_assembly.lens + + +class TestInputRegistration: + def test_graph_input_order_follows_the_inputs_dict(self): + model = _built_model(with_input=True) + inputs = {"Current function [A]": 0.5} + + lowering = RustModelLowering(model, inputs) + + assert lowering.input_name_order == list(inputs) + + def test_sensitivity_indices_point_into_the_input_order(self): + model = _built_model(with_input=True) + inputs = {"Current function [A]": 0.5} + model.calculate_sensitivities = ["Current function [A]"] + + lowering = RustModelLowering(model, inputs) + indices, names = lowering.sensitivity_indices(model.calculate_sensitivities) + + assert indices == [0] + assert names == ["Current function [A]"] + + def test_sensitivity_names_not_supplied_as_inputs_are_dropped(self): + model = _built_model(with_input=True) + inputs = {"Current function [A]": 0.5} + + lowering = RustModelLowering(model, inputs) + indices, names = lowering.sensitivity_indices(["Not an input"]) + + assert indices == [] + assert names == [] + + +class TestStateResidual: + def test_dae_residual_spans_every_state(self): + model = _built_model() + lowering = RustModelLowering(model, {}) + lowering.state_residual() + + compiled = lowering.compile() + + assert compiled.n_states == model.len_rhs_and_alg + + def test_compile_before_state_residual_is_rejected(self): + lowering = RustModelLowering(_built_model(), {}) + + with pytest.raises(pybamm.SolverError, match=r"state residual"): + lowering.compile() + + +class TestGenericEvaluators: + """The slots BaseSolver's backend-agnostic helpers read on the native path.""" + + def test_rhs_evaluator_spans_the_differential_block_only(self): + """A DAE's residual is wider than ydot0, so the rhs needs its own root.""" + model = _built_model(options={"surface form": "algebraic"}) + assert model.len_alg > 0, "expected a DAE" + lowering = RustModelLowering(model, {}) + lowering.state_residual() + + rhs = lowering.rhs_evaluator() + y0 = np.asarray(model.concatenated_initial_conditions.evaluate()).reshape(-1) + values = np.asarray(rhs(0.0, y0, np.array([]))).reshape(-1) + + assert values.shape[0] == model.len_rhs + + def test_rhs_evaluator_matches_the_residual_it_shares_nodes_with(self): + model = _built_model(options={"surface form": "algebraic"}) + lowering = RustModelLowering(model, {}) + lowering.state_residual() + rhs = lowering.rhs_evaluator() + compiled = lowering.compile() + + y0 = np.asarray(model.concatenated_initial_conditions.evaluate()).reshape(-1) + residual = np.asarray(compiled.rhs(0.0, y0, np.array([]))).reshape(-1) + + np.testing.assert_array_equal( + np.asarray(rhs(0.0, y0, np.array([]))).reshape(-1), + residual[: model.len_rhs], + ) + + def test_bound_events_are_views_not_copies(self): + """Views onto the roots the fused root-finding tape is built from.""" + model = _built_model() + lowering = RustModelLowering(model, {}) + lowering.state_residual() + lowering.termination_events() + compiled = lowering.compile() + + lowering.bind_generic_evaluators(compiled) + + assert [ + evaluator._cf for evaluator in model.terminate_events_eval + ] == compiled.events + + def test_bound_events_follow_model_events_order(self): + model = _built_model() + lowering = RustModelLowering(model, {}) + lowering.state_residual() + lowering.termination_events() + lowering.bind_generic_evaluators(lowering.compile()) + + events = _termination_events(model) + assert len(model.terminate_events_eval) == len(events) + y0 = np.asarray(model.concatenated_initial_conditions.evaluate()).reshape(-1) + for event, evaluator in zip(events, model.terminate_events_eval, strict=True): + np.testing.assert_allclose( + float(np.asarray(evaluator(0.0, y0, np.array([]))).ravel()[0]), + float(np.asarray(event.expression.evaluate(0.0, y0, inputs={})).item()), + rtol=1e-12, + ) + + @pytest.mark.parametrize( + "solver_class", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver] + ) + def test_every_native_solver_binds_both_slots(self, solver_class): + """An unbound events slot would silently skip the event-violation check.""" + model = _built_model() + events = _termination_events(model) + assert events, "expected events the check could silently skip" + + solver_class().set_up(model, inputs={}, t_eval=np.linspace(0, 100, 3)) + + assert model.rhs_eval is not None + assert len(model.terminate_events_eval) == len(events) + + +class TestOneGraphPerSetUp: + """One lowering seam per ``set_up``: every native evaluator off one graph.""" + + @staticmethod + def _count_graphs(monkeypatch, solver, model): + from pybamm.solvers import rust_lowering + + built = [] + original = rust_lowering.rust_graph_with_inputs + + def counting(*args, **kwargs): + built.append(None) + return original(*args, **kwargs) + + monkeypatch.setattr(rust_lowering, "rust_graph_with_inputs", counting) + solver.set_up(model, inputs={}, t_eval=np.linspace(0, 100, 3)) + return len(built) + + @pytest.mark.parametrize( + "solver_class", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver] + ) + def test_native_set_up_builds_one_graph_beside_the_initial_conditions( + self, monkeypatch, solver_class + ): + """Two graphs: the initial conditions, and the solver's shared lowering. + + Anything more means an expression the shared graph already holds was + lowered a second time. + """ + model = _built_model() + assert _termination_events(model), "expected events that could be re-lowered" + + assert self._count_graphs(monkeypatch, solver_class(), model) == 2 + + def test_casadi_path_still_lowers_per_group(self, monkeypatch): + """The skip is native-only; a CasADi model builds no Rust graph at all.""" + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + sim = pybamm.Simulation(model) + sim.build() + built = sim.built_model + built.convert_to_format = "casadi" + + count = self._count_graphs(monkeypatch, pybamm.IDAKLUSolver(), built) + + assert count == 0 + assert built.rhs_eval is not None + assert len(built.terminate_events_eval) == len(_termination_events(built)) diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index b81cb19afe..089594a67e 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -931,6 +931,17 @@ def test_pickle_first_states_across_processes(self, tmp_path): assert "DATA" in save_result.stdout assert save_result.stdout == load_result.stdout + def test_solution_save_load_rust_backend(self, tmp_path): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=pybamm.IDAKLUSolver()) + sol = sim.solve([0, 3600]) + expected = sol["Voltage [V]"].data + path = tmp_path / "sol.pkl" + sol.save(str(path)) + loaded = pybamm.load(str(path)) + np.testing.assert_allclose(loaded["Voltage [V]"].data, expected) + def test_solution_evals_with_inputs(self): model = pybamm.lithium_ion.SPM() geometry = model.default_geometry @@ -994,7 +1005,13 @@ def test_discrete_data_sum_errors(self): @pytest.mark.parametrize( "solver_class,use_post_sum,use_output_var", _solver_classes ) - def test_discrete_data_sum(self, solver_class, use_post_sum, use_output_var): + def test_discrete_data_sum( + self, solver_class, use_post_sum, use_output_var, monkeypatch + ): + if solver_class is pybamm.CasadiSolver: + monkeypatch.setattr( + pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi" + ) model = pybamm.BaseModel(name="test_model") c = pybamm.Variable("c") model.rhs = {c: -2 * c} @@ -1111,7 +1128,13 @@ def test_discrete_data_sum(self, solver_class, use_post_sum, use_output_var): @pytest.mark.parametrize( "solver_class,use_post_sum,use_output_var", _solver_classes ) - def test_explicit_time_integral(self, solver_class, use_post_sum, use_output_var): + def test_explicit_time_integral( + self, solver_class, use_post_sum, use_output_var, monkeypatch + ): + if solver_class is pybamm.CasadiSolver: + monkeypatch.setattr( + pybamm.BaseModel, "_DEFAULT_CONVERT_TO_FORMAT", "casadi" + ) times = np.linspace(0, 1, 10) c = pybamm.Variable("c") if solver_class == pybamm.IDAKLUSolver: diff --git a/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py b/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py index 7b0e308e0f..d2baa4255b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py +++ b/packages/pybamm/tests/unit/test_solvers/test_store_first_last.py @@ -114,6 +114,7 @@ def test_composes_with_output_variables(self): def test_non_idaklu_warns_and_no_ops(self): model = _build_simple_dae_model() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(store_first_last=True) t_eval = np.linspace(0, 1, 11) @@ -125,6 +126,7 @@ def test_non_idaklu_warns_and_no_ops(self): def test_non_idaklu_warning_fires_once_per_instance(self): model = _build_simple_dae_model() + model.convert_to_format = "casadi" solver = pybamm.CasadiSolver(store_first_last=True) t_eval = np.linspace(0, 1, 5) diff --git a/packages/pybammsolvers/CMakeLists.txt b/packages/pybammsolvers/CMakeLists.txt index 579692239d..002272795f 100644 --- a/packages/pybammsolvers/CMakeLists.txt +++ b/packages/pybammsolvers/CMakeLists.txt @@ -124,6 +124,9 @@ pybind11_add_module(idaklu src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionTypes.hpp + # IDAKLU expressions - Rust backend (header-only; FFI resolved via dlsym) + src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp + src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h # IDAKLU expressions - concrete implementations ${IDAKLU_EXPR_CASADI_SOURCE_FILES} ) diff --git a/packages/pybammsolvers/pyproject.toml b/packages/pybammsolvers/pyproject.toml index 0857991235..747b88dbe7 100644 --- a/packages/pybammsolvers/pyproject.toml +++ b/packages/pybammsolvers/pyproject.toml @@ -47,6 +47,7 @@ wheel.packages = ["src/pybammsolvers"] wheel.exclude = ["pybammsolvers/idaklu.cpp", "pybammsolvers/idaklu_source/**"] editable.mode = "redirect" editable.rebuild = true +editable.verbose = false # SUNDIALS/SuiteSparse are NOT vendored (would add ~400MB); source installs # need either system packages or `git submodule update --init`. diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp index 1440d4d2d6..8db146bf4f 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp @@ -15,6 +15,7 @@ #include "idaklu_source/IdakluJax.hpp" #include "idaklu_source/common.hpp" #include "idaklu_source/Expressions/Casadi/CasadiFunctions.hpp" +#include "idaklu_source/Expressions/Rust/RustFunctions.hpp" #include "idaklu_source/sundials_error_handler.hpp" #include "idaklu_source/reduce.hpp" #include "idaklu_source/StandaloneNewtonSolver.hpp" @@ -119,8 +120,97 @@ IDAKLUSolverGroup *create_casadi_solver_group( ); } + // No owner: CasADi's expression sets hold their own functions. return new IDAKLUSolverGroup( - std::move(solvers), number_of_states, number_of_parameters); + std::move(solvers), number_of_states, number_of_parameters, py::none()); +} + +IDAKLUSolverGroup *create_rust_solver_group( + py::object rust_evaluators, + int number_of_states, + int number_of_inputs, + int number_of_sens_params, + int number_of_events, + int number_of_algebraic_states, + const std::vector &output_lens, + const np_array_int &jac_colptrs, + const np_array_int &jac_rowvals, + int jac_nnz, + const np_array_int &alg_jac_rowvals, + const np_array_int &alg_jac_colvals, + int alg_jac_nnz, + np_array rhs_alg_id, + np_array atol_np, + double rel_tol, + py::dict py_opts +) { + auto setup_opts = SetupOptions(py_opts); + auto solver_opts = SolverOptions(py_opts); + + // Resolve the dlsym table here, before any parallel region. C++ magic + // statics are thread-safe but a *throwing* initialiser is retried by the + // next thread, so a missing symbol would raise once per solver thread. + rust_ffi(); + + const int num_solvers = setup_opts.num_solvers; + const int pool_size = static_cast(py::len(rust_evaluators)); + if (pool_size != num_solvers) { + throw std::invalid_argument( + "rust_evaluators holds " + std::to_string(pool_size) + + " evaluator(s) but num_solvers is " + std::to_string(num_solvers) + + "; each solver needs its own evaluator."); + } + + auto colptrs_ptr = jac_colptrs.unchecked<1>(); + auto rowvals_ptr = jac_rowvals.unchecked<1>(); + std::vector colptrs(colptrs_ptr.data(0), + colptrs_ptr.data(0) + colptrs_ptr.size()); + std::vector rowvals(rowvals_ptr.data(0), + rowvals_ptr.data(0) + rowvals_ptr.size()); + auto alg_rowvals_ptr = alg_jac_rowvals.unchecked<1>(); + auto alg_colvals_ptr = alg_jac_colvals.unchecked<1>(); + std::vector alg_rowvals( + alg_rowvals_ptr.data(0), + alg_rowvals_ptr.data(0) + alg_rowvals_ptr.size() + ); + std::vector alg_colvals( + alg_colvals_ptr.data(0), + alg_colvals_ptr.data(0) + alg_colvals_ptr.size() + ); + + std::vector> solvers; + solvers.reserve(num_solvers); + for (int i = 0; i < num_solvers; i++) { + // One evaluator per solver: they share the compiled tape through Rust's + // Arc and differ only in the scratch each one writes through. + uintptr_t ptr_int = rust_evaluators.attr("as_ptr")(i).cast(); + void* model_ptr = reinterpret_cast(ptr_int); + + auto functions = std::make_unique( + model_ptr, number_of_states, number_of_inputs, number_of_sens_params, + number_of_algebraic_states, number_of_events, jac_nnz, colptrs, rowvals, + alg_jac_nnz, alg_rowvals, alg_colvals, output_lens, setup_opts + ); + + solvers.emplace_back( + std::unique_ptr( + create_idaklu_solver( + std::move(functions), + number_of_sens_params, + jac_colptrs, jac_rowvals, jac_nnz, + 0, 0, + number_of_events, + rhs_alg_id, atol_np, rel_tol, + number_of_inputs, + solver_opts, setup_opts + ) + ) + ); + } + + return new IDAKLUSolverGroup( + std::move(solvers), number_of_states, number_of_sens_params, + std::move(rust_evaluators)); } PYBIND11_MAKE_OPAQUE(std::vector); @@ -143,6 +233,7 @@ PYBIND11_MODULE(idaklu, m) py::arg("y0"), py::arg("yp0"), py::arg("inputs"), + py::arg("pbar") = np_array(), py::arg("logger") = py::none(), py::return_value_policy::take_ownership); @@ -174,6 +265,27 @@ PYBIND11_MODULE(idaklu, m) py::arg("alg_jac"), py::return_value_policy::take_ownership); + m.def("create_rust_solver_group", &create_rust_solver_group, + "Create a Rust-backed IDAKLU solver group", + py::arg("rust_evaluators"), + py::arg("number_of_states"), + py::arg("number_of_inputs"), + py::arg("number_of_sens_params"), + py::arg("number_of_events"), + py::arg("n_alg"), + py::arg("output_lens"), + py::arg("jac_colptrs"), + py::arg("jac_rowvals"), + py::arg("jac_nnz"), + py::arg("alg_jac_rowvals"), + py::arg("alg_jac_colvals"), + py::arg("alg_jac_nnz"), + py::arg("rhs_alg_id"), + py::arg("atol"), + py::arg("rtol"), + py::arg("options"), + py::return_value_policy::take_ownership); + m.def("observe", &observe, "Observe variables", py::arg("ts"), @@ -263,6 +375,15 @@ PYBIND11_MODULE(idaklu, m) py::arg("atol"), py::arg("rtol"), py::arg("step_tol"), py::arg("max_iter"), py::arg("max_backtracks"), py::arg("eps_newt"), py::arg("use_sparse")) + .def(py::init, std::vector, + std::vector, sunrealtype, sunrealtype, + int, int, sunrealtype, bool>(), + py::arg("rust_model"), py::arg("n_rhs"), py::arg("n_alg"), + py::arg("jac_rows"), py::arg("jac_cols"), + py::arg("atol"), py::arg("rtol"), py::arg("step_tol"), + py::arg("max_iter"), py::arg("max_backtracks"), + py::arg("eps_newt"), py::arg("use_sparse")) .def("solve", &StandaloneNewtonSolver::solve, py::arg("t"), py::arg("y0"), py::arg("inputs"), py::return_value_policy::move) @@ -272,6 +393,17 @@ PYBIND11_MODULE(idaklu, m) py::class_(m, "Function"); + py::class_(m, "SolverStats") + .def_readonly("number_of_steps", &IDAKLUStats::nsteps) + .def_readonly("number_of_residual_evaluations", &IDAKLUStats::nrevals) + .def_readonly("number_of_linear_solver_setups", &IDAKLUStats::nlinsetups) + .def_readonly("number_of_error_test_failures", &IDAKLUStats::netfails) + .def_readonly("number_of_nonlinear_solver_iterations", &IDAKLUStats::nniters) + .def_readonly("number_of_nonlinear_solver_fails", &IDAKLUStats::nncfails) + .def_readonly("number_of_jacobian_evaluations", &IDAKLUStats::njevals) + .def_readonly("number_of_linear_iterations", &IDAKLUStats::nliters) + .def_readonly("number_of_linear_convergence_failures", &IDAKLUStats::nlcfails); + py::class_(m, "solution") .def_readwrite("t", &Solution::t) .def_readwrite("y", &Solution::y) @@ -279,6 +411,7 @@ PYBIND11_MODULE(idaklu, m) .def_readwrite("yS", &Solution::yS) .def_readwrite("ypS", &Solution::ypS) .def_readwrite("y_term", &Solution::y_term) - .def_readwrite("flag", &Solution::flag); + .def_readwrite("flag", &Solution::flag) + .def_readonly("stats", &Solution::stats); } diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp index 6b87e78ec6..a552cb75d7 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp @@ -12,6 +12,11 @@ class Expression { */ Expression() = default; + /** + * @brief Virtual destructor (required for safe polymorphic deletion via base pointer) + */ + virtual ~Expression() = default; + /** * @brief Evaluation operator (for use after setting input and output data references) */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp index eabc59012a..44d27e75fe 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp @@ -10,6 +10,14 @@ class ExpressionSet { public: + // Compile-time marker: does this ExprSet project output-variable + // sensitivities natively? Default false; RustFunctions overrides to true. + static constexpr bool kNativeOutputSensitivities = false; + + // Compile-time marker: does this ExprSet evaluate output variables over a + // batch of points in one call? Default false; RustFunctions overrides. + static constexpr bool kNativeBatchedOutputs = false; + /** * @brief Constructor */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp new file mode 100644 index 0000000000..107c815360 --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp @@ -0,0 +1,599 @@ +#ifndef PYBAMM_IDAKLU_RUST_FUNCTIONS_HPP +#define PYBAMM_IDAKLU_RUST_FUNCTIONS_HPP + +#include "../Base/Expression.hpp" +#include "../Base/ExpressionSet.hpp" +#include "../../Options.hpp" +#include +#include +#include +#include +#include + +// Rust FFI declarations - single source of truth +#include "pybamm_rust_ffi.h" + +/** + * @brief Common base for the Rust FFI Expression adapters. + * + * Every adapter writes a single fixed-length output block and differs only in + * which FFI entry point it forwards to, so the shape and sparsity interface + * lives here. `m_rows`/`m_cols` stay empty unless the adapter emits COO data, + * matching the CasADi adapters' contract for dense-vector outputs. + */ +class RustExpression : public Expression { +public: + /** + * @param model Opaque Rust model handle + * @param out_len Elements written per evaluation (nnz for COO adapters) + * @param n_args m_arg slots the FFI signature uses + * @param n_res m_res slots the FFI signature writes + */ + RustExpression(void* model, int out_len, int n_args = 3, int n_res = 1) + : m_model(model), m_out_len(out_len) { + m_arg.resize(n_args); + m_res.resize(n_res); + } + + expr_int out_shape(int k) override { return m_out_len; } + expr_int nnz() override { return m_out_len; } + expr_int nnz_out() override { return m_out_len; } + const std::vector& get_row() override { return m_rows; } + const std::vector& get_col() override { return m_cols; } + +protected: + void* m_model; + int m_out_len; + std::vector m_rows; + std::vector m_cols; +}; + +/** + * @brief RustRhsExpression: evaluates the RHS function via Rust FFI + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs, m_res[0]=output + */ +class RustRhsExpression : public RustExpression { +public: + RustRhsExpression(void* model, int n_states) + : RustExpression(model, n_states) {} + + void operator()() override { + PYBAMM_RUST_CALL(eval_rhs, *m_arg[0], m_arg[1], m_arg[2], m_res[0], m_model); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + PYBAMM_RUST_CALL(eval_rhs, *inputs[0], inputs[1], inputs[2], results[0], m_model); + } +}; + +/** + * @brief RustJacExpression: assembles Jacobian matrix via Rust FFI + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs, m_arg[3]=cj, m_res[0]=jac_data + */ +class RustJacExpression : public RustExpression { +public: + RustJacExpression(void* model, int nnz, + std::vector rows, + std::vector cols) + : RustExpression(model, nnz, 4) { + m_rows = std::move(rows); + m_cols = std::move(cols); + } + + void operator()() override { + PYBAMM_RUST_CALL(jac_assemble, *m_arg[0], m_arg[1], m_arg[2], *m_arg[3], m_res[0], m_model); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + // cj is injected by the SUNDIALS callback via m_arg[3] and is not part + // of the inputs vector, so the cj read stays as *m_arg[3] even here. + PYBAMM_RUST_CALL(jac_assemble, *inputs[0], inputs[1], inputs[2], *m_arg[3], results[0], m_model); + } +}; + +/** + * @brief RustJacActionExpression: computes Jacobian-vector product via Rust FFI + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs, m_arg[3]=v, m_res[0]=Jv + */ +class RustJacActionExpression : public RustExpression { +public: + RustJacActionExpression(void* model, int n_states) + : RustExpression(model, n_states, 4) {} + + void operator()() override { + PYBAMM_RUST_CALL(jac_action, *m_arg[0], m_arg[1], m_arg[2], m_arg[3], m_res[0], m_model); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + PYBAMM_RUST_CALL(jac_action, *inputs[0], inputs[1], inputs[2], inputs[3], results[0], m_model); + } +}; + +/** + * @brief RustMassActionExpression: computes mass-matrix-vector product via Rust FFI + * + * ABI: m_arg[0]=v, m_res[0]=Mv + */ +class RustMassActionExpression : public RustExpression { +public: + RustMassActionExpression(void* model, int n_states) + : RustExpression(model, n_states, 1) {} + + void operator()() override { + PYBAMM_RUST_CALL(mass_action, m_arg[0], m_res[0], m_model); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + PYBAMM_RUST_CALL(mass_action, inputs[0], results[0], m_model); + } +}; + +/** + * @brief RustEventsExpression: event evaluator backed by Rust FFI. + * + * Evaluates all events and writes concatenated results to output. + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs, m_res[0]=output (length total_event_len) + */ +class RustEventsExpression : public RustExpression { +public: + /// `total_event_len == 0` (no events) makes this an inert placeholder. + RustEventsExpression(void* model, int total_event_len) + : RustExpression(model, total_event_len) {} + + void operator()() override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(events_eval, *m_arg[0], m_arg[1], m_arg[2], m_res[0], m_model); + } + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(events_eval, *inputs[0], inputs[1], inputs[2], results[0], m_model); + } + } +}; + +/** + * @brief RustSensExpression: forward-sensitivity evaluator backed by Rust FFI. + * + * Computes ∂f/∂p_i for every configured sensitivity parameter and writes the + * result into `m_res[i]`. Constructed with `n_sens_params == 0` it acts as + * an empty placeholder, matching the previous behaviour for non-sensitivity + * solves. + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs; m_res[i]=∂f/∂p_i (length n_states) + */ +class RustSensExpression : public RustExpression { +public: + /// One result slot per sensitivity parameter; sundials_functions.inl + /// repoints them at the resvalS N_Vectors before each call. + RustSensExpression(void* model, int n_states, int n_sens_params) + : RustExpression(model, n_states, 3, n_sens_params), + m_n_sens_params(n_sens_params), + m_columns(static_cast(n_states) * n_sens_params) {} + + void operator()() override { + eval_columns(m_arg[0], m_arg[1], m_arg[2], m_res); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + eval_columns(inputs[0], inputs[1], inputs[2], results); + } + + // out_shape stays n_states (the per-parameter block), but a model without + // sensitivities writes nothing at all. + expr_int nnz() override { return m_n_sens_params == 0 ? 0 : m_out_len; } + expr_int nnz_out() override { return nnz(); } + +private: + /* + * `sens_eval_all` runs the tape's shared primal section once and then one + * tangent-only sweep per parameter, where per-parameter `sens_eval` repeats + * the primal every time (58% of the DFN tangent tape). It writes the + * columns contiguously, so they are scattered out to the per-parameter + * SUNDIALS buffers, which are separate N_Vectors. + */ + void eval_columns(const sunrealtype* t, const sunrealtype* y, + const sunrealtype* inputs, + const std::vector& results) { + if (m_n_sens_params == 0) { + return; + } + PYBAMM_RUST_CALL(sens_eval_all, *t, y, inputs, m_columns.data(), m_model); + const size_t n_states = static_cast(m_out_len); + for (int i = 0; i < m_n_sens_params; i++) { + std::memcpy(results[i], m_columns.data() + i * n_states, + n_states * sizeof(sunrealtype)); + } + } + + int m_n_sens_params; + std::vector m_columns; +}; + +/** + * @brief RustOutputExpression: evaluates a single output variable via Rust FFI. + * + * IDAKLUSolverOpenMP iterates `functions->var_fcns` whenever output variables + * are saved (save_outputs_only path), invoking the parameterized + * `operator()(inputs, results)` overload with `inputs = {&t, y, inputs_data}`. + * We forward straight into `rust_output_eval` for the configured `var_idx`. + * + * ABI: m_arg[0]=t, m_arg[1]=y, m_arg[2]=inputs; m_res[0]=output (length m_out_len) + */ +class RustOutputExpression : public RustExpression { +public: + RustOutputExpression(void* model, int var_idx, int out_len) + : RustExpression(model, out_len), m_var_idx(var_idx) {} + + void operator()() override { + int written = 0; + PYBAMM_RUST_CALL(output_eval, *m_arg[0], m_arg[1], m_arg[2], m_var_idx, + m_res[0], &written, m_model); + check_written(written); + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + int written = 0; + PYBAMM_RUST_CALL(output_eval, *inputs[0], inputs[1], inputs[2], m_var_idx, + results[0], &written, m_model); + check_written(written); + } + +private: + /* The caller sized the buffer from `m_out_len`, so a short write leaves the + tail holding the previous step's values. */ + void check_written(int written) const { + if (written != m_out_len) { + throw std::runtime_error( + std::string("pybammsolvers: Rust output variable ") + + std::to_string(m_var_idx) + " wrote " + std::to_string(written) + + " elements, expected " + std::to_string(m_out_len) + "."); + } + } + + int m_var_idx; +}; + +/** + * @brief RustAlgResExpression: evaluates algebraic residuals via Rust FFI. + * + * Constructed with `n_alg == 0` it behaves like an empty placeholder so the + * solver naturally falls back to the full-system IC path. + */ +class RustAlgResExpression : public RustExpression { +public: + RustAlgResExpression(void* model, int n_alg) + : RustExpression(model, n_alg) {} + + void operator()() override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(alg_res, *m_arg[0], m_arg[1], m_arg[2], m_res[0], m_model); + } + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(alg_res, *inputs[0], inputs[1], inputs[2], results[0], m_model); + } + } +}; + +/** + * @brief RustAlgJacExpression: assembles algebraic Jacobians via Rust FFI. + * + * The output ordering follows the COO `(row, col)` metadata passed in at + * construction time. With `m_nnz == 0` it acts as an empty placeholder. + */ +class RustAlgJacExpression : public RustExpression { +public: + RustAlgJacExpression(void* model, int nnz, + std::vector rows, + std::vector cols) + : RustExpression(model, nnz) { + m_rows = std::move(rows); + m_cols = std::move(cols); + } + + void operator()() override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(alg_jac_assemble, *m_arg[0], m_arg[1], m_arg[2], m_res[0], m_model); + } + } + + void operator()(const std::vector& inputs, + const std::vector& results) override { + if (m_out_len > 0) { + PYBAMM_RUST_CALL(alg_jac_assemble, *inputs[0], inputs[1], inputs[2], results[0], m_model); + } + } +}; + +/** + * @brief Shared base for the standalone algebraic Newton solve adapters. + * + * The Newton solver calls with the convention `F(t, y_alg, [y_diff; inputs])`: + * m_arg[0]=&t, m_arg[1]=y_alg, m_arg[2]=[y_diff; inputs]. The Rust FFI instead + * expects the full state `g(t, y_full, inputs)`, so both adapters gather + * y_diff and y_alg into a full-state buffer and forward the inputs slice that + * follows the y_diff block. + */ +class RustNewtonExpression : public RustExpression { +public: + RustNewtonExpression(void* model, int out_len, int n_rhs, int n_alg) + : RustExpression(model, out_len), m_n_rhs(n_rhs), m_n_alg(n_alg), + m_y_full(static_cast(n_rhs) + n_alg, 0.0) {} + + // Declaring the two-argument overload here would otherwise hide the + // zero-argument one that the subclasses define and this one delegates to. + using Expression::operator(); + + void operator()(const std::vector& inputs, + const std::vector& results) override { + m_arg.assign(inputs.begin(), inputs.end()); + m_res.assign(results.begin(), results.end()); + (*this)(); + } + +protected: + /// Gather [y_diff; y_alg] into m_y_full; returns the trailing inputs slice. + const double* gather_full_state() { + std::memcpy(m_y_full.data(), m_arg[2], m_n_rhs * sizeof(double)); // y_diff + std::memcpy(m_y_full.data() + m_n_rhs, m_arg[1], m_n_alg * sizeof(double)); // y_alg + return m_arg[2] + m_n_rhs; + } + + int m_n_rhs; + int m_n_alg; + std::vector m_y_full; +}; + +/** + * @brief RustNewtonResExpression: Newton residual adapter. + */ +class RustNewtonResExpression : public RustNewtonExpression { +public: + RustNewtonResExpression(void* model, int n_rhs, int n_alg) + : RustNewtonExpression(model, n_alg, n_rhs, n_alg) {} + + void operator()() override { + const double* inputs = gather_full_state(); + PYBAMM_RUST_CALL(alg_res, *m_arg[0], m_y_full.data(), inputs, m_res[0], m_model); + } +}; + +/** + * @brief RustNewtonJacExpression: Newton Jacobian adapter. + * + * `alg_jac_assemble` fills the output in the COO order of the (rows, cols) + * sparsity passed at construction, so BuildSparseResources consumes it + * unchanged. + */ +class RustNewtonJacExpression : public RustNewtonExpression { +public: + RustNewtonJacExpression(void* model, int n_rhs, int n_alg, int nnz, + std::vector rows, + std::vector cols) + : RustNewtonExpression(model, nnz, n_rhs, n_alg) { + m_rows = std::move(rows); + m_cols = std::move(cols); + } + + void operator()() override { + const double* inputs = gather_full_state(); + PYBAMM_RUST_CALL(alg_jac_assemble, *m_arg[0], m_y_full.data(), inputs, m_res[0], m_model); + } +}; + +/** + * @brief RustFunctions: ExpressionSet implementation for Rust-based models + * + * Wraps Rust FFI functions in the Expression interface expected by IDAKLU solver. + * + * Uses direct members (not unique_ptr) to ensure valid pointers can be passed + * to the base class constructor. In C++, base classes are initialized before + * members, but taking &member gives a valid address that will be populated + * by the time it's dereferenced during actual solving. + */ +class RustFunctions : public ExpressionSet { +public: + /** + * @brief Construct RustFunctions from a Rust model handle + * + * @param rust_model Opaque pointer to Rust Model + * @param n_states Number of state variables + * @param n_inputs Number of input parameters + * @param n_sens_params Number of forward-sensitivity parameters + * @param n_alg Number of algebraic states + * @param n_events Number of event functions + * @param nnz Number of non-zeros in Jacobian + * @param colptrs CSC column pointers + * @param rowvals CSC row indices + * @param alg_jac_nnz Number of non-zeros in the algebraic Jacobian + * @param alg_rowvals Algebraic Jacobian COO row indices + * @param alg_colvals Algebraic Jacobian COO column indices + * @param output_lens Output length of each output variable + * @param options Solver setup options + */ + RustFunctions( + void* rust_model, + int n_states, + int n_inputs, + int n_sens_params, + int n_alg, + int n_events, + int nnz, + const std::vector& colptrs, + const std::vector& rowvals, + int alg_jac_nnz, + const std::vector& alg_rowvals, + const std::vector& alg_colvals, + const std::vector& output_lens, + const SetupOptions& options + ) : + tmp_state_vector(n_states), + tmp_sparse_jacobian_data(nnz), + // Expression members - use direct members, not unique_ptr. Each owns + // its COO index vectors, so the conversions are passed as temporaries. + m_rhs(rust_model, n_states), + m_jac(rust_model, nnz, + convert_to_expr_int(rowvals), + csc_rowvals_to_coo_cols(colptrs, rowvals, nnz)), + m_jac_action(rust_model, n_states), + m_mass_action(rust_model, n_states), + m_events(rust_model, + n_events > 0 ? PYBAMM_RUST_VALUE(total_event_len, rust_model) : 0), + m_sens(rust_model, n_states, n_sens_params), + m_alg_res(rust_model, n_alg), + m_alg_jac(rust_model, alg_jac_nnz, + convert_to_expr_int(alg_rowvals), + convert_to_expr_int(alg_colvals)), + // Base class - addresses of direct members are valid even before init + ExpressionSet( + static_cast(&m_rhs), + static_cast(&m_jac), + nnz, + 0, // jac_bandwidth_lower (not used for sparse) + 0, // jac_bandwidth_upper (not used for sparse) + np_array_int(), // empty, we store directly + np_array_int(), // empty, we store directly + n_inputs, // inputs_length + static_cast(&m_jac_action), + static_cast(&m_mass_action), + static_cast(&m_sens), + static_cast(&m_events), + n_states, + n_events, // number of events + n_sens_params, // n_parameters (forward sensitivities) + options, + static_cast(&m_alg_res), + static_cast(&m_alg_jac) + ) + { + // Retain the Rust model handle for native output-sensitivity projection. + m_model = rust_model; + // Store sparsity pattern in base class members + jac_times_cjmass_colptrs = colptrs; + jac_times_cjmass_rowvals = rowvals; + // Allocate the inputs vector that the SUNDIALS callback layer + // populates each step (sundials_functions.inl sets m_arg[2] to + // inputs.data()). The base ExpressionSet ctor takes inputs_length + // but doesn't size the vector itself; we mirror what CasadiFunctions + // does in its own ctor body. + inputs.resize(n_inputs); + + // Construct one RustOutputExpression per requested output variable. + // unique_ptr keeps the Expression heap-stable across vector growth so + // the raw pointers we hand to ExpressionSet::var_fcns stay valid. + m_output_fcns.reserve(output_lens.size()); + for (int i = 0; i < static_cast(output_lens.size()); ++i) { + m_output_fcns.emplace_back( + std::make_unique( + rust_model, i, output_lens[i] + ) + ); + var_fcns.push_back(static_cast(m_output_fcns.back().get())); + } + } + + sunrealtype* get_tmp_state_vector() override { + return tmp_state_vector.data(); + } + + sunrealtype* get_tmp_sparse_jacobian_data() override { + return tmp_sparse_jacobian_data.data(); + } + + // Compile-time marker: this ExprSet projects output sensitivities natively, + // so the solver takes the native branch instead of the CasADi sparse loop. + static constexpr bool kNativeOutputSensitivities = true; + + // Compile-time marker: output variables can be evaluated over a batch of + // trajectory points in one call, amortising interpreter dispatch. + static constexpr bool kNativeBatchedOutputs = true; + + // Batch-evaluate every output variable over k staged points. + // ts: [k]; ys: [k * n_states], each point contiguous; + // out: [k * total_output_len], each point's stacked outputs contiguous. + void eval_outputs_batch( + const sunrealtype* ts, const sunrealtype* ys, int k, sunrealtype* out) { + PYBAMM_RUST_CALL(output_eval_batch, ts, ys, k, inputs.data(), out, m_model); + } + + // Project state sensitivities onto output-variable sensitivities via Rust FFI. + // yS_flat: [n_sens_params * n_states]; out: [n_sens_params * total_output_len]. + void project_output_sensitivities( + double t, sunrealtype* y, sunrealtype* inputs, + const sunrealtype* yS_flat, sunrealtype* out) { + PYBAMM_RUST_CALL(output_sens_project, t, y, inputs, yS_flat, out, m_model); + } + +private: + // Rust model handle, retained for native output-sensitivity projection. + void* m_model = nullptr; + + std::vector tmp_state_vector; + std::vector tmp_sparse_jacobian_data; + + // Expression members - direct members (not unique_ptr) so that addresses + // can be safely passed to base class constructor + RustRhsExpression m_rhs; + RustJacExpression m_jac; + RustJacActionExpression m_jac_action; + RustMassActionExpression m_mass_action; + RustEventsExpression m_events; + RustSensExpression m_sens; + RustAlgResExpression m_alg_res; + RustAlgJacExpression m_alg_jac; + // Output-variable expressions, one per configured output. Stored via + // unique_ptr so the Expression* references in `var_fcns` (base class) + // remain valid even if the vector reallocates. + std::vector> m_output_fcns; + + /** + * @brief Convert int64_t vector to expr_int vector + */ + static std::vector convert_to_expr_int(const std::vector& v) { + std::vector result; + result.reserve(v.size()); + for (const auto& val : v) { + result.push_back(static_cast(val)); + } + return result; + } + + /** + * @brief Convert CSC rowvals to COO column indices + * + * In CSC format, colptrs[j] to colptrs[j+1] gives the range of entries in column j. + * For COO format, we need the column index for each entry. + */ + static std::vector csc_rowvals_to_coo_cols( + const std::vector& colptrs, + const std::vector& rowvals, + int nnz + ) { + std::vector cols(nnz); + int n_cols = static_cast(colptrs.size()) - 1; + for (int col = 0; col < n_cols; ++col) { + for (int64_t k = colptrs[col]; k < colptrs[col + 1]; ++k) { + cols[k] = static_cast(col); + } + } + return cols; + } +}; + +#endif // PYBAMM_IDAKLU_RUST_FUNCTIONS_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h new file mode 100644 index 0000000000..042c45827b --- /dev/null +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h @@ -0,0 +1,264 @@ +#ifndef PYBAMM_RUST_FFI_H +#define PYBAMM_RUST_FFI_H + +#include + +/* Return codes */ +#define PYBAMM_SUCCESS 0 +#define PYBAMM_ERROR_NULL -1 +#define PYBAMM_ERROR_PANIC -2 +#define PYBAMM_ERROR_INVALID_PARAM -3 +#define PYBAMM_ERROR_INVALID_OUTPUT -4 +#define PYBAMM_ERROR_BUFFER_SMALL -5 +#define PYBAMM_ERROR_NO_SENS -6 +#define PYBAMM_ERROR_NO_OUTPUTS -7 +#define PYBAMM_ERROR_NO_ALG -8 +#define PYBAMM_ERROR_NO_EVENTS -9 + +/* ABI contract version. Must equal the Rust core's RUST_ABI_VERSION; the + * rust_ffi() resolver throws on mismatch. Bump in lockstep with ffi.rs. */ +#define PYBAMM_RUST_ABI_VERSION 1 + +#ifdef __cplusplus + +#if defined(_WIN32) +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#ifndef NOMINMAX +#define NOMINMAX +#endif +// PSAPI_VERSION 2 binds EnumProcessModules to its kernel32 export +// (K32EnumProcessModules), so no psapi.lib import library is needed. +#ifndef PSAPI_VERSION +#define PSAPI_VERSION 2 +#endif +#include +#include +#else +#include +#endif + +#include +#include +#include + +/* + * Rust FFI access — runtime symbol resolution (host-plugin model). + * + * pybammsolvers cannot link against PyBaMM's Rust core: PyBaMM depends on + * pybammsolvers, so the dependency only flows one way, and the Rust FFI entry + * points are compiled into PyBaMM's `pybamm.rust._core` Python extension rather + * than a standalone shared library. We therefore do NOT declare these functions as + * `extern "C"` (which would leave undefined symbols in idaklu and break loading + * the module unless the extension happened to be loaded first). Instead we + * resolve them at runtime via `find_symbol` below (`dlsym(RTLD_DEFAULT, ...)` + * on POSIX, a loaded-module walk on Windows). + * + * By the time a Rust-backed solver group is constructed, `pybamm.rust._core` is + * already imported in-process (the caller hands us a live CompiledModel), so + * its exported symbols are visible to the dynamic linker. This keeps idaklu + * free of undefined symbols — it loads cleanly standalone (the CasADi path is + * unaffected) and ships in release wheels, with the Rust path resolved lazily + * the first time it is used. + * + * Note (Linux): symbols are only visible to `dlsym(RTLD_DEFAULT, ...)` if the + * extension was loaded with global visibility. CPython defaults to RTLD_LOCAL + * on Linux, so `pybamm/rust/__init__.py` adds RTLD_GLOBAL around the import + * there (`_global_symbol_visibility`). On macOS the symbols are globally + * visible by default. + * + * Note (Windows): there is no process-global symbol namespace, so + * `find_symbol` walks the loaded-module list with `GetProcAddress` instead. + * The Rust extension is a cdylib, whose `#[no_mangle]` entry points sit in + * its DLL export table, so the walk finds them once the module is loaded. + */ + +/* + * Function-pointer typedefs for the Rust `extern "C"` entry points IDAKLU + * calls. This is deliberately a subset of `ffi.rs`, not a mirror of it: every + * name here is resolved eagerly and a missing symbol aborts the whole Rust + * path, so binding an entry point nothing calls would couple IDAKLU to Rust + * exports it does not need. Model metadata (state/param/output counts, + * sparsity, algebraic ids) reaches C++ as explicit `create_rust_solver_group` + * arguments, read on the Python side through pyo3. + */ +using rust_eval_rhs_t = int (*)(double, const double*, const double*, double*, void*); +using rust_jac_assemble_t = int (*)(double, const double*, const double*, double, double*, void*); +using rust_jac_action_t = int (*)(double, const double*, const double*, const double*, double*, void*); +using rust_mass_action_t = int (*)(const double*, double*, void*); +using rust_sens_eval_all_t = int (*)(double, const double*, const double*, double*, void*); +using rust_output_sens_project_t = + int (*)(double, const double*, const double*, const double*, double*, void*); +using rust_output_eval_t = int (*)(double, const double*, const double*, int, double*, int*, void*); +using rust_output_eval_batch_t = + int (*)(const double*, const double*, int, const double*, double*, void*); +using rust_alg_res_t = int (*)(double, const double*, const double*, double*, void*); +using rust_alg_jac_assemble_t = int (*)(double, const double*, const double*, double*, void*); +using rust_total_event_len_t = int (*)(const void*); +using rust_events_eval_t = int (*)(double, const double*, const double*, double*, void*); +using rust_abi_version_t = uint32_t (*)(); + +/* Resolved table of Rust FFI entry points. */ +struct RustFfi { + rust_eval_rhs_t eval_rhs; + rust_jac_assemble_t jac_assemble; + rust_jac_action_t jac_action; + rust_mass_action_t mass_action; + rust_sens_eval_all_t sens_eval_all; + rust_output_sens_project_t output_sens_project; + rust_output_eval_t output_eval; + rust_output_eval_batch_t output_eval_batch; + rust_alg_res_t alg_res; + rust_alg_jac_assemble_t alg_jac_assemble; + rust_total_event_len_t total_event_len; + rust_events_eval_t events_eval; + rust_abi_version_t abi_version; +}; + +namespace pybamm_rust_detail { + +/* Address of an exported symbol in any module loaded in this process, or + * nullptr. POSIX has this as `dlsym(RTLD_DEFAULT, ...)`; Windows has no + * process-global lookup, so every loaded module's export table is tried in + * turn (`K32EnumProcessModules` is the kernel32 alias of the psapi function, + * so no extra import library is linked). */ +inline void* find_symbol(const char* name) { +#if defined(_WIN32) + HANDLE process = GetCurrentProcess(); + DWORD bytes_needed = 0; + if (K32EnumProcessModules(process, nullptr, 0, &bytes_needed) == 0) { + return nullptr; + } + std::vector modules(bytes_needed / sizeof(HMODULE)); + if (K32EnumProcessModules(process, modules.data(), bytes_needed, &bytes_needed) == 0) { + return nullptr; + } + // The module list can shrink between the size query and the fill. + modules.resize(bytes_needed / sizeof(HMODULE)); + for (HMODULE module : modules) { + if (FARPROC sym = GetProcAddress(module, name)) { + return reinterpret_cast(sym); + } + } + return nullptr; +#else + return dlsym(RTLD_DEFAULT, name); +#endif +} + +template +inline Fn load_symbol(const char* name) { + void* sym = find_symbol(name); + if (sym == nullptr) { + throw std::runtime_error( + std::string("pybammsolvers: Rust FFI symbol '") + name + + "' could not be resolved. Import `pybamm.rust` before constructing a " + "Rust-backed IDAKLU solver: it loads pybamm.rust._core and, on Linux, " + "adds RTLD_GLOBAL so these symbols reach dlsym(RTLD_DEFAULT)."); + } + return reinterpret_cast(sym); +} + +/* Spelling of a Rust FFI status code, for error messages. */ +inline const char* status_name(int status) { + switch (status) { + case PYBAMM_SUCCESS: return "SUCCESS"; + case PYBAMM_ERROR_NULL: return "ERROR_NULL"; + case PYBAMM_ERROR_PANIC: return "ERROR_PANIC"; + case PYBAMM_ERROR_INVALID_PARAM: return "ERROR_INVALID_PARAM"; + case PYBAMM_ERROR_INVALID_OUTPUT: return "ERROR_INVALID_OUTPUT"; + case PYBAMM_ERROR_BUFFER_SMALL: return "ERROR_BUFFER_SMALL"; + case PYBAMM_ERROR_NO_SENS: return "ERROR_NO_SENS"; + case PYBAMM_ERROR_NO_OUTPUTS: return "ERROR_NO_OUTPUTS"; + case PYBAMM_ERROR_NO_ALG: return "ERROR_NO_ALG"; + case PYBAMM_ERROR_NO_EVENTS: return "ERROR_NO_EVENTS"; + default: return "unrecognised status"; + } +} + +[[noreturn]] inline void throw_ffi_error(const char* name, int status) { + throw std::runtime_error( + std::string("pybammsolvers: Rust FFI call '") + name + "' failed with " + + status_name(status) + " (" + std::to_string(status) + + "). The output buffer may be unwritten, so the evaluation cannot be trusted."); +} + +/* + * Call a status-returning entry point, throwing unless it reports success. + * + * A panic caught at the Rust boundary yields ERROR_PANIC with the output buffer + * potentially unwritten, so dropping the status would feed the previous step's + * stale values to SUNDIALS as a valid evaluation. Throwing matches the CasADi + * path, whose evaluation errors reach IDAKLUSolverGroup's handler the same way. + */ +template +inline void checked_call(const char* name, Fn fn, Args... args) { + const int status = fn(args...); + if (status != PYBAMM_SUCCESS) { + throw_ffi_error(name, status); + } +} + +/* As above for entry points returning a non-negative count rather than a status. */ +template +inline int checked_value(const char* name, Fn fn, Args... args) { + const int value = fn(args...); + if (value < 0) { + throw_ffi_error(name, value); + } + return value; +} + +} // namespace pybamm_rust_detail + +/* Checked call to `RustFfi::entry`; see `pybamm_rust_detail::checked_call`. */ +#define PYBAMM_RUST_CALL(entry, ...) \ + ::pybamm_rust_detail::checked_call(#entry, ::rust_ffi().entry, __VA_ARGS__) + +/* Checked read of a count-returning `RustFfi::entry`. */ +#define PYBAMM_RUST_VALUE(entry, ...) \ + ::pybamm_rust_detail::checked_value(#entry, ::rust_ffi().entry, __VA_ARGS__) + +/* + * Resolve and cache the Rust FFI table. The first call performs the dlsym + * lookups (thread-safe via the function-local static); subsequent calls return + * the cached table, so evaluation hot paths only pay an indirect call. + */ +inline const RustFfi& rust_ffi() { + using pybamm_rust_detail::load_symbol; + static const RustFfi table = [] { + RustFfi t{}; + t.eval_rhs = load_symbol("pybamm_rust_eval_rhs"); + t.jac_assemble = load_symbol("pybamm_rust_jac_assemble"); + t.jac_action = load_symbol("pybamm_rust_jac_action"); + t.mass_action = load_symbol("pybamm_rust_mass_action"); + t.sens_eval_all = load_symbol("pybamm_rust_sens_eval_all"); + t.output_sens_project = + load_symbol("pybamm_rust_output_sens_project"); + t.output_eval = load_symbol("pybamm_rust_output_eval"); + t.output_eval_batch = + load_symbol("pybamm_rust_output_eval_batch"); + t.alg_res = load_symbol("pybamm_rust_alg_res"); + t.alg_jac_assemble = load_symbol("pybamm_rust_alg_jac_assemble"); + t.total_event_len = load_symbol("pybamm_rust_total_event_len"); + t.events_eval = load_symbol("pybamm_rust_events_eval"); + t.abi_version = load_symbol("pybamm_rust_abi_version"); + if (t.abi_version() != PYBAMM_RUST_ABI_VERSION) { + throw std::runtime_error( + std::string("pybammsolvers: Rust FFI ABI version mismatch. " + "pybammsolvers was built for version ") + + std::to_string(PYBAMM_RUST_ABI_VERSION) + + " but the loaded PyBaMM Rust core reports version " + + std::to_string(t.abi_version()) + + ". Rebuild pybammsolvers and the PyBaMM Rust extension from the " + "same source tree."); + } + return t; + }(); + return table; +} + +#endif /* __cplusplus */ + +#endif /* PYBAMM_RUST_FFI_H */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp index 7c7aaa4380..a11fb4410f 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp @@ -34,6 +34,7 @@ class IDAKLUSolver const sunrealtype *y0, const sunrealtype *yp0, const sunrealtype *inputs, + const sunrealtype *pbar, bool save_adaptive_steps, bool save_interp_steps ) = 0; diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index 2ebc851a9a..46b1038589 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -1,7 +1,8 @@ #include "IDAKLUSolverGroup.hpp" #include +#include #include -#include +#include std::vector IDAKLUSolverGroup::solve( np_array t_eval_np, @@ -9,6 +10,7 @@ std::vector IDAKLUSolverGroup::solve( np_array y0_np, np_array yp0_np, np_array inputs, + np_array pbar, py::object logger) { DEBUG("IDAKLUSolverGroup::solve"); @@ -96,16 +98,31 @@ std::vector IDAKLUSolverGroup::solve( "inputs has wrong number of rows. Expected " + std::to_string(number_of_groups) + " but got " + std::to_string(inputs.shape()[0])); - const std::size_t solves_per_thread = number_of_groups / m_solvers.size(); - const std::size_t remainder_solves = number_of_groups % m_solvers.size(); + // pbar is optional: an empty array leaves IDAS at its unit default. + const bool has_pbar = pbar.size() > 0; + if (has_pbar) { + if (pbar.ndim() != 2) + throw std::domain_error("pbar has wrong number of dimensions. Expected 2 but got " + std::to_string(pbar.ndim())); + if (pbar.shape()[0] != number_of_groups) + throw std::domain_error( + "pbar has wrong number of rows. Expected " + std::to_string(number_of_groups) + + " but got " + std::to_string(pbar.shape()[0])); + if (pbar.shape()[1] != number_of_parameters) + throw std::domain_error( + "pbar has wrong number of cols. Expected " + std::to_string(number_of_parameters) + + " but got " + std::to_string(pbar.shape()[1])); + } const sunrealtype *y0 = y0_np.data(); const sunrealtype *yp0 = yp0_np.data(); const sunrealtype *inputs_data = inputs.data(); + const sunrealtype *pbar_data = has_pbar ? pbar.data() : nullptr; std::vector results(number_of_groups); - std::optional exception_message; + // One slot per input set, so the rethrow below can name which sets failed. + // Distinct indices need no synchronisation; empty means the set succeeded. + std::vector errors(number_of_groups); // Python exceptions carry their own type, which the string path below loses std::exception_ptr python_exception; @@ -113,27 +130,39 @@ std::vector IDAKLUSolverGroup::solve( // may log directly or must buffer until flush_logs(). set_loggers(logger); - omp_set_num_threads(m_solvers.size()); - #pragma omp parallel for - for (int i = 0; i < m_solvers.size(); i++) { - try { - for (int j = 0; j < solves_per_thread; j++) { - const std::size_t index = i * solves_per_thread + j; - const sunrealtype *y = y0 + index * y0_np.shape(1); - const sunrealtype *yp = yp0 + index * yp0_np.shape(1); - const sunrealtype *input = inputs_data + index * inputs.shape(1); - results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, save_adaptive_steps, save_interp_steps); - } - } catch (py::error_already_set &) { - #pragma omp critical - { - python_exception = std::current_exception(); - } - } catch (std::exception &e) { - // If an exception is thrown, we need to catch it and rethrow it outside the parallel region - #pragma omp critical - { - exception_message = std::string(e.what()); + // Scoped to this region rather than omp_set_num_threads, which would mutate a + // process-wide setting; never more threads than sets, so no team member idles + // and every thread id indexes a distinct solver. + const int team_size = std::max(1, std::min(m_solvers.size(), number_of_groups)); + // Self-scheduled, not a static block partition: a heterogeneous sweep (current + // sets that terminate early on events) otherwise leaves threads idle through + // the tail, and a static split also leaves n % num_solvers solves serial. + // Plain atomic, not schedule(dynamic) or omp atomic capture: the macOS wheels' + // libomp lacks __kmpc_dispatch_deinit, and MSVC wants -openmp:llvm for capture. + // The first round is pre-assigned by thread id, so the GIL holder always takes + // a set and streams its diagnostics rather than leaving every set to a worker. + std::atomic next_group{team_size}; + #pragma omp parallel num_threads(team_size) + { + const int thread = omp_get_thread_num(); + for (int i = thread; i < number_of_groups; + i = next_group.fetch_add(1, std::memory_order_relaxed)) { + const sunrealtype *y = y0 + i * y0_np.shape(1); + const sunrealtype *yp = yp0 + i * yp0_np.shape(1); + const sunrealtype *input = inputs_data + i * inputs.shape(1); + const sunrealtype *scales = pbar_data ? pbar_data + i * number_of_parameters : nullptr; + try { + results[i] = m_solvers[thread]->solve( + t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps); + } catch (py::error_already_set &) { + #pragma omp critical + { + if (!python_exception) { + python_exception = std::current_exception(); + } + } + } catch (std::exception &e) { + errors[i] = e.what(); } } } @@ -145,22 +174,21 @@ std::vector IDAKLUSolverGroup::solve( std::rethrow_exception(python_exception); } - if (exception_message.has_value()) { - py::set_error(PyExc_ValueError, exception_message->c_str()); - throw py::error_already_set(); + std::string failures; + for (int i = 0; i < number_of_groups; i++) { + if (errors[i].empty()) { + continue; + } + if (!failures.empty()) { + failures += "; "; + } + failures += "input set " + std::to_string(i) + ": " + errors[i]; } - - // Runs on this thread, so these solves log directly rather than buffering - for (int i = 0; i < remainder_solves; i++) { - const std::size_t index = number_of_groups - remainder_solves + i; - const sunrealtype *y = y0 + index * y0_np.shape(1); - const sunrealtype *yp = yp0 + index * yp0_np.shape(1); - const sunrealtype *input = inputs_data + index * inputs.shape(1); - results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, save_adaptive_steps, save_interp_steps); + if (!failures.empty()) { + py::set_error(PyExc_ValueError, failures.c_str()); + throw py::error_already_set(); } - flush_logs(); - // create solutions (needs to be serial as we're using the Python GIL) std::vector solutions(number_of_groups); for (int i = 0; i < number_of_groups; i++) { diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp index 3ba0239a78..8b44b52a07 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp @@ -13,11 +13,20 @@ class IDAKLUSolverGroup /** * @brief Default constructor + * + * `rust_owner` keeps whatever Python object owns the memory the solvers' + * expression sets point at (the Rust evaluator pool) alive for at least as + * long as the group, instead of relying on the caller to outlive us. + * pybind11 destroys the group with the GIL held, so its destructor is safe. + * It is required rather than defaulted, so a new construction site has to say + * what owns its memory; CasADi's expression sets own theirs and pass none. */ - IDAKLUSolverGroup(std::vector> solvers, int number_of_states, int number_of_parameters): + IDAKLUSolverGroup(std::vector> solvers, int number_of_states, int number_of_parameters, + py::object rust_owner): m_solvers(std::move(solvers)), number_of_states(number_of_states), - number_of_parameters(number_of_parameters) + number_of_parameters(number_of_parameters), + m_rust_owner(std::move(rust_owner)) {} // no copy constructor (unique_ptr cannot be copied) @@ -37,6 +46,7 @@ class IDAKLUSolverGroup np_array y0_np, np_array yp0_np, np_array inputs, + np_array pbar = np_array(), py::object logger = py::none()); @@ -56,6 +66,7 @@ class IDAKLUSolverGroup std::vector> m_solvers; int number_of_states; int number_of_parameters; + py::object m_rust_owner; }; #endif // PYBAMM_IDAKLU_SOLVER_GROUP_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp index c5b3198c19..ed18ef81d7 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.hpp @@ -3,6 +3,7 @@ #include "IDAKLUSolver.hpp" #include "common.hpp" +#include #include #include // For std::make_unique using std::vector; @@ -97,6 +98,12 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver vector res_dvar_dp; // Reused scratch for sparse->dense output sensitivity scatter (outputs-only mode) vector dvar_dp_dense; + // Staged (t, y) points for batched output evaluation (Rust ExprSet only): + // interpolation points accumulate here and flush through one batched FFI + // call instead of one tape walk per point. + vector stage_out_t_; + vector stage_out_y_; + size_t stage_out_first_ = 0; // i_save_ index of the first staged point bool const sensitivity; // cppcheck-suppress unusedStructMember bool const save_outputs_only; // cppcheck-suppress unusedStructMember bool save_hermite; // cppcheck-suppress unusedStructMember @@ -188,6 +195,9 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver sunrealtype *yp_val_ = nullptr; vector yS_val_; vector ypS_val_; + // |p| per sensitivity parameter, so IDAS weights the scaled sensitivity + // pbar*yS like a state. Empty leaves IDAS at its pbar = 1 default. + std::vector sens_scales_; SUNContext sunctx; @@ -223,6 +233,7 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver const sunrealtype *y0, const sunrealtype *yp0, const sunrealtype *inputs, + const sunrealtype *pbar, bool save_adaptive_steps, bool save_interp_steps ) override; @@ -296,6 +307,14 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ void ReinitializeIntegrator(const sunrealtype& t_val); + /** + * @brief Hand IDAS the per-parameter scales held in sens_scales_. + * + * Must run after every IDASensInit/IDASensReInit, both of which reset pbar + * to 1.0. + */ + void ApplySensitivityScales(); + /** * @brief Set a consistent initialization for the system of equations. */ @@ -365,7 +384,8 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver const std::vector &t_eval, const sunrealtype *y0, const sunrealtype *yp0, - const sunrealtype *inputs + const sunrealtype *inputs, + const sunrealtype *pbar ); /** @@ -459,11 +479,29 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ void SetStepOutput(sunrealtype &tval); + /** + * @brief Evaluate and store any staged batched-output points + */ + void FlushOutputBatch(); + /** * @brief Save the output function sensitivities at the requested time */ void SetStepOutputSensitivities(sunrealtype &tval); + /** + * @brief Save the output function sensitivities via native projection. + * Used when ExprSet::kNativeOutputSensitivities is true (Rust core), writing + * to the same flat yS positions as the CasADi SetStepOutputSensitivities. + */ + void NativeSetStepOutputSensitivities(sunrealtype &tval); + + // Reusable scratch for the native output-sensitivity projection (avoids + // per-step allocation). yS_flat_: [n_params * n_states] gather buffer; + // out_sens_flat_: [n_params * length_of_return_vector] projection output. + std::vector yS_flat_; + std::vector out_sens_flat_; + /** * @brief Save Hermite interpolation derivatives at the requested time */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl index b3959e4a21..d5deea3ee3 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl @@ -330,6 +330,7 @@ SolutionData IDAKLUSolverOpenMP::solve( const sunrealtype *y0, const sunrealtype *yp0, const sunrealtype *inputs, + const sunrealtype *pbar, bool save_adaptive_steps, bool save_interp_steps ) @@ -344,7 +345,7 @@ SolutionData IDAKLUSolverOpenMP::solve( // setup InitializeSolveStorage(number_of_evals, t_interp.size()); - SetupInitialState(t_eval, y0, yp0, inputs); + SetupInitialState(t_eval, y0, yp0, inputs, pbar); sunrealtype t0 = t_eval.front(); sunrealtype tf = t_eval.back(); @@ -464,6 +465,10 @@ void IDAKLUSolverOpenMP::InitializeSolveStorage( length_of_return_vector = ReturnVectorLength(); i_save_ = 0; + // Drop any staged batched-output points from a previous solve. + stage_out_t_.clear(); + stage_out_y_.clear(); + stage_out_first_ = 0; // Allocate output arrays. Pre-allocate 64 elements for initial storage. int est = std::max(n_evals + n_interps, 64); @@ -515,7 +520,8 @@ void IDAKLUSolverOpenMP::SetupInitialState( const std::vector &t_eval, const sunrealtype *y0, const sunrealtype *yp0, - const sunrealtype *inputs + const sunrealtype *inputs, + const sunrealtype *pbar ) { DEBUG("IDAKLUSolver::SetupInitialState"); @@ -524,6 +530,17 @@ void IDAKLUSolverOpenMP::SetupInitialState( functions->inputs[i] = inputs[i]; } + // Sanitised once per solve, then re-applied unchanged after each reinit. + // IDAS rejects a zero pbar, and a zero parameter has no scale of its own. + sens_scales_.clear(); + if (sensitivity && pbar != nullptr) { + sens_scales_.reserve(number_of_parameters); + for (int i = 0; i < number_of_parameters; i++) { + sunrealtype const scale = std::abs(pbar[i]); + sens_scales_.push_back((std::isfinite(scale) && scale > 0.0) ? scale : 1.0); + } + } + // Setup SUNDIALS vector pointers (member state) y_val_ = N_VGetArrayPointer(yy); yp_val_ = N_VGetArrayPointer(yyp); @@ -630,10 +647,8 @@ void IDAKLUSolverOpenMP::HandleBreakpoint( i_eval++; t_eval_next = t_eval[i_eval]; CheckErrors(IDASetStopTime(ida_mem, t_eval_next), "IDASetStopTime"); - if (solver_opts.print_stats) { - // Save stats before reinitializing (reinit resets IDA counters) - SaveStats(); - } + // Save stats before reinitializing (reinit resets IDA counters) + SaveStats(); // Reinitialize the solver to deal with the discontinuity at t = t_val ReinitializeIntegrator(t_val); @@ -648,11 +663,16 @@ template SolutionData IDAKLUSolverOpenMP::BuildSolutionData(int retval) { DEBUG("IDAKLUSolver::BuildSolutionData"); + SaveStats(); if (solver_opts.print_stats) { - SaveStats(); CaptureStats(); } + // Evaluate any tail of staged batched-output points before y is finalized. + if (save_outputs_only) { + FlushOutputBatch(); + } + // Finalize output arrays if (use_knot_reduction_) { knot_reducer->Finalize(); @@ -694,7 +714,8 @@ SolutionData IDAKLUSolverOpenMP::BuildSolutionData(int retval) { arg_sens0, arg_sens1, arg_sens2, - save_hermite + save_hermite, + accumulated_stats ); } @@ -780,9 +801,20 @@ void IDAKLUSolverOpenMP::ReinitializeIntegrator(const sunrealtype& t_va CheckErrors(IDAReInit(ida_mem, t_val, yy, yyp), "IDAReInit"); if (sensitivity) { CheckErrors(IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, yypS), "IDASensReInit"); + ApplySensitivityScales(); } } +template +void IDAKLUSolverOpenMP::ApplySensitivityScales() { + if (sens_scales_.empty()) { + return; + } + CheckErrors( + IDASetSensParams(ida_mem, nullptr, sens_scales_.data(), nullptr), + "IDASetSensParams"); +} + template void IDAKLUSolverOpenMP::ConsistentInitialization( const sunrealtype& t_val, @@ -957,6 +989,11 @@ void IDAKLUSolverOpenMP::SetStepFullSensitivities( } } +// Staged points per batched output evaluation: amortises per-instruction +// interpreter dispatch while bounding the extra state copy at +// kOutputBatchPoints * number_of_states. +static constexpr size_t kOutputBatchPoints = 128; + template void IDAKLUSolverOpenMP::SetStepOutput( sunrealtype &tval @@ -964,6 +1001,23 @@ void IDAKLUSolverOpenMP::SetStepOutput( DEBUG("IDAKLUSolver::SetStepOutput"); // FLAT STORAGE: Write output variables to y[i_save_ * stride + j] + if constexpr (ExprSet::kNativeBatchedOutputs) { + // The sensitivity path reads yS_val_ at the point, so it stays per-point. + if (!sensitivity) { + if (stage_out_t_.empty()) { + stage_out_first_ = i_save_; + stage_out_t_.reserve(kOutputBatchPoints); + stage_out_y_.reserve(kOutputBatchPoints * number_of_states); + } + stage_out_t_.push_back(tval); + stage_out_y_.insert(stage_out_y_.end(), y_val_, y_val_ + number_of_states); + if (stage_out_t_.size() == kOutputBatchPoints) { + FlushOutputBatch(); + } + return; + } + } + sunrealtype* y_dest = &y[i_save_ * length_of_return_vector]; size_t j = 0; for (auto& var_fcn : functions->var_fcns) { @@ -974,7 +1028,29 @@ void IDAKLUSolverOpenMP::SetStepOutput( } if (sensitivity) { - SetStepOutputSensitivities(tval); + if constexpr (ExprSet::kNativeOutputSensitivities) { + NativeSetStepOutputSensitivities(tval); + } else { + SetStepOutputSensitivities(tval); + } + } +} + +template +void IDAKLUSolverOpenMP::FlushOutputBatch() { + if constexpr (ExprSet::kNativeBatchedOutputs) { + const size_t k = stage_out_t_.size(); + if (k == 0) { + return; + } + DEBUG("IDAKLUSolver::FlushOutputBatch"); + // The staged points wrote nothing yet; their y rows are contiguous from + // stage_out_first_, so one batched call fills them all. + functions->eval_outputs_batch( + stage_out_t_.data(), stage_out_y_.data(), static_cast(k), + &y[stage_out_first_ * length_of_return_vector]); + stage_out_t_.clear(); + stage_out_y_.clear(); } } @@ -1045,6 +1121,35 @@ void IDAKLUSolverOpenMP::SetStepOutputSensitivities( } } +template +void IDAKLUSolverOpenMP::NativeSetStepOutputSensitivities( + sunrealtype &tval +) { + DEBUG("IDAKLUSolver::NativeSetStepOutputSensitivities"); + const size_t n_params = static_cast(number_of_parameters); + const size_t stride = length_of_return_vector; + const size_t yS_base = i_save_ * n_params * stride; + + // Gather per-param state sensitivities into a contiguous [n_params * n_states]. + yS_flat_.resize(n_params * number_of_states); + for (size_t p = 0; p < n_params; ++p) { + std::copy(yS_val_[p], yS_val_[p] + number_of_states, + yS_flat_.begin() + p * number_of_states); + } + // Project: out[p * stride + o] = d(output_o)/d(p_p). + out_sens_flat_.resize(n_params * stride); + functions->project_output_sensitivities( + tval, y_val_, functions->inputs.data(), + yS_flat_.data(), out_sens_flat_.data()); + // Transpose into the shared flat yS: the projection emits (param, output), + // yS is (output, param), as the CasADi path above writes it. + for (size_t p = 0; p < n_params; ++p) { + for (size_t o = 0; o < stride; ++o) { + yS[yS_base + o * n_params + p] = out_sens_flat_[p * stride + o]; + } + } +} + template void IDAKLUSolverOpenMP::SetStepHermite( sunrealtype &tval diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp index ab7cbf15f5..e6b3e7b238 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp @@ -2,6 +2,7 @@ #define PYBAMM_IDAKLU_SOLUTION_HPP #include "common.hpp" +#include "IDAKLUStats.hpp" /** * @brief Solution class @@ -18,9 +19,10 @@ class Solution * @brief Constructor */ Solution(int retval, np_array t_np, np_array y_np, np_array yp_np, - np_array yS_np, np_array ypS_np, np_array y_term_np) + np_array yS_np, np_array ypS_np, np_array y_term_np, + IDAKLUStats stats) : flag(retval), t(t_np), y(y_np), yp(yp_np), yS(yS_np), ypS(ypS_np), - y_term(y_term_np) + y_term(y_term_np), stats(stats) { } @@ -36,6 +38,7 @@ class Solution np_array yS; np_array ypS; np_array y_term; + IDAKLUStats stats; }; #endif // PYBAMM_IDAKLU_SOLUTION_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp index 2cb7a8b565..620b8d786e 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp @@ -48,7 +48,8 @@ class SolutionData ptrdiff_t arg_sens0, ptrdiff_t arg_sens1, ptrdiff_t arg_sens2, - bool save_hermite) + bool save_hermite, + IDAKLUStats stats) : flag(flag), t_vec(std::move(t)), y_vec(std::move(y)), @@ -59,7 +60,8 @@ class SolutionData arg_sens0(arg_sens0), arg_sens1(arg_sens1), arg_sens2(arg_sens2), - save_hermite(save_hermite) + save_hermite(save_hermite), + stats(stats) {} ~SolutionData() = default; @@ -81,7 +83,8 @@ class SolutionData vector_to_numpy_3d(std::move(yS_vec), arg_sens0, arg_sens1, arg_sens2), vector_to_numpy_3d(std::move(ypS_vec), save_hermite ? arg_sens0 : 0, arg_sens1, arg_sens2), - vector_to_numpy(std::move(yterm_vec)) + vector_to_numpy(std::move(yterm_vec)), + stats ); } @@ -97,6 +100,7 @@ class SolutionData ptrdiff_t arg_sens1 = 0; ptrdiff_t arg_sens2 = 0; bool save_hermite = false; + IDAKLUStats stats; }; #endif // PYBAMM_IDAKLU_SOLUTION_DATA_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp index 8f6a6f649f..de4afc2323 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp @@ -1,16 +1,19 @@ #include "StandaloneNewtonSolver.hpp" +#include "Expressions/Rust/RustFunctions.hpp" + #include // ────────────────────── StandaloneAlgebraicSystem ────────────────────── StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( - casadi::Function res_fn, - casadi::Function jac_fn, + std::unique_ptr res_fn, + std::unique_ptr jac_fn, + int n_vars, bool use_sparse) - : res_cf_(res_fn), - jac_cf_(jac_fn), - n_vars_(static_cast(res_fn.nnz_out(0))), + : res_(std::move(res_fn)), + jac_(std::move(jac_fn)), + n_vars_(n_vars), use_sparse_(use_sparse), sunctx_(nullptr), J_(nullptr), LS_(nullptr), res_nvec_(nullptr), delta_nvec_(nullptr) @@ -20,7 +23,7 @@ StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( res_nvec_ = N_VNew_Serial(n_vars_, sunctx_); delta_nvec_ = N_VNew_Serial(n_vars_, sunctx_); - int jac_nnz = static_cast(jac_cf_.nnz_out()); + int jac_nnz = static_cast(jac_->nnz_out()); jac_buf_.resize(jac_nnz > 0 ? jac_nnz : n_vars_ * n_vars_); if (use_sparse_ && jac_nnz > 0) { @@ -31,6 +34,17 @@ StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( } } +StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( + casadi::Function res_fn, + casadi::Function jac_fn, + bool use_sparse) + : StandaloneAlgebraicSystem( + std::make_unique(res_fn), + std::make_unique(jac_fn), + static_cast(res_fn.nnz_out(0)), + use_sparse) +{} + StandaloneAlgebraicSystem::~StandaloneAlgebraicSystem() { if (res_nvec_) N_VDestroy(res_nvec_); if (delta_nvec_) N_VDestroy(delta_nvec_); @@ -42,22 +56,22 @@ StandaloneAlgebraicSystem::~StandaloneAlgebraicSystem() { void StandaloneAlgebraicSystem::eval_residual( sunrealtype t, const sunrealtype* y, sunrealtype* res) { - res_cf_.m_arg[0] = &t; - res_cf_.m_arg[1] = y; - res_cf_.m_arg[2] = inputs_.data(); - res_cf_.m_res[0] = res; - res_cf_(); + res_->m_arg[0] = &t; + res_->m_arg[1] = y; + res_->m_arg[2] = inputs_.data(); + res_->m_res[0] = res; + (*res_)(); } int StandaloneAlgebraicSystem::solve_linear( sunrealtype t, const sunrealtype* y, sunrealtype* res, sunrealtype* delta) { - jac_cf_.m_arg[0] = &t; - jac_cf_.m_arg[1] = y; - jac_cf_.m_arg[2] = inputs_.data(); - jac_cf_.m_res[0] = jac_buf_.data(); - jac_cf_(); + jac_->m_arg[0] = &t; + jac_->m_arg[1] = y; + jac_->m_arg[2] = inputs_.data(); + jac_->m_res[0] = jac_buf_.data(); + (*jac_)(); if (use_sparse_) { sunrealtype* mat_data = SUNSparseMatrix_Data(J_); @@ -88,8 +102,8 @@ int StandaloneAlgebraicSystem::solve_linear( } void StandaloneAlgebraicSystem::BuildSparseResources(int jac_nnz) { - const auto& rows = jac_cf_.get_row(); - const auto& cols = jac_cf_.get_col(); + const auto& rows = jac_->get_row(); + const auto& cols = jac_->get_col(); int nnz_total = jac_nnz; // Build CSC from COO (same O(nnz) algorithm as AlgebraicICBuilder) @@ -132,9 +146,9 @@ void StandaloneAlgebraicSystem::BuildSparseResources(int jac_nnz) { } void StandaloneAlgebraicSystem::BuildDenseResources() { - const auto& rows = jac_cf_.get_row(); - const auto& cols = jac_cf_.get_col(); - int nnz_total = static_cast(jac_cf_.nnz_out()); + const auto& rows = jac_->get_row(); + const auto& cols = jac_->get_col(); + int nnz_total = static_cast(jac_->nnz_out()); // Build CSC structure for dense scatter std::vector col_count(n_vars_ + 1, 0); @@ -188,6 +202,33 @@ StandaloneNewtonSolver::StandaloneNewtonSolver( y_work_(n_vars_) {} +StandaloneNewtonSolver::StandaloneNewtonSolver( + std::uintptr_t rust_model_ptr, + int n_rhs, + int n_alg, + const std::vector& jac_rows, + const std::vector& jac_cols, + const std::vector& atol, + sunrealtype rtol, + sunrealtype step_tol, + int max_iter, + int max_backtracks, + sunrealtype epsNewt, + bool use_sparse) + : system_( + std::make_unique( + reinterpret_cast(rust_model_ptr), n_rhs, n_alg), + std::make_unique( + reinterpret_cast(rust_model_ptr), n_rhs, n_alg, + static_cast(jac_rows.size()), jac_rows, jac_cols), + n_alg, + use_sparse), + solver_(system_, system_.n_vars(), atol.data(), rtol, step_tol, + max_iter, max_backtracks, epsNewt), + n_vars_(system_.n_vars()), + y_work_(n_vars_) +{} + std::pair StandaloneNewtonSolver::solve( sunrealtype t, const np_array& y0_np, diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp index b4e1c6e5be..9d382f2ca1 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp @@ -4,18 +4,38 @@ #include "NonlinearSolver.hpp" #include "Expressions/Casadi/CasadiFunctions.hpp" #include "common.hpp" +#include #include #include /** - * @brief NonlinearSystem backed by standalone CasadiFunctions + SUNLinSol. + * @brief NonlinearSystem backed by standalone Expression functions + SUNLinSol. * * Self-contained: no dependency on IDAKLUSolverOpenMP or IDA memory. * Residual signature: F(t, y_alg, inputs) -> res (n_vars outputs) * Jacobian signature: J(t, y_alg, inputs) -> data (COO or dense) + * + * Polymorphic over the Expression base, so any backend (casadi, Rust, ...) + * that implements the Expression interface can drive the Newton solve. */ class StandaloneAlgebraicSystem : public NonlinearSystem { public: + /** + * @brief Expression-based constructor (the general entry point). + * @param res_fn Residual function F(t, y, inputs) -> res + * @param jac_fn Jacobian function J(t, y, inputs) -> data (COO) + * @param n_vars Number of algebraic variables (residual output length) + * @param use_sparse Whether to use a sparse (KLU) linear solver + */ + StandaloneAlgebraicSystem( + std::unique_ptr res_fn, + std::unique_ptr jac_fn, + int n_vars, + bool use_sparse); + + /** + * @brief Casadi convenience constructor; delegates to the Expression ctor. + */ StandaloneAlgebraicSystem( casadi::Function res_fn, casadi::Function jac_fn, @@ -39,8 +59,8 @@ class StandaloneAlgebraicSystem : public NonlinearSystem { void BuildSparseResources(int jac_nnz); void BuildDenseResources(); - CasadiFunction res_cf_; - CasadiFunction jac_cf_; + std::unique_ptr res_; + std::unique_ptr jac_; int n_vars_; bool use_sparse_; @@ -77,6 +97,27 @@ class StandaloneNewtonSolver { sunrealtype epsNewt, bool use_sparse); + /** + * @brief Rust-backed constructor. + * + * Builds Newton residual/Jacobian adapters over the Rust FFI from an opaque + * Rust model handle, then constructs the algebraic system and solver exactly + * as the casadi ctor does. + */ + StandaloneNewtonSolver( + std::uintptr_t rust_model_ptr, + int n_rhs, + int n_alg, + const std::vector& jac_rows, + const std::vector& jac_cols, + const std::vector& atol, + sunrealtype rtol, + sunrealtype step_tol, + int max_iter, + int max_backtracks, + sunrealtype epsNewt, + bool use_sparse); + /** * @brief Solve F(t, y, inputs) = 0 starting from y0. * @return (success, y_solution) with zero-copy numpy output. diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl index a96afa7bcd..89234d88d9 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/sundials_functions.inl @@ -157,6 +157,9 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, // create pointer to jac data, column pointers, and row values sunrealtype *jac_data; + // dense SUNMatrix + sparse-emitting expression (rust): eval into the tmp + // buffer, then scatter below via the stored CSC structure. + bool dense_needs_scatter = false; if (p_python_functions->setup_opts.using_sparse_matrix) { jac_data = SUNSparseMatrix_Data(JJ); @@ -166,7 +169,11 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, } else { - jac_data = SUNDenseMatrix_Data(JJ); + const sunindextype n_dense = SUNDenseMatrix_Rows(JJ) * SUNDenseMatrix_Columns(JJ); + dense_needs_scatter = p_python_functions->jac_times_cjmass->nnz_out() < n_dense; + jac_data = dense_needs_scatter + ? p_python_functions->get_tmp_sparse_jacobian_data() + : SUNDenseMatrix_Data(JJ); } DEBUG_VECTORn(yy, 100); @@ -254,6 +261,22 @@ int jacobian_eval(sunrealtype tt, sunrealtype cj, N_Vector yy, N_Vector yp, } else throw std::runtime_error("Unknown matrix format detected (Expected CSC or CSR)"); } + else if (dense_needs_scatter) + { + // scatter the sparse-emitting expression's CSC values into the dense matrix + auto jac_colptrs = p_python_functions->jac_times_cjmass_colptrs.data(); + auto jac_rowvals = p_python_functions->jac_times_cjmass_rowvals.data(); + sunrealtype *dense = SUNDenseMatrix_Data(JJ); + const sunindextype n_rows = SUNDenseMatrix_Rows(JJ); + const sunindextype n_cols = SUNDenseMatrix_Columns(JJ); + std::memset(dense, 0, static_cast(n_rows * n_cols) * sizeof(sunrealtype)); + for (sunindextype col_ij = 0; col_ij < n_cols; col_ij++) { + for (auto data_i = jac_colptrs[col_ij]; data_i < jac_colptrs[col_ij + 1]; data_i++) { + const auto row_ij = jac_rowvals[data_i]; + dense[col_ij * n_rows + row_ij] = jac_data[data_i]; + } + } + } return (0); } diff --git a/packages/pybammsolvers/tests/test_integration.py b/packages/pybammsolvers/tests/test_integration.py index d1cd95690c..748ddc12a7 100644 --- a/packages/pybammsolvers/tests/test_integration.py +++ b/packages/pybammsolvers/tests/test_integration.py @@ -232,3 +232,37 @@ def test_parallel_solver_group_uses_multiple_solvers( for idx, sol in enumerate(solutions): expected = model_y0 * np.exp(-decay_constants[idx] * sol.t) np.testing.assert_allclose(sol.y, expected, rtol=1e-5, atol=1e-8) + + +class TestSensitivityScales: + """The optional ``pbar`` argument carrying IDAS's sensitivity scales.""" + + pytestmark = pytest.mark.integration + + @staticmethod + def _args(solver_data): + t_eval = solver_data["model"]["t_eval"] + return ( + t_eval, + t_eval, + solver_data["y0"], + solver_data["yp0"], + solver_data["inputs"], + ) + + def test_an_empty_pbar_matches_omitting_it(self, exponential_decay_solver): + solver = exponential_decay_solver["solver"] + args = self._args(exponential_decay_solver) + + without = solver.solve(*args)[0] + with_empty = solver.solve(*args, np.empty((0, 0)))[0] + + np.testing.assert_array_equal(without.t, with_empty.t) + np.testing.assert_array_equal(without.y, with_empty.y) + + def test_a_misshapen_pbar_is_rejected(self, exponential_decay_solver): + # A silently ignored pbar would leave the scales at 1 with no warning. + # This fixture has no sens parameters, so any pbar is too wide. + solver = exponential_decay_solver["solver"] + with pytest.raises(Exception, match="pbar has wrong number of cols"): + solver.solve(*self._args(exponential_decay_solver), np.ones((1, 3))) diff --git a/pyproject.toml b/pyproject.toml index 43dd004f23..87b2b26e28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,8 @@ [tool.uv.workspace] members = ["packages/*"] +# The Rust crate is excluded as uv would glob it and error on the absent pyproject. +exclude = ["packages/pybamm-rust"] [tool.uv.sources] pybammsolvers = { workspace = true } @@ -22,6 +24,16 @@ no-build-isolation-package = ["pybammsolvers"] [dependency-groups] dev = [ "nox[uv]", + # Builds the Rust extension. The floor is the version this was verified + # against, not the technical minimum (<1.9.4 does not set + # PYO3_BUILD_EXTENSION_MODULE at all). + "maturin>=1.14.1,<2.0", + # Imported by packages/pybamm/hatch_build.py, which has unit tests. + "hatchling>=1.31.0", + # Only for `mypy.stubtest`, which pins pybamm/rust/_core.pyi against the + # built extension. + "mypy>=1.14", + "scipy-stubs", "pre-commit", # pybammsolvers build toolchain, required in the venv because the solver is # installed with --no-build-isolation (above) and rebuilds itself on import. diff --git a/uv.lock b/uv.lock index 5e39c6eb5e..bda2d9f5a7 100644 --- a/uv.lock +++ b/uv.lock @@ -5,17 +5,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version < '3.11'", ] @@ -29,11 +29,15 @@ members = [ [manifest.dependency-groups] dev = [ { name = "cmake", specifier = ">=4.4.2" }, + { name = "hatchling", specifier = ">=1.31.0" }, + { name = "maturin", specifier = ">=1.14.1,<2.0" }, + { name = "mypy", specifier = ">=1.14" }, { name = "ninja" }, { name = "nox", extras = ["uv"] }, { name = "pre-commit" }, { name = "pybind11", specifier = ">=3.0.1" }, { name = "scikit-build-core", specifier = ">=1.0.3" }, + { name = "scipy-stubs" }, ] [[package]] @@ -168,6 +172,69 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/c9/d7977eaacb9df673210491da99e6a247e93df98c715fc43fd136ce1d3d33/arrow-1.4.0-py3-none-any.whl", hash = "sha256:749f0769958ebdc79c173ff0b0670d59051a535fa26e8eba02953dc19eb43205", size = 68797, upload-time = "2025-10-18T17:46:45.663Z" }, ] +[[package]] +name = "ast-serialize" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/a9/11851c3e02a3fea2ddc9932d1fdc7d2edaeecc0d2e11bc5f2a7fde2b0934/ast_serialize-0.8.0.tar.gz", hash = "sha256:6c37c43e4004dfb42d321ddedc569dc17ff4259296f3af577c9ea46a809bc010", size = 845638, upload-time = "2026-08-07T11:29:02.152Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/16/6e520b57cd8c75914b38c670ad4593d13c22911e4306cc7165dab8b0789b/ast_serialize-0.8.0-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:3d822605fa7bb326ef868d25fafced7fc660fa46d9b90c02ea86d5e2f5d325f7", size = 863924, upload-time = "2026-08-07T11:27:34.579Z" }, + { url = "https://files.pythonhosted.org/packages/03/e1/48802de9b22a2bcad42ec80601a17e3f69172fe4f590e6311bcc2b323aeb/ast_serialize-0.8.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:2efa40b068197d5efb62655b43baadb842ed71c4958cccd3e8b86a35726f0119", size = 1177662, upload-time = "2026-08-07T11:27:36.196Z" }, + { url = "https://files.pythonhosted.org/packages/38/d4/323438db76bded3a1f3523a3167b8325916b2ddceb2107a330c6ec9fcf4d/ast_serialize-0.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:db1b957291bca08c7e72f43a12357b2948e20775d970e3fc3dac0aa3160ab725", size = 1167072, upload-time = "2026-08-07T11:27:37.646Z" }, + { url = "https://files.pythonhosted.org/packages/77/82/53c5400b54144b56de8ed7f957fd1ccd97e42482009292ab46121d15f8dd/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fdc0d5b18ff8fb364e87923e47c0a91d0d69dbcaeaa274591f7fd26892cc3a3a", size = 1225497, upload-time = "2026-08-07T11:27:39.225Z" }, + { url = "https://files.pythonhosted.org/packages/44/5f/36c07327a8b91303fbf1382c7c3e8a2902072dbe1b9546138a5288e75ff0/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9da7330f3e235bf7da89b8d39205c6350fc0c08a85379743f2df9fff87d6d980", size = 1227101, upload-time = "2026-08-07T11:27:40.799Z" }, + { url = "https://files.pythonhosted.org/packages/9d/48/5adf5c67addc7ddb328122208c6d375a84cf154984f412b4087330a157bd/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f3186969ee66a9863b00acc6523ace44c56974eecb348a7ea4b228d9f0b80e19", size = 1424001, upload-time = "2026-08-07T11:27:42.708Z" }, + { url = "https://files.pythonhosted.org/packages/38/a1/70074dd3869d2b0e934f91891d8d6b734361cd3b80f85ca7ece2e668ecdd/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:40a57b73731be45da4fa41430c4d5dc94a24b3a4faba7b9e069978c0402064ea", size = 1245545, upload-time = "2026-08-07T11:27:44.4Z" }, + { url = "https://files.pythonhosted.org/packages/e3/be/53b9c0a8a6399950c2e3546bdfab96d2b299d5b114b47eb94fd3c49c4054/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5075b9da3ef807eda752502446dfecea3b381c4900b7e27a5d5f4f899eb39951", size = 1248961, upload-time = "2026-08-07T11:27:45.781Z" }, + { url = "https://files.pythonhosted.org/packages/eb/13/3651d3812548a2bda15e26e5dd51aadb48cf682d0865370255fcf0e367dd/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:293cc1c5bfa741f8e3fbe8175b9c07beee487c9a6fdbb25a5acad9f1df2d30a9", size = 1243877, upload-time = "2026-08-07T11:27:47.325Z" }, + { url = "https://files.pythonhosted.org/packages/21/a0/521f0bf000f675e9312a4aae2c8ba7a992405d072a85c485e08fd59433b9/ast_serialize-0.8.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:e0910c3442a75216dde0f102d854ba2aaa71d2482e0ee213630b9bf29584fba3", size = 1293903, upload-time = "2026-08-07T11:27:49.264Z" }, + { url = "https://files.pythonhosted.org/packages/b1/7e/402fc902568aa2ee65865a3e151f000db0153da8ce6b1be4c9c349025f8d/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:43dd6d596879bb1cb8a12cc9dae7bb10090a39a35883026c24f82488a195619a", size = 1401070, upload-time = "2026-08-07T11:27:50.947Z" }, + { url = "https://files.pythonhosted.org/packages/ff/7c/97d4b66c057f1706fc8be6dd532cc77c988794357c8f4ffdb6adabb39562/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:8c9d537f59e936392cfd3597789d1390304dd659efc3c486ce7f40fb6b8a9f53", size = 1502602, upload-time = "2026-08-07T11:27:52.364Z" }, + { url = "https://files.pythonhosted.org/packages/89/6f/72cc3b71562001bba46e898ccfbf1844f7939b3e28912736206102f2e5a8/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:f0190a33d7f97c65e9069f7a7f40499eea6b5cbe260c558378109caf20ce934b", size = 1495848, upload-time = "2026-08-07T11:27:53.803Z" }, + { url = "https://files.pythonhosted.org/packages/a0/53/d6f629d1e49308b2f363dae028baa213ec222c9106fa1f7f0d1f7b41499a/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:77308ae6c5cf5264cc0f01a7c556ec77a9e68eb1f61b093534d698139fdc3b14", size = 1556556, upload-time = "2026-08-07T11:27:55.342Z" }, + { url = "https://files.pythonhosted.org/packages/ee/22/340f35dd8dfc6d412d53dc20699ca014b8d228db923e8ed4759c512b162c/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:8d53a23f27e1ed3a36b2d26fd2a1a6228c8e85a1ed62ff7cdb44bd610769f20a", size = 1417822, upload-time = "2026-08-07T11:27:56.712Z" }, + { url = "https://files.pythonhosted.org/packages/11/29/6dde5c13fbebc051d3a6df4ec0a6fd1d5359333cc1193f7f609f3410b4d8/ast_serialize-0.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ffa5e7cb08f96fed9121f77b224151e41caf88feab9d652bb46c78202b6fbeda", size = 1445153, upload-time = "2026-08-07T11:27:58.275Z" }, + { url = "https://files.pythonhosted.org/packages/62/c5/f473a8ed030f7a0ca24b9849cca184677a50c053867a7b808c2e1289bbd3/ast_serialize-0.8.0-cp314-cp314t-win32.whl", hash = "sha256:fa70ed4dea0bb18b30a1789c77baa701d0ef30c474f2ccabdea61e25623a8827", size = 1063711, upload-time = "2026-08-07T11:27:59.793Z" }, + { url = "https://files.pythonhosted.org/packages/23/63/39e171fcd38ca057c2e1979d5ee81ac7a3502784abe3d83df7454f7a0978/ast_serialize-0.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d8b3c8eee4c1baef9d4e84d2a59a805501617127be42615cb48970b15b0892b6", size = 1103740, upload-time = "2026-08-07T11:28:01.405Z" }, + { url = "https://files.pythonhosted.org/packages/21/1c/d00762b399e7726d68d0a088cc946e3a4c60f1c6176f557608f672f627f3/ast_serialize-0.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:ac4f0a83c55a9b782f79ad55a5247b7db123c1db405959791c2ef886e9710c9f", size = 1076021, upload-time = "2026-08-07T11:28:02.947Z" }, + { url = "https://files.pythonhosted.org/packages/4c/11/911210c3c78923273a9211a2b6cfc4c8aa723b30dab3e1c8d19afb983b40/ast_serialize-0.8.0-cp315-abi3.abi3t-macosx_10_12_x86_64.whl", hash = "sha256:86b8a1e6d90467345356098b040150e82fbc26d24a7a202224b13dc1f6264ca0", size = 1177715, upload-time = "2026-08-07T11:28:04.654Z" }, + { url = "https://files.pythonhosted.org/packages/77/89/6282881c8587606638db153cbe21e1e0c4d1f3970dee1aa0610a1c62a026/ast_serialize-0.8.0-cp315-abi3.abi3t-macosx_11_0_arm64.whl", hash = "sha256:39e92ff8e8cb45947fe9007174b2950e1fb098e6abd00266a13cd3bcf6675068", size = 1169347, upload-time = "2026-08-07T11:28:06.1Z" }, + { url = "https://files.pythonhosted.org/packages/97/78/a9f846a03a340ff3728c915f23338ca742742f3292700559cdb3ad999b1e/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c85d8d18db5b2dfcb3b7e38a4d600ca35504c0ed8a6f75cd1c811e4ffe248a15", size = 1225916, upload-time = "2026-08-07T11:28:07.654Z" }, + { url = "https://files.pythonhosted.org/packages/c0/15/aba6ef8a988a6eceb6f0359589aac509e29ae2dba67fd9bfd5af0c3f13e7/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9830ff7e764f74d9eefb01170c61a9f0fd2c027dac5fcb72e064decd57d56371", size = 1227135, upload-time = "2026-08-07T11:28:09.504Z" }, + { url = "https://files.pythonhosted.org/packages/94/29/3f63d696ea7c5b8abadcecc3505be51bd900daaccc522ed8322fa5b05a93/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6479d9722a4cd21b578f5478074c41e6169f04811996ec881655560f703a5bba", size = 1425040, upload-time = "2026-08-07T11:28:11.044Z" }, + { url = "https://files.pythonhosted.org/packages/e2/5d/0aac338604ff59df5774d4304307898982252f325ff7cafe31d52fedcb65/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a63bed264e818cd83eec11feed0f50aa162542b91132ef58afebc857182763a5", size = 1246278, upload-time = "2026-08-07T11:28:12.519Z" }, + { url = "https://files.pythonhosted.org/packages/23/ca/9f1ef795bb724719532bd86dbec11e5b66857d3fbe9b6772baec0191a6ed/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9d187197d234aa45d6cfa2b096be5f666e8cc2e7eb3722d0ab8926293cf5720c", size = 1250029, upload-time = "2026-08-07T11:28:13.896Z" }, + { url = "https://files.pythonhosted.org/packages/dc/25/5e061372d2ed953b9ba3b9c4f73de3b8e9234cda3f6c088db4686801d0e1/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_31_riscv64.whl", hash = "sha256:2d39a56282cfcc0d8eeea37267c754be59c98d48505c23b1dae5c6011f3813dd", size = 1243575, upload-time = "2026-08-07T11:28:15.37Z" }, + { url = "https://files.pythonhosted.org/packages/a8/c1/ae7da218053120635a4ca802366c69f707203641af95372eeb83f70dfd52/ast_serialize-0.8.0-cp315-abi3.abi3t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f7cc5f10386994c0f4844f1e6d6a97127e9b478660eb6dec2b257644f0acab64", size = 1294396, upload-time = "2026-08-07T11:28:16.813Z" }, + { url = "https://files.pythonhosted.org/packages/2e/89/271d1f49c5269fcddcc789ea3f25be401f6723fc1138aeda539f4d05516d/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_aarch64.whl", hash = "sha256:6102f2f985c2e542be85cd857678ec9356fefa792b93cadfadd31139f5696f27", size = 1401987, upload-time = "2026-08-07T11:28:18.333Z" }, + { url = "https://files.pythonhosted.org/packages/55/be/4e7d77fcf571ac7cb5cf7115a20c36642bd7d29473b45dfaaefeb9618f90/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_armv7l.whl", hash = "sha256:3a8660fe66667b76a6e9dccd1d33e66b229fde3b308db991c041609226c005b6", size = 1502904, upload-time = "2026-08-07T11:28:20.039Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ae/ed1de2db7e019d4236fbc164ffa5ef9a6022a300a342bbf142d21b7c141e/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_i686.whl", hash = "sha256:e7266307e5fba39836edb79def8608887af48820508bff3c5f2941e1e04d1534", size = 1496967, upload-time = "2026-08-07T11:28:21.734Z" }, + { url = "https://files.pythonhosted.org/packages/92/89/5fea507fae5c5f18b7dc7f95e5c00956574b8c717b8fd2049c504fab0b18/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_ppc64le.whl", hash = "sha256:4ca7e6fd1ad845d1cc649dc2ecd499db2f8f46af5bf8da7b70dd858774cc038b", size = 1559041, upload-time = "2026-08-07T11:28:23.194Z" }, + { url = "https://files.pythonhosted.org/packages/42/71/478d69df21b64e064554a68134c94be304270316ca676a94e63c389a636a/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_riscv64.whl", hash = "sha256:2880350b13d3eae69a0d70bc1fb6c9bfaca4dbd0e20ba8cd1aa483080b56ff06", size = 1417367, upload-time = "2026-08-07T11:28:24.601Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2d/8962dc8d5b3a9dc27b36f9db199afa25264c741505469d9ec10ffbfd2ba7/ast_serialize-0.8.0-cp315-abi3.abi3t-musllinux_1_2_x86_64.whl", hash = "sha256:ab0f9a59f7d63d0d441b56b9a818b273705264352d5115cfee12e940e816d958", size = 1446178, upload-time = "2026-08-07T11:28:26.152Z" }, + { url = "https://files.pythonhosted.org/packages/4f/22/14d2ad4fd1d1bcd0dc687ca268e0630069f45162496260c0efb70ee0ea72/ast_serialize-0.8.0-cp315-abi3.abi3t-win32.whl", hash = "sha256:0485a25ef519c62e749ee3c1ad8070e591b380d67226349eb5a70b228dc1ac4a", size = 1063811, upload-time = "2026-08-07T11:28:27.864Z" }, + { url = "https://files.pythonhosted.org/packages/18/1d/84a327c0202a41aa5fdba3ade33904d6d8f3b9e6806fa83568d835395850/ast_serialize-0.8.0-cp315-abi3.abi3t-win_amd64.whl", hash = "sha256:bd84d60bca7079e741be4ac5dbe237751a59d7f6f9f0126b11880d63822cbe16", size = 1105518, upload-time = "2026-08-07T11:28:29.691Z" }, + { url = "https://files.pythonhosted.org/packages/8c/92/74556dec52fde85a2ad84ed159991b916241043788609c15d8b77e14570b/ast_serialize-0.8.0-cp315-abi3.abi3t-win_arm64.whl", hash = "sha256:057769b5921336eb2d9124f2a731b42ed05ffdac559b840dbdf6f3937cf153dc", size = 1076319, upload-time = "2026-08-07T11:28:31.282Z" }, + { url = "https://files.pythonhosted.org/packages/d9/e3/6142e920fec6ef7bccabd8c24ed8ed99f8bdc6cb8b065e1df7c6a3b2d667/ast_serialize-0.8.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:e1bd223df0f6c96b396975fa604cb33bce53d9b4a0185490be4c4a289f7c9c87", size = 1184007, upload-time = "2026-08-07T11:28:34.654Z" }, + { url = "https://files.pythonhosted.org/packages/a6/e9/6e8be8df02b35d85e2b8809f7f1cfa290bdf5882b55127a539d049482db0/ast_serialize-0.8.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:ddd3b61f45c132da66c5476b281891e08c1fd87fbdabe8a6973e1622efc85f06", size = 1177588, upload-time = "2026-08-07T11:28:36.318Z" }, + { url = "https://files.pythonhosted.org/packages/8c/80/7e0fd2e2e2aba257820db4a8657c4c356844d36b914b20a4af294bcfb902/ast_serialize-0.8.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f9caa63fad8241257ae401b5ff0a64026c6adb36b8e86cbe8782d9ea505daf6", size = 1234575, upload-time = "2026-08-07T11:28:37.772Z" }, + { url = "https://files.pythonhosted.org/packages/b0/6a/3bae0af06f9b1bae3001c44d64215f5b567877e7aae9ffd45db11c3a7647/ast_serialize-0.8.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3926fa117b5e65019853a2969966d11c7175af377a3425991f3fe73784412405", size = 1236015, upload-time = "2026-08-07T11:28:39.14Z" }, + { url = "https://files.pythonhosted.org/packages/6f/c4/ce2d41a1bc22508e82618901f7e10f2a5e2f9556553fea90624daf9875e2/ast_serialize-0.8.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:485f1113af805e9e170b95ef993ca3fbd4f89c04bab25c58b4fc632d854801ab", size = 1432808, upload-time = "2026-08-07T11:28:40.664Z" }, + { url = "https://files.pythonhosted.org/packages/1a/90/f5058f209756dd70e958b7538aaa82d25d24944baf9ec8ae6f27b06fcacc/ast_serialize-0.8.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3ccebbed24f1281062d5852353c72c47502955926cfcb8345ffb3a44d87ff3d3", size = 1256251, upload-time = "2026-08-07T11:28:42.223Z" }, + { url = "https://files.pythonhosted.org/packages/bf/32/7f77ea87fa0836daab706ed5cb7f903bb25fa26a77439011aee626af11d8/ast_serialize-0.8.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:252f883290d1cdb728eb7fe1d9a7221b88af5a329aae0bc91ddee4dafb820331", size = 1258574, upload-time = "2026-08-07T11:28:43.751Z" }, + { url = "https://files.pythonhosted.org/packages/eb/5a/75b82ad2725b5e8e8c742732f9e76c6738a292d0709e1f60d10a973730b4/ast_serialize-0.8.0-cp39-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:96abc072ad29db8d02194afd47d68987322622787daceae82398d7b69f3ba2e6", size = 1254075, upload-time = "2026-08-07T11:28:45.28Z" }, + { url = "https://files.pythonhosted.org/packages/4e/54/8c20ed4eea805516a3fd23dd4a721ce28c64f50f0e4b359969f60a8c97a6/ast_serialize-0.8.0-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:9118ad3e369727060b2696fc4078f250ecffca4248ba87f537f55cea9f9dce06", size = 1301018, upload-time = "2026-08-07T11:28:46.851Z" }, + { url = "https://files.pythonhosted.org/packages/cb/5b/9f14430f12fe830b656fb38f8e2e05ee13b02a88967660bef46af0ab22a8/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f359df4bd921918af8bebd142a376c77511d7151cc8ba852760b587b5a4a54f3", size = 1409951, upload-time = "2026-08-07T11:28:48.312Z" }, + { url = "https://files.pythonhosted.org/packages/2d/3d/084882eca93c842bd4262591a071ec7f825340644035e51501208cc5a8d4/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:e94f9121d13fa36cbf21314783c77d05ae3a0868decd18cf5233fdcc6de49ac8", size = 1509544, upload-time = "2026-08-07T11:28:49.847Z" }, + { url = "https://files.pythonhosted.org/packages/ce/73/ea84852096c2036c61cc0b2f97b90242207419f534dc671060ee1c8e05cb/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:54f95b486018d262bcb387a9afd96f0da74508b442762b80c769454a6fbb3ee3", size = 1505671, upload-time = "2026-08-07T11:28:51.239Z" }, + { url = "https://files.pythonhosted.org/packages/cb/88/287b9a5300c1f2f651d259f670931b63110adc265b7613c885b44c5bc53d/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4c38b915511e32bc718c49dbce98ff9af36bac0ad6a604f58000cd5e3aecdba7", size = 1563685, upload-time = "2026-08-07T11:28:53.112Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f3/1bc3a79afcf0c2a8d2c37182d0d659d1545a9d7f7f6dc9cf3e63d6c17135/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:9a2ef9cf12f2de4f1028c42c1dd7d775255e0fb3e5bb48896c97e35ef52366fe", size = 1427977, upload-time = "2026-08-07T11:28:54.418Z" }, + { url = "https://files.pythonhosted.org/packages/5c/cd/440c798957e14e31776bfeb024d8fafe0bb1d5b89c51c2f067e69938f7b0/ast_serialize-0.8.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6f18048fe9f6dd266bd577cdec48bdcecb74faaa01fe941324435483b013ed2a", size = 1454335, upload-time = "2026-08-07T11:28:55.968Z" }, + { url = "https://files.pythonhosted.org/packages/4f/4a/587eb36dcc240a54c8660f599464516b469ecad96f0dbdb6bccbedb50745/ast_serialize-0.8.0-cp39-abi3-win32.whl", hash = "sha256:31883542dd6c94d178f5db3d32fbd69c5eb88b3a7c018e7ac8cc0c45195ddbed", size = 1068858, upload-time = "2026-08-07T11:28:57.541Z" }, + { url = "https://files.pythonhosted.org/packages/5f/a4/3e887bbd92164e183cb6e412c6a3e9198ddd446d7fe405958293ef5ef49c/ast_serialize-0.8.0-cp39-abi3-win_amd64.whl", hash = "sha256:861794565b06337005c1447ef23103a3d5a627d08bdc827870d00d0b28ef5f51", size = 1111839, upload-time = "2026-08-07T11:28:59Z" }, + { url = "https://files.pythonhosted.org/packages/25/6c/b400476d3ceba681ab929787edc9554f6d88fcc69435eb681b00fc0457a5/ast_serialize-0.8.0-cp39-abi3-win_arm64.whl", hash = "sha256:b2a5978662fd4db463dfb4b974d2b10ac6430b98f5333aabc7051909df3561d0", size = 1083655, upload-time = "2026-08-07T11:29:00.349Z" }, +] + [[package]] name = "asttokens" version = "3.0.1" @@ -680,17 +747,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -982,17 +1049,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] sdist = { url = "https://files.pythonhosted.org/packages/ae/b6/03bb70946330e88ffec97aefd3ea75ba575cb2e762061e0e62a213befee8/docutils-0.22.4.tar.gz", hash = "sha256:4db53b1fde9abecbb74d91230d32ab626d94f6badfc575d6db9194a49df29968", size = 2291750, upload-time = "2025-12-18T19:00:26.443Z" } @@ -1123,6 +1190,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, ] +[[package]] +name = "hatchling" +version = "1.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "pathspec" }, + { name = "pluggy" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "trove-classifiers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/e2/dfa73fe78f773018dcaebc6d09b819bc10d328ff5a6b4a66efa1e3d71f52/hatchling-1.31.0.tar.gz", hash = "sha256:6b48ad4068a482ed7239b3a8215bc55b47aad3345d58dfc94e553c5d2d46211b", size = 57208, upload-time = "2026-07-08T01:48:32.237Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/e2/2c0af0a52d16be74a4f194564fcdc417521ed863e9b65e4bc9052dacba6f/hatchling-1.31.0-py3-none-any.whl", hash = "sha256:aac80bec8b6fe35e8480f1c335be8910fa210a0e6f735a139be205dadcacb544", size = 77747, upload-time = "2026-07-08T01:48:31.024Z" }, +] + [[package]] name = "httpcore" version = "1.0.9" @@ -1267,17 +1350,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -1570,7 +1653,7 @@ wheels = [ [[package]] name = "jupyter-server" -version = "2.20.0" +version = "2.18.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -1593,9 +1676,9 @@ dependencies = [ { name = "traitlets" }, { name = "websocket-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6b/dc/db3a582633170186f8c8b31298d7eb26ad0eb031a1f53476c258b64eed05/jupyter_server-2.20.0.tar.gz", hash = "sha256:b5778ba337d8015a3dc2b80803ecdd5ac18d3797fddf61a50ea5fb472b4ebe14", size = 756523, upload-time = "2026-06-17T12:09:09.435Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/ec/9302cec1ccacdd33c1b1312ac31681c8975cae56c626d783ab49edf9c681/jupyter_server-2.18.0.tar.gz", hash = "sha256:568b27bce4320a53c3eebf1bdcbee9acf48a8ab7f66ec83d900ca9909d4fb770", size = 751152, upload-time = "2026-05-04T13:39:29.685Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f3/71/8c002223e873a870f5c41dc69b0a7c922301123e4a31d5d01ecb700aef77/jupyter_server-2.20.0-py3-none-any.whl", hash = "sha256:c3b67c93c471e947c18b5026f04f21614218adb706df8f48227d3ee8e0a7cdcc", size = 393143, upload-time = "2026-06-17T12:09:07.234Z" }, + { url = "https://files.pythonhosted.org/packages/cf/f9/050312d92072ddb9ce14c11171804c07435790c98d4350935a780d9e10c2/jupyter_server-2.18.0-py3-none-any.whl", hash = "sha256:69a5397a039d689da81a45955f9b23e95ee167f6d8a8d64372fb616f2aac650a", size = 391687, upload-time = "2026-05-04T13:39:27.549Z" }, ] [[package]] @@ -1613,7 +1696,7 @@ wheels = [ [[package]] name = "jupyterlab" -version = "4.5.10" +version = "4.5.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "async-lru" }, @@ -1630,11 +1713,10 @@ dependencies = [ { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "tornado" }, { name = "traitlets" }, - { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/74/24/621aa20ec0d2fe72f52095bda0fc1be7738ac21aabe4129ff623140d5cdf/jupyterlab-4.5.10.tar.gz", hash = "sha256:77e8d80b78be59b2eaba2154562e21caa6e79c2f1281d6f486584f7144ee2f47", size = 23998879, upload-time = "2026-07-21T12:43:27.324Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/22/8440ec827762146e7cdecf04335bd348795899d29dc6ae82238707353a2c/jupyterlab-4.5.7.tar.gz", hash = "sha256:55a9822c4754da305f41e113452c68383e214dcf96de760146af89ce5d5117b0", size = 23992763, upload-time = "2026-04-29T16:43:51.328Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/c9/940f95f17ee4e413ad252bf8d4f2ee9a341f18cfeda87775fef3d7847321/jupyterlab-4.5.10-py3-none-any.whl", hash = "sha256:5967ca61e692e67a2f30b5a2b901c941dc6ce56c0b0e357bc6d34fed5ec095f6", size = 12452502, upload-time = "2026-07-21T12:43:23.542Z" }, + { url = "https://files.pythonhosted.org/packages/3d/aa/537b8f7d80e799af19af35fb3ddfc970b951088a13c57dd9387dcfbb7f61/jupyterlab-4.5.7-py3-none-any.whl", hash = "sha256:fba4cb0e2c44a52859669d8c98b45de029d5e515f8407bf8534d2a8fc5f0964d", size = 12450123, upload-time = "2026-04-29T16:43:46.639Z" }, ] [[package]] @@ -1799,6 +1881,105 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b5/40/23569737873cc9637fd488606347e9dd92b9fa37ba4fcda1f98ee5219a97/latexcodec-3.0.1-py3-none-any.whl", hash = "sha256:a9eb8200bff693f0437a69581f7579eb6bca25c4193515c09900ce76451e452e", size = 18532, upload-time = "2025-06-17T18:47:30.726Z" }, ] +[[package]] +name = "librt" +version = "0.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/36/9b/356320fbae2ac8467e21c5e73e1389c80468e4998c62cc7d3536cc51b614/librt-0.15.0.tar.gz", hash = "sha256:4e66cbe84437497d951b799d3e1551291b6fb3d643820a7014b3655d57a59162", size = 214338, upload-time = "2026-08-07T10:49:42.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/48/12/e2e9ca532cf5a0e08c9489826c4a35c6958c92ba0313fda70e8c6c3912be/librt-0.15.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e1a49adf16a7c9d9646816c2946135527197b6fcf4347c7b8b761cf1bfbf4489", size = 148673, upload-time = "2026-08-07T10:46:22.569Z" }, + { url = "https://files.pythonhosted.org/packages/6d/7c/02005e23478bd5950618d9712e0fd2b4c511657857f3efd8ba6a5feabcdd/librt-0.15.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:81a398f45b45a59200e13cd5ad1ae1d3f44334de98b148331afe2cdfee701c52", size = 153547, upload-time = "2026-08-07T10:46:23.931Z" }, + { url = "https://files.pythonhosted.org/packages/a0/90/d8848a735f5642077fc4b3b4bebcdb08edf10178e3add45597f5201a368f/librt-0.15.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4eafbaff06b9563f8b1c850621ce51605de05208e09d4d71ce490bc972b7b9e8", size = 494355, upload-time = "2026-08-07T10:46:25.122Z" }, + { url = "https://files.pythonhosted.org/packages/e1/0b/8604f41ea02feace490e9e405a338a15f9905369f55b239a9ce31c946f24/librt-0.15.0-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:b0411b4066db926b80258c60dcb0e6db4c9cee312eab45b7e8866b17ddf9ada1", size = 485459, upload-time = "2026-08-07T10:46:26.447Z" }, + { url = "https://files.pythonhosted.org/packages/a0/ac/84153bda1ce0da609182527ab92b40d961809e544eefdc5a1c2422971416/librt-0.15.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:febb1ce6cac545a54e6b769982824e955a700fdd9fbf3a08a3d82c990968b57d", size = 498398, upload-time = "2026-08-07T10:46:27.701Z" }, + { url = "https://files.pythonhosted.org/packages/2c/3a/5ca6cd282b2c244bec8ec84102e09773264e9c02891d56ab3a8f0e4d7083/librt-0.15.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b230acc1c3bfe2d6f2627ba2b95dc92e58aa494600e9722d0e6ccbc931e59702", size = 515474, upload-time = "2026-08-07T10:46:28.9Z" }, + { url = "https://files.pythonhosted.org/packages/73/d3/bd34110234779eb843c6ed66aba7c9b2091d3dd85989f1fb9922f564cb7a/librt-0.15.0-cp310-cp310-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:6da110e5f314c19ab8478464d02ae18808ae73d522c15260fa4918acdcd64da9", size = 509484, upload-time = "2026-08-07T10:46:30.124Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6c/43c3f7f071d71631a7daa3b835ef2168ea39f20692d81464d4e47fbaa6d6/librt-0.15.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:eab9208b00ca55bf75983ec99f7bf13acc746a36102e98953addaad7f7ea1e1b", size = 532534, upload-time = "2026-08-07T10:46:31.511Z" }, + { url = "https://files.pythonhosted.org/packages/c5/1c/b854adf036ea817c40408873a5b794d65a91d9f0f39826f2ad2a2d5d7f48/librt-0.15.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:6c013cd3a1721e69e14380ada97eaa4b7b0cdf1c6b96fa765d4ea47c875088db", size = 537087, upload-time = "2026-08-07T10:46:32.734Z" }, + { url = "https://files.pythonhosted.org/packages/25/5c/c9a890e244e7dd725d3bd8b560e41f0aec787eaf343b46956a290ab7b841/librt-0.15.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:567b1c430f8bd560e689421468278ac5941bab4a05303b5d95b6ae10db03f451", size = 536575, upload-time = "2026-08-07T10:46:33.965Z" }, + { url = "https://files.pythonhosted.org/packages/5f/c5/c8e70b60b704299555f55db468eb46b1c81bfc60201ffbfe20407d89870c/librt-0.15.0-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:29c4cab9df457b19672c39be7f384ebb2bc925c4e2684b8780c222b43eb36389", size = 517142, upload-time = "2026-08-07T10:46:35.577Z" }, + { url = "https://files.pythonhosted.org/packages/56/d1/767a90c41f5d381b3195bc88ac0ec4afda35777c9c781e1f9848fedd965e/librt-0.15.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:bccbd8e5b0bffb7106cf18eb1baa3d7194b1cebb3b4b1cdbd4bdb19382a6ee6c", size = 558714, upload-time = "2026-08-07T10:46:36.829Z" }, + { url = "https://files.pythonhosted.org/packages/f9/b4/3c0624b8dc8301ab808f2b3a910995bcabe28df070fb9a0e5505ae997dae/librt-0.15.0-cp310-cp310-win32.whl", hash = "sha256:8ae493ed5f659a7761c43d42f183db514536073ded9bcf671d2d1df47e29a07e", size = 104426, upload-time = "2026-08-07T10:46:38.594Z" }, + { url = "https://files.pythonhosted.org/packages/31/98/e91c0382304bedb2db9c6801897319a9dcb68daac5e975819b562362f20d/librt-0.15.0-cp310-cp310-win_amd64.whl", hash = "sha256:bc25fb356d0c7810bb49ff3df908ad1fda6995d660ab099ded69244ed7ab6053", size = 125057, upload-time = "2026-08-07T10:46:40.052Z" }, + { url = "https://files.pythonhosted.org/packages/59/52/06790ced2ac7117f890c21bda43c39c958ec82aa665c0718e821d33ff939/librt-0.15.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:823b92cf3c18ecd08afc70c42473888b41b6e8ef5046f3b82c05c154a2fa3d22", size = 148039, upload-time = "2026-08-07T10:46:41.165Z" }, + { url = "https://files.pythonhosted.org/packages/e7/1d/8e150b7fc449a1f33c8a760965cc1f43b14fc1577d9d0b50ab2701420e74/librt-0.15.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c70bc1b602cf59917e8f0c7a2cbc8bcc6fbc14d5486136b00707a79619121d63", size = 153067, upload-time = "2026-08-07T10:46:42.418Z" }, + { url = "https://files.pythonhosted.org/packages/51/87/a162bc5a66a35599dc619ecb215145f4de7d68e886b479b6d12593139f7c/librt-0.15.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:814ff83a25b5fce8b9c80c4dd803153fb5c5599fc74db9e022466938368957ef", size = 493087, upload-time = "2026-08-07T10:46:43.657Z" }, + { url = "https://files.pythonhosted.org/packages/e5/3a/aeea1fc620cf48060d3065b37614edbf97043c099d0f50782bc8ca61d897/librt-0.15.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:57f5eeb6ad4c180de583b1038e61fe5fbd9796bb69a8a1c1a0c7ddbec4c8c60f", size = 485608, upload-time = "2026-08-07T10:46:45.038Z" }, + { url = "https://files.pythonhosted.org/packages/52/ff/fe571ad416f0856fd0d5578ffc2e6dc531891e586e36b647bcf50569cab8/librt-0.15.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:82909c8f7eb9952656b65d3147afde4cf8e6d5a991eebc86418b5e65843b0ab8", size = 498723, upload-time = "2026-08-07T10:46:46.35Z" }, + { url = "https://files.pythonhosted.org/packages/0f/e1/7a65eb5dedb1f00aebd948cdd8e17add48bf066cab3514e9daf84ab45a6c/librt-0.15.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f779070399f991400fc451719e0ea388eb7de313388bada2c127a35de05f798a", size = 516002, upload-time = "2026-08-07T10:46:47.599Z" }, + { url = "https://files.pythonhosted.org/packages/5f/45/59832b0ebfbd08c2742e6ece372ceb53f18bf1faef5d33c8daf3abebf749/librt-0.15.0-cp311-cp311-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bac89069bc496ebdf4f79ebb57bbd10d0b214c8454225deb672d91002bd17e18", size = 508607, upload-time = "2026-08-07T10:46:48.873Z" }, + { url = "https://files.pythonhosted.org/packages/ea/0d/37fa73f3b43ebd8259f91ae9102a15e5a54e65d581e48dea72df3e81d7a4/librt-0.15.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e0d00c708fb2f5822b152429b1ac80a58dbbbc3f6c232c4d13a3f7fcf2ea5b4c", size = 530422, upload-time = "2026-08-07T10:46:50.45Z" }, + { url = "https://files.pythonhosted.org/packages/26/02/e046c6fe7a5881ac34623242192f484426ba8a75595fd18f22c53a3f530f/librt-0.15.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:6c6624fe268625869485553dd7cc1daf30d22558215bb2a4ff16f67a9801a31a", size = 534303, upload-time = "2026-08-07T10:46:51.693Z" }, + { url = "https://files.pythonhosted.org/packages/95/32/d5e6d861ab0366f3edf74f887ab0c9eb9f535aaf01d32b80b4f734daa179/librt-0.15.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f56b397858a23dacf35ede366ed2212fdc03a6a57a1ad36468ad6e9dc5fac091", size = 536084, upload-time = "2026-08-07T10:46:52.951Z" }, + { url = "https://files.pythonhosted.org/packages/2a/de/d69d725513fe53fc90c6d7a1f86e4428939bad2fb905b17fe4c18d413dde/librt-0.15.0-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:4388184646efe2054911c5b00a1077d6d1ee86a95b7e8ba96dc7850a809f3f40", size = 514307, upload-time = "2026-08-07T10:46:54.194Z" }, + { url = "https://files.pythonhosted.org/packages/36/93/f8aded0d6682b4f25820fa86e0690f87f01df9fd7bd09ddb04d9167ad021/librt-0.15.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:97335f59082f9fe2ce6c2a9cc6433a0114bbb6cd4d5c09dd76c95c68b9f9a8b0", size = 557686, upload-time = "2026-08-07T10:46:55.443Z" }, + { url = "https://files.pythonhosted.org/packages/74/09/ffeb6bdeb6cd862b4272fddc8ad05f938dd25d020ed517e631813917d80a/librt-0.15.0-cp311-cp311-win32.whl", hash = "sha256:83380ffde38062a2e9bb55d83e74474f6614665528b98a6928720fc006dfffbb", size = 104917, upload-time = "2026-08-07T10:46:56.605Z" }, + { url = "https://files.pythonhosted.org/packages/96/28/7e2313a3ffbf0b4de7ba3da58a09e488507b4bd1ea2b5e69378354a23415/librt-0.15.0-cp311-cp311-win_amd64.whl", hash = "sha256:f75720477ee05d509a310e856cacc8d909adc182f7b91193c207bcc26d7ee6db", size = 125886, upload-time = "2026-08-07T10:46:57.729Z" }, + { url = "https://files.pythonhosted.org/packages/39/9e/04b8c3cde014ef255ee785730425268354543acc38902093a40afa0dc164/librt-0.15.0-cp311-cp311-win_arm64.whl", hash = "sha256:256237037a3ab001ae8d9803b2d43562a4c3aa38739843694349e4d5ebb0fd56", size = 111885, upload-time = "2026-08-07T10:46:58.787Z" }, + { url = "https://files.pythonhosted.org/packages/ba/39/99c25030e782bdfb7a21be8c05254806a2e4bbb05c8d50c2a2130acbfa05/librt-0.15.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e87bc679f86a99aa3b26e3c78eeb821a247c9a28eae48eaafcc32c3bf4c3bb9e", size = 151021, upload-time = "2026-08-07T10:47:00.057Z" }, + { url = "https://files.pythonhosted.org/packages/14/43/f4b1bd1b2888798a1409808889a25ea1ba49eaabce7d681ed27734c2df9d/librt-0.15.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:71599e011ac880e8e45d46047d714871894c7d4ab6f25626f8d4f89da21f368d", size = 155267, upload-time = "2026-08-07T10:47:01.311Z" }, + { url = "https://files.pythonhosted.org/packages/0c/db/3ad9c965c72f1e1d6beeec44ec10a54e17be8ae042fbb4baade16cbadced/librt-0.15.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c802434092b769b1d613ed2e13fac15fbfce1934a74bd10283b03c0fae231cd1", size = 503136, upload-time = "2026-08-07T10:47:02.45Z" }, + { url = "https://files.pythonhosted.org/packages/4b/07/5888a6d76acd62ebce66c61b74d94e9370b9c32929f111e487bb6546f8ed/librt-0.15.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:5500eeae393a184d14e1f35645962c27129d20c81afa4069e6ef826ebc2b3aaa", size = 496670, upload-time = "2026-08-07T10:47:03.675Z" }, + { url = "https://files.pythonhosted.org/packages/29/39/ab57cc2f5b276156da02bb7f5a8921bada1cb1993ffec99acf811c602c23/librt-0.15.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6ecfc32dfb46fb7b565bcd6abf9412acf978775a998273d22888a6d7953730dd", size = 513688, upload-time = "2026-08-07T10:47:04.981Z" }, + { url = "https://files.pythonhosted.org/packages/a7/b9/bdbb0b648b5c2befb031f4c6f3b1dd857415e8fb492a25a3c764a6681e6c/librt-0.15.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cc46cfd15022e35084355478c9ac809d90b1152222706ac9a7655ec21df6fa", size = 531904, upload-time = "2026-08-07T10:47:06.211Z" }, + { url = "https://files.pythonhosted.org/packages/93/26/473c2e4b6c104e9e58e27ce95fc8005c8bd4fc36cae4f254371125a92db8/librt-0.15.0-cp312-cp312-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d5f51401d102c885b9ca509e62c79b1dbff286e1b9b047fde6f763780789356d", size = 524427, upload-time = "2026-08-07T10:47:07.592Z" }, + { url = "https://files.pythonhosted.org/packages/26/60/03b3abb82b41714671b907bf6989b228e31e6a8af52dec82b5b0728dc250/librt-0.15.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:cc30523e3f1a23fb7511cc659834a0d01a1042bb9de359bc1c131cc4ec6c9656", size = 543155, upload-time = "2026-08-07T10:47:08.866Z" }, + { url = "https://files.pythonhosted.org/packages/f2/0e/9bb1f0a4affbd0a1888f4f79dc03ed2a299d9a2c26c59ab2a97dcbf11903/librt-0.15.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:59fe030d8ae4a57e3fb7756bf35a858de74e04066fc8555c53d0af979132af81", size = 546890, upload-time = "2026-08-07T10:47:10.327Z" }, + { url = "https://files.pythonhosted.org/packages/dc/84/6937a280d461f7de6e031ffb02edc2b7c3c90d49d630565ce8ff27cbc5f2/librt-0.15.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:5a6526a2a956bbb1e4ae3568c82e650fc99119c66bb011ea60715744955a2b4d", size = 555163, upload-time = "2026-08-07T10:47:11.798Z" }, + { url = "https://files.pythonhosted.org/packages/bc/95/2a2853c1ee014bf102116e7f897a04beeaeb2461b45b79af98bdfb95f1ef/librt-0.15.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:85ea21ec6730194d67156b0e0b5430ccb1d61f8b8b907e39b37f9812b74a13f0", size = 535812, upload-time = "2026-08-07T10:47:13.279Z" }, + { url = "https://files.pythonhosted.org/packages/c9/4c/cf9601c1b4c5f09280acd5d83abdb2e68527a2be8257136eb42304218622/librt-0.15.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1e47b8ba865d7ede071a91a7163073bbaeb72541f1ef8a07d512c45c7b5007f2", size = 573688, upload-time = "2026-08-07T10:47:14.727Z" }, + { url = "https://files.pythonhosted.org/packages/47/6d/9ac7cbec46189a7625af4b5acbd25f10d827f4141b2002181848c8418923/librt-0.15.0-cp312-cp312-win32.whl", hash = "sha256:a5207ec414d1c4a2a7231b2086970dc036f94293cdf338190984958a013a42f1", size = 106138, upload-time = "2026-08-07T10:47:15.973Z" }, + { url = "https://files.pythonhosted.org/packages/38/d0/2ae99c83be86ce23f925ac1aeeedc777e97f427c4a8d190c70d0a16e9a87/librt-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:73b30cfa976659b3917c8f6153bdb0591c6a9ec6583599fd24a689b690622022", size = 126974, upload-time = "2026-08-07T10:47:17.049Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ef/dd24f9635c730b86b87587967dda7516b1845e8b17684603d31607fed598/librt-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:a54cf9e0ef47b96af580849db5471142200568ce1e02cbf416addab551369570", size = 112292, upload-time = "2026-08-07T10:47:18.222Z" }, + { url = "https://files.pythonhosted.org/packages/e7/42/467b53a601b406ccd7b97c1fd54b59cb34f9185ad5ce7e9d5c3c4e8961c8/librt-0.15.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:db13ca398005abcbe538deda87b686d9bd08b7001cf40c4c06b444960ae10a26", size = 151029, upload-time = "2026-08-07T10:47:19.312Z" }, + { url = "https://files.pythonhosted.org/packages/3e/e6/36c2299b7a94b84fdd01220d8a777a71be5be0925bb0dbdf71c0a06a34d9/librt-0.15.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:aa1f1995789dca3698bc550aaceb09a51bd5df0a057ff84ff15296cd1975b801", size = 155194, upload-time = "2026-08-07T10:47:20.398Z" }, + { url = "https://files.pythonhosted.org/packages/c9/b6/ed5071f9325845e670bd36012757419767fbf56af77ed483077b9e4db541/librt-0.15.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:55456ea87d8df21808446d03817be2f65e20391c1c615d9187440dff28cd08dc", size = 502568, upload-time = "2026-08-07T10:47:21.652Z" }, + { url = "https://files.pythonhosted.org/packages/7f/81/6450c67c3615d87704bcbc21323fafc69c799b06a044c447529f725d4b01/librt-0.15.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:5a86a5a08c2235316bdb359d5dbb6ce0abfca7fac06363103e2c5af571d92f95", size = 496153, upload-time = "2026-08-07T10:47:22.925Z" }, + { url = "https://files.pythonhosted.org/packages/e1/d6/5f52b722bc75076954b3bfd49be15ea362df4d580c6fb315d0f617100d30/librt-0.15.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e56b6a368529bed262da40ce13f8fef590db0479819cca84f16a1f01ac356d0b", size = 513336, upload-time = "2026-08-07T10:47:24.213Z" }, + { url = "https://files.pythonhosted.org/packages/8d/e2/c08fd1d36ce63ea5a12b85c5d37f4550b5f86a692167e41e5a74222607ae/librt-0.15.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:234d8d394721fa0d786af15ebf1f3fb7f3ed82fd1cd0cde45c2f247b5d4281d2", size = 531661, upload-time = "2026-08-07T10:47:25.507Z" }, + { url = "https://files.pythonhosted.org/packages/3f/d8/d9482fcbeb177b9eb87bb3899eeb3b42be690313c652f9e146b1d0681fb2/librt-0.15.0-cp313-cp313-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d8363d7accb0286ac3a0e633f396e93800dafb8150494505daf9515bbda591f3", size = 524487, upload-time = "2026-08-07T10:47:26.79Z" }, + { url = "https://files.pythonhosted.org/packages/10/cc/075171517b41f861753034fbb151b42cfc83bcc853849f24f5e66fd60ccf/librt-0.15.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:0f0ee3644d951f31055ad07d77d92520e84505dd7a432cc4cd501dd70ee06785", size = 543201, upload-time = "2026-08-07T10:47:27.999Z" }, + { url = "https://files.pythonhosted.org/packages/b0/03/42c2330f37eeb475b6affeedd06518f60035f323af3a839335e3fc9fef2d/librt-0.15.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:2cfd1a81a648806e6a7717be4cc4d1bb392fa229752bf8444ba365e381e984d6", size = 546467, upload-time = "2026-08-07T10:47:29.396Z" }, + { url = "https://files.pythonhosted.org/packages/57/1e/1ad4c5638f7e64d8560328bd25c54b409a661bdb6ff254b38ff90744288d/librt-0.15.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:a6cd22c9da0d866558e46a041f1cc0c2bbb26b61b137b2347fa834c332e1d101", size = 555139, upload-time = "2026-08-07T10:47:30.815Z" }, + { url = "https://files.pythonhosted.org/packages/49/41/39fa7d15db1204cd1cbe6514680fbdc243adf754a0885061308f43afc013/librt-0.15.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:6d5225ef8801e4ea5e482fa9b5dfb891dd9ef6f6d870f1f25d449ca2c70ac218", size = 536050, upload-time = "2026-08-07T10:47:32.222Z" }, + { url = "https://files.pythonhosted.org/packages/1e/88/c6dcf0dd8e26dc0c9a499a2abab8646c86dcaf9ecea9524cb46d3686331a/librt-0.15.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d28a05796b99f749bf8794f17ba9ba1612d0076b802e9cfc62c554634e9ce3b", size = 573700, upload-time = "2026-08-07T10:47:33.527Z" }, + { url = "https://files.pythonhosted.org/packages/1b/9b/ab54c71a7918a7c34fa5327fb61390a77446a07a146fbfb1165250a61035/librt-0.15.0-cp313-cp313-pyemscripten_2025_0_wasm32.whl", hash = "sha256:2067ff438048cead9d223ca5675bae2a25e520a7c3e6c1498bf9c6892d22caab", size = 82194, upload-time = "2026-08-07T10:47:34.835Z" }, + { url = "https://files.pythonhosted.org/packages/8d/b2/4f9a243bb892395f3becb80789ade13771701091f9f07ab8230247953ba8/librt-0.15.0-cp313-cp313-win32.whl", hash = "sha256:1cd3b721f24c206398b9e26da3c3a9c011e6e89d06f318ba8ebefc30f1003890", size = 106231, upload-time = "2026-08-07T10:47:36.251Z" }, + { url = "https://files.pythonhosted.org/packages/bf/af/64aff4885a40b93132382f2c314647d722574605416504379184ef3045ea/librt-0.15.0-cp313-cp313-win_amd64.whl", hash = "sha256:f395a4a9a03ac062dbe9a9f82e0c720502e590a38feee6a757bc82e9c63afbd8", size = 126996, upload-time = "2026-08-07T10:47:37.453Z" }, + { url = "https://files.pythonhosted.org/packages/27/83/335bccf6c7cb9028cb0b54aead27d9ece3f01f83bc6baa2abace5da655c1/librt-0.15.0-cp313-cp313-win_arm64.whl", hash = "sha256:0a15cb554761247d84a3ec0cbdf4078d70725384f0e4662c0fa3b26266eb60ad", size = 112188, upload-time = "2026-08-07T10:47:38.729Z" }, + { url = "https://files.pythonhosted.org/packages/a8/93/949053fb462eecc4a9a5ee770a81f4b40be7b79538b245545d4aebc6b58b/librt-0.15.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f5de7feedc56337a088eb15cd9fafa9938367362221d8cc62c642b7f94821993", size = 149833, upload-time = "2026-08-07T10:47:39.86Z" }, + { url = "https://files.pythonhosted.org/packages/61/ca/8281aa6cd560a3420e4497729f6b704b53be3eeaaef82d5aeadddaf7441f/librt-0.15.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6c0eb900c0e91f4aebe680845242e614f1864edfd44106380d0752ac29522bf8", size = 154088, upload-time = "2026-08-07T10:47:41.065Z" }, + { url = "https://files.pythonhosted.org/packages/dd/02/1a1662dceaba6a086360891448d5ce9a7d3555976cae59a31a39d744b9c7/librt-0.15.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8c9a650a188e38bac005048cbe6342e81407782944d01934540ab75e417df21", size = 494215, upload-time = "2026-08-07T10:47:42.388Z" }, + { url = "https://files.pythonhosted.org/packages/69/84/99211619dc656370a3740c33d2b0b6d5a3fb1e73689314f6ed477a397dc4/librt-0.15.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:92bfed8deec93df30286b9fe9e3b1dd17329cc076a192b4ee5ec223841d54953", size = 491173, upload-time = "2026-08-07T10:47:43.683Z" }, + { url = "https://files.pythonhosted.org/packages/d4/aa/5448d0b05f4579b635d3899176817ebf561af0e57bacd425b5b1887264c1/librt-0.15.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ec4b19788f835711a2072f9dbe6b03b3bf32ed1f0fb30cf399bdd59d9f0c33fa", size = 505512, upload-time = "2026-08-07T10:47:45.314Z" }, + { url = "https://files.pythonhosted.org/packages/95/82/01940e40b83c43a546c4a3c896cf34ca272a9690899d55914e4827b3dcce/librt-0.15.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d4c7bacb70930f3d0a56f4ecf1be474a1f0d941b01dd73b756f3c256d42cb879", size = 523073, upload-time = "2026-08-07T10:47:46.66Z" }, + { url = "https://files.pythonhosted.org/packages/88/fa/759c0030f3ee371439eb26de34fc745807caf0abb878af7af4b8b7c3dd3d/librt-0.15.0-cp314-cp314-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3e79f05e4a08b4d880342673312bbc895b56df7765605796f15902eb5367d3ae", size = 515080, upload-time = "2026-08-07T10:47:48.319Z" }, + { url = "https://files.pythonhosted.org/packages/0b/27/894e072228fcb159703c655da69f8cd10dbed489c36e3df7dd032a2483be/librt-0.15.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a417149c0cba4d50b61e992e5a15e69eaf96746609b461cc4ed168aeef6b79dd", size = 534164, upload-time = "2026-08-07T10:47:49.875Z" }, + { url = "https://files.pythonhosted.org/packages/98/a3/0078e91c1f36f8815db17827de15650b9a3fe56c55fbf998c854b34e40d3/librt-0.15.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:da7a94d6a3411f579d72aa3e3bc5fbca7ed4549f3dbd7e5de3aa567333374285", size = 540616, upload-time = "2026-08-07T10:47:51.408Z" }, + { url = "https://files.pythonhosted.org/packages/86/33/81a29b796dd52a45e9ef7974c7732926e8f10f15b8d2be505665979f896d/librt-0.15.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:856f743ae607f2c1380eccb566c0038a9fb3eabf0fc2be2704d76d9f73557239", size = 545890, upload-time = "2026-08-07T10:47:52.818Z" }, + { url = "https://files.pythonhosted.org/packages/05/82/8be1baa1350e5d30cfd70ae79d0a6f4dc5862ef47f7bb2808aabc9bb86e5/librt-0.15.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:779a6e7c894737e5983e7790a9c78c4000c30e23c9aada08081bdbea53b0fa60", size = 523287, upload-time = "2026-08-07T10:47:54.165Z" }, + { url = "https://files.pythonhosted.org/packages/c6/4f/d1be6a01a35c20ef734e0e44113f87d4af756a9354a89dcfbe3b4f8af5e1/librt-0.15.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:96bb17dbe8bab3c0954fbebfc69ed395599de75b6bbc35e3270a878e15d4dd65", size = 565868, upload-time = "2026-08-07T10:47:55.566Z" }, + { url = "https://files.pythonhosted.org/packages/67/88/649cfa33f5825927b160610f670bdab012a64d627eddb94fa795ea4292fd/librt-0.15.0-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:7220697efaa6e5348fc3d18ee7f8563d4bfecd9872b37ffb915bfc1d08840622", size = 81619, upload-time = "2026-08-07T10:47:56.886Z" }, + { url = "https://files.pythonhosted.org/packages/22/31/8e88a8d5e48fc8d1a817787fb6811dfff6499acd6c8683dd83934aa6ede0/librt-0.15.0-cp314-cp314-win32.whl", hash = "sha256:f54598964d357b1c5ab77cf5d92f21e598fe0e23cdbe9618480807f81b4eba15", size = 100138, upload-time = "2026-08-07T10:47:58.093Z" }, + { url = "https://files.pythonhosted.org/packages/80/92/20fd6c4b6a1b1a564b076d55cd3d427d8428217d7638dc25a654cc4791d4/librt-0.15.0-cp314-cp314-win_amd64.whl", hash = "sha256:3ff5893a2c23d886aa9ce786de5ac6ddc74aeeaf90743682b74d920e117d2e28", size = 121258, upload-time = "2026-08-07T10:47:59.564Z" }, + { url = "https://files.pythonhosted.org/packages/fc/28/6af430b44d9ebb897b865a3c363b6dcace51357be2347cc0f8f869656a86/librt-0.15.0-cp314-cp314-win_arm64.whl", hash = "sha256:3722a099730704c9a3d70c879fc0f51daec25fe5f1555672d97bc595abeafb95", size = 106467, upload-time = "2026-08-07T10:48:01.097Z" }, + { url = "https://files.pythonhosted.org/packages/7e/aa/b42bb798942ced219f6d63b27e07f91237887a8d0bd0921666db79a13790/librt-0.15.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:38c0c7d4b6fc06c3324b3f9162c8391bfc4fd9dde53afe1033ce7edb48d5a714", size = 159523, upload-time = "2026-08-07T10:48:02.442Z" }, + { url = "https://files.pythonhosted.org/packages/75/03/1b53cd4ef904e73b1d828a5f90143bf94a2967d7cfff0b9ccf93e12aa9b4/librt-0.15.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8b2fdd7ead3c995c37940a790690660d0ca006c302db26cc51933f6766866fc3", size = 161638, upload-time = "2026-08-07T10:48:03.725Z" }, + { url = "https://files.pythonhosted.org/packages/ac/c4/9f9c9fba097d49e9e694c2b4dc331df31884645ecbc58a93b4b5fc69d2c5/librt-0.15.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2fde98cf1fc4bac144ce23c2c4c017b924ba714509ea9334977b0b27050c837d", size = 701795, upload-time = "2026-08-07T10:48:05.135Z" }, + { url = "https://files.pythonhosted.org/packages/4c/05/0966840bda0380c8ae167b9043c6230202941cc90ea29c48e096964c765e/librt-0.15.0-cp314-cp314t-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:e3b461183c5fa7681b48560f91515f53a953122fb30c71e07abc67d7ddf58c38", size = 682147, upload-time = "2026-08-07T10:48:06.555Z" }, + { url = "https://files.pythonhosted.org/packages/18/af/1c47ca573c30ea47d195aec26133af522fea1104afaace028d7b32247ea8/librt-0.15.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4bbcc257e3babea20a91715c361b24554ec4e8f51aa578568afc230799fe1a19", size = 696397, upload-time = "2026-08-07T10:48:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/2e/0f/1aed6223d4f9f9d1171a8596ff100ea4c3f7699fea7a4ba657c3e60daa6c/librt-0.15.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b845b8d48088fad0cadc84be4b8fda63203be7e9237b71015b3925443c1f35ab", size = 722542, upload-time = "2026-08-07T10:48:09.569Z" }, + { url = "https://files.pythonhosted.org/packages/c6/22/9e3a929aea456c97d69e6ef3884efea56d4807f97399471cc946baebd8af/librt-0.15.0-cp314-cp314t-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b30e600e8f337b9bd7f39b86d9fdfedc73cc46e3d0f745931a23a234220bb7e2", size = 729709, upload-time = "2026-08-07T10:48:11.129Z" }, + { url = "https://files.pythonhosted.org/packages/e9/1b/c327ef6018e3a9ca0b8e7c5eddeeb331ba8f9b76c24e126d37d0f6d62faf/librt-0.15.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:64b0c8c35aa4c4ed79896359f3e0b285cbe4e610042106500da4811c322cc108", size = 752891, upload-time = "2026-08-07T10:48:12.558Z" }, + { url = "https://files.pythonhosted.org/packages/d7/d1/d5f1ea02c56930087009e39db9b70660a663e76c730b27b925d786718457/librt-0.15.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:0da0d94cb802f32a0524653e7201f2cef72d5f700a5407678f5290483d4fcd08", size = 745301, upload-time = "2026-08-07T10:48:14.55Z" }, + { url = "https://files.pythonhosted.org/packages/d9/3c/5f7c585d15ebb2250c73e7c0ee4e9e47be72c65d520c07ddbcdc62037674/librt-0.15.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:4a6369168d371207339b1e50d4532b06a7121586141f82599505a3f315751d47", size = 747921, upload-time = "2026-08-07T10:48:16.453Z" }, + { url = "https://files.pythonhosted.org/packages/7f/52/1443a446486eba966bcbca1696b472e4f210320ec42f490a47f48fbf0fdc/librt-0.15.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:c434e072557ade9cbc642d052c89d031efe47d5c9614523619d0d74a02378e81", size = 727561, upload-time = "2026-08-07T10:48:18.089Z" }, + { url = "https://files.pythonhosted.org/packages/79/91/2270a9380f11725cf83ce1925a5e32dd1dde2be9bba597f25c10a38644e7/librt-0.15.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c7eec6a42018bc1d45763b1c162d3d2bf7c3b9a1b0ed30d3e91dcba390efefcc", size = 774417, upload-time = "2026-08-07T10:48:19.611Z" }, + { url = "https://files.pythonhosted.org/packages/9e/3b/f4b1548d4f5b99186737fe27aec238e9823e8d5d23bf4df007c030689dc5/librt-0.15.0-cp314-cp314t-win32.whl", hash = "sha256:6912fa5e635d74529ac7cdb1bdf6ca3af4453da8d1edbe0110ee1cb4ad407ebf", size = 104381, upload-time = "2026-08-07T10:48:21.048Z" }, + { url = "https://files.pythonhosted.org/packages/80/b6/134afad262def1de04c0843c376d02135f1168af43f22e09a52bd8394727/librt-0.15.0-cp314-cp314t-win_amd64.whl", hash = "sha256:8e11699ed745931c395acd3621b07062e0f840efa6935aad87a64ed0995f0915", size = 127034, upload-time = "2026-08-07T10:48:22.561Z" }, + { url = "https://files.pythonhosted.org/packages/99/5f/1b6846b20572bd699c9e9ec321a5f781845bee477df2aa2a43b28bc40119/librt-0.15.0-cp314-cp314t-win_arm64.whl", hash = "sha256:5d2a91724463bfed4f573cd7a9fdc856d2e230d0c0e5a61416a93481dccd8605", size = 110827, upload-time = "2026-08-07T10:48:23.804Z" }, +] + [[package]] name = "markdown-it-py" version = "3.0.0" @@ -1822,17 +2003,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -2015,6 +2196,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/33/ee4519fa02ed11a94aef9559552f3b17bb863f2ecfe1a35dc7f548cde231/matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76", size = 9516, upload-time = "2025-10-23T09:00:20.675Z" }, ] +[[package]] +name = "maturin" +version = "1.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e7/b3/addd877f871fb1860d46d3a4f206ecb10b946c85846805e6367631926fd3/maturin-1.14.1.tar.gz", hash = "sha256:9d6577a62cd08e0ceba7a0db06fb098e0c9b1b3429bad747a4f3a18215a1b3df", size = 369637, upload-time = "2026-06-19T05:19:49.774Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/f0/97c5a5bd9c71653a066c0976a484eaaae50b9369557838a4176b7b0bdaa5/maturin-1.14.1-py3-none-linux_armv6l.whl", hash = "sha256:522292398945442cdafa9daeb2271b2340fbde57027b818f923f88eab04174f8", size = 10207496, upload-time = "2026-06-19T05:19:09.321Z" }, + { url = "https://files.pythonhosted.org/packages/fe/83/294bca639b0e052f1e2f65199b3db258780c7d4e31408b934c9c974a1379/maturin-1.14.1-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:ffe5ad71f21d1e6603c4dd75f7fee34adf5ed5ebcebb692886549888ebb329ed", size = 19680113, upload-time = "2026-06-19T05:19:13.43Z" }, + { url = "https://files.pythonhosted.org/packages/43/b6/79c881410a3b1c187f7eb3d407aecae646c6a4433d630d72200359015e83/maturin-1.14.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f3306078070c1508fd715b9116070cbcaff5959024272a9f1e6f5cb29768b86c", size = 10169205, upload-time = "2026-06-19T05:19:16.615Z" }, + { url = "https://files.pythonhosted.org/packages/93/9d/44b6f26dcb7f7a04c5501ac2dbb6ca1490150682baa525ca5860504f9eab/maturin-1.14.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:cd457cd88961156e26379e1155bd287cc0ec1c8b2f1582b0660fb31b87c8842d", size = 10188098, upload-time = "2026-06-19T05:19:19.736Z" }, + { url = "https://files.pythonhosted.org/packages/1a/bd/9c0d5d6983905ce2c9edaa073a7e89355a9cf7f396988e05d32f1c37785d/maturin-1.14.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:dfc54ae32e6fcb18302193ab9a30b0b25eefffba994ae13238974805533ef75e", size = 10627576, upload-time = "2026-06-19T05:19:22.713Z" }, + { url = "https://files.pythonhosted.org/packages/e5/33/b096412bd6a7cb399652b260666f901adf88a687181a6dbd6a3f89f0a94e/maturin-1.14.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:a131d912b5267e640bc96d70f4914e10590aed64082ec9abacba7cea52004224", size = 10085181, upload-time = "2026-06-19T05:19:25.69Z" }, + { url = "https://files.pythonhosted.org/packages/56/8d/08c3bf469c38a23c9e6c877e338193001eb604d010fedc08341974e38528/maturin-1.14.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:be18fc568fb76884c0205456336892a75105ec398e6b667cd777c6268bd06d69", size = 10026363, upload-time = "2026-06-19T05:19:28.904Z" }, + { url = "https://files.pythonhosted.org/packages/3a/a4/c4d1a92839f8745ab4aab988a7db884a79d6d710bd3b286fcf9316dece1a/maturin-1.14.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:994a0c8ba3ad8a92b3a9ee1b02645d200d610216b15cff5102b0fe65e8e08666", size = 13321347, upload-time = "2026-06-19T05:19:32.411Z" }, + { url = "https://files.pythonhosted.org/packages/b3/fa/170f04624d03fd07d2a8b1b67de83a127af93aef9eaa425839553347297b/maturin-1.14.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:be80866363e605d137991b491a741a84cde9ae350183c4c85f49690ca9aaaa65", size = 10877609, upload-time = "2026-06-19T05:19:35.448Z" }, + { url = "https://files.pythonhosted.org/packages/61/ad/1ae2e1d0ded282bf2c55ac13f0811d87deb425e200ae64a15785675dede9/maturin-1.14.1-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:5282dffd4b539d2be245f4e5b1a5ab6bc1033b58f4a4872f5833f9d43c954aa4", size = 10417316, upload-time = "2026-06-19T05:19:38.28Z" }, + { url = "https://files.pythonhosted.org/packages/fb/27/bf677183920718da49cd7982d6a3ffc440aad8919329f571d189f81b7bdf/maturin-1.14.1-py3-none-win32.whl", hash = "sha256:1a04de0a20188f95c721b5702eed18140bdcccb28c386797093eca3f62f4d4e0", size = 8931293, upload-time = "2026-06-19T05:19:41.183Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/585adeb9167b08d3cdff0032a938b0e72655c92003df4f52c3f696a1bcc2/maturin-1.14.1-py3-none-win_amd64.whl", hash = "sha256:3c9f94640ecc4895e94abaf834a0684430032c865b2748a36c12461fd9252fdd", size = 10314067, upload-time = "2026-06-19T05:19:44.389Z" }, + { url = "https://files.pythonhosted.org/packages/51/d4/dac8c0720ae246be1700afb6fbdbbea20fe35b13f6570b2f70faa005df77/maturin-1.14.1-py3-none-win_arm64.whl", hash = "sha256:15cea8fcb3ba47dd636f50092bb34baea8b04ac777392f23e6bf8a9a61efb894", size = 9718943, upload-time = "2026-06-19T05:19:47.49Z" }, +] + [[package]] name = "mdit-py-plugins" version = "0.5.0" @@ -2053,14 +2258,14 @@ wheels = [ [[package]] name = "mistune" -version = "3.3.0" +version = "3.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/84/9c/1939635275ec7258e2b43b00dafabc36d89ad11aa7838d375dc1b0e561cb/mistune-3.3.0.tar.gz", hash = "sha256:3074ec4c61b384abe725128e4dcbb483f5a09cc4632012505cdee655d3a113b9", size = 110936, upload-time = "2026-06-21T13:11:39.458Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ca/84/620cc3f7e3adf6f5067e10f4dbae71295d8f9e16d5d3f9ef97c40f2f592c/mistune-3.2.1.tar.gz", hash = "sha256:7c8e5501d38bac1582e067e46c8343f17d57ea1aaa735823f3aba1fd59c88a28", size = 98003, upload-time = "2026-05-03T14:33:22.312Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/76/b90f9d48d43fbd80a79a20d3eab2e5109859c7a56dc663b23187385898f3/mistune-3.3.0-py3-none-any.whl", hash = "sha256:a758e578acda49d8195f9a860b132dae2cf7bf409381393b1c4e6e489a65397b", size = 61250, upload-time = "2026-06-21T13:11:37.938Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7f/a946aa4f8752b37102b41e64dca18a1976ac705c3a0d1dfe74d820a02552/mistune-3.2.1-py3-none-any.whl", hash = "sha256:78cdb0ba5e938053ccf63651b352508d2efa9411dc8810bfb05f2dc5140c0048", size = 53749, upload-time = "2026-05-03T14:33:20.551Z" }, ] [[package]] @@ -2117,6 +2322,60 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] +[[package]] +name = "mypy" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ast-serialize" }, + { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/82/6a/878cc1097d4035f82bd516658d0c528d2a9955bc7b363afcbd0b07fea11b/mypy-2.3.1.tar.gz", hash = "sha256:47c1b1207258513a9d93495f69c8be9de73916186f0e52703e8c461b7a623419", size = 3992554, upload-time = "2026-08-15T03:03:38.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/b9/de8f67e12d721cdcc8ba6cfc440b989a4ba4dfabe4402ae94dfdd8bb30a4/mypy-2.3.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:57a936373fc690c43a8cd7e7e12a35148e4ec5aa7698ad7fc0a9f918bdc5be41", size = 14015541, upload-time = "2026-08-15T03:01:53.104Z" }, + { url = "https://files.pythonhosted.org/packages/f1/8a/9e746ab012c67ed8ea3232a613716c306ee8c0b5682c80d8103b4f04568e/mypy-2.3.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d00d769056bde2f4e69c175071eba45cfb44fa1ed92bdfbfe64a93e0543b0cf0", size = 14248142, upload-time = "2026-08-15T03:02:43.201Z" }, + { url = "https://files.pythonhosted.org/packages/f7/5c/c99ff2d8d0e2c53393e32dfe22d9aa43a5d959d30db46c786dafd24527d3/mypy-2.3.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2166b29228835e1f88ff411e96639e6ca3c7fdde84b62ec211f70f86b4051167", size = 15193309, upload-time = "2026-08-15T03:01:28.714Z" }, + { url = "https://files.pythonhosted.org/packages/64/39/124638f745243faae1ff4b37d5426fe41c0f0454535edc82fe8102b56a3c/mypy-2.3.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:83d36c2924df7426333abe7faf4724a7e1aab0d9fd41625e81b4683034b80c13", size = 15498246, upload-time = "2026-08-15T03:02:46.29Z" }, + { url = "https://files.pythonhosted.org/packages/b2/83/31c0781e243836505c0fb5f4e865487d6df1023e4ad959f4ebd4b84a0226/mypy-2.3.1-cp310-cp310-win_amd64.whl", hash = "sha256:f12fdb70459d0060dea40b29e52163a961b156106d68d57882a6a9f648983a53", size = 11155028, upload-time = "2026-08-15T03:01:39.08Z" }, + { url = "https://files.pythonhosted.org/packages/a0/ab/bc2eb0129e72d7d7d93d5e981a78084a9abefda7efa732a7e02f97d6e27d/mypy-2.3.1-cp310-cp310-win_arm64.whl", hash = "sha256:e099200a1b1b1223a4951f0a90cbff1b8c91b250ba599dab1f7217a628144d90", size = 10151438, upload-time = "2026-08-15T03:02:19.04Z" }, + { url = "https://files.pythonhosted.org/packages/a4/be/c624d4241484f37dc62839e177ab607a9b8b3e96f0866544ca99e8e41d51/mypy-2.3.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:94f04929f1c44c35fb0061e912087edaf504acede963a4a7d00680bd089d8531", size = 13936739, upload-time = "2026-08-15T03:03:26.475Z" }, + { url = "https://files.pythonhosted.org/packages/53/84/e3cf72f90dce5960871c82551c8fba6da05fc1018f79be41c047bd126bdd/mypy-2.3.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f5d716048611e85ca9eefb2e1baa5d73ede389b5820ded260ea27c757d667af8", size = 14166460, upload-time = "2026-08-15T03:01:50.565Z" }, + { url = "https://files.pythonhosted.org/packages/4a/ff/6b97d58aa0f79a5ab9b472db1f6d6df1b11a51d74d0c08ab3760d3a613ba/mypy-2.3.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b091a455111214cb5c9d54a57b9618e9a49f9fe2a42e4e1ac86e9d104ed96ce8", size = 15100476, upload-time = "2026-08-15T03:03:12.079Z" }, + { url = "https://files.pythonhosted.org/packages/da/f0/cbb4b7d2ae3ac635f6b4f2d9b04070b8a92edf50da599d3b39e5ed109001/mypy-2.3.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:df12e20c9efd614738c71b390007ecd0181125afc4ccafca04d78a1d2eed2c01", size = 15347826, upload-time = "2026-08-15T03:03:02.856Z" }, + { url = "https://files.pythonhosted.org/packages/5f/10/91dcdc6f8d43fc08e6a06ab1f9732f3abaaf835ac1b2e67b9dff56910855/mypy-2.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:52eaf3a155f35cf80b40220288c861eb45f14a2340c1f6cbfbdb0feff32879d1", size = 11142615, upload-time = "2026-08-15T03:03:36.316Z" }, + { url = "https://files.pythonhosted.org/packages/3d/8a/28d54535bf4b9aa43b2d8918c2ef660378b9f66b23d78dcee052744ae622/mypy-2.3.1-cp311-cp311-win_arm64.whl", hash = "sha256:9b4eacbee8a69836c06eff6d0dd4e134a07c2b047755b30c08625fe214f322c6", size = 10141145, upload-time = "2026-08-15T03:03:07.406Z" }, + { url = "https://files.pythonhosted.org/packages/85/da/d6effc4f808a842d91edc22535dc9e799d2ff6e91449168b7f47a0771f54/mypy-2.3.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a32bbbb940af990d3be0b8af321c7b6815bb1b3b48142fe7459b9cc5f58959ff", size = 14047547, upload-time = "2026-08-15T03:02:57.707Z" }, + { url = "https://files.pythonhosted.org/packages/e4/e6/478229701dab76f26485fc8ff5d6f241f393da22447400bbc56f6946aebe/mypy-2.3.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ff715e45b2231a8e85de1d163d1b42791e4d7aab8f5145f85fee1b710b735aff", size = 14216515, upload-time = "2026-08-15T03:01:26.496Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fe/7c42327a3b21e84681f691982cbfe43f334a3685f3b683b72c376476c4fa/mypy-2.3.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:858fc57d3d91fa728e33e7ad71def60fc6272694607b306cd3292db53ae39080", size = 15307789, upload-time = "2026-08-15T03:03:31.62Z" }, + { url = "https://files.pythonhosted.org/packages/59/f4/7e597edbe01b5a56fa958ce541302dcaabfed979966f1dffedbea0ea0fc2/mypy-2.3.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:851833db876e7b650f93719c74b7879a08e338979c96054fdfc3bfd90a486355", size = 15548831, upload-time = "2026-08-15T03:03:15.55Z" }, + { url = "https://files.pythonhosted.org/packages/a3/52/cb31e084bc0314a1e384bdd677a4b80e55af04ccac077545e2238b9d320a/mypy-2.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:4c5095a327483591c94e0c8d3ef9e50d4ab1369b541eae007c1f23bc2a41f6bb", size = 11226359, upload-time = "2026-08-15T03:03:29.002Z" }, + { url = "https://files.pythonhosted.org/packages/7a/47/88fcf6217b43fa2da81a8c2611370af18141536a4f0294bbf98b457d456d/mypy-2.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:bbfe022634a2a195406bd469e888d2eaf193b02ba7e607391cd7640374aaae3b", size = 10214707, upload-time = "2026-08-15T03:02:48.807Z" }, + { url = "https://files.pythonhosted.org/packages/de/cf/862010ee800ca9c2bd0c4c0dacf0f092e5411824a09b8f97ad4be8fe250e/mypy-2.3.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:114dff494000f18bd10d5d95d84b8567b26da60279ecbe838131841df20e635d", size = 13964542, upload-time = "2026-08-15T03:02:21.43Z" }, + { url = "https://files.pythonhosted.org/packages/75/5a/3f3a2107b41e3e92e617e25daaee121413b91e9784bea733131ed4fecc5d/mypy-2.3.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c8637731bb5eee3671eb2c3200827aa3564ed8a9309ecee4d1afe77e6d031bdb", size = 14168922, upload-time = "2026-08-15T03:03:00.351Z" }, + { url = "https://files.pythonhosted.org/packages/8b/41/04dc4fe7e63d7820fa4eff272e95157d30cbea921388f3ab3fe77794cd0b/mypy-2.3.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1c80fbc405ed8020f5ff3802dc18cf060197bcdd3fbdd6a26ef2fd34dfdd5226", size = 15244791, upload-time = "2026-08-15T03:02:31.089Z" }, + { url = "https://files.pythonhosted.org/packages/96/fc/c3053b26b9054949285aa868cb6af8c10e7591541cacd79c5dcc06a1fcf9/mypy-2.3.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:84081f538ce27375045c02e3d7f81bd11d853400621ae245d87ce7b6c420ec74", size = 15501627, upload-time = "2026-08-15T03:03:34.128Z" }, + { url = "https://files.pythonhosted.org/packages/70/4e/d77daab008bbc4e5001374d7928f4a260d28f0e6747af444fc4763f7a310/mypy-2.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:e9144ac16fde007096f9563eb2041b4433c2d705c4218edeb79e7e9d01035ee6", size = 11243961, upload-time = "2026-08-15T03:02:11.952Z" }, + { url = "https://files.pythonhosted.org/packages/f0/f8/7eb68c136e4abd30569fe31ef2bfcb7eceae9952cab80017c04cd09f5d0c/mypy-2.3.1-cp313-cp313-win_arm64.whl", hash = "sha256:77ad9529e67dca28e511f5cd5671436584ce91f6d3bac159a353158187b986ac", size = 10213219, upload-time = "2026-08-15T03:02:26.361Z" }, + { url = "https://files.pythonhosted.org/packages/be/c4/42a49d44aeff804edf1b19acce0b49e8bd1a9c57dee9605dd8d980aa43d7/mypy-2.3.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:192abaedf75da1bc0b1cef104927e70ec49c1ef0031cc4825c7ee10a438ed24d", size = 13986778, upload-time = "2026-08-15T03:01:33.69Z" }, + { url = "https://files.pythonhosted.org/packages/45/13/9331fd2dfed7194d66c5304072894a8be3e51e9deda6863c1eceaa35a43d/mypy-2.3.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bf678dffd16efcda2c15cbd30e9ecc0081388e29ea23687a88e686ed92638dc3", size = 14188467, upload-time = "2026-08-15T03:02:40.554Z" }, + { url = "https://files.pythonhosted.org/packages/78/f7/f4a34edab45667c5465855dc585a20e87978ffa8aee711445b7239d120c6/mypy-2.3.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e036f06b41630f4c8a1d48f9ac6aa26acc65f8be089973f5519da643318f03f", size = 15225538, upload-time = "2026-08-15T03:03:09.761Z" }, + { url = "https://files.pythonhosted.org/packages/40/05/534b3590757bd05794f73e07f6666c2a77b8597ffed795c94ce570096aa0/mypy-2.3.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:71af9c8a894e862b58e92abb08e53b05a384a1e5e5d6dc7cda59126211a53d82", size = 15480805, upload-time = "2026-08-15T03:01:41.134Z" }, + { url = "https://files.pythonhosted.org/packages/55/da/bdfba852e2562f599624af5bb7d29e36b0b4f526f2b8bac85efe0dd1803d/mypy-2.3.1-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:3c80cd23d85368bdd9f37d5231dfd97d35bcbf5bf41af96ef3a9b078ad1957f9", size = 7761712, upload-time = "2026-08-15T03:02:36.008Z" }, + { url = "https://files.pythonhosted.org/packages/98/31/60fc64a74cdba4f2a5d642d32317993e479163e1ac7d91b695e5d15e2264/mypy-2.3.1-cp314-cp314-win_amd64.whl", hash = "sha256:4956f34d145e145562a0a0bf367f642bbc85c04ec2baf47ae015947c3169a85d", size = 11423968, upload-time = "2026-08-15T03:02:06.931Z" }, + { url = "https://files.pythonhosted.org/packages/a9/23/eb5950b24cd26ba3b78f87707a275568d633c77dae8e61c9661be6055ca6/mypy-2.3.1-cp314-cp314-win_arm64.whl", hash = "sha256:cfb12e360242d23d91f5e978d94f58ea66acf5804c4fb6f2f794a20d4cb1b595", size = 10399323, upload-time = "2026-08-15T03:02:33.671Z" }, + { url = "https://files.pythonhosted.org/packages/82/c7/f80f4e46c0b9a00eb5f78a79d49dda8bdf56a5230f7257fb33e76be04da7/mypy-2.3.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e5f1c50bb05b64e2026b52867e8d21106f01313c744a2c4ecc34c90d12e8d6e2", size = 15121308, upload-time = "2026-08-15T03:01:46.053Z" }, + { url = "https://files.pythonhosted.org/packages/5d/74/9b04f17c7074cc5188f02fb63a2ca1d43fedf479e84fe3091c39061a1d7f/mypy-2.3.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:667196b352f4cf304ded4c10f90cfc179263a1acfb3cdcfa984bdfd340d498bc", size = 15536590, upload-time = "2026-08-15T03:01:35.941Z" }, + { url = "https://files.pythonhosted.org/packages/26/04/c837ef6208e567774e2ed1f863f8ba6ec4817b1b6dd426315e5d559b6ec9/mypy-2.3.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9c53e395c12cad2c6d4b67d5da7c6057638a132d85c08b73646b18f802a0045", size = 16791074, upload-time = "2026-08-15T03:01:31.073Z" }, + { url = "https://files.pythonhosted.org/packages/37/68/48730230afa45192d5bd429a6a2ff24a6f8dedda90fdf2b221792b54518f/mypy-2.3.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:18162b128c3f9c703cd35f5537446900b0d21a2549aa7a95d21380d2ef643fb0", size = 17069183, upload-time = "2026-08-15T03:02:28.566Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ea/ca23fc9c20eeda09a15c9cbcf50015d0e73f409f6ead059e42aa69a608ff/mypy-2.3.1-cp314-cp314t-win_amd64.whl", hash = "sha256:30c0477d4aab7b7f39c8397dc877f2c96b9fe5588ec379f372c56eb63d599f63", size = 12154679, upload-time = "2026-08-15T03:02:04.809Z" }, + { url = "https://files.pythonhosted.org/packages/3b/67/8d982126034990869466f73b8db80dcb2234a7ac39b4dad093e047a79835/mypy-2.3.1-cp314-cp314t-win_arm64.whl", hash = "sha256:6941ab3619377bc3f32ca02876b07d27f216f5201604b664d3937ea0fdd23bb4", size = 10969159, upload-time = "2026-08-15T03:02:38.152Z" }, + { url = "https://files.pythonhosted.org/packages/8e/41/9675c7a1e78edecfba0b79e587a52594c56e189368261dc7b3a7fffb9527/mypy-2.3.1-py3-none-any.whl", hash = "sha256:6ed5c7e3419083268e5c9258bd1c1ef91af44a9e89374dbcaf37b775716e72eb", size = 2754338, upload-time = "2026-08-15T03:02:53.4Z" }, +] + [[package]] name = "mypy-extensions" version = "1.1.0" @@ -2154,17 +2413,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -2441,17 +2700,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] sdist = { url = "https://files.pythonhosted.org/packages/57/fd/0005efbd0af48e55eb3c7208af93f2862d4b1a56cd78e84309a2d959208d/numpy-2.4.2.tar.gz", hash = "sha256:659a6107e31a83c4e33f763942275fd278b21d095094044eb35569e86a21ddae", size = 20723651, upload-time = "2026-01-31T23:13:10.135Z" } @@ -2529,6 +2788,49 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/e5/b7d20451657664b07986c2f6e3be564433f5dcaf3482d68eaecd79afaf03/numpy-2.4.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:be71bf1edb48ebbbf7f6337b5bfd2f895d1902f6335a5830b20141fc126ffba0", size = 12502577, upload-time = "2026-01-31T23:13:07.08Z" }, ] +[[package]] +name = "numpy-typing-compat" +version = "20251206.2.4" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" } }, +] +sdist = { url = "https://files.pythonhosted.org/packages/42/5f/29fd5f29b0a5d96e2def96ecba3112fc330ecd16e8c97c2b332563c5e201/numpy_typing_compat-20251206.2.4.tar.gz", hash = "sha256:59882d23aaff054a2536da80564012cdce33487657be4d79c5925bb8705fcabc", size = 5011, upload-time = "2025-12-06T20:02:04.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/7c/5c2892e6bc0628a2ccf4e938e1e2db22794657ccb374672d66e20d73839e/numpy_typing_compat-20251206.2.4-py3-none-any.whl", hash = "sha256:a82e723bd20efaa4cf2886709d4264c144f1f2b609bda83d1545113b7e47a5b5", size = 6300, upload-time = "2025-12-06T20:01:57.578Z" }, +] + +[[package]] +name = "numpy-typing-compat" +version = "20260602.2.4" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" } }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fc/73/e331473d3db84a8e8883ac07bfd63a8ce9eb7196acbb672bda1f5b8d3294/numpy_typing_compat-20260602.2.4.tar.gz", hash = "sha256:e4eb661f312a7ad5805677967d5879e04fd7b97627fe910121ce7b1f43aa748c", size = 4603, upload-time = "2026-06-02T15:52:38.572Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/a8/94811eedac4cef5ef7df4b24e06715fa371724782e86a4573f5d172c8473/numpy_typing_compat-20260602.2.4-py3-none-any.whl", hash = "sha256:78d33917d5f6921f8d1c549db347a5b8d9768853e36796b08208c81f5b620977", size = 5879, upload-time = "2026-06-02T15:52:33.214Z" }, +] + [[package]] name = "opt-einsum" version = "3.4.0" @@ -2538,6 +2840,76 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, ] +[[package]] +name = "optype" +version = "0.9.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/3c/9d59b0167458b839273ad0c4fc5f62f787058d8f5aed7f71294963a99471/optype-0.9.3.tar.gz", hash = "sha256:5f09d74127d316053b26971ce441a4df01f3a01943601d3712dd6f34cdfbaf48", size = 96143, upload-time = "2025-03-31T17:00:08.392Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/d8/ac50e2982bdc2d3595dc2bfe3c7e5a0574b5e407ad82d70b5f3707009671/optype-0.9.3-py3-none-any.whl", hash = "sha256:2935c033265938d66cc4198b0aca865572e635094e60e6e79522852f029d9e8d", size = 84357, upload-time = "2025-03-31T17:00:06.464Z" }, +] + +[[package]] +name = "optype" +version = "0.17.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9b/86/e6f1f6f3487492dfcf3b7a2d4e2534d27af6ac05b364b276706906c34865/optype-0.17.1.tar.gz", hash = "sha256:07bfa32b795dea28fba8605a6288d36370d072f25183fb9c29b5a90f4b6f5638", size = 53572, upload-time = "2026-05-17T22:13:28.725Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/d4/c6a2b043e33f0dd012486dcebe0593585588d400175d22aad42049c88321/optype-0.17.1-py3-none-any.whl", hash = "sha256:82f2508ca31cb21e53a41648482d890fe1f5c6cb153720551af41161555adaf1", size = 65954, upload-time = "2026-05-17T22:13:27.549Z" }, +] + +[package.optional-dependencies] +numpy = [ + { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy-typing-compat", version = "20251206.2.4", source = { registry = "https://pypi.org/simple" } }, +] + +[[package]] +name = "optype" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/99/51/51dc9b1009e020f44703933d4d1ee3429c647c026ce7806b37ee2b257998/optype-0.18.0.tar.gz", hash = "sha256:ea10dee61b15ca299ed0d97025d362585c4dfc5481159bb999a1d0d414bbcb04", size = 59967, upload-time = "2026-06-07T22:13:17.534Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/91/2b064a117cb2593bc3eb04ee994134ca2a6bf2162c92961769947e3258cc/optype-0.18.0-py3-none-any.whl", hash = "sha256:91822ed8516e7a4f225ba53d30f291776c35f31553fb952d7cda98286679c5a6", size = 73410, upload-time = "2026-06-07T22:13:16.324Z" }, +] + +[package.optional-dependencies] +numpy = [ + { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy-typing-compat", version = "20260602.2.4", source = { registry = "https://pypi.org/simple" } }, +] + [[package]] name = "overrides" version = "7.7.0" @@ -2628,17 +3000,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -2738,84 +3110,109 @@ wheels = [ [[package]] name = "pillow" -version = "12.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/c2/669d88644cddb1485bd9534e63e8cf476c8e51cb3c3a1297677023505c0e/pillow-12.3.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:6c0016e7b354317c4e9e525b937ac8596c38d2d232b419529b9cd7a1cd46e39a", size = 5392418, upload-time = "2026-07-01T11:53:27.808Z" }, - { url = "https://files.pythonhosted.org/packages/6b/ba/3762f376a2948e3036488d773a146e0ae6ecc2ca03ac20e2615bd0b2ba02/pillow-12.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:bcc33feacfaefce60c12fd500a277533bdc02b10a19f7f6d348763d8140bbba7", size = 4785287, upload-time = "2026-07-01T11:53:29.761Z" }, - { url = "https://files.pythonhosted.org/packages/07/50/b5d688cc9c52d4482f3d5bcab6ce20bc2a74a85d2343841c907444a3be2c/pillow-12.3.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5594fc43d548a7ed94949d139aa1341b270f1863f11cfd37f5a6c8b778a6b67f", size = 6253754, upload-time = "2026-07-01T11:53:32.298Z" }, - { url = "https://files.pythonhosted.org/packages/4e/89/36f4cd76cf4baf05c50ababb976249153f18c959171c7f6ba09a6f217260/pillow-12.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f0606c8bf2cdefea14a43530f7657cbbb7ecf1c4222512492ef4a4434a9501ec", size = 6925605, upload-time = "2026-07-01T11:53:34.487Z" }, - { url = "https://files.pythonhosted.org/packages/eb/c0/4de58cf6633b9e3a6061ef4be6fb91fc3c90b812ece886f531e3c523d777/pillow-12.3.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:85f998ea1848bc6757289e739cfbdda3a04adfd58b02fc018ce54d754a5ce468", size = 6327788, upload-time = "2026-07-01T11:53:36.433Z" }, - { url = "https://files.pythonhosted.org/packages/87/3c/14d53682a19550dbbaf3b598f807d5457646c510805a44c7d7891cd1cd1a/pillow-12.3.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:25b9b82bb22e6e2b3cd07b39c68b7b862001226cb3dff7130d1cb914121b39ed", size = 7036288, upload-time = "2026-07-01T11:53:38.712Z" }, - { url = "https://files.pythonhosted.org/packages/38/1d/36279e3c77efe034e4cc2b0393ee74ffdb5a62391dacbf9b916154f5f0b8/pillow-12.3.0-cp310-cp310-win32.whl", hash = "sha256:37dc8f7bbb66efe481bb60defacef820c950c24713fb44962ed6aa2a50966de1", size = 6472396, upload-time = "2026-07-01T11:53:40.781Z" }, - { url = "https://files.pythonhosted.org/packages/48/7c/8fa0039574c476d7c6fa57dd7c32a130436877c6ec1e5ce1cc8ec44878c1/pillow-12.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:300557495eb45ebb8aec96c2da9c4be642fbf7cd937278b4013ba894ea8eb0eb", size = 7226887, upload-time = "2026-07-01T11:53:42.764Z" }, - { url = "https://files.pythonhosted.org/packages/fa/17/e324be141d173c1c919428066c3259f21c1b8982e564e01a4a81e96dbdcf/pillow-12.3.0-cp310-cp310-win_arm64.whl", hash = "sha256:514435a37670e3e5e08f3945b68718b6ed329bb84367777e16f9f4dfe1e61a0f", size = 2568039, upload-time = "2026-07-01T11:53:45.372Z" }, - { url = "https://files.pythonhosted.org/packages/fb/c8/0a78b0e02d7ac54bc03e5321c9220da52f0c2ea83b21f7c40e7f3169c502/pillow-12.3.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:00808c5e14ef63ac5161091d242999076604ff74b883423a11e5d7bbb38bf756", size = 5392415, upload-time = "2026-07-01T11:53:47.162Z" }, - { url = "https://files.pythonhosted.org/packages/b2/5b/a02d30018abd97ced9f5a6c63d28597694a00d066516b9c1c6de45859fc9/pillow-12.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:37d6d0a00072fd2948eb22bce7e1475f34569d90c87c59f7a2ec59541b77f7a6", size = 4785266, upload-time = "2026-07-01T11:53:49.079Z" }, - { url = "https://files.pythonhosted.org/packages/c8/98/766667a4be768150a202836acd9fad19c06824ca86c4286d3cf6b274964e/pillow-12.3.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bcb46e2f9feff8d06323983bd83ed00c201fdcab3d74973e7072a889b3979fcd", size = 6263814, upload-time = "2026-07-01T11:53:51.32Z" }, - { url = "https://files.pythonhosted.org/packages/3b/2d/ede717bc1144f63886c21fd349bb95860b0d1a21149ff16f2bb362b612b6/pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23d27a3e0307ec2244cc51e7287b919aa68d097504ebe19df4e76a98a3eea5bd", size = 6934408, upload-time = "2026-07-01T11:53:53.487Z" }, - { url = "https://files.pythonhosted.org/packages/a3/48/9c58b685e69d49c31af6c8eb9012055fab7e665785165c84796e2c73ce72/pillow-12.3.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4f883547d4b7f0495ebe7056b0cc2aea76094e7a4abc8e933540f3271df27d9c", size = 6337160, upload-time = "2026-07-01T11:53:55.457Z" }, - { url = "https://files.pythonhosted.org/packages/ff/fa/dc2a5c0ba6df93f67c31d34b808b7ce440b40cdbf96f0b81cde1d1e6fa93/pillow-12.3.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:236ff70b9312fb68943c703aa842ca6a758abfa45ac187a5e7c1452e96ef72b5", size = 7045172, upload-time = "2026-07-01T11:53:57.736Z" }, - { url = "https://files.pythonhosted.org/packages/86/a5/444817a4d4c4c2417df00513086ca196f388d8f9ef40c2e4ccd1ad1af54b/pillow-12.3.0-cp311-cp311-win32.whl", hash = "sha256:10e41f0fbf1eec8cfd234b8fe17a4caac7c9d0db4c204d3c173a8f9f6ef3232b", size = 6472232, upload-time = "2026-07-01T11:53:59.767Z" }, - { url = "https://files.pythonhosted.org/packages/63/c6/4bad1b18d132a50b27e1365e1ab163616f7a5bb56d330f66f9d1d9d4f9d4/pillow-12.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:8e95e1385e4998ae9694eeaa4730ba5457ff61185b3a55e2e7bea0880aef452a", size = 7233653, upload-time = "2026-07-01T11:54:02.066Z" }, - { url = "https://files.pythonhosted.org/packages/fd/16/00f91ab7760dc842f5aad55217e80fc4a7067a0604535249bc8a2d6d9870/pillow-12.3.0-cp311-cp311-win_arm64.whl", hash = "sha256:ebaea975e03d3141d9d3a507df75c9b3ec90fa9d2ffd07567b3a978d9d790b26", size = 2568195, upload-time = "2026-07-01T11:54:04.622Z" }, - { url = "https://files.pythonhosted.org/packages/37/bf/fb3ebff8ddcb76aac5a01389251bbbb9519922a9b520d8247c1ca864a25d/pillow-12.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965", size = 5345969, upload-time = "2026-07-01T11:54:06.397Z" }, - { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, - { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, - { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, - { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, - { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, - { url = "https://files.pythonhosted.org/packages/20/20/25e0f4dc178a6bc0696793720055519a0de89e7661dae886992decbd2f81/pillow-12.3.0-cp312-cp312-win32.whl", hash = "sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f", size = 6472684, upload-time = "2026-07-01T11:54:19.839Z" }, - { url = "https://files.pythonhosted.org/packages/45/89/da2f7971a317f83d807fdd4065c0af40208e59e692cc43d315a71a0e96d1/pillow-12.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09", size = 7227137, upload-time = "2026-07-01T11:54:22.025Z" }, - { url = "https://files.pythonhosted.org/packages/de/47/4845a0a6c0dbf1db8456bd9fc791f13c5ced7ced20606d08a0aacfd25b49/pillow-12.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510", size = 2568267, upload-time = "2026-07-01T11:54:24.051Z" }, - { url = "https://files.pythonhosted.org/packages/9d/ac/31fb64e1e7efb5a4b50cd3d92049ba89ac6e4d8d3bb6a74e15048ca3353e/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:21900ce7ba264168cd50defae43cd75d25c833ad4ad6e73ffc5596d12e25ac89", size = 4161684, upload-time = "2026-07-01T11:54:25.934Z" }, - { url = "https://files.pythonhosted.org/packages/87/b4/9805e23d2b4d77842b468513841fda254ee42f0289d25088340e4ff46e2d/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:4e8c2a84d977f50b9daed6eeaf3baef67d00d5d74d932288f02cb94518ee3ace", size = 4255487, upload-time = "2026-07-01T11:54:27.935Z" }, - { url = "https://files.pythonhosted.org/packages/df/39/ecf519435a200c693fe053a6ee4d835b41cf963a4dfc2551c4e637cb2a71/pillow-12.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:ae26d61dfa7a47befdc7572b521024e8745f3d809bd95ca9505a7bba9ef849ec", size = 3696433, upload-time = "2026-07-01T11:54:29.813Z" }, - { url = "https://files.pythonhosted.org/packages/42/92/2fc3ffad878ae8dd5469ec1bc8eb83b71f48e13efdf68f02709003982a32/pillow-12.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7a743ff716f746fc19a9557f60dab1600d4613255f8a7aeb3cdde4db7eb15a66", size = 5345889, upload-time = "2026-07-01T11:54:31.97Z" }, - { url = "https://files.pythonhosted.org/packages/10/76/8803c13605b763d33d156c4678fc77f8443389c0c51c8aef707bb02015f4/pillow-12.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d69141514cc30b774ceea5e3ed3a6635c8d8a96edf664689b890f4089111fb35", size = 4780109, upload-time = "2026-07-01T11:54:34.026Z" }, - { url = "https://files.pythonhosted.org/packages/1f/01/e18aff37cb0b4aac47ac90f016d347a49aca667ef97f190b06ac2aabc928/pillow-12.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f7401aebd7f581d7f83a439d87d474999317ee099218e5ad25d125290990ba65", size = 6263736, upload-time = "2026-07-01T11:54:36.131Z" }, - { url = "https://files.pythonhosted.org/packages/f7/62/de5bdd77d935331f4f802edc11e4d82950f642caad6cb2f949837b8560e2/pillow-12.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0847a763afefb695bc912d7c131e7e0632d4edc1d8698f58ddabec8e46b8b6d3", size = 6937129, upload-time = "2026-07-01T11:54:38.216Z" }, - { url = "https://files.pythonhosted.org/packages/70/4d/105627a13300c5e0df1d174230b32fd1273062c96f7745fd552b945d1e1d/pillow-12.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:571b9fcb07b97ef3a492028fb3d2dc0993ca23a06138b0315286566d29ef718a", size = 6339562, upload-time = "2026-07-01T11:54:40.354Z" }, - { url = "https://files.pythonhosted.org/packages/6b/1d/f13de01a553988ab895ba1c722e06cf3144d4f57656fd5b81b6d881f1179/pillow-12.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:756c768d0c9c2955feb7a56c37ea24aea2e369f8d36a88da270b6a9f19e62b5e", size = 7049439, upload-time = "2026-07-01T11:54:42.489Z" }, - { url = "https://files.pythonhosted.org/packages/c9/f9/066794cca041b969964f779ee5fa66a9498bbf34248ac39c5d7954e4198f/pillow-12.3.0-cp313-cp313-win32.whl", hash = "sha256:a876864214e136f0eb367788dbd7df045f4806801518e2cfe9e13229cfe06d8f", size = 6473287, upload-time = "2026-07-01T11:54:44.9Z" }, - { url = "https://files.pythonhosted.org/packages/a6/9b/7a58e61d62be561da3a356fe2384d4059a6345fc130e23ef1c36a5b81d24/pillow-12.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:1cca606cd25738df4ed873d5ad46bbdb3d83b5cbca291f6b4ff13a4df6b0bbe8", size = 7239691, upload-time = "2026-07-01T11:54:47.141Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b0/c4ed4f0ef8f8fa5ee8351537db6650bb8189f7e118842978dd6589065692/pillow-12.3.0-cp313-cp313-win_arm64.whl", hash = "sha256:b629de27fda84b42cde7edef0d85f13b958b47f6e9bbcbba9b673c562a89bd8b", size = 2568185, upload-time = "2026-07-01T11:54:49.137Z" }, - { url = "https://files.pythonhosted.org/packages/dc/01/001f65b68192f0228cc1dbbc8d2530ab5d58b61037ba0587f946fea607cd/pillow-12.3.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:9cf95fe4d0f84c82d282745d9bb08ad9f926efa00be4697e767b814ce40d4330", size = 4161736, upload-time = "2026-07-01T11:54:51.156Z" }, - { url = "https://files.pythonhosted.org/packages/1a/d2/0219746d0fd16fc8a84498e79452375be3797d3ce4044596ce565164b84f/pillow-12.3.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:8728f216dcdb6e6d555cf971cb34076139ad74b31fc2c14da4fafc741c5f6217", size = 4255435, upload-time = "2026-07-01T11:54:53.414Z" }, - { url = "https://files.pythonhosted.org/packages/c8/02/8d0bc62ef0302318c46ff2a512822d2610e81c7aa46c9b3abe6cbaca5ad0/pillow-12.3.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:a45650e8ce7fafffd731db8550230db6b0d306d181a90b67d3e6bca2f1990930", size = 3696262, upload-time = "2026-07-01T11:54:55.739Z" }, - { url = "https://files.pythonhosted.org/packages/85/e2/73c77d218410b14f5f2d565e8a998d5317b7b9c75368d29985139f7a46f0/pillow-12.3.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ba54cfebe86920a559a7c4d6b9050791c20513650a1952ebe3368c7dc70306f8", size = 5350344, upload-time = "2026-07-01T11:54:57.657Z" }, - { url = "https://files.pythonhosted.org/packages/c7/da/32c752228ae345f489e3a42499d817b6c3996da7e8a3bc7a04fc806b243b/pillow-12.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e158cb00350dc278f3b91551101aa7d12415a66ebf2c91d8d5ac14e56ddd3ad0", size = 4780131, upload-time = "2026-07-01T11:54:59.713Z" }, - { url = "https://files.pythonhosted.org/packages/b1/9d/8b2c807dbef61a5197c047afe99823787eb66f63daf9fb2432f91d6f0462/pillow-12.3.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9aeb04d6aef139de265b29683e119b638208f88cf73cdd1658aa07221165321", size = 6263757, upload-time = "2026-07-01T11:55:01.778Z" }, - { url = "https://files.pythonhosted.org/packages/5c/44/c85361f65dbe00eea8576ee467c768d25129989efb76e94f205e9ca9bb46/pillow-12.3.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:251bf95b67017e27b13d82f5b326234ca62d70f9cf4c2b9032de2358a3b12c7b", size = 6936962, upload-time = "2026-07-01T11:55:03.93Z" }, - { url = "https://files.pythonhosted.org/packages/18/7e/e483414b35800b86b6f08dbbc7803fb5cd52c4d6f897f47d53ea2c7e6f65/pillow-12.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fe3cca2e4e8a592be0f269a1ca4835c25199d9f3ce815c8491048f785b0a0198", size = 6339171, upload-time = "2026-07-01T11:55:05.989Z" }, - { url = "https://files.pythonhosted.org/packages/f0/f4/68c491844841ede6bed70189546b3ee9731cf9f2cbad396faff5e1ccba45/pillow-12.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:23aceaa007d6172b02c277f0cd359c79492bbb14f7072b4ede9fbcaf20648130", size = 7048116, upload-time = "2026-07-01T11:55:08.131Z" }, - { url = "https://files.pythonhosted.org/packages/a3/34/77f3f793fed8efc7d243f21b33c5a3f0d1c97ee70346d3db855587e155ff/pillow-12.3.0-cp314-cp314-win32.whl", hash = "sha256:af8d94b0db561cf68b88a267c5c44b49e134f525d0dc2cb7ed413a66bc23559a", size = 6467209, upload-time = "2026-07-01T11:55:10.408Z" }, - { url = "https://files.pythonhosted.org/packages/f1/e0/492879f69d94f91f60fc8cd05ba03650e9520afebb2fb7aa12777d7c7f38/pillow-12.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:fdafc9cce40277e0f7a0feabce0ee50dd2fa1800f3b38015e51296b5e814048d", size = 7237707, upload-time = "2026-07-01T11:55:12.745Z" }, - { url = "https://files.pythonhosted.org/packages/c9/ac/6b11f2875f1c2ac040d84e1bbf9cf22a88038f901ca1037898b280b38365/pillow-12.3.0-cp314-cp314-win_arm64.whl", hash = "sha256:e91206ee562682b51b98ef4b26a6ef48fd84e15fd4c4bc5ec768eb641d206838", size = 2565995, upload-time = "2026-07-01T11:55:14.736Z" }, - { url = "https://files.pythonhosted.org/packages/52/69/c2208e56af9bfc1913afb24020297a691eb1d4ef688474c8a04913f65e04/pillow-12.3.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:164b31cd1a0490ab6efae01aa5df49da7061be0af1b30e035b6e9a1bfe34ee6e", size = 5352503, upload-time = "2026-07-01T11:55:17.076Z" }, - { url = "https://files.pythonhosted.org/packages/07/70/e5686d753e898a45d778ff1718dba8516ead6ab6b95d85fc8c4b70650cf2/pillow-12.3.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5afb51d599ea772b8365ae807ae557f18bccfe46ab261fd1c2a9ed700fc6eb17", size = 4782956, upload-time = "2026-07-01T11:55:19.448Z" }, - { url = "https://files.pythonhosted.org/packages/d5/37/25c6692f06927ee973ff18c8d9ee98ad0b4d84ee67a09610c2dd1447958e/pillow-12.3.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3edce1d53195db527e0191f84b71d02022de0540bf43a16ed734ed7537b07385", size = 6322855, upload-time = "2026-07-01T11:55:21.613Z" }, - { url = "https://files.pythonhosted.org/packages/cc/91/420637fcb8f1bc11029e403b4538e6694744428d8246118e45719f944556/pillow-12.3.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf16ba1b4d0b6b7c8e534936632270cf70eb00dbe09005bc345b2677b726855c", size = 6989642, upload-time = "2026-07-01T11:55:24.006Z" }, - { url = "https://files.pythonhosted.org/packages/10/08/b94d7811281ccf0d143a1cf768d1c49e1e54af63e7b708ab2ee3eb87face/pillow-12.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:24870b09b224f7ae3c39ed07d10e819d06f8720bc551847b1d623832b5b0e28d", size = 6391281, upload-time = "2026-07-01T11:55:26.252Z" }, - { url = "https://files.pythonhosted.org/packages/d2/87/24233f785f55474dc02ce3e739c5528a77e3a862e9333d1dd7a25cc31f70/pillow-12.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:30f2aa603c41533cc25c05acd0da21636e84a315768feb631c937177db558931", size = 7096716, upload-time = "2026-07-01T11:55:28.318Z" }, - { url = "https://files.pythonhosted.org/packages/23/26/fcb2f6e37175b04f53570b59937867e2b80ee1685e744023153028fc14f9/pillow-12.3.0-cp314-cp314t-win32.whl", hash = "sha256:4b0a7fe987b14c31ebda6083f74f22b561fd3739bc0ac51e019622e3d72668c7", size = 6474125, upload-time = "2026-07-01T11:55:30.956Z" }, - { url = "https://files.pythonhosted.org/packages/90/de/3634abee5f1c9e13c56787b7d5517b0ba8d6de51700b95578cf338349c9f/pillow-12.3.0-cp314-cp314t-win_amd64.whl", hash = "sha256:962864dc93511324d51ddbb5b9f8731bf71675b93ca612a07441896f4688fb8c", size = 7242939, upload-time = "2026-07-01T11:55:34.044Z" }, - { url = "https://files.pythonhosted.org/packages/ce/2a/fd13f8eb24de5714a6eb444a3d67e2842c6c576e159a43793adf23051351/pillow-12.3.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0740a512dc522224c77d9aa5a8d70d8b7d73fb91f2c21125d8d025d3b8990e45", size = 2567506, upload-time = "2026-07-01T11:55:35.988Z" }, - { url = "https://files.pythonhosted.org/packages/75/18/2e8b40223153ccbc60df07f9e8928dc0c76202aa4e55ae9f53962b6510d6/pillow-12.3.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:b3c777e849237620b022f7f297dd67705f9f5cf1685f09f02e46f93e92725468", size = 5302510, upload-time = "2026-07-01T11:56:25.736Z" }, - { url = "https://files.pythonhosted.org/packages/46/3e/51fabf59d5ab801ceab709453d3ab6b180083496579549de4c45ced6528a/pillow-12.3.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:b343699e8308bdc51978310e1c959c584e7869cc8c40780058c87da7781a1e94", size = 4736058, upload-time = "2026-07-01T11:56:28.041Z" }, - { url = "https://files.pythonhosted.org/packages/bf/20/22fe9384b7949e25fb1293bcfc84fb82590ff4ea6b37c95b24d26d793d86/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fbd139c8447d25dd750ab79ee274cc5e1fe80fc56340ab10b18a195e1b6eca3e", size = 5237776, upload-time = "2026-07-01T11:56:30.263Z" }, - { url = "https://files.pythonhosted.org/packages/08/14/f6ba68107680ffa74b39985f3f30884e41318fbc4250caa423c79b4788bb/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7e480451b9fa137494bccd3a7d69adbe8ac65a87d97be61e11f1b1050a5bac3", size = 5860358, upload-time = "2026-07-01T11:56:32.68Z" }, - { url = "https://files.pythonhosted.org/packages/36/54/0169bc772ec491108b62f644f8ecf1fe5d8ae5ebafde2ee2142210166903/pillow-12.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a", size = 7231786, upload-time = "2026-07-01T11:56:35.046Z" }, +version = "12.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/21/c2bcdd5906101a30244eaffc1b6e6ce71a31bd0742a01eb89e660ebfac2d/pillow-12.2.0.tar.gz", hash = "sha256:a830b1a40919539d07806aa58e1b114df53ddd43213d9c8b75847eee6c0182b5", size = 46987819, upload-time = "2026-04-01T14:46:17.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/aa/d0b28e1c811cd4d5f5c2bfe2e022292bd255ae5744a3b9ac7d6c8f72dd75/pillow-12.2.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:a4e8f36e677d3336f35089648c8955c51c6d386a13cf6ee9c189c5f5bd713a9f", size = 5354355, upload-time = "2026-04-01T14:42:15.402Z" }, + { url = "https://files.pythonhosted.org/packages/27/8e/1d5b39b8ae2bd7650d0c7b6abb9602d16043ead9ebbfef4bc4047454da2a/pillow-12.2.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e589959f10d9824d39b350472b92f0ce3b443c0a3442ebf41c40cb8361c5b97", size = 4695871, upload-time = "2026-04-01T14:42:18.234Z" }, + { url = "https://files.pythonhosted.org/packages/f0/c5/dcb7a6ca6b7d3be41a76958e90018d56c8462166b3ef223150360850c8da/pillow-12.2.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a52edc8bfff4429aaabdf4d9ee0daadbbf8562364f940937b941f87a4290f5ff", size = 6269734, upload-time = "2026-04-01T14:42:20.608Z" }, + { url = "https://files.pythonhosted.org/packages/ea/f1/aa1bb13b2f4eba914e9637893c73f2af8e48d7d4023b9d3750d4c5eb2d0c/pillow-12.2.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:975385f4776fafde056abb318f612ef6285b10a1f12b8570f3647ad0d74b48ec", size = 8076080, upload-time = "2026-04-01T14:42:23.095Z" }, + { url = "https://files.pythonhosted.org/packages/a1/2a/8c79d6a53169937784604a8ae8d77e45888c41537f7f6f65ed1f407fe66d/pillow-12.2.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bd9c0c7a0c681a347b3194c500cb1e6ca9cab053ea4d82a5cf45b6b754560136", size = 6382236, upload-time = "2026-04-01T14:42:25.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/42/bbcb6051030e1e421d103ce7a8ecadf837aa2f39b8f82ef1a8d37c3d4ebc/pillow-12.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:88d387ff40b3ff7c274947ed3125dedf5262ec6919d83946753b5f3d7c67ea4c", size = 7070220, upload-time = "2026-04-01T14:42:28.68Z" }, + { url = "https://files.pythonhosted.org/packages/3f/e1/c2a7d6dd8cfa6b231227da096fd2d58754bab3603b9d73bf609d3c18b64f/pillow-12.2.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:51c4167c34b0d8ba05b547a3bb23578d0ba17b80a5593f93bd8ecb123dd336a3", size = 6493124, upload-time = "2026-04-01T14:42:31.579Z" }, + { url = "https://files.pythonhosted.org/packages/5f/41/7c8617da5d32e1d2f026e509484fdb6f3ad7efaef1749a0c1928adbb099e/pillow-12.2.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:34c0d99ecccea270c04882cb3b86e7b57296079c9a4aff88cb3b33563d95afaa", size = 7194324, upload-time = "2026-04-01T14:42:34.615Z" }, + { url = "https://files.pythonhosted.org/packages/2d/de/a777627e19fd6d62f84070ee1521adde5eeda4855b5cf60fe0b149118bca/pillow-12.2.0-cp310-cp310-win32.whl", hash = "sha256:b85f66ae9eb53e860a873b858b789217ba505e5e405a24b85c0464822fe88032", size = 6376363, upload-time = "2026-04-01T14:42:37.19Z" }, + { url = "https://files.pythonhosted.org/packages/e7/34/fc4cb5204896465842767b96d250c08410f01f2f28afc43b257de842eed5/pillow-12.2.0-cp310-cp310-win_amd64.whl", hash = "sha256:673aa32138f3e7531ccdbca7b3901dba9b70940a19ccecc6a37c77d5fdeb05b5", size = 7083523, upload-time = "2026-04-01T14:42:39.62Z" }, + { url = "https://files.pythonhosted.org/packages/2d/a0/32852d36bc7709f14dc3f64f929a275e958ad8c19a6deba9610d458e28b3/pillow-12.2.0-cp310-cp310-win_arm64.whl", hash = "sha256:3e080565d8d7c671db5802eedfb438e5565ffa40115216eabb8cd52d0ecce024", size = 2463318, upload-time = "2026-04-01T14:42:42.063Z" }, + { url = "https://files.pythonhosted.org/packages/68/e1/748f5663efe6edcfc4e74b2b93edfb9b8b99b67f21a854c3ae416500a2d9/pillow-12.2.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:8be29e59487a79f173507c30ddf57e733a357f67881430449bb32614075a40ab", size = 5354347, upload-time = "2026-04-01T14:42:44.255Z" }, + { url = "https://files.pythonhosted.org/packages/47/a1/d5ff69e747374c33a3b53b9f98cca7889fce1fd03d79cdc4e1bccc6c5a87/pillow-12.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:71cde9a1e1551df7d34a25462fc60325e8a11a82cc2e2f54578e5e9a1e153d65", size = 4695873, upload-time = "2026-04-01T14:42:46.452Z" }, + { url = "https://files.pythonhosted.org/packages/df/21/e3fbdf54408a973c7f7f89a23b2cb97a7ef30c61ab4142af31eee6aebc88/pillow-12.2.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f490f9368b6fc026f021db16d7ec2fbf7d89e2edb42e8ec09d2c60505f5729c7", size = 6280168, upload-time = "2026-04-01T14:42:49.228Z" }, + { url = "https://files.pythonhosted.org/packages/d3/f1/00b7278c7dd52b17ad4329153748f87b6756ec195ff786c2bdf12518337d/pillow-12.2.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8bd7903a5f2a4545f6fd5935c90058b89d30045568985a71c79f5fd6edf9b91e", size = 8088188, upload-time = "2026-04-01T14:42:51.735Z" }, + { url = "https://files.pythonhosted.org/packages/ad/cf/220a5994ef1b10e70e85748b75649d77d506499352be135a4989c957b701/pillow-12.2.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3997232e10d2920a68d25191392e3a4487d8183039e1c74c2297f00ed1c50705", size = 6394401, upload-time = "2026-04-01T14:42:54.343Z" }, + { url = "https://files.pythonhosted.org/packages/e9/bd/e51a61b1054f09437acfbc2ff9106c30d1eb76bc1453d428399946781253/pillow-12.2.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e74473c875d78b8e9d5da2a70f7099549f9eb37ded4e2f6a463e60125bccd176", size = 7079655, upload-time = "2026-04-01T14:42:56.954Z" }, + { url = "https://files.pythonhosted.org/packages/6b/3d/45132c57d5fb4b5744567c3817026480ac7fc3ce5d4c47902bc0e7f6f853/pillow-12.2.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:56a3f9c60a13133a98ecff6197af34d7824de9b7b38c3654861a725c970c197b", size = 6503105, upload-time = "2026-04-01T14:42:59.847Z" }, + { url = "https://files.pythonhosted.org/packages/7d/2e/9df2fc1e82097b1df3dce58dc43286aa01068e918c07574711fcc53e6fb4/pillow-12.2.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:90e6f81de50ad6b534cab6e5aef77ff6e37722b2f5d908686f4a5c9eba17a909", size = 7203402, upload-time = "2026-04-01T14:43:02.664Z" }, + { url = "https://files.pythonhosted.org/packages/bd/2e/2941e42858ebb67e50ae741473de81c2984e6eff7b397017623c676e2e8d/pillow-12.2.0-cp311-cp311-win32.whl", hash = "sha256:8c984051042858021a54926eb597d6ee3012393ce9c181814115df4c60b9a808", size = 6378149, upload-time = "2026-04-01T14:43:05.274Z" }, + { url = "https://files.pythonhosted.org/packages/69/42/836b6f3cd7f3e5fa10a1f1a5420447c17966044c8fbf589cc0452d5502db/pillow-12.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:6e6b2a0c538fc200b38ff9eb6628228b77908c319a005815f2dde585a0664b60", size = 7082626, upload-time = "2026-04-01T14:43:08.557Z" }, + { url = "https://files.pythonhosted.org/packages/c2/88/549194b5d6f1f494b485e493edc6693c0a16f4ada488e5bd974ed1f42fad/pillow-12.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:9a8a34cc89c67a65ea7437ce257cea81a9dad65b29805f3ecee8c8fe8ff25ffe", size = 2463531, upload-time = "2026-04-01T14:43:10.743Z" }, + { url = "https://files.pythonhosted.org/packages/58/be/7482c8a5ebebbc6470b3eb791812fff7d5e0216c2be3827b30b8bb6603ed/pillow-12.2.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2d192a155bbcec180f8564f693e6fd9bccff5a7af9b32e2e4bf8c9c69dbad6b5", size = 5308279, upload-time = "2026-04-01T14:43:13.246Z" }, + { url = "https://files.pythonhosted.org/packages/d8/95/0a351b9289c2b5cbde0bacd4a83ebc44023e835490a727b2a3bd60ddc0f4/pillow-12.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3f40b3c5a968281fd507d519e444c35f0ff171237f4fdde090dd60699458421", size = 4695490, upload-time = "2026-04-01T14:43:15.584Z" }, + { url = "https://files.pythonhosted.org/packages/de/af/4e8e6869cbed569d43c416fad3dc4ecb944cb5d9492defaed89ddd6fe871/pillow-12.2.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:03e7e372d5240cc23e9f07deca4d775c0817bffc641b01e9c3af208dbd300987", size = 6284462, upload-time = "2026-04-01T14:43:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/e9/9e/c05e19657fd57841e476be1ab46c4d501bffbadbafdc31a6d665f8b737b6/pillow-12.2.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b86024e52a1b269467a802258c25521e6d742349d760728092e1bc2d135b4d76", size = 8094744, upload-time = "2026-04-01T14:43:20.716Z" }, + { url = "https://files.pythonhosted.org/packages/2b/54/1789c455ed10176066b6e7e6da1b01e50e36f94ba584dc68d9eebfe9156d/pillow-12.2.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7371b48c4fa448d20d2714c9a1f775a81155050d383333e0a6c15b1123dda005", size = 6398371, upload-time = "2026-04-01T14:43:23.443Z" }, + { url = "https://files.pythonhosted.org/packages/43/e3/fdc657359e919462369869f1c9f0e973f353f9a9ee295a39b1fea8ee1a77/pillow-12.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:62f5409336adb0663b7caa0da5c7d9e7bdbaae9ce761d34669420c2a801b2780", size = 7087215, upload-time = "2026-04-01T14:43:26.758Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f8/2f6825e441d5b1959d2ca5adec984210f1ec086435b0ed5f52c19b3b8a6e/pillow-12.2.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:01afa7cf67f74f09523699b4e88c73fb55c13346d212a59a2db1f86b0a63e8c5", size = 6509783, upload-time = "2026-04-01T14:43:29.56Z" }, + { url = "https://files.pythonhosted.org/packages/67/f9/029a27095ad20f854f9dba026b3ea6428548316e057e6fc3545409e86651/pillow-12.2.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc3d34d4a8fbec3e88a79b92e5465e0f9b842b628675850d860b8bd300b159f5", size = 7212112, upload-time = "2026-04-01T14:43:32.091Z" }, + { url = "https://files.pythonhosted.org/packages/be/42/025cfe05d1be22dbfdb4f264fe9de1ccda83f66e4fc3aac94748e784af04/pillow-12.2.0-cp312-cp312-win32.whl", hash = "sha256:58f62cc0f00fd29e64b29f4fd923ffdb3859c9f9e6105bfc37ba1d08994e8940", size = 6378489, upload-time = "2026-04-01T14:43:34.601Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7b/25a221d2c761c6a8ae21bfa3874988ff2583e19cf8a27bf2fee358df7942/pillow-12.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:7f84204dee22a783350679a0333981df803dac21a0190d706a50475e361c93f5", size = 7084129, upload-time = "2026-04-01T14:43:37.213Z" }, + { url = "https://files.pythonhosted.org/packages/10/e1/542a474affab20fd4a0f1836cb234e8493519da6b76899e30bcc5d990b8b/pillow-12.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:af73337013e0b3b46f175e79492d96845b16126ddf79c438d7ea7ff27783a414", size = 2463612, upload-time = "2026-04-01T14:43:39.421Z" }, + { url = "https://files.pythonhosted.org/packages/4a/01/53d10cf0dbad820a8db274d259a37ba50b88b24768ddccec07355382d5ad/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:8297651f5b5679c19968abefd6bb84d95fe30ef712eb1b2d9b2d31ca61267f4c", size = 4100837, upload-time = "2026-04-01T14:43:41.506Z" }, + { url = "https://files.pythonhosted.org/packages/0f/98/f3a6657ecb698c937f6c76ee564882945f29b79bad496abcba0e84659ec5/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:50d8520da2a6ce0af445fa6d648c4273c3eeefbc32d7ce049f22e8b5c3daecc2", size = 4176528, upload-time = "2026-04-01T14:43:43.773Z" }, + { url = "https://files.pythonhosted.org/packages/69/bc/8986948f05e3ea490b8442ea1c1d4d990b24a7e43d8a51b2c7d8b1dced36/pillow-12.2.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:766cef22385fa1091258ad7e6216792b156dc16d8d3fa607e7545b2b72061f1c", size = 3640401, upload-time = "2026-04-01T14:43:45.87Z" }, + { url = "https://files.pythonhosted.org/packages/34/46/6c717baadcd62bc8ed51d238d521ab651eaa74838291bda1f86fe1f864c9/pillow-12.2.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5d2fd0fa6b5d9d1de415060363433f28da8b1526c1c129020435e186794b3795", size = 5308094, upload-time = "2026-04-01T14:43:48.438Z" }, + { url = "https://files.pythonhosted.org/packages/71/43/905a14a8b17fdb1ccb58d282454490662d2cb89a6bfec26af6d3520da5ec/pillow-12.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56b25336f502b6ed02e889f4ece894a72612fe885889a6e8c4c80239ff6e5f5f", size = 4695402, upload-time = "2026-04-01T14:43:51.292Z" }, + { url = "https://files.pythonhosted.org/packages/73/dd/42107efcb777b16fa0393317eac58f5b5cf30e8392e266e76e51cff28c3d/pillow-12.2.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f1c943e96e85df3d3478f7b691f229887e143f81fedab9b20205349ab04d73ed", size = 6280005, upload-time = "2026-04-01T14:43:54.242Z" }, + { url = "https://files.pythonhosted.org/packages/a8/68/b93e09e5e8549019e61acf49f65b1a8530765a7f812c77a7461bca7e4494/pillow-12.2.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:03f6fab9219220f041c74aeaa2939ff0062bd5c364ba9ce037197f4c6d498cd9", size = 8090669, upload-time = "2026-04-01T14:43:57.335Z" }, + { url = "https://files.pythonhosted.org/packages/4b/6e/3ccb54ce8ec4ddd1accd2d89004308b7b0b21c4ac3d20fa70af4760a4330/pillow-12.2.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5cdfebd752ec52bf5bb4e35d9c64b40826bc5b40a13df7c3cda20a2c03a0f5ed", size = 6395194, upload-time = "2026-04-01T14:43:59.864Z" }, + { url = "https://files.pythonhosted.org/packages/67/ee/21d4e8536afd1a328f01b359b4d3997b291ffd35a237c877b331c1c3b71c/pillow-12.2.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eedf4b74eda2b5a4b2b2fb4c006d6295df3bf29e459e198c90ea48e130dc75c3", size = 7082423, upload-time = "2026-04-01T14:44:02.74Z" }, + { url = "https://files.pythonhosted.org/packages/78/5f/e9f86ab0146464e8c133fe85df987ed9e77e08b29d8d35f9f9f4d6f917ba/pillow-12.2.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:00a2865911330191c0b818c59103b58a5e697cae67042366970a6b6f1b20b7f9", size = 6505667, upload-time = "2026-04-01T14:44:05.381Z" }, + { url = "https://files.pythonhosted.org/packages/ed/1e/409007f56a2fdce61584fd3acbc2bbc259857d555196cedcadc68c015c82/pillow-12.2.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1e1757442ed87f4912397c6d35a0db6a7b52592156014706f17658ff58bbf795", size = 7208580, upload-time = "2026-04-01T14:44:08.39Z" }, + { url = "https://files.pythonhosted.org/packages/23/c4/7349421080b12fb35414607b8871e9534546c128a11965fd4a7002ccfbee/pillow-12.2.0-cp313-cp313-win32.whl", hash = "sha256:144748b3af2d1b358d41286056d0003f47cb339b8c43a9ea42f5fea4d8c66b6e", size = 6375896, upload-time = "2026-04-01T14:44:11.197Z" }, + { url = "https://files.pythonhosted.org/packages/3f/82/8a3739a5e470b3c6cbb1d21d315800d8e16bff503d1f16b03a4ec3212786/pillow-12.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:390ede346628ccc626e5730107cde16c42d3836b89662a115a921f28440e6a3b", size = 7081266, upload-time = "2026-04-01T14:44:13.947Z" }, + { url = "https://files.pythonhosted.org/packages/c3/25/f968f618a062574294592f668218f8af564830ccebdd1fa6200f598e65c5/pillow-12.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:8023abc91fba39036dbce14a7d6535632f99c0b857807cbbbf21ecc9f4717f06", size = 2463508, upload-time = "2026-04-01T14:44:16.312Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a4/b342930964e3cb4dce5038ae34b0eab4653334995336cd486c5a8c25a00c/pillow-12.2.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:042db20a421b9bafecc4b84a8b6e444686bd9d836c7fd24542db3e7df7baad9b", size = 5309927, upload-time = "2026-04-01T14:44:18.89Z" }, + { url = "https://files.pythonhosted.org/packages/9f/de/23198e0a65a9cf06123f5435a5d95cea62a635697f8f03d134d3f3a96151/pillow-12.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:dd025009355c926a84a612fecf58bb315a3f6814b17ead51a8e48d3823d9087f", size = 4698624, upload-time = "2026-04-01T14:44:21.115Z" }, + { url = "https://files.pythonhosted.org/packages/01/a6/1265e977f17d93ea37aa28aa81bad4fa597933879fac2520d24e021c8da3/pillow-12.2.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:88ddbc66737e277852913bd1e07c150cc7bb124539f94c4e2df5344494e0a612", size = 6321252, upload-time = "2026-04-01T14:44:23.663Z" }, + { url = "https://files.pythonhosted.org/packages/3c/83/5982eb4a285967baa70340320be9f88e57665a387e3a53a7f0db8231a0cd/pillow-12.2.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d362d1878f00c142b7e1a16e6e5e780f02be8195123f164edf7eddd911eefe7c", size = 8126550, upload-time = "2026-04-01T14:44:26.772Z" }, + { url = "https://files.pythonhosted.org/packages/4e/48/6ffc514adce69f6050d0753b1a18fd920fce8cac87620d5a31231b04bfc5/pillow-12.2.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c727a6d53cb0018aadd8018c2b938376af27914a68a492f59dfcaca650d5eea", size = 6433114, upload-time = "2026-04-01T14:44:29.615Z" }, + { url = "https://files.pythonhosted.org/packages/36/a3/f9a77144231fb8d40ee27107b4463e205fa4677e2ca2548e14da5cf18dce/pillow-12.2.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd8c21c98c5cc60653bcb311bef2ce0401642b7ce9d09e03a7da87c878289d4", size = 7115667, upload-time = "2026-04-01T14:44:32.773Z" }, + { url = "https://files.pythonhosted.org/packages/c1/fc/ac4ee3041e7d5a565e1c4fd72a113f03b6394cc72ab7089d27608f8aaccb/pillow-12.2.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9f08483a632889536b8139663db60f6724bfcb443c96f1b18855860d7d5c0fd4", size = 6538966, upload-time = "2026-04-01T14:44:35.252Z" }, + { url = "https://files.pythonhosted.org/packages/c0/a8/27fb307055087f3668f6d0a8ccb636e7431d56ed0750e07a60547b1e083e/pillow-12.2.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dac8d77255a37e81a2efcbd1fc05f1c15ee82200e6c240d7e127e25e365c39ea", size = 7238241, upload-time = "2026-04-01T14:44:37.875Z" }, + { url = "https://files.pythonhosted.org/packages/ad/4b/926ab182c07fccae9fcb120043464e1ff1564775ec8864f21a0ebce6ac25/pillow-12.2.0-cp313-cp313t-win32.whl", hash = "sha256:ee3120ae9dff32f121610bb08e4313be87e03efeadfc6c0d18f89127e24d0c24", size = 6379592, upload-time = "2026-04-01T14:44:40.336Z" }, + { url = "https://files.pythonhosted.org/packages/c2/c4/f9e476451a098181b30050cc4c9a3556b64c02cf6497ea421ac047e89e4b/pillow-12.2.0-cp313-cp313t-win_amd64.whl", hash = "sha256:325ca0528c6788d2a6c3d40e3568639398137346c3d6e66bb61db96b96511c98", size = 7085542, upload-time = "2026-04-01T14:44:43.251Z" }, + { url = "https://files.pythonhosted.org/packages/00/a4/285f12aeacbe2d6dc36c407dfbbe9e96d4a80b0fb710a337f6d2ad978c75/pillow-12.2.0-cp313-cp313t-win_arm64.whl", hash = "sha256:2e5a76d03a6c6dcef67edabda7a52494afa4035021a79c8558e14af25313d453", size = 2465765, upload-time = "2026-04-01T14:44:45.996Z" }, + { url = "https://files.pythonhosted.org/packages/bf/98/4595daa2365416a86cb0d495248a393dfc84e96d62ad080c8546256cb9c0/pillow-12.2.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:3adc9215e8be0448ed6e814966ecf3d9952f0ea40eb14e89a102b87f450660d8", size = 4100848, upload-time = "2026-04-01T14:44:48.48Z" }, + { url = "https://files.pythonhosted.org/packages/0b/79/40184d464cf89f6663e18dfcf7ca21aae2491fff1a16127681bf1fa9b8cf/pillow-12.2.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:6a9adfc6d24b10f89588096364cc726174118c62130c817c2837c60cf08a392b", size = 4176515, upload-time = "2026-04-01T14:44:51.353Z" }, + { url = "https://files.pythonhosted.org/packages/b0/63/703f86fd4c422a9cf722833670f4f71418fb116b2853ff7da722ea43f184/pillow-12.2.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:6a6e67ea2e6feda684ed370f9a1c52e7a243631c025ba42149a2cc5934dec295", size = 3640159, upload-time = "2026-04-01T14:44:53.588Z" }, + { url = "https://files.pythonhosted.org/packages/71/e0/fb22f797187d0be2270f83500aab851536101b254bfa1eae10795709d283/pillow-12.2.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2bb4a8d594eacdfc59d9e5ad972aa8afdd48d584ffd5f13a937a664c3e7db0ed", size = 5312185, upload-time = "2026-04-01T14:44:56.039Z" }, + { url = "https://files.pythonhosted.org/packages/ba/8c/1a9e46228571de18f8e28f16fabdfc20212a5d019f3e3303452b3f0a580d/pillow-12.2.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:80b2da48193b2f33ed0c32c38140f9d3186583ce7d516526d462645fd98660ae", size = 4695386, upload-time = "2026-04-01T14:44:58.663Z" }, + { url = "https://files.pythonhosted.org/packages/70/62/98f6b7f0c88b9addd0e87c217ded307b36be024d4ff8869a812b241d1345/pillow-12.2.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:22db17c68434de69d8ecfc2fe821569195c0c373b25cccb9cbdacf2c6e53c601", size = 6280384, upload-time = "2026-04-01T14:45:01.5Z" }, + { url = "https://files.pythonhosted.org/packages/5e/03/688747d2e91cfbe0e64f316cd2e8005698f76ada3130d0194664174fa5de/pillow-12.2.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7b14cc0106cd9aecda615dd6903840a058b4700fcb817687d0ee4fc8b6e389be", size = 8091599, upload-time = "2026-04-01T14:45:04.5Z" }, + { url = "https://files.pythonhosted.org/packages/f6/35/577e22b936fcdd66537329b33af0b4ccfefaeabd8aec04b266528cddb33c/pillow-12.2.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8cbeb542b2ebc6fcdacabf8aca8c1a97c9b3ad3927d46b8723f9d4f033288a0f", size = 6396021, upload-time = "2026-04-01T14:45:07.117Z" }, + { url = "https://files.pythonhosted.org/packages/11/8d/d2532ad2a603ca2b93ad9f5135732124e57811d0168155852f37fbce2458/pillow-12.2.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4bfd07bc812fbd20395212969e41931001fd59eb55a60658b0e5710872e95286", size = 7083360, upload-time = "2026-04-01T14:45:09.763Z" }, + { url = "https://files.pythonhosted.org/packages/5e/26/d325f9f56c7e039034897e7380e9cc202b1e368bfd04d4cbe6a441f02885/pillow-12.2.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:9aba9a17b623ef750a4d11b742cbafffeb48a869821252b30ee21b5e91392c50", size = 6507628, upload-time = "2026-04-01T14:45:12.378Z" }, + { url = "https://files.pythonhosted.org/packages/5f/f7/769d5632ffb0988f1c5e7660b3e731e30f7f8ec4318e94d0a5d674eb65a4/pillow-12.2.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:deede7c263feb25dba4e82ea23058a235dcc2fe1f6021025dc71f2b618e26104", size = 7209321, upload-time = "2026-04-01T14:45:15.122Z" }, + { url = "https://files.pythonhosted.org/packages/6a/7a/c253e3c645cd47f1aceea6a8bacdba9991bf45bb7dfe927f7c893e89c93c/pillow-12.2.0-cp314-cp314-win32.whl", hash = "sha256:632ff19b2778e43162304d50da0181ce24ac5bb8180122cbe1bf4673428328c7", size = 6479723, upload-time = "2026-04-01T14:45:17.797Z" }, + { url = "https://files.pythonhosted.org/packages/cd/8b/601e6566b957ca50e28725cb6c355c59c2c8609751efbecd980db44e0349/pillow-12.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:4e6c62e9d237e9b65fac06857d511e90d8461a32adcc1b9065ea0c0fa3a28150", size = 7217400, upload-time = "2026-04-01T14:45:20.529Z" }, + { url = "https://files.pythonhosted.org/packages/d6/94/220e46c73065c3e2951bb91c11a1fb636c8c9ad427ac3ce7d7f3359b9b2f/pillow-12.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:b1c1fbd8a5a1af3412a0810d060a78b5136ec0836c8a4ef9aa11807f2a22f4e1", size = 2554835, upload-time = "2026-04-01T14:45:23.162Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ab/1b426a3974cb0e7da5c29ccff4807871d48110933a57207b5a676cccc155/pillow-12.2.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:57850958fe9c751670e49b2cecf6294acc99e562531f4bd317fa5ddee2068463", size = 5314225, upload-time = "2026-04-01T14:45:25.637Z" }, + { url = "https://files.pythonhosted.org/packages/19/1e/dce46f371be2438eecfee2a1960ee2a243bbe5e961890146d2dee1ff0f12/pillow-12.2.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:d5d38f1411c0ed9f97bcb49b7bd59b6b7c314e0e27420e34d99d844b9ce3b6f3", size = 4698541, upload-time = "2026-04-01T14:45:28.355Z" }, + { url = "https://files.pythonhosted.org/packages/55/c3/7fbecf70adb3a0c33b77a300dc52e424dc22ad8cdc06557a2e49523b703d/pillow-12.2.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c0a9f29ca8e79f09de89293f82fc9b0270bb4af1d58bc98f540cc4aedf03166", size = 6322251, upload-time = "2026-04-01T14:45:30.924Z" }, + { url = "https://files.pythonhosted.org/packages/1c/3c/7fbc17cfb7e4fe0ef1642e0abc17fc6c94c9f7a16be41498e12e2ba60408/pillow-12.2.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1610dd6c61621ae1cf811bef44d77e149ce3f7b95afe66a4512f8c59f25d9ebe", size = 8127807, upload-time = "2026-04-01T14:45:33.908Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c3/a8ae14d6defd2e448493ff512fae903b1e9bd40b72efb6ec55ce0048c8ce/pillow-12.2.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a34329707af4f73cf1782a36cd2289c0368880654a2c11f027bcee9052d35dd", size = 6433935, upload-time = "2026-04-01T14:45:36.623Z" }, + { url = "https://files.pythonhosted.org/packages/6e/32/2880fb3a074847ac159d8f902cb43278a61e85f681661e7419e6596803ed/pillow-12.2.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e9c4f5b3c546fa3458a29ab22646c1c6c787ea8f5ef51300e5a60300736905e", size = 7116720, upload-time = "2026-04-01T14:45:39.258Z" }, + { url = "https://files.pythonhosted.org/packages/46/87/495cc9c30e0129501643f24d320076f4cc54f718341df18cc70ec94c44e1/pillow-12.2.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:fb043ee2f06b41473269765c2feae53fc2e2fbf96e5e22ca94fb5ad677856f06", size = 6540498, upload-time = "2026-04-01T14:45:41.879Z" }, + { url = "https://files.pythonhosted.org/packages/18/53/773f5edca692009d883a72211b60fdaf8871cbef075eaa9d577f0a2f989e/pillow-12.2.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f278f034eb75b4e8a13a54a876cc4a5ab39173d2cdd93a638e1b467fc545ac43", size = 7239413, upload-time = "2026-04-01T14:45:44.705Z" }, + { url = "https://files.pythonhosted.org/packages/c9/e4/4b64a97d71b2a83158134abbb2f5bd3f8a2ea691361282f010998f339ec7/pillow-12.2.0-cp314-cp314t-win32.whl", hash = "sha256:6bb77b2dcb06b20f9f4b4a8454caa581cd4dd0643a08bacf821216a16d9c8354", size = 6482084, upload-time = "2026-04-01T14:45:47.568Z" }, + { url = "https://files.pythonhosted.org/packages/ba/13/306d275efd3a3453f72114b7431c877d10b1154014c1ebbedd067770d629/pillow-12.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6562ace0d3fb5f20ed7290f1f929cae41b25ae29528f2af1722966a0a02e2aa1", size = 7225152, upload-time = "2026-04-01T14:45:50.032Z" }, + { url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" }, + { url = "https://files.pythonhosted.org/packages/4e/b7/2437044fb910f499610356d1352e3423753c98e34f915252aafecc64889f/pillow-12.2.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0538bd5e05efec03ae613fd89c4ce0368ecd2ba239cc25b9f9be7ed426b0af1f", size = 5273969, upload-time = "2026-04-01T14:45:55.538Z" }, + { url = "https://files.pythonhosted.org/packages/f6/f4/8316e31de11b780f4ac08ef3654a75555e624a98db1056ecb2122d008d5a/pillow-12.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:394167b21da716608eac917c60aa9b969421b5dcbbe02ae7f013e7b85811c69d", size = 4659674, upload-time = "2026-04-01T14:45:58.093Z" }, + { url = "https://files.pythonhosted.org/packages/d4/37/664fca7201f8bb2aa1d20e2c3d5564a62e6ae5111741966c8319ca802361/pillow-12.2.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5d04bfa02cc2d23b497d1e90a0f927070043f6cbf303e738300532379a4b4e0f", size = 5288479, upload-time = "2026-04-01T14:46:01.141Z" }, + { url = "https://files.pythonhosted.org/packages/49/62/5b0ed78fce87346be7a5cfcfaaad91f6a1f98c26f86bdbafa2066c647ef6/pillow-12.2.0-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0c838a5125cee37e68edec915651521191cef1e6aa336b855f495766e77a366e", size = 7032230, upload-time = "2026-04-01T14:46:03.874Z" }, + { url = "https://files.pythonhosted.org/packages/c3/28/ec0fc38107fc32536908034e990c47914c57cd7c5a3ece4d8d8f7ffd7e27/pillow-12.2.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a6c9fa44005fa37a91ebfc95d081e8079757d2e904b27103f4f5fa6f0bf78c0", size = 5355404, upload-time = "2026-04-01T14:46:06.33Z" }, + { url = "https://files.pythonhosted.org/packages/5e/8b/51b0eddcfa2180d60e41f06bd6d0a62202b20b59c68f5a132e615b75aecf/pillow-12.2.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25373b66e0dd5905ed63fa3cae13c82fbddf3079f2c8bf15c6fb6a35586324c1", size = 6002215, upload-time = "2026-04-01T14:46:08.83Z" }, + { url = "https://files.pythonhosted.org/packages/bc/60/5382c03e1970de634027cee8e1b7d39776b778b81812aaf45b694dfe9e28/pillow-12.2.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:bfa9c230d2fe991bed5318a5f119bd6780cda2915cca595393649fc118ab895e", size = 7080946, upload-time = "2026-04-01T14:46:11.734Z" }, ] [[package]] name = "platformdirs" -version = "4.11.0" +version = "4.10.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/78/9b/560e4be8e26f6fd133a03630a8df0c663b9e8d61b4ade152b72005aec83b/platformdirs-4.11.0.tar.gz", hash = "sha256:0555d18370482847566ffabcaa53ad7c6c1c29f195989ae1ed634a05f76ea1e0", size = 31953, upload-time = "2026-07-21T13:09:36.565Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/47/e4501f49c178ae1d9f4a75073fda4204f52647993f075a9db4d14930e0c5/platformdirs-4.10.0.tar.gz", hash = "sha256:31e761a6a0ca04faf7353ea759bdba55652be214725111e5aac52dfa29d4bef7", size = 31224, upload-time = "2026-05-28T03:32:53.587Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7d/68/d8d58938dfb1370b266a1a729e6d77a985be23689a0496498ee17b2cbf90/platformdirs-4.11.0-py3-none-any.whl", hash = "sha256:360ccded2b7fce0af0ff80cc8f5942a1c5d99b0e856033acb030bfc634709e74", size = 23247, upload-time = "2026-07-21T13:09:35.422Z" }, + { url = "https://files.pythonhosted.org/packages/81/e6/cd9575ac904136b3cbf7aa7ee819ef86eedb7274e46f230e94ea4342e729/platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a", size = 22743, upload-time = "2026-05-28T03:32:52.175Z" }, ] [[package]] @@ -2858,7 +3255,7 @@ wheels = [ [[package]] name = "pre-commit" -version = "4.6.1" +version = "4.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cfgv" }, @@ -2867,9 +3264,9 @@ dependencies = [ { name = "pyyaml" }, { name = "virtualenv" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/25/3a/ddb78f32a0814e66b18a099377a106a2dcdce92d86a034d69d65df9b256e/pre_commit-4.6.1.tar.gz", hash = "sha256:03e809865c7d178b9979d06c761fcbfe6808fdaded8581a745bb110e52050421", size = 198646, upload-time = "2026-07-21T20:56:58.225Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/22/2de9408ac81acbb8a7d05d4cc064a152ccf33b3d480ebe0cd292153db239/pre_commit-4.6.0.tar.gz", hash = "sha256:718d2208cef53fdc38206e40524a6d4d9576d103eb16f0fec11c875e7716e9d9", size = 198525, upload-time = "2026-04-21T20:31:41.613Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/49/bc925106abcdac498074f2cbe6137e94e09f418dd2b7775df5b577dc0313/pre_commit-4.6.1-py2.py3-none-any.whl", hash = "sha256:0e3b2942510d1fb34eec167a3ec57331bf8442122f1153a9fb8b58f5c49b2717", size = 226186, upload-time = "2026-07-21T20:56:57.064Z" }, + { url = "https://files.pythonhosted.org/packages/80/6e/4b28b62ecb6aae56769c34a8ff1d661473ec1e9519e2d5f8b2c150086b26/pre_commit-4.6.0-py2.py3-none-any.whl", hash = "sha256:e2cf246f7299edcabcf15f9b0571fdce06058527f0a06535068a86d38089f29b", size = 226472, upload-time = "2026-04-21T20:31:40.092Z" }, ] [[package]] @@ -2996,7 +3393,9 @@ tqdm = [ [package.dev-dependencies] dev = [ + { name = "hatchling" }, { name = "hypothesis" }, + { name = "mypy" }, { name = "nbmake" }, { name = "nox" }, { name = "pre-commit" }, @@ -3007,6 +3406,9 @@ dev = [ { name = "pytest-snapshot" }, { name = "pytest-xdist" }, { name = "ruff" }, + { name = "scipy-stubs", version = "1.15.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy-stubs", version = "1.17.1.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "scipy-stubs", version = "1.18.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] docs = [ { name = "ipykernel" }, @@ -3065,8 +3467,10 @@ provides-extras = ["all", "bpx", "cite", "examples", "jax", "plot", "pydiffsol", [package.metadata.requires-dev] dev = [ + { name = "hatchling", specifier = ">=1.31.0" }, { name = "hypothesis" }, { name = "importlib-metadata", marker = "python_full_version < '3.10'" }, + { name = "mypy", specifier = ">=1.14" }, { name = "nbmake" }, { name = "nox" }, { name = "pre-commit" }, @@ -3077,6 +3481,7 @@ dev = [ { name = "pytest-snapshot" }, { name = "pytest-xdist" }, { name = "ruff" }, + { name = "scipy-stubs" }, ] docs = [ { name = "ipykernel" }, @@ -3370,7 +3775,7 @@ wheels = [ [[package]] name = "pytest" -version = "9.1.1" +version = "9.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, @@ -3381,9 +3786,9 @@ dependencies = [ { name = "pygments" }, { name = "tomli", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, + { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" }, ] [[package]] @@ -4008,17 +4413,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -4088,6 +4493,64 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/a5/df8f46ef7da168f1bc52cd86e09a9de5c6f19cc1da04454d51b7d4f43408/scipy-1.17.0-cp314-cp314t-win_arm64.whl", hash = "sha256:031121914e295d9791319a1875444d55079885bbae5bdc9c5e0f2ee5f09d34ff", size = 25246266, upload-time = "2026-01-10T21:30:45.923Z" }, ] +[[package]] +name = "scipy-stubs" +version = "1.15.3.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "optype", version = "0.9.3", source = { registry = "https://pypi.org/simple" } }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0b/5f/35c43bd7d412add4adcd68475702571b2489b50c40b6564f808b2355e452/scipy_stubs-1.15.3.0.tar.gz", hash = "sha256:e8f76c9887461cf9424c1e2ad78ea5dac71dd4cbb383dc85f91adfe8f74d1e17", size = 275699, upload-time = "2025-05-08T16:58:35.139Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/42/cd8dc81f8060de1f14960885ad5b2d2651f41de8b93d09f3f919d6567a5a/scipy_stubs-1.15.3.0-py3-none-any.whl", hash = "sha256:a251254cf4fd6e7fb87c55c1feee92d32ddbc1f542ecdf6a0159cdb81c2fb62d", size = 459062, upload-time = "2025-05-08T16:58:33.356Z" }, +] + +[[package]] +name = "scipy-stubs" +version = "1.17.1.5" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "optype", version = "0.17.1", source = { registry = "https://pypi.org/simple" }, extra = ["numpy"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/30/7a2e621918d1317ab972f797161131f2635648ad5d92baf0695dd009e4f9/scipy_stubs-1.17.1.5.tar.gz", hash = "sha256:284b1dd1dd46107a614971d170030d310cd88b2ac6b483f85285ee0ff87720bd", size = 399933, upload-time = "2026-05-25T21:34:33.6Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/26/d4bc2ba3427a623f79a6c10c8f427c7a55b56eb8b3eddc369319d97f741b/scipy_stubs-1.17.1.5-py3-none-any.whl", hash = "sha256:58ebf054a86c000c72e8982e121c4ead0d3d9ba7a6c38aa5fa71b07f96a427fd", size = 607388, upload-time = "2026-05-25T21:34:32.073Z" }, +] + +[[package]] +name = "scipy-stubs" +version = "1.18.0.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "optype", version = "0.18.0", source = { registry = "https://pypi.org/simple" }, extra = ["numpy"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a0/c6/c4c80bd13c0eec99121922c50cd1c72600672307b28e4ba67f636d6b298d/scipy_stubs-1.18.0.1.tar.gz", hash = "sha256:761cb35106200c90a37034d3c8120d3899d673de1e5879d3a2ac1d6816c62327", size = 410862, upload-time = "2026-07-12T21:09:29.993Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/0c/78404528783677077f072288724ff8ef8381bbc219b110e3147e60361b6d/scipy_stubs-1.18.0.1-py3-none-any.whl", hash = "sha256:704408c5f03a33924c9fed6dc9e25926dc331ed508d52f1beece0188525cebad", size = 616217, upload-time = "2026-07-12T21:09:28.45Z" }, +] + [[package]] name = "send2trash" version = "2.1.0" @@ -4099,11 +4562,11 @@ wheels = [ [[package]] name = "setuptools" -version = "83.0.0" +version = "80.10.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/95/faf61eb8363f26aa7e1d762267a8d602a1b26d4f3a1e758e92cb3cb8b054/setuptools-80.10.2.tar.gz", hash = "sha256:8b0e9d10c784bf7d262c4e5ec5d4ec94127ce206e8738f29a437945fbc219b70", size = 1200343, upload-time = "2026-01-25T22:38:17.252Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, + { url = "https://files.pythonhosted.org/packages/94/b8/f1f62a5e3c0ad2ff1d189590bfa4c46b4f3b6e49cef6f26c6ee4e575394d/setuptools-80.10.2-py3-none-any.whl", hash = "sha256:95b30ddfb717250edb492926c92b5221f7ef3fbcc2b07579bcd4a27da21d0173", size = 1064234, upload-time = "2026-01-25T22:38:15.216Z" }, ] [[package]] @@ -4135,11 +4598,11 @@ wheels = [ [[package]] name = "soupsieve" -version = "2.8.4" +version = "2.8.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/47/2c/0a5f6f8ee0d5589e48c7640213ed5175d52cf540a06725b628cc1a45d6ce/soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e", size = 121110, upload-time = "2026-05-24T13:55:57.154Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/ae/2d9c981590ed9999a0d91755b47fc74f74de286b0f5cee14c9269041e6c4/soupsieve-2.8.3.tar.gz", hash = "sha256:3267f1eeea4251fb42728b6dfb746edc9acaffc4a45b27e19450b676586e8349", size = 118627, upload-time = "2026-01-20T04:27:02.457Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" }, + { url = "https://files.pythonhosted.org/packages/46/2c/1462b1d0a634697ae9e55b3cecdcb64788e8b7d63f54d923fcd0bb140aed/soupsieve-2.8.3-py3-none-any.whl", hash = "sha256:ed64f2ba4eebeab06cc4962affce381647455978ffc1e36bb79a545b91f45a95", size = 37016, upload-time = "2026-01-20T04:27:01.012Z" }, ] [[package]] @@ -4215,13 +4678,13 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -4276,17 +4739,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -4340,17 +4803,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [ @@ -4679,6 +5142,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" }, ] +[[package]] +name = "trove-classifiers" +version = "2026.6.1.19" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c2/e3/7ca82ee24c82d344584abd5b8637b3bd056f2900226e8d82fc22f1184b92/trove_classifiers-2026.6.1.19.tar.gz", hash = "sha256:c5132b4b61a829d11cfbd2d72e97f20a45ed6edb95e45c5efdeb5e00836b2745", size = 17059, upload-time = "2026-06-01T19:41:34.649Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/a4/81502f486f01db95bc8320646a8a12511f5e556cb63d5e224d91816605c4/trove_classifiers-2026.6.1.19-py3-none-any.whl", hash = "sha256:ab4c4ec93cc4a4e7815fa759906e05e6bb3f2fbd92ea0f897288c6a43efd15b3", size = 14211, upload-time = "2026-06-01T19:41:33.434Z" }, +] + [[package]] name = "typing-extensions" version = "4.16.0" @@ -5035,17 +5507,17 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "(python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "(python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", ] dependencies = [