add ability to end the lattice at the beginning of an element

virtual_accelerator/bmad/factory.py

@@ -5,7 +5,11 @@
 import warnings
 
 
-from virtual_accelerator.bmad.variables import get_all_element_types, get_variables
+from virtual_accelerator.bmad.variables import (
+    get_all_element_types,
+    get_normalized_element_names,
+    get_variables,
+)
 from virtual_accelerator.utils.optional_dependencies import import_optional
 from virtual_accelerator.utils.variables import get_element_attr_mapping
 
@@ -40,8 +44,33 @@ def build_bmad_model(
     custom_beam_path: str | None,
     custom_tao_commands: list[str] | None = None,
     custom_aliases: dict[str, str] | None = None,
+    end_mode: str = "end",
 ):
-    """Build a lattice-specific LUMEBmadModel from a shared implementation."""
+    """
+
+    Build a lattice-specific LUMEBmadModel from a shared implementation
+
+    Parameters
+    ----------
+    spec : BmadModelSpec
+        Specification for the Bmad model to be built.
+    start_element : str
+        Name of the starting element in the lattice.
+    end_element : str
+        Name of the ending element in the lattice.
+    track_beam : bool
+        Whether to enable beam tracking.
+    custom_beam_path : str | None
+        Path to a custom beam file, if any.
+    custom_tao_commands : list[str] | None, optional
+        List of custom Tao commands to apply, by default None.
+    custom_aliases : dict[str, str] | None, optional
+        Dictionary of custom element aliases, by default None.
+    end_mode : str, optional
+        Mode for determining the end of the lattice slice, by default "end".
+
+
+    """
 
     _check_optional_modules(
         [
@@ -61,6 +90,20 @@ def build_bmad_model(
     init_file = os.path.join(lattice_root, spec.tao_init_relpath)
     tao = Tao(f"-init {init_file} -noplot -slice_lattice {start_element}:{end_element}")
 
+    # modify the end element if end_mode is "beginning"
+    if end_mode not in ["beginning", "end"]:
+        raise ValueError(f"Invalid end_mode: {end_mode}. Must be 'beginning' or 'end'.")
+
+    if end_mode == "beginning":
+        # stop tao at second to last element
+        normalized_element_list = get_normalized_element_names(tao)
+        new_end_element = normalized_element_list[
+            normalized_element_list.index(end_element) - 1
+        ]
+        tao = Tao(
+            f"-init {init_file} -noplot -slice_lattice {start_element}:{new_end_element}"
+        )
+
     # set tracking to start_element
     tao.cmd(f"set beam track_start = {start_element}")
 

virtual_accelerator/models/cu_hxr.py

@@ -22,7 +22,11 @@
 
 
 def get_cu_hxr_bmad_model(
-    start_element="OTR2", end_element="END", track_beam=False, custom_beam_path=None
+    start_element="OTR2",
+    end_element="END",
+    track_beam=False,
+    custom_beam_path=None,
+    end_mode="end",
 ):
     """
     Get the LUMEBmadModel for the CU_HXR lattice from OTR2 to END.
@@ -37,6 +41,8 @@ def get_cu_hxr_bmad_model(
         Whether to enable beam tracking in the model. Default is False.
     custom_beam_path: str, optional
         Path to custom beam file for tracking. If None, will use default design beam. Default is None.
+    end_mode: str, optional
+        The mode for determining the end of the lattice element. Must be either "beginning" or "end". Default is "end".
 
 
     Returns
@@ -61,6 +67,7 @@ def get_cu_hxr_bmad_model(
         end_element=end_element,
         track_beam=track_beam,
         custom_beam_path=custom_beam_path,
+        end_mode=end_mode,
         custom_tao_commands=[
             "set bmad_com lr_wakes_on=false",
             "set bmad_com sr_wakes_on=false",

virtual_accelerator/tests/test_cu_hxr.py

@@ -203,6 +203,20 @@ def test_screen_pvs_match_tao_lattice(self):
         )
         assert_screen_image_pvs_match_tao_lattice(model, screen_attrs=SCREEN_PV_ATTRS)
 
+    def test_end_mode(self):
+        model = get_cu_hxr_bmad_model(
+            end_element="OTR4", track_beam=False, end_mode="beginning"
+        )
+        assert model.get("name")[-1] == "DB00A"
+
+        model = get_cu_hxr_bmad_model(
+            start_element="YAG03",
+            end_element="TCAV0",
+            track_beam=False,
+            end_mode="beginning",
+        )
+        assert model.get("name")[-1] == "DH06"
+
 
 class TestCUHXRCheetah:
     pytestmark = [