modularity.py

import torch
from collections import deque
from probability import entropy_to_temperature, discrete_entropy, log2prob
from tqdm import trange, tqdm


ADJACENCY_EPS = 1e-15


def is_valid_adjacency_matrix(adj:torch.Tensor, enforce_sym=False, enforce_no_self=False, enforce_binary=False) -> bool:
    valid = torch.all(adj >= 0.) and torch.mean(adj) > ADJACENCY_EPS
    if enforce_binary:
        valid = valid and torch.all(torch.logical_or(adj == 0., adj == 1.))
    if enforce_no_self:
        valid = valid and torch.all(adj.diag() == 0.)
    if enforce_sym:
        valid = valid and torch.allclose(adj, adj.T)
    return valid


def soft_num_clusters(clusters:torch.Tensor):
    marginal_frac = clusters.mean(dim=0)
    plogp = -marginal_frac * torch.log(marginal_frac)
    # plogp cannot mathematically be less than zero or greater than 1/e; anywhere this happens is a case of floating
    # point instability, and we just throw it out
    plogp[torch.isnan(plogp)] = 0.
    plogp[plogp < 0] = 0.
    plogp[plogp > 0.3678794412] = 0.
    entropy = torch.sum(plogp)
    return torch.exp(entropy)


def girvan_newman(adj:torch.Tensor, clusters:torch.Tensor):
    """Compute GN modularity statistic on adjacency matrix A using cluster assignments P.

    This version follows the equations in the GN paper fairly literally. The first line computes the cluster-cluster
    connectivity matrix 'e', which is then summed along rows or columns to get expected connectivity.
    """
    e = (clusters.T @ adj @ clusters) / torch.sum(adj)
    e_sum_row = torch.sum(e, dim=0)
    e_sum_col = torch.sum(e, dim=1)
    return torch.trace(e) - torch.sum(e_sum_col * e_sum_row)


def girvan_newman_sym(adj:torch.Tensor, clusters:torch.Tensor):
    """Compute GN modularity statistic on symmetric adjacency matrix A using cluster assignments P.

    This version makes two assumptions: that A is symmetric, and that the rows of P sum to 1. These assumptions are
    exploited to reduce redundant calculations for a slight speedup over girvan_newman().
    """
    uB = adj @ clusters
    # Since P sums to 1 along rows, this is equivalent to sum(A) but on a much smaller matrix
    B = uB / torch.sum(uB)
    # The following is equivalent to Trace(P'*A*P) but without transposing and without computing off-diagonal elements
    # that would be dropped by trace anyway
    cluster_connectivity = torch.sum(clusters * B)
    # When A is symmetric, we don't need to compute the expected connectivity separately for rows and columns... just
    # compute once and multiply by itself for the trace. In the GN paper, this corresponds to computing the expected
    # connectivity statistic row- or column-wise (a_i or a_j).
    Brow = torch.sum(B, dim=0)
    expected_connectivity = torch.sum(Brow * Brow)
    # GN modularity statistic is the actual cluster connectivity minus the baseline 'expected' connectivity
    return cluster_connectivity - expected_connectivity


def sparsify(adj:torch.Tensor, fraction:float) -> torch.Tensor:
    """Sparsify and binarize a continuous adjacency/association matrix, keeping the top 'fraction' percent of elements

    Note: 'dead' rows/columns that are all zeros to begin with don't count towards sparsity, and neither does the main
    diagonal, which is set to zero regardless.
    """
    # First, remove the diagonal
    adj = adj - adj.diag().diag()
    # Select all off-diagonal elements
    i, j = torch.tril_indices(*adj.size(), offset=1)
    off_diag = adj[i, j]
    # Adjust fraction such that it is relative to the total number of nonzero elements
    frac_nonzero = (off_diag != 0).float().mean()
    # Quit early if there aren't any nonzero elements
    if frac_nonzero == 0.:
        return adj.new_zeros(adj.size())
    fraction = fraction * frac_nonzero
    # Find whatever value lives at the 'fraction'th quantile of all off-diagonal values
    cutoff = torch.quantile(off_diag[off_diag != 0], 1-fraction)
    # Set to 0 or 1 based on the cutoff
    return torch.where(adj < cutoff, 0., 1.)


def spectral_modularity(adj, max_clusters=None) -> torch.Tensor:
    """Spectral algorithm to quickly find an approximate maximum for the Girvan-Newman modularity score.

    See Newman, M. E. J. (2006). Modularity and community structure in networks. PNAS, 103(23), 8577–8582.
    """
    n = adj.size(0)
    max_clusters = max_clusters or n

    # Quit early if the adjacency matrix is all zeros (otherwise SVD below will error): return all-zeros for clustering
    if adj.mean() < ADJACENCY_EPS:
        return adj.new_zeros(n, max_clusters)

    # Normalize by sum of all 'edges' / AKA convert to probability assuming all A>=0
    adj = adj / adj.sum()
    # Compute degree of each 'vertex' / AKA compute marginal probability
    A1 = adj.sum(dim=1, keepdims=True)
    is_dead = A1.flatten() < ADJACENCY_EPS
    # Compute modularity matrix 'B': connectivity minus expected random connectivity
    B = adj - A1 * A1.T

    def gn_score(clusters):
        # Local helper function to compute Girvan-Newman modularity score using intermediate matrix 'B'
        return torch.sum(B * (clusters @ clusters.T))

    # Initialize everything into a single cluster, pruning out 'dead' units right away
    clusters = adj.new_zeros(n, max_clusters)
    clusters[~is_dead, 0] = 1.
    best_score = gn_score(clusters)
    # Iteratively subdivide until modularity score is not improved. This is a greedy method that takes any high-level
    # split that improves total modularity score. Each split is a branch of a binary tree, so after 'l' splits there
    # may be as many as 2^l modules, but some branches may be pruned early. This 'tree' is traversed breadth-first.
    # Tree traversal simply means keeping track of which column to try splitting next, which we keep track of in a
    # FIFO queue.
    queue, next_avail_col = deque([0]), 1
    while len(queue) > 0 and next_avail_col < max_clusters:
        col = queue.popleft()
        # Which variables are we working with here
        mask = clusters[:, col] == 1.
        # Skip this division if only one element left
        if mask.sum() <= 1:
            continue
        # Isolate the submatrix of B just containing these variables
        Bsub = B[mask, :][:, mask]
        # Compute single top eigenvectors for this sub-matrix
        _, _, v = torch.svd(Bsub)
        # Skip this division if the leading eigenvector does not contain both (+) and (-) elements. Otherwise, propose a
        # split based on the first eigenvector
        if torch.all(v[:, 0] >= 0.) or torch.all(v[:, 0] <= 0.):
            continue
        # Try new subdivision out: use current 'col' for (+) side of v, and use 'next_avail_col' for (-) side. In rare
        # cases, v will have exact zeros on nodes that are not 'dead' according to the EPS check above. The cluster
        # assignments for such nodes does not affect the score anyway (|v| is what matters), so we use >= in the (+)
        # cluster to include the zeros there arbitrarily.
        clusters[mask, col] = (v[:, 0] >= 0).float()
        clusters[mask, next_avail_col] = (v[:, 0] < 0).float()
        subdiv_score = gn_score(clusters)
        # If improved, keep it and push potential further subdivisions onto the queue
        if subdiv_score > best_score:
            queue.extend([col, next_avail_col])
            next_avail_col += 1
            best_score = subdiv_score
        # If not improved, undo change to 'clusters'
        else:
            clusters[:, col] = mask.float()
            clusters[:, next_avail_col] = 0.

    # Sanity check: only those units initially pruned as 'dead' should be missing
    assert torch.all((clusters.sum(dim=1) == 0.) == is_dead)

    return clusters


def monte_carlo_modularity(adj, clusters=None, max_clusters=None, steps=10000, target_entropy=0.15, device='cpu'):
    """Optimize GN modularity by initializing using spectral method, then repeatedly shuffling values around with
    probability proportional to how good the GN score is after that shuffle. Return the best-scoring clustering

    Each step, the temperature of the multinomial distribution is selected (scores are scaled) to match the given
    target entropy, which effectively sets the number of 'choices' for moves on each step. Testing suggests that
    an entropy of 0.1 (1.1ish options per step) is a good value.
    """
    adj = adj.to(device)

    assert is_valid_adjacency_matrix(adj, enforce_sym=True), \
        "monte_carlo_modularity expects a symmetric matrix of associations"

    if clusters is None:
        clusters = spectral_modularity(adj, max_clusters=max_clusters)

    max_clusters = clusters.size(1)
    best_score, best_clusters = girvan_newman_sym(adj, clusters).cpu(), clusters.clone()
    scores_history = torch.zeros(steps)
    temperature_history = torch.zeros(steps)
    entropy_history = torch.zeros(steps)
    # Some elements are initially not clustered ("dead") – don't bother moving these around
    is_dead = clusters.sum(dim=1) == 0.
    # Quit early if there are no units to move around
    if torch.all(is_dead):
        return clusters, torch.tensor(float('nan')), torch.tensor(float('nan')), torch.tensor(float('nan'))

    # At each step, we sample a new cluster with probability proportional to exp(score/temperature).
    # Keep track of the temperature from the previous iteration since it's a good guess for the
    # next iteration.
    temperature = 1.0

    # Pick a random item to shuffle around in each step
    ishuffle = torch.multinomial((~is_dead).float(), num_samples=steps, replacement=True)
    for i, idx in tqdm(enumerate(ishuffle), desc=f'MC shuffles', leave=False, total=steps):
        # Consider all possible re-assignments, but only consider creating a new column at most 1 time
        scores, did_new_col = float('-inf')*torch.ones(max_clusters), False
        for j in range(max_clusters):
            is_new_col = clusters[:, j].sum() == 0.
            if is_new_col and did_new_col:
                continue
            else:
                did_new_col = did_new_col or is_new_col
                # Re-assign vertex 'idx' to cluster 'j' and see what the score would be
                clusters[idx, :] = 0.
                clusters[idx, j] = 1.
                scores[j] = girvan_newman_sym(adj, clusters).cpu()
                # Track best-scoring clusters we've seen across all proposals
                if scores[j] > best_score:
                    best_score = scores[j]
                    best_clusters = clusters.clone()
        # Viewing scores as an energy, select temperature to match the target entropy
        temperature = entropy_to_temperature(scores, target_entropy, init_t=temperature, eps=0.01)
        # Keep a record of temperatures and entropies along the way
        temperature_history[i], entropy_history[i] = temperature, discrete_entropy(scores/temperature)
        # Pick randomly among decent-scoring options
        choice = torch.multinomial(log2prob(scores/temperature), num_samples=1)
        clusters[idx, :] = 0.
        clusters[idx, choice] = 1.
        scores_history[i] = scores[choice]

    # Sanity check: only those units initially pruned as 'dead' should be missing
    assert torch.all((best_clusters.sum(dim=1) == 0.) == is_dead)

    return best_clusters.cpu(), scores_history.cpu(), temperature_history.cpu(), entropy_history.cpu()


def gradient_ascent_modularity(adj, max_k=None, steps=5000, num_init=100):
    n = adj.size(0)
    max_k = max_k or n
    best_init, best_score = None, torch.tensor(float('-inf'))
    for _ in range(num_init):
        log_clusters = torch.randn(n, max_k, device=adj.device)
        tmp_score = girvan_newman_sym(adj, torch.softmax(log_clusters, dim=1))
        if tmp_score > best_score:
            best_score = tmp_score
            best_init = log_clusters
    log_clusters = best_init
    log_clusters.requires_grad_(True)
    scores = torch.zeros(steps, device=adj.device)
    opt = torch.optim.Adam([log_clusters], lr=0.005)
    for s in trange(steps, desc='GN gradient ascent'):
        opt.zero_grad()
        gn = girvan_newman_sym(adj, torch.softmax(log_clusters, dim=1))
        (-gn).backward()
        opt.step()

        with torch.no_grad():
            scores[s] = gn.item()

    return torch.softmax(log_clusters.detach(), dim=1).cpu(), scores


def greedy_alignment(cluster1, cluster2):
    c1c2 = cluster1.T @ cluster2
    n = c1c2.size()[0]
    # Greedily pick out max (r, c) indices of the c1c2 product, then remove row r and column c
    row_sort, col_sort = list(range(n)), list(range(n))
    for i in range(n):
        # Find 1D index of the max element of c1c2 alignments
        best = torch.argmax(c1c2).item()
        # Convert to 2D row, col indices and save 'sort' array in descending order of maxima
        row_sort[i], col_sort[i] = best // n, best % n
        # Overwrite this row, col pair with '-inf' so next iteration gets the next-highest value
        c1c2[row_sort[i], :], c1c2[:, col_sort[i]] = float('-inf'), float('-inf')
    return row_sort, col_sort


def alignment_score(cluster1, cluster2):
    n1, n2 = cluster1.sum(), cluster2.sum()
    row_sort, col_sort = greedy_alignment(cluster1, cluster2)
    return torch.sum(cluster1[:, row_sort] * cluster2[:, col_sort]) / min(n1, n2)


def shuffled_alignment_score(cluster1, cluster2, n_shuffle=10000):
    n = cluster1.size()[0]
    return torch.tensor([alignment_score(cluster1[torch.randperm(n), :], cluster2) for _ in trange(n_shuffle, leave=False)])


def cluster_id(cluster):
    is_dead = torch.all(cluster==0, dim=1)
    idx = torch.argmax(cluster, dim=1)
    idx[is_dead] = -1
    return idx


def sort_by_cluster(cluster, remove_dead=False):
    isort = torch.tensor(cluster).argmax(dim=1).argsort()
    if remove_dead:
        isort = torch.tensor([idx for idx in isort if cluster[idx, :].sum() != 0.])
    return isort