Atomizer/atomizer-field/tests/test_synthetic.py

"""
test_synthetic.py
Synthetic tests with known analytical solutions

Tests basic functionality without real FEA data:
- Model can be created
- Forward pass works
- Loss functions compute correctly
- Predictions have correct shape
"""

import torch
import numpy as np
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from neural_models.field_predictor import create_model
from neural_models.physics_losses import create_loss_function
from torch_geometric.data import Data


def test_model_creation():
    """
    Test 1: Can we create the model?

    Expected: Model instantiates with correct number of parameters
    """
    print("    Creating GNN model...")

    config = {
        'node_feature_dim': 12,
        'edge_feature_dim': 5,
        'hidden_dim': 64,
        'num_layers': 4,
        'dropout': 0.1
    }

    model = create_model(config)

    # Count parameters
    num_params = sum(p.numel() for p in model.parameters())

    print(f"    Model created: {num_params:,} parameters")

    return {
        'status': 'PASS',
        'message': f'Model created successfully ({num_params:,} params)',
        'metrics': {'parameters': num_params}
    }


def test_forward_pass():
    """
    Test 2: Can model process data?

    Expected: Forward pass completes without errors
    """
    print("    Testing forward pass...")

    # Create model
    config = {
        'node_feature_dim': 12,
        'edge_feature_dim': 5,
        'hidden_dim': 64,
        'num_layers': 4,
        'dropout': 0.1
    }

    model = create_model(config)
    model.eval()

    # Create dummy data
    num_nodes = 100
    num_edges = 300

    x = torch.randn(num_nodes, 12)  # Node features
    edge_index = torch.randint(0, num_nodes, (2, num_edges))  # Connectivity
    edge_attr = torch.randn(num_edges, 5)  # Edge features
    batch = torch.zeros(num_nodes, dtype=torch.long)  # Batch assignment

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, batch=batch)

    # Forward pass
    with torch.no_grad():
        results = model(data, return_stress=True)

    # Check outputs
    assert 'displacement' in results, "Missing displacement output"
    assert 'stress' in results, "Missing stress output"
    assert 'von_mises' in results, "Missing von Mises output"

    # Check shapes
    assert results['displacement'].shape == (num_nodes, 6), f"Wrong displacement shape: {results['displacement'].shape}"
    assert results['stress'].shape == (num_nodes, 6), f"Wrong stress shape: {results['stress'].shape}"
    assert results['von_mises'].shape == (num_nodes,), f"Wrong von Mises shape: {results['von_mises'].shape}"

    print(f"    Displacement shape: {results['displacement'].shape} [OK]")
    print(f"    Stress shape: {results['stress'].shape} [OK]")
    print(f"    Von Mises shape: {results['von_mises'].shape} [OK]")

    return {
        'status': 'PASS',
        'message': 'Forward pass successful',
        'metrics': {
            'num_nodes': num_nodes,
            'displacement_shape': list(results['displacement'].shape),
            'stress_shape': list(results['stress'].shape)
        }
    }


def test_loss_computation():
    """
    Test 3: Do loss functions work?

    Expected: All loss types compute without errors
    """
    print("    Testing loss functions...")

    # Create dummy predictions and targets
    num_nodes = 100

    predictions = {
        'displacement': torch.randn(num_nodes, 6),
        'stress': torch.randn(num_nodes, 6),
        'von_mises': torch.abs(torch.randn(num_nodes))
    }

    targets = {
        'displacement': torch.randn(num_nodes, 6),
        'stress': torch.randn(num_nodes, 6)
    }

    loss_types = ['mse', 'relative', 'physics', 'max']
    loss_values = {}

    for loss_type in loss_types:
        loss_fn = create_loss_function(loss_type)
        losses = loss_fn(predictions, targets)

        assert 'total_loss' in losses, f"Missing total_loss for {loss_type}"
        assert not torch.isnan(losses['total_loss']), f"NaN loss for {loss_type}"
        assert not torch.isinf(losses['total_loss']), f"Inf loss for {loss_type}"

        loss_values[loss_type] = losses['total_loss'].item()
        print(f"    {loss_type.upper()} loss: {loss_values[loss_type]:.6f} [OK]")

    return {
        'status': 'PASS',
        'message': 'All loss functions working',
        'metrics': loss_values
    }


def test_batch_processing():
    """
    Test 4: Can model handle batches?

    Expected: Batch processing works correctly
    """
    print("    Testing batch processing...")

    config = {
        'node_feature_dim': 12,
        'edge_feature_dim': 5,
        'hidden_dim': 64,
        'num_layers': 4,
        'dropout': 0.1
    }

    model = create_model(config)
    model.eval()

    # Create batch of 3 graphs
    graphs = []
    for i in range(3):
        num_nodes = 50 + i * 10  # Different sizes
        num_edges = 150 + i * 30

        x = torch.randn(num_nodes, 12)
        edge_index = torch.randint(0, num_nodes, (2, num_edges))
        edge_attr = torch.randn(num_edges, 5)
        batch = torch.full((num_nodes,), i, dtype=torch.long)

        graphs.append(Data(x=x, edge_index=edge_index, edge_attr=edge_attr, batch=batch))

    # Process batch
    total_nodes = sum(g.x.shape[0] for g in graphs)

    with torch.no_grad():
        for i, graph in enumerate(graphs):
            results = model(graph, return_stress=True)
            print(f"    Graph {i+1}: {graph.x.shape[0]} nodes -> predictions [OK]")

    return {
        'status': 'PASS',
        'message': 'Batch processing successful',
        'metrics': {'num_graphs': len(graphs), 'total_nodes': total_nodes}
    }


def test_gradient_flow():
    """
    Test 5: Do gradients flow correctly?

    Expected: Gradients computed without errors
    """
    print("    Testing gradient flow...")

    config = {
        'node_feature_dim': 12,
        'edge_feature_dim': 5,
        'hidden_dim': 64,
        'num_layers': 4,
        'dropout': 0.1
    }

    model = create_model(config)
    model.train()

    # Create dummy data
    num_nodes = 50
    num_edges = 150

    x = torch.randn(num_nodes, 12)
    edge_index = torch.randint(0, num_nodes, (2, num_edges))
    edge_attr = torch.randn(num_edges, 5)
    batch = torch.zeros(num_nodes, dtype=torch.long)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, batch=batch)

    # Forward pass
    results = model(data, return_stress=True)

    # Compute loss
    targets = {
        'displacement': torch.randn(num_nodes, 6),
        'stress': torch.randn(num_nodes, 6)
    }

    loss_fn = create_loss_function('mse')
    losses = loss_fn(results, targets)

    # Backward pass
    losses['total_loss'].backward()

    # Check gradients
    has_grad = sum(1 for p in model.parameters() if p.grad is not None)
    total_params = sum(1 for _ in model.parameters())

    print(f"    Parameters with gradients: {has_grad}/{total_params} [OK]")

    assert has_grad == total_params, f"Not all parameters have gradients"

    return {
        'status': 'PASS',
        'message': 'Gradients computed successfully',
        'metrics': {
            'parameters_with_grad': has_grad,
            'total_parameters': total_params
        }
    }


if __name__ == "__main__":
    print("\nRunning synthetic tests...\n")

    tests = [
        ("Model Creation", test_model_creation),
        ("Forward Pass", test_forward_pass),
        ("Loss Computation", test_loss_computation),
        ("Batch Processing", test_batch_processing),
        ("Gradient Flow", test_gradient_flow)
    ]

    passed = 0
    failed = 0

    for name, test_func in tests:
        print(f"[TEST] {name}")
        try:
            result = test_func()
            if result['status'] == 'PASS':
                print(f"  [PASS]\n")
                passed += 1
            else:
                print(f"  [FAIL]: {result['message']}\n")
                failed += 1
        except Exception as e:
            print(f"  [FAIL]: {str(e)}\n")
            failed += 1

    print(f"\nResults: {passed} passed, {failed} failed")