Atomizer/tests/unit/test_training_data_export.py

"""
Test script for training data export functionality.

Creates a simple beam optimization study with training data export enabled
to verify end-to-end functionality of AtomizerField training data collection.
"""

import json
import shutil
from pathlib import Path

# Configuration for test study with training data export
test_config = {
    "study_name": "training_data_export_test",
    "sim_file": "examples/Models/Circular Plate/Circular_Plate.sim",
    "fem_file": "examples/Models/Circular Plate/Circular_Plate_fem1.fem",

    "design_variables": [
        {
            "name": "thickness",
            "expression_name": "thickness",
            "min": 2.0,
            "max": 8.0
        },
        {
            "name": "radius",
            "expression_name": "radius",
            "min": 80.0,
            "max": 120.0
        }
    ],

    "objectives": [
        {
            "name": "max_stress",
            "type": "minimize",
            "extractor": {
                "type": "result_parameter",
                "parameter_name": "Max Von Mises Stress"
            }
        },
        {
            "name": "mass",
            "type": "minimize",
            "extractor": {
                "type": "expression",
                "expression_name": "mass"
            }
        }
    ],

    "constraints": [
        {
            "name": "stress_limit",
            "type": "less_than",
            "value": 300.0,
            "extractor": {
                "type": "result_parameter",
                "parameter_name": "Max Von Mises Stress"
            }
        }
    ],

    "optimization": {
        "algorithm": "NSGA-II",
        "n_trials": 10,
        "population_size": 4
    },

    # Enable training data export
    "training_data_export": {
        "enabled": True,
        "export_dir": "atomizer_field_training_data/test_study_001"
    },

    "version": "1.0"
}

def main():
    """Run test study with training data export."""

    # Create study directory
    study_dir = Path("studies/training_data_export_test")
    study_dir.mkdir(parents=True, exist_ok=True)

    setup_dir = study_dir / "1_setup"
    setup_dir.mkdir(exist_ok=True)

    results_dir = study_dir / "2_results"
    results_dir.mkdir(exist_ok=True)

    # Save workflow config
    config_path = setup_dir / "workflow_config.json"
    with open(config_path, 'w') as f:
        json.dump(test_config, f, indent=2)

    print("=" * 80)
    print("TRAINING DATA EXPORT TEST STUDY")
    print("=" * 80)
    print(f"\nStudy created: {study_dir}")
    print(f"Config saved: {config_path}")
    print(f"\nTraining data will be exported to:")
    print(f"  {test_config['training_data_export']['export_dir']}")
    print(f"\nNumber of trials: {test_config['optimization']['n_trials']}")
    print("\n" + "=" * 80)
    print("To run the test:")
    print(f"  cd {study_dir}")
    print("  python run_optimization.py")
    print("=" * 80)

    # Create run_optimization.py in study directory
    run_script = study_dir / "run_optimization.py"
    run_script_content = '''"""Run optimization for training data export test."""
import sys
from pathlib import Path

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

from optimization_engine.core.runner import OptimizationRunner

def main():
    """Run the optimization."""
    config_path = Path(__file__).parent / "1_setup" / "workflow_config.json"
    runner = OptimizationRunner(config_path)
    runner.run()

if __name__ == "__main__":
    main()
'''

    with open(run_script, 'w') as f:
        f.write(run_script_content)

    print(f"\nRun script created: {run_script}")
    print("\nTest study setup complete!")

if __name__ == "__main__":
    main()