iplotx 0.2.1__py3-none-any.whl → 0.3.1__py3-none-any.whl

iplotx/ingest/providers/tree/cogent3.py

@@ -1,112 +1,41 @@
 from typing import (
+    Any,
     Optional,
     Sequence,
 )
-from collections.abc import Hashable
-from operator import attrgetter
-import numpy as np
-import pandas as pd
-
-from ....typing import (
-    TreeType,
-    LayoutType,
-)
 from ...typing import (
     TreeDataProvider,
-    TreeData,
-)
-from ...heuristics import (
-    normalise_tree_layout,
 )
 
 
 class Cogent3DataProvider(TreeDataProvider):
-    def __call__(
-        self,
-        tree: TreeType,
-        layout: str | LayoutType,
-        orientation: str = "horizontal",
-        directed: bool | str = False,
-        vertex_labels: Optional[
-            Sequence[str] | dict[Hashable, str] | pd.Series | bool
-        ] = None,
-        edge_labels: Optional[Sequence[str] | dict] = None,
-    ) -> TreeData:
-        """Create tree data object for iplotx from cogent3.core.tree.PhyloNode classes."""
-
-        root_fun = lambda tree: tree.root()
-        preorder_fun = lambda tree: tree.preorder()
-        postorder_fun = lambda tree: tree.postorder()
-        children_fun = attrgetter("children")
-        branch_length_fun = attrgetter("length")
-        leaves_fun = lambda tree: tree.tips()
-
-        tree_data = {
-            "root": root_fun(tree),
-            "leaves": leaves_fun(tree),
-            "rooted": True,
-            "directed": directed,
-            "ndim": 2,
-            "layout_name": layout,
-        }
+    def preorder(self) -> Sequence[Any]:
+        return self.tree.preorder()
 
-        # Add vertex_df including layout
-        tree_data["vertex_df"] = normalise_tree_layout(
-            layout,
-            tree=tree,
-            orientation=orientation,
-            root_fun=root_fun,
-            preorder_fun=preorder_fun,
-            postorder_fun=postorder_fun,
-            children_fun=children_fun,
-            branch_length_fun=branch_length_fun,
-        )
-        if layout in ("radial",):
-            tree_data["layout_coordinate_system"] = "polar"
-        else:
-            tree_data["layout_coordinate_system"] = "cartesian"
+    def postorder(self) -> Sequence[Any]:
+        return self.tree.postorder()
 
-        # Add edge_df
-        edge_data = {"_ipx_source": [], "_ipx_target": []}
-        for node in preorder_fun(tree):
-            for child in node.children:
-                if directed == "parent":
-                    edge_data["_ipx_source"].append(child)
-                    edge_data["_ipx_target"].append(node)
-                else:
-                    edge_data["_ipx_source"].append(node)
-                    edge_data["_ipx_target"].append(child)
-        edge_df = pd.DataFrame(edge_data)
-        tree_data["edge_df"] = edge_df
+    def get_leaves(self) -> Sequence[Any]:
+        return self.tree.tips()
 
-        # Add vertex labels
-        if vertex_labels is None:
-            vertex_labels = False
-        if np.isscalar(vertex_labels) and vertex_labels:
-            tree_data["vertex_df"]["label"] = [
-                x.name for x in tree_data["vertices"].index
-            ]
-        elif not np.isscalar(vertex_labels):
-            # If a dict-like object is passed, it can be incomplete (e.g. only the leaves):
-            # we fill the rest with empty strings which are not going to show up in the plot.
-            if isinstance(vertex_labels, pd.Series):
-                vertex_labels = dict(vertex_labels)
-            if isinstance(vertex_labels, dict):
-                for vertex in tree_data["vertex_df"].index:
-                    if vertex not in vertex_labels:
-                        vertex_labels[vertex] = ""
-            tree_data["vertex_df"]["label"] = pd.Series(vertex_labels)
+    @staticmethod
+    def get_children(node: Any) -> Sequence[Any]:
+        return node.children
 
-        return tree_data
+    @staticmethod
+    def get_branch_length(node: Any) -> Optional[float]:
+        return node.length
 
-    def check_dependencies(self) -> bool:
+    @staticmethod
+    def check_dependencies() -> bool:
         try:
             import cogent3
         except ImportError:
             return False
         return True
 
-    def tree_type(self):
+    @staticmethod
+    def tree_type():
         from cogent3.core.tree import PhyloNode
 
         return PhyloNode

iplotx/ingest/providers/tree/ete4.py

@@ -1,112 +1,44 @@
 from typing import (
+    Any,
     Optional,
     Sequence,
 )
-from collections.abc import Hashable
-from operator import attrgetter
-import numpy as np
-import pandas as pd
+from functools import partialmethod
 
-from ....typing import (
-    TreeType,
-    LayoutType,
-)
 from ...typing import (
     TreeDataProvider,
-    TreeData,
-)
-from ...heuristics import (
-    normalise_tree_layout,
 )
 
 
 class Ete4DataProvider(TreeDataProvider):
-    def __call__(
-        self,
-        tree: TreeType,
-        layout: str | LayoutType,
-        orientation: str = "horizontal",
-        directed: bool | str = False,
-        vertex_labels: Optional[
-            Sequence[str] | dict[Hashable, str] | pd.Series | bool
-        ] = None,
-        edge_labels: Optional[Sequence[str] | dict] = None,
-    ) -> TreeData:
-        """Create tree data object for iplotx from ete4.core.tre.Tree classes."""
-
-        root_fun = attrgetter("root")
-        preorder_fun = lambda tree: tree.traverse("preorder")
-        postorder_fun = lambda tree: tree.traverse("postorder")
-        children_fun = attrgetter("children")
-        branch_length_fun = lambda node: node.dist if node.dist is not None else 1.0
-        leaves_fun = lambda tree: tree.leaves()
-
-        tree_data = {
-            "root": tree.root,
-            "leaves": leaves_fun(tree),
-            "rooted": True,
-            "directed": directed,
-            "ndim": 2,
-            "layout_name": layout,
-        }
+    def _traverse(self, order: str) -> Any:
+        """Traverse the tree."""
+        return self.tree.traverse(order)
 
-        # Add vertex_df including layout
-        tree_data["vertex_df"] = normalise_tree_layout(
-            layout,
-            tree=tree,
-            orientation=orientation,
-            root_fun=root_fun,
-            preorder_fun=preorder_fun,
-            postorder_fun=postorder_fun,
-            children_fun=children_fun,
-            branch_length_fun=branch_length_fun,
-        )
-        if layout in ("radial",):
-            tree_data["layout_coordinate_system"] = "polar"
-        else:
-            tree_data["layout_coordinate_system"] = "cartesian"
+    preorder = partialmethod(_traverse, order="preorder")
+    postorder = partialmethod(_traverse, order="postorder")
 
-        # Add edge_df
-        edge_data = {"_ipx_source": [], "_ipx_target": []}
-        for node in preorder_fun(tree):
-            for child in children_fun(node):
-                if directed == "parent":
-                    edge_data["_ipx_source"].append(child)
-                    edge_data["_ipx_target"].append(node)
-                else:
-                    edge_data["_ipx_source"].append(node)
-                    edge_data["_ipx_target"].append(child)
-        edge_df = pd.DataFrame(edge_data)
-        tree_data["edge_df"] = edge_df
+    def get_leaves(self) -> Sequence[Any]:
+        return self.tree.leaves()
 
-        # Add vertex labels
-        if vertex_labels is None:
-            vertex_labels = False
-        if np.isscalar(vertex_labels) and vertex_labels:
-            tree_data["vertex_df"]["label"] = [
-                x.name for x in tree_data["vertices"].index
-            ]
-        elif not np.isscalar(vertex_labels):
-            # If a dict-like object is passed, it can be incomplete (e.g. only the leaves):
-            # we fill the rest with empty strings which are not going to show up in the plot.
-            if isinstance(vertex_labels, pd.Series):
-                vertex_labels = dict(vertex_labels)
-            if isinstance(vertex_labels, dict):
-                for vertex in tree_data["vertex_df"].index:
-                    if vertex not in vertex_labels:
-                        vertex_labels[vertex] = ""
-            tree_data["vertex_df"]["label"] = pd.Series(vertex_labels)
+    @staticmethod
+    def get_children(node: Any) -> Sequence[Any]:
+        return node.children
 
-        return tree_data
+    @staticmethod
+    def get_branch_length(node: Any) -> Optional[float]:
+        return node.dist
 
-    def check_dependencies(self) -> bool:
+    @staticmethod
+    def check_dependencies() -> bool:
         try:
             from ete4 import Tree
         except ImportError:
             return False
         return True
 
-    def tree_type(self):
+    @staticmethod
+    def tree_type():
         from ete4 import Tree
 
         return Tree

iplotx/ingest/providers/tree/skbio.py

@@ -1,112 +1,41 @@
 from typing import (
+    Any,
     Optional,
     Sequence,
 )
-from collections.abc import Hashable
-from operator import attrgetter
-import numpy as np
-import pandas as pd
-
-from ....typing import (
-    TreeType,
-    LayoutType,
-)
 from ...typing import (
     TreeDataProvider,
-    TreeData,
-)
-from ...heuristics import (
-    normalise_tree_layout,
 )
 
 
 class SkbioDataProvider(TreeDataProvider):
-    def __call__(
-        self,
-        tree: TreeType,
-        layout: str | LayoutType,
-        orientation: str = "horizontal",
-        directed: bool | str = False,
-        vertex_labels: Optional[
-            Sequence[str] | dict[Hashable, str] | pd.Series | bool
-        ] = None,
-        edge_labels: Optional[Sequence[str] | dict] = None,
-    ) -> TreeData:
-        """Create tree data object for iplotx from skbio.tree.TreeNode classes."""
-
-        root_fun = lambda tree: tree.root()
-        preorder_fun = lambda tree: tree.preorder()
-        postorder_fun = lambda tree: tree.postorder()
-        children_fun = attrgetter("children")
-        branch_length_fun = attrgetter("length")
-        leaves_fun = lambda tree: tree.tips()
-
-        tree_data = {
-            "root": root_fun(tree),
-            "leaves": leaves_fun(tree),
-            "rooted": True,
-            "directed": directed,
-            "ndim": 2,
-            "layout_name": layout,
-        }
+    def preorder(self) -> Sequence[Any]:
+        return self.tree.preorder()
 
-        # Add vertex_df including layout
-        tree_data["vertex_df"] = normalise_tree_layout(
-            layout,
-            tree=tree,
-            orientation=orientation,
-            root_fun=root_fun,
-            preorder_fun=preorder_fun,
-            postorder_fun=postorder_fun,
-            children_fun=children_fun,
-            branch_length_fun=branch_length_fun,
-        )
-        if layout in ("radial",):
-            tree_data["layout_coordinate_system"] = "polar"
-        else:
-            tree_data["layout_coordinate_system"] = "cartesian"
+    def postorder(self) -> Sequence[Any]:
+        return self.tree.postorder()
 
-        # Add edge_df
-        edge_data = {"_ipx_source": [], "_ipx_target": []}
-        for node in preorder_fun(tree):
-            for child in children_fun(node):
-                if directed == "parent":
-                    edge_data["_ipx_source"].append(child)
-                    edge_data["_ipx_target"].append(node)
-                else:
-                    edge_data["_ipx_source"].append(node)
-                    edge_data["_ipx_target"].append(child)
-        edge_df = pd.DataFrame(edge_data)
-        tree_data["edge_df"] = edge_df
+    def get_leaves(self) -> Sequence[Any]:
+        return self.tree.tips()
 
-        # Add vertex labels
-        if vertex_labels is None:
-            vertex_labels = False
-        if np.isscalar(vertex_labels) and vertex_labels:
-            tree_data["vertex_df"]["label"] = [
-                x.name for x in tree_data["vertices"].index
-            ]
-        elif not np.isscalar(vertex_labels):
-            # If a dict-like object is passed, it can be incomplete (e.g. only the leaves):
-            # we fill the rest with empty strings which are not going to show up in the plot.
-            if isinstance(vertex_labels, pd.Series):
-                vertex_labels = dict(vertex_labels)
-            if isinstance(vertex_labels, dict):
-                for vertex in tree_data["vertex_df"].index:
-                    if vertex not in vertex_labels:
-                        vertex_labels[vertex] = ""
-            tree_data["vertex_df"]["label"] = pd.Series(vertex_labels)
+    @staticmethod
+    def get_children(node: Any) -> Sequence[Any]:
+        return node.children
 
-        return tree_data
+    @staticmethod
+    def get_branch_length(node: Any) -> Optional[float]:
+        return node.length
 
-    def check_dependencies(self) -> bool:
+    @staticmethod
+    def check_dependencies() -> bool:
         try:
             from skbio import TreeNode
         except ImportError:
             return False
         return True
 
-    def tree_type(self):
+    @staticmethod
+    def tree_type():
         from skbio import TreeNode
 
         return TreeNode

iplotx/ingest/typing.py

@@ -11,14 +11,20 @@ from typing import (
     Protocol,
     Optional,
     Sequence,
+    Any,
+    Iterable,
 )
 from collections.abc import Hashable
+import numpy as np
 import pandas as pd
 from ..typing import (
     GraphType,
     LayoutType,
     TreeType,
 )
+from .heuristics import (
+    normalise_tree_layout,
+)
 
 
 class NetworkData(TypedDict):
@@ -44,15 +50,13 @@ class NetworkDataProvider(Protocol):
         """Create network data object for iplotx from any provider."""
         raise NotImplementedError("Network data providers must implement this method.")
 
-    def check_dependencies(
-        self,
-    ):
+    @staticmethod
+    def check_dependencies():
         """Check whether the dependencies for this provider are installed."""
         raise NotImplementedError("Network data providers must implement this method.")
 
-    def graph_type(
-        self,
-    ):
+    @staticmethod
+    def graph_type():
         """Return the graph type from this provider to check for instances."""
         raise NotImplementedError("Network data providers must implement this method.")
 
@@ -63,11 +67,12 @@ class TreeData(TypedDict):
     rooted: bool
     directed: bool | str
     root: Optional[Hashable]
-    leaves: list[Hashable]
+    leaf_df: pd.DataFrame
     vertex_df: dict[Hashable, tuple[float, float]]
     edge_df: dict[Hashable, Sequence[tuple[float, float]]]
     layout_coordinate_system: str
     layout_name: str
+    orientation: str
     ndim: int
     tree_library: NotRequired[str]
 
@@ -75,26 +80,224 @@ class TreeData(TypedDict):
 class TreeDataProvider(Protocol):
     """Protocol for tree data ingestion provider for iplotx."""
 
-    def __call__(
+    def __init__(
         self,
         tree: TreeType,
-        layout: str | LayoutType,
-        orientation: Optional[str] = None,
-        directed: bool | str = False,
-        vertex_labels: Optional[Sequence[str] | dict[Hashable, str] | pd.Series] = None,
-        edge_labels: Optional[Sequence[str] | dict] = None,
-    ) -> TreeData:
-        """Create tree data object for iplotx from any provider."""
-        raise NotImplementedError("Tree data providers must implement this method.")
+    ) -> None:
+        """Initialize the provider with the tree type.
 
-    def check_dependencies(
-        self,
-    ):
+        Parameters:
+            tree: The tree type that this provider will handle.
+        """
+        self.tree = tree
+
+    @staticmethod
+    def check_dependencies():
         """Check whether the dependencies for this provider are installed."""
         raise NotImplementedError("Tree data providers must implement this method.")
 
-    def tree_type(
-        self,
-    ):
+    @staticmethod
+    def tree_type():
         """Return the tree type from this provider to check for instances."""
         raise NotImplementedError("Tree data providers must implement this method.")
+
+    def is_rooted(self) -> bool:
+        """Get whether the tree is rooted.
+
+        Returns:
+            A boolean indicating whether the tree is rooted.
+
+        Note: This is a default implemntation that can be overridden by the provider
+        if they support unrooted trees (e.g. Biopython).
+        """
+        return True
+
+    def get_root(self) -> Any:
+        """Get the tree root in a provider-specific data structure.
+
+        Returns:
+            The root of the tree.
+
+        Note: This is a default implemntation that can be overridden by the provider.
+        """
+        root_attr = self.tree.root
+        if callable(root_attr):
+            return root_attr()
+        else:
+            return root_attr
+
+    def get_leaves(self) -> Sequence[Any]:
+        """Get the tree leaves/tips in a provider-specific data structure.
+
+        Returns:
+            The leaves or tips of the tree.
+        """
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    def preorder(self) -> Iterable[Any]:
+        """Preorder (DFS - parent first) iteration over the tree.
+
+        Returns:
+            An iterable of nodes in preorder traversal.
+        """
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    def postorder(self) -> Iterable[Any]:
+        """Postorder (DFS - child first) iteration over the tree.
+
+        Returns:
+            An iterable of nodes in preorder traversal.
+        """
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    @staticmethod
+    def get_children(
+        node: Any,
+    ) -> Sequence[Any]:
+        """Get the children of a node.
+
+        Parameters:
+            node: The node to get the children from.
+        Returns:
+            A sequence of children nodes.
+        """
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    @staticmethod
+    def get_branch_length(
+        node: Any,
+    ) -> Optional[float]:
+        """Get the length of the branch to this node.
+
+        Parameters:
+            node: The node to get the branch length from.
+        Returns:
+            The branch length to the node.
+        """
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    def get_branch_length_default_to_one(
+        self,
+        node: Any,
+    ) -> float:
+        """Get the length of the branch to this node, defaulting to 1.0 if not available.
+
+        Parameters:
+            node: The node to get the branch length from.
+        Returns:
+            The branch length to the node, defaulting to 1.0 if not available.
+        """
+        branch_length = self.get_branch_length(node)
+        return branch_length if branch_length is not None else 1.0
+
+    def __call__(
+        self,
+        layout: str | LayoutType,
+        orientation: Optional[str],
+        layout_style: Optional[dict[str, int | float | str]] = None,
+        directed: bool | str = False,
+        vertex_labels: Optional[
+            Sequence[str] | dict[Hashable, str] | pd.Series | bool
+        ] = None,
+        edge_labels: Optional[Sequence[str] | dict] = None,
+        leaf_labels: Optional[Sequence[str] | dict[Hashable, str] | pd.Series] = None,
+    ) -> TreeData:
+        """Create tree data object for iplotx from ete4.core.tre.Tree classes."""
+
+        if layout_style is None:
+            layout_style = {}
+
+        if orientation is None:
+            if layout == "horizontal":
+                orientation = "right"
+            elif layout == "vertical":
+                orientation = "descending"
+            elif layout == "radial":
+                orientation = "clockwise"
+
+        tree_data = {
+            "root": self.get_root(),
+            "rooted": self.is_rooted(),
+            "directed": directed,
+            "ndim": 2,
+            "layout_name": layout,
+            "orientation": orientation,
+        }
+
+        # Add vertex_df including layout
+        tree_data["vertex_df"] = normalise_tree_layout(
+            layout,
+            orientation=orientation,
+            root=tree_data["root"],
+            preorder_fun=self.preorder,
+            postorder_fun=self.postorder,
+            children_fun=self.get_children,
+            branch_length_fun=self.get_branch_length_default_to_one,
+            **layout_style,
+        )
+        if layout in ("radial",):
+            tree_data["layout_coordinate_system"] = "polar"
+        else:
+            tree_data["layout_coordinate_system"] = "cartesian"
+
+        # Add edge_df
+        edge_data = {"_ipx_source": [], "_ipx_target": []}
+        for node in self.preorder():
+            for child in self.get_children(node):
+                if directed == "parent":
+                    edge_data["_ipx_source"].append(child)
+                    edge_data["_ipx_target"].append(node)
+                else:
+                    edge_data["_ipx_source"].append(node)
+                    edge_data["_ipx_target"].append(child)
+        edge_df = pd.DataFrame(edge_data)
+        tree_data["edge_df"] = edge_df
+
+        # Add leaf_df
+        tree_data["leaf_df"] = pd.DataFrame(index=self.get_leaves())
+
+        # Add vertex labels
+        if vertex_labels is None:
+            vertex_labels = False
+        if np.isscalar(vertex_labels) and vertex_labels:
+            tree_data["vertex_df"]["label"] = [
+                x.name for x in tree_data["vertex_df"].index
+            ]
+        elif not np.isscalar(vertex_labels):
+            # If a dict-like object is passed, it can be incomplete (e.g. only the leaves):
+            # we fill the rest with empty strings which are not going to show up in the plot.
+            if isinstance(vertex_labels, pd.Series):
+                vertex_labels = dict(vertex_labels)
+            if isinstance(vertex_labels, dict):
+                for vertex in tree_data["vertex_df"].index:
+                    if vertex not in vertex_labels:
+                        vertex_labels[vertex] = ""
+            tree_data["vertex_df"]["label"] = pd.Series(vertex_labels)
+
+        # Add leaf labels
+        if leaf_labels is None:
+            leaf_labels = False
+        if np.isscalar(leaf_labels) and leaf_labels:
+            tree_data["leaf_labels"]["label"] = [
+                # FIXME: this is likely broken
+                x.name
+                for x in tree_data["leaf_df"].index
+            ]
+        elif not np.isscalar(leaf_labels):
+            # Leaves are already in the dataframe in a certain order, so sequences are allowed
+            if isinstance(leaf_labels, (list, tuple, np.ndarray)):
+                leaf_labels = {
+                    leaf: label
+                    for leaf, label in zip(tree_data["leaf_df"].index, leaf_labels)
+                }
+            # If a dict-like object is passed, it can be incomplete (e.g. only the leaves):
+            # we fill the rest with empty strings which are not going to show up in the plot.
+            if isinstance(leaf_labels, pd.Series):
+                leaf_labels = dict(leaf_labels)
+            if isinstance(leaf_labels, dict):
+                for leaf in tree_data["leaf_df"].index:
+                    if leaf not in leaf_labels:
+                        leaf_labels[leaf] = ""
+            tree_data["leaf_df"]["label"] = pd.Series(leaf_labels)
+
+        return tree_data