phylogenie 2.1.23__tar.gz → 2.1.24__tar.gz

{phylogenie-2.1.23 → phylogenie-2.1.24}/PKG-INFO

@@ -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 → phylogenie-2.1.24}/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-2.1.23 → phylogenie-2.1.24}/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-2.1.23 → phylogenie-2.1.24}/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-2.1.24/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-2.1.24/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-2.1.23/phylogenie/io.py → phylogenie-2.1.24/phylogenie/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-2.1.24/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 → phylogenie-2.1.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "phylogenie"
-version = "2.1.23"
+version = "2.1.24"
 description = "Generate phylogenetic datasets with minimal setup effort"
 authors = ["Gabriele Marino <gabmarino.8601@gmail.com>"]
 readme = "README.md"