diff --git a/mbtrack2/tracking/beam_ion_effects.py b/mbtrack2/tracking/beam_ion_effects.py index 5e49ca50bad420935cb9df601373f1ca51d77323..db05c8d95bf195e52a72e99fc41e4fd5fd224fb8 100644 --- a/mbtrack2/tracking/beam_ion_effects.py +++ b/mbtrack2/tracking/beam_ion_effects.py @@ -547,6 +547,9 @@ class BeamIonElement(Element): self.generate_ions = generate_ions self.beam_ion_interaction = beam_ion_interaction self.ion_drift = ion_drift + + self.init_asbefore = True + self.switch_new = True # interfaces for apertures and montiors self.apertures = [] @@ -641,7 +644,7 @@ class BeamIonElement(Element): raise ValueError( f"The implementation for required beam-ion interaction model \ {field_model} is not implemented") - if isinstance(second_beam, IonParticles): + if isinstance(second_beam, IonParticles) and self.switch_new: sb_mx, sb_my, sb_stdx, sb_stdy = second_beam.mean_std_xy else: sb_mx, sb_stdx = ( @@ -664,7 +667,7 @@ class BeamIonElement(Element): ) elif field_model == "strong": - if isinstance(first_beam, IonParticles): + if isinstance(first_beam, IonParticles) and self.switch_new: fb_mx, fb_my = first_beam.mean_xy else: fb_mx, fb_my = ( @@ -788,7 +791,11 @@ class BeamIonElement(Element): if not empty_bucket and self.generate_ions: self.generate_new_ions(electron_bunch=electron_bunch) - + + if self.init_asbefore: + self.init_asbefore = False + self.ion_beam["charge"][0] = self.ion_beam["charge"][1] + for aperture in self.apertures: aperture.track(self.ion_beam)