diff --git a/favites_lite.py b/favites_lite.py index 59d9dba..6fd966a 100755 --- a/favites_lite.py +++ b/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 ===") diff --git a/plugins/ancestral_sequence/generate_seq.py b/plugins/ancestral_sequence/generate_seq.py index dcda0fa..95bf0b0 100644 --- a/plugins/ancestral_sequence/generate_seq.py +++ b/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_')})) diff --git a/plugins/common.py b/plugins/common.py index 8026b61..6ac53f4 100644 --- a/plugins/common.py +++ b/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 diff --git a/plugins/contact_network/ngg.py b/plugins/contact_network/ngg.py index 36b0aa4..f3e1a65 100644 --- a/plugins/contact_network/ngg.py +++ b/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) diff --git a/plugins/mutation_rates/common_treeswift.py b/plugins/mutation_rates/common_treeswift.py index d1b1328..492003b 100644 --- a/plugins/mutation_rates/common_treeswift.py +++ b/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: diff --git a/plugins/sample_times/specific_time.py b/plugins/sample_times/specific_time.py index b9b4b1b..cedfd1b 100644 --- a/plugins/sample_times/specific_time.py +++ b/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) diff --git a/plugins/sample_times/state_entry.py b/plugins/sample_times/state_entry.py index 3545631..85929aa 100644 --- a/plugins/sample_times/state_entry.py +++ b/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) diff --git a/plugins/sample_times/time_windows.py b/plugins/sample_times/time_windows.py index ff670ed..e6fde9e 100644 --- a/plugins/sample_times/time_windows.py +++ b/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) diff --git a/plugins/sequence_evolution/seqgen.py b/plugins/sequence_evolution/seqgen.py index ab3919a..84211ce 100644 --- a/plugins/sequence_evolution/seqgen.py +++ b/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) diff --git a/plugins/transmission_network/gemf.py b/plugins/transmission_network/gemf.py index d9a56b2..e61d1ec 100644 --- a/plugins/transmission_network/gemf.py +++ b/plugins/transmission_network/gemf.py @@ -21,22 +21,12 @@ def gemf_determine_initial_states(params, input_cn_fn, initial_states_fn): f.close() # write infected states -def gemf_write_infected_states(model, infected_states_fn): - infected_states = { - "Granich": ['I1', 'I2', 'I3', 'I4', 'A1', 'A2', 'A3', 'A4'], - "Hethcote-Yorke": ['MIS', 'MIA', 'FIS', 'FIA'], - "SAIS": ['I'], - "SAAPPHIIRE": ['E', 'P1', 'P2', 'I1', 'I2', 'A1', 'A2'], - "SAPHIRE": ['E', 'P', 'I', 'A'], - "SEIR": ['I'], - "SI": ['I'], - "SIR": ['I'], - "SIS": ['I'], - } - if model not in infected_states: +def gemf_write_infected_states(model, infected_states_fn, GLOBAL): + TN_MODELS = GLOBAL["MODELS"]["Transmission Network"] + if model not in TN_MODELS: error("Invalid GEMF model: %s" % model) f = open(infected_states_fn, 'w') - for s in infected_states[model]: + for s in TN_MODELS[model]["INF_STATES"]: f.write('%s\n' % s) f.close() @@ -54,7 +44,7 @@ def gemf_write_transition_rates(params, transition_rates_fn): f.close() # simulate an epidemic using GEMF_FAVITES -def gemf_favites(model, params, out_fn, config, verbose=True): +def gemf_favites(model, params, out_fn, config, GLOBAL, verbose=True): # set things up intermediate_gemf = "%s/GEMF_FAVITES" % out_fn['intermediate']; makedirs(intermediate_gemf) initial_states_fn = "%s/initial_states.tsv" % intermediate_gemf @@ -64,7 +54,7 @@ def gemf_favites(model, params, out_fn, config, verbose=True): # run gemf_determine_initial_states(params, out_fn['contact_network'], initial_states_fn) - gemf_write_infected_states(model, infected_states_fn) + gemf_write_infected_states(model, infected_states_fn, GLOBAL) gemf_write_transition_rates(params, transition_rates_fn) command = [ 'GEMF_FAVITES.py', @@ -91,21 +81,21 @@ def gemf_favites(model, params, out_fn, config, verbose=True): print_log("All State Transitions written to: %s" % out_fn['all_state_transitions']) # model-specific functions -def gemf_favites_granich(params, out_fn, config, verbose=True): - gemf_favites("Granich", params, out_fn, config, verbose=verbose) -def gemf_favites_hethcote_yorke(params, out_fn, config, verbose=True): - gemf_favites("Hethcote-Yorke", params, out_fn, config, verbose=verbose) -def gemf_favites_sais(params, out_fn, config, verbose=True): - gemf_favites("SAIS", params, out_fn, config, verbose=verbose) -def gemf_favites_saapphiire(params, out_fn, config, verbose=True): - gemf_favites("SAAPPHIIRE", params, out_fn, config, verbose=verbose) -def gemf_favites_saphire(params, out_fn, config, verbose=True): - gemf_favites("SAPHIRE", params, out_fn, config, verbose=verbose) -def gemf_favites_seir(params, out_fn, config, verbose=True): - gemf_favites("SEIR", params, out_fn, config, verbose=verbose) -def gemf_favites_si(params, out_fn, config, verbose=True): - gemf_favites("SI", params, out_fn, config, verbose=verbose) -def gemf_favites_sir(params, out_fn, config, verbose=True): - gemf_favites("SIR", params, out_fn, config, verbose=verbose) -def gemf_favites_sis(params, out_fn, config, verbose=True): - gemf_favites("SIS", params, out_fn, config, verbose=verbose) +def gemf_favites_granich(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Granich et al. (2008)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_hethcote_yorke(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Hethcote and Yorke (1984)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_sais(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Susceptible-Alert-Infected-Susceptible (SAIS)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_saapphiire(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("SAAPPHIIRE", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_saphire(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("SAPHIRE", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_seir(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Susceptible-Exposed-Infected-Removed (SEIR)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_si(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Susceptible-Infected (SI)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_sir(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Susceptible-Infected-Removed (SIR)", params, out_fn, config, GLOBAL, verbose=verbose) +def gemf_favites_sis(params, out_fn, config, GLOBAL, verbose=True): + gemf_favites("Susceptible-Infected-Susceptible (SIS)", params, out_fn, config, GLOBAL, verbose=verbose) diff --git a/plugins/viral_phylogeny_seeds/common_treesap.py b/plugins/viral_phylogeny_seeds/common_treesap.py index 8664bb3..f329e22 100644 --- a/plugins/viral_phylogeny_seeds/common_treesap.py +++ b/plugins/viral_phylogeny_seeds/common_treesap.py @@ -12,7 +12,7 @@ error("Unable to import treeswift. Install with: pip install treeswift") # sample a seed tree with TreeSAP -def treesap_seed(model, params, out_fn, config, verbose=True): +def treesap_seed(model, params, out_fn, config, GLOBAL, verbose=True): chain_trees = [read_tree_newick(l) for l in open(out_fn['viral_phylogeny_all_chains_time']).read().strip().splitlines()] if model == "Coalescent (Neutral)": tree = coalescent_const_pop_tree(100., len(chain_trees), continuous=True) @@ -36,11 +36,11 @@ def treesap_seed(model, params, out_fn, config, verbose=True): print_log("Viral Phylogeny (Time) written to: %s" % out_fn['viral_phylogeny_time']) # model-specific -def treesap_coalescent_const_pop(params, out_fn, config, verbose=True): - treesap_seed("Coalescent (Neutral)", params, out_fn, config, verbose=verbose) -def treesap_dualbirth(params, out_fn, config, verbose=True): - treesap_seed("Dual-Birth", params, out_fn, config, verbose=verbose) -def treesap_nonhom_yule(params, out_fn, config, verbose=True): - treesap_seed("Non-Homogeneous Yule", params, out_fn, config, verbose=verbose) -def treesap_yule(params, out_fn, config, verbose=True): - treesap_seed("Yule", params, out_fn, config, verbose=verbose) +def treesap_coalescent_const_pop(params, out_fn, config, GLOBAL, verbose=True): + treesap_seed("Coalescent (Neutral)", params, out_fn, config, GLOBAL, verbose=verbose) +def treesap_dualbirth(params, out_fn, config, GLOBAL, verbose=True): + treesap_seed("Dual-Birth", params, out_fn, config, GLOBAL, verbose=verbose) +def treesap_nonhom_yule(params, out_fn, config, GLOBAL, verbose=True): + treesap_seed("Non-Homogeneous Yule", params, out_fn, config, GLOBAL, verbose=verbose) +def treesap_yule(params, out_fn, config, GLOBAL, verbose=True): + treesap_seed("Yule", params, out_fn, config, GLOBAL, verbose=verbose) diff --git a/plugins/viral_phylogeny_trans/coatran.py b/plugins/viral_phylogeny_trans/coatran.py index 0948e1b..7c0f15d 100644 --- a/plugins/viral_phylogeny_trans/coatran.py +++ b/plugins/viral_phylogeny_trans/coatran.py @@ -8,7 +8,7 @@ error("Unable to import treeswift. Install with: pip install treeswift") # simulate a coalescent viral phylogeny using CoaTran -def coatran(exe, params, out_fn, config, verbose=True): +def coatran(exe, params, out_fn, config, GLOBAL, verbose=True): if exe in {'coatran_inftime', 'coatran_transtree'}: command = [exe, out_fn['transmission_network'], out_fn['sample_times']] elif exe == 'coatran_constant': @@ -29,9 +29,9 @@ def coatran(exe, params, out_fn, config, verbose=True): print_log("Transmission Chain Viral Phylogenies (Time) written to: %s" % out_fn['viral_phylogeny_all_chains_time']) # model-specific plugins -def coatran_constant(params, out_fn, config, verbose=True): - coatran("coatran_constant", params, out_fn, config, verbose=verbose) -def coatran_inftime(params, out_fn, config, verbose=True): - coatran("coatran_inftime", params, out_fn, config, verbose=verbose) -def coatran_transtree(params, out_fn, config, verbose=True): - coatran("coatran_transtree", params, out_fn, config, verbose=verbose) +def coatran_constant(params, out_fn, config, GLOBAL, verbose=True): + coatran("coatran_constant", params, out_fn, config, GLOBAL, verbose=verbose) +def coatran_inftime(params, out_fn, config, GLOBAL, verbose=True): + coatran("coatran_inftime", params, out_fn, config, GLOBAL, verbose=verbose) +def coatran_transtree(params, out_fn, config, GLOBAL, verbose=True): + coatran("coatran_transtree", params, out_fn, config, GLOBAL, verbose=verbose)