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

iplotx/label.py

@@ -1,3 +1,7 @@
+"""
+Module for label collection in iplotx.
+"""
+
 from typing import (
     Optional,
     Sequence,
@@ -26,13 +30,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],
         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 +60,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()
 
@@ -83,12 +108,13 @@ class LabelCollection(mpl.artist.Artist):
         self._labelartists = arts
         self._margins = np.array(margins)
 
-    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):
+    def get_offsets(self) -> np.ndarray:
+        """Get the positions (offsets) of the labels."""
         return self._offsets
 
     def _adjust_offsets_for_margins(self, offsets, dpi=72.0):
@@ -100,13 +126,22 @@ class LabelCollection(mpl.artist.Artist):
             offsets = trans_inv(trans(offsets) + margins)
         return offsets
 
-    def set_offsets(self, offsets):
-        """Set positions (offsets) of the labels."""
+    def set_offsets(self, offsets) -> 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):
             art.set_position((offset[0], offset[1]))
 
-    def set_rotations(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.
+        """
         for art, rotation in zip(self._labelartists, rotations):
             rot_deg = 180.0 / np.pi * rotation
             # Force the font size to be upwards
@@ -114,7 +149,7 @@ class LabelCollection(mpl.artist.Artist):
             art.set_rotation(rot_deg)
 
     @_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

@@ -18,7 +18,8 @@ def compute_tree_layout(
     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".
+        orientation: The orientation of the layout, e.g. "right", "left", "descending", or
+            "ascending".
 
     Returns:
         A layout dictionary with node positions.
@@ -71,7 +72,6 @@ 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]
         for child in children_fun(node):
             bl = branch_length_fun(child)
             if bl is None:
@@ -93,7 +93,7 @@ def _horizontal_tree_layout(
     layout = _horizontal_tree_layout_right(tree, **kwargs)
 
     if orientation == "left":
-        for key, value in layout.items():
+        for key in layout:
             layout[key][0] *= -1
     return layout
 

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

@@ -1,6 +1,5 @@
 from typing import (
     Any,
-    Never,
     Optional,
     Sequence,
 )
@@ -23,7 +22,7 @@ style_leaves = (
     "zorder",
     "tension",
     "looptension",
-    "lloopmaxangle",
+    "loopmaxangle",
     "rotate",
     "marker",
     "waypoints",
@@ -37,6 +36,15 @@ style_leaves = (
     "ports",
 )
 
+# These properties are not allowed to be rotated (global throughout the graph).
+# This might change in the future as the API improves.
+nonrotating_leaves = (
+    "offset",
+    "looptension",
+    "loopmaxangle",
+    "vertexpadding",
+)
+
 
 default = {
     "vertex": {
@@ -247,7 +255,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 +284,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 +334,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 +346,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 +359,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,6 +2,7 @@ from typing import (
     Optional,
     Sequence,
 )
+from collections.abc import Hashable
 
 import numpy as np
 import pandas as pd
@@ -26,6 +27,9 @@ from .edge import (
     EdgeCollection,
     make_stub_patch as make_undirected_edge_patch,
 )
+from .label import (
+    LabelCollection,
+)
 from .network import (
     _update_from_internal,
 )
@@ -50,13 +54,31 @@ class TreeArtist(mpl.artist.Artist):
         layout="horizontal",
         orientation="right",
         directed: bool | str = False,
-        vertex_labels: Optional[list | dict | pd.Series] = None,
+        vertex_labels: Optional[
+            bool | list[str] | dict[Hashable, str] | pd.Series
+        ] = None,
         edge_labels: Optional[Sequence] = 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.
+            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,
@@ -80,13 +102,23 @@ class TreeArtist(mpl.artist.Artist):
         self._add_vertices()
         self._add_edges()
 
-    def get_children(self):
+    def get_children(self) -> tuple[mpl.artist.Artist]:
+        """Get the children of this artist.
+
+        Returns:
+            The artists for vertices and edges.
+        """
         return (self._vertices, self._edges)
 
-    def set_figure(self, figure):
-        super().set_figure(figure)
+    def set_figure(self, fig) -> None:
+        """Set the figure for this artist and its children.
+
+        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)
 
     def get_offset_transform(self):
         """Get the offset transform (for vertices/edges)."""
@@ -132,43 +164,43 @@ class TreeArtist(mpl.artist.Artist):
         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_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 _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_edges(self) -> None:
         """Add edges to the network artist.
 
         NOTE: UndirectedEdgeCollection and ArrowCollection are both subclasses of
@@ -256,11 +288,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,
@@ -270,7 +297,7 @@ class TreeArtist(mpl.artist.Artist):
             self._edges.set_array(colorarray)
 
     @_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/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.2.1"

iplotx/vertex.py

@@ -6,7 +6,6 @@ from typing import (
     Optional,
     Sequence,
     Any,
-    Never,
 )
 import warnings
 import numpy as np
@@ -70,20 +69,21 @@ class VertexCollection(PatchCollection):
         self._index = layout.index
         self._style = style
         self._labels = labels
+        self._layout = layout
+        self._layout_coordinate_system = layout_coordinate_system
 
         # Create patches from structured data
-        patches, offsets, sizes, kwargs2 = self._init_vertex_patches(
-            layout,
-            layout_coordinate_system=layout_coordinate_system,
-        )
+        patches, sizes, kwargs2 = self._init_vertex_patches()
 
         kwargs.update(kwargs2)
-        kwargs["offsets"] = offsets
         kwargs["match_original"] = True
 
         # Pass to PatchCollection constructor
         super().__init__(patches, *args, **kwargs)
 
+        # Set offsets in coordinate system
+        self._update_offsets_from_layout()
+
         # Compute _transforms like in _CollectionWithScales for dpi issues
         self.set_sizes(sizes)
 
@@ -100,7 +100,7 @@ class VertexCollection(PatchCollection):
             children.append(self._label_collection)
         return tuple(children)
 
-    def set_figure(self, fig) -> Never:
+    def set_figure(self, fig) -> None:
         """Set the figure for this artist and all children."""
         super().set_figure(fig)
         self.set_sizes(self._sizes, self.get_figure(root=True).dpi)
@@ -123,7 +123,7 @@ class VertexCollection(PatchCollection):
         """Get vertex sizes (max of width and height), scaled by dpi."""
         return self._transforms[:, 0, 0]
 
-    def set_sizes(self, sizes, dpi=72.0):
+    def set_sizes(self, sizes, dpi: float = 72.0) -> None:
         """Set vertex sizes.
 
         This rescales the current vertex symbol/path linearly, using this
@@ -146,9 +146,70 @@ class VertexCollection(PatchCollection):
     get_size = get_sizes
     set_size = set_sizes
 
-    def _init_vertex_patches(
-        self, vertex_layout_df, layout_coordinate_system="cartesian"
-    ):
+    def get_layout(self) -> pd.DataFrame:
+        """Get the vertex layout.
+
+        Returns:
+            The vertex layout as a DataFrame.
+        """
+        return self._layout
+
+    def get_layout_coordinate_system(self) -> str:
+        """Get the layout coordinate system.
+
+        Returns:
+            Name of the layout coordinate system, e.g. "cartesian" or "polar".
+        """
+        return self._layout_coordinate_system
+
+    def get_offsets(self, ignore_layout: bool = True) -> np.ndarray:
+        """Get the vertex offsets.
+
+        Parameters:
+            ignore_layout: If True, return the matplotlib Artist._offsets directly, ignoring the
+                layout coordinate system. If False, it's equivalent to get_layout().values.
+
+        Returns:
+            The vertex offsets as a 2D numpy array.
+
+        Note: It is best for users to *not* ignore the layout coordinate system, as it may lead
+        to inconsistencies. However, some internal matplotlib functions require the default
+        signature of this function to look at the vanilla offsets, hence the default parameters.
+        """
+        if not ignore_layout:
+            return self.get_layout().values
+        else:
+            return self._offsets
+
+    def _update_offsets_from_layout(self) -> None:
+        """Update offsets in matplotlib coordinates from the layout DataFrame."""
+        if self._layout_coordinate_system == "cartesian":
+            self._offsets = self._layout.values
+        elif self._layout_coordinate_system == "polar":
+            # Convert polar coordinates (r, theta) to cartesian (x, y)
+            r = self._layout.iloc[:, 0].values
+            theta = self._layout.iloc[:, 1].values
+            if self._offsets is None:
+                self._offsets = np.zeros((len(r), 2))
+            self._offsets[:, 0] = r * np.cos(theta)
+            self._offsets[:, 1] = r * np.sin(theta)
+        else:
+            raise ValueError(
+                f"Layout coordinate system not supported: {self._layout_coordinate_system}."
+            )
+
+    def set_offsets(self, offsets: np.ndarray) -> None:
+        """Set the vertex positions/offsets in layout coordinates.
+
+        Parameters:
+            offsets: Array of coordinates in the layout coordinate system. For polar layouts,
+                these should be in the form of (r, theta) pairs.
+        """
+        self._layout.values[:] = offsets
+        self._update_offsets_from_layout()
+        self.stale = True
+
+    def _init_vertex_patches(self):
         style = self._style or {}
         if "cmap" in style:
             cmap_fun = _build_cmap_fun(
@@ -164,29 +225,16 @@ class VertexCollection(PatchCollection):
                     "No labels found, cannot resize vertices based on labels."
                 )
                 style["size"] = get_style("default.vertex")["size"]
-            else:
-                vertex_labels = self._labels
 
         if "cmap" in style:
             colorarray = []
         patches = []
-        offsets = []
         sizes = []
-        for i, (vid, row) in enumerate(vertex_layout_df.iterrows()):
-            # Centre of the vertex
-            offset = list(row.values)
-
-            # Transform to cartesian coordinates if needed
-            if layout_coordinate_system == "polar":
-                r, theta = offset
-                offset = [r * np.cos(theta), r * np.sin(theta)]
-
-            offsets.append(offset)
-
-            if style.get("size") == "label":
+        for i, (vid, row) in enumerate(self._layout.iterrows()):
+            if style.get("size", 20) == "label":
                 # NOTE: it's ok to overwrite the dict here
                 style["size"] = _get_label_width_height(
-                    str(vertex_labels[vid]), **style.get("label", {})
+                    str(self._labels[vid]), **style.get("label", {})
                 )
 
             stylei = rotate_style(style, index=i, key=vid)
@@ -207,7 +255,7 @@ class VertexCollection(PatchCollection):
             kwargs["cmap"] = style["cmap"]
             kwargs["norm"] = norm
 
-        return patches, offsets, sizes, kwargs
+        return patches, sizes, kwargs
 
     def _compute_label_collection(self):
         transform = self.get_offset_transform()
@@ -230,6 +278,12 @@ class VertexCollection(PatchCollection):
         )
 
     def get_labels(self):
+        """Get the vertex labels.
+
+        Returns:
+            The artist with the LabelCollection.
+        """
+
         if hasattr(self, "_label_collection"):
             return self._label_collection
         else:
@@ -280,7 +334,8 @@ def make_patch(
 
     # Size of vertices is determined in self._transforms, which scales with dpi, rather than here,
     # so normalise by the average dimension (btw x and y) to keep the ratio of the marker.
-    # If you check in get_sizes, you will see that rescaling also happens with the max of width and height.
+    # If you check in get_sizes, you will see that rescaling also happens with the max of width
+    # and height.
     size = np.asarray(size, dtype=float)
     size_max = size.max()
     if size_max > 0: