387 lines
14 KiB
Python
387 lines
14 KiB
Python
|
"""
|
||
|
|
Unit tests for TrainingDataExporter
|
||
|
|
|
||
|
|
Tests the training data export functionality for AtomizerField.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import json
|
||
|
|
import tempfile
|
||
|
|
import shutil
|
||
|
|
from pathlib import Path
|
||
|
|
from optimization_engine.training_data_exporter import TrainingDataExporter, create_exporter_from_config
|
||
|
|
|
||
|
|
|
||
|
|
class TestTrainingDataExporter:
|
||
|
|
"""Test suite for TrainingDataExporter class."""
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def temp_dir(self):
|
||
|
|
"""Create a temporary directory for tests."""
|
||
|
|
temp_path = Path(tempfile.mkdtemp())
|
||
|
|
yield temp_path
|
||
|
|
# Cleanup after test
|
||
|
|
if temp_path.exists():
|
||
|
|
shutil.rmtree(temp_path)
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def sample_exporter(self, temp_dir):
|
||
|
|
"""Create a sample exporter for testing."""
|
||
|
|
return TrainingDataExporter(
|
||
|
|
export_dir=temp_dir / "test_export",
|
||
|
|
study_name="test_study",
|
||
|
|
design_variable_names=["thickness", "width", "length"],
|
||
|
|
objective_names=["max_stress", "mass"],
|
||
|
|
constraint_names=["stress_constraint"],
|
||
|
|
metadata={"test_key": "test_value"}
|
||
|
|
)
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def sample_simulation_files(self, temp_dir):
|
||
|
|
"""Create sample .dat and .op2 files for testing."""
|
||
|
|
sim_dir = temp_dir / "sim_files"
|
||
|
|
sim_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
dat_file = sim_dir / "test_model.dat"
|
||
|
|
op2_file = sim_dir / "test_model.op2"
|
||
|
|
|
||
|
|
# Create dummy files with some content
|
||
|
|
dat_file.write_text("$ Nastran input deck\nGRID,1,0,0.0,0.0,0.0\n")
|
||
|
|
op2_file.write_bytes(b"DUMMY OP2 BINARY DATA" * 100)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"dat_file": dat_file,
|
||
|
|
"op2_file": op2_file
|
||
|
|
}
|
||
|
|
|
||
|
|
def test_exporter_initialization(self, temp_dir):
|
||
|
|
"""Test that exporter initializes correctly."""
|
||
|
|
export_dir = temp_dir / "test_export"
|
||
|
|
|
||
|
|
exporter = TrainingDataExporter(
|
||
|
|
export_dir=export_dir,
|
||
|
|
study_name="beam_study",
|
||
|
|
design_variable_names=["thickness", "width"],
|
||
|
|
objective_names=["stress", "mass"]
|
||
|
|
)
|
||
|
|
|
||
|
|
assert exporter.export_dir == export_dir
|
||
|
|
assert exporter.study_name == "beam_study"
|
||
|
|
assert exporter.design_variable_names == ["thickness", "width"]
|
||
|
|
assert exporter.objective_names == ["stress", "mass"]
|
||
|
|
assert exporter.trial_count == 0
|
||
|
|
assert len(exporter.exported_trials) == 0
|
||
|
|
assert export_dir.exists()
|
||
|
|
|
||
|
|
def test_readme_creation(self, sample_exporter):
|
||
|
|
"""Test that README.md is created."""
|
||
|
|
readme_path = sample_exporter.export_dir / "README.md"
|
||
|
|
assert readme_path.exists()
|
||
|
|
|
||
|
|
content = readme_path.read_text()
|
||
|
|
assert "test_study" in content
|
||
|
|
assert "AtomizerField Training Data" in content
|
||
|
|
assert "thickness" in content
|
||
|
|
assert "max_stress" in content
|
||
|
|
|
||
|
|
def test_export_trial_success(self, sample_exporter, sample_simulation_files):
|
||
|
|
"""Test successful trial export."""
|
||
|
|
design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0}
|
||
|
|
results = {
|
||
|
|
"objectives": {"max_stress": 245.3, "mass": 1.25},
|
||
|
|
"constraints": {"stress_constraint": -54.7},
|
||
|
|
"max_displacement": 1.23
|
||
|
|
}
|
||
|
|
|
||
|
|
success = sample_exporter.export_trial(
|
||
|
|
trial_number=1,
|
||
|
|
design_variables=design_vars,
|
||
|
|
results=results,
|
||
|
|
simulation_files=sample_simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
assert success
|
||
|
|
assert sample_exporter.trial_count == 1
|
||
|
|
assert len(sample_exporter.exported_trials) == 1
|
||
|
|
|
||
|
|
# Check directory structure
|
||
|
|
trial_dir = sample_exporter.export_dir / "trial_0001"
|
||
|
|
assert trial_dir.exists()
|
||
|
|
assert (trial_dir / "input").exists()
|
||
|
|
assert (trial_dir / "output").exists()
|
||
|
|
assert (trial_dir / "input" / "model.bdf").exists()
|
||
|
|
assert (trial_dir / "output" / "model.op2").exists()
|
||
|
|
assert (trial_dir / "metadata.json").exists()
|
||
|
|
|
||
|
|
def test_export_trial_metadata(self, sample_exporter, sample_simulation_files):
|
||
|
|
"""Test that metadata.json is created correctly."""
|
||
|
|
design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0}
|
||
|
|
results = {
|
||
|
|
"objectives": {"max_stress": 245.3, "mass": 1.25},
|
||
|
|
"constraints": {"stress_constraint": -54.7},
|
||
|
|
"max_displacement": 1.23
|
||
|
|
}
|
||
|
|
|
||
|
|
sample_exporter.export_trial(
|
||
|
|
trial_number=42,
|
||
|
|
design_variables=design_vars,
|
||
|
|
results=results,
|
||
|
|
simulation_files=sample_simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
metadata_path = sample_exporter.export_dir / "trial_0042" / "metadata.json"
|
||
|
|
assert metadata_path.exists()
|
||
|
|
|
||
|
|
with open(metadata_path, 'r') as f:
|
||
|
|
metadata = json.load(f)
|
||
|
|
|
||
|
|
assert metadata["trial_number"] == 42
|
||
|
|
assert metadata["atomizer_study"] == "test_study"
|
||
|
|
assert metadata["design_parameters"] == design_vars
|
||
|
|
assert metadata["results"]["objectives"] == {"max_stress": 245.3, "mass": 1.25}
|
||
|
|
assert metadata["results"]["constraints"] == {"stress_constraint": -54.7}
|
||
|
|
assert metadata["results"]["max_displacement"] == 1.23
|
||
|
|
assert "timestamp" in metadata
|
||
|
|
|
||
|
|
def test_export_trial_missing_dat_file(self, sample_exporter, temp_dir):
|
||
|
|
"""Test export fails gracefully when .dat file is missing."""
|
||
|
|
design_vars = {"thickness": 3.5}
|
||
|
|
results = {"objectives": {"max_stress": 245.3}}
|
||
|
|
|
||
|
|
# Create only .op2 file
|
||
|
|
sim_dir = temp_dir / "sim_files"
|
||
|
|
sim_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
op2_file = sim_dir / "test_model.op2"
|
||
|
|
op2_file.write_bytes(b"DUMMY OP2 DATA")
|
||
|
|
|
||
|
|
simulation_files = {
|
||
|
|
"dat_file": sim_dir / "nonexistent.dat", # Does not exist
|
||
|
|
"op2_file": op2_file
|
||
|
|
}
|
||
|
|
|
||
|
|
success = sample_exporter.export_trial(
|
||
|
|
trial_number=1,
|
||
|
|
design_variables=design_vars,
|
||
|
|
results=results,
|
||
|
|
simulation_files=simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
assert not success
|
||
|
|
assert sample_exporter.trial_count == 0
|
||
|
|
|
||
|
|
def test_export_trial_missing_op2_file(self, sample_exporter, temp_dir):
|
||
|
|
"""Test export fails gracefully when .op2 file is missing."""
|
||
|
|
design_vars = {"thickness": 3.5}
|
||
|
|
results = {"objectives": {"max_stress": 245.3}}
|
||
|
|
|
||
|
|
# Create only .dat file
|
||
|
|
sim_dir = temp_dir / "sim_files"
|
||
|
|
sim_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
dat_file = sim_dir / "test_model.dat"
|
||
|
|
dat_file.write_text("$ Nastran input")
|
||
|
|
|
||
|
|
simulation_files = {
|
||
|
|
"dat_file": dat_file,
|
||
|
|
"op2_file": sim_dir / "nonexistent.op2" # Does not exist
|
||
|
|
}
|
||
|
|
|
||
|
|
success = sample_exporter.export_trial(
|
||
|
|
trial_number=1,
|
||
|
|
design_variables=design_vars,
|
||
|
|
results=results,
|
||
|
|
simulation_files=simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
assert not success
|
||
|
|
assert sample_exporter.trial_count == 0
|
||
|
|
|
||
|
|
def test_export_multiple_trials(self, sample_exporter, sample_simulation_files):
|
||
|
|
"""Test exporting multiple trials."""
|
||
|
|
for i in range(1, 6):
|
||
|
|
design_vars = {"thickness": 2.0 + i * 0.5, "width": 40.0 + i * 5, "length": 180.0 + i * 10}
|
||
|
|
results = {
|
||
|
|
"objectives": {"max_stress": 200.0 + i * 10, "mass": 1.0 + i * 0.1}
|
||
|
|
}
|
||
|
|
|
||
|
|
success = sample_exporter.export_trial(
|
||
|
|
trial_number=i,
|
||
|
|
design_variables=design_vars,
|
||
|
|
results=results,
|
||
|
|
simulation_files=sample_simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
assert success
|
||
|
|
|
||
|
|
assert sample_exporter.trial_count == 5
|
||
|
|
assert len(sample_exporter.exported_trials) == 5
|
||
|
|
|
||
|
|
# Verify all trial directories exist
|
||
|
|
for i in range(1, 6):
|
||
|
|
trial_dir = sample_exporter.export_dir / f"trial_{i:04d}"
|
||
|
|
assert trial_dir.exists()
|
||
|
|
assert (trial_dir / "metadata.json").exists()
|
||
|
|
|
||
|
|
def test_finalize(self, sample_exporter, sample_simulation_files):
|
||
|
|
"""Test finalize creates study_summary.json."""
|
||
|
|
# Export a few trials
|
||
|
|
for i in range(1, 4):
|
||
|
|
design_vars = {"thickness": 2.0 + i * 0.5, "width": 50.0, "length": 200.0}
|
||
|
|
results = {"objectives": {"max_stress": 200.0 + i * 10, "mass": 1.2}}
|
||
|
|
sample_exporter.export_trial(i, design_vars, results, sample_simulation_files)
|
||
|
|
|
||
|
|
# Finalize
|
||
|
|
sample_exporter.finalize()
|
||
|
|
|
||
|
|
# Check study_summary.json
|
||
|
|
summary_path = sample_exporter.export_dir / "study_summary.json"
|
||
|
|
assert summary_path.exists()
|
||
|
|
|
||
|
|
with open(summary_path, 'r') as f:
|
||
|
|
summary = json.load(f)
|
||
|
|
|
||
|
|
assert summary["study_name"] == "test_study"
|
||
|
|
assert summary["total_trials"] == 3
|
||
|
|
assert summary["design_variables"] == ["thickness", "width", "length"]
|
||
|
|
assert summary["objectives"] == ["max_stress", "mass"]
|
||
|
|
assert summary["constraints"] == ["stress_constraint"]
|
||
|
|
assert "export_timestamp" in summary
|
||
|
|
assert summary["metadata"] == {"test_key": "test_value"}
|
||
|
|
|
||
|
|
def test_create_exporter_from_config_disabled(self):
|
||
|
|
"""Test that None is returned when export is disabled."""
|
||
|
|
config = {
|
||
|
|
"training_data_export": {
|
||
|
|
"enabled": False,
|
||
|
|
"export_dir": "/some/path"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
exporter = create_exporter_from_config(config)
|
||
|
|
assert exporter is None
|
||
|
|
|
||
|
|
def test_create_exporter_from_config_missing_export_dir(self):
|
||
|
|
"""Test that None is returned when export_dir is missing."""
|
||
|
|
config = {
|
||
|
|
"training_data_export": {
|
||
|
|
"enabled": True
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
exporter = create_exporter_from_config(config)
|
||
|
|
assert exporter is None
|
||
|
|
|
||
|
|
def test_create_exporter_from_config_success(self, temp_dir):
|
||
|
|
"""Test successful exporter creation from config."""
|
||
|
|
config = {
|
||
|
|
"study_name": "beam_optimization",
|
||
|
|
"training_data_export": {
|
||
|
|
"enabled": True,
|
||
|
|
"export_dir": str(temp_dir / "export")
|
||
|
|
},
|
||
|
|
"design_variables": [
|
||
|
|
{"name": "thickness"},
|
||
|
|
{"parameter": "width"},
|
||
|
|
{"other_key": "value"} # Should use "var_2"
|
||
|
|
],
|
||
|
|
"objectives": [
|
||
|
|
{"name": "max_stress"},
|
||
|
|
{"name": "mass"}
|
||
|
|
],
|
||
|
|
"constraints": [
|
||
|
|
{"name": "stress_limit"}
|
||
|
|
],
|
||
|
|
"version": "1.0",
|
||
|
|
"optimization": {
|
||
|
|
"algorithm": "NSGA-II",
|
||
|
|
"n_trials": 100
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
exporter = create_exporter_from_config(config)
|
||
|
|
|
||
|
|
assert exporter is not None
|
||
|
|
assert exporter.study_name == "beam_optimization"
|
||
|
|
assert exporter.design_variable_names == ["thickness", "width", "var_2"]
|
||
|
|
assert exporter.objective_names == ["max_stress", "mass"]
|
||
|
|
assert exporter.constraint_names == ["stress_limit"]
|
||
|
|
assert exporter.study_metadata["atomizer_version"] == "1.0"
|
||
|
|
assert exporter.study_metadata["optimization_algorithm"] == "NSGA-II"
|
||
|
|
assert exporter.study_metadata["n_trials"] == 100
|
||
|
|
|
||
|
|
def test_create_exporter_from_config_defaults(self, temp_dir):
|
||
|
|
"""Test exporter creation with minimal config."""
|
||
|
|
config = {
|
||
|
|
"training_data_export": {
|
||
|
|
"enabled": True,
|
||
|
|
"export_dir": str(temp_dir / "export")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
exporter = create_exporter_from_config(config)
|
||
|
|
|
||
|
|
assert exporter is not None
|
||
|
|
assert exporter.study_name == "unnamed_study"
|
||
|
|
assert exporter.design_variable_names == []
|
||
|
|
assert exporter.objective_names == []
|
||
|
|
assert exporter.constraint_names == []
|
||
|
|
|
||
|
|
def test_file_content_preservation(self, sample_exporter, temp_dir):
|
||
|
|
"""Test that file contents are preserved during copy."""
|
||
|
|
# Create files with specific content
|
||
|
|
sim_dir = temp_dir / "sim_files"
|
||
|
|
sim_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
dat_content = "$ Nastran Input Deck\nGRID,1,0,0.0,0.0,0.0\nGRID,2,0,1.0,0.0,0.0\n"
|
||
|
|
op2_content = b"\x00\x01\x02\x03NASTRAN_BINARY_DATA\xFF\xFE\xFD\xFC" * 50
|
||
|
|
|
||
|
|
dat_file = sim_dir / "model.dat"
|
||
|
|
op2_file = sim_dir / "model.op2"
|
||
|
|
dat_file.write_text(dat_content)
|
||
|
|
op2_file.write_bytes(op2_content)
|
||
|
|
|
||
|
|
simulation_files = {"dat_file": dat_file, "op2_file": op2_file}
|
||
|
|
|
||
|
|
# Export trial
|
||
|
|
sample_exporter.export_trial(
|
||
|
|
trial_number=1,
|
||
|
|
design_variables={"thickness": 3.0},
|
||
|
|
results={"objectives": {"stress": 200.0}},
|
||
|
|
simulation_files=simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
# Verify copied files have same content
|
||
|
|
exported_dat = sample_exporter.export_dir / "trial_0001" / "input" / "model.bdf"
|
||
|
|
exported_op2 = sample_exporter.export_dir / "trial_0001" / "output" / "model.op2"
|
||
|
|
|
||
|
|
assert exported_dat.read_text() == dat_content
|
||
|
|
assert exported_op2.read_bytes() == op2_content
|
||
|
|
|
||
|
|
def test_trial_numbering_format(self, sample_exporter, sample_simulation_files):
|
||
|
|
"""Test that trial directories use correct numbering format."""
|
||
|
|
# Test various trial numbers
|
||
|
|
trial_numbers = [1, 9, 10, 99, 100, 999, 1000]
|
||
|
|
|
||
|
|
for trial_num in trial_numbers:
|
||
|
|
sample_exporter.export_trial(
|
||
|
|
trial_number=trial_num,
|
||
|
|
design_variables={"thickness": 3.0},
|
||
|
|
results={"objectives": {"stress": 200.0}},
|
||
|
|
simulation_files=sample_simulation_files
|
||
|
|
)
|
||
|
|
|
||
|
|
expected_dir = sample_exporter.export_dir / f"trial_{trial_num:04d}"
|
||
|
|
assert expected_dir.exists()
|
||
|
|
|
||
|
|
# Verify specific formatting
|
||
|
|
assert (sample_exporter.export_dir / "trial_0001").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_0009").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_0010").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_0099").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_0100").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_0999").exists()
|
||
|
|
assert (sample_exporter.export_dir / "trial_1000").exists()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
pytest.main([__file__, "-v"])
|