class Segment(object):
"""
Simulate one fixed segment, i.e. sequence + range.
:ivar Madx madx:
:ivar list elements:
:ivar dict twiss_args:
"""
_columns = [
'name', 'l', 'angle', 'k1l',
's',
'x', 'y',
'betx','bety',
'alfx', 'alfy',
]
# TODO: extend list of merge-columns
_mixin_columns = [
'x', 'y',
'betx','bety',
'alfx', 'alfy',
]
def __init__(self, session, sequence, range, beam, twiss_args,
show_element_indicators):
"""
:param Session session:
:param str sequence:
:param tuple range:
"""
self.hook = HookCollection(
update=None,
remove=None,
show_element_indicators=None,
)
self.session = session
self.sequence = session.madx.sequences[sequence]
self.start, self.stop = self.parse_range(range)
self.range = (normalize_range_name(self.start.name),
normalize_range_name(self.stop.name))
self._beam = beam
self._twiss_args = twiss_args
self._show_element_indicators = show_element_indicators
self._use_beam(beam)
raw_elements = self.sequence.elements
# TODO: provide uncached version of elements with units:
self.elements = list(map(
session.utool.dict_add_unit, raw_elements))
# TODO: self.hook.create(self)
self.twiss()
@property
def madx(self):
return self.session.madx
@property
def utool(self):
return self.session.utool
@property
def data(self):
return {
'sequence': self.sequence.name,
'range': self.range,
'beam': self.beam,
'twiss': self.twiss_args,
}
def get_element_info(self, element):
"""Get :class:`ElementInfo` from element name or index."""
if isinstance(element, ElementInfo):
return element
if isinstance(element, (basestring, dict)):
element = self.sequence.elements.index(element)
element_data = self.session.utool.dict_add_unit(
self.sequence.elements[element])
if element < 0:
element += len(self.sequence.elements)
return ElementInfo(element_data['name'], element, element_data['at'])
def parse_range(self, range):
"""Convert a range str/tuple to a tuple of :class:`ElementInfo`."""
if isinstance(range, basestring):
range = range.split('/')
start_name, stop_name = range
return (self.get_element_info(start_name),
self.get_element_info(stop_name))
def destroy(self):
self.session.segment = None
self.hook.remove()
@property
def show_element_indicators(self):
return self._show_element_indicators
@show_element_indicators.setter
def show_element_indicators(self, show):
if show == self._show_element_indicators:
return
self._show_element_indicators = show
self.hook.show_element_indicators()
@property
def twiss_args(self):
return self._twiss_args
@twiss_args.setter
def twiss_args(self, twiss_args):
self._twiss_args = twiss_args
self.twiss()
@property
def beam(self):
"""Get the beam parameter dictionary."""
return self._beam
@beam.setter
def beam(self, beam):
"""Set beam from a parameter dictionary."""
self._beam = beam
self._use_beam(beam)
self.twiss()
def _use_beam(self, beam):
beam = self.utool.dict_strip_unit(beam)
beam = dict(beam, sequence=self.sequence.name)
self.madx.command.beam(**beam)
def element_by_position(self, pos):
"""Find optics element by longitudinal position."""
if pos is None:
return None
for elem in self.elements:
at, L = elem['at'], elem['l']
if pos >= at and pos <= at+L:
return elem
return None
def get_element_index(self, elem):
"""Get element index by it name."""
return self.sequence.elements.index(elem)
def get_twiss(self, elem, name):
"""Return beam envelope at element."""
element = self.get_element_info(elem)
if not self.contains(element):
return None
return self.tw[name][element.index - self.start.index]
def contains(self, element):
return (self.start.index <= element.index and
self.stop.index >= element.index)
def twiss(self):
"""Recalculate TWISS parameters."""
results = self.raw_twiss()
# Update TWISS results
self.tw = self.utool.dict_add_unit(results)
self.summary = self.utool.dict_add_unit(results.summary)
# data post processing
self.tw['s'] += self.start.at
self.pos = self.tw['s']
self.tw['envx'] = (self.tw['betx'] * self.summary['ex'])**0.5
self.tw['envy'] = (self.tw['bety'] * self.summary['ey'])**0.5
# Create aliases for x,y that have non-empty common prefix. The goal
# is to make the config file entries less awkward that hold this
# prefix:
self.tw['posx'] = self.tw['x']
self.tw['posy'] = self.tw['y']
self.hook.update()
def _get_twiss_args(self, **kwargs):
twiss_init = self.utool.dict_strip_unit(self.twiss_args)
twiss_args = {
'sequence': self.sequence.name,
'range': self.range,
'columns': self._columns,
'twiss_init': twiss_init,
}
twiss_args.update(kwargs)
return twiss_args
def raw_twiss(self, **kwargs):
return self.madx.twiss(**self._get_twiss_args(**kwargs))
def get_transfer_map(self, beg_elem, end_elem):
"""
Get the transfer matrix R(i,j) between the two elements.
This requires a full twiss call, so don't do it too often.
"""
info = self.get_element_info
twiss_args = self._get_twiss_args()
twiss_args['range_'] = (info(beg_elem).name, info(end_elem).name)
twiss_args['tw_range'] = twiss_args.pop('range')
return self.madx.get_transfer_map_7d(**twiss_args)
