diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index 169959bc..9e66c0ef 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -10,6 +10,7 @@ from . import initialization from . import utils from . import boris +from . import pusher from . import evolve from .solvers import vector_potential @@ -32,4 +33,4 @@ from .diagnostics import vtk from .boundary_conditions import boundaryconditions -from .boundary_conditions import grid_and_stencil \ No newline at end of file +from .boundary_conditions import grid_and_stencil diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index a4d7754c..f30aae27 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -50,10 +50,10 @@ def run_PyPIC3D(config_file): ##################################### INITIALIZE SIMULATION ################################################ - loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic = initialize_simulation(config_file) + loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic, particle_pusher = initialize_simulation(config_file) # initialize the simulation - jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic')) + jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher')) dt = world['dt'] output_dir = simulation_parameters['output_dir'] @@ -156,6 +156,7 @@ def run_PyPIC3D(config_file): J_func, solver, relativistic=relativistic, + particle_pusher=particle_pusher, ) # time loop to update the particles and fields diff --git a/PyPIC3D/boris.py b/PyPIC3D/boris.py index a23aeec1..69112043 100644 --- a/PyPIC3D/boris.py +++ b/PyPIC3D/boris.py @@ -1,313 +1,6 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -from PyPIC3D.boundary_conditions.grid_and_stencil import ( - BC_PERIODIC, - axis_has_active_cells, - collapse_axis_stencil, - prepare_particle_axis_stencil, +from PyPIC3D.pusher.boris import ( + boris_single_particle, + interpolate_field_to_particles, + relativistic_boris_single_particle, ) -from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights - -@jit -def particle_push(particles, E, B, grid, staggered_grid, dt, constants, periodic=True, relativistic=True): - """ - Updates the velocities of particles using the Boris algorithm. - - Args: - particles (Particles): The particles to be updated. - E (tuple): Electric field components (Ex, Ey, Ez). - B (tuple): Magnetic field components (Bx, By, Bz). - grid (Grid): The grid on which the fields are defined. - staggered_grid (Grid): The staggered grid for field interpolation. - dt (float): The time step for the update. - constants (dict): Dictionary containing physical constants. Must include: - - 'C': Speed of light in vacuum (m/s). - periodic (tuple): A tuple of three booleans indicating whether each dimension is periodic. - Default is (True, True, True) for fully periodic domains. - - Returns: - Particles: The particles with updated velocities. - """ - q = particles.get_charge() - m = particles.get_mass() - x, y, z = particles.get_forward_position() - vx, vy, vz = particles.get_velocity() - # get the charge, mass, position, and velocity of the particles - - shape_factor = particles.get_shape() - # get the shape factor of the particles - - ################## INTERPOLATION GRIDS ################################### - Ex_grid = staggered_grid[0], grid[1], grid[2] - Ey_grid = grid[0], staggered_grid[1], grid[2] - Ez_grid = grid[0], grid[1], staggered_grid[2] - # create the grids for the electric field components - Bx_grid = grid[0], staggered_grid[1], staggered_grid[2] - By_grid = staggered_grid[0], grid[1], staggered_grid[2] - Bz_grid = staggered_grid[0], staggered_grid[1], grid[2] - # create the staggered grids for the magnetic field components - ########################################################################## - - - ################## INTERPOLATE FIELDS TO PARTICLE POSITIONS ############## - Ex, Ey, Ez = E - # unpack the electric field components - efield_atx = interpolate_field_to_particles(Ex, x, y, z, Ex_grid, shape_factor, ghost_cells=True) - efield_aty = interpolate_field_to_particles(Ey, x, y, z, Ey_grid, shape_factor, ghost_cells=True) - efield_atz = interpolate_field_to_particles(Ez, x, y, z, Ez_grid, shape_factor, ghost_cells=True) - # calculate the electric field at the particle positions on the Yee-staggered component grids - Bx, By, Bz = B - # unpack the magnetic field components - bfield_atx = interpolate_field_to_particles(Bx, x, y, z, Bx_grid, shape_factor, ghost_cells=True) - bfield_aty = interpolate_field_to_particles(By, x, y, z, By_grid, shape_factor, ghost_cells=True) - bfield_atz = interpolate_field_to_particles(Bz, x, y, z, Bz_grid, shape_factor, ghost_cells=True) - # calculate the magnetic field at the particle positions on the Yee-staggered component grids - ######################################################################### - - - #################### BORIS ALGORITHM #################################### - if jnp.ndim(q) == 0: - boris_vmap = jax.vmap( - boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), - ) - relativistic_boris_vmap = jax.vmap( - relativistic_boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), - ) - else: - boris_vmap = jax.vmap( - boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), - ) - relativistic_boris_vmap = jax.vmap( - relativistic_boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), - ) - # vectorize the Boris algorithm for batch processing - - newvx, newvy, newvz = jax.lax.cond( - relativistic == True, - lambda _: relativistic_boris_vmap(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants), - lambda _: boris_vmap(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants), - operand=None - ) - # apply the Boris algorithm to update the velocities of the particles - ######################################################################### - - - particles.set_velocity(newvx, newvy, newvz) - # set the new velocities of the particles - return particles - -@jit -def boris_single_particle(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants): - """ - Updates the velocity of a single particle using the Boris algorithm. - Args: - x (float): Initial x position of the particle. - y (float): Initial y position of the particle. - z (float): Initial z position of the particle. - vx (float): Initial x component of the particle's velocity. - vy (float): Initial y component of the particle's velocity. - vz (float): Initial z component of the particle's velocity. - efield_atx (float): x component of the electric field at the particle's position. - efield_aty (float): y component of the electric field at the particle's position. - efield_atz (float): z component of the electric field at the particle's position. - bfield_atx (float): x component of the magnetic field at the particle's position. - bfield_aty (float): y component of the magnetic field at the particle's position. - bfield_atz (float): z component of the magnetic field at the particle's position. - q (float): Charge of the particle. - m (float): Mass of the particle. - dt (float): Time step for the update. - Returns: - tuple: Updated velocity components (vx, vy, vz) of the particle. - """ - - v = jnp.array([vx, vy, vz]) - # convert v into an array - - vminus = v + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) - # get v minus vector - - t = q*dt/(2*m)*jnp.array([bfield_atx, bfield_aty, bfield_atz]) - # calculate the t vector - vprime = vminus + jnp.cross(vminus, t) - # calculate the v prime vector - - s = 2*t / (1 + t[0]**2 + t[1]**2 + t[2]**2) - # calculate the s vector - vplus = vminus + jnp.cross(vprime, s) - # calculate the v plus vector - - newv = vplus + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) - # calculate the new velocity - return newv[0], newv[1], newv[2] - - - -@jit -def relativistic_boris_single_particle(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants): - """ - Perform a single step of the relativistic Boris algorithm for a charged particle. - - This function calculates the updated velocity of a charged particle under the influence - of electric and magnetic fields using the relativistic Boris algorithm. The algorithm - ensures energy conservation and is widely used in particle-in-cell (PIC) simulations. - - Args: - vx (float): Initial velocity of the particle in the x-direction (m/s). - vy (float): Initial velocity of the particle in the y-direction (m/s). - vz (float): Initial velocity of the particle in the z-direction (m/s). - efield_atx (float): Electric field at the particle's position in the x-direction (V/m). - efield_aty (float): Electric field at the particle's position in the y-direction (V/m). - efield_atz (float): Electric field at the particle's position in the z-direction (V/m). - bfield_atx (float): Magnetic field at the particle's position in the x-direction (T). - bfield_aty (float): Magnetic field at the particle's position in the y-direction (T). - bfield_atz (float): Magnetic field at the particle's position in the z-direction (T). - q (float): Charge of the particle (Coulombs). - m (float): Mass of the particle (kg). - dt (float): Time step for the simulation (seconds). - constants (dict): Dictionary containing physical constants. Must include: - - 'C': Speed of light in vacuum (m/s). - - Returns: - tuple: Updated velocity components of the particle in the x, y, and z directions (m/s). - """ - - - C = constants['C'] - # speed of light - - v = jnp.array([vx, vy, vz]) - # convert v into an array - - gamma = 1 / jnp.sqrt( 1 - ( (v[0]**2 + v[1]**2 + v[2]**2) / C**2 ) ) - # define the gamma factor - - uminus = v * gamma + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) - # get the u minus vector - - gamma_minus = jnp.sqrt( 1 + ( (uminus[0]**2 + uminus[1]**2 + uminus[2]**2) / C**2 ) ) - # define the gamma minus factor - - t = q*dt/(2*m)*jnp.array([bfield_atx, bfield_aty, bfield_atz]) / gamma_minus - # calculate the t vector - uprime = uminus + jnp.cross(uminus, t) - # calculate the u prime vector - - s = 2*t / (1 + t[0]**2 + t[1]**2 + t[2]**2) - # calculate the s vector - uplus = uminus + jnp.cross(uprime, s) - # calculate the u plus vector - - newu = uplus + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) - # calculate the new velocity - - new_gamma = jnp.sqrt( 1 + ( (newu[0]**2 + newu[1]**2 + newu[2]**2) / C**2 ) ) - # define the new gamma factor - - newv = newu / new_gamma - # convert the relativistic velocity back to the lab frame - - return newv[0], newv[1], newv[2] - - -@partial(jit, static_argnames=("ghost_cells",)) -def interpolate_field_to_particles(field, x, y, z, grid, shape_factor, ghost_cells=False): - """ - Interpolate a Yee-grid field component to particle positions using PIC shape functions. - - Args: - field (ndarray): Field component values on the supplied component grid. - x, y, z (ndarray): Particle coordinates. - grid (tuple): Component grid coordinates (x_grid, y_grid, z_grid). - shape_factor (int): Particle shape order (1 -> first order, 2 -> second order). - - Returns: - ndarray: Interpolated field value at each particle position. - """ - x_grid, y_grid, z_grid = grid - Nx = len(x_grid) - Ny = len(y_grid) - Nz = len(z_grid) - # grid point counts for each direction - x_active = axis_has_active_cells(Nx, ghost_cells=ghost_cells) - y_active = axis_has_active_cells(Ny, ghost_cells=ghost_cells) - z_active = axis_has_active_cells(Nz, ghost_cells=ghost_cells) - - _, _, deltax, xpts = prepare_particle_axis_stencil( - x, x_grid, Nx, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells - ) - _, _, deltay, ypts = prepare_particle_axis_stencil( - y, y_grid, Ny, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells - ) - _, _, deltaz, zpts = prepare_particle_axis_stencil( - z, z_grid, Nz, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells - ) - dx = x_grid[1] - x_grid[0] if Nx > 1 else 1.0 - dy = y_grid[1] - y_grid[0] if Ny > 1 else 1.0 - dz = z_grid[1] - z_grid[0] if Nz > 1 else 1.0 - # determine the distance from the closest grid nodes - - x_weights, y_weights, z_weights = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(deltax, deltay, deltaz, dx, dy, dz), - lambda _: get_second_order_weights(deltax, deltay, deltaz, dx, dy, dz), - operand=None, - ) - x_weights = jnp.asarray(x_weights) - y_weights = jnp.asarray(y_weights) - z_weights = jnp.asarray(z_weights) - # compute the shape function weights for the particles and convert them to arrays - - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - - # Keep full shape-factor computation but collapse inactive axes to an - # effective stencil size of 1 to avoid redundant interpolation work. - if x_active: - xpts_eff = xpts - x_weights_eff = x_weights - else: - xpts_eff, x_weights_eff = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=ghost_cells) - - if y_active: - ypts_eff = ypts - y_weights_eff = y_weights - else: - ypts_eff, y_weights_eff = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=ghost_cells) - - if z_active: - zpts_eff = zpts - z_weights_eff = z_weights - else: - zpts_eff, z_weights_eff = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=ghost_cells) - - def stencil_contribution(stencil_idx): - i, j, k = stencil_idx - return ( - field[xpts_eff[i, ...], ypts_eff[j, ...], zpts_eff[k, ...]] - * x_weights_eff[i, ...] - * y_weights_eff[j, ...] - * z_weights_eff[k, ...] - ) - # define a function to compute the contribution from each point in the effective stencil - - ii, jj, kk = jnp.meshgrid( - jnp.arange(xpts_eff.shape[0]), - jnp.arange(ypts_eff.shape[0]), - jnp.arange(zpts_eff.shape[0]), - indexing="ij", - ) - stencil_indicies = jnp.stack([ii.ravel(), jj.ravel(), kk.ravel()], axis=1) - # build effective stencil indices with shape (Sx*Sy*Sz, 3) - - interpolated_field = jnp.sum(jax.vmap(stencil_contribution)(stencil_indicies), axis=0) - # sum the contributions from all stencil points to get the final interpolated field value at each particle position - - return interpolated_field +from PyPIC3D.pusher.particle_push import particle_push diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index ed4f9aa1..f375106c 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -77,6 +77,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): x_active = axis_has_active_cells(Nx, ghost_cells=True) y_active = axis_has_active_cells(Ny, ghost_cells=True) z_active = axis_has_active_cells(Nz, ghost_cells=True) + # determine which directions are not quasi-1D. for species in particles: q = species.get_charge() @@ -84,18 +85,22 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): vx, vy, vz = species.get_velocity() shape_factor = species.get_shape() N_particles = species.get_number_of_particles() + # get particle information old_x = x - vx * dt old_y = y - vy * dt old_z = z - vz * dt + # compute old positions by backstepping along the velocity x0 = compute_particle_anchor(x, grid[0], shape_factor) y0 = compute_particle_anchor(y, grid[1], shape_factor) z0 = compute_particle_anchor(z, grid[2], shape_factor) + # compute the nearest grid point to the particle's current position old_x0 = compute_particle_anchor(old_x, grid[0], shape_factor) old_y0 = compute_particle_anchor(old_y, grid[1], shape_factor) old_z0 = compute_particle_anchor(old_z, grid[2], shape_factor) + # compute the nearest grid point to the particle's old position deltax = particle_axis_offset(x, x0, grid[0]) deltay = particle_axis_offset(y, y0, grid[1]) @@ -103,15 +108,19 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): old_deltax = particle_axis_offset(old_x, old_x0, grid[0]) old_deltay = particle_axis_offset(old_y, old_y0, grid[1]) old_deltaz = particle_axis_offset(old_z, old_z0, grid[2]) + # compute the particle's offset from the nearest grid point in each direction at both time steps shift_x = x0 - old_x0 shift_y = y0 - old_y0 shift_z = z0 - old_z0 + # compute how many grid points the particles have shifted in each direction. offsets = jnp.asarray([-2, -1, 0, 1, 2], dtype=x0.dtype) xpts = build_axis_stencil_points(x0, Nx, bc_x, offsets) ypts = build_axis_stencil_points(y0, Ny, bc_y, offsets) zpts = build_axis_stencil_points(z0, Nz, bc_z, offsets) + # build the 5-point Esirkepov stencil in each direction. + xpts = lax.cond( bc_x == BC_PERIODIC, lambda pts: _remap_staggered_periodic_ghosts(pts, x, Nx, world["x_wind"]), @@ -130,6 +139,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): lambda pts: pts, operand=zpts, ) + # remap out-of-domain stencil points into the ghost cells across the periodic seam, if needed. xw, yw, zw = jax.lax.cond( shape_factor == 1, @@ -143,6 +153,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): lambda _: get_second_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), operand=None, ) + # compute the particle shape weights at the current and old time steps tmp = jnp.zeros_like(xw[0]) @@ -153,13 +164,16 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): oxw = [tmp, oxw[0], oxw[1], oxw[2], tmp] oyw = [tmp, oyw[0], oyw[1], oyw[2], tmp] ozw = [tmp, ozw[0], ozw[1], ozw[2], tmp] + # pad the weights with zeros so they can be rolled. oxw = _roll_old_weights_to_new_frame(oxw, shift_x) oyw = _roll_old_weights_to_new_frame(oyw, shift_y) ozw = _roll_old_weights_to_new_frame(ozw, shift_z) + # shift the old weights into the new-cell frame so they can be differenced with the current weights. xpts, xw, oxw = _collapse_esirkepov_axis(xpts, xw, oxw, Nx) ypts, yw, oyw = _collapse_esirkepov_axis(ypts, yw, oyw, Ny) zpts, zw, ozw = _collapse_esirkepov_axis(zpts, zw, ozw, Nz) + # remove dummy stencil points for non 3D simulations if x_active and y_active and z_active: Wx_, Wy_, Wz_ = get_3D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles) @@ -187,6 +201,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=1) elif z_active and (not x_active) and (not y_active): Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=2) + # calculate the Esirkepov weights for the active dimensions dJx = jax.lax.cond( x_active, @@ -208,6 +223,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): lambda _: q * vz / (dx * dy * dz) * jnp.ones(N_particles), operand=None, ) + # compute the current contribution prefactors for each active dimension according to Esirkepov's formula Fx = dJx * Wx_ Fy = dJy * Wy_ @@ -252,6 +268,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): for j in range(5): for k in range(5): Jz = Jz.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Fz[i, j, k, :], mode="drop") + # deposit the current contributions onto the grid. Jx = fold_ghost_cells(Jx, bc_x, bc_y, bc_z) Jy = fold_ghost_cells(Jy, bc_x, bc_y, bc_z) @@ -260,6 +277,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): Jx = update_ghost_cells(Jx, bc_x, bc_y, bc_z) Jy = update_ghost_cells(Jy, bc_x, bc_y, bc_z) Jz = update_ghost_cells(Jz, bc_x, bc_y, bc_z) + # update ghost cells according to boundary conditions to prepare for field solve return (Jx, Jy, Jz) diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index cef9ee21..5355e8f3 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -2,7 +2,6 @@ from jax import jit import jax.numpy as jnp from functools import partial -from jax import lax from PyPIC3D.boundary_conditions.grid_and_stencil import ( axis_has_active_cells, @@ -14,21 +13,6 @@ from PyPIC3D.utils import bilinear_filter -def _remap_staggered_periodic_ghosts(points, position, axis_size, wind): - """Route out-of-domain staggered deposits into the ghost that folds across the seam.""" - points = jnp.where( - position[jnp.newaxis, ...] > 0.5 * wind, - jnp.where(points == axis_size - 1, 0, points), - points, - ) - points = jnp.where( - position[jnp.newaxis, ...] < -0.5 * wind, - jnp.where(points == 0, axis_size - 1, points), - points, - ) - return points - - @partial(jit, static_argnames=("filter",)) def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): """Compute current density (Jx,Jy,Jz) by depositing particle velocities.""" @@ -42,28 +26,31 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): Jx, Jy, Jz = J Nx, Ny, Nz = Jx.shape - x_active = axis_has_active_cells(Nx, ghost_cells=True) - y_active = axis_has_active_cells(Ny, ghost_cells=True) - z_active = axis_has_active_cells(Nz, ghost_cells=True) + # get the grid dimensions for the current density arrays bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] + # determine the boundary conditions for each axis to handle ghost cells correctly Jx = Jx.at[:, :, :].set(0) Jy = Jy.at[:, :, :].set(0) Jz = Jz.at[:, :, :].set(0) + # initialize current density arrays to zero before deposition for species in particles: shape_factor = species.get_shape() charge = species.get_charge() dq = charge / (dx * dy * dz) + # get the particle information needed for deposition, including charge and shape factor x, y, z = species.get_forward_position() vx, vy, vz = species.get_velocity() + # get the particle velocities and positions at the forward time step (t + dt) for deposition x = x - vx * world["dt"] / 2 y = y - vy * world["dt"] / 2 z = z - vz * world["dt"] / 2 + # half step back the particle positions to align with the time-centered current deposition scheme x, x0, deltax_node, xpts = prepare_particle_axis_stencil( x, @@ -74,6 +61,10 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): wind=world["x_wind"], ghost_cells=True, ) + # prepare the particle axis stencil for the x-axis, which includes determining the grid points that + # contribute to the deposition and the corresponding weights based on the particle's position and + # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. + y, y0, deltay_node, ypts = prepare_particle_axis_stencil( y, grid[1], @@ -82,7 +73,11 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): bc_y, wind=world["y_wind"], ghost_cells=True, - ) + ) + # prepare the particle axis stencil for the y-axis, which includes determining the grid points that + # contribute to the deposition and the corresponding weights based on the particle's position and + # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. + z, z0, deltaz_node, zpts = prepare_particle_axis_stencil( z, grid[2], @@ -91,11 +86,15 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): bc_z, wind=world["z_wind"], ghost_cells=True, - ) + ) + # prepare the particle axis stencil for the z-axis, which includes determining the grid points that + # contribute to the deposition and the corresponding weights based on the particle's position and + # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. deltax_face = (x - grid[0][0]) - (x0 + 0.5) * dx deltay_face = (y - grid[1][0]) - (y0 + 0.5) * dy deltaz_face = (z - grid[2][0]) - (z0 + 0.5) * dz + # compute the distances from the particles to the face-centered grid points. x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( shape_factor == 1, @@ -103,6 +102,7 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): lambda _: get_second_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), operand=None, ) + # compute the node-centered weights for deposition x_weights_face, y_weights_face, z_weights_face = jax.lax.cond( shape_factor == 1, @@ -110,28 +110,13 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): lambda _: get_second_order_weights(deltax_face, deltay_face, deltaz_face, dx, dy, dz), operand=None, ) + # compute the face-centered weights for deposition xpts = jnp.asarray(xpts) ypts = jnp.asarray(ypts) zpts = jnp.asarray(zpts) - xpts = jax.lax.cond( - bc_x == 0, - lambda pts: _remap_staggered_periodic_ghosts(pts, x, Nx, world["x_wind"]), - lambda pts: pts, - operand=xpts, - ) - ypts = jax.lax.cond( - bc_y == 0, - lambda pts: _remap_staggered_periodic_ghosts(pts, y, Ny, world["y_wind"]), - lambda pts: pts, - operand=ypts, - ) - zpts = jax.lax.cond( - bc_z == 0, - lambda pts: _remap_staggered_periodic_ghosts(pts, z, Nz, world["z_wind"]), - lambda pts: pts, - operand=zpts, - ) + # Boundary ownership stays in fold_ghost_cells: deposition may write + # into ghost cells, and folding maps those deposits back to the domain. x_weights_face = jnp.asarray(x_weights_face) y_weights_face = jnp.asarray(y_weights_face) @@ -140,12 +125,15 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): x_weights_node = jnp.asarray(x_weights_node) y_weights_node = jnp.asarray(y_weights_node) z_weights_node = jnp.asarray(z_weights_node) + xpts, x_weights_node = collapse_axis_stencil(xpts, x_weights_node, Nx, ghost_cells=True) _, x_weights_face = collapse_axis_stencil(xpts, x_weights_face, Nx, ghost_cells=True) ypts, y_weights_node = collapse_axis_stencil(ypts, y_weights_node, Ny, ghost_cells=True) _, y_weights_face = collapse_axis_stencil(ypts, y_weights_face, Ny, ghost_cells=True) zpts, z_weights_node = collapse_axis_stencil(zpts, z_weights_node, Nz, ghost_cells=True) _, z_weights_face = collapse_axis_stencil(zpts, z_weights_face, Nz, ghost_cells=True) + # Collapse stencils for inactive (single-cell) axes so deposition uses the + # reduced effective stencil; ghost-cell folding is handled later. xpts_eff = xpts ypts_eff = ypts @@ -164,6 +152,7 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): indexing="ij", ) combos = jnp.stack([ii.ravel(), jj.ravel(), kk.ravel()], axis=1) + # get the list of indices to map over def idx_and_dJ_values(idx): i, j, k = idx @@ -191,10 +180,12 @@ def idx_and_dJ_values(idx): return ix, iy, iz, valx, valy, valz ix, iy, iz, valx, valy, valz = jax.vmap(idx_and_dJ_values)(combos) + # use vectorized mapping to compute the current depositions. Jx = Jx.at[(ix, iy, iz)].add(valx, mode="drop") Jy = Jy.at[(ix, iy, iz)].add(valy, mode="drop") Jz = Jz.at[(ix, iy, iz)].add(valz, mode="drop") + # deposit the currents onto the faces. Jx = fold_ghost_cells(Jx, bc_x, bc_y, bc_z) Jy = fold_ghost_cells(Jy, bc_x, bc_y, bc_z) @@ -217,8 +208,10 @@ def filter_func(J_, filter): Jx = filter_func(Jx, filter) Jy = filter_func(Jy, filter) Jz = filter_func(Jz, filter) + # apply any filters Jx = update_ghost_cells(Jx, bc_x, bc_y, bc_z) Jy = update_ghost_cells(Jy, bc_x, bc_y, bc_z) Jz = update_ghost_cells(Jz, bc_x, bc_y, bc_z) + # update the ghost cells return (Jx, Jy, Jz) diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 5cdacff5..6d308e40 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -6,7 +6,7 @@ from jax import jit from functools import partial -from PyPIC3D.boris import ( +from PyPIC3D.pusher.particle_push import ( particle_push ) @@ -22,13 +22,13 @@ E_from_A, B_from_A, update_vector_potential ) -@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic")) -def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True): +@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) +def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True, particle_pusher="boris"): """ Advances the simulation by one time step for an electrostatic Particle-In-Cell (PIC) loop. This function performs the following steps: - 1. Pushes all particles using the current electric and magnetic fields (typically with the Boris algorithm). + 1. Pushes all particles using the current electric and magnetic fields with the selected particle pusher. 2. Updates particle positions. 3. Solves for the new electric field using the Poisson equation, updating the field and potential. 4. Packs the updated fields and returns the new particle and field states. @@ -54,8 +54,8 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], E, B, vertex_grid, center_grid, world['dt'], constants, relativistic=relativistic) - # use boris push for particle velocities + particles[i] = particle_push(particles[i], E, B, vertex_grid, center_grid, world['dt'], constants, relativistic=relativistic, particle_pusher=particle_pusher) + # use the selected particle pusher for particle velocities particles[i].update_position() # update the particle positions @@ -74,12 +74,12 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu return particles, fields -@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic")) -def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True): +@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) +def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True, particle_pusher="boris"): """ Advance an electrodynamic Particle-In-Cell (PIC) system by one time step. This routine performs, in order: - 1) Particle push using the Boris algorithm (optionally relativistic), + 1) Particle push using the selected particle pusher, 2) Particle position update, 3) Current deposition onto the grid via the provided `J_func`, 4) Field update (E then B) using curl operators and boundary conditions, @@ -129,8 +129,8 @@ def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_f ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], E, B, center_grid, vertex_grid, world['dt'], constants, relativistic=relativistic) - # use boris push for particle velocities + particles[i] = particle_push(particles[i], E, B, center_grid, vertex_grid, world['dt'], constants, relativistic=relativistic, particle_pusher=particle_pusher) + # use the selected particle pusher for particle velocities particles[i].update_position() # update the particle positions @@ -154,8 +154,8 @@ def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_f return particles, fields -@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic")) -def time_loop_vector_potential(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True): +@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) +def time_loop_vector_potential(particles, fields, world, constants, curl_func, J_func, solver, relativistic=True, particle_pusher="boris"): """ Advance a PIC (Particle-In-Cell) simulation by one time step using a vector-potential formulation for the electromagnetic fields. @@ -217,8 +217,8 @@ def time_loop_vector_potential(particles, fields, world, constants, curl_func, J ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], E, B, center_grid, vertex_grid, world['dt'], constants, relativistic=relativistic) - # use boris push for particle velocities + particles[i] = particle_push(particles[i], E, B, center_grid, vertex_grid, world['dt'], constants, relativistic=relativistic, particle_pusher=particle_pusher) + # use the selected particle pusher for particle velocities particles[i].update_position() # update the particle positions diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 5b8a4df6..10d24bd0 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -48,6 +48,8 @@ time_loop_electrodynamic, time_loop_electrostatic, time_loop_vector_potential ) +from PyPIC3D.pusher.particle_push import validate_particle_pusher + from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.solvers.vector_potential import initialize_vector_potential @@ -121,6 +123,7 @@ def default_parameters(): "Nt": None, # number of time steps "electrostatic": False, # boolean for electrostatic simulation "relativistic": True, # boolean for relativistic simulation + "particle_pusher": "boris", # particle pusher: boris, higuera_cary "benchmark": False, # boolean for using the profiler "verbose": False, # boolean for printing verbose output "GPUs": False, # boolean for using GPUs @@ -202,6 +205,8 @@ def initialize_simulation(toml_file): electrostatic = simulation_parameters['electrostatic'] solver = simulation_parameters['solver'] relativistic = simulation_parameters['relativistic'] + particle_pusher = simulation_parameters['particle_pusher'] + validate_particle_pusher(particle_pusher) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] # set the simulation parameters @@ -363,6 +368,8 @@ def initialize_simulation(toml_file): print("Relativistic simulation") else: print("Non-relativistic simulation") + print(f"Using {particle_pusher} particle pusher") + # print the selected particle pusher if electrostatic: print("Using electrostatic solver") @@ -408,7 +415,7 @@ def initialize_simulation(toml_file): return evolve_loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, \ - solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic + solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic, particle_pusher diff --git a/PyPIC3D/pusher/__init__.py b/PyPIC3D/pusher/__init__.py new file mode 100644 index 00000000..6e579869 --- /dev/null +++ b/PyPIC3D/pusher/__init__.py @@ -0,0 +1,16 @@ +from PyPIC3D.pusher.boris import ( + boris_single_particle, + interpolate_field_to_particles, + relativistic_boris_single_particle, +) +from PyPIC3D.pusher.higuera_cary import ( + gamma_from_u, + higuera_cary_single_particle, + u_from_v, + velocity_from_u, +) +from PyPIC3D.pusher.particle_push import ( + VALID_PARTICLE_PUSHERS, + particle_push, + validate_particle_pusher, +) diff --git a/PyPIC3D/pusher/boris.py b/PyPIC3D/pusher/boris.py new file mode 100644 index 00000000..f06f4274 --- /dev/null +++ b/PyPIC3D/pusher/boris.py @@ -0,0 +1,221 @@ +import jax +from jax import jit +import jax.numpy as jnp +from functools import partial + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + BC_PERIODIC, + axis_has_active_cells, + collapse_axis_stencil, + prepare_particle_axis_stencil, +) +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights + +@jit +def boris_single_particle(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants): + """ + Updates the velocity of a single particle using the Boris algorithm. + Args: + x (float): Initial x position of the particle. + y (float): Initial y position of the particle. + z (float): Initial z position of the particle. + vx (float): Initial x component of the particle's velocity. + vy (float): Initial y component of the particle's velocity. + vz (float): Initial z component of the particle's velocity. + efield_atx (float): x component of the electric field at the particle's position. + efield_aty (float): y component of the electric field at the particle's position. + efield_atz (float): z component of the electric field at the particle's position. + bfield_atx (float): x component of the magnetic field at the particle's position. + bfield_aty (float): y component of the magnetic field at the particle's position. + bfield_atz (float): z component of the magnetic field at the particle's position. + q (float): Charge of the particle. + m (float): Mass of the particle. + dt (float): Time step for the update. + Returns: + tuple: Updated velocity components (vx, vy, vz) of the particle. + """ + + v = jnp.array([vx, vy, vz]) + # convert v into an array + + vminus = v + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) + # get v minus vector + + t = q*dt/(2*m)*jnp.array([bfield_atx, bfield_aty, bfield_atz]) + # calculate the t vector + vprime = vminus + jnp.cross(vminus, t) + # calculate the v prime vector + + s = 2*t / (1 + t[0]**2 + t[1]**2 + t[2]**2) + # calculate the s vector + vplus = vminus + jnp.cross(vprime, s) + # calculate the v plus vector + + newv = vplus + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) + # calculate the new velocity + return newv[0], newv[1], newv[2] + + + +@jit +def relativistic_boris_single_particle(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants): + """ + Perform a single step of the relativistic Boris algorithm for a charged particle. + + This function calculates the updated velocity of a charged particle under the influence + of electric and magnetic fields using the relativistic Boris algorithm. The algorithm + ensures energy conservation and is widely used in particle-in-cell (PIC) simulations. + + Args: + vx (float): Initial velocity of the particle in the x-direction (m/s). + vy (float): Initial velocity of the particle in the y-direction (m/s). + vz (float): Initial velocity of the particle in the z-direction (m/s). + efield_atx (float): Electric field at the particle's position in the x-direction (V/m). + efield_aty (float): Electric field at the particle's position in the y-direction (V/m). + efield_atz (float): Electric field at the particle's position in the z-direction (V/m). + bfield_atx (float): Magnetic field at the particle's position in the x-direction (T). + bfield_aty (float): Magnetic field at the particle's position in the y-direction (T). + bfield_atz (float): Magnetic field at the particle's position in the z-direction (T). + q (float): Charge of the particle (Coulombs). + m (float): Mass of the particle (kg). + dt (float): Time step for the simulation (seconds). + constants (dict): Dictionary containing physical constants. Must include: + - 'C': Speed of light in vacuum (m/s). + + Returns: + tuple: Updated velocity components of the particle in the x, y, and z directions (m/s). + """ + + + C = constants['C'] + # speed of light + + v = jnp.array([vx, vy, vz]) + # convert v into an array + + gamma = 1 / jnp.sqrt( 1 - ( (v[0]**2 + v[1]**2 + v[2]**2) / C**2 ) ) + # define the gamma factor + + uminus = v * gamma + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) + # get the u minus vector + + gamma_minus = jnp.sqrt( 1 + ( (uminus[0]**2 + uminus[1]**2 + uminus[2]**2) / C**2 ) ) + # define the gamma minus factor + + t = q*dt/(2*m)*jnp.array([bfield_atx, bfield_aty, bfield_atz]) / gamma_minus + # calculate the t vector + uprime = uminus + jnp.cross(uminus, t) + # calculate the u prime vector + + s = 2*t / (1 + t[0]**2 + t[1]**2 + t[2]**2) + # calculate the s vector + uplus = uminus + jnp.cross(uprime, s) + # calculate the u plus vector + + newu = uplus + q*dt/(2*m)*jnp.array([efield_atx, efield_aty, efield_atz]) + # calculate the new velocity + + new_gamma = jnp.sqrt( 1 + ( (newu[0]**2 + newu[1]**2 + newu[2]**2) / C**2 ) ) + # define the new gamma factor + + newv = newu / new_gamma + # convert the relativistic velocity back to the lab frame + + return newv[0], newv[1], newv[2] + + +@partial(jit, static_argnames=("ghost_cells",)) +def interpolate_field_to_particles(field, x, y, z, grid, shape_factor, ghost_cells=False): + """ + Interpolate a Yee-grid field component to particle positions using PIC shape functions. + + Args: + field (ndarray): Field component values on the supplied component grid. + x, y, z (ndarray): Particle coordinates. + grid (tuple): Component grid coordinates (x_grid, y_grid, z_grid). + shape_factor (int): Particle shape order (1 -> first order, 2 -> second order). + + Returns: + ndarray: Interpolated field value at each particle position. + """ + x_grid, y_grid, z_grid = grid + Nx = len(x_grid) + Ny = len(y_grid) + Nz = len(z_grid) + # grid point counts for each direction + x_active = axis_has_active_cells(Nx, ghost_cells=ghost_cells) + y_active = axis_has_active_cells(Ny, ghost_cells=ghost_cells) + z_active = axis_has_active_cells(Nz, ghost_cells=ghost_cells) + + _, _, deltax, xpts = prepare_particle_axis_stencil( + x, x_grid, Nx, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells + ) + _, _, deltay, ypts = prepare_particle_axis_stencil( + y, y_grid, Ny, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells + ) + _, _, deltaz, zpts = prepare_particle_axis_stencil( + z, z_grid, Nz, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells + ) + dx = x_grid[1] - x_grid[0] if Nx > 1 else 1.0 + dy = y_grid[1] - y_grid[0] if Ny > 1 else 1.0 + dz = z_grid[1] - z_grid[0] if Nz > 1 else 1.0 + # determine the distance from the closest grid nodes + + x_weights, y_weights, z_weights = jax.lax.cond( + shape_factor == 1, + lambda _: get_first_order_weights(deltax, deltay, deltaz, dx, dy, dz), + lambda _: get_second_order_weights(deltax, deltay, deltaz, dx, dy, dz), + operand=None, + ) + x_weights = jnp.asarray(x_weights) + y_weights = jnp.asarray(y_weights) + z_weights = jnp.asarray(z_weights) + # compute the shape function weights for the particles and convert them to arrays + + xpts = jnp.asarray(xpts) + ypts = jnp.asarray(ypts) + zpts = jnp.asarray(zpts) + + # Keep full shape-factor computation but collapse inactive axes to an + # effective stencil size of 1 to avoid redundant interpolation work. + if x_active: + xpts_eff = xpts + x_weights_eff = x_weights + else: + xpts_eff, x_weights_eff = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=ghost_cells) + + if y_active: + ypts_eff = ypts + y_weights_eff = y_weights + else: + ypts_eff, y_weights_eff = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=ghost_cells) + + if z_active: + zpts_eff = zpts + z_weights_eff = z_weights + else: + zpts_eff, z_weights_eff = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=ghost_cells) + + def stencil_contribution(stencil_idx): + i, j, k = stencil_idx + return ( + field[xpts_eff[i, ...], ypts_eff[j, ...], zpts_eff[k, ...]] + * x_weights_eff[i, ...] + * y_weights_eff[j, ...] + * z_weights_eff[k, ...] + ) + # define a function to compute the contribution from each point in the effective stencil + + ii, jj, kk = jnp.meshgrid( + jnp.arange(xpts_eff.shape[0]), + jnp.arange(ypts_eff.shape[0]), + jnp.arange(zpts_eff.shape[0]), + indexing="ij", + ) + stencil_indicies = jnp.stack([ii.ravel(), jj.ravel(), kk.ravel()], axis=1) + # build effective stencil indices with shape (Sx*Sy*Sz, 3) + + interpolated_field = jnp.sum(jax.vmap(stencil_contribution)(stencil_indicies), axis=0) + # sum the contributions from all stencil points to get the final interpolated field value at each particle position + + return interpolated_field diff --git a/PyPIC3D/pusher/higuera_cary.py b/PyPIC3D/pusher/higuera_cary.py new file mode 100644 index 00000000..bb9c0b7e --- /dev/null +++ b/PyPIC3D/pusher/higuera_cary.py @@ -0,0 +1,113 @@ +import jax.numpy as jnp +from jax import jit + + +@jit +def gamma_from_u(u, constants): + """ + Calculate gamma from the relativistic velocity-like momentum u = gamma * v. + """ + C = constants["C"] + # speed of light + + return jnp.sqrt(1.0 + jnp.dot(u, u) / C**2) + + +@jit +def velocity_from_u(u, constants): + """ + Convert the relativistic velocity-like momentum u = gamma * v to velocity. + """ + gamma = gamma_from_u(u, constants) + # compute gamma from u + + return u / gamma + + +@jit +def u_from_v(v, constants): + """ + Convert velocity to the relativistic velocity-like momentum u = gamma * v. + """ + C = constants["C"] + # speed of light + + gamma = 1.0 / jnp.sqrt(1.0 - jnp.dot(v, v) / C**2) + # compute gamma from velocity + + return gamma * v + + +@jit +def gamma_higuera_cary(u_e, beta, constants): + """ + Calculate the Higuera-Cary gamma used in the magnetic rotation. + """ + C = constants["C"] + # speed of light + + beta_squared = jnp.dot(beta, beta) + u_star = jnp.dot(u_e, beta) / C + sigma = gamma_from_u(u_e, constants) ** 2 - beta_squared + # compute the scalar quantities for the Higuera-Cary gamma expression + + return jnp.sqrt((sigma + jnp.sqrt(sigma**2 + 4.0 * (beta_squared + u_star**2))) / 2.0) + + +@jit +def higuera_cary_single_particle(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants): + """ + Updates the velocity of a single particle using the Higuera-Cary algorithm. + Args: + vx (float): Initial x component of the particle velocity. + vy (float): Initial y component of the particle velocity. + vz (float): Initial z component of the particle velocity. + efield_atx (float): x component of the electric field at the particle. + efield_aty (float): y component of the electric field at the particle. + efield_atz (float): z component of the electric field at the particle. + bfield_atx (float): x component of the magnetic field at the particle. + bfield_aty (float): y component of the magnetic field at the particle. + bfield_atz (float): z component of the magnetic field at the particle. + q (float): Charge of the particle. + m (float): Mass of the particle. + dt (float): Time step for the update. + constants (dict): Dictionary containing the speed of light. + Returns: + tuple: Updated velocity components (vx, vy, vz) of the particle. + """ + + v = jnp.array([vx, vy, vz]) + # convert v into an array + + u = u_from_v(v, constants) + # convert velocity into relativistic velocity-like momentum + + E = jnp.array([efield_atx, efield_aty, efield_atz]) + B = jnp.array([bfield_atx, bfield_aty, bfield_atz]) + # convert the fields into vectors + + half_charge_timestep = q * dt / (2.0 * m) + # compute the half time-step charge-to-mass factor + + epsilon = half_charge_timestep * E + beta = half_charge_timestep * B + # compute the electric and magnetic update vectors + + u_e = u + epsilon + # apply the first half electric-field kick + + gamma_next = gamma_higuera_cary(u_e, beta, constants) + t = beta / gamma_next + # calculate the Higuera-Cary rotation vector + + s = 1.0 / (1.0 + jnp.dot(t, t)) + u_m = s * (u_e + jnp.dot(u_e, t) * t + jnp.cross(u_e, t)) + # rotate the momentum using the Higuera-Cary magnetic update + + newu = u_m + epsilon + jnp.cross(u_m, t) + # apply the second electric-field kick and final rotation correction + + newv = velocity_from_u(newu, constants) + # convert the relativistic velocity-like momentum back to velocity + + return newv[0], newv[1], newv[2] diff --git a/PyPIC3D/pusher/particle_push.py b/PyPIC3D/pusher/particle_push.py new file mode 100644 index 00000000..fad6e554 --- /dev/null +++ b/PyPIC3D/pusher/particle_push.py @@ -0,0 +1,126 @@ +import jax +from jax import jit +import jax.numpy as jnp +from functools import partial + +from PyPIC3D.pusher.boris import ( + boris_single_particle, + interpolate_field_to_particles, + relativistic_boris_single_particle, +) +from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle + + +VALID_PARTICLE_PUSHERS = ("boris", "higuera_cary") + + +def validate_particle_pusher(particle_pusher): + """ + Validate the particle pusher name before entering jitted simulation code. + """ + if particle_pusher not in VALID_PARTICLE_PUSHERS: + raise ValueError(f"particle_pusher must be one of {VALID_PARTICLE_PUSHERS}, got {particle_pusher!r}") + + +@partial(jit, static_argnames=("periodic", "relativistic", "particle_pusher")) +def particle_push(particles, E, B, grid, staggered_grid, dt, constants, periodic=True, relativistic=True, particle_pusher="boris"): + """ + Updates particle velocities using the selected particle pusher. + + Args: + particles (Particles): The particles to be updated. + E (tuple): Electric field components (Ex, Ey, Ez). + B (tuple): Magnetic field components (Bx, By, Bz). + grid (Grid): The grid on which the fields are defined. + staggered_grid (Grid): The staggered grid for field interpolation. + dt (float): The time step for the update. + constants (dict): Dictionary containing physical constants. + periodic (bool): Kept for compatibility with the existing call signature. + relativistic (bool): Selects relativistic or non-relativistic Boris. + particle_pusher (str): Particle pusher name: "boris" or "higuera_cary". + + Returns: + Particles: The particles with updated velocities. + """ + q = particles.get_charge() + m = particles.get_mass() + x, y, z = particles.get_forward_position() + vx, vy, vz = particles.get_velocity() + # get the charge, mass, position, and velocity of the particles + + shape_factor = particles.get_shape() + # get the shape factor of the particles + + ################## INTERPOLATION GRIDS ################################### + Ex_grid = staggered_grid[0], grid[1], grid[2] + Ey_grid = grid[0], staggered_grid[1], grid[2] + Ez_grid = grid[0], grid[1], staggered_grid[2] + # create the grids for the electric field components + Bx_grid = grid[0], staggered_grid[1], staggered_grid[2] + By_grid = staggered_grid[0], grid[1], staggered_grid[2] + Bz_grid = staggered_grid[0], staggered_grid[1], grid[2] + # create the staggered grids for the magnetic field components + ########################################################################## + + ################## INTERPOLATE FIELDS TO PARTICLE POSITIONS ############## + Ex, Ey, Ez = E + # unpack the electric field components + efield_atx = interpolate_field_to_particles(Ex, x, y, z, Ex_grid, shape_factor, ghost_cells=True) + efield_aty = interpolate_field_to_particles(Ey, x, y, z, Ey_grid, shape_factor, ghost_cells=True) + efield_atz = interpolate_field_to_particles(Ez, x, y, z, Ez_grid, shape_factor, ghost_cells=True) + # calculate the electric field at the particle positions on the Yee-staggered component grids + Bx, By, Bz = B + # unpack the magnetic field components + bfield_atx = interpolate_field_to_particles(Bx, x, y, z, Bx_grid, shape_factor, ghost_cells=True) + bfield_aty = interpolate_field_to_particles(By, x, y, z, By_grid, shape_factor, ghost_cells=True) + bfield_atz = interpolate_field_to_particles(Bz, x, y, z, Bz_grid, shape_factor, ghost_cells=True) + # calculate the magnetic field at the particle positions on the Yee-staggered component grids + ######################################################################### + + #################### PARTICLE PUSHER #################################### + if jnp.ndim(q) == 0: + boris_vmap = jax.vmap( + boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), + ) + relativistic_boris_vmap = jax.vmap( + relativistic_boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), + ) + higuera_cary_vmap = jax.vmap( + higuera_cary_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), + ) + else: + boris_vmap = jax.vmap( + boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + relativistic_boris_vmap = jax.vmap( + relativistic_boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + higuera_cary_vmap = jax.vmap( + higuera_cary_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + # vectorize the selected particle pusher for batch processing + + if particle_pusher == "boris": + newvx, newvy, newvz = jax.lax.cond( + relativistic == True, + lambda _: relativistic_boris_vmap(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants), + lambda _: boris_vmap(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants), + operand=None + ) + # apply the Boris algorithm to update particle velocities + elif particle_pusher == "higuera_cary": + newvx, newvy, newvz = higuera_cary_vmap(vx, vy, vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, q, m, dt, constants) + # apply the Higuera-Cary algorithm to update particle velocities + else: + raise ValueError(f"Unknown particle_pusher: {particle_pusher}") + ######################################################################### + + particles.set_velocity(newvx, newvy, newvz) + # set the new velocities of the particles + return particles diff --git a/PyPIC3D/solvers/vector_potential.py b/PyPIC3D/solvers/vector_potential.py index 56a9528d..9dd72c11 100644 --- a/PyPIC3D/solvers/vector_potential.py +++ b/PyPIC3D/solvers/vector_potential.py @@ -7,7 +7,7 @@ centered_finite_difference_divergence, centered_finite_difference_gradient ) -from PyPIC3D.boris import interpolate_field_to_particles +from PyPIC3D.pusher.boris import interpolate_field_to_particles from PyPIC3D.utils import digital_filter @@ -320,4 +320,4 @@ def interpolate_field(field, grid, target_grid, interpolation_order=1): interp_field = interp_flat.reshape(X_target.shape) # Reshape back to the target grid shape - return interp_field \ No newline at end of file + return interp_field diff --git a/demos/convergence_testing/interpolation_convergence_test.py b/demos/convergence_testing/interpolation_convergence_test.py index ade24627..fb9fc90e 100644 --- a/demos/convergence_testing/interpolation_convergence_test.py +++ b/demos/convergence_testing/interpolation_convergence_test.py @@ -1,7 +1,7 @@ import jax.numpy as jnp from scipy import stats -from PyPIC3D.boris import interpolate_field_to_particles +from PyPIC3D.pusher.boris import interpolate_field_to_particles from functools import partial diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 81c38125..359addce 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -14,6 +14,7 @@ electrostatic = false cfl = 0.99 shape_factor = 2 current_calculation = "j_from_rhov" +fast_backend = "default" alpha = 1.0 relativistic = true diff --git a/pyproject.toml b/pyproject.toml index e283bb21..9979253a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,8 @@ dependencies = [ "openpmd-api", ] -[tool.setuptools] -packages = ["PyPIC3D"] +[tool.setuptools.packages.find] +include = ["PyPIC3D*"] [project.scripts] PyPIC3D = "PyPIC3D.__main__:main" diff --git a/tests/boris_test.py b/tests/boris_test.py index eff8b5e9..96c02714 100644 --- a/tests/boris_test.py +++ b/tests/boris_test.py @@ -6,7 +6,9 @@ from functools import partial -from PyPIC3D.boris import boris_single_particle, interpolate_field_to_particles +from PyPIC3D.pusher.particle_push import validate_particle_pusher +from PyPIC3D.pusher.boris import boris_single_particle, interpolate_field_to_particles +from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle from PyPIC3D.utils import mae, convergence_test jax.config.update("jax_enable_x64", True) @@ -140,6 +142,73 @@ def measure_xz_radius(xs, zs): #print(f"Average radius: {avg_radius}") self.assertTrue( jnp.isclose(avg_radius, 1.28, atol=0.5) ) + def test_higuera_cary_zero_field_preserves_velocity(self): + vx, vy, vz = higuera_cary_single_particle( + 0.2, -0.1, 0.05, + 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # zero fields should not change the particle velocity + self.assertTrue(jnp.allclose(jnp.array([vx, vy, vz]), jnp.array([0.2, -0.1, 0.05]))) + + def test_higuera_cary_magnetic_field_preserves_speed(self): + vx0, vy0, vz0 = 0.2, 0.0, 0.0 + vx, vy, vz = higuera_cary_single_particle( + vx0, vy0, vz0, + 0.0, 0.0, 0.0, + 0.0, 0.0, 1.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # magnetic rotation should conserve speed + old_speed = jnp.sqrt(vx0**2 + vy0**2 + vz0**2) + new_speed = jnp.sqrt(vx**2 + vy**2 + vz**2) + self.assertTrue(jnp.allclose(new_speed, old_speed, rtol=1e-12, atol=1e-12)) + + def test_higuera_cary_electric_field_accelerates_along_field(self): + vx, vy, vz = higuera_cary_single_particle( + 0.0, 0.0, 0.0, + 1.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # electric field in x should accelerate the particle in x only + self.assertGreater(vx, 0.0) + self.assertTrue(jnp.allclose(jnp.array([vy, vz]), jnp.array([0.0, 0.0]))) + + def test_higuera_cary_vectorizes_with_per_particle_charge_and_mass(self): + higuera_cary_vmap = jax.vmap( + higuera_cary_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + newvx, newvy, newvz = higuera_cary_vmap( + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([1.0, 1.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([1.0, 2.0]), + jnp.array([1.0, 1.0]), + 0.1, + {"C": 10.0}, + ) + # per-particle charge and mass should work through the vectorized path + self.assertEqual(newvx.shape, (2,)) + self.assertGreater(newvx[1], newvx[0]) + self.assertTrue(jnp.allclose(newvy, jnp.zeros(2))) + self.assertTrue(jnp.allclose(newvz, jnp.zeros(2))) + + def test_validate_particle_pusher_rejects_unknown_name(self): + with self.assertRaises(ValueError): + validate_particle_pusher("not_a_pusher") + def test_trilinear_convergence(self): diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 0f262eee..4665549a 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -20,6 +20,8 @@ def test_setup_write_dir(self): def test_default_parameters(self): plotting, sim, const = default_parameters() self.assertIn('Nx', sim) + self.assertIn('particle_pusher', sim) + self.assertEqual(sim['particle_pusher'], 'boris') self.assertIn('eps', const) self.assertIn('plotfields', plotting) # check that the default parameters contain expected keys diff --git a/tests/particle_test.py b/tests/particle_test.py index 82e69946..1b8a93ba 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -572,9 +572,10 @@ def test_J_from_rhov_periodic_boundary_crossing(self): Jx_int = J[0][1:-1, 1:-1, 1:-1] self.assertAlmostEqual(float(jnp.sum(Jx_int) * dx * dy * dz), 0.5) - # With opposite-side periodic fold, all current lands at the last - # interior face for this near-boundary particle. - self.assertGreater(float(jnp.sum(Jx_int[-1, :, :])), 0.0) + # The interior portion stays on the last physical face; only the + # +x ghost deposit folds to the first physical face. + self.assertAlmostEqual(float(jnp.sum(Jx_int[0, :, :]) * dx * dy * dz), 0.05) + self.assertAlmostEqual(float(jnp.sum(Jx_int[-1, :, :]) * dx * dy * dz), 0.45) def test_rho_reduced_dimensions_periodic(self): Nx, Ny, Nz = 16, 1, 1