merge queue: checking #1829 on main (d53b008)

stdpopsim/engines.py

@@ -1,11 +1,13 @@
+import copy
 import logging
+import math
 import warnings
 
 import attr
 import msprime
-import stdpopsim
 import numpy as np
-import math
+
+import stdpopsim
 
 logger = logging.getLogger(__name__)
 
@@ -145,7 +147,7 @@ class _MsprimeEngine(Engine):
         )
     ]
 
-    # We default to the first model in the list.
+    # We default to ``hudson`` if no model is specified.
     model_class_map = {
         "hudson": msprime.StandardCoalescent,
         "dtwf": msprime.DiscreteTimeWrightFisher,
@@ -166,40 +168,73 @@ class _MsprimeEngine(Engine):
 
     @property
     def supported_models(self):
+        # NOTE: this list does not include models without a string alias.
+        # Only msprime models with a string alias are supported in the CLI
+        # (and this method is only referenced in the CLI logic).
         return list(self.model_class_map.keys())
 
-    def _convert_model_spec(self, model_str, model_changes):
+    def _convert_model_spec(self, model, model_changes):
         """
         Convert the specified model specification into a form suitable
-        for sim_ancestry. The model param is a string or None. The
-        model_changes is either None or list of (time, model_str) tuples.
+        for sim_ancestry. The model param is a string, msprime.AncestryModel
+        instance, or None. The model_changes is either None or list of
+        (time, string | msprime.AncestryModel) tuples.
         Also return the appropriate extra citations.
         """
+
+        def _from_model_spec(model):
+            # We default to the Hudson model if no model is specified
+            if model is None:
+                return _from_model_spec("hudson")
+            elif isinstance(model, str):
+                if model not in self.model_class_map:
+                    raise ValueError(f"Unrecognised model '{model}'")
+                return self.model_class_map[model](duration=None)
+            elif isinstance(model, msprime.AncestryModel):
+                model = copy.deepcopy(model)
+                model.duration = None
+                return model
+            raise TypeError("`model` must be a string or msprime.AncestryModel")
+
         citations = []
-        if model_str is None:
-            model_str = "hudson"
-        else:
-            if model_str not in self.model_class_map:
-                raise ValueError(f"Unrecognised model '{model_str}'")
-            if model_str in self.model_citations:
-                citations.extend(self.model_citations[model_str])
+        model = _from_model_spec(model)
+        if model.name in self.model_citations:
+            citations.extend(self.model_citations[model.name])
 
-        if model_changes is None:
-            model = model_str
-        else:
+        if model_changes is not None:
             model_list = []
             last_t = 0
-            last_model = model_str
-            for t, model in model_changes:
-                if model not in self.supported_models:
-                    raise ValueError(f"Unrecognised model '{model}'")
-                if model in self.model_citations:
-                    citations.extend(self.model_citations[model])
+            last_model = model
+
+            def _from_tuple(change):
+                if len(change) != 2:
+                    raise ValueError(
+                        "`model_changes` must be a list of (time, model) tuples"
+                    )
+                t, next_model = change
+                try:
+                    t = float(t)
+                    assert t >= 0 and np.isfinite(t)
+                except (ValueError, AssertionError, TypeError):
+                    raise ValueError(
+                        f"Specified time {t} is not a valid non-negative number"
+                    )
+                if next_model is None:
+                    raise ValueError("Unrecognised model: `None`")
+                return t, _from_model_spec(next_model)
+
+            for change in model_changes:
+                t, next_model = _from_tuple(change)
                 duration = t - last_t
-                model_list.append(self.model_class_map[last_model](duration=duration))
-                last_model = model
+                if duration < 0:
+                    raise ValueError("Tuples in `model_changes` must be sorted by time")
+                last_model.duration = duration
+                model_list.append(last_model)
+                last_model = next_model
+                if last_model.name in self.model_citations:
+                    citations.extend(self.model_citations[last_model.name])
                 last_t = t
-            model_list.append(self.model_class_map[last_model](duration=None))
+            model_list.append(last_model)
             model = model_list
 
         return model, citations
@@ -222,12 +257,15 @@ def simulate(
         for all engines.
 
         :param msprime_model: The msprime simulation model to be used.
-            One of ``hudson``, ``dtwf``, ``smc``, or ``smc_prime``.
+            One of ``hudson``, ``dtwf``, ``smc``, or ``smc_prime``, or a
+            ``msprime.AncestryModel`` instance.
             See msprime API documentation for details.
-        :type msprime_model: str
+        :type msprime_model: str or msprime.AncestryModel
         :param msprime_change_model: A list of (time, model) tuples, which
-            changes the simulation model to the new model at the time specified.
-        :type msprime_change_model: list of (float, str) tuples
+            changes the simulation model to the new model at the time
+            specified. Each model may be a string as in ``msprime_model``
+            or a ``msprime.AncestryModel`` instance.
+        :type msprime_change_model: list of (float, str or msprime.AncestryModel) tuples
         :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

tests/test_engines.py

@@ -2,10 +2,13 @@
 Tests for simulation engine infrastructure.
 """
 
-import stdpopsim
+import json
+
 import msprime
-import pytest
 import numpy as np
+import pytest
+
+import stdpopsim
 
 
 class TestEngineAPI:
@@ -167,3 +170,176 @@ def test_msprime_bad_samples(self):
                 contig=contig,
                 samples=samples,
             )
+
+    def check_execution_order(self, expected, msprime_model, msprime_change_model=None):
+        """
+        Small helper function to verify that the model execution order
+        recorded in the simulated tree sequence provenance matches `expected`.
+        Each entry in `expected` is a (ClassName, duration) tuple.
+        """
+        species = stdpopsim.get_species("AraTha")
+        engine = stdpopsim.get_engine("msprime")
+        kwargs = dict(
+            demographic_model=stdpopsim.PiecewiseConstantSize(species.population_size),
+            contig=species.get_contig(length=100),
+            samples={"pop_0": 5},
+            msprime_model=msprime_model,
+            msprime_change_model=msprime_change_model,
+        )
+        ts = engine.simulate(**kwargs)
+        prov = ts.provenance(0)
+        record = json.loads(prov.record)
+        models = record["parameters"]["model"]
+        # Flexible provenance parsing
+        if isinstance(models, dict):
+            models = [models]
+        actual = []
+        for model in models:
+            name = model["__class__"].replace("msprime.ancestry.", "")
+            duration = model["duration"]
+            actual.append((name, duration))
+        assert expected == actual
+
+    def test_check_execution_order(self):
+        # If no `msprime_change_model` is provided, we expect a single
+        # msprime ancestry model with duration=None
+        # msprime_model=None defaults to StandardCoalescent
+        self.check_execution_order([("StandardCoalescent", None)], None)
+        # string alias
+        self.check_execution_order([("StandardCoalescent", None)], "hudson")
+        # `msprime.AncestryModel`
+        self.check_execution_order(
+            [("StandardCoalescent", None)], msprime.StandardCoalescent()
+        )
+        self.check_execution_order([("SMCK", None)], msprime.SMCK(k=1))
+
+        # Test that msprime_change_model results in the correct sequence of
+        # msprime ancestry models with the corresponding durations
+        expected = [
+            ("DiscreteTimeWrightFisher", 20),
+            ("StandardCoalescent", None),
+        ]
+        # Input as string aliases
+        self.check_execution_order(expected, "dtwf", [(20, "hudson")])
+        # Input as `msprime.AncestryModel`
+        self.check_execution_order(
+            expected,
+            msprime.DiscreteTimeWrightFisher(),
+            [(20, msprime.StandardCoalescent())],
+        )
+        # Or a combination of both
+        model_changes = [
+            (20, msprime.SMCK(k=1)),
+            (21, msprime.SMCK(k=0)),
+            (100, "hudson"),
+        ]
+        expected = [
+            ("DiscreteTimeWrightFisher", 20),
+            ("SMCK", 21 - 20),
+            ("SMCK", 100 - 21),
+            ("StandardCoalescent", None),
+        ]
+        self.check_execution_order(expected, "dtwf", model_changes)
+
+    def test_msprime_ancestry_is_immutable(self):
+        engine = stdpopsim.get_engine("msprime")
+        species = stdpopsim.get_species("HomSap")
+        kwargs = dict(
+            demographic_model=species.get_demographic_model("AshkSub_7G19"),
+            contig=species.get_contig("chr1"),
+            samples={"YRI": 5, "CHB": 5, "CEU": 5},
+            dry_run=True,
+        )
+        # Test that we do not mutate the input model
+        model = msprime.SMCK(k=1, duration=10)
+        engine.simulate(
+            **kwargs,
+            msprime_model=model,
+            msprime_change_model=[(12, model), (99, model)],
+        )
+        assert model.duration == 10
+        engine.simulate(**kwargs, msprime_model=model)
+        assert model.duration == 10, "model should not be mutated"
+
+    def test_msprime_ancestry_model_errors(self):
+        engine = stdpopsim.get_engine("msprime")
+        species = stdpopsim.get_species("HomSap")
+        kwargs = dict(
+            demographic_model=species.get_demographic_model("AshkSub_7G19"),
+            contig=species.get_contig("chr1"),
+            samples={"YRI": 5, "CHB": 5, "CEU": 5},
+            dry_run=True,
+        )
+
+        # Fail if user provides an invalid msprime_model
+        class MyClass:
+            pass
+
+        invalid_models = [
+            "notvalid",
+            100,
+            MyClass(),
+            # Passing msprime_model as a list is not supported (although
+            # it would be a valid input for `msprime`).
+            # Instead, user should use the msprime_change_model argument.
+            [
+                msprime.DiscreteTimeWrightFisher(duration=10),
+                msprime.SMCK(k=1),
+            ],
+        ]
+
+        for invalid in invalid_models:
+            with pytest.raises((TypeError, ValueError)):
+                engine.simulate(**kwargs, msprime_model=invalid)
+
+        # Fail if user provides an invalid msprime_change_model
+        with pytest.raises(
+            ValueError,
+            match=r"`model_changes` must be a list of \(time, model\) tuples",
+        ):
+            engine.simulate(
+                **kwargs,
+                msprime_model="dtwf",
+                msprime_change_model=[
+                    ("hudson"),
+                ],
+            )
+
+        # Note: (10, None) is not a valid input msprime_change_model but
+        # msprime_model=None is
+        invalid_change_models = invalid_models + [None]
+        for invalid in invalid_change_models:
+            with pytest.raises((TypeError, ValueError)):
+                engine.simulate(
+                    **kwargs,
+                    msprime_model="dtwf",
+                    msprime_change_model=[
+                        (10, invalid),
+                    ],
+                )
+
+        invalid_times = [-1, None, "hudson", np.inf]
+        for invalid in invalid_times:
+            with pytest.raises(
+                ValueError,
+                match=r"Specified time .* is not a valid non-negative number",
+            ):
+                engine.simulate(
+                    **kwargs,
+                    msprime_model="dtwf",
+                    msprime_change_model=[
+                        (invalid, "hudson"),
+                    ],
+                )
+        with pytest.raises(
+            ValueError,
+            match=r"Tuples in `model_changes` must be sorted",
+        ):
+            engine.simulate(
+                **kwargs,
+                msprime_model="dtwf",
+                msprime_change_model=[
+                    (2, "hudson"),
+                    (1, "dtwf"),
+                ],
+            )