Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion PyPIC3D/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from . import initialization
from . import utils
from . import boris
from . import pusher
from . import evolve

from .solvers import vector_potential
Expand All @@ -32,4 +33,4 @@
from .diagnostics import vtk

from .boundary_conditions import boundaryconditions
from .boundary_conditions import grid_and_stencil
from .boundary_conditions import grid_and_stencil
5 changes: 3 additions & 2 deletions PyPIC3D/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Expand Down Expand Up @@ -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

Expand Down
317 changes: 5 additions & 312 deletions PyPIC3D/boris.py
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading