utils.py

# -*- coding: utf-8 -*-
import numpy as np
import torch
import networkx as nx
from tqdm import tqdm
from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score
from scipy.optimize import linear_sum_assignment


######################################## Clustering Metrics ########################################

def cluster_acc(y_pred, y_true):
    """Clustering accuracy with Hungarian algorithm"""
    y_true = y_true.astype(np.int64)
    y_pred = y_pred.astype(np.int64)

    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=np.int64)

    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    row_ind, col_ind = linear_sum_assignment(w.max() - w)
    return w[row_ind, col_ind].sum() / y_pred.size


def nmi(y_pred, y_true):
    """Normalized Mutual Information"""
    return normalized_mutual_info_score(y_true, y_pred, average_method='arithmetic')


def ari(y_pred, y_true):
    """Adjusted Rand Index"""
    return adjusted_rand_score(y_true, y_pred)


######################################## Graph Embedding Utils ########################################

def compute_position_and_hop_ids(adj_matrix, k=15, max_hop=99):
    """Compute position IDs and hop distance IDs from adjacency matrix"""
    if isinstance(adj_matrix, torch.Tensor):
        adj_dense = adj_matrix.cpu().numpy()
    else:
        adj_dense = adj_matrix

    N = adj_dense.shape[0]

    # Select Top-k neighbors by weight
    neighbor_indices_list = []
    for i in range(N):
        neighbors = adj_dense[i, :]
        top_k_idx = np.argsort(neighbors)[::-1][:k]
        neighbor_indices_list.append(top_k_idx)

    neighbor_indices = np.array(neighbor_indices_list)

    # Position IDs
    init_pos_ids = np.zeros((N, k+1), dtype=np.int64)
    for i in range(1, k+1):
        init_pos_ids[:, i] = i

    # Hop distance IDs using BFS
    hop_dis_ids = np.zeros((N, k+1), dtype=np.int64)

    edges = [(i, j) for i in range(N) for j in range(N) if adj_dense[i, j] > 0]
    G = nx.Graph()
    G.add_nodes_from(range(N))
    G.add_edges_from(edges)

    for node in tqdm(range(N), desc="Computing hop distances", leave=False):
        neighbors = neighbor_indices[node]
        for i, neighbor in enumerate(neighbors):
            try:
                hop = nx.shortest_path_length(G, source=node, target=neighbor)
            except nx.NetworkXNoPath:
                hop = max_hop
            hop_dis_ids[node, i+1] = min(hop, max_hop)

    return (torch.from_numpy(init_pos_ids).long(),
            torch.from_numpy(hop_dis_ids).long(),
            torch.from_numpy(neighbor_indices).long())


def construct_subgraph_features(X, neighbor_indices, k):
    """Construct subgraph features: [center node + Top-k neighbors]"""
    if isinstance(X, np.ndarray):
        X = torch.from_numpy(X).float()

    N, F = X.shape
    subgraph_features = torch.zeros(N, k+1, F, device=X.device)
    subgraph_features[:, 0, :] = X

    if neighbor_indices.device != X.device:
        neighbor_indices = neighbor_indices.to(X.device)

    subgraph_features[:, 1:, :] = X[neighbor_indices]
    return subgraph_features


def prepare_dual_view_data(X, adj1, adj2, k=15):
    """Prepare dual-view input data"""
    print("  - Processing view 1...")
    pos_ids_1, hop_ids_1, neighbor_idx_1 = compute_position_and_hop_ids(adj1, k=k)
    raw_features_1 = construct_subgraph_features(X, neighbor_idx_1, k=k)

    print("  - Processing view 2...")
    pos_ids_2, hop_ids_2, neighbor_idx_2 = compute_position_and_hop_ids(adj2, k=k)
    raw_features_2 = construct_subgraph_features(X, neighbor_idx_2, k=k)

    return {
        'raw_embeddings': raw_features_1,
        'raw_embeddings2': raw_features_2,
        'neighbor_indices': neighbor_idx_1,
        'neighbor_indices2': neighbor_idx_2,
        'int_embeddings': pos_ids_1,
        'hop_embeddings': hop_ids_1,
        'int_embeddings2': pos_ids_2,
        'hop_embeddings2': hop_ids_2
    }