feat(config): AtomizerSpec v2.0 Pydantic models, validators, and tests
Config Layer: - spec_models.py: Pydantic models for AtomizerSpec v2.0 - spec_validator.py: Semantic validation with detailed error reporting Extractors: - custom_extractor_loader.py: Runtime custom extractor loading - spec_extractor_builder.py: Build extractors from spec definitions Tools: - migrate_to_spec_v2.py: CLI tool for batch migration Tests: - test_migrator.py: Migration tests - test_spec_manager.py: SpecManager service tests - test_spec_api.py: REST API tests - test_mcp_tools.py: MCP tool tests - test_e2e_unified_config.py: End-to-end config tests
This commit is contained in:
328
optimization_engine/extractors/spec_extractor_builder.py
Normal file
328
optimization_engine/extractors/spec_extractor_builder.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
Spec Extractor Builder
|
||||
|
||||
Builds result extractors from AtomizerSpec v2.0 configuration.
|
||||
Combines builtin extractors with custom Python extractors.
|
||||
|
||||
P3.10: Integration with optimization runner
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from optimization_engine.extractors.custom_extractor_loader import (
|
||||
CustomExtractor,
|
||||
CustomExtractorContext,
|
||||
CustomExtractorLoader,
|
||||
load_custom_extractors,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Builtin Extractor Registry
|
||||
# ============================================================================
|
||||
|
||||
# Map of builtin extractor types to their extraction functions
|
||||
BUILTIN_EXTRACTORS = {}
|
||||
|
||||
|
||||
def _register_builtin_extractors():
|
||||
"""Lazily register builtin extractors to avoid circular imports."""
|
||||
global BUILTIN_EXTRACTORS
|
||||
if BUILTIN_EXTRACTORS:
|
||||
return
|
||||
|
||||
try:
|
||||
# Zernike OPD (recommended for mirrors)
|
||||
from optimization_engine.extractors.extract_zernike_figure import (
|
||||
ZernikeOPDExtractor,
|
||||
)
|
||||
BUILTIN_EXTRACTORS['zernike_opd'] = ZernikeOPDExtractor
|
||||
|
||||
# Mass extractors
|
||||
from optimization_engine.extractors.bdf_mass_extractor import extract_mass_from_bdf
|
||||
BUILTIN_EXTRACTORS['mass'] = extract_mass_from_bdf
|
||||
|
||||
from optimization_engine.extractors.extract_mass_from_expression import (
|
||||
extract_mass_from_expression,
|
||||
)
|
||||
BUILTIN_EXTRACTORS['mass_expression'] = extract_mass_from_expression
|
||||
|
||||
# Displacement
|
||||
from optimization_engine.extractors.extract_displacement import extract_displacement
|
||||
BUILTIN_EXTRACTORS['displacement'] = extract_displacement
|
||||
|
||||
# Stress
|
||||
from optimization_engine.extractors.extract_von_mises_stress import extract_solid_stress
|
||||
BUILTIN_EXTRACTORS['stress'] = extract_solid_stress
|
||||
|
||||
from optimization_engine.extractors.extract_principal_stress import (
|
||||
extract_principal_stress,
|
||||
)
|
||||
BUILTIN_EXTRACTORS['principal_stress'] = extract_principal_stress
|
||||
|
||||
# Frequency
|
||||
from optimization_engine.extractors.extract_frequency import extract_frequency
|
||||
BUILTIN_EXTRACTORS['frequency'] = extract_frequency
|
||||
|
||||
# Temperature
|
||||
from optimization_engine.extractors.extract_temperature import extract_temperature
|
||||
BUILTIN_EXTRACTORS['temperature'] = extract_temperature
|
||||
|
||||
# Strain energy
|
||||
from optimization_engine.extractors.extract_strain_energy import (
|
||||
extract_strain_energy,
|
||||
extract_total_strain_energy,
|
||||
)
|
||||
BUILTIN_EXTRACTORS['strain_energy'] = extract_strain_energy
|
||||
BUILTIN_EXTRACTORS['total_strain_energy'] = extract_total_strain_energy
|
||||
|
||||
# SPC forces
|
||||
from optimization_engine.extractors.extract_spc_forces import (
|
||||
extract_spc_forces,
|
||||
extract_total_reaction_force,
|
||||
)
|
||||
BUILTIN_EXTRACTORS['spc_forces'] = extract_spc_forces
|
||||
BUILTIN_EXTRACTORS['reaction_force'] = extract_total_reaction_force
|
||||
|
||||
logger.debug(f"Registered {len(BUILTIN_EXTRACTORS)} builtin extractors")
|
||||
|
||||
except ImportError as e:
|
||||
logger.warning(f"Some builtin extractors unavailable: {e}")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Spec Extractor Builder
|
||||
# ============================================================================
|
||||
|
||||
class SpecExtractorBuilder:
|
||||
"""
|
||||
Builds extraction functions from AtomizerSpec extractor definitions.
|
||||
"""
|
||||
|
||||
def __init__(self, spec: Dict[str, Any]):
|
||||
"""
|
||||
Initialize builder with an AtomizerSpec.
|
||||
|
||||
Args:
|
||||
spec: AtomizerSpec dictionary
|
||||
"""
|
||||
self.spec = spec
|
||||
self.custom_loader = CustomExtractorLoader()
|
||||
self._extractors: Dict[str, Callable] = {}
|
||||
self._custom_extractors: Dict[str, CustomExtractor] = {}
|
||||
|
||||
# Register builtin extractors
|
||||
_register_builtin_extractors()
|
||||
|
||||
def build(self) -> Dict[str, Callable]:
|
||||
"""
|
||||
Build all extractors from the spec.
|
||||
|
||||
Returns:
|
||||
Dictionary of extractor_id -> extraction_function
|
||||
"""
|
||||
for ext_def in self.spec.get('extractors', []):
|
||||
ext_id = ext_def.get('id', 'unknown')
|
||||
|
||||
if ext_def.get('builtin', True):
|
||||
# Builtin extractor
|
||||
extractor_func = self._build_builtin_extractor(ext_def)
|
||||
else:
|
||||
# Custom extractor
|
||||
extractor_func = self._build_custom_extractor(ext_def)
|
||||
|
||||
if extractor_func:
|
||||
self._extractors[ext_id] = extractor_func
|
||||
else:
|
||||
logger.warning(f"Failed to build extractor: {ext_id}")
|
||||
|
||||
return self._extractors
|
||||
|
||||
def _build_builtin_extractor(self, ext_def: Dict[str, Any]) -> Optional[Callable]:
|
||||
"""
|
||||
Build a builtin extractor function.
|
||||
|
||||
Args:
|
||||
ext_def: Extractor definition from spec
|
||||
|
||||
Returns:
|
||||
Callable extraction function or None
|
||||
"""
|
||||
ext_type = ext_def.get('type', '')
|
||||
ext_id = ext_def.get('id', '')
|
||||
config = ext_def.get('config', {})
|
||||
outputs = ext_def.get('outputs', [])
|
||||
|
||||
# Get base extractor
|
||||
base_extractor = BUILTIN_EXTRACTORS.get(ext_type)
|
||||
if base_extractor is None:
|
||||
logger.warning(f"Unknown builtin extractor type: {ext_type}")
|
||||
return None
|
||||
|
||||
# Create configured wrapper
|
||||
def create_extractor_wrapper(base, cfg, outs):
|
||||
"""Create a wrapper that applies config and extracts specified outputs."""
|
||||
def wrapper(op2_path: str, **kwargs) -> Dict[str, float]:
|
||||
"""Execute extractor and return outputs dict."""
|
||||
try:
|
||||
# Handle class-based extractors (like ZernikeOPDExtractor)
|
||||
if isinstance(base, type):
|
||||
# Instantiate with config
|
||||
instance = base(
|
||||
inner_radius=cfg.get('inner_radius_mm', 0),
|
||||
n_modes=cfg.get('n_modes', 21),
|
||||
**{k: v for k, v in cfg.items()
|
||||
if k not in ['inner_radius_mm', 'n_modes']}
|
||||
)
|
||||
raw_result = instance.extract(op2_path, **kwargs)
|
||||
else:
|
||||
# Function-based extractor
|
||||
raw_result = base(op2_path, **cfg, **kwargs)
|
||||
|
||||
# Map to output names
|
||||
result = {}
|
||||
if isinstance(raw_result, dict):
|
||||
# Use output definitions to select values
|
||||
for out_def in outs:
|
||||
out_name = out_def.get('name', '')
|
||||
source = out_def.get('source', out_name)
|
||||
if source in raw_result:
|
||||
result[out_name] = float(raw_result[source])
|
||||
elif out_name in raw_result:
|
||||
result[out_name] = float(raw_result[out_name])
|
||||
|
||||
# If no outputs defined, return all
|
||||
if not outs:
|
||||
result = {k: float(v) for k, v in raw_result.items()
|
||||
if isinstance(v, (int, float))}
|
||||
elif isinstance(raw_result, (int, float)):
|
||||
# Single value - use first output name or 'value'
|
||||
out_name = outs[0]['name'] if outs else 'value'
|
||||
result[out_name] = float(raw_result)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Extractor failed: {e}")
|
||||
return {out['name']: float('nan') for out in outs}
|
||||
|
||||
return wrapper
|
||||
|
||||
return create_extractor_wrapper(base_extractor, config, outputs)
|
||||
|
||||
def _build_custom_extractor(self, ext_def: Dict[str, Any]) -> Optional[Callable]:
|
||||
"""
|
||||
Build a custom Python extractor function.
|
||||
|
||||
Args:
|
||||
ext_def: Extractor definition with function source
|
||||
|
||||
Returns:
|
||||
Callable extraction function or None
|
||||
"""
|
||||
ext_id = ext_def.get('id', 'custom')
|
||||
func_def = ext_def.get('function', {})
|
||||
|
||||
if not func_def.get('source'):
|
||||
logger.error(f"Custom extractor {ext_id} has no source code")
|
||||
return None
|
||||
|
||||
try:
|
||||
custom_ext = CustomExtractor(
|
||||
extractor_id=ext_id,
|
||||
name=ext_def.get('name', 'Custom'),
|
||||
function_name=func_def.get('name', 'extract'),
|
||||
code=func_def.get('source', ''),
|
||||
outputs=ext_def.get('outputs', []),
|
||||
dependencies=func_def.get('dependencies', []),
|
||||
)
|
||||
custom_ext.compile()
|
||||
self._custom_extractors[ext_id] = custom_ext
|
||||
|
||||
# Create wrapper function
|
||||
def create_custom_wrapper(extractor):
|
||||
def wrapper(op2_path: str, bdf_path: str = None,
|
||||
params: Dict[str, float] = None,
|
||||
working_dir: str = None, **kwargs) -> Dict[str, float]:
|
||||
context = CustomExtractorContext(
|
||||
op2_path=op2_path,
|
||||
bdf_path=bdf_path,
|
||||
working_dir=working_dir,
|
||||
params=params or {}
|
||||
)
|
||||
return extractor.execute(context)
|
||||
return wrapper
|
||||
|
||||
return create_custom_wrapper(custom_ext)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to build custom extractor {ext_id}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Convenience Functions
|
||||
# ============================================================================
|
||||
|
||||
def build_extractors_from_spec(spec: Union[Dict[str, Any], Path, str]) -> Dict[str, Callable]:
|
||||
"""
|
||||
Build extraction functions from an AtomizerSpec.
|
||||
|
||||
Args:
|
||||
spec: AtomizerSpec dict, or path to spec JSON file
|
||||
|
||||
Returns:
|
||||
Dictionary of extractor_id -> extraction_function
|
||||
|
||||
Example:
|
||||
extractors = build_extractors_from_spec("atomizer_spec.json")
|
||||
results = extractors['E1']("model.op2")
|
||||
"""
|
||||
if isinstance(spec, (str, Path)):
|
||||
with open(spec) as f:
|
||||
spec = json.load(f)
|
||||
|
||||
builder = SpecExtractorBuilder(spec)
|
||||
return builder.build()
|
||||
|
||||
|
||||
def get_extractor_outputs(spec: Dict[str, Any], extractor_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get output definitions for an extractor.
|
||||
|
||||
Args:
|
||||
spec: AtomizerSpec dictionary
|
||||
extractor_id: ID of the extractor
|
||||
|
||||
Returns:
|
||||
List of output definitions [{name, units, description}, ...]
|
||||
"""
|
||||
for ext in spec.get('extractors', []):
|
||||
if ext.get('id') == extractor_id:
|
||||
return ext.get('outputs', [])
|
||||
return []
|
||||
|
||||
|
||||
def list_available_builtin_extractors() -> List[str]:
|
||||
"""
|
||||
List all available builtin extractor types.
|
||||
|
||||
Returns:
|
||||
List of extractor type names
|
||||
"""
|
||||
_register_builtin_extractors()
|
||||
return list(BUILTIN_EXTRACTORS.keys())
|
||||
|
||||
|
||||
__all__ = [
|
||||
'SpecExtractorBuilder',
|
||||
'build_extractors_from_spec',
|
||||
'get_extractor_outputs',
|
||||
'list_available_builtin_extractors',
|
||||
'BUILTIN_EXTRACTORS',
|
||||
]
|
||||
Reference in New Issue
Block a user