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

iplotx/network.py

@@ -128,12 +128,12 @@ class NetworkArtist(mpl.artist.Artist):
         if len(layout) == 0:
             return mpl.transforms.Bbox([[0, 0], [1, 1]])
 
-        if self._vertices is not None:
-            bbox = self._vertices.get_datalim(transData)
-
-        if self._edges is not None:
-            edge_bbox = self._edges.get_datalim(transData)
-            bbox = mpl.transforms.Bbox.union([bbox, edge_bbox])
+        bbox = mpl.transforms.Bbox.union(
+            [
+                self._vertices.get_datalim(transData),
+                self._edges.get_datalim(transData),
+            ]
+        )
 
         bbox = bbox.expanded(sw=(1.0 + pad), sh=(1.0 + pad))
         return bbox
@@ -162,6 +162,9 @@ class NetworkArtist(mpl.artist.Artist):
 
         self._vertices = VertexCollection(
             layout=self.get_layout(),
+            layout_coordinate_system=self._ipx_internal_data.get(
+                "layout_coordinate_system", "cartesian"
+            ),
             style=get_style(".vertex"),
             labels=self._get_label_series("vertex"),
             transform=self.get_transform(),
@@ -238,8 +241,6 @@ class NetworkArtist(mpl.artist.Artist):
             edgepatches,
             vertex_ids=adjacent_vertex_ids,
             vertex_collection=self._vertices,
-            layout=self.get_layout(),
-            layout_coordinate_system="cartesian",
             labels=labels,
             transform=self.get_offset_transform(),
             style=edge_style,

iplotx/plotting.py

@@ -123,9 +123,10 @@ def network(
 def tree(
     tree: Optional[TreeType] = None,
     layout: str | LayoutType = "horizontal",
-    orientation: str = "right",
+    orientation: Optional[str] = None,
     directed: bool | str = False,
     vertex_labels: Optional[list | dict | pd.Series] = None,
+    leaf_labels: Optional[list | dict | pd.Series] = None,
     ax: Optional[mpl.axes.Axes] = None,
     style: str | dict | Sequence[str | dict] = "tree",
     title: Optional[str] = None,
@@ -138,8 +139,9 @@ def tree(
     Parameters:
         tree: The tree to plot. Can be a BioPython.Phylo.Tree object.
         layout: The layout to use for plotting.
-        orientation: The orientation of the horizontal layout. Can be "right" or "left". Defaults to
-            "right".
+        orientation: The orientation of the layout. Can be "right" or "left". Defaults to
+            "right" for horizontal layout, "descending" or "ascending" for vertical layout,
+            and "clockwise" or "anticlockwise" for radial layout.
         directed: If False, donot draw arrows. If True or "child", draw arrows from parent to child
             node. If "parent", draw arrows the other way around.
 
@@ -160,6 +162,7 @@ def tree(
             transform=mpl.transforms.IdentityTransform(),
             offset_transform=ax.transData,
             vertex_labels=vertex_labels,
+            leaf_labels=leaf_labels,
         )
         ax.add_artist(artist)
 

iplotx/style.py

@@ -1,6 +1,5 @@
 from typing import (
     Any,
-    Never,
     Optional,
     Sequence,
 )
@@ -23,7 +22,9 @@ style_leaves = (
     "zorder",
     "tension",
     "looptension",
-    "lloopmaxangle",
+    "loopmaxangle",
+    "paralleloffset",
+    "offset",
     "rotate",
     "marker",
     "waypoints",
@@ -35,6 +36,17 @@ style_leaves = (
     "hmargin",
     "vmargin",
     "ports",
+    "extend",
+)
+
+# These properties are not allowed to be rotated (global throughout the graph).
+# This might change in the future as the API improves.
+nonrotating_leaves = (
+    "paralleloffset",
+    "looptension",
+    "loopmaxangle",
+    "vertexpadding",
+    "extend",
 )
 
 
@@ -56,7 +68,7 @@ default = {
         "linestyle": "-",
         "color": "black",
         "curved": False,
-        "offset": 3,
+        "paralleloffset": 3,
         "tension": 1,
         "looptension": 4,
         "loopmaxangle": 60,
@@ -103,6 +115,7 @@ hollow["vertex"]["edgecolor"] = "black"
 hollow["vertex"]["linewidth"] = 1.5
 hollow["vertex"]["marker"] = "r"
 hollow["vertex"]["size"] = "label"
+hollow["vertex"]["label"]["color"] = "black"
 
 tree = copy_with_deep_values(default)
 tree["vertex"]["size"] = 0
@@ -115,8 +128,8 @@ tree["vertex"]["label"]["bbox"] = {
 }
 tree["vertex"]["label"]["color"] = "black"
 tree["vertex"]["label"]["size"] = 12
-tree["vertex"]["label"]["horizontalalignment"] = "left"
-tree["vertex"]["label"]["hmargin"] = 5
+tree["vertex"]["label"]["verticalalignment"] = "center"
+tree["vertex"]["label"]["hmargin"] = 10
 
 
 styles = {
@@ -137,13 +150,17 @@ def get_stylename():
     return str(stylename)
 
 
-def get_style(name: str = "") -> dict[str, Any]:
+def get_style(name: str = "", *args) -> dict[str, Any]:
     """Get a *deep copy* of the chosen style.
 
     Parameters:
         name: The name of the style to get. If empty, the current style is returned.
             Substyles can be obtained by using a dot notation, e.g. "default.vertex".
             If "name" starts with a dot, it means a substyle of the current style.
+        *args: A single argument is accepted. If present, this value (usually a
+            dictionary) is returned if the queried style is not found. For example,
+            get_style(".nonexistent") raises an Exception but
+            get_style("nonexistent", {}) does not, returning an empty dict instead.
     Returns:
         The requected style or substyle.
 
@@ -152,6 +169,9 @@ def get_style(name: str = "") -> dict[str, Any]:
         useful for hashables that change hash upon copying, such as Biopython's
         tree nodes.
     """
+    if len(args) > 1:
+        raise ValueError("get_style() accepts at most one additional argument.")
+
     namelist = name.split(".")
     style = styles
     for i, namei in enumerate(namelist):
@@ -165,6 +185,8 @@ def get_style(name: str = "") -> dict[str, Any]:
             # which will not fail unless the uder tries to enter it
             elif namei not in style_leaves:
                 style = {}
+            elif len(args) > 0:
+                return args[0]
             else:
                 raise KeyError(f"Style not found: {name}")
 
@@ -247,7 +269,7 @@ def use(style: Optional[str | dict | Sequence] = None, **kwargs):
         raise
 
 
-def reset() -> Never:
+def reset() -> None:
     """Reset to default style."""
     global current
     current = copy_with_deep_values(styles["default"])
@@ -276,7 +298,7 @@ def context(style: Optional[str | dict | Sequence] = None, **kwargs):
 
 def unflatten_style(
     style_flat: dict[str, str | dict | int | float],
-) -> Never:
+) -> None:
     """Convert a flat or semi-flat style into a fully structured dict.
 
     Parameters:
@@ -326,7 +348,7 @@ def rotate_style(
     style: dict[str, Any],
     index: Optional[int] = None,
     key: Optional[Hashable] = None,
-    props: Sequence[str] = style_leaves,
+    props: Optional[Sequence[str]] = None,
 ) -> dict[str, Any]:
     """Rotate leaves of a style for a certain index or key.
 
@@ -338,7 +360,8 @@ def rotate_style(
         props: The properties to rotate, usually all leaf properties.
 
     Returns:
-        A style with rotated leaves, which describes the properties of a single element (e.g. vertex).
+        A style with rotated leaves, which describes the properties of a single element (e.g.
+        vertex).
 
     Example:
         >>> style = {'vertex': {'size': [10, 20]}}
@@ -350,6 +373,9 @@ def rotate_style(
             "At least one of 'index' or 'key' must be provided to rotate_style."
         )
 
+    if props is None:
+        props = tuple(prop for prop in style_leaves if prop not in nonrotating_leaves)
+
     style = copy_with_deep_values(style)
 
     for prop in props:

iplotx/tree.py

@@ -2,12 +2,15 @@ from typing import (
     Optional,
     Sequence,
 )
+from collections.abc import Hashable
+from collections import defaultdict
 
 import numpy as np
 import pandas as pd
 import matplotlib as mpl
 
 from .style import (
+    context,
     get_style,
     rotate_style,
 )
@@ -18,6 +21,7 @@ from .utils.matplotlib import (
 )
 from .ingest import (
     ingest_tree_data,
+    data_providers,
 )
 from .vertex import (
     VertexCollection,
@@ -26,6 +30,12 @@ from .edge import (
     EdgeCollection,
     make_stub_patch as make_undirected_edge_patch,
 )
+from .label import (
+    LabelCollection,
+)
+from .cascades import (
+    CascadeCollection,
+)
 from .network import (
     _update_from_internal,
 )
@@ -47,23 +57,49 @@ class TreeArtist(mpl.artist.Artist):
     def __init__(
         self,
         tree,
-        layout="horizontal",
-        orientation="right",
+        layout: Optional[str] = "horizontal",
+        orientation: Optional[str] = None,
         directed: bool | str = False,
-        vertex_labels: Optional[list | dict | pd.Series] = None,
-        edge_labels: Optional[Sequence] = None,
+        vertex_labels: Optional[
+            bool | list[str] | dict[Hashable, str] | pd.Series
+        ] = None,
+        edge_labels: Optional[Sequence | dict[Hashable, str] | pd.Series] = None,
+        leaf_labels: Optional[Sequence | dict[Hashable, str]] | pd.Series = None,
         transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
         offset_transform: Optional[mpl.transforms.Transform] = None,
     ):
-        self.tree = tree
+        """Initialize the TreeArtist.
+
+        Parameters:
+            tree: The tree to plot.
+            layout: The layout to use for the tree. Can be "horizontal", "vertical", or "radial".
+            orientation: The orientation of the tree layout. Can be "right" or "left" (for
+                horizontal and radial layouts) and "descending" or "ascending" (for vertical
+                layouts).
+            directed: Whether the tree is directed. Can be a boolean or a string with the
+                following choices: "parent" or "child".
+            vertex_labels: Labels for the vertices. Can be a list, dictionary, or pandas Series.
+            edge_labels: Labels for the edges. Can be a sequence of strings.
+            leaf_labels: Labels for the leaves. Can be a sequence of strings or a pandas Series.
+                These labels are positioned at the depth of the deepest leaf. If you want to
+                label leaves next to each leaf independently of how deep they are, use
+                the "vertex_labels" parameter instead - usually as a dict with the leaves
+                as keys and the labels as values.
+            transform: The transform to apply to the tree artist. This is usually the identity.
+            offset_transform: The offset transform to apply to the tree artist. This is
+                usually `ax.transData`.
+        """
 
+        self.tree = tree
         self._ipx_internal_data = ingest_tree_data(
             tree,
             layout,
             orientation=orientation,
             directed=directed,
+            layout_style=get_style(".layout", {}),
             vertex_labels=vertex_labels,
             edge_labels=edge_labels,
+            leaf_labels=leaf_labels,
         )
 
         super().__init__()
@@ -79,14 +115,52 @@ class TreeArtist(mpl.artist.Artist):
 
         self._add_vertices()
         self._add_edges()
+        self._add_leaf_vertices()
+
+        # NOTE: cascades need to be created after leaf vertices in case
+        # they are requested to wrap around them.
+        if "cascade" in self.get_vertices().get_style():
+            self._add_cascades()
+
+    def get_children(self) -> tuple[mpl.artist.Artist]:
+        """Get the children of this artist.
+
+        Returns:
+            The artists for vertices and edges.
+        """
+        children = [self._vertices, self._edges]
+        if hasattr(self, "_leaf_vertices"):
+            children.append(self._leaf_vertices)
+        if hasattr(self, "_cascades"):
+            children.append(self._cascades)
+        return tuple(children)
 
-    def get_children(self):
-        return (self._vertices, self._edges)
+    def set_figure(self, fig) -> None:
+        """Set the figure for this artist and its children.
 
-    def set_figure(self, figure):
-        super().set_figure(figure)
+        Parameters:
+            fig: the figure to set for this artist and its children.
+        """
+        super().set_figure(fig)
         for child in self.get_children():
-            child.set_figure(figure)
+            child.set_figure(fig)
+
+        # At the end, if there are cadcades with extent depending on
+        # leaf edges, we should update them
+        self._update_cascades_extent()
+
+    def _update_cascades_extent(self) -> None:
+        """Update cascades if extent depends on leaf labels."""
+        if not hasattr(self, "_cascades"):
+            return
+
+        style_cascade = self.get_vertices().get_style()["cascade"]
+        extend_to_labels = style_cascade.get("extend", False) == "leaf_labels"
+        if not extend_to_labels:
+            return
+
+        maxdepth = self._get_maxdepth_leaf_labels()
+        self._cascades.set_maxdepth(maxdepth)
 
     def get_offset_transform(self):
         """Get the offset transform (for vertices/edges)."""
@@ -105,6 +179,16 @@ class TreeArtist(mpl.artist.Artist):
         if kind == "vertex":
             layout = self._ipx_internal_data["vertex_df"][layout_columns]
             return layout
+        elif kind == "leaf":
+            leaves = self._ipx_internal_data["leaf_df"].index
+            layout = self._ipx_internal_data["vertex_df"][layout_columns]
+            # NOTE: workaround for a pandas bug
+            idxs = []
+            for i, vid in enumerate(layout.index):
+                if vid in leaves:
+                    idxs.append(i)
+            layout = layout.iloc[idxs]
+            return layout
 
         elif kind == "edge":
             return self._ipx_internal_data["edge_df"][layout_columns]
@@ -129,46 +213,160 @@ class TreeArtist(mpl.artist.Artist):
         edge_bbox = self._edges.get_datalim(transData)
         bbox = mpl.transforms.Bbox.union([bbox, edge_bbox])
 
+        if hasattr(self, "_cascades"):
+            cascades_bbox = self._cascades.get_datalim(transData)
+            bbox = mpl.transforms.Bbox.union([bbox, cascades_bbox])
+
+        if hasattr(self, "_leaf_vertices"):
+            leaf_labels_bbox = self._leaf_vertices.get_datalim(transData)
+            bbox = mpl.transforms.Bbox.union([bbox, leaf_labels_bbox])
+
         bbox = bbox.expanded(sw=(1.0 + pad), sh=(1.0 + pad))
         return bbox
 
-    def _get_label_series(self, kind):
+    def _get_label_series(self, kind: str) -> Optional[pd.Series]:
         if "label" in self._ipx_internal_data[f"{kind}_df"].columns:
             return self._ipx_internal_data[f"{kind}_df"]["label"]
         else:
             return None
 
-    def get_vertices(self):
+    def get_vertices(self) -> VertexCollection:
         """Get VertexCollection artist."""
         return self._vertices
 
-    def get_edges(self):
+    def get_edges(self) -> EdgeCollection:
         """Get EdgeCollection artist."""
         return self._edges
 
-    def get_vertex_labels(self):
+    def get_leaf_vertices(self) -> Optional[VertexCollection]:
+        """Get leaf VertexCollection artist."""
+        if hasattr(self, "_leaf_vertices"):
+            return self._leaf_vertices
+        return None
+
+    def get_vertex_labels(self) -> LabelCollection:
         """Get list of vertex label artists."""
         return self._vertices.get_labels()
 
-    def get_edge_labels(self):
+    def get_edge_labels(self) -> LabelCollection:
         """Get list of edge label artists."""
         return self._edges.get_labels()
 
-    def _add_vertices(self):
+    def get_leaf_labels(self) -> Optional[LabelCollection]:
+        if hasattr(self, "_leaf_vertices"):
+            return self._leaf_vertices.get_labels()
+        return None
+
+    def _add_vertices(self) -> None:
         """Add vertices to the tree."""
         self._vertices = VertexCollection(
             layout=self.get_layout(),
-            style=get_style(".vertex"),
-            labels=self._get_label_series("vertex"),
             layout_coordinate_system=self._ipx_internal_data.get(
                 "layout_coordinate_system",
                 "catesian",
             ),
+            style=get_style(".vertex"),
+            labels=self._get_label_series("vertex"),
             transform=self.get_transform(),
             offset_transform=self.get_offset_transform(),
         )
 
-    def _add_edges(self):
+    def _add_leaf_vertices(self) -> None:
+        """Add invisible deep vertices as leaf label anchors."""
+        leaf_layout = self.get_layout("leaf").copy()
+        # Set all to max depth
+        depth_idx = int(self._ipx_internal_data["layout_name"] == "vertical")
+        leaf_layout.iloc[:, depth_idx] = leaf_layout.iloc[:, depth_idx].max()
+
+        # Set invisible vertices with visible labels
+        layout_name = self._ipx_internal_data["layout_name"]
+        orientation = self._ipx_internal_data["orientation"]
+        if layout_name == "radial":
+            ha = "auto"
+        elif orientation in ("left", "ascending"):
+            ha = "right"
+        else:
+            ha = "left"
+
+        leaf_vertex_style = {
+            "size": 0,
+            "label": {
+                "verticalalignment": "center",
+                "horizontalalignment": ha,
+                "hmargin": 5,
+                "bbox": {
+                    "facecolor": (1, 1, 1, 0),
+                },
+            },
+        }
+        with context({"vertex": leaf_vertex_style}):
+            leaf_vertex_style = get_style(".vertex")
+            self._leaf_vertices = VertexCollection(
+                layout=leaf_layout,
+                layout_coordinate_system=self._ipx_internal_data.get(
+                    "layout_coordinate_system",
+                    "catesian",
+                ),
+                style=leaf_vertex_style,
+                labels=self._get_label_series("leaf"),
+                transform=self.get_transform(),
+                offset_transform=self.get_offset_transform(),
+            )
+
+    def _add_cascades(self) -> None:
+        """Add cascade patches."""
+        # NOTE: If leaf labels are present and the cascades are requested to wrap around them,
+        # we have to compute the max extend of the cascades from the leaf labels.
+        maxdepth = None
+        style_cascade = self.get_vertices().get_style()["cascade"]
+        extend_to_labels = style_cascade.get("extend", False) == "leaf_labels"
+        has_leaf_labels = self.get_leaf_labels() is not None
+        if extend_to_labels and not has_leaf_labels:
+            raise ValueError("Cannot extend cascades: no leaf labels.")
+
+        if extend_to_labels and has_leaf_labels:
+            maxdepth = self._get_maxdepth_leaf_labels()
+
+        self._cascades = CascadeCollection(
+            tree=self.tree,
+            layout=self.get_layout(),
+            layout_name=self._ipx_internal_data["layout_name"],
+            orientation=self._ipx_internal_data["orientation"],
+            style=style_cascade,
+            provider=data_providers["tree"][self._ipx_internal_data["tree_library"]],
+            transform=self.get_offset_transform(),
+            maxdepth=maxdepth,
+        )
+
+    def _get_maxdepth_leaf_labels(self):
+        layout_name = self.get_layout_name()
+        if layout_name == "radial":
+            maxdepth = 0
+            # These are the text boxes, they must all be included
+            bboxes = self.get_leaf_labels().get_datalims_children(
+                self.get_offset_transform()
+            )
+            for bbox in bboxes:
+                r1 = np.linalg.norm([bbox.xmax, bbox.ymax])
+                r2 = np.linalg.norm([bbox.xmax, bbox.ymin])
+                r3 = np.linalg.norm([bbox.xmin, bbox.ymax])
+                r4 = np.linalg.norm([bbox.xmin, bbox.ymin])
+                maxdepth = max(maxdepth, r1, r2, r3, r4)
+        else:
+            orientation = self.get_orientation()
+            bbox = self.get_leaf_labels().get_datalim(self.get_offset_transform())
+            if (layout_name, orientation) == ("horizontal", "right"):
+                maxdepth = bbox.xmax
+            elif layout_name == "horizontal":
+                maxdepth = bbox.xmin
+            elif (layout_name, orientation) == ("vertical", "descending"):
+                maxdepth = bbox.ymin
+            elif layout_name == "vertical":
+                maxdepth = bbox.ymax
+
+        return maxdepth
+
+    def _add_edges(self) -> None:
         """Add edges to the network artist.
 
         NOTE: UndirectedEdgeCollection and ArrowCollection are both subclasses of
@@ -227,7 +425,7 @@ class TreeArtist(mpl.artist.Artist):
                 if layout_name == "horizontal":
                     waypointsi = "x0y1"
                 elif layout_name == "vertical":
-                    waypointsi = "y0y0"
+                    waypointsi = "y0x1"
                 elif layout_name == "radial":
                     waypointsi = "r0a1"
                 else:
@@ -256,11 +454,6 @@ class TreeArtist(mpl.artist.Artist):
             edgepatches,
             vertex_ids=adjacent_vertex_ids,
             vertex_collection=self._vertices,
-            layout=self.get_layout(kind="vertex"),
-            layout_coordinate_system=self._ipx_internal_data.get(
-                "layout_coordinate_system",
-                "cartesian",
-            ),
             labels=labels,
             transform=self.get_offset_transform(),
             style=edge_style,
@@ -269,17 +462,32 @@ class TreeArtist(mpl.artist.Artist):
         if "cmap" in edge_style:
             self._edges.set_array(colorarray)
 
+    def get_layout_name(self) -> str:
+        """Get the layout name."""
+        return self._ipx_internal_data["layout_name"]
+
+    def get_orientation(self) -> Optional[str]:
+        """Get the orientation of the tree layout."""
+        return self._ipx_internal_data.get("orientation", None)
+
     @_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
 
-        # FIXME: Callbacks on stale vertices/edges??
+        # At the end, if there are cadcades with extent depending on
+        # leaf edges, we should update them
+        self._update_cascades_extent()
 
         # NOTE: looks like we have to manage the zorder ourselves
-        # this is kind of funny actually
+        # this is kind of funny actually. Btw we need to ensure
+        # that cascades are drawn behind (earlier than) vertices
+        # and edges at equal zorder because it looks better that way.
+        z_suborder = defaultdict(int)
+        if hasattr(self, "_cascades"):
+            z_suborder[self._cascades] = -1
         children = list(self.get_children())
-        children.sort(key=lambda x: x.zorder)
+        children.sort(key=lambda x: (x.zorder, z_suborder[x]))
         for art in children:
             art.draw(renderer)

iplotx/typing.py

@@ -1,8 +1,17 @@
+"""
+Typing hints for iplotx.
+
+Some of these types are legit, others are just Any but renamed to improve readability throughout
+the codebase.
+"""
+
 from typing import (
     Union,
     Sequence,
     Any,
+    TypeVar,
 )
+from collections.abc import Hashable
 import numpy as np
 import pandas as pd
 
@@ -34,3 +43,13 @@ GroupingType = Union[
     # igraph.clustering.Cover,
     # igraph.clustering.VertexCover,
 ]
+
+
+T = TypeVar("T")
+LeafProperty = Union[
+    T,
+    Sequence[T],
+    dict[Hashable, T],
+]
+
+Pair = tuple[T, T]

iplotx/version.py

@@ -2,4 +2,4 @@
 iplotx version information module.
 """
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"