PDGrapher/tests/test_gene_graph_encoder.py

"""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()