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

Issue:
kwargs for `simulate` functions
I recently noticed that there's no option to set `record_migrations=True` when using `simulate()`.

Since this might be an issue if people want to use other options such as `record_full_arg `, I propose we use **kwargs with engine.simulate(), which would then be passed to the slim and msprime engines.

I'm happy to implement this if others agree it's a good idea!


Discussion / hints:

Fix patch:
diff --git a/stdpopsim/engines.py b/stdpopsim/engines.py
index c9864966c..a423a905a 100644
--- a/stdpopsim/engines.py
+++ b/stdpopsim/engines.py
@@ -65,14 +65,16 @@ class Engine:
 
     def simulate(
         self,
-        demographic_model=None,
-        contig=None,
-        samples=None,
+        demographic_model,
+        contig,
+        samples,
+        *,
         seed=None,
         dry_run=False,
     ):
         """
-        Simulates the model for the specified contig and samples.
+        Simulates the model for the specified contig and samples. ``demographic_model``,
+        ``contig``, and ``samples`` must be specified.
 
         :param demographic_model: The demographic model to simulate.
         :type demographic_model: :class:`.DemographicModel`
@@ -126,13 +128,15 @@ class _MsprimeEngine(Engine):
 
     def simulate(
         self,
-        demographic_model=None,
-        contig=None,
-        samples=None,
+        demographic_model,
+        contig,
+        samples,
+        *,
         seed=None,
         msprime_model=None,
         msprime_change_model=None,
         dry_run=False,
+        **kwargs,
     ):
         """
         Simulate the demographic model using msprime.
@@ -149,6 +153,7 @@ def simulate(
         :param dry_run: If True, ``end_time=0`` is passed to :meth:`msprime.simulate()`
             to initialise the simulation and then immediately return.
         :type dry_run: bool
+        :param \\**kwargs: Further arguments passed to :meth:`msprime.simulate()`
         """
         if msprime_model is None:
             msprime_model = self.supported_models[0]
@@ -169,6 +174,13 @@ def simulate(
                     self.citations.extend(self.model_citations[model])
             demographic_events.sort(key=lambda x: x.time)
 
+        if "random_seed" in kwargs.keys():
+            if seed is None:
+                seed = kwargs["random_seed"]
+                del kwargs["random_seed"]
+            else:
+                raise ValueError("Cannot set both seed and random_seed")
+
         ts = msprime.simulate(
             samples=samples,
             recombination_map=contig.recombination_map,
@@ -179,6 +191,7 @@ def simulate(
             random_seed=seed,
             model=msprime_model,
             end_time=0 if dry_run else None,
+            **kwargs,
         )
 
         if contig.inclusion_mask is not None:
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index 70e78930e..6a268c864 100644
--- a/stdpopsim/slim_engine.py
+++ b/stdpopsim/slim_engine.py
@@ -1013,9 +1013,10 @@ def get_version(self):
 
     def simulate(
         self,
-        demographic_model=None,
-        contig=None,
-        samples=None,
+        demographic_model,
+        contig,
+        samples,
+        *,
         seed=None,
         mutation_types=None,
         extended_events=None,


Test patch:
diff --git a/tests/test_engines.py b/tests/test_engines.py
index 5c7315b3d..26b01ba6e 100644
--- a/tests/test_engines.py
+++ b/tests/test_engines.py
@@ -4,6 +4,7 @@
 import unittest
 
 import stdpopsim
+import msprime
 
 
 class TestEngineAPI(unittest.TestCase):
@@ -49,7 +50,8 @@ def test_get_engine(self):
 
     def test_abstract_base_class(self):
         e = stdpopsim.Engine()
-        self.assertRaises(NotImplementedError, e.simulate)
+        with self.assertRaises(NotImplementedError):
+            e.simulate(None, None, None)
         self.assertRaises(NotImplementedError, e.get_version)
 
 
@@ -68,3 +70,34 @@ def test_simulate_nonexistent_param(self):
             engine.simulate(**good_kwargs)
             with self.assertRaises(TypeError):
                 engine.simulate(**bad_kwargs)
+
+    def test_required_params(self):
+        species = stdpopsim.get_species("HomSap")
+        model = species.get_demographic_model("AshkSub_7G19")
+        contig = (species.get_contig("chr1"),)
+        for engine in stdpopsim.all_engines():
+            with self.assertRaises(TypeError):
+                engine.simulate(model, contig)
+
+    def test_msprime_kwargs(self):
+        species = stdpopsim.get_species("HomSap")
+        model = species.get_demographic_model("AshkSub_7G19")
+        contig = species.get_contig("chr22", length_multiplier=0.01)
+        samples = model.get_samples(10)
+        engine = stdpopsim.get_engine("msprime")
+        sim_arg = engine.simulate(
+            model, contig, samples, record_full_arg=True, random_seed=1
+        )
+        assert any(msprime.NODE_IS_RE_EVENT == sim_arg.tables.nodes.flags)
+
+    def test_msprime_seed(self):
+        species = stdpopsim.get_species("HomSap")
+        model = species.get_demographic_model("AshkSub_7G19")
+        contig = species.get_contig("chr22", length_multiplier=0.01)
+        samples = model.get_samples(10)
+        engine = stdpopsim.get_engine("msprime")
+        with self.assertRaises(ValueError):
+            engine.simulate(model, contig, samples, seed=1, random_seed=1)
+        sim_seed = engine.simulate(model, contig, samples, seed=1)
+        sim_random_seed = engine.simulate(model, contig, samples, random_seed=1)
+        self.assertEquals(sim_seed.tables.edges, sim_random_seed.tables.edges)