phylogenie 2.1.12__py3-none-any.whl → 2.1.14__py3-none-any.whl

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/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,45 +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)
-            fig = ax.get_figure()
-            fig.colorbar(sm, ax=ax, label=color_by)  # pyright: ignore
+            ax.get_figure().colorbar(sm, ax=ax)  # pyright: ignore
+
         else:
             raise ValueError(
                 f"Unknown coloring method: {coloring}. Choices are {list(Coloring)}."
@@ -69,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/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/treesimulator/features.py

@@ -4,31 +4,41 @@ 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
+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[str, str]:
-    return {node.name: get_node_state(node.name) for node in tree}
+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[str, int]:
-    return {node.name: get_mutation_id(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):
-    STATE = "state"
-    MUTATION = "mutation"
-    N_TIPS = "n_tips"
-    TIME = "time"
+    DEPTH = "depth"
+    DEPTH_LEVEL = "depth_level"
     HEIGHT = "height"
+    HEIGHT_LEVEL = "height_level"
+    MUTATION = "mutation"
+    N_LEAVES = "n_leaves"
+    STATE = "state"
 
 
 FEATURES_EXTRACTORS = {
-    Feature.STATE: _get_states,
+    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_TIPS: get_n_tips,
-    Feature.TIME: get_times,
-    Feature.HEIGHT: get_heights,
+    Feature.N_LEAVES: get_node_leaf_counts,
+    Feature.STATE: _get_states,
 }
 
 
@@ -36,4 +46,4 @@ 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])
+            node.set(feature.value, feature_maps[node])

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/utils.py

@@ -1,28 +1,59 @@
 from phylogenie.tree import Tree
 
 
-def get_n_tips(tree: Tree) -> dict[str, int]:
-    n_tips: dict[str, int] = {}
+def get_node_leaf_counts(tree: Tree) -> dict[Tree, int]:
+    n_leaves: dict[Tree, 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
+        n_leaves[node] = sum(n_leaves[child] for child in node.children) or 1
+    return n_leaves
 
 
-def get_times(tree: Tree) -> dict[str, float]:
-    times: dict[str, float] = {}
+def get_node_depth_levels(tree: Tree) -> dict[Tree, int]:
+    depth_levels: dict[Tree, int] = {}
     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
+        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_heights(tree: Tree) -> dict[str, int]:
-    heights: dict[str, int] = {}
+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.name] = 0
+            heights[node] = 0
         else:
-            heights[node.name] = 1 + max(heights[child.name] for child in node.children)
+            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.12.dist-info → phylogenie-2.1.14.dist-info}/METADATA

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

{phylogenie-2.1.12.dist-info → phylogenie-2.1.14.dist-info}/RECORD

@@ -1,6 +1,6 @@
 phylogenie/__init__.py,sha256=FQE_BakIftKaCnMaCZ4qpoSUH0JJPh2igAZLl_5pA28,2296
 phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
-phylogenie/generators/alisim.py,sha256=G7p6tfcDWncg4xl3NhEXxgikcZDUi_RdE33zWn_CvzY,2704
+phylogenie/generators/alisim.py,sha256=0aCLuGInifWgAvfh7zARWSKF4EMw3TjlPXMLSECui0k,2783
 phylogenie/generators/configs.py,sha256=WFoeKpgj9ZQIom7BKqwpgXbriiQGg3jFBMLoD8KButk,1073
 phylogenie/generators/dataset.py,sha256=pPwW9yxm9fkU0PPllFq8EsPlqau8tth-4OatbA_hEHo,2120
 phylogenie/generators/factories.py,sha256=TuVFQWRjq33Hewjw_Lp8tQ0l_IPtqYDyQCNJhtiHpw8,7882
@@ -10,26 +10,26 @@ 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=hjftVhapi9AoD2_e200E82rVtYHY5dNptvyMKYcbnQI,2999
+phylogenie/plot.py,sha256=5InWceLwjZV2qVYv2jaTJPs_0gfAEaKYGQt_jnLpaz0,3615
 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=bWQqbr8CbWZZHKQr-iav4gU2sLbnGRXuBk5_ZJZEZBU,3345
+phylogenie/tree.py,sha256=PjyfsjYWadvIQSOxyIO0oaRVN2pecXEnUR-k1mM4Vic,4159
 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
 phylogenie/treesimulator/events/core.py,sha256=RF7oHzAjkU675PnczaVc66d9gNrHBL-IhmVHtcy7MKE,7949
 phylogenie/treesimulator/events/mutations.py,sha256=erEvgfiv_X3G_DwK9Hqu-fAR8otupfwq66cp5tRZamM,3591
-phylogenie/treesimulator/features.py,sha256=f1t7DE_dw8MhD-FAzQr0j8pxHdvLENy5AQeAjJyOjOQ,1082
+phylogenie/treesimulator/features.py,sha256=Wj1rjbOAHHE9XqhrLND2GMgqTB_4M9MdAD5OWoUoatc,1356
 phylogenie/treesimulator/gillespie.py,sha256=LZHB2Ko147E78LoUCtN_BN7NYO1xhMYRy5PUZbN93c0,5283
-phylogenie/treesimulator/model.py,sha256=Ct0lfn6maKtjuFxivWx1MbFHvH3Y-fiJ0XXMdkN3Cak,5775
+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=tVorfhMuRkxogeNu45aBv4SDBflHg4faiIpnHNJPKPY,871
-phylogenie-2.1.12.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.12.dist-info/METADATA,sha256=65IaQc4XK0qJH_CwrqkdHtVpU2u8mxQLZxuzpDIljK8,5477
-phylogenie-2.1.12.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.12.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.12.dist-info/RECORD,,
+phylogenie/utils.py,sha256=l7kJ0mkdhplz_hTUseBxaGpEKngsWUOSiCMyLMIWYZQ,1936
+phylogenie-2.1.14.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.14.dist-info/METADATA,sha256=elXtIXUWfGalXgdhTgretrdYVQwOWwiSfHUiZicf5WQ,5477
+phylogenie-2.1.14.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.14.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.14.dist-info/RECORD,,