diff --git a/c/tskit/core.c b/c/tskit/core.c index 5a8ed6d9ac..027e4f147f 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -330,6 +330,11 @@ tsk_strerror_internal(int err) "values for any single site. " "(TSK_ERR_MUTATION_TIME_HAS_BOTH_KNOWN_AND_UNKNOWN)"; break; + case TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME: + ret = "Some mutation times are marked 'unknown' for a method that requires " + "no unknown times. (Use compute_mutation_times to add times?) " + "(TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME)"; + break; /* Migration errors */ case TSK_ERR_UNSORTED_MIGRATIONS: diff --git a/c/tskit/core.h b/c/tskit/core.h index 45a33dd8b7..5769062e24 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -500,6 +500,11 @@ the edge on which it occurs, and wasn't TSK_UNKNOWN_TIME. A single site had a mixture of known mutation times and TSK_UNKNOWN_TIME */ #define TSK_ERR_MUTATION_TIME_HAS_BOTH_KNOWN_AND_UNKNOWN -509 +/** +Some mutations have TSK_UNKNOWN_TIME in an algorithm where that's +disallowed (use compute_mutation_times?). +*/ +#define TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME -510 /** @} */ /** diff --git a/c/tskit/trees.c b/c/tskit/trees.c index bb7002386f..fc2cef9c90 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -6956,7 +6956,6 @@ tsk_treeseq_extend_edges_iter( out_parent[edges->child[e]] = edges->parent[e]; } - for (tj = tree_pos.out.start; tj != tree_pos.out.stop; tj += direction) { e = tree_pos.out.order[tj]; if (out_parent[edges->child[e]] == TSK_NULL) { @@ -7014,7 +7013,8 @@ tsk_treeseq_extend_edges_iter( } near_side[e_in] = there; while (c != p) { - for (ex_out = edges_out_head; ex_out != NULL; ex_out = ex_out->next) { + for (ex_out = edges_out_head; ex_out != NULL; + ex_out = ex_out->next) { e_out = ex_out->edge; if (edges->child[e_out] == c) { break; @@ -7053,9 +7053,9 @@ tsk_treeseq_extend_edges_iter( return ret; } - static int -tsk_treeseq_slide_mutation_nodes_up(const tsk_treeseq_t *self, tsk_mutation_table_t *mutations) +tsk_treeseq_slide_mutation_nodes_up( + const tsk_treeseq_t *self, tsk_mutation_table_t *mutations) { int ret = 0; bool valid; @@ -7096,9 +7096,13 @@ tsk_treeseq_slide_mutation_nodes_up(const tsk_treeseq_t *self, tsk_mutation_tabl e = tree_pos.in.order[tj]; parent[edges->child[e]] = edges->parent[e]; } - while (next_mut < (tsk_id_t) mutations->num_rows && - sites_position[mutations->site[next_mut]] < right) { + while (next_mut < (tsk_id_t) mutations->num_rows + && sites_position[mutations->site[next_mut]] < right) { t = mutations->time[next_mut]; + if (tsk_is_unknown_time(t)) { + ret = TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME; + goto out; + } c = mutations->node[next_mut]; p = parent[c]; while (p != TSK_NULL && nodes_time[p] <= t) { @@ -7119,7 +7123,6 @@ tsk_treeseq_slide_mutation_nodes_up(const tsk_treeseq_t *self, tsk_mutation_tabl return ret; } - int TSK_WARN_UNUSED tsk_treeseq_extend_edges(const tsk_treeseq_t *self, int max_iter, tsk_flags_t TSK_UNUSED(options), tsk_treeseq_t *output) @@ -7177,7 +7180,8 @@ tsk_treeseq_extend_edges(const tsk_treeseq_t *self, int max_iter, } /* Remap mutation nodes */ - ret = tsk_mutation_table_copy(&self->tables->mutations, &tables.mutations, TSK_NO_INIT); + ret = tsk_mutation_table_copy( + &self->tables->mutations, &tables.mutations, TSK_NO_INIT); if (ret != 0) { goto out; } diff --git a/python/tests/test_extend_edges.py b/python/tests/test_extend_edges.py index 94ffae0dc4..f130bc19dd 100644 --- a/python/tests/test_extend_edges.py +++ b/python/tests/test_extend_edges.py @@ -2,6 +2,7 @@ import numpy as np import pytest +import _tskit import tests.test_wright_fisher as wf import tskit from tests import tsutil @@ -53,7 +54,9 @@ def _slide_mutation_nodes_up(ts, mutations): while valid: left, right = tree_pos.interval - for j in range(tree_pos.out_range.start, tree_pos.out_range.stop, tskit.FORWARD): + for j in range( + tree_pos.out_range.start, tree_pos.out_range.stop, tskit.FORWARD + ): e = tree_pos.out_range.order[j] parent[edges.child[e]] = -1 @@ -61,8 +64,10 @@ def _slide_mutation_nodes_up(ts, mutations): e = tree_pos.in_range.order[j] parent[edges.child[e]] = edges.parent[e] - while (next_mut < mutations.num_rows and - ts.sites_position[mutations.site[next_mut]] < right): + while ( + next_mut < mutations.num_rows + and ts.sites_position[mutations.site[next_mut]] < right + ): t = mutations.time[next_mut] c = new_nodes[next_mut] p = parent[c] @@ -84,12 +89,13 @@ def _slide_mutation_nodes_up(ts, mutations): valid = tree_pos.next() # in C the node column can be edited in place - new_mutations = tskit.MutationTable() + new_mutations = mutations.copy() + new_mutations.clear() for mut, n in zip(mutations, new_nodes): new_mutations.append(mut.replace(node=n)) return new_mutations - + def _extend(ts, forwards=True): # `degree` will record the degree of each node in the tree we'd get if @@ -132,7 +138,7 @@ def _extend(ts, forwards=True): # if an edge from p->c has been extended, entirely replacing # another edge from p'->c, then both edges may be in edges_out, # and we only want to include the *first* one. - for e, x in edges_out: + for e, _ in edges_out: out_parent[edges.child[e]] = -1 tmp = [] for e, x in edges_out: @@ -168,7 +174,7 @@ def _extend(ts, forwards=True): # validate out_parent array for c, p in enumerate(out_parent): foundit = False - for e, x in edges_out: + for e, _ in edges_out: if edges.child[e] == c: assert edges.parent[e] == p foundit = True @@ -359,6 +365,18 @@ def test_runs(self): ts = msprime.simulate(5, mutation_rate=1.0, random_seed=126) self.verify_extend_edges(ts) + def test_unknown_times(self): + ts = msprime.simulate(5, mutation_rate=1.0, random_seed=126) + tables = ts.dump_tables() + tables.mutations.clear() + for mut in ts.mutations(): + tables.mutations.append(mut.replace(time=tskit.UNKNOWN_TIME)) + ts = tables.tree_sequence() + with pytest.raises( + _tskit.LibraryError, match="TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME" + ): + _ = ts.extend_edges() + def test_max_iter(self): ts = msprime.simulate(5, random_seed=126) with pytest.raises(ValueError, match="max_iter"): @@ -381,12 +399,12 @@ def test_simple_ex(self): # 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 # # Result: - # - # 6 6 6 6 - # +-+-+ +-+-+ +-+-+ +-+-+ - # 7 8 7 8 7 8 7 8 - # | | ++-+ | | +-++ | | - # 4 5 4 | 5 4 | 5 4 5 + # + # 6 6 6 6 + # +-+-+ +-+-+ +-+-+ +-+-+ + # 7 8 7 8 7 8 7 8 + # | | ++-+ | | +-++ | | + # 4 5 4 | 5 4 | 5 4 5 # +++ +++ +++ | | | | +++ +++ +++ # 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 @@ -436,7 +454,7 @@ def test_simple_ex(self): for t in ets.trees(): print(".....") print(t.interval) - print(t.draw(format='ascii')) + print(t.draw(format="ascii")) assert ts.num_edges == 18 assert ets.num_edges == 12 for t in ets.trees(): @@ -453,22 +471,14 @@ def test_wright_fisher(self): tables = wf.wf_sim(N=5, ngens=20, num_loci=100, deep_history=False, seed=3) tables.sort() tables.simplify() - ts = msprime.sim_mutations( - tables.tree_sequence(), - rate=0.01, - random_seed=888 - ) + ts = msprime.sim_mutations(tables.tree_sequence(), rate=0.01, random_seed=888) self.verify_extend_edges(ts, max_iter=1) self.verify_extend_edges(ts) def test_wright_fisher_unsimplified(self): tables = wf.wf_sim(N=6, ngens=22, num_loci=100, deep_history=False, seed=4) tables.sort() - ts = msprime.sim_mutations( - tables.tree_sequence(), - rate=0.01, - random_seed=888 - ) + ts = msprime.sim_mutations(tables.tree_sequence(), rate=0.01, random_seed=888) self.verify_extend_edges(ts, max_iter=1) self.verify_extend_edges(ts) @@ -476,11 +486,7 @@ def test_wright_fisher_with_history(self): tables = wf.wf_sim(N=8, ngens=15, num_loci=100, deep_history=True, seed=5) tables.sort() tables.simplify() - ts = msprime.sim_mutations( - tables.tree_sequence(), - rate=0.01, - random_seed=888 - ) + ts = msprime.sim_mutations(tables.tree_sequence(), rate=0.01, random_seed=888) self.verify_extend_edges(ts, max_iter=1) self.verify_extend_edges(ts) @@ -502,8 +508,12 @@ class TestExamples: """ def check(self, ts): - lib_ts = ts.extend_edges() + if np.any(tskit.is_unknown_time(ts.mutations_time)): + tables = ts.dump_tables() + tables.compute_mutation_times() + ts = tables.tree_sequence() py_ts = extend_edges(ts) + lib_ts = ts.extend_edges() lib_ts.tables.assert_equals(py_ts.tables) @pytest.mark.parametrize("ts", get_example_tree_sequences())