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

iplotx/network.py

@@ -1,50 +1,39 @@
-from typing import Union, Sequence
-import warnings
+from typing import Optional, Sequence
 import numpy as np
 import pandas as pd
 import matplotlib as mpl
-from matplotlib.transforms import Affine2D
 
 from .typing import (
     GraphType,
     LayoutType,
 )
-from .styles import (
+from .style import (
     get_style,
     rotate_style,
 )
-from .heuristics import (
-    network_library,
-    normalise_layout,
-    detect_directedness,
-)
 from .utils.matplotlib import (
     _stale_wrapper,
     _forwarder,
-    _get_label_width_height,
+    _build_cmap_fun,
+)
+from .ingest import (
+    ingest_network_data,
 )
 from .vertex import (
     VertexCollection,
-    make_patch as make_vertex_patch,
 )
-from .edge.undirected import (
-    UndirectedEdgeCollection,
+from .edge import (
+    EdgeCollection,
     make_stub_patch as make_undirected_edge_patch,
 )
-from .edge.directed import (
-    DirectedEdgeCollection,
-    make_arrow_patch,
-)
 
 
 @_forwarder(
     (
         "set_clip_path",
         "set_clip_box",
-        "set_transform",
         "set_snap",
         "set_sketch_params",
-        "set_figure",
         "set_animated",
         "set_picker",
     )
@@ -53,50 +42,62 @@ class NetworkArtist(mpl.artist.Artist):
     def __init__(
         self,
         network: GraphType,
-        layout: LayoutType = None,
-        vertex_labels: Union[None, list, dict, pd.Series] = None,
-        edge_labels: Union[None, Sequence] = None,
+        layout: Optional[LayoutType] = None,
+        vertex_labels: Optional[list | dict | pd.Series] = None,
+        edge_labels: Optional[Sequence] = None,
+        transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
+        offset_transform: Optional[mpl.transforms.Transform] = None,
     ):
         """Network container artist that groups all plotting elements.
 
         Parameters:
-            network (networkx.Graph or igraph.Graph): The network to plot.
-            layout (array-like): The layout of the network. If None, this function will attempt to
+            network: The network to plot.
+            layout: The layout of the network. If None, this function will attempt to
                 infer the layout from the network metadata, using heuristics. If that fails, an
                 exception will be raised.
-            vertex_labels (list, dict, or pandas.Series): The labels for the vertices. If None, no vertex labels
+            vertex_labels: The labels for the vertices. If None, no vertex labels
                 will be drawn. If a list, the labels are taken from the list. If a dict, the keys
                 should be the vertex IDs and the values should be the labels.
-            elge_labels (sequence): The labels for the edges. If None, no edge labels will be drawn.
-        """
-        super().__init__()
+            edge_labels: The labels for the edges. If None, no edge labels will be drawn.
 
+        """
         self.network = network
-        self._ipx_internal_data = _create_internal_data(
+        self._ipx_internal_data = ingest_network_data(
             network,
             layout,
             vertex_labels=vertex_labels,
             edge_labels=edge_labels,
         )
-        self._clear_state()
 
-    def _clear_state(self):
-        self._vertices = None
-        self._edges = None
-        self._vertex_labels = []
-        self._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):
-        artists = []
-        # Collect edges first. This way vertices are on top of edges,
-        # since vertices are drawn later. That is what most people expect.
-        if self._edges is not None:
-            artists.append(self._edges)
-        if self._vertices is not None:
-            artists.append(self._vertices)
-        artists.extend(self._edge_labels)
-        artists.extend(self._vertex_labels)
-        return tuple(artists)
+        return (self._vertices, self._edges)
+
+    def set_figure(self, fig):
+        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_vertices(self):
         """Get VertexCollection artist."""
@@ -108,13 +109,13 @@ class NetworkArtist(mpl.artist.Artist):
 
     def get_vertex_labels(self):
         """Get list of vertex label artists."""
-        return self._vertex_labels
+        return self._vertices.get_labels()
 
     def get_edge_labels(self):
         """Get list of edge label artists."""
-        return self._edge_labels
+        return self._edges.get_labels()
 
-    def get_datalim(self, transData, pad=0.05):
+    def get_datalim(self, transData, pad=0.15):
         """Get limits on x/y axes based on the graph layout data.
 
         Parameters:
@@ -122,228 +123,101 @@ class NetworkArtist(mpl.artist.Artist):
             pad (float): Padding to add to the limits. Default is 0.05.
                 Units are a fraction of total axis range before padding.
         """
-        # FIXME: transData works here, but it's probably kind of broken in general
-        import numpy as np
-
-        layout_columns = [
-            f"_ipx_layout_{i}" for i in range(self._ipx_internal_data["ndim"])
-        ]
-        layout = self._ipx_internal_data["vertex_df"][layout_columns].values
+        layout = self.get_layout().values
 
         if len(layout) == 0:
-            mins = np.array([0, 0])
-            maxs = np.array([1, 1])
-            return (mins, maxs)
+            return mpl.transforms.Bbox([[0, 0], [1, 1]])
 
-        # Use the layout as a base, and expand using bboxes from other artists
-        mins = np.min(layout, axis=0).astype(float)
-        maxs = np.max(layout, axis=0).astype(float)
-
-        # NOTE: unlike other Collections, the vertices are basically a
-        # PatchCollection with an offset transform using transData. Therefore,
-        # care should be taken if one wants to include it here
         if self._vertices is not None:
-            trans = transData.transform
-            trans_inv = transData.inverted().transform
-            verts = self._vertices
-            for path, offset in zip(verts.get_paths(), verts._offsets):
-                bbox = path.get_extents()
-                mins = np.minimum(mins, trans_inv(bbox.min + trans(offset)))
-                maxs = np.maximum(maxs, trans_inv(bbox.max + trans(offset)))
+            bbox = self._vertices.get_datalim(transData)
 
         if self._edges is not None:
-            for path in self._edges.get_paths():
-                bbox = path.get_extents()
-                mins = np.minimum(mins, bbox.min)
-                maxs = np.maximum(maxs, bbox.max)
-
-        if hasattr(self, "_groups") and self._groups is not None:
-            for path in self._groups.get_paths():
-                bbox = path.get_extents()
-                mins = np.minimum(mins, bbox.min)
-                maxs = np.maximum(maxs, bbox.max)
-
-        # 5% padding, on each side
-        pad = (maxs - mins) * pad
-        mins -= pad
-        maxs += pad
+            edge_bbox = self._edges.get_datalim(transData)
+            bbox = mpl.transforms.Bbox.union([bbox, edge_bbox])
 
-        return mpl.transforms.Bbox([mins, maxs])
+        bbox = bbox.expanded(sw=(1.0 + pad), sh=(1.0 + pad))
+        return bbox
 
-    def _add_vertices(self):
-        """Draw the vertices"""
-        vertex_style = get_style(".vertex")
+    def autoscale_view(self, tight=False):
+        """Recompute data limits from this artist and set autoscale based on them."""
+        bbox = self.get_datalim(self.axes.transData)
+        self.axes.update_datalim(bbox)
+        self.axes.autoscale_view(tight=tight)
 
+    def get_layout(self):
         layout_columns = [
             f"_ipx_layout_{i}" for i in range(self._ipx_internal_data["ndim"])
         ]
         vertex_layout_df = self._ipx_internal_data["vertex_df"][layout_columns]
-        if vertex_style.get("size") == "label":
-            if "label" not in self._ipx_internal_data["vertex_df"].columns:
-                warnings.warn(
-                    "No labels found, cannot resize vertices based on labels."
-                )
-                vertex_style["size"] = get_style("default.vertex")["size"]
-            else:
-                vertex_labels = self._ipx_internal_data["vertex_df"]["label"]
-
-        # FIXME:: this would be better off in the VertexCollection itself, like we do for groups
-        offsets = []
-        patches = []
-        for i, (vid, row) in enumerate(vertex_layout_df.iterrows()):
-            # Centre of the vertex
-            offsets.append(list(row[layout_columns].values))
-
-            if vertex_style.get("size") == "label":
-                # NOTE: it's ok to overwrite the dict here
-                vertex_style["size"] = _get_label_width_height(
-                    vertex_labels[vid], **vertex_style.get("label", {})
-                )
-
-            vertex_stylei = rotate_style(vertex_style, index=i, id=vid)
-
-            # Shape of the vertex (Patch)
-            art = make_vertex_patch(**vertex_stylei)
-            patches.append(art)
-
-        art = VertexCollection(
-            patches,
-            offsets=offsets if offsets else None,
-            offset_transform=self.axes.transData,
-            transform=Affine2D(),
-            match_original=True,
-        )
-        self._vertices = art
-
-    def _add_vertex_labels(self):
-        """Draw vertex labels."""
-        label_style = get_style(".vertex.label")
-        forbidden_props = ["hpadding", "vpadding"]
-        for prop in forbidden_props:
-            if prop in label_style:
-                del label_style[prop]
-
-        texts = []
-        vertex_labels = self._ipx_internal_data["vertex_df"]["label"]
-        for offset, label in zip(self._vertices._offsets, vertex_labels):
-            text = mpl.text.Text(
-                offset[0],
-                offset[1],
-                label,
-                transform=self.axes.transData,
-                **label_style,
-            )
-            texts.append(text)
-        self._vertex_labels = texts
+        return vertex_layout_df
 
-    def _add_edges(self):
-        """Draw the edges."""
-        if "labels" in self._ipx_internal_data["edge_df"].columns:
-            labels = self._ipx_internal_data["edge_df"]["labels"]
+    def _get_label_series(self, kind):
+        if "label" in self._ipx_internal_data[f"{kind}_df"].columns:
+            return self._ipx_internal_data[f"{kind}_df"]["label"]
         else:
-            labels = None
-
-        if self._ipx_internal_data["directed"]:
-            return self._add_directed_edges(labels=labels)
-        return self._add_undirected_edges(labels=labels)
-
-    def _add_directed_edges(self, labels=None):
-        """Draw directed edges."""
-        edge_style = get_style(".edge")
-        arrow_style = get_style(".arrow")
-
-        layout_columns = [
-            f"_ipx_layout_{i}" for i in range(self._ipx_internal_data["ndim"])
-        ]
-        vertex_layout_df = self._ipx_internal_data["vertex_df"][layout_columns]
-        edge_df = self._ipx_internal_data["edge_df"].set_index(
-            ["_ipx_source", "_ipx_target"]
-        )
+            return None
 
-        # This contains the patches for vertices, for edge shortening and such
-        vertex_paths = self._vertices._paths
-        vertex_indices = pd.Series(
-            np.arange(len(vertex_layout_df)), index=vertex_layout_df.index
+    def _add_vertices(self):
+        """Add vertices to the network artist."""
+
+        self._vertices = VertexCollection(
+            layout=self.get_layout(),
+            style=get_style(".vertex"),
+            labels=self._get_label_series("vertex"),
+            transform=self.get_transform(),
+            offset_transform=self.get_offset_transform(),
         )
 
-        edgepatches = []
-        arrowpatches = []
-        adjacent_vertex_ids = []
-        adjecent_vertex_centers = []
-        adjecent_vertex_paths = []
-        for i, (vid1, vid2) in enumerate(edge_df.index):
-            # Get the vertices for this edge
-            vcenter1 = vertex_layout_df.loc[vid1, layout_columns].values
-            vcenter2 = vertex_layout_df.loc[vid2, layout_columns].values
-            vpath1 = vertex_paths[vertex_indices[vid1]]
-            vpath2 = vertex_paths[vertex_indices[vid2]]
-
-            edge_stylei = rotate_style(edge_style, index=i, id=(vid1, vid2))
+    def _add_edges(self):
+        """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.
+        """
 
-            # 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))
-            adjecent_vertex_centers.append((vcenter1, vcenter2))
-            adjecent_vertex_paths.append((vpath1, vpath2))
+        labels = self._get_label_series("edge")
+        edge_style = get_style(".edge")
 
-            arrow_patch = make_arrow_patch(
-                **arrow_style,
+        if "cmap" in edge_style:
+            cmap_fun = _build_cmap_fun(
+                edge_style["color"],
+                edge_style["cmap"],
             )
-            arrowpatches.append(arrow_patch)
-
-        adjacent_vertex_ids = np.array(adjacent_vertex_ids)
-        adjecent_vertex_centers = np.array(adjecent_vertex_centers)
-        # NOTE: the paths might have different number of sides, so it cannot be recast
-
-        # TODO:: deal with "ports" a la graphviz
-
-        art = DirectedEdgeCollection(
-            edges=edgepatches,
-            arrows=arrowpatches,
-            labels=labels,
-            vertex_ids=adjacent_vertex_ids,
-            vertex_paths=adjecent_vertex_paths,
-            vertex_centers=adjecent_vertex_centers,
-            transform=self.axes.transData,
-            style=edge_style,
-        )
-        self._edges = art
-
-    def _add_undirected_edges(self, labels=None):
-        """Draw undirected edges."""
-        edge_style = get_style(".edge")
+        else:
+            cmap_fun = None
 
-        layout_columns = [
-            f"_ipx_layout_{i}" for i in range(self._ipx_internal_data["ndim"])
-        ]
-        vertex_layout_df = self._ipx_internal_data["vertex_df"][layout_columns]
         edge_df = self._ipx_internal_data["edge_df"].set_index(
             ["_ipx_source", "_ipx_target"]
         )
 
-        # This contains the patches for vertices, for edge shortening and such
-        vertex_paths = self._vertices._paths
-        vertex_indices = pd.Series(
-            np.arange(len(vertex_layout_df)), index=vertex_layout_df.index
-        )
-
+        if "cmap" in edge_style:
+            colorarray = []
         edgepatches = []
         adjacent_vertex_ids = []
-        adjecent_vertex_centers = []
-        adjecent_vertex_paths = []
         for i, (vid1, vid2) in enumerate(edge_df.index):
-            # Get the vertices for this edge
-            vcenter1 = vertex_layout_df.loc[vid1, layout_columns].values
-            vcenter2 = vertex_layout_df.loc[vid2, layout_columns].values
-            vpath1 = vertex_paths[vertex_indices[vid1]]
-            vpath2 = vertex_paths[vertex_indices[vid2]]
-
-            edge_stylei = rotate_style(edge_style, index=i, id=(vid1, vid2))
+            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"])
 
             # These are not the actual edges drawn, only stubs to establish
             # the styles which are then fed into the dynamic, optimised
@@ -353,155 +227,62 @@ class NetworkArtist(mpl.artist.Artist):
             )
             edgepatches.append(patch)
             adjacent_vertex_ids.append((vid1, vid2))
-            adjecent_vertex_centers.append((vcenter1, vcenter2))
-            adjecent_vertex_paths.append((vpath1, vpath2))
-
-        adjacent_vertex_ids = np.array(adjacent_vertex_ids)
-        adjecent_vertex_centers = np.array(adjecent_vertex_centers)
-        # NOTE: the paths might have different number of sides, so it cannot be recast
 
-        # TODO:: deal with "ports" a la graphviz
+        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
 
-        art = UndirectedEdgeCollection(
+        self._edges = EdgeCollection(
             edgepatches,
-            labels=labels,
             vertex_ids=adjacent_vertex_ids,
-            vertex_paths=adjecent_vertex_paths,
-            vertex_centers=adjecent_vertex_centers,
-            transform=self.axes.transData,
+            vertex_collection=self._vertices,
+            layout=self.get_layout(),
+            layout_coordinate_system="cartesian",
+            labels=labels,
+            transform=self.get_offset_transform(),
             style=edge_style,
+            directed=self._ipx_internal_data["directed"],
         )
-        self._edges = art
-
-    def _process(self):
-        self._clear_state()
-
-        # TODO: some more things might be plotted before this
-
-        # NOTE: we plot vertices first to get size etc. for edge shortening
-        # but when the mpl engine runs down all children artists for actual
-        # drawing it uses get_children() to get the order. Whatever is last
-        # in that order will get drawn on top (vis-a-vis zorder).
-        self._add_vertices()
-        self._add_edges()
-        if "label" in self._ipx_internal_data["vertex_df"].columns:
-            self._add_vertex_labels()
-
-        # TODO: callbacks for stale vertices/edges
-
-        # Forward mpl properties to children
-        # TODO sort out all of the things that need to be forwarded
-        for child in self.get_children():
-            # set the figure & axes on child, this ensures each artist
-            # down the hierarchy knows where to draw
-            if hasattr(child, "set_figure"):
-                child.set_figure(self.figure)
-            child.axes = self.axes
-
-            # forward the clippath/box to the children need this logic
-            # because mpl exposes some fast-path logic
-            clip_path = self.get_clip_path()
-            if clip_path is None:
-                clip_box = self.get_clip_box()
-                child.set_clip_box(clip_box)
-            else:
-                child.set_clip_path(clip_path)
+        if "cmap" in edge_style:
+            self._edges.set_array(colorarray)
 
     @_stale_wrapper
-    def draw(self, renderer, *args, **kwds):
+    def draw(self, renderer):
         """Draw each of the children, with some buffering mechanism."""
+        if not self.get_children():
+            self._add_vertices()
+            self._add_edges()
+
         if not self.get_visible():
             return
 
-        if not self.get_children():
-            self._process()
+        # 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, *args, **kwds)
-
-
-# INTERNAL ROUTINES
-def _create_internal_data(
-    network,
-    layout=None,
-    vertex_labels=None,
-    edge_labels=None,
-):
-    """Create internal data for the network."""
-    nl = network_library(network)
-    directed = detect_directedness(network)
-
-    if nl == "networkx":
-        # Vertices are indexed by node ID
-        vertex_df = normalise_layout(layout).loc[pd.Index(network.nodes)]
-        ndim = vertex_df.shape[1]
-        vertex_df.columns = [f"_ipx_layout_{i}" for i in range(ndim)]
-
-        # Vertex labels
-        if vertex_labels is not None:
-            if len(vertex_labels) != len(vertex_df):
-                raise ValueError(
-                    "Vertex labels must be the same length as the number of vertices."
-                )
-            vertex_df["label"] = vertex_labels
-
-        # Edges are a list of tuples, because of multiedges
-        tmp = []
-        for u, v, d in network.edges.data():
-            row = {"_ipx_source": u, "_ipx_target": v}
-            row.update(d)
-            tmp.append(row)
-        edge_df = pd.DataFrame(tmp)
-        del tmp
-
-        # Edge labels
-        if edge_labels is not None:
-            if len(edge_labels) != len(edge_df):
-                raise ValueError(
-                    "Edge labels must be the same length as the number of edges."
-                )
-            edge_df["labels"] = edge_labels
-
-    else:
-        # Vertices are ordered integers, no gaps
-        vertex_df = normalise_layout(layout)
-        ndim = vertex_df.shape[1]
-        vertex_df.columns = [f"_ipx_layout_{i}" for i in range(ndim)]
-
-        # Vertex labels
-        if vertex_labels is not None:
-            if len(vertex_labels) != len(vertex_df):
-                raise ValueError(
-                    "Vertex labels must be the same length as the number of vertices."
-                )
-            vertex_df["label"] = vertex_labels
-
-        # Edges are a list of tuples, because of multiedges
-        tmp = []
-        for edge in network.es:
-            row = {"_ipx_source": edge.source, "_ipx_target": edge.target}
-            row.update(edge.attributes())
-            tmp.append(row)
-        edge_df = pd.DataFrame(tmp)
-        del tmp
-
-        # Edge labels
-        if edge_labels is not None:
-            if len(edge_labels) != len(edge_df):
-                raise ValueError(
-                    "Edge labels must be the same length as the number of edges."
-                )
-            edge_df["labels"] = edge_labels
-
-    internal_data = {
-        "vertex_df": vertex_df,
-        "edge_df": edge_df,
-        "directed": directed,
-        "network_library": nl,
-        "ndim": ndim,
-    }
-    return internal_data
+            art.draw(renderer)
+
+
+def _update_from_internal(style, row, kind):
+    """Update single vertex/edge style from internal data."""
+    if "color" in row:
+        style["color"] = row["color"]
+    if "facecolor" in row:
+        style["facecolor"] = row["facecolor"]
+    if "edgecolor" in row:
+        if kind == "vertex":
+            style["edgecolor"] = row["edgecolor"]
+        else:
+            style["color"] = row["edgecolor"]
+
+    if "linewidth" in row:
+        style["linewidth"] = row["linewidth"]
+    if "linestyle" in row:
+        style["linestyle"] = row["linestyle"]
+    if "alpha" in row:
+        style["alpha"] = row["alpha"]