Callbacks on EnsembleProblem rather then on ODEProblem within Ensemble


I have an EnsembleProblem, or at least I want to model a system of equations with that. I want the ensemble member’s parameters to modify based on the solutions of some ensemble members every
N timesteps. It is like a discretized coupling of subsets of the members, based on some conditions of their state (which members are coupled is a priory NOT known).
I hoped to model this with Discrete callbacks of an EnsembleProblem but the callbacks seem to act on each member independently. But the callbacks would have to communicate. I would need the “integrator” in the callback options to be the EnsembleProblem rather than a member of the ensemble. As far as I understand this is not possible because each member is solved independently and there is not necessarily a common timestepping. That might be the same reason why init(EnsembleProblem, … ) is also not defined?

The alternative is of course to write all ensemble members in one system, which might be substantially slower since my system will be large O(7e3), or more.

Maybe someone has advice on how to approach this?


You can just define a non-local variable which you access in the callback condition and affect! functions. What did you try? Doing it like that should be rather straightforward and similar to some of the tutorials which flip a global parameter. You just need to make sure you write thread-safe code if you’re multithreading it.

What I end up doing is basically reinventing the ensemble idea for my purposes, which is probably not what you suggest, but in this “ensemble” of systems, I have to replace the state vector routinely. So it is not really an ensemble of one ODE that samples the parameter space, but just a bunch of ODEs that follow the same system (the same physics) but can have very different states.

1.) define an array of integrators that hold my (large) set of ODEs.
2.) init each member
3.) advance each ODE using Distributed.pmap by timestep DT.
4.) evaluate the state vectors of all members, come up with new initial conditions, reinit!() or set_u!() each state vector
→ to 3.)

Nearly all steps are simply parallelizable, while each ODE is on its own dt (they can be very stiff). The evaluation of all state vectors is the main bottleneck since I have to loop (or map) through all integrator objects. This probably lacks efficiency, but it works for now.

I haven’t tried your suggestion yet. I am actually not sure if I understand, maybe you can elaborate?

Instead of init and reinit!, set_u! everywhere, they can just be all one solve call that syncs inside of a DiscreteCallback. That probably won’t make too much of a difference in the end though, since what you did isn’t equivalent to building the ensemble idea, instead it’s equivalent to reinventing the DiscreteCallback system (which is just implemented via checking conditions and then doing a few modifications to the integrator before eventually resetting a few FSAL caches and going some more).