Add minimal testing for GANOrchestrator and implement simple discriminator fixture

Author: wli51

tests/engine/conftest.py

@@ -112,6 +112,43 @@ def forward(self, x):
     return MultiOutputConv()
 
 
+@pytest.fixture
+def simple_discriminator():
+    """
+    Simple discriminator model for GAN testing.
+    Takes concatenated input/target stack (B, 6, H, W) -> outputs score (B, 1)
+    Uses conv + global average pooling + linear layer.
+    """
+    class SimpleDiscriminator(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Conv2d(
+                in_channels=6,  # stacked input + target
+                out_channels=16,
+                kernel_size=3,
+                padding=1,
+                bias=True
+            )
+            self.pool = nn.AdaptiveAvgPool2d(1)  # Global average pooling
+            self.fc = nn.Linear(16, 1)  # Output single score
+        
+        def forward(self, x):
+            x = self.conv(x)
+            x = torch.relu(x)
+            x = self.pool(x)  # (B, 16, 1, 1)
+            x = x.flatten(1)  # (B, 16)
+            x = self.fc(x)    # (B, 1)
+            return x
+    
+    return SimpleDiscriminator()
+
+
+@pytest.fixture
+def random_stack():
+    """Random stack tensor (batch=2, channels=6, height=8, width=8) for discriminator."""
+    return torch.randn(2, 6, 8, 8)
+
+
 @pytest.fixture
 def sample_inputs():
     """Create sample inputs for loss computation."""
@@ -208,3 +245,31 @@ def forward_pass_context_eval(forward_group, random_input, random_target, torch_
         inputs=random_input.to(torch_device),
         targets=random_target.to(torch_device),
     )
+
+
+@pytest.fixture
+def disc_optimizer(simple_discriminator):
+    """Create an Adam optimizer for the discriminator model."""
+    import torch.optim as optim
+    return optim.Adam(simple_discriminator.parameters(), lr=1e-3)
+
+
+@pytest.fixture
+def discriminator_forward_group(simple_discriminator, disc_optimizer, torch_device):
+    """Create a DiscriminatorForwardGroup with the simple discriminator and optimizer."""
+    from virtual_stain_flow.engine.forward_groups import DiscriminatorForwardGroup
+    return DiscriminatorForwardGroup(
+        discriminator=simple_discriminator,
+        optimizer=disc_optimizer,
+        device=torch_device,
+    )
+
+
+@pytest.fixture
+def gan_orchestrator(forward_group, discriminator_forward_group):
+    """Create a GANOrchestrator with generator and discriminator forward groups."""
+    from virtual_stain_flow.engine.orchestrators import GANOrchestrator
+    return GANOrchestrator(
+        generator_fg=forward_group,
+        discriminator_fg=discriminator_forward_group,
+    )

tests/engine/test_forward_group.py

@@ -5,7 +5,8 @@
 
 from virtual_stain_flow.engine.forward_groups import (
     AbstractForwardGroup,
-    GeneratorForwardGroup
+    GeneratorForwardGroup,
+    DiscriminatorForwardGroup
 )
 from virtual_stain_flow.engine.names import INPUTS, TARGETS, PREDS, GENERATOR_MODEL
 
@@ -128,3 +129,52 @@ def test_forward_output_arity_mismatch(self, multi_output_model, random_input, r
 
         with pytest.raises(ValueError, match="Model returned 2 outputs.*output_keys expects 1"):
             forward_group(train=False, inputs=random_input, targets=random_target)
+
+
+class TestDiscriminatorForwardGroup:
+    """Test DiscriminatorForwardGroup functionality."""
+
+    def test_forward_train_mode(self, simple_discriminator, random_stack):
+        """Test that discriminator is set to train mode when train=True."""
+        forward_group = DiscriminatorForwardGroup(
+            device=torch.device("cpu"),
+            discriminator=simple_discriminator
+        )
+        
+        ctx = forward_group(train=True, stack=random_stack)
+        
+        assert forward_group.model.training is True
+        assert ctx["p"].requires_grad is True
+
+    def test_forward_eval_mode(self, simple_discriminator, random_stack):
+        """Test that discriminator is set to eval mode when train=False."""
+        forward_group = DiscriminatorForwardGroup(
+            device=torch.device("cpu"),
+            discriminator=simple_discriminator
+        )
+        
+        ctx = forward_group(train=False, stack=random_stack)
+        
+        assert forward_group.model.training is False
+        assert ctx["p"].requires_grad is False
+
+    def test_optimizer_zero_grad(self, simple_discriminator, disc_optimizer, random_stack):
+        """Test that optimizer.zero_grad() is called when train=True."""
+        forward_group = DiscriminatorForwardGroup(
+            device=torch.device("cpu"),
+            discriminator=simple_discriminator,
+            optimizer=disc_optimizer
+        )
+        
+        # Manually create some gradients
+        dummy_loss = sum(p.sum() for p in forward_group.model.parameters())
+        dummy_loss.backward()
+        
+        # Check that gradients exist
+        assert any(p.grad is not None for p in forward_group.model.parameters())
+        
+        # Forward should zero gradients
+        _ = forward_group(train=True, stack=random_stack)
+        
+        # Gradients should be None (set_to_none=True)
+        assert all(p.grad is None for p in forward_group.model.parameters())

tests/engine/test_orchestrator.py

@@ -0,0 +1,63 @@
+"""Tests for GANOrchestrator."""
+
+import torch
+
+from virtual_stain_flow.engine.names import INPUTS, TARGETS, PREDS
+
+
+class TestGANOrchestrator:
+    """Test GANOrchestrator functionality."""
+
+    def test_discriminator_forward(self, gan_orchestrator, random_input, random_target):
+        """Test that _discriminator_forward produces correct context with real and fake stacks."""
+        ctx = gan_orchestrator._discriminator_forward(
+            train=False,
+            inputs=random_input,
+            targets=random_target
+        )
+        
+        # Check that generator outputs are present
+        assert INPUTS in ctx
+        assert TARGETS in ctx
+        assert PREDS in ctx
+        
+        # Check that discriminator outputs for real and fake are present
+        assert "real_stack" in ctx
+        assert "fake_stack" in ctx
+        assert "p_real_as_real" in ctx
+        assert "p_fake_as_real" in ctx
+        
+        # Verify shapes
+        batch_size = random_input.shape[0]
+        assert ctx["p_real_as_real"].shape[0] == batch_size
+        assert ctx["p_fake_as_real"].shape[0] == batch_size
+        
+        # Verify real_stack is concatenation of inputs and targets
+        expected_real_stack = torch.cat([ctx[INPUTS], ctx[TARGETS]], dim=1)
+        assert torch.allclose(ctx["real_stack"], expected_real_stack)
+        
+        # Verify fake_stack is concatenation of inputs and preds
+        expected_fake_stack = torch.cat([ctx[INPUTS], ctx[PREDS]], dim=1)
+        assert torch.allclose(ctx["fake_stack"], expected_fake_stack)
+
+    def test_generator_forward(self, gan_orchestrator, random_input, random_target):
+        """Test that _generator_forward produces correct context with generator outputs and discriminator score."""
+        ctx = gan_orchestrator._generator_forward(
+            train=False,
+            inputs=random_input,
+            targets=random_target
+        )
+        
+        # Check that generator outputs are present
+        assert INPUTS in ctx
+        assert TARGETS in ctx
+        assert PREDS in ctx
+        
+        # Check that discriminator score for fake is present
+        assert "p_fake_as_real" in ctx
+        
+        # Verify shapes
+        batch_size = random_input.shape[0]
+        assert ctx[PREDS].shape[0] == batch_size
+        assert ctx["p_fake_as_real"].shape[0] == batch_size
+        assert ctx["p_fake_as_real"].shape[1] == 1  # Single score output