Merge pull request #2388 from jeromekelleher/more-internal-samples-2

Some tests for no-leaf-samples

Author: jeromekelleher

CHANGELOG.md

@@ -1,7 +1,12 @@
 # Changelog
 
 
-## [1.3.5] - 2025-XX-XX
+## [1.4.0] - 2025-XX-XX
+
+**New features**
+
+- The `FixedPedigree` simulation model now supports internal samples ({issue}`1855`,
+  {pr}`2321`, {pr}`2326`, {pr}`2388`, {user}`abureau`, {user}`jeromekelleher`).
 
 **Breaking changes**:
 
@@ -33,7 +38,7 @@ Bugfix release for issues with Dirac and Beta coalescent models.
 
 - Add `record_provenance` argument to `sim_mutations` ({issue}`2272`, {pr}`2273`, {user}`peterlharp`)
 
-- Add support and wheels for Python3.12, drop Python3.8 
+- Add support and wheels for Python3.12, drop Python3.8
 
 ## [1.3.1] - 2024-02-09
 
@@ -48,7 +53,7 @@ Bugfix release for issues with Dirac and Beta coalescent models.
 - Add a `MicrosatMutationModel` mutation model class, that
   represents a generalized interface for constructing mutational
   models appropriate to STRs. In addition 3 specific microsat models
-  are added `SMM`, `TPM`, and `EL2`.  
+  are added `SMM`, `TPM`, and `EL2`.
   ({issue}`2013`, {user}`andrewkern`).
 
 - Raise an error if `log_arg_likelihood` is called on a recombinant tree

lib/msprime.c

@@ -3007,6 +3007,25 @@ msp_pedigree_add_sample_ancestry(msp_t *self, tsk_id_t node_id)
     return ret;
 }
 
+static void
+msp_skip_pedigree_non_samples(msp_t *self)
+{
+    individual_t *ind = NULL;
+    pedigree_t *pedigree = &self->pedigree;
+    const tsk_id_t num_individuals = (tsk_id_t) pedigree->num_individuals;
+    const tsk_flags_t *nodes_flags = self->tables->nodes.flags;
+    tsk_id_t j, node;
+
+    for (j = 0; j < num_individuals; j++) {
+        ind = pedigree->visit_order[j];
+        node = ind->nodes[0];
+        if ((nodes_flags[node] & TSK_NODE_IS_SAMPLE) != 0) {
+            break;
+        }
+    }
+    pedigree->next_individual = j;
+}
+
 static int MSP_WARN_UNUSED
 msp_pedigree_initialise(msp_t *self)
 {
@@ -3057,6 +3076,7 @@ msp_pedigree_initialise(msp_t *self)
     counter->value = 0;
 
     self->pedigree.next_individual = 0;
+    msp_skip_pedigree_non_samples(self);
 out:
     return ret;
 }

lib/tests/test_pedigrees.c

@@ -204,6 +204,54 @@ verify_pedigree(double recombination_rate, unsigned long seed,
     tsk_table_collection_free(&tables);
 }
 
+static void
+verify_pedigree_event_by_event(double recombination_rate, unsigned long seed,
+    tsk_size_t num_individuals, tsk_id_t *parents, double *time, tsk_flags_t *is_sample,
+    tsk_id_t *population, uint32_t additional_nodes)
+{
+    int ret, status1, status2;
+    int ploidy = 2;
+    msp_t msp1, msp2;
+    tsk_table_collection_t tables1, tables2;
+    gsl_rng *rng1 = safe_rng_alloc();
+    gsl_rng *rng2 = safe_rng_alloc();
+
+    ret = build_pedigree_sim(&msp1, &tables1, rng1, 100, ploidy, num_individuals,
+        parents, time, is_sample, population);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = build_pedigree_sim(&msp2, &tables2, rng2, 100, ploidy, num_individuals,
+        parents, time, is_sample, population);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = msp_initialise(&msp1);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = msp_initialise(&msp2);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = msp_run(&msp1, DBL_MAX, UINT32_MAX);
+    CU_ASSERT_FATAL(ret == MSP_EXIT_MODEL_COMPLETE || ret == MSP_EXIT_COALESCENCE);
+    status1 = ret;
+
+    ret = MSP_EXIT_MAX_EVENTS;
+    while (ret == MSP_EXIT_MAX_EVENTS) {
+        ret = msp_run(&msp2, DBL_MAX, 1);
+        CU_ASSERT_FATAL(ret >= 0);
+        msp_verify(&msp2, 0);
+    }
+    CU_ASSERT_FATAL(ret == MSP_EXIT_MODEL_COMPLETE || ret == MSP_EXIT_COALESCENCE);
+    status2 = ret;
+    CU_ASSERT_EQUAL(status1, status2);
+
+    CU_ASSERT_TRUE(tsk_table_collection_equals(&tables1, &tables2, 0));
+
+    tsk_table_collection_free(&tables1);
+    msp_free(&msp1);
+    tsk_table_collection_free(&tables2);
+    msp_free(&msp2);
+    gsl_rng_free(rng1);
+    gsl_rng_free(rng2);
+}
+
 static void
 test_trio(void)
 {
@@ -850,6 +898,29 @@ test_internal_samples(void)
 
     verify_pedigree(0, 1, 3, parents, time, is_sample, NULL, 0);
     verify_pedigree(0.1, 1, 3, parents, time, is_sample, NULL, 0);
+
+    verify_pedigree_event_by_event(0, 1, 3, parents, time, is_sample, NULL, 0);
+    verify_pedigree_event_by_event(0.1, 1, 3, parents, time, is_sample, NULL, 0);
+}
+
+static void
+test_no_leaf_samples(void)
+{
+    tsk_id_t *parents = deep_n10_parents;
+    double *time = deep_n10_time;
+    size_t num_inds = sizeof(deep_n10_time) / sizeof(*deep_n10_time);
+    tsk_flags_t is_sample[num_inds];
+    size_t j;
+
+    for (j = 0; j < num_inds; j++) {
+        is_sample[j] = time[j] == 1;
+    }
+
+    verify_pedigree(0, 1, num_inds, parents, time, is_sample, NULL, 0);
+    verify_pedigree(0.1, 1, num_inds, parents, time, is_sample, NULL, 0);
+
+    verify_pedigree_event_by_event(0, 1, num_inds, parents, time, is_sample, NULL, 0);
+    verify_pedigree_event_by_event(0.1, 1, num_inds, parents, time, is_sample, NULL, 0);
 }
 
 static void
@@ -895,6 +966,7 @@ test_trio_same_pop(void)
     tsk_id_t population[] = { 2, 2, 2 };
 
     verify_pedigree(0, 1, 3, parents, time, NULL, population, 0);
+    verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0);
 }
 
 static void
@@ -906,6 +978,7 @@ test_trio_child_different_pop(void)
     tsk_id_t population[] = { 2, 2, 1 };
 
     verify_pedigree(0, 1, 3, parents, time, NULL, population, 0);
+    verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0);
 }
 
 int
@@ -918,6 +991,7 @@ main(int argc, char **argv)
         { "test_sibs", test_sibs },
         { "test_ancient_sample", test_ancient_sample },
         { "test_internal_samples", test_internal_samples },
+        { "test_no_leaf_samples", test_no_leaf_samples },
         { "test_large_family", test_large_family },
         { "test_unrelated_n3", test_unrelated_n3 },
         { "test_very_deep_n2", test_very_deep_n2 },
@@ -933,7 +1007,6 @@ main(int argc, char **argv)
         { "test_replicates_exit_coalescence", test_replicates_exit_coalescence },
         { "test_errors", test_errors },
         { "test_combined_with_other_models", test_combined_with_other_models },
-
         { "test_trio_same_pop", test_trio_same_pop },
         { "test_trio_child_different_pop", test_trio_child_different_pop },
         CU_TEST_INFO_NULL,

tests/test_algorithms.py

@@ -436,6 +436,21 @@ def test_pedigree_internal_sample(self, r):
         input_tables.nodes.assert_equals(output_tables.nodes[: len(input_tables.nodes)])
         assert len(output_tables.edges) >= 2
 
+    @pytest.mark.parametrize("r", [0, 0.1, 1])
+    def test_pedigree_only_internal_samples(self, r):
+        input_tables = simulate_pedigree(
+            num_founders=4, num_generations=4, sample_gen=[1]
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ts_path = pathlib.Path(tmpdir) / "pedigree.trees"
+            input_tables.dump(ts_path)
+            ts = self.run_script(f"0 --from-ts {ts_path} -r {r} --model=fixed_pedigree")
+        output_tables = ts.dump_tables()
+        input_tables.individuals.assert_equals(output_tables.individuals)
+        input_tables.nodes.assert_equals(output_tables.nodes[: len(input_tables.nodes)])
+        assert np.all(ts.nodes_time[ts.samples()] == 1)
+        assert len(output_tables.edges) >= 2
+
     @pytest.mark.parametrize("num_founders", [1, 2, 20])
     def test_one_gen_pedigree(self, num_founders):
         tables = simulate_pedigree(num_founders=num_founders, num_generations=1)

tests/test_pedigree.py

@@ -567,6 +567,30 @@ def test_shallow_internal(self, num_founders, recombination_rate):
         )
         self.verify(tables, recombination_rate)
 
+    @pytest.mark.parametrize("num_founders", [2, 3, 5])
+    @pytest.mark.parametrize("recombination_rate", [0, 0.01])
+    def test_internal_and_leaf_samples(self, num_founders, recombination_rate):
+        tables = simulate_pedigree(
+            num_founders=num_founders,
+            num_children_prob=[0, 0, 1],
+            num_generations=5,
+            sequence_length=100,
+            sample_gen=[0, 2],
+        )
+        self.verify(tables, recombination_rate)
+
+    @pytest.mark.parametrize("num_founders", [2, 3, 5])
+    @pytest.mark.parametrize("recombination_rate", [0, 0.01])
+    def test_no_leaf_samples(self, num_founders, recombination_rate):
+        tables = simulate_pedigree(
+            num_founders=num_founders,
+            num_children_prob=[0, 0, 1],
+            num_generations=5,
+            sequence_length=100,
+            sample_gen=[1],
+        )
+        self.verify(tables, recombination_rate)
+
     @pytest.mark.parametrize("num_founders", [2, 3, 10, 20])
     @pytest.mark.parametrize("recombination_rate", [0, 0.01])
     def test_deep(self, num_founders, recombination_rate):
@@ -836,8 +860,11 @@ def verify(self, input_tables, recombination_rate=0):
             recombination_rate=recombination_rate,
             random_seed=1,
         )
+        # print(ts1.tables)
+        # print(ts1.draw_text())
         sim.run(event_chunk=1)
         output_tables = tskit.TableCollection.fromdict(sim.tables.asdict())
+        # print(output_tables)
         output_tables.assert_equals(ts1.tables, ignore_provenance=True)
         return ts1