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

iplotx/tree.py

@@ -0,0 +1,312 @@
+from typing import (
+    Optional,
+    Sequence,
+)
+from collections.abc import Hashable
+
+import numpy as np
+import pandas as pd
+import matplotlib as mpl
+
+from .style import (
+    get_style,
+    rotate_style,
+)
+from .utils.matplotlib import (
+    _stale_wrapper,
+    _forwarder,
+    _build_cmap_fun,
+)
+from .ingest import (
+    ingest_tree_data,
+)
+from .vertex import (
+    VertexCollection,
+)
+from .edge import (
+    EdgeCollection,
+    make_stub_patch as make_undirected_edge_patch,
+)
+from .label import (
+    LabelCollection,
+)
+from .network import (
+    _update_from_internal,
+)
+
+
+@_forwarder(
+    (
+        "set_clip_path",
+        "set_clip_box",
+        "set_snap",
+        "set_sketch_params",
+        "set_animated",
+        "set_picker",
+    )
+)
+class TreeArtist(mpl.artist.Artist):
+    """Artist for plotting trees."""
+
+    def __init__(
+        self,
+        tree,
+        layout="horizontal",
+        orientation="right",
+        directed: bool | str = False,
+        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,
+    ):
+        """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,
+            orientation=orientation,
+            directed=directed,
+            vertex_labels=vertex_labels,
+            edge_labels=edge_labels,
+        )
+
+        super().__init__()
+
+        # This is usually the identity (which scales poorly with dpi)
+        self.set_transform(transform)
+
+        # This is usually transData
+        self.set_offset_transform(offset_transform)
+
+        zorder = get_style(".network").get("zorder", 1)
+        self.set_zorder(zorder)
+
+        self._add_vertices()
+        self._add_edges()
+
+    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, 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(fig)
+
+    def get_offset_transform(self):
+        """Get the offset transform (for vertices/edges)."""
+        return self._offset_transform
+
+    def set_offset_transform(self, offset_transform):
+        """Set the offset transform (for vertices/edges)."""
+        self._offset_transform = offset_transform
+
+    def get_layout(self, kind="vertex"):
+        """Get vertex or edge layout."""
+        layout_columns = [
+            f"_ipx_layout_{i}" for i in range(self._ipx_internal_data["ndim"])
+        ]
+
+        if kind == "vertex":
+            layout = self._ipx_internal_data["vertex_df"][layout_columns]
+            return layout
+
+        elif kind == "edge":
+            return self._ipx_internal_data["edge_df"][layout_columns]
+        else:
+            raise ValueError(f"Unknown layout kind: {kind}. Use 'vertex' or 'edge'.")
+
+    def get_datalim(self, transData, pad=0.15):
+        """Get limits on x/y axes based on the graph layout data.
+
+        Parameters:
+            transData (Transform): The transform to use for the data.
+            pad (float): Padding to add to the limits. Default is 0.05.
+                Units are a fraction of total axis range before padding.
+        """
+        layout = self.get_layout().values
+
+        if len(layout) == 0:
+            return mpl.transforms.Bbox([[0, 0], [1, 1]])
+
+        bbox = self._vertices.get_datalim(transData)
+
+        edge_bbox = self._edges.get_datalim(transData)
+        bbox = mpl.transforms.Bbox.union([bbox, edge_bbox])
+
+        bbox = bbox.expanded(sw=(1.0 + pad), sh=(1.0 + pad))
+        return bbox
+
+    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) -> VertexCollection:
+        """Get VertexCollection artist."""
+        return self._vertices
+
+    def get_edges(self) -> EdgeCollection:
+        """Get EdgeCollection artist."""
+        return self._edges
+
+    def get_vertex_labels(self) -> LabelCollection:
+        """Get list of vertex label artists."""
+        return self._vertices.get_labels()
+
+    def get_edge_labels(self) -> LabelCollection:
+        """Get list of edge label artists."""
+        return self._edges.get_labels()
+
+    def _add_vertices(self) -> None:
+        """Add vertices to the tree."""
+        self._vertices = VertexCollection(
+            layout=self.get_layout(),
+            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) -> None:
+        """Add edges to the network artist.
+
+        NOTE: UndirectedEdgeCollection and ArrowCollection are both subclasses of
+        PatchCollection. When used with a cmap/norm, they set their facecolor
+        according to the cmap, even though most likely we only want the edgecolor
+        set that way. It can make for funny looking plots that are not uninteresting
+        but mostly niche at this stage. Therefore we sidestep the whole cmap thing
+        here.
+        """
+
+        labels = self._get_label_series("edge")
+        edge_style = get_style(".edge")
+
+        if "cmap" in edge_style:
+            cmap_fun = _build_cmap_fun(
+                edge_style["color"],
+                edge_style["cmap"],
+            )
+        else:
+            cmap_fun = None
+
+        edge_df = self._ipx_internal_data["edge_df"].set_index(
+            ["_ipx_source", "_ipx_target"]
+        )
+
+        if "cmap" in edge_style:
+            colorarray = []
+        edgepatches = []
+        adjacent_vertex_ids = []
+        waypoints = []
+        for i, (vid1, vid2) in enumerate(edge_df.index):
+            edge_stylei = rotate_style(edge_style, index=i, key=(vid1, vid2))
+
+            # FIXME:: Improve this logic. We have three layers of priority:
+            # 1. Explicitely set in the style of "plot"
+            # 2. Internal through network attributes
+            # 3. Default styles
+            # Because 1 and 3 are merged as a style context on the way in,
+            # it's hard to squeeze 2 in the middle. For now, we will assume
+            # the priority order is 2-1-3 instead (internal property is
+            # highest priority).
+            # This is also why we cannot shift this logic further into the
+            # EdgeCollection class, which is oblivious of NetworkArtist's
+            # internal data. In fact, one would argue this needs to be
+            # pushed outwards to deal with the wrong ordering.
+            _update_from_internal(edge_stylei, edge_df.iloc[i], kind="edge")
+
+            if cmap_fun is not None:
+                colorarray.append(edge_stylei["color"])
+                edge_stylei["color"] = cmap_fun(edge_stylei["color"])
+
+            # Tree layout determines waypoints
+            waypointsi = edge_stylei.pop("waypoints", None)
+            if waypointsi is None:
+                layout_name = self._ipx_internal_data["layout_name"]
+                if layout_name == "horizontal":
+                    waypointsi = "x0y1"
+                elif layout_name == "vertical":
+                    waypointsi = "y0y0"
+                elif layout_name == "radial":
+                    waypointsi = "r0a1"
+                else:
+                    waypointsi = "none"
+            waypoints.append(waypointsi)
+
+            # These are not the actual edges drawn, only stubs to establish
+            # the styles which are then fed into the dynamic, optimised
+            # factory (the collection) below
+            patch = make_undirected_edge_patch(
+                **edge_stylei,
+            )
+            edgepatches.append(patch)
+            adjacent_vertex_ids.append((vid1, vid2))
+
+        if "cmap" in edge_style:
+            vmin = np.min(colorarray)
+            vmax = np.max(colorarray)
+            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
+            edge_style["norm"] = norm
+
+        edge_style["waypoints"] = waypoints
+
+        # NOTE: Trees are directed is their "directed" property is True, "child", or "parent"
+        self._edges = EdgeCollection(
+            edgepatches,
+            vertex_ids=adjacent_vertex_ids,
+            vertex_collection=self._vertices,
+            labels=labels,
+            transform=self.get_offset_transform(),
+            style=edge_style,
+            directed=bool(self._ipx_internal_data["directed"]),
+        )
+        if "cmap" in edge_style:
+            self._edges.set_array(colorarray)
+
+    @_stale_wrapper
+    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??
+
+        # NOTE: looks like we have to manage the zorder ourselves
+        # this is kind of funny actually
+        children = list(self.get_children())
+        children.sort(key=lambda x: x.zorder)
+        for art in children:
+            art.draw(renderer)

iplotx/typing.py

@@ -1,41 +1,55 @@
-from typing import Union, Sequence
-from numpy import ndarray
-from pandas import DataFrame
-
-from .importing import igraph, networkx
-
-
-igraphGraph = igraph.Graph if igraph is None else None
-if networkx is not None:
-    from networkx import Graph as networkxGraph
-    from networkx import DiGraph as networkxDiGraph
-    from networkx import MultiGraph as networkxMultiGraph
-    from networkx import MultiDiGraph as networkxMultiDiGraph
-
-    networkxOmniGraph = Union[
-        networkxGraph, networkxDiGraph, networkxMultiGraph, networkxMultiDiGraph
-    ]
-else:
-    networkxOmniGraph = None
-
-if igraphGraph is not None and networkxOmniGraph is not None:
-    GraphType = Union[igraphGraph, networkxOmniGraph]
-elif igraphGraph is not None:
-    GraphType = igraphGraph
-else:
-    GraphType = networkxOmniGraph
-
-LayoutType = Union[str, Sequence[Sequence[float]], ndarray, DataFrame]
-
-if (igraph is not None) and (networkx is not None):
-    # networkx returns generators of sets, igraph has its own classes
-    # additionally, one can put list of memberships
-    GroupingType = Union[
-        Sequence[set],
-        igraph.clustering.Clustering,
-        igraph.clustering.VertexClustering,
-        igraph.clustering.Cover,
-        igraph.clustering.VertexCover,
-        Sequence[int],
-        Sequence[str],
-    ]
+"""
+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
+
+
+# NOTE: GraphType is supposed to indicate any kind of graph object that is accepted by
+# iplotx's functions, e.g. igraph.Graph or networkx.Graph and subclasses. It is not
+# quite possible to really statically type it because providers can add their own
+# types - together with protocols to process them - at runtime.
+# Nonetheless, for increased readibility we define separately-named types in this
+# module to be used throughout the codebase.
+GraphType = Any
+TreeType = Any
+
+# NOTE: The commented ones are not a mistake: they are supported but cannot be
+# statically typed if the user has no igraph installed (it's a soft dependency).
+LayoutType = Union[
+    str,
+    Sequence[Sequence[float]],
+    np.ndarray,
+    pd.DataFrame,
+    # igraph.Layout,
+]
+GroupingType = Union[
+    Sequence[set],
+    Sequence[int],
+    Sequence[str],
+    # igraph.clustering.Clustering,
+    # igraph.clustering.VertexClustering,
+    # igraph.clustering.Cover,
+    # igraph.clustering.VertexCover,
+]
+
+
+T = TypeVar("T")
+LeafProperty = Union[
+    T,
+    Sequence[T],
+    dict[Hashable, T],
+]
+
+Pair = tuple[T, T]

iplotx/utils/geometry.py

@@ -1,4 +1,4 @@
-from math import tan, atan2
+from math import atan2
 import numpy as np
 
 
@@ -21,22 +21,56 @@ def _evaluate_cubic_bezier(points, t):
     )
 
 
+def _evaluate_cubic_bezier_derivative(points, t):
+    """Evaluate the derivative of a cubic Bezier curve at t."""
+    p0, p1, p2, p3 = points
+    # (dx / dt, dy / dt) is the parametric gradient
+    # to get the angle from this, one can just atanh(dy/dt, dx/dt)
+    # This is equivalent to computing the actual bezier curve
+    # at low t, of course, which is the geometric interpretation
+    # (obviously, division by t is irrelenant)
+    return (
+        3 * p0 * (1 - t) ** 2
+        + 3 * p1 * (1 - t) * (-3 * t + 1)
+        + 3 * p2 * t * (2 - 3 * t)
+        + 3 * p3 * t**2
+    )
+
+
 def convex_hull(points):
-    """Compute the convex hull of a set of 2D points."""
-    from ..importing import igraph
+    """Compute the convex hull of a set of 2D points.
+
+    This is guaranteed to return the vertices clockwise.
 
+    (Therefore, (v[i+1] - v[i]) rotated *left* by pi/2 points *outwards* of the convex hull.)
+    """
     points = np.asarray(points)
+    if len(points) < 3:
+        return np.arange(len(points))
+
+    hull_idx = None
 
     # igraph's should be faster in 2D
-    if igraph is not None:
+    try:
+        import igraph
+
         hull_idx = igraph.convex_hull(list(points))
-    else:
-        try:
-            from scipy.spatial import ConvexHull
+    except ImportError:
+        pass
 
-            hull_idx = ConvexHull(points).vertices
-        except ImportError:
-            hull_idx = _convex_hull_Graham_scan(points)
+    # Otherwise, try scipy
+    try:
+        from scipy.spatial import ConvexHull
+
+        # NOTE: scipy guarantees counterclockwise ordering in 2D
+        # https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.ConvexHull.html
+        hull_idx = ConvexHull(points).vertices[::-1]
+    except ImportError:
+        pass
+
+    # Last resort: our own Graham scan
+    if hull_idx is None:
+        hull_idx = _convex_hull_Graham_scan(points)
 
     return hull_idx
 
@@ -45,11 +79,10 @@ def convex_hull(points):
 # Compared to that C implementation, this is a bit more vectorised and messes less with memory as usual when
 # optimising Python/numpy code
 def _convex_hull_Graham_scan(points):
-    """Compute the indices for the convex hull of a set of 2D points using Graham's scan algorithm."""
-    if len(points) < 4:
-        # NOTE: for an exact triangle, this does not guarantee chirality. Should be ok anyway
-        return np.arange(len(points))
+    """Compute the indices for the convex hull of a set of 2D points using Graham's scan algorithm.
 
+    NOTE: This works from 3 points upwards, guaranteed clockwise.
+    """
     points = np.asarray(points)
 
     # Find pivot (bottom left corner)
@@ -136,9 +169,13 @@ def _convex_hull_Graham_scan(points):
 
 
 def _compute_group_path_with_vertex_padding(
+    hull,
     points,
     transform,
     vertexpadding=10,
+    points_per_curve=30,
+    # TODO: check how dpi affects this
+    dpi=72.0,
 ):
     """Offset path for a group based on vertex padding.
 
@@ -146,82 +183,92 @@ def _compute_group_path_with_vertex_padding(
 
     # NOTE: this would look better as a cubic Bezier, but ok for now.
     """
+    # Short form
+    ppc = points_per_curve
+
+    # No padding, set degenerate path
+    if vertexpadding == 0:
+        for j, point in enumerate(hull):
+            points[ppc * j : ppc * (j + 1)] = point
+        points[-1] = points[0]
+        return points
 
     # Transform into figure coordinates
     trans = transform.transform
     trans_inv = transform.inverted().transform
     points = trans(points)
 
-    # Find the vertex centers, to recompute the offsets from scratch
-    # Independent whether this is a first call or a later draw,
-    # finding the vertex center can be done at once
-    # 0.         .->.vcenter
-    #  |         |  ^
-    #  |         |  |
-    # 1.--.2     .--.
-    # singleton group
-    s2 = 0.96
-    if len(points) == 9:
-        points[:] = 0.5 * (points[0] + points[4])
-        points[0] += np.array([0, -1]) * vertexpadding
-        points[1] += np.array([-s2, -s2]) * vertexpadding
-        points[2] += np.array([-1, 0]) * vertexpadding
-        points[3] += np.array([-s2, s2]) * vertexpadding
-        points[4] += np.array([0, 1]) * vertexpadding
-        points[5] += np.array([s2, s2]) * vertexpadding
-        points[6] += np.array([1, 0]) * vertexpadding
-        points[7] += np.array([s2, -s2]) * vertexpadding
-        points[8] += np.array([0, -1]) * vertexpadding
-    else:
-        # doublet group are a bit different from triangles+
-        if len(points) == 11:
-            # points per vertex
-            ppv = 5
-            points[:-1:ppv] = 0.5 * (points[:-1:ppv] + points[ppv - 1 : -1 : ppv])
-        else:
-            ppv = 3
-            points[:-1:ppv] = (
-                points[:-1:ppv] + points[ppv - 1 : -1 : ppv] - points[1:-1:ppv]
-            )
-        for j in range(1, ppv):
-            points[j:-1:ppv] = points[:-1:ppv]
+    # Singleton: draw a circle around it
+    if len(hull) == 1:
+
+        # NOTE: linspace is double inclusive, which covers CLOSEPOLY
+        thetas = np.linspace(
+            -np.pi,
+            np.pi,
+            len(points),
+        )
+        # NOTE: dpi scaling might need to happen here
+        perimeter = vertexpadding * np.vstack([np.cos(thetas), np.sin(thetas)]).T
+        return trans_inv(trans(hull[0]) + perimeter)
+
+    # Doublet: draw two semicircles
+    if len(hull) == 2:
+
+        # Unit vector connecting the two points
+        dv = trans(hull[0]) - trans(hull[1])
+        dv = dv / np.sqrt((dv**2).sum())
+
+        # Draw a semicircle
+        angles = np.linspace(-0.5 * np.pi, 0.5 * np.pi, 30)
+        vs = np.array([np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)])
+        vs = vs.T.reshape((len(angles), 2, 2))
+        vs = np.matmul(dv, vs)
+
+        # NOTE: dpi scaling might need to happen here
+        semicircle1 = vertexpadding * vs
+        semicircle2 = vertexpadding * np.matmul(vs, -np.diag((1, 1)))
+
+        # Put it together
+        vs1 = trans_inv(trans(hull[0]) + semicircle1)
+        vs2 = trans_inv(trans(hull[1]) + semicircle2)
+        points[:ppc] = vs1
+        points[ppc:-1] = vs2
         points[-1] = points[0]
+        return points
 
-        # Compute all shift vectors by diff, arctan2, then add 90 degrees, tan, norm
-        # This maintains chirality
-        # NOTE: the last point is just going back to the beginning, this
-        # is a quirk or how mpl's closed paths work
-
-        # Normalised diff
-        vpoints = points[:-1:ppv].copy()
-        vpoints[0] -= points[-2]
-        vpoints[1:] -= points[:-1:ppv][:-1]
-        vpoints = (vpoints.T / np.sqrt((vpoints**2).sum(axis=1))).T
-
-        # Rotate by 90 degrees
-        vpads = vpoints @ np.array([[0, 1], [-1, 0]])
-
-        # Permute diff for the end
-        vpads_perm = np.zeros_like(vpads)
-        vpads_perm[:-1] = vpads[1:]
-        vpads_perm[-1] = vpads[0]
-
-        # Shift the points
-        if ppv == 3:
-            points[:-1:ppv] += vpads * vertexpadding
-            points[1:-1:ppv] += (vpads + vpads_perm) * vertexpadding
-            points[2:-1:ppv] += vpads_perm * vertexpadding
-        else:
-            points[:-1:ppv] += vpads * vertexpadding
-            points[1:-1:ppv] += (vpads + vpoints) * vertexpadding
-            points[2:-1:ppv] += vpoints * vertexpadding
-            points[3:-1:ppv] += (vpads_perm + vpoints) * vertexpadding
-            points[4:-1:ppv] += vpads_perm * vertexpadding
+    # At least three points, i.e. a nondegenerate convex hull
+    nsides = len(hull)
+    for i, point1 in enumerate(hull):
+        point0 = hull[i - 1]
+        point2 = hull[(i + 1) % nsides]
 
-    # mpl's quirky closed-path thing
-    points[-1] = points[0]
+        # NOTE: this can be optimised by computing things once
+        # unit vector to previous point
+        dv0 = trans(point1) - trans(point0)
+        dv0 = dv0 / np.sqrt((dv0**2).sum())
 
-    # Transform back to data coordinates
-    points = trans_inv(points)
+        # unit vector to next point
+        dv2 = trans(point2) - trans(point1)
+        dv2 = dv2 / np.sqrt((dv2**2).sum())
 
+        # span the angles
+        theta0 = atan2(dv0[1], dv0[0])
+        theta2 = atan2(dv2[1], dv2[0])
+
+        # The worst that can happen is that we go exactly backwards, i.e. theta2 == theta0 + np.pi
+        # if it's more than that, we are on the inside of the convex hull due to the periodicity of atan2
+        if theta2 - theta0 > np.pi:
+            theta2 -= 2 * np.pi
+
+        # angles is from the point of view of the first vector, dv0
+        angles = np.linspace(theta0 + np.pi / 2, theta2 + np.pi / 2, ppc)
+        vs = np.array([np.cos(angles), np.sin(angles)]).T
+
+        # NOTE: dpi scaling might need to happen here
+        chunkcircle = vertexpadding * vs
+
+        vs1 = trans_inv(trans(point1) + chunkcircle)
+        points[i * ppc : (i + 1) * ppc] = vs1
+
+    points[-1] = points[0]
     return points

iplotx/utils/internal.py

@@ -0,0 +1,3 @@
+def _make_layout_columns(ndim):
+    columns = [f"_ipx_layout_{i}" for i in range(ndim)]
+    return columns