Added a PIC-solver option to BeamIonElement() and TransverseSpaceCharge() classes.

mbtrack2/tracking/beam_ion_effects.py

@@ -16,6 +16,7 @@ import numpy as np
 from numpy.random import choice, normal, uniform
 from scipy.constants import c, e
 
+from mbtrack2.pic.pic2d import PICSolver2D
 from mbtrack2.tracking.aperture import ElipticalAperture
 from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.emfields import _efieldn_mit, get_displaced_efield
@@ -416,7 +417,7 @@ class BeamIonElement(Element):
         self.ion_field_model = ion_field_model
         self.ion_element_length = ion_element_length
         self.generate_method = generate_method
-        self.n_ion_macroparticles_per_bunch = 30
+        self.n_ion_macroparticles_per_bunch = n_ion_macroparticles_per_bunch
         self.ion_beam_monitor_name = ion_beam_monitor_name
         self.ion_beam = IonParticles(
             mp_number=1,
@@ -559,28 +560,14 @@ class BeamIonElement(Element):
                                               sb_stdx, sb_stdy, sb_mx, sb_my)
 
         elif field_model == "PIC":
-            from PyPIC import FFT_OpenBoundary
-            from PyPIC import geom_impact_ellip as ellipse
-            qe = e
-            Dx = 0.1 * sb_stdx
-            Dy = 0.1 * sb_stdy
-            x_aper = 10 * sb_stdx
-            y_aper = 10 * sb_stdy
-            chamber = ellipse.ellip_cham_geom_object(x_aper=x_aper,
-                                                     y_aper=y_aper)
-            picFFT = FFT_OpenBoundary.FFT_OpenBoundary(
-                x_aper=chamber.x_aper,
-                y_aper=chamber.y_aper,
-                dx=Dx,
-                dy=Dy,
-                fftlib="pyfftw",
-            )
-            nel_part = 0 * second_beam["x"] + 1.0
-            picFFT.scatter(second_beam["x"], second_beam["y"], nel_part)
-            picFFT.solve()
-            en_x, en_y = picFFT.gather(first_beam["x"], first_beam["y"])
-            en_x /= qe * second_beam["x"].shape[0]
-            en_y /= qe * second_beam["x"].shape[0]
+            solver = PICSolver2D(grid_size=(-10 * sb_stdx, 10 * sb_stdx,
+                                            -10 * sb_stdy, 10 * sb_stdy),
+                                 cell_size=(0.1 * sb_stdx, 0.1 * sb_stdy))
+            solver.poisson_solver_2d_igf()
+            en_x, en_y = solver.interpolate_electric_field(
+                first_beam["x"], first_beam["y"])
+        else:
+            raise ValueError("Field model is not defined correctly.")
         return en_x, en_y
 
     def _get_new_beam_momentum(self,

mbtrack2/tracking/spacecharge.py

@@ -3,6 +3,7 @@ Module for transverse space charge calculations.
 """
 from scipy.constants import c
 
+from mbtrack2.pic.pic2d import PICSolver2D
 from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.emfields import _efieldn_mit, get_displaced_efield
 
@@ -41,7 +42,7 @@ class TransverseSpaceCharge(Element):
     ratio_threshold = 1e-3
     absolute_threshold = 1e-10
 
-    def __init__(self, ring, interaction_length, n_bins=100):
+    def __init__(self, ring, interaction_length, n_bins=100, method='BE'):
         """
         Initialize the SpaceCharge object.
 
@@ -58,6 +59,12 @@ class TransverseSpaceCharge(Element):
         self.n_bins = n_bins
         self.interaction_length = interaction_length
         self.efieldn = _efieldn_mit
+        self.method = method
+        if method == "PIC":
+            self.solver = PICSolver2D(
+                grid_size=(-10 * ring.sigma()[0], 10 * ring.sigma()[0],
+                           -10 * ring.sigma()[1], 10 * ring.sigma()[1]),
+                cell_size=(0.1 * ring.sigma()[0], 0.1 * ring.sigma()[1]))
 
     @Element.parallel
     def track(self, bunch):
@@ -72,8 +79,7 @@ class TransverseSpaceCharge(Element):
         """
         prefactor = self.interaction_length / (self.ring.E0 *
                                                self.ring.gamma**2)
-        (bins, sorted_index, profile,
-         center) = bunch.binning(n_bin=self.n_bins)
+        (bins, sorted_index, profile, _) = bunch.binning(n_bin=self.n_bins)
         dz = (bins[1] - bins[0]) * c
         charge_density = bunch.charge_per_mp * profile / dz
         for bin_index in range(self.n_bins - 1):
@@ -86,12 +92,24 @@ class TransverseSpaceCharge(Element):
             if len(x) != 0 and len(y) != 0:
                 mean_x, std_x = x.mean(), x.std()
                 mean_y, std_y = y.mean(), y.std()
-
+                if self.method == "BE":
                     en_x, en_y = get_displaced_efield(self.efieldn,
                                                       bunch['x'][particle_ids],
                                                       bunch['y'][particle_ids],
-                                                  std_x, std_y, mean_x, mean_y)
-
+                                                      std_x, std_y, mean_x,
+                                                      mean_y)
+                elif self.method == "PIC":
+                    self.solver.update_grid_igf(
+                        grid_size=(-10 * std_x, 10 * std_x, -10 * std_y,
+                                   10 * std_y),
+                        cell_size=(0.1 * std_x, 0.1 * std_y))
+                    self.solver.poisson_solver_2d_igf()
+                    en_x, en_y = self.solver.interpolate_electric_field(
+                        bunch['x'][particle_ids], bunch['y'][particle_ids])
+                else:
+                    raise ValueError(
+                        "Method for finding electic field is not defined properly. Should be 'BE' or 'PIC'."
+                    )
                 kicks_x = prefactor * en_x * charge_density[bin_index]
                 kicks_y = prefactor * en_y * charge_density[bin_index]
                 bunch['xp'][particle_ids] += kicks_x