Skip to content

Commit

Permalink
Merge pull request #40 from roman-martin/feature_arex
Browse files Browse the repository at this point in the history 
Add arex/arey to available alignment errors
  • Loading branch information
@rdemaria
rdemaria authored Aug 25, 2020
2 parents 714ca8b + 79d9a41 commit 4ab545d
Show file tree
Hide file tree
Showing 3 changed files with 229 additions and 10 deletions.
48 changes: 38 additions & 10 deletions pysixtrack/line.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
# missing access to particles._m:
deg2rad = np.pi / 180.


class Line(Element):
_description = [
("elements", "", "List of elements", ()),
Expand Down Expand Up @@ -357,22 +358,37 @@ def find_element_ids(self, element_name):
"""
# will raise error if element not present:
idx_el = self.element_names.index(element_name)
idx_after_el = idx_el + 1
if self.element_names[idx_after_el] == element_name + "_aperture":
idx_after_el += 1
try:
# if aperture marker is present
idx_after_el = self.element_names.index(element_name + "_aperture") + 1
except ValueError:
# if aperture marker is not present
idx_after_el = idx_el + 1
return idx_el, idx_after_el

def add_offset_error_to(self, element_name, dx=0, dy=0):
idx_el, idx_after_el = self.find_element_ids(element_name)
if not dx and not dy:
return
xyshift = elements.XYShift(dx=dx, dy=dy)
inv_xyshift = elements.XYShift(dx=-dx, dy=-dy)
self.insert_element(idx_el, xyshift, element_name + "_offset_in")
self.insert_element(
idx_after_el + 1, inv_xyshift, element_name + "_offset_out"
)

def add_aperture_offset_error_to(self, element_name, arex=0, arey=0):
idx_el, idx_after_el = self.find_element_ids(element_name)
idx_el_aper = idx_after_el - 1
if not self.element_names[idx_el_aper] == element_name + "_aperture":
# it is allowed to provide arex/arey without providing an aperture
print('Info: Element', element_name, ': arex/y provided without aperture -> arex/y ignored')
return
xyshift = elements.XYShift(dx=arex, dy=arey)
inv_xyshift = elements.XYShift(dx=-arex, dy=-arey)
self.insert_element(idx_el_aper, xyshift, element_name + "_aperture_offset_in")
self.insert_element(
idx_after_el + 1, inv_xyshift, element_name + "_aperture_offset_out"
)

def add_tilt_error_to(self, element_name, angle):
'''Alignment error of transverse rotation around s-axis.
The element corresponding to the given `element_name`
Expand All @@ -384,8 +400,6 @@ def add_tilt_error_to(self, element_name, angle):
by `angle` as well.
'''
idx_el, idx_after_el = self.find_element_ids(element_name)
if not angle:
return
element = self.elements[self.element_names.index(element_name)]
if isinstance(element, elements.Multipole) and (
element.hxl or element.hyl):
Expand Down Expand Up @@ -450,7 +464,7 @@ def apply_madx_errors(self, error_table):
if error_type == "name":
continue
if any(error_table[error_type]):
if error_type in ["dx", "dy", "dpsi"]:
if error_type in ["dx", "dy", "dpsi", "arex", "arey"]:
# available alignment error
continue
elif error_type[:1] == "k" and error_type[-1:] == "l":
Expand Down Expand Up @@ -478,14 +492,28 @@ def apply_madx_errors(self, error_table):
dy = error_table["dy"][i_line]
except KeyError:
dy = 0
self.add_offset_error_to(element_name, dx, dy)
if dx or dy:
self.add_offset_error_to(element_name, dx, dy)

# add tilt
try:
dpsi = error_table["dpsi"][i_line]
except KeyError:
dpsi = 0
if dpsi:
self.add_tilt_error_to(element_name, angle=dpsi / deg2rad)

# add aperture-only offset
try:
arex = error_table["arex"][i_line]
except KeyError:
arex = 0
try:
arey = error_table["arey"][i_line]
except KeyError:
pass
arey = 0
if arex or arey:
self.add_aperture_offset_error_to(element_name, arex, arey)

# add multipole error
knl = [
Expand Down
2 changes: 2 additions & 0 deletions tests/test_line.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,15 @@ def test_line():
assert len(line) == n_elements

line.add_offset_error_to(multipole_name, dx=0, dy=0)
n_elements += 2
assert len(line) == n_elements

line.add_offset_error_to(multipole_name, dx=0.2, dy=-0.003)
n_elements += 2
assert len(line) == n_elements

line.add_tilt_error_to(multipole_name, angle=0)
n_elements += 2
assert len(line) == n_elements

line.add_tilt_error_to(multipole_name, angle=0.1)
Expand Down
189 changes: 189 additions & 0 deletions tests/test_madx_import.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,3 +107,192 @@ def test_madx_import():
)
else:
raise ValueError("mode not understood")


def test_error_import():
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)

from cpymad.madx import Madx

madx = Madx()

madx.input('''
MQ1: Quadrupole, K1:=KQ1, L=1.0, apertype=CIRCLE, aperture={0.04};
MQ2: Quadrupole, K1:=KQ2, L=1.0, apertype=CIRCLE, aperture={0.04};
MQ3: Quadrupole, K1:=0.0, L=1.0, apertype=CIRCLE, aperture={0.04};
KQ1 = 0.02;
KQ2 = -0.02;
testseq: SEQUENCE, l = 20.0;
MQ1, at = 5;
MQ2, at = 12;
MQ3, at = 18;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
Select, flag=makethin, pattern="MQ1", slice=2;
makethin, sequence=testseq;
use, sequence=testseq;
!---assign misalignments and field errors
select, flag = error, clear;
select, flag = error, pattern = "MQ1";
ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
select, flag = error, clear;
select, flag = error, pattern = "MQ2";
ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
select, flag = error, clear;
select, flag = error, pattern = "MQ3";
ealign, dx = 0.00, dy = 0.00, arex = 0.00, arey = 0.00, dpsi = 0.00;
efcomp, DKN = {0.0, 0.0, 0.001, 0.002}, DKS = {0.0, 0.0, 0.003, 0.004};
select, flag = error, full;
''')
seq = madx.sequence.testseq
# store already applied errors:
madx.command.esave(file='lattice_errors.err')
madx.command.readtable(file='lattice_errors.err', table="errors")
os.remove('lattice_errors.err')
errors = madx.table.errors

pysixtrack_line = pysixtrack.Line.from_madx_sequence(seq, install_apertures=True)
pysixtrack_line.apply_madx_errors(errors)
madx.input('stop;')

expected_element_num = (
2 # start and end marker
+ 6 # drifts (including drift between MQ1 slices)
+ 3 + 2 # quadrupoles + MQ1 slices
+ 3 + 2 # corresponding aperture elements
+ 2*(3+1) # dx/y in/out for MQ1 slices and MQ2
+ 2 # tilt in/out for MQ2
+ 2*3 # arex/y in/out for MQ1 slices
)
assert len(pysixtrack_line) == expected_element_num

expected_element_order = [
pysixtrack.elements.Drift, # start marker
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in of MQ1 1st slice
pysixtrack.elements.Multipole, # MQ1 1st slice
pysixtrack.elements.XYShift, # arex/y in for MQ1 1st slice
pysixtrack.elements.LimitEllipse, # MQ1 1st slice aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 1st slice
pysixtrack.elements.XYShift, # dx/y out for MQ1 1st slice
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ1 marker
pysixtrack.elements.Drift, # MQ1 marker
pysixtrack.elements.XYShift, # arex/y in for MQ1 marker
pysixtrack.elements.LimitEllipse, # MQ1 marker aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 marker
pysixtrack.elements.XYShift, # dx/y out for MQ1 marker
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ1 2nd slice
pysixtrack.elements.Multipole, # MQ1 2nd slice
pysixtrack.elements.XYShift, # arex/y in for MQ1 2nd slice
pysixtrack.elements.LimitEllipse, # MQ1 2nd slice aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 2nd slice
pysixtrack.elements.XYShift, # dx/y out for MQ1 2nd slice
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ2
pysixtrack.elements.SRotation, # tilt in for MQ2
pysixtrack.elements.Multipole, # MQ2
pysixtrack.elements.LimitEllipse, # MQ2 aperture
pysixtrack.elements.SRotation, # tilt out for MQ2
pysixtrack.elements.XYShift, # dx/y out for MQ2
pysixtrack.elements.Drift,
pysixtrack.elements.Multipole, # MQ3
pysixtrack.elements.LimitEllipse, # MQ3 aperture
pysixtrack.elements.Drift,
pysixtrack.elements.Drift # end marker
]
for element, expected_element in zip(pysixtrack_line.elements,
expected_element_order):
assert isinstance(element, expected_element)

idx_MQ3 = pysixtrack_line.element_names.index('mq3')
MQ3 = pysixtrack_line.elements[idx_MQ3]
assert abs(MQ3.knl[2] - 0.001) < 1e-14
assert abs(MQ3.knl[3] - 0.002) < 1e-14
assert abs(MQ3.ksl[2] - 0.003) < 1e-14
assert abs(MQ3.ksl[3] - 0.004) < 1e-14


def test_neutral_errors():
# make sure that some misaligned drifts do not influence particle
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)

from cpymad.madx import Madx

madx = Madx()

madx.input('''
T1: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T2: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T3: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
KQ1 = 0.02;
KQ2 = -0.02;
testseq: SEQUENCE, l = 20.0;
T1, at = 5;
T2, at = 12;
T3, at = 18;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
Select, flag=makethin, pattern="MQ1", slice=2;
makethin, sequence=testseq;
use, sequence=testseq;
!---misalign collimators
select, flag = error, clear;
select, flag = error, pattern = "T1";
ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
select, flag = error, clear;
select, flag = error, pattern = "T2";
ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
select, flag = error, clear;
select, flag = error, pattern = "T3";
ealign, dx = 0.02, dy = 0.01, arex = 0.03, arey = 0.02, dpsi = 0.1;
select, flag = error, full;
''')
seq = madx.sequence.testseq
# store already applied errors:
madx.command.esave(file='lattice_errors.err')
madx.command.readtable(file='lattice_errors.err', table="errors")
os.remove('lattice_errors.err')
errors = madx.table.errors

pysixtrack_line = pysixtrack.Line.from_madx_sequence(seq, install_apertures=True)
pysixtrack_line.apply_madx_errors(errors)
madx.input('stop;')

initial_x = 0.025
initial_y = -0.015

particle = pysixtrack.Particles()
particle.x = initial_x
particle.y = initial_y
particle.state = 1

pysixtrack_line.track(particle)

assert abs(particle.x-initial_x) < 1e-14
assert abs(particle.y-initial_y) < 1e-14

0 comments on commit 4ab545d

Please sign in to comment.