Skip to content

Reuse one TS across the trajectories of an ensemble - #326

Merged
singhharsh1708 merged 8 commits into
mainfrom
ensemble-reuse
Oct 25, 2026
Merged

singhharsh1708 merged 8 commits into
mainfrom
ensemble-reuse

Conversation

@singhharsh1708

@singhharsh1708 singhharsh1708 commented Oct 24, 2026 •

Copy link
Copy Markdown
Member

PETScDiffEq's solve builds and destroys its TS on every call, so in an EnsembleProblem each trajectory pays the full setup. #319 made reinit! keep the TS when nothing structural changed. This uses that inside ensembles: EnsembleSerial() and EnsembleThreads() solves now restart one kept TS for every trajectory that differs only in u0, p or tspan, and set up a new one otherwise, and PETScDiffEq.reusing(f) gives the same to a loop of solves such as a parameter sweep or a fit. SciMLBase's ensemble code is unchanged; the solves it makes simply find the kept TS.

A kept TS is reused only when the problem type, algorithm and keywords are the same, the state has the same size and the span runs the same way; keyword arrays, algorithm options, the mass matrix and differential_vars must also equal copies taken at setup, the Jacobian prototype's sparsity pattern must match, and f must be the same function. A prob_func that changes the mass matrix in place, swaps f or changes the sparsity gets a fresh setup. Failed or retried solves, callbacks, tstops, Float32, distributed or DM solves and the methods #319 does not reuse are never kept. A SplitODEProblem with an additive ARKIMEX tableau is reused too; the SplitFunction's scratch cache, which SciMLBase builds anew for each solve, is not compared. Warnings the setup prints are printed again for every solve that reuses a kept TS, as independent solves would print them. Kept TSs are registered in LIVE_HANDLES and destroyed when the outermost reusing returns. _reuse! gained a tssolve flag, because TSSolve resets PETSc's reject and iteration counts itself.

Every trajectory is bitwise identical to an independent solve (t, u, retcode and all stats). The new testset checks this for TSRK, TSRosW and BDF under p, u0 and tspan changes, in-place p, structural changes, a callback, failing trajectories and a 64-trajectory EnsembleThreads run at -t 4 (139 tests), and over 30 repeated ensembles the live heap stays flat with no kept handles left. Per trajectory on Lotka-Volterra with default tolerances (minimum of 3 runs, N = 1000 and 10000): TSRK("5dp") 65.2 and 70.3 µs to 30.7 and 31.3 µs, TSRosW("ra34pw2") 451 and 443 µs to 389 and 382 µs, TSImplicit("bdf") 397 and 390 µs to 343 and 326 µs. EnsembleThreads at -t 4 gives the same speedup, as solves still run one at a time under the lock. On HIRES and the 1000-state Brusselator the change is within noise, since setup is a small part of each solve there.

AI Disclosure: Claude assisted with the code. Testing and review by @singhharsh1708.

@singhharsh1708
singhharsh1708 merged commit 8a075f7 into main Oct 25, 2026
31 of 33 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant