PDGrapher/tests/test_population_batch.py

"""Tests for population-level data layer: batch, dataset, collate."""
import sys
import os

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

import pytest
import torch
from torch.utils.data import DataLoader

from gidflow.data import (
    PopulationPerturbationBatch,
    SyntheticPopulationPerturbationDataset,
    population_collate_fn,
)


# ---------------------------------------------------------------------------
# PopulationPerturbationBatch
# ---------------------------------------------------------------------------

class TestPopulationPerturbationBatch:

    def _make_batch(self, B=2, Ns=4, Nt=6, G=8):
        return PopulationPerturbationBatch(
            source_cells=torch.randn(B, Ns, G),
            target_cells=torch.randn(B, Nt, G),
            perturbation=torch.zeros(B, G),
            source_mask=torch.ones(B, Ns, dtype=torch.bool),
            target_mask=torch.ones(B, Nt, dtype=torch.bool),
        )

    def test_shapes(self):
        b = self._make_batch()
        assert b.source_cells.shape == (2, 4, 8)
        assert b.target_cells.shape == (2, 6, 8)
        assert b.perturbation.shape == (2, 8)
        assert b.source_mask.shape  == (2, 4)
        assert b.target_mask.shape  == (2, 6)

    def test_to_device_cpu(self):
        b = self._make_batch()
        b2 = b.to(torch.device("cpu"))
        assert b2.source_cells.device.type == "cpu"

    def test_optional_fields_default_none(self):
        b = self._make_batch()
        assert b.context is None
        assert b.metadata is None


# ---------------------------------------------------------------------------
# SyntheticPopulationPerturbationDataset
# ---------------------------------------------------------------------------

class TestSyntheticPopulationPerturbationDataset:

    @pytest.fixture
    def ds(self):
        return SyntheticPopulationPerturbationDataset(
            num_conditions=16,
            num_genes=32,
            min_cells=4,
            max_cells=12,
            perturbation_size=2,
            seed=0,
        )

    def test_length(self, ds):
        assert len(ds) == 16

    def test_item_keys(self, ds):
        item = ds[0]
        assert "source_cells" in item
        assert "target_cells" in item
        assert "perturbation" in item

    def test_gene_dim(self, ds):
        item = ds[0]
        assert item["source_cells"].shape[-1] == 32
        assert item["target_cells"].shape[-1] == 32
        assert item["perturbation"].shape == (32,)

    def test_cell_counts_in_range(self, ds):
        for item in ds:
            ns = item["source_cells"].shape[0]
            nt = item["target_cells"].shape[0]
            assert 4 <= ns <= 12
            assert 4 <= nt <= 12

    def test_perturbation_is_multihot(self, ds):
        for item in ds:
            pert = item["perturbation"]
            assert ((pert == 0) | (pert == 1)).all()
            assert int(pert.sum().item()) == 2

    def test_reproducible(self):
        ds1 = SyntheticPopulationPerturbationDataset(num_conditions=4, num_genes=8, seed=7)
        ds2 = SyntheticPopulationPerturbationDataset(num_conditions=4, num_genes=8, seed=7)
        assert torch.allclose(ds1[0]["source_cells"], ds2[0]["source_cells"])

    def test_different_seeds_differ(self):
        ds1 = SyntheticPopulationPerturbationDataset(num_conditions=4, num_genes=16, seed=1)
        ds2 = SyntheticPopulationPerturbationDataset(num_conditions=4, num_genes=16, seed=2)
        # Perturbations are same size [G] regardless of cell count
        assert not torch.allclose(ds1[0]["perturbation"], ds2[0]["perturbation"])


# ---------------------------------------------------------------------------
# population_collate_fn
# ---------------------------------------------------------------------------

class TestPopulationCollateFn:

    def _make_items(self, G=16):
        """Items with deliberately different cell counts."""
        return [
            {"source_cells": torch.randn(3, G), "target_cells": torch.randn(5, G), "perturbation": torch.zeros(G)},
            {"source_cells": torch.randn(7, G), "target_cells": torch.randn(2, G), "perturbation": torch.ones(G)},
        ]

    def test_output_type(self):
        items = self._make_items()
        batch = population_collate_fn(items)
        assert isinstance(batch, PopulationPerturbationBatch)

    def test_padded_shapes(self):
        items = self._make_items(G=16)
        batch = population_collate_fn(items)
        B, G = 2, 16
        assert batch.source_cells.shape == (B, 7, G)   # max(3, 7)
        assert batch.target_cells.shape == (B, 5, G)   # max(5, 2)
        assert batch.perturbation.shape == (B, G)
        assert batch.source_mask.shape  == (B, 7)
        assert batch.target_mask.shape  == (B, 5)

    def test_mask_values(self):
        items = self._make_items(G=8)
        batch = population_collate_fn(items)
        # item 0 has 3 source cells → first 3 real, rest padding
        assert batch.source_mask[0, :3].all()
        assert not batch.source_mask[0, 3:].any()
        # item 1 has 7 source cells → all real
        assert batch.source_mask[1].all()

    def test_padding_is_zero(self):
        items = self._make_items(G=8)
        batch = population_collate_fn(items)
        # padded positions for item 0 (source cells 3..6)
        assert (batch.source_cells[0, 3:] == 0).all()

    def test_max_cells_truncation(self):
        items = self._make_items(G=8)
        batch = population_collate_fn(items, max_source_cells=4, max_target_cells=3)
        assert batch.source_cells.shape[1] == 4
        assert batch.target_cells.shape[1] == 3

    def test_dataloader_integration(self):
        ds = SyntheticPopulationPerturbationDataset(num_conditions=8, num_genes=16, seed=0)
        loader = DataLoader(ds, batch_size=4, collate_fn=population_collate_fn, shuffle=False)
        batch = next(iter(loader))
        assert isinstance(batch, PopulationPerturbationBatch)
        assert batch.source_cells.shape[0] == 4
        assert batch.source_cells.shape[2] == 16

    def test_perturbation_stacked_correctly(self):
        G = 8
        items = [
            {"source_cells": torch.randn(2, G), "target_cells": torch.randn(2, G),
             "perturbation": torch.tensor([1.0, 0, 0, 0, 0, 0, 0, 0])},
            {"source_cells": torch.randn(2, G), "target_cells": torch.randn(2, G),
             "perturbation": torch.tensor([0.0, 1, 0, 0, 0, 0, 0, 0])},
        ]
        batch = population_collate_fn(items)
        assert batch.perturbation[0, 0] == 1.0
        assert batch.perturbation[1, 1] == 1.0