nystrom-ncut 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

nystrom_ncut/propagation_utils.py

@@ -96,7 +96,6 @@ def distance_from_features(
         D = D / (2 * features.var(dim=0).sum())
     else:
         raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
-
     return D
 
 
@@ -184,13 +183,37 @@ def propagate_knn(
     V_list = []
     for _v in torch.chunk(inp_features, n_chunks, dim=0):
         _v = _v.to(device)
-        _A = affinity_from_features(subgraph_features, _v, affinity_focal_gamma, distance).mT
 
-        if knn is not None:
-            mask = torch.full_like(_A, True, dtype=torch.bool)
-            mask[torch.arange(len(_v))[:, None], _A.topk(knn, dim=-1, largest=True).indices] = False
-            _A[mask] = 0.0
-        _A = F.normalize(_A, p=1, dim=-1)
+        # _A = affinity_from_features(subgraph_features, _v, affinity_focal_gamma, distance).mT
+        # if knn is not None:
+        #     mask = torch.full_like(_A, True, dtype=torch.bool)
+        #     mask[torch.arange(len(_v))[:, None], _A.topk(knn, dim=-1, largest=True).indices] = False
+        #     _A[mask] = 0.0
+        # _A = F.normalize(_A, p=1, dim=-1)
+
+        if distance == 'cosine':
+            _A = _v @ subgraph_features.T
+        elif distance == 'euclidean':
+            _A = - torch.cdist(_v, subgraph_features, p=2)
+        elif distance == 'rbf':
+            _A = - torch.cdist(_v, subgraph_features, p=2) ** 2
+        else:
+            raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
+
+        # keep topk KNN for each row
+        topk_sim, topk_idx = _A.topk(knn, dim=-1, largest=True)
+        row_id = torch.arange(topk_idx.shape[0], device=_A.device)[:, None].expand(
+            -1, topk_idx.shape[1]
+        )
+        _A = torch.sparse_coo_tensor(
+            torch.stack([row_id, topk_idx], dim=-1).reshape(-1, 2).T,
+            topk_sim.reshape(-1),
+            size=(_A.shape[0], _A.shape[1]),
+            device=_A.device,
+        )
+        _A = _A.to_dense().to(dtype=subgraph_output.dtype)
+        _D = _A.sum(-1)
+        _A /= _D[:, None]
 
         _V = _A @ subgraph_output
         if move_output_to_cpu:

{nystrom_ncut-0.0.5.dist-info → nystrom_ncut-0.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.0.5
+Version: 0.0.6
 Summary: Normalized Cut and Nyström Approximation
 Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
 Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/

{nystrom_ncut-0.0.5.dist-info → nystrom_ncut-0.0.6.dist-info}/RECORD

@@ -2,10 +2,10 @@ nystrom_ncut/__init__.py,sha256=Cww-_OsyQHLKpgw_Wh28_tUOvIMMr7Ey8w-tH7v99xQ,452
 nystrom_ncut/common.py,sha256=qdR_JwknT9H1Cv5LopwdwZfORFx-O8MLiRI6ZF1Qohc,558
 nystrom_ncut/ncut_pytorch.py,sha256=wRQXUPBOW2_vutocKf0J19HrFVkBYQePAYUEfotLfx4,11701
 nystrom_ncut/nystrom.py,sha256=HbwON9pLW3gEZvOmbDJwkQNHolOo1EBvwBPeh2p2uJE,8833
-nystrom_ncut/propagation_utils.py,sha256=mD6rZ_mwYjYXs1cp5ZaTK0FrJ4YhyCdoIUrdGRP9k-M,12119
+nystrom_ncut/propagation_utils.py,sha256=OCqnv7P9kzDlwqeJzNWpJ3TdTEpk7AD0rJhX8MazZYs,13061
 nystrom_ncut/visualize_utils.py,sha256=QmBatlX7Q-ZWF_iJ1zFDnPHFuofz3tCmtoNeeoMPw3U,18558
-nystrom_ncut-0.0.5.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.0.5.dist-info/METADATA,sha256=n9zlRYBD02k478INScrj9V9rZ1mhXTylcMjkmQDgl1A,6058
-nystrom_ncut-0.0.5.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-nystrom_ncut-0.0.5.dist-info/top_level.txt,sha256=j7g_j0S048EvguFFnGgD5Ewd3r2H6klsxd5A4dd-wHw,13
-nystrom_ncut-0.0.5.dist-info/RECORD,,
+nystrom_ncut-0.0.6.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.0.6.dist-info/METADATA,sha256=FD53Ov3g9u4tbBP_Sxxd2hf1yUdg_Hy3ShWq_xGOZFA,6058
+nystrom_ncut-0.0.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.0.6.dist-info/top_level.txt,sha256=j7g_j0S048EvguFFnGgD5Ewd3r2H6klsxd5A4dd-wHw,13
+nystrom_ncut-0.0.6.dist-info/RECORD,,

{nystrom_ncut-0.0.5.dist-info → nystrom_ncut-0.0.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.7.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any