Skip to content

add keep_unary for pedigrees and dtwf #2145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
GertjanBisschop wants to merge 1 commit into from
Closed

add keep_unary for pedigrees and dtwf #2145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 31 additions & 6 deletions algorithms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1140,6 +1140,10 @@ def dtwf_generation(self):
for h in H:
segments_to_merge = len(h)
if segments_to_merge == 1:
if self.store_unary:
# FIXME: should be flagged as a PASS_TRHOUGH_EVENT
new_node_id = self.store_node(pop_idx)
self.store_arg_edges(h[0][1], new_node_id)
h = []
elif segments_to_merge >= 2:
for _, individual in h:
Expand Down Expand Up @@ -1247,8 +1251,8 @@ def dtwf_climb_pedigree(self):
for ploid in range(ind.ploidy):
self.process_pedigree_common_ancestors(ind, ploid)

def store_arg_edges(self, segment, u=None):
if u is None:
def store_arg_edges(self, segment, u=-1):
if u == -1:
u = len(self.tables.nodes) - 1
# Store edges pointing to current node to the left
x = segment
Expand Down Expand Up @@ -1724,6 +1728,7 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1):
pop = self.P[pop_id]
defrag_required = False
coalescence = False
store_edge = False
alpha = None
z = None
merged_head = None
Expand Down Expand Up @@ -1802,9 +1807,20 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1):
self.set_segment_mass(alpha)
z = alpha
if self.full_arg:
store_edge = True
if not coalescence:
self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)
self.store_arg_edges(z)
new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)

elif self.store_unary:
if new_node_id != -1:
store_edge = True
elif self.model == "dtwf":
# PASS_THROUGH_EVENT has been dealt with in merge_n_ancestors
new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT)
store_edge = True

if store_edge:
self.store_arg_edges(z, new_node_id)
if defrag_required:
self.defrag_segment_chain(z)
if coalescence:
Expand Down Expand Up @@ -1853,6 +1869,8 @@ def merge_two_ancestors(self, population_index, label, x, y):
z = None
coalescence = False
defrag_required = False
store_edge = False
u = -1
while x is not None or y is not None:
alpha = None
if x is None or y is None:
Expand Down Expand Up @@ -1939,10 +1957,17 @@ def merge_two_ancestors(self, population_index, label, x, y):
z = alpha

if self.full_arg:
store_edge = True
if not coalescence:
u = self.store_node(population_index, flags=msprime.NODE_IS_CA_EVENT)
self.store_arg_edges(z, u)
elif self.store_unary and coalescence:
elif self.store_unary:
if u != -1:
store_edge = True
elif self.model == "dtwf":
u = self.store_node(population_index)
store_edge = True

if store_edge:
self.store_arg_edges(z, u)
if defrag_required:
self.defrag_segment_chain(z)
Expand Down
71 changes: 61 additions & 10 deletions lib/msprime.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2784,6 +2784,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l
int ret = 0;
bool coalescence = false;
bool defrag_required = false;
bool store_edge = false;
tsk_id_t v;
double l, r, l_min, r_max;
avl_node_t *node;
Expand Down Expand Up @@ -2938,26 +2939,39 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l
}
}
if (self->store_full_arg) {
store_edge = true;
if (!coalescence) {
ret = msp_store_node(
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
if (ret < 0) {
goto out;
}
}
// is specified new_node_id impossible when recording full_arg?
ret = msp_store_arg_edges(self, z, TSK_NULL);
if (ret != 0) {
goto out;
}

} else {
if (self->store_unary && coalescence) {
ret = msp_store_arg_edges(self, z, new_node_id);
if (ret < 0) {
goto out;
if (self->store_unary) {
if (new_node_id != TSK_NULL) {
store_edge = true;
} else {
if (self->model.type == MSP_MODEL_DTWF) {
new_node_id = msp_store_node(
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
ret = (int) new_node_id;
if (ret < 0) {
goto out;
}
store_edge = true;
}
}
}
}

if (store_edge) {
ret = msp_store_arg_edges(self, z, new_node_id);
if (ret != 0) {
goto out;
}
}
if (defrag_required) {
ret = msp_defrag_segment_chain(self, z);
if (ret != 0) {
Expand Down Expand Up @@ -3020,6 +3034,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
int ret = MSP_ERR_GENERIC;
bool coalescence = false;
bool defrag_required = false;
bool store_edge = false;
uint32_t j, h;
double l, r, r_max, next_l, l_min;
avl_node_t *node;
Expand Down Expand Up @@ -3180,14 +3195,35 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
}
}
if (self->store_full_arg) {
store_edge = true;
if (!coalescence) {
ret = msp_store_node(
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, individual);
if (ret < 0) {
goto out;
}
}
ret = msp_store_arg_edges(self, z, TSK_NULL);
} else {
if (self->store_unary) {
if (new_node_id != TSK_NULL) {
store_edge = true;
} else {
if (self->model.type == MSP_MODEL_DTWF) {
// PASS_THROUGH_EVENT have been dealth with in merge_n_ancestors
new_node_id = msp_store_node(self, MSP_NODE_IS_CA_EVENT, self->time,
population_id, individual);
ret = (int) new_node_id;
if (ret < 0) {
goto out;
}
store_edge = true;
}
}
}
}

if (store_edge) {
ret = msp_store_arg_edges(self, z, new_node_id);
if (ret != 0) {
goto out;
}
Expand Down Expand Up @@ -3246,6 +3282,21 @@ msp_merge_n_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,

if (num_common_ancestors == 1) {
merged_head = msp_priority_queue_pop(self, Q);
if (self->store_unary) {
if (self->model.type == MSP_MODEL_DTWF) {
new_node_id = msp_store_node(
/*FIXME: not a CA_EVENT but PASS_THROUGH_EVENT*/
self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL);
ret = (int) new_node_id;
if (ret < 0) {
goto out;
}
}
ret = msp_store_arg_edges(self, merged_head, new_node_id);
if (ret != 0) {
goto out;
}
}
} else if (num_common_ancestors >= 2) {
msp_remove_individuals_from_population(self, Q);
if (num_common_ancestors == 2) {
Expand Down
45 changes: 45 additions & 0 deletions lib/tests/test_ancestry.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -824,6 +824,49 @@ test_dtwf_multi_locus_simulation(void)
tsk_table_collection_free(&tables);
}

static void
test_dtwf_multi_locus_simulation_unary(void)
{
int ret;
const char *model_name;
uint32_t n = 100;
msp_t msp;
tsk_size_t num_edges, num_edges_simple;
gsl_rng *rng = safe_rng_alloc();
tsk_table_collection_t tables;

ret = build_sim(&msp, &tables, rng, 100, 1, NULL, n);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_initialise(&msp);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_recombination_rate(&msp, 1);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = msp_set_simulation_model_dtwf(&msp);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_population_configuration(&msp, 0, n, 0, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_store_unary(&msp, true);
CU_ASSERT_EQUAL(ret, 0);
model_name = msp_get_model_name(&msp);
CU_ASSERT_STRING_EQUAL(model_name, "dtwf");

ret = msp_run(&msp, DBL_MAX, UINT32_MAX);
CU_ASSERT_EQUAL(ret, 0);
msp_verify(&msp, 0);

// verify whether there is at least one unary node
num_edges = tables.edges.num_rows;
ret = tsk_table_collection_simplify(&tables, NULL, 0, 0, NULL);
CU_ASSERT_EQUAL_FATAL(ret, 0);
num_edges_simple = tables.edges.num_rows;
CU_ASSERT_TRUE(num_edges_simple < num_edges);

ret = msp_free(&msp);
CU_ASSERT_EQUAL(ret, 0);
gsl_rng_free(rng);
tsk_table_collection_free(&tables);
}

static void
test_dtwf_deterministic(void)
{
Expand Down Expand Up @@ -3573,6 +3616,8 @@ main(int argc, char **argv)

{ "test_dtwf_single_locus_simulation", test_dtwf_single_locus_simulation },
{ "test_dtwf_multi_locus_simulation", test_dtwf_multi_locus_simulation },
{ "test_dtwf_multi_locus_simulation_unary",
test_dtwf_multi_locus_simulation_unary },
{ "test_dtwf_deterministic", test_dtwf_deterministic },
{ "test_dtwf_simultaneous_historical_samples",
test_dtwf_simultaneous_historical_samples },
Expand Down
Loading