Example-61: Line indexing and slicing

# In this example sequence indexing is detailed

# Indexing can be used to extract individual elements or a sequence of elements (list of elements)
# Slicing syntax is used for sequence extraction
# Note, the first element is included, while the final one is excluded
# But both edges are kept if names are used for slicing

# Import

import matplotlib
from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
matplotlib.rcParams['text.usetex'] = True

from plotly import graph_objects

import torch

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.bpm import BPM
from model.library.line import Line

from model.command.layout import Layout

# Define simple FODO based lattice

B1 = BPM('B1', direction='inverse')
B2 = BPM('B2', direction='forward')

QF1 = Quadrupole('QF1', 0.5, +0.20)
QF2 = Quadrupole('QF2', 0.5, +0.20)

QD1 = Quadrupole('QD1', 0.5, -0.19)
QD2 = Quadrupole('QD2', 0.5, -0.19)

BM1 = Dipole('BM1', 3.50, torch.pi/4.0)
BM2 = Dipole('BM2', 3.50, torch.pi/4.0)

SF1 = Sextupole('SF1', 0.25)
SF2 = Sextupole('SF2', 0.25)

SD1 = Sextupole('SD1', 0.25)
SD2 = Sextupole('SD2', 0.25)

DR1 = Drift('DR1', 0.25)
DR2 = Drift('DR2', 0.25)
DR3 = Drift('DR3', 0.25)
DR4 = Drift('DR4', 0.25)
DR5 = Drift('DR5', 0.25)
DR6 = Drift('DR6', 0.25)
DR7 = Drift('DR7', 0.25)
DR8 = Drift('DR8', 0.25)

FODO = Line('FODO',
            [B1, QF1, DR1, SF1, DR2, BM1, DR3, SD1, DR4, QD1, QD2, DR5, SD2, DR6, BM2, DR7, SF2, DR8, QF2, B2],
            propagate=True,
            dp=0.0,
            exact=False,
            output=False,
            matrix=False)

# Print FODO line

print(FODO)

BPM(name="B1", direction="inverse", dp=0.0)
Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR3", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR4", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD1", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD2", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR5", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SD2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR6", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM2", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR7", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SF2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="B2", direction="forward", dp=0.0)

# Getting element by index (position in the first level sequence)

for i in range(len(FODO)):
    print(f'{i:02d}', FODO[i])

00 BPM(name="B1", direction="inverse", dp=0.0)
01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
04 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
05 Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
06 Drift(name="DR3", length=0.25, dp=0.0, exact=False, ns=1, order=0)
07 Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
08 Drift(name="DR4", length=0.25, dp=0.0, exact=False, ns=1, order=0)
09 Quadrupole(name="QD1", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
10 Quadrupole(name="QD2", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
11 Drift(name="DR5", length=0.25, dp=0.0, exact=False, ns=1, order=0)
12 Sextupole(name="SD2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
13 Drift(name="DR6", length=0.25, dp=0.0, exact=False, ns=1, order=0)
14 Dipole(name="BM2", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
15 Drift(name="DR7", length=0.25, dp=0.0, exact=False, ns=1, order=0)
16 Sextupole(name="SF2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
17 Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
18 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
19 BPM(name="B2", direction="forward", dp=0.0)

# Note, indexing is circular

print(FODO[1 - len(FODO)])
print(FODO[1])
print(FODO[1 + len(FODO)])

Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)

# Getting subsequence by indices (start and stop first level positions)
# Note, the first element is included, while the last one is excluded

start, stop, step = 1, 4, 1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

start, stop, step = 1, 4 + 1, 1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)

01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
04 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)

# Forward sequence (step > 0)
# Note, wraps around for stop < step

start, stop, step = 1, 5 + 1, 1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

start, stop, step = 17, 3 + 1, 1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
04 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
05 Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)

17 Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
18 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
19 BPM(name="B2", direction="forward", dp=0.0)
00 BPM(name="B1", direction="inverse", dp=0.0)
01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)

# Inverse sequence (step < 0)
# Note, wraps around for stop > step

start, stop, step = 5, 2 - 1, -1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

start, stop, step = 2, 16 + 1, -1
for element in FODO[start:stop:step]:
    print(f'{FODO.position(element.name):02d}', element)
print()

05 Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
04 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
03 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)

02 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
01 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
00 BPM(name="B1", direction="inverse", dp=0.0)
19 BPM(name="B2", direction="forward", dp=0.0)
18 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)

# Iteration
# Note, iteration propagates through nested sequences if any

for element in FODO:
    print(element)

BPM(name="B1", direction="inverse", dp=0.0)
Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR3", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR4", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD1", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD2", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR5", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SD2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR6", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM2", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR7", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SF2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="B2", direction="forward", dp=0.0)

# Getting element by name (first occurance in the first level sequence)
# Note, this returns first occurance of the element with selected name
# To avoid confusion, use elements with unique names

for name in FODO.names:
    print(name, FODO[name])

B1 BPM(name="B1", direction="inverse", dp=0.0)
QF1 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR1 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SF1 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR2 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
BM1 Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
DR3 Drift(name="DR3", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SD1 Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR4 Drift(name="DR4", length=0.25, dp=0.0, exact=False, ns=1, order=0)
QD1 Quadrupole(name="QD1", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
QD2 Quadrupole(name="QD2", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR5 Drift(name="DR5", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SD2 Sextupole(name="SD2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR6 Drift(name="DR6", length=0.25, dp=0.0, exact=False, ns=1, order=0)
BM2 Dipole(name="BM2", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
DR7 Drift(name="DR7", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SF2 Sextupole(name="SF2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR8 Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
QF2 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
B2 BPM(name="B2", direction="forward", dp=0.0)

# Name based slicing
# Note, name based slicing includes edges and wraps around (use with step equal to 1 or -1)

for element in FODO['QF1':'SD1']:
    print(element.name, element)
print()

for element in FODO['QF2':'DR1']:
    print(element.name, element)
print()

for element in FODO['QF1':'SD1':-1]:
    print(element.name, element)
print()

QF1 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR1 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SF1 Sextupole(name="SF1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR2 Drift(name="DR2", length=0.25, dp=0.0, exact=False, ns=1, order=0)
BM1 Dipole(name="BM1", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
DR3 Drift(name="DR3", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SD1 Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)

QF2 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
B2 BPM(name="B2", direction="forward", dp=0.0)
B1 BPM(name="B1", direction="inverse", dp=0.0)
QF1 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR1 Drift(name="DR1", length=0.25, dp=0.0, exact=False, ns=1, order=0)

QF1 Quadrupole(name="QF1", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
B1 BPM(name="B1", direction="inverse", dp=0.0)
B2 BPM(name="B2", direction="forward", dp=0.0)
QF2 Quadrupole(name="QF2", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR8 Drift(name="DR8", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SF2 Sextupole(name="SF2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR7 Drift(name="DR7", length=0.25, dp=0.0, exact=False, ns=1, order=0)
BM2 Dipole(name="BM2", length=3.5, angle=0.7853981633974493, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
DR6 Drift(name="DR6", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SD2 Sextupole(name="SD2", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
DR5 Drift(name="DR5", length=0.25, dp=0.0, exact=False, ns=1, order=0)
QD2 Quadrupole(name="QD2", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
QD1 Quadrupole(name="QD1", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
DR4 Drift(name="DR4", length=0.25, dp=0.0, exact=False, ns=1, order=0)
SD1 Sextupole(name="SD1", length=0.25, ms=0.0, dp=0.0, exact=False, ns=1, order=0)