Repository navigation
Reuse one TS across the trajectories of an ensemble - #326
Merged
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PETScDiffEq's
solvebuilds and destroys its TS on every call, so in anEnsembleProblemeach trajectory pays the full setup. #319 madereinit!keep the TS when nothing structural changed. This uses that inside ensembles:EnsembleSerial()andEnsembleThreads()solves now restart one kept TS for every trajectory that differs only inu0,portspan, and set up a new one otherwise, andPETScDiffEq.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_varsmust also equal copies taken at setup, the Jacobian prototype's sparsity pattern must match, andfmust be the same function. A prob_func that changes the mass matrix in place, swapsfor 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. ASplitODEProblemwith an additive ARKIMEX tableau is reused too; theSplitFunction'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 inLIVE_HANDLESand destroyed when the outermostreusingreturns._reuse!gained atssolveflag, 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
EnsembleThreadsrun 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 4gives 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.