diff --git a/virtual_accelerator/bmad/cu_transformer.py b/virtual_accelerator/bmad/cu_transformer.py index df670df..468d1da 100644 --- a/virtual_accelerator/bmad/cu_transformer.py +++ b/virtual_accelerator/bmad/cu_transformer.py @@ -157,6 +157,10 @@ def get_tao_property(self, tao: Tao, control_name: str): return self.screen_attributes[element_name]["bins"][1] elif attr == "RESOLUTION": return self.screen_attributes[element_name]["resolution"] + elif attr == "X": + return tao.ele(element_name).orbit.x * 1e3 # convert from m to mm + elif attr == "Y": + return tao.ele(element_name).orbit.y * 1e3 # convert from m to mm else: return ele_attr[attr] diff --git a/virtual_accelerator/bmad/variables.py b/virtual_accelerator/bmad/variables.py index 84303cd..74a6eb3 100644 --- a/virtual_accelerator/bmad/variables.py +++ b/virtual_accelerator/bmad/variables.py @@ -111,6 +111,10 @@ def get_variables( # Only query TAO for elements we actually need element_type = tao.ele_head(element_name)["key"] + # handle BPMS by name since the device type is just Monitor, but we want to treat it as a BPM + if "BPM" in device_name: + element_type = "BPM" + # Apply element type mappings element_type = ELEMENT_TYPE_MAPPING.get(element_type, element_type) diff --git a/virtual_accelerator/tests/_bmad_model_test_utils.py b/virtual_accelerator/tests/_bmad_model_test_utils.py index d42d19d..3d21c6c 100644 --- a/virtual_accelerator/tests/_bmad_model_test_utils.py +++ b/virtual_accelerator/tests/_bmad_model_test_utils.py @@ -185,3 +185,71 @@ def assert_screen_image_pvs_in_supported_variables( "Screen image PVs missing from model.supported_variables: " + ", ".join(missing_image_pvs) ) + + +def assert_bpm_pvs_match_tao_lattice( + model, + bpm_attrs: tuple[str, ...] = ("X", "Y"), +) -> None: + """ + Verify that mapped BPM elements expose expected BPM PVs. + + Parameters + ---------- + model : LUMEModel + The LUME BMAD model instance to check. + bpm_attrs : tuple[str, ...] + BPM PV attributes expected for each BPM device. + + Raises + ------ + AssertionError + If BPM lattice elements are missing PV prefix mappings, or if expected + BPM PVs are missing from ``model.supported_variables``. + """ + element_names = model.tao.lat_list("*", "ele.name") + bpm_elements = { + element_name.split("#")[0] + for element_name in element_names + if element_name not in ("BEGINNING", "END") + and "BPM" in element_name.split("#")[0] + } + + if not bpm_elements: + return + + pv_prefix_by_element = { + element_name + for _, element_name in model.transformer.control_name_to_bmad.items() + } + + missing_mapping = sorted( + element_name + for element_name in bpm_elements + if element_name not in pv_prefix_by_element + ) + assert not missing_mapping, "BPM elements missing PV prefix mapping: " + ", ".join( + missing_mapping + ) + + pv_prefix_by_element = { + element_name: pv_prefix + for pv_prefix, element_name in model.transformer.control_name_to_bmad.items() + } + + supported_variable_names = set(model.supported_variables) + missing_bpm_pvs = {} + + for element_name in sorted(bpm_elements): + pv_prefix = pv_prefix_by_element[element_name] + expected_pvs = {f"{pv_prefix}:{attr}" for attr in bpm_attrs} + absent_pvs = sorted(expected_pvs - supported_variable_names) + if absent_pvs: + missing_bpm_pvs[element_name] = absent_pvs + + assert not missing_bpm_pvs, ( + "BPM PVs missing from model.supported_variables: " + + "; ".join( + f"{element}: {', '.join(pvs)}" for element, pvs in missing_bpm_pvs.items() + ) + ) diff --git a/virtual_accelerator/tests/test_cu_hxr.py b/virtual_accelerator/tests/test_cu_hxr.py index 2d7de40..9fac3b4 100644 --- a/virtual_accelerator/tests/test_cu_hxr.py +++ b/virtual_accelerator/tests/test_cu_hxr.py @@ -5,6 +5,7 @@ from virtual_accelerator.tests._bmad_model_test_utils import ( HAS_BMAD_DEPS, TEST_BEAM_PATH, + assert_bpm_pvs_match_tao_lattice, assert_bmad_model_initialization, assert_bmad_model_twiss_outputs, assert_bmad_model_track_beam_custom_path, @@ -104,6 +105,10 @@ def test_vkicker_pvs_match_tao_lattice(self): model = get_cu_hxr_bmad_model() assert_magnet_pvs_match_tao_lattice(model, "VKicker") + def test_bpm_pvs_match_tao_lattice(self): + model = get_cu_hxr_bmad_model() + assert_bpm_pvs_match_tao_lattice(model) + @pytest.mark.skipif( (not HAS_CHEETAH_DEPS) or (not HAS_LCLS_LATTICE), diff --git a/virtual_accelerator/tests/test_facet.py b/virtual_accelerator/tests/test_facet.py index 206e236..123c4d5 100644 --- a/virtual_accelerator/tests/test_facet.py +++ b/virtual_accelerator/tests/test_facet.py @@ -4,6 +4,7 @@ from virtual_accelerator.tests._bmad_model_test_utils import ( HAS_BMAD_DEPS, TEST_BEAM_PATH, + assert_bpm_pvs_match_tao_lattice, assert_bmad_model_initialization, assert_bmad_model_track_beam_custom_path, assert_bmad_model_twiss_outputs, @@ -75,3 +76,7 @@ def test_hkicker_pvs_match_tao_lattice(self): def test_vkicker_pvs_match_tao_lattice(self): model = get_facet_bmad_model() assert_magnet_pvs_match_tao_lattice(model, "VKicker") + + def test_bpm_pvs_match_tao_lattice(self): + model = get_facet_bmad_model() + assert_bpm_pvs_match_tao_lattice(model)