Initial numba edge diffs

Author: benjeffery

.github/workflows/tests.yml

@@ -126,6 +126,14 @@ jobs:
             python -m pytest -x --cov=tskit --cov-report=xml --cov-branch -n2 --durations=20 tests
           fi
 
+      - name: Run numba tests
+        working-directory: python
+        run: |
+          source ~/.profile
+          conda activate anaconda-client-env
+          pip install numba
+          python -m pytest -x --only-numba-tests --cov=tskit.numba --cov-report=xml --cov-branch -n2 --durations=20 tests
+  
       - name: Upload coverage to Codecov
         uses: codecov/[email protected]
         with:

python/tests/conftest.py

@@ -64,13 +64,20 @@ def pytest_addoption(parser):
         default=False,
         help="To help debugging, draw lines around the plotboxes in SVG output files",
     )
+    parser.addoption(
+        "--only-numba-tests",
+        action="store_true",
+        default=False,
+        help="Only run tests marked with @pytest.mark.numba",
+    )
 
 
 def pytest_configure(config):
     """
     Add docs on the "slow" marker
     """
     config.addinivalue_line("markers", "slow: mark test as slow to run")
+    config.addinivalue_line("markers", "numba: mark test as a Numba test")
 
 
 def pytest_collection_modifyitems(config, items):
@@ -79,6 +86,18 @@ def pytest_collection_modifyitems(config, items):
         for item in items:
             if "slow" in item.keywords:
                 item.add_marker(skip_slow)
+    if config.getoption("--only-numba-tests"):
+        only_numba = pytest.mark.skip(reason="--only-numba-tests specified")
+        for item in items:
+            if "numba" not in item.keywords:
+                item.add_marker(only_numba)
+    else:
+        numba_tests_skipped = pytest.mark.skip(
+            reason="--only-numba-tests not specified, skipping numba tests"
+        )
+        for item in items:
+            if "numba" in item.keywords:
+                item.add_marker(numba_tests_skipped)
 
 
 @fixture

python/tests/test_jit.py

@@ -0,0 +1,36 @@
+import itertools
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import tests.tsutil as tsutil
+
+
+def test_numba_import_error():
+    # Mock numba as not available
+    with patch.dict(sys.modules, {"numba": None}):
+        with pytest.raises(ImportError, match="pip install numba"):
+            import tskit.jit.numba  # noqa: F401
+
+
+@pytest.mark.numba
+@pytest.mark.parametrize("ts", tsutil.get_example_tree_sequences())
+def test_correct_trees_forward(ts):
+    import tskit.jit.numba as jit_numba
+
+    numba_ts = jit_numba.numba_tree_sequence(ts)
+    in_index = ts.indexes_edge_insertion_order
+    out_index = ts.indexes_edge_removal_order
+    for numba_edge_diff, edge_diff in itertools.zip_longest(
+        numba_ts.edge_diffs(), ts.edge_diffs()
+    ):
+        assert edge_diff.interval == numba_edge_diff.interval
+        for edge_in_index, edge in itertools.zip_longest(
+            range(*numba_edge_diff.edges_in_index_range), edge_diff.edges_in
+        ):
+            assert edge.id == in_index[edge_in_index]
+        for edge_out_index, edge in itertools.zip_longest(
+            range(*numba_edge_diff.edges_out_index_range), edge_diff.edges_out
+        ):
+            assert edge.id == out_index[edge_out_index]

python/tskit/jit/__init__.py

@@ -0,0 +1,83 @@
+from dataclasses import dataclass
+
+
+try:
+    import numba
+except ImportError:
+    raise ImportError(
+        "Numba is not installed. Please install it with `pip install numba` "
+        "or `conda install numba` to use the tskit.jit.numba module."
+    )
+
+
+# Decorator that makes a jited dataclass by removing certain methods
+# that are not compatible with Numba's JIT compilation.
+def jitdataclass(cls):
+    dc_cls = dataclass(cls, eq=False, match_args=False)
+    del dc_cls.__dataclass_params__
+    del dc_cls.__dataclass_fields__
+    del dc_cls.__repr__
+    del dc_cls.__replace__
+    return numba.experimental.jitclass(dc_cls)
+
+
+@jitdataclass
+class NumbaEdgeDiff:
+    interval: numba.types.UniTuple(numba.float64, 2)
+    edges_in_index_range: numba.types.UniTuple(numba.int32, 2)
+    edges_out_index_range: numba.types.UniTuple(numba.int32, 2)
+
+
+@jitdataclass
+class NumbaTreeSequence:
+    num_edges: numba.int64
+    sequence_length: numba.float64
+    edges_left: numba.float64[:]
+    edges_right: numba.float64[:]
+    indexes_edge_insertion_order: numba.int32[:]
+    indexes_edge_removal_order: numba.int32[:]
+
+    def edge_diffs(self, include_terminal=False):
+        left = 0.0
+        j = 0
+        k = 0
+        edges_left = self.edges_left
+        edges_right = self.edges_right
+        in_order = self.indexes_edge_insertion_order
+        out_order = self.indexes_edge_removal_order
+
+        while j < self.num_edges or left < self.sequence_length:
+            in_start = j
+            out_start = k
+
+            while k < self.num_edges and edges_right[out_order[k]] == left:
+                k += 1
+            while j < self.num_edges and edges_left[in_order[j]] == left:
+                j += 1
+            in_end = j
+            out_end = k
+
+            right = self.sequence_length
+            if j < self.num_edges:
+                right = min(right, edges_left[in_order[j]])
+            if k < self.num_edges:
+                right = min(right, edges_right[out_order[k]])
+
+            yield NumbaEdgeDiff((left, right), (in_start, in_end), (out_start, out_end))
+
+            left = right
+
+        # Handle remaining edges that haven't been processed
+        if include_terminal:
+            yield NumbaEdgeDiff((left, right), (j, j), (k, self.num_edges))
+
+
+def numba_tree_sequence(ts):
+    return NumbaTreeSequence(
+        num_edges=ts.num_edges,
+        sequence_length=ts.sequence_length,
+        edges_left=ts.edges_left,
+        edges_right=ts.edges_right,
+        indexes_edge_insertion_order=ts.indexes_edge_insertion_order,
+        indexes_edge_removal_order=ts.indexes_edge_removal_order,
+    )