diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py index 3d63d329266a485f79f29055309af85ed90d8ee7..b524c36efa79f9a9b875356fff1de16d63ec73a5 100644 --- a/mbtrack2/impedance/resistive_wall.py +++ b/mbtrack2/impedance/resistive_wall.py @@ -7,7 +7,7 @@ import numpy as np from scipy.constants import c, epsilon_0, mu_0 from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction -from mbtrack2.tracking.particles_electromagnetic_fields import _wofz +from mbtrack2.tracking.emfields import _wofz def skin_depth(frequency, rho, mu_r=1, epsilon_r=1): diff --git a/mbtrack2/tracking/beam_ion_effects.py b/mbtrack2/tracking/beam_ion_effects.py index 458aa1705a855d98de4ff8aa514f467d67616244..da354ba011dec6978d5b7d4d45a131a4abcba1d8 100644 --- a/mbtrack2/tracking/beam_ion_effects.py +++ b/mbtrack2/tracking/beam_ion_effects.py @@ -18,12 +18,9 @@ from scipy.constants import c, e from mbtrack2.tracking.aperture import ElipticalAperture from mbtrack2.tracking.element import Element +from mbtrack2.tracking.emfields import _efieldn_mit, get_displaced_efield from mbtrack2.tracking.monitors import Monitor from mbtrack2.tracking.particles import Beam, Bunch -from mbtrack2.tracking.particles_electromagnetic_fields import ( - _efieldn_mit, - get_displaced_efield, -) class IonMonitor(Monitor, metaclass=ABCMeta): @@ -356,6 +353,9 @@ class BeamIonElement(Element): The number of records in the built-in ion beam monitor. Should be number of turns times number of bunches because the monitor records every turn after each bunch passage. n_ion_macroparticles_per_bunch : int, optional The number of ion macroparticles generated per electron bunch passed. Defaults to 30. + ion_beam_monitor_name : str, optional + If provided, the name of the ion monitor output file. It must end with an extension '.hdf5'. + If None, no ion monitor file is generated. use_ion_phase_space_monitor : bool, optional Whether to use the ion phase space monitor. generate_method : str, optional @@ -431,7 +431,7 @@ class BeamIonElement(Element): if self.ion_beam_monitor_name: warnings.warn( - 'BeamIonMonitor.ion_beam_monitor_name.close() should be called at the end of tracking', + 'BeamIonMonitor.beam_monitor.close() should be called at the end of tracking', UserWarning, stacklevel=2) self.beam_monitor = IonMonitor( @@ -467,13 +467,13 @@ class BeamIonElement(Element): @wraps(track) def track_wrapper(*args, **kwargs): if isinstance(args[1], Beam): - if Beam.switch_mpi: + self = args[0] + beam = args[1] + if beam.mpi_switch: warnings.warn( 'Tracking through beam-ion element is performed sequentially. Bunches are not parallelized.', UserWarning, stacklevel=2) - self = args[0] - beam = args[1] for bunch in beam.bunch_list: track(self, bunch, *args[2:], **kwargs) else: @@ -611,9 +611,9 @@ class BeamIonElement(Element): The new y momentum of the first beam. """ - en_x, en_y = BeamIonElement._get_efields(first_beam, - second_beam, - field_model=field_model) + en_x, en_y = self._get_efields(first_beam, + second_beam, + field_model=field_model) kicks_x = prefactor * en_x kicks_y = prefactor * en_y new_xp = first_beam["xp"] + kicks_x diff --git a/mbtrack2/tracking/particles_electromagnetic_fields.py b/mbtrack2/tracking/emfields.py similarity index 98% rename from mbtrack2/tracking/particles_electromagnetic_fields.py rename to mbtrack2/tracking/emfields.py index 3780ebbe9193e3ae7f3cbcd4cd147381f2eb4821..a451471c4f7e66c41cd4f431b66c7a73fb86f407 100644 --- a/mbtrack2/tracking/particles_electromagnetic_fields.py +++ b/mbtrack2/tracking/emfields.py @@ -104,8 +104,8 @@ def efieldn_gauss_round(x, y, sig_x, sig_y): """ r_squared = x*x + y*y sig_r = sig_x - amplitude = (1 - np.exp(-r_squared / - (2*sig_r*sig_r))) / (2*pi*epsilon_0*r_squared) + amplitude = -np.expm1(-r_squared / + (2*sig_r*sig_r)) / (2*pi*epsilon_0*r_squared) return x * amplitude, y * amplitude diff --git a/mbtrack2/tracking/spacecharge.py b/mbtrack2/tracking/spacecharge.py index 3ecb779859145e733620a42b03b1e9d15daf8dca..661bc10d7445f164e08bede7c7b37105b996fb80 100644 --- a/mbtrack2/tracking/spacecharge.py +++ b/mbtrack2/tracking/spacecharge.py @@ -4,10 +4,7 @@ Module for transverse space charge calculations. from scipy.constants import c from mbtrack2.tracking.element import Element -from mbtrack2.tracking.particles_electromagnetic_fields import ( - _efieldn_mit, - get_displaced_efield, -) +from mbtrack2.tracking.emfields import _efieldn_mit, get_displaced_efield class TransverseSpaceCharge(Element):