phylogenie 2.1.13__tar.gz → 2.1.14__tar.gz

{phylogenie-2.1.13 → phylogenie-2.1.14}/PKG-INFO

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

{phylogenie-2.1.13 → phylogenie-2.1.14}/phylogenie/generators/alisim.py

@@ -9,6 +9,7 @@ from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import data, string
 from phylogenie.generators.trees import TreeDatasetGeneratorConfig
 from phylogenie.io import dump_newick
+from phylogenie.utils import get_node_depths
 
 MSAS_DIRNAME = "MSAs"
 TREES_DIRNAME = "trees"
@@ -71,8 +72,9 @@ class AliSimDatasetGenerator(DatasetGenerator):
                     "Tree simulation timed out, retrying with different parameters..."
                 )
 
+        times = get_node_depths(tree)
         for leaf in tree.get_leaves():
-            leaf.name += f"|{leaf.get_time()}"
+            leaf.name += f"|{times[leaf]}"
         dump_newick(tree, f"{tree_filename}.nwk")
 
         self._generate_one_from_tree(msa_filename, f"{tree_filename}.nwk", rng, d)

{phylogenie-2.1.13 → phylogenie-2.1.14}/phylogenie/plot.py

@@ -6,7 +6,7 @@ import matplotlib.pyplot as plt
 from matplotlib.axes import Axes
 
 from phylogenie.tree import Tree
-from phylogenie.utils import get_times
+from phylogenie.utils import get_node_depths
 
 
 class Coloring(str, Enum):
@@ -22,44 +22,69 @@ def plot_tree(
     coloring: str | Coloring | None = None,
     cmap: str | None = None,
     show_legend: bool = True,
+    x_feature: str | None = None,
 ) -> Axes:
     if ax is None:
         ax = plt.gca()
 
-    xs = get_times(tree)
+    xs = (
+        get_node_depths(tree)
+        if x_feature is None
+        else {node: node.get(x_feature) for node in tree}
+    )
     ys = {node: i for i, node in enumerate(tree.inorder_traversal())}
 
     if color_by is not None:
-        features = set(node.get(color_by) for node in tree)
+        features = set(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."
                 )
-            cmap = "tab20" if cmap is None else cmap
-            colormap = plt.get_cmap(cmap, len(features))
+            colormap = plt.get_cmap("tab20" if cmap is None else cmap)
             feature_colors = {
                 f: mcolors.to_hex(colormap(i)) for i, f in enumerate(features)
             }
-            colors = {node: feature_colors[node.get(color_by)] for node in tree}
-            legend_handles = [
-                mpatches.Patch(color=feature_colors[f], label=str(f)) for f in 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 features
+                ]
+                if any(color_by not in node.features for node in tree):
+                    legend_handles.append(
+                        mpatches.Patch(color=default_color, label="None")
+                    )
                 ax.legend(handles=legend_handles, title=color_by)  # pyright: ignore
-        elif coloring in {Coloring.CONTINUOUS}:
-            cmap = "viridis" if cmap is None else cmap
-            values = list(map(float, features))
-            norm = mcolors.Normalize(vmin=min(values), vmax=max(values))
-            colormap = plt.get_cmap(cmap)
-            colors = {node: colormap(norm(float(node.get(color_by)))) for node in tree}
+
+        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
+            }
 
             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)}."
@@ -68,11 +93,11 @@ def plot_tree(
         colors = {node: default_color for node in tree}
 
     for node in tree:
-        x1, y1 = xs[node.name], ys[node]
+        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.name], ys[node.parent]
+        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
 

{phylogenie-2.1.13 → phylogenie-2.1.14}/phylogenie/tree.py

@@ -23,12 +23,49 @@ class Tree:
         return self._features.copy()
 
     @property
-    def time_to_parent(self) -> float:
-        if self.parent is None and self.branch_length is None:
+    def depth(self) -> float:
+        if self.parent is None:
             return 0.0
         if self.branch_length is None:
             raise ValueError(f"Branch length of node {self.name} is not set.")
-        return self.branch_length
+        return self.parent.depth + self.branch_length
+
+    @property
+    def depth_level(self) -> int:
+        if self.parent is None:
+            return 0
+        return self.parent.depth_level + 1
+
+    @property
+    def height(self) -> float:
+        if self.is_leaf():
+            return 0.0
+        if any(child.branch_length is None for child in self.children):
+            raise ValueError(
+                f"Branch length of one or more children of node {self.name} is not set."
+            )
+        return max(
+            child.branch_length + child.height  # pyright: ignore
+            for child in self.children
+        )
+
+    @property
+    def height_level(self) -> int:
+        if self.is_leaf():
+            return 0
+        return 1 + max(child.height_level for child in self.children)
+
+    def set(self, key: str, value: Any) -> None:
+        self._features[key] = value
+
+    def update_features(self, features: dict[str, Any]) -> None:
+        self._features.update(features)
+
+    def get(self, key: str) -> Any:
+        return self._features[key]
+
+    def delete(self, key: str) -> None:
+        del self._features[key]
 
     def add_child(self, child: "Tree") -> "Tree":
         child._parent = self
@@ -75,23 +112,12 @@ class Tree:
         return not self.children
 
     def get_leaves(self) -> tuple["Tree", ...]:
-        return tuple(node for node in self if not node.children)
+        return tuple(node for node in self if node.is_leaf())
 
-    def get_time(self) -> float:
-        parent_time = 0 if self.parent is None else self.parent.get_time()
-        return self.time_to_parent + parent_time
-
-    def set(self, key: str, value: Any) -> None:
-        self._features[key] = value
-
-    def update_features(self, features: dict[str, Any]) -> None:
-        self._features.update(features)
-
-    def get(self, key: str) -> Any:
-        return self._features[key]
-
-    def delete(self, key: str) -> None:
-        del self._features[key]
+    def ladderize(self, feature: str) -> None:
+        self._children.sort(key=lambda x: x.get(feature))
+        for child in self.children:
+            child.ladderize(feature)
 
     def copy(self):
         new_tree = Tree(self.name, self.branch_length)
@@ -103,5 +129,8 @@ class Tree:
     def __iter__(self) -> Iterator["Tree"]:
         return self.preorder_traversal()
 
+    def __len__(self) -> int:
+        return sum(1 for _ in self)
+
     def __repr__(self) -> str:
         return f"TreeNode(name='{self.name}', branch_length={self.branch_length}, features={self.features})"

phylogenie-2.1.14/phylogenie/treesimulator/features.py

@@ -0,0 +1,49 @@
+from collections.abc import Iterable
+from enum import Enum
+
+from phylogenie.tree import Tree
+from phylogenie.treesimulator.events import get_mutation_id
+from phylogenie.treesimulator.model import get_node_state
+from phylogenie.utils import (
+    get_node_depth_levels,
+    get_node_depths,
+    get_node_height_levels,
+    get_node_heights,
+    get_node_leaf_counts,
+)
+
+
+def _get_states(tree: Tree) -> dict[Tree, str]:
+    return {node: get_node_state(node.name) for node in tree}
+
+
+def _get_mutations(tree: Tree) -> dict[Tree, int]:
+    return {node: get_mutation_id(node.name) for node in tree}
+
+
+class Feature(str, Enum):
+    DEPTH = "depth"
+    DEPTH_LEVEL = "depth_level"
+    HEIGHT = "height"
+    HEIGHT_LEVEL = "height_level"
+    MUTATION = "mutation"
+    N_LEAVES = "n_leaves"
+    STATE = "state"
+
+
+FEATURES_EXTRACTORS = {
+    Feature.DEPTH: get_node_depths,
+    Feature.DEPTH_LEVEL: get_node_depth_levels,
+    Feature.HEIGHT: get_node_heights,
+    Feature.HEIGHT_LEVEL: get_node_height_levels,
+    Feature.MUTATION: _get_mutations,
+    Feature.N_LEAVES: get_node_leaf_counts,
+    Feature.STATE: _get_states,
+}
+
+
+def set_features(tree: Tree, features: Iterable[Feature]) -> None:
+    for feature in features:
+        feature_maps = FEATURES_EXTRACTORS[feature](tree)
+        for node in tree:
+            node.set(feature.value, feature_maps[node])

{phylogenie-2.1.13 → phylogenie-2.1.14}/phylogenie/treesimulator/model.py

@@ -89,9 +89,7 @@ class Model:
     def _set_branch_length(self, node: Tree, time: float) -> None:
         if node.branch_length is not None:
             raise ValueError(f"Branch length of node {node.name} is already set.")
-        node.branch_length = (
-            time if node.parent is None else time - node.parent.get_time()
-        )
+        node.branch_length = time if node.parent is None else time - node.parent.depth
 
     def _stem(self, individual: Individual, time: float) -> None:
         self._set_branch_length(individual.node, time)

phylogenie-2.1.14/phylogenie/utils.py

@@ -0,0 +1,59 @@
+from phylogenie.tree import Tree
+
+
+def get_node_leaf_counts(tree: Tree) -> dict[Tree, int]:
+    n_leaves: dict[Tree, int] = {}
+    for node in tree.postorder_traversal():
+        n_leaves[node] = sum(n_leaves[child] for child in node.children) or 1
+    return n_leaves
+
+
+def get_node_depth_levels(tree: Tree) -> dict[Tree, int]:
+    depth_levels: dict[Tree, int] = {}
+    for node in tree:
+        if node.parent is None:
+            depth_levels[node] = 0
+        else:
+            depth_levels[node] = depth_levels[node.parent] + 1
+    return depth_levels
+
+
+def get_node_depths(tree: Tree) -> dict[Tree, float]:
+    depths: dict[Tree, float] = {}
+    for node in tree:
+        if node.parent is None:
+            depths[node] = 0
+        else:
+            if node.branch_length is None:
+                raise ValueError(f"Branch length of node {node.name} is not set.")
+            depths[node] = depths[node.parent] + node.branch_length
+    return depths
+
+
+def get_node_height_levels(tree: Tree) -> dict[Tree, int]:
+    height_levels: dict[Tree, int] = {}
+    for node in tree.postorder_traversal():
+        if node.is_leaf():
+            height_levels[node] = 0
+        else:
+            height_levels[node] = max(
+                1 + height_levels[child] for child in node.children
+            )
+    return height_levels
+
+
+def get_node_heights(tree: Tree) -> dict[Tree, float]:
+    heights: dict[Tree, float] = {}
+    for node in tree.postorder_traversal():
+        if node.is_leaf():
+            heights[node] = 0
+        else:
+            if any(child.branch_length is None for child in node.children):
+                raise ValueError(
+                    f"Branch length of one or more children of node {node.name} is not set."
+                )
+            heights[node] = max(
+                child.branch_length + heights[child]  # pyright: ignore
+                for child in node.children
+            )
+    return heights

{phylogenie-2.1.13 → phylogenie-2.1.14}/pyproject.toml

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

phylogenie-2.1.13/phylogenie/treesimulator/features.py

@@ -1,39 +0,0 @@
-from collections.abc import Iterable
-from enum import Enum
-
-from phylogenie.tree import Tree
-from phylogenie.treesimulator.events import get_mutation_id
-from phylogenie.treesimulator.model import get_node_state
-from phylogenie.utils import get_heights, get_n_tips, get_times
-
-
-def _get_states(tree: Tree) -> dict[str, str]:
-    return {node.name: get_node_state(node.name) for node in tree}
-
-
-def _get_mutations(tree: Tree) -> dict[str, int]:
-    return {node.name: get_mutation_id(node.name) for node in tree}
-
-
-class Feature(str, Enum):
-    STATE = "state"
-    MUTATION = "mutation"
-    N_TIPS = "n_tips"
-    TIME = "time"
-    HEIGHT = "height"
-
-
-FEATURES_EXTRACTORS = {
-    Feature.STATE: _get_states,
-    Feature.MUTATION: _get_mutations,
-    Feature.N_TIPS: get_n_tips,
-    Feature.TIME: get_times,
-    Feature.HEIGHT: get_heights,
-}
-
-
-def set_features(tree: Tree, features: Iterable[Feature]) -> None:
-    for feature in features:
-        feature_maps = FEATURES_EXTRACTORS[feature](tree)
-        for node in tree:
-            node.set(feature.value, feature_maps[node.name])

phylogenie-2.1.13/phylogenie/utils.py

@@ -1,28 +0,0 @@
-from phylogenie.tree import Tree
-
-
-def get_n_tips(tree: Tree) -> dict[str, int]:
-    n_tips: dict[str, int] = {}
-    for node in tree.postorder_traversal():
-        n_tips[node.name] = (
-            1 if node.is_leaf() else sum(n_tips[child.name] for child in node.children)
-        )
-    return n_tips
-
-
-def get_times(tree: Tree) -> dict[str, float]:
-    times: dict[str, float] = {}
-    for node in tree:
-        parent_time = 0 if node.parent is None else times[node.parent.name]
-        times[node.name] = node.time_to_parent + parent_time
-    return times
-
-
-def get_heights(tree: Tree) -> dict[str, int]:
-    heights: dict[str, int] = {}
-    for node in tree.postorder_traversal():
-        if node.is_leaf():
-            heights[node.name] = 0
-        else:
-            heights[node.name] = 1 + max(heights[child.name] for child in node.children)
-    return heights