Merge pull request #990 from izabelcavassim/patch-3

Adding lognormal distribution to DFEs

stdpopsim/ext/selection.py

@@ -18,8 +18,6 @@ def __attrs_post_init__(self):
         if self.dominance_coeff < 0:
             raise ValueError(f"Invalid dominance coefficient {self.dominance_coeff}.")
 
-        #       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
         # distribution_args here, and add unit tests.
         if self.distribution_type == "f":
@@ -41,12 +39,11 @@ def __attrs_post_init__(self):
                 )
         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:
+            if len(self.distribution_args) != 1:
                 raise ValueError(
                     "Exponentially-distributed sel. coefs. (distribution_type='e') "
-                    "use a (mean) parameterisation, requiring mean > 0."
+                    "use a (mean) parameterisation."
                 )
         elif self.distribution_type == "n":
             # An normally-distributed fitness effect (mean, standard deviation).
@@ -68,19 +65,34 @@ def __attrs_post_init__(self):
                     "Weibull-distributed sel. coef. (distribution_type='w') "
                     "use a (scale, shape) parameterisation, requiring parameters > 0."
                 )
+        elif self.distribution_type == "l":
+            # An lognormally-distributed fitness effect (logmean, sdmean).
+            # See Eidos documentation for rlnorm().
+            if len(self.distribution_args) != 2 or self.distribution_args[1] <= 0:
+                raise ValueError(
+                    "Lognormally-distributed sel. coefs. (distribution_type='l') "
+                    "use a (logmean, logsd) parameterisation, requiring logsd > 0."
+                )
+            self.distribution_type = "s"
+            # dealing with lognormal distribution
+            # (adding instead of multiplying the mean):
+            logmean = self.distribution_args[0]
+            logsd = self.distribution_args[1]
+            self.distribution_args = f'"return rlnorm(1, {logmean} + log(Q), {logsd});"'
         else:
             raise ValueError(
                 f"{self.distribution_type} is not a supported distribution type"
             )
 
-        # The index of the param in the distribution_args list that should be
+        # The index(s) of the param in the distribution_args list that should be
         # multiplied by Q when using --slim-scaling-factor Q.
         self.Q_scaled_index = {
-            "e": 0,  # mean
-            "f": 0,  # fixed value
-            "g": 0,  # mean
-            "n": 1,  # standard deviation
-            "w": 0,  # scale
+            "e": [0],  # mean
+            "f": [0],  # fixed value
+            "g": [0],  # mean
+            "n": [0, 1],  # mean and sd
+            "w": [0],  # scale
+            "s": [],  # script types should just printout arguments
         }[self.distribution_type]
 
 

stdpopsim/slim_engine.py

@@ -848,9 +848,16 @@ def matrix2str(
 
     # Mutation types.
     for i, m in enumerate(mutation_types):
-        distrib_args = [str(arg) for arg in m.distribution_args]
-        distrib_args[m.Q_scaled_index] = "Q * " + distrib_args[m.Q_scaled_index]
-        distrib_args = ", ".join(distrib_args)
+        if len(m.Q_scaled_index) >= 1:
+            distrib_args = [str(arg) for arg in m.distribution_args]
+            for j in m.Q_scaled_index:
+                distrib_args[m.Q_scaled_index[j]] = (
+                    "Q * " + distrib_args[m.Q_scaled_index[j]]
+                )
+            distrib_args = ", ".join(distrib_args)
+        # dealing with distributions given by "s" Eidos script:
+        else:
+            distrib_args = m.distribution_args
         printsc(
             f"    initializeMutationType({i}, {m.dominance_coeff}, "
             + f'"{m.distribution_type}", {distrib_args});'

tests/test_slim_engine.py

@@ -792,6 +792,26 @@ def test_unweighted_mutations_are_not_simulated_by_slim(self):
         )
         assert ts.num_sites == 0
 
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_that_DFE_lognornal_produce_sites(self):
+        contig = get_test_contig()
+        contig.mutation_types = [
+            stdpopsim.ext.MutationType(),
+            stdpopsim.ext.MutationType(
+                distribution_type="l", distribution_args=[0.01, 0.2]
+            ),
+        ]
+        contig.genomic_element_types[0].proportions = [0.5]
+        ts = slim_simulate_no_recap(
+            demographic_model=self.model,
+            contig=contig,
+            samples=self.samples,
+            slim_scaling_factor=10,
+            slim_burn_in=0.1,
+            dry_run=True,
+        )
+        assert ts.num_sites > 0
+
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_weighted_mutations_are_simulated_by_slim(self):
         contig = get_test_contig()
@@ -880,7 +900,7 @@ def test_distribution_type_e(self):
             )
 
     def test_bad_distribution_args_e(self):
-        for distribution_args in ([], [0], [-0.1], [0.1, 0.4, 0.5]):
+        for distribution_args in ([], [0, 1], [0.1, 0.4, 0.5]):
             with pytest.raises(ValueError):
                 stdpopsim.ext.MutationType(
                     distribution_type="e", distribution_args=distribution_args
@@ -912,6 +932,19 @@ def test_bad_distribution_args_w(self):
                     distribution_type="w", distribution_args=distribution_args
                 )
 
+    def test_distribution_type_l(self):
+        for distribution_args in ([-0.1, 0.2], [0.1, 0.1], [50, 50]):
+            stdpopsim.ext.MutationType(
+                distribution_type="l", distribution_args=distribution_args
+            )
+
+    def test_bad_distribution_args_l(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="l", distribution_args=distribution_args
+                )
+
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
 class TestDrawMutation(PiecewiseConstantSizeMixin):