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

Issue:
should SLiM follow msprime in doing Jukes-Cantor mutations?
Since msprime 1.0, `sim_mutations( )` has by default produced JC mutations (and hence, nucelotide mutations). Questions:

1. Do we want to do JC for msprime mutations? (I think yes?)
2. Do we want to do JC for *SLiM* mutations? (I also think yes?)
3. Do we want to convert the output of SLiM simulations so that the ancestral and derived states are nucleotides rather than lists of SLiM IDs (which produce technically invalid and at least weird VCFs)? We'd do this with [`pyslim.convert_alleles( )`](https://tskit.dev/pyslim/docs/latest/python_api.html#pyslim.convert_alleles). (I'm not sure? maybe by default but with an optional parameter?) This would mean that the tree sequences that get output *can't* be loaded back into SLiM (unless there's an optional parameter that is toggled).


Discussion / hints:
We discussed this in the meeting this week, and I think the answer is "yes" to all of these. Do you see any issue with this, @bhaller? @jeromekelleher?
I'd say yes - the majority of users won't know anything about SLiM and will expect to get VCFs out the other end with ACGT alleles.
It generally makes sense to me.  The only gotcha is that an ancestral state of A and a mutation back to A are not the same thing in SLiM, nor are two different mutational lineages to T at the same position that arose independently.  So you may want to think about various "uniquing down" strategies if you want stdpopsim to really make complete biological sense when selection is involved – IBS rather than IBD.  The SLiM manual has quite a bit of discussion about this in version 4.0, including example recipes for various types of uniquing down, etc.  In my opinion it's a weakness in SLiM's design, and maybe I will try to address it at some point.  For now, users have to worry about it.  Besides that, yes – getting ACGT out from stdpopsim by default makes good sense.
Hm - looking at the details, it's going to quite annoying to actually shift the script to being a nucleotide model. HOWEVER, since all we want is Jukes-Cantor mutations, we could equivalently just call [`pyslim.generate_nucleotides`](https://tskit.dev/pyslim/docs/latest/python_api.html#pyslim.generate_nucleotides) on the tree sequence afterwards. I think this is definitely what we want to do, since changing the SLiM code would also require re-doing how the testing code works; probably in the future someone will want to mess around with nucleotide mutations more, and can get into the details at that point.

My only question about this is how slow the calls to `generate_nucleotides` and `convert_alleles` will be - they both loop over the mutations and do nontrivial things in python.
Hmm.  At that point, it sounds to me like maybe this ought to be optional, or just left up to the user to do if they want to.  If it's a time-consuming operation that the user may or may not want, and if it's easy for them to do it themselves, then maybe just leave it up to them, and document how they would do it?  Or at a minimum I'd suggest making it the default behavior but having a way to turn it off, for users who either (a) don't even want nucleotides, or (b) don't want to pay the price in terms of speed?
We should see how time-consuming it is, but yes.
Hmm, `generate_nucleotides` and `convert_alleles` seem prohibitively expensive to be used by default ... for example,
```python
import stdpopsim
import pyslim
import time

engine = stdpopsim.get_engine("slim")
species = stdpopsim.get_species("HomSap")
contig = species.get_contig("chr1", genetic_map="HapMapII_GRCh37")
model = stdpopsim.PiecewiseConstantSize(100)
samples = model.get_samples(10)

start_sim = time.time()
ts = engine.simulate(
    demographic_model=model,
    contig=contig,
    samples=samples,
    slim_burn_in=0.1,
)
print(f"engine.simulate {time.time() - start_sim} seconds")

start_gen = time.time()
ts = pyslim.generate_nucleotides(ts)
print(f"pyslim.generate_nucleotides {time.time() - start_gen} seconds")

start_conv = time.time()
ts = pyslim.convert_alleles(ts)
print(f"pyslim.convert_alleles {time.time() - start_conv} seconds")

print(f"num_mutations {ts.num_mutations}")
```

```python
#engine.simulate 2.703385829925537 seconds
#pyslim.generate_nucleotides 533.7081677913666 seconds
#pyslim.convert_alleles 520.3803863525391 seconds
#num_mutations 3288
```

OMG. Well, this is probably something we should look at doing more efficiently, somehow. Let's leave this open.
OK, over at https://github.com/tskit-dev/pyslim/pull/299 we're making `convert_alleles` go fast. And, if we use nucleotide models in SLiM, then we don't need `generate_nucleotides`. (I'm sure the latter can go faster, but we don't yet need it to?) I'll do a pyslim release with the faster version once it's sorted.

Fix patch:
diff --git a/requirements/CI/requirements.txt b/requirements/CI/requirements.txt
index 6e2201c49..9dd8e31ff 100644
--- a/requirements/CI/requirements.txt
+++ b/requirements/CI/requirements.txt
@@ -14,7 +14,7 @@ msprime==1.2.0
 attrs==21.4.0
 appdirs==1.4.4
 humanize==4.4.0
-pyslim==1.0
+pyslim==1.0.1
 pandas==1.3.4
 numpy==1.21.5; python_version=='3.7'
 numpy==1.22.3; python_version>='3.8'
@@ -23,3 +23,4 @@ zarr==2.11.3
 biopython==1.79
 demes==0.2.2
 demesdraw==0.3.1
+kastore==0.3.2
diff --git a/requirements/development.txt b/requirements/development.txt
index 383ceeeb6..2618f2850 100644
--- a/requirements/development.txt
+++ b/requirements/development.txt
@@ -18,7 +18,7 @@ attrs
 appdirs
 humanize
 pre-commit
-pyslim>=0.600
+pyslim>=1.0.1
 pandas
 numpy
 scikit-allel
@@ -27,3 +27,4 @@ biopython
 demes==0.2.0
 demesdraw==0.2.0
 boto3
+kastore
diff --git a/stdpopsim/cli.py b/stdpopsim/cli.py
index c81cefe98..89b245c82 100644
--- a/stdpopsim/cli.py
+++ b/stdpopsim/cli.py
@@ -593,6 +593,14 @@ def add_simulate_species_parser(parser, species):
             "For msprime, seeds must be > 0 and < 2^32."
         ),
     )
+    species_parser.add_argument(
+        "--keep-mutation-ids-as-alleles",
+        action="store_true",
+        help=(
+            "If set, mutations in SLiM simulations will have alleles coded as "
+            "numeric ids rather than random nucleotides."
+        ),
+    )
 
     model_help = (
         "Specify a simulation model. If no model is specified, a single population "
@@ -724,10 +732,16 @@ def run_simulation(args):
                 f"Cannot sample from more than {model.num_sampling_populations} "
                 "populations"
             )
+        if args.keep_mutation_ids_as_alleles and args.engine != "slim":
+            exit(
+                "Option --keep-mutation-ids-as-alleles only applies to the SLiM "
+                "engine."
+            )
 
         samples = stdpopsim.utils.parse_population_sample_pairs(args.samples)
         if isinstance(samples, list):
-            # FIXME: for back-compatibility with deprecated positional sample counts
+            # TODO: For back-compatibility with deprecated positional sample
+            # counts, eventually remove
             samples = model.get_samples(*samples)
         else:
             assert isinstance(samples, dict)
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index 4d7ae0838..58189b6c1 100644
--- a/stdpopsim/slim_engine.py
+++ b/stdpopsim/slim_engine.py
@@ -1197,6 +1197,12 @@ def matrix2str(
             if not m.convert_to_substitution:
                 # T is the default for WF simulations.
                 printsc(f"    m{mid}.convertToSubstitution = F;")
+            # TODO: when msprime.SLiMMutationModel supports stacking policy,
+            # set policy such that there's at most a single mutation per-site
+            # and individual
+            # printsc(f"    m{mid}.mutationStackGroup = 0;")
+            # printsc(f"    m{mid}.mutationStackPolicy = 'l';")
+        mut_types = ", ".join([str(mt) for mt in mut_type_list])
         mut_types = ", ".join([str(mt) for mt in mut_type_list])
         # We will assign a proportion of 0.0 to mutation types that are neutral
         # unless all the mutation types in a DFE are neutral. In such case,
@@ -1462,6 +1468,7 @@ def simulate(
         verbosity=None,
         logfile=None,
         logfile_interval=100,
+        keep_mutation_ids_as_alleles=False,
         _recap_and_rescale=True,
     ):
         """
@@ -1496,6 +1503,10 @@ def simulate(
             (currently, only mean and SD of fitness values per population).
             Defaults to None, meaning "do not log".
         :param logfile_interval: How often to write to the log file, in generations.
+        :param keep_mutation_ids_as_alleles: If true, alleles will be coded by integer
+            mutation ids (assigned by SLiM) rather than by randomly-generated
+            nucleotides.
+        :type keep_mutation_ids_as_alleles: bool
         """
 
         if slim_scaling_factor <= 0:
@@ -1585,7 +1596,13 @@ def script_file_f():
         ts = _add_dfes_to_metadata(ts, contig)
         if _recap_and_rescale:
             ts = self._recap_and_rescale(
-                ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events
+                ts,
+                seed,
+                recap_epoch,
+                contig,
+                slim_scaling_factor,
+                keep_mutation_ids_as_alleles,
+                extended_events,
             )
 
         if contig.inclusion_mask is not None:
@@ -1665,11 +1682,19 @@ def _simplify_remembered(self, ts):
         return ts.simplify(samples=list(nodes), filter_populations=False)
 
     def _recap_and_rescale(
-        self, ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events=None
+        self,
+        ts,
+        seed,
+        recap_epoch,
+        contig,
+        slim_scaling_factor,
+        keep_mutation_ids_as_alleles,
+        extended_events=None,
     ):
         """
         Apply post-SLiM transformations to ``ts``. This rescales node times,
-        does recapitation, simplification, and adds neutral mutations.
+        does recapitation, simplification, adds neutral mutations, and converts
+        alleles to nucleotides.
         """
         # Times come from SLiM generation numbers, which may have been
         # divided by a scaling factor for computational tractability.
@@ -1763,10 +1788,18 @@ def _get_msp_rate_map(breaks, is_this_dfe, rate):
                         rates[is_this_dfe] = rate
                         return msprime.RateMap(position=breaks, rate=rates)
 
-                    # Add mutations to recapitated part of trees.
+                    # Use msprime.SLiMMutationModel rather than msprime.JC69
+                    # for neutral DFEs.  This ensures that there will be a
+                    # 'selection_coef' key in the mutation metadata (e.g. the
+                    # mutation metadata structure will be consistent across the
+                    # tree sequence).
+                    # TODO: set stacking policy to "l" when supported
                     model = msprime.SLiMMutationModel(
-                        type=mt["slim_mutation_type_id"], next_id=_get_next_id(ts)
+                        type=mt["slim_mutation_type_id"],
+                        next_id=_get_next_id(ts),
                     )
+
+                    # Add mutations to recapitated part of trees.
                     start_time = None
                     end_time = None
                     if dfe["id"] == "recapitation":
@@ -1790,6 +1823,16 @@ def _get_msp_rate_map(breaks, is_this_dfe, rate):
                         keep=True,
                         random_seed=mut_seed,
                     )
+
+        if not keep_mutation_ids_as_alleles:
+            nuc_seed = rng.randrange(1, 2**32)
+            ts = pyslim.convert_alleles(
+                pyslim.generate_nucleotides(
+                    ts,
+                    seed=nuc_seed,
+                )
+            )
+
         return ts
 
     def recap_and_rescale(
@@ -1801,6 +1844,7 @@ def recap_and_rescale(
         extended_events=None,
         slim_scaling_factor=1.0,
         seed=None,
+        keep_mutation_ids_as_alleles=False,
         **kwargs,
     ):
         """
@@ -1854,7 +1898,13 @@ def recap_and_rescale(
             )
 
         ts = self._recap_and_rescale(
-            ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events
+            ts,
+            seed,
+            recap_epoch,
+            contig,
+            slim_scaling_factor,
+            keep_mutation_ids_as_alleles,
+            extended_events,
         )
         return ts
 


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 2fdd1d6a6..a4c5e34a6 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -49,7 +49,7 @@ jobs:
         uses: actions/cache@v2
         env:
           # Increase this to reset the cache if the key hasn't changed.
-          CACHE_NUM: 0
+          CACHE_NUM: 2
         with:
           path: |
             ${{ steps.find-conda.outputs.CONDA }}/envs/${{ env.CONDA_ENV_NAME }}
diff --git a/tests/test_dfes.py b/tests/test_dfes.py
index 81b67ef1d..54b16b305 100644
--- a/tests/test_dfes.py
+++ b/tests/test_dfes.py
@@ -679,10 +679,11 @@ def test_simulation_runs(self):
             for j, md in zip(
                 mut.derived_state.split(","), mut.metadata["mutation_list"]
             ):
+                uid = f"{mut.id}_{j}"
                 if md["selection_coeff"] != 0.0:
                     nonneutral[mut.id] = True
-                if j not in mut_info:
-                    mut_info[int(j)] = md
+                if uid not in mut_info:
+                    mut_info[uid] = md
 
         num_nonneutral = sum(nonneutral)
         nonneutral_positions = ts.tables.sites.position[
diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py
index cbae1e522..2db8025c7 100644
--- a/tests/test_slim_engine.py
+++ b/tests/test_slim_engine.py
@@ -26,9 +26,9 @@
 def count_mut_types(ts):
     mut_info = {}
     for mut in ts.mutations():
-        for j, md in zip(mut.derived_state.split(","), mut.metadata["mutation_list"]):
-            if j not in mut_info:
-                mut_info[int(j)] = md
+        for j, md in enumerate(mut.metadata["mutation_list"]):
+            uid = f"{mut.id}_{j}"
+            mut_info[uid] = md
     num_neutral = sum([mut_info[j]["selection_coeff"] == 0.0 for j in mut_info])
     return [num_neutral, abs(len(mut_info) - num_neutral)]
 
@@ -281,6 +281,74 @@ def test_assert_min_version(self):
             engine._assert_min_version("3.5", engine.slim_path())
             engine._assert_min_version("3.6", None)
 
+    def test_stacked_mutations(self):
+        # Verify that `count_mut_types` works with stacked mutations, after
+        # these have been converted to nucleotides. When it's possible to set
+        # the stacking policy to "l" in SLiMMutationModel, this test will break
+        # and can be removed.
+        engine = stdpopsim.get_engine("slim")
+        species = stdpopsim.get_species("HomSap")
+        contig = species.get_contig(length=10)
+        model = stdpopsim.PiecewiseConstantSize(100)
+        samples = {"pop_0": 5}
+        contig.mutation_rate = 1e-2
+        while True:
+            ts = engine.simulate(
+                demographic_model=model,
+                contig=contig,
+                samples=samples,
+                slim_burn_in=0,
+                keep_mutation_ids_as_alleles=False,
+            )
+            is_stacked = [len(m.metadata["mutation_list"]) > 1 for m in ts.mutations()]
+            if any(is_stacked):
+                break
+        count_mut_types(ts)
+
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_allele_codes(self):
+        engine = stdpopsim.get_engine("slim")
+        species = stdpopsim.get_species("AraTha")
+        contig = species.get_contig("5", left=0, right=100000)
+        contig.add_dfe(
+            intervals=[[0, contig.length // 2]],
+            DFE=stdpopsim.DFE(
+                id="test",
+                description="test",
+                long_description="test",
+                mutation_types=[
+                    stdpopsim.MutationType(
+                        distribution_type="n",
+                        distribution_args=[0, 0.01],
+                    )
+                ],
+            ),
+        )
+        model = stdpopsim.PiecewiseConstantSize(species.population_size)
+        samples = {"pop_0": 5}
+        # nucleotides
+        ts = engine.simulate(
+            demographic_model=model,
+            contig=contig,
+            samples=samples,
+            slim_scaling_factor=10,
+            slim_burn_in=0,
+        )
+        for mut in ts.mutations():
+            assert mut.derived_state in "ACGT"
+        # slim mutation ids (may be stacked)
+        ts = engine.simulate(
+            demographic_model=model,
+            contig=contig,
+            samples=samples,
+            keep_mutation_ids_as_alleles=True,
+            slim_scaling_factor=10,
+            slim_burn_in=0,
+        )
+        for mut in ts.mutations():
+            alleles = mut.derived_state.split(",")
+            assert all([x.isnumeric() for x in alleles])
+
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
 class TestCLI:
@@ -561,6 +629,14 @@ def test_errors(self, tmp_path):
         with pytest.raises(SystemExit, match="A DFE bed file and a DFE annotation"):
             capture_output(stdpopsim.cli.stdpopsim_main, cmd)
 
+        # keep mutation ids as alleles without slim engine
+        cmd = (
+            f"-q -e msprime HomSap -c chr22 -s 1234 -l 0.01 -o {fname} "
+            f"-d OutOfAfrica_3G09 --keep-mutation-ids-as-alleles YRI:5"
+        ).split()
+        with pytest.raises(SystemExit, match="only applies to the SLiM engine"):
+            capture_output(stdpopsim.cli.stdpopsim_main, cmd)
+
     @mock.patch("stdpopsim.slim_engine._SLiMEngine.get_version", return_value="64.64")
     def test_dry_run(self, _mocked_get_version):
         # --dry-run should run slim, but not create an output file.