Improve code coverage in test

pygmtools/jittor_backend.py

@@ -420,7 +420,7 @@ def _comp_aff_score(x, k):
                 idx[i].append(ix[0].item() if ix.shape[0]>1 else ix.item())
         idx = jt.Var(idx)
 
-        assert jt.all(score_combo_max >= score_ori), jt.min(score_combo_max - score_ori)
+        assert jt.all(score_combo_max + 1e-4 >= score_ori), jt.min(score_combo_max - score_ori)
         X_upt = X_combo[mask1, mask2, idx, :, :]
         X = X_upt * X_mask + X_upt.transpose(0, 1).transpose(2, 3) * X_mask.transpose(0, 1) + X * (1 - X_mask - X_mask.transpose(0, 1))
         assert jt.all(X.transpose(0, 1).transpose(2, 3) == X)

pygmtools/linear_solvers.py

@@ -647,7 +647,7 @@ def sinkhorn(s, n1=None, n2=None, unmatch1=None, unmatch2=None,
             unmatch1 = _unsqueeze(unmatch1, 0, backend)
             unmatch2 = _unsqueeze(unmatch2, 0, backend)
         if not _check_shape(unmatch1, 2, backend) or not _check_shape(unmatch2, 2, backend):
-            raise ValueError(f'the input argument unmatch1 and unmatch2 are illegal. They should be 2-dim'
+            raise ValueError(f'the input arguments unmatch1 and unmatch2 are illegal. They should be 2-dim'
                              f'for batched input, and 1-dim for non-batched input.')
         if not all((_get_shape(unmatch1, backend)[1] == _get_shape(s, backend)[1],
                     _get_shape(unmatch2, backend)[1] == _get_shape(s, backend)[2],
@@ -1141,7 +1141,7 @@ def hungarian(s, n1=None, n2=None, unmatch1=None, unmatch2=None,
             unmatch1 = _unsqueeze(unmatch1, 0, backend)
             unmatch2 = _unsqueeze(unmatch2, 0, backend)
         if not _check_shape(unmatch1, 2, backend) or not _check_shape(unmatch2, 2, backend):
-            raise ValueError(f'the input argument unmatch1 and unmatch2 are illegal. They should be 2-dim'
+            raise ValueError(f'the input arguments unmatch1 and unmatch2 are illegal. They should be 2-dim'
                              f'for batched input, and 1-dim for non-batched input.')
         if not all((_get_shape(unmatch1, backend)[1] == _get_shape(s, backend)[1],
                     _get_shape(unmatch2, backend)[1] == _get_shape(s, backend)[2],

pygmtools/neural_solvers.py

@@ -15,7 +15,9 @@
 
 import importlib
 import pygmtools
-from pygmtools.utils import NOT_IMPLEMENTED_MSG, _check_shape, _get_shape, _unsqueeze, _squeeze, _check_data_type
+import numpy as np
+from pygmtools.utils import NOT_IMPLEMENTED_MSG, from_numpy, \
+    _check_shape, _get_shape, _unsqueeze, _squeeze, _check_data_type
 from pygmtools.classic_solvers import __check_gm_arguments
 
 
@@ -239,6 +241,8 @@ def pca_gm(feat1, feat2, A1, A2, n1=None, n2=None,
 
         if all([_check_shape(_, 2, backend) for _ in (feat1, feat2, A1, A2)]):
             feat1, feat2, A1, A2 = [_unsqueeze(_, 0, backend) for _ in (feat1, feat2, A1, A2)]
+            if type(n1) is int: n1 = from_numpy(np.array([n1]), backend=backend)
+            if type(n2) is int: n2 = from_numpy(np.array([n2]), backend=backend)
             non_batched_input = True
         elif all([_check_shape(_, 3, backend) for _ in (feat1, feat2, A1, A2)]):
             non_batched_input = False
@@ -255,6 +259,8 @@ def pca_gm(feat1, feat2, A1, A2, n1=None, n2=None,
             raise ValueError(
                 f'the input dimensions do not match. Got feat1:{_get_shape(feat1, backend)}, '
                 f'feat2:{_get_shape(feat2, backend)}, A1:{_get_shape(A1, backend)}, A2:{_get_shape(A2, backend)}!')
+    if n1 is not None: _check_data_type(n1, 'n1', backend)
+    if n2 is not None: _check_data_type(n2, 'n2', backend)
 
     args = (feat1, feat2, A1, A2, n1, n2, in_channel, hidden_channel, out_channel, num_layers, sk_max_iter, sk_tau,
            network, pretrain)
@@ -500,6 +506,8 @@ def ipca_gm(feat1, feat2, A1, A2, n1=None, n2=None,
 
         if all([_check_shape(_, 2, backend) for _ in (feat1, feat2, A1, A2)]):
             feat1, feat2, A1, A2 = [_unsqueeze(_, 0, backend) for _ in (feat1, feat2, A1, A2)]
+            if type(n1) is int: n1 = from_numpy(np.array([n1]), backend=backend)
+            if type(n2) is int: n2 = from_numpy(np.array([n2]), backend=backend)
             non_batched_input = True
         elif all([_check_shape(_, 3, backend) for _ in (feat1, feat2, A1, A2)]):
             non_batched_input = False
@@ -516,6 +524,8 @@ def ipca_gm(feat1, feat2, A1, A2, n1=None, n2=None,
             raise ValueError(
                 f'the input dimensions do not match. Got feat1:{_get_shape(feat1, backend)}, '
                 f'feat2:{_get_shape(feat2, backend)}, A1:{_get_shape(A1, backend)}, A2:{_get_shape(A2, backend)}!')
+    if n1 is not None: _check_data_type(n1, 'n1', backend)
+    if n2 is not None: _check_data_type(n2, 'n2', backend)
 
     args = (feat1, feat2, A1, A2, n1, n2, in_channel, hidden_channel, out_channel, num_layers, cross_iter,
             sk_max_iter, sk_tau, network, pretrain)
@@ -774,6 +784,8 @@ def cie(feat_node1, feat_node2, A1, A2, feat_edge1, feat_edge2, n1=None, n2=None
                 and all([_check_shape(_, 3, backend) for _ in (feat_edge1, feat_edge2)]):
             feat_node1, feat_node2, A1, A2, feat_edge1, feat_edge2 =\
                 [_unsqueeze(_, 0, backend) for _ in (feat_node1, feat_node2, A1, A2, feat_edge1, feat_edge2)]
+            if type(n1) is int: n1 = from_numpy(np.array([n1]), backend=backend)
+            if type(n2) is int: n2 = from_numpy(np.array([n2]), backend=backend)
             non_batched_input = True
         elif all([_check_shape(_, 3, backend) for _ in (feat_node1, feat_node2, A1, A2)]) \
                 and all([_check_shape(_, 4, backend) for _ in (feat_edge1, feat_edge2)]):
@@ -798,6 +810,8 @@ def cie(feat_node1, feat_node2, A1, A2, feat_edge1, feat_edge2, n1=None, n2=None
                 f'the input dimensions do not match. Got feat_node1:{_get_shape(feat_node1, backend)}, '
                 f'feat_node2:{_get_shape(feat_node2, backend)}, A1:{_get_shape(A1, backend)}, A2:{_get_shape(A2, backend)},'
                 f'feat_edge1:{_get_shape(feat_edge1, backend)}, feat_edge2:{_get_shape(feat_edge2, backend)}!')
+    if n1 is not None: _check_data_type(n1, 'n1', backend)
+    if n2 is not None: _check_data_type(n2, 'n2', backend)
 
     args = (feat_node1, feat_node2, A1, A2, feat_edge1, feat_edge2, n1, n2,
             in_node_channel, in_edge_channel, hidden_channel, out_channel, num_layers,

pygmtools/numpy_backend.py

@@ -443,7 +443,7 @@ def _comp_aff_score(x, k):
         idx = np.argmax(score_combo,axis=-1)
         score_combo = np.max(score_combo, axis=-1)
 
-        assert np.all(score_combo >= score_ori), np.min(score_combo - score_ori)
+        assert np.all(score_combo + 1e-4 >= score_ori), np.min(score_combo - score_ori)
         X_upt = X_combo[mask1, mask2, idx, :, :]
         X = X_upt * X_mask + X_upt.swapaxes(0,1).swapaxes(2,3) * X_mask.swapaxes(0,1) + X * (1 - X_mask - X_mask.swapaxes(0, 1))
         assert np.all(X.swapaxes(0,1).swapaxes(2,3) == X)
@@ -764,14 +764,16 @@ def gamgm_real(
                 UUt = np.matmul(U_hung, U_hung.transpose())
                 cluster_weight = np.repeat(cluster_M, ns.astype('i4'), axis=0)
                 cluster_weight = np.repeat(cluster_weight, ns.astype('i4'), axis=1)
-                quad = np.linalg.multi_dot(supA, UUt * cluster_weight, supA, U_hung) * quad_weight * 2
+                quad = np.linalg.multi_dot((supA, UUt * cluster_weight, supA, U_hung)) * quad_weight * 2
                 unary = np.matmul(supW * cluster_weight, U_hung)
-                max_vals = (unary + quad).max(axis=1).values
+                max_vals = (unary + quad).max(axis=1)
                 U = U * (unary + quad > outlier_thresh)
                 if verbose:
                     print(f'hungarian #iter={i}/{max_iter} '
-                          f'unary+quad score thresh={outlier_thresh:.3f}, #>thresh={np.sum(max_vals > outlier_thresh)}/{max_vals.shape[0]}'
-                          f' min:{max_vals.min():.4f}, mean:{max_vals.mean():.4f}, median:{max_vals.median():.4f}, max:{max_vals.max():.4f}')
+                          f'unary+quad score thresh={outlier_thresh:.3f}, '
+                          f'#>thresh={np.sum(max_vals > outlier_thresh)}/{max_vals.shape[0]} '
+                          f'min:{max_vals.min():.4f}, mean:{max_vals.mean():.4f}, '
+                          f'median:{np.median(max_vals):.4f}, max:{max_vals.max():.4f}')
 
             if np.linalg.norm(np.matmul(U, U.T) - lastUUt) < converge_thresh:
                 break

pygmtools/paddle_backend.py

@@ -431,7 +431,7 @@ def _comp_aff_score(x, k):
         idx = paddle.argmax(score_combo, axis=-1)
         score_combo = paddle.max(score_combo, axis=-1)
 
-        # assert paddle.all(score_combo >= score_ori), paddle.min(score_combo - score_ori)
+        assert paddle.all(score_combo + 1e-4 >= score_ori), paddle.min(score_combo - score_ori)
         X_upt = X_combo[mask1, mask2, idx]
         X = X_upt * X_mask + X_upt.transpose((1, 0, 3, 2))* X_mask.transpose((1, 0, 2, 3)) + X * (1 - X_mask - X_mask.transpose((1, 0, 2, 3)))
         assert paddle.all(X.transpose((1, 0, 3, 2)) == X)
@@ -749,8 +749,12 @@ def gamgm_real(
                 U = U * (unary + quad > outlier_thresh)
                 if verbose:
                     print(f'hungarian #iter={i}/{max_iter} '
-                          f'unary+quad score thresh={outlier_thresh:.3f}, #>thresh={paddle.sum(max_vals > outlier_thresh)}/{max_vals.shape[0]}'
-                          f' min:{max_vals.min():.4f}, mean:{max_vals.mean():.4f}, median:{max_vals.median():.4f}, max:{max_vals.max():.4f}')
+                          f'unary+quad score thresh={outlier_thresh:.3f}, '
+                          f'#>thresh={paddle.sum(max_vals > outlier_thresh).numpy().squeeze()}/{max_vals.shape[0]} '
+                          f'min:{max_vals.min().numpy().squeeze():.4f}, '
+                          f'mean:{max_vals.mean().numpy().squeeze():.4f}, '
+                          f'median:{max_vals.median().numpy().squeeze():.4f}, '
+                          f'max:{max_vals.max().numpy().squeeze():.4f}')
 
             if paddle.linalg.norm(paddle.mm(U, U.t()) - lastUUt) < converge_thresh:
                 break

pygmtools/pytorch_backend.py

@@ -431,7 +431,7 @@ def _comp_aff_score(x, k):
 
         score_combo, idx = torch.max(score_combo, dim=-1)
 
-        assert torch.all(score_combo >= score_ori), torch.min(score_combo - score_ori)
+        assert torch.all(score_combo + 1e-4 >= score_ori), torch.min(score_combo - score_ori)
         X_upt = X_combo[mask1, mask2, idx, :, :]
         X = X_upt * X_mask + X_upt.transpose(0, 1).transpose(2, 3) * X_mask.transpose(0, 1) + X * (1 - X_mask - X_mask.transpose(0, 1))
         assert torch.all(X.transpose(0, 1).transpose(2, 3) == X)

tests/test_classic_solvers.py

@@ -142,14 +142,16 @@ def _test_classic_solver_on_linear_assignment(num_nodes1, num_nodes2, node_feat_
         F1, F2, X_gt = (pygm.utils.build_batch(_) for _ in (F1, F2, X_gt))
         if batch_size > 1:
             F1, F2, n1, n2, X_gt = data_to_numpy(F1, F2, n1, n2, X_gt)
+            if unmatch:
+                unmatch1, unmatch2 = (pygm.utils.build_batch(_) for _ in (unmatch1, unmatch2))
+                unmatch1, unmatch2 = data_to_numpy(unmatch1, unmatch2)
         else:
             F1, F2, n1, n2, X_gt = data_to_numpy(
                 F1.squeeze(0), F2.squeeze(0), n1, n2, X_gt.squeeze(0)
             )
-
-        if unmatch:
-            unmatch1, unmatch2 = (pygm.utils.build_batch(_) for _ in (unmatch1, unmatch2))
-            unmatch1, unmatch2 = data_to_numpy(unmatch1, unmatch2)
+            if unmatch:
+                unmatch1, unmatch2 = (pygm.utils.build_batch(_) for _ in (unmatch1, unmatch2))
+                unmatch1, unmatch2 = data_to_numpy(unmatch1.squeeze(0), unmatch2.squeeze(0))
 
         last_X = None
         for working_backend in backends:
@@ -200,7 +202,7 @@ def test_hungarian():
     # non-batched input
     _test_classic_solver_on_linear_assignment([10], [30], 10, pygm.hungarian, {
         'nproc': [1],
-        'outlier_num': [0]
+        'outlier_num': [0, 5]
     }, ['pytorch', 'numpy', 'paddle' ,'jittor','tensorflow'])
 
 
@@ -245,6 +247,7 @@ def test_sinkhorn():
         'max_iter': [500],
         'batched_operation': [True],
         'dummy_row': [True],
+        'outlier_num': [0, 5]
     }, ['pytorch', 'numpy', 'paddle', 'jittor','tensorflow'])
 
     _test_classic_solver_on_linear_assignment(*args1)

tests/test_multi_graph_solvers.py

@@ -203,6 +203,7 @@ def test_gamgm():
             'node_aff_fn': [functools.partial(pygm.utils.gaussian_aff_fn, sigma=.1)],
             'verbose': [True],
             'n_univ': [10],
+            'outlier_thresh': [0., 0.1],
             'ns': [np.array([num_nodes] * (num_graphs // 2) + [num_nodes-1] * (num_graphs - num_graphs // 2))],
         }, ['pytorch', 'numpy', 'paddle', 'jittor'])
 

tests/test_neural_solvers.py

@@ -74,6 +74,8 @@ def _test_neural_solver_on_isomorphic_graphs(graph_num_nodes, node_feat_dim, sol
         for working_backend in backends:
             pygm.BACKEND = working_backend
             _A1, _A2, _F1, _F2, _EF1, _EF2, _n1, _n2 = data_from_numpy(A1, A2, F1, F2, EF1, EF2, n1, n2)
+            if batch_size == 1:
+                _n1, _n2 = _n1.item(), _n2.item()
 
             if mode == 'lawler-qap':
                 if batch_size > 1:
@@ -82,7 +84,6 @@ def _test_neural_solver_on_isomorphic_graphs(graph_num_nodes, node_feat_dim, sol
                     _K = pygm.utils.build_aff_mat(_F1, _edge1, _conn1, _F2, _edge2, _conn2, _n1, _ne1, _n2, _ne2,
                                                   **aff_param_dict)
                 else:
-                    _n1, _n2 = _n1.item(), _n2.item()
                     _conn1, _edge1 = pygm.utils.dense_to_sparse(_A1)
                     _conn2, _edge2 = pygm.utils.dense_to_sparse(_A2)
                     _K = pygm.utils.build_aff_mat(_F1, _edge1, _conn1, _F2, _edge2, _conn2, _n1, None, _n2, None,
@@ -112,6 +113,9 @@ def _test_neural_solver_on_isomorphic_graphs(graph_num_nodes, node_feat_dim, sol
                 f"params: {';'.join([k + '=' + str(v) for k, v in aff_param_dict.items()])};" \
                 f"{';'.join([k + '=' + str(v) for k, v in solver_param_dict.items()])}"
 
+            if 'pretrain' in solver_param_dict and solver_param_dict['pretrain'] is None:
+                _X1 = pygm.hungarian(_X1, _n1, _n2)
+
             if last_X is not None:
                 assert np.abs(pygm.utils.to_numpy(_X1) - last_X).sum() < 5e-3, \
                     f"Incorrect GM solution for {working_backend}, " \
@@ -124,6 +128,7 @@ def _test_neural_solver_on_isomorphic_graphs(graph_num_nodes, node_feat_dim, sol
                                   f"params: {';'.join([k + '=' + str(v) for k, v in aff_param_dict.items()])};" \
                                   f"{';'.join([k + '=' + str(v) for k, v in solver_param_dict.items()])}"
 
+
 def test_pca_gm():
     _test_neural_solver_on_isomorphic_graphs(list(range(10, 30, 2)), 1024, pygm.pca_gm, 'individual-graphs', {
         'pretrain': ['voc', 'willow', 'voc-all'],
@@ -134,6 +139,13 @@ def test_pca_gm():
         'pretrain': ['voc'],
     }, ['pytorch', 'numpy', 'jittor'])
 
+    # test more layers
+    _test_neural_solver_on_isomorphic_graphs([10], 1024, pygm.pca_gm, 'individual-graphs', {
+        'num_layers': [3],
+        'pretrain': [None],
+    }, ['pytorch', 'numpy', 'jittor'])
+
+
 def test_ipca_gm():
     _test_neural_solver_on_isomorphic_graphs(list(range(10, 30, 2)), 1024, pygm.ipca_gm, 'individual-graphs', {
         'pretrain': ['voc', 'willow'],
@@ -144,6 +156,13 @@ def test_ipca_gm():
         'pretrain': ['voc'],
     }, ['pytorch', 'numpy', 'jittor'])
 
+    # test more layers
+    _test_neural_solver_on_isomorphic_graphs([10], 1024, pygm.ipca_gm, 'individual-graphs', {
+        'num_layers': [3],
+        'pretrain': [None],
+    }, ['pytorch', 'numpy', 'jittor'])
+
+
 def test_cie():
     _test_neural_solver_on_isomorphic_graphs(list(range(10, 30, 2)), 1024, pygm.cie, 'individual-graphs-edge', {
             'pretrain': ['voc', 'willow'],
@@ -154,6 +173,12 @@ def test_cie():
         'pretrain': ['voc'],
     }, ['pytorch', 'numpy', 'jittor'])
 
+    # test more layers
+    _test_neural_solver_on_isomorphic_graphs([10], 1024, pygm.cie, 'individual-graphs-edge', {
+        'num_layers': [3],
+        'pretrain': [None],
+    }, ['pytorch', 'numpy', 'jittor'])
+
 def test_ngm():
     _test_neural_solver_on_isomorphic_graphs(list(range(10, 30, 2)), 1024, pygm.ngm, 'lawler-qap', {
         'edge_aff_fn': [functools.partial(pygm.utils.gaussian_aff_fn, sigma=1.), pygm.utils.inner_prod_aff_fn],
@@ -168,8 +193,9 @@ def test_ngm():
         'pretrain': ['voc'],
     }, ['pytorch', 'numpy', 'jittor'])
 
+
 if __name__ == '__main__':
-    # test_pca_gm()
-    # test_ipca_gm()
-    # test_cie()
+    test_pca_gm()
+    test_ipca_gm()
+    test_cie()
     test_ngm()