phylogenie 2.1.23__py3-none-any.whl → 2.1.24__py3-none-any.whl

phylogenie/__init__.py

@@ -11,7 +11,7 @@ from phylogenie.generators import (
     FBDTreeDatasetGenerator,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import dump_newick, load_fasta, load_newick
+from phylogenie.io import dump_newick, load_fasta, load_newick, load_nexus
 from phylogenie.msa import MSA
 from phylogenie.skyline import (
     SkylineMatrix,
@@ -100,6 +100,7 @@ __all__ = [
     "generate_trees",
     "simulate_tree",
     "dump_newick",
+    "load_nexus",
     "load_fasta",
     "load_newick",
     "MSA",

phylogenie/draw.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from itertools import islice
 from typing import Any
 
 import matplotlib.colors as mcolors
@@ -25,7 +26,7 @@ def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color])
 
     xs = (
         get_node_depth_levels(tree)
-        if any(node.branch_length is None for node in tree)
+        if any(node.branch_length is None for node in islice(tree, 1, None))
         else get_node_depths(tree)
     )
     ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}

phylogenie/generators/dataset.py

@@ -56,7 +56,7 @@ class DatasetGenerator(ABC, StrictBaseModel):
             for i in range(n_samples)
         )
         df = pd.DataFrame(
-            [r for r in tqdm(jobs, total=n_samples, desc=f"Generating {data_dir}...")]
+            [j for j in tqdm(jobs, f"Generating {data_dir}...", n_samples)]
         )
         df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
 

phylogenie/io/__init__.py

@@ -0,0 +1,5 @@
+from phylogenie.io.fasta import load_fasta
+from phylogenie.io.newick import dump_newick, load_newick
+from phylogenie.io.nexus import load_nexus
+
+__all__ = ["load_fasta", "load_newick", "dump_newick", "load_nexus"]

phylogenie/io/fasta.py

@@ -0,0 +1,25 @@
+from typing import Callable
+
+from phylogenie.msa import MSA, Sequence
+
+
+def load_fasta(
+    fasta_file: str, extract_time_from_id: Callable[[str], float] | None = None
+) -> MSA:
+    sequences: list[Sequence] = []
+    with open(fasta_file, "r") as f:
+        for line in f:
+            if not line.startswith(">"):
+                raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
+            id = line[1:].strip()
+            time = None
+            if extract_time_from_id is not None:
+                time = extract_time_from_id(id)
+            elif "|" in id:
+                try:
+                    time = float(id.split("|")[-1])
+                except ValueError:
+                    pass
+            chars = next(f).strip()
+            sequences.append(Sequence(id, chars, time))
+    return MSA(sequences)

phylogenie/{io.py → io/newick.py}

@@ -1,17 +1,17 @@
 import re
-from typing import Callable
 
-from phylogenie.msa import MSA, Sequence
 from phylogenie.tree import Tree
 
 
-def _parse_newick(newick: str) -> Tree:
+def parse_newick(newick: str, translations: dict[str, str] | None = None) -> Tree:
     newick = newick.strip()
+    newick = re.sub(r"^\[\&[^\]]*\]", "", newick).strip()
+
     stack: list[list[Tree]] = []
     current_children: list[Tree] = []
     current_nodes: list[Tree] = []
     i = 0
-    while i < len(newick):
+    while True:
 
         def _read_chars(stoppers: list[str]) -> str:
             nonlocal i
@@ -29,7 +29,10 @@ def _parse_newick(newick: str) -> Tree:
             i += 1
             continue
 
-        current_node = Tree(_read_chars(["[", ":", ",", ")", ";"]))
+        name = _read_chars([":", "[", ",", ")", ";"])
+        if translations is not None and name in translations:
+            name = translations[name]
+        current_node = Tree(name)
 
         if newick[i] == "[":
             i += 1
@@ -62,17 +65,15 @@ def _parse_newick(newick: str) -> Tree:
 
         i += 1
 
-    raise ValueError("Newick string should end with ';'")
-
 
 def load_newick(filepath: str) -> Tree | list[Tree]:
     with open(filepath, "r") as file:
-        trees = [_parse_newick(newick) for newick in file]
+        trees = [parse_newick(newick) for newick in file]
     return trees[0] if len(trees) == 1 else trees
 
 
-def _to_newick(tree: Tree) -> str:
-    children_newick = ",".join([_to_newick(child) for child in tree.children])
+def to_newick(tree: Tree) -> str:
+    children_newick = ",".join([to_newick(child) for child in tree.children])
     newick = tree.name
     if tree.features:
         reprs = {k: repr(v).replace("'", '"') for k, v in tree.features.items()}
@@ -99,26 +100,4 @@ def dump_newick(trees: Tree | list[Tree], filepath: str) -> None:
         trees = [trees]
     with open(filepath, "w") as file:
         for t in trees:
-            file.write(_to_newick(t) + ";\n")
-
-
-def load_fasta(
-    fasta_file: str, extract_time_from_id: Callable[[str], float] | None = None
-) -> MSA:
-    sequences: list[Sequence] = []
-    with open(fasta_file, "r") as f:
-        for line in f:
-            if not line.startswith(">"):
-                raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
-            id = line[1:].strip()
-            time = None
-            if extract_time_from_id is not None:
-                time = extract_time_from_id(id)
-            elif "|" in id:
-                try:
-                    time = float(id.split("|")[-1])
-                except ValueError:
-                    pass
-            chars = next(f).strip()
-            sequences.append(Sequence(id, chars, time))
-    return MSA(sequences)
+            file.write(to_newick(t) + ";\n")

phylogenie/io/nexus.py

@@ -0,0 +1,43 @@
+import re
+from collections.abc import Iterator
+
+from phylogenie.io.newick import parse_newick
+from phylogenie.tree import Tree
+
+
+def _parse_translate_block(lines: Iterator[str]) -> dict[str, str]:
+    translations: dict[str, str] = {}
+    for line in lines:
+        match = re.match(r"\s*(\d+)\s+['\"]?([^'\",;]+)['\"]?", line)
+        if match is None:
+            if ";" in line:
+                return translations
+            else:
+                raise ValueError(f"Invalid translate line: {line.strip()}")
+        translations[match.group(1)] = match.group(2)
+    raise ValueError("Translate block not terminated with ';'")
+
+
+def _parse_trees_block(lines: Iterator[str]) -> dict[str, Tree]:
+    trees: dict[str, Tree] = {}
+    translations = {}
+    for line in lines:
+        line = line.strip()
+        if line.upper() == "TRANSLATE":
+            translations = _parse_translate_block(lines)
+        elif line.upper() == "END;":
+            return trees
+        else:
+            match = re.match(r"^TREE\s*\*?\s+(\S+)\s*=\s*(.+)$", line, re.IGNORECASE)
+            if match is None:
+                raise ValueError(f"Invalid tree line. Expected 'TREE name = newick'")
+            trees[match.group(1)] = parse_newick(match.group(2), translations)
+    return trees
+
+
+def load_nexus(nexus_file: str) -> dict[str, Tree]:
+    with open(nexus_file, "r") as f:
+        for line in f:
+            if line.strip().upper() == "BEGIN TREES;":
+                return _parse_trees_block(f)
+    raise ValueError("No TREES block found in the NEXUS file.")

{phylogenie-2.1.23.dist-info → phylogenie-2.1.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.23
+Version: 2.1.24
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com

{phylogenie-2.1.23.dist-info → phylogenie-2.1.24.dist-info}/RECORD

@@ -1,13 +1,16 @@
-phylogenie/__init__.py,sha256=8ulA-U7-WnBLNsqRYqPOPrJG8X4ZzyEU02oaDbR_Hxs,2849
-phylogenie/draw.py,sha256=WCjr_HCn-wCuxwkDhHA9Ijl1Sm6av6FBLkyR0MK2bVM,5271
+phylogenie/__init__.py,sha256=KV83hJ153WAwCIWEPWXtSsqdI_ncojxiXzt1X_lC7SU,2879
+phylogenie/draw.py,sha256=37eQDK8TKmL8bxSzXvi7jNvrR7ulj7MByR5MA1FtyyE,5317
 phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
 phylogenie/generators/alisim.py,sha256=0aCLuGInifWgAvfh7zARWSKF4EMw3TjlPXMLSECui0k,2783
 phylogenie/generators/configs.py,sha256=WFoeKpgj9ZQIom7BKqwpgXbriiQGg3jFBMLoD8KButk,1073
-phylogenie/generators/dataset.py,sha256=pPwW9yxm9fkU0PPllFq8EsPlqau8tth-4OatbA_hEHo,2120
+phylogenie/generators/dataset.py,sha256=kY92diePr2IjiLejHLixJoYRc-2LpM-GBt3wkX9SYvA,2109
 phylogenie/generators/factories.py,sha256=TuVFQWRjq33Hewjw_Lp8tQ0l_IPtqYDyQCNJhtiHpw8,7882
 phylogenie/generators/trees.py,sha256=6tHS013RNRyIsObeYq2Kx9it7Yc0TgiMkeAHbskHzAM,10748
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
-phylogenie/io.py,sha256=vUG2yVtoV98tNHut46uSuB3VPj6s64VDhUp1EaSk0o0,4084
+phylogenie/io/__init__.py,sha256=gtRYtDdZSTlWCj3I4vmMJSAs93jdz5RySkCakD3sxlQ,214
+phylogenie/io/fasta.py,sha256=IWtNb_RQLR6kvS0G826wB9SodkCGfugddoUHx78Yrec,837
+phylogenie/io/newick.py,sha256=vw0fafh1LL0SXQifIIa1TQ7g5KgTCIAX6vzq-bUOrKE,3396
+phylogenie/io/nexus.py,sha256=IKbV8lJ_Q053iYJ7JzVQPCUqSkSfmiRpUchFTrLHZuE,1551
 phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
 phylogenie/models.py,sha256=pCg9ob0RpLUHwM49x4knKxL4FNPr3-EU_6zMXsvxtAg,370
 phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
@@ -28,8 +31,8 @@ phylogenie/treesimulator/model.py,sha256=Nyg6R8XmMwZMSw1-dII81sU9uU7tDe-NMs8v1qK
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=GknvAFXyiaWeeYJ8Lk5d6E2VHT-xW6ONEojYbtJYiB8,476
 phylogenie/utils.py,sha256=ehVk_2kvjW8Q_EyM2kxBPHYiK-KlPmZQx7JeVN6Fh-E,5419
-phylogenie-2.1.23.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.23.dist-info/METADATA,sha256=ym-NHK9hkvBR4kH_KGCnLnckV_iEzvFk0fdMzrkVudQ,5477
-phylogenie-2.1.23.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.23.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.23.dist-info/RECORD,,
+phylogenie-2.1.24.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.24.dist-info/METADATA,sha256=18OWTAKMES6VnsF454GNsr7NRDZ2AQNPprpkcxDfbZk,5477
+phylogenie-2.1.24.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.24.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.24.dist-info/RECORD,,