phylogenie 3.1.8__py3-none-any.whl → 3.1.10__py3-none-any.whl

phylogenie/draw.py

@@ -32,6 +32,8 @@ def draw_tree(
     ax: Axes | None = None,
     colors: Color | dict[Tree, Color] = "black",
     backward_time: bool = False,
+    branch_kwargs: dict[str, Any] | None = None,
+    sampled_ancestor_kwargs: dict[str, Any] | None = None,
 ) -> Axes:
     """
     Draw a phylogenetic tree with colored branches.
@@ -46,6 +48,10 @@ def draw_tree(
         A single color for all branches or a dictionary mapping each node to a color.
     backward_time : bool, optional
         If True, the x-axis is inverted to represent time going backward.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
+    sampled_ancestor_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to highlight sampled ancestors.
 
     Returns
     -------
@@ -54,6 +60,12 @@ def draw_tree(
     """
     if ax is None:
         ax = plt.gca()
+    if branch_kwargs is None:
+        branch_kwargs = {}
+    if sampled_ancestor_kwargs is None:
+        sampled_ancestor_kwargs = {}
+    if "marker" not in sampled_ancestor_kwargs:
+        sampled_ancestor_kwargs["marker"] = "o"
 
     if not isinstance(colors, dict):
         colors = {node: colors for node in tree}
@@ -66,20 +78,33 @@ def draw_tree(
         else get_node_depths(tree)
     )
 
-    ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}
+    leaves = tree.get_leaves()
+    ys: dict[Tree, float] = {
+        node: i
+        for i, node in enumerate(leaves)
+        if node.parent is None or node.branch_length != 0
+    }
     for node in tree.postorder_traversal():
         if node.is_internal():
-            ys[node] = sum(ys[child] for child in node.children) / len(node.children)
+            children = [child for child in node.children if child.branch_length != 0]
+            ys[node] = sum(ys[child] for child in children) / len(children)
+    for leaf in leaves:
+        if leaf.parent is not None and leaf.branch_length == 0:
+            ys[leaf] = ys[leaf.parent]
 
     if tree.branch_length is not None:
         xmin = xs[tree] + tree.branch_length if backward_time else 0
-        ax.hlines(y=ys[tree], xmin=xmin, xmax=xs[tree], color=colors[tree])  # pyright: ignore
+        ax.hlines(  # pyright: ignore
+            y=ys[tree], xmin=xmin, xmax=xs[tree], color=colors[tree], **branch_kwargs
+        )
     for node in tree:
         x1, y1 = xs[node], ys[node]
+        if node.parent is not None and node.branch_length == 0:
+            ax.plot(x1, y1, color=colors[node], **sampled_ancestor_kwargs)  # pyright: ignore
         for child in node.children:
             x2, y2 = xs[child], ys[child]
-            ax.hlines(y=y2, xmin=x1, xmax=x2, color=colors[child])  # pyright: ignore
-            ax.vlines(x=x1, ymin=y1, ymax=y2, color=colors[child])  # pyright: ignore
+            ax.hlines(y=y2, xmin=x1, xmax=x2, color=colors[child], **branch_kwargs)  # pyright: ignore
+            ax.vlines(x=x1, ymin=y1, ymax=y2, color=colors[child], **branch_kwargs)  # pyright: ignore
 
     if backward_time:
         ax.invert_xaxis()
@@ -117,6 +142,7 @@ def draw_dated_tree(
     calibration_nodes: tuple[CalibrationNode, CalibrationNode],
     ax: Axes | None = None,
     colors: Color | dict[Tree, Color] = "black",
+    branch_kwargs: dict[str, Any] | None = None,
 ) -> Axes:
     """
     Draw a phylogenetic tree with branches positioned according to calibrated dates.
@@ -131,6 +157,8 @@ def draw_dated_tree(
         The matplotlib Axes to draw on. If None, uses the current Axes.
     colors : Color | dict[Tree, Color], optional
         A single color for all branches or a dictionary mapping each node to a color.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
 
     Returns
     -------
@@ -139,6 +167,8 @@ def draw_dated_tree(
     """
     if ax is None:
         ax = plt.gca()
+    if branch_kwargs is None:
+        branch_kwargs = {}
 
     if not isinstance(colors, dict):
         colors = {node: colors for node in tree}
@@ -160,6 +190,7 @@ def draw_dated_tree(
             xmin=mdates.date2num(origin_date),  # pyright: ignore
             xmax=mdates.date2num(xs[tree]),  # pyright: ignore
             color=colors[tree],
+            **branch_kwargs,
         )
     for node in tree:
         x1, y1 = xs[node], ys[node]
@@ -170,8 +201,15 @@ def draw_dated_tree(
                 xmin=mdates.date2num(x1),  # pyright: ignore
                 xmax=mdates.date2num(x2),  # pyright: ignore
                 color=colors[child],
+                **branch_kwargs,
+            )
+            ax.vlines(  # pyright: ignore
+                x=mdates.date2num(x1),  # pyright: ignore
+                ymin=y1,
+                ymax=y2,
+                color=colors[child],
+                **branch_kwargs,
             )
-            ax.vlines(x=mdates.date2num(x1), ymin=y1, ymax=y2, color=colors[child])  # pyright: ignore
 
     ax.xaxis.set_major_locator(mdates.AutoDateLocator())
     ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
@@ -260,6 +298,8 @@ def draw_colored_tree_categorical(
     show_legend: bool = True,
     labels: dict[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    sampled_ancestor_kwargs: dict[str, Any] | None = None,
 ):
     """
     Draw a phylogenetic tree with branches colored based on categorical metadata.
@@ -284,6 +324,10 @@ def draw_colored_tree_categorical(
         A mapping from category values to labels for the legend.
     legend_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the legend.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
+    sampled_ancestor_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to highlight sampled ancestors.
 
     Returns
     -------
@@ -300,7 +344,14 @@ def draw_colored_tree_categorical(
         labels=labels,
         legend_kwargs=legend_kwargs,
     )
-    return draw_tree(tree=tree, ax=ax, colors=colors, backward_time=backward_time)
+    return draw_tree(
+        tree=tree,
+        ax=ax,
+        colors=colors,
+        backward_time=backward_time,
+        branch_kwargs=branch_kwargs,
+        sampled_ancestor_kwargs=sampled_ancestor_kwargs,
+    )
 
 
 def draw_colored_dated_tree_categorical(
@@ -313,6 +364,7 @@ def draw_colored_dated_tree_categorical(
     show_legend: bool = True,
     labels: dict[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
 ) -> Axes:
     """
     Draw a dated phylogenetic tree with branches colored based on categorical metadata.
@@ -337,6 +389,8 @@ def draw_colored_dated_tree_categorical(
         A mapping from category values to labels for the legend.
     legend_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the legend.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
 
     Returns
     -------
@@ -354,7 +408,11 @@ def draw_colored_dated_tree_categorical(
         legend_kwargs=legend_kwargs,
     )
     return draw_dated_tree(
-        tree=tree, calibration_nodes=calibration_nodes, ax=ax, colors=colors
+        tree=tree,
+        calibration_nodes=calibration_nodes,
+        ax=ax,
+        colors=colors,
+        branch_kwargs=branch_kwargs,
     )
 
 
@@ -382,7 +440,7 @@ def _init_colored_tree_continuous(
     vmin: float | None = ...,
     vmax: float | None = ...,
     *,
-    show_hist: Literal[True],
+    show_hist: Literal[True] = True,
     hist_kwargs: dict[str, Any] | None = ...,
     hist_axes_kwargs: dict[str, Any] | None = ...,
 ) -> tuple[Axes, dict[Tree, Color], Axes]: ...
@@ -394,7 +452,6 @@ def _init_colored_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
-    *,
     show_hist: bool = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
@@ -478,6 +535,8 @@ def draw_colored_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    sampled_ancestor_kwargs: dict[str, Any] | None = None,
     *,
     show_hist: Literal[False],
     hist_kwargs: dict[str, Any] | None = None,
@@ -493,8 +552,9 @@ def draw_colored_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
-    *,
-    show_hist: Literal[True],
+    branch_kwargs: dict[str, Any] | None = None,
+    sampled_ancestor_kwargs: dict[str, Any] | None = None,
+    show_hist: Literal[True] = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
 ) -> tuple[Axes, Axes]: ...
@@ -507,6 +567,8 @@ def draw_colored_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    sampled_ancestor_kwargs: dict[str, Any] | None = None,
     show_hist: bool = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
@@ -532,6 +594,10 @@ def draw_colored_tree_continuous(
         The minimum value for normalization. If None, uses the minimum of the data.
     vmax : float | None, optional
         The maximum value for normalization. If None, uses the maximum of the data.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
+    sampled_ancestor_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to highlight sampled ancestors.
     show_hist : bool, optional
         Whether to display a histogram of the continuous values.
     hist_kwargs : dict[str, Any] | None, optional
@@ -559,7 +625,12 @@ def draw_colored_tree_continuous(
             hist_axes_kwargs=hist_axes_kwargs,
         )
         return draw_tree(
-            tree=tree, ax=ax, colors=colors, backward_time=backward_time
+            tree=tree,
+            ax=ax,
+            colors=colors,
+            backward_time=backward_time,
+            branch_kwargs=branch_kwargs,
+            sampled_ancestor_kwargs=sampled_ancestor_kwargs,
         ), hist_ax
 
     ax, colors = _init_colored_tree_continuous(
@@ -574,7 +645,14 @@ def draw_colored_tree_continuous(
         hist_kwargs=hist_kwargs,
         hist_axes_kwargs=hist_axes_kwargs,
     )
-    return draw_tree(tree=tree, ax=ax, colors=colors, backward_time=backward_time)
+    return draw_tree(
+        tree=tree,
+        ax=ax,
+        colors=colors,
+        backward_time=backward_time,
+        branch_kwargs=branch_kwargs,
+        sampled_ancestor_kwargs=sampled_ancestor_kwargs,
+    )
 
 
 @overload
@@ -587,6 +665,7 @@ def draw_colored_dated_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
     *,
     show_hist: Literal[False],
     hist_kwargs: dict[str, Any] | None = None,
@@ -602,8 +681,8 @@ def draw_colored_dated_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
-    *,
-    show_hist: Literal[True],
+    branch_kwargs: dict[str, Any] | None = None,
+    show_hist: Literal[True] = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
 ) -> tuple[Axes, Axes]: ...
@@ -616,6 +695,7 @@ def draw_colored_dated_tree_continuous(
     colormap: str | Colormap = "viridis",
     vmin: float | None = None,
     vmax: float | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
     show_hist: bool = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
@@ -641,6 +721,8 @@ def draw_colored_dated_tree_continuous(
         The minimum value for normalization. If None, uses the minimum of the data.
     vmax : float | None, optional
         The maximum value for normalization. If None, uses the maximum of the data.
+    branch_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the branch drawing functions.
     show_hist : bool, optional
         Whether to display a histogram of the continuous values.
     hist_kwargs : dict[str, Any] | None, optional
@@ -672,6 +754,7 @@ def draw_colored_dated_tree_continuous(
             calibration_nodes=calibration_nodes,
             ax=ax,
             colors=colors,
+            branch_kwargs=branch_kwargs,
         ), hist_ax
 
     ax, colors = _init_colored_tree_continuous(
@@ -691,4 +774,5 @@ def draw_colored_dated_tree_continuous(
         calibration_nodes=calibration_nodes,
         ax=ax,
         colors=colors,
+        branch_kwargs=branch_kwargs,
     )

phylogenie/io/fasta.py

@@ -1,3 +1,4 @@
+from datetime import date
 from pathlib import Path
 from typing import Callable
 
@@ -5,7 +6,8 @@ from phylogenie.msa import MSA, Sequence
 
 
 def load_fasta(
-    fasta_file: str | Path, extract_time_from_id: Callable[[str], float] | None = None
+    fasta_file: str | Path,
+    extract_time_from_id: Callable[[str], float | date] | None = None,
 ) -> MSA:
     sequences: list[Sequence] = []
     with open(fasta_file, "r") as f:
@@ -17,10 +19,14 @@ def load_fasta(
             if extract_time_from_id is not None:
                 time = extract_time_from_id(id)
             elif "|" in id:
+                last_metadata = id.split("|")[-1]
                 try:
-                    time = float(id.split("|")[-1])
+                    time = float(last_metadata)
                 except ValueError:
-                    pass
+                    try:
+                        time = date.fromisoformat(last_metadata)
+                    except ValueError:
+                        pass
             chars = next(f).strip()
             sequences.append(Sequence(id, chars, time))
     return MSA(sequences)

phylogenie/msa.py

@@ -1,5 +1,6 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
+from datetime import date
 
 import numpy as np
 
@@ -8,7 +9,7 @@ import numpy as np
 class Sequence:
     id: str
     chars: str
-    time: float | None = None
+    time: float | date | None = None
 
 
 class MSA:
@@ -25,8 +26,8 @@ class MSA:
         return [sequence.id for sequence in self.sequences]
 
     @property
-    def times(self) -> list[float]:
-        times: list[float] = []
+    def times(self) -> list[float | date]:
+        times: list[float | date] = []
         for sequence in self:
             if sequence.time is None:
                 raise ValueError(f"Time is not set for sequence {sequence.id}.")

{phylogenie-3.1.8.dist-info → phylogenie-3.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.8
+Version: 3.1.10
 Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{phylogenie-3.1.8.dist-info → phylogenie-3.1.10.dist-info}/RECORD

@@ -1,8 +1,8 @@
 phylogenie/__init__.py,sha256=redo-SdvWwQPlyqOfdSW0TkkZf-aP98kLRCUIg4pB8c,3304
-phylogenie/draw.py,sha256=J_ijON8O848f-r2VnOu8MbSR_vAcvO8zPq5u5Vq5J_U,23323
+phylogenie/draw.py,sha256=ZZvOQrmW_X9wUiOl5h2nvNBVX7BFH014EJ7evEI2orA,26840
 phylogenie/main.py,sha256=ry3B3HiwibZG3_qB58T5UhWy5dp6neYUtSqzL9LrSkA,1698
 phylogenie/mixins.py,sha256=wMwqP6zkqME9eMyzx5FS6-p9X8yW09jIC8jge8pHlkk,907
-phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
+phylogenie/msa.py,sha256=RGXxmo1WX9wi5NpEQTxMuH4zk8vM0f8A7Knm0YUJ9AI,2111
 phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=p694PBe_tk25A6N8vGGWxuqoDtt3nHFUsIYJrwR_76Y,494
@@ -14,7 +14,7 @@ phylogenie/generators/factories.py,sha256=2mTFdFbbLyV3v79JaOEVtqLOmxQHaOUv1S-Y3v
 phylogenie/generators/trees.py,sha256=8dO1CkU34E6mmMAHrYqiLV_VA8r54cSEOo-UzoHiN20,10467
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
 phylogenie/io/__init__.py,sha256=3v_bxv9RVeQ3KZzxNFeV9KLVxaC_Whf7rgWLnBKoEr0,95
-phylogenie/io/fasta.py,sha256=kx9uVATLpzpXdhhNNMvMpB5Vdwh9CTWepTGachvEwV4,1154
+phylogenie/io/fasta.py,sha256=1Bjm_70fkqTA5rCsBY17iMW-lyW1smM-wjq3dl93sFc,1371
 phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
 phylogenie/skyline/matrix.py,sha256=v4SitY7VbXprqlqQckjWTzW5hwRmCyIF595R6IJMxWw,9268
 phylogenie/skyline/parameter.py,sha256=TVqkqirGXNN-VP8hnIJACPkOxUan6LkGa5o_JcPfwbY,4834
@@ -33,9 +33,9 @@ phylogenie/treesimulator/events/mutations.py,sha256=8Nqa2fg7fwaVNe5XSkGDSwp9pIKQ
 phylogenie/treesimulator/io/__init__.py,sha256=rfP-zp8SP8baq5_4dPAr10WH0W6KfoMCxdTZDCSXtzE,185
 phylogenie/treesimulator/io/newick.py,sha256=8Pr_jixByPOaVch18w-rFt62HYy0U97YMu0H-QSwIy0,3449
 phylogenie/treesimulator/io/nexus.py,sha256=zqT9dzj413z_s0hqp3Cdq5NMO6lv-zuuaJlaqzaqaB8,1847
-phylogenie-3.1.8.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-3.1.8.dist-info/METADATA,sha256=0lQMRPORgVWh0cp4OYt1RjCkYtCdy2P3ugiBmqB2oHQ,5194
-phylogenie-3.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-phylogenie-3.1.8.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
-phylogenie-3.1.8.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
-phylogenie-3.1.8.dist-info/RECORD,,
+phylogenie-3.1.10.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-3.1.10.dist-info/METADATA,sha256=sdHbD8Gx14dS0TkKkN9dpmZkTKgt0ib__qtTX2noIo0,5195
+phylogenie-3.1.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+phylogenie-3.1.10.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
+phylogenie-3.1.10.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
+phylogenie-3.1.10.dist-info/RECORD,,