iplotx 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

iplotx/ingest/providers/tree/cogent3.py

@@ -0,0 +1,112 @@
+from typing import (
+    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,
+        }
+
+        # 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"
+
+        # 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
+
+        # 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)
+
+        return tree_data
+
+    def check_dependencies(self) -> bool:
+        try:
+            import cogent3
+        except ImportError:
+            return False
+        return True
+
+    def tree_type(self):
+        from cogent3.core.tree import PhyloNode
+
+        return PhyloNode

iplotx/ingest/providers/tree/ete4.py

@@ -0,0 +1,112 @@
+from typing import (
+    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 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,
+        }
+
+        # 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"
+
+        # 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
+
+        # 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)
+
+        return tree_data
+
+    def check_dependencies(self) -> bool:
+        try:
+            from ete4 import Tree
+        except ImportError:
+            return False
+        return True
+
+    def tree_type(self):
+        from ete4 import Tree
+
+        return Tree

iplotx/ingest/providers/tree/skbio.py

@@ -0,0 +1,112 @@
+from typing import (
+    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,
+        }
+
+        # 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"
+
+        # 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
+
+        # 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)
+
+        return tree_data
+
+    def check_dependencies(self) -> bool:
+        try:
+            from skbio import TreeNode
+        except ImportError:
+            return False
+        return True
+
+    def tree_type(self):
+        from skbio import TreeNode
+
+        return TreeNode

iplotx/ingest/typing.py

@@ -0,0 +1,100 @@
+"""
+Typing module for data/object ingestion. This module described the abstract data types that providers need to comply with to be compatible with iplotx.
+
+Networkx and trees are treated separately for practical reasons: many tree analysis libraries rely heavily on recursive data structures, which do not
+work as well on general networks.
+"""
+
+from typing import (
+    NotRequired,
+    TypedDict,
+    Protocol,
+    Optional,
+    Sequence,
+)
+from collections.abc import Hashable
+import pandas as pd
+from ..typing import (
+    GraphType,
+    LayoutType,
+    TreeType,
+)
+
+
+class NetworkData(TypedDict):
+    """Network data structure for iplotx."""
+
+    directed: bool
+    vertex_df: pd.DataFrame
+    edge_df: pd.DataFrame
+    ndim: int
+    network_library: NotRequired[str]
+
+
+class NetworkDataProvider(Protocol):
+    """Protocol for network data ingestion provider for iplotx."""
+
+    def __call__(
+        self,
+        network: GraphType,
+        layout: Optional[LayoutType] = None,
+        vertex_labels: Optional[Sequence[str] | dict[Hashable, str] | pd.Series] = None,
+        edge_labels: Optional[Sequence[str] | dict] = None,
+    ) -> NetworkData:
+        """Create network data object for iplotx from any provider."""
+        raise NotImplementedError("Network data providers must implement this method.")
+
+    def check_dependencies(
+        self,
+    ):
+        """Check whether the dependencies for this provider are installed."""
+        raise NotImplementedError("Network data providers must implement this method.")
+
+    def graph_type(
+        self,
+    ):
+        """Return the graph type from this provider to check for instances."""
+        raise NotImplementedError("Network data providers must implement this method.")
+
+
+class TreeData(TypedDict):
+    """Tree data structure for iplotx."""
+
+    rooted: bool
+    directed: bool | str
+    root: Optional[Hashable]
+    leaves: list[Hashable]
+    vertex_df: dict[Hashable, tuple[float, float]]
+    edge_df: dict[Hashable, Sequence[tuple[float, float]]]
+    layout_coordinate_system: str
+    layout_name: str
+    ndim: int
+    tree_library: NotRequired[str]
+
+
+class TreeDataProvider(Protocol):
+    """Protocol for tree data ingestion provider for iplotx."""
+
+    def __call__(
+        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.")
+
+    def check_dependencies(
+        self,
+    ):
+        """Check whether the dependencies for this provider are installed."""
+        raise NotImplementedError("Tree data providers must implement this method.")
+
+    def tree_type(
+        self,
+    ):
+        """Return the tree type from this provider to check for instances."""
+        raise NotImplementedError("Tree data providers must implement this method.")

iplotx/label.py

@@ -0,0 +1,127 @@
+from typing import (
+    Optional,
+    Sequence,
+)
+import numpy as np
+import matplotlib as mpl
+
+from .style import (
+    rotate_style,
+    copy_with_deep_values,
+)
+from .utils.matplotlib import (
+    _stale_wrapper,
+    _forwarder,
+)
+
+
+@_forwarder(
+    (
+        "set_clip_path",
+        "set_clip_box",
+        "set_snap",
+        "set_sketch_params",
+        "set_animated",
+        "set_picker",
+    )
+)
+class LabelCollection(mpl.artist.Artist):
+    def __init__(
+        self,
+        labels: Sequence[str],
+        style: Optional[dict[str, dict]] = None,
+        offsets: Optional[np.ndarray] = None,
+        transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
+    ):
+        self._labels = labels
+        self._offsets = offsets if offsets is not None else np.zeros((len(labels), 2))
+        self._style = style
+        super().__init__()
+
+        self.set_transform(transform)
+        self._create_artists()
+
+    def get_children(self):
+        return tuple(self._labelartists)
+
+    def set_figure(self, figure):
+        super().set_figure(figure)
+        for child in self.get_children():
+            child.set_figure(figure)
+        self._update_offsets(dpi=figure.dpi)
+
+    def _get_margins_with_dpi(self, dpi=72.0):
+        return self._margins * dpi / 72.0
+
+    def _create_artists(self):
+        style = copy_with_deep_values(self._style) if self._style is not None else {}
+        transform = self.get_transform()
+
+        margins = []
+
+        forbidden_props = ["rotate"]
+        for prop in forbidden_props:
+            if prop in style:
+                del style[prop]
+
+        arts = []
+        for i, (anchor_id, label) in enumerate(self._labels.items()):
+            stylei = rotate_style(style, index=i, key=anchor_id)
+            # Margins are handled separately
+            hmargin = stylei.pop("hmargin", 0.0)
+            vmargin = stylei.pop("vmargin", 0.0)
+            margins.append((hmargin, vmargin))
+
+            art = mpl.text.Text(
+                self._offsets[i][0],
+                self._offsets[i][1],
+                label,
+                transform=transform,
+                **stylei,
+            )
+            arts.append(art)
+        self._labelartists = arts
+        self._margins = np.array(margins)
+
+    def _update_offsets(self, dpi=72.0):
+        """Update offsets including margins."""
+        offsets = self._adjust_offsets_for_margins(self._offsets, dpi=dpi)
+        self.set_offsets(offsets)
+
+    def get_offsets(self):
+        return self._offsets
+
+    def _adjust_offsets_for_margins(self, offsets, dpi=72.0):
+        margins = self._get_margins_with_dpi(dpi=dpi)
+        if (margins != 0).any():
+            transform = self.get_transform()
+            trans = transform.transform
+            trans_inv = transform.inverted().transform
+            offsets = trans_inv(trans(offsets) + margins)
+        return offsets
+
+    def set_offsets(self, offsets):
+        """Set positions (offsets) of the labels."""
+        self._offsets = np.asarray(offsets)
+        for art, offset in zip(self._labelartists, self._offsets):
+            art.set_position((offset[0], offset[1]))
+
+    def set_rotations(self, rotations):
+        for art, rotation in zip(self._labelartists, rotations):
+            rot_deg = 180.0 / np.pi * rotation
+            # Force the font size to be upwards
+            rot_deg = ((rot_deg + 90) % 180) - 90
+            art.set_rotation(rot_deg)
+
+    @_stale_wrapper
+    def draw(self, renderer):
+        """Draw each of the children, with some buffering mechanism."""
+        if not self.get_visible():
+            return
+
+        self._update_offsets(dpi=renderer.dpi)
+
+        # We should manage zorder ourselves, but we need to compute
+        # the new offsets and angles of arrows from the edges before drawing them
+        for art in self.get_children():
+            art.draw(renderer)

iplotx/layout.py

@@ -0,0 +1,139 @@
+"""
+Layout functions, currently limited to trees.
+"""
+
+from collections.abc import Hashable
+
+import numpy as np
+
+
+def compute_tree_layout(
+    tree,
+    layout: str,
+    orientation: str,
+    **kwargs,
+) -> dict[Hashable, list[float]]:
+    """Compute the layout for a tree.
+
+    Parameters:
+        tree: The tree to compute the layout for.
+        layout: The name of the layout, e.g. "horizontal" or "radial".
+        orientation: The orientation of the layout, e.g. "right", "left", "descending", or "ascending".
+
+    Returns:
+        A layout dictionary with node positions.
+    """
+
+    if layout == "radial":
+        layout_dict = _circular_tree_layout(tree, orientation=orientation, **kwargs)
+    elif layout == "horizontal":
+        layout_dict = _horizontal_tree_layout(tree, orientation=orientation, **kwargs)
+    elif layout == "vertical":
+        layout_dict = _vertical_tree_layout(tree, orientation=orientation, **kwargs)
+    else:
+        raise ValueError(f"Tree layout not available: {layout}")
+
+    return layout_dict
+
+
+def _horizontal_tree_layout_right(
+    tree,
+    root_fun: callable,
+    preorder_fun: callable,
+    postorder_fun: callable,
+    children_fun: callable,
+    branch_length_fun: callable,
+) -> dict[Hashable, list[float]]:
+    """Build a tree layout horizontally, left to right.
+
+    The strategy is the usual one:
+    1. Compute the y values for the leaves, from 0 upwards.
+    2. Compute the y values for the internal nodes, bubbling up (postorder).
+    3. Set the x value for the root as 0.
+    4. Compute the x value of all nodes, trickling down (BFS/preorder).
+    5. Compute the edges from the end nodes.
+    """
+    layout = {}
+
+    # Set the y values for vertices
+    i = 0
+    for node in postorder_fun(tree):
+        children = children_fun(node)
+        if len(children) == 0:
+            layout[node] = [None, i]
+            i += 1
+        else:
+            layout[node] = [
+                None,
+                np.mean([layout[child][1] for child in children]),
+            ]
+
+    # Set the x values for vertices
+    layout[root_fun(tree)][0] = 0
+    for node in preorder_fun(tree):
+        x0, y0 = layout[node]
+        for child in children_fun(node):
+            bl = branch_length_fun(child)
+            if bl is None:
+                bl = 1.0
+            layout[child][0] = layout[node][0] + bl
+
+    return layout
+
+
+def _horizontal_tree_layout(
+    tree,
+    orientation="right",
+    **kwargs,
+) -> dict[Hashable, list[float]]:
+    """Horizontal tree layout."""
+    if orientation not in ("right", "left"):
+        raise ValueError("Orientation must be 'right' or 'left'.")
+
+    layout = _horizontal_tree_layout_right(tree, **kwargs)
+
+    if orientation == "left":
+        for key, value in layout.items():
+            layout[key][0] *= -1
+    return layout
+
+
+def _vertical_tree_layout(
+    tree,
+    orientation="descending",
+    **kwargs,
+) -> dict[Hashable, list[float]]:
+    """Vertical tree layout."""
+    sign = 1 if orientation == "descending" else -1
+    layout = _horizontal_tree_layout(tree, **kwargs)
+    for key, value in layout.items():
+        # Invert x and y
+        layout[key] = value[::-1]
+        # Orient vertically
+        layout[key][1] *= sign
+    return layout
+
+
+def _circular_tree_layout(
+    tree,
+    orientation="right",
+    starting_angle=0,
+    angular_span=360,
+    **kwargs,
+) -> dict[Hashable, list[float]]:
+    """Circular tree layout."""
+    # Short form
+    th = starting_angle * np.pi / 180
+    th_span = angular_span * np.pi / 180
+    sign = 1 if orientation == "right" else -1
+
+    layout = _horizontal_tree_layout_right(tree, **kwargs)
+    ymax = max(point[1] for point in layout.values())
+    for key, (x, y) in layout.items():
+        r = x
+        theta = sign * th_span * y / (ymax + 1) + th
+        # We export r and theta to ensure theta does not
+        # modulo 2pi if we take the tan and then arctan later.
+        layout[key] = (r, theta)
+
+    return layout