File size: 6,633 Bytes
07fcdfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """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
|