New implementation for one node type normalization seems to work.

src/icosagon/normalize.py

@@ -46,11 +46,15 @@ def norm_adj_mat_one_node_type_sparse(adj_mat):
     adj_mat = adj_mat.coalesce()
     indices = adj_mat.indices()
     values = adj_mat.values()
-    degrees = torch.zeros(adj_mat.shape[0])
-    degrees = degrees.index_add(0, indices[0], values.to(degrees.dtype))
-    print('degrees:', degrees)
-    print('values:', values)
-    values = values.to(degrees.dtype) / degrees[indices[0]]
+    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
@@ -68,8 +72,12 @@ def norm_adj_mat_one_node_type_dense(adj_mat):
         raise ValueError('adj_mat must be a square matrix')
 
     adj_mat = adj_mat + torch.eye(adj_mat.shape[0], dtype=adj_mat.dtype)
-    degrees = adj_mat.sum(1).view(-1, 1).to(torch.float32)
-    adj_mat = adj_mat.to(degrees.dtype) / degrees
+    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
 

tests/icosagon/test_normalize.py

@@ -6,6 +6,7 @@ import decagon_pytorch.normalize
 import torch
 import pytest
 import numpy as np
+from math import sqrt
 
 
 def test_add_eye_sparse_01():
@@ -53,7 +54,7 @@ def test_norm_adj_mat_one_node_type_sparse_02():
     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)
+    assert torch.all(adj_mat_sparse.to_dense() - adj_mat_dense < 0.000001)
 
 
 def test_norm_adj_mat_one_node_type_dense_01():
@@ -68,28 +69,42 @@ def test_norm_adj_mat_one_node_type_dense_02():
         [1, 0, 1, 0], # 3
         [1, 1, 0, 1], # 4
         [0, 0, 1, 0]  # 2
+    #    3  3  4  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]
+    expect_denom = np.array([
+        [ 3,           3,           sqrt(3)*2,   sqrt(6)   ],
+        [ 3,           3,           sqrt(3)*2,   sqrt(6)   ],
+        [ sqrt(3)*2,   sqrt(3)*2,   4,           sqrt(2)*2 ],
+        [ sqrt(6),     sqrt(6),     sqrt(2)*2,   2         ]
     ], dtype=np.float32)
+    expect = (adj_mat.detach().cpu().numpy().astype(np.float32) + np.eye(4)) / expect_denom
+    # 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)
+    assert np.all(res - expect < 0.000001)
 
 
-@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 = torch.rand((10, 10))
+    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 = (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)
+    assert np.all(adj_mat_dec - adj_mat_ico < 0.000001)