PDGrapher/tests/test_sparse_planner.py

"""Tests for P3: SparsePlanner, PopulationEncoder, GapEncoder, PopulationGIDModel,
target metrics, and perturbation eval metrics."""
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

import pytest
import torch

from gidflow.models import (
    PopulationEncoder,
    GapEncoder,
    SparsePlanner,
    PopulationGIDModel,
)
from gidflow.metrics import (
    recall_at_k, precision_at_k, ndcg_at_k, mrr, jaccard_topk,
    compute_all_target_metrics,
    pearson_r_mean, pearson_r_topk_de, de_direction_agreement,
    compute_all_perturbation_metrics,
    mmd_rbf_projected,
)
from gidflow.data import SyntheticPopulationPerturbationDataset, population_collate_fn
from torch.utils.data import DataLoader


# ---------------------------------------------------------------------------
# PopulationEncoder
# ---------------------------------------------------------------------------
class TestPopulationEncoder:

    def test_output_shape_with_var(self):
        enc = PopulationEncoder(num_genes=32, hidden_dim=64, output_dim=16, use_var=True)
        cells = torch.randn(3, 10, 32)
        z = enc(cells)
        assert z.shape == (3, 16)

    def test_output_shape_no_var(self):
        enc = PopulationEncoder(num_genes=32, hidden_dim=64, output_dim=16, use_var=False)
        z = enc(torch.randn(2, 8, 32))
        assert z.shape == (2, 16)

    def test_mask_accepted(self):
        enc = PopulationEncoder(32, 64, 16, use_var=True)
        cells = torch.randn(2, 8, 32)
        mask  = torch.ones(2, 8, dtype=torch.bool)
        mask[0, 6:] = False
        z = enc(cells, mask)
        assert z.shape == (2, 16)
        assert not z.isnan().any()

    def test_gradient_flows(self):
        enc = PopulationEncoder(16, 32, 8)
        x   = torch.randn(2, 5, 16, requires_grad=True)
        z   = enc(x)
        z.mean().backward()
        assert x.grad is not None


# ---------------------------------------------------------------------------
# GapEncoder
# ---------------------------------------------------------------------------
class TestGapEncoder:

    def test_output_shape(self):
        gap = GapEncoder(input_dim=16, hidden_dim=32, output_dim=24)
        z_s = torch.randn(3, 16)
        z_t = torch.randn(3, 16)
        z_g = gap(z_s, z_t)
        assert z_g.shape == (3, 24)

    def test_gradient_flows(self):
        gap = GapEncoder(16, 32, 24)
        z_s = torch.randn(2, 16, requires_grad=True)
        z_t = torch.randn(2, 16, requires_grad=True)
        gap(z_s, z_t).mean().backward()
        assert z_s.grad is not None


# ---------------------------------------------------------------------------
# SparsePlanner
# ---------------------------------------------------------------------------
class TestSparsePlanner:

    @pytest.fixture
    def planner(self):
        return SparsePlanner(gap_dim=32, num_genes=64, hidden_dim=64)

    def test_score_shape(self, planner):
        z_gap = torch.randn(4, 32)
        scores = planner(z_gap)
        assert scores.shape == (4, 64)

    def test_soft_mask_range(self, planner):
        z_gap = torch.randn(2, 32)
        mask  = planner.get_soft_mask(z_gap)
        assert (mask >= 0).all() and (mask <= 1).all()

    def test_topk_mask_binary(self, planner):
        z_gap = torch.randn(2, 32)
        mask  = planner.get_topk_mask(z_gap, k=5)
        assert ((mask == 0) | (mask == 1)).all()
        assert mask.sum(dim=-1).eq(5).all()

    def test_ste_topk_gradient(self, planner):
        z_gap = torch.randn(2, 32, requires_grad=True)
        mask  = planner.get_ste_topk_mask(z_gap, k=5)
        mask.sum().backward()
        assert z_gap.grad is not None
        # Hard mask but gradient is not zero
        assert z_gap.grad.abs().sum() > 0

    def test_ste_mask_is_binary(self, planner):
        z_gap = torch.randn(2, 32)
        mask  = planner.get_ste_topk_mask(z_gap, k=3)
        assert ((mask == 0) | (mask == 1)).all()


# ---------------------------------------------------------------------------
# PopulationGIDModel (integration)
# ---------------------------------------------------------------------------
class TestPopulationGIDModel:

    @pytest.fixture
    def model(self):
        return PopulationGIDModel(
            num_genes=32,
            encoder_hidden=64, encoder_output=16,
            gap_hidden=32, gap_output=32,
            planner_hidden=32,
            response_hidden=64, response_pert_dim=32,
            n_layers=1,
            planner_topk=3,
        )

    def test_forward_keys(self, model):
        B, N, G = 2, 6, 32
        src = torch.randn(B, N, G)
        tgt = torch.randn(B, N, G)
        out = model(src, tgt)
        assert "target_scores" in out
        assert "target_mask" in out
        assert "pred_cells" in out

    def test_shapes(self, model):
        B, Ns, Nt, G = 2, 8, 6, 32
        src = torch.randn(B, Ns, G)
        tgt = torch.randn(B, Nt, G)
        out = model(src, tgt)
        assert out["target_scores"].shape == (B, G)
        assert out["pred_cells"].shape == (B, Ns, G)

    def test_teacher_forcing(self, model):
        B, N, G = 2, 5, 32
        src  = torch.randn(B, N, G)
        tgt  = torch.randn(B, N, G)
        pert = torch.zeros(B, G); pert[:, :3] = 1.0
        out  = model(src, tgt, true_perturbation=pert)
        assert out["pred_cells"].shape == (B, N, G)

    def test_no_nan(self, model):
        src = torch.randn(2, 6, 32)
        tgt = torch.randn(2, 6, 32)
        out = model(src, tgt)
        assert not out["target_scores"].isnan().any()
        assert not out["pred_cells"].isnan().any()

    def test_end_to_end_gradient(self, model):
        src = torch.randn(2, 5, 32, requires_grad=False)
        tgt = torch.randn(2, 5, 32)
        out = model(src, tgt)
        out["pred_cells"].mean().backward()
        # Check at least one param has gradient
        grads = [p.grad for p in model.parameters() if p.grad is not None]
        assert len(grads) > 0

    def test_predict_targets_shape(self, model):
        model.eval()
        src = torch.randn(2, 5, 32)
        tgt = torch.randn(2, 5, 32)
        mask = model.predict_targets(src, tgt, topk=3)
        assert mask.shape == (2, 32)
        assert mask.sum(dim=-1).eq(3).all()


# ---------------------------------------------------------------------------
# Target metrics
# ---------------------------------------------------------------------------
class TestTargetMetrics:

    def _make_scores_targets(self, B=4, G=64, k_true=3):
        targets = torch.zeros(B, G)
        for b in range(B):
            idx = torch.randperm(G)[:k_true]
            targets[b, idx] = 1.0
        # Perfect scores: true targets get high score
        scores = torch.rand(B, G)
        scores[targets.bool()] += 10.0
        return scores, targets

    def test_recall_at_k_perfect(self):
        scores, targets = self._make_scores_targets(k_true=3)
        r = recall_at_k(scores, targets, k=3)
        assert r.item() > 0.9

    def test_recall_at_k_random(self):
        torch.manual_seed(7)
        targets = torch.zeros(100, 1000); targets[:, :3] = 1.0
        scores  = torch.rand(100, 1000)
        r = recall_at_k(scores, targets, k=3)
        # Expected ~3/1000 * 3 ≈ 0.009 → very low
        assert r.item() < 0.1

    def test_ndcg_perfect_vs_random(self):
        scores, targets = self._make_scores_targets()
        ndcg_good = ndcg_at_k(scores, targets, k=5).item()
        scores_rand = torch.rand_like(scores)
        ndcg_rand = ndcg_at_k(scores_rand, targets, k=5).item()
        assert ndcg_good > ndcg_rand

    def test_mrr_perfect(self):
        B, G = 3, 32
        targets = torch.zeros(B, G); targets[:, 0] = 1.0
        scores  = torch.zeros(B, G); scores[:, 0] = 10.0
        r = mrr(scores, targets)
        assert abs(r.item() - 1.0) < 0.01

    def test_compute_all_returns_keys(self):
        scores, targets = self._make_scores_targets()
        out = compute_all_target_metrics(scores, targets, ks=(1, 5))
        for k in ["recall@1", "recall@5", "ndcg@1", "ndcg@5", "mrr"]:
            assert k in out


# ---------------------------------------------------------------------------
# Perturbation eval metrics
# ---------------------------------------------------------------------------
class TestPerturbationEval:

    def _make_populations(self, B=2, N=20, G=16):
        src  = torch.randn(B, N, G)
        tgt  = src + torch.randn(B, 1, G) * 0.5   # shifted
        pred = tgt + torch.randn(B, N, G) * 0.1   # close to true
        rand = torch.randn(B, N, G) + 5.0          # far from true
        return src, tgt, pred, rand

    def test_pearson_r_perfect(self):
        B, N, G = 2, 10, 16
        cells = torch.randn(B, N, G)
        r = pearson_r_mean(cells, cells)
        assert r.item() > 0.99

    def test_pearson_r_good_better_than_random(self):
        src, tgt, pred, rand = self._make_populations()
        r_good = pearson_r_mean(pred, tgt).item()
        r_rand = pearson_r_mean(rand, tgt).item()
        assert r_good > r_rand

    def test_pearson_r_topk_de(self):
        src, tgt, pred, _ = self._make_populations(G=32)
        r = pearson_r_topk_de(pred, tgt, src, topk=8)
        assert -1.0 <= r.item() <= 1.0

    def test_de_direction_agreement_perfect(self):
        src  = torch.randn(2, 10, 16)
        tgt  = src + 1.0   # all positive shift
        pred = src + 1.0   # identical shift
        agr  = de_direction_agreement(pred, tgt, src)
        assert agr.item() > 0.95

    def test_compute_all_returns_keys(self):
        src, tgt, pred, _ = self._make_populations()
        out = compute_all_perturbation_metrics(pred, tgt, src)
        assert "pearson_r_mean" in out
        assert "de_direction_agreement" in out
        assert "nmse" in out


# ---------------------------------------------------------------------------
# mmd_rbf_projected
# ---------------------------------------------------------------------------
class TestMMDProjected:

    def test_same_dist_low_mmd(self):
        torch.manual_seed(0)
        x = torch.randn(2, 20, 128)
        mmd = mmd_rbf_projected(x, x, n_components=16)
        assert mmd.item() < 0.5

    def test_different_dist_higher_mmd(self):
        torch.manual_seed(1)
        x = torch.randn(2, 20, 128)
        y = torch.randn(2, 20, 128) + 5.0
        mmd_proj = mmd_rbf_projected(x, y, n_components=16)
        mmd_zero = mmd_rbf_projected(x, x, n_components=16)
        assert mmd_proj.item() > mmd_zero.item()

    def test_shape(self):
        x = torch.randn(3, 15, 64)
        y = torch.randn(3, 12, 64)
        mmd = mmd_rbf_projected(x, y, n_components=8)
        assert mmd.shape == torch.Size([])