| """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, |
| ) |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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) |
| |
| assert not torch.allclose(ds1[0]["perturbation"], ds2[0]["perturbation"]) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
| assert batch.target_cells.shape == (B, 5, G) |
| 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) |
| |
| assert batch.source_mask[0, :3].all() |
| assert not batch.source_mask[0, 3:].any() |
| |
| assert batch.source_mask[1].all() |
|
|
| def test_padding_is_zero(self): |
| items = self._make_items(G=8) |
| batch = population_collate_fn(items) |
| |
| 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 |
|
|
