GID-Flow / PDGrapher /tests /test_gene_graph_encoder.py

Upload GID-Flow project snapshot (deduped: code + key artifacts)

07fcdfe verified 4 months ago

9.2 kB

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

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