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

Issue:
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.
Do not simulate neutral mutations
This is a merge of #1151 and #1152 by @mufernando, with minor changes. Closes #1149.

Here's where we're at:

This PR makes it so SLiM does not simulate neutral mutation types, leaving those for msprime.
This is verified, in part, by saving information about mutation types to top-level metadata, when `verbosity=3`, in `ts.metadata['SLiM']['user_metadata']`.

We would *also* like to have stdpopsim store information in top-level metadata (e.g., under `ts.metadata['stdpopsim']`); for instance, which mutation types go with which DFE? @mufernando made a start at this in #1152. But, that's *separate* because (1) we'll be sticking that information in a different slot in metadata (SLiM can only put stuff in `["SLiM"]["user_metadata"]`); (2) we want to put *more* information than what we have easily available within SLiM (e.g., names of DFEs); (3) our use of metadata here is us asking SLiM to tell us exactly what it's doing, for unit testing, and that's just a different thing than the goal of putting useful information for downstream users in metadata. (But - the SLiM-produced metadata will totally be useful for testing the stdpopsim metadata, when we get that in.)

So - I think we should merge this, and build on #1152 to fix up #1148. (And, I think this was @mufernando's intention?)

Could you have a look at the justification and the code, @mufernando?


Discussion / hints:

# [Codecov](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium) Report
> Merging [#1157](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium) (b11c41c) into [main](https://codecov.io/gh/popsim-consortium/stdpopsim/commit/93ad9b3b4d8d428ef0eae0b865683d0b5e1d83f8?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium) (93ad9b3) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157/graphs/tree.svg?width=650&height=150&src=pr&token=AOkcaVfvhX&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium)](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium)

```diff
@@           Coverage Diff           @@
##             main    #1157   +/-   ##
=======================================
  Coverage   99.54%   99.54%           
=======================================
  Files          95       95           
  Lines        3081     3086    +5     
  Branches      395      397    +2     
=======================================
+ Hits         3067     3072    +5     
  Misses          6        6           
  Partials        8        8           
```


| [Impacted Files](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium) | Coverage Δ | |
|---|---|---|
| [stdpopsim/dfe.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium#diff-c3RkcG9wc2ltL2RmZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/slim\_engine.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium#diff-c3RkcG9wc2ltL3NsaW1fZW5naW5lLnB5) | `99.72% <100.00%> (+<0.01%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium). Last update [93ad9b3...b11c41c](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/1157?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=popsim-consortium).

thanks for organizing this!! the code looks good to me heh.

Yes, that was the intention with those two PRs. I'll rebase #1152 and we can continue the discussion there.

Fix patch:
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index 4ec751fe4..27157c0b0 100644
--- a/stdpopsim/slim_engine.py
+++ b/stdpopsim/slim_engine.py
@@ -57,6 +57,7 @@
 import msprime
 import pyslim
 import tskit
+import json
 
 logger = logging.getLogger(__name__)
 
@@ -602,12 +603,62 @@
 
 """
 
+_raw_stdpopsim_top_level_schema = {
+    "stdpopsim": {
+        "description": "Top-level metadata for a tree sequence produced by stdposim using the SLiM engine",  # noqa: E501
+        "type": "object",
+        "properties": {
+            "DFEs": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "id": {"type": "string"},
+                        "description": {"type": "string"},
+                        "long_description": {"type": "string"},
+                        "mutation_types": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "dominance_coeff": {"type": "number"},
+                                    "distribution_type": {"type": "string"},
+                                    "distribution_args": {
+                                        "type": "array",
+                                        "items": {"type": ["number", "string"]},
+                                    },
+                                    "convert_to_substitution": {"type": "boolean"},
+                                    "Q_scaled_index": {
+                                        "type": "array",
+                                        "items": {"type": "number"},
+                                    },
+                                    "slim_mutation_type_id": {"type": "integer"},
+                                    "is_neutral": {"type": "boolean"},
+                                },
+                            },
+                        },
+                        "proportions": {"type": "array", "items": {"type": "number"}},
+                        "citations": {
+                            "type": "array",
+                            "items": {"type": ["string", "object"]},
+                        },
+                        "intervals": {
+                            "type": "array",
+                            "items": {"type": "array", "minItems": 2, "maxItems": 2},
+                        },
+                        "start_time": {"type": "number"},
+                        "end_time": {"type": "number"},
+                    },
+                },
+            }
+        },
+    }
+}
 
-def get_msp_and_slim_mutation_rate_maps(contig):
+
+def get_slim_mutation_rate_map(contig):
     """
-    Returns a :class:`msprime.RateMap` with the mutation rate map for overlay
-    of neutral mutations and a tuple with the breakpoints and rates of mutations
-    for SLiM.
+    Returns a tuple with the breakpoints and rates of mutations for SLiM.
     """
     breaks, dfe_labels = contig.dfe_breakpoints()  # beware -1 labels
     slim_fractions = np.array(
@@ -625,11 +676,7 @@ def get_msp_and_slim_mutation_rate_maps(contig):
     )  # append 0 for the -1 labels
     slim_rates = contig.mutation_rate * slim_fractions[dfe_labels]
     slim_breaks = breaks[1:] - 1  # SLiM has inclusive right endpoints
-
-    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))
+    return (slim_breaks, slim_rates)
 
 
 def slim_array_string(iterable, indent, width=80):
@@ -648,6 +695,74 @@ def slim_array_string(iterable, indent, width=80):
     )
 
 
+def _enum_dfe_and_intervals(contig):
+    """
+    Returns an iterator which, for every `DFE` in the `contig`, returns a tuple
+    containing a numeric id for said DFE, the DFE object and the list of
+    intervals associated with the DFE. We need to assign each DFE a numeric id
+    both in the slim_makescript and in the _add_dfes_to_metadata, so with this
+    function we can ensure these two steps assigns numeric ids to the DFEs in the
+    same way.
+    """
+    assert len(contig.dfe_list) == len(contig.interval_list)
+    for i, (d, ints) in enumerate(zip(contig.dfe_list, contig.interval_list)):
+        yield i, d, ints
+
+
+def _dfe_to_mtypes(contig):
+    """
+    Assigns a mutation type id to each of the mutation types contained in this
+    `Contig`. This function will return a dictionary with `len(contig.dfe_list)`
+    elements, in which the position of the DFE in `contig.dfe_list` is the key.
+    For each DFE, the dictionary holds a list of tuples that contain the an
+    assigned mutation type id (used in  SLiM) and the `MutationType` object.
+    This is necessary so that we use the same numeric ids to the mutation types
+    in the _add_dfes_to_metadata and slim_makescript steps.
+    """
+    mid = 0
+    dfe_to_mtypes = {}
+    for i, d, _ in _enum_dfe_and_intervals(contig):
+        dfe_to_mtypes[i] = []
+        for mt in d.mutation_types:
+            dfe_to_mtypes[i].append((mid, mt))
+            mid += 1
+    return dfe_to_mtypes
+
+
+def _add_dfes_to_metadata(ts, contig):
+    """
+    Adds the DFEs to the top-level metadata of a tree sequence using
+    information in the `contig`.
+    """
+
+    def _get_json(obj):
+        # recursively flattens an object to a dictionary.
+        return json.loads(
+            json.dumps(obj, default=lambda o: getattr(o, "__dict__", str(o)))
+        )
+
+    tables = ts.dump_tables()
+    schema = tables.metadata_schema.asdict()
+    metadata = tables.metadata
+    schema["properties"].update(_raw_stdpopsim_top_level_schema)
+    dfes = _get_json(contig.dfe_list)
+    dfe_to_mtypes = _dfe_to_mtypes(contig)
+    for i, d, ints in _enum_dfe_and_intervals(contig):
+        dfes[i]["intervals"] = ints.tolist()
+        dfes[i]["proportions"] = d.proportions
+        for j, (mid, mt) in enumerate(dfe_to_mtypes[i]):
+            add = {
+                "slim_mutation_type_id": mid,
+                "is_neutral": mt.is_neutral,
+            }
+            dfes[i]["mutation_types"][j].update(add)
+    new_metadata = {"stdpopsim": {"DFEs": dfes}}
+    metadata.update(new_metadata)
+    tables.metadata_schema = tskit.MetadataSchema(schema)
+    tables.metadata = metadata
+    return tables.tree_sequence()
+
+
 def msprime_rm_to_slim_rm(recombination_map):
     """
     Convert recombination map from start position coords to end position coords.
@@ -671,7 +786,6 @@ def slim_makescript(
     slim_rate_map,
 ):
 
-    mutation_types = contig.mutation_types()
     pop_names = [pop.name for pop in demographic_model.model.populations]
     # Use copies of these so that the time frobbing below doesn't have
     # side-effects in the caller's model.
@@ -780,14 +894,17 @@ def fix_time(event):
     fitness_callbacks = []
     condition_on_allele_frequency = []
     op_id = stdpopsim.ext.ConditionOnAlleleFrequency.op_id
+    slim_mutation_ids = []
+    for x in _dfe_to_mtypes(contig).values():
+        slim_mutation_ids.extend([u[0] for u in x])
     for ee in extended_events:
         if hasattr(ee, "mutation_type_id"):
             mt_id = getattr(ee, "mutation_type_id")
             cls_name = ee.__class__.__name__
-            if not (0 <= mt_id < len(mutation_types)):
+            if mt_id not in slim_mutation_ids:
                 raise ValueError(
-                    f"Invalid {cls_name} {ee.mutation_type_id} {len(mutation_types)} "
-                    f"event with mutation type id {mt_id}.{contig}"
+                    f"Invalid {cls_name} event with mutation type id "
+                    f"{mt_id} on {contig} with valid IDs {slim_mutation_ids}."
                 )
         if hasattr(ee, "start_time") and hasattr(ee, "end_time"):
             # Now that GenerationAfter times have been accounted for, we can
@@ -919,11 +1036,12 @@ def matrix2str(
         return "".join(s)
 
     # Genomic element types.
-    next_mutation_type_id = 0
-    for j, d in enumerate(contig.dfe_list):
+    dfe_mtypes = _dfe_to_mtypes(contig)
+    for j, d, ints in _enum_dfe_and_intervals(contig):
         # Mutation types.
         mut_type_list = []
-        for m in d.mutation_types:
+        assert len(dfe_mtypes[j]) == len(d.mutation_types)
+        for mid, m in dfe_mtypes[j]:
             if len(m.Q_scaled_index) >= 1:
                 distrib_args = [str(arg) for arg in m.distribution_args]
                 for k in m.Q_scaled_index:
@@ -934,8 +1052,6 @@ def matrix2str(
             # dealing with distributions given by "s" Eidos script,
             else:
                 distrib_args = "; ".join([f'"{a}"' for a in m.distribution_args])
-            mid = next_mutation_type_id
-            next_mutation_type_id += 1
             mut_type_list.append(mid)
             printsc(
                 f"    initializeMutationType({mid}, {m.dominance_coeff}, "
@@ -945,6 +1061,12 @@ 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])
+        # 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,
+        # SLiM would not allow all mutation types in a genomic element type to
+        # have 0.0 proportion. And the proportions in SLiM won't matter anyway,
+        # because all the intervals in this fully neutral DFE will have 0.0
+        # mutation rate anyway.
         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)
@@ -958,11 +1080,11 @@ def matrix2str(
                 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)
+        if ints.shape[0] > 0:
+            element_starts = slim_array_string(ints[:, 0], indent)
             # stdpopsim intervals are 0-based left inclusive, right exclusive, but
             # SLiM intervals are right inclusive
-            element_ends = slim_array_string(contig.interval_list[j][:, 1] - 1, indent)
+            element_ends = slim_array_string(ints[:, 1] - 1, indent)
             printsc(
                 f"    initializeGenomicElement({j}, {element_starts}, {element_ends});"
             )
@@ -1190,6 +1312,7 @@ def simulate(
         slim_burn_in=10.0,
         dry_run=False,
         verbosity=None,
+        _recap_and_rescale=True,
     ):
         """
         Simulate the demographic model using SLiM.
@@ -1239,7 +1362,7 @@ def simulate(
 
         # figuring out mutations in SLiM and neutral mutations to be added
         # by msprime later
-        msp_rate_map, slim_rate_map = get_msp_and_slim_mutation_rate_maps(contig)
+        slim_rate_map = get_slim_mutation_rate_map(contig)
 
         # TODO: remove this after a release or two. See #745.
         self._warn_zigzag(demographic_model)
@@ -1289,9 +1412,11 @@ def script_file_f():
 
             ts = tskit.load(ts_file.name)
 
-        ts = self._recap_and_rescale(
-            ts, seed, recap_epoch, contig, slim_scaling_factor, msp_rate_map
-        )
+        ts = _add_dfes_to_metadata(ts, contig)
+        if _recap_and_rescale:
+            ts = self._recap_and_rescale(
+                ts, seed, recap_epoch, contig, slim_scaling_factor
+            )
 
         if contig.inclusion_mask is not None:
             ts = stdpopsim.utils.mask_tree_sequence(ts, contig.inclusion_mask, False)
@@ -1369,27 +1494,49 @@ 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, msp_rate_map
-    ):
+    def _recap_and_rescale(self, ts, seed, recap_epoch, contig, slim_scaling_factor):
         """
         Apply post-SLiM transformations to ``ts``. This rescales node times,
         does recapitation, simplification, and adds neutral mutations.
         """
-        # Node times come from SLiM generation numbers, which may have been
+        # Times come from SLiM generation numbers, which may have been
         # divided by a scaling factor for computational tractability.
         tables = ts.dump_tables()
         for table in (tables.nodes, tables.migrations, tables.mutations):
             table.time *= slim_scaling_factor
-        ts_metadata = tables.metadata
-        ts_metadata["SLiM"]["generation"] *= slim_scaling_factor
-        tables.metadata = ts_metadata
+        metadata = tables.metadata
+        metadata["SLiM"]["generation"] *= slim_scaling_factor
+        # Finding what slim id to use in recap DFE
+        max_id = -1
+        for dfe in metadata["stdpopsim"]["DFEs"]:
+            for mt in dfe["mutation_types"]:
+                max_id = max(mt["slim_mutation_type_id"], max_id)
+        recap_dfe = {
+            "id": "recapitation",
+            "description": "DFE used in recapitation",
+            "long_description": "neutral mutations added with msprime.sim_mutations "
+            "to the recapited portion of the tree sequence",
+            "mutation_types": [
+                {
+                    "dominance_coeff": 0.5,
+                    "distribution_type": "f",
+                    "distribution_args": [0],
+                    "convert_to_substitution": False,
+                    "Q_scaled_index": [0],
+                    "slim_mutation_type_id": max_id + 1,
+                    "is_neutral": True,
+                }
+            ],
+            "proportions": [1.0],
+            "citations": [],
+            "intervals": [[0, ts.sequence_length]],
+        }
+        metadata["stdpopsim"]["DFEs"].append(recap_dfe)
+        tables.metadata = metadata
         ts = tables.tree_sequence()
-        slim_gen = ts.metadata["SLiM"]["generation"]
 
         rng = random.Random(seed)
-        s1, s2 = rng.randrange(1, 2 ** 32), rng.randrange(1, 2 ** 32)
-
+        s0 = rng.randrange(1, 2 ** 32)
         population_configurations = [
             msprime.PopulationConfiguration(
                 initial_size=pop.start_size,
@@ -1408,44 +1555,58 @@ def _recap_and_rescale(
             initial_state=ts,
             demography=demography,
             recombination_rate=contig.recombination_map,
-            random_seed=s1,
+            random_seed=s0,
         )
-
         ts = self._simplify_remembered(ts)
 
-        # Figuring out SLiM mutation id metadata
-        max_id = -1
-        for mut in ts.mutations():
-            for d in mut.derived_state.split(","):
-                max_id = max(max_id, int(d))
-
-        # Creating a SLiM mutation model
-        # TODO: fix this up to do the right types? or, share them?
-        model = msprime.SLiMMutationModel(
-            type=len(contig.mutation_types()), next_id=max_id + 1
-        )
-
-        # Add mutations to recapitated part of trees.
-        ts = msprime.sim_mutations(
-            ts,
-            rate=contig.mutation_rate,
-            model=model,
-            start_time=slim_gen,
-            keep=True,
-            random_seed=s2,
-        )
-
-        # Adding neutral mutations to the SLiM part
-        s3 = rng.randrange(1, 2 ** 32)
-        ts = msprime.sim_mutations(
-            ts,
-            rate=msp_rate_map,
-            model=model,
-            end_time=slim_gen,
-            keep=True,
-            random_seed=s3,
-        )
-
+        # Adding neutral mutations to simulation and recapitation periods
+        breaks, dfe_labels = contig.dfe_breakpoints()  # beware -1 labels
+        for i, dfe in enumerate(ts.metadata["stdpopsim"]["DFEs"]):
+            assert len(dfe["proportions"]) == len(dfe["mutation_types"])
+            for prop, mt in zip(dfe["proportions"], dfe["mutation_types"]):
+                if mt["is_neutral"]:
+                    mut_seed = rng.randrange(1, 2 ** 32)
+                    # Figuring out SLiM mutation id metadata
+
+                    def _get_next_id(ts):
+                        max_id = -1
+                        for mut in ts.mutations():
+                            for d in mut.derived_state.split(","):
+                                max_id = max(max_id, int(d))
+                        return max_id + 1
+
+                    def _get_msp_rate_map(breaks, is_this_dfe, rate):
+                        rates = np.zeros(shape=is_this_dfe.shape)
+                        rates[is_this_dfe] = rate
+                        return msprime.RateMap(position=breaks, rate=rates)
+
+                    # Add mutations to recapitated part of trees.
+                    model = msprime.SLiMMutationModel(
+                        type=mt["slim_mutation_type_id"], next_id=_get_next_id(ts)
+                    )
+                    start_time = None
+                    end_time = None
+                    if dfe["id"] == "recapitation":
+                        start_time = metadata["SLiM"]["generation"]
+                        msp_rate_map = _get_msp_rate_map(
+                            np.array([0, contig.length]),
+                            np.array([True]),
+                            contig.mutation_rate,
+                        )
+                    else:
+                        msp_rate_map = _get_msp_rate_map(
+                            breaks, dfe_labels == i, prop * contig.mutation_rate
+                        )
+                        end_time = metadata["SLiM"]["generation"]
+                    ts = msprime.sim_mutations(
+                        ts,
+                        rate=msp_rate_map,
+                        model=model,
+                        start_time=start_time,
+                        end_time=end_time,
+                        keep=True,
+                        random_seed=mut_seed,
+                    )
         return ts
 
     def recap_and_rescale(
@@ -1482,7 +1643,7 @@ def recap_and_rescale(
             should be consulted to determine if the behaviour is appropriate
             for your case.
         """
-        msp_rate_map, slim_rate_map = get_msp_and_slim_mutation_rate_maps(contig)
+        slim_rate_map = get_slim_mutation_rate_map(contig)
 
         with open(os.devnull, "w") as script_file:
             recap_epoch = slim_makescript(
@@ -1497,9 +1658,7 @@ def recap_and_rescale(
                 slim_rate_map,
             )
 
-        ts = self._recap_and_rescale(
-            ts, seed, recap_epoch, contig, slim_scaling_factor, msp_rate_map
-        )
+        ts = self._recap_and_rescale(ts, seed, recap_epoch, contig, slim_scaling_factor)
         return ts
 
 


Test patch:
diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py
index 03e45669c..f3c315d1a 100644
--- a/tests/test_slim_engine.py
+++ b/tests/test_slim_engine.py
@@ -2,7 +2,6 @@
 Tests for SLiM simulation engine.
 """
 import os
-import re
 import io
 import sys
 import tempfile
@@ -23,30 +22,6 @@
 slim_path = os.environ.get("SLIM", "slim")
 
 
-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
-    tmp_trees_file = match.group(1)
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        script_file = os.path.join(tmpdir, "script.slim")
-        trees_file = os.path.join(tmpdir, "out.trees")
-        # Write out the script with a new location for the trees_file.
-        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, verbosity=verbosity)
-        ts = tskit.load(trees_file)
-    return ts
-
-
 def count_mut_types(ts):
     mut_info = {}
     for mut in ts.mutations():
@@ -217,7 +192,7 @@ def test_recap_and_rescale(self):
                 slim_burn_in=0,
                 seed=seed,
             )
-            ts2_headless = slim_simulate_no_recap(
+            ts2_headless = engine.simulate(
                 demographic_model=model,
                 contig=contig,
                 samples=samples,
@@ -225,6 +200,7 @@ def test_recap_and_rescale(self):
                 slim_scaling_factor=10,
                 slim_burn_in=0,
                 seed=seed,
+                _recap_and_rescale=False,
             )
             ts2 = engine.recap_and_rescale(
                 ts2_headless,
@@ -942,7 +918,7 @@ def verify_rates_match(self, match_rate, intervals, slim_rates, slim_ends):
             for rate, end in zip(slim_rates, slim_ends):
                 # SLiM's endpoints are inclusive
                 if left <= end:
-                    assert rate == match_rate
+                    assert np.isclose(rate, match_rate)
                 if right - 1 <= end:
                     break
             assert right - 1 <= end
@@ -974,6 +950,41 @@ def verify_mutation_rates(self, contig, ts):
             )
         self.verify_rates_match(0.0, empty_intervals, slim_rates, slim_ends)
 
+    def verify_dfes_metadata(self, ts, ge_types):
+        def _slim_proportions(dfe):
+            is_dfe_neutral = True
+            for mt in dfe["mutation_types"]:
+                is_dfe_neutral &= mt["is_neutral"]
+            # Recall a mutation type will have proportion 0.0 in SLiM if it is
+            # neutral and not all the mutation types in the corresponding DFE are
+            # neutral. SLiM does not allow a genomicelementtype to have all 0.0
+            # proportions.
+            return [
+                prop if (not mt["is_neutral"]) or is_dfe_neutral else 0.0
+                for prop, mt in zip(dfe["proportions"], dfe["mutation_types"])
+            ]
+
+        for i, dfe in enumerate(ts.metadata["stdpopsim"]["DFEs"]):
+            if dfe["id"] == "recapitation":
+                continue
+            ge = self.slim_metadata_key0(ge_types, str(i))
+            int_ends = []
+            int_starts = []
+            for inter in dfe["intervals"]:
+                int_starts.append(inter[0])
+                # in SLiM ends are inclusive
+                int_ends.append(inter[1] - 1)
+            ge_from_meta = {
+                "intervalEnds": int_ends,
+                "intervalStarts": int_starts,
+                "mutationFractions": _slim_proportions(dfe),
+                "mutationTypes": [
+                    mt["slim_mutation_type_id"] for mt in dfe["mutation_types"]
+                ],
+            }
+            ge = self.slim_metadata_key0(ge_types, str(i))
+            assert ge == ge_from_meta
+
     def verify_genomic_elements(self, contig, ts):
         # compare information about genomic elements as recorded
         # by SLiM itself in metadata to what we think it should be
@@ -983,6 +994,7 @@ def verify_genomic_elements(self, contig, ts):
         ge_types = self.slim_metadata_key0(
             ts.metadata["SLiM"]["user_metadata"], "genomicElementTypes"
         )
+        self.verify_dfes_metadata(ts, ge_types)
         assert len(contig.dfe_list) == len(ge_types)
         for j, (dfe, intervals) in enumerate(
             zip(contig.dfe_list, contig.interval_list)
@@ -1119,6 +1131,7 @@ def test_same_dfes(self):
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_slim_produces_mutations(self):
+        engine = stdpopsim.get_engine("slim")
         contig = get_test_contig()
         dfe = stdpopsim.DFE(
             id="test",
@@ -1133,43 +1146,48 @@ def test_slim_produces_mutations(self):
             proportions=[0.5, 0.5],
         )
         contig.add_dfe(np.array([[0, contig.length]], dtype="int"), dfe)
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
             verbosity=3,  # to get metadata output
+            _recap_and_rescale=False,
         )
         assert ts.num_sites > 0
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_no_neutral_mutations_are_simulated_by_slim(self):
+        engine = stdpopsim.get_engine("slim")
         contig = get_test_contig()
         contig.mutation_rate = 10 * contig.mutation_rate
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
             verbosity=3,
+            _recap_and_rescale=False,
         )
         assert ts.num_sites == 0
         self.verify_genomic_elements(contig, ts)
         self.verify_mutation_rates(contig, ts)
+        assert len(ts.metadata["stdpopsim"]["DFEs"]) == len(contig.dfe_list)
 
         contig.dfe_list[0].mutation_types = [
             stdpopsim.MutationType() for i in range(10)
         ]
         contig.dfe_list[0].proportions = [1 / 10 for i in range(10)]
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
             verbosity=3,
+            _recap_and_rescale=False,
         )
         assert ts.num_sites == 0
         self.verify_genomic_elements(contig, ts)
@@ -1177,15 +1195,17 @@ def test_no_neutral_mutations_are_simulated_by_slim(self):
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_neutral_dfe_slim_proportions(self):
+        engine = stdpopsim.get_engine("slim")
         contig = get_test_contig()
         Q = 10
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=Q,
             slim_burn_in=0.1,
             verbosity=3,
+            _recap_and_rescale=False,
         )
         ge_types = self.slim_metadata_key0(
             ts.metadata["SLiM"]["user_metadata"], "genomicElementTypes"
@@ -1210,13 +1230,14 @@ def test_neutral_dfe_slim_proportions(self):
                 proportions=[0.9, 0.1],
             ),
         )
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             slim_scaling_factor=Q,
             slim_burn_in=0.1,
             verbosity=3,
+            _recap_and_rescale=False,
         )
         ge_types = self.slim_metadata_key0(
             ts.metadata["SLiM"]["user_metadata"], "genomicElementTypes"
@@ -1365,6 +1386,7 @@ def test_save_point_creation(self):
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_drawn_mutation_not_lost(self):
+        engine = stdpopsim.get_engine("slim")
         ct = self.contig
         ct.mutation_rate = 0.0
         mt = [
@@ -1396,18 +1418,20 @@ def test_drawn_mutation_not_lost(self):
                 allele_frequency=0,
             ),
         ]
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=ct,
             samples=self.samples,
             extended_events=extended_events,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
+            _recap_and_rescale=False,
         )
         assert ts.num_mutations == 1
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_drawn_mutation_is_lost(self):
+        engine = stdpopsim.get_engine("slim")
         ct = self.contig
         ct.mutation_rate = 0.0
         mt = [
@@ -1439,18 +1463,20 @@ def test_drawn_mutation_is_lost(self):
                 allele_frequency=0,
             ),
         ]
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=ct,
             samples=self.samples,
             extended_events=extended_events,
             slim_scaling_factor=10,
             slim_burn_in=0.1,
+            _recap_and_rescale=False,
         )
         assert ts.num_mutations == 0
 
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_drawn_mutation_meets_AF_threshold(self):
+        engine = stdpopsim.get_engine("slim")
         ct = self.contig
         ct.mutation_rate = 0.0
         mt = [
@@ -1484,7 +1510,7 @@ def test_drawn_mutation_meets_AF_threshold(self):
                     allele_frequency=af_threshold,
                 ),
             ]
-            ts = slim_simulate_no_recap(
+            ts = engine.simulate(
                 demographic_model=self.model,
                 contig=ct,
                 samples=self.samples,
@@ -1492,6 +1518,7 @@ def test_drawn_mutation_meets_AF_threshold(self):
                 slim_scaling_factor=10,
                 slim_burn_in=0.1,
                 seed=seed,
+                _recap_and_rescale=False,
             )
             assert ts.num_mutations == 1
             assert self.allele_frequency(ts) >= af_threshold
@@ -1641,6 +1668,7 @@ def test_no_drawn_mutation(self):
 class TestFixedSelectionCoefficient(PiecewiseConstantSizeMixin):
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_drawn_mutation_has_correct_selection_coeff(self):
+        engine = stdpopsim.get_engine("slim")
         contig = self.contig
         contig.dfe_list[0].mutation_types = [
             stdpopsim.MutationType(),
@@ -1671,13 +1699,14 @@ def test_drawn_mutation_has_correct_selection_coeff(self):
             ),
         ]
         scaling_factor = 10
-        ts = slim_simulate_no_recap(
+        ts = engine.simulate(
             demographic_model=self.model,
             contig=contig,
             samples=self.samples,
             extended_events=extended_events,
             slim_scaling_factor=scaling_factor,
             slim_burn_in=0.1,
+            _recap_and_rescale=False,
         )
         assert ts.num_mutations == 1
         mut = next(ts.mutations())
@@ -1698,6 +1727,7 @@ class TestChangeMutationFitness(PiecewiseConstantSizeMixin):
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     @pytest.mark.skip("can't get this to complete?")
     def test_positive_mutation_meets_AF_threshold(self):
+        engine = stdpopsim.get_engine("slim")
         for af_threshold, seed in zip((0.5, 1), (1, 2)):
             extended_events = [
                 stdpopsim.ext.DrawMutation(
@@ -1743,7 +1773,7 @@ def test_positive_mutation_meets_AF_threshold(self):
             ct.dfe_list[0].mutation_types = mt
             print(ct.dfe_list)
             print(ct.mutation_types())
-            ts = slim_simulate_no_recap(
+            ts = engine.simulate(
                 demographic_model=self.model,
                 contig=ct,
                 samples=self.samples,
@@ -1751,6 +1781,7 @@ def test_positive_mutation_meets_AF_threshold(self):
                 slim_scaling_factor=10,
                 slim_burn_in=0.1,
                 seed=seed,
+                _recap_and_rescale=False,
             )
             assert ts.num_mutations == 1
             assert self.allele_frequency(ts) >= af_threshold