diff --git a/mbtrack2/tracking/beam_ion_effects.py b/mbtrack2/tracking/beam_ion_effects.py index cb566c10e1f069422856c743425a6feb4ce07f7c..c9633c330f0ef4f7c600cf18967a3182fb6d82c3 100644 --- a/mbtrack2/tracking/beam_ion_effects.py +++ b/mbtrack2/tracking/beam_ion_effects.py @@ -464,13 +464,13 @@ class BeamIonElement(Element): @wraps(track) def track_wrapper(*args, **kwargs): if isinstance(args[1], Beam): - if Beam.switch_mpi: + self = args[0] + beam = args[1] + if beam.mpi_switch: warnings.warn( 'Tracking through beam-ion element is performed sequentially. Bunches are not parallelized.', UserWarning, stacklevel=2) - self = args[0] - beam = args[1] for bunch in beam.bunch_list: track(self, bunch, *args[2:], **kwargs) else: @@ -608,9 +608,9 @@ class BeamIonElement(Element): The new y momentum of the first beam. """ - en_x, en_y = BeamIonElement._get_efields(first_beam, - second_beam, - field_model=field_model) + en_x, en_y = self._get_efields(first_beam, + second_beam, + field_model=field_model) kicks_x = prefactor * en_x kicks_y = prefactor * en_y new_xp = first_beam["xp"] + kicks_x