Remove dataclass, add ts properties

Author: benjeffery

python/tests/test_jit.py

@@ -84,3 +84,54 @@ def coalescent_nodes_python(ts):
     C2 = _coalescent_nodes_numba(numba_ts, ts.num_nodes, ts.edges_parent)
 
     np.testing.assert_array_equal(C1, C2)
+
+
+def test_numba_tree_sequence_properties(ts_fixture):
+    """
+    Test that NumbaTreeSequence properties have correct contents and dtypes.
+    """
+    ts = ts_fixture
+    import tskit.jit.numba as jit_numba
+
+    numba_ts = jit_numba.numba_tree_sequence(ts)
+
+    assert numba_ts.num_edges == ts.num_edges
+    assert numba_ts.sequence_length == ts.sequence_length
+    np.testing.assert_array_equal(numba_ts.edges_left, ts.edges_left)
+    np.testing.assert_array_equal(numba_ts.edges_right, ts.edges_right)
+    np.testing.assert_array_equal(numba_ts.edges_parent, ts.edges_parent)
+    np.testing.assert_array_equal(numba_ts.edges_child, ts.edges_child)
+    assert numba_ts.edges_left.dtype == np.float64
+    assert numba_ts.edges_right.dtype == np.float64
+    assert numba_ts.edges_parent.dtype == np.int32
+    assert numba_ts.edges_child.dtype == np.int32
+    np.testing.assert_array_equal(numba_ts.nodes_time, ts.nodes_time)
+    np.testing.assert_array_equal(numba_ts.nodes_flags, ts.nodes_flags)
+    np.testing.assert_array_equal(numba_ts.nodes_population, ts.nodes_population)
+    np.testing.assert_array_equal(numba_ts.nodes_individual, ts.nodes_individual)
+    assert numba_ts.nodes_time.dtype == np.float64
+    assert numba_ts.nodes_flags.dtype == np.uint32
+    assert numba_ts.nodes_population.dtype == np.int32
+    assert numba_ts.nodes_individual.dtype == np.int32
+    np.testing.assert_array_equal(numba_ts.individuals_flags, ts.individuals_flags)
+    assert numba_ts.individuals_flags.dtype == np.uint32
+    np.testing.assert_array_equal(numba_ts.sites_position, ts.sites_position)
+    assert numba_ts.sites_position.dtype == np.float64
+    np.testing.assert_array_equal(numba_ts.mutations_site, ts.mutations_site)
+    np.testing.assert_array_equal(numba_ts.mutations_node, ts.mutations_node)
+    np.testing.assert_array_equal(numba_ts.mutations_parent, ts.mutations_parent)
+    np.testing.assert_array_equal(numba_ts.mutations_time, ts.mutations_time)
+    assert numba_ts.mutations_site.dtype == np.int32
+    assert numba_ts.mutations_node.dtype == np.int32
+    assert numba_ts.mutations_parent.dtype == np.int32
+    assert numba_ts.mutations_time.dtype == np.float64
+    np.testing.assert_array_equal(
+        numba_ts.indexes_edge_insertion_order, ts.indexes_edge_insertion_order
+    )
+    np.testing.assert_array_equal(
+        numba_ts.indexes_edge_removal_order, ts.indexes_edge_removal_order
+    )
+    assert numba_ts.indexes_edge_insertion_order.dtype == np.int32
+    assert numba_ts.indexes_edge_removal_order.dtype == np.int32
+    assert numba_ts.breakpoints.dtype == np.float64
+    np.testing.assert_array_equal(numba_ts.breakpoints, ts.breakpoints(as_array=True))

python/tskit/jit/numba.py

@@ -1,6 +1,3 @@
-from dataclasses import dataclass
-
-
 try:
     import numba
 except ImportError:
@@ -10,45 +7,92 @@
     )
 
 
-# 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)
-    del dc_cls.__dataclass_params__
-    del dc_cls.__dataclass_fields__
-    del dc_cls.__repr__
-    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)
-
-
-@jitdataclass
+tree_sequence_spec = [
+    ("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[:]),
+    ("individuals_flags", numba.uint32[:]),
+    ("nodes_time", numba.float64[:]),
+    ("nodes_flags", numba.uint32[:]),
+    ("nodes_population", numba.int32[:]),
+    ("nodes_individual", numba.int32[:]),
+    ("edges_parent", numba.int32[:]),
+    ("edges_child", numba.int32[:]),
+    ("sites_position", numba.float64[:]),
+    ("mutations_site", numba.int32[:]),
+    ("mutations_node", numba.int32[:]),
+    ("mutations_parent", numba.int32[:]),
+    ("mutations_time", numba.float64[:]),
+    ("breakpoints", numba.float64[:]),
+]
+
+
+@numba.experimental.jitclass(tree_sequence_spec)
 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 __init__(
+        self,
+        num_edges,
+        sequence_length,
+        edges_left,
+        edges_right,
+        indexes_edge_insertion_order,
+        indexes_edge_removal_order,
+        individuals_flags,
+        nodes_time,
+        nodes_flags,
+        nodes_population,
+        nodes_individual,
+        edges_parent,
+        edges_child,
+        sites_position,
+        mutations_site,
+        mutations_node,
+        mutations_parent,
+        mutations_time,
+        breakpoints,
+    ):
+        self.num_edges = num_edges
+        self.sequence_length = sequence_length
+        self.edges_left = edges_left
+        self.edges_right = edges_right
+        self.indexes_edge_insertion_order = indexes_edge_insertion_order
+        self.indexes_edge_removal_order = indexes_edge_removal_order
+        self.individuals_flags = individuals_flags
+        self.nodes_time = nodes_time
+        self.nodes_flags = nodes_flags
+        self.nodes_population = nodes_population
+        self.nodes_individual = nodes_individual
+        self.edges_parent = edges_parent
+        self.edges_child = edges_child
+        self.sites_position = sites_position
+        self.mutations_site = mutations_site
+        self.mutations_node = mutations_node
+        self.mutations_parent = mutations_parent
+        self.mutations_time = mutations_time
+        self.breakpoints = breakpoints
 
     def tree_position(self):
         return NumbaTreePosition(self, (0, 0), (0, 0), (0, 0))
 
 
-@jitdataclass
+tree_position_spec = [
+    ("ts", NumbaTreeSequence.class_type.instance_type),
+    ("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)),
+]
+
+
+@numba.experimental.jitclass(tree_position_spec)
 class NumbaTreePosition:
-    ts: NumbaTreeSequence
-    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)
+    def __init__(self, ts, interval, edges_in_index_range, edges_out_index_range):
+        self.ts = ts
+        self.interval = interval
+        self.edges_in_index_range = edges_in_index_range
+        self.edges_out_index_range = edges_out_index_range
 
     def next(self):  # noqa: A003
         M = self.ts.num_edges
@@ -87,4 +131,17 @@ def numba_tree_sequence(ts):
         edges_right=ts.edges_right,
         indexes_edge_insertion_order=ts.indexes_edge_insertion_order,
         indexes_edge_removal_order=ts.indexes_edge_removal_order,
+        individuals_flags=ts.individuals_flags,
+        nodes_time=ts.nodes_time,
+        nodes_flags=ts.nodes_flags,
+        nodes_population=ts.nodes_population,
+        nodes_individual=ts.nodes_individual,
+        edges_parent=ts.edges_parent,
+        edges_child=ts.edges_child,
+        sites_position=ts.sites_position,
+        mutations_site=ts.mutations_site,
+        mutations_node=ts.mutations_node,
+        mutations_parent=ts.mutations_parent,
+        mutations_time=ts.mutations_time,
+        breakpoints=ts.breakpoints(as_array=True),
     )