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

iplotx/edge/arrow.py

@@ -1,16 +1,171 @@
+"""
+Module for edge arrows in iplotx.
+"""
+
 import numpy as np
 import matplotlib as mpl
 from matplotlib.patches import PathPatch
 
+from ..style import (
+    get_style,
+    rotate_style,
+)
+
+
+class EdgeArrowCollection(mpl.collections.PatchCollection):
+    """Collection of arrow patches for plotting directed edgs."""
+
+    _factor = 1.0
+
+    def __init__(
+        self,
+        edge_collection,
+        *args,
+        transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
+        **kwargs,
+    ) -> None:
+        """Initialize the edge arrow collection.
+
+        Parameters:
+            edge_collection: The edge collection to which these arrows belong.
+            transform: The transform to apply to the arrows. This related to the arrow size
+                scaling, not the arrow tip position which is controlled by set_offset_transform.
+        """
+
+        self._edge_collection = edge_collection
+        self._style = get_style(".edge.arrow")
+
+        patches, sizes = self._create_artists()
+
+        if "cmap" in self._edge_collection._style:
+            kwargs["cmap"] = self._edge_collection._style["cmap"]
+            kwargs["norm"] = self._edge_collection._style["norm"]
+
+        super().__init__(
+            patches,
+            offsets=np.zeros((len(patches), 2)),
+            offset_transform=self.get_offset_transform(),
+            transform=transform,
+            match_original=True,
+            *args,
+            **kwargs,
+        )
+        self._angles = np.zeros(len(self._paths))
+
+        # Compute _transforms like in _CollectionWithScales for dpi issues
+        self.set_sizes(sizes)
+
+    def get_sizes(self):
+        """Get arrow sizes (max of width and height), not scaled by dpi."""
+        return self._sizes
+
+    def get_sizes_dpi(self):
+        """Get arrow sizes (max of width and height) scaled by dpi."""
+        return self._transforms[:, 0, 0]
+
+    def set_sizes(self, sizes, dpi=72.0):
+        """Set vertex sizes.
+
+        This rescales the current vertex symbol/path linearly, using this
+        value as the largest of width and height.
+
+        @param sizes: A sequence of vertex sizes or a single size.
+        """
+        if sizes is None:
+            self._sizes = np.array([])
+            self._transforms = np.empty((0, 3, 3))
+        else:
+            self._sizes = np.asarray(sizes)
+            self._transforms = np.zeros((len(self._sizes), 3, 3))
+            scale = self._sizes * dpi / 72.0 * self._factor
+            self._transforms[:, 0, 0] = scale
+            self._transforms[:, 1, 1] = scale
+            self._transforms[:, 2, 2] = 1.0
+        self.stale = True
+
+    get_size = get_sizes
+    set_size = set_sizes
+
+    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)
+        for child in self.get_children():
+            child.set_figure(fig)
+
+    def get_offset_transform(self):
+        """Get offset transform for the edge arrows. This sets the tip of each arrow."""
+        return self._edge_collection.get_transform()
+
+    get_size = get_sizes
+    set_size = set_sizes
+
+    def _create_artists(self):
+        style = self._style if self._style is not None else {}
+
+        patches = []
+        sizes = []
+        for i in range(len(self._edge_collection._vertex_ids)):
+            stylei = rotate_style(style, index=i)
+            if ("facecolor" not in stylei) and ("color" not in stylei):
+                stylei["facecolor"] = self._edge_collection.get_edgecolors()[i][:3]
+            if ("edgecolor" not in stylei) and ("color" not in stylei):
+                stylei["edgecolor"] = self._edge_collection.get_edgecolors()[i][:3]
+            if "alpha" not in stylei:
+                stylei["alpha"] = self._edge_collection.get_edgecolors()[i][3]
+            if "linewidth" not in stylei:
+                stylei["linewidth"] = self._edge_collection.get_linewidths()[i]
+
+            patch, size = make_arrow_patch(
+                **stylei,
+            )
+            patches.append(patch)
+            sizes.append(size)
+
+        return patches, sizes
+
+    def set_array(self, A):
+        """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."""
+        # 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)
+        has_facecolor = self._facecolors[:, 3] > 0
+        self._facecolors[has_facecolor] = colors[has_facecolor]
+
+    @mpl.artist.allow_rasterization
+    def draw(self, renderer):
+        self.set_sizes(self._sizes, self.get_figure(root=True).dpi)
+        super().draw(renderer)
+
 
 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)
 
+    # Normalise by the max size, this is taken care of in _transforms
+    # subsequently in a way that is nice to dpi scaling
+    size_max = max(width, height)
+    if size_max > 0:
+        height /= size_max
+        width /= size_max
+
     if marker == "|>":
-        codes = ["MOVETO", "LINETO", "LINETO"]
+        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, -width * 0.5], [0, 0]]),
+            np.array(
+                [
+                    [-height, width * 0.5],
+                    [-height, -width * 0.5],
+                    [0, 0],
+                    [-height, width * 0.5],
+                ]
+            ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
@@ -25,8 +180,10 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
             closed=False,
         )
     elif marker == ">>":
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
         overhang = kwargs.pop("overhang", 0.25)
-        codes = ["MOVETO", "LINETO", "LINETO", "LINETO"]
+        codes = ["MOVETO", "LINETO", "LINETO", "LINETO", "CLOSEPOLY"]
         path = mpl.path.Path(
             np.array(
                 [
@@ -34,14 +191,17 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
                     [-height, -width * 0.5],
                     [-height * (1.0 - overhang), 0],
                     [-height, width * 0.5],
+                    [0, 0],
                 ]
             ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
     elif marker == ")>":
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
         overhang = kwargs.pop("overhang", 0.25)
-        codes = ["MOVETO", "LINETO", "CURVE3", "CURVE3"]
+        codes = ["MOVETO", "LINETO", "CURVE3", "CURVE3", "CLOSEPOLY"]
         path = mpl.path.Path(
             np.array(
                 [
@@ -49,6 +209,7 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
                     [-height, -width * 0.5],
                     [-height * (1.0 - overhang), 0],
                     [-height, width * 0.5],
+                    [0, 0],
                 ]
             ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
@@ -70,8 +231,37 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=False,
         )
+    elif marker == "|":
+        kwargs["facecolor"] = "none"
+        if "color" in kwargs:
+            kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO", "LINETO"]
+        path = mpl.path.Path(
+            np.array([[-height, width * 0.5], [-height, -width * 0.5]]),
+            codes=[getattr(mpl.path.Path, x) for x in codes],
+            closed=False,
+        )
+    elif marker == "s":
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO", "LINETO", "LINETO", "LINETO", "CLOSEPOLY"]
+        path = mpl.path.Path(
+            np.array(
+                [
+                    [-height, width * 0.5],
+                    [-height, -width * 0.5],
+                    [0, -width * 0.5],
+                    [0, width * 0.5],
+                    [-height, width * 0.5],
+                ]
+            ),
+            codes=[getattr(mpl.path.Path, x) for x in codes],
+            closed=True,
+        )
     elif marker == "d":
-        codes = ["MOVETO", "LINETO", "LINETO", "LINETO"]
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO", "LINETO", "LINETO", "LINETO", "CLOSEPOLY"]
         path = mpl.path.Path(
             np.array(
                 [
@@ -79,13 +269,16 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
                     [-height, 0],
                     [-height * 0.5, -width * 0.5],
                     [0, 0],
+                    [-height * 0.5, width * 0.5],
                 ]
             ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
     elif marker == "p":
-        codes = ["MOVETO", "LINETO", "LINETO", "LINETO"]
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO", "LINETO", "LINETO", "LINETO", "CLOSEPOLY"]
         path = mpl.path.Path(
             np.array(
                 [
@@ -93,13 +286,16 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
                     [0, 0],
                     [0, -width],
                     [-height, -width],
+                    [-height, 0],
                 ]
             ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
     elif marker == "q":
-        codes = ["MOVETO", "LINETO", "LINETO", "LINETO"]
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO", "LINETO", "LINETO", "LINETO", "CLOSEPOLY"]
         path = mpl.path.Path(
             np.array(
                 [
@@ -107,16 +303,30 @@ def make_arrow_patch(marker: str = "|>", width: float = 8, **kwargs):
                     [0, 0],
                     [0, width],
                     [-height, width],
+                    [-height, 0],
+                ]
+            ),
+            codes=[getattr(mpl.path.Path, x) for x in codes],
+            closed=True,
+        )
+    elif marker == "none":
+        if "color" in kwargs:
+            kwargs["facecolor"] = kwargs["edgecolor"] = kwargs.pop("color")
+        codes = ["MOVETO"]
+        path = mpl.path.Path(
+            np.array(
+                [
+                    [0, 0],
                 ]
             ),
             codes=[getattr(mpl.path.Path, x) for x in codes],
             closed=True,
         )
+    else:
+        raise ValueError(f"Arrow marker not found: {marker}.")
 
     patch = PathPatch(
         path,
         **kwargs,
     )
-    return patch
-
-    raise KeyError(f"Unknown marker: {marker}")
+    return patch, size_max

iplotx/edge/geometry.py

@@ -0,0 +1,392 @@
+"""
+Support module with geometry- and path-related functions for edges.
+"""
+
+from typing import Optional
+from math import atan2, tan, pi
+import numpy as np
+import matplotlib as mpl
+
+from ..typing import (
+    Pair,
+)
+from .ports import _get_port_unit_vector
+
+
+def _compute_loops_per_angle(nloops, angles):
+    if len(angles) == 0:
+        return [(0, 2 * pi, nloops)]
+
+    angles_sorted_closed = list(sorted(angles))
+    angles_sorted_closed.append(angles_sorted_closed[0] + 2 * pi)
+    deltas = np.diff(angles_sorted_closed)
+
+    # Now we have the deltas and the total number of loops
+    # 1. Assign all loops to the largest wedge
+    idx_dmax = deltas.argmax()
+    if nloops == 1:
+        return [
+            (
+                angles_sorted_closed[idx_dmax],
+                angles_sorted_closed[idx_dmax + 1],
+                nloops,
+            )
+        ]
+
+    # 2. Check if any other wedges are larger than this
+    # If not, we are done (this is the algo in igraph)
+    dsplit = deltas[idx_dmax] / nloops
+    if (deltas > dsplit).sum() < 2:
+        return [
+            (
+                angles_sorted_closed[idx_dmax],
+                angles_sorted_closed[idx_dmax + 1],
+                nloops,
+            )
+        ]
+
+        # 3. Check how small the second-largest wedge would become
+    idx_dsort = np.argsort(deltas)
+    return [
+        (
+            angles_sorted_closed[idx_dmax],
+            angles_sorted_closed[idx_dmax + 1],
+            nloops - 1,
+        ),
+        (
+            angles_sorted_closed[idx_dsort[-2]],
+            angles_sorted_closed[idx_dsort[-2] + 1],
+            1,
+        ),
+    ]
+
+
+def _get_shorter_edge_coords(vpath, vsize, theta):
+    # Bound theta from -pi to pi (why is that not guaranteed?)
+    theta = (theta + pi) % (2 * pi) - pi
+
+    # Size zero vertices need no shortening
+    if vsize == 0:
+        return np.array([0, 0])
+
+    for i in range(len(vpath)):
+        v1 = vpath.vertices[i]
+        v2 = vpath.vertices[(i + 1) % len(vpath)]
+        theta1 = atan2(*((v1)[::-1]))
+        theta2 = atan2(*((v2)[::-1]))
+
+        # atan2 ranges ]-3.14, 3.14]
+        # so it can be that theta1 is -3 and theta2 is +3
+        # therefore we need two separate cases, one that cuts at pi and one at 0
+        cond1 = theta1 <= theta <= theta2
+        cond2 = (
+            (theta1 + 2 * pi) % (2 * pi)
+            <= (theta + 2 * pi) % (2 * pi)
+            <= (theta2 + 2 * pi) % (2 * pi)
+        )
+        if cond1 or cond2:
+            break
+    else:
+        raise ValueError("Angle for patch not found")
+
+    # The edge meets the patch of the vertex on the v1-v2 size,
+    # at angle theta from the center
+    mtheta = tan(theta)
+    if v2[0] == v1[0]:
+        xe = v1[0]
+    else:
+        m12 = (v2[1] - v1[1]) / (v2[0] - v1[0])
+        xe = (v1[1] - m12 * v1[0]) / (mtheta - m12)
+    ye = mtheta * xe
+    ve = np.array([xe, ye])
+    return ve * vsize
+
+
+def _fix_parallel_edges_straight(
+    paths,
+    indices,
+    indices_inv,
+    trans,
+    trans_inv,
+    offset=3,
+):
+    """Offset parallel edges along the same path."""
+    ntot = len(indices) + len(indices_inv)
+
+    # This is straight so two vertices anyway
+    # NOTE: all paths will be the same, which is why we need to offset them
+    vs, ve = trans(paths[indices[0]].vertices)
+
+    # Move orthogonal to the line
+    fracs = (vs - ve) / np.sqrt(((vs - ve) ** 2).sum()) @ np.array([[0, 1], [-1, 0]])
+
+    # NOTE: for now treat both direction the same
+    for i, idx in enumerate(indices + indices_inv):
+        # Offset the path
+        paths[idx].vertices = trans_inv(
+            trans(paths[idx].vertices) + fracs * offset * (i - ntot / 2)
+        )
+
+
+def _compute_loop_path(
+    vcoord_fig,
+    vpath,
+    vsize,
+    angle1,
+    angle2,
+    trans_inv,
+    looptension,
+):
+    # Shorten at starting angle
+    start = _get_shorter_edge_coords(vpath, vsize, angle1) + vcoord_fig
+    # Shorten at end angle
+    end = _get_shorter_edge_coords(vpath, vsize, angle2) + vcoord_fig
+
+    aux1 = (start - vcoord_fig) * looptension + vcoord_fig
+    aux2 = (end - vcoord_fig) * looptension + vcoord_fig
+
+    vertices = np.vstack(
+        [
+            start,
+            aux1,
+            aux2,
+            end,
+        ]
+    )
+    codes = ["MOVETO"] + ["CURVE4"] * 3
+
+    # Offset to place and transform to data coordinates
+    vertices = trans_inv(vertices)
+    codes = [getattr(mpl.path.Path, x) for x in codes]
+    path = mpl.path.Path(
+        vertices,
+        codes=codes,
+    )
+    return path
+
+
+def _compute_edge_path_straight(
+    vcoord_data,
+    vpath_fig,
+    vsize_fig,
+    trans,
+    trans_inv,
+    **kwargs,
+):
+
+    # Coordinates in figure (default) coords
+    vcoord_fig = trans(vcoord_data)
+
+    points = []
+
+    # Angle of the straight line
+    theta = atan2(*((vcoord_fig[1] - vcoord_fig[0])[::-1]))
+
+    # 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]
+    )
+    points.append(ve)
+
+    codes = ["MOVETO", "LINETO"]
+    path = mpl.path.Path(
+        points,
+        codes=[getattr(mpl.path.Path, x) for x in codes],
+    )
+    path.vertices = trans_inv(path.vertices)
+    return path, (theta, theta + np.pi)
+
+
+def _compute_edge_path_waypoints(
+    waypoints,
+    vcoord_data,
+    vpath_fig,
+    vsize_fig,
+    trans,
+    trans_inv,
+    layout_coordinate_system: str = "cartesian",
+    points_per_curve: int = 30,
+    **kwargs,
+):
+
+    if waypoints in ("x0y1", "y0x1"):
+        assert layout_coordinate_system == "cartesian"
+
+        # Coordinates in figure (default) coords
+        vcoord_fig = trans(vcoord_data)
+
+        if waypoints == "x0y1":
+            waypoint = np.array([vcoord_fig[0][0], vcoord_fig[1][1]])
+        else:
+            waypoint = np.array([vcoord_fig[1][0], vcoord_fig[0][1]])
+
+        # Angles of the straight lines
+        theta0 = atan2(*((waypoint - vcoord_fig[0])[::-1]))
+        theta1 = atan2(*((waypoint - vcoord_fig[1])[::-1]))
+
+        # Shorten at starting vertex
+        vs = (
+            _get_shorter_edge_coords(vpath_fig[0], vsize_fig[0], theta0) + vcoord_fig[0]
+        )
+
+        # Shorten at end vertex
+        ve = (
+            _get_shorter_edge_coords(vpath_fig[1], vsize_fig[1], theta1) + vcoord_fig[1]
+        )
+
+        points = [vs, waypoint, ve]
+        codes = ["MOVETO", "LINETO", "LINETO"]
+        angles = (theta0, theta1)
+    elif waypoints == "r0a1":
+        assert layout_coordinate_system == "polar"
+
+        r0, alpha0 = vcoord_data[0]
+        r1, alpha1 = vcoord_data[1]
+        idx_inner = np.argmin([r0, r1])
+        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)]
+        )
+        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:
+        raise NotImplementedError(
+            f"Edge shortening with waypoints not implemented yet: {waypoints}.",
+        )
+
+    path = mpl.path.Path(
+        points,
+        codes=[getattr(mpl.path.Path, x) for x in codes],
+    )
+
+    path.vertices = trans_inv(path.vertices)
+    return path, angles
+
+
+def _compute_edge_path_curved(
+    tension,
+    vcoord_data,
+    vpath_fig,
+    vsize_fig,
+    trans,
+    trans_inv,
+    ports=(None, None),
+):
+    """Shorten the edge path along a cubic Bezier between the vertex centres.
+
+    The most important part is that the derivative of the Bezier at the start
+    and end point towards the vertex centres: people notice if they do not.
+    """
+
+    # Coordinates in figure (default) coords
+    vcoord_fig = trans(vcoord_data)
+
+    dv = vcoord_fig[1] - vcoord_fig[0]
+    edge_straight_length = np.sqrt((dv**2).sum())
+
+    auxs = [None, None]
+    for i in range(2):
+        if ports[i] is not None:
+            der = _get_port_unit_vector(ports[i], trans_inv)
+            auxs[i] = der * edge_straight_length * tension + vcoord_fig[i]
+
+    # Both ports defined, just use them and hope for the best
+    # Obviously, if the user specifies ports that make no sense,
+    # this is going to be a (technically valid) mess.
+    if all(aux is not None for aux in auxs):
+        pass
+
+    # If no ports are specified (the most common case), compute
+    # the Bezier and shorten it
+    elif all(aux is None for aux in auxs):
+        # Put auxs along the way
+        auxs = np.array(
+            [
+                vcoord_fig[0] + 0.33 * dv,
+                vcoord_fig[1] - 0.33 * dv,
+            ]
+        )
+        # Right rotation from the straight edge
+        dv_rot = -0.1 * dv @ np.array([[0, 1], [-1, 0]])
+        # Shift the auxs orthogonal to the straight edge
+        auxs += dv_rot * tension
+
+    # First port is defined
+    elif (auxs[0] is not None) and (auxs[1] is None):
+        auxs[1] = auxs[0]
+
+    # Second port is defined
+    else:
+        auxs[0] = auxs[1]
+
+    vs = [None, None]
+    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]
+        )
+
+    path = {
+        "vertices": [
+            vs[0],
+            auxs[0],
+            auxs[1],
+            vs[1],
+        ],
+        "codes": ["MOVETO"] + ["CURVE4"] * 3,
+    }
+
+    path = mpl.path.Path(
+        path["vertices"],
+        codes=[getattr(mpl.path.Path, x) for x in path["codes"]],
+    )
+
+    # Return to data transform
+    path.vertices = trans_inv(path.vertices)
+    return path, tuple(thetas)
+
+
+def _compute_edge_path(
+    *args,
+    tension: float = 0,
+    waypoints: str = "none",
+    ports: Pair[Optional[str]] = (None, None),
+    layout_coordinate_system: str = "cartesian",
+    **kwargs,
+):
+    """Compute the edge path in a few different ways."""
+    if (waypoints != "none") and (tension != 0):
+        raise ValueError("Waypoints not supported for curved edges.")
+
+    if waypoints != "none":
+        return _compute_edge_path_waypoints(
+            waypoints,
+            *args,
+            layout_coordinate_system=layout_coordinate_system,
+            **kwargs,
+        )
+
+    if tension == 0:
+        return _compute_edge_path_straight(*args, **kwargs)
+
+    return _compute_edge_path_curved(
+        tension,
+        *args,
+        ports=ports,
+        **kwargs,
+    )