take ordered file to mimic tbs in ms

Author: stsmall

project/sim_modules/readconfig.py

@@ -25,12 +25,19 @@ def read_config_sims(configFile, ms_path):
     contig_len = int(config.getfloat(sim, "contiglen"))
     num_loci = config.getint(sim, "loci")
 
-    effective_size = config.get(sim, "effective_population_size")
-    if "," in effective_size:
-        effective_size = list(map(float, effective_size.split(",")))
-        effective_size = list(map(int, effective_size))
+    ne_size = config.get(sim, "effective_population_size")
+    if "," in ne_size:
+        effective_size = list(map(float, ne_size.split(",")))
+        effective_size = list(map(int, ne_size))
+    elif ne_size[0].isalpha():
+        if os.path.exists(ne_size):
+            print(f"loading {ne_size} ...")
+            effective_size = list(np.loadtxt(ne_size))
+        else:
+            print(f"loading {os.path.join(config_path, ne_size)} ...")
+            effective_size = list(np.loadtxt(os.path.join(config_path, ne_size)))
     else:
-        effective_size = int(float(effective_size))
+        effective_size = int(float(ne_size))
 
     recomb_rate = config.get(sim, "recombination_rate")
     if recomb_rate:

project/sim_modules/sim_msprime.py

@@ -156,19 +156,15 @@ def run_simulation(param_df):
     # build demographic command line
     demo_events = demo_config(param_df)
 
-    ne0 = np.random.choice(scaled_Ne)
-    mu_t = np.random.choice(mu)
-    rec_t = np.random.choice(rec)
-    pfileout.write(f"{ne0}\t{mu_t}\t{rec_t}\n")
     # check demo
     if dry_run:
         checkDemo(pops, demo_events)
         return None
     else:
         trees = msp.simulate(
-                          Ne=ne0,
-                          recombination_rate=rec_t,
-                          mutation_rate=mu_t,
+                          Ne=param_df["ne_t"],
+                          recombination_rate=param_df["rec_t"],
+                          mutation_rate=param_df["mu_t"],
                           num_replicates=model_dt["loci"],
                           length=model_dt["contig_length"],
                           population_configurations=pops,
@@ -209,7 +205,8 @@ def checkDemo(pops, demo_events):
 
 
 def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
-                     outfile: str, nprocs: int, stats_config: str, dryrun: bool):
+                     outfile: str, nprocs: int, stats_config: str, dryrun: bool,
+                     order: bool):
     """Run code for simulating msprime.
 
     Parameters
@@ -260,61 +257,80 @@ def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
 
     # set mutation rate
     global mu
+    l_mu = np.nan
     mut_rate = model_dt["mutation_rate"]
     if type(mut_rate) is list:
         if len(mut_rate) == 2:
             low, high = mut_rate
             mu = np.random.uniform(low, high, sim_number)
         else:
             mu = mut_rate
+            if order:
+                l_mu = len(mu)
     else:
-        mu = [mut_rate]
+        mu = [mut_rate] * sim_number
 
     # set recombination rate
     global rec
+    l_rec = np.nan
     rec_rate = model_dt["recombination_rate"]
     if type(rec_rate) is list:
         if len(rec_rate) == 2:
             low, high = rec_rate
             rec = np.random.uniform(low, high, sim_number)
+            # rec = np.random.exponential(rec_rate, sim_number)
         else:
             rec = rec_rate
+            if order:
+                l_rec = len(rec)
     else:
-        # rec = np.random.exponential(rec_rate, sim_number)
-        rec = [rec_rate]
+        rec = [rec_rate] * sim_number
 
     # set ploidy
     global ploidy
     ploidy = model_dt["ploidy"]
 
     # set effective pop size
     global scaled_Ne
+    l_ne = np.nan
     effective_size = model_dt["eff_size"]
     if type(effective_size) is list:
         if len(effective_size) == 2:
             low, high = effective_size
             scaled_Ne = np.random.randint(low, high, sim_number) * ploidy
         else:
             scaled_Ne = effective_size
+            if order:
+                l_ne = len(scaled_Ne)
     else:
-        scaled_Ne = [effective_size * ploidy]
-    global pfileout
-    pfileout = open(f"{outfile}.ne_mu_rec.out", 'w')
+        scaled_Ne = [effective_size * ploidy] * sim_number
     # =========================================================================
     #  Main simulations
     # =========================================================================
     # set up generator fx for MP
+    if order:
+        l_min = np.nanmin([l_mu, l_rec, l_ne])
+        sim_number = int(l_min)
+        print(f"order requested, setting sim_number to shortest param file: {l_min}")
+
+    with open(f"{outfile}.ne_mu_rec.out", 'w') as pfile:
+        pfile.write("Ne\tmu\trec\n")
+        for i in range(sim_number):
+            pfile.write(f"{int(scaled_Ne[i])}\t{mu[i]}\t{rec[i]}\n")
+
     event = param_df["event"].values
     pops = param_df["pops"].values
     time_arr = list(zip(*param_df["time"].values))
     value_arr = list(zip(*param_df["value"].values))
-    param_gen = ({"time": time_arr[i], "event": event, "pops": pops, "value": value_arr[i]} for i in range(sim_number))
+    param_gen = ({"ne_t": scaled_Ne[i], "mu_t": mu[i], "rec_t": rec[i],
+                  "time": time_arr[i], "event": event, "pops": pops, "value": value_arr[i]} for i in range(sim_number))
+    # param_gen = ({"time": time_arr[i], "event": event, "pops": pops, "value": value_arr[i]} for i in range(sim_number))
     param_gen = list(param_gen)
+
     # check nprocs
     if nprocs > multiprocessing.cpu_count():  # check that there are not more requested than available
         print("not {nprocs} processors available, setting to {multiprocessing.cpu_count()}")
         nprocs = multiprocessing.cpu_count()
-
     # perform sims
     global stats_dt
     global header_len
@@ -349,4 +365,3 @@ def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
         pops_outfile.close()
     else:
         print("No stats file given with msprime")
-    pfileout.close()

run_sims.py

@@ -104,6 +104,8 @@ def parse_args(args_in):
     parser.add_argument("--set_priors", type=str,
                         help="provide a list of priors to be reused will overide"
                         " all other parameters")
+    parser.add_argument("--in_order", action="store_true",
+                        help="will draw from rec, mu, and Ne files in order")
     parser.add_argument("--nprocs", type=int, default=1,
                         help="processors for MP")
     parser.add_argument("--dryrun", action="store_true",
@@ -126,6 +128,7 @@ def main():
     outfile = args.out
     ploidy = args.ploidy
     priors_df = args.set_priors
+    order = args.in_order
     nprocs = args.nprocs
     dry_run = args.dryrun
     stats_config = args.stats_config
@@ -145,14 +148,15 @@ def main():
     # =========================================================================
     # parses the model file and draw all parameters
     demo_df, param_df = parse_model(model_file, sim_number)
+    # load and/or write priors
     if priors_df:
         assert os.path.exists(priors_df), "no priors file found"
         param_df = pd.read_csv(
             priors_df, header=[0, 1], index_col=0, skipinitialspace=True)
-    else:
-        # write priors
+    else:  # save priors
         priors_outfile = f"{sim_path}"
         write_params(param_df, priors_outfile, sim_number, dry_run)
+    # call sims
     if "discoal" in ms_path:
         simulate_discoal(ms_path, model_dt, demo_df, param_df, sim_number,
                          sim_path, nprocs, stats_config, dry_run)
@@ -167,7 +171,7 @@ def main():
                       sim_path, nprocs, stats_config, dry_run, args.approx)
     elif "msprime" in ms_path:
         simulate_msprime(model_dt, demo_df, param_df, sim_number,
-                         sim_path, nprocs, stats_config, dry_run)
+                         sim_path, nprocs, stats_config, dry_run, order)
 
 
 if __name__ == "__main__":