`Coevolve` aggregator for `VariableRateJump`

gzagatti commented 1 year ago

This PR introduces the Coevolve aggregator for VariableRateJump. This aggregator supports the SSAStepper() integrator for the efficient simulation of jump processes with time-varying rates.

Some background, in JumpProcesses.jl most of the algorithms are borrowed from the biochemistry literature. However, jump processes have also been extensively studied in the statistics literature where they go by point processes (see Daley and Vere-Jones [1]). In the point process literature, we find many of the simulation algorithms implemented in this library that go by other names. Among the most popular one is Ogata's algorithm [2] which is very similar to Gillespie's direct method. Ogata's algorithm as adapted in Daley [1] is a general-purpose algorithm for any point process that evolves through time. Farajtabar et al. [3] modifies Ogata's algorithm to use a priority queue in order to simulate a compound Hawkes process using a priority queue. Here I adapt Farajtabar's COEVOLVE algorithm to simulate any evolutionary point process, thus the name of the algorithm.

This algorithm requires user-defined bounds and a dependency graph. Therefore, I extend the VariableRateJump to accept three aditional keyword arguments lrate, urate and L. lrate(u, p, t) is a function that computes the lower-bound of the rate from t to t + L, urate computes the upper-bound, and L computes the interval for which the bounds are valid.

A toy example is given below

using JumpProcesses, Plots
L(u,p,t) = (1/p[1])*2
rate(u,p,t) = t*p[1]*u[1]
lrate(u, p, t) = rate(u, p, t)
urate(u,p,t) = rate(u, p, t + L(u,p,t))
affect!(integrator) = integrator.u[1] -= 1
vrj = VariableRateJump(rate, affect!; lrate=lrate, urate=urate, L=L)
u0 = [1_000.]
p = (1.0,)
tspan = (0., 10.)
prob = DiscreteProblem(u0, tspan, p)
jprob = JumpProblem(prob, Coevolve(), vrj; dep_graph=[[1]])
sol = solve(jprob, SSAStepper())
plot(sol, title="Trajectory", legend=false)

I have performed some benchmarks to evaluate the performance of the model. I ran the jump benchmarks of SciMLBenchmarks. Coevolve performs on par with the alternatives, with a fixed penalty over NRM. The QueueMethod basically reduces to the NRM when all the jumps are ConstantRateJump or MassActionJump, but I was not able to equate the NRM performance.

Benchmark	Direct	FRM	SortingDirect	NRM	DirectCR	RSSA	RSSACR	Coevolve
Diffusion CTRW	5.1 s		1.2 s	0.7 s	0.4 s	1.6 s	0.4 s	1 s
Multistate Model	0.1 s	0.2 s	0.1 s	0.2 s	0.2 s	0.1 s	0.1 s	0.4 s
Neg. Feed. Gene Expression	0.2 ms	0.3 ms	0.2 ms	0.4 ms	0.4 ms	0.4 ms	0.8 ms	0.4 ms
Marchetti Gene Expression	0.4 s	0.6 s	0.4s	0.9 s	0.0.9 s	0.6 s	0.9 s	1.3 s

More complex examples include the compound Hawkes process which was the main motivation for contributing this PR.

I have also developed a new benchmark using the compound Hawkes process and found significant improvements compared to using the ODE solver. Please check my branch of SciMLBenchmarks for more details. Note that the Direct method does not run within the allowed time when the number of nodes reaches above 40.

V	Direct (n)	Direct (median time)	Coevolve (n)	QueueMethod (median time)
1	50	124.5 μs	50	3.4 μs
10	50	28.6 ms	50	455.5 μs
20	50	223.8 ms	50	5.2 ms
30	50	715.1 ms	50	7.9 ms
40	31	715.1 ms	50	30.2 ms
50	15	4.1 s	50	61.7 ms
60	10	6.7 s	50	137.1 ms
70	6	10.9 s	50	202.9 ms
80	4	17.8 s	50	346.4 ms
90	3	27.24 s	50	642.0 ms

I have added unit tests and updated the documentation to reflect the changes introduced here.

As I completed this PR, I learnt about PR #252 which also extends VariableRateJump. The author implements a different algorithm which also uses rate bounds. While the algorithm implemented in that PR is analogous to rejection-based algorithms, the algorithm in this PR is analogous to next-reaction methods. So they both face different trade-offs.

I am looking forward to your feedback.

[1] D. J. Daley and D. Vere-Jones, An Introduction to the Theory of Point Processes: Volume I: Elementary Theory and Methods, 2nd ed. New York: Springer-Verlag, 2003. doi: 10.1007/b97277.

[2] Y. Ogata, “On Lewis’ simulation method for point processes,” IEEE Transactions on Information Theory, vol. 27, no. 1, pp. 23–31, Jan. 1981, doi: 10.1109/TIT.1981.1056305.

[3] M. Farajtabar, Y. Wang, M. Gomez-Rodriguez, S. Li, H. Zha, and L. Song, “COEVOLVE: a joint point process model for information diffusion and network evolution,” J. Mach. Learn. Res., vol. 18, no. 1, pp. 1305–1353, Jan. 2017, doi: 10.5555/3122009.3122050.

isaacsas commented 1 year ago

This looks great! Thanks also for so comprehensively commenting and adding docs/tests.

I’ll try to give this a comprehensive review later today.

coveralls commented 1 year ago

Pull Request Test Coverage Report for Build 3644727547

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.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

47 of 54 (87.04%) changed or added relevant lines in 3 files are covered.
60 unchanged lines in 7 files lost coverage.
Overall coverage increased (+0.7%) to 89.028%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/aggregators/aggregators.jl	3	4	75.0%
src/jumps.jl	19	22	86.36%
src/problem.jl	25	28	89.29%
<!--	Total:	47	54	87.04%	-->

Files with Coverage Reduction	New Missed Lines	%
src/extended_jump_array.jl	1	27.54%
src/aggregators/prioritytable.jl	1	86.86%
src/coupling.jl	2	78.18%
src/spatial/spatial_massaction_jump.jl	3	85.45%
src/SSA_stepper.jl	4	94.12%
src/jumps.jl	22	86.19%
src/problem.jl	27	71.26%
<!--	Total:	60	-->

Totals
Change from base Build 3544760886:	0.7%
Covered Lines:	2069
Relevant Lines:	2324

💛 - Coveralls

gzagatti commented 1 year ago

I renamed QueueMethod to Coevolve to better reflect the source of the algorithm. Also, I felt QueueMethod was too generic and the Method suffix a bit redundant.

isaacsas commented 1 year ago

Sorry, grading took longer than expected today, so probably won’t get to looking this over till tomorrow afternoon (last day of classes here…).

gzagatti commented 1 year ago

Don't worry. Grading is always time consuming. Thanks for the update.

gzagatti commented 1 year ago

Thanks for your suggestions! See answers to your comments below. I will be pushing some changes later today. I am looking forward to hearing your reactions after I make the pushes.

Should L take the upper/lower/current rates as optional keywords to avoid having to recalculate them (when they are available already from the algorithm)?

Should we consider merging urate, lrate and L into one function that returns all three? I can argue for/against this, so I'd be interested in hearing your thoughts.

See my answer above on this.

Why the changes to JumpSets? Are you sure this doesn't break their working with and preservingTuples for Direct simulations?

See my answer above on this. But, so far I have not faced any issues with using the other aggregators and all tests pass.

What do you mean by preservingTuples for Direct simulations?

I'd prefer to keep ConstantRateJumps distinguished in the algorithms, and more generally just handle them and MassActionJumps specially. We can assume an ordering of MassActionJumps, ConstantRateJumps and VariableRateJumps for agrregators that support all three. In that case, couldn't you then just reuse the NRM time updates for them? I'd be fine with some refactoring of NRM to make that easier to reuse (or even considering merging NRM and Coevolve if we can make it functionally equivalent for systems without VariableRateJumps).

See my answer above on this. In summary, I think we should honour the current API that tells aggregator A will be used to solve the problem, so aggregator B should not be used.

With that being said, the main loop in Coevolve should reduce to the one in NRM when using MassActionJumps and/or ConstantRateJumps.

The main differences is that NRM re-uses the random number to re-compute the new time in update_dependent_rates at line 103:

update!(p.pq, rx, t + oldrate / cur_rates[rx] * (p.pq[rx] - t))

This come at the cost of indexing the mutable binary heap.

On the other hand, Coevolve always generate a new random number when updating the jump. Everything else in Coevolve should be the same since the loop short-circuits when evaluating the next time of any jump with an infinity time window L.

The solution then is to use the random number generated previously to re-compute s (the time to the next jump), which should get us closer to the performance of NRM. This should not be hard to implement.

Can you update some of the other tests that loop through most of the aggregators and add Coevolve to them? It'd be good to check that it works for mixes of MassActionJumps and ConstantRateJumps too.

I think I have now updated all tests that loop through aggregators to include Coevolve. Let me know otherwise.

gzagatti commented 1 year ago

I have made all the changes following your feedback. You can now look at the modifications for a second round of comments.

In particular, I have used the same NRM trick of re-using the random number when updating the dependent rates which brought some efficiency gains. I am still not beating the NRM performance but it's now much closer to it.

Benchmark	Direct	FRM	SortingDirect	NRM	DirectCR	RSSA	RSSACR	Coevolve
Diffusion CTRW	4.9 s		1.2 s	0.7 s	0.4 s	1.7 s	0.4 s	0.7 s
Multistate Model	0.1 s	0.2 s	0.1 s	0.2 s	0.2 s	0.1 s	0.1 s	0.3 s
Neg. Feed. Gene Expression	0.2 ms	0.3 ms	0.2 ms	0.3 ms	0.4 ms	0.4 ms	0.8 ms	0.4 ms
Marchetti Gene Expression	0.5 s	0.7 s	0.4 s	0.8 s	0.9 s	0.6 s	0.8 s	1.0 s

isaacsas commented 1 year ago

The new numbers look much better! However, I would still strongly prefer to drop the function creation in get_rates and just explicitly handle MassActionJumps and ConstantRateJumps rate/time updates, bypassing the rejection while loop / logic. As I said above, there is no reason the performance shouldn't be identical to NRM for them I would think.

gzagatti commented 1 year ago

I tried to make changes to the next_time function to avoid creating anonymous functions. I am afraid that we still hit the same performance as the last one I reported earlier as you can see below:

Benchmark	Direct	FRM	SortingDirect	NRM	DirectCR	RSSA	RSSACR	Coevolve
Diffusion CTRW	4.7 s		1.1 s	0.7 s	0.4 s	1.5 s	0.4 s	0.7 s
Multistate Model	0.1 s	0.2 s	0.1 s	0.2 s	0.2 s	0.1 s	0.1 s	0.3 s
Neg. Feed. Gene Expression	0.2 ms	0.3 ms	0.2 ms	0.3 ms	0.4 ms	0.4 ms	0.7 ms	0.4 ms
Marchetti Gene Expression	0.5 s	0.7 s	0.4 s	0.8 s	0.9 s	0.6 s	0.8 s	1.0 s

You previously mentioned that creating functions and capturing a variable could cause an overhead. I am wondering if in the previous iteration Julia was not simply compiling all this away to reach a performance as close to returning a number which is what is doing now. My understanding is that Julia really shines on doing these kind of smart compilations which allow us to write more natural code without that many branches.

With the current iteration, I feel the loop in next_time is not as neat as it was in the previous one, but it is marginally optimized.

Let me know what you think.

gzagatti commented 1 year ago

Yet, one more iteration. I have completely separated the treatment of each jump type in Coevolve. I believe you might prefer this last iteration because of that.

In terms of performance, we still obtain very similar values from the previous two iterations:

Benchmark	Direct	FRM	SortingDirect	NRM	DirectCR	RSSA	RSSACR	Coevolve
Diffusion CTRW	4.9 s		1.2 s	0.7 s	0.4 s	1.7 s	0.3 s	0.7 s
Multistate Model	0.1 s	0.2 s	0.1 s	0.2 s	0.2 s	0.1 s	0.1 s	0.3 s
Neg. Feed. Gene Expression	0.2 ms	0.2 ms	0.2 ms	0.3 ms	0.4 ms	0.4 ms	0.7 ms	0.4 ms
Marchetti Gene Expression	0.4 s	0.6 s	0.4 s	0.8 s	0.8 s	0.6 s	0.8 s	1.0 s

I'll post the performance for the Hawkes benchmark as soon as it is completed.

gzagatti commented 1 year ago

Here is the performance for the Hawkes benchmark, much improved from the first iteration:

V	Direct (n)	Direct (median time)	Coevolve (n)	QueueMethod (median time)
1	50	203.7 μs	50	2.6 μs
10	50	30.9 ms	50	451.1 μs
20	50	221.5 ms	50	2.2 ms
30	50	776.5 ms	50	10.7 ms
40	33	1.9 s	50	23.0 ms
50	16	3.8 s	50	46.0 ms
60	10	6.6 s	50	73.1 ms
70	6	11.7 s	50	120.4 ms
80	4	17.9 s	50	188.7 ms
90	3	26.8 s	50	258.4 ms

isaacsas commented 1 year ago

I'm perplexed why it isn't matching NRM now; I'll try to play around with it tonight. (Back to grading today unfortunately -- but then I should be done with grading till Thursday).

isaacsas commented 1 year ago

@gzagatti I made some tutorial wording changes and fixes to get it to build docs locally for preview purposes. Can you follow the directions at

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork

to enable maintainers to push to your fork's PR branch so I can push to the PR? Thanks!

gzagatti commented 1 year ago

The "Allow edits by maintainers" is activated. Let me know if you are still having trouble pushing to the fork branch.

isaacsas commented 1 year ago

Weird, pushing to the pr branch seems to update your repo but not update the PR. I've never had that happen before.

I'm just going to push to a new PR that can be for updates to the docs from this PR. (So we can finish this one as is and I'll just tweak the docs there and merge it after this merges.)

isaacsas commented 1 year ago

Nevermind, looks like it did work. Strange.

isaacsas commented 1 year ago

@gzagatti just to update you. I've gone through and updated a subset of the tutorials, along with updating some of the new doc strings. I also changed L to rateinterval as I think we need something more descriptive for this field. If you have a better name though I'm happy to change it again.

Stil TODO on my end:

Finish tweaking the docs I haven't yet gotten to.
Figure out how we want to handle JumpProblem. It seems quite a bit more complicated now as rewritten, so it would be nice to try to simplify it. Also, we'd like to make sure we can handle mixes of the two VariableRateJump types now (what I call general vs. bounded VariableRateJumps in the tutorial/docs I've updated). Adding a test for this would be good too (i.e. a mix of bounded and unbounded VariableRateJumps, and hence requiring being over an ODEProblem, but still using Coevolve for the bounded jumps).
(Future work): After this PR we should figure out why performance is below NRM for non-VariableRateJump systems. Given we are using FunctionWrappers there may also be a small gain to be had by keeping lrate=nothing when no lrate is specified, and skipping the call to the lrate function in the rejection sampling code.

A few minor comments for you:

I updated the docs so they can build locally. As they are dynamic it is good to try that out yourself to make sure everything looks ok and the links/examples all work. You just go in the doc folder, activate the Project.toml there, and call include("make.jl"). The resulting build directory should then have the local doc build.
Since we assume Julia 1.6, you don't need f(u; p = p) you can just say f(u; p) and if there is a variable with the name p it will get picked up for you.

isaacsas commented 1 year ago

@gzagatti do the changes I made look ok to you? Otherwise I think this is good to go (though some more tests mixing bounded and general VariableRateJumps, along with the other jump types would be good at some point).

isaacsas commented 1 year ago

All merged. Thanks @gzagatti this is a great PR.

gzagatti commented 1 year ago

You're welcome. Thanks for your support @isaacsas!

xiaomingfu2013 commented 1 year ago

Hi, thank you for the great work to support VariableRateJump type! Maybe I missed something, in https://github.com/SciML/JumpProcesses.jl/blob/2afd12369ef5e848543a9e315d2ab02c8f44a667/src/aggregators/coevolve.jl#L30 it seems to me that !isempty(rs) should check whether the JumpProblem has either ConstantRateJump or VariableRateJump. But in the construction of Coevolve aggregator, rates only refers to VariableRateJump https://github.com/SciML/JumpProcesses.jl/blob/2afd12369ef5e848543a9e315d2ab02c8f44a667/src/aggregators/coevolve.jl#L66 should the condition !isempty(rs) be changed to !isempty(urates) instead? Thank you!

isaacsas commented 1 year ago

Hi-, good catch. Could you submit a PR with that change?

SciML / JumpProcesses.jl