Merge branch 'develop' of gitlab.synchrotron-soleil.fr:PA/collective-effects/mbtrack2 into develop

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

tests/unit/tracking/test_beam_ion_effects.py

@@ -25,8 +25,10 @@ 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)
+        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)
@@ -37,11 +39,15 @@ 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)
 
     # 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)
+        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"]))
@@ -50,6 +56,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):
@@ -58,14 +66,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, small_bunch):
@@ -73,9 +81,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):
@@ -118,8 +126,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
@@ -199,11 +207,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'):
         
@@ -216,11 +219,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
@@ -239,9 +237,15 @@ class TestBeamIonElement:
                              [('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):
-        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"])
+        charge_gen = beam_ion.ion_beam.charge
+        
+        # loose half of electron bunch
+        small_bunch.alive[0:5] = False
+        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'), 
@@ -288,15 +292,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'):
@@ -314,12 +323,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.track(beam_ion.ion_beam)
     
         assert len(beam_ion.ion_beam["x"]) == 0
 
@@ -332,3 +343,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
+        
\ No newline at end of file