Fix CI, add extra test

Author: benjeffery

.github/workflows/tests.yml

@@ -126,14 +126,6 @@ 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/requirements/CI-complete/requirements.txt

@@ -6,6 +6,7 @@ lshmm==0.0.8
 msgpack==1.1.0
 msprime==1.3.3
 networkx==3.2.1
+numba==0.61.2
 portion==2.6.0
 pytest==8.3.5
 pytest-cov==6.0.0

python/requirements/CI-tests-pip/requirements.txt

@@ -11,4 +11,5 @@ networkx==3.2.1
 msgpack==1.1.0
 newick==1.10.0
 kastore==0.3.3
-jsonschema==4.23.0
+jsonschema==4.23.0
+numba>=0.60.0

python/tests/conftest.py

@@ -64,20 +64,13 @@ 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):
@@ -86,18 +79,6 @@ 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

@@ -2,6 +2,9 @@
 import sys
 from unittest.mock import patch
 
+import msprime
+import numba
+import numpy as np
 import pytest
 
 import tests.tsutil as tsutil
@@ -14,7 +17,6 @@ def test_numba_import_error():
             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
@@ -34,3 +36,49 @@ def test_correct_trees_forward(ts):
             range(*numba_edge_diff.edges_out_index_range), edge_diff.edges_out
         ):
             assert edge.id == out_index[edge_out_index]
+
+
+def test_using_from_jit_function():
+    """
+    Test that we can use the numba jit function from the tskit.jit module.
+    """
+    import tskit.jit.numba as jit_numba
+
+    ts = msprime.sim_ancestry(
+        samples=10, sequence_length=100, recombination_rate=1, random_seed=42
+    )
+
+    @numba.njit
+    def _coalescent_nodes_numba(numba_ts, num_nodes, edges_parent):
+        is_coalescent = np.zeros(num_nodes, dtype=np.int8)
+        num_children = np.zeros(num_nodes, dtype=np.int64)
+        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]
+                num_children[edges_parent[e]] -= 1
+            for j in range(*tree_pos.edges_in_index_range):
+                e = numba_ts.indexes_edge_insertion_order[j]
+                p = edges_parent[e]
+                num_children[p] += 1
+                if num_children[p] == 2:
+                    is_coalescent[p] = True
+        return is_coalescent
+
+    def coalescent_nodes_python(ts):
+        is_coalescent = np.zeros(ts.num_nodes, dtype=bool)
+        num_children = np.zeros(ts.num_nodes, dtype=int)
+        for _, edges_out, edges_in in ts.edge_diffs():
+            for e in edges_out:
+                num_children[e.parent] -= 1
+            for e in edges_in:
+                num_children[e.parent] += 1
+                if num_children[e.parent] == 2:
+                    # Num_children will always be exactly two once, even arity is greater
+                    is_coalescent[e.parent] = True
+        return is_coalescent
+
+    numba_ts = jit_numba.numba_tree_sequence(ts)
+    C1 = coalescent_nodes_python(ts)
+    C2 = _coalescent_nodes_numba(numba_ts, ts.num_nodes, ts.edges_parent)
+
+    np.testing.assert_array_equal(C1, C2)

python/tskit/jit/numba.py

@@ -13,11 +13,20 @@
 # 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)
+    dc_cls = dataclass(cls, eq=False)
     del dc_cls.__dataclass_params__
     del dc_cls.__dataclass_fields__
     del dc_cls.__repr__
-    del dc_cls.__replace__
+    try:
+        del dc_cls.__replace__
+    except AttributeError:
+        # __replace__ is not available in Python < 3.10
+        pass
+    try:
+        del dc_cls.__match_args__
+    except AttributeError:
+        # __match_args__ is not available in Python < 3.10
+        pass
     return numba.experimental.jitclass(dc_cls)