GID-Flow / PDGrapher /tests /test_sparse_planner.py

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

07fcdfe verified 4 months ago

10.7 kB

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

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