"""Tests for GeneGraphEncoder, graph_utils, and GIDModel with gene graph.""" import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import pytest import torch from gidflow.data import ( load_pdgrapher_edge_index, extract_subgraph, build_normalized_adjacency, ) from gidflow.models import GeneGraphEncoder, PopulationGIDModel PDG_EI_PATH = ( "/data/boom/Protein/PDGrapher/data/processed/" "torch_data/real_lognorm/edge_index_A375.pt" ) real_data_available = pytest.mark.skipif( not os.path.exists(PDG_EI_PATH), reason="PDGrapher edge_index not found" ) # --------------------------------------------------------------------------- # graph_utils # --------------------------------------------------------------------------- class TestGraphUtils: def _small_ei(self, G=20, E=60): src = torch.randint(0, G, (E,)) dst = torch.randint(0, G, (E,)) return torch.stack([src, dst]) def test_build_normalized_adjacency_shape(self): ei = self._small_ei(G=10, E=30) A = build_normalized_adjacency(ei, num_nodes=10) assert A.shape == (10, 10) def test_adjacency_nonneg_and_diagonal_positive(self): ei = self._small_ei(G=10, E=40) A = build_normalized_adjacency(ei, num_nodes=10) # All values ≥ 0 assert (A >= 0).all() # Diagonal (self-loops always added) should be positive assert (A.diag() > 0).all() def test_adjacency_undirected_is_symmetric(self): G = 10 src = torch.randint(0, G, (30,)); dst = torch.randint(0, G, (30,)) # Build undirected: include both directions ei = torch.stack([torch.cat([src, dst]), torch.cat([dst, src])]) A = build_normalized_adjacency(ei, num_nodes=G) assert torch.allclose(A, A.T, atol=1e-5) def test_extract_subgraph_reduces_edges(self): G = 10 full_genes = [f"G{i}" for i in range(G)] target_genes = ["G0", "G1", "G2", "G5"] # Create edges covering full graph ei = torch.tensor([[0,1,2,5,7,8],[1,2,5,0,8,9]]) sub_ei, mask = extract_subgraph(ei, full_genes, target_genes) # Only edges between target nodes should remain assert sub_ei.shape[0] == 2 assert sub_ei.max() < len(target_genes) def test_extract_subgraph_correct_remapping(self): full_genes = ["A", "B", "C", "D"] target_genes = ["B", "D"] # Add a B-D edge (idx 1 → idx 3) alongside A-B and C-D ei = torch.tensor([[0, 1, 2, 1, 3], [1, 0, 3, 3, 1]]) sub_ei, _ = extract_subgraph(ei, full_genes, target_genes) # B→D and D→B edges: B=0 in target, D=1 in target assert sub_ei.shape[1] > 0 assert sub_ei.max().item() < len(target_genes) @real_data_available def test_load_pdgrapher_edge_index(self): ei, _ = load_pdgrapher_edge_index(PDG_EI_PATH) assert ei.shape[0] == 2 assert ei.shape[1] == 303678 assert ei.max().item() == 10715 # --------------------------------------------------------------------------- # GeneGraphEncoder # --------------------------------------------------------------------------- class TestGeneGraphEncoder: def _make_encoder(self, G=50, out_dim=16): enc = GeneGraphEncoder( num_genes=G, node_feature_dim=16, hidden_dim=32, output_dim=out_dim, n_layers=2 ) ei = torch.randint(0, G, (2, 200)) enc.set_graph(ei) return enc def test_output_shape(self): enc = self._make_encoder(G=50, out_dim=16) emb = enc() assert emb.shape == (50, 16) def test_no_nan(self): enc = self._make_encoder(G=40, out_dim=8) emb = enc() assert not emb.isnan().any() assert not emb.isinf().any() def test_gradient_through_embedding(self): enc = self._make_encoder(G=30, out_dim=8) emb = enc() emb.sum().backward() # Gene embedding weights should have gradients assert enc.gene_embedding.weight.grad is not None def test_with_protein_features(self): G, F_p = 40, 4 enc = GeneGraphEncoder( num_genes=G, node_feature_dim=16, hidden_dim=32, output_dim=16, n_layers=2, protein_feature_dim=F_p ) ei = torch.randint(0, G, (2, 150)) enc.set_graph(ei) prot = torch.rand(G, F_p) emb = enc(prot) assert emb.shape == (G, 16) def test_different_graphs_give_different_embeddings(self): G = 30 enc = GeneGraphEncoder(G, 16, 32, 16, n_layers=1) ei1 = torch.randint(0, G, (2, 100)) ei2 = torch.randint(0, G, (2, 100)) enc.set_graph(ei1); emb1 = enc().detach() enc.set_graph(ei2); emb2 = enc().detach() # Different graphs → different embeddings (very likely with random edges) assert not torch.allclose(emb1, emb2) def test_set_graph_raises_before_forward(self): enc = GeneGraphEncoder(20, 8, 16, 8, n_layers=1) with pytest.raises(RuntimeError): enc() # no set_graph called @real_data_available def test_real_pdgrapher_graph(self): ei, _ = load_pdgrapher_edge_index(PDG_EI_PATH) enc = GeneGraphEncoder(10716, node_feature_dim=16, hidden_dim=32, output_dim=16, n_layers=1) enc.set_graph(ei) emb = enc() assert emb.shape == (10716, 16) assert not emb.isnan().any() # --------------------------------------------------------------------------- # PopulationGIDModel with gene graph # --------------------------------------------------------------------------- class TestGIDModelWithGeneGraph: @pytest.fixture def model_and_ei(self): G = 60 ei = torch.randint(0, G, (2, 250)) model = PopulationGIDModel( num_genes=G, encoder_hidden=32, encoder_output=16, gap_hidden=32, gap_output=32, planner_hidden=32, response_hidden=32, response_pert_dim=16, n_layers=1, planner_topk=2, encoder_use_var=False, use_gene_graph=True, gene_graph_node_dim=16, gene_graph_hidden_dim=32, gene_graph_output_dim=16, gene_graph_n_layers=2, ) model.set_gene_graph(ei) return model, ei def test_forward_shapes(self, model_and_ei): model, _ = model_and_ei B, G = 2, 60 src = torch.randn(B, 1, G) tgt = torch.randn(B, 1, G) out = model(src, tgt) assert out["target_scores"].shape == (B, G) assert out["pred_cells"].shape == (B, 1, G) def test_no_nan(self, model_and_ei): model, _ = model_and_ei src = torch.randn(2, 1, 60); tgt = torch.randn(2, 1, 60) out = model(src, tgt) assert not out["target_scores"].isnan().any() assert not out["pred_cells"].isnan().any() def test_gradient_flows_through_gcn(self, model_and_ei): model, _ = model_and_ei src = torch.randn(2, 1, 60); tgt = torch.randn(2, 1, 60) out = model(src, tgt) out["target_scores"].sum().backward() # GCN gene embedding should get gradient gcn_grad = model.gene_graph_encoder.gene_embedding.weight.grad assert gcn_grad is not None assert gcn_grad.abs().sum() > 0 def test_gcn_off_matches_no_graph(self): """Without graph, model should still forward correctly.""" G = 60 model_no_graph = PopulationGIDModel( num_genes=G, encoder_hidden=32, encoder_output=16, gap_hidden=32, gap_output=32, planner_hidden=32, response_hidden=32, response_pert_dim=16, n_layers=1, planner_topk=2, encoder_use_var=False, use_gene_graph=False, ) src = torch.randn(2, 1, G); tgt = torch.randn(2, 1, G) out = model_no_graph(src, tgt) assert out["target_scores"].shape == (2, G) def test_predict_targets_shape(self, model_and_ei): model, _ = model_and_ei model.eval() src = torch.randn(2, 1, 60); tgt = torch.randn(2, 1, 60) mask = model.predict_targets(src, tgt, topk=3) assert mask.shape == (2, 60) assert mask.sum(dim=-1).eq(3).all() def test_with_protein_features_and_graph(self): G, F_p = 50, 4 ei = torch.randint(0, G, (2, 200)) model = PopulationGIDModel( num_genes=G, encoder_hidden=32, encoder_output=16, gap_hidden=32, gap_output=32, planner_hidden=32, response_hidden=32, response_pert_dim=16, n_layers=1, planner_topk=1, encoder_use_var=False, protein_input_dim=F_p, protein_hidden_dim=16, protein_output_dim=8, use_gene_graph=True, gene_graph_node_dim=16, gene_graph_hidden_dim=32, gene_graph_output_dim=16, gene_graph_n_layers=1, ) model.set_gene_graph(ei) model.set_protein_features(torch.rand(G, F_p)) src = torch.randn(2, 1, G); tgt = torch.randn(2, 1, G) out = model(src, tgt) assert out["target_scores"].shape == (2, G) assert not out["target_scores"].isnan().any()