Utility Functions

SimFrame provides various utility functions for inspecting, modifying and saving lattices and lattice elements. This notebook will demonstrate some of this functionality.

For general tips on setting up SimFrame, see getting_started.

import sys
import numpy as np
import matplotlib.pyplot as plt

import SimulationFramework.Framework as fw  # noqa E402

# Define a new framework instance, in directory 'getting_started'.
#       "clean" will empty (delete everything!) in the directory if true
#       "verbose" will print a progressbar if true
framework = fw.Framework(
        directory="./getting_started",
        clean=False,
        verbose=False
    )

scaling = 4

Load the settings from a definitions file and check which codes correspond to each line.

framework.loadSettings("Lattices/clara400_v13_SP3.def")

print([f"{line}: {framework[line].code}" for line in framework.lines])

['generator: ASTRA', 'injector400: astra', 'S02: elegant', 'L02: elegant', 'S03: elegant', 'L03: elegant', 'S04: elegant', 'L4H: elegant', 'S05: elegant', 'VBC: elegant', 'S06: elegant', 'L04: elegant', 'S07: elegant', 'SP3: elegant']

Check the global parameters accessible to this instance of SimFrame and all of the frameworkLattice objects.

print(framework.global_parameters)

{'beam': {'filename': None, 'code': None},
 'GPTLICENSE': '',
 'delete_tracking_files': False,
 'astra_use_wsl': 1,
 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started',
 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./',
 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'}

framework.generator.load_defaults("clara_400_2ps_Gaussian")

print(framework.settings)

FrameworkSettings({'settingsFilename': '/home/xkc85723/Documents/masterlattice/MasterLattice/./Lattices/clara400_v13_SP3.def', 'files': {'injector400': {'code': 'ASTRA', 'charge': {'cathode': True, 'space_charge_mode': '2D', 'mirror_charge': True}, 'input': {'particle_definition': 'initial_distribution'}, 'output': {'zstart': 0, 'end_element': 'CLA-S02-SIM-APER-01'}}, 'S02': {'code': 'elegant', 'output': {'start_element': 'CLA-S02-SIM-APER-01', 'end_element': 'CLA-L02-SIM-APER-01'}}, 'L02': {'code': 'elegant', 'output': {'start_element': 'CLA-L02-SIM-APER-01', 'end_element': 'CLA-S03-SIM-APER-01'}, 'input': {}}, 'S03': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S03-SIM-APER-01', 'end_element': 'CLA-L03-SIM-APER-01'}}, 'L03': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L03-SIM-APER-01', 'end_element': 'CLA-S04-SIM-APER-01'}}, 'S04': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S04-SIM-APER-01', 'end_element': 'CLA-L4H-SIM-APER-01'}, 'groups': {'laser-heater': {'dipoleangle': 0.0001, 'width': 0.2, 'gap': 0.02}}}, 'L4H': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L4H-SIM-APER-01', 'end_element': 'CLA-S05-SIM-APER-01'}}, 'S05': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S05-SIM-APER-01', 'end_element': 'CLA-S05-SIM-MARK-01'}}, 'VBC': {'code': 'elegant', 'model': 'test', 'input': {}, 'output': {'start_element': 'CLA-S05-SIM-MARK-01', 'end_element': 'CLA-S06-SIM-MARK-01'}, 'groups': {'bunch_compressor': {'dipoleangle': 0.1185, 'width': 0.2, 'gap': 0.02}}}, 'S06': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-S06-SIM-MARK-01', 'end_element': 'CLA-L04-SIM-APER-01'}}, 'L04': {'code': 'elegant', 'input': {}, 'output': {'start_element': 'CLA-L04-SIM-APER-01', 'end_element': 'CLA-S07-SIM-APER-01'}}, 'S07': {'code': 'elegant', 'input': {}, 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'output': {'start_element': 'CLA-S07-SIM-APER-01', 'end_element': 'CLA-SP3-SIM-MARK-01'}}, 'SP3': {'code': 'elegant', 'input': {}, 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'output': {'start_element': 'CLA-SP3-SIM-MARK-01', 'end_element': 'CLA-SP3-DIA-FCUP-01'}}}, 'generator': {'code': 'astra', 'default': 'clara_400_3ps'}, 'elements': {'filename': ['YAML/Injector400.yaml', 'YAML/S02.yaml', 'YAML/L02.yaml', 'YAML/S03.yaml', 'YAML/L03.yaml', 'YAML/S04.yaml', 'YAML/L4H.yaml', 'YAML/S05.yaml', 'YAML/VBC.yaml', 'YAML/S06.yaml', 'YAML/L04.yaml', 'YAML/S07_SP3.yaml']}, 'groups': {'bunch_compressor': {'type': 'chicane', 'elements': ['CLA-VBC-MAG-DIP-01', 'CLA-VBC-MAG-DIP-02', 'CLA-VBC-MAG-DIP-03', 'CLA-VBC-MAG-DIP-04']}}, 'global': {'ASTRAsettings': {'High_res': True, 'N_red': 1, 'EmitS': True, 'PhaseS': True, 'TrackS': True, 'RefS': True, 'Track_All': True, 'Phase_Scan': False, 'Auto_Phase': True, 'check_ref_part': False, 'Zphase': 1, 'Zemit': 1050, 'H_max': 0.0007, 'H_min': 0.0007, 'title': 'trial', 'run_no': 1, 'Smooth_x': 2, 'Smooth_y': 2, 'Smooth_z': 2, 'Max_scale': 0.1, 'Max_count': 100, 'cell_var': 2, 'min_grid': 3.424657e-13, 'Lmagnetized': False}}})

framework.change_subdirectory('./utility_functions')

Accessing and modifying elements

Elements and their parameters can be retrieved either using getElement or as a key in the framework object itself.

print(framework.getElement("CLA-S02-MAG-QUAD-01"))
print(framework["CLA-S02-MAG-QUAD-01"])

objectname='CLA-S02-MAG-QUAD-01' objecttype='quadrupole' objectdefaults={} allowedkeywords=['subelement', 'sub_elements', 'length', 'online_model_name', 'controller_name', 'pv_root', 'pv_suffixes', 'field_integral_coefficients', 'global_rotation', 'centre', 'datum', 'start', 'end', 'buffer_start', 'buffer_end', 'buffer_start_length', 'buffer_end_length', 'position_errors', 'rotation_errors', 'tilt_angle', 'array_names', 'rotation', 'smooth', 'field_reference_position', 'parent', 'k1l', 'k1', 'field_definition', 'gradient', 'fringe_field_coefficient', 'field_type', 'scale_field', 'multipoles'] length=0.1007 position_errors=[0, 0, 0] rotation_errors=[0, 0, 0] global_rotation=[0.0, 0.0, 0.0] rotation=[0, 0, 0] starting_rotation=0.0 conversion_rules_elegant={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_file_name': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} conversion_rules_ocelot={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_definition_sdds': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} starting_offset=[0, 0, 0] subelement=False field_definition=None wakefield_definition=None k1l=-0.128366 n_kicks=4 field_reference_position='middle' fringe_field_coefficient=0.0 strength_errors=[0] gradient=0.0 scale_field=False multipoles=[0] smooth=2 centre=[0.0, 0.0, 3.52715] datum=[0.0, 0.0, 3.5775] global_parameters={'beam': {'filename': None, 'code': None}, 'GPTLICENSE': '', 'delete_tracking_files': False, 'astra_use_wsl': 1, 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started', 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./', 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'} PV='CLA-S02-MAG-QUAD-01' isr_enable=True sr_enable=True type='quadrupole' bore=0.037
objectname='CLA-S02-MAG-QUAD-01' objecttype='quadrupole' objectdefaults={} allowedkeywords=['subelement', 'sub_elements', 'length', 'online_model_name', 'controller_name', 'pv_root', 'pv_suffixes', 'field_integral_coefficients', 'global_rotation', 'centre', 'datum', 'start', 'end', 'buffer_start', 'buffer_end', 'buffer_start_length', 'buffer_end_length', 'position_errors', 'rotation_errors', 'tilt_angle', 'array_names', 'rotation', 'smooth', 'field_reference_position', 'parent', 'k1l', 'k1', 'field_definition', 'gradient', 'fringe_field_coefficient', 'field_type', 'scale_field', 'multipoles'] length=0.1007 position_errors=[0, 0, 0] rotation_errors=[0, 0, 0] global_rotation=[0.0, 0.0, 0.0] rotation=[0, 0, 0] starting_rotation=0.0 conversion_rules_elegant={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_file_name': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} conversion_rules_ocelot={'length': 'l', 'entrance_edge_angle': 'e1', 'exit_edge_angle': 'e2', 'edge_field_integral': 'fint', 'half_gap': 'hgap', 'horizontal_size': 'x_max', 'vertical_size': 'y_max', 'output_filename': 'filename', 'csr_bins': 'bins', 'hangle': 'hkick', 'vangle': 'vkick', 'csrdz': 'dz', 'field_definition_sdds': 'inputfile', 'change_momentum': 'change_p0', 'bunched_beam': 'bunched_beam_mode', 'current_bins': 'n_bins', 'lsc_enable': 'lsc', 'tilt_angle': 'tilt', 'isr_enable': 'isr', 'sr_enable': 'synch_rad', 'n_kicks': 'n_slices'} starting_offset=[0, 0, 0] subelement=False field_definition=None wakefield_definition=None k1l=-0.128366 n_kicks=4 field_reference_position='middle' fringe_field_coefficient=0.0 strength_errors=[0] gradient=0.0 scale_field=False multipoles=[0] smooth=2 centre=[0.0, 0.0, 3.52715] datum=[0.0, 0.0, 3.5775] global_parameters={'beam': {'filename': None, 'code': None}, 'GPTLICENSE': '', 'delete_tracking_files': False, 'astra_use_wsl': 1, 'master_subdir': '/home/xkc85723/Documents/simframe/examples/notebooks/getting_started', 'master_lattice_location': '/home/xkc85723/Documents/masterlattice/MasterLattice/./', 'simcodes_location': '/home/xkc85723/Documents/simcodes/SimCodes/./'} PV='CLA-S02-MAG-QUAD-01' isr_enable=True sr_enable=True type='quadrupole' bore=0.037

print(framework.getElement("CLA-S02-MAG-QUAD-01", "k1l"))
print(framework["CLA-S02-MAG-QUAD-01"].k1l)

-0.128366
-0.128366

All objects of a given type defined in the entire lattice can be accessed as follows.

(NB here cavity on-crest phase is zero and ASTRA convention is followed for the phase sign)

cavnames = [c["name"] for c in framework.getElementType("cavity")]
print({k: v for k, v in zip(cavnames, framework.getElementType("cavity", "phase"))})

{'CLA-HRG1-GUN-CAV-01': 9, 'CLA-L01-LIN-CAV-01': 16, 'CLA-L02-LIN-CAV-01': 23, 'CLA-L03-LIN-CAV-01': 8, 'CLA-L4H-LIN-CAV-01': -184, 'CLA-L04-LIN-CAV-01': -45}

Various options are available for changing a lattice element, or multiple elements of a given type simultaneously.

# Set all dipole angles in the lattice to zero and use `setElementType`
values = [0 for _ in framework.getElementType("dipole")]
dipnames = [d["name"] for d in framework.getElementType("dipole")]
framework.setElementType("dipole", "angle", values)
print({k: v for k, v in zip(dipnames, framework.getElementType("dipole", "angle"))})

# Set all dipole angles to 0.1 and use modifyElementType
value = 0.1
framework.modifyElementType("dipole", "angle", 0.1)
print({k: v for k, v in zip(dipnames, framework.getElementType("dipole", "angle"))})

{'CLA-VBC-MAG-DIP-01': 0.0, 'CLA-VBC-MAG-DIP-02': 0.0, 'CLA-VBC-MAG-DIP-03': 0.0, 'CLA-VBC-MAG-DIP-04': 0.0, 'CLA-SP3-MAG-DIP-01': 0.0}
{'CLA-VBC-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-02': 0.1, 'CLA-VBC-MAG-DIP-03': 0.1, 'CLA-VBC-MAG-DIP-04': 0.1, 'CLA-SP3-MAG-DIP-01': 0.1}

s02q1 = "CLA-S02-MAG-QUAD-01"
s02q2 = "CLA-S02-MAG-QUAD-02"
setattr(framework.elementObjects[s02q1], 'k1l', 1)
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}")
framework[s02q1].k1l = 2
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}")
framework.modifyElement(s02q1, "k1l", 3)
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}")
framework.modifyElement(s02q1, ["k1l", "n_kicks"], [3, 5])
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}")
framework.modifyElement(s02q1, {"k1l": 4, "n_kicks": 6})
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}")
framework.modifyElements([s02q1, s02q2], {"k1l": 5, "n_kicks": 7})
print(f"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}, \
{s02q2} k1l = {framework.getElement(s02q2, 'k1l')}, n_kicks = {framework[s02q2].n_kicks})

CLA-S02-MAG-QUAD-01 k1l = 1.0
CLA-S02-MAG-QUAD-01 k1l = 2.0
CLA-S02-MAG-QUAD-01 k1l = 3.0
CLA-S02-MAG-QUAD-01 k1l = 3.0, n_kicks = 5
CLA-S02-MAG-QUAD-01 k1l = 4.0, n_kicks = 6
CLA-S02-MAG-QUAD-01 k1l = 5.0, n_kicks = 7, CLA-S02-MAG-QUAD-02 k1l = 5.0, n_kicks = 7

framework.saveParametersFile("params.yaml", framework.getElementType("dipole"))