| """Tests for SimplePopulationResponseModel.""" |
| import sys |
| import os |
|
|
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) |
|
|
| import pytest |
| import torch |
|
|
| from gidflow.models import SimplePopulationResponseModel |
|
|
|
|
| class TestSimplePopulationResponseModel: |
|
|
| @pytest.fixture |
| def model(self): |
| return SimplePopulationResponseModel(num_genes=32, hidden_dim=64, perturbation_dim=32) |
|
|
| def test_output_shape(self, model): |
| B, N, G = 4, 10, 32 |
| src = torch.randn(B, N, G) |
| pert = torch.zeros(B, G) |
| pert[:, :3] = 1.0 |
| out = model(src, pert) |
| assert out.shape == (B, N, G), f"Expected ({B},{N},{G}), got {out.shape}" |
|
|
| def test_different_n_per_call(self, model): |
| """Model should handle varying N between calls.""" |
| for N in [1, 5, 32, 64]: |
| src = torch.randn(2, N, 32) |
| pert = torch.zeros(2, 32) |
| out = model(src, pert) |
| assert out.shape == (2, N, 32) |
|
|
| def test_no_nan(self, model): |
| src = torch.randn(3, 8, 32) |
| pert = torch.zeros(3, 32) |
| pert[:, 0] = 1.0 |
| out = model(src, pert) |
| assert not out.isnan().any(), "NaN in output" |
| assert not out.isinf().any(), "Inf in output" |
|
|
| def test_gradient_flows(self, model): |
| src = torch.randn(2, 6, 32, requires_grad=True) |
| pert = torch.randn(2, 32, requires_grad=True) |
| out = model(src, pert) |
| loss = out.mean() |
| loss.backward() |
| assert src.grad is not None |
| assert pert.grad is not None |
|
|
| def test_mask_accepted(self, model): |
| """Forward accepts source_mask without error.""" |
| B, N, G = 2, 8, 32 |
| src = torch.randn(B, N, G) |
| pert = torch.zeros(B, G) |
| mask = torch.ones(B, N, dtype=torch.bool) |
| out = model(src, pert, source_mask=mask) |
| assert out.shape == (B, N, G) |
|
|
| def test_gpu_if_available(self, model): |
| if not torch.cuda.is_available(): |
| pytest.skip("CUDA not available") |
| device = torch.device("cuda") |
| model = model.to(device) |
| B, N, G = 2, 5, 32 |
| src = torch.randn(B, N, G, device=device) |
| pert = torch.zeros(B, G, device=device) |
| out = model(src, pert) |
| assert out.device.type == "cuda" |
|
|
| def test_perturbation_changes_output(self, model): |
| """Different perturbations should give different predictions.""" |
| model.eval() |
| src = torch.randn(1, 8, 32) |
| pert1 = torch.zeros(1, 32); pert1[0, 0] = 1.0 |
| pert2 = torch.zeros(1, 32); pert2[0, 5] = 1.0 |
| with torch.no_grad(): |
| out1 = model(src, pert1) |
| out2 = model(src, pert2) |
| assert not torch.allclose(out1, out2), "Different perturbations gave identical output" |
|
|
| def test_batch_size_1(self, model): |
| out = model(torch.randn(1, 4, 32), torch.zeros(1, 32)) |
| assert out.shape == (1, 4, 32) |
|
|
