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)