Abstracting incremental algorithms and supporting numba #2778
jeromekelleher
started this conversation in
Ideas
Replies: 1 comment
-
This is the way to go. I've recently been stuck several times by "I can do this easily with trees but know that it won't be as fast as edge diffs". |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Problem
Something that has bothered me for a while is that our incremental algorithms for computing stats and so on have a lot of duplicated code that's quite subtle. We don't use the
edge_diffs
in C very much because the API is quite clumsy, it involves allocating quite a bit of memory, and it touches memory that we may not need at all for a given calculation. We've ended up duplicating code because it's fairly straighforward, and compiles well (e.g. here, here, here and more)We also have the problem that none of these algorithms now support left-right bounds, and if we were to (say) add support for parallelism to the general stats API, we'd have to update three separate, slightly different, versions of the incremental tree generation code.
Another problem is that we don't have a natural pattern for doing incremental algorithms using Numba (single tree algorithms work really well).
Proposed solution
We create a new C "class" that captures the state required to iterate over the trees. The crucial difference to the edge_diffs methods is that we only store the range of indexes into ordering arrays for the edges_out and edges_in, not the actual edges.
Here's a rough version done using numba to solve the problem in #2774
WARNING this code is not fully tested and probably has subtle logic problems!!!
We get the following results:
It's fast! As we get to larger and larger sample sizes, the time spent in Python using the edge_diffs approach becomes more important. I've included the time required to compute
ts.diversity()
here as a benchmark to give us a sense of how the numba code competes with C code doing something comparable. We can see that for large sample size the numba code is even faster (this is a much simpler algorithm, after all).The key pattern that we'll have for client code will be this:
Note that we could add support for starting at a particular position quite straighforwardly here (and maybe even add support for "direction", so you can go backwards too).
I think this should be as fast as the currently fully-compiled in versions because the call to
tree_pos.next()
should be inlineable, and it should lead to quite good cache behaviour. The numba results above seem to support this anyway.What about numba?
You can do efficient incremental algorithms now using numba by copying the code above into your project, but I guess it would be nice to make this generally available to people so we don't duplicate the code. Maybe we need a tskit-numba package or something?
Beta Was this translation helpful? Give feedback.
All reactions