Fix GRASP's dead similar-patient graph

pyhealth/models/grasp.py

@@ -254,8 +254,6 @@ def __init__(
         self.GCN.initialize_parameters()
         self.GCN_2 = GraphConvolution(self.hidden_dim, self.hidden_dim, bias=True)
         self.GCN_2.initialize_parameters()
-        self.A_mat = None
-
         self.bn = nn.BatchNorm1d(self.hidden_dim)
 
     def sample_gumbel(self, shape, eps=1e-20):
@@ -310,12 +308,16 @@ def grasp_encoder(
 
         centers, codes = cluster(hidden_t, self.cluster_num, input.device)
 
-        if self.A_mat is None:
+        # Build the similar-cluster kNN graph (the point of GRASP).
+        # k must be < number of nodes (cluster_num); fall back to 
+        # identity only when there are too few clusters to form a neighbor graph.
+        k = min(20, self.cluster_num - 1)
+        if k < 1:
             A_mat = np.eye(self.cluster_num)
         else:
             A_mat = kneighbors_graph(
                 np.array(centers.detach().cpu().numpy()),
-                20,
+                k,
                 mode="connectivity",
                 include_self=False,
             ).toarray()

tests/core/test_grasp.py

@@ -241,5 +241,40 @@ def test_batch_smaller_than_cluster_num(self):
         self.assertEqual(ret["y_prob"].shape[0], 1)
 
 
+    def test_similar_cluster_graph_is_built_not_identity(self):
+        """Regression test: the kNN similar-cluster graph must be built and fed
+        to the GCN, not the identity fallback.
+
+        Previously ``self.A_mat`` was set to None and never reassigned, so the
+        ``if self.A_mat is None`` branch always won and the GCN received
+        ``np.eye`` (identity) -- the kNN graph construction was dead code and
+        GRASP's core "similar patients" mechanism did nothing.
+        """
+        import numpy as np
+        from unittest.mock import patch
+        from pyhealth.models.grasp import GRASPLayer
+
+        x = torch.randn(8, 5, 6)
+        mask = torch.ones(8, 5)
+
+        def adj_fed_to_gcn(layer):
+            layer.eval()
+            with patch.object(layer.GCN, "forward", wraps=layer.GCN.forward) as spy:
+                with torch.no_grad():
+                    layer(x, mask=mask)
+            return spy.call_args[0][0].detach().cpu().numpy()
+
+        # cluster_num >= 2: a real kNN graph, not identity.
+        adj = adj_fed_to_gcn(GRASPLayer(input_dim=6, hidden_dim=8, cluster_num=3, block="GRU"))
+        self.assertFalse(np.allclose(adj, np.eye(3)), "GCN got identity; graph is dead")
+        self.assertEqual(float(np.diag(adj).sum()), 0.0)   # include_self=False, no self loops
+        self.assertGreater(float(adj.sum()), 0.0)          # has edges
+
+        # cluster_num == 1: genuine identity fallback (too few clusters for a graph).
+        adj1 = adj_fed_to_gcn(GRASPLayer(input_dim=6, hidden_dim=8, cluster_num=1, block="GRU"))
+        self.assertTrue(np.allclose(adj1, np.eye(1)))
+
+
+
 if __name__ == "__main__":
     unittest.main()