ComputerAssignments/CA6_systematic_generalization/tests/test_systematic_generalization.py

import pytest
import torch
import numpy as np
from typing import Dict, List, Any
import tempfile
import os
import json


from src.core.components import (
    Component,
    ComponentType,
    CompositeExpression,
    SystematicGeneralizationTask,
    ExpressionGenerator,
)
from src.models.neural_architectures import (
    ModularNetwork,
    AttentionComposer,
    GraphComposer,
    HierarchicalComposer,
    MetaLearningComposer,
)
from src.models.symbolic_systems import (
    Symbol,
    Compound,
    SymbolicReasoner,
    ProgramSynthesizer,
)
from src.models.neurosymbolic_systems import (
    NeuralModuleNetwork,
    NeuroSymbolicComposer,
    HybridReasoningSystem,
)
from src.datasets.dataset_implementations import (
    SCANDataset,
    ArithmeticReasoningDataset,
    VisualReasoningDataset,
)
from src.evaluation.evaluation_framework import (
    EvaluationMetrics,
    SystematicGeneralizationEvaluator,
)

class TestCoreComponents:

    def test_component_creation(self):
        component = Component("test", ComponentType.PRIMITIVE, 0)

        assert component.name == "test"
        assert component.type == ComponentType.PRIMITIVE
        assert component.arity == 0
        assert component.representation is not None
        assert isinstance(component.representation, torch.Tensor)

    def test_component_serialization(self):
        component = Component("test", ComponentType.OPERATOR, 2)
        component_dict = component.to_dict()

        assert component_dict["name"] == "test"
        assert component_dict["type"] == "operator"
        assert component_dict["arity"] == 2


        restored_component = Component.from_dict(component_dict)
        assert restored_component.name == component.name
        assert restored_component.type == component.type
        assert restored_component.arity == component.arity

    def test_composite_expression(self):
        components = [
            Component("a", ComponentType.PRIMITIVE),
            Component("+", ComponentType.OPERATOR, 2),
            Component("b", ComponentType.PRIMITIVE),
        ]

        expr = CompositeExpression(components, "a + b", 3, 1)

        assert len(expr) == 3
        assert expr.structure == "a + b"
        assert expr.complexity == 1
        assert expr.has_component_type(ComponentType.PRIMITIVE)
        assert expr.has_component_type(ComponentType.OPERATOR)
        assert not expr.has_component_type(ComponentType.MODIFIER)

    def test_expression_generator(self):
        primitives = [
            Component("x", ComponentType.PRIMITIVE),
            Component("y", ComponentType.PRIMITIVE),
        ]
        operators = [
            Component("+", ComponentType.OPERATOR, 2),
            Component("*", ComponentType.OPERATOR, 2),
        ]

        generator = ExpressionGenerator(primitives, operators)
        expressions = generator.generate_expressions(max_depth=2, max_examples=10)

        assert len(expressions) > 0
        assert all(isinstance(expr, CompositeExpression) for expr in expressions)
        assert all(expr.complexity >= 1 for expr in expressions)

class TestNeuralArchitectures:

    def test_modular_network(self):
        model = ModularNetwork(10, 5, embed_dim=32, hidden_dim=64)


        prim1 = torch.randint(0, 10, (2,))
        op = torch.randint(0, 5, (2,))
        prim2 = torch.randint(0, 10, (2,))

        output = model(prim1, op, prim2)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

    def test_attention_composer(self):
        model = AttentionComposer(vocab_size=50, embed_dim=32, hidden_dim=64)


        tokens = torch.randint(0, 50, (2, 10))
        output = model(tokens)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

    def test_graph_composer(self):
        model = GraphComposer(vocab_size=50, embed_dim=32)


        node_ids = torch.randint(0, 50, (2, 5))
        adj_matrix = torch.rand(2, 5, 5)

        output = model(node_ids, adj_matrix)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

    def test_hierarchical_composer(self):
        model = HierarchicalComposer(vocab_size=50, embed_dim=32, hidden_dim=64)


        tokens = torch.randint(0, 50, (2, 10))
        output = model(tokens)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

    def test_meta_learning_composer(self):
        model = MetaLearningComposer(vocab_size=50, embed_dim=32)


        support_examples = [
            (torch.randint(0, 50, (2,)), torch.rand(2)),
            (torch.randint(0, 50, (2,)), torch.rand(2)),
        ]
        query_examples = torch.randint(0, 50, (2, 3))

        output = model(support_examples, query_examples)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

class TestSymbolicSystems:

    def test_symbol_evaluation(self):
        x = Symbol("x")
        y = Symbol("y", 5)


        result_x = x.evaluate({"x": 10})
        result_y = y.evaluate()

        assert result_x == 10
        assert result_y == 5

    def test_compound_evaluation(self):
        x = Symbol("x")
        y = Symbol("y")
        const_2 = Symbol("2", 2)


        expr_add = Compound("+", x, const_2)
        result_add = expr_add.evaluate({"x": 5})
        assert result_add == 7


        expr_mul = Compound("*", x, y)
        result_mul = expr_mul.evaluate({"x": 3, "y": 4})
        assert result_mul == 12

    def test_symbolic_reasoner(self):
        reasoner = SymbolicReasoner()


        reasoner.add_rule(r"\((.+) \+ 0\)", r"\1", priority=10)
        reasoner.add_fact(Symbol("fact", True))


        expr = Compound("+", Symbol("x"), Symbol("0", 0))
        result = reasoner.evaluate_expression(expr, {"x": 5})

        assert result == 5

    def test_program_synthesizer(self):
        synthesizer = ProgramSynthesizer()


        synthesizer.add_primitive("+", lambda x, y: x + y, 2)
        synthesizer.add_primitive("square", lambda x: x * x, 1)


        examples = [(1, 1), (2, 4), (3, 9)]
        program = synthesizer.synthesize_program(examples, max_depth=2)

        if program:
            assert isinstance(program, (Symbol, Compound))


            result = program.evaluate({"x": 4})
            assert result == 16

class TestNeuroSymbolicSystems:

    def test_neural_module_network(self):
        model = NeuralModuleNetwork(vocab_size=50, embed_dim=32)


        program_tokens = torch.randint(0, 50, (2, 5))
        context = torch.randn(2, 32)

        output = model(program_tokens, context)

        assert output.shape == (2, 1)
        assert not torch.isnan(output).any()

    def test_neurosymbolic_composer(self):
        model = NeuroSymbolicComposer(symbol_vocab_size=50, embed_dim=32)


        symbol_seq = torch.randint(0, 50, (2, 6))
        output = model(symbol_seq)

        assert isinstance(output, dict)
        assert "values" in output
        assert "operations" in output
        assert "memory_state" in output
        assert output["values"].shape == (2, 6, 1)

    def test_hybrid_reasoning_system(self):
        model = HybridReasoningSystem(vocab_size=50, embed_dim=32)


        input_tokens = torch.randint(0, 50, (2, 4))
        output = model(input_tokens, use_symbolic=True)

        assert isinstance(output, dict)
        assert "neural_output" in output
        assert "symbol_probabilities" in output
        assert "confidence" in output

class TestDatasets:

    def test_scan_dataset(self):
        dataset = SCANDataset(split_type="simple", max_length=10)

        assert len(dataset) > 0
        assert dataset.vocab_size > 0


        item = dataset[0]
        assert "command" in item
        assert "actions" in item
        assert "command_text" in item
        assert "action_text" in item


        stats = dataset.get_statistics()
        assert "total_examples" in stats
        assert "vocab_size" in stats

    def test_arithmetic_dataset(self):
        dataset = ArithmeticReasoningDataset(number_range=(1, 20))

        assert len(dataset) > 0
        assert dataset.vocab_size > 0


        item = dataset[0]
        assert "expression" in item
        assert "result" in item
        assert "expr_text" in item


        stats = dataset.get_statistics()
        assert "total_examples" in stats
        assert "vocab_size" in stats

    def test_visual_dataset(self):
        dataset = VisualReasoningDataset()

        assert len(dataset) > 0
        assert dataset.vocab_size > 0


        item = dataset[0]
        assert "description" in item
        assert "description_text" in item


        stats = dataset.get_statistics()
        assert "total_examples" in stats
        assert "vocab_size" in stats

class TestEvaluationFramework:

    def test_evaluation_metrics(self):
        metrics = EvaluationMetrics(
            accuracy=0.85,
            precision=0.82,
            recall=0.88,
            f1_score=0.85,
            systematic_generalization_gap=0.15,
            compositional_accuracy=0.80,
            length_generalization_accuracy=0.75,
            complexity_generalization_accuracy=0.70,
            training_time=120.5,
            inference_time=0.05,
            memory_usage=512.0,
        )

        assert metrics.accuracy == 0.85
        assert metrics.systematic_generalization_gap == 0.15
        assert metrics.training_time == 120.5

    def test_evaluator_creation(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            evaluator = SystematicGeneralizationEvaluator(temp_dir)

            assert evaluator.results_dir == temp_dir
            assert len(evaluator.evaluation_history) == 0

class TestIntegration:

    def test_end_to_end_pipeline(self):


        model = ModularNetwork(10, 5, embed_dim=16, hidden_dim=32)


        dataset = SCANDataset(split_type="simple", max_length=5)


        item = dataset[0]


        prim1 = torch.randint(0, 10, (1,))
        op = torch.randint(0, 5, (1,))
        prim2 = torch.randint(0, 10, (1,))

        output = model(prim1, op, prim2)

        assert output.shape == (1, 1)
        assert not torch.isnan(output).any()

    def test_configuration_loading(self):
        from src.experiments.experiment_runner import ExperimentConfig


        config = ExperimentConfig(
            experiment_name="test_experiment",
            embed_dim=32,
            hidden_dim=64,
            num_epochs=10,
        )


        config_dict = config.to_dict()
        assert config_dict["experiment_name"] == "test_experiment"
        assert config_dict["embed_dim"] == 32


        restored_config = ExperimentConfig.from_dict(config_dict)
        assert restored_config.experiment_name == config.experiment_name
        assert restored_config.embed_dim == config.embed_dim


@pytest.fixture
def sample_model():
    return ModularNetwork(10, 5, embed_dim=16, hidden_dim=32)

@pytest.fixture
def sample_dataset():
    return SCANDataset(split_type="simple", max_length=5)

@pytest.fixture
def sample_config():
    from src.experiments.experiment_runner import ExperimentConfig

    return ExperimentConfig(
        experiment_name="test_experiment", embed_dim=16, hidden_dim=32, num_epochs=5
    )


class TestPerformance:

    def test_model_memory_usage(self):
        model = HierarchicalComposer(vocab_size=100, embed_dim=64, hidden_dim=128)


        total_params = sum(p.numel() for p in model.parameters())


        assert total_params < 1000000

    def test_dataset_loading_speed(self):
        import time

        start_time = time.time()
        dataset = SCANDataset(split_type="simple", max_length=10)
        load_time = time.time() - start_time


        assert load_time < 5.0
        assert len(dataset) > 0

if __name__ == "__main__":

    pytest.main([__file__, "-v"])