New implementation of normalize for one/two node types and sparse/dense - good.

src/icosagon/normalize.py

@@ -9,17 +9,37 @@ import scipy.sparse as sp
 import torch
 
 
-def add_eye_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
+def _check_tensor(adj_mat):
     if not isinstance(adj_mat, torch.Tensor):
         raise ValueError('adj_mat must be a torch.Tensor')
 
+
+def _check_sparse(adj_mat):
     if not adj_mat.is_sparse:
         raise ValueError('adj_mat must be sparse')
 
+
+def _check_dense(adj_mat):
+    if adj_mat.is_sparse:
+        raise ValueError('adj_mat must be dense')
+
+
+def _check_square(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')
 
+
+def _check_2d(adj_mat):
+    if len(adj_mat.shape) != 2:
+        raise ValueError('adj_mat must be a square matrix')
+
+
+def add_eye_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_sparse(adj_mat)
+    _check_square(adj_mat)
+
     adj_mat = adj_mat.coalesce()
     indices = adj_mat.indices()
     values = adj_mat.values()
@@ -36,12 +56,42 @@ def add_eye_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
     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')
+def norm_adj_mat_one_node_type_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_sparse(adj_mat)
+    _check_square(adj_mat)
 
     adj_mat = add_eye_sparse(adj_mat)
+    adj_mat = norm_adj_mat_two_node_types_sparse(adj_mat)
+
+    return adj_mat
+
+
+def norm_adj_mat_one_node_type_dense(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_dense(adj_mat)
+    _check_square(adj_mat)
+
+    adj_mat = adj_mat + torch.eye(adj_mat.shape[0], dtype=adj_mat.dtype)
+    adj_mat = norm_adj_mat_two_node_types_dense(adj_mat)
+
+    return adj_mat
+
+
+def norm_adj_mat_one_node_type(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_square(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_two_node_types_sparse(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_sparse(adj_mat)
+    _check_2d(adj_mat)
 
     adj_mat = adj_mat.coalesce()
     indices = adj_mat.indices()
@@ -50,28 +100,17 @@ def norm_adj_mat_one_node_type_sparse(adj_mat):
     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')
+def norm_adj_mat_two_node_types_dense(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_dense(adj_mat)
+    _check_2d(adj_mat)
 
-    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)
@@ -82,29 +121,11 @@ def norm_adj_mat_one_node_type_dense(adj_mat):
     return adj_mat
 
 
-def norm_adj_mat_one_node_type(adj_mat):
+def norm_adj_mat_two_node_types(adj_mat: torch.Tensor) -> torch.Tensor:
+    _check_tensor(adj_mat)
+    _check_2d(adj_mat)
+
     if adj_mat.is_sparse:
-        return norm_adj_mat_one_node_type_sparse(adj_mat)
+        return norm_adj_mat_two_node_types_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
+        return norm_adj_mat_two_node_types_dense(adj_mat)

tests/icosagon/test_normalize.py

@@ -1,7 +1,10 @@
 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
+    norm_adj_mat_one_node_type, \
+    norm_adj_mat_two_node_types_sparse, \
+    norm_adj_mat_two_node_types_dense, \
+    norm_adj_mat_two_node_types
 import decagon_pytorch.normalize
 import torch
 import pytest
@@ -108,3 +111,75 @@ def test_norm_adj_mat_one_node_type_dense_03():
     print('adj_mat_dec:', adj_mat_dec)
     print('adj_mat_ico:', adj_mat_ico)
     assert np.all(adj_mat_dec - adj_mat_ico < 0.000001)
+
+
+def test_norm_adj_mat_two_node_types_sparse_01():
+    adj_mat = torch.rand((10, 20))
+    adj_mat = (adj_mat > .5)
+    adj_mat = adj_mat.to_sparse()
+    _ = norm_adj_mat_two_node_types_sparse(adj_mat)
+
+
+def test_norm_adj_mat_two_node_types_sparse_02():
+    adj_mat_dense = torch.rand((10, 20))
+    adj_mat_dense = (adj_mat_dense > .5)
+    adj_mat_sparse = adj_mat_dense.to_sparse()
+    adj_mat_sparse = norm_adj_mat_two_node_types_sparse(adj_mat_sparse)
+    adj_mat_dense = norm_adj_mat_two_node_types_dense(adj_mat_dense)
+    assert torch.all(adj_mat_sparse.to_dense() - adj_mat_dense < 0.000001)
+
+
+def test_norm_adj_mat_two_node_types_dense_01():
+    adj_mat = torch.rand((10, 20))
+    adj_mat = (adj_mat > .5)
+    _ = norm_adj_mat_two_node_types_dense(adj_mat)
+
+
+def test_norm_adj_mat_two_node_types_dense_02():
+    adj_mat = torch.tensor([
+        [0, 1, 1, 0], # 2
+        [1, 0, 1, 0], # 2
+        [1, 1, 0, 1], # 3
+        [0, 0, 1, 0] # 1
+    #    2  2  3  1
+    ])
+    expect_denom = np.array([
+        [ 2,           2,           sqrt(6),   sqrt(2)   ],
+        [ 2,           2,           sqrt(6),   sqrt(2)   ],
+        [ sqrt(6),     sqrt(6),     3,         sqrt(3)   ],
+        [ sqrt(2),     sqrt(2),     sqrt(3),   1         ]
+    ], dtype=np.float32)
+    expect = adj_mat.detach().cpu().numpy().astype(np.float32) / expect_denom
+    res = decagon_pytorch.normalize.norm_adj_mat_two_node_types(adj_mat)
+    res = res.todense().astype(np.float32)
+    print('res:', res)
+    print('expect:', expect)
+    assert np.all(res - expect < 0.000001)
+
+
+def test_norm_adj_mat_two_node_types_dense_03():
+    adj_mat = torch.tensor([
+        [0, 1, 1,  0,  0],
+        [1, 0, 1,  0,  1],
+        [1, 1, 0, .5, .5],
+        [0, 0, .5, 0,  1],
+        [0, 1, .5, 1,  0]
+    ])
+    adj_mat_dec = decagon_pytorch.normalize.norm_adj_mat_two_node_types(adj_mat)
+    adj_mat_ico = norm_adj_mat_two_node_types_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 < 0.000001)
+
+
+def test_norm_adj_mat_two_node_types_dense_04():
+    adj_mat = torch.rand((10, 20))
+    adj_mat_dec = decagon_pytorch.normalize.norm_adj_mat_two_node_types(adj_mat)
+    adj_mat_ico = norm_adj_mat_two_node_types_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 < 0.000001)