File indexing completed on 2026-11-09 09:29:52
0001 """Geometry utilities for the ePIC B0 tracker AID2E example."""
0002
0003 from pathlib import Path
0004 import shutil
0005 import xml.etree.ElementTree as ET
0006 from typing import Any, Dict
0007
0008
0009 B0_CONSTANTS = {
0010 "B0TrackerLayer1_zstart",
0011 "B0TrackerLayer2_zstart",
0012 "B0TrackerLayer3_zstart",
0013 "B0TrackerLayer4_zstart",
0014 }
0015
0016
0017 def geometry_expressions(design_point: Dict[str, Any]) -> Dict[str, str]:
0018 """
0019 Convert the AID2E B0 design point into the XML expressions used by ePIC.
0020
0021 z1 is the first-layer offset and dz2/dz3/dz4 are successive
0022 longitudinal offsets, matching the original B0 optimization.
0023 """
0024 z1 = float(design_point["b0_tracker.z1"])
0025 dz2 = float(design_point["b0_tracker.dz2"])
0026 dz3 = float(design_point["b0_tracker.dz3"])
0027 dz4 = float(design_point["b0_tracker.dz4"])
0028
0029 return {
0030 "B0TrackerLayer1_zstart":
0031 f"B0Tracker_length/2.0+({z1})",
0032 "B0TrackerLayer2_zstart":
0033 f"B0Tracker_length/2.0+({z1})+({dz2})",
0034 "B0TrackerLayer3_zstart":
0035 f"B0Tracker_length/2.0+({z1})+({dz2})+({dz3})",
0036 "B0TrackerLayer4_zstart":
0037 f"B0Tracker_length/2.0+({z1})+({dz2})+({dz3})+({dz4})",
0038 }
0039
0040
0041 def edit_b0_geometry(
0042 xml_path: str | Path,
0043 design_point: Dict[str, Any],
0044 ) -> Path:
0045 """
0046 Modify the four B0 tracker layer z-start constants in an XML file.
0047 """
0048 xml_path = Path(xml_path)
0049
0050 tree = ET.parse(xml_path)
0051 root = tree.getroot()
0052
0053 expressions = geometry_expressions(design_point)
0054 found = set()
0055
0056 for constant in root.findall(".//constant"):
0057 name = constant.get("name")
0058
0059 if name in expressions:
0060 constant.set("value", expressions[name])
0061 found.add(name)
0062
0063 missing = B0_CONSTANTS - found
0064
0065 if missing:
0066 raise RuntimeError(
0067 "B0 constants not found in XML: "
0068 + ", ".join(sorted(missing))
0069 )
0070
0071 tree.write(
0072 xml_path,
0073 encoding="utf-8",
0074 xml_declaration=True,
0075 )
0076
0077 return xml_path
0078
0079
0080 def create_b0_geometry(
0081 *,
0082 design_point: Dict[str, Any],
0083 source_xml: str | Path,
0084 output_xml: str | Path,
0085 ) -> Path:
0086 """
0087 Copy the nominal ePIC B0 tracker XML and apply one AID2E design point.
0088 """
0089 source_xml = Path(source_xml)
0090 output_xml = Path(output_xml)
0091
0092 if not source_xml.exists():
0093 raise FileNotFoundError(
0094 f"Nominal B0 XML does not exist: {source_xml}"
0095 )
0096
0097 output_xml.parent.mkdir(parents=True, exist_ok=True)
0098
0099 shutil.copy2(source_xml, output_xml)
0100
0101 return edit_b0_geometry(
0102 xml_path=output_xml,
0103 design_point=design_point,
0104 )
0105
0106
0107 def read_b0_geometry(xml_path: str | Path) -> Dict[str, str]:
0108 """
0109 Read back the four optimized B0 constants.
0110
0111 Useful for validation/tests.
0112 """
0113 tree = ET.parse(xml_path)
0114 root = tree.getroot()
0115
0116 values = {}
0117
0118 for constant in root.findall(".//constant"):
0119 name = constant.get("name")
0120
0121 if name in B0_CONSTANTS:
0122 values[name] = constant.get("value")
0123
0124 missing = B0_CONSTANTS - set(values)
0125
0126 if missing:
0127 raise RuntimeError(
0128 "B0 constants not found in XML: "
0129 + ", ".join(sorted(missing))
0130 )
0131
0132 return values
