Utility Functions

SIMBA 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 SIMBA, see getting_started.

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

import simba.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",
        master_lattice=os.environ["LATTICE_LOCATION"],
        generator_defaults="clara.yaml",
        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_combined.def")

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

['generator: ASTRA', 'injector400: astra', 'Linac: elegant', 'FEBE: elegant']

Check the global parameters accessible to this instance of SIMBA 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/simba/docs/source/examples/notebooks/getting_started', 'master_lattice': '/home/xkc85723/Documents/laura-lattices/CLARA/./', 'simcodes_location': '/home/xkc85723/Documents/venv/lib/python3.12/site-packages/SimCodes/'}

framework.generator.load_defaults("clara_400_2ps_Gaussian")

print(framework.settings)

FrameworkSettings({'settingsFilename': '/home/xkc85723/Documents/laura-lattices/CLARA/./Lattices/clara400_v13_combined.def', 'global': {}, 'generator': {'code': 'astra', 'default': 'clara_400_3ps'}, 'section': '/home/xkc85723/Documents/laura-lattices/CLARA/sections.yaml', 'element_list': '/home/xkc85723/Documents/laura-lattices/CLARA/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']}}, 'elements': {}, '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'}}, 'Linac': {'code': 'elegant', 'output': {'start_element': 'CLA-S02-SIM-APER-01', 'end_element': 'CLA-FEA-SIM-START-01'}}, 'FEBE': {'code': 'elegant', 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'input': {}, 'output': {'start_element': 'CLA-FEA-SIM-START-01', 'end_element': 'CLA-FED-SIM-DUMP-01-START'}}}, 'layout': '/home/xkc85723/Documents/laura-lattices/CLARA/layouts.yaml'})

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"])

name='CLA-S02-MAG-QUAD-01' hardware_class='Magnet' hardware_type='Quadrupole' hardware_model='Generic' machine_area='S02' virtual_name='' alias=[] subelement=False CASCADING_RULES={} physical=PhysicalElement(middle=Position(x=0.0, y=0.0, z=3.52715239) datum=Position(x=0.0, y=0.0, z=3.59149478) length=0.12868478212775) electrical=ElectricalElement(minI=-53.0, maxI=53.0, read_tolerance=0.1) manufacturer=ManufacturerElement(manufacturer='Quadrupole Type 1', serial_number='13248', hardware_class='Quadrupole') simulation=MagnetSimulationElement(field_definition=None, wakefield_definition=None, field_reference_position=None, scale_field=False, n_kicks=7, smooth=2, edge_field_integral=0.5, edge1_effects=1, edge2_effects=1, sr_enable=True, integration_order=4, nonlinear=1, smoothing_half_width=1, edge_order=2, csr_bins=100, deltaL=0.0, csr_enable=True, isr_enable=True, field_amplitude=0.0) controls=ControlsInformation(variables={'readback': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:RBV', dtype=<class 'float'>, protocol='CA', units='N/A', description='Readback of CLA-S02-MAG-QUAD-01', read_only=True), 'setpoint': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:SP', dtype=<class 'float'>, protocol='CA', units='N/A', description='Setpoint of CLA-S02-MAG-QUAD-01', read_only=True), 'state': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:STATE', dtype=<class 'int'>, protocol='CA', units='N/A', description='State of CLA-S02-MAG-QUAD-01', read_only=False)}) degauss=DegaussableElement(tolerance=0.5, values=[52.89, -52.89, 31.8, -31.8, 21.2, -21.2, 10.6, -10.6, 5.3, -5.3, 0.0], steps=11) magnetic=Quadrupole_Magnet(skew=False, length=0.12868478212775, multipoles=Multipoles(K1L=Multipole(order=1 normal=5 skew=0.0 radius=0.0)), systematic_multipoles=Multipoles(), random_multipoles=Multipoles(), field_integral_coefficients=FieldIntegral(coefficients=[-2.23133410405682e-10, 4.5196171252132e-08, -3.46208258004659e-06, 0.000111195870210961, 0.0238129337415767, 0.00981229429460256]), linear_saturation_coefficients=LinearSaturationFit(m=0.025500836847330708, I_max=15.575435287034225, f=0.0, a=-1.3755963213967208e-05, I0=942.4766071889139, d=12.215564413215553, L=128.68477053001592), settle_time=45.0, entrance_edge_angle=0.0, exit_edge_angle=0.0, gap=0.032, bore=0.037, plane='horizontal', width=0.2, tilt=0.0, edge_field_integral=0.5, fringe_field_coefficient=0.0, half_gap=0.016)
name='CLA-S02-MAG-QUAD-01' hardware_class='Magnet' hardware_type='Quadrupole' hardware_model='Generic' machine_area='S02' virtual_name='' alias=[] subelement=False CASCADING_RULES={} physical=PhysicalElement(middle=Position(x=0.0, y=0.0, z=3.52715239) datum=Position(x=0.0, y=0.0, z=3.59149478) length=0.12868478212775) electrical=ElectricalElement(minI=-53.0, maxI=53.0, read_tolerance=0.1) manufacturer=ManufacturerElement(manufacturer='Quadrupole Type 1', serial_number='13248', hardware_class='Quadrupole') simulation=MagnetSimulationElement(field_definition=None, wakefield_definition=None, field_reference_position=None, scale_field=False, n_kicks=7, smooth=2, edge_field_integral=0.5, edge1_effects=1, edge2_effects=1, sr_enable=True, integration_order=4, nonlinear=1, smoothing_half_width=1, edge_order=2, csr_bins=100, deltaL=0.0, csr_enable=True, isr_enable=True, field_amplitude=0.0) controls=ControlsInformation(variables={'readback': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:RBV', dtype=<class 'float'>, protocol='CA', units='N/A', description='Readback of CLA-S02-MAG-QUAD-01', read_only=True), 'setpoint': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:SP', dtype=<class 'float'>, protocol='CA', units='N/A', description='Setpoint of CLA-S02-MAG-QUAD-01', read_only=True), 'state': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:STATE', dtype=<class 'int'>, protocol='CA', units='N/A', description='State of CLA-S02-MAG-QUAD-01', read_only=False)}) degauss=DegaussableElement(tolerance=0.5, values=[52.89, -52.89, 31.8, -31.8, 21.2, -21.2, 10.6, -10.6, 5.3, -5.3, 0.0], steps=11) magnetic=Quadrupole_Magnet(skew=False, length=0.12868478212775, multipoles=Multipoles(K1L=Multipole(order=1 normal=5 skew=0.0 radius=0.0)), systematic_multipoles=Multipoles(), random_multipoles=Multipoles(), field_integral_coefficients=FieldIntegral(coefficients=[-2.23133410405682e-10, 4.5196171252132e-08, -3.46208258004659e-06, 0.000111195870210961, 0.0238129337415767, 0.00981229429460256]), linear_saturation_coefficients=LinearSaturationFit(m=0.025500836847330708, I_max=15.575435287034225, f=0.0, a=-1.3755963213967208e-05, I0=942.4766071889139, d=12.215564413215553, L=128.68477053001592), settle_time=45.0, entrance_edge_angle=0.0, exit_edge_angle=0.0, gap=0.032, bore=0.037, plane='horizontal', width=0.2, tilt=0.0, edge_field_integral=0.5, fringe_field_coefficient=0.0, half_gap=0.016)

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

5
5

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("rfcavity")]
print({k: v for k, v in zip(cavnames, framework.getElementType("rfcavity", "phase"))})

{'CLA-L02-LIN-CAV-01': 23.0, 'CLA-L4H-LIN-CAV-01': -184.0, 'CLA-L01-LIN-CAV-01': 16.0, 'CLA-L03-LIN-CAV-01': 8.0, 'CLA-L04-LIN-CAV-01': 45.0, 'CLA-HRG1-GUN-CAV-01': 9.0}

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-FMS-MAG-DIP-01': 0, 'CLA-VBC-MAG-DIP-01': 0, 'CLA-FMS-MAG-PMD-03': 0, 'CLA-FMS-MAG-DIP-05': 0, 'CLA-FMS-MAG-PMD-02': 0, 'CLA-FMS-MAG-DIP-06': 0, 'CLA-FMS-MAG-DIP-07': 0, 'CLA-FEA-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-03': 0, 'CLA-FEA-MAG-DIP-02': 0, 'CLA-VBC-MAG-DIP-03': 0, 'CLA-FED-MAG-DIP-01': 0, 'CLA-FMS-MAG-PMD-01': 0, 'CLA-VBC-MAG-DIP-02': 0, 'CLA-SP2-MAG-DIP-01': 0, 'CLA-VBC-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-08': 0, 'CLA-FMS-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-02': 0, 'CLA-SP3-MAG-DIP-01': 0, 'CLA-FEA-MAG-DIP-01': 0, 'CLA-FEA-MAG-DIP-03': 0, 'CLA-C2V-MAG-DIP-01': 0}
{'CLA-FMS-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-01': 0.1, 'CLA-FMS-MAG-PMD-03': 0.1, 'CLA-FMS-MAG-DIP-05': 0.1, 'CLA-FMS-MAG-PMD-02': 0.1, 'CLA-FMS-MAG-DIP-06': 0.1, 'CLA-FMS-MAG-DIP-07': 0.1, 'CLA-FEA-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-03': 0.1, 'CLA-FEA-MAG-DIP-02': 0.1, 'CLA-VBC-MAG-DIP-03': 0.1, 'CLA-FED-MAG-DIP-01': 0.1, 'CLA-FMS-MAG-PMD-01': 0.1, 'CLA-VBC-MAG-DIP-02': 0.1, 'CLA-SP2-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-08': 0.1, 'CLA-FMS-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-02': 0.1, 'CLA-SP3-MAG-DIP-01': 0.1, 'CLA-FEA-MAG-DIP-01': 0.1, 'CLA-FEA-MAG-DIP-03': 0.1, 'CLA-C2V-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
CLA-S02-MAG-QUAD-01 k1l = 2
CLA-S02-MAG-QUAD-01 k1l = 3
CLA-S02-MAG-QUAD-01 k1l = 3, n_kicks = 5
CLA-S02-MAG-QUAD-01 k1l = 4, n_kicks = 6
CLA-S02-MAG-QUAD-01 k1l = 5, n_kicks = 7,CLA-S02-MAG-QUAD-02 k1l = 5, n_kicks = 7

import shutil
shutil.rmtree("./utility_functions/")
os.remove('dipoles.yaml')