implemented IBS effect and covered with unit tests

ocelot/cpbd/physics_proc.py

@@ -1,6 +1,6 @@
 from ocelot.cpbd.io import save_particle_array
-from ocelot.common.globals import h_eV_s, m_e_eV, m_e_GeV, ro_e, speed_of_light
-from ocelot.cpbd.beam import Twiss, beam_matching
+from ocelot.common.globals import h_eV_s, m_e_eV, m_e_GeV, ro_e, speed_of_light, q_e
+from ocelot.cpbd.beam import Twiss, beam_matching, global_slice_analysis, s_to_cur, get_envelope
 from ocelot.utils.acc_utils import slice_bunching
 from ocelot.common.ocelog import *
 from ocelot.cpbd.beam import ParticleArray
@@ -638,3 +638,118 @@ def apply(self, p_array, dz=0):
         p_array.E = self.Eref
         p_array.p()[:] = p_new[:]
 
+
+class IBS(PhysProc):
+    """
+    Intrabeam Scattering (IBS) Physics Process.
+
+    This class models the intrabeam scattering process in particle beams. Two methods are implemented based on different formulations:
+
+    1. **Huang Method:** Based on Z. Huang's work: *Intrabeam Scattering in an X-ray FEL Driver* (LCLS-TN-02-8, 2002).
+       URL: https://www-ssrl.slac.stanford.edu/lcls/technotes/LCLS-TN-02-8.pdf
+    2. **Nagaitsev Method:** Based on S. Nagaitsev's work: *Intrabeam scattering formulas for fast numerical evaluation*
+       (PRAB 8, 064403, 2005). URL: https://journals.aps.org/prab/abstract/10.1103/PhysRevSTAB.8.064403
+
+    The methods can be selected using `self.method`, which accepts "Huang" or "Nagaitsev".
+
+    Key assumptions and features:
+    - A **round beam approximation** is used, where `sigma_xy = (sigma_x + sigma_y) / 2`.
+    - Beam parameters are calculated within a slice of width ±0.5 `sigma_tau` relative to the slice with maximum current
+      (`self.slice = "Imax"` by default).
+    - The Coulomb logarithm (`self.Clog`) is a configurable parameter, defaulting to 8 (based on Z. Huang's work).
+    - A factor of 2 is applied to the Nagaitsev formula to align high-energy results with the Huang formula.
+
+    Attributes:
+    - `method` (str): Selected method for IBS calculation ("Huang" or "Nagaitsev").
+    - `Clog` (float): Coulomb logarithm (default: 8) is used constant if update_Clog = False
+    - `update_Clog` (bool): recalculate Clog on each step using Huang's formula without cut off.
+    - `bounds` (list): Range of bounds for slice analysis in units of `sigma_tau` (default: `[-0.5, 0.5]`).
+    - `slice` (str): Reference slice ("Imax" for maximum current by default).
+
+    Methods:
+    - `get_beam_params(p_array)`: Computes beam parameters such as `sigma_xy`, `sigma_z`, and normalized emittance.
+    - `apply(p_array, dz)`: Applies the IBS process to a particle array over a specified distance.
+
+    Raises:
+    - `ValueError`: If an invalid method is specified.
+    """
+
+    def __init__(self, method="Huang"):
+        PhysProc.__init__(self)
+        self.method = method
+        self.Clog = 8
+        self.update_Clog = True
+        self.bounds = [-0.5, 0.5]
+        self.slice = "Imax"
+
+        self.emit_n = None
+        self.sigma_xy = None
+        self.sigma_z = None
+        self.sigma_dgamma0 = None
+
+    def get_beam_params(self, p_array):
+        """
+        Compute beam parameters such as sigma_xy, sigma_z, and normalized emittance.
+
+        :param p_array: ParticleArray
+            Input particle array containing particle properties.
+        """
+        tws = get_envelope(p_array, bounds=self.bounds, slice=self.slice)
+        self.sigma_z = np.std(p_array.tau())
+        self.sigma_xy = (np.sqrt(tws.xx) + np.sqrt(tws.yy)) / 2.
+        self.emit_n = (tws.emit_xn + tws.emit_yn) / 2.
+        pc = np.sqrt(p_array.E ** 2 - m_e_GeV ** 2)
+        self.sigma_dgamma0 = np.sqrt(tws.pp)*pc / m_e_GeV
+
+    def estimate_Clog(self):
+        """
+        Used formula from Hunag's paper without cut offs
+
+        :return: float - Coulomb logarithm
+        """
+        Clog = np.log(self.emit_n * self.emit_n / (ro_e * self.sigma_xy))
+        return Clog
+
+    def apply(self, p_array, dz):
+        """
+        Apply the IBS process to a particle array over a given path length.
+
+        :param p_array: ParticleArray
+            The particle array to be processed.
+        :param dz: float
+            The path length over which the IBS process is applied.
+
+        Raises:
+        - ValueError: If an invalid method is specified.
+        """
+        # Number of particles
+        Nb = np.sum(p_array.q_array) / q_e
+
+        # Update beam parameters
+        self.get_beam_params(p_array)
+
+        # estimate C log
+        if self.update_Clog:
+            Clog = self.estimate_Clog()
+        else:
+            Clog = self.Clog
+
+        # Particle energy and momentum
+        gamma = p_array.E / m_e_GeV
+        pc = np.sqrt(p_array.E**2 - m_e_GeV**2)
+
+        # Calculate sigma_gamma based on the selected method
+        if self.method == "Huang":
+            sigma_gamma = np.sqrt(Clog * ro_e**2 * Nb / (self.sigma_xy * self.emit_n * self.sigma_z) * dz / 4)
+
+        elif self.method == "Nagaitsev":
+            xi = (self.sigma_dgamma0 * self.sigma_xy / (gamma * self.emit_n))**2
+            F = (1 - xi**0.25) * np.log(xi + 1) / xi
+            sigma_gamma = np.sqrt(Clog / 4 * ro_e**2 * Nb / (self.sigma_xy * self.emit_n * self.sigma_z) * F * dz)
+        else:
+            raise ValueError(f"Invalid method '{self.method}'. Choose 'Huang' or 'Nagaitsev'.")
+
+        # Energy spread and update particle momenta
+        sigma_e = sigma_gamma * m_e_GeV / pc
+        p_array.p()[:] += np.random.randn(p_array.n) * sigma_e
+

unit_tests/ebeam_test/phys_proc/phys_proc_conf.py

@@ -8,6 +8,7 @@
 from ocelot.utils.acc_utils import chicane_RTU
 from ocelot.cpbd.sc import LSC
 from ocelot.cpbd.wake3D import s2current
+from ocelot.cpbd.physics_proc import IBS
 
 """Lattice elements definition"""
 

unit_tests/ebeam_test/phys_proc/phys_proc_test.py

@@ -479,6 +479,65 @@ def test_phase_space_aperture(lattice, p_array, parameter, update_ref_values=Fal
     assert check_result(result1 + result2)
 
 
+@pytest.mark.parametrize('parameter', [0, 1, 2, 3])
+def test_ibs(lattice, p_array, parameter, update_ref_values=False):
+    """
+    test PhysicsProc WakeTable
+
+    0 - tracking of the electron beam with positive energy chirp trough undulator
+    1 - tracking of the electron beam with negative energy chirp trough undulator
+    """
+    emit_x = 1e-6  # m
+    sigma_x = 250e-6  # m
+    sigma_z = 1e-3  # m
+    E = 0.1
+    gamma = E / m_e_GeV
+    parray_init = generate_parray(sigma_x=sigma_x, sigma_px=emit_x / gamma / sigma_x, sigma_y=sigma_x,
+                                  sigma_py=emit_x / gamma / sigma_x,
+                                  sigma_tau=sigma_z, sigma_p=1e-3 / 100, chirp=0.00, charge=0.25e-9, nparticles=2000000,
+                                  energy=E,
+                                  tau_trunc=None, tws=None, shape="gauss")
+    parray = copy.deepcopy(parray_init)
+    ibs = IBS()
+    ibs.Clog = 8
+    if parameter == 0:
+        ibs.method = "Nagaitsev"
+        ibs.update_Clog = True
+    if parameter == 1:
+        ibs.method = "Nagaitsev"
+        ibs.update_Clog = False
+    if parameter == 2:
+        ibs.method = "Huang"
+        ibs.update_Clog = True
+    else:
+        ibs.method = "Huang"
+        ibs.update_Clog = False
+    tws = get_envelope(parray, bounds=[-0.5, 0.5], slice="Imax")
+
+    pc = np.sqrt(parray.E ** 2 - m_e_GeV ** 2)
+
+    sigma_p = np.sqrt(tws.pp) * pc
+
+    Sigma_e_ocl = [sigma_p * 1e6]
+    s_ocl = [1]
+    s_pos = 1
+    for i in range(29):
+        ibs.apply(parray, dz=1)
+
+        tws = get_envelope(parray, bounds=[-0.5, 0.5], slice="Imax")
+        sigma_p = np.sqrt(tws.pp) * pc
+        Sigma_e_ocl.append(sigma_p * 1e6)
+        s_pos += 1
+        s_ocl.append(s_pos)
+    B = np.column_stack((s_ocl, Sigma_e_ocl))
+    if update_ref_values:
+        return numpy2json(B)
+
+    B_ref = json2numpy(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + str(parameter) + '.json'))
+
+    result = check_matrix(B, B_ref, tolerance=1.0e-2, tolerance_type='relative', assert_info=' current - ')
+    assert check_result(result)
+
 def setup_module(module):
 
     f = open(pytest.TEST_RESULTS_FILE, 'a')
@@ -525,12 +584,14 @@ def test_update_ref_values(lattice, p_array, cmdopt):
     update_functions.append('test_track_spontan_rad_effects')
     update_functions.append("test_dechirper_offaxis")
     update_functions.append("test_phase_space_aperture")
+    update_functions.append("test_ibs")
 
     update_function_parameters = {}
     update_function_parameters['test_track_smooth'] = [0, 1]
     update_function_parameters['test_track_aperture'] = [0, 1, 2]
     update_function_parameters['test_track_ellipt_aperture'] = [0, 1, 2]
     update_function_parameters['test_phase_space_aperture'] = [0, 1, 2]
+    update_function_parameters['test_ibs'] = [0, 1, 2, 3]
 
     parameter = update_function_parameters[cmdopt] if cmdopt in update_function_parameters.keys() else ['']
 

unit_tests/ebeam_test/phys_proc/ref_results/test_ibs0.json

@@ -0,0 +1 @@
+[{"0": 1.0, "1": 1.0003445810416542}, {"0": 2.0, "1": 1.0132264582504005}, {"0": 3.0, "1": 1.0261418348455744}, {"0": 4.0, "1": 1.038839600921569}, {"0": 5.0, "1": 1.051433924828487}, {"0": 6.0, "1": 1.0637691897022543}, {"0": 7.0, "1": 1.0759164336677531}, {"0": 8.0, "1": 1.0877526417501853}, {"0": 9.0, "1": 1.0996930020999416}, {"0": 10.0, "1": 1.1111914024224365}, {"0": 11.0, "1": 1.1227622635463177}, {"0": 12.0, "1": 1.1342028375672764}, {"0": 13.0, "1": 1.1450911445886964}, {"0": 14.0, "1": 1.1565376880003095}, {"0": 15.0, "1": 1.1679669259529066}, {"0": 16.0, "1": 1.1789980627090422}, {"0": 17.0, "1": 1.1899054119658774}, {"0": 18.0, "1": 1.2007118798803684}, {"0": 19.0, "1": 1.211069534012699}, {"0": 20.0, "1": 1.22181904712469}, {"0": 21.0, "1": 1.2325695487418533}, {"0": 22.0, "1": 1.2429672837146108}, {"0": 23.0, "1": 1.2531420379287972}, {"0": 24.0, "1": 1.2632975093964207}, {"0": 25.0, "1": 1.2736676968584473}, {"0": 26.0, "1": 1.2836114078375374}, {"0": 27.0, "1": 1.2938845568653625}, {"0": 28.0, "1": 1.3039373571748745}, {"0": 29.0, "1": 1.3138357800093694}, {"0": 30.0, "1": 1.3235118415419846}]

unit_tests/ebeam_test/phys_proc/ref_results/test_ibs1.json

@@ -0,0 +1 @@
+[{"0": 1.0, "1": 0.9997025703785725}, {"0": 2.0, "1": 1.012395903762503}, {"0": 3.0, "1": 1.0251126088319444}, {"0": 4.0, "1": 1.0376056618449674}, {"0": 5.0, "1": 1.0499048940657663}, {"0": 6.0, "1": 1.06225375386254}, {"0": 7.0, "1": 1.0740669180810403}, {"0": 8.0, "1": 1.0862370443699896}, {"0": 9.0, "1": 1.0979075093330777}, {"0": 10.0, "1": 1.1098043464756342}, {"0": 11.0, "1": 1.121319362226039}, {"0": 12.0, "1": 1.132884804762504}, {"0": 13.0, "1": 1.144198017262989}, {"0": 14.0, "1": 1.1554318737881544}, {"0": 15.0, "1": 1.1664870514941097}, {"0": 16.0, "1": 1.1776092358083947}, {"0": 17.0, "1": 1.1885954628388211}, {"0": 18.0, "1": 1.199411196630008}, {"0": 19.0, "1": 1.2100810668961404}, {"0": 20.0, "1": 1.2206573312518325}, {"0": 21.0, "1": 1.2312257284671644}, {"0": 22.0, "1": 1.241806660485859}, {"0": 23.0, "1": 1.2523288788252158}, {"0": 24.0, "1": 1.2626353888570516}, {"0": 25.0, "1": 1.2728536428570845}, {"0": 26.0, "1": 1.2829232218752427}, {"0": 27.0, "1": 1.29329744289136}, {"0": 28.0, "1": 1.3030613265594209}, {"0": 29.0, "1": 1.3129054994902707}, {"0": 30.0, "1": 1.322368482665442}]

unit_tests/ebeam_test/phys_proc/ref_results/test_ibs2.json

@@ -0,0 +1 @@
+[{"0": 1.0, "1": 0.9998725663150522}, {"0": 2.0, "1": 1.0223093672819514}, {"0": 3.0, "1": 1.044489859534837}, {"0": 4.0, "1": 1.0665051846968936}, {"0": 5.0, "1": 1.087622957458697}, {"0": 6.0, "1": 1.108630064808506}, {"0": 7.0, "1": 1.1292257039728177}, {"0": 8.0, "1": 1.149530533920015}, {"0": 9.0, "1": 1.16927330577773}, {"0": 10.0, "1": 1.189006692964033}, {"0": 11.0, "1": 1.2082081415488588}, {"0": 12.0, "1": 1.2265741595986799}, {"0": 13.0, "1": 1.2451696400193075}, {"0": 14.0, "1": 1.263567835036204}, {"0": 15.0, "1": 1.2815975529900088}, {"0": 16.0, "1": 1.299439824803275}, {"0": 17.0, "1": 1.3168757154359425}, {"0": 18.0, "1": 1.3343874934934168}, {"0": 19.0, "1": 1.351039740496107}, {"0": 20.0, "1": 1.3676166982793703}, {"0": 21.0, "1": 1.3842125314468827}, {"0": 22.0, "1": 1.400807921453688}, {"0": 23.0, "1": 1.416748713155206}, {"0": 24.0, "1": 1.4325431350905102}, {"0": 25.0, "1": 1.4481730809665256}, {"0": 26.0, "1": 1.4639326081421193}, {"0": 27.0, "1": 1.479480820344132}, {"0": 28.0, "1": 1.4946609183628063}, {"0": 29.0, "1": 1.5096723391311508}, {"0": 30.0, "1": 1.5245596506212855}]

unit_tests/ebeam_test/phys_proc/ref_results/test_ibs3.json

@@ -0,0 +1 @@
+[{"0": 1.0, "1": 0.9997366834643983}, {"0": 2.0, "1": 1.0128116957460225}, {"0": 3.0, "1": 1.0255280772289488}, {"0": 4.0, "1": 1.0380603533558583}, {"0": 5.0, "1": 1.0502672024469313}, {"0": 6.0, "1": 1.0623469884300116}, {"0": 7.0, "1": 1.074269895880912}, {"0": 8.0, "1": 1.0860745441892916}, {"0": 9.0, "1": 1.0977684088851027}, {"0": 10.0, "1": 1.1096538520734391}, {"0": 11.0, "1": 1.1213326628111062}, {"0": 12.0, "1": 1.1330678189367918}, {"0": 13.0, "1": 1.144540165744393}, {"0": 14.0, "1": 1.1559297948264675}, {"0": 15.0, "1": 1.1672267365600004}, {"0": 16.0, "1": 1.1781902814931788}, {"0": 17.0, "1": 1.1891458979140237}, {"0": 18.0, "1": 1.19967018490185}, {"0": 19.0, "1": 1.2105303479325875}, {"0": 20.0, "1": 1.2209897685203805}, {"0": 21.0, "1": 1.2315624567594723}, {"0": 22.0, "1": 1.2420008342560112}, {"0": 23.0, "1": 1.2522537838660226}, {"0": 24.0, "1": 1.2626586925321777}, {"0": 25.0, "1": 1.272933815016061}, {"0": 26.0, "1": 1.2830779924354359}, {"0": 27.0, "1": 1.2933493056384662}, {"0": 28.0, "1": 1.3032151019868898}, {"0": 29.0, "1": 1.313314500259722}, {"0": 30.0, "1": 1.3230406742204555}]