phylogenie 2.1.27__py3-none-any.whl → 2.1.29__py3-none-any.whl

phylogenie/draw.py

@@ -57,6 +57,7 @@ def draw_tree(
     vmin: float | None = None,
     vmax: float | None = None,
     show_legend: bool = True,
+    labels: dict[Any, Any] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
     show_hist: bool = True,
     hist_kwargs: dict[str, Any] | None = None,
@@ -98,7 +99,10 @@ def draw_tree(
 
         if show_legend:
             legend_handles = [
-                mpatches.Patch(color=feature_colors[f], label=str(f))
+                mpatches.Patch(
+                    color=feature_colors[f],
+                    label=str(f) if labels is None else labels[f],
+                )
                 for f in feature_colors
             ]
             if any(color_by not in node.features for node in tree):

phylogenie/io/fasta.py

@@ -1,10 +1,11 @@
+from pathlib import Path
 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
+    fasta_file: str | Path, extract_time_from_id: Callable[[str], float] | None = None
 ) -> MSA:
     sequences: list[Sequence] = []
     with open(fasta_file, "r") as f:

phylogenie/io/newick.py

@@ -1,4 +1,5 @@
 import re
+from pathlib import Path
 
 from phylogenie.tree import Tree
 
@@ -66,7 +67,7 @@ def parse_newick(newick: str, translations: dict[str, str] | None = None) -> Tre
         i += 1
 
 
-def load_newick(filepath: str) -> Tree | list[Tree]:
+def load_newick(filepath: str | Path) -> Tree | list[Tree]:
     with open(filepath, "r") as file:
         trees = [parse_newick(newick) for newick in file]
     return trees[0] if len(trees) == 1 else trees
@@ -95,7 +96,7 @@ def to_newick(tree: Tree) -> str:
     return newick
 
 
-def dump_newick(trees: Tree | list[Tree], filepath: str) -> None:
+def dump_newick(trees: Tree | list[Tree], filepath: str | Path) -> None:
     if isinstance(trees, Tree):
         trees = [trees]
     with open(filepath, "w") as file:

phylogenie/io/nexus.py

@@ -1,5 +1,6 @@
 import re
 from collections.abc import Iterator
+from pathlib import Path
 
 from phylogenie.io.newick import parse_newick
 from phylogenie.tree import Tree
@@ -35,7 +36,7 @@ def _parse_trees_block(lines: Iterator[str]) -> dict[str, Tree]:
     return trees
 
 
-def load_nexus(nexus_file: str) -> dict[str, Tree]:
+def load_nexus(nexus_file: str | Path) -> dict[str, Tree]:
     with open(nexus_file, "r") as f:
         for line in f:
             if line.strip().upper() == "BEGIN TREES;":

phylogenie/treesimulator/gillespie.py

@@ -1,6 +1,6 @@
-import os
 import time
 from collections.abc import Iterable, Sequence
+from pathlib import Path
 from typing import Any
 
 import joblib
@@ -106,7 +106,7 @@ def simulate_tree(
 
 
 def generate_trees(
-    output_dir: str,
+    output_dir: str | Path,
     n_trees: int,
     events: Sequence[Event],
     min_tips: int = 1,
@@ -119,7 +119,13 @@ def generate_trees(
     n_jobs: int = -1,
     timeout: float = np.inf,
 ) -> pd.DataFrame:
-    def _simulate_tree(seed: int) -> tuple[Tree, dict[str, Any]]:
+    if isinstance(output_dir, str):
+        output_dir = Path(output_dir)
+    if output_dir.exists():
+        raise FileExistsError(f"Output directory {output_dir} already exists")
+    output_dir.mkdir(parents=True)
+
+    def _simulate_tree(i: int, seed: int) -> dict[str, Any]:
         while True:
             try:
                 tree, metadata = simulate_tree(
@@ -132,27 +138,21 @@ def generate_trees(
                     seed=seed,
                     timeout=timeout,
                 )
+                metadata["file_id"] = i
                 if node_features is not None:
                     set_features(tree, node_features)
-                return (tree, metadata)
+                dump_newick(tree, output_dir / f"{i}.nwk")
+                return metadata
             except TimeoutError:
                 print("Simulation timed out, retrying with a different seed...")
             seed += 1
 
-    if os.path.exists(output_dir):
-        raise FileExistsError(f"Output directory {output_dir} already exists")
-    os.makedirs(output_dir)
-
     rng = default_rng(seed)
     jobs = joblib.Parallel(n_jobs=n_jobs, return_as="generator_unordered")(
-        joblib.delayed(_simulate_tree)(seed=int(rng.integers(2**32)))
-        for _ in range(n_trees)
+        joblib.delayed(_simulate_tree)(i=i, seed=int(rng.integers(2**32)))
+        for i in range(n_trees)
     )
 
-    df: list[dict[str, Any]] = []
-    for i, (tree, metadata) in tqdm(
-        enumerate(jobs), total=n_trees, desc=f"Generating trees in {output_dir}..."
-    ):
-        df.append({"file_id": i} | metadata)
-        dump_newick(tree, os.path.join(output_dir, f"{i}.nwk"))
-    return pd.DataFrame(df)
+    return pd.DataFrame(
+        [md for md in tqdm(jobs, f"Generating trees in {output_dir}...", n_trees)]
+    )

{phylogenie-2.1.27.dist-info → phylogenie-2.1.29.dist-info}/METADATA

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

{phylogenie-2.1.27.dist-info → phylogenie-2.1.29.dist-info}/RECORD

@@ -1,5 +1,5 @@
 phylogenie/__init__.py,sha256=FdNHTw1U1eYDz8w_EOvk9R0rdOXvCMAgbm7ilXs6Fog,2861
-phylogenie/draw.py,sha256=37eQDK8TKmL8bxSzXvi7jNvrR7ulj7MByR5MA1FtyyE,5317
+phylogenie/draw.py,sha256=jM24oCGkU66IsEm56IAcq0pCSVXC8ltyYJQizNwDJl8,5451
 phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
 phylogenie/generators/alisim.py,sha256=1YQLpOG_Bpn9YqExQqEu-wz1MDGCbpPmTdhrBb6TbWc,2820
 phylogenie/generators/configs.py,sha256=zjNnljBBC-d6DTjWjieE8-HnKeE58LAy3moBYiD6NFo,1042
@@ -8,9 +8,9 @@ phylogenie/generators/factories.py,sha256=qSz-s7ltQwUf7R6TPZ-5xFyOnnSHkKn0u5Iuke
 phylogenie/generators/trees.py,sha256=zsmL3_YUUwtn-_NrGSvwZVSGIm9kpf8SPHxoQCIT_hE,10168
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
 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/io/fasta.py,sha256=CUFO06m7wClprarsMheZojM4Os2NQf3ALYUXSWzfNL0,869
+phylogenie/io/newick.py,sha256=4j7m9B11erVJijye0GfpbWomk2ttlBXdyHNoBc_fvIg,3435
+phylogenie/io/nexus.py,sha256=KGf0rH8sAVAuPnmhokj2XbQTBsiuhPujJflRixUunGY,1583
 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
@@ -27,13 +27,13 @@ phylogenie/treesimulator/events/contact_tracing.py,sha256=QqinLfQmBBi6qF3nFE_Qsi
 phylogenie/treesimulator/events/core.py,sha256=bhgQgi5L-oaHsoWJmUOsTTzWxi0POYxVLoF-KrC8AGQ,8179
 phylogenie/treesimulator/events/mutations.py,sha256=fGxxeXhiFrMAaS-s_vd42aVBZ6XQiRjpW9uJUATNk1E,3752
 phylogenie/treesimulator/features.py,sha256=nNfv0kPudAZnrnufIVuFCR_VEJSP-6fZ2L-VG5_-Fqo,1366
-phylogenie/treesimulator/gillespie.py,sha256=_O1TqK-LOjH6k8ZSos-CbVQ7wJZPzf1Awv_1DXt0NYE,5730
+phylogenie/treesimulator/gillespie.py,sha256=3xPeqkSitP71Qa2a8d8vMkdV4CCdLy1rRCdDsNctAug,5732
 phylogenie/treesimulator/model.py,sha256=lhDwmBFQ8Qh8qVGZPgED0vehtPC3DE7_CgCV_8rPB-A,4641
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=p694PBe_tk25A6N8vGGWxuqoDtt3nHFUsIYJrwR_76Y,494
 phylogenie/utils.py,sha256=ehVk_2kvjW8Q_EyM2kxBPHYiK-KlPmZQx7JeVN6Fh-E,5419
-phylogenie-2.1.27.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.27.dist-info/METADATA,sha256=UUZTU3YJ5yrAzQZ8EiiVsTmhuZzKNiGU_roQazkf3Mo,5477
-phylogenie-2.1.27.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.27.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.27.dist-info/RECORD,,
+phylogenie-2.1.29.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.29.dist-info/METADATA,sha256=mVsGaofekWG_o0Mu2CqYu7RZsEL7MKsY3m5w7wv8hAw,5477
+phylogenie-2.1.29.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.29.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.29.dist-info/RECORD,,