Add test flags

mbtrack2/tracking/beam_ion_effects.py

@@ -548,6 +548,9 @@ class BeamIonElement(Element):
         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 = []
         self.monitors = []
@@ -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 = (
@@ -789,6 +792,10 @@ 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)