def model_func(params):
Nanc, N_1F, r_2, N_2F, T_1, T_2 = params
model = momi.DemographicModel(N_e=1, gen_time=1, muts_per_gen=1e-8)
model.add_size_param('Nanc')
model.add_size_param("N_1F")
model.add_growth_param('r_2')
model.add_size_param("N_2F")
model.add_time_param("T_1")
model.add_time_param("T_2")
model.add_leaf("YRI", N="N_1F")
model.add_leaf("CEU", N="N_2F", g='r_2')
model.add_leaf("CHB", N="N_2F", g=0)
model.move_lineages("CHB", "CEU", t="T_1", N="N_2F", g="r_2")
model.move_lineages("CEU",
"YRI",
t=lambda params: 2 * params.T_1,
N="N_1F")
model.set_size("YRI",
N='Nanc',
g=0,
t=lambda params: params.T_2 + 2 * params.T_1)
model.set_params({
'Nanc': Nanc,
'N_1F': N_1F,
'r_2': r_2,
'N_2F': N_2F,
'T_1': T_1,
'T_2': T_2,
})
return model
def _simulate(self,
name,
N_e=1e6,
tau=20000,
epsilon=10,
num_replicates=100,
verbose=False):
model = momi.DemographicModel(N_e=N_e)
model.add_leaf(name)
## epsilon > 1 is bottleneck backwards in time
## epsilon < 1 is expansion
## epsilon == 1 is constant size
model.set_size(name, t=tau, N=N_e * epsilon)
sampled_n_dict = {name: self.paramsdict["nsamps"]}
if verbose: print(sampled_n_dict)
ac = model.simulate_data(
length=self.paramsdict["length"],
num_replicates=num_replicates,
recoms_per_gen=self.paramsdict["recoms_per_gen"],
muts_per_gen=self._sample_mu(),
sampled_n_dict=sampled_n_dict)
try:
sfs = ac.extract_sfs(n_blocks=1)
except ValueError:
## If _sample_mu() returns zero, or a very small value with respect to
## sequence length, Ne, and tau, then you can get a case where there
## are no snps in the data, and constructing the sfs freaks.
raise PTAError(
"Can't extract SFS from a simulation with no variation. Check that muts_per_gen looks reasonable."
)
return sfs
def model():
# nu1F, nu2B, nu2F, !m!, Tp, T
model = momi.DemographicModel(N_e=1, gen_time=25, muts_per_gen=2.35e-8)
model.add_size_param('N_A')
model.add_size_param("nu_1F", lower=1e-2, upper=100)
# model.add_size_param("nu_2B", lower=1e-2, upper=100)
model.add_growth_param('g_2')
model.add_size_param("nu_2F", lower=1e-2, upper=100)
model.add_time_param("tp", lower=0, upper=5)
model.add_time_param("t", lower=0, upper=5)
model.add_leaf("CEU", N=lambda x: x.nu_2F * x.N_A, g='g_2')
model.add_leaf("YRI", N=lambda x: x.nu_1F * x.N_A)
model.move_lineages("CEU",
"YRI",
t=lambda x: x.t * 50 * x.N_A,
N=lambda x: x.nu_1F * x.N_A)
model.set_size("YRI", N='N_A', g=0, t=lambda x: (x.t + x.tp) * 50 * x.N_A)
#model.set_params(randomize=True)
# [1.880, 0.0724, 1.764, !0.930!, 0.363, 0.112]
model.set_params({
'N_A': 7300,
'nu_1F': 1.880,
'g_2': np.log(1.764 / 0.0724) / (0.112 * 50 * 7300),
'nu_2F': 1.764,
'tp': 0.363,
't': 0.112
})
return model
def twopop_migration(sfs, pop_list, main_pop):
if len(pop_list) != 2:
sys.exit('pop_list does not contain 2 populations for 2pop model')
tp = pop_list
tp = [i for i in tp if i != main_pop]
new_pop = tp[0]
TWOMIGRATIONEVS = momi.DemographicModel(N_e=1e5)
TWOMIGRATIONEVS.set_data(sfs)
##
TWOMIGRATIONEVS.add_time_param("tdiv", upper=1e6)
TWOMIGRATIONEVS.add_time_param("tmig_" + main_pop + "_" + new_pop,
upper_constraints=['tdiv'])
##
TWOMIGRATIONEVS.add_size_param("n_" + main_pop, upper=1e7)
TWOMIGRATIONEVS.add_size_param("n_" + new_pop, upper=1e5)
##
TWOMIGRATIONEVS.add_leaf(main_pop, N="n_" + main_pop)
TWOMIGRATIONEVS.add_leaf(new_pop, N="n_" + new_pop)
##
TWOMIGRATIONEVS.add_pulse_param("mfrac_" + main_pop + "_" + new_pop,
upper=0.5)
##
TWOMIGRATIONEVS.move_lineages(new_pop, main_pop, t="tdiv")
TWOMIGRATIONEVS.move_lineages(new_pop,
main_pop,
t="tmig_" + main_pop + "_" + new_pop,
p="mfrac_" + main_pop + "_" + new_pop)
##
TWOMIGRATIONEVS.optimize(method="TNC")
return (TWOMIGRATIONEVS)
def threepop_w_migration_in_beach(sfs, pop_list, main_pop):
tp = pop_list
tp = [i for i in tp if i != main_pop]
pair1 = main_pop + "_" + tp[0]
pair2 = main_pop + "_" + tp[1]
THREEMIGBEACH = momi.DemographicModel(N_e=1e5)
THREEMIGBEACH.set_data(sfs)
THREEMIGBEACH.add_time_param("tdiv_" + pair1)
THREEMIGBEACH.add_time_param("tdiv_" + pair2)
THREEMIGBEACH.add_time_param(
"tmig_" + tp[0] + "_" + tp[1],
upper_constraints=["tdiv_" + pair1, "tdiv_" + pair2],
upper=1e5)
THREEMIGBEACH.add_leaf(main_pop)
THREEMIGBEACH.add_leaf(tp[0])
THREEMIGBEACH.add_leaf(tp[1])
THREEMIGBEACH.add_size_param("n_" + main_pop, 1e7)
THREEMIGBEACH.add_size_param("n_" + tp[0], 1e5)
THREEMIGBEACH.add_size_param("n_" + tp[1], 1e5)
THREEMIGBEACH.add_pulse_param("mfrac_" + tp[0] + "_" + tp[1])
THREEMIGBEACH.move_lineages(tp[0], main_pop, t="tdiv_" + pair1)
THREEMIGBEACH.move_lineages(tp[1], main_pop, t="tdiv_" + pair1)
THREEMIGBEACH.move_lineages(tp[0],
tp[1],
t="tmig_" + tp[0] + "_" + tp[1],
p="mfrac_" + tp[0] + "_" + tp[1])
THREEMIGBEACH.optimize(method="TNC")
return (THREEMIGBEACH)
def simple_five_pop_demo(x=np.random.normal(size=30)):
assert len(x) == 30
# make all params positive
x = np.exp(x)
# # allow negative growth rates
# for i in range(15,20):
# x[i] = np.log(x[i])
# # make times increasing
# for i in range(1,15):
# x[i] = x[i] + x[i-1]
t = np.cumsum(x[:15])
# allow negative growth rates
g = np.log(x[15:20])
model = momi.DemographicModel(1.0, .25)
for pop in range(1, 6):
model.add_leaf(pop)
model.set_size(5, t[0], g=g[0])
model.set_size(4, t[1], g=g[1])
model.set_size(3, t[2], g=g[2])
model.set_size(2, t[3], g=g[3])
model.set_size(1, t[4], g=g[4])
model.move_lineages(5, 4, t=t[5], N=x[20])
model.set_size(3, t=t[6], N=x[21])
model.set_size(2, t=t[7], N=x[22])
model.set_size(1, t[8], N=x[23])
model.move_lineages(4, 3, t[9], N=x[24])
model.set_size(2, t[10], N=x[25])
model.set_size(1, t[11], N=x[26])
model.move_lineages(3, 2, t[12], N=x[27])
model.set_size(1, t[13], N=x[28])
model.move_lineages(2, 1, t[14], N=x[29])
return model
def model_func(params):
"""
N_anc: Ancestral population size
N_1F: Size of ancestral population after growth and YRI population
N_2B: Bottleneck size of Eurasian population
nu_2F_gen: Final size of CEU population (in genetic units)
r2: Growth rate for CEU population
N_3F: Final size of CHB population
r3: Growth rate for CHB population
T1: Time of epoch between ancestral size growth and first split
T2: Time of second epoch between two splits
T3: Time of last epoch, time of last split
"""
N_Anc, N_1F, N_2B, nu_2F_gen, r2, N_3F, r3, T1, T2, T3 = params
model = momi.DemographicModel(N_e=1e5)
# momi2 uses backward in time models
# 1. Set leafs at time = 0
model.add_leaf("YRI", N=N_1F)
# we have nu_2F_gen is in genetic units so we translate it
model.add_leaf("CEU", N=nu_2F_gen * N_Anc, g=r2)
model.add_leaf("CHB", N=N_3F, g=r3)
# 2. Merge CHB to CEU (T3 time ago) and set size to bottleneck with g=0
model.move_lineages("CHB", "CEU", t=T3, N=N_1F, g=0)
# 3. Merge CEU to YRI (T3+T2 time ago)
model.move_lineages("CEU", "YRI", t=T3 + T2)
# 4. Change size of ancestral (YRI) population to N_anc (T3+T2+T1 time ago)
model.set_size("YRI", N=N_Anc, g=0, t=T3 + T2 + T1)
return model
def test_stochastic_jointime_inference(sampled_n=(5, 5, 5),
folded=False,
add_n=0,
use_theta=False,
theta=.1,
num_runs=10000):
t0 = random.uniform(.25, 2.5)
t1 = t0 + random.uniform(.5, 5.0)
num_bases = 1e3
theta = theta / num_bases
model = momi.DemographicModel(1, muts_per_gen=theta)
model.add_leaf(1)
model.add_leaf(2)
model.add_leaf(3)
#model.add_parameter("join_time", t0, scaled_lower=0.0, scaled_upper=t1)
#model.add_param("join_time", x0=t0, upper_x=t1)
model.add_time_param("join_time", t0, upper=t1)
model.move_lineages(1, 2, t="join_time")
model.move_lineages(2, 3, t=t1)
sampled_pops = (1, 2, 3)
data = model.simulate_data(num_bases,
0,
num_runs,
sampled_n_dict=dict(zip(sampled_pops,
sampled_n)))
sfs = data.extract_sfs(1)
assert sfs.n_snps() > 0
sfs = sfs._copy(sampled_n=np.array(sampled_n) + add_n)
if folded:
sfs = sfs.fold()
print((t0, t1))
#prim_log_lik = momi.likelihood._raw_log_lik
#prim_log_lik.reset_grad_count()
#assert not prim_log_lik.num_grad_calls()
if not use_theta:
model.set_mut_rate(None)
#model.set_x(random.uniform(0, t1), "join_time")
model.set_params({"join_time": random.uniform(0, t1)}, scaled=True)
model.set_data(sfs)
res = model.stochastic_optimize(n_minibatches=10,
num_iters=100,
svrg_epoch=10)
# make sure autograd is calling the rearranged gradient
#assert bool(prim_log_lik.num_grad_calls())
print(res.jac)
#assert abs(res.x - t0) / t0 < .05
#assert (model.get_params()["join_time"] - t0) / t0 < .05
assert (res.parameters["join_time"] - t0) / t0 < .05
def simple_three_pop_demo(t0, t1):
model = momi.DemographicModel(1., .25)
model.add_leaf(1)
model.add_leaf(2)
model.add_leaf(3)
model.move_lineages(1, 2, t0)
model.move_lineages(2, 3, t0 + t1)
return model
def test_P():
t1 = np.random.exponential(.25)
t2 = np.random.exponential(.25) + t1
t3 = np.random.exponential(.5) + t2
p1 = np.random.uniform(0, 1)
p2 = np.random.uniform(0, 1)
i = np.random.choice([0, 1])
j = 1 - i
demo0 = momi.DemographicModel(1.0, .25)
demo1 = momi.DemographicModel(1.0, .25)
for d in (demo0, demo1):
d.add_leaf(0)
d.add_leaf(1)
d.move_lineages(0, 1, t3)
demo0.move_lineages(0, 1, t=t1, p=p1)
demo0.move_lineages(i, j, t=t2, p=p2)
demo1.move_lineages(0, 'x', t=t1, p=p1)
demo1.move_lineages('x', 1, t=t1)
demo1.move_lineages(i, 'y', t=t2, p=p2)
demo1.move_lineages('y', j, t=t2)
demo0 = demo0._get_demo({0:5,1:6})
demo1 = demo1._get_demo({0:5,1:6})
#root_event = ('-ej', t3, 0, 1)
#pulse_events0 = [('-ep', t1, 0, 1, p1),
# ('-ep', t2, i, j, p2)]
#pulse_events1 = [('-ep', t1, 0, 'x', p1), ('-ej', t1, 'x', 1),
# ('-ep', t2, i, 'y', p2), ('-ej', t2, 'y', j)]
#demo0 = make_demography(pulse_events0 + [root_event],
# (0, 1), (5, 6))
#demo1 = make_demography(pulse_events1 + [root_event],
# (0, 1), (5, 6))
p = 20
vecs = [np.random.normal(size=(p, n + 1)) for n in demo0.sampled_n]
vals0, vals1 = [expected_sfs_tensor_prod(vecs, d)
for d in (demo0, demo1)]
assert np.allclose(vals0, vals1)
def exp_growth_model(x=np.random.normal(size=3)):
t, g, g2 = x
t, g2 = np.exp(t), np.exp(g2)
model = momi.DemographicModel(1.0, .25)
model.add_leaf(0, g=g)
model.set_size(0, t=t, g=g2)
model.set_size(0, t=3 * t, g=0)
return model
def simple_two_pop_demo(x=np.random.normal(size=4)):
x = [np.exp(xi) for xi in x]
model = momi.DemographicModel(1., .25)
model.add_leaf(1)
model.add_leaf(0)
model.set_size(1, t=0.0, N=x[1])
model.set_size(0, t=0.0, N=x[2])
model.move_lineages(0, 1, t=x[0])
model.set_size(1, t=x[0], N=x[3])
return model
def test_events_before_sample():
n_events = 4
t = [0.0]
for i in range(n_events):
t += [np.random.exponential(1. / float(n_events)) + t[-1]]
t = t[1:]
demo0 = momi.DemographicModel(1.0, .25)
demo1 = momi.DemographicModel(1.0, .25)
p = np.random.uniform(0, 1)
for d in (demo0, demo1):
d.add_leaf("a")
d.add_leaf("b", t=t[3])
d.move_lineages("a", "b", t=t[0], p=p)
demo0.set_size("c", t=0, N=10, g=1)
demo0.move_lineages("a", "c", t=t[1])
demo0.move_lineages("c", "b", t=t[2])
demo1.set_size("a", t=t[1], N=10*np.exp(-t[1]), g=1.0)
demo1.move_lineages("a", "b", t=t[2])
demo0, demo1 = [d._get_demo({"a":7,"b":5}) for d in (demo0, demo1)]
#events = [('-ep', t[0], 'a', 'b', np.random.uniform(0, 1))]
#demo0 = make_demography(events + [('-en', 0.0, 'c', 10.0), ('-eg', 0.0, 'c', 1.0),
# ('-ej', t[1], 'a', 'c'),
# ('-ej', t[2], 'c', 'b')],
# sampled_pops=('a', 'b'), sampled_n=(7, 5),
# sampled_t=(0., t[3]))
#demo1 = make_demography(events + [('-en', t[1], 'a', 10.0 * np.exp(-t[1])), ('-eg', t[1], 'a', 1.0),
# ('-ej', t[2], 'a', 'b')],
# sampled_pops=('a', 'b'), sampled_n=(7, 5),
# sampled_t=(0., t[3]))
vecs = [np.random.normal(size=(10, n + 1)) for n in demo0.sampled_n]
val0, val1 = [expected_sfs_tensor_prod(vecs, d) for d in (demo0, demo1)]
assert np.allclose(val0, val1)
def piecewise_constant_demo(x=np.random.normal(size=15)):
assert x.shape[0] % 2 == 1
model = momi.DemographicModel(1.0, .25)
model.add_leaf(0, N=np.exp(x[0]))
#events_list = [('-en', 0., 0, 1e4 * np.exp(x[0]))]
prev_time = 0.0
for i in range(int((x.shape[0] - 1) / 2)):
prev_time = np.exp(x[2 * i + 1]) + prev_time
N = np.exp(x[2 * i + 2])
model.set_size(0, t=prev_time, N=N)
return model
def model_func(params):
"""
Some model
"""
nuB, nuF, TB, TF = params
model = momi.DemographicModel(N_e=1, gen_time=1, muts_per_gen=1e-8)
model.add_leaf("Pop1", N=nuF)
model.set_size("Pop1", N=nuB, g=0, t=TF)
model.set_size("Pop1", N=10000, g=0, t=TB)
return model
def simple_admixture_demo(x=np.random.normal(size=7)):
t = np.cumsum(np.exp(x[:5]))
p = 1.0 / (1.0 + np.exp(x[5:]))
ret = momi.DemographicModel(1., .25)
ret.add_leaf("b")
ret.add_leaf("a")
ret.move_lineages("a", 2, t[1], p=1. - p[1])
ret.move_lineages("a", 3, t[0], p=1. - p[0])
ret.move_lineages(2, 3, t[2])
ret.move_lineages(3, "b", t[3])
ret.move_lineages("a", "b", t[4])
return ret
def twopop_split(sfs, pop_list, main_pop):
if len(pop_list) != 2:
sys.exit('pop_list does not contain 2 populations for 2pop model')
tp = pop_list
tp = [i for i in tp if i != main_pop]
new_pop = tp[0]
TWOSPLIT = momi.DemographicModel(N_e=1e5)
TWOSPLIT.set_data(sfs)
TWOSPLIT.add_time_param("tdiv", upper=1e6)
TWOSPLIT.add_leaf(main_pop)
TWOSPLIT.add_leaf(new_pop)
TWOSPLIT.move_lineages(new_pop, main_pop, t="tdiv")
TWOSPLIT.optimize(method="TNC")
return (TWOSPLIT)
def model_func(params):
nu21, nu22, t2, nu11, Nanc, t1 = params
model = momi.DemographicModel(N_e=1, gen_time=1, muts_per_gen=2.35e-08)
model.add_size_param('Nanc', lower=1.0, upper=100000000.0)
model.add_time_param('t1', lower=0.2, upper=10000000.0)
model.add_size_param('nu11', lower=1.0, upper=100000000.0)
model.add_pulse_param('s1', lower=0.001, upper=0.999)
model.add_time_param('t2', lower=0.2, upper=10000000.0)
model.add_size_param('nu21', lower=1.0, upper=100000000.0)
model.add_size_param('nu22', lower=1.0, upper=100000000.0)
model.add_leaf(
pop_name='YRI',
t=0,
N='nu21',
g=0,
)
model.add_leaf(
pop_name='CEU',
t=0,
N='nu22',
g=0,
)
model.move_lineages(
pop_from='CEU',
pop_to='YRI',
t='t2',
p=1,
N='nu11',
g=lambda params: np.log(params.nu11 / params.Nanc) / params.t1,
)
model.set_size(
pop_name='YRI',
t=lambda params: params.t2 + params.t1,
N='Nanc',
g=0,
)
model.set_params({
'nu21': nu21,
'nu22': nu22,
't2': t2,
'nu11': nu11,
'Nanc': Nanc,
't1': t1,
})
return model
def threepop_split_from_main(sfs, pop_list, main_pop):
tp = pop_list
tp = [i for i in tp if i != main_pop]
pair1 = main_pop + "_" + tp[0]
pair2 = main_pop + "_" + tp[1]
THREEPSFM = momi.DemographicModel(N_e=1e5)
THREEPSFM.set_data(sfs)
THREEPSFM.add_time_param("tdiv_" + pair1)
THREEPSFM.add_time_param("tdiv_" + pair2)
THREEPSFM.add_leaf(main_pop)
THREEPSFM.add_leaf(tp[0])
THREEPSFM.add_leaf(tp[1])
THREEPSFM.move_lineages(tp[0], main_pop, t="tdiv_" + pair1)
THREEPSFM.move_lineages(tp[1], main_pop, t="tdiv_" + pair2)
THREEPSFM.optimize(method="TNC")
return (THREEPSFM)
def simple_admixture_3pop(x=None):
if x is None:
x = np.random.normal(size=7)
t = np.cumsum(np.exp(x[:5]))
p = 1.0 / (1.0 + np.exp(x[5:]))
model = momi.DemographicModel(1., .25)
model.add_leaf("b")
model.add_leaf("a")
model.add_leaf("c")
model.move_lineages("a", "c", t[1], p=1. - p[1])
model.move_lineages("a", "d", t[0], p=1. - p[0])
model.move_lineages("c", "d", t[2])
model.move_lineages("d", "b", t[3])
model.move_lineages("a", "b", t[4])
return model
def test_archaic_and_pairwisediffs():
#logging.basicConfig(level=logging.DEBUG)
theta = 1
N_e = 1.0
join_time = 1.0
num_runs = 1000
def logit(p):
return np.log(p / (1. - p))
def expit(x):
return 1. / (1. + np.exp(-x))
n_bases = 1000
model = momi.DemographicModel(
N_e, muts_per_gen=theta/4./N_e/n_bases)
model.add_time_param(
"sample_t", random.uniform(0.001, join_time-.001) / join_time,
upper=join_time)
model.add_size_param("N", 1.0)
model.add_leaf("a", N="N")
model.add_leaf("b", t="sample_t", N="N")
model.move_lineages("a", "b", join_time)
data = model.simulate_data(length=n_bases,
recoms_per_gen=0,
num_replicates=num_runs,
sampled_n_dict={"a": 2, "b": 2})
model.set_data(data.extract_sfs(1),
use_pairwise_diffs=False,
mem_chunk_size=-1)
true_params = np.array(list(model.get_params().values()))
#model.set_x([logit(random.uniform(.001, join_time-.001) / join_time),
model.set_params([
logit(random.uniform(.001, join_time-.001) / join_time),
random.uniform(-1, 1)],
scaled=True)
res = model.optimize(method="trust-ncg", hessp=True)
inferred_params = np.array(list(model.get_params().values()))
assert np.max(np.abs(np.log(true_params / inferred_params))) < .2
def threepop_sizes(sfs, pop_list, main_pop):
tp = pop_list
tp = [i for i in tp if i != main_pop]
pair1 = main_pop + "_" + tp[0]
pair2 = main_pop + "_" + tp[1]
THREESIZES = momi.DemographicModel(N_e=1e5)
THREESIZES.set_data(sfs)
THREESIZES.add_time_param("tdiv_" + pair1)
THREESIZES.add_time_param("tdiv_" + pair2)
THREESIZES.add_leaf(main_pop)
THREESIZES.add_leaf(tp[0])
THREESIZES.add_leaf(tp[1])
THREESIZES.add_size_param("n_" + main_pop, 1e7)
THREESIZES.add_size_param("n_" + tp[0], 1e5)
THREESIZES.add_size_param("n_" + tp[1], 1e5)
THREESIZES.move_lineages(tp[0], main_pop, t="tdiv_" + pair1)
THREESIZES.move_lineages(tp[1], main_pop, t="tdiv_" + pair2)
THREESIZES.optimize(method="TNC")
return (THREESIZES)
def random_tree_demo(num_leaf_pops, lins_per_pop):
#events_list = []
sampled_pops = list(range(1, num_leaf_pops + 1))
model = momi.DemographicModel(1.0, .25)
for p in sampled_pops:
model.add_leaf(p, N=random.expovariate(1.0))
roots = list(sampled_pops)
#for i in roots:
# events_list += [('-en', 0.0, i, random.expovariate(1.0))]
t = 0.0
while len(roots) > 1:
i, j = random.sample(roots, 2)
t += random.expovariate(1.0)
#events_list += [('-ej', t, i, j),
# ('-en', t, j, random.expovariate(1.0))]
model.move_lineages(i, j, t, N=random.expovariate(1.0))
roots.remove(i)
#return make_demo_hist(events_list, sampled_pops, [lins_per_pop] * num_leaf_pops)
return model
def to_demography(self, print_events = False):
""" convert a networkx graph into a momi demography object """
model = momi.DemographicModel(N_e=1e4, gen_time=25,
muts_per_gen=1.25e-8)
leaves = self.get_leaf_nodes()
for l in leaves:
model.add_leaf(l, t=0)
self.set_leaves_attribute()
events = self.get_events()
leaves_dict = nx.get_node_attributes(self, 'leaves')
admixture_proportions = self.get_admixture_proportions()
for e in events:
n, t = e[0], e[1]
if self.nodes[n]['type'] == 'merge':
model.add_time_param(n)
children = list(self.successors(n))
model.move_lineages(leaves_dict[children[0]], leaves_dict[children[1]], t = n)
model.set_params({n: t})
if print_events:
print(f'move from {leaves_dict[children[0]]} to ' +
f'{leaves_dict[children[1]]} at t = {t:.2f}')
elif self.nodes[n]['type'] == 'admixture':
pass
#model.set_params({n: t})
else:
raise ValueError('event can only be admixture or merge')
for e in admixture_proportions:
p = admixture_proportions[e]
t = self.get_event_time(e[0])
admixed_edge_name = f'{e[0]}_{e[1]}'
model.add_time_param(admixed_edge_name)
model.add_pulse_param(f'{admixed_edge_name}_proportion')
model.move_lineages(leaves_dict[e[1]],
leaves_dict[e[0]], t = admixed_edge_name, p = f'{admixed_edge_name}_proportion')
model.set_params({admixed_edge_name: t, f'{admixed_edge_name}_proportion': p})
if print_events:
print(f'move from {leaves_dict[e[1]]} to ' +
f'{leaves_dict[e[0]]} at t = {t:.2f}' +
f' and proportion = {p}')
return model
def simple_nea_admixture_demo(N_chb_bottom,
N_chb_top,
pulse_t,
pulse_p,
ej_chb,
ej_yri,
sampled_n=(14, 10)):
ej_chb = pulse_t + ej_chb
ej_yri = ej_chb + ej_yri
G_chb = -np.log(N_chb_top / N_chb_bottom) / ej_chb
model = momi.DemographicModel(1., .25)
model.add_leaf("yri")
model.add_leaf("chb")
model.set_size("chb", 0., N=N_chb_bottom, g=G_chb)
model.move_lineages("chb", "nea", t=pulse_t, p=pulse_p)
model.move_lineages("chb", "yri", t=ej_chb)
model.move_lineages("yri", "nea", t=ej_yri)
return model
def test_underflow_robustness(folded):
num_runs = 1000
sampled_pops = (1, 2, 3)
sampled_n = (5, 5, 5)
n_bases = int(1e3)
demo = momi.DemographicModel(1.0, .25, muts_per_gen=2.5 / n_bases)
for p in sampled_pops:
demo.add_leaf(p)
demo.add_time_param("t0")
demo.add_time_param("t1", lower_constraints=["t0"])
demo.move_lineages(1, 2, "t0")
demo.move_lineages(2, 3, "t1")
true_params = np.array([0.5, 0.7])
demo.set_params(true_params)
data = demo.simulate_data(
length=n_bases,
recoms_per_gen=0.0,
num_replicates=num_runs,
sampled_n_dict=dict(zip(sampled_pops, sampled_n)))
sfs = data.extract_sfs(1)
if folded:
sfs = sfs.fold()
demo.set_data(sfs)
demo.set_params({"t0": 0.1, "t1": 100.0})
optimize_res = demo.optimize()
print(optimize_res)
inferred_params = np.array(list(demo.get_params().values()))
error = (true_params - inferred_params) / true_params
print("# Truth:\n", true_params)
print("# Inferred:\n", inferred_params)
print("# Max Relative Error: %f" % max(abs(error)))
print("# Relative Error:", "\n", error)
assert max(abs(error)) < .1
def model_func(params):
Nanc, N_1F, r_2, N_2F, Tp, T = params
model = momi.DemographicModel(N_e=1, gen_time=1, muts_per_gen=1e-8)
model.add_size_param('N_Anc')
model.add_size_param("N_1F")
model.add_growth_param('r_2')
model.add_size_param("N_2F")
model.add_time_param("Tp")
model.add_time_param("T")
model.add_leaf("YRI", N="N_1F")
model.add_leaf("CEU", N="N_2F", g='r_2')
model.move_lineages("CEU", "YRI", t="T", N="N_1F")
model.set_size("YRI", N='N_Anc', g=0, t="Tp")
model.set_params({
'N_Anc': Nanc,
'N_1F': N_1F,
'r_2': r_2,
'N_2F': N_2F,
'Tp': Tp,
'T': T,
})
return model
#!/home/sergio/miniconda2/envs/momi-py36/bin/ python
import momi
import logging
import pickle
logging.basicConfig(level=logging.INFO, filename="log_model3_platy.log")
sfs = momi.Sfs.load("platy_sfs.gz")
#Model 3 - isolation with migration without population size changes
platy_model3 = momi.DemographicModel(N_e=1e5,
gen_time=2.35,
muts_per_gen=2.5e-9)
platy_model3.set_data(sfs)
platy_model3.add_time_param("tdiv_AF_CEP", lower=5e3, upper=3e5)
platy_model3.add_time_param("tmig_AF_CEP",
upper_constraints=["tdiv_AF_CEP"],
lower=5e3)
platy_model3.add_pulse_param("mfrac_AF_CEP", upper=.2)
platy_model3.add_time_param("tmig_CEP_AF",
upper_constraints=["tdiv_AF_CEP"],
lower=5e3)
platy_model3.add_pulse_param("mfrac_CEP_AF", upper=.2)
platy_model3.add_leaf("AF", g=1e-5, N=4.88e5)
platy_model3.add_leaf("CEP", g=1e-5, N=1.28e5)
platy_model3.set_size("CEP", t="tdiv_AF_CEP", g=0)
## because only looking at this particular split, elected to change ranges to get
print("\nPRIORS")
print("MUT RATE: 2.21e-9")
print("ANCESTRAL NE: 300,000")
print("GEN TIME: 1")
print("DIV TIME RANGE: 500,000 to 2,500,000")
print("NE RANGE: 1,000 to 2,000,000")
print("MIGRATION RANGE: 0 to 0.5")
print("MIGRATION DATE RANGE: 0 to 1,000,000 [25,000 for SECC model]\n\n")
print("begin setting up models\n")
##### PURE ISOLATION MODEL #####
print("\nPure Isolation model (base model)")
pure_isolation_model = momi.DemographicModel(
N_e=300000, muts_per_gen=2.21e-9,
gen_time=1) ## why tho -- can you give it something?
pure_isolation_model.set_data(sfs)
## set up divergence times
pure_isolation_model.add_time_param("tdiv_sc", lower=500000, upper=2500000)
## set up effective population size
pure_isolation_model.add_size_param(
"ne_s", lower=1000,
upper=2000000) ## this is from Brian's paper on cardinals
pure_isolation_model.add_size_param("ne_c", lower=1000, upper=2000000)
## set up populations and phylogeny
pure_isolation_model.add_leaf("Son", N="ne_s")
pure_isolation_model.add_leaf("Chi", N="ne_c")
pure_isolation_model.move_lineages("Son", "Chi", t="tdiv_sc")
pure_isolation_model.set_params({
'tdiv_sc': 807803.1164182945,
import momi
from scipy.optimize import differential_evolution
data = momi.sfs_from_dadi("YRI.CEU.CHB.fs")
model = momi.DemographicModel(N_e=1e5, muts_per_gen=2.35e-8)
model.set_data(data, length=4.04*(10**6))
model.add_size_param("nuA")
model.add_size_param("nuEu")
model.add_size_param("nuAf")
model.add_size_param("nuAs")
model.add_size_param("nuB")
model.add_growth_param('gAs')
model.add_growth_param('gEu')
model.add_time_param("TEuAs")
model.add_time_param("TB")
model.add_time_param("TAf")
model.add_leaf("CEU", N="nuEu")
model.add_leaf("YRI", N="nuAf")
model.add_leaf("CHB", N="nuAs")
# set growth rates
model.set_size("CEU", N="nuEu", t=0, g='gEu') #lambda params: np.log(params.nuEu/params.nuEu0)/params.TEuAs)
model.set_size("CHB", N="nuAs", t=0, g='gAs') #lambda params: np.log(params.nuAs/params.nuAs0)/params.TEuAs)
# Eu and As coalescence