Skip to content

Commit

Permalink
Merge pull request #60 from roman-martin/mad_loader_errors
Browse files Browse the repository at this point in the history 
loader of MAD-X errors
  • Loading branch information
@rdemaria
rdemaria authored Oct 29, 2020
2 parents 4ab545d + 82b94d0 commit 8e1d432
Show file tree
Hide file tree
Showing 3 changed files with 232 additions and 107 deletions.
145 changes: 62 additions & 83 deletions pysixtrack/line.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,7 @@ def from_madx_sequence(
exact_drift=False,
drift_threshold=1e-6,
install_apertures=False,
apply_madx_errors=False,
):

line = cls(elements=[], element_names=[])
Expand All @@ -342,6 +343,9 @@ def from_madx_sequence(
):
line.append_element(el, el_name)

if apply_madx_errors:
line._apply_madx_errors(sequence)

return line

# error handling (alignment, multipole orders, ...):
Expand All @@ -366,7 +370,7 @@ def find_element_ids(self, element_name):
idx_after_el = idx_el + 1
return idx_el, idx_after_el

def add_offset_error_to(self, element_name, dx=0, dy=0):
def _add_offset_error_to(self, element_name, dx=0, dy=0):
idx_el, idx_after_el = self.find_element_ids(element_name)
xyshift = elements.XYShift(dx=dx, dy=dy)
inv_xyshift = elements.XYShift(dx=-dx, dy=-dy)
Expand All @@ -375,7 +379,7 @@ def add_offset_error_to(self, element_name, dx=0, dy=0):
idx_after_el + 1, inv_xyshift, element_name + "_offset_out"
)

def add_aperture_offset_error_to(self, element_name, arex=0, arey=0):
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":
Expand All @@ -389,7 +393,7 @@ def add_aperture_offset_error_to(self, element_name, arex=0, arey=0):
idx_after_el + 1, inv_xyshift, element_name + "_aperture_offset_out"
)

def add_tilt_error_to(self, element_name, angle):
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`
gets wrapped by SRotation elements with rotation angle
Expand Down Expand Up @@ -418,7 +422,7 @@ def add_tilt_error_to(self, element_name, angle):
self.insert_element(idx_el, srot, element_name + "_tilt_in")
self.insert_element(idx_after_el + 1, inv_srot, element_name + "_tilt_out")

def add_multipole_error_to(self, element_name, knl=[], ksl=[]):
def _add_multipole_error_to(self, element_name, knl=[], ksl=[]):
# will raise error if element not present:
assert element_name in self.element_names
element = self.elements[self.element_names.index(element_name)]
Expand All @@ -435,97 +439,72 @@ def add_multipole_error_to(self, element_name, knl=[], ksl=[]):
for i, component in enumerate(ksl):
element.ksl[i] += component

def apply_madx_errors(self, error_table):
"""Applies MAD-X error_table (with multipole errors,
dx and dy offset errors and dpsi tilt errors)
to existing elements in this Line instance.
def _apply_madx_errors(self, madx_sequence):
"""Applies errors from MAD-X sequence to existing
elements in this Line instance.
Return error_table names which were not found in the
elements of this Line instance (and thus not treated).
Return names of MAD-X elements with existing align_errors
or field_errors which were not found in the elements of
this Line instance (and thus not treated).
Example via cpymad:
madx = cpymad.madx.Madx()
# (...set up lattice and errors in cpymad...)
seq = madx.sequence.some_lattice
# store already applied errors:
madx.command.esave(file='lattice_errors.err')
madx.command.readtable(
file='lattice_errors.err', table="errors")
errors = madx.table.errors
pysixtrack_line = Line.from_madx_sequence(seq)
pysixtrack_line.apply_madx_errors(errors)
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
apply_madx_errors=True
)
"""
max_multipole_err = 0
# check for errors in table which cannot be treated yet:
for error_type in error_table.keys():
if error_type == "name":
continue
if any(error_table[error_type]):
if error_type in ["dx", "dy", "dpsi", "arex", "arey"]:
# available alignment error
continue
elif error_type[:1] == "k" and error_type[-1:] == "l":
# available multipole error
order = int("".join(c for c in error_type if c.isdigit()))
max_multipole_err = max(max_multipole_err, order)
else:
print(
f'Warning: MAD-X error type "{error_type}"'
" not implemented yet."
)

elements_not_found = []
for i_line, element_name in enumerate(error_table["name"]):
for element, element_name in zip(
madx_sequence.expanded_elements,
madx_sequence.expanded_element_names()
):
if element_name not in self.element_names:
elements_not_found.append(element_name)
continue
if element.align_errors or element.field_errors:
elements_not_found.append(element_name)
continue

# add offset
try:
dx = error_table["dx"][i_line]
except KeyError:
dx = 0
try:
dy = error_table["dy"][i_line]
except KeyError:
dy = 0
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:
arey = 0
if arex or arey:
self.add_aperture_offset_error_to(element_name, arex, arey)

# add multipole error
knl = [
error_table[f"k{o}l"][i_line]
for o in range(max_multipole_err + 1)
]
ksl = [
error_table[f"k{o}sl"][i_line]
for o in range(max_multipole_err + 1)
]
if any(knl) or any(ksl):
self.add_multipole_error_to(element_name, knl, ksl)
if element.align_errors:
# add offset
dx = element.align_errors.dx
dy = element.align_errors.dy
if dx or dy:
self._add_offset_error_to(element_name, dx, dy)

# add tilt
dpsi = element.align_errors.dpsi
if dpsi:
self._add_tilt_error_to(element_name, angle=dpsi / deg2rad)

# add aperture-only offset
arex = element.align_errors.arex
arey = element.align_errors.arey
if arex or arey:
self._add_aperture_offset_error_to(element_name, arex, arey)

# check for errors which cannot be treated yet:
for error_type in dir(element.align_errors):
if not error_type[0] == '_' and \
error_type not in ['dx', 'dy', 'dpsi', 'arex',
'arey', 'count', 'index']:
print(
f'Warning: MAD-X error type "{error_type}"'
" not implemented yet."
)

if element.field_errors:
# add multipole error
if any(element.field_errors.dkn) or \
any(element.field_errors.dks):
knl = element.field_errors.dkn
ksl = element.field_errors.dks
knl = knl[:np.amax(np.where(knl)) + 1] # delete trailing zeros
ksl = ksl[:np.amax(np.where(ksl)) + 1] # to keep order low
self._add_multipole_error_to(element_name, knl, ksl)

return elements_not_found

Expand Down
12 changes: 6 additions & 6 deletions tests/test_line.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,24 +34,24 @@ def test_line():
n_elements += 1
assert len(line) == n_elements

line.add_offset_error_to(multipole_name, dx=0, dy=0)
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)
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)
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)
line._add_tilt_error_to(multipole_name, angle=0.1)
n_elements += 2
assert len(line) == n_elements

line.add_multipole_error_to(multipole_name, knl=[0, 0.1], ksl=[-0.03, 0.01])
# line.add_multipole_error_to(drift_exact,knl=[0,0.1],ksl=[-0.03,0.01])
line._add_multipole_error_to(multipole_name, knl=[0, 0.1], ksl=[-0.03, 0.01])
# line._add_multipole_error_to(drift_exact,knl=[0,0.1],ksl=[-0.03,0.01])

line_dict = line.to_dict()
line = pysixtrack.Line.from_dict(line_dict)
Expand Down
Loading

0 comments on commit 8e1d432

Please sign in to comment.