Repository: tskit-dev/msprime (language: python)

Issue:
Enable gene conversion for SMC_K model
The general idea of how this should be done is already incorporated in `algorithms.py`, what's left is implementing this in C and adding a fair amount of tests.


Discussion / hints:

Fix patch:
diff --git a/algorithms.py b/algorithms.py
index 484ff7488..5925b16b2 100644
--- a/algorithms.py
+++ b/algorithms.py
@@ -840,6 +840,7 @@ def __init__(
         gene_conversion_rate=0.0,
         gene_conversion_length=1,
         discrete_genome=True,
+        hull_offset=None,
     ):
         # Must be a square matrix.
         N = len(migration_matrix)
@@ -900,9 +901,7 @@ def __init__(
 
         # set hull_offset for smc_k, deviates from actual pattern
         # implemented using `ParametricAncestryModel()`
-        self.hull_offset = None
-        if self.model == "smc_k":
-            self.hull_offset = 0.0
+        self.hull_offset = hull_offset
 
         self.initialise(tables.tree_sequence())
 
@@ -1054,7 +1053,7 @@ def alloc_hull(self, left, right, lineage_head):
         alpha = lineage_head
         hull = self.hull_stack.pop()
         hull.left = left
-        hull.right = min(right + self.hull_offset, self.L)
+        hull.right = right
         while alpha.prev is not None:
             alpha = alpha.prev
         assert alpha is not None
@@ -1799,9 +1798,7 @@ def hudson_recombination_event(self, label, return_heads=False):
             self.P[pop].reset_hull_right(label, lhs_hull, rhs_right, lhs_hull.right)
 
             # create hull for alpha
-            alpha_hull = self.alloc_hull(
-                alpha.left, rhs_right - self.hull_offset, alpha
-            )
+            alpha_hull = self.alloc_hull(alpha.left, rhs_right, alpha)
             self.P[alpha.population].add_hull(label, alpha_hull)
 
         self.set_segment_mass(alpha)
@@ -1947,10 +1944,9 @@ def wiuf_gene_conversion_within_event(self, label):
         else:
             # rbp lies beyond segment chain, regular recombination logic applies
             if insert_alpha and self.model == "smc_k":
-                assert reset_right >= 0
-                self.P[pop].reset_hull_right(
-                    label, hull, hull_right + self.hull_offset, reset_right
-                )
+                assert reset_right > 0
+                reset_right = min(reset_right + self.hull_offset, self.L)
+                self.P[pop].reset_hull_right(label, hull, hull.right, reset_right)
 
         #        y            z
         #  |  ========== ... ===== |
@@ -1965,6 +1961,7 @@ def wiuf_gene_conversion_within_event(self, label):
         if new_individual_head is not None:
             if self.model == "smc_k":
                 assert hull_left < hull_right
+                hull_right = min(self.L, hull_right + self.hull_offset)
                 hull = self.alloc_hull(hull_left, hull_right, new_individual_head)
                 self.P[new_individual_head.population].add_hull(
                     new_individual_head.label, hull
@@ -2037,12 +2034,12 @@ def wiuf_gene_conversion_left_event(self, label):
 
         if self.model == "smc_k":
             # lhs logic is identical to the lhs recombination event
-            rhs_right = lhs_hull.right
-            lhs_hull.right = right + self.hull_offset
-            self.P[pop].reset_hull_right(label, lhs_hull, rhs_right, lhs_hull.right)
+            lhs_old_right = lhs_hull.right
+            lhs_new_right = min(self.L, right + self.hull_offset)
+            self.P[pop].reset_hull_right(label, lhs_hull, lhs_old_right, lhs_new_right)
 
             # rhs
-            hull = self.alloc_hull(alpha.left, rhs_right, alpha)
+            hull = self.alloc_hull(alpha.left, lhs_old_right, alpha)
             self.P[pop].add_hull(label, hull)
 
         self.set_segment_mass(alpha)
@@ -2818,6 +2815,7 @@ def run_simulate(args):
         gene_conversion_rate=gc_rate,
         gene_conversion_length=mean_tract_length,
         discrete_genome=args.discrete,
+        hull_offset=args.offset,
     )
     ts = s.simulate(args.end_time)
     ts.dump(args.output_file)
@@ -2904,6 +2902,7 @@ def add_simulator_arguments(parser):
         help="The delta_t value for selective sweeps",
     )
     parser.add_argument("--model", default="hudson")
+    parser.add_argument("--offset", type=float, default=0.0)
     parser.add_argument(
         "--from-ts",
         "-F",
diff --git a/lib/msprime.c b/lib/msprime.c
index 5df950d19..aeceb7248 100644
--- a/lib/msprime.c
+++ b/lib/msprime.c
@@ -918,8 +918,8 @@ msp_alloc_hull(msp_t *self, double left, double right, segment_t *lineage)
     tsk_bug_assert(hull->id < UINT_MAX * self->num_populations);
 
     hull->left = left;
-    hull->right = GSL_MIN(
-        right + self->model.params.smc_k_coalescent.hull_offset, self->sequence_length);
+    tsk_bug_assert(right <= self->sequence_length);
+    hull->right = right;
     hull->lineage = lineage;
     hull->count = 0;
     hull->insertion_order = UINT64_MAX;
@@ -1344,7 +1344,10 @@ msp_insert_hull(msp_t *self, hull_t *hull)
 
     /* insert hullend into state */
     hullend = msp_alloc_hullend(self, hull->right, u->label);
-    tsk_bug_assert(hullend != NULL);
+    if (hullend == NULL) {
+        ret = MSP_ERR_NO_MEMORY;
+        goto out;
+    }
     node = msp_alloc_avl_node(self);
     if (node == NULL) {
         ret = MSP_ERR_NO_MEMORY;
@@ -1876,7 +1879,7 @@ msp_verify_hulls(msp_t *self)
     hull_t *hull, hull_a, hull_b;
     hullend_t *hullend;
     fenwick_t *coal_mass_index;
-    double pos;
+    double pos, hull_right;
     uint32_t N;
     uint64_t io;
 
@@ -1887,11 +1890,12 @@ msp_verify_hulls(msp_t *self)
             count = 0;
             avl = &self->populations[population_id].ancestors[label_id];
             /* generate all possible lineage pairs */
-            if (msp_get_num_population_ancestors(self, population_id) == 0) {
+            if (avl_count(avl) == 0) {
                 continue;
             }
             for (a = avl->head; a->next != NULL; a = a->next) {
                 x = (segment_t *) a->item;
+                hull_right = x->hull->right;
                 hull_a.left = x->left;
                 while (x->next != NULL) {
                     x = x->next;
@@ -1899,6 +1903,7 @@ msp_verify_hulls(msp_t *self)
                 hull_a.right
                     = GSL_MIN(x->right + self->model.params.smc_k_coalescent.hull_offset,
                         self->sequence_length);
+                tsk_bug_assert(hull_a.right == hull_right);
                 for (b = a->next; b != NULL; b = b->next) {
                     y = (segment_t *) b->item;
                     hull_b.left = y->left;
@@ -3254,6 +3259,10 @@ msp_recombination_event(msp_t *self, label_id_t label, segment_t **lhs, segment_
 
         /* create new hull for alpha */
         rhs_hull = msp_alloc_hull(self, alpha->left, rhs_right, alpha);
+        if (rhs_hull == NULL) {
+            ret = MSP_ERR_NO_MEMORY;
+            goto out;
+        }
         ret = msp_insert_hull(self, rhs_hull);
         if (ret != 0) {
             goto out;
@@ -3308,6 +3317,9 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
     segment_t *x, *y, *alpha, *head, *tail, *z, *new_individual_head;
     double left_breakpoint, right_breakpoint, tl;
     bool insert_alpha;
+    hull_t *hull = NULL;
+    double reset_right = 0.0;
+    double tract_hull_left, tract_hull_right;
 
     tsk_bug_assert(self->gc_mass_index != NULL);
     self->num_gc_events++;
@@ -3338,6 +3350,11 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
         return 0;
     }
 
+    if (self->model.type == MSP_MODEL_SMC_K) {
+        hull = segment_get_hull(y);
+        tsk_bug_assert(hull != NULL);
+    }
+
     /* Process left break */
     insert_alpha = true;
     if (left_breakpoint <= y->left) {
@@ -3355,6 +3372,7 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
             insert_alpha = false;
         } else {
             x->next = NULL;
+            reset_right = x->right;
         }
         y->prev = NULL;
         alpha = y;
@@ -3383,6 +3401,7 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
         y->right = left_breakpoint;
         msp_set_segment_mass(self, y);
         tail = y;
+        reset_right = left_breakpoint;
 
         if (!msp_has_breakpoint(self, left_breakpoint)) {
             ret = msp_insert_breakpoint(self, left_breakpoint);
@@ -3395,7 +3414,10 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
 
     // Find the segment z that the right breakpoint falls in
     z = alpha;
+    tract_hull_left = z->left;
+    tract_hull_right = 0.0;
     while (z != NULL && right_breakpoint >= z->right) {
+        tract_hull_right = z->right;
         z = z->next;
     }
 
@@ -3424,6 +3446,7 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
             }
             z->right = right_breakpoint;
             z->next = NULL;
+            tract_hull_right = right_breakpoint;
             msp_set_segment_mass(self, z);
 
             if (!msp_has_breakpoint(self, right_breakpoint)) {
@@ -3452,6 +3475,16 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
         }
         head->prev = tail;
         msp_set_segment_mass(self, head);
+    } else {
+        // rbp lies beyond segment chain, regular recombination logic applies
+        if (insert_alpha && self->model.type == MSP_MODEL_SMC_K) {
+            tsk_bug_assert(reset_right > 0);
+            reset_right
+                = GSL_MIN(reset_right + self->model.params.smc_k_coalescent.hull_offset,
+                    self->sequence_length);
+            msp_reset_hull_right(
+                self, hull, hull->right, reset_right, y->population, y->label);
+        }
     }
 
     //        y            z
@@ -3466,6 +3499,22 @@ msp_gene_conversion_event(msp_t *self, label_id_t label)
         new_individual_head = head;
     }
     if (new_individual_head != NULL) {
+        if (self->model.type == MSP_MODEL_SMC_K) {
+            tsk_bug_assert(tract_hull_left < tract_hull_right);
+            tract_hull_right = GSL_MIN(
+                tract_hull_right + self->model.params.smc_k_coalescent.hull_offset,
+                self->sequence_length);
+            hull = msp_alloc_hull(
+                self, tract_hull_left, tract_hull_right, new_individual_head);
+            if (hull == NULL) {
+                ret = MSP_ERR_NO_MEMORY;
+                goto out;
+            }
+            ret = msp_insert_hull(self, hull);
+            if (ret != 0) {
+                goto out;
+            }
+        }
         ret = msp_insert_individual(self, new_individual_head);
     } else {
         self->num_noneffective_gc_events++;
@@ -3718,6 +3767,10 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l
             r += self->model.params.smc_k_coalescent.hull_offset;
             hull = msp_alloc_hull(
                 self, merged_head->left, GSL_MIN(r, self->sequence_length), merged_head);
+            if (hull == NULL) {
+                ret = MSP_ERR_NO_MEMORY;
+                goto out;
+            }
             ret = msp_insert_hull(self, hull);
             if (ret != 0) {
                 goto out;
@@ -4151,6 +4204,10 @@ msp_insert_root_segments(msp_t *self, const segment_t *head, segment_t **new_hea
                 if (self->state != MSP_STATE_NEW) {
                     /* correct hull->right is set at the end */
                     hull = msp_alloc_hull(self, head->left, copy->right, copy);
+                    if (hull == NULL) {
+                        ret = MSP_ERR_NO_MEMORY;
+                        goto out;
+                    }
                 }
             }
         } else {
@@ -4490,7 +4547,10 @@ msp_hulls_init_counts(msp_t *self, population_id_t pop, label_id_t label)
 
         /* insert hullend into hulls_right */
         hullend = msp_alloc_hullend(self, hull->right, (label_id_t) label);
-        tsk_bug_assert(hullend != NULL);
+        if (hullend == NULL) {
+            ret = MSP_ERR_NO_MEMORY;
+            goto out;
+        }
         node_right = msp_alloc_avl_node(self);
         if (node_right == NULL) {
             ret = MSP_ERR_NO_MEMORY;
@@ -4533,8 +4593,13 @@ msp_initialise_smc_k(msp_t *self)
                     seg = seg->next;
                 }
                 /* insert into hulls_left */
+                right += self->model.params.smc_k_coalescent.hull_offset;
+                right = GSL_MIN(right, self->sequence_length);
                 hull = msp_alloc_hull(self, left, right, head);
-                tsk_bug_assert(hull != NULL);
+                if (hull == NULL) {
+                    ret = MSP_ERR_NO_MEMORY;
+                    goto out;
+                }
                 h_node = msp_alloc_avl_node(self);
                 if (h_node == NULL) {
                     ret = MSP_ERR_NO_MEMORY;
@@ -4957,9 +5022,12 @@ msp_gene_conversion_left_event(msp_t *self, label_id_t label)
     int ret = 0;
     const double gc_left_total = msp_get_total_gc_left(self);
     double h = gsl_rng_uniform(self->rng) * gc_left_total;
-    double tl, bp;
+    double tl, bp, lhs_old_right, lhs_new_right;
     segment_t *y, *x, *alpha;
+    hull_t *rhs_hull;
+    hull_t *lhs_hull = NULL;
 
+    lhs_hull = NULL;
     y = msp_find_gc_left_individual(self, label, h);
     assert(y != NULL);
 
@@ -4989,6 +5057,10 @@ msp_gene_conversion_left_event(msp_t *self, label_id_t label)
     tsk_bug_assert(y != NULL);
     self->num_gc_events++;
     x = y->prev;
+    if (self->model.type == MSP_MODEL_SMC_K) {
+        lhs_hull = segment_get_hull(y);
+        tsk_bug_assert(lhs_hull != NULL);
+    }
 
     if (y->left < bp) {
         //  x          y
@@ -5033,6 +5105,30 @@ msp_gene_conversion_left_event(msp_t *self, label_id_t label)
         // Ensure y points to the last segment left of the break for full ARG recording
         y = x;
     }
+    lhs_new_right = y->right;
+
+    if (self->model.type == MSP_MODEL_SMC_K) {
+        // lhs logic is identical to the lhs recombination event
+        lhs_old_right = lhs_hull->right;
+        lhs_new_right
+            = GSL_MIN(lhs_new_right + self->model.params.smc_k_coalescent.hull_offset,
+                self->sequence_length);
+        msp_reset_hull_right(
+            self, lhs_hull, lhs_old_right, lhs_new_right, y->population, y->label);
+
+        // rhs
+        tsk_bug_assert(alpha->left < lhs_old_right);
+        rhs_hull = msp_alloc_hull(self, alpha->left, lhs_old_right, alpha);
+        if (rhs_hull == NULL) {
+            ret = MSP_ERR_NO_MEMORY;
+            goto out;
+        }
+        ret = msp_insert_hull(self, rhs_hull);
+        if (ret != 0) {
+            goto out;
+        }
+    }
+
     msp_set_segment_mass(self, alpha);
     tsk_bug_assert(alpha->prev == NULL);
     ret = msp_insert_individual(self, alpha);
diff --git a/msprime/ancestry.py b/msprime/ancestry.py
index 31f94184f..552feddfc 100644
--- a/msprime/ancestry.py
+++ b/msprime/ancestry.py
@@ -833,7 +833,6 @@ def _parse_sim_ancestry(
     models = _parse_model_arg(model)
     is_dtwf = isinstance(models[0], DiscreteTimeWrightFisher)
     is_pedigree = any(isinstance(model, FixedPedigree) for model in models)
-    is_smck = any(isinstance(model, SmcKApproxCoalescent) for model in models)
 
     if record_full_arg:
         if coalescing_segments_only is not None:
@@ -938,12 +937,6 @@ def _parse_sim_ancestry(
     if discrete_genome and math.floor(sequence_length) != sequence_length:
         raise ValueError("Must have integer sequence length with discrete_genome=True")
 
-    if is_smck:
-        if gene_conversion_rate is not None or gene_conversion_tract_length is not None:
-            raise ValueError(
-                "Gene conversion has not been implemented yet for smc_k models."
-            )
-
     recombination_map = _parse_rate_map(
         recombination_rate, sequence_length, "recombination"
     )


Test patch:
diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c
index 321654809..1c4380256 100644
--- a/lib/tests/test_ancestry.c
+++ b/lib/tests/test_ancestry.c
@@ -3941,6 +3941,7 @@ test_setup_smc_k_plus(void)
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     ret = msp_set_simulation_model_smc_k(&msp, hull_offset);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL(msp.model.params.smc_k_coalescent.hull_offset, hull_offset);
 
     model = msp_get_model(&msp)->type;
     CU_ASSERT_EQUAL(model, MSP_MODEL_SMC_K);
@@ -4195,6 +4196,71 @@ test_fenwick_rebuild_smc_k(void)
     tsk_table_collection_free(&tables);
 }
 
+static void
+run_smc_k_gc_simulation(
+    double gc_rate, double tract_length, bool discrete_genome, double offset)
+{
+
+    int ret;
+    uint32_t n = 10;
+    sample_t *samples = malloc(n * sizeof(sample_t));
+    msp_t msp;
+    gsl_rng *rng = safe_rng_alloc();
+    tsk_table_collection_t tables;
+    tsk_treeseq_t ts;
+    long seed = 144;
+    double recombination_rate = 0.0;
+
+    gsl_rng_set(rng, seed);
+
+    CU_ASSERT_FATAL(&msp != NULL);
+    CU_ASSERT_FATAL(samples != NULL);
+    memset(samples, 0, n * sizeof(sample_t));
+    ret = build_sim(&msp, &tables, rng, 100, 1, samples, n);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = msp_set_recombination_rate(&msp, 0.01);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL_FATAL(msp_set_discrete_genome(&msp, discrete_genome), 0);
+    CU_ASSERT_EQUAL_FATAL(msp_set_recombination_rate(&msp, recombination_rate), 0);
+    CU_ASSERT_EQUAL_FATAL(msp_set_gene_conversion_rate(&msp, gc_rate), 0);
+    CU_ASSERT_EQUAL_FATAL(msp_set_gene_conversion_tract_length(&msp, tract_length), 0);
+    ret = msp_initialise(&msp);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = msp_set_simulation_model_smc_k(&msp, offset);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = msp_run(&msp, DBL_MAX, ULONG_MAX);
+    CU_ASSERT_EQUAL(ret, 0);
+    msp_verify(&msp, 0);
+    CU_ASSERT_TRUE(msp_is_completed(&msp));
+    CU_ASSERT_TRUE(msp_get_num_gene_conversion_events(&msp) > 0);
+
+    /* Make sure we can build a tree sequence out of the result */
+    ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    tsk_treeseq_free(&ts);
+
+    ret = msp_free(&msp);
+    CU_ASSERT_EQUAL(ret, 0);
+    gsl_rng_free(rng);
+    free(samples);
+    tsk_table_collection_free(&tables);
+}
+
+static void
+test_smc_k_gc(void)
+{
+    double tract_lengths[] = { 1.0, 1.3333, 5, 10 };
+    size_t j;
+
+    for (j = 0; j < sizeof(tract_lengths) / sizeof(double); j++) {
+        run_smc_k_gc_simulation(1.0, tract_lengths[j], true, 0.0);
+        run_smc_k_gc_simulation(1.0, tract_lengths[j], false, 0.75);
+        run_smc_k_gc_simulation(1.0, tract_lengths[j], true, 5.0);
+        run_smc_k_gc_simulation(1.0, tract_lengths[j], false, 3.1);
+    }
+}
+
 int
 main(int argc, char **argv)
 {
@@ -4294,6 +4360,7 @@ main(int argc, char **argv)
         { "test_mixed_model_smc_k", test_mixed_model_smc_k },
         { "test_mixed_model_smc_k_large", test_mixed_model_smc_k_large },
         { "test_fenwick_rebuild_smc_k", test_fenwick_rebuild_smc_k },
+        { "test_smc_k_gc", test_smc_k_gc },
         CU_TEST_INFO_NULL,
     };
 
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 4388f0c24..0a95abd78 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -443,6 +443,11 @@ def test_smck(self):
         for tree in ts.trees():
             assert tree.num_roots == 1
 
+        ts = self.run_script("10 -L 1000 -r 0.01 --model smc_k --offset 0.50")
+        assert ts.num_trees > 1
+        for tree in ts.trees():
+            assert tree.num_roots == 1
+
         ts = self.run_script("10 -L 1000 -d -r 0.01 --model smc_k -p 2 -g 0.1")
         assert ts.num_trees > 1
         for tree in ts.trees():
@@ -452,3 +457,8 @@ def test_smck(self):
         assert ts.num_trees > 1
         for tree in ts.trees():
             assert tree.num_roots == 1
+
+        ts = self.run_script("10 -L 1000 -c 0.04 2  --model smc_k --offset 0.75")
+        assert ts.num_trees > 1
+        for tree in ts.trees():
+            assert tree.num_roots == 1
diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py
index 11c057311..08f895e76 100644
--- a/tests/test_ancestry.py
+++ b/tests/test_ancestry.py
@@ -2959,7 +2959,8 @@ def test_model_switch_high_rec(self):
             assert tree.num_roots == 1
 
     @pytest.mark.parametrize("hull_offset", [2, 0.5, 1e-6, 2.133])
-    def test_smc_k_plus(self, hull_offset):
+    @pytest.mark.parametrize("discrete_genome", [True, False])
+    def test_smc_k_plus(self, hull_offset, discrete_genome):
         tss = msprime.sim_ancestry(
             samples=10,
             population_size=10_000,
@@ -2968,6 +2969,7 @@ def test_smc_k_plus(self, hull_offset):
             recombination_rate=1e-5,
             sequence_length=100,
             num_replicates=10,
+            discrete_genome=discrete_genome,
         )
         for ts in tss:
             assert ts.num_trees > 1
@@ -2990,15 +2992,21 @@ def test_two_pops(self):
         for tree in ts.trees():
             assert tree.num_roots == 1
 
-    def test_gc(self):
-        with pytest.raises(
-            ValueError,
-            match="Gene conversion has not been implemented yet for smc_k models.",
-        ):
-            _ = msprime.sim_ancestry(
-                samples=10,
-                model=msprime.SmcKApproxCoalescent(),
-                sequence_length=100,
-                gene_conversion_rate=0.5,
-                gene_conversion_tract_length=5,
-            )
+    @pytest.mark.parametrize("hull_offset", [2, 0.5, 1e-6, 2.133])
+    @pytest.mark.parametrize("discrete_genome", [True, False])
+    @pytest.mark.parametrize("recombination_rate", [0.0, 0.1])
+    def test_gc(self, hull_offset, discrete_genome, recombination_rate):
+        tss = msprime.sim_ancestry(
+            samples=10,
+            model=msprime.SmcKApproxCoalescent(hull_offset=hull_offset),
+            sequence_length=100,
+            gene_conversion_rate=1.0,
+            gene_conversion_tract_length=5,
+            num_replicates=5,
+            discrete_genome=discrete_genome,
+            recombination_rate=recombination_rate,
+        )
+        for ts in tss:
+            assert ts.num_trees > 1
+            for tree in ts.trees():
+                assert tree.num_roots == 1