Skip to content

Embed avl into lineage struct #2385

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

Merged
jeromekelleher merged 5 commits into from
Jul 24, 2025
Merged

Embed avl into lineage struct #2385

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a1c80f7
Update copyright years
jeromekelleher Jul 24, 2025
581f83f
Embed avl_node in lineage struct and remove alloc/frees
jeromekelleher Jul 24, 2025
814e191
Remove 3 more AVL searches from smc_k
jeromekelleher Jul 24, 2025
9f2dbe5
Update a DTWF test
jeromekelleher Jul 24, 2025
8e2ab4f
Remove some redundant AVL searches
jeromekelleher Jul 24, 2025
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
63 changes: 16 additions & 47 deletions lib/msprime.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
** Copyright (C) 2015-2024 University of Oxford
** Copyright (C) 2015-2025 University of Oxford
**
** This file is part of msprime.
**
Expand Down Expand Up @@ -1470,19 +1470,14 @@ msp_insert_individual(msp_t *self, lineage_t *lin)

tsk_bug_assert(lin != NULL);
tsk_bug_assert(lin->head != NULL);
node = msp_alloc_avl_node(self);
if (node == NULL) {
ret = MSP_ERR_NO_MEMORY;
goto out;
}
node = &lin->avl_node;
avl_init_node(node, lin);
node = avl_insert_node(msp_get_lineage_population(self, lin), node);
tsk_bug_assert(node != NULL);
tsk_bug_assert(node == &lin->avl_node);

if (self->model.type == MSP_MODEL_SMC_K) {
ret = msp_insert_individual_hull(self, lin);
}
out:
return ret;
}

Expand All @@ -1491,12 +1486,11 @@ msp_remove_individual(msp_t *self, lineage_t *lin)
{
avl_node_t *node;
avl_tree_t *pop;

tsk_bug_assert(lin != NULL);
pop = msp_get_lineage_population(self, lin);
node = avl_search(pop, lin);
tsk_bug_assert(node != NULL);
node = &lin->avl_node;
avl_unlink_node(pop, node);
msp_free_avl_node(self, node);
msp_free_lineage(self, lin);
}

Expand Down Expand Up @@ -1688,11 +1682,8 @@ msp_verify_segments(msp_t *self, bool verify_breakpoints)
label_segments == object_heap_get_num_allocated(&self->segment_heap[k]));
tsk_bug_assert(
label_hulls == object_heap_get_num_allocated(&self->hull_heap[k]));
/* tsk_bug_assert( */
/* label_hulls == object_heap_get_num_allocated(&self->hullend_heap[k])); */
}
total_avl_nodes = msp_get_num_ancestors(self) + avl_count(&self->breakpoints)
+ avl_count(&self->overlap_counts)
total_avl_nodes = avl_count(&self->breakpoints) + avl_count(&self->overlap_counts)
+ avl_count(&self->non_empty_populations);
tsk_bug_assert(msp_get_num_ancestors(self)
== object_heap_get_num_allocated(&self->lineage_heap));
Expand All @@ -1705,8 +1696,7 @@ msp_verify_segments(msp_t *self, bool verify_breakpoints)
}
tsk_bug_assert(total_avl_nodes + pedigree_avl_nodes
== object_heap_get_num_allocated(&self->avl_node_heap));
tsk_bug_assert(total_avl_nodes - msp_get_num_ancestors(self)
- avl_count(&self->non_empty_populations)
tsk_bug_assert(total_avl_nodes - avl_count(&self->non_empty_populations)
== object_heap_get_num_allocated(&self->node_mapping_heap));
if (self->recomb_mass_index != NULL) {
msp_verify_segment_index(
Expand Down Expand Up @@ -2682,7 +2672,6 @@ msp_move_individual(msp_t *self, avl_node_t *node, avl_tree_t *source,

ind = (lineage_t *) node->item;
avl_unlink_node(source, node);
msp_free_avl_node(self, node);
if (self->model.type == MSP_MODEL_SMC_K) {
msp_remove_hull(self, ind);
}
Expand Down Expand Up @@ -4146,8 +4135,7 @@ msp_merge_n_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
tsk_bug_assert(u->lineage != NULL);
if (u->lineage->population != population_id) {
current_pop = &self->populations[u->lineage->population];
avl_node = avl_search(&current_pop->ancestors[label], u->lineage);
tsk_bug_assert(avl_node != NULL);
avl_node = &u->lineage->avl_node;
ret = msp_move_individual(
self, avl_node, &current_pop->ancestors[label], population_id, label);
if (ret != 0) {
Expand Down Expand Up @@ -4229,7 +4217,6 @@ msp_reset_memory_state(msp_t *self)
u = v;
}
avl_unlink_node(&pop->ancestors[label], node);
msp_free_avl_node(self, node);
if (lin->hull != NULL) {
msp_remove_hull(self, lin);
}
Expand Down Expand Up @@ -5816,11 +5803,7 @@ static int
msp_change_label(msp_t *self, lineage_t *lin, label_id_t label)
{
avl_tree_t *pop = &self->populations[lin->population].ancestors[lin->label];
avl_node_t *node;

/* Find the this individual in the AVL tree. */
node = avl_search(pop, lin);
tsk_bug_assert(node != NULL);
avl_node_t *node = &lin->avl_node;
return msp_move_individual(self, node, pop, lin->population, label);
}

Expand Down Expand Up @@ -7162,7 +7145,6 @@ msp_simple_bottleneck(msp_t *self, demographic_event_t *event)
lin = (lineage_t *) node->item;
u = lin->head;
avl_unlink_node(pop, node);
msp_free_avl_node(self, node);
msp_free_lineage(self, lin);
q_node = msp_alloc_avl_node(self);
if (q_node == NULL) {
Expand Down Expand Up @@ -7314,7 +7296,6 @@ msp_instantaneous_bottleneck(msp_t *self, demographic_event_t *event)
* set for the root at u */
lin = (lineage_t *) avl_nodes[j]->item;
avl_unlink_node(pop, avl_nodes[j]);
msp_free_avl_node(self, avl_nodes[j]);
set_node = msp_alloc_avl_node(self);
if (set_node == NULL) {
ret = MSP_ERR_NO_MEMORY;
Expand Down Expand Up @@ -7601,9 +7582,7 @@ msp_std_common_ancestor_event(
tsk_bug_assert(node != NULL);
} else {
self->num_ca_events++;
msp_free_avl_node(self, x_node);
msp_free_lineage(self, x_lin);
msp_free_avl_node(self, y_node);
msp_free_lineage(self, y_lin);
ret = msp_merge_two_ancestors(self, population_id, label, x, y, TSK_NULL, NULL);
}
Expand Down Expand Up @@ -7647,14 +7626,13 @@ msp_smc_k_common_ancestor_event(
double random_mass, num_pairs, remaining_mass;
size_t hull_id;
fenwick_t *coal_mass_index;
avl_tree_t *avl;
avl_tree_t *ancestors;
avl_node_t *x_node, *y_node, *search;
hull_t *x_hull, *y_hull = NULL;
lineage_t *x_lin, *y_lin;
segment_t *x, *y;

/* find first hull */
/* FIX ME: clean up the various type castings */
coal_mass_index = &self->populations[population_id].coal_mass_index[label];
num_pairs = fenwick_get_total(coal_mass_index);
random_mass = gsl_ran_flat(self->rng, 0, num_pairs);
Expand All @@ -7666,9 +7644,7 @@ msp_smc_k_common_ancestor_event(
remaining_mass = fenwick_get_cumulative_sum(coal_mass_index, hull_id) - random_mass;

/* find second hull */
avl = &self->populations[population_id].hulls_left[label];
search = avl_search(avl, x_hull);
tsk_bug_assert(search != NULL);
search = &x_hull->left_avl_node;
for (search = search->prev; remaining_mass >= 0; search = search->prev) {
tsk_bug_assert(search != NULL);
y_hull = (hull_t *) search->item;
Expand All @@ -7682,19 +7658,15 @@ msp_smc_k_common_ancestor_event(
msp_remove_hull(self, y_lin);

/* retrieve ancestors linked to both hulls */
avl = &self->populations[population_id].ancestors[label];
ancestors = &self->populations[population_id].ancestors[label];
x = x_lin->head;
x_node = avl_search(avl, x_lin);
tsk_bug_assert(x_node != NULL);
avl_unlink_node(avl, x_node);
x_node = &x_lin->avl_node;
avl_unlink_node(ancestors, x_node);
y = y_lin->head;
y_node = avl_search(avl, y_lin);
tsk_bug_assert(y_node != NULL);
avl_unlink_node(avl, y_node);
y_node = &y_lin->avl_node;
avl_unlink_node(ancestors, y_node);

self->num_ca_events++;
msp_free_avl_node(self, x_node);
msp_free_avl_node(self, y_node);
msp_free_lineage(self, x_lin);
msp_free_lineage(self, y_lin);
ret = msp_merge_two_ancestors(self, population_id, label, x, y, TSK_NULL, NULL);
Expand Down Expand Up @@ -7781,9 +7753,7 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t
y = y_lin->head;
avl_unlink_node(ancestors, y_node);
self->num_ca_events++;
msp_free_avl_node(self, x_node);
msp_free_lineage(self, x_lin);
msp_free_avl_node(self, y_node);
msp_free_lineage(self, y_lin);
ret = msp_merge_two_ancestors(self, pop_id, label, x, y, TSK_NULL, NULL);
} else {
Expand Down Expand Up @@ -7947,7 +7917,6 @@ msp_multi_merger_common_ancestor_event(
lin = (lineage_t *) node->item;
u = lin->head;
avl_unlink_node(ancestors, node);
msp_free_avl_node(self, node);
msp_free_lineage(self, lin);

q_node = msp_alloc_avl_node(self);
Expand Down
3 changes: 2 additions & 1 deletion lib/msprime.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
** Copyright (C) 2015-2024 University of Oxford
** Copyright (C) 2015-2025 University of Oxford
**
** This file is part of msprime.
**
Expand Down Expand Up @@ -91,6 +91,7 @@ typedef struct lineage_t_t {
label_id_t label;
segment_t *head;
segment_t *tail;
avl_node_t avl_node;
struct hull_t_t *hull;
} lineage_t;

Expand Down
16 changes: 13 additions & 3 deletions lib/tests/test_ancestry.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
** Copyright (C) 2016-2024 University of Oxford
** Copyright (C) 2016-2025 University of Oxford
**
** This file is part of msprime.
**
Expand Down Expand Up @@ -795,7 +795,7 @@ test_dtwf_multi_locus_simulation(void)
long seed = 10;
double migration_matrix[] = { 0, 0.1, 0.1, 0 };
const char *model_name;
size_t num_ca_events, num_re_events;
size_t num_ca_events, num_re_events, num_mig_events;
double t;
tsk_table_collection_t tables;
msp_t msp;
Expand All @@ -809,6 +809,8 @@ test_dtwf_multi_locus_simulation(void)
CU_ASSERT_EQUAL_FATAL(msp_set_recombination_rate(&msp, 0.1), 0);
ret = msp_set_population_configuration(&msp, 0, n, 0, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_population_configuration(&msp, 1, n, 0, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_migration_matrix(&msp, 4, migration_matrix);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_store_migrations(&msp, true);
Expand All @@ -819,12 +821,14 @@ test_dtwf_multi_locus_simulation(void)
CU_ASSERT_STRING_EQUAL(model_name, "dtwf");

ret = msp_run(&msp, DBL_MAX, ULONG_MAX);
CU_ASSERT_EQUAL(ret, 0);
msp_verify(&msp, 0);
num_ca_events = msp_get_num_common_ancestor_events(&msp);
num_re_events = msp_get_num_recombination_events(&msp);
num_mig_events = tables.migrations.num_rows;
CU_ASSERT_TRUE(num_ca_events > 0);
CU_ASSERT_TRUE(num_re_events > 0);
CU_ASSERT_EQUAL(ret, 0);
CU_ASSERT_TRUE(num_mig_events > 0);
msp_free(&msp);
tsk_table_collection_free(&tables);

Expand All @@ -838,7 +842,12 @@ test_dtwf_multi_locus_simulation(void)
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_population_configuration(&msp, 0, n, 0, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_population_configuration(&msp, 1, n, 0, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_migration_matrix(&msp, 4, migration_matrix);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_set_store_migrations(&msp, true);
CU_ASSERT_EQUAL(ret, 0);
ret = msp_initialise(&msp);
CU_ASSERT_EQUAL(ret, 0);
t = 1;
Expand All @@ -854,6 +863,7 @@ test_dtwf_multi_locus_simulation(void)
CU_ASSERT_EQUAL(ret, 0);
CU_ASSERT_TRUE(num_ca_events == msp_get_num_common_ancestor_events(&msp));
CU_ASSERT_TRUE(num_re_events == msp_get_num_recombination_events(&msp));
CU_ASSERT_TRUE(num_mig_events == tables.migrations.num_rows);

ret = msp_free(&msp);
CU_ASSERT_EQUAL(ret, 0);
Expand Down
Loading