Merge branch 'develop' into 'fast_wakepotential'

# Conflicts:
#   tests/unit/tracking/test_beam_ion_effects.py
#   tests/unit/tracking/test_synchrotron.py

CITATION.cff

+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

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

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,

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))

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:
@@ -436,6 +436,8 @@ def transverse_map_sector_generator(ring, positions):
         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
     -------
     sectors : list
@@ -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])

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.

tests/unit/tracking/test_beam_ion_effects.py

@@ -31,38 +31,30 @@ 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"]))
@@ -71,6 +63,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):
@@ -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,10 +89,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):
@@ -148,8 +139,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
@@ -221,8 +212,8 @@ class TestElipticalAperture:
 
 @pytest.fixture
 def generate_beam_ion(demo_ring):
-
-    def generate(ion_mass=1.67e-27,
+    def generate(
+            ion_mass=1.67e-27,
             ion_charge=1.6e-19,
             ionization_cross_section=1e-22,
             residual_gas_density=1e50,
@@ -230,11 +221,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'):
         
@@ -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
@@ -273,13 +254,17 @@ class TestBeamIonElement:
     @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):
+    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 np.isclose(beam_ion.ion_beam.charge, charge_gen*1.5)
+
     @pytest.mark.parametrize('ion_field_model, electron_field_model',
                              [('weak', 'weak'), ('weak', 'strong'),
                               ('strong', 'weak'), ('strong', '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,13 +353,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.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
 
@@ -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
+        

tests/unit/tracking/test_synchrotron.py

@@ -143,6 +143,10 @@ class TestSynchrotron:
         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