chg_package/tests/test_chg.py

"""
Unit tests for CHG algorithm
"""

import unittest
import numpy as np
from chg_algorithm import CHG, CHGOptimizer


class TestCHG(unittest.TestCase):
    
    def setUp(self):
        """Set up test fixtures"""
        self.model = CHG(input_dim=2, hidden_dim=8, num_heads=2)
        self.X_train = np.random.randn(20, 2)
        self.y_train = np.random.randn(20)
        self.X_test = np.random.randn(10, 2)
    
    def test_model_initialization(self):
        """Test model initialization"""
        self.assertEqual(self.model.input_dim, 2)
        self.assertEqual(self.model.hidden_dim, 8)
        self.assertEqual(self.model.num_heads, 2)
        self.assertEqual(self.model.head_dim, 4)
    
    def test_parameter_shapes(self):
        """Test parameter matrix shapes"""
        self.assertEqual(self.model.W_q.shape, (2, 8))
        self.assertEqual(self.model.W_k.shape, (2, 8))
        self.assertEqual(self.model.W_v.shape, (2, 8))
        self.assertEqual(self.model.W_heads.shape, (2, 1))
    
    def test_covariance_computation(self):
        """Test covariance matrix computation"""
        K = self.model._compute_covariance(self.X_train, self.X_train)
        
        # Check shape
        self.assertEqual(K.shape, (20, 20))
        
        # Check symmetry (approximately)
        self.assertTrue(np.allclose(K, K.T, atol=1e-6))
        
        # Check positive semi-definiteness
        eigenvals = np.linalg.eigvals(K)
        self.assertTrue(np.all(eigenvals >= -1e-6))
    
    def test_fit_predict(self):
        """Test fit and predict functionality"""
        pred_mean, pred_var = self.model.fit_predict(
            self.X_train, self.y_train, self.X_test
        )
        
        # Check output shapes
        self.assertEqual(pred_mean.shape, (10,))
        self.assertEqual(pred_var.shape, (10,))
        
        # Check variance is positive
        self.assertTrue(np.all(pred_var > 0))
    
    def test_log_marginal_likelihood(self):
        """Test log marginal likelihood computation"""
        lml = self.model.log_marginal_likelihood(self.X_train, self.y_train)
        
        # Should return a finite number
        self.assertTrue(np.isfinite(lml))
        self.assertIsInstance(lml, float)
    
    def test_layer_norm(self):
        """Test layer normalization"""
        x = np.random.randn(5, 8)
        gamma = np.ones(8)
        beta = np.zeros(8)
        
        normalized = self.model._layer_norm(x, gamma, beta)
        
        # Check shape preservation
        self.assertEqual(normalized.shape, x.shape)
        
        # Check normalization (mean ≈ 0, std ≈ 1)
        mean = np.mean(normalized, axis=-1)
        std = np.std(normalized, axis=-1)
        self.assertTrue(np.allclose(mean, 0, atol=1e-6))
        self.assertTrue(np.allclose(std, 1, atol=1e-6))
    
    def test_gelu_activation(self):
        """Test GELU activation function"""
        x = np.array([-2, -1, 0, 1, 2])
        result = self.model._gelu(x)
        
        # Check shape preservation
        self.assertEqual(result.shape, x.shape)
        
        # Check monotonicity
        self.assertTrue(np.all(np.diff(result) > 0))
        
        # Check specific values
        self.assertAlmostEqual(result[2], 0.0, places=6)  # GELU(0) = 0


class TestCHGOptimizer(unittest.TestCase):
    
    def setUp(self):
        """Set up test fixtures"""
        self.model = CHG(input_dim=2, hidden_dim=6, num_heads=2)
        self.optimizer = CHGOptimizer(self.model, learning_rate=0.01)
        self.X = np.random.randn(15, 2)
        self.y = np.random.randn(15)
    
    def test_optimizer_initialization(self):
        """Test optimizer initialization"""
        self.assertEqual(self.optimizer.lr, 0.01)
        self.assertIs(self.optimizer.model, self.model)
    
    def test_gradient_computation(self):
        """Test gradient computation"""
        gradients = self.optimizer.compute_gradients(self.X, self.y)
        
        # Check that gradients are computed for all parameters
        expected_params = ['W_q', 'W_k', 'W_v', 'W_ff1', 'W_ff2', 'W_heads', 'scale']
        for param in expected_params:
            self.assertIn(param, gradients)
            
            # Check gradient shape matches parameter shape
            param_shape = getattr(self.model, param).shape
            self.assertEqual(gradients[param].shape, param_shape)
    
    def test_optimization_step(self):
        """Test optimization step"""
        # Store initial parameter values
        initial_params = {}
        for param in ['W_q', 'W_k', 'W_v', 'W_heads', 'scale']:
            initial_params[param] = getattr(self.model, param).copy()
        
        # Perform optimization step
        self.optimizer.step(self.X, self.y)
        
        # Check that parameters have been updated
        for param in initial_params:
            updated_param = getattr(self.model, param)
            self.assertFalse(np.allclose(initial_params[param], updated_param))


class TestExperiment(unittest.TestCase):
    
    def test_run_chg_experiment(self):
        """Test the complete experiment function"""
        from chg_algorithm import run_chg_experiment
        
        # Redirect stdout to capture print statements
        import io
        import sys
        captured_output = io.StringIO()
        sys.stdout = captured_output
        
        try:
            model, pred_mean, pred_var = run_chg_experiment()
            
            # Restore stdout
            sys.stdout = sys.__stdout__
            
            # Check outputs
            self.assertIsInstance(model, CHG)
            self.assertEqual(pred_mean.shape, (25,))
            self.assertEqual(pred_var.shape, (25,))
            self.assertTrue(np.all(pred_var > 0))
            
            # Check that performance metrics were printed
            output = captured_output.getvalue()
            self.assertIn("CHG Performance", output)
            self.assertIn("RMSE", output)
            self.assertIn("MAE", output)
            
        finally:
            # Ensure stdout is restored even if test fails
            sys.stdout = sys.__stdout__


if __name__ == "__main__":
    unittest.main()