Finish moving population to Lineage in Python

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,15 +1831,15 @@ 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)
             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[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