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

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

iplotx/label.py

@@ -1,8 +1,13 @@
+"""
+Module for label collection in iplotx.
+"""
+
 from typing import (
     Optional,
     Sequence,
 )
 import numpy as np
+import pandas as pd
 import matplotlib as mpl
 
 from .style import (
@@ -26,13 +31,28 @@ from .utils.matplotlib import (
     )
 )
 class LabelCollection(mpl.artist.Artist):
+    """Collection of labels for iplotx with styles.
+
+    NOTE: This class is not a subclass of `mpl.collections.Collection`, although in some ways items
+    behaves like one. It is named LabelCollection quite literally to indicate it contains a list of
+    labels for vertices, edges, etc.
+    """
+
     def __init__(
         self,
-        labels: Sequence[str],
+        labels: pd.Series,
         style: Optional[dict[str, dict]] = None,
         offsets: Optional[np.ndarray] = None,
         transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
-    ):
+    ) -> None:
+        """Initialize a collection of labels.
+
+        Parameters:
+            labels: A sequence of labels to be displayed.
+            style: A dictionary of styles to apply to the labels. The keys are style properties.
+            offsets: A sequence of offsets for each label, specifying the position of the label.
+            transform: A transform to apply to the labels. This is usually ax.transData.
+        """
         self._labels = labels
         self._offsets = offsets if offsets is not None else np.zeros((len(labels), 2))
         self._style = style
@@ -41,19 +61,25 @@ class LabelCollection(mpl.artist.Artist):
         self.set_transform(transform)
         self._create_artists()
 
-    def get_children(self):
+    def get_children(self) -> tuple[mpl.artist.Artist]:
+        """Get the children of this artist, which are the label artists."""
         return tuple(self._labelartists)
 
-    def set_figure(self, figure):
-        super().set_figure(figure)
+    def set_figure(self, fig) -> None:
+        """Set the figure of this artist.
+
+        Parameters:
+            fig: The figure to set.
+        """
+        super().set_figure(fig)
         for child in self.get_children():
-            child.set_figure(figure)
-        self._update_offsets(dpi=figure.dpi)
+            child.set_figure(fig)
+        self._update_offsets(dpi=fig.dpi)
 
-    def _get_margins_with_dpi(self, dpi=72.0):
+    def _get_margins_with_dpi(self, dpi: float = 72.0) -> np.ndarray:
         return self._margins * dpi / 72.0
 
-    def _create_artists(self):
+    def _create_artists(self) -> None:
         style = copy_with_deep_values(self._style) if self._style is not None else {}
         transform = self.get_transform()
 
@@ -72,6 +98,11 @@ class LabelCollection(mpl.artist.Artist):
             vmargin = stylei.pop("vmargin", 0.0)
             margins.append((hmargin, vmargin))
 
+            # Initially, ignore autoalignment since we do not know the
+            # rotations
+            if stylei.get("horizontalalignment") == "auto":
+                stylei["horizontalalignment"] = "center"
+
             art = mpl.text.Text(
                 self._offsets[i][0],
                 self._offsets[i][1],
@@ -82,14 +113,20 @@ class LabelCollection(mpl.artist.Artist):
             arts.append(art)
         self._labelartists = arts
         self._margins = np.array(margins)
+        self._rotations = np.zeros(len(self._labels))
 
-    def _update_offsets(self, dpi=72.0):
+    def _update_offsets(self, dpi: float = 72.0) -> None:
         """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
+        self.set_offsets(self._offsets, dpi=dpi)
+
+    def get_offsets(self, with_margins: bool = False) -> np.ndarray:
+        """Get the positions (offsets) of the labels."""
+        if not with_margins:
+            return self._offsets
+        else:
+            return np.array(
+                [art.get_position() for art in self._labelartists],
+            )
 
     def _adjust_offsets_for_margins(self, offsets, dpi=72.0):
         margins = self._get_margins_with_dpi(dpi=dpi)
@@ -97,24 +134,69 @@ class LabelCollection(mpl.artist.Artist):
             transform = self.get_transform()
             trans = transform.transform
             trans_inv = transform.inverted().transform
-            offsets = trans_inv(trans(offsets) + margins)
+            rotations = self.get_rotations()
+            vrot = [np.cos(rotations), np.sin(rotations)]
+
+            margins_rot = np.empty_like(margins)
+            margins_rot[:, 0] = margins[:, 0] * vrot[0] - margins[:, 1] * vrot[1]
+            margins_rot[:, 1] = margins[:, 0] * vrot[1] + margins[:, 1] * vrot[0]
+            offsets = trans_inv(trans(offsets) + margins_rot)
         return offsets
 
-    def set_offsets(self, offsets):
-        """Set positions (offsets) of the labels."""
+    def set_offsets(self, offsets, dpi: float = 72.0) -> None:
+        """Set positions (offsets) of the labels.
+
+        Parameters:
+            offsets: A sequence of offsets for each label, specifying the position of the label.
+        """
         self._offsets = np.asarray(offsets)
-        for art, offset in zip(self._labelartists, self._offsets):
+        offsets_with_margins = self._adjust_offsets_for_margins(offsets, dpi=dpi)
+        for art, offset in zip(self._labelartists, offsets_with_margins):
             art.set_position((offset[0], offset[1]))
 
-    def set_rotations(self, rotations):
+    def get_rotations(self) -> np.ndarray:
+        """Get the rotations of the labels in radians."""
+        return self._rotations
+
+    def set_rotations(self, rotations: Sequence[float]) -> None:
+        """Set the rotations of the labels.
+
+        Parameters:
+            rotations: A sequence of rotations in radians for each label.
+        """
+        self._rotations = np.asarray(rotations)
+        ha = self._style.get("horizontalalignment", "center")
         for art, rotation in zip(self._labelartists, rotations):
             rot_deg = 180.0 / np.pi * rotation
             # Force the font size to be upwards
+            if ha == "auto":
+                if -90 <= rot_deg < 90:
+                    art.set_horizontalalignment("left")
+                else:
+                    art.set_horizontalalignment("right")
             rot_deg = ((rot_deg + 90) % 180) - 90
             art.set_rotation(rot_deg)
 
+    def get_datalim(self, transData=None) -> mpl.transforms.Bbox:
+        """Get the data limits of the labels."""
+        bboxes = self.get_datalims_children(transData=transData)
+        bbox = mpl.transforms.Bbox.union(bboxes)
+        return bbox
+
+    def get_datalims_children(self, transData=None) -> Sequence[mpl.transforms.Bbox]:
+        """Get the data limits of the children of this artist."""
+        if transData is None:
+            transData = self.get_transform()
+        trans_inv = transData.inverted().transform_bbox
+        bboxes = []
+        for art in self._labelartists:
+            bbox_fig = art.get_bbox_patch().get_extents()
+            bbox_data = trans_inv(bbox_fig)
+            bboxes.append(bbox_data)
+        return bboxes
+
     @_stale_wrapper
-    def draw(self, renderer):
+    def draw(self, renderer) -> None:
         """Draw each of the children, with some buffering mechanism."""
         if not self.get_visible():
             return

iplotx/layout.py

@@ -2,34 +2,48 @@
 Layout functions, currently limited to trees.
 """
 
-from collections.abc import Hashable
+from typing import Any
+from collections.abc import (
+    Hashable,
+    Callable,
+)
 
 import numpy as np
 
 
 def compute_tree_layout(
-    tree,
     layout: str,
     orientation: str,
+    root: Any,
+    preorder_fun: Callable,
+    postorder_fun: Callable,
+    children_fun: Callable,
+    branch_length_fun: Callable,
     **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".
+        layout: The name of the layout, e.g. "horizontal", "vertial", or "radial".
+        orientation: The orientation of the layout, e.g. "right", "left", "descending",
+            "ascending", "clockwise", "anticlockwise".
 
     Returns:
         A layout dictionary with node positions.
     """
+    kwargs["root"] = root
+    kwargs["preorder_fun"] = preorder_fun
+    kwargs["postorder_fun"] = postorder_fun
+    kwargs["children_fun"] = children_fun
+    kwargs["branch_length_fun"] = branch_length_fun
+    kwargs["orientation"] = orientation
 
     if layout == "radial":
-        layout_dict = _circular_tree_layout(tree, orientation=orientation, **kwargs)
+        layout_dict = _radial_tree_layout(**kwargs)
     elif layout == "horizontal":
-        layout_dict = _horizontal_tree_layout(tree, orientation=orientation, **kwargs)
+        layout_dict = _horizontal_tree_layout(**kwargs)
     elif layout == "vertical":
-        layout_dict = _vertical_tree_layout(tree, orientation=orientation, **kwargs)
+        layout_dict = _vertical_tree_layout(**kwargs)
     else:
         raise ValueError(f"Tree layout not available: {layout}")
 
@@ -37,12 +51,11 @@ def compute_tree_layout(
 
 
 def _horizontal_tree_layout_right(
-    tree,
-    root_fun: callable,
-    preorder_fun: callable,
-    postorder_fun: callable,
-    children_fun: callable,
-    branch_length_fun: callable,
+    root: Any,
+    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.
 
@@ -57,7 +70,7 @@ def _horizontal_tree_layout_right(
 
     # Set the y values for vertices
     i = 0
-    for node in postorder_fun(tree):
+    for node in postorder_fun():
         children = children_fun(node)
         if len(children) == 0:
             layout[node] = [None, i]
@@ -69,9 +82,8 @@ def _horizontal_tree_layout_right(
             ]
 
     # Set the x values for vertices
-    layout[root_fun(tree)][0] = 0
-    for node in preorder_fun(tree):
-        x0, y0 = layout[node]
+    layout[root][0] = 0
+    for node in preorder_fun():
         for child in children_fun(node):
             bl = branch_length_fun(child)
             if bl is None:
@@ -82,7 +94,6 @@ def _horizontal_tree_layout_right(
 
 
 def _horizontal_tree_layout(
-    tree,
     orientation="right",
     **kwargs,
 ) -> dict[Hashable, list[float]]:
@@ -90,22 +101,21 @@ def _horizontal_tree_layout(
     if orientation not in ("right", "left"):
         raise ValueError("Orientation must be 'right' or 'left'.")
 
-    layout = _horizontal_tree_layout_right(tree, **kwargs)
+    layout = _horizontal_tree_layout_right(**kwargs)
 
     if orientation == "left":
-        for key, value in layout.items():
+        for key in layout:
             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)
+    sign = -1 if orientation == "descending" else 1
+    layout = _horizontal_tree_layout(**kwargs)
     for key, value in layout.items():
         # Invert x and y
         layout[key] = value[::-1]
@@ -114,24 +124,35 @@ def _vertical_tree_layout(
     return layout
 
 
-def _circular_tree_layout(
-    tree,
-    orientation="right",
-    starting_angle=0,
-    angular_span=360,
+def _radial_tree_layout(
+    orientation: str = "right",
+    start: float = 180,
+    span: float = 360,
     **kwargs,
-) -> dict[Hashable, list[float]]:
-    """Circular tree layout."""
+) -> dict[Hashable, tuple[float, float]]:
+    """Radial tree layout.
+
+    Parameters:
+        orientation: Whether the layout fans out towards the right (clockwise) or left
+            (anticlockwise).
+        start: The starting angle in degrees, default is -180 (left).
+        span: The angular span in degrees, default is 360 (full circle). When this is
+            360, it leaves a small gap at the end to ensure the first and last leaf
+            are not overlapping.
+    Returns:
+        A dictionary with the radial layout.
+    """
     # Short form
-    th = starting_angle * np.pi / 180
-    th_span = angular_span * np.pi / 180
-    sign = 1 if orientation == "right" else -1
+    th = start * np.pi / 180
+    th_span = span * np.pi / 180
+    pad = int(span == 360)
+    sign = -1 if orientation in ("right", "clockwise") else 1
 
-    layout = _horizontal_tree_layout_right(tree, **kwargs)
+    layout = _horizontal_tree_layout_right(**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
+        theta = sign * th_span * y / (ymax + pad) + 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)