329 lines
11 KiB
Python
329 lines
11 KiB
Python
|
"""
|
||
|
|
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',
|
||
|
|
]
|