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

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,162 @@
+"""
+Module for label collection in iplotx.
+"""
+
+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):
+    """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
+        super().__init__()
+
+        self.set_transform(transform)
+        self._create_artists()
+
+    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, 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(fig)
+        self._update_offsets(dpi=fig.dpi)
+
+    def _get_margins_with_dpi(self, dpi: float = 72.0) -> np.ndarray:
+        return self._margins * dpi / 72.0
+
+    def _create_artists(self) -> None:
+        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: 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) -> np.ndarray:
+        """Get the positions (offsets) of the labels."""
+        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) -> 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: 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
+            rot_deg = ((rot_deg + 90) % 180) - 90
+            art.set_rotation(rot_deg)
+
+    @_stale_wrapper
+    def draw(self, renderer) -> None:
+        """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):
+        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 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)
+    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