Repository: popsim-consortium/stdpopsim (language: python)

Issue:
SLiM and msprime engines handle population sizes differently for haploid species
In #1361, ploidy was added as an attribute of species, but the SLiM engine was left untouched (aside from setting the number of sampled individuals appropriately). However, this means that population size is treated differently between the msprime and SLiM engines for haploid species:
```python
import stdpopsim
import numpy as np

contig = stdpopsim.Contig.basic_contig(length=1, ploidy=1)
model = stdpopsim.PiecewiseConstantSize(1000)
reps = 100 

engine = stdpopsim.get_engine("msprime")
tmrca_msp = []
for _ in range(reps):
    ts_msp = engine.simulate(model, contig, {"pop_0":2})
    tmrca_msp.append(ts_msp.first().tmrca(0,1))
print(f"pair TMRCA (msprime): {np.mean(tmrca_msp)} +/- {np.std(tmrca_msp)/np.sqrt(reps)}")

engine = stdpopsim.get_engine("slim")
tmrca_slim = []
for _ in range(reps):
    ts_slim = engine.simulate(model, contig, {"pop_0":2})
    tmrca_slim.append(ts_slim.first().tmrca(0,1))
print(f"pair TMRCA (slim): {np.mean(tmrca_slim)} +/- {np.std(tmrca_slim)/np.sqrt(reps)}")      

# pair TMRCA (msprime): 1000.116049495669 +/- 93.04547569055264
# pair TMRCA (slim): 2274.7033717826785 +/- 192.0079953821991
```
The solution should be to rescale the demographic model's population size inside `slim_makescript` if the contig is haploid. But, this raises the question, are the population sizes for EscCol / StrAga in the catalog actually # haploids, or are they off by a factor of 2? 


Discussion / hints:

Fix patch:
diff --git a/docs/development.rst b/docs/development.rst
index c347cd9dc..a53374ba4 100644
--- a/docs/development.rst
+++ b/docs/development.rst
@@ -575,6 +575,7 @@ with a brief discussion of possible courses of action to take when components ha
 
 4. The **effective population size** should represent the historical average effective population size,
    and should ideally produce simulated data that matches the average observed genetic diversity in that species.
+   Population size is defined as the number of individuals, regardless of ploidy.
    However, this will often not capture features of genetic variation that are caused by recent changes in population size and the presence of population structure.
    To capture those, one should also provide a demographic model (or multiple models) for the species
    (see `Adding a new demographic model`_).
@@ -1171,6 +1172,7 @@ The final three attributes
 describe the inferred demographic history that you wish to code.
 This history consists of ancestral population size changes,
 migration rates, split times, and admixture events.
+Note that population size is defined as the number of individuals, regardless of ploidy.
 These attributes should be coded using the standard format of ``msprime``.
 If this is your first time specifying a demographic model using ``msprime``,
 then we highly recommend that you take some time to read through its
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index 6de3adf99..69b134d3e 100644
--- a/stdpopsim/slim_engine.py
+++ b/stdpopsim/slim_engine.py
@@ -900,7 +900,10 @@ def fix_time(event):
     growth_rates = np.empty(shape=(dd.num_populations, len(epochs)), dtype=float)
     for j, epoch in enumerate(epochs):
         for i, pop in enumerate(epoch.populations):
-            N[i, j] = int(pop.end_size)
+            # SLiM simulates a diploid population, so rescale population size
+            # depending on contig ploidy. If/when the SLiM simulation is
+            # converted to a haploid model, this rescaling should be removed.
+            N[i, j] = int(pop.end_size * contig.ploidy / 2)
             growth_rates[i, j] = pop.growth_rate
 
     admixture_pulses = []
@@ -1693,8 +1696,8 @@ def _recap_and_rescale(
     ):
         """
         Apply post-SLiM transformations to ``ts``. This rescales node times,
-        does recapitation, simplification, adds neutral mutations, and converts
-        alleles to nucleotides.
+        does recapitation, simplification, adds neutral mutations, converts
+        alleles to nucleotides, and rebuilds the individual table for haploids.
         """
         # Times come from SLiM generation numbers, which may have been
         # divided by a scaling factor for computational tractability.
@@ -1759,10 +1762,14 @@ def _recap_and_rescale(
             migration_matrix=recap_epoch.migration_matrix,
         )
 
+        # `recap_epoch` contains population sizes from the demographic model,
+        # that are the number of individuals regardless of ploidy. Thus,
+        # ploidy must be set here, as is done in `_MsprimeEngine.simulate`
         ts = msprime.sim_ancestry(
             initial_state=ts,
             demography=demography,
             recombination_rate=contig.recombination_map,
+            ploidy=contig.ploidy,
             random_seed=s0,
         )
         ts = self._simplify_remembered(ts)
@@ -1833,6 +1840,10 @@ def _get_msp_rate_map(breaks, is_this_dfe, rate):
                 )
             )
 
+        # If haploid, rebuild individual table so each sample is an individual
+        if contig.ploidy == 1:
+            ts = stdpopsim.utils.haploidize_individuals(ts)
+
         return ts
 
     def recap_and_rescale(
diff --git a/stdpopsim/utils.py b/stdpopsim/utils.py
index d6e8c106d..8fadb8a9d 100644
--- a/stdpopsim/utils.py
+++ b/stdpopsim/utils.py
@@ -308,3 +308,22 @@ def parse_population_sample_pairs(pop_sample_pairs):
             "Sample specification must be in the form "
             "<population_name:number_of_samples>"
         )
+
+
+def haploidize_individuals(ts):
+    """
+    Rebuild IndividualsTable of tree sequence such that each individual is
+    referenced by only one sample, by duplicating individuals that are
+    referenced by multiple samples in the original tree sequence.
+    """
+    tables = ts.dump_tables()
+    old_inds = tables.individuals.copy()
+    tables.individuals.clear()
+    # Make sure we don't refer to any individuals that are not samples
+    node_indiv = np.zeros_like(ts.nodes_individual) - 1
+    for new_idx, node in enumerate(ts.samples()):
+        old_idx = ts.nodes_individual[node]
+        tables.individuals.append(old_inds[old_idx])
+        node_indiv[node] = new_idx
+    tables.nodes.individual = node_indiv
+    return tables.tree_sequence()


Test patch:
diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py
index 6226fde5b..d0e30a562 100644
--- a/tests/test_slim_engine.py
+++ b/tests/test_slim_engine.py
@@ -11,6 +11,7 @@
 from numpy.testing import assert_array_equal
 import collections
 import re
+import logging
 
 import pytest
 import tskit
@@ -2885,3 +2886,104 @@ def test_errors(self):
             stdpopsim.ext.selection_coeff_from_mutation("foo", next(ts.mutations()))
         with pytest.raises(ValueError, match="must be a"):
             stdpopsim.ext.selection_coeff_from_mutation(ts, "bar")
+
+
+@pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
+class TestPloidy:
+    """
+    Test that population sizes used in SLiM engine are scaled correctly
+    with regards to ploidy.
+
+    - Check that (diploid) population size used in SLiM simulation equals
+    model population size * ploidy / 2.
+
+    - Check that `recap_epoch` returned by `slim_makescript`, used to parameterize
+    recapitation with msprime, contains population size from the original model
+    (e.g. # of individuals).
+
+    - Check that individuals in tree sequence are haploid/diploid.
+    """
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    @pytest.mark.usefixtures("caplog")
+    def test_slim_population_size_haploid(self, caplog):
+        N = 100
+        engine = stdpopsim.get_engine("slim")
+        contig = stdpopsim.Contig.basic_contig(length=1000, ploidy=1)
+        model = stdpopsim.PiecewiseConstantSize(N)
+        with caplog.at_level(logging.DEBUG):
+            engine.simulate(model, contig, samples={"pop_0": 2}, verbosity=2)
+        log_str = " ".join([rec.getMessage() for rec in caplog.records])
+        # match: "1: p = sim.addSubpop(0, <SLiM population size>);"
+        extract_ne = re.compile(".+1: p = sim.addSubpop\\(0, ([0-9]+)\\).+")
+        match = extract_ne.match(log_str)
+        assert match is not None
+        (Nslim,) = match.groups()
+        assert int(Nslim) == int(N / 2)
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    @pytest.mark.usefixtures("caplog")
+    def test_slim_population_size_diploid(self, caplog):
+        N = 100
+        engine = stdpopsim.get_engine("slim")
+        contig = stdpopsim.Contig.basic_contig(length=1000, ploidy=2)
+        model = stdpopsim.PiecewiseConstantSize(N)
+        with caplog.at_level(logging.DEBUG):
+            engine.simulate(model, contig, samples={"pop_0": 2}, verbosity=2)
+        log_str = " ".join([rec.getMessage() for rec in caplog.records])
+        # match: "1: p = sim.addSubpop(0, <SLiM population size>);"
+        extract_ne = re.compile(".+1: p = sim.addSubpop\\(0, ([0-9]+)\\).+")
+        match = extract_ne.match(log_str)
+        assert match is not None
+        (Nslim,) = match.groups()
+        assert int(Nslim) == int(N)
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_recap_population_size(self):
+        N = 100
+        model = stdpopsim.PiecewiseConstantSize(N)
+        for ploidy in [1, 2]:
+            contig = stdpopsim.Contig.basic_contig(length=1000, ploidy=ploidy)
+            sample_sets = model.get_sample_sets({"pop_0": 2}, ploidy=contig.ploidy)
+            rate_map = stdpopsim.get_slim_mutation_rate_map(contig)
+            with open(os.devnull, "w") as scriptfile:
+                recap_epoch = stdpopsim.slim_makescript(
+                    scriptfile,
+                    "unused",
+                    model,
+                    contig,
+                    sample_sets,
+                    extended_events=None,
+                    scaling_factor=1,
+                    burn_in=0,
+                    slim_rate_map=rate_map,
+                )
+            assert int(recap_epoch.populations[0].start_size) == N
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_individual_ploidy(self):
+        N = 100
+        model = stdpopsim.PiecewiseConstantSize(N)
+        engine = stdpopsim.get_engine("slim")
+        for ploidy in [1, 2]:
+            contig = stdpopsim.Contig.basic_contig(length=1000, ploidy=ploidy)
+            ts = engine.simulate(model, contig, samples={"pop_0": 2})
+            assert ts.num_individuals == 2
+            assert ts.num_samples == 2 * ploidy
+            individual = ts.tables.nodes.individual
+            assert len(np.unique(individual[: ts.num_samples])) == ts.num_individuals
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_haploidize_individuals(self):
+        N = 100
+        model = stdpopsim.PiecewiseConstantSize(N)
+        engine = stdpopsim.get_engine("slim")
+        contig = stdpopsim.Contig.basic_contig(length=1000, ploidy=2)
+        ts = engine.simulate(model, contig, samples={"pop_0": 3})
+        ts_hap = stdpopsim.utils.haploidize_individuals(ts)
+        assert ts_hap.num_individuals == ts.num_individuals * 2
+        for i, j in zip(ts.samples(), ts_hap.samples()):
+            assert i == j
+            ind_i = ts.nodes_individual[i]
+            ind_j = ts_hap.nodes_individual[j]
+            assert ts.tables.individuals[ind_i] == ts_hap.tables.individuals[ind_j]