SciML / JumpProcesses.jl

Build and simulate jump equations like Gillespie simulations and jump diffusions with constant and state-dependent rates and mix with differential equations and scientific machine learning (SciML)
https://docs.sciml.ai/JumpProcesses/stable/
Other
140 stars 35 forks source link

CoevolveSynced #315

Closed gzagatti closed 1 year ago

gzagatti commented 1 year ago

This PR introduces a new aggregator CoevolveSynced.

A known drawback of Coevolve is that it requires to look ahead during the thinning routine. This limits the interoperability of Coevolve with ODEProblems. It turns out that we can sync the thinning routine with the stepper which removes the need to look ahead. The only requirement is that the local bound of the rate computed at time t given u be valid from t to t + rateinterval which is a very mild requirement.

The PR modifies Coevolve as following. Instead of computing the next jump time, it only computes the next candidate time. The candidate is then pushed to the priority queue which is passed to the stepper. The stepper steps at the earliest candidate time and determines whether this time should be accepted or rejected. If the time is rejected, the jump is not computed. Otherwise, the jump is recorded and dependencies are updated. This procedure should be equivalent to Coevolve.

The algorithm only modifies the handling of VariableRateJump while the other jumps are handled the same way as in Coevolve. Despite of that, there seems to be a small penalty as CoevolveSynced runs slightly slower than Coevolve when no VariableRateJumps are involved. I am wondering if there is some obvious efficiencies that I have missed out.

Below you can find the SciMLBenchmarks:

image Diffusion model

image Mendes multivariate

image Negative feedback gene expression model

image Negative feedback Marchetti model

image Multivariate Hawkes model

In terms of correctness, CoevolveSynced produces similar statistics to Coevolve and QQ-plots are comparable as you can see below:

image QQ-plots Multivariate Hawkes model

While we review this PR, I would like to benchmark CoevolveSynced with PDMPCHV for the synapse model discussed in this Discourse thread which is of very practical value.

rveltz commented 1 year ago

A known drawback of Coevolve is that it requires to look ahead during the thinning routine. This limits the interoperability of Coevolve with ODEProblems. It turns out that we can sync the thinning routine with the stepper which removes the need to look ahead. The only requirement is that the local bound of the rate computed at time t given u be valid from t to t + rateinterval which is a very mild requirement.

It seems to indicate that you want to update the rate function according to the previous (stochastic ) history of the current trajectory. Is it what you want to do in the end?

coveralls commented 1 year ago

Pull Request Test Coverage Report for Build 4773569602

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

Details


Changes Missing Coverage Covered Lines Changed/Added Lines %
src/solve.jl 0 1 0.0%
src/aggregators/coevolve.jl 66 72 91.67%
<!-- Total: 68 75 90.67% -->
Files with Coverage Reduction New Missed Lines %
src/JumpProcesses.jl 1 0.0%
src/aggregators/coevolve.jl 2 93.6%
src/solve.jl 5 57.14%
src/SSA_stepper.jl 6 93.75%
src/aggregators/aggregators.jl 7 66.67%
<!-- Total: 21 -->
Totals Coverage Status
Change from base Build 4746294877: 0.6%
Covered Lines: 2306
Relevant Lines: 2605

💛 - Coveralls
gzagatti commented 1 year ago

It seems to indicate that you want to update the rate function according to the previous (stochastic ) history of the current trajectory. Is it what you want to do in the end?

What do you mean by the "previous (stochastic ) history of the current trajectory"? Is that vector u? If that's the case that's correct. The algorithm has a priority queue which holds the candidate times for the next jump. The stepper will stop at the earliest candidate, say time t corresponding to jump k. At time t, the state of the model is u. It then proceeds as following:

  1. If t is longer than rateinterval which was previously computed. The candidate t gets automatically rejected. We then re-compute another candidate using the upper round function, t, u and p.
  2. Otherwise, we determine whether t will be accepted or rejected. To do that, we evaluate the rate at time t, state u and parameters p. We accept the time if the computed rate is more than rng from the upper bound which was previously computed. Otherwise, we reject the time and re-compute another candidate.
  3. If time t is accepted, we update all candidates that depend on k.

Does it make sense now?

gzagatti commented 1 year ago

With the goal of developing a new benchmark for jump problems coupled with continuous ones, I have started working on the synapse model discussed in this Discourse thread.

You can find a discussion thread with regards to this development in rveltz/SynapseElife#1.

In particular, I have raised some questions about the best design pattern for this sort of problem in my latest comment. @isaacsas and @ChrisRackauckas when and if you have time please help with your inputs.

isaacsas commented 1 year ago

@gzagatti let me know when you are finished updating. Hopefully we can merge this shortly and then get the paper finalized.

gzagatti commented 1 year ago

@isaacsas I'm done with the review.

For some reason, JuliaFormatter wanted to format most of the files in the repo, so I didn't apply it. I'll let you do it once we're satisfied with the changes.

isaacsas commented 1 year ago

Don't worry about the formater failing. There was a style change so almost every file will need updating. I don't want to mix those changes up with your PR...

Going to take a last pass through the PR this afternoon, and then hopefully we will be good to go (modulo the comment I made on the other PR about that test/issue).

isaacsas commented 1 year ago

Ok so three issues to settle:

  1. Does the early return in SSAStepper that skips the discrete callbacks when rejecting events make sense? I feel like we should still always check user discrete callbacks as that is kind of the promise of how discrete callbacks work (i.e. they are checked each iteration of the time-stepper).
  2. Whether the u_modified behavior is strictly correct. If @ChrisRackauckas doesn't get back to us soon I'll assume it is ok and proceed.
  3. We need to merge the test and check the issue in the other PR about not evaluating urate and lrate beyond tstop.
gzagatti commented 1 year ago

With regard to your points:

  1. I believe you are talking about line 266 in SSA_stepper.jl. Before Coevolve, every step of the stepper performed a jump so save_everystep=true and save_positions=(false, true)|(true, false)|(true, true) were all equivalent. In fact, we see from line 178 of the same file when initializing the stepper:
    save_everystep = any(cb.save_positions)

Also, I don't see in the stepper code an instruction to save before jumps. From my understanding the code only saves after the jump even if save_positions=(true, false).

With Coevolve, save_everystep and save_positions are not equivalent. Therefore, the setting save_everystep=false and save_positions=(false,true) should only save right after jump times. We should then handle save_positions and save_everystep separately.

In this case, the struct might need to be refactored with an entry for save_positions. Also, the docstring of save_everystep needs to be modified since it currently reads:

    """Whether to save every time a jump occurs."""
    save_everystep::Bool

I was not that comfortable with modifying the stepper. You might have a better idea on how to handle all the saving corner cases, so please let me know how to proceed.

  1. Agreed. The main idea is to avoid expensive re-computation of derivatives when the candidate time is rejected. Nothing in the problem should change. The only thing that needs re-computation is a new candidate time. I took inspiration from the explanation of u_modified! in the documentation and the example on the AutoAbstol tutorial.

  2. How do you want to merge the test? Are you planning to first approve that PR, or should I just merge the test from the PR?

isaacsas commented 1 year ago

1 isn't just about saving. apply_discrete_callback! checks each callback's condition and calls its affect if the condition is true. The SciML docs state the condition is checked after every timestep. The PR changes this behavior and promise currently (if t is being increased I would define that as a timestep). (As an aside, if I recall right, not saving twice is because with SSAStepper the interpolation is piecewise constant. With an ODE solver I believe one does need to save twice since the time interpolation is at least linear.)

For 3. if that PR would result in an incorrect algorithm, and it would be substantial work to fix it, I guess it would be better to just add an analogous test here and make sure the issue is fixed here. I would've liked to release a fix on the older Coevolve for anyone actively using it, but if it would be a major PR to fix it then probably it makes more sense to just merge this PR and release.

gzagatti commented 1 year ago
  1. I have no problem with save_everystep following the model time. However, we need to ensure that the setting save_everystep=false and save_positions=(false,true) will only save at jump time (not time steps). It is otherwise, very difficult to simulate the model to save exactly at the jumps and nothing else. I faced this exact problem when writing the synapse benchmark in here. My suggestion that we treat save_everystep and save_positions separately thus stand.

I didn't quite understand your last point on the interpolation. However, my understanding is that save_positions=(true, false) should save right before the jump, which means that we should just get the value of u before the jump.

  1. Ok, I'll add just add the test in this PR. It will indeed be less work.
isaacsas commented 1 year ago

SSAStepper always just saves after the jump as long as one of the save_positions components is true. This is because there is no reason to save before and after since the solution object sol = solve(...) uses piecewise constant interpolation to be able to exactly evaluate sol(t) for any t (assuming one saves the post-jump value every step). Perhaps this needs better documentation, but this is the intended behavior.

For an ODE time stepper sol(t) is evaluated by using polynomial interpolation from times the timestepper stopped at that bracket t. Having the pre and post jump values becomes necessary to accurately interpolate the solution to time t.

gzagatti commented 1 year ago

Thanks for clarifying on the interpolation. It's much clearer now.

I think it might be worth adding some explanation about it on the documentation. Users might indeed be expecting a behavior similar to the ODE case, such that save_positions=(true, true) would return two values at t. (Or at least that's what I understand from the ODE documentation).

gzagatti commented 1 year ago

I have pushed the changes to the saving behavior that you originally suggested above. In addition to that, I have included some documentation on it to guide library users.

I'm glad we had the discussion above. Thanks to it, the saving behavior of the stepper is much clearer to me.

I hope that's all you need to get the PR merged.

isaacsas commented 1 year ago

@gzagatti all merged, thanks for all your efforts on this!

gzagatti commented 1 year ago

@isaacsas it's always a pleasure working with you. I learn a lot!