Perf tweaks to diversity (~20% drop in time)

Author: benjeffery

docs/numba.md

@@ -183,37 +183,44 @@ JIT compilation - using simple loops and fixed-size arrays with minimal object a
             # Process the outgoing edges
             for j in range(tree_index.out_range.start, tree_index.out_range.stop):
                 h = tree_index.out_range.order[j]
-                u = edge_child[h]
+                child = edge_child[h]
+                child_parent = edge_parent[h]
 
-                running_sum -= branch_length[u] * summary[u]
-                parent[u] = -1
-                branch_length[u] = 0.0
+                running_sum -= branch_length[child] * summary[child]
+                parent[child] = -1
+                branch_length[child] = 0.0
 
-                u = edge_parent[h]
+                u = child_parent
+                parent_u = parent[u]
                 while u != -1:
                     running_sum -= branch_length[u] * summary[u]
-                    state[u] -= state[edge_child[h]]
+                    state[u] -= state[child]
                     summary[u] = state[u] * (n - state[u]) * two_over_denom
                     running_sum += branch_length[u] * summary[u]
-                    u = parent[u]
+                    u = parent_u
+                    if u != -1:
+                        parent_u = parent[u]
 
             # Process the incoming edges
             for j in range(tree_index.in_range.start, tree_index.in_range.stop):
                 h = tree_index.in_range.order[j]
-                u = edge_child[h]
-                v = edge_parent[h]
+                child = edge_child[h]
+                child_parent = edge_parent[h]
 
-                parent[u] = v
-                branch_length[u] = node_times[v] - node_times[u]
-                running_sum += branch_length[u] * summary[u]
+                parent[child] = child_parent
+                branch_length[child] = node_times[child_parent] - node_times[child]
+                running_sum += branch_length[child] * summary[child]
 
-                u = v
+                u = child_parent
+                parent_u = parent[u]
                 while u != -1:
                     running_sum -= branch_length[u] * summary[u]
-                    state[u] += state[edge_child[h]]
+                    state[u] += state[child]
                     summary[u] = state[u] * (n - state[u]) * two_over_denom
                     running_sum += branch_length[u] * summary[u]
-                    u = parent[u]
+                    u = parent_u
+                    if u != -1:
+                        parent_u = parent[u]
 
             result += running_sum * (
                 tree_index.interval[1] - tree_index.interval[0]

python/tests/test_jit.py

@@ -207,37 +207,44 @@ def diversity(numba_ts):
             # Process the outgoing edges
             for j in range(tree_index.out_range.start, tree_index.out_range.stop):
                 h = tree_index.out_range.order[j]
-                u = edge_child[h]
+                child = edge_child[h]
+                child_parent = edge_parent[h]
 
-                running_sum -= branch_length[u] * summary[u]
-                parent[u] = -1
-                branch_length[u] = 0.0
+                running_sum -= branch_length[child] * summary[child]
+                parent[child] = -1
+                branch_length[child] = 0.0
 
-                u = edge_parent[h]
+                u = child_parent
+                parent_u = parent[u]
                 while u != -1:
                     running_sum -= branch_length[u] * summary[u]
-                    state[u] -= state[edge_child[h]]
+                    state[u] -= state[child]
                     summary[u] = state[u] * (n - state[u]) * two_over_denom
                     running_sum += branch_length[u] * summary[u]
-                    u = parent[u]
+                    u = parent_u
+                    if u != -1:
+                        parent_u = parent[u]
 
             # Process the incoming edges
             for j in range(tree_index.in_range.start, tree_index.in_range.stop):
                 h = tree_index.in_range.order[j]
-                u = edge_child[h]
-                v = edge_parent[h]
+                child = edge_child[h]
+                child_parent = edge_parent[h]
 
-                parent[u] = v
-                branch_length[u] = node_times[v] - node_times[u]
-                running_sum += branch_length[u] * summary[u]
+                parent[child] = child_parent
+                branch_length[child] = node_times[child_parent] - node_times[child]
+                running_sum += branch_length[child] * summary[child]
 
-                u = v
+                u = child_parent
+                parent_u = parent[u]
                 while u != -1:
                     running_sum -= branch_length[u] * summary[u]
-                    state[u] += state[edge_child[h]]
+                    state[u] += state[child]
                     summary[u] = state[u] * (n - state[u]) * two_over_denom
                     running_sum += branch_length[u] * summary[u]
-                    u = parent[u]
+                    u = parent_u
+                    if u != -1:
+                        parent_u = parent[u]
 
             result += running_sum * (tree_index.interval[1] - tree_index.interval[0])
 

python/tskit/jit/numba.py

@@ -17,7 +17,10 @@
 
 FORWARD = 1  #: Direction constant for forward tree traversal
 REVERSE = -1  #: Direction constant for reverse tree traversal
+
+# Retrieve these here to avoid lookups in tight loops
 NODE_IS_SAMPLE = tskit.NODE_IS_SAMPLE
+NULL = tskit.NULL
 
 edge_range_spec = [
     ("start", numba.int32),
@@ -98,7 +101,7 @@ class NumbaTreeIndex:
     def __init__(self, ts):
         self.ts = ts
         self.index = -1
-        self.direction = tskit.NULL
+        self.direction = NULL
         self.interval = (0, 0)
         self.in_range = NumbaEdgeRange(0, 0, np.zeros(0, dtype=np.int32))
         self.out_range = NumbaEdgeRange(0, 0, np.zeros(0, dtype=np.int32))