Refactored how GEMF plugins determine infected states

favites_lite.py

@@ -107,7 +107,7 @@ def validate_config(config):
             error("Step not implemented yet: %s" % step)
         if model not in PLUGIN_FUNCTIONS[step]:
             error("%s model not implemented yet: %s" % (step, model))
-        PLUGIN_FUNCTIONS[step][model](params, out_fn, config, verbose=verbose)
+        PLUGIN_FUNCTIONS[step][model](params, out_fn, config, GLOBAL, verbose=verbose)
     end_time = time()
     if verbose:
         print_log(); print_log("=== Completion ===")

plugins/ancestral_sequence/generate_seq.py

@@ -3,7 +3,7 @@
 from random import shuffle, uniform
 
 # random sequence with exact frequencies
-def exact_freqs(params, out_fn, config, verbose=True):
+def exact_freqs(params, out_fn, config, GLOBAL, verbose=True):
     probs = [params['p_%s' % n] for n in 'ACGT']; p_A, p_C, p_G, p_T = probs; k = params['k']
     if not check_props(probs):
         error("Invalid base frequencies: %s" % str({k:params[k] for k in params if k.startswith('p_')}))
@@ -15,7 +15,7 @@ def exact_freqs(params, out_fn, config, verbose=True):
         print_log("Ancestral Sequence written to: %s" % out_fn['ancestral_seq'])
 
 # random sequence with die roll at each position
-def seq_die_roll(params, out_fn, config, verbose=True):
+def seq_die_roll(params, out_fn, config, GLOBAL, verbose=True):
     probs = [params['p_%s' % n] for n in 'ACGT']; p_A, p_C, p_G, p_T = probs; k = params['k']
     if not check_props(probs):
         error("Invalid base frequencies: %s" % str({k:params[k] for k in params if k.startswith('p_')}))

plugins/common.py

@@ -5,7 +5,7 @@
 ZERO_THRESH = 0.00000000001
 
 # dummy plugin function
-def DUMMY_PLUGIN_FUNC(params, out_fn, config, verbose=True):
+def DUMMY_PLUGIN_FUNC(params, out_fn, config, GLOBAL, verbose=True):
     pass
 
 # return the current time as a string

plugins/contact_network/ngg.py

@@ -3,7 +3,7 @@
 from subprocess import call
 
 # simulate contact network using NiemaGraphGen
-def ngg(exe, params, out_fn, config, verbose=True):
+def ngg(exe, params, out_fn, config, GLOBAL, verbose=True):
     if exe == 'ngg_barabasi_albert':
         command = [exe, str(params['n']), str(params['m'])]
     elif exe == 'ngg_barbell':
@@ -28,21 +28,21 @@ def ngg(exe, params, out_fn, config, verbose=True):
     f.close()
     if verbose:
         print_log("Contact Network written to: %s" % out_fn['contact_network'])
-def ngg_barabasi_albert(params, out_fn, config, verbose=True):
-    ngg("ngg_barabasi_albert", params, out_fn, config, verbose=verbose)
-def ngg_barbell(params, out_fn, config, verbose=True):
-    ngg("ngg_barbell", params, out_fn, config, verbose=verbose)
-def ngg_complete(params, out_fn, config, verbose=True):
-    ngg("ngg_complete", params, out_fn, config, verbose=verbose)
-def ngg_cycle(params, out_fn, config, verbose=True):
-    ngg("ngg_cycle", params, out_fn, config, verbose=verbose)
-def ngg_empty(params, out_fn, config, verbose=True):
-    ngg("ngg_empty", params, out_fn, config, verbose=verbose)
-def ngg_erdos_renyi(params, out_fn, config, verbose=True):
-    ngg("ngg_erdos_renyi", params, out_fn, config, verbose=verbose)
-def ngg_newman_watts_strogatz(params, out_fn, config, verbose=True):
-    ngg("ngg_newman_watts_strogatz", params, out_fn, config, verbose=verbose)
-def ngg_path(params, out_fn, config, verbose=True):
-    ngg("ngg_path", params, out_fn, config, verbose=verbose)
-def ngg_ring_lattice(params, out_fn, config, verbose=True):
-    ngg("ngg_ring_lattice", params, out_fn, config, verbose=verbose)
+def ngg_barabasi_albert(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_barabasi_albert", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_barbell(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_barbell", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_complete(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_complete", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_cycle(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_cycle", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_empty(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_empty", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_erdos_renyi(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_erdos_renyi", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_newman_watts_strogatz(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_newman_watts_strogatz", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_path(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_path", params, out_fn, config, GLOBAL, verbose=verbose)
+def ngg_ring_lattice(params, out_fn, config, GLOBAL, verbose=True):
+    ngg("ngg_ring_lattice", params, out_fn, config, GLOBAL, verbose=verbose)

plugins/mutation_rates/common_treeswift.py

@@ -16,7 +16,7 @@
     error("Unable to import treeswift. Install with: pip install treeswift")
 
 # Autocorrelated Constant Increment
-def treeswift_autocorr_const_inc(params, out_fn, config, verbose=True):
+def treeswift_autocorr_const_inc(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time'])
     p = params['p']; p_over_2 = p/2; delta = params['delta']; R_min = params['R_min']; R_max = params['R_max']
     for node in tree.traverse_preorder():
@@ -39,7 +39,7 @@ def treeswift_autocorr_const_inc(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Autocorrelated Exponential
-def treeswift_autocorr_exp(params, out_fn, config, verbose=True):
+def treeswift_autocorr_exp(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time'])
     for node in tree.traverse_preorder():
         if node.is_root():
@@ -53,7 +53,7 @@ def treeswift_autocorr_exp(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Autocorrelated Gamma
-def treeswift_autocorr_gamma(params, out_fn, config, verbose=True):
+def treeswift_autocorr_gamma(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); v = params['v']
     for node in tree.traverse_preorder():
         if node.is_root():
@@ -71,7 +71,7 @@ def treeswift_autocorr_gamma(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Autocorrelated Log-Normal
-def treeswift_autocorr_lognorm(params, out_fn, config, verbose=True):
+def treeswift_autocorr_lognorm(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); v = params['v']
     for node in tree.traverse_preorder():
         if node.is_root():
@@ -85,7 +85,7 @@ def treeswift_autocorr_lognorm(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Constant
-def treeswift_constant(params, out_fn, config, verbose=True):
+def treeswift_constant(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); r = params['rate']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -95,7 +95,7 @@ def treeswift_constant(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Chi-Squared
-def treeswift_chisq(params, out_fn, config, verbose=True):
+def treeswift_chisq(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); k = params['k']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -105,7 +105,7 @@ def treeswift_chisq(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Exponential
-def treeswift_exp(params, out_fn, config, verbose=True):
+def treeswift_exp(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); b = params['beta']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -115,7 +115,7 @@ def treeswift_exp(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # F
-def treeswift_f_dist(params, out_fn, config, verbose=True):
+def treeswift_f_dist(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); d1 = params['d1']; d2 = params['d2']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -125,7 +125,7 @@ def treeswift_f_dist(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Gamma
-def treeswift_gamma(params, out_fn, config, verbose=True):
+def treeswift_gamma(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); k = params['k']; theta = params['theta']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -135,7 +135,7 @@ def treeswift_gamma(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Inverse Gaussian (Wald)
-def treeswift_inverse_gaussian(params, out_fn, config, verbose=True):
+def treeswift_inverse_gaussian(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); mu = params['mu']; lam = params['lambda']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -145,7 +145,7 @@ def treeswift_inverse_gaussian(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Log-Normal
-def treeswift_lognorm(params, out_fn, config, verbose=True):
+def treeswift_lognorm(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); mu = params['mu']; sigma = params['sigma']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -155,7 +155,7 @@ def treeswift_lognorm(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Noncentral Chi-Squared
-def treeswift_noncentral_chisq(params, out_fn, config, verbose=True):
+def treeswift_noncentral_chisq(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); k = params['k']; lam = params['lambda']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -165,7 +165,7 @@ def treeswift_noncentral_chisq(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Noncentral F
-def treeswift_noncentral_f(params, out_fn, config, verbose=True):
+def treeswift_noncentral_f(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); d1 = params['d1']; d2 = params['d2']; lam = params['lambda']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -175,7 +175,7 @@ def treeswift_noncentral_f(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Pareto
-def treeswift_pareto(params, out_fn, config, verbose=True):
+def treeswift_pareto(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); alpha = params['alpha']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -185,7 +185,7 @@ def treeswift_pareto(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Power
-def treeswift_power(params, out_fn, config, verbose=True):
+def treeswift_power(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); c = params['c']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -195,7 +195,7 @@ def treeswift_power(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Rayleigh
-def treeswift_rayleigh(params, out_fn, config, verbose=True):
+def treeswift_rayleigh(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); sigma = params['sigma']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -205,7 +205,7 @@ def treeswift_rayleigh(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Triangular
-def treeswift_triangular(params, out_fn, config, verbose=True):
+def treeswift_triangular(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); a = params['a']; b = params['b']; c = params['c']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -215,7 +215,7 @@ def treeswift_triangular(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Truncated Normal
-def treeswift_truncnorm(params, out_fn, config, verbose=True):
+def treeswift_truncnorm(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); mu = params['mu']; sigma = params['sigma']; a = params['a']; b = params['b']
     nodes = [node for node in tree.traverse_preorder() if node.edge_length is not None]
     rates = truncnorm.rvs(a=a, b=b, loc=mu, scale=sigma, size=len(nodes))
@@ -226,7 +226,7 @@ def treeswift_truncnorm(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Uniform
-def treeswift_uniform(params, out_fn, config, verbose=True):
+def treeswift_uniform(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); a = params['a']; b = params['b']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:
@@ -236,7 +236,7 @@ def treeswift_uniform(params, out_fn, config, verbose=True):
         print_log("Viral Phylogeny (Mutations) written to: %s" % out_fn['viral_phylogeny_mut'])
 
 # Weibull
-def treeswift_weibull(params, out_fn, config, verbose=True):
+def treeswift_weibull(params, out_fn, config, GLOBAL, verbose=True):
     tree = read_tree_newick(out_fn['viral_phylogeny_time']); shape = params['shape']; scale = params['scale']
     for node in tree.traverse_preorder():
         if node.edge_length is not None:

plugins/sample_times/specific_time.py

@@ -2,7 +2,7 @@
 from .. import *
 
 # sample individuals at a specific time
-def specific_time(params, out_fn, config, verbose=True):
+def specific_time(params, out_fn, config, GLOBAL, verbose=True):
     states_to_sample = {s.strip() for s in params['sampled_states'].split(',')}
     nodes_to_sample = set()
     for l in open(out_fn['all_state_transitions']):
@@ -22,6 +22,6 @@ def specific_time(params, out_fn, config, verbose=True):
         print_log("Sample Times written to: %s" % out_fn['sample_times'])
 
 # sample individuals at end time
-def end_time(params, out_fn, config, verbose=True):
+def end_time(params, out_fn, config, GLOBAL, verbose=True):
     p = params.copy(); p['sample_time'] = config["Transmission Network"]['param']['duration']
-    specific_time(p, out_fn, config, verbose=verbose)
+    specific_time(p, out_fn, config, GLOBAL, verbose=verbose)

plugins/sample_times/state_entry.py

@@ -25,7 +25,7 @@ def state_entry(params, out_fn, mode, verbose=True):
     f.close()
     if verbose:
         print_log("Sample Times written to: %s" % out_fn['sample_times'])
-def state_entry_first(params, out_fn, config, verbose=True):
+def state_entry_first(params, out_fn, config, GLOBAL, verbose=True):
     state_entry(params, out_fn, "first", verbose=verbose)
-def state_entry_all(params, out_fn, config, verbose=True):
+def state_entry_all(params, out_fn, config, GLOBAL, verbose=True):
     state_entry(params, out_fn, "all", verbose=verbose)

plugins/sample_times/time_windows.py

@@ -11,7 +11,7 @@
     error("Unable to import scipy. Install with: pip install scipy")
 
 # sample individuals according to some distribution in their time windows
-def time_windows(model, params, out_fn, config, verbose=True):
+def time_windows(model, params, out_fn, config, GLOBAL, verbose=True):
     end_time = config["Transmission Network"]['param']['duration']
     states_to_sample = {s.strip() for s in params['sampled_states'].split(',')}
     windows = dict() # windows[node] = list of [state, start, end]
@@ -52,11 +52,11 @@ def time_windows(model, params, out_fn, config, verbose=True):
     f.close()
 
 # model-specific functions
-def time_windows_trunc_expon(params, out_fn, config, verbose=True):
-    time_windows("Truncated Exponential", params, out_fn, config, verbose=verbose)
-def time_windows_trunc_gamma(params, out_fn, config, verbose=True):
-    time_windows("Truncated Gamma", params, out_fn, config, verbose=verbose)
-def time_windows_trunc_normal(params, out_fn, config, verbose=True):
-    time_windows("Truncated Normal", params, out_fn, config, verbose=verbose)
-def time_windows_uniform(params, out_fn, config, verbose=True):
-    time_windows("Uniform", params, out_fn, config, verbose=verbose)
+def time_windows_trunc_expon(params, out_fn, config, GLOBAL, verbose=True):
+    time_windows("Truncated Exponential", params, out_fn, config, GLOBAL, verbose=verbose)
+def time_windows_trunc_gamma(params, out_fn, config, GLOBAL, verbose=True):
+    time_windows("Truncated Gamma", params, out_fn, config, GLOBAL, verbose=verbose)
+def time_windows_trunc_normal(params, out_fn, config, GLOBAL, verbose=True):
+    time_windows("Truncated Normal", params, out_fn, config, GLOBAL, verbose=verbose)
+def time_windows_uniform(params, out_fn, config, GLOBAL, verbose=True):
+    time_windows("Uniform", params, out_fn, config, GLOBAL, verbose=verbose)

plugins/sequence_evolution/seqgen.py

@@ -8,7 +8,7 @@
     error("Unable to import treeswift. Install with: pip install treeswift")
 
 # run seq-gen
-def seqgen(mode, params, out_fn, config, verbose=True):
+def seqgen(mode, params, out_fn, config, GLOBAL, verbose=True):
     treestr = open(out_fn['viral_phylogeny_mut']).read().strip().replace('[&R] ','')
     if ',' not in treestr: # if one-node tree, add DUMMY 0-length leaf
         treestr = "(DUMMY:0,%s);" % treestr.replace('(','').replace(')','')[:-1]
@@ -52,9 +52,9 @@ def seqgen(mode, params, out_fn, config, verbose=True):
         print_log("Sequences written to: %s" % out_fn['sequences'])
 
 # model-specific functions
-def seqgen_gtr(params, out_fn, config, verbose=True):
-    seqgen("GTR", params, out_fn, config, verbose=verbose)
-def seqgen_gtr_codon(params, out_fn, config, verbose=True):
-    seqgen("GTR+Codon", params, out_fn, config, verbose=verbose)
-def seqgen_gtr_gamma(params, out_fn, config, verbose=True):
-    seqgen("GTR+G", params, out_fn, config, verbose=verbose)
+def seqgen_gtr(params, out_fn, config, GLOBAL, verbose=True):
+    seqgen("GTR", params, out_fn, config, GLOBAL, verbose=verbose)
+def seqgen_gtr_codon(params, out_fn, config, GLOBAL, verbose=True):
+    seqgen("GTR+Codon", params, out_fn, config, GLOBAL, verbose=verbose)
+def seqgen_gtr_gamma(params, out_fn, config, GLOBAL, verbose=True):
+    seqgen("GTR+G", params, out_fn, config, GLOBAL, verbose=verbose)