Backtracking tests and moving modules around for better organization

src/virtual_stain_flow/models/blocks/__init__.py

@@ -0,0 +1,34 @@
+from .blocks import (
+    AbstractBlock,
+    Conv2DConvNeXtBlock,
+    Conv2DNormActBlock,
+)
+from .up_down_blocks import (
+	Bilinear2DUpsampleBlock,
+	Conv2DDownBlock,
+	ConvTrans2DUpBlock,
+	IdentityBlock,
+	MaxPool2DDownBlock,
+	PixelShuffle2DUpBlock,
+)
+from .utils import (
+    ActivationType,
+    get_activation,
+    get_norm
+)
+
+
+__all__ = [
+    "AbstractBlock",
+    "Conv2DConvNeXtBlock",
+    "Conv2DNormActBlock",
+    "Bilinear2DUpsampleBlock",
+    "Conv2DDownBlock",
+    "ConvTrans2DUpBlock",
+    "IdentityBlock",
+    "MaxPool2DDownBlock",
+    "PixelShuffle2DUpBlock",
+    "ActivationType",
+    "get_activation",
+    "get_norm",
+]

src/virtual_stain_flow/models/up_down_blocks.py → src/virtual_stain_flow/models/blocks/up_down_blocks.py

@@ -51,23 +51,19 @@ class AbstractDownBlock(AbstractBlock):
     def __init__(self, in_channels, out_channels, num_units, **kwargs):
         super().__init__(in_channels, out_channels, num_units, **kwargs)
 
-    @property
     def out_h(self, in_h: int) -> int:
         return in_h // 2
 
-    @property
     def out_w(self, in_w: int) -> int:
         return in_w // 2
 
 class AbstractUpBlock(AbstractBlock):
     def __init__(self, in_channels, out_channels, num_units, **kwargs):
         super().__init__(in_channels, out_channels, num_units, **kwargs)
 
-    @property
     def out_h(self, in_h: int) -> int:
         return in_h * 2
 
-    @property
     def out_w(self, in_w: int) -> int:
         return in_w * 2
 

src/virtual_stain_flow/models/model.py

@@ -21,7 +21,7 @@
 
 from .encoder import Encoder
 from .decoder import Decoder
-from .utils import (
+from .blocks import (
     get_activation,
     ActivationType
 )

src/virtual_stain_flow/models/stages.py

@@ -29,7 +29,7 @@
     AbstractBlock,
     Conv2DNormActBlock
 )
-from .up_down_blocks import (
+from .blocks import (
     Conv2DDownBlock,
     ConvTrans2DUpBlock
 )

src/virtual_stain_flow/models/unet.py

@@ -13,18 +13,18 @@
 import torch
 import torch.nn as nn
 
-from .utils import ActivationType
 from .factory import qualname
 from .model import BaseGeneratorModel
 from .encoder import Encoder
 from .decoder import Decoder
-from .blocks import Conv2DNormActBlock
-from .up_down_blocks import (
+from .blocks import (
+    Conv2DNormActBlock,
     IdentityBlock,
     Conv2DDownBlock,
     MaxPool2DDownBlock,
     ConvTrans2DUpBlock,
-    Bilinear2DUpsampleBlock
+    Bilinear2DUpsampleBlock,
+    ActivationType
 )
 
 class UNet(BaseGeneratorModel):

src/virtual_stain_flow/models/unext.py

@@ -15,14 +15,15 @@
 import torch
 import torch.nn as nn
 
-from .utils import ActivationType
 from .factory import qualname
 from .model import BaseGeneratorModel
 from .decoder import Decoder
-from .blocks import Conv2DConvNeXtBlock, Conv2DNormActBlock
-from .up_down_blocks import (
+from .blocks import (
+    Conv2DConvNeXtBlock, 
+    Conv2DNormActBlock,
     ConvTrans2DUpBlock,
-    PixelShuffle2DUpBlock
+    PixelShuffle2DUpBlock,
+    ActivationType
 )
 
 class ConvNeXtUNet(BaseGeneratorModel):

tests/models/test_compute_blocks.py

@@ -0,0 +1,176 @@
+import timm
+import torch
+import pytest
+import torch.nn as nn
+
+from virtual_stain_flow.models.blocks import (
+	AbstractBlock,
+	Conv2DConvNeXtBlock,
+	Conv2DNormActBlock,
+)
+
+
+class DummyBlock(AbstractBlock):
+	def forward(self, x):
+		return x
+
+
+class TestAbstractBlock:
+	@pytest.mark.parametrize(
+		("parameter", "value", "exception", "message"),
+		[
+			("in_channels", "3", TypeError, "Expected in_channels to be int"),
+			("in_channels", 0, ValueError, "Expected in_channels to be positive"),
+			("in_channels", -1, ValueError, "Expected in_channels to be positive"),
+			("out_channels", "8", TypeError, "Expected out_channels to be int"),
+			("out_channels", 0, ValueError, "Expected out_channels to be positive"),
+			("out_channels", -1, ValueError, "Expected out_channels to be positive"),
+			("num_units", "2", TypeError, "Expected num_units to be int"),
+			("num_units", 0, ValueError, "Expected num_units to be positive"),
+			("num_units", -1, ValueError, "Expected num_units to be positive"),
+		],
+	)
+	def test_rejects_invalid_channel_and_unit_values(
+		self, parameter, value, exception, message
+	):
+		kwargs = {"in_channels": 3, "out_channels": 8, "num_units": 2}
+		kwargs[parameter] = value
+
+		with pytest.raises(exception, match=message):
+			DummyBlock(**kwargs)
+
+
+class TestConv2DConvNeXtBlock:
+	@pytest.mark.parametrize(
+		("in_channels", "out_channels"),
+		[(1, 2), (3, 8), (8, 3)],
+	)
+	def test_convnext_block_produces_configured_output_channels(
+		self, in_channels, out_channels
+	):
+		block = Conv2DConvNeXtBlock(
+			in_channels=in_channels, out_channels=out_channels
+		)
+
+		output = block(torch.randn(2, in_channels, 16, 16))
+
+		assert output.shape == (2, out_channels, 16, 16)
+
+	def test_convnext_block_defaults_output_channels_to_input_channels(self):
+		block = Conv2DConvNeXtBlock(in_channels=3)
+
+		output = block(torch.randn(2, 3, 16, 16))
+
+		assert output.shape == (2, 3, 16, 16)
+
+	@pytest.mark.parametrize("num_units", [1, 2, 3])
+	def test_convnext_block_appends_configured_number_of_units(self, num_units):
+		block = Conv2DConvNeXtBlock(
+			in_channels=3, out_channels=8, num_units=num_units
+		)
+
+		convnext_units = [
+			layer
+			for layer in block.network
+			if isinstance(layer, timm.models.convnext.ConvNeXtBlock)
+		]
+
+		assert len(convnext_units) == num_units
+
+	def test_convnext_block_exposes_abstract_block_metadata_and_dimensions(self):
+		input_tensor = torch.randn(2, 3, 12, 18)
+		block = Conv2DConvNeXtBlock(in_channels=3, out_channels=8, num_units=2)
+
+		output = block(input_tensor)
+
+		assert block.in_channels == input_tensor.shape[1]
+		assert block.out_channels == output.shape[1]
+		assert block.num_units == 2
+		assert block.out_h(input_tensor.shape[2]) == output.shape[2]
+		assert block.out_w(input_tensor.shape[3]) == output.shape[3]
+
+
+class TestConv2DNormActBlock:
+	@pytest.mark.parametrize(
+		("in_channels", "out_channels"),
+		[(1, 2), (3, 8), (8, 3)],
+	)
+	def test_norm_act_block_produces_configured_output_channels(
+		self, in_channels, out_channels
+	):
+		block = Conv2DNormActBlock(
+			in_channels=in_channels, out_channels=out_channels
+		)
+
+		output = block(torch.randn(2, in_channels, 16, 16))
+
+		assert output.shape == (2, out_channels, 16, 16)
+
+	def test_norm_act_block_defaults_output_channels_to_input_channels(self):
+		block = Conv2DNormActBlock(in_channels=3)
+
+		output = block(torch.randn(2, 3, 16, 16))
+
+		assert output.shape == (2, 3, 16, 16)
+
+	@pytest.mark.parametrize("num_units", [1, 2, 3])
+	def test_norm_act_block_appends_configured_number_of_convolutions(
+		self, num_units
+	):
+		block = Conv2DNormActBlock(
+			in_channels=3, out_channels=8, num_units=num_units
+		)
+
+		convolutions = [
+			layer for layer in block.network if isinstance(layer, nn.Conv2d)
+		]
+
+		assert len(convolutions) == num_units
+
+	@pytest.mark.parametrize(
+		("norm_type", "expected_type"),
+		[
+			("batch", nn.BatchNorm2d),
+			("layer", nn.GroupNorm),
+			("none", nn.Identity),
+		],
+	)
+	def test_norm_act_block_uses_configured_normalization(
+		self, norm_type, expected_type
+	):
+		block = Conv2DNormActBlock(in_channels=3, norm_type=norm_type)
+
+		normalization_layers = block.network[1::3]
+
+		assert all(isinstance(layer, expected_type) for layer in normalization_layers)
+
+	@pytest.mark.parametrize(
+		("act_type", "expected_type"),
+		[
+			("relu", nn.ReLU),
+			("gelu", nn.GELU),
+			("sigmoid", nn.Sigmoid),
+			("softmax", nn.Softmax),
+			("none", nn.Identity),
+		],
+	)
+	def test_norm_act_block_uses_configured_activation(
+		self, act_type, expected_type
+	):
+		block = Conv2DNormActBlock(in_channels=3, act_type=act_type)
+
+		activation_layers = block.network[2::3]
+
+		assert all(isinstance(layer, expected_type) for layer in activation_layers)
+
+	def test_norm_act_block_exposes_abstract_block_metadata_and_dimensions(self):
+		input_tensor = torch.randn(2, 3, 12, 18)
+		block = Conv2DNormActBlock(in_channels=3, out_channels=8, num_units=2)
+
+		output = block(input_tensor)
+
+		assert block.in_channels == input_tensor.shape[1]
+		assert block.out_channels == output.shape[1]
+		assert block.num_units == 2
+		assert block.out_h(input_tensor.shape[2]) == output.shape[2]
+		assert block.out_w(input_tensor.shape[3]) == output.shape[3]

tests/models/test_up_down_blocks.py

@@ -0,0 +1,50 @@
+import pytest
+import torch
+
+from virtual_stain_flow.models.blocks import (
+	Bilinear2DUpsampleBlock,
+	Conv2DDownBlock,
+	ConvTrans2DUpBlock,
+	IdentityBlock,
+	MaxPool2DDownBlock,
+	PixelShuffle2DUpBlock,
+)
+
+
+class TestUpDownBlocks:
+	@pytest.mark.parametrize(
+		("block_type", "kwargs", "in_channels", "expected_out_channels", "scale"),
+		[
+			(IdentityBlock, {}, 3, 3, 1),
+			(IdentityBlock, {"out_channels": 8}, 3, 3, 1),
+			(Conv2DDownBlock, {}, 3, 6, 0.5),
+			(Conv2DDownBlock, {"out_channels": 8}, 3, 8, 0.5),
+			(MaxPool2DDownBlock, {}, 3, 3, 0.5),
+			(MaxPool2DDownBlock, {"out_channels": 8}, 3, 3, 0.5),
+			(ConvTrans2DUpBlock, {}, 4, 2, 2),
+			(ConvTrans2DUpBlock, {"out_channels": 3}, 4, 3, 2),
+			(PixelShuffle2DUpBlock, {}, 4, 4, 2),
+			(PixelShuffle2DUpBlock, {"out_channels": 8}, 4, 4, 2),
+			(Bilinear2DUpsampleBlock, {}, 3, 3, 2),
+			(Bilinear2DUpsampleBlock, {"out_channels": 8}, 3, 3, 2),
+		],
+	)
+	def test_output_channels_and_spatial_dimensions(
+		self, block_type, kwargs, in_channels, expected_out_channels, scale
+	):
+		input_tensor = torch.randn(2, in_channels, 12, 18)
+		block = block_type(in_channels=in_channels, **kwargs)
+
+		output = block(input_tensor)
+
+		expected_height = int(input_tensor.shape[2] * scale)
+		expected_width = int(input_tensor.shape[3] * scale)
+		assert block.out_channels == expected_out_channels
+		assert output.shape == (
+			input_tensor.shape[0],
+			expected_out_channels,
+			expected_height,
+			expected_width,
+		)
+		assert block.out_h(input_tensor.shape[2]) == output.shape[2]
+		assert block.out_w(input_tensor.shape[3]) == output.shape[3]