Skip to content

Towards a JumpProblem with minimal mutating state #554

@isaacsas

Description

@isaacsas

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:

  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:

  1. MassActionJumps are modified to store a function that maps parameters to unscaled rate constants (or scaled or both, that is a minor detail).
  2. 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).
  3. 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.
  4. 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.

  1. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions