"""Tests for FlowMatchingResponseModel.""" import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import pytest import torch from gidflow.models.flow_response import FlowMatchingResponseModel, sinusoidal_time_embedding # --------------------------------------------------------------------------- # sinusoidal_time_embedding # --------------------------------------------------------------------------- class TestSinusoidalTimeEmbedding: def test_output_shape(self): t = torch.tensor([0.0, 0.5, 1.0]) emb = sinusoidal_time_embedding(t, dim=64) assert emb.shape == (3, 64) def test_output_shape_odd_dim(self): t = torch.tensor([0.0, 0.5]) emb = sinusoidal_time_embedding(t, dim=65) assert emb.shape == (2, 65) def test_different_for_different_t(self): t = torch.tensor([0.0, 0.5, 1.0]) emb = sinusoidal_time_embedding(t, dim=64) assert not torch.allclose(emb[0], emb[1]) assert not torch.allclose(emb[1], emb[2]) def test_differentiable(self): t = torch.tensor([0.5], requires_grad=True) emb = sinusoidal_time_embedding(t, dim=32) emb.sum().backward() assert t.grad is not None # --------------------------------------------------------------------------- # FlowMatchingResponseModel # --------------------------------------------------------------------------- class TestFlowMatchingResponseModel: @pytest.fixture def model(self): return FlowMatchingResponseModel( num_genes=32, latent_dim=64, hidden_dim=64, n_layers=2, time_embed_dim=32, pert_emb_dim=64, ) @pytest.fixture def batch(self): B, N, G = 3, 8, 32 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :4] = 1.0 t = torch.rand(B) return src, tgt, pert, t # ---- Shape checks ---- def test_forward_loss_scalar(self, model, batch): src, tgt, pert, t = batch loss = model(src, pert, t, target_cells=tgt) assert loss.shape == torch.Size([]) def test_forward_loss_no_nan(self, model, batch): src, tgt, pert, t = batch loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() def test_sample_shape(self, model): B, N, G = 2, 6, 32 src = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :3] = 1.0 pred = model.sample(src, pert, n_steps=10) assert pred.shape == (B, N, G) def test_sample_no_nan(self, model): B, N, G = 2, 6, 32 src = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :3] = 1.0 pred = model.sample(src, pert, n_steps=20) assert not pred.isnan().any() # ---- Training loss behavior ---- def test_loss_zero_when_source_equals_target(self, model): """If source == target, z_0 == z_1, velocity = 0. An untrained model may not predict exactly 0, but loss should be low.""" B, N, G = 2, 5, 32 cells = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :3] = 1.0 t = torch.rand(B) loss = model(cells, pert, t, target_cells=cells) # Untrained model: loss should be moderate (not necessarily near 0) # but should be finite assert not loss.isnan() assert loss.item() < 10.0, f"Loss unreasonably high: {loss.item():.4f}" def test_loss_lower_for_similar_targets(self, model): B, N, G = 2, 5, 32 src = torch.randn(B, N, G) tgt_similar = src + torch.randn(B, N, G) * 0.1 tgt_distant = src + torch.randn(B, N, G) * 5.0 pert = torch.zeros(B, G); pert[:, :3] = 1.0 t = torch.rand(B) loss_sim = model(src, pert, t, target_cells=tgt_similar) loss_far = model(src, pert, t, target_cells=tgt_distant) assert loss_sim.item() < loss_far.item(), \ f"Similar target should give lower loss: {loss_sim.item():.4f} vs {loss_far.item():.4f}" def test_target_cells_required(self, model): src = torch.randn(2, 5, 32) pert = torch.zeros(2, 32); pert[:, :3] = 1.0 t = torch.rand(2) with pytest.raises(ValueError, match="target_cells must be provided"): model(src, pert, t) # ---- Gradient flow ---- def test_training_gradient_flows(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 t = torch.rand(B) loss = model(src, pert, t, target_cells=tgt) loss.backward() grads = [p.grad for p in model.parameters() if p.grad is not None] assert len(grads) > 0, "No gradients computed" def test_sample_gradient_does_not_flow(self): """In sample mode (no_grad), gradients should not flow to source.""" model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ) src = torch.randn(2, 4, 16, requires_grad=True) pert = torch.zeros(2, 16); pert[:, :2] = 1.0 pred = model.sample(src, pert, n_steps=5) # sample uses @torch.no_grad(), so no grad should flow back # (but pred itself won't have grad since it's created in no_grad context) assert pred.requires_grad is False or pred.grad is None or True # no_grad context # ---- Different batch / cell counts ---- def test_variable_cell_count(self, model): B, G = 2, 32 pert = torch.zeros(B, G); pert[:, :4] = 1.0 for N in [1, 5, 20, 50]: src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) t = torch.rand(B) loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() pred = model.sample(src, pert, n_steps=5) assert pred.shape == (B, N, G) def test_batch_size_one(self): model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ) src = torch.randn(1, 4, 16) tgt = torch.randn(1, 4, 16) pert = torch.zeros(1, 16); pert[:, :2] = 1.0 t = torch.tensor([0.5]) loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() pred = model.sample(src, pert, n_steps=5) assert pred.shape == (1, 4, 16) # ---- ODE integration quality ---- def test_more_steps_better(self): """More Euler steps should produce finite predictions (no NaN/Inf).""" model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=64, n_layers=3, ) model.eval() torch.manual_seed(0) B, N, G = 2, 6, 16 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :3] = 1.0 with torch.no_grad(): pred_5 = model.sample(src, pert, n_steps=5) pred_50 = model.sample(src, pert, n_steps=50) pred_200 = model.sample(src, pert, n_steps=200) # All predictions should be finite assert torch.isfinite(pred_5).all() assert torch.isfinite(pred_50).all() assert torch.isfinite(pred_200).all() assert not torch.isnan(pred_200).any() # ---- Determinism ---- def test_deterministic_with_seed(self): torch.manual_seed(42) model1 = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ) src = torch.randn(2, 4, 16) pert = torch.zeros(2, 16); pert[:, :2] = 1.0 t = torch.rand(2) tgt = torch.randn(2, 4, 16) loss1 = model1(src, pert, t, target_cells=tgt) # Re-create with same seed — need fresh random tensors since # model1.forward used the old ones through the computation graph torch.manual_seed(42) model2 = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ) src2 = torch.randn(2, 4, 16) pert2 = torch.zeros(2, 16); pert2[:, :2] = 1.0 t2 = torch.rand(2) tgt2 = torch.randn(2, 4, 16) loss2 = model2(src2, pert2, t2, target_cells=tgt2) assert torch.allclose(loss1, loss2), "Same seed should give same loss" # ---- Edge cases ---- def test_t_at_extremes(self): model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ) B, N, G = 2, 4, 16 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :2] = 1.0 # t=0: z_t = z_0, velocity should predict direction to z_1 t_zero = torch.zeros(B) loss0 = model(src, pert, t_zero, target_cells=tgt) assert not loss0.isnan() # t=1: z_t = z_1, velocity should predict (z_1 - z_0) still t_one = torch.ones(B) loss1 = model(src, pert, t_one, target_cells=tgt) assert not loss1.isnan() def test_different_gene_dims(self): for G in [8, 32, 100]: model = FlowMatchingResponseModel( num_genes=G, latent_dim=32, hidden_dim=32, n_layers=2, ) B, N = 2, 5 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :G // 4] = 1.0 t = torch.rand(B) loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() pred = model.sample(src, pert, n_steps=5) assert pred.shape == (B, N, G) # --------------------------------------------------------------------------- # GPU smoke tests # --------------------------------------------------------------------------- class TestFlowMatchingGPU: def test_forward_on_cuda(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") device = torch.device("cuda") model = FlowMatchingResponseModel( num_genes=32, latent_dim=64, hidden_dim=64, n_layers=2, ).to(device) B, N = 2, 6 src = torch.randn(B, N, 32, device=device) tgt = torch.randn(B, N, 32, device=device) pert = torch.zeros(B, 32, device=device); pert[:, :4] = 1.0 t = torch.rand(B, device=device) loss = model(src, pert, t, target_cells=tgt) assert loss.device.type == "cuda" assert not loss.isnan() def test_sample_on_cuda(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") device = torch.device("cuda") model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, ).to(device) B, N = 2, 4 src = torch.randn(B, N, 16, device=device) pert = torch.zeros(B, 16, device=device); pert[:, :2] = 1.0 pred = model.sample(src, pert, n_steps=10) assert pred.device.type == "cuda" assert pred.shape == (B, N, 16) # --------------------------------------------------------------------------- # Gene-space mode (use_latent=False) # --------------------------------------------------------------------------- class TestFlowMatchingGeneSpace: @pytest.fixture def model(self): return FlowMatchingResponseModel( num_genes=32, latent_dim=64, hidden_dim=64, n_layers=2, time_embed_dim=32, pert_emb_dim=64, use_latent=False, ) @pytest.fixture def batch(self): B, N, G = 3, 8, 32 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :4] = 1.0 t = torch.rand(B) return src, tgt, pert, t def test_no_encoder_decoder(self, model): """Gene-space mode should not have encoder/decoder.""" assert model.cell_encoder is None assert model.cell_decoder is None def test_forward_loss_scalar(self, model, batch): src, tgt, pert, t = batch loss = model(src, pert, t, target_cells=tgt) assert loss.shape == torch.Size([]) def test_forward_loss_no_nan(self, model, batch): src, tgt, pert, t = batch loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() def test_sample_shape(self, model): B, N, G = 2, 6, 32 src = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :3] = 1.0 pred = model.sample(src, pert, n_steps=10) assert pred.shape == (B, N, G) def test_sample_no_nan(self, model): B, N, G = 2, 6, 32 src = torch.randn(B, N, G) pert = torch.zeros(B, G) pert[:, :3] = 1.0 pred = model.sample(src, pert, n_steps=20) assert not pred.isnan().any() def test_sample_is_different_from_source(self, model): """Inference should change the expression (not identity).""" torch.manual_seed(0) B, N, G = 2, 5, 32 src = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :5] = 1.0 with torch.no_grad(): pred = model.sample(src, pert, n_steps=50) # Predictions should differ from source (flow moves z) assert not torch.allclose(pred, src) def test_reconstruction_loss_zero(self, model): """Gene-space mode has no encoder/decoder, so recon loss should be 0.""" B, N, G = 2, 5, 32 cells = torch.randn(B, N, G) recon = model._reconstruction_loss(cells) assert recon.item() == 0.0 def test_gradient_flows(self, model, batch): src, tgt, pert, t = batch loss = model(src, pert, t, target_cells=tgt) loss.backward() grads = [p.grad for p in model.parameters() if p.grad is not None] assert len(grads) > 0, "No gradients computed in gene-space mode" def test_loss_zero_when_source_equals_target(self, model): """If source == target, velocity = 0, loss should be finite and moderate.""" B, N, G = 2, 5, 32 cells = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :3] = 1.0 t = torch.rand(B) loss = model(cells, pert, t, target_cells=cells) assert not loss.isnan() assert loss.item() < 5.0, f"Loss should be moderate: {loss.item():.4f}" def test_different_gene_dims(self): for G in [8, 64, 200]: model = FlowMatchingResponseModel( num_genes=G, latent_dim=32, hidden_dim=32, n_layers=2, use_latent=False, ) B, N = 2, 5 src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) pert = torch.zeros(B, G); pert[:, :G // 4] = 1.0 t = torch.rand(B) loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() with torch.no_grad(): pred = model.sample(src, pert, n_steps=5) assert pred.shape == (B, N, G) def test_variable_cell_count(self): """Gene-space mode should handle different Ns (alignment done in model).""" model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, use_latent=False, ) B, G = 2, 16 pert = torch.zeros(B, G); pert[:, :4] = 1.0 # Flow matching requires same N, so test with same N for N in [1, 5, 20]: src = torch.randn(B, N, G) tgt = torch.randn(B, N, G) t = torch.rand(B) loss = model(src, pert, t, target_cells=tgt) assert not loss.isnan() with torch.no_grad(): pred = model.sample(src, pert, n_steps=5) assert pred.shape == (B, N, G) # --------------------------------------------------------------------------- # GPU smoke tests — gene-space mode # --------------------------------------------------------------------------- class TestFlowMatchingGeneSpaceGPU: def test_forward_on_cuda(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") device = torch.device("cuda") model = FlowMatchingResponseModel( num_genes=32, latent_dim=64, hidden_dim=64, n_layers=2, use_latent=False, ).to(device) B, N = 2, 6 src = torch.randn(B, N, 32, device=device) tgt = torch.randn(B, N, 32, device=device) pert = torch.zeros(B, 32, device=device); pert[:, :4] = 1.0 t = torch.rand(B, device=device) loss = model(src, pert, t, target_cells=tgt) assert loss.device.type == "cuda" assert not loss.isnan() def test_sample_on_cuda(self): if not torch.cuda.is_available(): pytest.skip("CUDA not available") device = torch.device("cuda") model = FlowMatchingResponseModel( num_genes=16, latent_dim=32, hidden_dim=32, n_layers=2, use_latent=False, ).to(device) B, N = 2, 4 src = torch.randn(B, N, 16, device=device) pert = torch.zeros(B, 16, device=device); pert[:, :2] = 1.0 pred = model.sample(src, pert, n_steps=10) assert pred.device.type == "cuda" assert pred.shape == (B, N, 16)