decagon-pytorch/src/icosagon/normalize.py

#
# Copyright (C) Stanislaw Adaszewski, 2020
# License: GPLv3
#


import numpy as np
import scipy.sparse as sp
import torch


def add_eye_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
    if not isinstance(adj_mat, torch.Tensor):
        raise ValueError('adj_mat must be a torch.Tensor')

    if not adj_mat.is_sparse:
        raise ValueError('adj_mat must be sparse')

    if len(adj_mat.shape) != 2 or \
        adj_mat.shape[0] != adj_mat.shape[1]:
        raise ValueError('adj_mat must be a square matrix')

    adj_mat = adj_mat.coalesce()
    indices = adj_mat.indices()
    values = adj_mat.values()

    eye_indices = torch.arange(adj_mat.shape[0], dtype=indices.dtype).view(1, -1)
    eye_indices = torch.cat((eye_indices, eye_indices), 0)
    eye_values = torch.ones(adj_mat.shape[0], dtype=values.dtype)

    indices = torch.cat((indices, eye_indices), 1)
    values = torch.cat((values, eye_values), 0)

    adj_mat = torch.sparse_coo_tensor(indices=indices, values=values, size=adj_mat.shape)

    return adj_mat


def norm_adj_mat_one_node_type_sparse(adj_mat):
    if len(adj_mat.shape) != 2 or \
        adj_mat.shape[0] != adj_mat.shape[1]:
        raise ValueError('adj_mat must be a square matrix')

    adj_mat = add_eye_sparse(adj_mat)

    adj_mat = adj_mat.coalesce()
    indices = adj_mat.indices()
    values = adj_mat.values()
    degrees_row = torch.zeros(adj_mat.shape[0])
    degrees_row = degrees_row.index_add(0, indices[0], values.to(degrees_row.dtype))
    degrees_col = torch.zeros(adj_mat.shape[1])
    degrees_col = degrees_col.index_add(0, indices[1], values.to(degrees_col.dtype))
    # degrees_row = torch.sqrt(degrees_row)
    # degrees_col = torch.sqrt(degrees_col)
    # print('degrees:', degrees)
    # print('values:', values)
    values = values.to(degrees_row.dtype) / torch.sqrt(degrees_row[indices[0]] * degrees_col[indices[1]])
    adj_mat = torch.sparse_coo_tensor(indices=indices, values=values, size=adj_mat.shape)

    return adj_mat


def norm_adj_mat_one_node_type_dense(adj_mat):
    if not isinstance(adj_mat, torch.Tensor):
        raise ValueError('adj_mat must be a torch.Tensor')

    if adj_mat.is_sparse:
        raise ValueError('adj_mat must be dense')

    if len(adj_mat.shape) != 2 or \
        adj_mat.shape[0] != adj_mat.shape[1]:
        raise ValueError('adj_mat must be a square matrix')

    adj_mat = adj_mat + torch.eye(adj_mat.shape[0], dtype=adj_mat.dtype)
    degrees_row = adj_mat.sum(1).view(-1, 1).to(torch.float32)
    degrees_col = adj_mat.sum(0).view(1, -1).to(torch.float32)
    degrees_row = torch.sqrt(degrees_row)
    degrees_col = torch.sqrt(degrees_col)
    adj_mat = adj_mat.to(degrees_row.dtype) / degrees_row
    adj_mat = adj_mat / degrees_col

    return adj_mat


def norm_adj_mat_one_node_type(adj_mat):
    if adj_mat.is_sparse:
        return norm_adj_mat_one_node_type_sparse(adj_mat)
    else:
        return norm_adj_mat_one_node_type_dense(adj_mat)


# def norm_adj_mat_one_node_type(adj):
#     adj = sp.coo_matrix(adj)
#     assert adj.shape[0] == adj.shape[1]
#     adj_ = adj + sp.eye(adj.shape[0])
#     rowsum = np.array(adj_.sum(1))
#     degree_mat_inv_sqrt = np.power(rowsum, -0.5).flatten()
#     degree_mat_inv_sqrt = sp.diags(degree_mat_inv_sqrt)
#     adj_normalized = adj_.dot(degree_mat_inv_sqrt).transpose().dot(degree_mat_inv_sqrt)
#     return adj_normalized


def norm_adj_mat_two_node_types(adj):
    adj = sp.coo_matrix(adj)
    rowsum = np.array(adj.sum(1))
    colsum = np.array(adj.sum(0))
    rowdegree_mat_inv = sp.diags(np.nan_to_num(np.power(rowsum, -0.5)).flatten())
    coldegree_mat_inv = sp.diags(np.nan_to_num(np.power(colsum, -0.5)).flatten())
    adj_normalized = rowdegree_mat_inv.dot(adj).dot(coldegree_mat_inv).tocoo()
    return adj_normalized