Merge branch 'feature-ions_update' into 'develop'

Update BeamIonElement

See merge request !33

mbtrack2/tracking/beam_ion_effects.py

@@ -6,7 +6,6 @@ IonMonitor
 IonAperture
 IonParticles
 """
-import warnings
 from abc import ABCMeta
 from functools import wraps
 from itertools import count
@@ -36,14 +35,18 @@ class IonMonitor(Monitor, metaclass=ABCMeta):
     total_size : int
         The total number of steps to be simulated.
     file_name : str, optional
-        The name of the HDF5 file to store the data. If not provided, a new file will be created. Defaults to None.
+        The name of the HDF5 file to store the data. If not provided, a new 
+        file will be created. 
+        Defaults to None.
 
     Methods
     -------
-    monitor_init(group_name, save_every, buffer_size, total_size, dict_buffer, dict_file, file_name=None, dict_dtype=None)
+    monitor_init(group_name, save_every, buffer_size, total_size, dict_buffer, 
+                 dict_file, file_name=None, dict_dtype=None)
         Initialize the monitor object.
     track(bunch)
         Tracking method for the element.
+
     Raises
     ------
     ValueError
@@ -56,16 +59,16 @@ class IonMonitor(Monitor, metaclass=ABCMeta):
     def __init__(self, save_every, buffer_size, total_size, file_name=None):
         group_name = "IonData_" + str(next(self._n_monitors))
         dict_buffer = {
-            "mean": (6, buffer_size),
-            "std": (6, buffer_size),
+            "mean": (4, buffer_size),
+            "std": (4, buffer_size),
             "charge": (buffer_size, ),
-            "charge_per_mp": (buffer_size, ),
+            "mp_number": (buffer_size, ),
         }
         dict_file = {
-            "mean": (6, total_size),
-            "std": (6, total_size),
+            "mean": (4, total_size),
+            "std": (4, total_size),
             "charge": (total_size, ),
-            "charge_per_mp": (buffer_size, ),
+            "mp_number": (buffer_size, ),
         }
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name)
@@ -96,13 +99,18 @@ class IonMonitor(Monitor, metaclass=ABCMeta):
         total_size : int
             The total number of steps to be simulated.
         dict_buffer : dict
-            A dictionary containing the names and sizes of the attribute buffers.
+            A dictionary containing the names and sizes of the attribute 
+            buffers.
         dict_file : dict
-            A dictionary containing the names and shapes of the datasets to be created.
+            A dictionary containing the names and shapes of the datasets to be 
+            created.
         file_name : str, optional
-            The name of the HDF5 file to store the data. If not provided, a new file will be created. Defaults to None.
+            The name of the HDF5 file to store the data. If not provided, a new 
+            file will be created. 
+            Defaults to None.
         dict_dtype : dict, optional
-            A dictionary containing the names and data types of the datasets. Defaults to None.
+            A dictionary containing the names and data types of the datasets. 
+            Defaults to None.
 
         Raises
         ------
@@ -157,8 +165,10 @@ class IonAperture(ElipticalAperture):
     """
     Class representing an ion aperture.
 
-    Inherits from ElipticalAperture. Unlike in ElipticalAperture, ions are removed from IonParticles instead of just being flagged as not "alive".
-    For beam-ion simulations there are too many lost particles and it is better to remove them.
+    Inherits from ElipticalAperture. Unlike in ElipticalAperture, ions are 
+    removed from IonParticles instead of just being flagged as not "alive".
+    For beam-ion simulations there are too many lost particles and it is better 
+    to remove them.
 
     Attributes
     ----------
@@ -208,15 +218,33 @@ class IonParticles(Bunch):
     ring : Synchrotron class object
         The ring object representing the accelerator ring.
     track_alive : bool, optional
-        Flag indicating whether to track the alive particles. Default is False.
+        Flag indicating whether to track the alive particles. 
+        Default is False.
     alive : bool, optional
-        Flag indicating whether the particles are alive. Default is True.
+        Flag indicating whether the particles are alive. 
+        Default is True.
+        
+    Attributes
+    ----------
+    charge : float
+        Bunch charge in [C].
+    mean : array of shape (4,)
+        Charge-weighted mean values for x, xp, y and yp coordinates.
+    mean_xy : tuple of 2 float
+        Charge-weighted mean values for x and y coordinates.
+    std : array of shape (4,)
+        Charge-weighted std values for x, xp, y and yp coordinates.
+    mean_std_xy : tuple of 4 float
+        Charge-weighted mean and std values for x and y coordinates.
+        
     Methods:
     --------
-    generate_as_a_distribution(electron_bunch)
-        Generates the particle positions based on a normal distribution, taking distribution parameters from an electron bunch.
-    generate_from_random_samples(electron_bunch)
-        Generates the particle positions and times based on random samples from electron positions.
+    generate_as_a_distribution(electron_bunch, charge)
+        Generates the particle positions based on a normal distribution, taking 
+        distribution parameters from an electron bunch.
+    generate_from_random_samples(electron_bunch, charge)
+        Generates the particle positions and times based on random samples from 
+        electron positions.
     """
 
     def __init__(self,
@@ -242,6 +270,7 @@ class IonParticles(Bunch):
             "yp": np.zeros((mp_number, ), dtype=np.float64),
             "tau": np.zeros((mp_number, ), dtype=np.float64),
             "delta": np.zeros((mp_number, ), dtype=np.float64),
+            "charge": np.zeros((mp_number, ), dtype=np.float64),
         }
         self.charge_per_mp = 0
 
@@ -254,14 +283,73 @@ class IonParticles(Bunch):
     def mp_number(self, value):
         self._mp_number = int(value)
 
-    def generate_as_a_distribution(self, electron_bunch):
+    @property
+    def charge(self):
+        """Bunch charge in [C]"""
+        return self["charge"].sum()
+
+    def _mean_weighted(self, coord):
+        if self.charge == 0:
+            return np.zeros_like(coord, dtype=float)
+        else:
+            mean = [[np.average(self[name], weights=self["charge"])]
+                    for name in coord]
+            return np.squeeze(np.array(mean))
+
+    def _mean_std_weighted(self, coord):
+        if self.charge == 0:
+            return np.zeros_like(coord,
+                                 dtype=float), np.zeros_like(coord,
+                                                             dtype=float)
+        else:
+            mean = [[np.average(self[name], weights=self["charge"])]
+                    for name in coord]
+            var = []
+            for i, name in enumerate(coord):
+                var.append(
+                    np.average((self[name] - mean[i])**2,
+                               weights=self["charge"]))
+            var = np.array(var)
+            return np.squeeze(np.array(mean)), np.squeeze(np.sqrt(var))
+
+    @property
+    def mean(self):
+        """Return charge-weighted mean values for x, xp, y and yp coordinates."""
+        coord = ["x", "xp", "y", "yp"]
+        return self._mean_weighted(coord)
+
+    @property
+    def mean_xy(self):
+        """Return charge-weighted mean values for x and y coordinates."""
+        coord = ["x", "y"]
+        mean = self._mean_weighted(coord)
+        return mean[0], mean[1]
+
+    @property
+    def std(self):
+        """Return charge-weighted std values for x, xp, y and yp coordinates."""
+        coord = ["x", "xp", "y", "yp"]
+        _, std = self._mean_std_weighted(coord)
+        return std
+
+    @property
+    def mean_std_xy(self):
+        """Return charge-weighted mean and std values for x and y coordinates."""
+        coord = ["x", "y"]
+        mean, std = self._mean_std_weighted(coord)
+        return mean[0], mean[1], std[0], std[1]
+
+    def generate_as_a_distribution(self, electron_bunch, charge):
         """
-        Generates the particle positions based on a normal distribution, taking distribution parameters from an electron bunch.
+        Generates the particle positions based on a normal distribution, taking 
+        distribution parameters from an electron bunch.
 
         Parameters:
         ----------
         electron_bunch : Bunch
             An instance of the Bunch class representing the electron bunch.
+        charge : float
+            Total ion charge generated in [C].
         """
         if electron_bunch.is_empty:
             raise ValueError("Electron bunch is empty.")
@@ -289,15 +377,19 @@ class IonParticles(Bunch):
             high=self.ion_element_length / c,
             size=self.mp_number,
         )
+        self["charge"] = np.ones((self.mp_number, )) * charge / self.mp_number
 
-    def generate_from_random_samples(self, electron_bunch):
+    def generate_from_random_samples(self, electron_bunch, charge):
         """
-        Generates the particle positions and times based on random samples from electron positions in the bunch.
+        Generates the particle positions and times based on random samples from 
+        electron positions in the bunch.
 
         Parameters:
         ----------
         electron_bunch : Bunch
             An instance of the Bunch class representing the electron bunch.
+        charge : float
+            Total ion charge generated in [C].
         """
         if electron_bunch.is_empty:
             raise ValueError("Electron bunch is empty.")
@@ -316,6 +408,7 @@ class IonParticles(Bunch):
             high=self.ion_element_length / c,
             size=self.mp_number,
         )
+        self["charge"] = np.ones((self.mp_number, )) * charge / self.mp_number
 
     def __add__(self, new_particles):
         self.mp_number += new_particles.mp_number
@@ -331,46 +424,70 @@ class BeamIonElement(Element):
     """
     Represents an element for simulating beam-ion interactions.
     
+    Apertures and monitors for the ion beam (instances of IonAperture and 
+    IonMonitor) can be added to tracking after the BeamIonElement object has 
+    been initialized by using:
+        beam_ion.apertures.append(aperture)
+        beam_ion.monitors.append(monitor)
+        
+    If the a IonMonitor object is used to record the ion beam, at the end of 
+    tracking the user should close the monitor, for example by calling:
+        beam_ion.monitors[0].close()
+
     Parameters
     ----------
     ion_mass : float
-        The mass of the ions in kg.
+        The mass of the ions in [kg].
     ion_charge : float
-        The charge of the ions in Coulomb.
+        The charge of the ions in [C].
     ionization_cross_section : float
-        The cross section of ionization in meters^2.
+        The cross section of ionization in [m^2].
     residual_gas_density : float
-        The residual gas density in meters^-3.
+        The residual gas density in [m^-3].
     ring : instance of Synchrotron()
         The ring.
-    ion_field_model : str
-        The ion field model, the options are 'weak' (acts on each macroparticle), 'strong' (acts on c.m.), 'PIC'.
+    ion_field_model : {'weak', 'strong', 'PIC'}
+        The ion field model used to update electron beam coordinates:
+            - 'weak': ion field acts on each macroparticle.
+            - 'strong': ion field acts on electron bunch c.m. 
+            - 'PIC': a PIC solver is used to get the ion electric field and the
+            result is interpolated on the electron bunch coordinates.
+        For both 'weak' and 'strong' models, the electric field is computed 
+        using the Bassetti-Erskine formula [1], so assuming a Gaussian beam 
+        distribution.
         For 'PIC' the PyPIC package is required.
-    electron_field_model : str
-        The electron field model, the options are 'weak', 'strong', 'PIC'.
+    electron_field_model : {'weak', 'strong', 'PIC'}
+        The electron field model, defined in the same way as ion_field_model.
     ion_element_length : float
-        The length of the beam-ion interaction region. For example, if only a single interaction point is used this should be equal to ring.L. 
-    x_radius : float
-        The x radius of the aperture.
-    y_radius : float
-        The y radius of the aperture.
-    n_steps : int
-        The number of records in the built-in ion beam monitor. Should be number of turns times number of bunches because the monitor records every turn after each bunch passage.
+        The length of the beam-ion interaction region. For example, if only a 
+        single interaction point is used this should be equal to ring.L. 
     n_ion_macroparticles_per_bunch : int, optional
-        The number of ion macroparticles generated per electron bunch passed. Defaults to 30.
-    ion_beam_monitor_name : str, optional
-         If provided, the name of the ion monitor output file. It must end with an extension '.hdf5'.
-         If None, no ion monitor file is generated.
-    use_ion_phase_space_monitor : bool, optional
-        Whether to use the ion phase space monitor.
-    generate_method : str, optional
-        The method to generate the ion macroparticles, the options are 'distribution', 'samples'. Defaults to 'distribution'. 
-        'distribution' generates a distribution statistically equivalent to the distribution of electrons. 
-        'samples' generates ions from random samples of electron positions.
+        The number of ion macroparticles generated per electron bunch passage. 
+        Defaults to 30.
+    generate_method : {'distribution', 'samples'}, optional
+        The method to generate the ion macroparticles:
+            - 'distribution' generates a distribution statistically equivalent 
+            to the distribution of electrons. 
+            - 'samples' generates ions from random samples of electron 
+            positions.
+        Defaults to 'distribution'.
+    generate_ions : bool, optional
+        If True, generate ions during BeamIonElement.track calls.
+        Default is True.
+    beam_ion_interaction : bool, optional
+        If True, update both beam and ion beam coordinate due to beam-ion 
+        interaction during BeamIonElement.track calls.
+        Default is True.
+    ion_drift : bool, optional
+        If True, update ion beam coordinate due to drift time between bunches.
+        Default is True.
 
     Methods
     -------
-    __init__(ion_mass, ion_charge, ionization_cross_section, residual_gas_density, ring, ion_field_model, electron_field_model, ion_element_length, n_steps, x_radius, y_radius, ion_beam_monitor_name=None, use_ion_phase_space_monitor=False, n_ion_macroparticles_per_bunch=30, generate_method='distribution')
+    __init__(ion_mass, ion_charge, ionization_cross_section, 
+             residual_gas_density, ring, ion_field_model, electron_field_model, 
+             ion_element_length, n_ion_macroparticles_per_bunch=30, 
+             generate_method='distribution')
         Initializes the BeamIonElement object.
     parallel(track)
         Defines the decorator @parallel to handle tracking of Beam() objects.
@@ -386,9 +503,13 @@ class BeamIonElement(Element):
     Raises
     ------
     UserWarning
-        If the BeamIonMonitor object is used, the user should call the close() method at the end of tracking.
-    NotImplementedError
-        If the ion phase space monitor is used.
+        If the BeamIonMonitor object is used, the user should call the close() 
+        method at the end of tracking.
+        
+    References
+    ----------
+    [1] : Bassetti, M., & Erskine, G. A. (1980). Closed expression for the 
+    electrical field of a two-dimensional Gaussian charge (No. ISR-TH-80-06).
     """
 
     def __init__(self,
@@ -400,16 +521,11 @@ class BeamIonElement(Element):
                  ion_field_model,
                  electron_field_model,
                  ion_element_length,
-                 n_steps,
-                 x_radius,
-                 y_radius,
-                 ion_beam_monitor_name=None,
-                 use_ion_phase_space_monitor=False,
                  n_ion_macroparticles_per_bunch=30,
-                 generate_method='distribution'):
-        if use_ion_phase_space_monitor:
-            raise NotImplementedError(
-                "Ion phase space monitor is not implemented.")
+                 generate_method='distribution',
+                 generate_ions=True,
+                 beam_ion_interaction=True,
+                 ion_drift=True):
         self.ring = ring
         self.bunch_spacing = ring.L / ring.h
         self.ion_mass = ion_mass
@@ -423,30 +539,18 @@ class BeamIonElement(Element):
         if not self.generate_method in ["distribution", "samples"]:
             raise ValueError("Wrong generate_method.")
         self.n_ion_macroparticles_per_bunch = n_ion_macroparticles_per_bunch
-        self.ion_beam_monitor_name = ion_beam_monitor_name
         self.ion_beam = IonParticles(
             mp_number=1,
             ion_element_length=self.ion_element_length,
             ring=self.ring)
-        self.ion_beam["x"] = 0
-        self.ion_beam["xp"] = 0
-        self.ion_beam["y"] = 0
-        self.ion_beam["yp"] = 0
-        self.ion_beam["tau"] = 0
-        self.ion_beam["delta"] = 0
-
-        if self.ion_beam_monitor_name:
-            warnings.warn(
-                'BeamIonMonitor.beam_monitor.close() should be called at the end of tracking',
-                UserWarning,
-                stacklevel=2)
-            self.beam_monitor = IonMonitor(
-                1,
-                int(n_steps / 10),
-                n_steps,
-                file_name=self.ion_beam_monitor_name)
-
-        self.aperture = IonAperture(X_radius=x_radius, Y_radius=y_radius)
+
+        self.generate_ions = generate_ions
+        self.beam_ion_interaction = beam_ion_interaction
+        self.ion_drift = ion_drift
+
+        # interfaces for apertures and montiors
+        self.apertures = []
+        self.monitors = []
 
     def parallel(track):
         """
@@ -514,16 +618,17 @@ class BeamIonElement(Element):
 
     def _get_efields(self, first_beam, second_beam, field_model):
         """
-        Calculates the electromagnetic field of the first beam acting on the second beam for a given field model.
+        Calculates the electromagnetic field of the second beam acting on the 
+        first beam for a given field model.
     
         Parameters
         ----------
         first_beam : IonParticles or Bunch
-            The first beam, represented as an instance of IonParticles() or Bunch().
+            The first beam, which is being acted on.
         second_beam : IonParticles or Bunch
-            The second beam, represented as an instance of IonParticles() or Bunch().
-        field_model : str, optional
-            The field model used for the interaction. Options are 'weak', 'strong', or 'PIC'.
+            The second beam, which is generating an electric field.
+        field_model : {'weak', 'strong', 'PIC'}
+            The field model used for the interaction.
     
         Returns
         -------
@@ -534,8 +639,11 @@ class BeamIonElement(Element):
         """
         if not field_model in ["weak", "strong", "PIC"]:
             raise ValueError(
-                f"The implementation for required beam-ion interaction model {field_model} is not implemented"
-            )
+                f"The implementation for required beam-ion interaction model \
+                {field_model} is not implemented")
+        if isinstance(second_beam, IonParticles):
+            sb_mx, sb_my, sb_stdx, sb_stdy = second_beam.mean_std_xy
+        else:
             sb_mx, sb_stdx = (
                 second_beam["x"].mean(),
                 second_beam["x"].std(),
@@ -556,6 +664,9 @@ class BeamIonElement(Element):
             )
 
         elif field_model == "strong":
+            if isinstance(first_beam, IonParticles):
+                fb_mx, fb_my = first_beam.mean_xy
+            else:
                 fb_mx, fb_my = (
                     first_beam["x"].mean(),
                     first_beam["y"].mean(),
@@ -594,18 +705,19 @@ class BeamIonElement(Element):
                                prefactor,
                                field_model="strong"):
         """
-        Calculates the new momentum of the first beam due to the interaction with the second beam.
+        Calculates the new momentum of the first beam due to the interaction 
+        with the second beam.
         
         Parameters
         ----------
         first_beam : IonParticles or Bunch
-            The first beam, represented as an instance of IonParticles() or Bunch().
+            The first beam, which is being acted on.
         second_beam : IonParticles or Bunch
-            The second beam, represented as an instance of IonParticles() or Bunch().
+            The second beam, which is generating an electric field.
         prefactor : float
             A scaling factor applied to the calculation of the new momentum.
-        field_model : str
-            The field model used for the interaction. Options are 'weak', 'strong', or 'PIC'.
+        field_model : {'weak', 'strong', 'PIC'}
+            The field model used for the interaction.
             Default is "strong".
         
         Returns
@@ -635,7 +747,7 @@ class BeamIonElement(Element):
 
         Parameters
         ----------
-        electron_bunch : ElectronBunch
+        electron_bunch : Bunch
             The electron bunch used to generate new ions.
 
         Returns
@@ -647,18 +759,16 @@ class BeamIonElement(Element):
             ion_element_length=self.ion_element_length,
             ring=self.ring,
         )
+        new_ion_charge = (electron_bunch.charge *
+                          self.ionization_cross_section *
+                          self.residual_gas_density * self.ion_element_length)
         if self.generate_method == 'distribution':
             new_ion_particles.generate_as_a_distribution(
-                electron_bunch=electron_bunch)
+                electron_bunch=electron_bunch, charge=new_ion_charge)
         elif self.generate_method == 'samples':
             new_ion_particles.generate_from_random_samples(
-                electron_bunch=electron_bunch)
+                electron_bunch=electron_bunch, charge=new_ion_charge)
         self.ion_beam += new_ion_particles
-        self.ion_beam.charge_per_mp = (electron_bunch.charge *
-                                       self.ionization_cross_section *
-                                       self.residual_gas_density *
-                                       self.ion_element_length /
-                                       self.n_ion_macroparticles_per_bunch)
 
     @parallel
     def track(self, electron_bunch):
@@ -676,15 +786,16 @@ class BeamIonElement(Element):
         else:
             empty_bucket = False
 
-        if not empty_bucket:
+        if not empty_bucket and self.generate_ions:
             self.generate_new_ions(electron_bunch=electron_bunch)
 
-        self.aperture.track(self.ion_beam)
+        for aperture in self.apertures:
+            aperture.track(self.ion_beam)
 
-        if self.ion_beam_monitor_name is not None:
-            self.beam_monitor.track(self.ion_beam)
+        for monitor in self.monitors:
+            monitor.track(self.ion_beam)
 
-        if not empty_bucket:
+        if not empty_bucket and self.beam_ion_interaction:
             prefactor_to_ion_field = -self.ion_beam.charge / (self.ring.E0)
             prefactor_to_electron_field = -electron_bunch.charge * (
                 e / (self.ion_mass * c**2))
@@ -704,4 +815,5 @@ class BeamIonElement(Element):
             self._update_beam_momentum(electron_bunch, new_xp_electrons,
                                        new_yp_electrons)
 
+        if self.ion_drift:
             self.track_ions_in_a_drift(drift_length=self.bunch_spacing)

tests/unit/tracking/test_beam_ion_effects.py

@@ -25,8 +25,10 @@ class TestIonParticles:
 
     # Generate particle distribution using electron bunch parameters via generate_as_a_distribution()
     def test_generate_as_distribution(self, generate_ion_particles, large_bunch):
-        ions = generate_ion_particles(mp_number=1e5)
-        ions.generate_as_a_distribution(large_bunch)
+        mp_number = 1e5
+        charge = 1e-9
+        ions = generate_ion_particles(mp_number)
+        ions.generate_as_a_distribution(large_bunch, charge)
     
         assert np.isclose(ions["x"].mean(), large_bunch["x"].mean(), rtol=0.1, atol=1e-5)
         assert np.isclose(ions["x"].std(), large_bunch["x"].std(), rtol=0.1, atol=1e-5)
@@ -37,11 +39,15 @@ class TestIonParticles:
         assert np.all(ions["delta"] == 0)
         assert np.all(ions["tau"] >= -ions.ion_element_length)
         assert np.all(ions["tau"] <= ions.ion_element_length)
+        assert np.isclose(ions.charge, charge)
+        assert np.isclose(ions["charge"].mean(), charge/mp_number, rtol=0.01, atol=1e-9)
 
     # Generate random particle samples from electron bunch via generate_from_random_samples()
     def test_generate_from_random_samples(self, generate_ion_particles, large_bunch):
-        ions = generate_ion_particles(mp_number=1e5)
-        ions.generate_from_random_samples(large_bunch)
+        mp_number = 1e5
+        charge = 1e-9
+        ions = generate_ion_particles(mp_number)
+        ions.generate_from_random_samples(large_bunch, charge)
     
         assert np.all(np.isin(ions["x"], large_bunch["x"]))
         assert np.all(np.isin(ions["y"], large_bunch["y"]))
@@ -50,6 +56,8 @@ class TestIonParticles:
         assert np.all(ions["delta"] == 0)
         assert np.all(ions["tau"] >= -ions.ion_element_length)
         assert np.all(ions["tau"] <= ions.ion_element_length)
+        assert np.isclose(ions.charge, charge)
+        assert np.isclose(ions["charge"].mean(), charge/mp_number, rtol=0.01, atol=1e-9)
         
     # Add two IonParticles instances together and verify combined particle arrays
     def test_add_ion_particles(self, generate_ion_particles):
@@ -58,14 +66,14 @@ class TestIonParticles:
     
         combined = ions1 + ions2
         assert combined.mp_number == 150
-        assert all(combined[coord].shape == (150,) for coord in ["x","xp","y","yp","tau","delta"])
+        assert all(combined[coord].shape == (150,) for coord in ["x","xp","y","yp","tau","delta","charge"])
         assert np.all(combined.alive == True)
 
     # Initialize with alive=False and verify all particles marked as dead
     def test_init_dead_particles(self, generate_ion_particles):
         ions = generate_ion_particles(alive=False)
         assert np.all(ions.alive == False)
-        assert all(ions[coord].shape == (1,) for coord in ["x","xp","y","yp","tau","delta"])
+        assert all(ions[coord].shape == (1,) for coord in ["x","xp","y","yp","tau","delta","charge"])
 
     # Generate distributions with electron bunch containing no particles
     def test_generate_from_empty_bunch(self, generate_ion_particles, small_bunch):
@@ -73,9 +81,9 @@ class TestIonParticles:
         ions = generate_ion_particles()
     
         with pytest.raises(ValueError):
-            ions.generate_as_a_distribution(small_bunch)
+            ions.generate_as_a_distribution(small_bunch, 1e-9)
         with pytest.raises(ValueError):
-            ions.generate_from_random_samples(small_bunch)
+            ions.generate_from_random_samples(small_bunch, 1e-9)
             
 @pytest.fixture
 def generate_ion_monitor(tmp_path):
@@ -118,8 +126,8 @@ class TestIonMonitor:
     def test_data_structures_initialization(self, generate_ion_monitor):
         monitor = generate_ion_monitor()
     
-        assert monitor.mean.shape == (6, 10)
-        assert monitor.std.shape == (6, 10)
+        assert monitor.mean.shape == (4, 10)
+        assert monitor.std.shape == (4, 10)
         assert monitor.charge.shape == (10,)
         assert monitor.time.shape == (10,)
         assert monitor.time.dtype == int
@@ -199,11 +207,6 @@ def generate_beam_ion(demo_ring):
             ion_field_model="strong",
             electron_field_model="strong",
             ion_element_length=demo_ring.L,
-            n_steps=int(demo_ring.h*10),
-            x_radius=0.1,
-            y_radius=0.1,
-            ion_beam_monitor_name=None,
-            use_ion_phase_space_monitor=False,
             n_ion_macroparticles_per_bunch=30,
             generate_method='samples'):
         
@@ -216,11 +219,6 @@ def generate_beam_ion(demo_ring):
             ion_field_model=ion_field_model,
             electron_field_model=electron_field_model, 
             ion_element_length=ion_element_length,
-            n_steps=n_steps,
-            x_radius=x_radius,
-            y_radius=y_radius,
-            ion_beam_monitor_name=ion_beam_monitor_name,
-            use_ion_phase_space_monitor=use_ion_phase_space_monitor,
             n_ion_macroparticles_per_bunch=n_ion_macroparticles_per_bunch,
             generate_method=generate_method)
         return beam_ion
@@ -239,9 +237,15 @@ class TestBeamIonElement:
                              [('weak','weak'), ('weak','strong'), 
                               ('strong','weak'), ('strong', 'strong')])
     def test_track_bunch_partially_lost(self, generate_beam_ion, small_bunch, ion_field_model, electron_field_model):
-        small_bunch.alive[0:5] = False
         beam_ion = generate_beam_ion(ion_field_model=ion_field_model, electron_field_model=electron_field_model)
         assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+        charge_gen = beam_ion.ion_beam.charge
+        
+        # loose half of electron bunch
+        small_bunch.alive[0:5] = False
+        assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+        
+        assert np.isclose(beam_ion.ion_beam.charge, charge_gen*1.5)
 
     @pytest.mark.parametrize('ion_field_model, electron_field_model',
                              [('weak','weak'), ('weak','strong'), 
@@ -288,15 +292,20 @@ class TestBeamIonElement:
         assert np.allclose(beam_ion.ion_beam["y"], initial_y + drift_length)
 
     # Monitor records ion beam data at specified intervals when enabled
-    def test_monitor_recording(self, generate_beam_ion, small_bunch, tmp_path):
-        monitor_file = str(tmp_path / "test_monitor.hdf5")
-        with pytest.warns(UserWarning):
-            beam_ion = generate_beam_ion(ion_beam_monitor_name=monitor_file)
+    def test_monitor_recording(self, 
+                               generate_beam_ion, 
+                               small_bunch, 
+                               generate_ion_monitor,
+                               tmp_path):
+        file_name=tmp_path / "test_monitor.hdf5"
+        monitor = generate_ion_monitor(file_name=file_name)
+        beam_ion = generate_beam_ion()
+        beam_ion.monitors.append(monitor)
     
         beam_ion.track(small_bunch)
     
-        assert os.path.exists(monitor_file)
-        with hp.File(monitor_file, 'r') as f:
+        assert os.path.exists(file_name)
+        with hp.File(file_name, 'r') as f:
             cond = False
             for key in f.keys():
                 if key.startswith('IonData'):
@@ -314,12 +323,14 @@ class TestBeamIonElement:
     # Boundary conditions at aperture edges
     def test_aperture_boundary(self, generate_beam_ion, small_bunch):
         x_radius = 0.001
-        beam_ion = generate_beam_ion(x_radius=x_radius, y_radius=x_radius)
+        aperture = IonAperture(X_radius=x_radius, Y_radius=x_radius)
+        beam_ion = generate_beam_ion()
+        beam_ion.apertures.append(aperture)
     
         beam_ion.generate_new_ions(small_bunch)
     
         beam_ion.ion_beam["x"] = np.ones_like(beam_ion.ion_beam["x"]) * (x_radius * 1.1)
-        beam_ion.aperture.track(beam_ion.ion_beam)
+        beam_ion.track(beam_ion.ion_beam)
     
         assert len(beam_ion.ion_beam["x"]) == 0
 
@@ -332,3 +343,23 @@ class TestBeamIonElement:
         beam_ion.clear_ions()
         assert len(beam_ion.ion_beam["x"]) == 1
         assert beam_ion.ion_beam["x"][0] == 0
+        
+    # Tracking with a pre-set cloud
+    def test_track_preset_cloud(self, generate_beam_ion, small_bunch, generate_ion_particles):
+        beam_ion = generate_beam_ion()
+        assert beam_ion.ion_beam.charge == 0
+        
+        preset_ion_particles = generate_ion_particles(mp_number=1000)
+        preset_charge = 1e-9
+        preset_ion_particles.generate_as_a_distribution(small_bunch, preset_charge)
+        beam_ion.ion_beam += preset_ion_particles
+        assert np.isclose(beam_ion.ion_beam.charge, preset_charge)
+        
+        beam_ion.generate_ions = False
+        assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+        assert np.isclose(beam_ion.ion_beam.charge, preset_charge)
+        
+        beam_ion.generate_ions = True
+        assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+        assert beam_ion.ion_beam.charge > preset_charge
+        
\ No newline at end of file