Merge pull request #981 from izabelcavassim/patch-3

Adding other DFE distributions (e, n, and w)

stdpopsim/ext/selection.py

@@ -18,7 +18,6 @@ def __attrs_post_init__(self):
         if self.dominance_coeff < 0:
             raise ValueError(f"Invalid dominance coefficient {self.dominance_coeff}.")
 
-        # TODO: Add more distribution types like "e", "n", and "w".
         #       We probably shouldn't support "s" because it takes an
         #       arbitrary Eidos code string as an argument.
         # To add a new distribution type: validate the
@@ -37,9 +36,38 @@ def __attrs_post_init__(self):
             # See Eidos documentation for rgamma().
             if len(self.distribution_args) != 2 or self.distribution_args[1] <= 0:
                 raise ValueError(
-                    "Gamma-distributed mutation types (distribution_type='g') "
+                    "Gamma-distributed sel. coefs. (distribution_type='g') "
                     "use a (mean, shape) parameterisation, requiring shape > 0."
                 )
+        elif self.distribution_type == "e":
+            # An exponentially-distributed fitness effect (mean).
+            # A negative value for the mean is permitted.
+            # See Eidos documentation for rexp().
+            if len(self.distribution_args) != 1 or self.distribution_args[0] <= 0:
+                raise ValueError(
+                    "Exponentially-distributed sel. coefs. (distribution_type='e') "
+                    "use a (mean) parameterisation, requiring mean > 0."
+                )
+        elif self.distribution_type == "n":
+            # An normally-distributed fitness effect (mean, standard deviation).
+            # See Eidos documentation for rnorm().
+            if len(self.distribution_args) != 2 or self.distribution_args[1] <= 0:
+                raise ValueError(
+                    "Normally-distributed sel. coefs. (distribution_type='n') "
+                    "use a (mean, sd) parameterisation, requiring sd > 0."
+                )
+        elif self.distribution_type == "w":
+            # A Weibull-distributed fitness effect (scale, shape).
+            # See Eidos documentation for rweibull().
+            if (
+                len(self.distribution_args) != 2
+                or self.distribution_args[0] <= 0
+                or self.distribution_args[1] <= 0
+            ):
+                raise ValueError(
+                    "Weibull-distributed sel. coef. (distribution_type='w') "
+                    "use a (scale, shape) parameterisation, requiring parameters > 0."
+                )
         else:
             raise ValueError(
                 f"{self.distribution_type} is not a supported distribution type"

tests/test_slim_engine.py

@@ -873,6 +873,45 @@ def test_bad_distribution_args_g(self):
                     distribution_type="g", distribution_args=distribution_args
                 )
 
+    def test_distribution_type_e(self):
+        for distribution_args in ([0.1], [10], [5000]):
+            stdpopsim.ext.MutationType(
+                distribution_type="e", distribution_args=distribution_args
+            )
+
+    def test_bad_distribution_args_e(self):
+        for distribution_args in ([], [0], [-0.1], [0.1, 0.4, 0.5]):
+            with pytest.raises(ValueError):
+                stdpopsim.ext.MutationType(
+                    distribution_type="e", distribution_args=distribution_args
+                )
+
+    def test_distribution_type_n(self):
+        for distribution_args in ([-0.1, 0.2], [0.1, 0.1], [50, 50]):
+            stdpopsim.ext.MutationType(
+                distribution_type="n", distribution_args=distribution_args
+            )
+
+    def test_bad_distribution_args_n(self):
+        for distribution_args in ([], [0.1, -1], [0.1, 0.4, 0.5], [0.1]):
+            with pytest.raises(ValueError):
+                stdpopsim.ext.MutationType(
+                    distribution_type="n", distribution_args=distribution_args
+                )
+
+    def test_distribution_type_w(self):
+        for distribution_args in ([0.1, 0.2], [0.1, 0.1], [50, 50]):
+            stdpopsim.ext.MutationType(
+                distribution_type="w", distribution_args=distribution_args
+            )
+
+    def test_bad_distribution_args_w(self):
+        for distribution_args in ([], [-0.1, 1], [0.1, -1], [0.1, 0.4, 0.5], [0.1]):
+            with pytest.raises(ValueError):
+                stdpopsim.ext.MutationType(
+                    distribution_type="w", distribution_args=distribution_args
+                )
+
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
 class TestDrawMutation(PiecewiseConstantSizeMixin):