phylogenie 3.1.9__tar.gz → 3.1.10__tar.gz

{phylogenie-3.1.9/src/phylogenie.egg-info → phylogenie-3.1.10}/PKG-INFO

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

{phylogenie-3.1.9 → phylogenie-3.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "phylogenie"
-version = "3.1.9"
+version = "3.1.10"
 description = "Generate phylogenetic datasets with minimal setup effort"
 readme = "README.md"
 requires-python = ">=3.10"

{phylogenie-3.1.9 → phylogenie-3.1.10}/src/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,
     )
 
 
@@ -477,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,
@@ -492,6 +552,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[True] = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
@@ -505,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,
@@ -530,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
@@ -557,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(
@@ -572,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
@@ -585,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,
@@ -600,6 +681,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[True] = True,
     hist_kwargs: dict[str, Any] | None = None,
     hist_axes_kwargs: dict[str, Any] | None = None,
@@ -613,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,
@@ -638,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
@@ -669,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(
@@ -688,4 +774,5 @@ def draw_colored_dated_tree_continuous(
         calibration_nodes=calibration_nodes,
         ax=ax,
         colors=colors,
+        branch_kwargs=branch_kwargs,
     )

{phylogenie-3.1.9 → phylogenie-3.1.10/src/phylogenie.egg-info}/PKG-INFO

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