New tskit.jit module

[jeromekelleher]

We've talked about having a numba support for tskit in various places, so let's add a new optional module "tskit.jit" which provides basic utilities for numba. This would be an optional dependency, and have an entirely separate namespace. E.g

import tskit.jit as tsjit
import numba as nb

@nb.njit
def algorithm(tree_pos)
  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]
            # Do something with out edge ID e
        for j in range(tree_pos.in_range[0], tree_pos.in_range[1]):
            e = tree_pos.edge_insertion_order[j]
            # Do something with in edge ID e

tree_pos = tsjit.TreePosition(ts)
algorithm(tree_pos)

This should enable a lot of downstream algorithms and reduce code duplication.

[jeromekelleher]

So, to be clear, this would not be installed by default and would only be included in the install if we did pip install tskit[jit] or something similar.

[gregorgorjanc]

Is this so you don't depend on numba by default?

[jeromekelleher]

Yep. It's a tricky dependency, not needed by downstream things like msprime.

[benjeffery]

See also #2778

[benjeffery]

Fixed in several PRs e.g. #3225, #3256, #3255