utilities/workflows/geometry_utils.py

0001 """Helpers for modifying detector geometry in workflows.
0002
0003 These utilities facilitate modifying detector geometries
0004 based on design parameters.
0005
0006 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0007 Homepage: https://aid2e.github.io/aid2e-framework
0008 Repository: https://github.com/aid2e/AID2E-framework.git
0009
0010 """
0011
0012 from typing import Any, Dict, List, Tuple
0013 import xml.etree.ElementTree as ET
0014
0015
0016 def modify_xml_files(modifications: Dict[str, List[Tuple[str, str, str, Any]]]) -> None:
0017     """Edit XML files
0018
0019     Helper function to apply a list of modifications to XML
0020     files (e.g. DD4hep compact files).
0021
0022     Args:
0023       modifications: dictionary mapping file path -> list of
0024                      (xml_path, attribute, unit, new_value)
0025     """
0026     for file, parameters in modifications.items():
0027
0028         tree = ET.parse(file)
0029         for parameter in parameters:
0030
0031             path, attribute, units, value = parameter
0032             element = tree.getroot().find(path)
0033
0034             if units != '':
0035                 element.set(attribute, "{}*{}".format(value, units))
0036             else:
0037                 element.set(attribute, "{}".format(value))
0038
0039         # save edits and exit
0040         tree.write(file)