diff --git a/xtrack/slicing.py b/xtrack/slicing.py index 5ed1ebd56..d68a0d757 100644 --- a/xtrack/slicing.py +++ b/xtrack/slicing.py @@ -280,20 +280,10 @@ def _slice_element(self, name, element, chosen_slicing, _edge_markers=True) -> O name=name, ) - if _edge_markers: - - if isinstance(element, xt.Replica): - element = element.resolve(self._line) - _buffer = element._buffer - - entry_marker, exit_marker = f'{name}_entry', f'{name}_exit' - self._line._element_dict[entry_marker] = xt.Marker(_buffer=_buffer) - self._line._element_dict[exit_marker] = xt.Marker(_buffer=_buffer) - slices_to_add = [entry_marker] + slices_to_add + [exit_marker] - - # Handle aperture if isinstance(element, xt.Replica): element = element.resolve(self._line) + + # Handle aperture if (hasattr(element, 'name_associated_aperture') and element.name_associated_aperture is not None): new_slices_to_add = [] @@ -308,7 +298,19 @@ def _slice_element(self, name, element, chosen_slicing, _edge_markers=True) -> O new_slices_to_add += [aper_name] aper_index += 1 new_slices_to_add += [nn] - slices_to_add = new_slices_to_add + # Final aperture + aper_name = f'{name}_aper..{aper_index}' + self._line._element_dict[aper_name] = xt.Replica( + parent_name=element.name_associated_aperture) + + slices_to_add = new_slices_to_add + [aper_name] + + if _edge_markers: + _buffer = element._buffer + entry_marker, exit_marker = f'{name}_entry', f'{name}_exit' + self._line._element_dict[entry_marker] = xt.Marker(_buffer=_buffer) + self._line._element_dict[exit_marker] = xt.Marker(_buffer=_buffer) + slices_to_add = [entry_marker] + slices_to_add + [exit_marker] return slices_to_add