diff --git a/algorithms.py b/algorithms.py index e17a8ed19..fdf4b1fd4 100644 --- a/algorithms.py +++ b/algorithms.py @@ -885,6 +885,7 @@ def __init__( gene_conversion_length=1, discrete_genome=True, hull_offset=None, + stop_at_local_mrca=True, ): # Must be a square matrix. N = len(migration_matrix) @@ -906,6 +907,7 @@ def __init__( self.migration_matrix = migration_matrix self.num_labels = num_labels self.num_populations = N + self.stop_at_local_mrca = stop_at_local_mrca self.max_segments = max_segments self.coalescing_segments_only = coalescing_segments_only self.additional_nodes = msprime.NodeType(additional_nodes) @@ -1253,7 +1255,7 @@ def simulate(self, end_time): if self.model == "hudson": self.hudson_simulate(end_time) elif self.model == "dtwf": - self.dtwf_simulate() + self.dtwf_simulate(end_time) elif self.model == "fixed_pedigree": self.pedigree_simulate() elif self.model == "single_sweep": @@ -1524,11 +1526,11 @@ def pedigree_simulate(self): self.pedigree = Pedigree(self.tables) self.dtwf_climb_pedigree() - def dtwf_simulate(self): + def dtwf_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. """ - while self.ancestors_remain(): + while self.ancestors_remain() and self.t < end_time: self.t += 1 self.verify() self.dtwf_generation() @@ -2394,6 +2396,7 @@ def merge_two_ancestors(self, population_index, label, x, y, u=-1): new_lineage = self.alloc_lineage(None, population_index, label=label) coalescence = False defrag_required = False + min_overlap = 0 while x is not None or y is not None: alpha = None @@ -2435,12 +2438,12 @@ def merge_two_ancestors(self, population_index, label, x, y, u=-1): j = self.S.floor_key(r_max) self.S[r_max] = self.S[j] # Update the number of extant segments. - if self.S[left] == 2: + if self.S[left] == min_overlap: self.S[left] = 0 right = self.S.succ_key(left) else: right = left - while right < r_max and self.S[right] != 2: + while right < r_max and self.S[right] != min_overlap: self.S[right] -= 1 right = self.S.succ_key(right) alpha = self.alloc_segment( @@ -2860,6 +2863,7 @@ def run_simulate(args): gene_conversion_length=mean_tract_length, discrete_genome=args.discrete, hull_offset=args.offset, + stop_at_local_mrca=args.stop_at_local_mrca, ) ts = s.simulate(args.end_time) ts.dump(args.output_file) @@ -2946,6 +2950,7 @@ def add_simulator_arguments(parser): help="The delta_t value for selective sweeps", ) parser.add_argument("--model", default="hudson") + parser.add_argument("--stop-at-local-mrca", default=True, type=bool) parser.add_argument("--offset", type=float, default=0.0) parser.add_argument( "--from-ts",