Planted partition model

planted-partition-model.py

import graph_tool.all as gt
import numpy as np
from scipy.sparse.csgraph import connected_components


def generate_PPI_network(Cave, mixing_rate, N, q):
    """
    Generate PPI network. 
    
    Cave: int. The average degree. 
    mixing_rate: float. The mixing of communities, with range (0,1]. 
                The mixing_rate = 1 generates an Erdos-Renyi random graph, and mixing_rate~0 generates well-separated communities. 
                It is defined as the ratio p_out / pave, where pout is the probability of inter-community edges, and 
                pave is the average edge probability (the density of edges in the network).
  
    return: net, membership
        net: scipy.csr_matrix representation of the adjacency matrix of the generated network. 
        membership: numpy array of membership IDs of the nodes in the network.
      
    """
    memberships = np.sort(np.arange(N) % q)

    q = int(np.max(memberships) + 1)
    N = len(memberships)
    U = sparse.csr_matrix((np.ones(N), (np.arange(N), memberships)), shape=(N, q))

    Cout = np.maximum(1, mixing_rate * Cave)
    Cin = q * Cave - (q - 1) * Cout
    pout = Cout / N
    pin = Cin / N

    Nk = np.array(U.sum(axis=0)).reshape(-1)

    P = np.ones((q, q)) * pout + np.eye(q) * (pin - pout)
    probs = np.diag(Nk) @ P @ np.diag(Nk)
    gt_params = {
        "b": memberships,
        "probs": probs,
        "micro_degs": False,
        "in_degs": np.ones_like(memberships) * Cave,
        "out_degs": np.ones_like(memberships) * Cave,
    }

    # Generate the network until the degree sequence
    # satisfied the thresholds
    while True:
        g = gt.generate_sbm(**gt_params)

        A = gt.adjacency(g).T

        A.data = np.ones_like(A.data)
        # check if the graph is connected
        if connected_components(A)[0] == 1:
            break
    return A, memberships