Merge pull request #24 from andrewkern/li_stephan_model

added Li and Stephan model, changes to example

dmel_example.py

@@ -7,18 +7,61 @@
 import stdpopsim
 from stdpopsim import drosophila_melanogaster
 
-chrom = drosophila_melanogaster.genome.chromosomes["chrX"]
+chrom = drosophila_melanogaster.genome.chromosomes["chr2R"]
 recomb_map = chrom.recombination_map()
 
+print("Testing Li and Stephan (2006) model")
+model = drosophila_melanogaster.LiStephanTwoPopulation()
+model.debug()
+
+samples = [msprime.Sample(population=0, time=0),
+           msprime.Sample(population=0, time=0),
+           msprime.Sample(population=1, time=0),
+           msprime.Sample(population=1, time=0)]
+
+
+ts = msprime.simulate(
+    samples=samples,
+    # recombination rate placeholder for quick runtime
+    recombination_rate=1e-09,
+    mutation_rate=chrom.mean_mutation_rate,
+    **model.asdict())
+print("simulated:", ts.num_trees, ts.num_sites)
+
+print("====================\n")
+print("Testing Sheehan and Song (2016) model")
 model = drosophila_melanogaster.SheehanSongThreeEpoch()
+
 model.debug()
 
-samples = [msprime.Sample(population=0, time=0),msprime.Sample(population=0,
-                                                               time=0)]
+samples = [msprime.Sample(population=0, time=0),
+           msprime.Sample(population=0, time=0)
+          ]
 
 ts = msprime.simulate(
     samples=samples,
+    # recombination rate placeholder for quick runtime
+    recombination_rate=1e-09,
+    mutation_rate=chrom.mean_mutation_rate,
+    **model.asdict())
+print("simulated:", ts.num_trees, ts.num_sites)
+
+print("====================\n")
+print("Testing Sheehan and Song (2016) model on chrom2R map.")
+print("This will take a while- perhaps 2 hours\n")
+model = drosophila_melanogaster.SheehanSongThreeEpoch()
+
+model.debug()
+
+samples = [msprime.Sample(population=0, time=0),
+           msprime.Sample(population=0, time=0)
+          ]
+
+ts = msprime.simulate(
+    samples=samples,
+    # actual recombination map; slow
     recombination_map=chrom.recombination_map(),
     mutation_rate=chrom.mean_mutation_rate,
     **model.asdict())
 print("simulated:", ts.num_trees, ts.num_sites)
+

stdpopsim/drosophila_melanogaster.py

@@ -111,3 +111,47 @@ def __init__(self):
                 time=t_2, initial_size=N_A, population_id=0)
         ]
         self.migration_matrix = [[0]]
+
+
+class LiStephanTwoPopulation(models.Model):
+    """
+    two population model of Li and Stephan (2006) for
+    African and European divergence
+
+    .. todo:: document this model, including the original publications
+        and clear information about what the different population indexes
+        mean.
+
+    """
+
+    def __init__(self):
+
+        # African Parameter values from "Demographic History of the African
+        # Population" section
+        N_A0 = 8.603e06
+        t_A0 = 6000  # generations
+        N_A1 = N_A0 / 5.0
+        # European Parameter values from "Demo History of Euro Population"
+        N_E0 = 1.075e06
+        N_E1 = 2200
+        t_AE = 1580  # generations
+        t_E1 = 1580 - 34
+        self.population_configurations = [
+            msprime.PopulationConfiguration(initial_size=N_A0),
+            msprime.PopulationConfiguration(initial_size=N_E0)
+        ]
+        self.demographic_events = [
+            # Size change at Euro bottleneck
+            msprime.PopulationParametersChange(
+                time=t_E1, initial_size=N_E1, population_id=1),
+            # Split
+            msprime.MassMigration(
+                time=t_AE, source=1, destination=0, proportion=1.0),
+            # African bottleneck
+            msprime.PopulationParametersChange(
+                time=t_A0, initial_size=N_A1, population_id=0)
+        ]
+        self.migration_matrix = [
+            [0, 0],
+            [0, 0],
+        ]

tests/test_drosophila_melanogaster.py

@@ -56,3 +56,24 @@ def test_debug_runs(self):
         model.debug(output)
         s = output.getvalue()
         self.assertGreater(len(s), 0)
+
+
+class TestLiStephanTwoPopulation(unittest.TestCase):
+    """
+    Basic tests for the LiStephanTwoPopulation model.
+    """
+
+    def test_simulation_runs(self):
+        model = drosophila_melanogaster.LiStephanTwoPopulation()
+        samples = [msprime.Sample(population=0, time=0),
+                   msprime.Sample(population=1, time=0)]
+        ts = msprime.simulate(
+            samples=samples, **model.asdict())
+        self.assertEqual(ts.num_populations, 2)
+
+    def test_debug_runs(self):
+        model = drosophila_melanogaster.LiStephanTwoPopulation()
+        output = io.StringIO()
+        model.debug(output)
+        s = output.getvalue()
+        self.assertGreater(len(s), 0)