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)