diff --git a/algorithms.py b/algorithms.py index c69b95685..d9850104d 100644 --- a/algorithms.py +++ b/algorithms.py @@ -1003,6 +1003,7 @@ def initialise(self, ts): seg = self.alloc_segment( left, right, root, population, tail ) + seg.lineage = root_lineages[root] tail.next = seg root_segments_tail[root] = seg self.S[self.L] = -1 @@ -1013,18 +1014,18 @@ def initialise(self, ts): if lineage is not None: seg = lineage.head left_end = seg.left - self.add_lineage(lineage) while seg is not None: self.set_segment_mass(seg) seg = seg.next + self.add_lineage(lineage) if self.model == "smc_k": for node in range(ts.num_nodes): - seg = root_segments_head[node] - if seg is not None: + lineage = root_lineages[node] + if lineage is not None: + seg = lineage.head left_end = seg.left - pop = seg.population - lineage = seg.lineage + pop = lineage.population label = lineage.label right_end = root_segments_tail[node].right new_hull = self.alloc_hull(left_end, right_end, lineage) @@ -1830,7 +1831,7 @@ def hudson_recombination_event(self, label, return_heads=False): right_lineage = self.alloc_lineage(alpha, left_lineage.population, label=label) if self.model == "smc_k": # modify original hull - pop = alpha.population + 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) @@ -1838,7 +1839,7 @@ def hudson_recombination_event(self, label, return_heads=False): # create hull for alpha alpha_hull = self.alloc_hull(alpha.left, rhs_right, right_lineage) - self.P[alpha.population].add_hull(label, alpha_hull) + self.P[pop].add_hull(label, alpha_hull) self.set_segment_mass(alpha) self.add_lineage(right_lineage) @@ -2006,7 +2007,7 @@ def wiuf_gene_conversion_within_event(self, label): assert hull_left < hull_right hull_right = min(self.L, hull_right + self.hull_offset) hull = self.alloc_hull(hull_left, hull_right, lineage) - self.P[new_individual_head.population].add_hull(lineage.label, hull) + self.P[lineage.population].add_hull(lineage.label, hull) self.add_lineage(lineage) def wiuf_gene_conversion_left_event(self, label): @@ -2093,11 +2094,9 @@ def hudson_recombination_event_sweep_phase(self, label, sweep_site, pop_freq): Implements a recombination event in during a selective sweep. """ left_lin, right_lin = self.hudson_recombination_event(label, return_heads=True) - lhs = left_lin.head - rhs = right_lin.head r = random.random() - if sweep_site < rhs.left: + if sweep_site < right_lin.head.left: if r < 1.0 - pop_freq: # move rhs to other population self.P[right_lin.population].remove_individual( @@ -2249,11 +2248,9 @@ def bottleneck_event(self, pop_id, label, intensity): self.merge_ancestors(H, pop_id, label) def store_additional_nodes_edges(self, flag, new_node_id, z): - # FIXME if self.additional_nodes.value & flag > 0: if new_node_id == -1: - self.store_node(population_index) - new_node_id = len(self.tables.nodes) - 1 + new_node_id = self.store_node(z.lineage.population) self.update_node_flag(new_node_id, flag) self.store_arg_edges(z, new_node_id) return new_node_id