edgebundling.py

import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.collections import LineCollection
from itertools import count
from heapq import heappush, heappop
from collections import defaultdict
import time
import pandas as pd
from datashader.bundling import hammer_bundle  # New import for hammer bundling

###############################################################################
# Minimal AbstractBundling base class (refactored from .abstractBundling import)
###############################################################################
class AbstractBundling:
    def __init__(self, G: nx.Graph):
        self.G = G

    def bundle(self):
        raise NotImplementedError("Subclasses should implement 'bundle'.")

###############################################################################
# Simple SplineC placeholder (refactoring out the nx2ipe dependency)
###############################################################################
class SplineC:
    def __init__(self, points):
        self.points = points

###############################################################################
# A base SpannerBundling class that SpannerBundlingNoSP depends on
###############################################################################
class SpannerBundling(AbstractBundling):
    """
    S-EPB. Implementation
    
    weightFactor: kappa value that sets the bundling strength
    distortion: t value that sets the maximum allowed stretch/distortion
    numWorkers: number of workers that process biconnected components 
    """
    def __init__(self, G: nx.Graph, weightFactor=2, distortion=2, numWorkers=1):
        super().__init__(G)
        self.distortion = distortion
        self.weightFactor = weightFactor
        self.mode = "greedy"
        self.name = None
        self.numWorkers = numWorkers

    @property
    def name(self):
        return f"SEPB_d_{self.distortion}_w_{self.weightFactor}_{self.mode}"

    @name.setter
    def name(self, value):
        self._name = value

    def bundle(self):
        # Default does nothing
        return 0.0

    def process(self, component):
        # Default does nothing
        pass

    def spanner(self, g, k):
        # Default does nothing
        return None

###############################################################################
# The requested SpannerBundlingNoSP class
###############################################################################
class SpannerBundlingNoSP(SpannerBundling):
    """
    S-EPB where instead of computing single source shortest paths we reuse
    shortest paths during the spanner construction. 
    """
    def __init__(self, G: nx.Graph, weightFactor=2, distortion=2):
        super().__init__(G)
        self.distortion = distortion
        self.weightFactor = weightFactor
        self.mode = "reuse"

    def bundle(self):
        """
        Executes the bundling process on all biconnected components.
        Returns the total time for bundling.
        """
        t_start = time.process_time()

        if nx.is_directed(self.G):
            # Convert to undirected for the biconnected components
            GG = self.G.to_undirected(as_view=True)
            components = nx.biconnected_components(GG)
        else:
            components = nx.biconnected_components(self.G)

        to_process = []
        for nodes in components:
            if len(nodes) > 2:
                subg = self.G.subgraph(nodes).copy()
                to_process.append(subg)

        # Sort the components from largest to smallest
        to_process = sorted(to_process, key=lambda x: len(x.nodes()), reverse=True)

        # Process each component
        for comp in to_process:
            self.process(comp)

        return time.process_time() - t_start

    def process(self, component):
        """
        Process a component: build a spanner, then for each edge not in
        the spanner, store a 'path' and create a Spline if possible.
        """
        T = self.spanner(component, self.distortion)

        # Mark edges in T as 'Spanning'
        for u, v, data in T.edges(data=True):
            data["weight"] = np.power(data["dist"], self.weightFactor)

        for u, v in T.edges():
            self.G[u][v]["Layer"] = "Spanning"
            self.G[u][v]["Stroke"] = "blue"

        # For edges not in T, build a spline from the stored path
        for u, v, data in component.edges(data=True):
            if T.has_edge(u, v):
                continue

            path = data.get("path", [])
            if len(path) < 1:
                continue

            spline_points = []
            current = path[0]
            for nxt in path[1:-1]:
                x = component.nodes[nxt].get("X", component.nodes[nxt].get("x", 0))
                y = component.nodes[nxt].get("Y", component.nodes[nxt].get("y", 0))
                spline_points.append((x, y))
                current = nxt

            self.G[u][v]["Spline"] = SplineC(spline_points)
            self.G[u][v]["Layer"] = "Bundled"
            self.G[u][v]["Stroke"] = "purple"

        return

    def spanner(self, g, k):
        """
        Create a spanner and store the shortest path in edge['path'] when the
        edge is not added to the spanner.
        """
        if nx.is_directed(g):
            spanner = nx.DiGraph()
        else:
            spanner = nx.Graph()

        edges = sorted(g.edges(data=True), key=lambda t: t[2].get("dist", 1))

        for u, v, data in edges:
            if u not in spanner.nodes:
                spanner.add_edge(u, v, dist=data["dist"])
                continue
            if v not in spanner.nodes:
                spanner.add_edge(u, v, dist=data["dist"])
                continue

            pred, pathLength = nx.dijkstra_predecessor_and_distance(
                spanner, u, weight="dist", cutoff=k * data["dist"]
            )

            # If v is in pathLength, we store the path in data['path']
            if v in pathLength:
                # reconstruct path from v back to u
                path = []
                nxt = v
                while nxt != u:
                    path.append(nxt)
                    nxt = pred[nxt][0]
                # remove the first node (==v) because we typically want just intermediate
                path = path[1:]
                path.reverse()

                data["path"] = path
            else:
                spanner.add_edge(u, v, dist=data["dist"])

        return spanner

###############################################################################
# Function to plot only the bundled edges (with optional color gradient)
###############################################################################
def plot_bundled_edges_only(G, edge_gradient=False, node_colors=None, ax=None, **plot_kwargs):
    """
    Plots only the edges whose 'Layer' is 'Bundled' (or user-defined).
    Nodes are plotted for reference in black.
    
    Parameters:
        G: NetworkX graph
        title: Plot title
        edge_gradient: If True, color edges with gradient
        node_colors: Dictionary of node colors
        ax: Optional matplotlib axis to plot on. If None, creates new figure.
        **plot_kwargs: Additional keyword arguments passed to LineCollection
    """
    # Use provided axis or create new one
    if ax is None:
        plt.figure(figsize=(8, 8))
        ax = plt.gca()

    # 1. Extract positions
    pos = {}
    for node, data in G.nodes(data=True):
        x = data.get('X', data.get('x', 0))
        y = data.get('Y', data.get('y', 0))
        pos[node] = (x, y)

    # 2. Assign or retrieve node colors. If your graph doesn't already have
    #    some color-coded attribute, you can define them here. 
    #    For example, let's just fix them to green for demonstration:
    # node_colors = {}
    # for node in G.nodes():
    #     node_colors[node] = (0.0, 0.5, 0.0, 1.0)  # RGBA

    # 3. Build up segments (and possibly per-segment colors) for the edges
    def binomial(n, k):
        """Compute the binomial coefficient (n choose k)."""
        coeff = 1
        for i in range(1, k + 1):
            coeff *= (n - i + 1) / i
        return coeff

    def approxBezier(points, n=50):
        """
        Compute and return n points along a Bezier curve defined by control points.
        """
        X, Y = [], []
        m = len(points) - 1
        binom_vals = [binomial(m, i) for i in range(m + 1)]
        t_values = np.linspace(0, 1, n)
        for t in t_values:
            pX, pY = 0.0, 0.0
            for i, p in enumerate(points):
                coeff = binom_vals[i] * ((1 - t) ** (m - i)) * (t ** i)
                pX += coeff * p[0]
                pY += coeff * p[1]
            X.append(pX)
            Y.append(pY)
        return np.column_stack([X, Y])

    edge_segments = []
    edge_colors = []

    for u, v, data in G.edges(data=True):
        if data.get("Layer", None) != "Bundled":
            # Skip edges not marked as bundled
            continue

        # (a) Gather the control points
        if "Spline" in data and data["Spline"] is not None:
            spline_obj = data["Spline"]
            control_points = list(spline_obj.points)
            # Add the start/end for completeness
            control_points = [pos[u]] + control_points + [pos[v]]
        else:
            # fallback to a straight line
            control_points = [pos[u], pos[v]]

        # (b) Approximate a curve from these control points
        #     We always subdivide if edge_gradient is True. 
        #     If not gradient-based, only subdivide for an actual curve.
        do_subdivide = edge_gradient or (len(control_points) > 2)
        if do_subdivide:
            curve_points = approxBezier(control_points, n=50)
        else:
            curve_points = np.array(control_points)

        # (c) If we're using gradient, we break it into small segments, each with a color
        if edge_gradient:
            c_u = np.array(node_colors[u])  # RGBA for source node
            c_v = np.array(node_colors[v])  # RGBA for target node
            num_pts = len(curve_points)
            for i in range(num_pts - 1):
                p0 = curve_points[i]
                p1 = curve_points[i + 1]
                # fraction along the curve
                t = i / max(1, (num_pts - 2))  
                seg_color = (1 - t) * c_u + t * c_v  # linear interpolation in RGBA
                edge_segments.append([p0, p1])
                edge_colors.append(seg_color)
        else:
            # Single color for the entire edge
            if len(curve_points) > 1:
                edge_segments.append([curve_points[0], curve_points[-1]])
                edge_colors.append((0.5, 0.0, 0.5, 0.9))  # purple RGBA

    # 4. Plot
    # Remove the plt.figure() call since we're using the provided axis

    # Set default values for LineCollection
    lc_kwargs = {
        'linewidths': 1,
        'alpha': 0.9
    }
    
    # If colors weren't explicitly passed and we calculated edge_colors, use them
    if 'colors' not in plot_kwargs and edge_colors:
        lc_kwargs['colors'] = edge_colors
    
    # Update with user-provided kwargs
    lc_kwargs.update(plot_kwargs)
    
    # Create the LineCollection with all parameters
    lc = LineCollection(edge_segments, **lc_kwargs)
    ax.add_collection(lc)

    # The nodes in black
    # node_positions = np.array([pos[n] for n in G.nodes()])
    # ax.scatter(node_positions[:, 0], node_positions[:, 1], color="black", s=20, alpha=0.8)

    # ax.set_aspect('equal')
    # Remove plt.show() since we want to allow further additions to the plot

###############################################################################
# Convenience function to run SpannerBundlingNoSP on a graph and plot results
###############################################################################
def run_and_plot_spanner_bundling_no_sp(G, weightFactor=2, distortion=2, edge_gradient=False, node_colors=None, ax=None, **plot_kwargs):
    """
    Create an instance of SpannerBundlingNoSP, run .bundle(), and
    plot only the bundled edges. Pass edge_gradient=True to see 
    color-gradient edges.
    
    Additional keyword arguments are passed to the LineCollection for edge styling.
    """
    bundler = SpannerBundlingNoSP(G, weightFactor=weightFactor, distortion=distortion)
    bundler.bundle()
    plot_bundled_edges_only(G, 
                          edge_gradient=edge_gradient,
                          node_colors=node_colors,
                          ax=ax,
                          **plot_kwargs)

def run_hammer_bundling(G, accuracy=500, advect_iterations=50, batch_size=20000, 
                       decay=0.01, initial_bandwidth=1.1, iterations=4, 
                       max_segment_length=0.016, min_segment_length=0.008, 
                       tension=1.2):
    """
    Run hammer bundling on a NetworkX graph and return the bundled paths.
    """
    # Create nodes DataFrame
    nodes = []
    node_to_index = {}
    for i, (node, attr) in enumerate(G.nodes(data=True)):
        x = attr.get('X', attr.get('x', 0))
        y = attr.get('Y', attr.get('y', 0))
        nodes.append({'node': node, 'x': x, 'y': y})
        node_to_index[node] = i
    nodes_df = pd.DataFrame(nodes)
    
    # Create edges DataFrame
    edges = []
    for u, v in G.edges():
        edges.append({'source': node_to_index[u], 'target': node_to_index[v]})
    edges_df = pd.DataFrame(edges)
    
    # Apply hammer bundling
    bundled_paths = hammer_bundle(nodes_df, edges_df,
                                accuracy=accuracy,
                                advect_iterations=advect_iterations,
                                batch_size=batch_size,
                                decay=decay,
                                initial_bandwidth=initial_bandwidth,
                                iterations=iterations,
                                max_segment_length=max_segment_length,
                                min_segment_length=min_segment_length,
                                tension=tension)
    
    # Convert bundled paths to a format compatible with our plotting function
    paths = []
    current_path = []
    edge_index = 0
    
    for _, row in bundled_paths.iterrows():
        if pd.isna(row['x']) or pd.isna(row['y']):
            if current_path:
                # Get source and target nodes for this edge
                source_idx = edges_df.iloc[edge_index]['source']
                target_idx = edges_df.iloc[edge_index]['target']
                source_node = nodes_df.iloc[source_idx]['node']
                target_node = nodes_df.iloc[target_idx]['node']
                
                paths.append((source_node, target_node, current_path))
                current_path = []
                edge_index += 1
        else:
            current_path.append((row['x'], row['y']))
    
    if current_path:  # Handle the last path
        source_idx = edges_df.iloc[edge_index]['source']
        target_idx = edges_df.iloc[edge_index]['target']
        source_node = nodes_df.iloc[source_idx]['node']
        target_node = nodes_df.iloc[target_idx]['node']
        paths.append((source_node, target_node, current_path))
    
    return paths

def plot_bundled_edges(G, bundled_paths, edge_gradient=False, node_colors=None, ax=None, **plot_kwargs):
    """
    Generic plotting function that works with both bundling methods.
    
    Parameters:
        G: NetworkX graph
        bundled_paths: List of (source, target, path_points) tuples
        edge_gradient: If True, color edges with gradient
        node_colors: Dictionary of node colors
        ax: Optional matplotlib axis
        **plot_kwargs: Additional styling arguments
    """
    if ax is None:
        plt.figure(figsize=(8, 8))
        ax = plt.gca()

    def approxBezier(points, n=50):
        """Compute points along a Bezier curve."""
        points = np.array(points)
        t = np.linspace(0, 1, n)
        return np.array([(1-t)*points[:-1] + t*points[1:] for t in t]).reshape(-1, 2)

    edge_segments = []
    edge_colors = []

    for source, target, path_points in bundled_paths:
        points = np.array(path_points)
        
        if edge_gradient:
            # Create segments with gradient colors
            c_u = np.array(node_colors[source])
            c_v = np.array(node_colors[target])
            num_pts = len(points)
            
            for i in range(num_pts - 1):
                p0, p1 = points[i], points[i + 1]
                t = i / max(1, (num_pts - 2))
                seg_color = (1 - t) * c_u + t * c_v
                edge_segments.append([p0, p1])
                edge_colors.append(seg_color)
        else:
            # Single color for the entire path
            for i in range(len(points) - 1):
                edge_segments.append([points[i], points[i + 1]])
                edge_colors.append((0.5, 0.0, 0.5, 0.9))

    # Plot edges
    lc_kwargs = {'linewidths': 1, 'alpha': 0.9}
    if edge_colors:
        lc_kwargs['colors'] = edge_colors
    lc_kwargs.update(plot_kwargs)
    
    lc = LineCollection(edge_segments, **lc_kwargs)
    ax.add_collection(lc)
    ax.autoscale()

def run_and_plot_bundling(G, method='hammer', edge_gradient=False, node_colors=None, ax=None, 
                         bundling_params=None, **plot_kwargs):
    """
    Unified function to run and plot different bundling methods.
    
    Parameters:
        G: NetworkX graph
        method: 'spanner' or 'hammer'
        bundling_params: dict of parameters specific to the bundling method
        Other parameters same as plot_bundled_edges
    """
    bundling_params = bundling_params or {}
    
    if method == 'spanner':
        bundler = SpannerBundlingNoSP(G, **bundling_params)
        bundler.bundle()
        
        # Extract bundled paths from SpannerBundling format
        bundled_paths = []
        for u, v, data in G.edges(data=True):
            if data.get("Layer") == "Bundled" and "Spline" in data:
                spline_points = data["Spline"].points
                pos_u = (G.nodes[u].get('X', G.nodes[u].get('x', 0)), 
                        G.nodes[u].get('Y', G.nodes[u].get('y', 0)))
                pos_v = (G.nodes[v].get('X', G.nodes[v].get('x', 0)), 
                        G.nodes[v].get('Y', G.nodes[v].get('y', 0)))
                path = [pos_u] + list(spline_points) + [pos_v]
                bundled_paths.append((u, v, path))
                
    elif method == 'hammer':
        bundled_paths = run_hammer_bundling(G, **bundling_params)
    else:
        raise ValueError(f"Unknown bundling method: {method}")
    
    plot_bundled_edges(G, bundled_paths, edge_gradient, node_colors, ax, **plot_kwargs)