Skip to content

Commit 836510b

Browse files
GertjanBisschopGertjanBisschop
committed
first pass
1 parent a932851 commit 836510b

7 files changed

Lines changed: 316 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: 56 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,38 @@ 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+
ret = msp_store_node(
2958+
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
2959+
if (ret != 0) {
2960+
goto out;
2961+
}
2962+
store_edge = true;
2963+
}
29582964
}
29592965
}
29602966
}
2967+
2968+
if (store_edge) {
2969+
ret = msp_store_arg_edges(self, z, new_node_id);
2970+
if (ret != 0) {
2971+
goto out;
2972+
}
2973+
}
29612974
if (defrag_required) {
29622975
ret = msp_defrag_segment_chain(self, z);
29632976
if (ret != 0) {
@@ -3020,6 +3033,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
30203033
int ret = MSP_ERR_GENERIC;
30213034
bool coalescence = false;
30223035
bool defrag_required = false;
3036+
bool store_edge = false;
30233037
uint32_t j, h;
30243038
double l, r, r_max, next_l, l_min;
30253039
avl_node_t *node;
@@ -3180,14 +3194,34 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
31803194
}
31813195
}
31823196
if (self->store_full_arg) {
3197+
store_edge = true;
31833198
if (!coalescence) {
31843199
ret = msp_store_node(
31853200
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, individual);
31863201
if (ret < 0) {
31873202
goto out;
31883203
}
31893204
}
3190-
ret = msp_store_arg_edges(self, z, TSK_NULL);
3205+
} else {
3206+
if (self->store_unary) {
3207+
if (new_node_id != TSK_NULL) {
3208+
store_edge = true;
3209+
} else {
3210+
if (self->model.type == MSP_MODEL_DTWF) {
3211+
// PASS_THROUGH_EVENT have been dealth with in merge_n_ancestors
3212+
ret = msp_store_node(self, MSP_NODE_IS_CA_EVENT, self->time,
3213+
population_id, individual);
3214+
if (ret < 0) {
3215+
goto out;
3216+
}
3217+
store_edge = true;
3218+
}
3219+
}
3220+
}
3221+
}
3222+
3223+
if (store_edge) {
3224+
ret = msp_store_arg_edges(self, z, new_node_id);
31913225
if (ret != 0) {
31923226
goto out;
31933227
}
@@ -3246,6 +3280,18 @@ msp_merge_n_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
32463280

32473281
if (num_common_ancestors == 1) {
32483282
merged_head = msp_priority_queue_pop(self, Q);
3283+
if (self->store_unary && self->model.type == MSP_MODEL_DTWF) {
3284+
ret = msp_store_node(
3285+
/*FIXME: not a CA_EVENT but PASS_THROUGH_EVENT*/
3286+
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
3287+
if (ret < 0) {
3288+
goto out;
3289+
}
3290+
ret = msp_store_arg_edges(self, merged_head, TSK_NULL);
3291+
if (ret < 0) {
3292+
goto out;
3293+
}
3294+
}
32493295
} else if (num_common_ancestors >= 2) {
32503296
msp_remove_individuals_from_population(self, Q);
32513297
if (num_common_ancestors == 2) {

‎lib/tests/test_ancestry.c‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -824,6 +824,47 @@ 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 = 10;
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, 0.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_store_unary(&msp, true);
847+
CU_ASSERT_EQUAL(ret, 0);
848+
model_name = msp_get_model_name(&msp);
849+
CU_ASSERT_STRING_EQUAL(model_name, "dtwf");
850+
851+
ret = msp_run(&msp, DBL_MAX, UINT32_MAX);
852+
CU_ASSERT_EQUAL(ret, 0);
853+
msp_verify(&msp, 0);
854+
855+
// verify whether there is at least one unary node
856+
num_edges = tables.edges.num_rows;
857+
ret = tsk_table_collection_simplify(&tables, NULL, 0, 0, NULL);
858+
CU_ASSERT_EQUAL_FATAL(ret, 0);
859+
num_edges_simple = tables.edges.num_rows;
860+
CU_ASSERT_TRUE(num_edges_simple < num_edges);
861+
862+
ret = msp_free(&msp);
863+
CU_ASSERT_EQUAL(ret, 0);
864+
gsl_rng_free(rng);
865+
tsk_table_collection_free(&tables);
866+
}
867+
827868
static void
828869
test_dtwf_deterministic(void)
829870
{
@@ -3573,6 +3614,8 @@ main(int argc, char **argv)
35733614

35743615
{ "test_dtwf_single_locus_simulation", test_dtwf_single_locus_simulation },
35753616
{ "test_dtwf_multi_locus_simulation", test_dtwf_multi_locus_simulation },
3617+
{ "test_dtwf_multi_locus_simulation_unary",
3618+
test_dtwf_multi_locus_simulation_unary },
35763619
{ "test_dtwf_deterministic", test_dtwf_deterministic },
35773620
{ "test_dtwf_simultaneous_historical_samples",
35783621
test_dtwf_simultaneous_historical_samples },

0 commit comments

Comments
 (0)