diff --git a/algorithms.py b/algorithms.py index f39d7012e..ce0006aef 100644 --- a/algorithms.py +++ b/algorithms.py @@ -1016,57 +1016,12 @@ def initialise(self, ts): lineage = root_lineages[node] if lineage is not None: seg = lineage.head - left_end = seg.left while seg is not None: self.set_segment_mass(seg) lineage.tail = seg seg = seg.next self.add_lineage(lineage) - if self.model == "smc_k": - for node in range(ts.num_nodes): - lineage = root_lineages[node] - if lineage is not None: - seg = lineage.head - left_end = seg.left - pop = lineage.population - label = lineage.label - right_end = root_segments_tail[node].right - new_hull = self.alloc_hull(left_end, right_end, lineage) - # insert Hull - floor = self.P[pop].hulls_left[label].floor_key(new_hull) - insertion_order = 0 - if floor is not None: - if floor.left == new_hull.left: - insertion_order = floor.insertion_order + 1 - new_hull.insertion_order = insertion_order - self.P[pop].hulls_left[label][new_hull] = -1 - - # initialise the correct coalesceable pairs count - for pop in self.P: - for label, ost_left in enumerate(pop.hulls_left): - avl = ost_left.avl - ost_right = pop.hulls_right[label] - count = 0 - for hull in avl.keys(): - floor = ost_right.floor_key(HullEnd(hull.left)) - num_ending_before_hull = 0 - if floor is not None: - num_ending_before_hull = ost_right.rank[floor] + 1 - num_pairs = count - num_ending_before_hull - avl[hull] = num_pairs - pop.coal_mass_index[label].set_value(hull.index, num_pairs) - # insert HullEnd - hull_end = HullEnd(hull.right) - floor = ost_right.floor_key(hull_end) - insertion_order = 0 - if floor is not None: - if floor.x == hull.right: - insertion_order = floor.insertion_order + 1 - hull_end.insertion_order = insertion_order - ost_right[hull_end] = -1 - count += 1 - def ancestors_remain(self): """ Returns True if the simulation is not finished, i.e., there is some ancestral @@ -1198,6 +1153,15 @@ def store_edge(self, left, right, parent, child): tskit.Edge(left=left, right=right, parent=parent, child=child) ) + def update_lineage_right(self, lineage): + if self.model == "smc_k": + # modify original hull + pop = lineage.population + hull = lineage.hull + old_right = hull.right + hull.right = min(lineage.tail.right + self.hull_offset, self.L) + self.P[pop].reset_hull_right(lineage.label, hull, old_right, hull.right) + def add_lineage(self, lineage): pop = lineage.population self.P[pop].add(lineage, lineage.label) @@ -1208,6 +1172,15 @@ def add_lineage(self, lineage): assert x.lineage == lineage x = x.next + if self.model == "smc_k": + head = lineage.head + assert head.prev is None + hull = self.alloc_hull(head.left, head.right, lineage) + right = lineage.tail.right + hull.right = min(right + self.hull_offset, self.L) + pop = self.P[lineage.population] + pop.add_hull(lineage.label, hull) + def finalise(self): """ Finalises the simulation returns an msprime tree sequence object. @@ -1836,22 +1809,11 @@ def hudson_recombination_event(self, label): left_lineage.tail = x lhs_tail = x + self.update_lineage_right(left_lineage) right_lineage = self.alloc_lineage(alpha, left_lineage.population, label=label) self.set_segment_mass(alpha) self.add_lineage(right_lineage) - if self.model == "smc_k": - # modify original hull - pop = left_lineage.population - lhs_hull = lhs_tail.get_hull() - rhs_right = lhs_hull.right - lhs_hull.right = min(lhs_tail.right + self.hull_offset, self.L) - self.P[pop].reset_hull_right(label, lhs_hull, rhs_right, lhs_hull.right) - - # create hull for alpha - alpha_hull = self.alloc_hull(alpha.left, rhs_right, right_lineage) - self.P[pop].add_hull(label, alpha_hull) - if self.additional_nodes.value & msprime.NODE_IS_RE_EVENT > 0: self.store_node(left_lineage.population, flags=msprime.NODE_IS_RE_EVENT) self.store_arg_edges(lhs_tail) @@ -1890,11 +1852,8 @@ def wiuf_gene_conversion_within_event(self, label): # lbp rbp return None self.num_gc_events += 1 - hull = y.get_hull() - assert (self.model == "smc_k") == (hull is not None) lineage = y.lineage pop = lineage.population - reset_right = -1 # Process left break insert_alpha = True @@ -1913,7 +1872,6 @@ def wiuf_gene_conversion_within_event(self, label): insert_alpha = False else: x.next = None - reset_right = x.right y.prev = None alpha = y tail = x @@ -1935,15 +1893,11 @@ def wiuf_gene_conversion_within_event(self, label): y.right = left_breakpoint self.set_segment_mass(y) tail = y - reset_right = left_breakpoint self.set_segment_mass(alpha) # Find the segment z that the right breakpoint falls in z = alpha - hull_left = z.left - hull_right = -1 while z is not None and right_breakpoint >= z.right: - hull_right = z.right z = z.next head = None @@ -1967,7 +1921,6 @@ def wiuf_gene_conversion_within_event(self, label): z.right = right_breakpoint z.next = None self.set_segment_mass(z) - hull_right = right_breakpoint else: # tail z # ====== @@ -1985,12 +1938,6 @@ def wiuf_gene_conversion_within_event(self, label): tail.next = head head.prev = tail self.set_segment_mass(head) - else: - # rbp lies beyond segment chain, regular recombination logic applies - if insert_alpha and self.model == "smc_k": - assert reset_right > 0 - reset_right = min(reset_right + self.hull_offset, self.L) - self.P[pop].reset_hull_right(label, hull, hull.right, reset_right) # y z # | ========== ... ===== | @@ -2005,12 +1952,8 @@ def wiuf_gene_conversion_within_event(self, label): if new_individual_head is not None: # FIXME when doing the smc_k update lineage.reset_segments() + self.update_lineage_right(lineage) new_lineage = self.alloc_lineage(new_individual_head, pop) - if self.model == "smc_k": - assert hull_left < hull_right - hull_right = min(self.L, hull_right + self.hull_offset) - hull = self.alloc_hull(hull_left, hull_right, new_lineage) - self.P[new_lineage.population].add_hull(new_lineage.label, hull) self.add_lineage(new_lineage) def wiuf_gene_conversion_left_event(self, label): @@ -2042,8 +1985,6 @@ def wiuf_gene_conversion_left_event(self, label): x = y.prev lineage = y.lineage pop = lineage.population - lhs_hull = y.get_hull() - assert (self.model == "smc_k") == (lhs_hull is not None) if y.left < bp: # x y # ===== =====|==== @@ -2061,7 +2002,6 @@ def wiuf_gene_conversion_left_event(self, label): y.next = None y.right = bp self.set_segment_mass(y) - right = y.right else: # x y # ===== | ========= @@ -2075,19 +2015,10 @@ def wiuf_gene_conversion_left_event(self, label): x.next = None y.prev = None alpha = y - right = x.right - - if self.model == "smc_k": - # lhs logic is identical to the lhs recombination event - lhs_old_right = lhs_hull.right - lhs_new_right = min(self.L, right + self.hull_offset) - self.P[pop].reset_hull_right(label, lhs_hull, lhs_old_right, lhs_new_right) - - # rhs - hull = self.alloc_hull(alpha.left, lhs_old_right, alpha) - self.P[pop].add_hull(label, hull) + # FIXME lineage.reset_segments() + self.update_lineage_right(lineage) self.set_segment_mass(alpha) assert alpha.prev is None @@ -2574,16 +2505,6 @@ def insert_merged_lineage( # assert tail == new_lineage.tail self.add_lineage(new_lineage) - if self.model == "smc_k": - merged_head = new_lineage.head - assert merged_head.prev is None - hull = self.alloc_hull(merged_head.left, merged_head.right, new_lineage) - while merged_head is not None: - right = merged_head.right - merged_head = merged_head.next - hull.right = min(right + self.hull_offset, self.L) - pop = self.P[new_lineage.population] - pop.add_hull(new_lineage.label, hull) return new_lineage def print_state(self, verify=False): diff --git a/lib/msprime.c b/lib/msprime.c index 35560d095..f53a7038d 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -462,10 +462,14 @@ 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. */ -static void +static int msp_reindex_segments(msp_t *self) { + int ret = 0; avl_node_t *node; avl_tree_t *population_ancestors; segment_t *seg; @@ -473,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]; @@ -481,9 +487,74 @@ 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) { + goto out; + } + } } } } +out: + return ret; +} + +/* Setup the various SMC K indexes either after a simulation model change + * or during msp_initialise */ +static int +msp_setup_smc_k_indexes(msp_t *self) +{ + int ret = 0; + label_id_t label; + size_t num_hulls, j; + population_t *pop; + + for (j = 0; j < self->num_populations; j++) { + pop = &self->populations[j]; + if (pop->hulls_left != NULL) { + msp_safe_free(pop->hulls_left); + pop->hulls_left = NULL; + } + if (pop->hulls_right != NULL) { + msp_safe_free(pop->hulls_right); + pop->hulls_right = NULL; + } + if (pop->coal_mass_index != NULL) { + for (label = 0; label < (label_id_t) self->num_labels; label++) { + fenwick_free(&pop->coal_mass_index[label]); + } + msp_safe_free(pop->coal_mass_index); + pop->coal_mass_index = NULL; + } + } + + 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; + } + } + } +out: + return ret; } /* Setup the mass indexes either after a simulation model change @@ -553,7 +624,14 @@ msp_setup_mass_indexes(msp_t *self) } } - msp_reindex_segments(self); + if (self->model.type == MSP_MODEL_SMC_K) { + ret = msp_setup_smc_k_indexes(self); + if (ret != 0) { + goto out; + } + } + + ret = msp_reindex_segments(self); out: return ret; } @@ -583,66 +661,6 @@ msp_alloc_memory_blocks_hulls(msp_t *self) return ret; } -/* Setup the mass indexes either after a simulation model change - * or during msp_initialise */ -static int -msp_setup_smc_k(msp_t *self) -{ - int ret = 0; - label_id_t label; - size_t num_hulls, j; - - /* Copied logic from msp_setup_mass_indexes */ - /* For simplicity, we always drop the mass 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; - } - if (self->populations[j].hulls_right != NULL) { - msp_safe_free(self->populations[j].hulls_right); - self->populations[j].hulls_right = NULL; - } - if (self->populations[j].coal_mass_index != NULL) { - for (label = 0; label < (label_id_t) self->num_labels; label++) { - fenwick_free(&self->populations[j].coal_mass_index[label]); - } - msp_safe_free(self->populations[j].coal_mass_index); - self->populations[j].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; - 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: - return ret; -} - static int msp_alloc_populations(msp_t *self) { @@ -935,6 +953,7 @@ msp_alloc_hull(msp_t *self, double left, double right, lineage_t *lineage) uint32_t j; tsk_bug_assert(lineage != NULL); + tsk_bug_assert(0 <= left && left <= right); label = lineage->label; if (object_heap_empty(&self->hull_heap[label])) { @@ -1331,9 +1350,9 @@ hullend_adjust_insertion_order(hullend_t *h, avl_node_t *node) } static inline avl_tree_t * -msp_get_segment_population(msp_t *self, segment_t *u) +msp_get_lineage_population(msp_t *self, lineage_t *lineage) { - return &self->populations[u->lineage->population].ancestors[u->lineage->label]; + return &self->populations[lineage->population].ancestors[lineage->label]; } static int MSP_WARN_UNUSED @@ -1481,6 +1500,29 @@ msp_remove_hull(msp_t *self, hull_t *hull) msp_free_hullend(self, query_ptr, label); } +static inline int MSP_WARN_UNUSED +msp_insert_individual_hull(msp_t *self, lineage_t *lin) +{ + int ret = 0; + hull_t *hull; + double hull_right; + const double hull_offset = self->model.params.smc_k_coalescent.hull_offset; + + hull_right = GSL_MIN(lin->tail->right + hull_offset, self->sequence_length); + hull = msp_alloc_hull(self, lin->head->left, hull_right, lin); + if (hull == NULL) { + ret = MSP_ERR_NO_MEMORY; + goto out; + } + lin->hull = hull; + ret = msp_insert_hull(self, hull); + if (ret != 0) { + goto out; + } +out: + return ret; +} + static inline int MSP_WARN_UNUSED msp_insert_individual(msp_t *self, lineage_t *lin) { @@ -1495,8 +1537,12 @@ msp_insert_individual(msp_t *self, lineage_t *lin) goto out; } avl_init_node(node, lin); - node = avl_insert_node(msp_get_segment_population(self, lin->head), node); + node = avl_insert_node(msp_get_lineage_population(self, lin), node); tsk_bug_assert(node != NULL); + + if (self->model.type == MSP_MODEL_SMC_K) { + ret = msp_insert_individual_hull(self, lin); + } out: return ret; } @@ -1507,7 +1553,7 @@ 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_segment_population(self, lin->head); + pop = msp_get_lineage_population(self, lin); node = avl_search(pop, lin); tsk_bug_assert(node != NULL); avl_unlink_node(pop, node); @@ -1709,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)); @@ -2015,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 @@ -2361,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); @@ -2671,12 +2727,9 @@ msp_move_individual(msp_t *self, avl_node_t *node, avl_tree_t *source, } ind->label = dest_label; } + ind->hull = new_hull; if (new_hull != NULL) { new_hull->lineage = ind; - ret = msp_insert_hull(self, new_hull); - if (ret != 0) { - goto out; - } } lineage_reset_segments(ind); ret = msp_insert_individual(self, ind); @@ -3299,7 +3352,7 @@ msp_recombination_event(msp_t *self, label_id_t label, lineage_t **lhs, lineage_ double breakpoint; lineage_t *left_lineage, *right_lineage; segment_t *x, *y, *alpha, *lhs_tail; - hull_t *lhs_hull, *rhs_hull; + hull_t *lhs_hull; double lhs_right, rhs_right; self->num_re_events++; @@ -3370,16 +3423,16 @@ msp_recombination_event(msp_t *self, label_id_t label, lineage_t **lhs, lineage_ msp_reset_hull_right( self, lhs_hull, rhs_right, lhs_right, left_lineage->population, label); - /* create new hull for alpha */ - rhs_hull = msp_alloc_hull(self, alpha->left, rhs_right, alpha->lineage); - if (rhs_hull == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - ret = msp_insert_hull(self, rhs_hull); - if (ret != 0) { - goto out; - } + /* /1* create new hull for alpha *1/ */ + /* rhs_hull = msp_alloc_hull(self, alpha->left, rhs_right, alpha->lineage); */ + /* if (rhs_hull == NULL) { */ + /* ret = MSP_ERR_NO_MEMORY; */ + /* goto out; */ + /* } */ + /* ret = msp_insert_hull(self, rhs_hull); */ + /* if (ret != 0) { */ + /* goto out; */ + /* } */ } if (self->additional_nodes & MSP_NODE_IS_RE_EVENT) { ret = msp_store_arg_recombination(self, lhs_tail, alpha); @@ -3428,7 +3481,6 @@ msp_gene_conversion_event(msp_t *self, label_id_t label) bool insert_alpha; hull_t *hull = NULL; double reset_right = 0.0; - double tract_hull_left, tract_hull_right; population_id_t population; tsk_bug_assert(self->gc_mass_index != NULL); @@ -3526,10 +3578,7 @@ msp_gene_conversion_event(msp_t *self, label_id_t label) // Find the segment z that the right breakpoint falls in z = alpha; - tract_hull_left = z->left; - tract_hull_right = 0.0; while (z != NULL && right_breakpoint >= z->right) { - tract_hull_right = z->right; z = z->next; } @@ -3558,7 +3607,6 @@ msp_gene_conversion_event(msp_t *self, label_id_t label) } z->right = right_breakpoint; z->next = NULL; - tract_hull_right = right_breakpoint; msp_set_segment_mass(self, z); if (!msp_has_breakpoint(self, right_breakpoint)) { @@ -3622,21 +3670,22 @@ msp_gene_conversion_event(msp_t *self, label_id_t label) if (ret != 0) { goto out; } - if (self->model.type == MSP_MODEL_SMC_K) { - tsk_bug_assert(tract_hull_left < tract_hull_right); - tract_hull_right = GSL_MIN( - tract_hull_right + self->model.params.smc_k_coalescent.hull_offset, - self->sequence_length); - hull = msp_alloc_hull(self, tract_hull_left, tract_hull_right, new_lineage); - if (hull == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - ret = msp_insert_hull(self, hull); - if (ret != 0) { - goto out; - } - } + /* if (self->model.type == MSP_MODEL_SMC_K) { */ + /* tsk_bug_assert(tract_hull_left < tract_hull_right); */ + /* tract_hull_right = GSL_MIN( */ + /* tract_hull_right + self->model.params.smc_k_coalescent.hull_offset, */ + /* self->sequence_length); */ + /* hull = msp_alloc_hull(self, tract_hull_left, tract_hull_right, + * new_lineage); */ + /* if (hull == NULL) { */ + /* ret = MSP_ERR_NO_MEMORY; */ + /* goto out; */ + /* } */ + /* ret = msp_insert_hull(self, hull); */ + /* if (ret != 0) { */ + /* goto out; */ + /* } */ + /* } */ } else { self->num_noneffective_gc_events++; } @@ -3696,10 +3745,7 @@ msp_insert_merged_ancestor(msp_t *self, lineage_t *new_lineage, tsk_id_t new_nod { int ret = 0; segment_t *z = new_lineage->tail; - segment_t *merged_head = new_lineage->head; segment_t *y; - hull_t *hull = NULL; - double r; if (coalescence) { ret = msp_conditional_compress_overlap_counts(self, l_min, r_max); @@ -3746,27 +3792,27 @@ msp_insert_merged_ancestor(msp_t *self, lineage_t *new_lineage, tsk_id_t new_nod if (ret_merged_head != NULL) { *ret_merged_head = new_lineage->head; } - if (self->model.type == MSP_MODEL_SMC_K) { - if (merged_head != NULL) { - y = merged_head; - r = 0; - while (y != NULL) { - r = y->right; - y = y->next; - } - r += self->model.params.smc_k_coalescent.hull_offset; - hull = msp_alloc_hull(self, merged_head->left, - GSL_MIN(r, self->sequence_length), merged_head->lineage); - if (hull == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - ret = msp_insert_hull(self, hull); - if (ret != 0) { - goto out; - } - } - } + /* if (self->model.type == MSP_MODEL_SMC_K) { */ + /* if (merged_head != NULL) { */ + /* y = merged_head; */ + /* r = 0; */ + /* while (y != NULL) { */ + /* r = y->right; */ + /* y = y->next; */ + /* } */ + /* r += self->model.params.smc_k_coalescent.hull_offset; */ + /* hull = msp_alloc_hull(self, merged_head->left, */ + /* GSL_MIN(r, self->sequence_length), merged_head->lineage); */ + /* if (hull == NULL) { */ + /* ret = MSP_ERR_NO_MEMORY; */ + /* goto out; */ + /* } */ + /* ret = msp_insert_hull(self, hull); */ + /* if (ret != 0) { */ + /* goto out; */ + /* } */ + /* } */ + /* } */ out: return ret; } @@ -4259,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; @@ -4297,14 +4344,19 @@ static int msp_insert_root_segments(msp_t *self, const segment_t *head, segment_t **new_head) { int ret = 0; - lineage_t *lineage = NULL; segment_t *copy, *prev; const segment_t *seg; const tsk_id_t *restrict node_population = self->tables->nodes.population; double breakpoints[2]; int j; const label_id_t label = 0; - hull_t *hull = NULL; + lineage_t *lineage + = msp_alloc_lineage(self, NULL, NULL, node_population[head->value], label); + + if (lineage == NULL) { + ret = MSP_ERR_NO_MEMORY; + goto out; + } prev = NULL; for (seg = head; seg != NULL; seg = seg->next) { @@ -4326,48 +4378,21 @@ msp_insert_root_segments(msp_t *self, const segment_t *head, segment_t **new_hea ret = MSP_ERR_NO_MEMORY; goto out; } - if (seg == head && new_head != NULL) { - *new_head = copy; - } copy->prev = prev; if (prev == NULL) { - lineage = msp_alloc_lineage( - self, copy, NULL, node_population[head->value], label); - if (lineage == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - ret = msp_insert_individual(self, lineage); - if (ret != 0) { - goto out; - } - if (self->model.type == MSP_MODEL_SMC_K) { - if (self->state != MSP_STATE_NEW) { - /* correct hull->right is set at the end */ - hull = msp_alloc_hull(self, head->left, copy->right, lineage); - if (hull == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - } - } + lineage->head = copy; } else { - lineage->tail = copy; prev->next = copy; } + lineage->tail = copy; + copy->lineage = lineage; msp_set_segment_mass(self, copy); prev = copy; } - /* insert hull into algorithm state */ - if (hull != NULL) { - tsk_bug_assert(self->model.type == MSP_MODEL_SMC_K); - hull->right - = GSL_MIN(prev->right + self->model.params.smc_k_coalescent.hull_offset, - self->sequence_length); - ret = msp_insert_hull(self, hull); - if (ret != 0) { - goto out; - } + ret = msp_insert_individual(self, lineage); + + if (new_head != NULL) { + *new_head = lineage->head; } out: return ret; @@ -4646,6 +4671,7 @@ msp_apply_fixed_events(msp_t *self, double time) return ret; } +#if 0 static int MSP_WARN_UNUSED msp_hulls_init_counts(msp_t *self, population_id_t pop, label_id_t label) { @@ -4684,6 +4710,7 @@ msp_hulls_init_counts(msp_t *self, population_id_t pop, label_id_t label) } value = visited - num_ending_before_hull; hull->count = value; + fenwick_set_value(coal_mass, hull->id, (double) value); visited++; @@ -4706,9 +4733,11 @@ msp_hulls_init_counts(msp_t *self, population_id_t pop, label_id_t label) out: return ret; } +#endif +#if 0 static int -msp_initialise_smc_k(msp_t *self) +msp_initialise_smc_k(msp_t *MSP_UNUSED(self)) { int ret = 0; population_id_t population_id; @@ -4762,6 +4791,7 @@ msp_initialise_smc_k(msp_t *self) out: return ret; } +#endif int msp_reset(msp_t *self) @@ -4946,17 +4976,6 @@ msp_initialise(msp_t *self) if (ret != 0) { goto out; } - /* SMC_K logic */ - 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; } @@ -5169,7 +5188,6 @@ msp_gene_conversion_left_event(msp_t *self, label_id_t label) population_id_t population; lineage_t *lineage, *new_lineage; segment_t *y, *x, *alpha; - hull_t *rhs_hull; hull_t *lhs_hull = NULL; lhs_hull = NULL; @@ -5285,18 +5303,6 @@ msp_gene_conversion_left_event(msp_t *self, label_id_t label) self->sequence_length); msp_reset_hull_right( self, lhs_hull, lhs_old_right, lhs_new_right, y->lineage->population, label); - - // rhs - tsk_bug_assert(alpha->left < lhs_old_right); - rhs_hull = msp_alloc_hull(self, alpha->left, lhs_old_right, alpha->lineage); - if (rhs_hull == NULL) { - ret = MSP_ERR_NO_MEMORY; - goto out; - } - ret = msp_insert_hull(self, rhs_hull); - if (ret != 0) { - goto out; - } } out: @@ -8442,17 +8448,6 @@ msp_set_simulation_model(msp_t *self, int model) goto out; } } - /* SMC_K logic */ - 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; @@ -8485,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 e165d0bc9..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) { @@ -3892,16 +3896,15 @@ test_setup_smc_k(void) CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_recombination_rate(&msp, 0.01); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_initialise(&msp); - CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_simulation_model_smc_k(&msp, 0.0); CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); model = msp_get_model(&msp)->type; CU_ASSERT_EQUAL(model, MSP_MODEL_SMC_K); model_name = msp_get_model_name(&msp); CU_ASSERT_STRING_EQUAL(model_name, "smc_k"); - while ((ret = msp_run(&msp, DBL_MAX, 1)) == MSP_EXIT_MAX_EVENTS) { msp_verify(&msp, 0); CU_ASSERT_FALSE(msp_is_completed(&msp)); @@ -3938,8 +3941,6 @@ test_setup_smc_k_plus(void) CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_recombination_rate(&msp, 0.01); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_initialise(&msp); - CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_simulation_model_smc_k(&msp, hull_offset); CU_ASSERT_EQUAL_FATAL(ret, 0); CU_ASSERT_EQUAL(msp.model.params.smc_k_coalescent.hull_offset, hull_offset); @@ -3949,6 +3950,9 @@ test_setup_smc_k_plus(void) model_name = msp_get_model_name(&msp); CU_ASSERT_STRING_EQUAL(model_name, "smc_k"); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); + while ((ret = msp_run(&msp, DBL_MAX, 1)) == MSP_EXIT_MAX_EVENTS) { msp_verify(&msp, 0); CU_ASSERT_FALSE(msp_is_completed(&msp)); @@ -3982,18 +3986,33 @@ test_reset_smc_k(void) CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_recombination_rate(&msp, 0.01); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_initialise(&msp); - CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_simulation_model_smc_k(&msp, 0.0); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_run(&msp, t, ULONG_MAX); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + 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); @@ -4225,10 +4244,12 @@ run_smc_k_gc_simulation( CU_ASSERT_EQUAL_FATAL(msp_set_recombination_rate(&msp, recombination_rate), 0); CU_ASSERT_EQUAL_FATAL(msp_set_gene_conversion_rate(&msp, gc_rate), 0); CU_ASSERT_EQUAL_FATAL(msp_set_gene_conversion_tract_length(&msp, tract_length), 0); - ret = msp_initialise(&msp); - CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_simulation_model_smc_k(&msp, offset); CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + msp_verify(&msp, 0); ret = msp_run(&msp, DBL_MAX, ULONG_MAX); CU_ASSERT_EQUAL(ret, 0);