plotting.py

import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as mcolors

def plot_initial_cycles(graph, positions, initial_cycles, ax=None, title="Initial Cycle Set"):
    """Plots the graph and highlights edges in the initial cycle set."""
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))

    nx.draw_networkx_nodes(graph, positions, node_size=5, node_color='gray', alpha=0.6, ax=ax)
    nx.draw_networkx_edges(graph, positions, edge_color='lightgray', alpha=0.5, ax=ax)

    # Collect all unique edges from the initial cycles
    highlight_edges = set()
    for cycle in initial_cycles:
        for edge in cycle:
            # Ensure edge nodes are in the graph before adding
            u, v = edge
            if u in graph and v in graph:
                 highlight_edges.add(tuple(sorted(edge))) # Use sorted tuple for consistency

    # Draw highlighted edges
    nx.draw_networkx_edges(graph, positions, edgelist=list(highlight_edges), edge_color='red', width=1.5, ax=ax)

    ax.set_title(title)
    ax.axis('off')
    return ax

def plot_final_rho(graph, positions, rho_star_map, edge_list, ax=None, title="Final Rho* Weights"):
    """Plots the graph with edges colored by rho* values."""
    if ax is None:
        fig, ax = plt.subplots(figsize=(12, 10)) # Slightly wider for colorbar

    rho_values = np.array([rho_star_map.get(edge, 0.0) for edge in edge_list])
    # Handle case where all rho are zero or constant
    min_rho = np.min(rho_values) if len(rho_values) > 0 else 0
    max_rho = np.max(rho_values) if len(rho_values) > 0 else 0

    if max_rho <= min_rho: # If all values are the same (or empty)
        cmap = cm.get_cmap('viridis')
        edge_colors = [cmap(0.5)] * len(edge_list) # Use a mid-range color
        norm = mcolors.Normalize(vmin=min_rho - 1e-6, vmax=max_rho + 1e-6) # Avoid division by zero
    else:
        cmap = cm.get_cmap('viridis') # Or 'plasma', 'inferno', 'magma'
        norm = mcolors.Normalize(vmin=min_rho, vmax=max_rho)
        edge_colors = [cmap(norm(rho_star_map.get(edge, 0.0))) for edge in edge_list]

    nx.draw_networkx_nodes(graph, positions, node_size=5, node_color='black', alpha=0.7, ax=ax)
    nx.draw_networkx_edges(graph, positions, edge_color=edge_colors, width=1.0, ax=ax)

    # Add a colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([]) # Necessary for standalone colorbar
    plt.colorbar(sm, ax=ax, shrink=0.8, label='rho* value')

    ax.set_title(title)
    ax.axis('off')
    return ax