diff --git a/mbtrack2/tracking/beam_ion_effects.py b/mbtrack2/tracking/beam_ion_effects.py
index a041fee621da8f0249c792213eb0bf5e28ea3901..717b55fefa29fa83120199404be083bca89a9195 100644
--- a/mbtrack2/tracking/beam_ion_effects.py
+++ b/mbtrack2/tracking/beam_ion_effects.py
@@ -56,16 +56,16 @@ class IonMonitor(Monitor, metaclass=ABCMeta):
def __init__(self, save_every, buffer_size, total_size, file_name=None):
group_name = "IonData_" + str(next(self._n_monitors))
dict_buffer = {
- "mean": (6, buffer_size),
- "std": (6, buffer_size),
+ "mean": (4, buffer_size),
+ "std": (4, buffer_size),
"charge": (buffer_size, ),
- "charge_per_mp": (buffer_size, ),
+ "mp_number": (buffer_size, ),
}
dict_file = {
- "mean": (6, total_size),
- "std": (6, total_size),
+ "mean": (4, total_size),
+ "std": (4, total_size),
"charge": (total_size, ),
- "charge_per_mp": (buffer_size, ),
+ "mp_number": (buffer_size, ),
}
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name)
@@ -242,6 +242,7 @@ class IonParticles(Bunch):
"yp": np.zeros((mp_number, ), dtype=np.float64),
"tau": np.zeros((mp_number, ), dtype=np.float64),
"delta": np.zeros((mp_number, ), dtype=np.float64),
+ "charge": np.zeros((mp_number, ), dtype=np.float64),
}
self.charge_per_mp = 0
@@ -253,8 +254,61 @@ class IonParticles(Bunch):
@mp_number.setter
def mp_number(self, value):
self._mp_number = int(value)
+
+ @property
+ def charge(self):
+ """Bunch charge in [C]"""
+ return self["charge"].sum()
+
+ def _mean_weighted(self, coord):
+ """
+ Return the mean position of alive particles for each coordinates.
+ """
+ if self.charge == 0:
+ return np.zeros_like(coord, dtype=float)
+ else:
+ mean = [[np.average(self[name], weights=self["charge"])] for name in coord]
+ return np.squeeze(np.array(mean))
- def generate_as_a_distribution(self, electron_bunch):
+ def _mean_std_weighted(self, coord):
+ """
+ Return the standard deviation of the position of alive
+ particles for each coordinates.
+ """
+ if self.charge == 0:
+ return np.zeros_like(coord, dtype=float), np.zeros_like(coord, dtype=float)
+ else:
+ mean = [[np.average(self[name], weights=self["charge"])] for name in coord]
+ var = []
+ for i, name in enumerate(coord):
+ var.append(np.average((self[name]-mean[i])**2, weights=self["charge"]))
+ var = np.array(var)
+ return np.squeeze(np.array(mean)), np.squeeze(np.sqrt(var))
+
+ @property
+ def mean(self):
+ coord = ["x","xp","y","yp"]
+ return self._mean_weighted(coord)
+
+ @property
+ def mean_weighted(self):
+ coord = ["x","y"]
+ mean = self._mean_weighted(coord)
+ return mean[0], mean[1]
+
+ @property
+ def std(self):
+ coord = ["x","xp","y","yp"]
+ _, std = self._mean_std_weighted(coord)
+ return std
+
+ @property
+ def mean_std_weighted(self):
+ coord = ["x","y"]
+ mean, std = self._mean_std_weighted(coord)
+ return mean[0], mean[1], std[0], std[1]
+
+ def generate_as_a_distribution(self, electron_bunch, charge):
"""
Generates the particle positions based on a normal distribution, taking distribution parameters from an electron bunch.
@@ -289,8 +343,9 @@ class IonParticles(Bunch):
high=self.ion_element_length / c,
size=self.mp_number,
)
+ self["charge"] = np.ones((self.mp_number, )) * charge / self.mp_number
- def generate_from_random_samples(self, electron_bunch):
+ def generate_from_random_samples(self, electron_bunch, charge):
"""
Generates the particle positions and times based on random samples from electron positions in the bunch.
@@ -316,6 +371,7 @@ class IonParticles(Bunch):
high=self.ion_element_length / c,
size=self.mp_number,
)
+ self["charge"] = np.ones((self.mp_number, )) * charge / self.mp_number
def __add__(self, new_particles):
self.mp_number += new_particles.mp_number
@@ -514,7 +570,7 @@ class BeamIonElement(Element):
def _get_efields(self, first_beam, second_beam, field_model):
"""
- Calculates the electromagnetic field of the first beam acting on the second beam for a given field model.
+ Calculates the electromagnetic field of the second beam acting on the first beam for a given field model.
Parameters
----------
@@ -536,14 +592,17 @@ class BeamIonElement(Element):
raise ValueError(
f"The implementation for required beam-ion interaction model {field_model} is not implemented"
)
- 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 isinstance(second_beam, IonParticles):
+ sb_mx, sb_my, sb_stdx, sb_stdy = second_beam.mean_std_weighted
+ 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,
@@ -556,10 +615,13 @@ class BeamIonElement(Element):
)
elif field_model == "strong":
- fb_mx, fb_my = (
- first_beam["x"].mean(),
- first_beam["y"].mean(),
- )
+ if isinstance(first_beam, IonParticles):
+ fb_mx, fb_my = first_beam.mean_weighted
+ 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)
@@ -647,18 +709,19 @@ class BeamIonElement(Element):
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)
+ electron_bunch=electron_bunch,
+ charge=new_ion_charge)
elif self.generate_method == 'samples':
new_ion_particles.generate_from_random_samples(
- electron_bunch=electron_bunch)
+ electron_bunch=electron_bunch,
+ charge=new_ion_charge)
self.ion_beam += new_ion_particles
- self.ion_beam.charge_per_mp = (electron_bunch.charge *
- self.ionization_cross_section *
- self.residual_gas_density *
- self.ion_element_length /
- self.n_ion_macroparticles_per_bunch)
@parallel
def track(self, electron_bunch):
diff --git a/tests/unit/tracking/test_beam_ion_effects.py b/tests/unit/tracking/test_beam_ion_effects.py
index 368d063b0ac15aeb93aba8a52e24ed574f59a5fa..8a556efe27e65fa8c77396aa39a139cfe4235122 100644
--- a/tests/unit/tracking/test_beam_ion_effects.py
+++ b/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,7 +56,9 @@ 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):
ions1 = generate_ion_particles(mp_number=100)
@@ -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