"""
Module implementing necessary functionalities for beam-ion interactions.
Classes:
BeamIonElement
IonMonitor
"""
from functools import wraps
from typing import Callable
from numpy.typing import NDArray
from scipy.constants import c, e
from mbtrack2.tracking.aperture import Aperture
from mbtrack2.tracking.element import Element
from mbtrack2.tracking.emfields import _efieldn_mit, get_displaced_efield
from mbtrack2.tracking.monitors.monitors import Monitor
from mbtrack2.tracking.particles import Beam, Bunch, IonParticles
from mbtrack2.utilities.synchrotron import Synchrotron
[docs]
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].
ion_charge : float
The charge of the ions in [C].
ionization_cross_section : float
The cross section of ionization in [m^2].
residual_gas_density : float
The residual gas density in [m^-3].
ring : instance of Synchrotron()
The ring.
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 : {'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.
n_ion_macroparticles_per_bunch : int, optional
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_ion_macroparticles_per_bunch=30,
generate_method='distribution')
Initializes the BeamIonElement object.
parallel(track)
Defines the decorator @parallel to handle tracking of Beam() objects.
clear_ions()
Clear the ion particles in the ion beam.
track_ions_in_a_drift(drift_length)
Tracks the ions in a drift.
generate_new_ions(electron_bunch)
Generate new ions based on the given electron bunch.
track(electron_bunch)
Beam-ion interaction kicks.
Raises
------
UserWarning
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).
"""
[docs]
def __init__(self,
ion_mass: float,
ion_charge: float,
ionization_cross_section: float,
residual_gas_density: float,
ring: Synchrotron,
ion_field_model: str,
electron_field_model: str,
ion_element_length: float,
n_ion_macroparticles_per_bunch: int = 30,
generate_method: str = 'distribution',
generate_ions: bool = True,
beam_ion_interaction: bool = True,
ion_drift: bool = True):
self.ring: Synchrotron = ring
self.bunch_spacing = ring.L / ring.h
self.ion_mass = ion_mass
self.ionization_cross_section = ionization_cross_section
self.residual_gas_density = residual_gas_density
self.ion_charge = ion_charge
self.electron_field_model = electron_field_model
self.ion_field_model = ion_field_model
self.ion_element_length = ion_element_length
self.generate_method = generate_method
if self.generate_method not in ["distribution", "samples"]:
raise ValueError("Wrong generate_method.")
self.n_ion_macroparticles_per_bunch: int = n_ion_macroparticles_per_bunch
self.ion_beam = IonParticles(
mp_number=1,
ion_element_length=self.ion_element_length,
ring=self.ring)
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 = []
[docs]
@staticmethod
def parallel(track: Callable) -> Callable:
"""
Defines the decorator @parallel which handle the embarrassingly
parallel case which happens when there is no bunch to bunch
interaction in the tracking routine.
Adding @Element.parallel allows to write the track method of the
Element subclass for a Bunch object instead of a Beam object.
Parameters
----------
track : function, method of an Element subclass
track method of an Element subclass which takes a Bunch object as
input
Returns
-------
track_wrapper: function, method of an Element subclass
track method of an Element subclass which takes a Beam object or a
Bunch object as input
"""
@wraps(track)
def track_wrapper(*args, **kwargs):
if isinstance(args[1], Beam):
self = args[0]
beam = args[1]
if beam.mpi_switch:
raise ValueError(
"Tracking through beam-ion element is performed sequentially."
)
for bunch in beam:
track(self, bunch, *args[2:], **kwargs)
else:
self = args[0]
bunch = args[1]
track(self, bunch, *args[2:], **kwargs)
return track_wrapper
[docs]
def clear_ions(self):
"""
Clear the ion particles in the ion beam.
"""
self.ion_beam.particles = IonParticles(
mp_number=1,
ion_element_length=self.ion_element_length,
ring=self.ring)
[docs]
def track_ions_in_a_drift(self, drift_length: float):
"""
Tracks the ions in a drift.
Parameters
----------
drift_length : float
The drift length in meters.
"""
drifted_ions_x = self.ion_beam["x"] + drift_length * self.ion_beam["xp"]
drifted_ions_y = self.ion_beam["y"] + drift_length * self.ion_beam["yp"]
self.ion_beam["x"] = drifted_ions_x
self.ion_beam["y"] = drifted_ions_y
[docs]
def _get_efields(self, first_beam: Bunch | IonParticles,
second_beam: Bunch | IonParticles,
field_model: str) -> tuple[NDArray, NDArray]:
"""
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, which is being acted on.
second_beam : IonParticles or Bunch
The second beam, which is generating an electric field.
field_model : {'weak', 'strong', 'PIC'}
The field model used for the interaction.
Returns
-------
en_x : numpy.ndarray
The x component of the electric field.
en_y : numpy.ndarray
The y component of the electric field.
"""
if not field_model in ["weak", "strong", "PIC"]:
raise ValueError(
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(),
)
if field_model == "weak":
en_x, en_y = get_displaced_efield(
_efieldn_mit,
first_beam["x"],
first_beam["y"],
sb_stdx,
sb_stdy,
sb_mx,
sb_my,
)
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(),
)
en_x, en_y = get_displaced_efield(_efieldn_mit, fb_mx, fb_my,
sb_stdx, sb_stdy, sb_mx, sb_my)
elif field_model == "PIC":
from PyPIC import FFT_OpenBoundary
from PyPIC import geom_impact_ellip as ellipse
qe = e
Dx = 0.1 * sb_stdx
Dy = 0.1 * sb_stdy
x_aper = 10 * sb_stdx
y_aper = 10 * sb_stdy
chamber = ellipse.ellip_cham_geom_object(x_aper=x_aper,
y_aper=y_aper)
picFFT = FFT_OpenBoundary.FFT_OpenBoundary(
x_aper=chamber.x_aper,
y_aper=chamber.y_aper,
dx=Dx,
dy=Dy,
fftlib="pyfftw",
)
nel_part = 0 * second_beam["x"] + 1.0
picFFT.scatter(second_beam["x"], second_beam["y"], nel_part)
picFFT.solve()
en_x, en_y = picFFT.gather(first_beam["x"], first_beam["y"])
en_x /= qe * second_beam["x"].shape[0]
en_y /= qe * second_beam["x"].shape[0]
else:
raise ValueError(
f"The implementation for required beam-ion interaction\
model {field_model} is not implemented")
return en_x, en_y
[docs]
def _get_new_beam_momentum(
self,
first_beam: Bunch | IonParticles,
second_beam: Bunch | IonParticles,
prefactor: float,
field_model: str = "strong") -> tuple[NDArray, NDArray]:
"""
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, which is being acted on.
second_beam : 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 : {'weak', 'strong', 'PIC'}
The field model used for the interaction.
Default is "strong".
Returns
-------
new_xp : numpy.ndarray
The new x momentum of the first beam.
new_yp : numpy.ndarray
The new y momentum of the first beam.
"""
en_x, en_y = self._get_efields(first_beam,
second_beam,
field_model=field_model)
kicks_x = prefactor * en_x
kicks_y = prefactor * en_y
new_xp = first_beam["xp"] + kicks_x
new_yp = first_beam["yp"] + kicks_y
return new_xp, new_yp
[docs]
def _update_beam_momentum(self, beam: Bunch, new_xp: NDArray,
new_yp: NDArray):
beam["xp"] = new_xp
beam["yp"] = new_yp
[docs]
def generate_new_ions(self, electron_bunch: Bunch):
"""
Generate new ions based on the given electron bunch.
Parameters
----------
electron_bunch : Bunch
The electron bunch used to generate new ions.
Returns
-------
None
"""
new_ion_particles = IonParticles(
mp_number=self.n_ion_macroparticles_per_bunch,
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, charge=new_ion_charge)
elif self.generate_method == 'samples':
new_ion_particles.generate_from_random_samples(
electron_bunch=electron_bunch, charge=new_ion_charge)
self.ion_beam += new_ion_particles
[docs]
@parallel
def track(self, electron_bunch: Bunch):
"""
Beam-ion interaction kicks.
Parameters
----------
electron_bunch : Bunch() or Beam() class object
An electron bunch to be interacted with.
"""
if electron_bunch.is_empty:
empty_bucket = True
else:
empty_bucket = False
if not empty_bucket and self.generate_ions:
self.generate_new_ions(electron_bunch=electron_bunch)
for aperture in self.apertures:
aperture.track(self.ion_beam)
for monitor in self.monitors:
monitor.track(self.ion_beam)
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))
new_xp_ions, new_yp_ions = self._get_new_beam_momentum(
self.ion_beam,
electron_bunch,
prefactor_to_electron_field,
field_model=self.electron_field_model,
)
new_xp_electrons, new_yp_electrons = self._get_new_beam_momentum(
electron_bunch,
self.ion_beam,
prefactor_to_ion_field,
field_model=self.ion_field_model,
)
self._update_beam_momentum(self.ion_beam, new_xp_ions, new_yp_ions)
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)