phylogenie 2.1.19__py3-none-any.whl → 2.1.21__py3-none-any.whl

phylogenie/__init__.py

@@ -1,3 +1,4 @@
+from phylogenie.draw import Coloring, draw_tree
 from phylogenie.generators import (
     AliSimDatasetGenerator,
     BDEITreeDatasetGenerator,
@@ -12,7 +13,6 @@ from phylogenie.generators import (
 )
 from phylogenie.io import dump_newick, load_fasta, load_newick
 from phylogenie.msa import MSA
-from phylogenie.plot import Coloring, plot_tree
 from phylogenie.skyline import (
     SkylineMatrix,
     SkylineMatrixCoercible,
@@ -46,6 +46,18 @@ from phylogenie.treesimulator import (
     get_FBD_events,
     simulate_tree,
 )
+from phylogenie.utils import (
+    compute_colless_index,
+    compute_mean_leaf_pairwise_distance,
+    compute_sackin_index,
+    get_distance,
+    get_mrca,
+    get_node_depth_levels,
+    get_node_depths,
+    get_node_height_levels,
+    get_node_heights,
+    get_node_leaf_counts,
+)
 
 __all__ = [
     "AliSimDatasetGenerator",
@@ -92,5 +104,15 @@ __all__ = [
     "load_newick",
     "MSA",
     "Coloring",
-    "plot_tree",
+    "draw_tree",
+    "compute_colless_index",
+    "compute_mean_leaf_pairwise_distance",
+    "compute_sackin_index",
+    "get_distance",
+    "get_mrca",
+    "get_node_depths",
+    "get_node_depth_levels",
+    "get_node_heights",
+    "get_node_height_levels",
+    "get_node_leaf_counts",
 ]

phylogenie/draw.py

@@ -0,0 +1,142 @@
+from enum import Enum
+from typing import Any
+
+import matplotlib.colors as mcolors
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # pyright: ignore
+
+from phylogenie.tree import Tree
+from phylogenie.utils import get_node_depth_levels, get_node_depths
+
+
+class Coloring(str, Enum):
+    DISCRETE = "discrete"
+    CONTINUOUS = "continuous"
+
+
+Color = str | tuple[float, float, float] | tuple[float, float, float, float]
+
+
+def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color]) -> Axes:
+    if not isinstance(colors, dict):
+        colors = {node: colors for node in tree}
+
+    xs = (
+        get_node_depth_levels(tree)
+        if any(node.branch_length is None for node in tree)
+        else get_node_depths(tree)
+    )
+    ys = {node: i for i, node in enumerate(tree.inorder_traversal())}
+
+    for node in tree:
+        x1, y1 = xs[node], ys[node]
+        if node.parent is None:
+            ax.hlines(y=y1, xmin=0, xmax=x1, color=colors[node])  # pyright: ignore
+            continue
+        x0, y0 = xs[node.parent], ys[node.parent]
+        ax.vlines(x=x0, ymin=y0, ymax=y1, color=colors[node])  # pyright: ignore
+        ax.hlines(y=y1, xmin=x0, xmax=x1, color=colors[node])  # pyright: ignore
+
+    ax.set_yticks([])  # pyright: ignore
+    return ax
+
+
+def draw_tree(
+    tree: Tree,
+    ax: Axes | None = None,
+    color_by: str | None = None,
+    coloring: str | Coloring | None = None,
+    default_color: Color = "black",
+    cmap: str | None = None,
+    show_legend: bool = True,
+    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,
+) -> Axes | tuple[Axes, Axes]:
+    if ax is None:
+        ax = plt.gca()
+
+    if color_by is None:
+        return _draw_colored_tree(tree, ax, colors=default_color)
+
+    features = [node.get(color_by) for node in tree if color_by in node.features]
+
+    if coloring is None:
+        coloring = (
+            Coloring.CONTINUOUS
+            if any(isinstance(f, float) for f in features)
+            else Coloring.DISCRETE
+        )
+
+    if coloring == Coloring.DISCRETE:
+        if any(isinstance(f, float) for f in features):
+            raise ValueError(
+                "Discrete coloring selected but feature values are not all categorical."
+            )
+
+        colormap = plt.get_cmap("tab20" if cmap is None else cmap)
+        feature_colors = {
+            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features))
+        }
+        colors = {
+            node: (
+                feature_colors[node.get(color_by)]
+                if color_by in node.features
+                else default_color
+            )
+            for node in tree
+        }
+
+        if show_legend:
+            legend_handles = [
+                mpatches.Patch(color=feature_colors[f], label=str(f))
+                for f in feature_colors
+            ]
+            if any(color_by not in node.features for node in tree):
+                legend_handles.append(mpatches.Patch(color=default_color, label="NA"))
+            ax.legend(handles=legend_handles)  # pyright: ignore
+
+        return _draw_colored_tree(tree, ax, colors)
+
+    if coloring == Coloring.CONTINUOUS:
+        vmin = min(features) if vmin is None else vmin
+        vmax = max(features) if vmax is None else vmax
+        norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
+        colormap = plt.get_cmap("viridis" if cmap is None else cmap)
+        colors = {
+            node: (
+                colormap(norm(float(node.get(color_by))))
+                if color_by in node.features
+                else default_color
+            )
+            for node in tree
+        }
+
+        if show_hist:
+            default_hist_axes_kwargs = {"width": "25%", "height": "25%"}
+            if hist_axes_kwargs is not None:
+                default_hist_axes_kwargs.update(hist_axes_kwargs)
+            hist_ax = inset_axes(ax, **default_hist_axes_kwargs)  # pyright: ignore
+
+            hist_kwargs = {} if hist_kwargs is None else hist_kwargs
+            _, bins, patches = hist_ax.hist(features, **hist_kwargs)  # pyright: ignore
+
+            for patch, b0, b1 in zip(  # pyright: ignore
+                patches, bins[:-1], bins[1:]  # pyright: ignore
+            ):
+                midpoint = (b0 + b1) / 2  # pyright: ignore
+                patch.set_facecolor(colormap(norm(midpoint)))  # pyright: ignore
+            return _draw_colored_tree(tree, ax, colors), hist_ax  # pyright: ignore
+
+        else:
+            sm = plt.cm.ScalarMappable(cmap=colormap, norm=norm)
+            ax.get_figure().colorbar(sm, ax=ax)  # pyright: ignore
+            return _draw_colored_tree(tree, ax, colors)
+
+    raise ValueError(
+        f"Unknown coloring method: {coloring}. Choices are {list(Coloring)}."
+    )

phylogenie/tree.py

@@ -73,6 +73,8 @@ class Tree:
         del self._features[key]
 
     def add_child(self, child: "Tree") -> "Tree":
+        if child.parent is not None:
+            raise ValueError(f"Node {child.name} already has a parent.")
         child._parent = self
         self._children.append(child)
         return self
@@ -81,10 +83,12 @@ class Tree:
         self._children.remove(child)
         child._parent = None
 
-    def set_parent(self, node: "Tree | None"):
-        self._parent = node
-        if node is not None:
-            node._children.append(self)
+    def set_parent(self, parent: "Tree | None"):
+        if self.parent is not None:
+            self.parent.remove_child(self)
+        self._parent = parent
+        if parent is not None:
+            parent._children.append(self)
 
     def inorder_traversal(self) -> Iterator["Tree"]:
         if self.is_leaf():
@@ -107,6 +111,12 @@ class Tree:
             yield from child.postorder_traversal()
         yield self
 
+    def iter_ancestors(self, stop: "Tree | None" = None) -> Iterator["Tree"]:
+        node = self
+        while node is not None and node is not stop:
+            yield node
+            node = node.parent
+
     def breadth_first_traversal(self) -> Iterator["Tree"]:
         queue: deque["Tree"] = deque([self])
         while queue:
@@ -126,6 +136,12 @@ class Tree:
     def get_leaves(self) -> tuple["Tree", ...]:
         return tuple(node for node in self if node.is_leaf())
 
+    def get_internal_nodes(self) -> tuple["Tree", ...]:
+        return tuple(node for node in self if not node.is_leaf())
+
+    def is_binary(self) -> bool:
+        return all(len(node.children) in (0, 2) for node in self)
+
     def ladderize(self, criterion: Callable[["Tree"], Any]) -> None:
         self._children.sort(key=criterion)
         for child in self.children:

phylogenie/utils.py

@@ -1,3 +1,6 @@
+from itertools import chain
+from math import comb
+
 from phylogenie.tree import Tree
 
 
@@ -57,3 +60,117 @@ def get_node_heights(tree: Tree) -> dict[Tree, float]:
                 for child in node.children
             )
     return heights
+
+
+def get_mrca(node1: Tree, node2: Tree) -> Tree:
+    node1_ancestors = set(node1.iter_ancestors())
+    for node2_ancestor in node2.iter_ancestors():
+        if node2_ancestor in node1_ancestors:
+            return node2_ancestor
+    raise ValueError(f"No common ancestor found between node {node1} and node {node2}.")
+
+
+def get_distance(node1: Tree, node2: Tree) -> float:
+    mrca = get_mrca(node1, node2)
+    path = list(chain(node1.iter_ancestors(stop=mrca), node2.iter_ancestors(stop=mrca)))
+    if any(node.branch_length is None for node in path):
+        return len(path)
+    return sum(node.branch_length for node in path)  # pyright: ignore
+
+
+def compute_sackin_index(tree: Tree, normalize: bool = False) -> float:
+    """
+    Compute the Sackin index of a tree.
+
+    Parameters
+    ----------
+    tree : Tree
+        The input tree.
+    normalize : bool, optional
+        Whether to normalize the index between 0 and 1, by default False.
+
+    Returns
+    -------
+    float
+        The Sackin index of the tree.
+
+    References
+    ----------
+    - Kwang-Tsao Shao and Robert R Sokal. Tree Balance. Systematic Zoology, 39(3):266, 1990.
+    """
+    leaves = tree.get_leaves()
+    depth_levels = get_node_depth_levels(tree)
+    sackin_index = sum(depth_levels[leaf] for leaf in leaves)
+    if normalize:
+        n = len(leaves)
+        max_sackin_index = (n + 2) * (n - 1) / 2
+        return (sackin_index - n) / (max_sackin_index - n)
+    return sackin_index
+
+
+def compute_colless_index(tree: Tree, normalize: bool = False) -> float:
+    """
+    Compute the Colless index of a binary tree.
+
+    Parameters
+    ----------
+    tree : Tree
+        The input binary tree.
+    normalize : bool, optional
+        Whether to normalize the index between 0 and 1, by default False.
+
+    Returns
+    -------
+    float
+        The Colless index of the tree.
+
+    References
+    ----------
+    - Donald H. Colless. Review of phylogenetics: the theory and practice of phylogenetic systematics. Systematic Zoology, 31(1):100–104, 1982.
+    """
+    if not tree.is_binary():
+        raise ValueError("Colless index is only defined for binary trees.")
+
+    internal_nodes = tree.get_internal_nodes()
+    if not internal_nodes:
+        raise ValueError(
+            "Tree must have at least one internal node to compute the Colless index."
+        )
+
+    colless_index = 0
+    leaf_counts = get_node_leaf_counts(tree)
+    for node in internal_nodes:
+        left, right = node.children
+        colless_index += abs(leaf_counts[left] - leaf_counts[right])
+    if normalize:
+        n_leaves = len(leaf_counts) - len(internal_nodes)
+        max_colless_index = comb(n_leaves, 2)
+        return colless_index / max_colless_index
+    return colless_index
+
+
+def compute_mean_leaf_pairwise_distance(tree: Tree) -> float:
+    """
+    Compute the mean pairwise distance between all pairs of leaves in the tree.
+
+    Parameters
+    ----------
+    tree : Tree
+        The input tree.
+
+    Returns
+    -------
+    float
+        The mean pairwise distance between all pairs of leaves in the tree.
+    """
+    leaves = tree.get_leaves()
+    n_leaves = len(leaves)
+    if n_leaves < 2:
+        return 0.0
+
+    total_distance = sum(
+        get_distance(leaves[i], leaves[j])
+        for i in range(n_leaves)
+        for j in range(i + 1, n_leaves)
+    )
+    return total_distance / comb(n_leaves, 2)

{phylogenie-2.1.19.dist-info → phylogenie-2.1.21.dist-info}/METADATA

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

{phylogenie-2.1.19.dist-info → phylogenie-2.1.21.dist-info}/RECORD

@@ -1,4 +1,5 @@
-phylogenie/__init__.py,sha256=FQE_BakIftKaCnMaCZ4qpoSUH0JJPh2igAZLl_5pA28,2296
+phylogenie/__init__.py,sha256=8ulA-U7-WnBLNsqRYqPOPrJG8X4ZzyEU02oaDbR_Hxs,2849
+phylogenie/draw.py,sha256=mXCBS5xtJh0zJza8K2QkmRg_YbeCINst0vh6YMv6HQg,4959
 phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
 phylogenie/generators/alisim.py,sha256=0aCLuGInifWgAvfh7zARWSKF4EMw3TjlPXMLSECui0k,2783
 phylogenie/generators/configs.py,sha256=WFoeKpgj9ZQIom7BKqwpgXbriiQGg3jFBMLoD8KButk,1073
@@ -10,13 +11,12 @@ phylogenie/io.py,sha256=nwy8DOknt0HqF9qMeFZHrCmSXpM5AGrU5oajwTtD6vY,3973
 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
-phylogenie/plot.py,sha256=5xwPbZWO0jX_spEVPH7UUnxc7evqQbhSPKHWpXHNXro,5121
 phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
 phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
 phylogenie/skyline/parameter.py,sha256=EM9qlPt0JhMBy3TbztM0dj24BaGNEy8KWKdTObDKhbI,4644
 phylogenie/skyline/vector.py,sha256=bJP7_FNX_Klt6wXqsyfj0KX3VNj6-dIhzCKSJuQcOV0,7115
-phylogenie/tree.py,sha256=qaHMEcnqEtjd8Jlo2wF1vw0Gi6tcMlXngK_6IKEaQkw,4568
+phylogenie/tree.py,sha256=mv8Jrb2jzK5q-vx5E1-mAQLagkP4aHgaL6tNpAOPfIE,5194
 phylogenie/treesimulator/__init__.py,sha256=yqS2vtYMhdWSXc9RAnX1dd4zAqSQweMLyVKTnJLfGTU,1106
 phylogenie/treesimulator/events/__init__.py,sha256=6zSgZ0MEUMvTK4yPlSolJnRWzCARLS-jYreTzh45mQo,1033
 phylogenie/treesimulator/events/contact_tracing.py,sha256=_nJ85yhgGkeruQgMHvGpDYoyhheBf8M4LgZWiWdi5dY,4801
@@ -27,9 +27,9 @@ phylogenie/treesimulator/gillespie.py,sha256=LZHB2Ko147E78LoUCtN_BN7NYO1xhMYRy5P
 phylogenie/treesimulator/model.py,sha256=Nyg6R8XmMwZMSw1-dII81sU9uU7tDe-NMs8v1qKE4_M,5746
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=GknvAFXyiaWeeYJ8Lk5d6E2VHT-xW6ONEojYbtJYiB8,476
-phylogenie/utils.py,sha256=faDicpAtjxM-QDwPKOv8AOxPe7skPUOGKmVWBZRzUbU,1990
-phylogenie-2.1.19.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.19.dist-info/METADATA,sha256=g2cXt1gj4UITtGZgm21eKyOdyKHZYULHDuIpU32Gi9o,5477
-phylogenie-2.1.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.19.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.19.dist-info/RECORD,,
+phylogenie/utils.py,sha256=ehVk_2kvjW8Q_EyM2kxBPHYiK-KlPmZQx7JeVN6Fh-E,5419
+phylogenie-2.1.21.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.21.dist-info/METADATA,sha256=De8Vqs1ngZOKgBdivLw1hWrtMDdtFWLRDW3M4-IAAyg,5477
+phylogenie-2.1.21.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.21.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.21.dist-info/RECORD,,

phylogenie/plot.py

@@ -1,136 +0,0 @@
-from enum import Enum
-from typing import Any
-
-import matplotlib.colors as mcolors
-import matplotlib.patches as mpatches
-import matplotlib.pyplot as plt
-from matplotlib.axes import Axes
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # pyright: ignore
-
-from phylogenie.tree import Tree
-from phylogenie.utils import get_node_depth_levels, get_node_depths
-
-
-class Coloring(str, Enum):
-    DISCRETE = "discrete"
-    CONTINUOUS = "continuous"
-
-
-def plot_tree(
-    tree: Tree,
-    ax: Axes | None = None,
-    color_by: str | None = None,
-    default_color: str = "black",
-    coloring: str | Coloring | None = None,
-    cmap: str | None = None,
-    show_legend: bool = True,
-    show_hist: bool = True,
-    hist_kwargs: dict[str, Any] | None = None,
-    hist_axes_kwargs: dict[str, Any] | None = None,
-) -> Axes | tuple[Axes, Axes]:
-    if ax is None:
-        ax = plt.gca()
-
-    xs = (
-        get_node_depth_levels(tree)
-        if any(node.branch_length is None for node in tree)
-        else get_node_depths(tree)
-    )
-    ys = {node: i for i, node in enumerate(tree.inorder_traversal())}
-
-    if color_by is not None:
-        features = [node.get(color_by) for node in tree if color_by in node.features]
-
-        if coloring is None and any(isinstance(f, float) for f in features):
-            coloring = Coloring.CONTINUOUS
-        elif coloring is None:
-            coloring = Coloring.DISCRETE
-
-        if coloring == Coloring.DISCRETE:
-            if any(isinstance(f, float) for f in features):
-                raise ValueError(
-                    "Discrete coloring selected but feature values are not all categorical."
-                )
-            colormap = plt.get_cmap("tab20" if cmap is None else cmap)
-            feature_colors = {
-                f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features))
-            }
-            colors = {
-                node: (
-                    feature_colors[node.get(color_by)]
-                    if color_by in node.features
-                    else default_color
-                )
-                for node in tree
-            }
-
-            if show_legend:
-                legend_handles = [
-                    mpatches.Patch(color=feature_colors[f], label=str(f))
-                    for f in feature_colors
-                ]
-                if any(color_by not in node.features for node in tree):
-                    legend_handles.append(
-                        mpatches.Patch(color=default_color, label="NA")
-                    )
-                ax.legend(handles=legend_handles, title=color_by)  # pyright: ignore
-
-        elif coloring == Coloring.CONTINUOUS:
-            norm = mcolors.Normalize(vmin=min(features), vmax=max(features))
-            colormap = plt.get_cmap("viridis" if cmap is None else cmap)
-            colors = {
-                node: (
-                    colormap(norm(float(node.get(color_by))))
-                    if color_by in node.features
-                    else default_color
-                )
-                for node in tree
-            }
-
-            if show_hist:
-                default_hist_axes_kwargs = {"width": "25%", "height": "25%"}
-                if hist_axes_kwargs is not None:
-                    default_hist_axes_kwargs.update(hist_axes_kwargs)
-                hist_ax = inset_axes(ax, **default_hist_axes_kwargs)  # pyright: ignore
-
-                hist_kwargs = {} if hist_kwargs is None else hist_kwargs
-                _, bins, patches = hist_ax.hist(  # pyright: ignore
-                    features, **hist_kwargs
-                )
-
-                for patch, b0, b1 in zip(  # pyright: ignore
-                    patches, bins[:-1], bins[1:]  # pyright: ignore
-                ):
-                    midpoint = (b0 + b1) / 2  # pyright: ignore
-                    patch.set_facecolor(colormap(norm(midpoint)))  # pyright: ignore
-            else:
-                sm = plt.cm.ScalarMappable(cmap=colormap, norm=norm)
-                ax.get_figure().colorbar(sm, ax=ax)  # pyright: ignore
-
-        else:
-            raise ValueError(
-                f"Unknown coloring method: {coloring}. Choices are {list(Coloring)}."
-            )
-    else:
-        colors = {node: default_color for node in tree}
-
-    for node in tree:
-        x1, y1 = xs[node], ys[node]
-        if node.parent is None:
-            ax.hlines(y=y1, xmin=0, xmax=x1, color=colors[node])  # pyright: ignore
-            continue
-        x0, y0 = xs[node.parent], ys[node.parent]
-        ax.vlines(x=x0, ymin=y0, ymax=y1, color=colors[node])  # pyright: ignore
-        ax.hlines(y=y1, xmin=x0, xmax=x1, color=colors[node])  # pyright: ignore
-
-    for node in tree:
-        x1, y1 = xs[node], ys[node]
-        if node.parent is None:
-            ax.hlines(y=y1, xmin=0, xmax=x1, color=colors[node])  # pyright: ignore
-            continue
-        x0, y0 = xs[node.parent], ys[node.parent]
-        ax.vlines(x=x0, ymin=y0, ymax=y1, color=colors[node])  # pyright: ignore
-        ax.hlines(y=y1, xmin=x0, xmax=x1, color=colors[node])  # pyright: ignore
-
-    ax.set_yticks([])  # pyright: ignore
-    return ax