graph_manipulator.py

import networkx as nx
import random

def generate_d_graph_chain(size, d):
    """ Generate a d-graph chain (succession of unit d-graphs).
        If you slice any 2*d+1 part of the graph, it will be a unit d-graph
        :param size The number of nodes in the chain (should not be less than 2*d+1)
        :param d The number of connection on the left and on the right for any node
        :return The d-graph chain
    """
    G = nx.Graph()

    for idx in range(size):
        # Create the node
        G.add_node(idx)

        # Link the node to d previous nodes
        for prev in range(max(0, idx-d), idx):
            G.add_edge(prev, idx)

    return G


def generate_approx_d_graph_chain(size, d_max, d_avg, rnd_seed=None):
    """ Generate an almost d-graph chain (succession of unit d-graphs). Almost because they are d-graphs in average
        with a coverage variation.
        :param size The number of nodes in the chain (should not be less than 2*d+1)
        :param d_max The max number of connection on the left and on the right for any node
        :param d_avg The average d value in the graph (ie 2*d average coverage)
        :param rnd_seed Fix the random seed for reproducibility
        :return The d-graph chain
    """
    # Reproducibility
    if rnd_seed:
        random.seed(rnd_seed)

    # Sample size computation
    sursample_needed = round(size * (2 * d_max - 2 * d_avg) / (4 * d_max))
    total_size = size + sursample_needed

    # Choose which idx to jump over
    to_skip = random.sample(range(total_size), sursample_needed)
    to_skip.sort()

    # Generate sequence
    G = nx.Graph()
    previous_nodes = [None]* d_max
    next_idx = 0
    for idx in range(total_size):
        if len(to_skip) > 0 and to_skip[0] == idx:
            to_skip.pop(0)
            previous_nodes.append(None)
        else:
            # Create the node
            G.add_node(next_idx)

            # link the node with previous ones
            for node_idx in previous_nodes:
                if node_idx is not None:
                    G.add_edge(next_idx, node_idx)

            # Update the state
            previous_nodes.append(next_idx)
            next_idx += 1

        # Remove oldest previous node
        previous_nodes.pop(0)

    return G


""" Merge 2 nodes of a graph G.
    The new node have edges from both of the previous nodes (dereplicated).
    If a link between node1 and node2 exist, it's discarded.
    @param G The graph to manipulate
    @param node1 First node to merge
    @param node2 Second node to merge
    @return The name of the new node in G
"""
def merge_nodes(G, node1, node2):
    # Create the new node
    new_node = f"{node1}_{node2}" if node1 < node2 else f"{node2}_{node1}"
    G.add_node(new_node)

    # List the edges to add
    neighbors = []
    for edge in G.edges(node1):
        neighbor = edge[0] if edge[0] != node1 else edge[1]

        if neighbor == node2:
            continue

        neighbors.append(neighbor)

    for edge in G.edges(node2):
        neighbor = edge[0] if edge[0] != node2 else edge[1]

        if neighbor == node1:
            continue

        neighbors.append(neighbor)

    # De-replicate neighbors
    neighbors = set(neighbors)

    # Add neighbors
    for neighbor in neighbors:
        G.add_edge(neighbor, new_node)

    # Remove previous nodes from the graph
    G.remove_node(node1)
    G.remove_node(node2)

    return new_node