""" 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', ]