Initial numba module

[benjeffery]

Part of #3135

[benjeffery]

I've done quite a bit of numba investigation and found a way to use dataclasses in numba code. This seems to come at very little performance cost compared to tuples and is a lot nicer. Using a generator also seems to work fine!

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 128 lines in your changes missing coverage. Please review.

Project coverage is 89.24%. Comparing base (c5bf634) to head (a3f3567).

Files with missing lines	Patch %	Lines
python/tskit/jit/numba.py	0.00%	128 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3225      +/-   ##
==========================================
- Coverage   89.61%   89.24%   -0.37%     
==========================================
  Files          28       29       +1     
  Lines       31983    32111     +128     
  Branches     5888     5898      +10     
==========================================
- Hits        28660    28656       -4     
- Misses       1888     2018     +130     
- Partials     1435     1437       +2     

Flag	Coverage Δ
c-tests	86.59% <ø> (ø)
lwt-tests	80.38% <ø> (ø)
python-c-tests	88.15% <ø> (ø)
python-tests	98.77% <ø> (-0.03%)	⬇️
python-tests-numpy1	51.57% <0.00%> (-0.87%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
python/tskit/jit/numba.py	0.00% <0.00%> (ø)

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[benjeffery]

Having some CI weirdness that I'm not yet able to recreate.

[benjeffery]

CI Fixed.

Here's some benchmarking with the "coalescent_nodes" method from #2778 on a TS with 12M edges:

Using ts.edge_diffs: 23.3s
Calculating edge diffs and coalescent nodes in a single numba.njit function: 0.085s
Using the classes here, calculating coalescent nodes in separate client numba.njit function: 0.093s

[jeromekelleher]

Shall we move the first commit into its own PR? It's cluttering up this one and making it hard to see the real changes.

[jeromekelleher]

I had imagined something lower level that was basically a copy of the TreePosition class from here: https://github.com/jeromekelleher/sc2ts/blob/7758245c3dc537aeec3b7cd6282241b65f8843dd/sc2ts/jit.py#L107

So, we don't try to provide Pythonic APIs, but just provide direct access to the edges out and edges in, which can be numba compiled like the example in the sc2ts code.

[benjeffery]

just provide direct access to the edges out and edges in

That's how this code works,
Your sc2ts code here:

    while tree_pos.next():
        for j in range(tree_pos.out_range[0], tree_pos.out_range[1]):
            e = tree_pos.edge_removal_order[j]
            c = edges_child[e]
            p = edges_parent[e]
            parent[c] = -1
            u = p
            while u != -1:
                num_samples[u] -= num_samples[c]
                u = parent[u]

becomes

    for tree_pos in numba_ts.edge_diffs():
        for j in range(*tree_pos.edges_out_index_range):
            e = numba_ts.indexes_edge_removal_order[j]
            c = edges_child[e]
            p = edges_parent[e]
            parent[c] = -1
            u = p
            while u != -1:
                num_samples[u] -= num_samples[c]
                u = parent[u]

It is still compiled, and 30% faster (for the coalesent nodes example)!

[jeromekelleher]

Ahh, I didn't spot that sorry. How is it faster then?

I do think we should just stick with the TreePosition interface though, because we want to support seeking backwards as well, and ultimately randomly. There's no point in adding a layer for indirection on top of that.

[benjeffery]

How is it faster then?

Mutating numpy arrays to maintain the state involves the following:

1. Creating a temporary list (build_list).
2. Performing bounds checks for the slice.
3. Copying the data from the list into the array's memory.

Whereas yielding lightweight immuatable objects is much more amenable to numba optimisation. We might be able to get the same gains by using native objects for the state rather than numpy arrays if you are set against iteration.

[jeromekelleher]

Let's talk it through in person - I don't have time to form an educated opinion I'm afraid.

[benjeffery]

I've tried to closely match the exisiting tsutil implementation with next and prev. Need to so some perf checks with this new code under numba.

[benjeffery]

New code looks just as fast, proceeding to add some more tests. Will merge this then before doign docs.

[benjeffery]

Getting some weird failures on Windows here, and coverage not counting for the new module, will fix.

I've added a stab at some docs.

[hyanwong]

Re docs, eventually we probably want a "high performance" tutorial with some of this stuff, but I can have a stab at that after 1.0. There's some comments here: tskit-dev/tutorials#150 (comment) and some code examples at tskit-dev/tutorials#63

[hyanwong]

I normally think of "tree traversal" as iterating through the tree structure itself. Do you mean "Iterating through trees" here? I can't see any pre/postorder traversal code here.

[benjeffery]

Good point, I was avoiding the word "iteration" not not confuse it with a Python iterator - but I'll change it back as this is more confusing!

[hyanwong]

Maybe just "moving between trees"?