From f74ace41ac27b5cf212d64089c9970cab09fd660 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 2 Sep 2026 12:31:04 +0100 Subject: [PATCH 1/4] Fix out-of-bounds write for causal ancestral alleles The node traversal allocated an array with one entry per node, but pushes the virtual root when the causal allele is the ancestral allele. The virtual root's ID is num_nodes, so the traversal wrote one element past the end of the array. Numba compiles with bounds checking disabled, so this failed silently. Move the compiled kernels into a new tstrait.jit module, sizing the traversal buffer from the tree arrays, which already include the virtual root, and dropping it from the returned values. The stack is now a preallocated array with the causal nodes passed in as an array, which removes numba.typed from the codebase along with the special case for an empty stack. Add tests/test_jit.py, which exercises the kernels through their py_func attribute. Running the untranslated Python gives us numpy's bounds checking and lets coverage measure the kernels, which reach 100% statement and branch coverage from small hand written examples. --- CHANGELOG.md | 6 ++ tests/test_jit.py | 158 +++++++++++++++++++++++++++++++++++++++ tstrait/genetic_value.py | 78 +++++-------------- tstrait/jit.py | 76 +++++++++++++++++++ 4 files changed, 258 insertions(+), 60 deletions(-) create mode 100644 tests/test_jit.py create mode 100644 tstrait/jit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ffcd1e..ae7c563 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,12 @@ In development edge effects, and edge, node, and individual genetic values {pr}`155` - Clarified that tstrait currently uses a site-mode effect model {pr}`155` +### Fix + +- Fix an out-of-bounds write in the node traversal when the causal allele is + the ancestral allele, in which case the virtual root is a causal node + {pr}`157` + ## [0.1.2] - 2026-03-03 Maintenance release diff --git a/tests/test_jit.py b/tests/test_jit.py new file mode 100644 index 0000000..d526cbc --- /dev/null +++ b/tests/test_jit.py @@ -0,0 +1,158 @@ +""" +Unit tests for the numba compiled kernels in ``tstrait.jit``. + +The kernels are exercised through ``unjit``, which returns the untranslated +Python function. Numba compiles with bounds checking disabled, so an out of +bounds access in the compiled code fails silently; running the Python gives us +numpy's bounds checking and lets coverage measure the kernels. The examples +here are small enough to check by hand. +""" + +import numpy as np +import pytest +import tskit + +from tstrait import jit + + +def unjit(func): + """ + Return the pure Python implementation of a numba jitted function. + """ + return getattr(func, "py_func", func) + + +compute_nodes_genetic_value = unjit(jit._compute_nodes_genetic_value) +accumulate_individual_values = unjit(jit._accumulate_individual_values) + +# A balanced binary tree on 4 leaves, as returned by +# tskit.Tree.generate_balanced(4). Node 7 is the virtual root, whose only +# child is the root of the tree. +# +# 6 +# / \ +# 4 5 +# / \ / \ +# 0 1 2 3 +# +NUM_NODES = 7 +VIRTUAL_ROOT = 7 +LEFT_CHILD = np.array([-1, -1, -1, -1, 0, 2, 4, 6], dtype=np.int32) +RIGHT_SIB = np.array([1, -1, 3, -1, 5, -1, -1, -1], dtype=np.int32) + + +def mutation_array(nodes): + """ + Return the ``has_mutation`` array in which the specified nodes are marked. + """ + has_mutation = np.zeros(NUM_NODES + 1, dtype=bool) + has_mutation[nodes] = True + return has_mutation + + +def compute(causal_nodes, mutation_nodes=(), effect_size=1): + return compute_nodes_genetic_value( + left_child_array=LEFT_CHILD, + right_sib_array=RIGHT_SIB, + causal_nodes=np.array(causal_nodes, dtype=np.int32), + has_mutation=mutation_array(list(mutation_nodes)), + effect_size=effect_size, + ) + + +def test_tree_arrays_match_tskit(): + """ + Guard against the hand written arrays above drifting from tskit. + """ + tree = tskit.Tree.generate_balanced(4) + assert tree.tree_sequence.num_nodes == NUM_NODES + assert tree.virtual_root == VIRTUAL_ROOT + np.testing.assert_array_equal(tree.left_child_array, LEFT_CHILD) + np.testing.assert_array_equal(tree.right_sib_array, RIGHT_SIB) + + +class TestComputeNodesGeneticValue: + def test_no_causal_nodes(self): + np.testing.assert_array_equal(compute([]), np.zeros(NUM_NODES)) + + def test_leaf(self): + # Node 0 has no children, so the sibling walk is never entered. + np.testing.assert_array_equal(compute([0]), [1, 0, 0, 0, 0, 0, 0]) + + def test_internal_node(self): + np.testing.assert_array_equal(compute([4]), [1, 1, 0, 0, 1, 0, 0]) + + def test_root(self): + np.testing.assert_array_equal(compute([6]), np.ones(NUM_NODES)) + + def test_virtual_root(self): + # The ancestral state is the causal allele, so the virtual root is + # causal. Its value is computed but is not part of the output. + value = compute([VIRTUAL_ROOT]) + assert len(value) == NUM_NODES + np.testing.assert_array_equal(value, np.ones(NUM_NODES)) + + def test_mutation_blocks_subtree(self): + # Node 4 carries a mutation, so it and its descendants keep the value + # of whatever allele that mutation introduced, ie. zero here. + np.testing.assert_array_equal(compute([6], [4]), [0, 0, 1, 1, 0, 1, 1]) + + def test_multiple_causal_nodes(self): + np.testing.assert_array_equal(compute([1, 5], [1, 5]), [0, 1, 1, 1, 0, 1, 0]) + + def test_effect_size(self): + np.testing.assert_array_equal( + compute([4], effect_size=-2.5), [-2.5, -2.5, 0, 0, -2.5, 0, 0] + ) + + @pytest.mark.parametrize("causal_nodes", [[], [0], [4], [6], [VIRTUAL_ROOT], [1, 5]]) + def test_output_excludes_virtual_root(self, causal_nodes): + assert len(compute(causal_nodes)) == NUM_NODES + + +class TestAccumulateIndividualValues: + def test_diploid(self): + # The node to individual map of tests.data.binary_tree, in which node + # 6 belongs to no individual. + nodes_individual = np.array([1, 1, 2, 2, 0, 0, tskit.NULL], dtype=np.int32) + nodes_genetic_value = np.array([1, 2, 3, 4, 5, 6, 7], dtype=float) + value = accumulate_individual_values(nodes_genetic_value, nodes_individual, 3) + np.testing.assert_array_equal(value, [11, 3, 7]) + + def test_repeated_nodes(self): + # Triploids: three nodes contribute to each individual. + nodes_individual = np.array([0, 1, 0, 1, 0, 1], dtype=np.int32) + nodes_genetic_value = np.array([1, 2, 4, 8, 16, 32], dtype=float) + value = accumulate_individual_values(nodes_genetic_value, nodes_individual, 2) + np.testing.assert_array_equal(value, [21, 42]) + + def test_no_individuals(self): + nodes_individual = np.full(NUM_NODES, tskit.NULL, dtype=np.int32) + value = accumulate_individual_values(np.zeros(NUM_NODES), nodes_individual, 0) + np.testing.assert_array_equal(value, []) + + +def test_jitted_matches_python(): + """ + The tests above all bypass the compiler, so check that the kernels still + compile and agree with the Python they were written as. + """ + causal_nodes = np.array([VIRTUAL_ROOT], dtype=np.int32) + has_mutation = mutation_array([4]) + nodes_individual = np.array([1, 1, 2, 2, 0, 0, tskit.NULL], dtype=np.int32) + + args = (LEFT_CHILD, RIGHT_SIB, causal_nodes, has_mutation, 0.5) + nodes_genetic_value = jit._compute_nodes_genetic_value(*args) + np.testing.assert_array_equal( + nodes_genetic_value, compute_nodes_genetic_value(*args) + ) + np.testing.assert_array_equal(nodes_genetic_value, [0, 0, 0.5, 0.5, 0, 0.5, 0.5]) + + individual_value = jit._accumulate_individual_values( + nodes_genetic_value, nodes_individual, 3 + ) + np.testing.assert_array_equal( + individual_value, + accumulate_individual_values(nodes_genetic_value, nodes_individual, 3), + ) + np.testing.assert_array_equal(individual_value, [0.5, 0, 1]) diff --git a/tstrait/genetic_value.py b/tstrait/genetic_value.py index 78fca11..287dc27 100644 --- a/tstrait/genetic_value.py +++ b/tstrait/genetic_value.py @@ -1,51 +1,11 @@ -import numba import numpy as np import pandas as pd import tskit +from . import jit from .base import _check_dataframe, _check_instance, _check_non_decreasing # noreorder -@numba.njit -def _compute_nodes_genetic_value( - left_child_array, - right_sib_array, - stack, - has_mutation, - num_nodes, - effect_size, -): # pragma: no cover - """ - Compute the node genetic values for the specified set of mutations - encoded in the stack. - """ - genetic_value = np.zeros(num_nodes) - while len(stack) > 0: - parent_node_id = stack.pop() - genetic_value[parent_node_id] = effect_size - child_node_id = left_child_array[parent_node_id] - while child_node_id != tskit.NULL: - if not has_mutation[child_node_id]: - stack.append(child_node_id) - child_node_id = right_sib_array[child_node_id] - return genetic_value - - -@numba.njit -def _accumulate_individual_values( - nodes_genetic_value, nodes_individual, num_nodes, num_individuals -): # pragma: no cover - """ - Accumulate the individual genetic values by summing their node contributions. - """ - individuals_genetic_value = np.zeros(num_individuals) - for u in range(num_nodes): - ind = nodes_individual[u] - if ind != tskit.NULL: - individuals_genetic_value[ind] += nodes_genetic_value[u] - return individuals_genetic_value - - def _accumulate_edge_values(nodes_genetic_value, nodes_edge, num_nodes, num_edges): """ Accumulate the edge genetic values by summing their node contributions. @@ -86,23 +46,21 @@ def _node_genetic_values(self, tree, site, causal_allele, effect_size): for m in site.mutations: state_transitions[m.node] = m.derived_state has_mutation[m.node] = True - stack = numba.typed.List() - for node, allele in state_transitions.items(): - if allele == causal_allele: - stack.append(node) - - if len(stack) == 0: - genetic_value = np.zeros(self.ts.num_nodes) - else: - genetic_value = _compute_nodes_genetic_value( - left_child_array=tree.left_child_array, - right_sib_array=tree.right_sib_array, - stack=stack, - has_mutation=has_mutation, - num_nodes=self.ts.num_nodes, - effect_size=effect_size, - ) - return genetic_value + causal_nodes = np.array( + [ + node + for node, allele in state_transitions.items() + if allele == causal_allele + ], + dtype=np.int32, + ) + return jit._compute_nodes_genetic_value( + left_child_array=tree.left_child_array, + right_sib_array=tree.right_sib_array, + causal_nodes=causal_nodes, + has_mutation=has_mutation, + effect_size=effect_size, + ) def _run(self, level): """ @@ -137,8 +95,8 @@ def _run(self, level): effect_size=data.effect_size, ) if level == "individual": - genetic_value = _accumulate_individual_values( - genetic_value, ts.nodes_individual, ts.num_nodes, ts.num_individuals + genetic_value = jit._accumulate_individual_values( + genetic_value, ts.nodes_individual, ts.num_individuals ) elif level == "edge": genetic_value = _accumulate_edge_values( diff --git a/tstrait/jit.py b/tstrait/jit.py new file mode 100644 index 0000000..20929c3 --- /dev/null +++ b/tstrait/jit.py @@ -0,0 +1,76 @@ +""" +Numba compiled kernels. + +Two conventions apply throughout this module: + +1. Node indexed arrays follow the tskit quintuply linked tree encoding, in + which the arrays have ``num_nodes + 1`` entries and the last entry + corresponds to the virtual root. + +2. Numba compiles with bounds checking disabled by default, so an out of + bounds access here fails silently rather than raising ``IndexError``. + Keeping the compiled functions in one module lets us unit test them + through their ``py_func`` attribute in ``tests/test_jit.py``, which runs + the untranslated Python and therefore gets numpy's bounds checking and + coverage measurement for free. +""" + +import numba +import numpy as np +import tskit + + +@numba.njit +def _compute_nodes_genetic_value( + left_child_array, + right_sib_array, + causal_nodes, + has_mutation, + effect_size, +): + """ + Compute the genetic value of each node by assigning ``effect_size`` to + every node in ``causal_nodes`` and to each of their descendants that is + not separated from them by a node carrying a mutation. + + The ``left_child_array``, ``right_sib_array`` and ``has_mutation`` inputs + are indexed by node ID and have ``num_nodes + 1`` entries. The + ``causal_nodes`` are distinct node IDs, and may include the virtual root, + which happens when the ancestral state is the causal allele. The returned + array has one entry per node and does not include the virtual root. + """ + num_nodes = len(left_child_array) - 1 + # One extra entry, so that the virtual root can be a causal node. + genetic_value = np.zeros(num_nodes + 1) + # Each node is pushed at most once, since causal nodes other than the + # virtual root all carry a mutation and so are never pushed as a child. + stack = np.empty(num_nodes + 1, dtype=np.int32) + stack_top = len(causal_nodes) + stack[:stack_top] = causal_nodes + while stack_top > 0: + stack_top -= 1 + parent_node_id = stack[stack_top] + genetic_value[parent_node_id] = effect_size + child_node_id = left_child_array[parent_node_id] + while child_node_id != tskit.NULL: + if not has_mutation[child_node_id]: + stack[stack_top] = child_node_id + stack_top += 1 + child_node_id = right_sib_array[child_node_id] + return genetic_value[:num_nodes] + + +@numba.njit +def _accumulate_individual_values( + nodes_genetic_value, nodes_individual, num_individuals +): + """ + Accumulate the individual genetic values by summing their node + contributions. + """ + individuals_genetic_value = np.zeros(num_individuals) + for u in range(len(nodes_individual)): + ind = nodes_individual[u] + if ind != tskit.NULL: + individuals_genetic_value[ind] += nodes_genetic_value[u] + return individuals_genetic_value From d46e3cf122a0973e2de4533165a320b17c7ed665 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 2 Sep 2026 13:24:33 +0100 Subject: [PATCH 2/4] Test the jit kernels over a range of tree topologies Replace the single hand written tree with trees from the tskit generators, covering polytomies at and below the root, ladders, and the degenerate cases in which the virtual root has no children, one child, or many: a single node, a tree sequence with no nodes, isolated samples, and a forest containing an unreachable node. Each topology is drawn in the class docstring so that the expected values can be checked against it. Run every test against both the compiled kernel and the Python it was written as, through the jit and nojit parameters of the fixtures. This replaces the smoke test comparing the two implementations to each other, since they are now both checked against fixed expectations. Use the documented tree sequences in tests/data.py for the accumulation kernel, rather than inventing node to individual maps, and add tests composing the two kernels the way genetic_value does. --- tests/test_jit.py | 500 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 386 insertions(+), 114 deletions(-) diff --git a/tests/test_jit.py b/tests/test_jit.py index d526cbc..9726122 100644 --- a/tests/test_jit.py +++ b/tests/test_jit.py @@ -1,11 +1,14 @@ """ Unit tests for the numba compiled kernels in ``tstrait.jit``. -The kernels are exercised through ``unjit``, which returns the untranslated -Python function. Numba compiles with bounds checking disabled, so an out of -bounds access in the compiled code fails silently; running the Python gives us -numpy's bounds checking and lets coverage measure the kernels. The examples -here are small enough to check by hand. +Each test runs twice, against the compiled kernel and against the Python it was +written as, through the ``jit`` and ``nojit`` parameters of the fixtures below. +Numba compiles with bounds checking disabled, so an out of bounds access in the +compiled code fails silently; running the Python gives us numpy's bounds +checking and lets coverage measure the kernels. + +The examples are small trees from the tskit generators, drawn in the class +docstrings so that the expected values can be checked against the topology. """ import numpy as np @@ -14,145 +17,414 @@ from tstrait import jit +from .data import ( + binary_tree, + diff_ind_tree, + triploid_tree, +) # noreorder -def unjit(func): + +def kernel(func, param): + """ + Return either the compiled kernel or the Python it was written as. + """ + return func if param == "jit" else func.py_func + + +@pytest.fixture(params=["jit", "nojit"]) +def node_genetic_value(request): + """ + Return a function computing the node genetic values for a tskit tree. + + The ``causal_nodes`` are the nodes at which the causal allele appears, + which includes the virtual root when the ancestral state is causal, and the + ``mutation_nodes`` are the nodes at which a mutation occurs. """ - Return the pure Python implementation of a numba jitted function. + func = kernel(jit._compute_nodes_genetic_value, request.param) + + def f(tree, causal_nodes, mutation_nodes=(), effect_size=1): + has_mutation = np.zeros(len(tree.left_child_array), dtype=bool) + has_mutation[list(mutation_nodes)] = True + return func( + left_child_array=tree.left_child_array, + right_sib_array=tree.right_sib_array, + causal_nodes=np.array(causal_nodes, dtype=np.int32), + has_mutation=has_mutation, + effect_size=effect_size, + ) + + return f + + +@pytest.fixture(params=["jit", "nojit"]) +def individual_genetic_value(request): """ - return getattr(func, "py_func", func) + Return a function accumulating node genetic values over the individuals of + a tree sequence. + """ + func = kernel(jit._accumulate_individual_values, request.param) + + def f(ts, nodes_genetic_value): + return func( + np.asarray(nodes_genetic_value, dtype=float), + ts.nodes_individual, + ts.num_individuals, + ) + return f -compute_nodes_genetic_value = unjit(jit._compute_nodes_genetic_value) -accumulate_individual_values = unjit(jit._accumulate_individual_values) -# A balanced binary tree on 4 leaves, as returned by -# tskit.Tree.generate_balanced(4). Node 7 is the virtual root, whose only -# child is the root of the tree. -# -# 6 -# / \ -# 4 5 -# / \ / \ -# 0 1 2 3 -# -NUM_NODES = 7 -VIRTUAL_ROOT = 7 -LEFT_CHILD = np.array([-1, -1, -1, -1, 0, 2, 4, 6], dtype=np.int32) -RIGHT_SIB = np.array([1, -1, 3, -1, 5, -1, -1, -1], dtype=np.int32) +def isolated_samples_tree(n): + """ + Return a tree of n isolated sample nodes, which are therefore all roots. + """ + tables = tskit.TableCollection(sequence_length=1) + for _ in range(n): + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) + return tables.tree_sequence().first() -def mutation_array(nodes): +def multiroot_tree(): """ - Return the ``has_mutation`` array in which the specified nodes are marked. + Return generate_balanced(4) with the edges into the root removed, so that + nodes 4 and 5 are both roots and node 6 is isolated. """ - has_mutation = np.zeros(NUM_NODES + 1, dtype=bool) - has_mutation[nodes] = True - return has_mutation + tables = tskit.Tree.generate_balanced(4).tree_sequence.dump_tables() + tables.edges.replace_with(tables.edges[tables.edges.parent != 6]) + return tables.tree_sequence().first() -def compute(causal_nodes, mutation_nodes=(), effect_size=1): - return compute_nodes_genetic_value( - left_child_array=LEFT_CHILD, - right_sib_array=RIGHT_SIB, - causal_nodes=np.array(causal_nodes, dtype=np.int32), - has_mutation=mutation_array(list(mutation_nodes)), - effect_size=effect_size, - ) +def empty_tree(): + """ + Return the tree of a tree sequence that has no nodes. + """ + return tskit.TableCollection(sequence_length=1).tree_sequence().first() -def test_tree_arrays_match_tskit(): +class TestBalancedBinaryTree: """ - Guard against the hand written arrays above drifting from tskit. + tskit.Tree.generate_balanced(4), in which node 7 is the virtual root:: + + 6 + +-+-+ + 4 5 + +++ +++ + 0 1 2 3 """ - tree = tskit.Tree.generate_balanced(4) - assert tree.tree_sequence.num_nodes == NUM_NODES - assert tree.virtual_root == VIRTUAL_ROOT - np.testing.assert_array_equal(tree.left_child_array, LEFT_CHILD) - np.testing.assert_array_equal(tree.right_sib_array, RIGHT_SIB) + @pytest.fixture + def tree(self): + return tskit.Tree.generate_balanced(4) -class TestComputeNodesGeneticValue: - def test_no_causal_nodes(self): - np.testing.assert_array_equal(compute([]), np.zeros(NUM_NODES)) + def test_no_causal_nodes(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, []), [0, 0, 0, 0, 0, 0, 0] + ) - def test_leaf(self): + def test_leaf(self, tree, node_genetic_value): # Node 0 has no children, so the sibling walk is never entered. - np.testing.assert_array_equal(compute([0]), [1, 0, 0, 0, 0, 0, 0]) + np.testing.assert_array_equal( + node_genetic_value(tree, [0]), [1, 0, 0, 0, 0, 0, 0] + ) - def test_internal_node(self): - np.testing.assert_array_equal(compute([4]), [1, 1, 0, 0, 1, 0, 0]) + def test_internal_node(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [4]), [1, 1, 0, 0, 1, 0, 0] + ) - def test_root(self): - np.testing.assert_array_equal(compute([6]), np.ones(NUM_NODES)) + def test_root(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [6]), [1, 1, 1, 1, 1, 1, 1] + ) - def test_virtual_root(self): + def test_virtual_root(self, tree, node_genetic_value): # The ancestral state is the causal allele, so the virtual root is - # causal. Its value is computed but is not part of the output. - value = compute([VIRTUAL_ROOT]) - assert len(value) == NUM_NODES - np.testing.assert_array_equal(value, np.ones(NUM_NODES)) + # causal. Its own value is computed but is not part of the output. + value = node_genetic_value(tree, [7]) + assert len(value) == 7 + np.testing.assert_array_equal(value, [1, 1, 1, 1, 1, 1, 1]) + + def test_mutation_on_internal_node(self, tree, node_genetic_value): + # Node 4 and its children carry whatever allele the mutation on node 4 + # introduced, so the traversal does not descend into them. + np.testing.assert_array_equal( + node_genetic_value(tree, [6], [4]), [0, 0, 1, 1, 0, 1, 1] + ) + + def test_mutation_on_leaf(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [6], [0]), [0, 1, 1, 1, 1, 1, 1] + ) + + def test_two_causal_nodes(self, tree, node_genetic_value): + # Both mutations are back to the causal allele, so both nodes are + # causal and each starts its own traversal. + np.testing.assert_array_equal( + node_genetic_value(tree, [1, 5], [1, 5]), [0, 1, 1, 1, 0, 1, 0] + ) + + def test_effect_size(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [4], effect_size=-2.5), + [-2.5, -2.5, 0, 0, -2.5, 0, 0], + ) + + +class TestStarTree: + """ + tskit.Tree.generate_star(5), in which node 6 is the virtual root. The + polytomy at the root makes the sibling walk iterate five times:: + + 5 + +-+-+-+-+ + 0 1 2 3 4 + """ + + @pytest.fixture + def tree(self): + return tskit.Tree.generate_star(5) + + def test_root(self, tree, node_genetic_value): + np.testing.assert_array_equal(node_genetic_value(tree, [5]), [1, 1, 1, 1, 1, 1]) + + def test_virtual_root(self, tree, node_genetic_value): + np.testing.assert_array_equal(node_genetic_value(tree, [6]), [1, 1, 1, 1, 1, 1]) + + def test_mutation_on_root(self, tree, node_genetic_value): + # The root is the virtual root's only child, so a mutation there stops + # the traversal immediately. + np.testing.assert_array_equal( + node_genetic_value(tree, [6], [5]), [0, 0, 0, 0, 0, 0] + ) + + def test_mutations_along_sibling_chain(self, tree, node_genetic_value): + # Mutations at both ends of the chain of children and in the middle. + np.testing.assert_array_equal( + node_genetic_value(tree, [5], [0, 2, 4]), [0, 1, 0, 1, 0, 1] + ) + + +class TestNonBinaryTree: + """ + tskit.Tree.generate_balanced(6, arity=3), in which node 10 is the virtual + root. The polytomy is below the root:: + + 9 + +---+---+ + 6 7 8 + +++ +++ +++ + 0 1 2 3 4 5 + """ - def test_mutation_blocks_subtree(self): - # Node 4 carries a mutation, so it and its descendants keep the value - # of whatever allele that mutation introduced, ie. zero here. - np.testing.assert_array_equal(compute([6], [4]), [0, 0, 1, 1, 0, 1, 1]) + @pytest.fixture + def tree(self): + return tskit.Tree.generate_balanced(6, arity=3) - def test_multiple_causal_nodes(self): - np.testing.assert_array_equal(compute([1, 5], [1, 5]), [0, 1, 1, 1, 0, 1, 0]) + def test_virtual_root(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [10]), [1, 1, 1, 1, 1, 1, 1, 1, 1, 1] + ) + + def test_middle_subtree(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [7]), [0, 0, 1, 1, 0, 0, 0, 1, 0, 0] + ) + + def test_mutation_on_middle_subtree(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [10], [7]), [1, 1, 0, 0, 1, 1, 1, 0, 1, 1] + ) + + def test_mutations_on_outer_subtrees(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [10], [6, 8]), [0, 0, 1, 1, 0, 0, 0, 1, 0, 1] + ) + + +class TestCombTree: + """ + tskit.Tree.generate_comb(5), in which node 9 is the virtual root. A ladder + is the deepest traversal for a given number of leaves:: + + 8 + +-+-+ + | 7 + | +-+-+ + | | 6 + | | +-++ + | | | 5 + | | | +++ + 0 1 2 3 4 + """ + + @pytest.fixture + def tree(self): + return tskit.Tree.generate_comb(5) + + def test_root(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [8]), [1, 1, 1, 1, 1, 1, 1, 1, 1] + ) + + def test_virtual_root(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [9]), [1, 1, 1, 1, 1, 1, 1, 1, 1] + ) + + def test_second_rung(self, tree, node_genetic_value): + # Everything below node 7, which is all of the tree except leaf 0 and + # the root. + np.testing.assert_array_equal( + node_genetic_value(tree, [7]), [0, 1, 1, 1, 1, 1, 1, 1, 0] + ) + + def test_mutation_near_root(self, tree, node_genetic_value): + np.testing.assert_array_equal( + node_genetic_value(tree, [8], [7]), [1, 0, 0, 0, 0, 0, 0, 0, 1] + ) - def test_effect_size(self): + def test_mutation_near_leaves(self, tree, node_genetic_value): np.testing.assert_array_equal( - compute([4], effect_size=-2.5), [-2.5, -2.5, 0, 0, -2.5, 0, 0] + node_genetic_value(tree, [8], [5]), [1, 1, 1, 0, 0, 0, 1, 1, 1] ) - @pytest.mark.parametrize("causal_nodes", [[], [0], [4], [6], [VIRTUAL_ROOT], [1, 5]]) - def test_output_excludes_virtual_root(self, causal_nodes): - assert len(compute(causal_nodes)) == NUM_NODES + +class TestDegenerateTrees: + """ + Trees that do not have a single root above a set of internal nodes. + """ + + def test_single_node(self, node_genetic_value): + # generate_balanced(1) is the single node 0, with virtual root 1. + tree = tskit.Tree.generate_balanced(1) + np.testing.assert_array_equal(node_genetic_value(tree, [1]), [1]) + np.testing.assert_array_equal(node_genetic_value(tree, [0]), [1]) + + def test_no_nodes(self, node_genetic_value): + # The virtual root is node 0 and there is nothing to return. + tree = empty_tree() + assert tree.virtual_root == 0 + np.testing.assert_array_equal(node_genetic_value(tree, [0]), []) + + def test_isolated_samples(self, node_genetic_value): + # Nodes 0, 1 and 2 are all roots, so the virtual root has three + # children: + # + # 0 1 2 + # + tree = isolated_samples_tree(3) + np.testing.assert_array_equal(node_genetic_value(tree, [3]), [1, 1, 1]) + + def test_isolated_samples_with_mutation(self, node_genetic_value): + tree = isolated_samples_tree(3) + np.testing.assert_array_equal(node_genetic_value(tree, [3], [1]), [1, 0, 1]) + + def test_multiple_roots(self, node_genetic_value): + # Nodes 4 and 5 are both roots and node 6 is isolated, so it is not + # reached from the virtual root: + # + # 4 5 + # +++ +++ + # 0 1 2 3 + # + tree = multiroot_tree() + assert tree.roots == [4, 5] + np.testing.assert_array_equal( + node_genetic_value(tree, [7]), [1, 1, 1, 1, 1, 1, 0] + ) + + def test_multiple_roots_one_blocked(self, node_genetic_value): + tree = multiroot_tree() + np.testing.assert_array_equal( + node_genetic_value(tree, [7], [4]), [0, 0, 1, 1, 0, 1, 0] + ) class TestAccumulateIndividualValues: - def test_diploid(self): - # The node to individual map of tests.data.binary_tree, in which node - # 6 belongs to no individual. - nodes_individual = np.array([1, 1, 2, 2, 0, 0, tskit.NULL], dtype=np.int32) - nodes_genetic_value = np.array([1, 2, 3, 4, 5, 6, 7], dtype=float) - value = accumulate_individual_values(nodes_genetic_value, nodes_individual, 3) - np.testing.assert_array_equal(value, [11, 3, 7]) - - def test_repeated_nodes(self): - # Triploids: three nodes contribute to each individual. - nodes_individual = np.array([0, 1, 0, 1, 0, 1], dtype=np.int32) - nodes_genetic_value = np.array([1, 2, 4, 8, 16, 32], dtype=float) - value = accumulate_individual_values(nodes_genetic_value, nodes_individual, 2) - np.testing.assert_array_equal(value, [21, 42]) - - def test_no_individuals(self): - nodes_individual = np.full(NUM_NODES, tskit.NULL, dtype=np.int32) - value = accumulate_individual_values(np.zeros(NUM_NODES), nodes_individual, 0) - np.testing.assert_array_equal(value, []) - - -def test_jitted_matches_python(): - """ - The tests above all bypass the compiler, so check that the kernels still - compile and agree with the Python they were written as. - """ - causal_nodes = np.array([VIRTUAL_ROOT], dtype=np.int32) - has_mutation = mutation_array([4]) - nodes_individual = np.array([1, 1, 2, 2, 0, 0, tskit.NULL], dtype=np.int32) - - args = (LEFT_CHILD, RIGHT_SIB, causal_nodes, has_mutation, 0.5) - nodes_genetic_value = jit._compute_nodes_genetic_value(*args) - np.testing.assert_array_equal( - nodes_genetic_value, compute_nodes_genetic_value(*args) - ) - np.testing.assert_array_equal(nodes_genetic_value, [0, 0, 0.5, 0.5, 0, 0.5, 0.5]) - - individual_value = jit._accumulate_individual_values( - nodes_genetic_value, nodes_individual, 3 - ) - np.testing.assert_array_equal( - individual_value, - accumulate_individual_values(nodes_genetic_value, nodes_individual, 3), - ) - np.testing.assert_array_equal(individual_value, [0.5, 0, 1]) + """ + Sum the node genetic values over the nodes of each individual, using the + tree sequences in tests/data.py. The node values are powers of two, so the + individual values identify the nodes that contributed to them. + """ + + def test_binary_tree(self, individual_genetic_value): + # Individual 0 is nodes 4 and 5, individual 1 is nodes 0 and 1, and + # individual 2 is nodes 2 and 3. Node 6 has no individual. + ts = binary_tree() + np.testing.assert_array_equal(ts.nodes_individual, [1, 1, 2, 2, 0, 0, -1]) + np.testing.assert_array_equal( + individual_genetic_value(ts, [1, 2, 4, 8, 16, 32, 64]), [48, 3, 12] + ) + + def test_diff_ind_tree(self, individual_genetic_value): + # The same tree, with the leaves paired up the other way around. + ts = diff_ind_tree() + np.testing.assert_array_equal(ts.nodes_individual, [1, 2, 1, 2, 0, 0, -1]) + np.testing.assert_array_equal( + individual_genetic_value(ts, [1, 2, 4, 8, 16, 32, 64]), [48, 5, 10] + ) + + def test_triploid_tree(self, individual_genetic_value): + # Two triploids: nodes 0, 2 and 4, and nodes 1, 3 and 5. + ts = triploid_tree() + np.testing.assert_array_equal(ts.nodes_individual, [0, 1, 0, 1, 0, 1, -1, -1]) + np.testing.assert_array_equal( + individual_genetic_value(ts, [1, 2, 4, 8, 16, 32, 64, 128]), [21, 42] + ) + + def test_no_individuals(self, individual_genetic_value): + ts = tskit.Tree.generate_balanced(4).tree_sequence + assert ts.num_individuals == 0 + np.testing.assert_array_equal( + individual_genetic_value(ts, [1, 2, 4, 8, 16, 32, 64]), [] + ) + + +class TestNodeAndIndividualValues: + """ + The two kernels composed, as they are used in tstrait.genetic_value, so + that the individual values can be traced back to the tree topology. The + tree of tests.data.binary_tree is:: + + 6 + +-+-+ + 4 5 + +++ +++ + 0 1 2 3 + + with individual 0 being nodes 4 and 5, individual 1 nodes 0 and 1, and + individual 2 nodes 2 and 3. + """ + + def test_internal_node(self, node_genetic_value, individual_genetic_value): + ts = binary_tree() + value = node_genetic_value(ts.first(), [4]) + np.testing.assert_array_equal(value, [1, 1, 0, 0, 1, 0, 0]) + # Individual 0 has one copy through node 4, individual 1 has two + # copies through nodes 0 and 1, and individual 2 has none. + np.testing.assert_array_equal(individual_genetic_value(ts, value), [1, 2, 0]) + + def test_virtual_root(self, node_genetic_value, individual_genetic_value): + # The causal allele is the ancestral state, so every node carries it + # and every diploid has two copies. + ts = binary_tree() + value = node_genetic_value(ts.first(), [7]) + np.testing.assert_array_equal(value, [1, 1, 1, 1, 1, 1, 1]) + np.testing.assert_array_equal(individual_genetic_value(ts, value), [2, 2, 2]) + + def test_triploid(self, node_genetic_value, individual_genetic_value): + # tests.data.triploid_tree, in which node 6 is the parent of nodes 3, + # 4 and 5 and node 7 is the root: + # + # 7 + # +-+-+---+ + # | | | 6 + # | | | +-+-+ + # 0 1 2 3 4 5 + # + ts = triploid_tree() + value = node_genetic_value(ts.first(), [6]) + np.testing.assert_array_equal(value, [0, 0, 0, 1, 1, 1, 1, 0]) + # Individual 0 is nodes 0, 2 and 4, and individual 1 is nodes 1, 3 + # and 5. + np.testing.assert_array_equal(individual_genetic_value(ts, value), [1, 2]) From 63c33d0a8f56d6d0fd8d549d196046b82c58447c Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 2 Sep 2026 13:26:20 +0100 Subject: [PATCH 3/4] Fixup refs in the CHANGELOG --- CHANGELOG.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ae7c563..29c6b93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,20 +8,20 @@ In development - Extended `genetic_value` with `level="individual"` (default), `level="node"`, and `level="edge"` to return genetic values for the corresponding entities. - {pr}`155` -- Added `edge_effect` to compute introduced effects on edges {pr}`155` + {pr}`189` +- Added `edge_effect` to compute introduced effects on edges {pr}`189` ### Documentation - Added a worked example relating causal-allele effects, - edge effects, and edge, node, and individual genetic values {pr}`155` -- Clarified that tstrait currently uses a site-mode effect model {pr}`155` + edge effects, and edge, node, and individual genetic values {pr}`189` +- Clarified that tstrait currently uses a site-mode effect model {pr}`189` ### Fix - Fix an out-of-bounds write in the node traversal when the causal allele is the ancestral allele, in which case the virtual root is a causal node - {pr}`157` + {pr}`192`, {issue}`191`. ## [0.1.2] - 2026-03-03 From 2ef7834a8a1b02dc01a9064885d742a95704d93a Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 2 Sep 2026 13:28:04 +0100 Subject: [PATCH 4/4] Bump docs build version to fix error --- .github/workflows/docs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 9cdc583..7cde629 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -13,4 +13,4 @@ permissions: jobs: Docs: - uses: tskit-dev/.github/.github/workflows/docs.yml@v19 + uses: tskit-dev/.github/.github/workflows/docs.yml@v21