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

Issue:
fixing neutral mutations
Fix to #1149

We were simulating neutral mutations within SLiM before. The fix was not as straightforward as @petrelharp pointed in the issue because we cannot have a genomic element type with all 0.0 proportions in SLiM. I added tests to this behavior (one of which is redundant, but is more obviously correct so I thought it wouldn't hurt? Also made some minor cosmetic changes throughout...
Dfes in ts metadata
The idea of this PR is to add metadata to tree sequences produced by the SLiM engine. I'm focusing on adding information about the DFEs onto the tree sequence and this will have two main uses:

- It will allow us to streamline the way we overlay neutral mutations with `sim_mutations`. The metadata in the ts will encode information about which mutation types SLiM did not simulate and we can simply iterate over those and add them with `sim_mutations`. Importantly, this will make so that the slim mutation type ids are the same as the ones in the original `Contig`, so we won't have to create new mutation types in the "recapitate and mutate" step.
- The metadata is also useful for downstream analyses.  For instance, this will allow people to figure out which intervals have which mutation types and at what proportions/rates. In #1148 @petrelharp describes our most pressing use case, that @xin-huang needs to know which mutations are neutral within exons to benchmark the DFE inference software.

PS: note this PR is based off of #1157, so some of the changes displayed here are from the other PR. as soon as that is merged I can rebase onto main.
don't simulate neutral mutations in SLiM
@mufernando points out that it looks like we are actually simulating neutral mutations twice, once with SLiM and once with msprime. The fix should be to change [this line](https://github.com/popsim-consortium/stdpopsim/blob/1c489090a6dde9414dc1b8825cdb2d0d540dbe3c/stdpopsim/slim_engine.py#L946) to something like:
```
        mut_props = ", ".join([str(prop) if mt.is_neutral() else 0.0 for prop, mt in d.proportions])
```

We also should have a test that catches this, by doing a SLiM simulation without recapitation and checking if any actual mutations are of the neutral types.

Ah - the test could be added to [this one](https://github.com/popsim-consortium/stdpopsim/blob/1c489090a6dde9414dc1b8825cdb2d0d540dbe3c/tests/test_slim_engine.py#L987), checking that all neutral mutation types have 0.0 proportion. If so, then we could go ahead and recapitate and still run the test.
Dfes in ts metadata
The idea of this PR is to add metadata to tree sequences produced by the SLiM engine. I'm focusing on adding information about the DFEs onto the tree sequence and this will have two main uses:

- It will allow us to streamline the way we overlay neutral mutations with `sim_mutations`. The metadata in the ts will encode information about which mutation types SLiM did not simulate and we can simply iterate over those and add them with `sim_mutations`. Importantly, this will make so that the slim mutation type ids are the same as the ones in the original `Contig`, so we won't have to create new mutation types in the "recapitate and mutate" step.
- The metadata is also useful for downstream analyses.  For instance, this will allow people to figure out which intervals have which mutation types and at what proportions/rates. In #1148 @petrelharp describes our most pressing use case, that @xin-huang needs to know which mutations are neutral within exons to benchmark the DFE inference software.

PS: note this PR is based off of #1157, so some of the changes displayed here are from the other PR. as soon as that is merged I can rebase onto main.
Dfes in ts metadata
The idea of this PR is to add metadata to tree sequences produced by the SLiM engine. I'm focusing on adding information about the DFEs onto the tree sequence and this will have two main uses:

- It will allow us to streamline the way we overlay neutral mutations with `sim_mutations`. The metadata in the ts will encode information about which mutation types SLiM did not simulate and we can simply iterate over those and add them with `sim_mutations`. Importantly, this will make so that the slim mutation type ids are the same as the ones in the original `Contig`, so we won't have to create new mutation types in the "recapitate and mutate" step.
- The metadata is also useful for downstream analyses.  For instance, this will allow people to figure out which intervals have which mutation types and at what proportions/rates. In #1148 @petrelharp describes our most pressing use case, that @xin-huang needs to know which mutations are neutral within exons to benchmark the DFE inference software.

PS: note this PR is based off of #1157, so some of the changes displayed here are from the other PR. as soon as that is merged I can rebase onto main.
store annotation information in metadata
Proposal (with @mufernando):
In the SLiM script we know what mutation type goes in each genomic region, but we don't use that information later when adding neutral mutations with msprime.
So, we should in fact store that information in top-level metadata, and use the information to tell msprime which mutation types to add.

To do this, we need to:

1. store in metadata, for each DFE, (a) the intervals it applies to, (b) and the mutation types it has, including the SLiM ID and whether it's being simulated by SLiM
2. 
2. when we add mutations with msprime, we no longer have to re-compute the msp_rate_map, instead we loop over the mutation types within each dfe, and, if it's neutral, we add the mutations with msprime.


Discussion / hints:
The only test failing is because of pyslim.
We need to change `ts.reference_sequence` to `ts.reference_sequence.data`.
Note some of the changes here are from my previous PR.
Looks good to me, after a few fixes!
I think this is ready to go now - I've rebased it, even - but, have a look, and let's see about the tests.
Sorry - did I mess things up?
No, I fucked up in merging your commits I think. Ideas to test the metadata?
test that neutral mutations have proportion 0 in SLiM, similar to this: https://github.com/popsim-consortium/stdpopsim/blob/1c489090a6dde9414dc1b8825cdb2d0d540dbe3c/tests/test_slim_engine.py#L987
Peter, it looks like this test you linked is doing the right thing. But none of these tests were being run because we were waiting for SLiM 3.7. See how you were computing the proportion of neutral mutations and excluding them from the mutation rate here: https://github.com/popsim-consortium/stdpopsim/blob/1c489090a6dde9414dc1b8825cdb2d0d540dbe3c/tests/test_slim_engine.py#L997
We need to change `ts.reference_sequence` to `ts.reference_sequence.data`.
Note some of the changes here are from my previous PR.
Looks good to me, after a few fixes!
I think this is ready to go now - I've rebased it, even - but, have a look, and let's see about the tests.
Sorry - did I mess things up?
No, I fucked up in merging your commits I think. Ideas to test the metadata?
We need to change `ts.reference_sequence` to `ts.reference_sequence.data`.
Note some of the changes here are from my previous PR.
Looks good to me, after a few fixes!
I think this is ready to go now - I've rebased it, even - but, have a look, and let's see about the tests.
Sorry - did I mess things up?
No, I fucked up in merging your commits I think. Ideas to test the metadata?


Fix patch:
diff --git a/stdpopsim/dfe.py b/stdpopsim/dfe.py
index 725dadf49..a3983f183 100644
--- a/stdpopsim/dfe.py
+++ b/stdpopsim/dfe.py
@@ -224,6 +224,10 @@ def __attrs_post_init__(self):
                     "mutation_types must be a list of MutationType objects."
                 )
 
+    @property
+    def is_neutral(self):
+        return all([m.is_neutral for m in self.mutation_types])
+
     def __str__(self):
         long_desc_lines = [
             line.strip()
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index 9f8060056..4ec751fe4 100644
--- a/stdpopsim/slim_engine.py
+++ b/stdpopsim/slim_engine.py
@@ -603,17 +603,6 @@
 """
 
 
-def msprime_rm_to_slim_rm(recombination_map):
-    """
-    Convert recombination map from start position coords to end position coords.
-    """
-    rates = recombination_map.rate.copy()
-    # replace missing values with 0 recombination rate
-    rates[recombination_map.missing] = 0
-    ends = [int(pos) - 1 for pos in recombination_map.position]
-    return rates, ends[1:]
-
-
 def get_msp_and_slim_mutation_rate_maps(contig):
     """
     Returns a :class:`msprime.RateMap` with the mutation rate map for overlay
@@ -626,11 +615,11 @@ def get_msp_and_slim_mutation_rate_maps(contig):
             sum(
                 [
                     p
-                    for p, mt in zip(g.proportions, g.mutation_types)
-                    if not mt.is_neutral
+                    for p, mt in zip(d.proportions, d.mutation_types)
+                    if (not mt.is_neutral)
                 ]
             )
-            for g in contig.dfe_list
+            for d in contig.dfe_list
         ]
         + [0]
     )  # append 0 for the -1 labels
@@ -640,7 +629,6 @@ def get_msp_and_slim_mutation_rate_maps(contig):
     msp_mutation_rate_map = msprime.RateMap(
         position=breaks, rate=contig.mutation_rate * (1 - slim_fractions[dfe_labels])
     )
-
     return (msp_mutation_rate_map, (slim_breaks, slim_rates))
 
 
@@ -660,6 +648,17 @@ def slim_array_string(iterable, indent, width=80):
     )
 
 
+def msprime_rm_to_slim_rm(recombination_map):
+    """
+    Convert recombination map from start position coords to end position coords.
+    """
+    rates = recombination_map.rate.copy()
+    # replace missing values with 0 recombination rate
+    rates[recombination_map.missing] = 0
+    ends = [int(pos) - 1 for pos in recombination_map.position]
+    return rates, ends[1:]
+
+
 def slim_makescript(
     script_file,
     trees_file,
@@ -946,10 +945,19 @@ def matrix2str(
                 # T is the default for WF simulations.
                 printsc(f"    m{mid}.convertToSubstitution = F;")
         mut_types = ", ".join([str(mt) for mt in mut_type_list])
-        mut_props = ", ".join([str(prop) for prop in d.proportions])
+        mut_props_list = [
+            prop if (not mt.is_neutral) or d.is_neutral else 0.0
+            for prop, mt in zip(d.proportions, d.mutation_types)
+        ]
+        mut_props = ", ".join(map(str, mut_props_list))
         printsc(
             f"    initializeGenomicElementType({j}, c({mut_types}), c({mut_props}));"
         )
+        if d.is_neutral:
+            printsc(
+                f"    // Note: genomic element type {j} is entirely neutral,"
+                " so will not be simulated by SLiM."
+            )
         if contig.interval_list[j].shape[0] > 0:
             element_starts = slim_array_string(contig.interval_list[j][:, 0], indent)
             # stdpopsim intervals are 0-based left inclusive, right exclusive, but


Test patch:
diff --git a/tests/test_genomes.py b/tests/test_genomes.py
index 4b128cc66..b8e50af3c 100644
--- a/tests/test_genomes.py
+++ b/tests/test_genomes.py
@@ -263,6 +263,7 @@ def test_basic_dfe(self):
                 assert a == b
             for a, b in zip(mt, dfe.mutation_types):
                 assert a == b
+        assert dfe.is_neutral
 
     def test_dfe_defaults(self):
         m1 = stdpopsim.MutationType()
@@ -278,6 +279,30 @@ def test_dfe_defaults(self):
         assert isinstance(dfe.proportions, list)
         assert len(dfe.proportions) == 1
         assert dfe.proportions[0] == 1
+        assert dfe.is_neutral
+
+    def test_dfe_is_neutral(self):
+        for neutral in (True, False):
+            for dist in ("f", "e"):
+                props = [0.3, 0.7]
+                if neutral:
+                    svals = [0.0, 0.0]
+                else:
+                    svals = [0.0, 0.1]
+                mt = [
+                    stdpopsim.MutationType(
+                        distribution_type=dist, distribution_args=[s]
+                    )
+                    for s in svals
+                ]
+                dfe = stdpopsim.DFE(
+                    id=0,
+                    description="test",
+                    long_description="test test",
+                    proportions=props,
+                    mutation_types=mt,
+                )
+                assert dfe.is_neutral is (neutral and dist == "f")
 
     @pytest.mark.usefixtures("capsys")
     def test_printing_dfe(self, capsys):
diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py
index 99e21255f..f05025abb 100644
--- a/tests/test_slim_engine.py
+++ b/tests/test_slim_engine.py
@@ -5,7 +5,6 @@
 import re
 import io
 import sys
-import itertools
 import tempfile
 import math
 from unittest import mock
@@ -14,7 +13,6 @@
 
 import pytest
 import tskit
-import pyslim
 import msprime
 
 import stdpopsim
@@ -24,14 +22,13 @@
 IS_WINDOWS = sys.platform.startswith("win")
 
 
-def slim_simulate_no_recap(seed=1234, **kwargs):
+def slim_simulate_no_recap(seed=1234, verbosity=2, **kwargs):
     """
     Return the tree sequence produced by SLiM, without recapitation, etc.
     """
     kwargs.update(slim_script=True)
     engine = stdpopsim.get_engine("slim")
     out, _ = capture_output(engine.simulate, **kwargs)
-
     # Find the name of the temporary trees_file in the script.
     match = re.search(r'"trees_file",\s*"([^"]*)"', out)
     assert match is not None
@@ -44,8 +41,8 @@ def slim_simulate_no_recap(seed=1234, **kwargs):
         out = out.replace(tmp_trees_file, trees_file)
         with open(script_file, "w") as f:
             f.write(out)
-        engine._run_slim(script_file, seed=seed)
-        ts = pyslim.load(trees_file)
+        engine._run_slim(script_file, seed=seed, verbosity=verbosity)
+        ts = tskit.load(trees_file)
     return ts
 
 
@@ -859,11 +856,7 @@ def allele_frequency(self, ts):
         """
         # surely there's a simpler way!
         assert ts.num_mutations == 1
-        alive = list(
-            itertools.chain.from_iterable(
-                ts.individual(i).nodes for i in ts.individuals_alive_at(0)
-            )
-        )
+        alive = ts.samples(time=0)
         mut = next(ts.mutations())
         tree = ts.at(ts.site(mut.site).position)
         have_mut = [u for u in alive if tree.is_descendant(u, mut.node)]
@@ -944,6 +937,7 @@ def test_slim_simulates_dfes(self):
             contig=contig,
             samples=self.samples,
             seed=123,
+            verbosity=3,
         )
         for j, mt in enumerate(self.example_mut_types):
             sites = np.where(
@@ -989,15 +983,15 @@ def verify_mutation_rates(self, contig, ts):
         Q = self.slim_metadata_key0(ts.metadata["SLiM"]["user_metadata"], "Q")
 
         for dfe, intervals in zip(contig.dfe_list, contig.interval_list):
-            prop_neutral = 0.0
+            prop_nonneutral = 0.0
             for mt, p in zip(dfe.mutation_types, dfe.proportions):
-                if mt.is_neutral:
-                    prop_neutral += p
-            mut_rate = contig.mutation_rate * (1 - prop_neutral) * Q
+                if not mt.is_neutral:
+                    prop_nonneutral += p
+            mut_rate = contig.mutation_rate * prop_nonneutral * Q
             self.verify_rates_match(mut_rate, intervals, slim_rates, slim_ends)
 
         # also verify any bits not covered by DFEs have no mutations
-        empty_intervals = np.empty((0, 2), dtype="int")
+        empty_intervals = np.array([[0, int(contig.length)]], dtype="int")
         for intervals in contig.interval_list:
             empty_intervals = stdpopsim.utils.mask_intervals(
                 empty_intervals,
@@ -1020,6 +1014,15 @@ def verify_genomic_elements(self, contig, ts):
         ):
             assert str(j) in ge_types
             ge = self.slim_metadata_key0(ge_types, str(j))
+            # checking that the neutral mutations have 0.0 proportion unless
+            # all the mutations are neutral in this dfe
+            for mt, dfe_prop, slim_prop in zip(
+                dfe.mutation_types, dfe.proportions, ge["mutationFractions"]
+            ):
+                if mt.is_neutral and not dfe.is_neutral:
+                    assert slim_prop == 0.0
+                else:
+                    assert slim_prop == dfe_prop
             # "+1" because SLiM's intervals are closed on both ends,
             # stdpopsim's are closed on the left, open on the right
             slim_intervals = np.column_stack(
@@ -1168,14 +1171,18 @@ def test_slim_produces_mutations(self):
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_no_neutral_mutations_are_simulated_by_slim(self):
         contig = get_test_contig()
+        contig.mutation_rate = 10 * contig.mutation_rate
         ts = slim_simulate_no_recap(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
+            verbosity=3,
         )
         assert ts.num_sites == 0
+        self.verify_genomic_elements(contig, ts)
+        self.verify_mutation_rates(contig, ts)
 
         contig.dfe_list[0].mutation_types = [
             stdpopsim.MutationType() for i in range(10)
@@ -1187,8 +1194,62 @@ def test_no_neutral_mutations_are_simulated_by_slim(self):
             samples=self.samples,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
+            verbosity=3,
         )
         assert ts.num_sites == 0
+        self.verify_genomic_elements(contig, ts)
+        self.verify_mutation_rates(contig, ts)
+
+    def test_neutral_dfe_slim_proportions(self):
+        contig = get_test_contig()
+        Q = 10
+        ts = slim_simulate_no_recap(
+            demographic_model=self.model,
+            contig=contig,
+            samples=self.samples,
+            slim_scaling_factor=Q,
+            slim_burn_in=0.1,
+            verbosity=3,
+        )
+        ge_types = self.slim_metadata_key0(
+            ts.metadata["SLiM"]["user_metadata"], "genomicElementTypes"
+        )
+        assert np.allclose(ge_types["0"][0]["mutationFractions"], [1.0])
+        assert np.allclose(
+            ts.metadata["SLiM"]["user_metadata"]["mutationRates"][0]["rates"], [0.0 * Q]
+        )
+        contig.add_dfe(
+            intervals=np.array([[0, contig.length]], dtype="int"),
+            DFE=stdpopsim.DFE(
+                id="neutral_sel",
+                description="neutral with some selected",
+                long_description="test",
+                mutation_types=[
+                    stdpopsim.MutationType(),
+                    stdpopsim.MutationType(
+                        distribution_type="e",
+                        distribution_args=[-0.01],
+                    ),
+                ],
+                proportions=[0.9, 0.1],
+            ),
+        )
+        ts = slim_simulate_no_recap(
+            demographic_model=self.model,
+            contig=contig,
+            samples=self.samples,
+            slim_scaling_factor=Q,
+            slim_burn_in=0.1,
+            verbosity=3,
+        )
+        ge_types = self.slim_metadata_key0(
+            ts.metadata["SLiM"]["user_metadata"], "genomicElementTypes"
+        )
+        assert np.allclose(ge_types["1"][0]["mutationFractions"], [0.0, 0.1])
+        assert np.allclose(
+            ts.metadata["SLiM"]["user_metadata"]["mutationRates"][0]["rates"],
+            [contig.mutation_rate * 0.1 * Q],
+        )
 
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")