Improve phylogeny speed

meteor/meteor.py

@@ -649,6 +649,12 @@ def get_arguments() -> Namespace:  # pragma: no cover
         type=isborned01,
         help="Minimum number of informative sites in the alignment (default: >= %(default)d).",
     )
+    tree_parser.add_argument(
+        "-r",
+        dest="gtr",
+        action="store_true",
+        help="Compute GTR model (default: False, slower).",
+    )
     tree_parser.add_argument(
         "-f",
         dest="format",
@@ -857,14 +863,7 @@ def main() -> None:  # pragma: no cover
             strain_detector = Strain(meteor, 100, 2, 2, 0.2, 1, 1, 0.2, 10, False)
             strain_detector.execute()
             meteor.tree_dir = Path(tmpdirname) / "tree"
-            trees = TreeBuilder(
-                meteor,
-                0.1,
-                4,
-                800,
-                600,
-                None,
-            )
+            trees = TreeBuilder(meteor, 0.1, 4, False, 800, 600, None)
             trees.execute()
     # Close logging
     logger.handlers[0].close()

meteor/phylogeny.py

@@ -48,6 +48,7 @@ class Phylogeny(Session):
     msp_file_list: list[Path]
     max_gap: float
     min_info_sites: int
+    gtr: bool
     tree_files: list[Path] = field(default_factory=list)
 
     def compute_site_info(self, sequences: Iterable[str]) -> list[float]:
@@ -171,8 +172,11 @@ def process_msp_file(
             )
             # cleaned_alignment = load_aligned_seqs(ali_file, moltype="dna")
             # d = EstimateDistances(cleaned_alignment, submodel=GTR())
-            d = EstimateDistances(aligned_seqs, submodel=GTR())
-            d.run(show_progress=False)
+            if self.gtr:
+                d = EstimateDistances(aligned_seqs, submodel=GTR())
+                d.run(show_progress=False)
+            else:
+                d = aligned_seqs.distance_matrix(calc="tn93", show_progress=False)
 
             # Create UPGMA Tree
             mycluster = upgma(d.get_pairwise_distances())

meteor/tests/test_phylogeny.py

@@ -27,7 +27,7 @@ def phylogeny_builder(datadir: Path, tmp_path: Path) -> Phylogeny:
     meteor.tree_dir.mkdir()
     meteor.threads = 1
     meteor.DEFAULT_GAP_CHAR = "?"
-    return Phylogeny(meteor, [Path(datadir / "msp_0864.fasta")], 0.5, 4)
+    return Phylogeny(meteor, [Path(datadir / "msp_0864.fasta")], 0.5, 4, False)
 
 
 def test_compute_site_info(phylogeny_builder: Phylogeny):

meteor/tests/test_treebuilder.py

@@ -29,7 +29,7 @@ def treebuilder_builder(datadir: Path, tmp_path: Path) -> TreeBuilder:
     meteor.tree_dir.mkdir()
     meteor.threads = 1
     meteor.strain_dir = datadir / "strain"
-    return TreeBuilder(meteor, 0.5, 4, 500, 500, None)
+    return TreeBuilder(meteor, 0.5, 4, False, 500, 500, None)
 
 
 def test_concatenate(treebuilder_builder: TreeBuilder, datadir: Path):

meteor/treebuilder.py

@@ -42,6 +42,7 @@ class TreeBuilder(Session):
     meteor: type[Component]
     max_gap: float
     min_info_sites: int
+    gtr: bool
     width: int
     height: int
     format: str | None
@@ -103,7 +104,7 @@ def execute(self) -> None:
         msp_file_list = self.concatenate(msp_file_dict)
         # Compute phylogenies
         phylogeny_process = Phylogeny(
-            self.meteor, msp_file_list, self.max_gap, self.min_info_sites
+            self.meteor, msp_file_list, self.max_gap, self.min_info_sites, self.gtr
         )
         phylogeny_process.execute()
         # Analyze tree data