2D Particle-in-cell solver

- Added a subpackage for a 2D Poisson solver with open boundary conditions, the returned field are normalised by the total deposited charge
- Some of the code (grid2d.f90) is written in fortran and has to be compiled with f2py to be used
- Added *.so compiled files to .gitignore

.gitignore

@@ -2,3 +2,4 @@
 *.ipynb_checkpoints*
 *.hdf5
 *.pyc
+*.so
\ No newline at end of file

mbtrack2/__init__.py

@@ -2,6 +2,7 @@
 __version__ = "0.7.0"
 from mbtrack2.impedance import *
 from mbtrack2.instability import *
+from mbtrack2.pic import *
 from mbtrack2.tracking import *
 from mbtrack2.utilities import *
 

mbtrack2/pic/__init__.py

+from mbtrack2.pic.grid2d import deposit_cic, interpolate_electric_field
+from mbtrack2.pic.pic2d import PICSolver2D
\ No newline at end of file

mbtrack2/pic/grid2d.f90

+subroutine deposit_cic(nx, ny, x_min, y_min, dx, dy, &
+    x, y, q, n_particles, rho)
+    implicit none
+
+    integer, intent(in) :: nx, ny       ! Number of grid points in x and y directions
+    real(8), intent(in) :: x_min, y_min ! Grid origin
+    real(8), intent(in) :: dx, dy       ! Grid spacing
+    real(8), intent(in) :: q ! Charge per macroparticle
+    integer, intent(in) :: n_particles  ! Number of particles
+    real(8), intent(in) :: x(n_particles), y(n_particles) ! Particle positions
+
+    real(8), intent(inout) :: rho(nx, ny)  ! Charge density array
+
+    integer :: i, i_grid, j_grid
+    real(8) :: wx, wy
+    do i = 1, n_particles
+
+        i_grid = int((x(i) - x_min) / dx)
+        j_grid = int((y(i) - y_min) / dy)
+
+        wx = ((x(i) - x_min) / dx) - dble(i_grid)
+        wy = ((y(i) - y_min) / dy) - dble(j_grid)
+
+        if (i_grid >= 0 .and. i_grid < nx-1 .and. j_grid >= 0 .and. j_grid < ny-1) then
+            rho(i_grid+1, j_grid+1) = rho(i_grid+1, j_grid+1) + q * (1.0d0 - wx) * (1.0d0 - wy)
+            rho(i_grid+2, j_grid+1) = rho(i_grid+2, j_grid+1) + q * wx * (1.0d0 - wy)
+            rho(i_grid+1, j_grid+2) = rho(i_grid+1, j_grid+2) + q * (1.0d0 - wx) * wy
+            rho(i_grid+2, j_grid+2) = rho(i_grid+2, j_grid+2) + q * wx * wy
+        end if
+    end do
+
+end subroutine deposit_cic
+
+subroutine interpolate_electric_field(nx, ny, n_particles, x_min, x_max, y_min, y_max, &
+    x_particles, y_particles, Ex_grid, Ey_grid, &
+    Ex_interp, Ey_interp)
+    implicit none
+
+    ! Input parameters
+    integer, intent(in) :: nx, ny, n_particles  ! Number of grid points and particles
+    real(8), intent(in) :: x_min, x_max, y_min, y_max  ! Domain bounds
+    real(8), intent(in) :: x_particles(n_particles), y_particles(n_particles)  ! Particle positions
+    real(8), intent(in) :: Ex_grid(nx, ny), Ey_grid(nx, ny)  ! Grid electric fields
+
+    ! Output parameters
+    real(8), intent(out) :: Ex_interp(n_particles), Ey_interp(n_particles)  ! Interpolated fields
+
+    ! Local variables
+    integer :: i, j, px
+    real(8) :: dx, dy, x_rel, y_rel, wx, wy, ex, ey
+
+    ! Calculate grid spacing
+    dx = (x_max - x_min) / dble(nx - 1)
+    dy = (y_max - y_min) / dble(ny - 1)
+
+    ! Loop over each particle to interpolate electric fields
+    do px = 1, n_particles
+        ! Find the indices of the lower-left grid point
+        i = int((x_particles(px) - x_min) / dx)
+        j = int((y_particles(px) - y_min) / dy)
+
+        ! Relative positions within the cell
+        x_rel = (x_particles(px) - (x_min + i * dx)) / dx
+        y_rel = (y_particles(px) - (y_min + j * dy)) / dy
+
+        ! Calculate weights for bilinear interpolation
+        wx = x_rel
+        wy = y_rel
+
+        ! Check bounds and perform bilinear interpolation
+        if (i >= 0 .and. i < nx-1 .and. j >= 0 .and. j < ny-1) then
+            ex = (1.0d0 - wx) * (1.0d0 - wy) * Ex_grid(i+1, j+1) + &
+            wx * (1.0d0 - wy) * Ex_grid(i+2, j+1) + &
+            (1.0d0 - wx) * wy * Ex_grid(i+1, j+2) + &
+            wx * wy * Ex_grid(i+2, j+2)
+
+            ey = (1.0d0 - wx) * (1.0d0 - wy) * Ey_grid(i+1, j+1) + &
+            wx * (1.0d0 - wy) * Ey_grid(i+2, j+1) + &
+            (1.0d0 - wx) * wy * Ey_grid(i+1, j+2) + &
+            wx * wy * Ey_grid(i+2, j+2)
+
+            Ex_interp(px) = ex
+            Ey_interp(px) = ey
+        else
+            Ex_interp(px) = 0.0d0
+            Ey_interp(px) = 0.0d0
+        end if
+    end do
+
+end subroutine interpolate_electric_field

mbtrack2/pic/pic2d.py

+import numpy as np
+from scipy.constants import epsilon_0
+
+import mbtrack2.pic.grid2d as grid2d
+
+
+class PICSolver2D:
+
+    def __init__(self, grid_size, cell_size):
+        # Define grid size and domain bounds based on the distribution's standard deviations
+        self.x_min, self.x_max, self.y_min, self.y_max = grid_size
+        self.dx, self.dy = cell_size
+
+        self.nx, self.ny = map(
+            lambda delta, d: int(np.floor(delta / d)),
+            [self.x_max - self.x_min, self.y_max - self.y_min],
+            [self.dx, self.dy])
+
+        self.rho = np.zeros((self.nx, self.ny), order='F', dtype=np.float64)
+        self.phi = np.zeros((self.nx, self.ny), dtype=np.float64)
+
+    def deposit_cic(self, x, y, q):
+        '''Implemented in fortran pic2d.f90'''
+        grid2d.deposit_cic(self.x_min, self.y_min, self.dx, self.dy, x, y, q,
+                           self.rho)
+        self.rho /= (epsilon_0 * self.dx * self.dy * self.rho.sum())
+
+
+# Original Python implementation (only to quickly check that fortran version is much faster before the merge request)
+#     def deposit_cic(self, bunch):
+#         for i in range(len(bunch)):
+#             x, y = bunch['x'][i], bunch['y'][i]
+#             q = bunch.charge_per_mp
+
+#             # Calculate the grid indices of the particle's position relative to the grid's origin
+#             i_grid = int((x - self.x_min) / self.dx)
+#             j_grid = int((y - self.y_min) / self.dy)
+
+#             # Fractional distance of the particle within the cell
+#             wx = ((x - self.x_min) / self.dx) - i_grid
+#             wy = ((y - self.y_min) / self.dy) - j_grid
+
+#             # Distribute the charge to the 4 nearest grid points using CIC interpolation
+#             if 0 <= i_grid < self.nx - 1 and 0 <= j_grid < self.ny - 1:
+#                 self.rho[i_grid, j_grid] += q * (1 - wx) * (1 - wy)
+#                 self.rho[i_grid + 1, j_grid] += q * wx * (1 - wy)
+#                 self.rho[i_grid, j_grid + 1] += q * (1 - wx) * wy
+#                 self.rho[i_grid + 1, j_grid + 1] += q * wx * wy
+
+    def reset_deposit(self):
+        self.rho = np.zeros((self.nx, self.ny), order='F', dtype=np.float64)
+
+    def poisson_solver_2d_fft(self):
+        """
+        Solves the 2D Poisson equation with open boundary conditions using FFT.
+        """
+
+        # Compute the Fourier transform of the charge density
+        rho_hat = np.fft.fft2(self.rho, s=(self.nx, self.ny))
+
+        # Create the wave numbers for the grid
+        kx = np.fft.fftfreq(self.nx, d=self.dx) * 2 * np.pi
+        ky = np.fft.fftfreq(self.ny, d=self.dy) * 2 * np.pi
+        kx, ky = np.meshgrid(kx, ky, indexing='ij')
+        # Compute the squared wave numbers
+        k_squared = kx*kx + ky*ky
+
+        # Avoid division by zero by setting zero frequency to a small value (pseudo-inverse)
+        k_squared[0, 0] = 1.0  # Prevents division by zero at (kx, ky) = (0, 0)
+
+        # Solve for the potential in Fourier space
+        phi_hat = rho_hat / k_squared
+
+        # Set the zero frequency component to zero (remove the mean to enforce open boundary conditions)
+        phi_hat[0, 0] = 0.0
+
+        # Transform back to real space
+        self.phi = np.fft.ifft2(phi_hat, s=(self.nx, self.ny)).real
+
+    def calculate_electric_field(self):
+        # Compute the electric field components from the potential
+        E_x, E_y = np.gradient(
+            -self.phi, self.dx,
+            self.dy)  # Negative gradient for the electric field
+        return E_x, E_y
+
+    def interpolate_electric_field(self, x_particles, y_particles):
+        """ Interpolates the electric field at particle positions. """
+        E_x, E_y = self.calculate_electric_field()
+        # Interpolation with scipy griddata is orders of magnitude slower than fortran implementation
+        E_x_interp, E_y_interp = grid2d.interpolate_electric_field(
+            self.x_min, self.x_max, self.y_min, self.y_max, x_particles,
+            y_particles, E_x, E_y)
+        return E_x_interp, E_y_interp