Merge branch 'develop' into prepare_v0.8.0

mbtrack2/tracking/element.py

@@ -444,6 +444,10 @@ def transverse_map_sector_generator(ring, positions):
     """
     N_sec = len(positions)
     sectors = []
+    if hasattr(ring, "adts") and ring.adts is not None:
+        adts = np.array([val / N_sec for val in ring.adts])
+    else:
+        adts = None
     if ring.optics.use_local_values:
         for i in range(N_sec):
             sectors.append(
@@ -455,49 +459,47 @@ def transverse_map_sector_generator(ring, positions):
                                     ring.optics.local_beta,
                                     ring.optics.local_dispersion,
                                     2 * np.pi * ring.tune / N_sec,
-                                    ring.chro / N_sec,
-                                    adts=ring.adts /
-                                    N_sec if ring.adts else None))
+                                    np.asarray(ring.chro) / N_sec,
+                                    adts=adts))
     else:
         import at
 
-        def _compute_chro(ring, pos, dp=1e-2, order=4):
+        def _compute_chro(ring, N_sec, dp=1e-2, order=4):
             lat = deepcopy(ring.optics.lattice)
             lat.append(at.Marker("END"))
-            fit, dpa, tune = at.physics.nonlinear.chromaticity(lat,
+            fit, _, _ = at.physics.nonlinear.chromaticity(lat,
                                                           method='linopt',
                                                           dpm=dp,
                                                           n_points=100,
                                                           order=order)
             Chrox, Chroy = fit[0, 1:], fit[1, 1:]
-            chrox = np.interp(pos, s, Chrox)
-            chroy = np.interp(pos, s, Chroy)
-
+            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 = _compute_chro(ring, N_sec)
         for i in range(N_sec):
             alpha0 = ring.optics.alpha(positions[i])
             beta0 = ring.optics.beta(positions[i])
             dispersion0 = ring.optics.dispersion(positions[i])
             mu0 = ring.optics.mu(positions[i])
-            chro0 = _compute_chro(ring, positions[i])
+            chro0 = _chro[:, i]
             if i != (N_sec - 1):
                 alpha1 = ring.optics.alpha(positions[i + 1])
                 beta1 = ring.optics.beta(positions[i + 1])
                 dispersion1 = ring.optics.dispersion(positions[i + 1])
                 mu1 = ring.optics.mu(positions[i + 1])
-                chro1 = _compute_chro(ring, positions[i + 1])
+                chro1 = _chro[:, i + 1]
             else:
                 alpha1 = ring.optics.alpha(positions[0])
                 beta1 = ring.optics.beta(positions[0])
                 dispersion1 = ring.optics.dispersion(positions[0])
                 mu1 = ring.optics.mu(ring.L)
-                chro1 = _compute_chro(ring, ring.L)
+                chro1 = _chro[:, -1]
             phase_diff = mu1 - mu0
             chro_diff = chro1 - chro0
             sectors.append(
-                TransverseMapSector(
-                    ring,
+                TransverseMapSector(ring,
                                     alpha0,
                                     beta0,
                                     dispersion0,
@@ -506,5 +508,5 @@ def transverse_map_sector_generator(ring, positions):
                                     dispersion1,
                                     phase_diff,
                                     chro_diff,
-                ))
+                                    adts=adts))
     return sectors

mbtrack2/tracking/particles.py

@@ -687,6 +687,8 @@ class Beam:
         Initialize beam with a given filling pattern and marco-particle number
         per bunch. Then initialize the different bunches with a 6D gaussian
         phase space.
+    init_bunch_list_mpi(bunch, filling_pattern)
+        Initialize a beam using MPI parallelisation with a Bunch per core.
     mpi_init()
         Switch on MPI parallelisation and initialise a Mpi object
     mpi_gather()
@@ -714,6 +716,8 @@ class Beam:
                 raise ValueError(("The length of the bunch list is {} ".format(
                     len(bunch_list)) + "but should be {}".format(self.ring.h)))
             self.bunch_list = bunch_list
+            self.update_filling_pattern()
+            self.update_distance_between_bunches()
 
     def __len__(self):
         """Return the number of (not empty) bunches"""
@@ -864,6 +868,40 @@ class Beam:
             for bunch in self.not_empty:
                 bunch.init_gaussian()
 
+    def init_bunch_list_mpi(self, bunch, filling_pattern):
+        """
+        Initialize a beam using MPI parallelisation with a Bunch per core.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+            The bunch given should probably depend on the mpi.rank so that each 
+            core can track a different bunch.
+            Example: beam.init_bunch_list_mpi(bunch_list[comm.rank], filling_pattern)
+        filling_pattern : array-like of bool of length ring.h
+            Filling pattern of the beam as a list or an array of bool.
+
+        """
+        filling_pattern = np.array(filling_pattern)
+
+        if len(filling_pattern) != self.ring.h:
+            raise ValueError(("The length of filling pattern is {} ".format(
+                len(filling_pattern)) +
+                              "but should be {}".format(self.ring.h)))
+
+        if filling_pattern.dtype != np.dtype("bool"):
+            raise TypeError("dtype {} should be bool.".format(
+                filling_pattern.dtype))
+
+        self.bunch_list = [
+            Bunch(self.ring, mp_number=1, alive=filling_pattern[i])
+            for i in range(self.ring.h)
+        ]
+        self.update_filling_pattern()
+        self.update_distance_between_bunches()
+        self.mpi_init()
+        self[self.mpi.bunch_num] = bunch
+
     def update_filling_pattern(self):
         """Update the beam filling pattern."""
         filling_pattern = []

tests/unit/tracking/test_element.py

@@ -374,3 +374,54 @@ class TestTransverseMap:
             assert not np.array_equal(initial_values[attr], current[attr])
             assert not np.array_equal(initial_values[attr], old[attr])
             np.testing.assert_allclose(current[attr], old[attr])
+
+class TestTransverseMapSectorGenerator:
+
+    # Generate sectors with local optics values when use_local_values is True
+    def test_local_optics_values(self, demo_ring):
+        positions = np.array([0, 25, 50, 75])
+        sectors = transverse_map_sector_generator(demo_ring, positions)
+    
+        assert len(sectors) == len(positions)
+        for sector in sectors:
+            assert np.array_equal(sector.alpha0, demo_ring.optics.local_alpha)
+            assert np.array_equal(sector.beta0, demo_ring.optics.local_beta)
+            assert np.array_equal(sector.dispersion0, demo_ring.optics.local_dispersion)
+
+    # Generate sectors with AT lattice optics when use_local_values is False
+    def test_at_lattice_optics(self, ring_with_at_lattice):
+        positions = np.array([0, 25, 50, 75])
+        sectors = transverse_map_sector_generator(ring_with_at_lattice, positions)
+    
+        assert len(sectors) == len(positions)
+        for i, sector in enumerate(sectors):
+            assert np.array_equal(sector.alpha0, ring_with_at_lattice.optics.alpha(positions[i]))
+            assert np.array_equal(sector.beta0, ring_with_at_lattice.optics.beta(positions[i]))
+
+    # Calculate phase differences between consecutive positions
+    def test_phase_differences(self, ring_with_at_lattice):
+        positions = np.array([0, 25, 50])
+        ring_with_at_lattice.optics.use_local_values = False
+        sectors = transverse_map_sector_generator(ring_with_at_lattice, positions)
+    
+        for i, sector in enumerate(sectors):
+            if i < len(positions)-1:
+                expected_phase = ring_with_at_lattice.optics.mu(positions[i+1]) - ring_with_at_lattice.optics.mu(positions[i])
+                assert np.allclose(sector.tune_diff * (2 * np.pi), expected_phase)
+
+    # Compute chromaticity differences between sectors
+    def test_chromaticity_differences(self, demo_ring):
+        positions = np.array([0, 50])
+        sectors = transverse_map_sector_generator(demo_ring, positions)
+    
+        expected_chro = np.asarray(demo_ring.chro) / len(positions)
+        assert np.allclose(sectors[0].chro_diff, expected_chro)
+
+    # Handle ADTS parameters correctly when provided
+    def test_adts_handling(self, demo_ring):
+        demo_ring.adts = np.array([1.0, 1.0, 1.0, 1.0])
+        positions = np.array([0, 50])
+        sectors = transverse_map_sector_generator(demo_ring, positions)
+    
+        expected_adts = demo_ring.adts / len(positions)
+        assert np.allclose(sectors[0].adts_poly[0].coef, expected_adts[0])
\ No newline at end of file

tests/unit/tracking/test_particle.py

@@ -305,3 +305,24 @@ class TestBeam:
         beam.mpi_close()
         assert not beam.mpi_switch
         assert beam.mpi is None
+        
+    def test_init_bunch_list(self, demo_ring):
+        filling_pattern = np.ones((demo_ring.h,), dtype=bool)
+        filling_pattern[5] = False
+        filling_pattern[8:10] = False
+        bunch_list = [Bunch(demo_ring, mp_number=1, alive=filling_pattern[i]) for i in range(demo_ring.h)]
+        beam = Beam(demo_ring, bunch_list)
+        assert len(beam) == filling_pattern.sum()
+        np.testing.assert_array_equal(beam.filling_pattern, filling_pattern)
+        assert beam.distance_between_bunches is not None
+        
+    def test_init_bunch_list_mpi(self, demo_ring, generate_bunch):
+        filling_pattern = np.zeros((demo_ring.h,), dtype=bool)
+        filling_pattern[0] = True
+        beam = Beam(demo_ring)
+        beam.init_bunch_list_mpi(generate_bunch(), filling_pattern)
+        
+        assert len(beam) == 1
+        np.testing.assert_array_equal(beam.filling_pattern, filling_pattern)
+        assert beam.distance_between_bunches[0] == demo_ring.h
+        
\ No newline at end of file