Refactor test folder mirroring the main package

Author: pietrobarbiero

tests/__init__.py

@@ -1,8 +1,7 @@
 """
-Comprehensive tests for torch_concepts.nn.modules.high.base.learner
+Comprehensive tests for torch_concepts.nn.modules.high
 
-Tests the BaseLearner class with metrics setup, optimizer configuration,
-and loss computation for binary and categorical concepts.
+Tests high-level model modules (CBM, CEM, CGM, etc.).
 """
 import unittest
 import torch
@@ -11,7 +10,6 @@
 from torch_concepts.annotations import Annotations, AxisAnnotation
 from torch_concepts.distributions import Delta
 from torch_concepts.nn.modules.high.base.learner import BaseLearner
-from torchmetrics import Accuracy, MeanSquaredError
 
 
 class MockLearner(BaseLearner):
@@ -621,5 +619,125 @@ def test_instantiate_metric_dict_with_non_dict(self):
         self.assertEqual(metrics, {})
 
 
+class TestHighLevelModels(unittest.TestCase):
+    """Test high-level model architectures."""
+
+    def setUp(self):
+        """Set up common test fixtures."""
+        # Create simple annotations for testing
+        concept_labels = ['color', 'shape', 'size']
+        task_labels = ['class1', 'class2']
+        self.annotations = Annotations({
+            1: AxisAnnotation(labels=concept_labels + task_labels)
+        })
+        self.variable_distributions = {
+            'color': Delta,
+            'shape': Delta,
+            'size': Delta,
+            'class1': Delta,
+            'class2': Delta
+        }
+
+    def test_cbm_placeholder(self):
+        """Placeholder test for CBM model."""
+        # CBM requires complex setup with inference strategies
+        # This is a placeholder to ensure the test file runs
+        self.assertTrue(True)
+
+    def test_cem_placeholder(self):
+        """Placeholder test for CEM model."""
+        # CEM requires complex setup with embeddings
+        # This is a placeholder to ensure the test file runs
+        self.assertTrue(True)
+
+
+class TestBatchValidation(unittest.TestCase):
+    """Test batch structure validation in BaseLearner."""
+    
+    def setUp(self):
+        """Create a mock learner instance for testing unpack_batch."""
+        # Create a mock learner that implements both _check_batch and unpack_batch
+        self.learner = type('MockLearner', (), {})()
+        # Bind both methods from BaseLearner
+        self.learner._check_batch = BaseLearner._check_batch.__get__(self.learner)
+        self.learner.unpack_batch = BaseLearner.unpack_batch.__get__(self.learner)
+    
+    def test_valid_batch_structure(self):
+        """Test that valid batch structure is accepted."""
+        valid_batch = {
+            'inputs': torch.randn(4, 10),
+            'concepts': torch.randn(4, 2)
+        }
+        inputs, concepts, transforms = self.learner.unpack_batch(valid_batch)
+        self.assertIsNotNone(inputs)
+        self.assertIsNotNone(concepts)
+        self.assertEqual(transforms, {})
+    
+    def test_batch_with_transforms(self):
+        """Test that batch with transforms is handled correctly."""
+        batch_with_transforms = {
+            'inputs': torch.randn(4, 10),
+            'concepts': torch.randn(4, 2),
+            'transforms': {'scaler': 'some_transform'}
+        }
+        inputs, concepts, transforms = self.learner.unpack_batch(batch_with_transforms)
+        self.assertIsNotNone(inputs)
+        self.assertIsNotNone(concepts)
+        self.assertEqual(transforms, {'scaler': 'some_transform'})
+    
+    def test_missing_inputs_key(self):
+        """Test that missing 'inputs' key raises KeyError."""
+        invalid_batch = {
+            'concepts': torch.randn(4, 2)
+        }
+        with self.assertRaises(KeyError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn('inputs', str(context.exception))
+        self.assertIn("missing required keys", str(context.exception))
+    
+    def test_missing_concepts_key(self):
+        """Test that missing 'concepts' key raises KeyError."""
+        invalid_batch = {
+            'inputs': torch.randn(4, 10)
+        }
+        with self.assertRaises(KeyError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn('concepts', str(context.exception))
+        self.assertIn("missing required keys", str(context.exception))
+    
+    def test_missing_both_keys(self):
+        """Test that missing both required keys raises KeyError."""
+        invalid_batch = {
+            'data': torch.randn(4, 10)
+        }
+        with self.assertRaises(KeyError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn("missing required keys", str(context.exception))
+    
+    def test_non_dict_batch(self):
+        """Test that non-dict batch raises TypeError."""
+        invalid_batch = torch.randn(4, 10)
+        with self.assertRaises(TypeError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn("Expected batch to be a dict", str(context.exception))
+    
+    def test_tuple_batch(self):
+        """Test that tuple batch raises TypeError."""
+        invalid_batch = (torch.randn(4, 10), torch.randn(4, 2))
+        with self.assertRaises(TypeError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn("Expected batch to be a dict", str(context.exception))
+    
+    def test_empty_dict_batch(self):
+        """Test that empty dict raises KeyError with helpful message."""
+        invalid_batch = {}
+        with self.assertRaises(KeyError) as context:
+            self.learner.unpack_batch(invalid_batch)
+        self.assertIn("missing required keys", str(context.exception))
+        self.assertIn("Found keys: []", str(context.exception))
+
+
 if __name__ == '__main__':
     unittest.main()
+
+

tests/test_delta.py renamed to tests/distributions/test_delta.py

@@ -0,0 +1,86 @@
+"""
+Comprehensive tests for torch_concepts.nn.modules.low.encoders
+
+Tests all encoder modules (linear, exogenous, selector, stochastic).
+"""
+import unittest
+import torch
+import torch.nn as nn
+from torch_concepts.nn.modules.low.encoders.exogenous import LinearZU
+
+
+class TestLinearZU(unittest.TestCase):
+    """Test LinearZU."""
+
+    def test_initialization(self):
+        """Test encoder initialization."""
+        encoder = LinearZU(
+            in_features=128,
+            out_features=10,
+            exogenous_size=16
+        )
+        self.assertEqual(encoder.in_features, 128)
+        self.assertEqual(encoder.out_features, 10)
+        self.assertEqual(encoder.exogenous_size, 16)
+
+    def test_forward_shape(self):
+        """Test forward pass output shape."""
+        encoder = LinearZU(
+            in_features=64,
+            out_features=5,
+            exogenous_size=8
+        )
+        embeddings = torch.randn(4, 64)
+        output = encoder(embeddings)
+        self.assertEqual(output.shape, (4, 5, 8))
+
+    def test_gradient_flow(self):
+        """Test gradient flow through encoder."""
+        encoder = LinearZU(
+            in_features=32,
+            out_features=3,
+            exogenous_size=4
+        )
+        embeddings = torch.randn(2, 32, requires_grad=True)
+        output = encoder(embeddings)
+        loss = output.sum()
+        loss.backward()
+        self.assertIsNotNone(embeddings.grad)
+
+    def test_different_embedding_sizes(self):
+        """Test various embedding sizes."""
+        for emb_size in [4, 8, 16, 32]:
+            encoder = LinearZU(
+                in_features=64,
+                out_features=5,
+                exogenous_size=emb_size
+            )
+            embeddings = torch.randn(2, 64)
+            output = encoder(embeddings)
+            self.assertEqual(output.shape, (2, 5, emb_size))
+
+    def test_encoder_output_dimension(self):
+        """Test output dimension calculation."""
+        encoder = LinearZU(
+            in_features=128,
+            out_features=10,
+            exogenous_size=16
+        )
+        self.assertEqual(encoder.out_endogenous_dim, 10)
+        self.assertEqual(encoder.out_encoder_dim, 10 * 16)
+
+    def test_leaky_relu_activation(self):
+        """Test that LeakyReLU is applied."""
+        encoder = LinearZU(
+            in_features=32,
+            out_features=3,
+            exogenous_size=4
+        )
+        embeddings = torch.randn(2, 32)
+        output = encoder(embeddings)
+        # Output should have passed through LeakyReLU
+        self.assertIsNotNone(output)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_nn_modules_high_learner_comprehensive.py renamed to tests/nn/modules/high/base/test_learner.py

@@ -0,0 +1,116 @@
+"""
+Comprehensive tests for torch_concepts.nn.modules.low.encoders
+
+Tests all encoder modules (linear, exogenous, selector, stochastic).
+"""
+import unittest
+import torch
+import torch.nn as nn
+from torch_concepts.nn.modules.low.encoders.linear import LinearZC, LinearUC
+
+
+class TestLinearZC(unittest.TestCase):
+    """Test LinearZC."""
+
+    def test_initialization(self):
+        """Test encoder initialization."""
+        encoder = LinearZC(
+            in_features=128,
+            out_features=10
+        )
+        self.assertEqual(encoder.in_features, 128)
+        self.assertEqual(encoder.out_features, 10)
+        self.assertIsInstance(encoder.encoder, nn.Sequential)
+
+    def test_forward_shape(self):
+        """Test forward pass output shape."""
+        encoder = LinearZC(
+            in_features=128,
+            out_features=10
+        )
+        embeddings = torch.randn(4, 128)
+        output = encoder(embeddings)
+        self.assertEqual(output.shape, (4, 10))
+
+    def test_gradient_flow(self):
+        """Test gradient flow through encoder."""
+        encoder = LinearZC(
+            in_features=64,
+            out_features=5
+        )
+        embeddings = torch.randn(2, 64, requires_grad=True)
+        output = encoder(embeddings)
+        loss = output.sum()
+        loss.backward()
+        self.assertIsNotNone(embeddings.grad)
+
+    def test_batch_processing(self):
+        """Test different batch sizes."""
+        encoder = LinearZC(
+            in_features=32,
+            out_features=5
+        )
+        for batch_size in [1, 4, 8]:
+            embeddings = torch.randn(batch_size, 32)
+            output = encoder(embeddings)
+            self.assertEqual(output.shape, (batch_size, 5))
+
+    def test_with_bias_false(self):
+        """Test encoder without bias."""
+        encoder = LinearZC(
+            in_features=32,
+            out_features=5,
+            bias=False
+        )
+        embeddings = torch.randn(2, 32)
+        output = encoder(embeddings)
+        self.assertEqual(output.shape, (2, 5))
+
+
+class TestLinearUC(unittest.TestCase):
+    """Test LinearUC."""
+
+    def test_initialization(self):
+        """Test encoder initialization."""
+        encoder = LinearUC(
+            in_features_exogenous=16,
+            n_exogenous_per_concept=2
+        )
+        self.assertEqual(encoder.n_exogenous_per_concept, 2)
+
+    def test_forward_shape(self):
+        """Test forward pass output shape."""
+        encoder = LinearUC(
+            in_features_exogenous=8,
+            n_exogenous_per_concept=2
+        )
+        # Input shape: (batch, concepts, in_features * n_exogenous_per_concept)
+        exog = torch.randn(4, 5, 16)  # 8 * 2 = 16
+        output = encoder(exog)
+        self.assertEqual(output.shape, (4, 5))
+
+    def test_single_exogenous_per_concept(self):
+        """Test with single exogenous per concept."""
+        encoder = LinearUC(
+            in_features_exogenous=10,
+            n_exogenous_per_concept=1
+        )
+        exog = torch.randn(3, 4, 10)
+        output = encoder(exog)
+        self.assertEqual(output.shape, (3, 4))
+
+    def test_gradient_flow(self):
+        """Test gradient flow."""
+        encoder = LinearUC(
+            in_features_exogenous=8,
+            n_exogenous_per_concept=2
+        )
+        exog = torch.randn(2, 3, 16, requires_grad=True)
+        output = encoder(exog)
+        loss = output.sum()
+        loss.backward()
+        self.assertIsNotNone(exog.grad)
+
+
+if __name__ == '__main__':
+    unittest.main()