decagon-pytorch/tests/icosagon/test_normalize.py

from icosagon.normalize import add_eye_sparse, \
    norm_adj_mat_one_node_type_sparse, \
    norm_adj_mat_one_node_type_dense, \
    norm_adj_mat_one_node_type
import decagon_pytorch.normalize
import torch
import pytest
import numpy as np


def test_add_eye_sparse_01():
    adj_mat_dense = torch.rand((10, 10))
    adj_mat_sparse = adj_mat_dense.to_sparse()

    adj_mat_dense += torch.eye(10)
    adj_mat_sparse = add_eye_sparse(adj_mat_sparse)

    assert torch.all(adj_mat_sparse.to_dense() == adj_mat_dense)


def test_add_eye_sparse_02():
    adj_mat_dense = torch.rand((10, 20))
    adj_mat_sparse = adj_mat_dense.to_sparse()

    with pytest.raises(ValueError):
        _ = add_eye_sparse(adj_mat_sparse)


def test_add_eye_sparse_03():
    adj_mat_dense = torch.rand((10, 10))

    with pytest.raises(ValueError):
        _ = add_eye_sparse(adj_mat_dense)


def test_add_eye_sparse_04():
    adj_mat_dense = np.random.rand(10, 10)

    with pytest.raises(ValueError):
        _ = add_eye_sparse(adj_mat_dense)


def test_norm_adj_mat_one_node_type_sparse_01():
    adj_mat = torch.rand((10, 10))
    adj_mat = (adj_mat > .5)
    adj_mat = adj_mat.to_sparse()
    _ = norm_adj_mat_one_node_type_sparse(adj_mat)


def test_norm_adj_mat_one_node_type_sparse_02():
    adj_mat_dense = torch.rand((10, 10))
    adj_mat_dense = (adj_mat_dense > .5)
    adj_mat_sparse = adj_mat_dense.to_sparse()
    adj_mat_sparse = norm_adj_mat_one_node_type_sparse(adj_mat_sparse)
    adj_mat_dense = norm_adj_mat_one_node_type_dense(adj_mat_dense)
    assert torch.all(adj_mat_sparse.to_dense() == adj_mat_dense)


def test_norm_adj_mat_one_node_type_dense_01():
    adj_mat = torch.rand((10, 10))
    adj_mat = (adj_mat > .5)
    _ = norm_adj_mat_one_node_type_dense(adj_mat)


def test_norm_adj_mat_one_node_type_dense_02():
    adj_mat = torch.tensor([
        [0, 1, 1, 0], # 3
        [1, 0, 1, 0], # 3
        [1, 1, 0, 1], # 4
        [0, 0, 1, 0]  # 2
    ])
    expect = np.array([
        [1/3, 1/3, 1/3, 0],
        [1/3, 1/3, 1/3, 0],
        [1/4, 1/4, 1/4, 1/4],
        [0, 0, 1/2, 1/2]
    ], dtype=np.float32)
    res = decagon_pytorch.normalize.norm_adj_mat_one_node_type(adj_mat)
    res = res.todense().astype(np.float32)
    print('res:', res)
    print('expect:', expect)
    assert torch.all(res == expect)


@pytest.mark.skip
def test_norm_adj_mat_one_node_type_dense_03():
    adj_mat = torch.rand((10, 10))
    adj_mat = (adj_mat > .5)
    adj_mat_dec = decagon_pytorch.normalize.norm_adj_mat_one_node_type(adj_mat)
    adj_mat_ico = norm_adj_mat_one_node_type_dense(adj_mat)
    adj_mat_dec = adj_mat_dec.todense()
    adj_mat_ico = adj_mat_ico.detach().cpu().numpy()
    print('adj_mat_dec:', adj_mat_dec)
    print('adj_mat_ico:', adj_mat_ico)
    assert np.all(adj_mat_dec == adj_mat_ico)