Merge pull request #2321 from jeromekelleher/internal-pedigree-sample…

…-first-pass

Initial steps for internal pedigree samples

algorithms.py

@@ -931,7 +931,12 @@ def __init__(
         # implemented using `ParametricAncestryModel()`
         self.hull_offset = hull_offset
 
-        self.initialise(tables.tree_sequence())
+        if model == "fixed_pedigree":
+            self.t = 0
+            self.S[0] = 0
+            self.S[self.L] = -1
+        else:
+            self.initialise(tables.tree_sequence())
 
         self.num_ca_events = 0
         self.num_re_events = 0
@@ -981,6 +986,7 @@ def initialise(self, ts):
         root_segments_tail = [None for _ in range(ts.num_nodes)]
         root_lineages = [None for _ in range(ts.num_nodes)]
         last_S = -1
+        start_time = np.inf
         for tree in ts.trees():
             left, right = tree.interval
             S = 0 if tree.num_roots == 1 else tree.num_roots
@@ -991,6 +997,7 @@ def initialise(self, ts):
             # any ancestral segments to the state.
             if tree.num_roots > 1:
                 for root in tree.roots:
+                    start_time = min(start_time, tree.time(root))
                     population = ts.node(root).population
                     if root_segments_head[root] is None:
                         seg = self.alloc_segment(left, right, root)
@@ -1014,7 +1021,7 @@ def initialise(self, ts):
         # Insert the segment chains into the algorithm state.
         for node in range(ts.num_nodes):
             lineage = root_lineages[node]
-            if lineage is not None:
+            if lineage is not None and ts.nodes_time[node] == start_time:
                 seg = lineage.head
                 left_end = seg.left
                 while seg is not None:
@@ -1616,6 +1623,16 @@ def process_pedigree_common_ancestors(self, ind, ploid):
         individual, then recombine and distribute the remaining ancestral
         material among its parent ploids.
         """
+        node = ind.nodes[ploid]
+        is_sample = (self.tables.nodes.flags[node] & tskit.NODE_IS_SAMPLE) > 0
+        if is_sample:
+            segment = self.alloc_segment(0, self.L, node)
+            lineage = self.alloc_lineage(segment, population=0)
+            self.add_lineage(lineage)
+            ind.add_common_ancestor(lineage.head, ploid=ploid)
+            for k in list(self.S.keys())[:-1]:
+                self.S[k] += 1
+
         common_ancestors = ind.common_ancestors[ploid]
         if len(common_ancestors) == 0:
             # No ancestral material inherited on this ploid of this individual
@@ -1632,7 +1649,6 @@ def process_pedigree_common_ancestors(self, ind, ploid):
         # monoploid genome for this ploid of this individual.
         # If any coalescences occur, they use the corresponding node ID.
         # FIXME update the population/label here
-        node = ind.nodes[ploid]
         genome = self.merge_ancestors(common_ancestors, 0, 0, node)
         if ind.parents[ploid] == tskit.NULL:
             # If this individual is a founder we need to make sure that all
@@ -1676,18 +1692,6 @@ def dtwf_climb_pedigree(self):
         Simulates transmission of ancestral material through a pre-specified
         pedigree
         """
-        assert self.num_populations == 1  # Single pop/pedigree for now
-        pop = self.P[0]
-
-        # Go through the extant lineages and gather the ancestral material
-        # into the corresponding pedigree individuals.
-        for lineage in pop.iter_ancestors():
-            u = lineage.head.node
-            node = self.tables.nodes[u]
-            assert node.individual != tskit.NULL
-            ind = self.pedigree.individuals[node.individual]
-            ind.add_common_ancestor(lineage.head, ploid=ind.nodes.index(u))
-
         # Visit pedigree individuals in time order.
         visit_order = sorted(self.pedigree.individuals, key=lambda x: (x.time, x.id))
         for ind in visit_order:

lib/msprime.c

@@ -1859,11 +1859,13 @@ msp_verify_overlaps(msp_t *self)
     int ok = overlap_counter_alloc(&counter, self->sequence_length, 0);
     tsk_bug_assert(ok == 0);
 
-    /* add in the overlaps for ancient samples */
-    for (j = self->next_sampling_event; j < self->num_sampling_events; j++) {
-        se = self->sampling_events[j];
-        for (u = self->root_segments[se.sample]; u != NULL; u = u->next) {
-            overlap_counter_increment_interval(&counter, u->left, u->right);
+    if (self->model.type != MSP_MODEL_WF_PED) {
+        /* add in the overlaps for ancient samples */
+        for (j = self->next_sampling_event; j < self->num_sampling_events; j++) {
+            se = self->sampling_events[j];
+            for (u = self->root_segments[se.sample]; u != NULL; u = u->next) {
+                overlap_counter_increment_interval(&counter, u->left, u->right);
+            }
         }
     }
 
@@ -2495,12 +2497,10 @@ msp_store_coalescence_edge(
     int ret = 0;
 
     if (self->model.type == MSP_MODEL_DTWF || self->model.type == MSP_MODEL_WF_PED) {
-        if (self->additional_nodes & MSP_NODE_IS_RE_EVENT) {
-            // parent and child can be equal
-            // don't store edges
-            if (parent == child) {
-                return ret;
-            }
+        // parent and child can be equal if we're storing RE events, or if we've
+        // got internal samples. Don't store edges
+        if (parent == child) {
+            return ret;
         }
     }
 
@@ -2914,13 +2914,30 @@ msp_pedigree_add_individual_common_ancestor(
 }
 
 static int MSP_WARN_UNUSED
-msp_pedigree_add_sample_ancestry(msp_t *self, segment_t *segment)
+msp_pedigree_add_sample_ancestry(msp_t *self, tsk_id_t node_id)
 {
     int ret = 0;
     tsk_size_t ploid;
-    tsk_id_t node_id = segment->value;
     tsk_id_t individual_id;
     individual_t *ind;
+    avl_node_t *a;
+    node_mapping_t *counter;
+    segment_t *seg
+        = msp_alloc_segment(self, 0, self->sequence_length, node_id, 0, 0, NULL, NULL);
+    lineage_t *lin = msp_alloc_lineage(self, seg, seg, 0, 0);
+
+    if (seg == NULL || lin == NULL) {
+        ret = MSP_ERR_NO_MEMORY;
+        goto out;
+    }
+    ret = msp_insert_individual(self, lin);
+    if (ret != 0) {
+        goto out;
+    }
+    for (a = self->overlap_counts.head; a != self->overlap_counts.tail; a = a->next) {
+        counter = (node_mapping_t *) a->item;
+        counter->value += 1;
+    }
 
     tsk_bug_assert(node_id < (tsk_id_t) self->tables->nodes.num_rows);
     individual_id = self->tables->nodes.individual[node_id];
@@ -2934,14 +2951,7 @@ msp_pedigree_add_sample_ancestry(msp_t *self, segment_t *segment)
         }
     }
     tsk_bug_assert(ploid < ind->ploidy);
-    if (avl_count(&ind->common_ancestors[ploid]) > 0) {
-        /* This is where we'd deal with the internal samples. What we'll probably
-         * need to do is to go through the ancestry after we've processed the
-         * individual and then pad out any gaps in the ancestry segments. */
-        ret = MSP_ERR_PEDIGREE_INTERNAL_SAMPLE;
-        goto out;
-    }
-    ret = msp_pedigree_add_individual_common_ancestor(self, ind->id, segment, ploid);
+    ret = msp_pedigree_add_individual_common_ancestor(self, ind->id, seg, ploid);
     if (ret != 0) {
         goto out;
     }
@@ -2955,9 +2965,11 @@ msp_pedigree_initialise(msp_t *self)
     int ret = 0;
     population_t *pop;
     lineage_t *lin;
+    segment_t *seg;
     avl_node_t *a;
     label_id_t label = 0;
     tsk_size_t j;
+    node_mapping_t *counter;
 
     if (self->next_demographic_event != NULL) {
         ret = MSP_ERR_UNSUPPORTED_OPERATION;
@@ -2981,14 +2993,21 @@ msp_pedigree_initialise(msp_t *self)
 
     for (j = 0; j < self->num_populations; j++) {
         pop = &self->populations[j];
+        /* Rather than messing about with how we initialise from trees, it's
+         * easier to just remove the lineages here, before we add them
+         * back later when dealing with samples in the pedigree. */
         for (a = pop->ancestors[label].head; a != NULL; a = a->next) {
             lin = (lineage_t *) a->item;
-            ret = msp_pedigree_add_sample_ancestry(self, lin->head);
-            if (ret != 0) {
-                goto out;
+            for (seg = lin->head; seg != NULL; seg = seg->next) {
+                msp_free_segment(self, seg);
             }
+            msp_remove_individual(self, lin);
         }
     }
+    tsk_bug_assert(avl_count(&self->overlap_counts) == 2);
+    counter = (node_mapping_t *) self->overlap_counts.head->item;
+    counter->value = 0;
+
     self->pedigree.next_individual = 0;
 out:
     return ret;
@@ -4589,31 +4608,6 @@ msp_apply_sampling_events(msp_t *self, double time)
     return ret;
 }
 
-static int MSP_WARN_UNUSED
-msp_pedigree_insert_ancient_samples(msp_t *self)
-{
-    int ret = 0;
-    sampling_event_t *se;
-    segment_t *root_seg, *new_head;
-
-    while (self->next_sampling_event < self->num_sampling_events
-           && self->sampling_events[self->next_sampling_event].time <= self->time) {
-        se = self->sampling_events + self->next_sampling_event;
-        root_seg = self->root_segments[se->sample];
-        ret = msp_insert_root_segments(self, root_seg, &new_head);
-        if (ret != 0) {
-            goto out;
-        }
-        ret = msp_pedigree_add_sample_ancestry(self, new_head);
-        if (ret != 0) {
-            goto out;
-        }
-        self->next_sampling_event++;
-    }
-out:
-    return ret;
-}
-
 static double
 msp_get_next_fixed_event_time(msp_t *self)
 {
@@ -5483,13 +5477,21 @@ msp_pedigree_process_common_ancestors(msp_t *self, individual_t *ind, tsk_size_t
 {
     int ret = 0;
     tsk_id_t node = ind->nodes[ploid];
+    bool is_sample = (self->tables->nodes.flags[node] & TSK_NODE_IS_SAMPLE) != 0;
     tsk_id_t parent = ind->parents[ploid];
     avl_tree_t *common_ancestors = &ind->common_ancestors[ploid];
     segment_t *genome, *parent_ancestry[MSP_MAX_PED_PLOIDY], *seg;
     const tsk_size_t ploidy = self->ploidy;
     tsk_size_t j;
     tsk_id_t *parent_nodes;
 
+    if (is_sample) {
+        ret = msp_pedigree_add_sample_ancestry(self, node);
+        if (ret != 0) {
+            goto out;
+        }
+    }
+
     /* FIXME - assuming 1 label here */
     ret = msp_merge_n_ancestors(
         self, common_ancestors, ind->population, 0, node, &genome);
@@ -5607,10 +5609,6 @@ msp_run_pedigree(msp_t *self, double max_time, unsigned long max_events)
         }
         tsk_bug_assert(ind->time >= self->time);
         self->time = ind->time;
-        ret = msp_pedigree_insert_ancient_samples(self);
-        if (ret != 0) {
-            goto out;
-        }
         for (ploid = 0; ploid < ind->ploidy; ploid++) {
             ret = msp_pedigree_process_common_ancestors(self, ind, ploid);
             if (ret != 0) {

lib/tests/test_pedigrees.c

@@ -553,8 +553,8 @@ test_replicates_ancient_samples(void)
     CU_ASSERT_EQUAL_FATAL(ret, 0);
 
     for (j = 0; j < 20; j++) {
-        msp_verify(&msp, 0);
         ret = msp_run(&msp, DBL_MAX, UINT32_MAX);
+
         CU_ASSERT_EQUAL_FATAL(ret, MSP_EXIT_MODEL_COMPLETE);
         msp_verify(&msp, 0);
         ret = msp_finalise_tables(&msp);
@@ -844,24 +844,12 @@ test_errors(void)
 static void
 test_internal_samples(void)
 {
-    int ret;
     tsk_id_t parents[] = { -1, -1, -1, -1, 0, 1 };
     double time[] = { 1.0, 1.0, 0.0 };
     tsk_flags_t is_sample[] = { true, true, true };
-    msp_t msp;
-    tsk_table_collection_t tables;
-    gsl_rng *rng = safe_rng_alloc();
-
-    ret = build_pedigree_sim(
-        &msp, &tables, rng, 100, 2, 3, parents, time, is_sample, NULL);
-    ret = msp_initialise(&msp);
-    CU_ASSERT_EQUAL_FATAL(ret, 0);
-    ret = msp_run(&msp, DBL_MAX, UINT32_MAX);
-    CU_ASSERT_EQUAL_FATAL(ret, MSP_ERR_PEDIGREE_INTERNAL_SAMPLE);
-    tsk_table_collection_free(&tables);
-    msp_free(&msp);
 
-    gsl_rng_free(rng);
+    verify_pedigree(0, 1, 3, parents, time, is_sample, NULL, 0);
+    verify_pedigree(0.1, 1, 3, parents, time, is_sample, NULL, 0);
 }
 
 static void

lib/util.c

@@ -183,12 +183,6 @@ msp_strerror_internal(int err)
             ret = "All individuals in the input pedigree must be associated with "
                   "exactly two parents (can be TSK_NULL, if not known)";
             break;
-        case MSP_ERR_PEDIGREE_INTERNAL_SAMPLE:
-            ret = "Samples that are internal nodes in the pedigree are not "
-                  "currently supported. Please comment on this GitHub issue if you "
-                  "would like to see this feature implemented: "
-                  "https://github.com/tskit-dev/msprime/issues/1855 ";
-            break;
 
         case MSP_ERR_BAD_PROPORTION:
             ret = "Proportion values must have 0 <= x <= 1";

lib/util.h

@@ -134,7 +134,6 @@
 #define MSP_ERR_PEDIGREE_TIME_TRAVEL                                -88
 #define MSP_ERR_PEDIGREE_IND_NOT_DIPLOID                            -89
 #define MSP_ERR_PEDIGREE_IND_NOT_TWO_PARENTS                        -90
-#define MSP_ERR_PEDIGREE_INTERNAL_SAMPLE                            -91
 
 /* clang-format on */
 /* This bit is 0 for any errors originating from tskit */

tests/test_algorithms.py

@@ -423,6 +423,21 @@ def test_pedigree_trio(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_internal_sample(self, r):
+        input_tables = simulate_pedigree(num_founders=4, num_generations=4)
+        flags = input_tables.nodes.flags
+        flags[8:12] = tskit.NODE_IS_SAMPLE  # Make inds 4 and 5 samples
+        input_tables.nodes.flags = flags
+        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 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)