iplotx 0.3.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

iplotx/__init__.py

@@ -11,6 +11,9 @@ from .plotting import (
     network,
     tree,
 )
+import iplotx.artists as artists
+import iplotx.style as style
+
 
 # Shortcut to iplotx.plotting.network
 plot = network
@@ -19,5 +22,7 @@ __all__ = [
     "network",
     "tree",
     "plot",
+    "artists",
+    "style",
     "__version__",
 ]

iplotx/artists.py

@@ -0,0 +1,24 @@
+"""
+All artists defined in iplotx.
+"""
+
+from .network import NetworkArtist
+from .tree import TreeArtist
+from .vertex import VertexCollection
+from .edge import EdgeCollection
+from .label import LabelCollection
+from .edge.arrow import EdgeArrowCollection
+from .edge.leaf import LeafEdgeCollection
+from .cascades import CascadeCollection
+
+
+___all__ = (
+    NetworkArtist,
+    TreeArtist,
+    VertexCollection,
+    EdgeCollection,
+    LeafEdgeCollection,
+    LabelCollection,
+    EdgeArrowCollection,
+    CascadeCollection,
+)

iplotx/cascades.py

@@ -2,6 +2,7 @@ from typing import (
     Any,
     Optional,
 )
+import warnings
 import numpy as np
 import pandas as pd
 
@@ -38,7 +39,7 @@ class CascadeCollection(mpl.collections.PatchCollection):
 
         # NOTE: there is a weird bug in pandas when using generic Hashable-s
         # with .loc. Seems like doing .T[...] works for individual index
-        # elements only though
+        # elements only though (i.e. using __getitem__ a la dict)
         def get_node_coords(node):
             return layout.T[node].values
 
@@ -54,9 +55,19 @@ class CascadeCollection(mpl.collections.PatchCollection):
         # These patches need at least a facecolor (usually) or an edgecolor
         # so it's safe to make a list from these
         nodes_unordered = set()
-        for prop in ("facecolor", "edgecolor"):
+        for prop in ("facecolor", "edgecolor", "linewidth", "linestyle"):
             if prop in style:
-                nodes_unordered |= set(style[prop].keys())
+                value = style[prop]
+                if isinstance(value, dict):
+                    nodes_unordered |= set(value.keys())
+
+        if len(nodes_unordered) == 0:
+            warnings.warn(
+                "No nodes found in the style for the cascading patches. "
+                "Please provide a style with at least one dict-like "
+                "specification among the following properties: 'facecolor', "
+                "'edgecolor', 'color', 'linewidth', or 'linestyle'.",
+            )
 
         # Draw the patches from the closest to the root (earlier drawing)
         # to the closer to the leaves (later drawing).
@@ -70,32 +81,15 @@ class CascadeCollection(mpl.collections.PatchCollection):
                 f"Cascading patches not implemented for layout: {layout_name}.",
             )
 
-        nleaves = sum(1 for leaf in provider(tree).get_leaves())
-        extend_mode = style.get("extend", False)
-        if extend_mode and (extend_mode != "leaf_labels"):
-            if layout_name == "horizontal":
-                if orientation == "right":
-                    maxdepth = layout.values[:, 0].max()
-                else:
-                    maxdepth = layout.values[:, 0].min()
-            elif layout_name == "vertical":
-                if orientation == "descending":
-                    maxdepth = layout.values[:, 1].min()
-                else:
-                    maxdepth = layout.values[:, 1].max()
-            elif layout_name == "radial":
-                # layout values are: r, theta
-                maxdepth = layout.values[:, 0].max()
         self._maxdepth = maxdepth
 
         cascading_patches = []
+        nleaves = sum(1 for leaf in provider(tree).get_leaves())
         for node in drawing_order:
             stylei = rotate_style(style, key=node)
             stylei.pop("extend", None)
             # Default alpha is 0.5 for simple colors
-            if isinstance(stylei.get("facecolor", None), str) and (
-                "alpha" not in stylei
-            ):
+            if isinstance(stylei.get("facecolor", None), str) and ("alpha" not in stylei):
                 stylei["alpha"] = 0.5
 
             provider_node = provider(node)
@@ -137,9 +131,7 @@ class CascadeCollection(mpl.collections.PatchCollection):
                 rmax = maxdepth if extend else leaves_coords[:, 0].max()
                 thetamin = leaves_coords[:, 1].min() - 0.5 * dtheta
                 thetamax = leaves_coords[:, 1].max() + 0.5 * dtheta
-                thetas = np.linspace(
-                    thetamin, thetamax, max(30, (thetamax - thetamin) // 3)
-                )
+                thetas = np.linspace(thetamin, thetamax, max(30, (thetamax - thetamin) // 3))
                 xs = list(rmin * np.cos(thetas)) + list(rmax * np.cos(thetas[::-1]))
                 ys = list(rmin * np.sin(thetas)) + list(rmax * np.sin(thetas[::-1]))
                 points = list(zip(xs, ys))
@@ -200,23 +192,22 @@ class CascadeCollection(mpl.collections.PatchCollection):
             for path in self.get_paths():
                 # Old radii
                 r2old = np.linalg.norm(path.vertices[-2])
-                path.vertices[(len(path.vertices) - 1) // 2 :] *= (
-                    self.get_maxdepth() / r2old
-                )
+                # Update the outer part of the wedge patch
+                path.vertices[(len(path.vertices) - 1) // 2 :] *= self.get_maxdepth() / r2old
             return
 
         if (layout_name, orientation) == ("horizontal", "right"):
             for path in self.get_paths():
                 path.vertices[[1, 2], 0] = self.get_maxdepth()
-        elif (layout_name, orientation) == ("horizontal", "right"):
+        elif (layout_name, orientation) == ("horizontal", "left"):
             for path in self.get_paths():
                 path.vertices[[0, 3], 0] = self.get_maxdepth()
         elif (layout_name, orientation) == ("vertical", "descending"):
             for path in self.get_paths():
-                path.vertices[[1, 2], 1] = self.get_maxdepth()
+                path.vertices[[0, 1], 1] = self.get_maxdepth()
         elif (layout_name, orientation) == ("vertical", "ascending"):
             for path in self.get_paths():
-                path.vertices[[0, 3], 1] = self.get_maxdepth()
+                path.vertices[[2, 3], 1] = self.get_maxdepth()
         else:
             raise ValueError(
                 f"Layout name and orientation not supported: {layout_name}, {orientation}."

iplotx/edge/__init__.py

@@ -121,8 +121,47 @@ class EdgeCollection(mpl.collections.PatchCollection):
                 transform=transform,
             )
 
+        if "split" in self._style:
+            self._add_subedges(
+                len(patches),
+                self._style["split"],
+            )
+
+    def _add_subedges(
+        self,
+        nedges,
+        style,
+    ):
+        """Add subedges to shadow the current edges."""
+        segments = [np.zeros((2, 2)) for i in range(nedges)]
+        kwargs = {
+            "linewidths": [],
+            "edgecolors": [],
+            "linestyles": [],
+        }
+        for i in range(nedges):
+            vids = self._vertex_ids[i]
+            stylei = rotate_style(style, index=i, key=vids, key2=vids[-1])
+            for key, values in kwargs.items():
+                # iplotx uses singular style properties
+                key = key.rstrip("s")
+                # "color" has higher priority than "edgecolor"
+                if (key == "edgecolor") and ("color" in stylei):
+                    val = stylei["color"]
+                else:
+                    val = stylei.get(key.rstrip("s"), getattr(self, f"get_{key}")()[i])
+                values.append(val)
+
+        self._subedges = mpl.collections.LineCollection(
+            segments,
+            transform=self.get_transform(),
+            **kwargs,
+        )
+
     def get_children(self) -> tuple:
         children = []
+        if hasattr(self, "_subedges"):
+            children.append(self._subedges)
         if hasattr(self, "_arrows"):
             children.append(self._arrows)
         if hasattr(self, "_label_collection"):
@@ -209,7 +248,7 @@ class EdgeCollection(mpl.collections.PatchCollection):
         index = pd.Series(
             np.arange(len(index)),
             index=index,
-        )
+        ).to_dict()
 
         voffsets = []
         vpaths = []
@@ -262,11 +301,13 @@ class EdgeCollection(mpl.collections.PatchCollection):
         trans_inv = transform.inverted().transform
 
         # 1. Make a list of vertices with loops, and store them for later
-        loop_vertex_dict = defaultdict(lambda: dict(indices=[], edge_angles=[]))
-        for i, (v1, v2) in enumerate(vids):
-            # Postpone loops (step 3)
-            if v1 == v2:
-                loop_vertex_dict[v1]["indices"].append(i)
+        # NOTE: vinfo["loops"] can be False when we want no loops (e.g. leaf edges)
+        if vinfo.get("loops", True):
+            loop_vertex_dict = defaultdict(lambda: dict(indices=[], edge_angles=[]))
+            for i, (v1, v2) in enumerate(vids):
+                # Postpone loops (step 3)
+                if v1 == v2:
+                    loop_vertex_dict[v1]["indices"].append(i)
 
         # 2. Make paths for non-loop edges
         # NOTE: keep track of parallel edges to offset them
@@ -274,7 +315,7 @@ class EdgeCollection(mpl.collections.PatchCollection):
         paths = []
         for i, (v1, v2) in enumerate(vids):
             # Postpone loops (step 3)
-            if v1 == v2:
+            if vinfo.get("loops", True) and (v1 == v2):
                 paths.append(None)
                 continue
 
@@ -296,10 +337,20 @@ class EdgeCollection(mpl.collections.PatchCollection):
                 tension = 0
                 ports = None
 
+            # False is a synonym for "none"
             waypoints = edge_stylei.get("waypoints", "none")
+            if waypoints is False or waypoints is np.False_:
+                waypoints = "none"
+            elif waypoints is True or waypoints is np.True_:
+                raise ValueError(
+                    "Could not determine automatically type of edge waypoints.",
+                )
             if waypoints != "none":
                 ports = edge_stylei.get("ports", (None, None))
 
+            if not isinstance(waypoints, str):
+                __import__("ipdb").set_trace()
+
             # Compute actual edge path
             path, angles = _compute_edge_path(
                 vcoord_data,
@@ -329,11 +380,23 @@ class EdgeCollection(mpl.collections.PatchCollection):
             if (offset != 0).any():
                 path.vertices[:] = trans_inv(trans(path.vertices) + offset)
 
+            # If splitting is active, split the path here, shedding off the last straight
+            # segment but only if waypoints were used
+            if hasattr(self, "_subedges") and waypoints != "none":
+                # NOTE: we are already in the middle of a redraw, so we can happily avoid
+                # causing stale of the subedges. They are already scheduled to be redrawn
+                # at the end of this function.
+                self._subedges._paths[i].vertices[:] = path.vertices[-2:].copy()
+                # NOTE: instead of shortening the path, we just make the last bit invisible
+                # that makes it easier on memory management etc.
+                path.vertices[-1] = path.vertices[-2]
+
             # Collect angles for this vertex, to be used for loops plotting below
-            if v1 in loop_vertex_dict:
-                loop_vertex_dict[v1]["edge_angles"].append(angles[0])
-            if v2 in loop_vertex_dict:
-                loop_vertex_dict[v2]["edge_angles"].append(angles[1])
+            if vinfo.get("loops", True):
+                if v1 in loop_vertex_dict:
+                    loop_vertex_dict[v1]["edge_angles"].append(angles[0])
+                if v2 in loop_vertex_dict:
+                    loop_vertex_dict[v2]["edge_angles"].append(angles[1])
 
             # Add the path for this non-loop edge
             paths.append(path)
@@ -360,41 +423,45 @@ class EdgeCollection(mpl.collections.PatchCollection):
                     )
 
         # 3. Deal with loops at the end
-        for vid, ldict in loop_vertex_dict.items():
-            vpath = vpaths[ldict["indices"][0]][0]
-            vsize = vsizes[ldict["indices"][0]][0]
-            vcoord_fig = trans(vcenters[ldict["indices"][0]][0])
-            nloops = len(ldict["indices"])
-            edge_angles = ldict["edge_angles"]
-
-            # The space between the existing angles is where we can fit the loops
-            # One loop we can fit in the largest wedge, multiple loops we need
-            nloops_per_angle = _compute_loops_per_angle(nloops, edge_angles)
-
-            idx = 0
-            for theta1, theta2, nloops in nloops_per_angle:
-                # Angular size of each loop in this wedge
-                delta = (theta2 - theta1) / nloops
-
-                # Iterate over individual loops
-                for j in range(nloops):
-                    thetaj1 = theta1 + j * delta + max(delta - loopmaxangle, 0) / 2
-                    thetaj2 = thetaj1 + min(delta, loopmaxangle)
-
-                    # Get the path for this loop
-                    path = _compute_loop_path(
-                        vcoord_fig,
-                        vpath,
-                        vsize,
-                        thetaj1,
-                        thetaj2,
-                        trans_inv,
-                        looptension=self._style.get("looptension", 2.5),
-                    )
-                    paths[ldict["indices"][idx]] = path
-                    idx += 1
+        if vinfo.get("loops", True):
+            for vid, ldict in loop_vertex_dict.items():
+                vpath = vpaths[ldict["indices"][0]][0]
+                vsize = vsizes[ldict["indices"][0]][0]
+                vcoord_fig = trans(vcenters[ldict["indices"][0]][0])
+                nloops = len(ldict["indices"])
+                edge_angles = ldict["edge_angles"]
+
+                # The space between the existing angles is where we can fit the loops
+                # One loop we can fit in the largest wedge, multiple loops we need
+                nloops_per_angle = _compute_loops_per_angle(nloops, edge_angles)
+
+                idx = 0
+                for theta1, theta2, nloops in nloops_per_angle:
+                    # Angular size of each loop in this wedge
+                    delta = (theta2 - theta1) / nloops
+
+                    # Iterate over individual loops
+                    for j in range(nloops):
+                        thetaj1 = theta1 + j * delta + max(delta - loopmaxangle, 0) / 2
+                        thetaj2 = thetaj1 + min(delta, loopmaxangle)
+
+                        # Get the path for this loop
+                        path = _compute_loop_path(
+                            vcoord_fig,
+                            vpath,
+                            vsize,
+                            thetaj1,
+                            thetaj2,
+                            trans_inv,
+                            looptension=self._style.get("looptension", 2.5),
+                        )
+                        paths[ldict["indices"][idx]] = path
+                        idx += 1
 
         self._paths = paths
+        # FIXME:??
+        # if hasattr(self, "_subedges"):
+        #    self._subedges.stale = True
 
     def _update_labels(self):
         if self._labels is None:
@@ -459,13 +526,13 @@ class EdgeCollection(mpl.collections.PatchCollection):
         if not self.get_visible():
             return
 
+        # This includes the subedges if present
         self._update_paths()
         # This sets the arrow offsets
         self._update_children()
 
         super().draw(renderer)
         for child in self.get_children():
-            # This sets the arrow sizes with dpi scaling
             child.draw(renderer)
 
     def get_ports(self) -> Optional[LeafProperty[Pair[Optional[str]]]]:
@@ -485,7 +552,8 @@ class EdgeCollection(mpl.collections.PatchCollection):
                 edge end.
         """
         if ports is None:
-            del self._style["ports"]
+            if "ports" in self._style:
+                del self._style["ports"]
         else:
             self._style["ports"] = ports
         self.stale = True
@@ -519,7 +587,8 @@ class EdgeCollection(mpl.collections.PatchCollection):
 
         """
         if tension is None:
-            del self._style["tension"]
+            if "tension" in self._style:
+                del self._style["tension"]
         else:
             self._style["tension"] = tension
         self.stale = True
@@ -579,7 +648,8 @@ class EdgeCollection(mpl.collections.PatchCollection):
             looptension: The tension to use for loops. If None, the default is 2.5.
         """
         if looptension is None:
-            del self._style["looptension"]
+            if "looptension" in self._style:
+                del self._style["looptension"]
         else:
             self._style["looptension"] = looptension
         self.stale = True
@@ -599,7 +669,8 @@ class EdgeCollection(mpl.collections.PatchCollection):
             offset: The offset in points for parallel straight edges. If None, the default is 3.
         """
         if offset is None:
-            del self._style["offset"]
+            if "offset" in self._style:
+                del self._style["offset"]
         else:
             self._style["offset"] = offset
         self.stale = True
@@ -629,6 +700,8 @@ def make_stub_patch(**kwargs):
         "offset",
         "paralleloffset",
         "cmap",
+        "norm",
+        "split",
     ]
     for prop in forbidden_props:
         if prop in kwargs:

iplotx/edge/arrow.py

@@ -2,6 +2,8 @@
 Module for edge arrows in iplotx.
 """
 
+from typing import Never
+
 import numpy as np
 import matplotlib as mpl
 from matplotlib.patches import PathPatch
@@ -124,12 +126,17 @@ class EdgeArrowCollection(mpl.collections.PatchCollection):
 
         return patches, sizes
 
-    def set_array(self, A):
+    def set_array(self, A: np.ndarray) -> Never:
         """Set the array for cmap/norm coloring, but keep the facecolors as set (usually 'none')."""
         raise ValueError("Setting an array for arrows directly is not supported.")
 
-    def set_colors(self, colors):
-        """Set arrow colors (edge and/or face) based on a colormap."""
+    def set_colors(self, colors: np.ndarray) -> None:
+        """Set arrow colors (edge and/or face) based on a colormap.
+
+        Parameters:
+            colors: Color array to apply. This must be an Nx3 or Nx4 vector of RGB or RGBA colors.
+                This function will NOT attempt to convert other color descriptions to RGB/RGBA.
+        """
         # NOTE: facecolors is always an array because we come from patches
         # It can have zero alpha (i.e. if we choose "none", or a hollow marker)
         self.set_edgecolor(colors)
@@ -145,6 +152,8 @@ class EdgeArrowCollection(mpl.collections.PatchCollection):
 def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
     """Make a patch of the given marker shape and size."""
     height = kwargs.pop("height", width * 1.3)
+    if height == "width":
+        height = width
 
     # Normalise by the max size, this is taken care of in _transforms
     # subsequently in a way that is nice to dpi scaling
@@ -169,6 +178,38 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
+    elif marker == "|\\":
+        codes = ["MOVETO", "LINETO", "LINETO", "CLOSEPOLY"]
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        path = mpl.path.Path(
+            np.array(
+                [
+                    [-height, width * 0.5],
+                    [-height, 0],
+                    [0, 0],
+                    [-height, width * 0.5],
+                ]
+            ),
+            codes=[getattr(mpl.path.Path, x) for x in codes],
+            closed=True,
+        )
+    elif marker == "|/":
+        codes = ["MOVETO", "LINETO", "LINETO", "CLOSEPOLY"]
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        path = mpl.path.Path(
+            np.array(
+                [
+                    [-height, 0],
+                    [-height, -width * 0.5],
+                    [0, 0],
+                    [-height, 0],
+                ]
+            ),
+            codes=[getattr(mpl.path.Path, x) for x in codes],
+            closed=True,
+        )
     elif marker == ">":
         kwargs["facecolor"] = "none"
         if "color" in kwargs:

iplotx/edge/geometry.py

@@ -96,6 +96,7 @@ def _get_shorter_edge_coords(vpath, vsize, theta):
         xe = v1[0]
     else:
         m12 = (v2[1] - v1[1]) / (v2[0] - v1[0])
+        print(m12, mtheta)
         xe = (v1[1] - m12 * v1[0]) / (mtheta - m12)
     ye = mtheta * xe
     ve = np.array([xe, ye])
@@ -171,25 +172,41 @@ def _compute_edge_path_straight(
     vsize_fig,
     trans,
     trans_inv,
+    layout_coordinate_system: str = "cartesian",
     **kwargs,
 ):
+    if layout_coordinate_system not in ("cartesian", "polar"):
+        raise ValueError(
+            f"Layout coordinate system not supported for straight edges: {layout_coordinate_system}.",
+        )
+
+    if layout_coordinate_system == "polar":
+        r0, theta0 = vcoord_data[0]
+        r1, theta1 = vcoord_data[1]
+        vcoord_data_cart = np.array(
+            [
+                [r0 * np.cos(theta0), r0 * np.sin(theta0)],
+                [r1 * np.cos(theta1), r1 * np.sin(theta1)],
+            ]
+        )
+    else:
+        vcoord_data_cart = vcoord_data
 
     # Coordinates in figure (default) coords
-    vcoord_fig = trans(vcoord_data)
+    vcoord_fig = trans(vcoord_data_cart)
 
     points = []
 
     # Angle of the straight line
     theta = atan2(*((vcoord_fig[1] - vcoord_fig[0])[::-1]))
+    print(vcoord_data_cart, vcoord_fig, theta)
 
     # Shorten at starting vertex
     vs = _get_shorter_edge_coords(vpath_fig[0], vsize_fig[0], theta) + vcoord_fig[0]
     points.append(vs)
 
     # Shorten at end vertex
-    ve = (
-        _get_shorter_edge_coords(vpath_fig[1], vsize_fig[1], theta + pi) + vcoord_fig[1]
-    )
+    ve = _get_shorter_edge_coords(vpath_fig[1], vsize_fig[1], theta + pi) + vcoord_fig[1]
     points.append(ve)
 
     codes = ["MOVETO", "LINETO"]
@@ -213,7 +230,6 @@ def _compute_edge_path_waypoints(
     ports: Pair[Optional[str]] = (None, None),
     **kwargs,
 ):
-
     if waypoints in ("x0y1", "y0x1"):
         assert layout_coordinate_system == "cartesian"
 
@@ -236,8 +252,7 @@ def _compute_edge_path_waypoints(
 
             # Shorten at vertex border
             vshorts[i] = (
-                _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], thetas[i])
-                + vcoord_fig[i]
+                _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], thetas[i]) + vcoord_fig[i]
             )
 
         # Shorten waypoints to keep the angles right
@@ -285,10 +300,7 @@ def _compute_edge_path_waypoints(
                 theta = atan2(*(_get_port_unit_vector(ports[i], trans_inv)[::-1]))
 
             # Shorten at vertex border
-            vshort = (
-                _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], theta)
-                + vcoord_fig[i]
-            )
+            vshort = _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], theta) + vcoord_fig[i]
             thetas.append(theta)
             vshorts.append(vshort)
 
@@ -305,16 +317,12 @@ def _compute_edge_path_waypoints(
         idx_outer = 1 - idx_inner
         alpha_outer = [alpha0, alpha1][idx_outer]
 
-        # FIXME: this is aware of chirality as stored by the layout function
         betas = np.linspace(alpha0, alpha1, points_per_curve)
         waypoints = [r0, r1][idx_inner] * np.vstack([np.cos(betas), np.sin(betas)]).T
-        endpoint = [r0, r1][idx_outer] * np.array(
-            [np.cos(alpha_outer), np.sin(alpha_outer)]
-        )
+        endpoint = [r0, r1][idx_outer] * np.array([np.cos(alpha_outer), np.sin(alpha_outer)])
         points = np.array(list(waypoints) + [endpoint])
         points = trans(points)
         codes = ["MOVETO"] + ["LINETO"] * len(waypoints)
-        # FIXME: same as previus comment
         angles = (alpha0 + pi / 2, alpha1)
 
     else:
@@ -391,10 +399,7 @@ def _compute_edge_path_curved(
     thetas = [None, None]
     for i in range(2):
         thetas[i] = atan2(*((auxs[i] - vcoord_fig[i])[::-1]))
-        vs[i] = (
-            _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], thetas[i])
-            + vcoord_fig[i]
-        )
+        vs[i] = _get_shorter_edge_coords(vpath_fig[i], vsize_fig[i], thetas[i]) + vcoord_fig[i]
 
     path = {
         "vertices": [
@@ -438,7 +443,11 @@ def _compute_edge_path(
         )
 
     if tension == 0:
-        return _compute_edge_path_straight(*args, **kwargs)
+        return _compute_edge_path_straight(
+            *args,
+            layout_coordinate_system=layout_coordinate_system,
+            **kwargs,
+        )
 
     return _compute_edge_path_curved(
         tension,