A big pain point in various places is that JumpProblem stores mutable state beyond u0 and p (i.e. aggregations, MassActionJumps, and rngs). This makes updating parameters more challenging, complicates remake, and also complicates EnsembleProblems.
If we are willing to allow a breaking release I think at least the rng and MassActionJumps can be fixed in a relatively clean way that we could do in a few phases:
Phase 1:
- We switch to a uniform, integrator stored
rng across SciML.
- Chris and I already discussed this, and I think I now have this mapped out for OrdinaryDiffEq, StochasticDiffEq, and JumpProcesses in a way that will not be breaking for the other two, but allow users to pass a
rng to solve that is stored in integrators. (Note this will also make it much easier to route rngs and handle their seeding for EnsembleProblems!).
Phase 2:
MassActionJumps are modified to store a function that maps parameters to unscaled rate constants (or scaled or both, that is a minor detail).
- Aggregations now store the numeric
scaled_rates, and recalculate them in their initialization from the current parameters and their stored MassActionJump. While this means they are recalculated every solve call I doubt this is a big cost. We could later add a way to communicate whether they really need to be recalculated (e.g. via a kwarg to initialization).
reset_aggregated_jumps is much cleaner as now it just reinitializes the aggregation, which it already does, and which then handles updating the parameters via reinitialization of the aggregation.
- We can drop the SymbolicIndexingInterface hook needed for MTK to initialize MassActionJumps.
Phase 3: This is the part I current feel probably isn't worth doing.
- We move construction of the aggregation and DiscreteCallback to init.
This means that JumpProblems no longer have lots of hidden mutating state that requires deepcopying all over the place. It makes EnsembleProblems and remake much easier to manage I think. However, this has a very big potential overhead as now unless someone uses the integrator interface and reinit!, we completely rebuild the aggregation on each call to solve. In our current EnsembleProblem setup, based around remaking problems and calling solve instead of using integrators, this just wouldn't work well and would add a ton of overhead. However, if EnsembleProblems were modified to support the integrator interface and use reinit! on each thread then maybe this would still be a good change (but we'd then need some kind of integrator_func as a prob_func alternative).
I think phase 1 and 2 are clear wins.
A big pain point in various places is that JumpProblem stores mutable state beyond
u0andp(i.e. aggregations, MassActionJumps, and rngs). This makes updating parameters more challenging, complicatesremake, and also complicatesEnsembleProblems.If we are willing to allow a breaking release I think at least the rng and MassActionJumps can be fixed in a relatively clean way that we could do in a few phases:
Phase 1:
rngacross SciML.rngtosolvethat is stored in integrators. (Note this will also make it much easier to routerngs and handle their seeding forEnsembleProblems!).Phase 2:
MassActionJumps are modified to store a function that maps parameters to unscaled rate constants (or scaled or both, that is a minor detail).scaled_rates, and recalculate them in their initialization from the current parameters and their stored MassActionJump. While this means they are recalculated every solve call I doubt this is a big cost. We could later add a way to communicate whether they really need to be recalculated (e.g. via a kwarg to initialization).reset_aggregated_jumpsis much cleaner as now it just reinitializes the aggregation, which it already does, and which then handles updating the parameters via reinitialization of the aggregation.Phase 3: This is the part I current feel probably isn't worth doing.
This means that JumpProblems no longer have lots of hidden mutating state that requires deepcopying all over the place. It makes EnsembleProblems and remake much easier to manage I think. However, this has a very big potential overhead as now unless someone uses the integrator interface and
reinit!, we completely rebuild the aggregation on each call tosolve. In our current EnsembleProblem setup, based around remaking problems and calling solve instead of using integrators, this just wouldn't work well and would add a ton of overhead. However, if EnsembleProblems were modified to support the integrator interface and usereinit!on each thread then maybe this would still be a good change (but we'd then need some kind ofintegrator_funcas aprob_funcalternative).I think phase 1 and 2 are clear wins.