Skip to content
Snippets Groups Projects
Commit 1fd72e16 authored by Alexis GAMELIN's avatar Alexis GAMELIN
Browse files

Improve docstrings and tests

parent 42b73361
No related branches found
No related tags found
1 merge request!33Update BeamIonElement
......@@ -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
......@@ -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,
......@@ -261,9 +289,6 @@ class IonParticles(Bunch):
return self["charge"].sum()
def _mean_weighted(self, coord):
"""
Return the mean position of alive particles for each coordinates.
"""
if self.charge == 0:
return np.zeros_like(coord, dtype=float)
else:
......@@ -272,10 +297,6 @@ class IonParticles(Bunch):
return np.squeeze(np.array(mean))
def _mean_std_weighted(self, coord):
"""
Return the standard deviation of the position of alive
particles for each coordinates.
"""
if self.charge == 0:
return np.zeros_like(coord,
dtype=float), np.zeros_like(coord,
......@@ -293,35 +314,42 @@ class IonParticles(Bunch):
@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_weighted(self):
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_weighted(self):
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.")
......@@ -353,12 +381,15 @@ class IonParticles(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.")
......@@ -392,47 +423,71 @@ class IonParticles(Bunch):
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.
......@@ -448,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,
......@@ -463,7 +522,10 @@ class BeamIonElement(Element):
electron_field_model,
ion_element_length,
n_ion_macroparticles_per_bunch=30,
generate_method='distribution'):
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
......@@ -481,16 +543,10 @@ class BeamIonElement(Element):
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
self.generate_ions = True
self.beam_ion_interaction = True
self.ion_drift = True
self.generate_ions = generate_ions
self.beam_ion_interaction = beam_ion_interaction
self.ion_drift = ion_drift
# interfaces for apertures and montiors
self.apertures = []
......@@ -562,16 +618,17 @@ class BeamIonElement(Element):
def _get_efields(self, first_beam, second_beam, field_model):
"""
Calculates the electromagnetic field of the second beam acting on the first 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
-------
......@@ -582,10 +639,10 @@ 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_weighted
sb_mx, sb_my, sb_stdx, sb_stdy = second_beam.mean_std_xy
else:
sb_mx, sb_stdx = (
second_beam["x"].mean(),
......@@ -608,7 +665,7 @@ class BeamIonElement(Element):
elif field_model == "strong":
if isinstance(first_beam, IonParticles):
fb_mx, fb_my = first_beam.mean_weighted
fb_mx, fb_my = first_beam.mean_xy
else:
fb_mx, fb_my = (
first_beam["x"].mean(),
......@@ -648,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
......@@ -689,7 +747,7 @@ class BeamIonElement(Element):
Parameters
----------
electron_bunch : ElectronBunch
electron_bunch : Bunch
The electron bunch used to generate new ions.
Returns
......
......@@ -237,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'),
......@@ -336,4 +342,24 @@ class TestBeamIonElement:
beam_ion.clear_ions()
assert len(beam_ion.ion_beam["x"]) == 1
assert beam_ion.ion_beam["x"][0] == 0
\ No newline at end of file
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
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment