census node per lineage, not segment: closes #2540

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.

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;
             }

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,17 +2978,18 @@ 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++;
         }
     }
     CU_ASSERT_TRUE(num_census_nodes > 1);
 
     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

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:
     """