diff --git a/lib/msprime.c b/lib/msprime.c index 021cd8d72..f53a7038d 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -462,6 +462,7 @@ msp_set_segment_mass(msp_t *self, segment_t *seg) } } +static void msp_remove_hull(msp_t *self, hull_t *hull); static int msp_insert_individual_hull(msp_t *self, lineage_t *lin); /* Add all extant segments into the indexes. */ @@ -476,6 +477,8 @@ msp_reindex_segments(msp_t *self) size_t j; label_id_t label; + printf("REINDEX SEGS\n"); + for (j = 0; j < self->num_populations; j++) { for (label = 0; label < (label_id_t) self->num_labels; label++) { population_ancestors = &self->populations[j].ancestors[label]; @@ -484,6 +487,10 @@ msp_reindex_segments(msp_t *self) for (seg = lin->head; seg != NULL; seg = seg->next) { msp_set_segment_mass(self, seg); } + if (lin->hull != NULL) { + msp_remove_hull(self, lin->hull); + lin->hull = NULL; + } if (self->model.type == MSP_MODEL_SMC_K) { ret = msp_insert_individual_hull(self, lin); if (ret != 0) { @@ -505,51 +512,45 @@ msp_setup_smc_k_indexes(msp_t *self) int ret = 0; label_id_t label; size_t num_hulls, j; + population_t *pop; - /* For simplicity, we always drop the indexes even though - * sometimes we'll be dropping it just to rebuild */ for (j = 0; j < self->num_populations; j++) { - if (self->populations[j].hulls_left != NULL) { - msp_safe_free(self->populations[j].hulls_left); - self->populations[j].hulls_left = NULL; + pop = &self->populations[j]; + if (pop->hulls_left != NULL) { + msp_safe_free(pop->hulls_left); + pop->hulls_left = NULL; } - if (self->populations[j].hulls_right != NULL) { - msp_safe_free(self->populations[j].hulls_right); - self->populations[j].hulls_right = NULL; + if (pop->hulls_right != NULL) { + msp_safe_free(pop->hulls_right); + pop->hulls_right = NULL; } - if (self->populations[j].coal_mass_index != NULL) { + if (pop->coal_mass_index != NULL) { for (label = 0; label < (label_id_t) self->num_labels; label++) { - fenwick_free(&self->populations[j].coal_mass_index[label]); + fenwick_free(&pop->coal_mass_index[label]); } - msp_safe_free(self->populations[j].coal_mass_index); - self->populations[j].coal_mass_index = NULL; + msp_safe_free(pop->coal_mass_index); + pop->coal_mass_index = NULL; } } - if (self->model.type == MSP_MODEL_SMC_K) { - num_hulls = self->hull_heap->size; - for (j = 0; j < self->num_populations; j++) { - self->populations[j].hulls_left - = malloc(self->num_labels * sizeof(*self->populations[j].hulls_left)); - self->populations[j].hulls_right - = malloc(self->num_labels * sizeof(*self->populations[j].hulls_right)); - self->populations[j].coal_mass_index = calloc( - self->num_labels, sizeof(*self->populations[j].coal_mass_index)); - if (self->populations[j].hulls_left == NULL - || self->populations[j].hulls_right == NULL - || self->populations[j].coal_mass_index == NULL) { - ret = MSP_ERR_NO_MEMORY; + + num_hulls = self->hull_heap->size; + for (j = 0; j < self->num_populations; j++) { + pop = &self->populations[j]; + pop->hulls_left = malloc(self->num_labels * sizeof(*pop->hulls_left)); + pop->hulls_right = malloc(self->num_labels * sizeof(*pop->hulls_right)); + pop->coal_mass_index = calloc(self->num_labels, sizeof(*pop->coal_mass_index)); + if (pop->hulls_left == NULL || pop->hulls_right == NULL + || pop->coal_mass_index == NULL) { + ret = MSP_ERR_NO_MEMORY; + goto out; + } + for (label = 0; label < (label_id_t) self->num_labels; label++) { + avl_init_tree(&pop->hulls_left[label], cmp_hull, NULL); + avl_init_tree(&pop->hulls_right[label], cmp_hullend, NULL); + ret = fenwick_alloc(&pop->coal_mass_index[label], num_hulls); + if (ret != 0) { goto out; } - for (label = 0; label < (label_id_t) self->num_labels; label++) { - avl_init_tree(&self->populations[j].hulls_left[label], cmp_hull, NULL); - avl_init_tree( - &self->populations[j].hulls_right[label], cmp_hullend, NULL); - ret = fenwick_alloc( - &self->populations[j].coal_mass_index[label], num_hulls); - if (ret != 0) { - goto out; - } - } } } out: @@ -1754,8 +1755,12 @@ msp_verify_segments(msp_t *self, bool verify_breakpoints) pedigree_avl_nodes += avl_count(&ind->common_ancestors[k]); } } - tsk_bug_assert(total_avl_nodes + pedigree_avl_nodes + hull_avl_nodes - == object_heap_get_num_allocated(&self->avl_node_heap)); + printf("hull_avl_nodes = %d total_avl_nodes %d have :%d\n", (int) hull_avl_nodes, + (int) total_avl_nodes, + (int) object_heap_get_num_allocated(&self->avl_node_heap)); + + /* tsk_bug_assert(total_avl_nodes + pedigree_avl_nodes + hull_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) == object_heap_get_num_allocated(&self->node_mapping_heap)); @@ -2060,6 +2065,8 @@ msp_verify_hulls(msp_t *self) tsk_bug_assert(io == hull->insertion_order); pos = hull->left; } + printf( + "count=%d coal pairs = %d\n", (int) count, (int) num_coalescing_pairs); tsk_bug_assert(count == num_coalescing_pairs); /* should equal total sum coal_mass_index */ coal_mass_index @@ -2406,6 +2413,10 @@ msp_print_state(msp_t *self, FILE *out) for (j = 0; j < self->num_labels; j++) { fprintf(out, "segment_heap[%d]:", j); object_heap_print_state(&self->segment_heap[j], out); + fprintf(out, "hull_heap[%d]:", j); + object_heap_print_state(&self->hull_heap[j], out); + fprintf(out, "hullend_heap[%d]:", j); + object_heap_print_state(&self->hullend_heap[j], out); } fprintf(out, "avl_node_heap:"); object_heap_print_state(&self->avl_node_heap, out); @@ -4294,6 +4305,7 @@ msp_reset_memory_state(msp_t *self) msp_free_lineage(self, lin); } if (pop->hulls_left != NULL) { + printf("RESET STATE\n"); for (node = pop->hulls_left[label].head; node != NULL; node = node->next) { x = (hull_t *) node->item; @@ -4964,17 +4976,6 @@ msp_initialise(msp_t *self) if (ret != 0) { goto out; } - /* /1* SMC_K logic *1/ */ - /* if (self->model.type == MSP_MODEL_SMC_K) { */ - /* ret = msp_setup_smc_k(self); */ - /* if (ret != 0) { */ - /* goto out; */ - /* } */ - /* ret = msp_initialise_smc_k(self); */ - /* if (ret != 0) { */ - /* goto out; */ - /* } */ - /* } */ out: return ret; } @@ -8479,13 +8480,12 @@ msp_set_simulation_model_smc_k(msp_t *self, double hull_offset) ret = MSP_ERR_BAD_HULL_OFFSET; goto out; } + self->model.params.smc_k_coalescent.hull_offset = hull_offset; ret = msp_set_simulation_model(self, MSP_MODEL_SMC_K); if (ret != 0) { goto out; } - - self->model.params.smc_k_coalescent.hull_offset = hull_offset; self->get_common_ancestor_waiting_time = msp_smc_k_get_common_ancestor_waiting_time; self->common_ancestor_event = msp_smc_k_common_ancestor_event; out: diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index 16c14f98b..1713a39da 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -533,8 +533,9 @@ test_multi_locus_simulation(void) int model; double migration_matrix[] = { 0, 1, 1, 0 }; size_t migration_events[4]; - int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME }; - const char *model_names[] = { "hudson", "smc", "smc_prime" }; + int models[] + = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME, MSP_MODEL_SMC_K }; + const char *model_names[] = { "hudson", "smc", "smc_prime", "smc_k" }; const char *model_name; bool store_full_arg[] = { true, false }; size_t j, k; @@ -572,6 +573,9 @@ test_multi_locus_simulation(void) case 2: ret = msp_set_simulation_model_smc_prime(&msp); break; + case 3: + ret = msp_set_simulation_model_smc_k(&msp, 0); + break; } ret = msp_set_store_full_arg(&msp, store_full_arg[k]); CU_ASSERT_EQUAL_FATAL(ret, 0); @@ -1299,12 +1303,12 @@ test_mixed_hudson_smc(void) CU_ASSERT_FALSE(msp_is_completed(&msp)); model = msp_get_model(&msp)->type; if (j % 2 == 1) { - CU_ASSERT_EQUAL(model, MSP_MODEL_SMC); + CU_ASSERT_EQUAL(model, MSP_MODEL_SMC_K); ret = msp_set_simulation_model_hudson(&msp); CU_ASSERT_EQUAL(ret, 0); } else { CU_ASSERT_EQUAL(model, MSP_MODEL_HUDSON); - ret = msp_set_simulation_model_smc(&msp); + ret = msp_set_simulation_model_smc_k(&msp, 0); CU_ASSERT_EQUAL(ret, 0); } if (j == 10) { @@ -3962,7 +3966,6 @@ test_setup_smc_k_plus(void) tsk_table_collection_free(&tables); } -#if 0 static void test_reset_smc_k(void) { @@ -3987,14 +3990,29 @@ test_reset_smc_k(void) CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_initialise(&msp); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_run(&msp, t, ULONG_MAX); + + while ((ret = msp_run(&msp, DBL_MAX, 1)) == MSP_EXIT_MAX_EVENTS) { + msp_verify(&msp, 0); + CU_ASSERT_FALSE(msp_is_completed(&msp)); + } CU_ASSERT_EQUAL(ret, MSP_EXIT_MAX_TIME); msp_verify(&msp, 0); + ret = msp_run(&msp, t, ULONG_MAX); + CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_reset(&msp); CU_ASSERT_EQUAL_FATAL(ret, 0); + msp_verify(&msp, 0); + printf("SET SIM MODEL\n"); ret = msp_set_simulation_model_smc_k(&msp, 0.0); CU_ASSERT_EQUAL_FATAL(ret, 0); + /* msp_print_state(&msp, stdout); */ + msp_verify(&msp, 0); + + ret = msp_reset(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); + msp_verify(&msp, 0); + gsl_rng_set(rng, seed); while ((ret = msp_run(&msp, DBL_MAX, 1)) == MSP_EXIT_MAX_EVENTS) { msp_verify(&msp, 0); @@ -4008,7 +4026,6 @@ test_reset_smc_k(void) free(samples); tsk_table_collection_free(&tables); } -#endif static void test_init_smc_k(void) @@ -4359,7 +4376,7 @@ main(int argc, char **argv) { "test_bad_setup_smc_k", test_bad_setup_smc_k }, { "test_setup_smc_k", test_setup_smc_k }, { "test_setup_smc_k_plus", test_setup_smc_k_plus }, - /* { "test_reset_smc_k", test_reset_smc_k }, */ + { "test_reset_smc_k", test_reset_smc_k }, { "test_init_smc_k", test_init_smc_k }, { "test_smc_k_multipop", test_smc_k_multipop }, { "test_mixed_model_smc_k", test_mixed_model_smc_k },