From 96be30d99ce083cf38a00f674c04f462d8d41995 Mon Sep 17 00:00:00 2001 From: peter Date: Fri, 28 Aug 2026 16:26:48 -0700 Subject: [PATCH] census node per lineage, not segment: closes #2540 and check for existing nodes --- CHANGELOG.md | 9 +++++++ lib/msprime.c | 41 ++++++++++++++++++++++---------- lib/tests/test_ancestry.c | 50 ++++++++++++++++++++++++++++++++++----- tests/test_demography.py | 44 +++++++++++++++++++++++++--------- 4 files changed, 115 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0244ac3aa..da87291d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,15 @@ In development. +**Bug fixes** + +- Previously, census events added a new census node to every distinct ancestral + segment, even when several of these segments were within the same lineage. + With this change, census events now assign a new census node to each + distinct lineage. This will in general reduce the number of census nodes, + and result in census nodes with more than one distinct segment (which + previously did not happen). ({pr}`2541`, {user}`petrelharp`) + ## [1.4.3] - 2026-08-31 Feature release, to prepare for the upcoming SLiM v6 release. diff --git a/lib/msprime.c b/lib/msprime.c index 32eee3b6d..5df9a98f0 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -7421,6 +7421,7 @@ msp_census_event(msp_t *self, demographic_event_t *event) lineage_t *lin; tsk_id_t i, j; tsk_id_t u; + tsk_node_table_t *nodes = &self->tables->nodes; for (i = 0; i < (int) self->num_populations; i++) { for (j = 0; j < (int) self->num_labels; j++) { @@ -7431,23 +7432,39 @@ msp_census_event(msp_t *self, demographic_event_t *event) while (node != NULL) { lin = (lineage_t *) node->item; - seg = lin->head; - while (seg != NULL) { - // Add an edge to the edge table. - ret = tsk_node_table_add_row(&self->tables->nodes, - MSP_NODE_IS_CEN_EVENT, event->time, i, TSK_NULL, NULL, 0); + + // Find whether this lineage already has a node at this time: + // if so it should only have one! (However, this is not + // sufficient for applying this to DTWF at integer times.) + u = TSK_NULL; + for (seg = lin->head; seg != NULL; seg = seg->next) { + if (nodes->time[seg->value] == event->time) { + u = seg->value; + break; + } + } + if (u == TSK_NULL) { + /* Add a node for this ancestor */ + ret = tsk_node_table_add_row( + nodes, MSP_NODE_IS_CEN_EVENT, event->time, i, TSK_NULL, NULL, 0); if (ret < 0) { goto out; } u = (tsk_id_t) ret; - // Add an edge joining the segment to the new node. - ret = msp_store_edge(self, seg->left, seg->right, u, seg->value); - if (ret != 0) { - goto out; + } else { + nodes->flags[u] |= MSP_NODE_IS_CEN_EVENT; + } + + for (seg = lin->head; seg != NULL; seg = seg->next) { + if (u != seg->value) { + // Add an edge joining the segment to the new node. + ret = msp_store_edge(self, seg->left, seg->right, u, seg->value); + if (ret != 0) { + goto out; + } + // Modify segment node id. + seg->value = u; } - // Modify segment node id. - seg->value = u; - seg = seg->next; } node = node->next; } diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index 2ecb39d26..37d0b8d65 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -2921,16 +2921,53 @@ test_census_event(void) uint32_t n = 10; msp_t msp; gsl_rng *rng = safe_rng_alloc(); - tsk_table_collection_t tables; + tsk_table_collection_t tables1, tables2; int num_census_nodes = 0; int i; + double census_time; - ret = build_sim(&msp, &tables, rng, 2, 1, NULL, n); + gsl_rng_set(rng, 123); + ret = build_sim(&msp, &tables1, rng, 2, 1, NULL, n); CU_ASSERT_EQUAL(ret, 0); ret = msp_set_recombination_rate(&msp, 1); /* Add a census event in at 0.5 generations. */ - ret = msp_add_census_event(&msp, 0.5); + census_time = 0.5; + ret = msp_add_census_event(&msp, census_time); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = msp_run(&msp, DBL_MAX, UINT32_MAX); + CU_ASSERT_EQUAL(ret, 0); + msp_verify(&msp, 0); + msp_print_state(&msp, _devnull); + + /* Check there is more than 1 node at the census time. */ + for (i = 0; i < tables1.nodes.num_rows; i++) { + if (tables1.nodes.time[i] == census_time) { + num_census_nodes++; + } + } + CU_ASSERT_TRUE(num_census_nodes > 1); + + ret = msp_free(&msp); + + /* Add a census event at a time where there's already a node */ + for (i = n; i < tables1.nodes.num_rows; i++) { + if (tables1.nodes.time[i] != census_time) { + census_time = tables1.nodes.time[i]; + break; + } + } + CU_ASSERT_TRUE(census_time != 0.5); + + gsl_rng_set(rng, 123); + ret = build_sim(&msp, &tables2, rng, 2, 1, NULL, n); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_set_recombination_rate(&msp, 1); + + ret = msp_add_census_event(&msp, census_time); CU_ASSERT_EQUAL(ret, 0); ret = msp_initialise(&msp); CU_ASSERT_EQUAL(ret, 0); @@ -2941,8 +2978,8 @@ test_census_event(void) msp_print_state(&msp, _devnull); /* Check there is more than 1 node at the census time. */ - for (i = 0; i < tables.nodes.num_rows; i++) { - if (tables.nodes.time[i] == 0.5) { + for (i = 0; i < tables2.nodes.num_rows; i++) { + if (tables2.nodes.time[i] == census_time) { num_census_nodes++; } } @@ -2950,8 +2987,9 @@ test_census_event(void) ret = msp_free(&msp); CU_ASSERT_EQUAL(ret, 0); + tsk_table_collection_free(&tables1); + tsk_table_collection_free(&tables2); gsl_rng_free(rng); - tsk_table_collection_free(&tables); } static void diff --git a/tests/test_demography.py b/tests/test_demography.py index 42b6b71af..7d14256e6 100644 --- a/tests/test_demography.py +++ b/tests/test_demography.py @@ -3347,20 +3347,22 @@ class TestCensusEvent: Tests of the census demographic event. """ - def verify(self, ts, census_time): + def verify(self, ts, census_time, unary=True): """ Verifies that a census event has been added correctly. """ - census_ids = np.where(ts.tables.nodes.flags == msprime.NODE_IS_CEN_EVENT)[0] - for u in census_ids: - assert ts.tables.nodes.time[u] == census_time + census_ids = np.where(ts.tables.nodes.time == census_time)[0] + assert np.all( + (ts.tables.nodes.flags[census_ids] & msprime.NODE_IS_CEN_EVENT) > 0 + ) assert len(census_ids) > 1 # Check that all samples have a census ancestor on each tree. for tree in ts.trees(): leaves = [] census_nodes = [u for u in census_ids if u in list(tree.nodes())] for node in census_nodes: - assert len(tree.children(node)) == 1 + if unary: + assert len(tree.children(node)) == 1 le = list(tree.leaves(node)) leaves += le leaves.sort() @@ -3413,12 +3415,17 @@ def test_census_at_existing_node_time(self): sample_size=3, random_seed=3, ) - with pytest.raises(_msprime.LibraryError): - msprime.simulate( - sample_size=3, - random_seed=3, - demographic_events=[demog_mod.CensusEvent(time=ts.tables.nodes.time[3])], - ) + n = ts.node(3).asdict() + census_time = n["time"] + ts = msprime.simulate( + sample_size=3, + random_seed=3, + demographic_events=[demog_mod.CensusEvent(time=census_time)], + ) + self.verify(ts, census_time, unary=False) + nn = ts.node(3).asdict() + nn["flags"] &= ~(msprime.NODE_IS_CEN_EVENT) + assert n == nn def test_migration_time_equals_census_time(self): census_time = 100 @@ -3487,6 +3494,21 @@ def test_dtwf_census_nodes_intermediate_time_ok(self): ) self.verify(ts, 1.1) + def test_census_at_end_time(self): + demography = msprime.Demography.isolated_model([100]) + end_time = 50 + demography.add_census(time=end_time) + ts = msprime.sim_ancestry( + samples=5, + demography=demography, + random_seed=1, + model="hudson", + end_time=end_time, + recombination_rate=1e-7, + sequence_length=1e6, + ) + self.verify(ts, end_time) + class TestPossibleLineagesOldStyle: """