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)