Xy screen variables

virtual_accelerator/bmad/variables.py

@@ -412,6 +412,8 @@ def get_screen_variables(
             screen_spec=screen_spec,
             index=0,  # need to reverse the order of the shape for the ArraySize0_RBV and ArraySize1_RBV variables since they are in row-major order
         ),
+        bmad_actions.BPMXVariable(name=f"{base_pv}:X", element_name=screen_name),
+        bmad_actions.BPMYVariable(name=f"{base_pv}:Y", element_name=screen_name),
     ]
 
     return variables

virtual_accelerator/cheetah/actions.py

@@ -5,6 +5,7 @@
 classes.
 """
 
+from cheetah.accelerator import Screen
 from lume_cheetah.actions import (
     CheetahReadOnlyEnumVariable,
     CheetahReadOnlyNDVariable,
@@ -314,3 +315,33 @@ def _get(self, simulator):
 
     def _set(self, simulator, value):
         super()._set(simulator, 1.0 if bool(value) else 0.0)
+
+
+class ScreenCentroidVariable(TorchScalarVariable, _ReadOnlyActionMixin):
+    """Read-only scalar for beam centroid multiplied by 1e3 (to convert to mm, mrad)."""
+
+    element_name: str
+    centroid_axis: str
+    unit: str
+
+    def _get(self, simulator):
+        element = getattr(simulator.segment, self.element_name)
+        if not isinstance(element, Screen):
+            raise ValueError(
+                f"Element {self.element_name!r} is not a Screen and cannot provide {self.centroid_axis.upper()} readback"
+            )
+        return getattr(element.get_read_beam(), self.centroid_axis).mean().item() * 1e3
+
+
+class ScreenXVariable(ScreenCentroidVariable):
+    """Read-only scalar for beam x centroid."""
+
+    centroid_axis: str = "x"
+    unit: str = "mm"
+
+
+class ScreenYVariable(ScreenCentroidVariable):
+    """Read-only scalar for beam y centroid."""
+
+    centroid_axis: str = "y"
+    unit: str = "mm"

virtual_accelerator/cheetah/variables.py

@@ -37,15 +37,13 @@
     "Image:ArraySize0_RBV": "ScreenImageArraySizeVariable",
     "RESOLUTION": "ScreenResolutionVariable",
     "IMAGE": "ScreenImageVariable",
-    "N_OF_ROW": "ScreenImageArraySizeVariable",
-    "N_OF_COL": "ScreenImageArraySizeVariable",
+    "X": "ScreenXVariable",
+    "Y": "ScreenYVariable",
 }
 
 SCREEN_ARRAY_SIZE_INDEX_BY_SUFFIX = {
     "Image:ArraySize1_RBV": 0,
-    "N_OF_ROW": 0,
     "Image:ArraySize0_RBV": 1,
-    "N_OF_COL": 1,
 }
 
 

virtual_accelerator/tests/_bmad_model_test_utils.py

@@ -8,6 +8,25 @@
 
 TEST_BEAM_PATH = os.path.join(Path(__file__).parent, "../bmad", "test_beam")
 
+DEFAULT_MAGNET_PV_ATTRS = (
+    "BCTRL",
+    "BACT",
+    "BDES",
+    "BMIN",
+    "BMAX",
+    "STATCTRLSUB.T",
+    "CTRL",
+)
+DEFAULT_SCREEN_PV_ATTRS = (
+    "Image:ArrayData",
+    "Image:ArraySize1_RBV",
+    "Image:ArraySize0_RBV",
+    "RESOLUTION",
+    "X",
+    "Y",
+)
+DEFAULT_BPM_PV_ATTRS = ("X", "Y", "TMIT")
+
 
 def _normalize_element_name(element_name: str) -> str:
     """Return an element name without any split-index suffix.
@@ -261,15 +280,7 @@ def assert_magnet_pvs_match_tao_lattice(
     model,
     element_key: str,
     excluded_elements: Iterable[str] = (),
-    element_attrs: tuple[str, ...] = (
-        "BCTRL",
-        "BACT",
-        "BDES",
-        "BMIN",
-        "BMAX",
-        "STATCTRLSUB.T",
-        "CTRL",
-    ),
+    element_attrs: tuple[str, ...] = DEFAULT_MAGNET_PV_ATTRS,
 ) -> None:
     """Assert magnet PV coverage for a Tao-backed model.
 
@@ -299,15 +310,7 @@ def assert_magnet_pvs_match_cheetah_segment(
     model,
     element_key: str,
     excluded_elements: Iterable[str] = (),
-    element_attrs: tuple[str, ...] = (
-        "BCTRL",
-        "BACT",
-        "BDES",
-        "BMIN",
-        "BMAX",
-        "STATCTRLSUB.T",
-        "CTRL",
-    ),
+    element_attrs: tuple[str, ...] = DEFAULT_MAGNET_PV_ATTRS,
 ) -> None:
     """Assert magnet PV coverage for a Cheetah-backed model.
 
@@ -339,15 +342,7 @@ def assert_magnet_pvs_match_lattice_elements(
     element_names: Sequence[str],
     element_keys: Sequence[str],
     excluded_elements: Iterable[str] = (),
-    element_attrs: tuple[str, ...] = (
-        "BCTRL",
-        "BACT",
-        "BDES",
-        "BMIN",
-        "BMAX",
-        "STATCTRLSUB.T",
-        "CTRL",
-    ),
+    element_attrs: tuple[str, ...] = DEFAULT_MAGNET_PV_ATTRS,
 ) -> None:
     """Assert magnet PV coverage from explicit lattice metadata sequences.
 
@@ -397,12 +392,7 @@ def assert_magnet_pvs_match_lattice_elements(
 def assert_screen_image_pvs_in_supported_variables(
     model,
     screen_elements: tuple[str, ...] | list[str] | None = None,
-    screen_attrs: tuple[str, ...] = (
-        "Image:ArrayData",
-        "Image:ArraySize1_RBV",
-        "Image:ArraySize0_RBV",
-        "RESOLUTION",
-    ),
+    screen_attrs: tuple[str, ...] = DEFAULT_SCREEN_PV_ATTRS,
 ) -> None:
     """
     Verify image-related PVs for screen elements are present in supported variables.
@@ -434,14 +424,9 @@ def assert_screen_image_pvs_in_supported_variables(
     )
 
 
-def assert_screen_image_pvs_match_tao_dump_locations(
+def assert_screen_image_pvs_match_tao_lattice(
     model,
-    screen_attrs: tuple[str, ...] = (
-        "Image:ArrayData",
-        "Image:ArraySize1_RBV",
-        "Image:ArraySize0_RBV",
-        "RESOLUTION",
-    ),
+    screen_attrs: tuple[str, ...] = DEFAULT_SCREEN_PV_ATTRS,
 ) -> None:
     """Assert screen image PV coverage using Tao ``dump_locations``.
 
@@ -461,12 +446,7 @@ def assert_screen_image_pvs_match_tao_dump_locations(
 
 def assert_screen_image_pvs_match_cheetah_segment(
     model,
-    screen_attrs: tuple[str, ...] = (
-        "Image:ArrayData",
-        "Image:ArraySize1_RBV",
-        "Image:ArraySize0_RBV",
-        "RESOLUTION",
-    ),
+    screen_attrs: tuple[str, ...] = DEFAULT_SCREEN_PV_ATTRS,
 ) -> None:
     """Assert screen image PV coverage for screen elements in a Cheetah segment.
 
@@ -492,7 +472,7 @@ def assert_screen_image_pvs_match_cheetah_segment(
 
 def assert_bpm_pvs_match_tao_lattice(
     model,
-    bpm_attrs: tuple[str, ...] = ("X", "Y", "TMIT"),
+    bpm_attrs: tuple[str, ...] = DEFAULT_BPM_PV_ATTRS,
 ) -> None:
     """
     Verify that mapped BPM elements expose expected BPM PVs.
@@ -527,7 +507,7 @@ def assert_bpm_pvs_match_tao_lattice(
 
 def assert_bpm_pvs_match_cheetah_segment(
     model,
-    bpm_attrs: tuple[str, ...] = ("X", "Y", "TMIT"),
+    bpm_attrs: tuple[str, ...] = DEFAULT_BPM_PV_ATTRS,
 ) -> None:
     """Assert BPM PV coverage for BPM elements in a Cheetah segment.
 
@@ -555,7 +535,7 @@ def assert_bpm_pvs_match_cheetah_segment(
 def assert_bpm_pvs_match_elements(
     model,
     bpm_elements: Iterable[str],
-    bpm_attrs: tuple[str, ...] = ("X", "Y", "TMIT"),
+    bpm_attrs: tuple[str, ...] = DEFAULT_BPM_PV_ATTRS,
 ) -> None:
     """Assert BPM PV coverage for an explicit BPM element name collection.
 

virtual_accelerator/tests/test_cu_hxr.py

@@ -25,7 +25,7 @@
     assert_magnet_pvs_match_cheetah_segment,
     assert_magnet_pvs_match_tao_lattice,
     assert_roundtrip_pv_get_set,
-    assert_screen_image_pvs_in_supported_variables,
+    assert_screen_image_pvs_match_tao_lattice,
 )
 
 CU_HXR_PROFMON_CONFIG_PATH = (
@@ -58,7 +58,7 @@ def test_initialization(self):
         model = get_cu_hxr_bmad_model(
             end_element="OTR4", track_beam=True, custom_beam_path=TEST_BEAM_PATH
         )
-        assert_screen_image_pvs_in_supported_variables(model)
+        assert_screen_image_pvs_match_tao_lattice(model)
 
         # test getting all of the supported variables to ensure no errors with screen variable setup
         _ = model.get(list(model.supported_variables))
@@ -82,7 +82,7 @@ def test_cu_hxr_twiss(self):
 
     def test_sub_lattice(self):
         model = get_cu_hxr_bmad_model("QE04#1", "OTR2")
-        assert len(model.supported_variables) < 40
+        assert len(model.supported_variables) < 50
 
         # test getting partial lattice with beam tracking
         model = get_cu_hxr_bmad_model(
@@ -150,6 +150,12 @@ def test_roundtrip_pv_get_set(self):
         )
         assert_roundtrip_pv_get_set(model)
 
+    def test_screen_pvs_match_tao_lattice(self):
+        model = get_cu_hxr_bmad_model(
+            custom_beam_path=TEST_BEAM_PATH, end_element="OTR4", track_beam=True
+        )
+        assert_screen_image_pvs_match_tao_lattice(model)
+
 
 class TestCUHXRCheetah:
     pytestmark = [