Skip to content

Commit a000f35

Browse files
GertjanBisschopGertjanBisschop
committed
first pass
1 parent d05b851 commit a000f35

7 files changed

Lines changed: 335 additions & 40 deletions

File tree

‎algorithms.py‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1140,6 +1140,10 @@ def dtwf_generation(self):
11401140
for h in H:
11411141
segments_to_merge = len(h)
11421142
if segments_to_merge == 1:
1143+
if self.store_unary:
1144+
# FIXME: should be flagged as a PASS_TRHOUGH_EVENT
1145+
new_node_id = self.store_node(pop_idx)
1146+
self.store_arg_edges(h[0][1], new_node_id)
11431147
h = []
11441148
elif segments_to_merge >= 2:
11451149
for _, individual in h:
@@ -1247,8 +1251,8 @@ def dtwf_climb_pedigree(self):
12471251
for ploid in range(ind.ploidy):
12481252
self.process_pedigree_common_ancestors(ind, ploid)
12491253

1250-
def store_arg_edges(self, segment, u=None):
1251-
if u is None:
1254+
def store_arg_edges(self, segment, u=-1):
1255+
if u == -1:
12521256
u = len(self.tables.nodes) - 1
12531257
# Store edges pointing to current node to the left
12541258
x = segment
@@ -1724,6 +1728,7 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1):
17241728
pop = self.P[pop_id]
17251729
defrag_required = False
17261730
coalescence = False
1731+
store_edge = False
17271732
alpha = None
17281733
z = None
17291734
merged_head = None
@@ -1802,9 +1807,20 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1):
18021807
self.set_segment_mass(alpha)
18031808
z = alpha
18041809
if self.full_arg:
1810+
store_edge = True
18051811
if not coalescence:
1806-
self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)
1807-
self.store_arg_edges(z)
1812+
new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)
1813+
1814+
elif self.store_unary:
1815+
if new_node_id != -1:
1816+
store_edge = True
1817+
elif self.model == "dtwf":
1818+
# PASS_THROUGH_EVENT has been dealt with in merge_n_ancestors
1819+
new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)
1820+
store_edge = True
1821+
1822+
if store_edge:
1823+
self.store_arg_edges(z, new_node_id)
18081824
if defrag_required:
18091825
self.defrag_segment_chain(z)
18101826
if coalescence:
@@ -1853,6 +1869,8 @@ def merge_two_ancestors(self, population_index, label, x, y):
18531869
z = None
18541870
coalescence = False
18551871
defrag_required = False
1872+
store_edge = False
1873+
u = -1
18561874
while x is not None or y is not None:
18571875
alpha = None
18581876
if x is None or y is None:
@@ -1939,10 +1957,17 @@ def merge_two_ancestors(self, population_index, label, x, y):
19391957
z = alpha
19401958

19411959
if self.full_arg:
1960+
store_edge = True
19421961
if not coalescence:
19431962
u = self.store_node(population_index, flags=msprime.NODE_IS_CA_EVENT)
1944-
self.store_arg_edges(z, u)
1945-
elif self.store_unary and coalescence:
1963+
elif self.store_unary:
1964+
if u != -1:
1965+
store_edge = True
1966+
elif self.model == "dtwf":
1967+
u = self.store_node(population_index)
1968+
store_edge = True
1969+
1970+
if store_edge:
19461971
self.store_arg_edges(z, u)
19471972
if defrag_required:
19481973
self.defrag_segment_chain(z)

‎lib/msprime.c‎

Lines changed: 61 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2784,6 +2784,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l
27842784
int ret = 0;
27852785
bool coalescence = false;
27862786
bool defrag_required = false;
2787+
bool store_edge = false;
27872788
tsk_id_t v;
27882789
double l, r, l_min, r_max;
27892790
avl_node_t *node;
@@ -2938,26 +2939,39 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l
29382939
}
29392940
}
29402941
if (self->store_full_arg) {
2942+
store_edge = true;
29412943
if (!coalescence) {
29422944
ret = msp_store_node(
29432945
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
29442946
if (ret < 0) {
29452947
goto out;
29462948
}
29472949
}
2948-
// is specified new_node_id impossible when recording full_arg?
2949-
ret = msp_store_arg_edges(self, z, TSK_NULL);
2950-
if (ret != 0) {
2951-
goto out;
2952-
}
2950+
29532951
} else {
2954-
if (self->store_unary && coalescence) {
2955-
ret = msp_store_arg_edges(self, z, new_node_id);
2956-
if (ret < 0) {
2957-
goto out;
2952+
if (self->store_unary) {
2953+
if (new_node_id != TSK_NULL) {
2954+
store_edge = true;
2955+
} else {
2956+
if (self->model.type == MSP_MODEL_DTWF) {
2957+
new_node_id = msp_store_node(
2958+
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
2959+
ret = (int) new_node_id;
2960+
if (ret < 0) {
2961+
goto out;
2962+
}
2963+
store_edge = true;
2964+
}
29582965
}
29592966
}
29602967
}
2968+
2969+
if (store_edge) {
2970+
ret = msp_store_arg_edges(self, z, new_node_id);
2971+
if (ret != 0) {
2972+
goto out;
2973+
}
2974+
}
29612975
if (defrag_required) {
29622976
ret = msp_defrag_segment_chain(self, z);
29632977
if (ret != 0) {
@@ -3020,6 +3034,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
30203034
int ret = MSP_ERR_GENERIC;
30213035
bool coalescence = false;
30223036
bool defrag_required = false;
3037+
bool store_edge = false;
30233038
uint32_t j, h;
30243039
double l, r, r_max, next_l, l_min;
30253040
avl_node_t *node;
@@ -3180,14 +3195,35 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
31803195
}
31813196
}
31823197
if (self->store_full_arg) {
3198+
store_edge = true;
31833199
if (!coalescence) {
31843200
ret = msp_store_node(
31853201
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, individual);
31863202
if (ret < 0) {
31873203
goto out;
31883204
}
31893205
}
3190-
ret = msp_store_arg_edges(self, z, TSK_NULL);
3206+
} else {
3207+
if (self->store_unary) {
3208+
if (new_node_id != TSK_NULL) {
3209+
store_edge = true;
3210+
} else {
3211+
if (self->model.type == MSP_MODEL_DTWF) {
3212+
// PASS_THROUGH_EVENT have been dealth with in merge_n_ancestors
3213+
new_node_id = msp_store_node(self, MSP_NODE_IS_CA_EVENT, self->time,
3214+
population_id, individual);
3215+
ret = (int) new_node_id;
3216+
if (ret < 0) {
3217+
goto out;
3218+
}
3219+
store_edge = true;
3220+
}
3221+
}
3222+
}
3223+
}
3224+
3225+
if (store_edge) {
3226+
ret = msp_store_arg_edges(self, z, new_node_id);
31913227
if (ret != 0) {
31923228
goto out;
31933229
}
@@ -3246,6 +3282,21 @@ msp_merge_n_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
32463282

32473283
if (num_common_ancestors == 1) {
32483284
merged_head = msp_priority_queue_pop(self, Q);
3285+
if (self->store_unary) {
3286+
if (self->model.type == MSP_MODEL_DTWF) {
3287+
new_node_id = msp_store_node(
3288+
/*FIXME: not a CA_EVENT but PASS_THROUGH_EVENT*/
3289+
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
3290+
ret = (int) new_node_id;
3291+
if (ret < 0) {
3292+
goto out;
3293+
}
3294+
}
3295+
ret = msp_store_arg_edges(self, merged_head, new_node_id);
3296+
if (ret != 0) {
3297+
goto out;
3298+
}
3299+
}
32493300
} else if (num_common_ancestors >= 2) {
32503301
msp_remove_individuals_from_population(self, Q);
32513302
if (num_common_ancestors == 2) {

‎lib/tests/test_ancestry.c‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -824,6 +824,49 @@ test_dtwf_multi_locus_simulation(void)
824824
tsk_table_collection_free(&tables);
825825
}
826826

827+
static void
828+
test_dtwf_multi_locus_simulation_unary(void)
829+
{
830+
int ret;
831+
const char *model_name;
832+
uint32_t n = 100;
833+
msp_t msp;
834+
tsk_size_t num_edges, num_edges_simple;
835+
gsl_rng *rng = safe_rng_alloc();
836+
tsk_table_collection_t tables;
837+
838+
ret = build_sim(&msp, &tables, rng, 100, 1, NULL, n);
839+
CU_ASSERT_EQUAL(ret, 0);
840+
ret = msp_initialise(&msp);
841+
CU_ASSERT_EQUAL(ret, 0);
842+
ret = msp_set_recombination_rate(&msp, 1);
843+
CU_ASSERT_EQUAL_FATAL(ret, 0);
844+
ret = msp_set_simulation_model_dtwf(&msp);
845+
CU_ASSERT_EQUAL(ret, 0);
846+
ret = msp_set_population_configuration(&msp, 0, n, 0, true);
847+
CU_ASSERT_EQUAL(ret, 0);
848+
ret = msp_set_store_unary(&msp, true);
849+
CU_ASSERT_EQUAL(ret, 0);
850+
model_name = msp_get_model_name(&msp);
851+
CU_ASSERT_STRING_EQUAL(model_name, "dtwf");
852+
853+
ret = msp_run(&msp, DBL_MAX, UINT32_MAX);
854+
CU_ASSERT_EQUAL(ret, 0);
855+
msp_verify(&msp, 0);
856+
857+
// verify whether there is at least one unary node
858+
num_edges = tables.edges.num_rows;
859+
ret = tsk_table_collection_simplify(&tables, NULL, 0, 0, NULL);
860+
CU_ASSERT_EQUAL_FATAL(ret, 0);
861+
num_edges_simple = tables.edges.num_rows;
862+
CU_ASSERT_TRUE(num_edges_simple < num_edges);
863+
864+
ret = msp_free(&msp);
865+
CU_ASSERT_EQUAL(ret, 0);
866+
gsl_rng_free(rng);
867+
tsk_table_collection_free(&tables);
868+
}
869+
827870
static void
828871
test_dtwf_deterministic(void)
829872
{
@@ -3573,6 +3616,8 @@ main(int argc, char **argv)
35733616

35743617
{ "test_dtwf_single_locus_simulation", test_dtwf_single_locus_simulation },
35753618
{ "test_dtwf_multi_locus_simulation", test_dtwf_multi_locus_simulation },
3619+
{ "test_dtwf_multi_locus_simulation_unary",
3620+
test_dtwf_multi_locus_simulation_unary },
35763621
{ "test_dtwf_deterministic", test_dtwf_deterministic },
35773622
{ "test_dtwf_simultaneous_historical_samples",
35783623
test_dtwf_simultaneous_historical_samples },

0 commit comments

Comments
 (0)