def test_moments_output(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=0)
G.add_edge('root','pop1')
dg = demography.DemoGraph(G)
fs = dg.SFS(theta=1.0, pop_ids=['pop1'], sample_sizes=[10])
self.assertTrue(np.allclose(moments.Demographics1D.snm([10]), fs))
G = test_graph()
dg = demography.DemoGraph(G)
fs = dg.SFS(theta=1, pop_ids=dg.leaves, sample_sizes=[10,10])
def test_frozen(self):
order = 2
Ne = 100
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=.1, frozen=True)
G.add_node('B', nu=1, T=.1, frozen=True)
G.add_edges_from([('root','A'),('root','B')])
dg = demography.DemoGraph(G, Ne=Ne)
gens = 10
Ns = [[Ne] * gens]
Tmrcas_init = demography.tmrcas.steady_state_tmrca(dg, Ne, order)
Tmrcas = demography.tmrcas.evolve_t(Tmrcas_init, order, Ns, [[1]], [True])
self.assertTrue(Tmrcas[1] == Tmrcas_init[1]+gens)
self.assertTrue(np.isclose(Tmrcas[0]-Tmrcas[1]**2,
Tmrcas_init[0]-Tmrcas_init[1]**2))
Tmrcas = dg.tmrca(['A','B'], order=order)
self.assertTrue(np.isclose(Tmrcas[0], Tmrcas[1]))
self.assertTrue(np.isclose(Tmrcas[0], Tmrcas[2]))
self.assertTrue(np.isclose(Tmrcas[3], Tmrcas[4]))
self.assertTrue(np.isclose(Tmrcas[3], Tmrcas[5]))
order = 2
G = nx.DiGraph()
Ne = 100
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=.1, frozen=True)
G.add_node('B', nu=1, T=.1)
G.add_edges_from([('root','A'),('root','B')])
dg = demography.DemoGraph(G, Ne=100)
Tmrcas = dg.tmrca(['A','B'], order=order)
self.assertTrue(np.isclose(Tmrcas[0], Tmrcas[1]))
self.assertTrue(np.isclose(Tmrcas[0], Tmrcas[2]+8400))
self.assertTrue(np.isclose(Tmrcas[3], Tmrcas[4]))
self.assertTrue(np.isclose(Tmrcas[3], Tmrcas[5]+20))
order = 2
G = nx.DiGraph()
Ne = 100
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=2, T=.1, frozen=True)
G.add_node('B', nu=2, T=.1)
G.add_edges_from([('root','A'),('root','B')])
dg = demography.DemoGraph(G, Ne=100)
Tmrcas = dg.tmrca(['A','B'], order=order)
self.assertTrue(np.isclose(Tmrcas[3], 2*Ne*(1+dg.G.nodes['A']['T'])))
self.assertTrue(np.isclose(Tmrcas[4], 2*Ne*(1+dg.G.nodes['A']['T'])))
self.assertTrue(Tmrcas[5] > 2*Ne)
def test_theta_computation(self):
G = example_graph_simple()
Ne = 1e4
u = 1e-8
theta = 4 * Ne * u
dg = demography.DemoGraph(G, Ne=Ne, mutation_rate=u)
self.assertTrue(np.isclose(dg.get_theta(), theta))
dg2 = demography.DemoGraph(G,
Ne=Ne,
mutation_rate=u,
sequence_length=5)
self.assertTrue(np.isclose(dg2.get_theta(), theta * 5))
dg3 = demography.DemoGraph(G)
self.assertTrue(np.isclose(dg3.get_theta(), 1))
def test_msprime_from_graph(self):
G = test_graph()
dg = demography.DemoGraph(G)
dg.Ne = 1e4
pc, mm, de = dg.msprime_inputs()
self.assertTrue(len(pc) == 7)
self.assertTrue(len(mm) == 7)
def dg_without_selfing():
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=1)
G.add_node('B', nu=1, T=1)
G.add_edges_from([('root','A'), ('root','B')])
return demography.DemoGraph(G)
def test_initial_sizes(self):
Ne = 7310
dg = demography.DemoGraph(ooa(), Ne=Ne)
(pop_config, mig_mat, demo_events) = dg.msprime_inputs()
N_t = demography.coalescence_rates.pop_sizes(pop_config, demo_events)
self.assertTrue(N_t[0][5](0) == Ne * 5.77)
self.assertTrue(N_t[0][0](0) == Ne)
def test_check_pts_sample_size(self):
G = test_graph()
dg = demography.DemoGraph(G)
self.assertRaises(AssertionError,
dg.SFS, ['pop1', 'pop2'], [30, 30],
pts=20,
engine='dadi')
def test_migration_matrix(self):
Ne = 7310
dg = demography.DemoGraph(ooa(), Ne=Ne)
(pop_config, mig_mat, demo_events) = dg.msprime_inputs()
ms = demography.coalescence_rates.migration_matrices(pop_config,
mig_mat, demo_events)
self.assertTrue(np.all(ms[max(ms.keys())] == 0))
def test_get_samples(self):
G = test_graph()
dg = demography.DemoGraph(G)
pop_ids = ['pop1','pop2']
ns = [10,20]
samples = dg.msprime_samples(pop_ids, ns)
self.assertTrue(len(samples) == np.sum(ns))
def dg_with_pulse():
Ne = 100
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=0.1, pulse={('B', 0.5, 0.1)})
G.add_node('B', nu=1, T=0.1)
G.add_edges_from([('root','A'), ('root','B')])
return demography.DemoGraph(G, Ne=Ne)
def example_three_split():
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=1)
G.add_node('pop2', nu=1, T=1)
G.add_node('pop3', nu=1, T=1)
G.add_edges_from([('root', 'pop1'), ('root', 'pop2'), ('root', 'pop3')])
return demography.DemoGraph(G)
def test_test_graph(self):
G = test_graph()
dg = demography.DemoGraph(G)
fs_dadi = dg.SFS(['pop1', 'pop2'], [4, 4],
engine='dadi',
pts=[10, 20, 30])
fs_moments = dg.SFS(['pop1', 'pop2'], [4, 4], engine='moments')
self.assertTrue(np.allclose(fs_dadi, fs_moments, 0.2))
def simple_graph():
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=.5)
G.add_node('pop2', nu=1, T=.3)
G.add_edges_from([('root', 'pop1'), ('root', 'pop2')])
dg = demography.DemoGraph(G)
return dg
def test_initial_migration_matrix(self):
G = ooa()
dg = demography.DemoGraph(G, Ne=7310)
pop_configs, mig_mat, demo_events = dg.msprime_inputs(Ne=7300)
contemp_pops = [3, 4, 5]
self.assertTrue(
np.all([
mig_mat[i][j] == 0 for i in range(6) for j in range(6)
if (i not in contemp_pops or j not in contemp_pops)
]))
def test_extending_frozen_pops(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=.5)
G.add_node('pop2', nu=1, T=.3)
G.add_edges_from([('root','pop1'),('root','pop2')])
dg = demography.DemoGraph(G)
dg_aug = demography.integration.augment_with_frozen(dg, ['pop1','pop2'])
acc_times = demography.util.get_accumulated_times(dg_aug)
self.assertTrue(np.all(np.array(list(acc_times.values()))-list(acc_times.values())[0] < 1e-12))
def test_integration_with_marginalize(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=2, T=.2)
G.add_node('pop2', nu=.5, T=.1)
G.add_edges_from([('root', 'pop1'), ('root', 'pop2')])
dg = demography.DemoGraph(G)
fs = dg.SFS(['pop1'], [10], engine='dadi', pts=30)
self.assertTrue(len(fs) == 11)
self.assertTrue(fs.pop_ids == ['pop1'])
def test_pop_configs(self):
G = test_graph()
dg = demography.DemoGraph(G)
Ne = 1e4
growth_rates = demography.msprime_functions.get_population_growth_rates(dg, Ne)
population_configurations, pop_indexes = demography.msprime_functions.get_population_configurations(dg, ['pop1','pop2'], growth_rates, Ne)
self.assertTrue(population_configurations[pop_indexes['root']].initial_size == Ne)
self.assertTrue(population_configurations[pop_indexes['pop1']].initial_size == Ne)
self.assertTrue(population_configurations[pop_indexes['pop2']].initial_size == 3.0*Ne)
self.assertTrue(population_configurations[pop_indexes['pop2']].growth_rate == growth_rates['pop2'])
def dg_with_marginalize():
Ne = 100
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=0.2)
G.add_node('B0', nu=1, T=0.1)
G.add_node('B', nu=1, T=0.1)
G.add_node('C', nu=1, T=0.05)
G.add_edges_from([('root','A'), ('root','B0'), ('B0','B'), ('B0','C')])
return demography.DemoGraph(G, Ne=Ne)
def test_dadi_pass(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=0.0)
G.add_edge('root', 'pop1')
dg = demography.DemoGraph(G)
fs = dg.SFS(['pop1'], [10], engine='dadi', pts=30)
self.assertTrue(len(fs) == 11)
fs = dg.SFS(['pop1'], [10], engine='dadi', pts=[50, 60, 70])
self.assertTrue(np.allclose(fs[1:-1], 1. / np.linspace(1, 9, 9),
0.005))
def test_pulses(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=1, T=.5)
G.add_node('pop2', nu=1, T=.5, pulse={('pop1', .1, .05), ('pop1', .5, 0.1)})
G.add_edges_from([('root','pop1'),('root','pop2')])
dg = demography.DemoGraph(G)
(present_pops, integration_times, nus, migration_matrices, frozen_pops,
selfing_rates, events) = demography.integration.get_moments_arguments(dg)
lineages = demography.integration.get_number_needed_lineages(dg, ['pop1','pop2'], [10, 10], events)
fs = dg.SFS(['pop1','pop2'], [10,10])
def test_frozen(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=2, T=.2)
G.add_node('pop2', nu=.5, T=.1)
G.add_edges_from([('root', 'pop1'), ('root', 'pop2')])
dg = demography.DemoGraph(G)
fs = dg.SFS(['pop1', 'pop2'], [10, 2], engine='dadi', pts=30)
self.assertTrue(fs.ndim == 2)
self.assertTrue(fs.pop_ids == ['pop1', 'pop2'])
self.assertTrue(np.all(fs.sample_sizes == [10, 2]))
def test_get_growth_rates(self):
G = test_graph()
dg = demography.DemoGraph(G)
Ne = 1e4
growth_rates = demography.msprime_functions.get_population_growth_rates(dg, Ne)
self.assertTrue(growth_rates['root'] == 0)
self.assertTrue(growth_rates['A'] == 0)
self.assertTrue(growth_rates['C'] == 0)
self.assertTrue(growth_rates['pop1'] == 0)
self.assertTrue(growth_rates['pop2'] != 0)
self.assertTrue(np.isclose(3, 0.5*np.exp(growth_rates['pop2'] * 2*Ne*dg.G.nodes['pop2']['T'])))
def test_all_lineages_moved(self):
y = np.ones(21)
Ne = 7310
dg = demography.DemoGraph(ooa(), Ne=Ne)
(pop_config, mig_mat, demo_events) = dg.msprime_inputs()
pulses = demography.coalescence_rates.pulse_events(pop_config,
demo_events)
for gen in sorted(pulses.keys()):
for pulse in pulses[gen]:
y = pulse.dot(y)
self.assertTrue(y[0] == 21)
self.assertTrue(np.all(y[1:] == 0))
def test_merge_two_pops_result(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('pop1', nu=2, T=.1)
G.add_node('pop2', nu=.5, T=.1)
G.add_node('pop3', nu=1, T=.1)
G.add_edges_from([('root', 'pop1'), ('root', 'pop2')])
G.add_weighted_edges_from([('pop1', 'pop3', .7), ('pop2', 'pop3', .3)])
dg = demography.DemoGraph(G)
fs_dadi = dg.SFS(['pop3'], [10], engine='dadi', pts=[60, 80, 100])
fs_moments = dg.SFS(['pop3'], [10], engine='moments')
self.assertTrue(np.allclose(fs_dadi, fs_moments, .006))
def test_population_split_two_pop(self):
G = nx.DiGraph()
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=1, T=0)
G.add_node('B', nu=1, T=0)
G.add_edges_from([('root', 'A'), ('root', 'B')])
dg = demography.DemoGraph(G)
fs = dg.SFS(['A', 'B'], [10, 10], engine='dadi', pts=20)
fs = dg.SFS(['A', 'B'], [10, 10], engine='dadi', pts=[20, 25, 30])
fs = dg.SFS(['A', 'B'], [10, 10], engine='dadi', pts=[60, 80, 100])
fs_moments = moments.Demographics2D.snm([10, 10])
self.assertTrue(np.allclose(fs_moments, fs, .02))
def two_epoch(params, selfing=None, Ne=1000):
[nu, T] = params
G = nx.DiGraph()
if selfing is None:
G.add_node('root', nu=1, T=0)
G.add_node('A', nu=nu, T=T)
else:
G.add_node('root', nu=1, T=0, selfing=selfing)
G.add_node('A', nu=nu, T=T, selfing=selfing)
G.add_edge('root','A')
dg = demography.DemoGraph(G, Ne=Ne)
return dg
def test_lineages_needed(self):
G = test_graph()
dg = demography.DemoGraph(G)
(present_pops, integration_times, nus, migration_matrices, frozen_pops,
selfing_rates, events) = demography.integration.get_moments_arguments(dg)
pop_ids = ['pop1','pop2']
sample_sizes = [10, 20]
lineages = demography.integration.get_number_needed_lineages(dg,
pop_ids, sample_sizes, events)
self.assertTrue(lineages['pop1'] == sample_sizes[0])
self.assertTrue(lineages['pop2'] == sample_sizes[1])
self.assertTrue(lineages['root'] == 40)
def bottleneck_growth(params=None):
if params is None:
(nuB, TB, nuF, TF) = (0.5, 0.2, 3.0, 0.05)
else:
(nuB, TB, nuF, TF) = params
G = initialize_graph()
G.add_node('pop_B', nu=nuB, T=TB)
G.add_node('pop', nu0=nuB, nuF=nuF, T=TF)
edges = [('root', 'pop_B'), ('pop_B', 'B')]
G.add_edges_from(edges)
return demography.DemoGraph(G)
def browning_america(params=None, Ne=7310):
if params is None:
(nuAf, TAf, nuB, TB, nuEu0, nuEuF, nuAs0, nuAsF, nuAdm0, nuAdmF, TEuAs,
TAdm, mAfB, mAfEu, mAfAs,
mEuAs) = (1.98, 0.265, 0.255, 0.0766, 0.141, 4.66, 0.0758, 6.27, 4.10,
7.48, 0.0621, 0.0008, 2.19, 0.366, 0.114, 0.455)
else:
(nuAf, TAf, nuB, TB, nuEu0, nuEuF, nuAs0, nuAsF, nuAdm0, nuAdmF, TEuAs,
TAdm, mAfB, mAfEu, mAfAs, mEuAs) = params
tol = 1e-10 # offset admixture pulses so they occur in the correct order
frac_Eu_pulse = (TEuAs + tol) / (TEuAs + TAdm)
frac_As_pulse = (TEuAs + 2 * tol) / (TEuAs + TAdm)
G = initialize_graph()
G.add_node('A', nu=nuAf, T=TAf)
G.add_node('Af0',
nu=nuAf,
T=TB + TEuAs,
m={
'B': mAfB,
'EUR': mAfEu,
'ASIA': mAfAs
})
G.add_node('B', nu=nuB, T=TB, m={'Af0': mAfB})
G.add_node('AFR', nu=nuAf, T=TAdm, m={'EUR': mAfEu, 'ASIA': mAfAs})
G.add_node('EUR',
nu0=nuEu0,
nuF=nuEuF,
T=TEuAs + TAdm,
m={
'Af0': mAfEu,
'AFR': mAfEu,
'ASIA': mEuAs
},
pulse={('ADMIX', frac_Eu_pulse, 2. / 3)})
G.add_node('ASIA',
nu0=nuAs0,
nuF=nuAsF,
T=TEuAs + TAdm,
m={
'Af0': mAfAs,
'AFR': mAfAs,
'EUR': mEuAs
},
pulse={('ADMIX', frac_As_pulse, 1. / 2)})
G.add_node('ADMIX', nu0=nuAdm0, nuF=nuAdmF, T=TAdm)
edges = [('root', 'A'), ('A', 'Af0'), ('A', 'B'), ('B', 'EUR'),
('B', 'ASIA'), ('Af0', 'ADMIX'), ('Af0', 'AFR')]
G.add_edges_from(edges)
dg = demography.DemoGraph(G, Ne=Ne)
return dg
def split_IM(params=None):
if params is None:
(nuA, nu1, nu2, T, m12, m21) = (1.0, 2.0, 3.0, 0.2, 0.5, 2.0)
else:
(nuA, nu1, nu2, T, m12, m21) = params
G = initialize_graph()
G.nodes['root']['nu'] = nuA
G.add_node('pop1', nu=nu1, T=T, m={'pop2': m12})
G.add_node('pop2', nu=nu2, T=T, m={'pop1': m21})
edges = [('root', 'pop1'), ('root', 'pop2')]
G.add_edges_from(edges)
return demography.DemoGraph(G)
