From 289302877b0ab189a01c3b795dbc3cd5ca819da2 Mon Sep 17 00:00:00 2001
From: Amine Ghozlane <amine.ghozlane@pasteur.fr>
Date: Tue, 26 Nov 2024 08:28:17 +0100
Subject: [PATCH] Add tn93 distance

---
 meteor/meteor.py                 |  8 ++++++++
 meteor/phylogeny.py              | 10 +++++++---
 meteor/tests/test_phylogeny.py   |  2 +-
 meteor/tests/test_treebuilder.py |  2 +-
 meteor/treebuilder.py            |  3 ++-
 5 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/meteor/meteor.py b/meteor/meteor.py
index 96192f1..559b54e 100644
--- a/meteor/meteor.py
+++ b/meteor/meteor.py
@@ -678,6 +678,12 @@ def get_arguments() -> Namespace:  # pragma: no cover
         required=True,
         help="Path to output directory.",
     )
+    tree_parser.add_argument(
+        "-r",
+        dest="gtr",
+        action='store_true',
+        help="Compute GTR genetic distance (slow).",
+    )
     tree_parser.add_argument(
         "--tmp",
         dest="tmp_path",
@@ -781,6 +787,7 @@ def main() -> None:  # pragma: no cover
             args.width,
             args.height,
             args.format,
+            args.gtr
         )
         trees.execute()
     # Run download catalogues
@@ -864,6 +871,7 @@ def main() -> None:  # pragma: no cover
                 800,
                 600,
                 None,
+                False
             )
             trees.execute()
     # Close logging
diff --git a/meteor/phylogeny.py b/meteor/phylogeny.py
index df33c2d..dc6f49a 100644
--- a/meteor/phylogeny.py
+++ b/meteor/phylogeny.py
@@ -47,6 +47,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]:
@@ -170,11 +171,14 @@ 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:
+                dists = EstimateDistances(aligned_seqs, submodel=GTR())
+                dists.run(show_progress=False)
+            else:
+                dists = aligned_seqs.distance_matrix(calc="tn93", show_progress=False)
 
             # Create UPGMA Tree
-            mycluster = upgma(d.get_pairwise_distances())
+            mycluster = upgma(dists.get_pairwise_distances())
             mycluster = mycluster.unrooted_deepcopy()
 
             with tree_file.open("w") as f:
diff --git a/meteor/tests/test_phylogeny.py b/meteor/tests/test_phylogeny.py
index 91acbc1..2783923 100644
--- a/meteor/tests/test_phylogeny.py
+++ b/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, True)
 
 
 def test_compute_site_info(phylogeny_builder: Phylogeny):
diff --git a/meteor/tests/test_treebuilder.py b/meteor/tests/test_treebuilder.py
index 27045be..fe95adc 100644
--- a/meteor/tests/test_treebuilder.py
+++ b/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, 500, 500, None, False)
 
 
 def test_concatenate(treebuilder_builder: TreeBuilder, datadir: Path):
diff --git a/meteor/treebuilder.py b/meteor/treebuilder.py
index bef96f0..9fd1169 100644
--- a/meteor/treebuilder.py
+++ b/meteor/treebuilder.py
@@ -45,6 +45,7 @@ class TreeBuilder(Session):
     width: int
     height: int
     format: str | None
+    gtr: bool
 
     def __post_init__(self) -> None:
         if self.format not in TreeBuilder.OUTPUT_FORMATS:
@@ -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