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

Issue:
Multiple merger polyploid scaling
Ping @jeromekelleher. Adding tests for polyploid scaling for multiple merger coalescents turned out to be easy, but also sent me down a bit of a rabbit hole. Summary of changes:

- The extra spaces on various lines in `msprime.c` which are unrelated to anything I've done (1396, 1899...) were put in by clang. Let me know if you want me to do something about them.
- For the Dirac coalescent, population growth at rate `g` should yield a coalescent growth rate of `2g`,essentially because the coalescence probability in the Moran model is `1/N^2` rather than `1/N`. We had overlooked that, but it is now fixed. 
- Ploidy wasn't being handled quite right in the Dirac coalescent either, but is now fixed. In the process, I got rid of some `if (ploidy == 1) {...} else {...}` statements which became redundant.
- For both the Dirac and Beta coalescents, I moved the allocation of the ancestor buckets `Q` later into the rejection sampler. It is now only done once a multiple merger event is definitely taking place.
- The Dirac coalescent test in `test_ancestry.py` started failing after the ploidy fix because the rate of multiple mergers wasn't high enough to guarantee one anymore. I increased it.
- I also noticed we were disallowing `psi = 1` in the Dirac coalescent for no reason. The documented (and sensible) range is `0 < psi <= 1`. I tweaked the boundary values and associated tests to make the code consistent with the documentation.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed5bd0ee3..b4030913f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## [1.3.3] - 2024-XX-XX
+
+**Bug fixes**:
+
+- Correct the Dirac coalescent time scaling with polyploidy and population growth.
+
 ## [1.3.2] - 2024-07-08
 
 - Add `record_provenance` argument to `sim_mutations` ({issue}`2272`, {pr}`2273`, {user}`peterlharp`)
diff --git a/lib/msprime.c b/lib/msprime.c
index b7b2b30ea..2316fb00b 100644
--- a/lib/msprime.c
+++ b/lib/msprime.c
@@ -7949,33 +7949,18 @@ msp_dirac_get_common_ancestor_waiting_time_from_rate(
     msp_t *self, population_t *pop, double lambda)
 {
     double ret = DBL_MAX;
-    double alpha = pop->growth_rate;
+    double alpha = 2.0 * pop->growth_rate;
     double t = self->time;
     double u, dt, z;
+    double pop_size = pop->initial_size;
 
     if (lambda > 0.0) {
         u = gsl_ran_exponential(self->rng, 1.0 / lambda);
         if (alpha == 0.0) {
-            if (self->ploidy == 1) {
-                ret = pop->initial_size * pop->initial_size * u;
-            } else {
-                /* For ploidy > 1 we assume N/2 two-parent families, so that the rate
-                 * with which 2 lineages belong to a common family is (2/N)^2 */
-                ret = pop->initial_size * pop->initial_size * u / 4.0;
-            }
+            ret = pop_size * pop_size * u;
         } else {
             dt = t - pop->start_time;
-            if (self->ploidy == 1) {
-                z = 1
-                    + alpha * pop->initial_size * pop->initial_size * exp(-alpha * dt)
-                          * u;
-            } else {
-                /* For ploidy > 1 we assume N/2 two-parent families, so that the rate
-                 * with which 2 lineages belong to a common family is (2/N)^2 */
-                z = 1
-                    + alpha * pop->initial_size * pop->initial_size * exp(-alpha * dt)
-                          * u / 4.0;
-            }
+            z = 1 + alpha * pop_size * pop_size * exp(-alpha * dt) * u;
             /* if z is <= 0 no coancestry can occur */
             if (z > 0) {
                 ret = log(z) / alpha;
@@ -7994,13 +7979,7 @@ msp_dirac_get_common_ancestor_waiting_time(
 {
     population_t *pop = &self->populations[pop_id];
     unsigned int n = (unsigned int) avl_count(&pop->ancestors[label]);
-    double c = self->model.params.dirac_coalescent.c;
-    double lambda = n * (n - 1.0) / 2.0;
-    if (self->ploidy == 1) {
-        lambda += c;
-    } else {
-        lambda += c / (2.0 * self->ploidy);
-    }
+    double lambda = gsl_sf_choose(n, 2) + self->model.params.dirac_coalescent.c;
 
     return msp_dirac_get_common_ancestor_waiting_time_from_rate(self, pop, lambda);
 }
@@ -8018,71 +7997,64 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t
     double nC2, p;
     double psi = self->model.params.dirac_coalescent.psi;
 
-    /* We assume haploid reproduction is single-parent, while all other ploidies
-     * are two-parent */
-    if (self->ploidy == 1) {
-        num_parental_copies = 1;
-    } else {
-        num_parental_copies = 2 * self->ploidy;
-    }
-    Q = tsk_malloc(num_parental_copies * sizeof(*Q));
-    if (Q == NULL) {
-        ret = MSP_ERR_NO_MEMORY;
-        goto out;
-    }
-
     ancestors = &self->populations[pop_id].ancestors[label];
     n = avl_count(ancestors);
     nC2 = gsl_sf_choose(n, 2);
-    if (self->ploidy == 1) {
-        p = (nC2 / (nC2 + self->model.params.dirac_coalescent.c));
-    } else {
-        p = (nC2 / (nC2 + self->model.params.dirac_coalescent.c / (2.0 * self->ploidy)));
-    }
+    p = nC2 / (nC2 + self->model.params.dirac_coalescent.c);
+
     if (gsl_rng_uniform(self->rng) < p) {
-        /* When 2 * ploidy parental chromosomes are available, Mendelian segregation
-         * results in a merger only 1 / (2 * ploidy) of the time. */
-        if (self->ploidy == 1
-            || gsl_rng_uniform(self->rng) < 1.0 / (2.0 * self->ploidy)) {
-            /* Choose x and y */
-            n = avl_count(ancestors);
-            j = (uint32_t) gsl_rng_uniform_int(self->rng, n);
-            x_node = avl_at(ancestors, j);
-            tsk_bug_assert(x_node != NULL);
-            x_lin = (lineage_t *) x_node->item;
-            x = x_lin->head;
-            avl_unlink_node(ancestors, x_node);
-            j = (uint32_t) gsl_rng_uniform_int(self->rng, n - 1);
-            y_node = avl_at(ancestors, j);
-            tsk_bug_assert(y_node != NULL);
-            y_lin = (lineage_t *) y_node->item;
-            y = y_lin->head;
-            avl_unlink_node(ancestors, y_node);
-            self->num_ca_events++;
-            msp_free_avl_node(self, x_node);
-            msp_free_lineage(self, x_lin);
-            msp_free_avl_node(self, y_node);
-            msp_free_lineage(self, y_lin);
-            ret = msp_merge_two_ancestors(self, pop_id, label, x, y, TSK_NULL, NULL);
-        }
+        /* Choose x and y */
+        n = avl_count(ancestors);
+        j = (uint32_t) gsl_rng_uniform_int(self->rng, n);
+        x_node = avl_at(ancestors, j);
+        tsk_bug_assert(x_node != NULL);
+        x_lin = (lineage_t *) x_node->item;
+        x = x_lin->head;
+        avl_unlink_node(ancestors, x_node);
+        j = (uint32_t) gsl_rng_uniform_int(self->rng, n - 1);
+        y_node = avl_at(ancestors, j);
+        tsk_bug_assert(y_node != NULL);
+        y_lin = (lineage_t *) y_node->item;
+        y = y_lin->head;
+        avl_unlink_node(ancestors, y_node);
+        self->num_ca_events++;
+        msp_free_avl_node(self, x_node);
+        msp_free_lineage(self, x_lin);
+        msp_free_avl_node(self, y_node);
+        msp_free_lineage(self, y_lin);
+        ret = msp_merge_two_ancestors(self, pop_id, label, x, y, TSK_NULL, NULL);
     } else {
-        for (j = 0; j < num_parental_copies; j++) {
-            avl_init_tree(&Q[j], cmp_segment_queue, NULL);
-        }
         num_participants = gsl_ran_binomial(self->rng, psi, n);
-        ret = msp_multi_merger_common_ancestor_event(
-            self, ancestors, Q, num_participants, num_parental_copies);
-        if (ret < 0) {
-            goto out;
-        }
-        /* All the lineages that have been assigned to the particular pots can now be
-         * merged.
-         */
-        for (j = 0; j < num_parental_copies; j++) {
-            ret = msp_merge_ancestors(self, &Q[j], pop_id, label, TSK_NULL, NULL);
+        if (num_participants > 1) {
+            /* We assume haploid reproduction is single-parent, while all other ploidies
+             * are two-parent */
+            if (self->ploidy == 1) {
+                num_parental_copies = 1;
+            } else {
+                num_parental_copies = 2 * self->ploidy;
+            }
+            Q = tsk_malloc(num_parental_copies * sizeof(*Q));
+            if (Q == NULL) {
+                ret = MSP_ERR_NO_MEMORY;
+                goto out;
+            }
+            for (j = 0; j < num_parental_copies; j++) {
+                avl_init_tree(&Q[j], cmp_segment_queue, NULL);
+            }
+            ret = msp_multi_merger_common_ancestor_event(
+                self, ancestors, Q, num_participants, num_parental_copies);
             if (ret < 0) {
                 goto out;
             }
+            /* All the lineages that have been assigned to the particular pots can now be
+             * merged.
+             */
+            for (j = 0; j < num_parental_copies; j++) {
+                ret = msp_merge_ancestors(self, &Q[j], pop_id, label, TSK_NULL, NULL);
+                if (ret < 0) {
+                    goto out;
+                }
+            }
         }
     }
 out:
@@ -8242,22 +8214,6 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l
     double truncation_point = beta_compute_truncation(self);
     double beta_x, u, increment;
 
-    /* We assume haploid reproduction is single-parent, while all other ploidies
-     * are two-parent */
-    if (self->ploidy == 1) {
-        num_parental_copies = 1;
-    } else {
-        num_parental_copies = 2 * self->ploidy;
-    }
-    Q = tsk_malloc(num_parental_copies * sizeof(*Q));
-    if (Q == NULL) {
-        ret = MSP_ERR_NO_MEMORY;
-        goto out;
-    }
-
-    for (j = 0; j < num_parental_copies; j++) {
-        avl_init_tree(&Q[j], cmp_segment_queue, NULL);
-    }
     ancestors = &self->populations[pop_id].ancestors[label];
     n = avl_count(ancestors);
     beta_x = ran_inc_beta(self->rng, 2.0 - alpha, alpha, truncation_point);
@@ -8295,6 +8251,22 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l
             num_participants = 2 + gsl_ran_binomial(self->rng, beta_x, n - 2);
         } while (gsl_rng_uniform(self->rng) > 1 / gsl_sf_choose(num_participants, 2));
 
+        /* We assume haploid reproduction is single-parent, while all other ploidies
+         * are two-parent */
+        if (self->ploidy == 1) {
+            num_parental_copies = 1;
+        } else {
+            num_parental_copies = 2 * self->ploidy;
+        }
+        Q = tsk_malloc(num_parental_copies * sizeof(*Q));
+        if (Q == NULL) {
+            ret = MSP_ERR_NO_MEMORY;
+            goto out;
+        }
+
+        for (j = 0; j < num_parental_copies; j++) {
+            avl_init_tree(&Q[j], cmp_segment_queue, NULL);
+        }
         ret = msp_multi_merger_common_ancestor_event(
             self, ancestors, Q, num_participants, num_parental_copies);
         if (ret < 0) {
diff --git a/msprime/_msprimemodule.c b/msprime/_msprimemodule.c
index 6b3954f02..b3de9398e 100644
--- a/msprime/_msprimemodule.c
+++ b/msprime/_msprimemodule.c
@@ -1199,8 +1199,8 @@ Simulator_parse_simulation_model(Simulator *self, PyObject *py_model)
             goto out;
         }
         c = PyFloat_AsDouble(value);
-        if (psi <= 0 || psi >= 1.0) {
-            PyErr_SetString(PyExc_ValueError, "Must have 0 < psi < 1");
+        if (psi <= 0 || psi > 1.0) {
+            PyErr_SetString(PyExc_ValueError, "Must have 0 < psi <= 1");
             goto out;
         }
         if (c < 0) {
diff --git a/verification.py b/verification.py
index dd7ca5b13..ea3a6ca40 100644
--- a/verification.py
+++ b/verification.py
@@ -3431,12 +3431,14 @@ def _run(self, pop_size, alpha, growth_rate, num_replicates=10000):
         logging.debug(f"running Beta growth for {pop_size} {alpha} {growth_rate}")
         b = growth_rate * (alpha - 1)
         model = (msprime.BetaCoalescent(alpha=alpha),)
-        ploidy = 2
-        a = 1 / (2 * ploidy * self.compute_beta_timescale(pop_size, alpha, ploidy))
-        name = f"N={pop_size}_alpha={alpha}_growth_rate={growth_rate}_ploidy={ploidy}"
-        self.compare_tmrca(
-            pop_size, growth_rate, model, num_replicates, a, b, ploidy, name
-        )
+        for ploidy in range(2, 7):
+            a = 1 / (2 * ploidy * self.compute_beta_timescale(pop_size, alpha, ploidy))
+            name = (
+                f"N={pop_size}_alpha={alpha}_growth_rate={growth_rate}_ploidy={ploidy}"
+            )
+            self.compare_tmrca(
+                pop_size, growth_rate, model, num_replicates, a, b, ploidy, name
+            )
         ploidy = 1
         a = 1 / self.compute_beta_timescale(pop_size, alpha, ploidy)
         name = f"N={pop_size}_alpha={alpha}_growth_rate={growth_rate}_ploidy={ploidy}"
@@ -3474,12 +3476,14 @@ def test_100000_11_001(self):
 class DiracGrowth(XiGrowth):
     def _run(self, pop_size, c, psi, growth_rate, num_replicates=10000):
         logging.debug(f"running Dirac growth for {pop_size} {c} {psi} {growth_rate}")
-        b = growth_rate
+        b = 2 * growth_rate
         model = (msprime.DiracCoalescent(psi=psi, c=c),)
-        p = 2
-        a = (1 + c * psi * psi / (2 * p)) / (pop_size * pop_size)
-        name = f"N={pop_size}_c={c}_psi={psi}_growth_rate={growth_rate}_ploidy={p}"
-        self.compare_tmrca(pop_size, growth_rate, model, num_replicates, a, b, p, name)
+        for p in range(2, 7):
+            a = (1 + c * psi * psi / (2 * p)) / (pop_size * pop_size)
+            name = f"N={pop_size}_c={c}_psi={psi}_growth_rate={growth_rate}_ploidy={p}"
+            self.compare_tmrca(
+                pop_size, growth_rate, model, num_replicates, a, b, p, name
+            )
         p = 1
         a = (1 + c * psi * psi) / (pop_size * pop_size)
         name = f"N={pop_size}_c={c}_psi={psi}_growth_rate={growth_rate}_ploidy={p}"


Test patch:
diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py
index 08f895e76..14eacb403 100644
--- a/tests/test_ancestry.py
+++ b/tests/test_ancestry.py
@@ -206,7 +206,7 @@ def test_multimerger_dirac(self):
             recombination_rate=0.1,
             record_full_arg=True,
             sequence_length=10,
-            model=msprime.DiracCoalescent(psi=0.5, c=5),
+            model=msprime.DiracCoalescent(psi=0.5, c=100),
         )
         self.verify(sim, multiple_mergers=True)
 
diff --git a/tests/test_lowlevel.py b/tests/test_lowlevel.py
index 8b33e4a52..73d2812e5 100644
--- a/tests/test_lowlevel.py
+++ b/tests/test_lowlevel.py
@@ -1353,7 +1353,7 @@ def test_dirac_simulation_model(self):
                 make_sim(model=model)
         with pytest.raises(ValueError):
             make_sim(model=get_simulation_model("dirac"))
-        for bad_psi in [-1, 0, -1e-6, 1, 1e6]:
+        for bad_psi in [-1, 0, -1e-6, 1 + 1e-6, 1e6]:
             with pytest.raises(ValueError):
                 make_sim(
                     model=get_simulation_model("dirac", c=1, psi=bad_psi),
@@ -1363,7 +1363,7 @@ def test_dirac_simulation_model(self):
                 make_sim(
                     model=get_simulation_model("dirac", psi=0.5, c=bad_c),
                 )
-        for psi in [0.99, 0.2, 1e-4]:
+        for psi in [1, 0.2, 1e-4]:
             for c in [5.0, 1e2, 1e-4]:
                 model = get_simulation_model("dirac", psi=psi, c=c)
                 sim = make_sim(model=model)