add bpms

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]

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)
 

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()
+        )
+    )

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),

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)