From 05487a54f610c53ca1796adcdd29c1bfbc7b11fc Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 29 May 2026 16:16:48 -0700 Subject: [PATCH 1/3] add BPM variable support --- virtual_accelerator/bmad/cu_transformer.py | 4 ++++ virtual_accelerator/bmad/variables.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/virtual_accelerator/bmad/cu_transformer.py b/virtual_accelerator/bmad/cu_transformer.py index df670df..afd1043 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) From 2fcb6d68f05d002849bfe9b05d4df87ddaca3b87 Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 29 May 2026 16:26:22 -0700 Subject: [PATCH 2/3] add bpm tests --- .../tests/_bmad_model_test_utils.py | 67 +++++++++++++++++++ virtual_accelerator/tests/test_cu_hxr.py | 5 ++ virtual_accelerator/tests/test_facet.py | 5 ++ 3 files changed, 77 insertions(+) diff --git a/virtual_accelerator/tests/_bmad_model_test_utils.py b/virtual_accelerator/tests/_bmad_model_test_utils.py index d42d19d..0efeb99 100644 --- a/virtual_accelerator/tests/_bmad_model_test_utils.py +++ b/virtual_accelerator/tests/_bmad_model_test_utils.py @@ -185,3 +185,70 @@ 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) From a784c39291c5845d0b512db2586ae23454b0811b Mon Sep 17 00:00:00 2001 From: Ryan Roussel Date: Fri, 29 May 2026 16:28:18 -0700 Subject: [PATCH 3/3] linting --- virtual_accelerator/bmad/cu_transformer.py | 4 ++-- virtual_accelerator/tests/_bmad_model_test_utils.py | 11 ++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/virtual_accelerator/bmad/cu_transformer.py b/virtual_accelerator/bmad/cu_transformer.py index afd1043..468d1da 100644 --- a/virtual_accelerator/bmad/cu_transformer.py +++ b/virtual_accelerator/bmad/cu_transformer.py @@ -158,9 +158,9 @@ def get_tao_property(self, tao: Tao, control_name: str): 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 + 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 + return tao.ele(element_name).orbit.y * 1e3 # convert from m to mm else: return ele_attr[attr] diff --git a/virtual_accelerator/tests/_bmad_model_test_utils.py b/virtual_accelerator/tests/_bmad_model_test_utils.py index 0efeb99..3d21c6c 100644 --- a/virtual_accelerator/tests/_bmad_model_test_utils.py +++ b/virtual_accelerator/tests/_bmad_model_test_utils.py @@ -224,10 +224,12 @@ def assert_bpm_pvs_match_tao_lattice( } missing_mapping = sorted( - element_name for element_name in bpm_elements if element_name not in pv_prefix_by_element + 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) + assert not missing_mapping, "BPM elements missing PV prefix mapping: " + ", ".join( + missing_mapping ) pv_prefix_by_element = { @@ -248,7 +250,6 @@ def assert_bpm_pvs_match_tao_lattice( 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() + f"{element}: {', '.join(pvs)}" for element, pvs in missing_bpm_pvs.items() ) )