Merge branch 'ions' into 'develop'

Ions

See merge request !18

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

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,7 +611,7 @@ class BeamIonElement(Element):
             The new y momentum of the first beam.
         """
 
-        en_x, en_y = BeamIonElement._get_efields(first_beam,
+        en_x, en_y = self._get_efields(first_beam,
                                        second_beam,
                                        field_model=field_model)
         kicks_x = prefactor * en_x

mbtrack2/tracking/particles_electromagnetic_fields.py → 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
 
 

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