diff --git a/CITATION.cff b/CITATION.cff
new file mode 100644
index 0000000000000000000000000000000000000000..1c6a56bc3885c651e681de0ba368ef21ec410ad7
--- /dev/null
+++ b/CITATION.cff
@@ -0,0 +1,26 @@
+cff-version: 1.2.0
+title: mbtrack2
+message: 'If you use this software, please cite it as below.'
+type: software
+authors:
+ - family-names: Gamelin
+ given-names: Alexis
+ - family-names: Foosang
+ given-names: Watanyu
+ - family-names: Yamomoto
+ given-names: Naoto
+ - family-names: Gubaidulin
+ given-names: Vadim
+ - family-names: Nagaoka
+ given-names: Ryutaro
+identifiers:
+ - type: doi
+ value: 10.5281/zenodo.14418989
+repository-code: >-
+ https://gitlab.synchrotron-soleil.fr/PA/collective-effects/mbtrack2
+abstract: >
+ mbtrack2 is a coherent object-oriented framework written
+ in python to work on collective effects in synchrotrons.
+license: BSD-3-Clause
+version: 0.8.0
+date-released: '2024-12-16'
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index 18257107f17c4a6026d7dd4e3d99187cd914bd71..fe1acd4431445d6a901e71b9bd0834e150c67f9d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,7 +2,5 @@ FROM gitlab-registry.synchrotron-soleil.fr/pa/collective-effects/python_mpi:late
LABEL name mbtrack2
USER dockeruser
WORKDIR '/home/dockeruser'
-ENV PATH ${PATH}:/home/dockeruser/miniconda3/bin:/home/dockeruser/miniconda3/condabin
-RUN pip3 install accelerator-toolbox==0.5.0
COPY --chown=dockeruser mbtrack2 /home/dockeruser/mbtrack2
ENV PYTHONPATH=/home/dockeruser/
\ No newline at end of file
diff --git a/mbtrack2/instability/__init__.py b/mbtrack2/instability/__init__.py
index f7a1516f07a5ba85219041527106425accf715ce..15fd68680eb6632db29e1eaf0a90369ddc7f4db3 100644
--- a/mbtrack2/instability/__init__.py
+++ b/mbtrack2/instability/__init__.py
@@ -7,6 +7,7 @@ from mbtrack2.instability.instabilities import (
mbi_threshold,
rwmbi_growth_rate,
rwmbi_threshold,
+ transverse_gaussian_space_charge_tune_shift,
)
from mbtrack2.instability.ions import (
fast_beam_ion,
diff --git a/mbtrack2/instability/instabilities.py b/mbtrack2/instability/instabilities.py
index 0a198eb98c90e9f55f16dc5a0805faa948889837..c67bc2205ce9b1de9642ed7c18f4210d128be374 100644
--- a/mbtrack2/instability/instabilities.py
+++ b/mbtrack2/instability/instabilities.py
@@ -454,3 +454,82 @@ def rwmbi_threshold(ring, beff, rho_material, plane='x'):
(1-frac_tune) * omega0) / (2*c*Z0*rho_material))**0.5
return Ith
+
+
+def transverse_gaussian_space_charge_tune_shift(ring, bunch_current, **kwargs):
+ """
+ Return the (maximum) transverse space charge tune shift for a Gaussian
+ bunch in the linear approximation, see Eq.(1) of [1].
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ Ring parameters.
+ bunch_current : float
+ Bunch current in [A].
+ sigma_s : float, optional
+ RMS bunch length in [s].
+ Default is ring.sigma_0.
+ emit_x : float, optional
+ Horizontal emittance in [m.rad].
+ Default is ring.emit[0].
+ emit_y : float, optional
+ Vertical emittance in [m.rad].
+ Default is ring.emit[1].
+ use_lattice : bool, optional
+ If True, use beta fonctions along the lattice.
+ If False, local values of beta fonctions are used.
+ Default is ring.optics.use_local_values.
+ n_points : int, optional
+ Number of points in the lattice to be considered if use_lattice ==
+ True. Default is 1000.
+ sigma_delta : float, optional
+ Relative energy spread.
+ Default is ring.sigma_delta.
+ gamma : float, optional
+ Relativistic Lorentz gamma.
+ Default is ring.gamma.
+
+ Returns
+ -------
+ tune_shift : array of shape (2,)
+ Horizontal and vertical space charge tune shift.
+
+ Reference
+ ---------
+ [1] : Antipov, S. A., Gubaidulin, V., Agapov, I., Garcia, E. C., &
+ Gamelin, A. (2024). Space Charge and Future Light Sources.
+ arXiv preprint arXiv:2409.08637.
+
+ """
+ sigma_s = kwargs.get('sigma_s', ring.sigma_0)
+ emit_x = kwargs.get('emit_x', ring.emit[0])
+ emit_y = kwargs.get('emit_y', ring.emit[1])
+ use_lattice = kwargs.get('use_lattice', not ring.optics.use_local_values)
+ sigma_delta = kwargs.get('sigma_delta', ring.sigma_delta)
+ gamma = kwargs.get('gamma', ring.gamma)
+ n_points = kwargs.get('n_points', 1000)
+ q = ring.particle.charge
+ m = ring.particle.mass
+ r_0 = 1 / (4*pi*epsilon_0) * q**2 / (m * c**2)
+ N = np.abs(bunch_current / ring.f0 / q)
+ sigma_z = sigma_s * c
+
+ if use_lattice:
+ s = np.linspace(0, ring.L, n_points)
+ beta = ring.optics.beta(s)
+ sig_x = (emit_x * beta[0] +
+ ring.optics.dispersion(s)[0]**2 * sigma_delta**2)**0.5
+ sig_y = (emit_y * beta[1] +
+ ring.optics.dispersion(s)[2]**2 * sigma_delta**2)**0.5
+ sig_xy = np.array([sig_x, sig_y])
+ return -r_0 * N / ((2 * pi)**1.5 * gamma**3 * sigma_z) * np.trapz(
+ beta / (sig_xy * (sig_y+sig_x)), s)
+ else:
+ beta = ring.optics.local_beta
+ sig_x = np.sqrt(beta[0] * emit_x)
+ sig_y = np.sqrt(beta[1] * emit_y)
+ sig_xy = np.array([sig_x, sig_y])
+ return -r_0 * N * ring.L / (
+ (2 * pi)**1.5 * gamma**3 * sigma_z) * beta / (sig_xy *
+ (sig_y+sig_x))
diff --git a/mbtrack2/tracking/beam_ion_effects.py b/mbtrack2/tracking/beam_ion_effects.py
index a041fee621da8f0249c792213eb0bf5e28ea3901..5e49ca50bad420935cb9df601373f1ca51d77323 100644
--- a/mbtrack2/tracking/beam_ion_effects.py
+++ b/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
@@ -330,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.
@@ -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,16 +639,19 @@ 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(),
+ )
+ sb_my, sb_stdy = (
+ second_beam["y"].mean(),
+ second_beam["y"].std(),
)
- sb_mx, sb_stdx = (
- second_beam["x"].mean(),
- second_beam["x"].std(),
- )
- sb_my, sb_stdy = (
- second_beam["y"].mean(),
- second_beam["y"].std(),
- )
if field_model == "weak":
en_x, en_y = get_displaced_efield(
_efieldn_mit,
@@ -556,10 +664,13 @@ class BeamIonElement(Element):
)
elif field_model == "strong":
- fb_mx, fb_my = (
- first_beam["x"].mean(),
- first_beam["y"].mean(),
- )
+ 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(),
+ )
en_x, en_y = get_displaced_efield(_efieldn_mit, fb_mx, fb_my,
sb_stdx, sb_stdy, sb_mx, sb_my)
@@ -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)
- self.track_ions_in_a_drift(drift_length=self.bunch_spacing)
+ if self.ion_drift:
+ self.track_ions_in_a_drift(drift_length=self.bunch_spacing)
diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 0f69f6cfb391776d5f747b2230fb9a9c374df895..1575a1973bc4a1ddedc33b0318f4f464f95ff051 100644
--- a/mbtrack2/tracking/element.py
+++ b/mbtrack2/tracking/element.py
@@ -415,7 +415,7 @@ class TransverseMap(TransverseMapSector):
2 * np.pi * ring.tune, ring.chro, ring.adts)
-def transverse_map_sector_generator(ring, positions):
+def transverse_map_sector_generator(ring, positions, **kwargs):
"""
Convenience function which generate a list of TransverseMapSector elements
from a ring:
@@ -435,6 +435,8 @@ def transverse_map_sector_generator(ring, positions):
of the TransverseMapSector elements.
The array should contain the initial position (s=0) but not the end
position (s=ring.L), so like position = np.array([0, pos1, pos2, ...]).
+
+ See at.physics.nonlinear.chromaticity for **kwargs
Returns
-------
@@ -463,8 +465,10 @@ def transverse_map_sector_generator(ring, positions):
adts=adts))
else:
import at
+ dp = kwargs.get('dp', 1e-2)
+ order = kwargs.get('order', 1)
- def _compute_chro(ring, N_sec, dp=1e-2, order=4):
+ def _compute_chro(ring, N_sec, dp, order):
lat = deepcopy(ring.optics.lattice)
lat.append(at.Marker("END"))
fit, _, _ = at.physics.nonlinear.chromaticity(lat,
@@ -472,12 +476,18 @@ def transverse_map_sector_generator(ring, positions):
dpm=dp,
n_points=100,
order=order)
- Chrox, Chroy = fit[0, 1:], fit[1, 1:]
- chrox = np.array([np.linspace(0, val, N_sec) for val in Chrox])
- chroy = np.array([np.linspace(0, val, N_sec) for val in Chroy])
- return np.array([chrox, chroy])
+ chro_xy = [
+ elem for pair in zip(fit[0, 1:], fit[1, 1:]) for elem in pair
+ ]
+ len_chro = int(order * 2)
+ _chro = np.zeros((len_chro, N_sec))
+ for i in range(len_chro):
+ chro_order_splited = np.linspace(0, chro_xy[i], N_sec)
+ _chro[i, :] = chro_order_splited
+
+ return _chro
- _chro = _compute_chro(ring, N_sec)
+ _chro = _compute_chro(ring, N_sec, dp, order)
for i in range(N_sec):
alpha0 = ring.optics.alpha(positions[i])
beta0 = ring.optics.beta(positions[i])
diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py
index 41458eff38988e909c1f3c90b2d4863c0a3e12ae..87717d06aec2dbed49728178571168faa5bc4bc0 100644
--- a/mbtrack2/tracking/synchrotron.py
+++ b/mbtrack2/tracking/synchrotron.py
@@ -373,23 +373,78 @@ class Synchrotron:
tuneS = phi / (2 * np.pi)
return tuneS
- def get_adts(self):
+ def get_adts(self,
+ xm=1e-4,
+ ym=1e-4,
+ npoints=9,
+ plot=False,
+ ax=None,
+ **kwargs):
"""
Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar
componenet from AT lattice.
+
+ Parameters
+ ----------
+ xm : float, optional
+ Maximum horizontal amplitude in [m].
+ Default is 1e-4.
+ ym : float, optional
+ Maximum vertical amplitude in [m].
+ Default is 1e-4.
+ npoints : int, optional
+ Number of points in each plane.
+ Default is 9.
+ plot : bool, optional
+ If True, plot the computed tune shift with amplitude.
+ Default is False.
+ ax : array of shape (2,) of matplotlib.axes.Axes, optional
+ Axes where to plot the figures.
+ Default is None.
+
+ Returns
+ -------
+ ax : array of shape (2,) of matplotlib.axes.Axes
+ Only if plot is True.
+
+ See at.physics.nonlinear.detuning for **kwargs
"""
import at
if self.optics.use_local_values:
raise ValueError("ADTS needs to be provided manualy as no AT" +
" lattice file is loaded.")
-
- det = at.physics.nonlinear.gen_detuning_elem(self.optics.lattice)
- coef_xx = np.array([det.A1 / 2, 0])
- coef_yx = np.array([det.A2 / 2, 0])
- coef_xy = np.array([det.A2 / 2, 0])
- coef_yy = np.array([det.A3 / 2, 0])
+ r0, r1, x, q_dx, y, q_dy = at.physics.nonlinear.detuning(
+ self.optics.lattice.radiation_off(copy=True),
+ npoints=npoints,
+ xm=xm,
+ ym=ym,
+ **kwargs)
+ coef_xx = np.array([r1[0][0] / 2, 0])
+ coef_yx = np.array([r1[0][1] / 2, 0])
+ coef_xy = np.array([r1[0][1] / 2, 0])
+ coef_yy = np.array([r1[1][1] / 2, 0])
self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
+ if plot:
+ if ax is None:
+ fig, ax = plt.subplots(2, 1)
+
+ ax[0].scatter(x * 1e6, q_dx[:, 0])
+ ax[0].scatter(y * 1e6, q_dy[:, 0])
+ ax[0].set_ylabel("Delta Tune X")
+ ax[0].set_xlabel("Amplitude [um]")
+ ax[0].legend(['dqx/dx', 'dqx/dy'])
+ ax[0].grid()
+
+ ax[1].scatter(x * 1e6, q_dx[:, 1])
+ ax[1].scatter(y * 1e6, q_dy[:, 1])
+ ax[1].set_ylabel("Delta Tune Y")
+ ax[1].set_xlabel("Amplitude [um]")
+ ax[1].legend(['dqy/dx', 'dqy/dy'])
+ ax[1].grid()
+
+ return ax
+
def get_chroma(self, order=4, dpm=0.02, n_points=100):
"""
Compute chromaticity (linear and nonlinear) from AT lattice and update the property.
diff --git a/tests/unit/tracking/test_beam_ion_effects.py b/tests/unit/tracking/test_beam_ion_effects.py
index 920ca6caa66bd3450b6a52f5b833476d70d64c44..6f3c393b5d2d3c15b8536d44fd8bc1bf94c7edaa 100644
--- a/tests/unit/tracking/test_beam_ion_effects.py
+++ b/tests/unit/tracking/test_beam_ion_effects.py
@@ -31,39 +31,31 @@ def generate_ion_particles(demo_ring):
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)
-
- 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)
- assert np.isclose(ions["y"].mean(),
- large_bunch["y"].mean(),
- rtol=0.1,
- atol=1e-5)
- assert np.isclose(ions["y"].std(),
- large_bunch["y"].std(),
- rtol=0.1,
- atol=1e-5)
+ def test_generate_as_distribution(self, generate_ion_particles, 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)
+ assert np.isclose(ions["y"].mean(), large_bunch["y"].mean(), rtol=0.1, atol=1e-5)
+ assert np.isclose(ions["y"].std(), large_bunch["y"].std(), rtol=0.1, atol=1e-5)
assert np.all(ions["xp"] == 0)
assert np.all(ions["yp"] == 0)
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)
-
+ def test_generate_from_random_samples(self, generate_ion_particles, 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"]))
assert np.all(ions["xp"] == 0)
@@ -71,7 +63,9 @@ 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):
ions1 = generate_ion_particles(mp_number=100)
@@ -79,16 +73,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,
@@ -97,11 +89,10 @@ 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):
@@ -147,11 +138,11 @@ class TestIonMonitor:
# Data structures are properly initialized with correct shapes and types
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.charge.shape == (10, )
- assert monitor.time.shape == (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
# Initialize monitor with total_size not divisible by buffer_size
@@ -221,23 +212,18 @@ class TestElipticalAperture:
@pytest.fixture
def generate_beam_ion(demo_ring):
-
- def generate(ion_mass=1.67e-27,
- ion_charge=1.6e-19,
- ionization_cross_section=1e-22,
- residual_gas_density=1e50,
- ring=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'):
-
+ def generate(
+ ion_mass=1.67e-27,
+ ion_charge=1.6e-19,
+ ionization_cross_section=1e-22,
+ residual_gas_density=1e50,
+ ring=demo_ring,
+ ion_field_model="strong",
+ electron_field_model="strong",
+ ion_element_length=demo_ring.L,
+ n_ion_macroparticles_per_bunch=30,
+ generate_method='samples'):
+
beam_ion = BeamIonElement(
ion_mass=ion_mass,
ion_charge=ion_charge,
@@ -247,11 +233,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
@@ -271,14 +252,18 @@ class TestBeamIonElement:
assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
@pytest.mark.parametrize('ion_field_model, electron_field_model',
- [('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):
+ [('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):
+ 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
- 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"])
+ 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'),
@@ -337,15 +322,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'):
@@ -363,14 +353,15 @@ 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.ion_beam["x"] = np.ones_like(beam_ion.ion_beam["x"]) * (x_radius * 1.1)
+ beam_ion.track(beam_ion.ion_beam)
+
assert len(beam_ion.ion_beam["x"]) == 0
# Ion clearing removes all particles as expected
@@ -382,3 +373,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
+
diff --git a/tests/unit/tracking/test_synchrotron.py b/tests/unit/tracking/test_synchrotron.py
index 252551f2a4ce1783182c15937b3138e18a31fac9..2a7dc4054cc84e8c5fac7fd79e95b257edf66dbe 100644
--- a/tests/unit/tracking/test_synchrotron.py
+++ b/tests/unit/tracking/test_synchrotron.py
@@ -142,7 +142,11 @@ class TestSynchrotron:
def test_get_adts(self, ring_with_at_lattice):
ring_with_at_lattice.get_adts()
assert ring_with_at_lattice.adts is not None
-
+
+ def test_get_adts_plot(self, ring_with_at_lattice):
+ axes = ring_with_at_lattice.get_adts(plot=True)
+ assert axes is not None
+
def test_get_chroma(self, ring_with_at_lattice):
ring_with_at_lattice.get_chroma()
assert len(ring_with_at_lattice.chro) == 8