diff --git a/virtual_accelerator/bmad/factory.py b/virtual_accelerator/bmad/factory.py index 247ab35..26889df 100644 --- a/virtual_accelerator/bmad/factory.py +++ b/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}") diff --git a/virtual_accelerator/models/cu_hxr.py b/virtual_accelerator/models/cu_hxr.py index 1a30aaa..42ce7ce 100644 --- a/virtual_accelerator/models/cu_hxr.py +++ b/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", diff --git a/virtual_accelerator/tests/test_cu_hxr.py b/virtual_accelerator/tests/test_cu_hxr.py index ae3fcd6..f25e328 100644 --- a/virtual_accelerator/tests/test_cu_hxr.py +++ b/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 = [