utilities/epic_utils/epic_design_config.py

0001 #!/usr/bin/env python3
0002 """
0003 ePIC-specific design configuration with XML modification support.
0004 Extends the base DesignConfig from the configurations module.
0005 """
0006
0007 from typing import Dict, List, Optional, Tuple, Any
0008 from pydantic import BaseModel, Field, RootModel, model_validator
0009 from pathlib import Path
0010 import yaml
0011 import os
0012 import re
0013
0014 from aid2e.utilities.configurations.base_models import BaseParameter
0015 from aid2e.utilities.configurations.design_config import DesignConfig, ParameterConstraint
0016
0017
0018 class EpicParameter(BaseParameter):
0019     """
0020     Parameter with XML modification capability for ePIC detector.
0021     Extends BaseParameter with XML path, file path, and unit information.
0022     """
0023     value: float
0024     bounds: Tuple[float, float]
0025     xml_path: str  # XPath to XML element, e.g., "//constant[@name='...']/@value"
0026     unit: Optional[str] = None  # e.g., "mm", "cm", "um"
0027
0028     @property
0029     def type(self) -> str:
0030         return "epic_range"
0031
0032
0033 class EpicParameterGroup(BaseModel):
0034     """
0035     Group of ePIC parameters that share the same XML file.
0036     """
0037     file_path: str  # Path to XML file, can include $DETECTOR_PATH
0038     parameters: Dict[str, EpicParameter]
0039
0040
0041 class EpicDesignParameters(RootModel[Dict[str, EpicParameterGroup]]):
0042     """Collection of ePIC parameter groups."""
0043
0044     @model_validator(mode="before")
0045     @classmethod
0046     def inject_qualified_names(cls, values: Dict[str, dict]):
0047         """
0048         Injects full qualified names like 'group.param' into each parameter.
0049         This ensures parameters are uniquely identified.
0050         """
0051         for group_name, group_data in values.items():
0052             param_dict = group_data.get("parameters", {})
0053             for param_name, param_data in param_dict.items():
0054                 if isinstance(param_data, dict) and "name" not in param_data:
0055                     param_data["name"] = f"{group_name}.{param_name}"
0056         return values
0057
0058
0059 class EpicDesignConfig(DesignConfig):
0060     """
0061     ePIC-specific design configuration with XML integration.
0062     Extends DesignConfig with XML modification capabilities and optimization groups.
0063
0064     Note: Uses 'epic_design_parameters' instead of 'design_parameters' to distinguish
0065     from generic configs in YAML files.
0066     """
0067     # Override to use ePIC-specific parameters
0068     design_parameters: Optional[Any] = None  # Set to None to avoid conflicts
0069     epic_design_parameters: EpicDesignParameters
0070     optimization_groups: Optional[Dict[str, List[str]]] = Field(default_factory=dict)
0071
0072     def get_flat_parameters(self) -> Dict[str, BaseParameter]:
0073         """Returns a flat dictionary of all parameters keyed by their qualified name."""
0074         flat = {}
0075         for group in self.epic_design_parameters.root.values():
0076             for param in group.parameters.values():
0077                 flat[param.name] = param
0078         return flat
0079
0080     def get_parameter_names(self) -> List[str]:
0081         """Get all parameter qualified names."""
0082         return list(self.get_flat_parameters().keys())
0083
0084     def get_xml_modifications(self, param_values: Optional[Dict[str, float]] = None) -> Dict[str, List[Tuple[str, str, float]]]:
0085         """
0086         Get XML modifications for given parameter values.
0087
0088         Args:
0089             param_values: Dictionary of qualified parameter names to values.
0090                          If None, uses default values from config.
0091
0092         Returns:
0093             Dictionary mapping file_path -> [(xml_path, unit, new_value), ...]
0094         """
0095         if param_values is None:
0096             # Use default values from config
0097             param_values = {name: param.value for name, param in self.get_flat_parameters().items()}
0098
0099         modifications = {}
0100
0101         for group_name, group in self.epic_design_parameters.root.items():
0102             # Expand environment variables in file path
0103             file_path = os.path.expandvars(group.file_path)
0104
0105             if file_path not in modifications:
0106                 modifications[file_path] = []
0107
0108             for param_name, param in group.parameters.items():
0109                 qualified_name = f"{group_name}.{param_name}"
0110                 if qualified_name in param_values:
0111                     new_value = param_values[qualified_name]
0112                     modifications[file_path].append((
0113                         param.xml_path,
0114                         param.unit or "",
0115                         new_value
0116                     ))
0117
0118         return modifications
0119
0120     def get_file_paths(self) -> List[str]:
0121         """Get all unique file paths referenced in the configuration."""
0122         return list(set(
0123             os.path.expandvars(group.file_path)
0124             for group in self.epic_design_parameters.root.values()
0125         ))
0126
0127     def get_optimization_group(self, group_name: str) -> Optional[List[str]]:
0128         """Get parameter names for a specific optimization group."""
0129         return self.optimization_groups.get(group_name) if self.optimization_groups else None
0130
0131     def get_all_optimization_groups(self) -> Dict[str, List[str]]:
0132         """Get all optimization groups."""
0133         return self.optimization_groups or {}
0134
0135
0136 class EpicDesignConfigLoader:
0137     """
0138     Loader for ePIC design configurations.
0139     Loads YAML files and instantiates EpicDesignConfig objects.
0140     """
0141
0142     @staticmethod
0143     def load(file_path: str) -> "EpicDesignConfig":
0144         """
0145         Load an ePIC design configuration from a YAML file.
0146
0147         Args:
0148             file_path: Path to the YAML configuration file
0149
0150         Returns:
0151             EpicDesignConfig instance
0152
0153         Raises:
0154             FileNotFoundError: If the file doesn't exist
0155             ValueError: If the file format is invalid
0156         """
0157         path = Path(file_path)
0158         if not path.exists():
0159             raise FileNotFoundError(f"Configuration file not found: {file_path}")
0160
0161         with open(path, 'r') as f:
0162             data = yaml.safe_load(f)
0163
0164         if 'epic_design_parameters' not in data:
0165             raise ValueError(f"Invalid configuration file format: {file_path}. Missing 'epic_design_parameters'.")
0166
0167         return EpicDesignConfig(**data)