updates for simulating and plotting sawtooth. useful for benchmarking.

Author: terhorst

smcpp/commands/plot.py

@@ -74,18 +74,24 @@ def main(self, args):
             if len(args.offsets) != len(args.model):
                 raise RuntimeError("Please specify one offset per model")
         for fn, off in zip_longest(args.model, args.offsets, fillvalue=None):
-            if not os.path.exists(fn):
-                sys.exit("File not found: %s" % fn)
-            res = json.load(open(fn, "rt"))
-            if args.step_function:
-                mod = res["model"]
-                klass = getattr(model, mod["class"])
-                m = klass.from_dict(mod)
-                a = m.stepwise_values().astype("float")
-                s = m.s
-                d = {"a": m.stepwise_values(), "s": m.s, "N0": mod["N0"]}
+            if fn in ["human", "sawtooth"]:
+                label = None
+                p = getattr(util, fn)
+                d = {k: p[k] for k in "abs"}
+                d['N0'] = p['N0']
             else:
-                d = res
+                if not os.path.exists(fn):
+                    sys.exit("File not found: %s" % fn)
+                res = json.load(open(fn, "rt"))
+                if args.step_function:
+                    mod = res["model"]
+                    klass = getattr(model, mod["class"])
+                    m = klass.from_dict(mod)
+                    a = m.stepwise_values().astype("float")
+                    s = m.s
+                    d = {"a": m.stepwise_values(), "s": m.s, "N0": mod["N0"]}
+                else:
+                    d = res
             d["g"] = args.g
             psfs.append((d, off or 0))
         fig, series = plotting.plot_psfs(

smcpp/plotting.py

@@ -70,7 +70,7 @@ def g(x, y, label, data=data, **kwargs):
             y = np.concatenate([y, [a[-1], a[-1]]])
             # if not logy:
             #     y *= 1e-3
-            series.append((fn, x, y, my_axplot, off, N0, g))
+            series.append((None, x, y, my_axplot, off, N0, g))
         elif "model" in d:
             cls = getattr(model, d["model"]["class"])
             mb = cls.from_dict(d["model"])

util/make_datasets.py

@@ -7,10 +7,9 @@
 sys.path.append("util/")
 import multiprocessing
 import argparse
-import cPickle as pickle
+import pickle
 import os.path
 import vcf
-import cStringIO
 from collections import Counter
 
 import scrm, smcpp.util

util/sim_sawtooth.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+
+import msprime as msp
+from smcpp.util import sawtooth
+import sys
+
+if __name__ == "__main__":
+    sawtooth_events = [
+        (.000582262, 1318.18),
+        (.00232905, -329.546),
+        (.00931919, 82.3865),
+        (.0372648, -20.5966),
+        (.149059, 5.14916),
+        (0.596236, 0.),
+    ]
+    N0 = 14312
+    de = (
+            [msp.PopulationParametersChange(time=0, initial_size=5 * N0)] + 
+            [
+                msp.PopulationParametersChange(
+                time=4 * t * N0,
+                growth_rate=g / (4 * N0))
+                for t, g in sawtooth_events
+                ])
+    msp.DemographyDebugger(demographic_events=de).print_history()
+    msp.simulate(int(sys.argv[1]), length=1e8, mutation_rate=1.25e-8, recombination_rate=1e-9, 
+            demographic_events=de).write_vcf(open(sys.argv[2], "wt"), ploidy=2)