diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 00000000..366ba48b --- /dev/null +++ b/.coveragerc @@ -0,0 +1,11 @@ +[run] +source = + PyPIC3D/PyPIC3D +omit = + PyPIC3D/tests/* + tests/* + */__pycache__/* + +[report] +skip_empty = True +show_missing = True diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 971b4dd9..733abc1a 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -8,7 +8,7 @@ on: jobs: test-coverage: - runs-on: ubuntu-latest + runs-on: self-hosted steps: - name: Checkout code uses: actions/checkout@v4 @@ -22,11 +22,11 @@ jobs: run: | python -m pip install --upgrade pip pip install . - pip install pytest coverage + pip install coverage - name: Run tests with coverage run: | - coverage run -m pytest + coverage run -m unittest tests/code_tests/*.py coverage xml - name: Upload coverage report diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index f1fe3bdc..728d65dd 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -13,7 +13,7 @@ permissions: jobs: build: - runs-on: ubuntu-latest + runs-on: self-hosted steps: - uses: actions/checkout@v4 @@ -27,4 +27,4 @@ jobs: pip install . - name: Test with unittest run: | - python -m unittest tests/*.py + python -m unittest tests/code_tests/*.py tests/physics_tests/*.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index 9e66c0ef..2abab663 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -6,21 +6,14 @@ import matplotlib.pyplot as plt # import external libraries -from . import errors from . import initialization from . import utils -from . import boris from . import pusher from . import evolve -from .solvers import vector_potential -from .solvers import pstd -from .solvers import fdtd -from .solvers import first_order_yee from .solvers import electrostatic_yee from .particles import particle_initialization -from .particles import species_class from .deposition import shapes from .deposition import Esirkepov @@ -30,7 +23,6 @@ from .diagnostics import plotting from .diagnostics import fluid_quantities from .diagnostics import openPMD -from .diagnostics import vtk from .boundary_conditions import boundaryconditions from .boundary_conditions import grid_and_stencil diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index f313f02a..773f3fe1 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -9,7 +9,6 @@ import time import jax from jax import block_until_ready -import jax.numpy as jnp from tqdm import tqdm #from memory_profiler import profile @@ -19,153 +18,192 @@ write_particles_phase_space, write_data ) -from PyPIC3D.diagnostics.openPMD import ( - write_openpmd_particles, write_openpmd_fields -) - -from PyPIC3D.diagnostics.vtk import ( - plot_field_slice_vtk, plot_vectorfield_slice_vtk, plot_vtk_particles +from PyPIC3D.diagnostics.async_writer import ( + create_async_tiled_openpmd_field_writer, + create_async_tiled_openpmd_particle_writer, + enqueue_openpmd_field_output, + enqueue_openpmd_particle_output, ) from PyPIC3D.utils import ( dump_parameters_to_toml, load_config_file, compute_energy, - setup_pmd_files, add_external_fields + setup_pmd_files, add_external_fields, compute_total_momentum ) from PyPIC3D.initialization import ( initialize_simulation ) -from PyPIC3D.diagnostics.fluid_quantities import ( - compute_mass_density -) - -from PyPIC3D.deposition.rho import compute_rho - # Importing functions from the PyPIC3D package ############################################################################################################ +def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): + if len(fields) < 8: + return + + overflow = fields[-1] + if bool(jax.device_get(overflow)): + raise RuntimeError("tiled particle tile capacity overflowed during periodic retile") + + 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, particle_pusher = initialize_simulation(config_file) + loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, solver, electrostatic, verbose, GPUs, Nt, relativistic, particle_pusher, species_config = initialize_simulation(config_file) # initialize the simulation - jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher')) - dt = world['dt'] output_dir = simulation_parameters['output_dir'] - vertex_grid = tuple(g[1:-1] for g in world['grids']['vertex']) - # unpack the physical interior of the vertex grid (strip ghost cell positions) + particle_species_names = simulation_parameters.get("particle_species_names") + + def loop_with_static_world( + particles, + species_config, + fields, + constants, + solver, + relativistic=True, + particle_pusher="boris", + ): + if electrostatic: + return loop( + particles, + species_config, + fields, + world, + constants, + solver, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + return loop( + particles, + species_config, + fields, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) - scalar_field_names = ["rho", "mass_density"] - vector_field_names = ["E", "B", "J"] + jit_loop = jax.jit( + loop_with_static_world, + static_argnames=('solver', 'relativistic', 'particle_pusher'), + ) E, B, J, rho, phi, external_fields, *rest = fields # unpack the fields total_E, total_B = add_external_fields(E, B, external_fields) # energy diagnostics use the fields seen by the particle pusher - e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants) + e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants, species_config=species_config) # Compute the energy of the system initial_energy = e_energy + b_energy + kinetic_energy - if plotting_parameters['plot_openpmd_fields']: setup_pmd_files( os.path.join(output_dir, "data"), "fields", ".h5") - if plotting_parameters['plot_openpmd_particles']: setup_pmd_files( os.path.join(output_dir, "data"), "particles", ".h5") + field_writer = None + particle_writer = None + if plotting_parameters['plot_openpmd_fields']: + setup_pmd_files( os.path.join(output_dir, "data"), "fields", ".h5") + field_writer = create_async_tiled_openpmd_field_writer( + world, + os.path.join(output_dir, "data"), + filename="fields", + file_extension=".h5", + queue_size=int(plotting_parameters.get("openpmd_field_queue_size", 2)), + ) + if plotting_parameters['plot_openpmd_particles']: + setup_pmd_files( os.path.join(output_dir, "data"), "particles", ".h5") + particle_writer = create_async_tiled_openpmd_particle_writer( + world, + constants, + os.path.join(output_dir, "data"), + filename="particles", + file_extension=".h5", + queue_size=int(plotting_parameters.get("openpmd_particle_queue_size", 2)), + ) # setup the openPMD files if needed ############################################################################################################ ###################################################### SIMULATION LOOP ##################################### - for t in tqdm(range(Nt)): - - # plot the data - if t % plotting_parameters['plotting_interval'] == 0: - - plot_num = t // plotting_parameters['plotting_interval'] - # determine the plot number - - E, B, J, rho, phi, external_fields, *rest = fields - # unpack the fields - - total_E, total_B = add_external_fields(E, B, external_fields) - # energy diagnostics use the fields seen by the particle pusher - e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants) - # Compute the energy of the system - write_data(f"{output_dir}/data/total_energy.txt", t * dt, e_energy + b_energy + kinetic_energy) - write_data(f"{output_dir}/data/energy_error.txt", t * dt, abs( initial_energy - (e_energy + b_energy + kinetic_energy)) / max(initial_energy, 1e-10)) - write_data(f"{output_dir}/data/electric_field_energy.txt", t * dt, e_energy) - write_data(f"{output_dir}/data/magnetic_field_energy.txt", t * dt, b_energy) - write_data(f"{output_dir}/data/kinetic_energy.txt", t * dt, kinetic_energy) - # Write the total energy to a file - total_momentum = sum(particle_species.momentum() for particle_species in particles) - # Total momentum of the particles - write_data(f"{output_dir}/data/total_momentum.txt", t * dt, total_momentum) - # Write the total momentum to a file - - # for species in particles: - # write_data(f"{output_dir}/data/{species.name}_kinetic_energy.txt", t * dt, species.kinetic_energy()) - - - if plotting_parameters['plot_phasespace']: - write_particles_phase_space(particles, t, output_dir) - - - - if plotting_parameters['plot_vtk_scalars']: - rho = compute_rho(particles, rho, world, constants) - # calculate the charge density based on the particle positions - mass_density = compute_mass_density(particles, rho, world) - # calculate the mass density based on the particle positions - - y_mid = world['Ny']//2 + 1 - # midplane index shifted by 1 for ghost cells - fields_mag = [rho[1:-1, y_mid, 1:-1], mass_density[1:-1, y_mid, 1:-1]] - plot_field_slice_vtk(fields_mag, scalar_field_names, 1, vertex_grid, t, "scalar_field", output_dir, world) - # Plot the scalar fields in VTK format - - - if plotting_parameters['plot_vtk_vectors']: - y_mid = world['Ny']//2 + 1 - # midplane index shifted by 1 for ghost cells - vector_field_slices = [ [E[0][1:-1, y_mid, 1:-1], E[1][1:-1, y_mid, 1:-1], E[2][1:-1, y_mid, 1:-1]], - [B[0][1:-1, y_mid, 1:-1], B[1][1:-1, y_mid, 1:-1], B[2][1:-1, y_mid, 1:-1]], - [J[0][1:-1, y_mid, 1:-1], J[1][1:-1, y_mid, 1:-1], J[2][1:-1, y_mid, 1:-1]]] - plot_vectorfield_slice_vtk(vector_field_slices, vector_field_names, 1, vertex_grid, t, 'vector_field', output_dir, world) - # Plot the vector fields in VTK format - - if plotting_parameters['plot_vtk_particles']: - plot_vtk_particles(particles, plot_num, output_dir) - # Plot the particles in VTK format - - if plotting_parameters['plot_openpmd_particles']: - write_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5") - # Write the particles in openPMD format - - if plotting_parameters['plot_openpmd_fields']: - write_openpmd_fields(fields, world, os.path.join(output_dir, "data"), plot_num, t, "fields", ".h5") - # Write the fields in openPMD format - - fields = (E, B, J, rho, phi, external_fields, *rest) - # repack the fields - - particles, fields = jit_loop( - particles, - fields, - world, - constants, - curl_func, - J_func, - solver, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - # time loop to update the particles and fields - - - return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world + loop_error = None + try: + for t in tqdm(range(Nt)): + + # plot the data + if t % plotting_parameters['plotting_interval'] == 0: + + plot_num = t // plotting_parameters['plotting_interval'] + # determine the plot number + + E, B, J, rho, phi, external_fields, *rest = fields + # unpack the fields + + total_E, total_B = add_external_fields(E, B, external_fields) + # energy diagnostics use the fields seen by the particle pusher + e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants, species_config=species_config) + # Compute the energy of the system + write_data(f"{output_dir}/data/total_energy.txt", t * dt, e_energy + b_energy + kinetic_energy) + write_data(f"{output_dir}/data/energy_error.txt", t * dt, abs( initial_energy - (e_energy + b_energy + kinetic_energy)) / max(initial_energy, 1e-10)) + write_data(f"{output_dir}/data/electric_field_energy.txt", t * dt, e_energy) + write_data(f"{output_dir}/data/magnetic_field_energy.txt", t * dt, b_energy) + write_data(f"{output_dir}/data/kinetic_energy.txt", t * dt, kinetic_energy) + # Write the total energy to a file + total_momentum = compute_total_momentum(particles, species_config=species_config) + # Total momentum of the particles + write_data(f"{output_dir}/data/total_momentum.txt", t * dt, total_momentum) + # Write the total momentum to a file + + if particle_writer is not None: + enqueue_openpmd_particle_output( + particle_writer, + particles, + world, + plot_num, + t, + species_config=species_config, + species_names=particle_species_names, + ) + # Queue tiled particle chunks for asynchronous openPMD output + + if field_writer is not None: + enqueue_openpmd_field_output(field_writer, fields, world, plot_num, t) + # Queue tiled field chunks for asynchronous openPMD output + + particles, fields = jit_loop( + particles, + species_config, + fields, + constants, + solver, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + # time loop to update the particles and fields + _raise_if_tiled_particles_overflowed(fields, simulation_parameters) + # fixed tile capacity overflow would silently drop particles; stop as + # soon as it is detected. + + except BaseException as exc: + loop_error = exc + raise + finally: + writer_error = None + for writer in (particle_writer, field_writer): + if writer is None: + continue + try: + writer.close(raise_errors=loop_error is None) + except BaseException as exc: + if writer_error is None: + writer_error = exc + if loop_error is None and writer_error is not None: + raise writer_error + + + return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world, species_config def main(): ###################### JAX SETTINGS ######################################################################## @@ -184,7 +222,7 @@ def main(): start = time.time() # start the timer - Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world = block_until_ready(run_PyPIC3D(toml_file)) + Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world, species_config = block_until_ready(run_PyPIC3D(toml_file)) # run the PyPIC3D simulation end = time.time() @@ -195,7 +233,7 @@ def main(): total_E, total_B = add_external_fields(E, B, external_fields) # energy diagnostics use the fields seen by the particle pusher - e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants) + e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants, species_config=species_config) # compute the energy of the system print(f"Final Electric Field Energy: {e_energy}") print(f"Final Magnetic Field Energy: {b_energy}") diff --git a/PyPIC3D/boris.py b/PyPIC3D/boris.py deleted file mode 100644 index 69112043..00000000 --- a/PyPIC3D/boris.py +++ /dev/null @@ -1,6 +0,0 @@ -from PyPIC3D.pusher.boris import ( - boris_single_particle, - interpolate_field_to_particles, - relativistic_boris_single_particle, -) -from PyPIC3D.pusher.particle_push import particle_push diff --git a/PyPIC3D/boundary_conditions/PML.py b/PyPIC3D/boundary_conditions/PML.py index 5f3572ab..87c7875c 100644 --- a/PyPIC3D/boundary_conditions/PML.py +++ b/PyPIC3D/boundary_conditions/PML.py @@ -1,8 +1,8 @@ import math +import jax import jax.numpy as jnp -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC @@ -128,6 +128,104 @@ def initialize_pml_state(world): return e_memory, b_memory +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("PML tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def _tile_scalar_profile(profile, tile_shape, num_guard_cells=2): + """ + Split one ghost-celled PML profile into compact tile-local profile arrays. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + Nx = int(profile.shape[0]) - 2 + Ny = int(profile.shape[1]) - 2 + Nz = int(profile.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + g = int(num_guard_cells) + + if g != 1: + profile_tiles = jnp.zeros( + (ntx, nty, ntz, tile_nx + 2 * g, tile_ny + 2 * g, tile_nz + 2 * g), + dtype=profile.dtype, + ) + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = tx * tile_nx + iy = ty * tile_ny + iz = tz * tile_nz + tile_with_one_guard = profile[ix:ix + tile_nx + 2, iy:iy + tile_ny + 2, iz:iz + tile_nz + 2] + profile_tiles = profile_tiles.at[ + tx, + ty, + tz, + g - 1:g + tile_nx + 1, + g - 1:g + tile_ny + 1, + g - 1:g + tile_nz + 1, + ].set(tile_with_one_guard) + return profile_tiles + + def tile_at(tx, ty, tz): + start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) + size = (tile_nx + 2, tile_ny + 2, tile_nz + 2) + return jax.lax.dynamic_slice(profile, start, size) + + return jnp.stack( + [ + jnp.stack( + [ + jnp.stack([tile_at(tx, ty, tz) for tz in range(ntz)], axis=0) + for ty in range(nty) + ], + axis=0, + ) + for tx in range(ntx) + ], + axis=0, + ) + + +def tile_pml_profiles(world, tile_shape): + """ + Tile the ghost-celled PML conductivity profiles. + + The leading tile axes match tiled field storage. The final three axes use + the same guard-cell depth as the field tiles. + """ + + _, _, _, _, profiles = world["pml"] + g = int(world.get("guard_cells", 2)) + return tuple(_tile_scalar_profile(profile, tile_shape, num_guard_cells=g) for profile in profiles) + + +def initialize_tiled_pml_state(world, tile_shape): + """ + Allocate tile-local ADE memory for stretched tiled Yee derivatives. + + The memory arrays have leading tile axes followed by the tile physical + interior. They intentionally do not store halos; the ADE terms are attached + to the derivatives after the tile halo exchange has already supplied the + stencil values. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx = _tile_axis_count(int(world["Nx"]), tile_nx) + nty = _tile_axis_count(int(world["Ny"]), tile_ny) + ntz = _tile_axis_count(int(world["Nz"]), tile_nz) + shape = (ntx, nty, ntz, tile_nx, tile_ny, tile_nz) + + e_memory = tuple(jnp.zeros(shape) for _ in range(6)) + b_memory = tuple(jnp.zeros(shape) for _ in range(6)) + tiled_profiles = tile_pml_profiles(world, tile_shape) + + return e_memory, b_memory, tiled_profiles + + def stretch_spatial_derivative(derivative, memory, sigma, dt): """ Apply one coordinate-stretched PML derivative. @@ -249,21 +347,93 @@ def apply_pml_to_b_curl(derivatives, world, pml_state): return (curl_x, curl_y, curl_z), (e_memory, b_memory) -def update_ghost_cells_for_pml(field, world): +def apply_tiled_pml_to_e_curl(derivatives, world, pml_state): + """ + Stretch tile-local B derivatives before assembling Ampere's-law curls. """ - Fill ghost cells without periodically wrapping across a PML wall. - The PML damping happens in the stretched derivatives above. This ghost-cell - rule only prevents a periodic stencil from feeding a wave back through the - opposite side of a PML-active axis. + dBz_dy, dBy_dz, dBx_dz, dBz_dx, dBy_dx, dBx_dy = derivatives + e_memory, b_memory, tiled_profiles = pml_state + ( + memory_dBz_dy, + memory_dBy_dz, + memory_dBx_dz, + memory_dBz_dx, + memory_dBy_dx, + memory_dBx_dy, + ) = e_memory + + sigma_x, sigma_y, sigma_z = tiled_profiles + g = int(world.get("guard_cells", 2)) + sigma_x = sigma_x[:, :, :, g:-g, g:-g, g:-g] + sigma_y = sigma_y[:, :, :, g:-g, g:-g, g:-g] + sigma_z = sigma_z[:, :, :, g:-g, g:-g, g:-g] + dt = world["dt"] + + dBz_dy, memory_dBz_dy = stretch_spatial_derivative(dBz_dy, memory_dBz_dy, sigma_y, dt) + dBy_dz, memory_dBy_dz = stretch_spatial_derivative(dBy_dz, memory_dBy_dz, sigma_z, dt) + dBx_dz, memory_dBx_dz = stretch_spatial_derivative(dBx_dz, memory_dBx_dz, sigma_z, dt) + dBz_dx, memory_dBz_dx = stretch_spatial_derivative(dBz_dx, memory_dBz_dx, sigma_x, dt) + dBy_dx, memory_dBy_dx = stretch_spatial_derivative(dBy_dx, memory_dBy_dx, sigma_x, dt) + dBx_dy, memory_dBx_dy = stretch_spatial_derivative(dBx_dy, memory_dBx_dy, sigma_y, dt) + + curl_x = dBz_dy - dBy_dz + curl_y = dBx_dz - dBz_dx + curl_z = dBy_dx - dBx_dy + + e_memory = ( + memory_dBz_dy, + memory_dBy_dz, + memory_dBx_dz, + memory_dBz_dx, + memory_dBy_dx, + memory_dBx_dy, + ) + + return (curl_x, curl_y, curl_z), (e_memory, b_memory, tiled_profiles) + + +def apply_tiled_pml_to_b_curl(derivatives, world, pml_state): + """ + Stretch tile-local E derivatives before assembling Faraday-law curls. """ - bc_x = world["boundary_conditions"]["x"] - bc_y = world["boundary_conditions"]["y"] - bc_z = world["boundary_conditions"]["z"] - _, pml_x, pml_y, pml_z, _ = world["pml"] - bc_x = jnp.where((pml_x) & (bc_x == BC_PERIODIC), BC_CONDUCTING, bc_x) - bc_y = jnp.where((pml_y) & (bc_y == BC_PERIODIC), BC_CONDUCTING, bc_y) - bc_z = jnp.where((pml_z) & (bc_z == BC_PERIODIC), BC_CONDUCTING, bc_z) + dEz_dy, dEy_dz, dEx_dz, dEz_dx, dEy_dx, dEx_dy = derivatives + e_memory, b_memory, tiled_profiles = pml_state + ( + memory_dEz_dy, + memory_dEy_dz, + memory_dEx_dz, + memory_dEz_dx, + memory_dEy_dx, + memory_dEx_dy, + ) = b_memory + + sigma_x, sigma_y, sigma_z = tiled_profiles + g = int(world.get("guard_cells", 2)) + sigma_x = sigma_x[:, :, :, g:-g, g:-g, g:-g] + sigma_y = sigma_y[:, :, :, g:-g, g:-g, g:-g] + sigma_z = sigma_z[:, :, :, g:-g, g:-g, g:-g] + dt = world["dt"] + + dEz_dy, memory_dEz_dy = stretch_spatial_derivative(dEz_dy, memory_dEz_dy, sigma_y, dt) + dEy_dz, memory_dEy_dz = stretch_spatial_derivative(dEy_dz, memory_dEy_dz, sigma_z, dt) + dEx_dz, memory_dEx_dz = stretch_spatial_derivative(dEx_dz, memory_dEx_dz, sigma_z, dt) + dEz_dx, memory_dEz_dx = stretch_spatial_derivative(dEz_dx, memory_dEz_dx, sigma_x, dt) + dEy_dx, memory_dEy_dx = stretch_spatial_derivative(dEy_dx, memory_dEy_dx, sigma_x, dt) + dEx_dy, memory_dEx_dy = stretch_spatial_derivative(dEx_dy, memory_dEx_dy, sigma_y, dt) + + curl_x = dEz_dy - dEy_dz + curl_y = dEx_dz - dEz_dx + curl_z = dEy_dx - dEx_dy + + b_memory = ( + memory_dEz_dy, + memory_dEy_dz, + memory_dEx_dz, + memory_dEz_dx, + memory_dEy_dx, + memory_dEx_dy, + ) - return update_ghost_cells(field, bc_x, bc_y, bc_z) + return (curl_x, curl_y, curl_z), (e_memory, b_memory, tiled_profiles) diff --git a/PyPIC3D/boundary_conditions/__init__.py b/PyPIC3D/boundary_conditions/__init__.py index 09aae75a..5eb29abc 100644 --- a/PyPIC3D/boundary_conditions/__init__.py +++ b/PyPIC3D/boundary_conditions/__init__.py @@ -7,10 +7,17 @@ ) from PyPIC3D.boundary_conditions.boundaryconditions import ( - update_ghost_cells, - fold_ghost_cells, - apply_conducting_bc, - apply_scalar_conducting_bc, apply_supergaussian_boundary_condition, apply_zero_boundary_condition, ) + +from PyPIC3D.boundary_conditions.ghost_cells import ( + update_tiled_ghost_cells, + update_tiled_vector_ghost_cells, + update_tiled_ghost_cells_for_pml, + update_tiled_vector_ghost_cells_for_pml, + apply_tiled_conducting_bc, + apply_tiled_scalar_conducting_bc, + fold_tiled_ghost_cells, + fold_tiled_vector_ghost_cells, +) diff --git a/PyPIC3D/boundary_conditions/boundaryconditions.py b/PyPIC3D/boundary_conditions/boundaryconditions.py index f8a22fe6..6725a030 100644 --- a/PyPIC3D/boundary_conditions/boundaryconditions.py +++ b/PyPIC3D/boundary_conditions/boundaryconditions.py @@ -1,238 +1,10 @@ # Christopher Woolford, April 6th 2026 -# Field-level boundary condition operations for ghost-celled 3D arrays. +# Field-level boundary condition operations for ordinary 3D arrays. -from jax import jit, lax +from jax import jit import jax.numpy as jnp # import external libraries -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -# import internal modules - - -@jit -def update_ghost_cells(field, bc_x, bc_y, bc_z): - """ - Fill ghost cells of a 3D field from its interior based on boundary conditions. - - Ghost cells provide the halo data needed by finite-difference stencils - that reach across the domain boundary. For periodic BCs, the ghost - layer mirrors the opposite interior edge. For conducting BCs, ghost - cells are zeroed (Dirichlet condition). - - The field is assumed to have shape (Nx+2, Ny+2, Nz+2) with one ghost - cell on each side per axis. Indices [0] and [-1] are ghost cells; - indices [1] and [-2] are the first and last interior cells. - - Args: - field (jnp.ndarray): 3D field array with shape (Nx+2, Ny+2, Nz+2). - bc_x (int): Boundary condition code for the x-axis (BC_PERIODIC or BC_CONDUCTING). - bc_y (int): Boundary condition code for the y-axis. - bc_z (int): Boundary condition code for the z-axis. - - Returns: - jnp.ndarray: Field with ghost cells filled. - """ - # x-axis ghost cells - # Periodic: ghost[0] <- interior[-2], ghost[-1] <- interior[1] - # Conducting: ghost[0] = ghost[-1] = 0 - field = lax.cond( - bc_x == BC_PERIODIC, - lambda f: f.at[0, :, :].set(f[-2, :, :]).at[-1, :, :].set(f[1, :, :]), - lambda f: f.at[0, :, :].set(0.0).at[-1, :, :].set(0.0), - operand=field - ) - - # y-axis ghost cells - field = lax.cond( - bc_y == BC_PERIODIC, - lambda f: f.at[:, 0, :].set(f[:, -2, :]).at[:, -1, :].set(f[:, 1, :]), - lambda f: f.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0), - operand=field - ) - - # z-axis ghost cells - field = lax.cond( - bc_z == BC_PERIODIC, - lambda f: f.at[:, :, 0].set(f[:, :, -2]).at[:, :, -1].set(f[:, :, 1]), - lambda f: f.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0), - operand=field - ) - - return field - - -@jit -def fold_ghost_cells(field, bc_x, bc_y, bc_z): - """ - Fold ghost cell deposits back into the interior after particle deposition. - - During charge/current deposition, particles near the boundary may - deposit into ghost cells. This function adds those ghost deposits to - the corresponding interior cells and then clears the ghost layer. - - For periodic BCs, ghost deposits are added to the opposite interior - edge (ghost[0] -> interior[1], ghost[-1] -> interior[-2]). - - For conducting BCs with reflecting particles, ghost deposits are reflected - back into the domain with opposite sign. - - In both cases, ghost cells are cleared to zero after folding. - - Args: - field (jnp.ndarray): 3D field array with shape (Nx+2, Ny+2, Nz+2). - bc_x (int): Boundary condition code for the x-axis (BC_PERIODIC or BC_CONDUCTING). - bc_y (int): Boundary condition code for the y-axis. - bc_z (int): Boundary condition code for the z-axis. - - Returns: - jnp.ndarray: Field with ghost cell deposits folded back into the interior. - """ - # x-axis: fold ghost deposits back to interior - field = lax.cond( - bc_x == BC_PERIODIC, - lambda f: f.at[1, :, :].add(f[-1, :, :]).at[-2, :, :].add(f[0, :, :]), - # add the ghost deposits to the opposite inner cell. - lambda f: f.at[1, :, :].add(-1 * f[0, :, :]).at[-2, :, :].add(-1 * f[-1, :, :]), - # for conducting BCs, the ghost deposits should be reflected back with opposite sign. - operand=field - ) - # if its periodic, the ghost deposits should be added to the other side. - # if its conducting, the ghost deposits should be reflected back with opposite sign. - # Clear x ghost cells after folding (both periodic and conducting) - field = field.at[0, :, :].set(0.0).at[-1, :, :].set(0.0) - - # y-axis: fold ghost deposits back to interior - field = lax.cond( - bc_y == BC_PERIODIC, - lambda f: f.at[:, 1, :].add(f[:, -1, :]).at[:, -2, :].add(f[:, 0, :]), - lambda f: f.at[:, 1, :].add(-1 * f[:, 0, :]).at[:, -2, :].add(-1 * f[:, -1, :]), - operand=field - ) - # Clear y ghost cells after folding - field = field.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0) - - # z-axis: fold ghost deposits back to interior - field = lax.cond( - bc_z == BC_PERIODIC, - lambda f: f.at[:, :, 1].add(f[:, :, -1]).at[:, :, -2].add(f[:, :, 0]), - lambda f: f.at[:, :, 1].add(-1 * f[:, :, 0]).at[:, :, -2].add(-1 * f[:, :, -1]), - operand=field - ) - # Clear z ghost cells after folding - field = field.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0) - - return field - - -@jit -def apply_conducting_bc(E, bc_x, bc_y, bc_z): - """ - Enforce conducting boundary conditions by zeroing tangential E at boundaries. - - On a perfectly conducting wall, the tangential electric field must vanish. - On a ghost-celled (Nx+2, Ny+2, Nz+2) grid, the first physical face is - at index [1] and the last at index [-2]. - - For a conducting wall normal to a given axis, the two tangential - components are zeroed at both boundary faces along that axis: - - Wall normal to x: zero Ey, Ez at x-boundaries [1] and [-2]. - - Wall normal to y: zero Ex, Ez at y-boundaries [1] and [-2]. - - Wall normal to z: zero Ex, Ey at z-boundaries [1] and [-2]. - - Args: - E (tuple): Electric field components (Ex, Ey, Ez), each with shape (Nx+2, Ny+2, Nz+2). - bc_x (int): Boundary condition code for the x-axis (BC_PERIODIC or BC_CONDUCTING). - bc_y (int): Boundary condition code for the y-axis. - bc_z (int): Boundary condition code for the z-axis. - - Returns: - tuple: Electric field components (Ex, Ey, Ez) with conducting BCs applied. - """ - Ex, Ey, Ez = E - - ### x-boundary conducting: zero tangential Ey, Ez ### - Ey = lax.cond( - bc_x == BC_CONDUCTING, - lambda f: f.at[1, :, :].set(0.0).at[-2, :, :].set(0.0), - lambda f: f, - operand=Ey - ) - Ez = lax.cond( - bc_x == BC_CONDUCTING, - lambda f: f.at[1, :, :].set(0.0).at[-2, :, :].set(0.0), - lambda f: f, - operand=Ez - ) - - ### y-boundary conducting: zero tangential Ex, Ez ### - Ex = lax.cond( - bc_y == BC_CONDUCTING, - lambda f: f.at[:, 1, :].set(0.0).at[:, -2, :].set(0.0), - lambda f: f, - operand=Ex - ) - Ez = lax.cond( - bc_y == BC_CONDUCTING, - lambda f: f.at[:, 1, :].set(0.0).at[:, -2, :].set(0.0), - lambda f: f, - operand=Ez - ) - - ### z-boundary conducting: zero tangential Ex, Ey ### - Ex = lax.cond( - bc_z == BC_CONDUCTING, - lambda f: f.at[:, :, 1].set(0.0).at[:, :, -2].set(0.0), - lambda f: f, - operand=Ex - ) - Ey = lax.cond( - bc_z == BC_CONDUCTING, - lambda f: f.at[:, :, 1].set(0.0).at[:, :, -2].set(0.0), - lambda f: f, - operand=Ey - ) - - return (Ex, Ey, Ez) - - -@jit -def apply_scalar_conducting_bc(field, bc_x, bc_y, bc_z): - """ - Enforce conducting boundary conditions on a scalar field (e.g. electric potential phi). - - Sets the field to zero at each conducting boundary face. On a - ghost-celled grid, boundary faces are at interior indices [1] and [-2]. - - Args: - field (jnp.ndarray): 3D scalar field with shape (Nx+2, Ny+2, Nz+2). - bc_x (int): Boundary condition code for the x-axis (BC_PERIODIC or BC_CONDUCTING). - bc_y (int): Boundary condition code for the y-axis. - bc_z (int): Boundary condition code for the z-axis. - - Returns: - jnp.ndarray: Scalar field with conducting BCs applied. - """ - field = lax.cond( - bc_x == BC_CONDUCTING, - lambda f: f.at[1, :, :].set(0.0).at[-2, :, :].set(0.0), - lambda f: f, - operand=field - ) - field = lax.cond( - bc_y == BC_CONDUCTING, - lambda f: f.at[:, 1, :].set(0.0).at[:, -2, :].set(0.0), - lambda f: f, - operand=field - ) - field = lax.cond( - bc_z == BC_CONDUCTING, - lambda f: f.at[:, :, 1].set(0.0).at[:, :, -2].set(0.0), - lambda f: f, - operand=field - ) - return field - - # No @jit: the Python for-loop is unrolled during JAX tracing by any # JIT-compiled caller, so explicit JIT here would cause tracing errors. def apply_supergaussian_boundary_condition(field, boundary_thickness, order, strength): diff --git a/PyPIC3D/boundary_conditions/ghost_cells.py b/PyPIC3D/boundary_conditions/ghost_cells.py new file mode 100644 index 00000000..96c67a96 --- /dev/null +++ b/PyPIC3D/boundary_conditions/ghost_cells.py @@ -0,0 +1,418 @@ +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC + + +def _reduced_tiled_axes(field_tiles, tile_shape, g): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx, nty, ntz = field_tiles.shape[:3] + return ( + int(ntx) == 1 and tile_nx == 1, + int(nty) == 1 and tile_ny == 1, + int(ntz) == 1 and tile_nz == 1, + ) + + +def _is_stacked_tiled_vector_field(field_tiles): + return hasattr(field_tiles, "ndim") and field_tiles.ndim == 7 and int(field_tiles.shape[0]) == 3 + + +def _stack_tiled_vector_field(field_tiles): + if _is_stacked_tiled_vector_field(field_tiles): + return field_tiles + return jnp.stack(field_tiles, axis=0) + + +def _unstack_tiled_vector_field(field_tiles): + if _is_stacked_tiled_vector_field(field_tiles): + return field_tiles[0], field_tiles[1], field_tiles[2] + return field_tiles + + +def _restore_tiled_vector_layout(stacked_tiles, original_tiles): + if _is_stacked_tiled_vector_field(original_tiles): + return stacked_tiles + return _unstack_tiled_vector_field(stacked_tiles) + + +def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Refresh tile halos using the field boundary conditions. + + ``field_tiles`` has shape + ``(ntx, nty, ntz, tile_nx + 2*g, tile_ny + 2*g, tile_nz + 2*g)``. + The last three axes are the tile-local ghost-celled field with + ``g = num_guard_cells`` guard cells on each side. The first three axes + identify the tile in the global tiling. + """ + + g = int(num_guard_cells) + if tile_shape is None: + tile_shape = tuple(int(width) for width in world["tile_shape"]) + reduced_x, reduced_y, reduced_z = _reduced_tiled_axes(field_tiles, tile_shape, g) + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + # x-halos: internal tile boundaries exchange neighboring interiors. At a + # conducting global wall, the exterior ghost face is zero, matching the + # assembled-field boundary convention. + if reduced_x: + lower_x = jnp.broadcast_to(field_tiles[:, :, :, g:g + 1, :, :], field_tiles[:, :, :, :g, :, :].shape) + upper_x = jnp.broadcast_to(field_tiles[:, :, :, g:g + 1, :, :], field_tiles[:, :, :, -g:, :, :].shape) + field_tiles = field_tiles.at[:, :, :, :g, :, :].set(lower_x) + field_tiles = field_tiles.at[:, :, :, -g:, :, :].set(upper_x) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, :g, :, :].set(0.0).at[:, :, :, -g:, :, :].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + else: + field_tiles = field_tiles.at[:, :, :, :g, :, :].set( + jnp.roll(field_tiles[:, :, :, -2 * g:-g, :, :], shift=1, axis=0) + ) + field_tiles = field_tiles.at[:, :, :, -g:, :, :].set( + jnp.roll(field_tiles[:, :, :, g:2 * g, :, :], shift=-1, axis=0) + ) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles.at[0, :, :, :g, :, :].set(0.0).at[-1, :, :, -g:, :, :].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + + # y-halos use the x-refreshed field so tile-edge/corner guard cells match + # the sequential ghost-cell convention used by assembled arrays. + if reduced_y: + lower_y = jnp.broadcast_to(field_tiles[:, :, :, :, g:g + 1, :], field_tiles[:, :, :, :, :g, :].shape) + upper_y = jnp.broadcast_to(field_tiles[:, :, :, :, g:g + 1, :], field_tiles[:, :, :, :, -g:, :].shape) + field_tiles = field_tiles.at[:, :, :, :, :g, :].set(lower_y) + field_tiles = field_tiles.at[:, :, :, :, -g:, :].set(upper_y) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, :, :g, :].set(0.0).at[:, :, :, :, -g:, :].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + else: + field_tiles = field_tiles.at[:, :, :, :, :g, :].set( + jnp.roll(field_tiles[:, :, :, :, -2 * g:-g, :], shift=1, axis=1) + ) + field_tiles = field_tiles.at[:, :, :, :, -g:, :].set( + jnp.roll(field_tiles[:, :, :, :, g:2 * g, :], shift=-1, axis=1) + ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles.at[:, 0, :, :, :g, :].set(0.0).at[:, -1, :, :, -g:, :].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + + # z-halos complete the edge and corner values after x/y have been refreshed. + if reduced_z: + lower_z = jnp.broadcast_to(field_tiles[:, :, :, :, :, g:g + 1], field_tiles[:, :, :, :, :, :g].shape) + upper_z = jnp.broadcast_to(field_tiles[:, :, :, :, :, g:g + 1], field_tiles[:, :, :, :, :, -g:].shape) + field_tiles = field_tiles.at[:, :, :, :, :, :g].set(lower_z) + field_tiles = field_tiles.at[:, :, :, :, :, -g:].set(upper_z) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, :, :, :g].set(0.0).at[:, :, :, :, :, -g:].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + else: + field_tiles = field_tiles.at[:, :, :, :, :, :g].set( + jnp.roll(field_tiles[:, :, :, :, :, -2 * g:-g], shift=1, axis=2) + ) + field_tiles = field_tiles.at[:, :, :, :, :, -g:].set( + jnp.roll(field_tiles[:, :, :, :, :, g:2 * g], shift=-1, axis=2) + ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, 0, :, :, :g].set(0.0).at[:, :, -1, :, :, -g:].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) + + return field_tiles + + +def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Refresh tile halos for a tiled vector field. + """ + + stacked_tiles = _stack_tiled_vector_field(field_tiles) + refreshed = jax.vmap( + lambda component: update_tiled_ghost_cells(component, world, num_guard_cells, tile_shape), + in_axes=0, + out_axes=0, + )(stacked_tiles) + return _restore_tiled_vector_layout(refreshed, field_tiles) + + +def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Refresh tile halos without wrapping across PML-active global walls. + """ + + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + _, pml_x, pml_y, pml_z, _ = world["pml"] + + bc_x = jnp.where((pml_x) & (bc_x == BC_PERIODIC), BC_CONDUCTING, bc_x) + bc_y = jnp.where((pml_y) & (bc_y == BC_PERIODIC), BC_CONDUCTING, bc_y) + bc_z = jnp.where((pml_z) & (bc_z == BC_PERIODIC), BC_CONDUCTING, bc_z) + + pml_world = { + "boundary_conditions": {"x": bc_x, "y": bc_y, "z": bc_z}, + "tile_shape": world["tile_shape"], + } + return update_tiled_ghost_cells(field_tiles, pml_world, num_guard_cells, tile_shape) + + +def update_tiled_vector_ghost_cells_for_pml(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Refresh tile halos for a vector field with PML-active exterior walls. + """ + + stacked_tiles = _stack_tiled_vector_field(field_tiles) + refreshed = jax.vmap( + lambda component: update_tiled_ghost_cells_for_pml(component, world, num_guard_cells, tile_shape), + in_axes=0, + out_axes=0, + )(stacked_tiles) + return _restore_tiled_vector_layout(refreshed, field_tiles) + + +def apply_tiled_conducting_bc(E_tiles, world, num_guard_cells=2): + """ + Zero tangential electric-field components on global conducting faces. + """ + + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + Ex, Ey, Ez = E_tiles + g = int(num_guard_cells) + lower = g + upper = -g - 1 + + Ey = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda f: f.at[0, :, :, lower, :, :].set(0.0).at[-1, :, :, upper, :, :].set(0.0), + lambda f: f, + operand=Ey, + ) + Ez = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda f: f.at[0, :, :, lower, :, :].set(0.0).at[-1, :, :, upper, :, :].set(0.0), + lambda f: f, + operand=Ez, + ) + + Ex = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda f: f.at[:, 0, :, :, lower, :].set(0.0).at[:, -1, :, :, upper, :].set(0.0), + lambda f: f, + operand=Ex, + ) + Ez = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda f: f.at[:, 0, :, :, lower, :].set(0.0).at[:, -1, :, :, upper, :].set(0.0), + lambda f: f, + operand=Ez, + ) + + Ex = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda f: f.at[:, :, 0, :, :, lower].set(0.0).at[:, :, -1, :, :, upper].set(0.0), + lambda f: f, + operand=Ex, + ) + Ey = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda f: f.at[:, :, 0, :, :, lower].set(0.0).at[:, :, -1, :, :, upper].set(0.0), + lambda f: f, + operand=Ey, + ) + + return Ex, Ey, Ez + + +def apply_tiled_scalar_conducting_bc(field_tiles, world, num_guard_cells=2): + """ + Zero scalar field values on global conducting faces. + """ + + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + g = int(num_guard_cells) + lower = g + upper = -g - 1 + + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda f: f.at[0, :, :, lower, :, :].set(0.0).at[-1, :, :, upper, :, :].set(0.0), + lambda f: f, + operand=field_tiles, + ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda f: f.at[:, 0, :, :, lower, :].set(0.0).at[:, -1, :, :, upper, :].set(0.0), + lambda f: f, + operand=field_tiles, + ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda f: f.at[:, :, 0, :, :, lower].set(0.0).at[:, :, -1, :, :, upper].set(0.0), + lambda f: f, + operand=field_tiles, + ) + + return field_tiles + + +def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Add tile-ghost deposits into owning interiors, then clear ghosts. + + Internal tile ghosts always belong to the neighboring physical tile + interior. At global conducting walls, exterior ghost deposits reflect into + the adjacent boundary cell with the same sign convention used by the + legacy global ghost-cell fold. + """ + + g = int(num_guard_cells) + if tile_shape is None: + tile_shape = tuple(int(width) for width in world["tile_shape"]) + reduced_x, reduced_y, reduced_z = _reduced_tiled_axes(field_tiles, tile_shape, g) + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + if reduced_x: + x_ghost_sum = ( + jnp.sum(field_tiles[:, :, :, :g, :, :], axis=3, keepdims=True) + + jnp.sum(field_tiles[:, :, :, -g:, :, :], axis=3, keepdims=True) + ) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, g:g + 1, :, :].add(-x_ghost_sum), + lambda tiles: tiles.at[:, :, :, g:g + 1, :, :].add(x_ghost_sum), + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :g, :, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, -g:, :, :].set(0.0) + else: + x_lower_ghost = field_tiles[:, :, :, :g, :, :] + x_upper_ghost = field_tiles[:, :, :, -g:, :, :] + field_tiles = field_tiles.at[:, :, :, -2 * g:-g, :, :].add( + jnp.roll(x_lower_ghost, shift=-1, axis=0) + ) + field_tiles = field_tiles.at[:, :, :, g:2 * g, :, :].add( + jnp.roll(x_upper_ghost, shift=1, axis=0) + ) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles + .at[-1, :, :, -2 * g:-g, :, :].add(-x_lower_ghost[0, :, :, :, :]) + .at[0, :, :, g:2 * g, :, :].add(-x_upper_ghost[-1, :, :, :, :]) + .at[0, :, :, g:2 * g, :, :].add(-x_lower_ghost[0, :, :, :, :]) + .at[-1, :, :, -2 * g:-g, :, :].add(-x_upper_ghost[-1, :, :, :, :]), + lambda tiles: tiles, + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :g, :, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, -g:, :, :].set(0.0) + + if reduced_y: + y_ghost_sum = ( + jnp.sum(field_tiles[:, :, :, :, :g, :], axis=4, keepdims=True) + + jnp.sum(field_tiles[:, :, :, :, -g:, :], axis=4, keepdims=True) + ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, :, g:g + 1, :].add(-y_ghost_sum), + lambda tiles: tiles.at[:, :, :, :, g:g + 1, :].add(y_ghost_sum), + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :, :g, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, -g:, :].set(0.0) + else: + y_lower_ghost = field_tiles[:, :, :, :, :g, :] + y_upper_ghost = field_tiles[:, :, :, :, -g:, :] + field_tiles = field_tiles.at[:, :, :, :, -2 * g:-g, :].add( + jnp.roll(y_lower_ghost, shift=-1, axis=1) + ) + field_tiles = field_tiles.at[:, :, :, :, g:2 * g, :].add( + jnp.roll(y_upper_ghost, shift=1, axis=1) + ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles + .at[:, -1, :, :, -2 * g:-g, :].add(-y_lower_ghost[:, 0, :, :, :]) + .at[:, 0, :, :, g:2 * g, :].add(-y_upper_ghost[:, -1, :, :, :]) + .at[:, 0, :, :, g:2 * g, :].add(-y_lower_ghost[:, 0, :, :, :]) + .at[:, -1, :, :, -2 * g:-g, :].add(-y_upper_ghost[:, -1, :, :, :]), + lambda tiles: tiles, + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :, :g, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, -g:, :].set(0.0) + + if reduced_z: + z_ghost_sum = ( + jnp.sum(field_tiles[:, :, :, :, :, :g], axis=5, keepdims=True) + + jnp.sum(field_tiles[:, :, :, :, :, -g:], axis=5, keepdims=True) + ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, :, :, :, g:g + 1].add(-z_ghost_sum), + lambda tiles: tiles.at[:, :, :, :, :, g:g + 1].add(z_ghost_sum), + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :, :, :g].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :, -g:].set(0.0) + else: + z_lower_ghost = field_tiles[:, :, :, :, :, :g] + z_upper_ghost = field_tiles[:, :, :, :, :, -g:] + field_tiles = field_tiles.at[:, :, :, :, :, -2 * g:-g].add( + jnp.roll(z_lower_ghost, shift=-1, axis=2) + ) + field_tiles = field_tiles.at[:, :, :, :, :, g:2 * g].add( + jnp.roll(z_upper_ghost, shift=1, axis=2) + ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles + .at[:, :, -1, :, :, -2 * g:-g].add(-z_lower_ghost[:, :, 0, :, :]) + .at[:, :, 0, :, :, g:2 * g].add(-z_upper_ghost[:, :, -1, :, :]) + .at[:, :, 0, :, :, g:2 * g].add(-z_lower_ghost[:, :, 0, :, :]) + .at[:, :, -1, :, :, -2 * g:-g].add(-z_upper_ghost[:, :, -1, :, :]), + lambda tiles: tiles, + operand=field_tiles, + ) + field_tiles = field_tiles.at[:, :, :, :, :, :g].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :, -g:].set(0.0) + + return field_tiles + + +def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): + """ + Fold tile-ghost deposits for a tiled vector field. + """ + + stacked_tiles = _stack_tiled_vector_field(field_tiles) + folded = jax.vmap( + lambda component: fold_tiled_ghost_cells(component, world, num_guard_cells, tile_shape), + in_axes=0, + out_axes=0, + )(stacked_tiles) + return _restore_tiled_vector_layout(folded, field_tiles) diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 3a812488..bf17d5e3 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -1,17 +1,14 @@ import jax -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 ( - BC_PERIODIC, - axis_has_active_cells, compute_particle_anchor, particle_axis_offset, ) -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.boundary_conditions.ghost_cells import update_tiled_vector_ghost_cells +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles def shift_old_stencil(old_w_list, shift): @@ -45,211 +42,282 @@ def collapse_redundant_axis(points, current_weights, old_weights, axis_active, e return collapsed_points, collapsed_current, collapsed_old -def clear_ghost_cells(field, axis): - """Clear the two ghost layers along one axis.""" - field_axis = jnp.moveaxis(field, axis, 0) - # move the axis to be cleared to the front for easier indexing - field_axis = field_axis.at[0, :, :].set(0.0) - field_axis = field_axis.at[1, :, :].set(0.0) - field_axis = field_axis.at[-2, :, :].set(0.0) - field_axis = field_axis.at[-1, :, :].set(0.0) - # clear the two ghost layers by setting them to zero - field = jnp.moveaxis(field_axis, 0, axis) - # move the axis back to its original position - return field +def _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=0): + if dim == 0: + Wx = jnp.stack([xw[i] - oxw[i] for i in range(5)], axis=0) + Wy = jnp.stack([(xw[i] + oxw[i]) / 2 for i in range(5)], axis=0) + Wz = jnp.stack([(xw[i] + oxw[i]) / 2 for i in range(5)], axis=0) + elif dim == 1: + Wy = jnp.stack([yw[j] - oyw[j] for j in range(5)], axis=0) + Wx = jnp.stack([(yw[j] + oyw[j]) / 2 for j in range(5)], axis=0) + Wz = jnp.stack([(yw[j] + oyw[j]) / 2 for j in range(5)], axis=0) + else: + Wz = jnp.stack([zw[k] - ozw[k] for k in range(5)], axis=0) + Wx = jnp.stack([(zw[k] + ozw[k]) / 2 for k in range(5)], axis=0) + Wy = jnp.stack([(zw[k] + ozw[k]) / 2 for k in range(5)], axis=0) + + return Wx, Wy, Wz + + +def _compact_2d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, null_dim=2): + def stack_plane(rows): + return jnp.stack([jnp.stack(row, axis=0) for row in rows], axis=0) + + if null_dim == 0: + Wx = stack_plane( + [ + [ + 1 / 3 * (yw[j] * zw[k] + oyw[j] * ozw[k]) + + 1 / 6 * (yw[j] * ozw[k] + oyw[j] * zw[k]) + for k in range(5) + ] + for j in range(5) + ] + ) + Wy = stack_plane( + [[1 / 2 * (yw[j] - oyw[j]) * (zw[k] + ozw[k]) for k in range(5)] for j in range(5)] + ) + Wz = stack_plane( + [[1 / 2 * (zw[k] - ozw[k]) * (yw[j] + oyw[j]) for k in range(5)] for j in range(5)] + ) + elif null_dim == 1: + Wx = stack_plane( + [[1 / 2 * (xw[i] - oxw[i]) * (zw[k] + ozw[k]) for k in range(5)] for i in range(5)] + ) + Wy = stack_plane( + [ + [ + 1 / 3 * (xw[i] * zw[k] + oxw[i] * ozw[k]) + + 1 / 6 * (xw[i] * ozw[k] + oxw[i] * zw[k]) + for k in range(5) + ] + for i in range(5) + ] + ) + Wz = stack_plane( + [[1 / 2 * (zw[k] - ozw[k]) * (xw[i] + oxw[i]) for k in range(5)] for i in range(5)] + ) + else: + Wx = stack_plane( + [[1 / 2 * (xw[i] - oxw[i]) * (yw[j] + oyw[j]) for j in range(5)] for i in range(5)] + ) + Wy = stack_plane( + [[1 / 2 * (yw[j] - oyw[j]) * (xw[i] + oxw[i]) for j in range(5)] for i in range(5)] + ) + Wz = stack_plane( + [ + [ + 1 / 3 * (xw[i] * yw[j] + oxw[i] * oyw[j]) + + 1 / 6 * (xw[i] * oyw[j] + oxw[i] * yw[j]) + for j in range(5) + ] + for i in range(5) + ] + ) + return Wx, Wy, Wz -def enforce_bc_along_axis(field, axis, bc, component_axis): - """ - Enforce boundary conditions along a specified axis of a field array. - - This function applies boundary conditions to ghost cells of a field by moving - the specified axis to the front for easier manipulation, applying the appropriate - boundary condition logic, and then moving the axis back to its original position. - - Parameters - ---------- - field : jax.Array - The field array to which boundary conditions will be applied. - Expected to have two ghost cells at indices 0, 1, -2, -1 along the specified axis. - axis : int - The axis along which to enforce boundary conditions (0, 1, or 2). - bc : str - The type of boundary condition to apply. Options are: - - "reflecting": Ghost cells are reflected across the boundary with optional sign flip - - "absorbing": Ghost cells are left unchanged (will be cleared afterward) - - "periodic": Ghost cells are folded by adding them to opposite physical layers (default) - component_axis : int - The axis corresponding to the field component direction. Used to determine - the sign for reflecting boundary conditions (sign = -1 if axis == component_axis, - else 1). - - Returns - ------- - jax.Array - The field array with boundary conditions enforced and ghost cells cleared. - - Notes - ----- - - Ghost cells are assumed to be at indices [0, 1, -2, -1] along the axis. - - The extended physical region is at indices [2:-2] along the axis. - - After applying boundary conditions, all ghost cells are cleared to zero. + +def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells, tile_shape): """ + Fold Esirkepov current deposits using the configured guard depth. + Internal tile faces are ownership transfers. At global particle boundaries, + periodic deposits wrap, reflecting deposits mirror with the Esirkepov + component sign, and absorbing deposits are discarded. + """ - field_axis = jnp.moveaxis(field, axis, 0) - # move the axis to be folded to the front for easier indexing. + g = int(num_guard_cells) + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx, nty, ntz = field_tiles.shape[:3] + reduced_axes = ( + int(ntx) == 1 and tile_nx == 1, + int(nty) == 1 and tile_ny == 1, + int(ntz) == 1 and tile_nz == 1, + ) + particle_bc = world.get("particle_boundary_conditions", {"x": 0, "y": 0, "z": 0}) - if bc == "reflecting": + def fold_axis(field_tiles, axis, bc, reduced_axis): + tiles = jnp.moveaxis(field_tiles, (axis, 3 + axis), (0, 3)) + lower_ghost = tiles[:, :, :, :g, :, :] + upper_ghost = tiles[:, :, :, -g:, :, :] sign = -1.0 if axis == component_axis else 1.0 - field_axis = field_axis.at[2, :, :].add(sign * field_axis[0, :, :]) - field_axis = field_axis.at[3, :, :].add(sign * field_axis[1, :, :]) - field_axis = field_axis.at[-4, :, :].add(sign * field_axis[-2, :, :]) - field_axis = field_axis.at[-3, :, :].add(sign * field_axis[-1, :, :]) - # if the boundary is reflecting, then reflect the ghost layers across the boundary - elif bc == "absorbing": - field_axis = field_axis - # if the boundary is absorbing, then neglect the ghost layers by leaving them as they are, which will be cleared to zero in the next step. - else: - # PERIODIC BC IS THE DEFAULT - field_axis = field_axis.at[-4, :, :].add(field_axis[0, :, :]) - field_axis = field_axis.at[-3, :, :].add(field_axis[1, :, :]) - field_axis = field_axis.at[2, :, :].add(field_axis[-2, :, :]) - field_axis = field_axis.at[3, :, :].add(field_axis[-1, :, :]) - # if the boundary is periodic, fold the ghost layers by adding them to the opposite physical layer. + if reduced_axis: + ghost_sum = ( + jnp.sum(lower_ghost, axis=3, keepdims=True) + + jnp.sum(upper_ghost, axis=3, keepdims=True) + ) - field = jnp.moveaxis(field_axis, 0, axis) - # move the axis back to its original position after folding the ghost layers according to the boundary condition. - field = clear_ghost_cells(field, axis) - # clear the ghost cells + def periodic_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(ghost_sum) - return field + def reflecting_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(sign * ghost_sum) -def ghost_cell_bc_esirkepov(field, bc_x, bc_y, bc_z, component_axis): - """ - Apply boundary conditions to ghost cells using Esirkepov method. - - Enforces boundary conditions along all three axes (x, y, z) for a field - by applying the appropriate boundary condition type along each spatial dimension. - - Parameters - ---------- - field : ndarray - The field array to which boundary conditions will be applied. - bc_x : str or callable - Boundary condition to enforce along the x-axis (axis=0). - bc_y : str or callable - Boundary condition to enforce along the y-axis (axis=1). - bc_z : str or callable - Boundary condition to enforce along the z-axis (axis=2). - component_axis : int - The axis index of the field component for which boundary conditions - are being applied. - - Returns - ------- - ndarray - The field array with boundary conditions applied along all three axes. - """ - field = enforce_bc_along_axis(field, axis=0, bc=bc_x, component_axis=component_axis) - # enforce the x boundary conditions along the x axis (axis=0) - field = enforce_bc_along_axis(field, axis=1, bc=bc_y, component_axis=component_axis) - # enforce the y boundary conditions along the y axis (axis=1) - field = enforce_bc_along_axis(field, axis=2, bc=bc_z, component_axis=component_axis) - # enforce the z boundary conditions along the z axis (axis=2) - return field + def absorbing_boundary(tiles): + return tiles + + tiles = jax.lax.switch(bc, (periodic_boundary, reflecting_boundary, absorbing_boundary), tiles) + tiles = tiles.at[:, :, :, :g, :, :].set(0.0) + tiles = tiles.at[:, :, :, -g:, :, :].set(0.0) + return jnp.moveaxis(tiles, (0, 3), (axis, 3 + axis)) + + tiles = tiles.at[:-1, :, :, -2 * g:-g, :, :].add(lower_ghost[1:, :, :, :, :, :]) + tiles = tiles.at[1:, :, :, g:2 * g, :, :].add(upper_ghost[:-1, :, :, :, :, :]) + + def periodic_boundary(tiles): + tiles = tiles.at[-1, :, :, -2 * g:-g, :, :].add(lower_ghost[0, :, :, :, :, :]) + tiles = tiles.at[0, :, :, g:2 * g, :, :].add(upper_ghost[-1, :, :, :, :, :]) + return tiles + + def reflecting_boundary(tiles): + tiles = tiles.at[0, :, :, g:2 * g, :, :].add(sign * lower_ghost[0, :, :, :, :, :]) + tiles = tiles.at[-1, :, :, -2 * g:-g, :, :].add(sign * upper_ghost[-1, :, :, :, :, :]) + return tiles + + def absorbing_boundary(tiles): + return tiles + + tiles = jax.lax.switch(bc, (periodic_boundary, reflecting_boundary, absorbing_boundary), tiles) + tiles = tiles.at[:, :, :, :g, :, :].set(0.0) + tiles = tiles.at[:, :, :, -g:, :, :].set(0.0) + return jnp.moveaxis(tiles, (0, 3), (axis, 3 + axis)) + + field_tiles = fold_axis(field_tiles, 0, particle_bc["x"], reduced_axes[0]) + field_tiles = fold_axis(field_tiles, 1, particle_bc["y"], reduced_axes[1]) + field_tiles = fold_axis(field_tiles, 2, particle_bc["z"], reduced_axes[2]) + return field_tiles -def eliminate_esirkepov_ghost_cells(field): - slices = [slice(1, -1), slice(1, -1), slice(1, -1)] - # remove one layer on each side, returning the solver's ordinary one-ghost current array. - return field[tuple(slices)] -def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): - """Esirkepov current deposition supporting 1D/2D/3D via inactive dims.""" - if grid is None: - grid = world["grids"]["center"] +def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells, tile_shape): + return tuple( + fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells, tile_shape) + for component_axis, component in enumerate(field_tiles) + ) + + +def Esirkepov_current( + particles: TiledParticles, + species_config: SpeciesConfig, + J, + constants, + world, + filter=None, +): + """ + Deposit Esirkepov current into tile-local current buffers. + + ``particles.x`` is the old particle position after the velocity push. The + future position used in Esirkepov's charge-conserving difference is + predicted locally as ``x + u*dt``. Particle retile ownership is applied by + the caller after deposition. + """ + + if not isinstance(particles, TiledParticles): + raise ValueError("Public Esirkepov_current requires TiledParticles.") + if filter not in (None, "none"): + raise ValueError("Esirkepov current filtering is not supported; use filter='none'.") + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + # get the tile shape + g = int(world["guard_cells"]) + # get the number of guard cells on the tiles + tiled_grid = world["grids"]["tiled_center_grid"] + # get the tile grid for the current deposition + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + # get spatial resolution + dt = world["dt"] + # get temporal resolution + shape_factor = world["shape_factor"] + # get shape factor Jx, Jy, Jz = J - Nx, Ny, Nz = Jx.shape - dx, dy, dz, dt = world["dx"], world["dy"], world["dz"], world["dt"] - bc_x = world["boundary_conditions"]["x"] - bc_y = world["boundary_conditions"]["y"] - bc_z = world["boundary_conditions"]["z"] - - Jx = Jx.at[:, :, :].set(0) - Jy = Jy.at[:, :, :].set(0) - Jz = Jz.at[:, :, :].set(0) - - 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. - - Nx_ext, Ny_ext, Nz_ext = Nx + 2, Ny + 2, Nz + 2 - # calculate the size of the extended grid with two ghost layers on each side for Esirkepov deposition. - - x_extended = jnp.concatenate([ grid[0][:1] - dx, grid[0], grid[0][-1:] + dx,]) - # add +- dx ghost points to the x grid axis for an extra ghost cell on each side - y_extended = jnp.concatenate([ grid[1][:1] - dy, grid[1], grid[1][-1:] + dy,]) - # add +- dy ghost points to the y grid axis for an extra ghost cell on each side - z_extended = jnp.concatenate([ grid[2][:1] - dz, grid[2], grid[2][-1:] + dz,]) - # add +- dz ghost points to the z grid axis for an extra ghost cell on each side - - extended_grid = (x_extended, y_extended, z_extended) - # build the extended grid tuple for passing to deposition functions - - for species in particles: - Jx_ext = jnp.zeros((Nx_ext, Ny_ext, Nz_ext), dtype=Jx.dtype) - Jy_ext = jnp.zeros_like(Jx_ext) - Jz_ext = jnp.zeros_like(Jx_ext) - # Deposit one species at a time so particle boundary labels can fold - # the extended ghost layers correctly. - - q = species.get_charge() - x, y, z = species.get_forward_position() - vx, vy, vz = species.get_velocity() - active = species.get_active_mask().astype(x.dtype) - 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, extended_grid[0], shape_factor) - y0 = compute_particle_anchor(y, extended_grid[1], shape_factor) - z0 = compute_particle_anchor(z, extended_grid[2], shape_factor) - # compute the nearest grid point to the particle's current position - - old_x0 = compute_particle_anchor(old_x, extended_grid[0], shape_factor) - old_y0 = compute_particle_anchor(old_y, extended_grid[1], shape_factor) - old_z0 = compute_particle_anchor(old_z, extended_grid[2], shape_factor) - # compute the nearest grid point to the particle's old position - - deltax = particle_axis_offset(x, x0, extended_grid[0]) - deltay = particle_axis_offset(y, y0, extended_grid[1]) - deltaz = particle_axis_offset(z, z0, extended_grid[2]) - old_deltax = particle_axis_offset(old_x, old_x0, extended_grid[0]) - old_deltay = particle_axis_offset(old_y, old_y0, extended_grid[1]) - old_deltaz = particle_axis_offset(old_z, old_z0, extended_grid[2]) - # compute the particle's offset from the nearest grid point in each direction at both time steps + # unpack current density + ntx, nty, ntz = Jx.shape[:3] + # get the number of tiles + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + # unpack the tile shapes + local_Nx = tile_nx + 2 * g + local_Ny = tile_ny + 2 * g + local_Nz = tile_nz + 2 * g + # compute the local shape including guard cells + + x_active = ntx * tile_nx > 1 + y_active = nty * tile_ny > 1 + z_active = ntz * tile_nz > 1 + # determine which axes are actually active and which ones are redundant + + Jx_template = jnp.zeros_like(Jx[0, 0, 0]) + Jy_template = jnp.zeros_like(Jy[0, 0, 0]) + Jz_template = jnp.zeros_like(Jz[0, 0, 0]) + # build a template array for the local J tiles + + species_weighted_charge = species_config.charge * species_config.weight + # compute the species weighted charge + + def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): + old_x = x_tile[..., 0].reshape(-1) + old_y = x_tile[..., 1].reshape(-1) + old_z = x_tile[..., 2].reshape(-1) + # get the old positions and reshape them as 1D arrays + vx = u_tile[..., 0].reshape(-1) + vy = u_tile[..., 1].reshape(-1) + vz = u_tile[..., 2].reshape(-1) + # get the velocities and reshape them as 1D arrays + active = active_tile.reshape(-1).astype(old_x.dtype) + # get the active mask and reshape it as a 1D array + q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) + # broadcast the species weighted charge to the shape of the active tile and reshape it as a 1D array + N_particles = old_x.shape[0] + # get the number of particles in the tile + update_x1 = jnp.broadcast_to(species_config.update_x[:, 0, jnp.newaxis], active_tile.shape).reshape(-1) + update_x2 = jnp.broadcast_to(species_config.update_x[:, 1, jnp.newaxis], active_tile.shape).reshape(-1) + update_x3 = jnp.broadcast_to(species_config.update_x[:, 2, jnp.newaxis], active_tile.shape).reshape(-1) + # determine which axes are updated for each particle and reshape them as 1D arrays + + x = old_x + jnp.where(update_x1, vx * dt, 0.0) + y = old_y + jnp.where(update_x2, vy * dt, 0.0) + z = old_z + jnp.where(update_x3, vz * dt, 0.0) + # step the particle positions forward in time using the velocity and dt, but only for the axes that are updated + + x_grid = tiled_grid[0][tx, ty, tz] + y_grid = tiled_grid[1][tx, ty, tz] + z_grid = tiled_grid[2][tx, ty, tz] + # get the local grid for the tile in each axis + + x0 = compute_particle_anchor(x, x_grid, shape_factor) + y0 = compute_particle_anchor(y, y_grid, shape_factor) + z0 = compute_particle_anchor(z, z_grid, shape_factor) + # get the particle anchor points for the new positions in each axis + old_x0 = compute_particle_anchor(old_x, x_grid, shape_factor) + old_y0 = compute_particle_anchor(old_y, y_grid, shape_factor) + old_z0 = compute_particle_anchor(old_z, z_grid, shape_factor) + # get the particle anchor points for the old positions in each axis + + deltax = particle_axis_offset(x, x0, x_grid) + deltay = particle_axis_offset(y, y0, y_grid) + deltaz = particle_axis_offset(z, z0, z_grid) + # compute the particle offsets from the anchor points for the new positions in each axis + old_deltax = particle_axis_offset(old_x, old_x0, x_grid) + old_deltay = particle_axis_offset(old_y, old_y0, y_grid) + old_deltaz = particle_axis_offset(old_z, old_z0, z_grid) + # compute the particle offsets from the anchor points for the old positions in each axis 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. + # get the difference between the new and old anchor points to determine how much the old weights need to be shifted to align with the new anchor points offsets = jnp.asarray([-2, -1, 0, 1, 2], dtype=x0.dtype) - # define the offsets for the 5-point Esirkepov stencil in each direction - xpts = x0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] ypts = y0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] zpts = z0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] - # Build the 5-point stencil in the two-ghost deposition buffer. + # compute the 5-point stencil indices for the new anchor points in each axis xw, yw, zw = jax.lax.cond( shape_factor == 1, @@ -263,149 +331,204 @@ 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 + # get the current weights for the new and old positions based on the shape factor tmp = jnp.zeros_like(xw[0]) - xw = [tmp, xw[0], xw[1], xw[2], tmp] yw = [tmp, yw[0], yw[1], yw[2], tmp] zw = [tmp, zw[0], zw[1], zw[2], tmp] - 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. + # build the 5 point stencil weights for the new and old positions, padding with zeros at the ghost points oxw = shift_old_stencil(oxw, shift_x) oyw = shift_old_stencil(oyw, shift_y) ozw = shift_old_stencil(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_redundant_axis(xpts, xw, oxw, x_active, Nx_ext) - ypts, yw, oyw = collapse_redundant_axis(ypts, yw, oyw, y_active, Ny_ext) - zpts, zw, ozw = collapse_redundant_axis(zpts, zw, ozw, z_active, Nz_ext) - # remove dummy stencil points for non 3D simulations - - if x_active and y_active and z_active: - # if full 3D, calculate full Esirkepov weights with the 3D formula - Wx_, Wy_, Wz_ = get_3D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles) - elif (x_active and y_active and (not z_active)) or (x_active and z_active and (not y_active)) or ( - y_active and z_active and (not x_active) - # if 2D, calculate Esirkepov weights with the 2D formula for the active plane - ): - null_dim = lax.cond( - not x_active, - lambda _: 0, - lambda _: lax.cond( - not y_active, - lambda _: 1, - lambda _: 2, - operand=None, - ), - operand=None, - ) + # shift the old weights to align with the new anchor points based on the computed shifts - Wx_, Wy_, Wz_ = get_2D_esirkepov_weights( - xw, yw, zw, oxw, oyw, ozw, N_particles, null_dim=null_dim - ) - elif x_active and (not y_active) and (not z_active): - # if 1D in x, calculate Esirkepov weights with the 1D formula for x and pad with zeros for y and z - Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=0) - elif y_active and (not x_active) and (not z_active): - # if 1D in y, calculate Esirkepov weights with the 1D formula for y and pad with zeros for x and z - 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): - # if 1D in z, calculate Esirkepov weights with the 1D formula for z and pad with zeros for x and y - 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 + xpts, xw, oxw = collapse_redundant_axis(xpts, xw, oxw, x_active, local_Nx) + ypts, yw, oyw = collapse_redundant_axis(ypts, yw, oyw, y_active, local_Ny) + zpts, zw, ozw = collapse_redundant_axis(zpts, zw, ozw, z_active, local_Nz) + # collapse any redundant axes (if the axis is inactive) to ensure that the weights and points are correctly aligned for deposition dJx = jax.lax.cond( x_active, - lambda _: active * (-(q / (dy * dz)) / dt) * jnp.ones(N_particles), - lambda _: active * q * vx / (dx * dy * dz) * jnp.ones(N_particles), + lambda _: active * (-(q / (dy * dz)) / dt), + lambda _: active * q * vx / (dx * dy * dz), operand=None, ) - dJy = jax.lax.cond( y_active, - lambda _: active * (-(q / (dx * dz)) / dt) * jnp.ones(N_particles), - lambda _: active * q * vy / (dx * dy * dz) * jnp.ones(N_particles), + lambda _: active * (-(q / (dx * dz)) / dt), + lambda _: active * q * vy / (dx * dy * dz), operand=None, ) - dJz = jax.lax.cond( z_active, - lambda _: active * (-(q / (dx * dy)) / dt) * jnp.ones(N_particles), - lambda _: active * q * vz / (dx * dy * dz) * jnp.ones(N_particles), + lambda _: active * (-(q / (dx * dy)) / dt), + lambda _: active * q * vz / (dx * dy * dz), operand=None, ) - # compute the current contribution prefactors for each active dimension according to Esirkepov's formula + # compute the local current contributions for each axis based on whether the axis is active or not, using the Esirkepov formula for charge-conserving current deposition. - Fx = dJx * Wx_ - Fy = dJy * Wy_ - Fz = dJz * Wz_ + tile_Jx = Jx_template + tile_Jy = Jy_template + tile_Jz = Jz_template + # initialize the local tile current arrays to zero - Jx_loc = jnp.zeros_like(Fx) - Jy_loc = jnp.zeros_like(Fy) - Jz_loc = jnp.zeros_like(Fz) + if x_active and y_active and z_active: + # if all three axes are active, compute the 3D Esirkepov weights and deposit the currents accordingly + Wx_, Wy_, Wz_ = get_3D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles) + Fx = dJx * Wx_ + Fy = dJy * Wy_ + Fz = dJz * Wz_ + Jx_loc = jnp.cumsum(Fx, axis=0) + Jy_loc = jnp.cumsum(Fy, axis=1) + Jz_loc = jnp.cumsum(Fz, axis=2) - Jx_loc = jnp.cumsum(Fx, axis=0) - Jy_loc = jnp.cumsum(Fy, axis=1) - Jz_loc = jnp.cumsum(Fz, axis=2) - if x_active: - for i in range(5): - for j in range(5): - for k in range(5): - Jx_ext = Jx_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Jx_loc[i, j, k, :], mode="drop") - else: for i in range(5): for j in range(5): for k in range(5): - Jx_ext = Jx_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Fx[i, j, k, :], mode="drop") - - if y_active: - for i in range(5): + ix = xpts[i, :] + iy = ypts[j, :] + iz = zpts[k, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Jx_loc[i, j, k, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Jy_loc[i, j, k, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Jz_loc[i, j, k, :], mode="drop") + elif (x_active and y_active and (not z_active)) or (x_active and z_active and (not y_active)) or ( + y_active and z_active and (not x_active) + ): # if two axes are active and one is inactive, compute the 2D Esirkepov weights and deposit the currents accordingly + if not x_active: + null_dim = 0 + elif not y_active: + null_dim = 1 + else: + null_dim = 2 + Wx_, Wy_, Wz_ = _compact_2d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, null_dim=null_dim) + Fx = dJx * Wx_ + Fy = dJy * Wy_ + Fz = dJz * Wz_ + + if null_dim == 0: + Jy_loc = jnp.cumsum(Fy, axis=0) + Jz_loc = jnp.cumsum(Fz, axis=1) for j in range(5): for k in range(5): - Jy_ext = Jy_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Jy_loc[i, j, k, :], mode="drop") - else: - for i in range(5): - for j in range(5): + ix = xpts[2, :] + iy = ypts[j, :] + iz = zpts[k, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Fx[j, k, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Jy_loc[j, k, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Jz_loc[j, k, :], mode="drop") + elif null_dim == 1: + Jx_loc = jnp.cumsum(Fx, axis=0) + Jz_loc = jnp.cumsum(Fz, axis=1) + for i in range(5): for k in range(5): - Jy_ext = Jy_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Fy[i, j, k, :], mode="drop") - - if z_active: + ix = xpts[i, :] + iy = ypts[2, :] + iz = zpts[k, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Jx_loc[i, k, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Fy[i, k, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Jz_loc[i, k, :], mode="drop") + else: + Jx_loc = jnp.cumsum(Fx, axis=0) + Jy_loc = jnp.cumsum(Fy, axis=1) + for i in range(5): + for j in range(5): + ix = xpts[i, :] + iy = ypts[j, :] + iz = zpts[2, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Jx_loc[i, j, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Jy_loc[i, j, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Fz[i, j, :], mode="drop") + elif x_active and (not y_active) and (not z_active): + # if only the x-axis is active, compute the 1D Esirkepov weights and deposit the currents accordingly + Wx_, Wy_, Wz_ = _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=0) + Fx = dJx * Wx_ + Fy = dJy * Wy_ + Fz = dJz * Wz_ + Jx_loc = jnp.cumsum(Fx, axis=0) for i in range(5): - for j in range(5): - for k in range(5): - Jz_ext = Jz_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Jz_loc[i, j, k, :], mode="drop") + ix = xpts[i, :] + iy = ypts[2, :] + iz = zpts[2, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Jx_loc[i, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Fy[i, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Fz[i, :], mode="drop") + elif y_active and (not x_active) and (not z_active): + # if only the y-axis is active, compute the 1D Esirkepov weights and deposit the currents accordingly + Wx_, Wy_, Wz_ = _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=1) + Fx = dJx * Wx_ + Fy = dJy * Wy_ + Fz = dJz * Wz_ + Jy_loc = jnp.cumsum(Fy, axis=0) + for j in range(5): + ix = xpts[2, :] + iy = ypts[j, :] + iz = zpts[2, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Fx[j, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Jy_loc[j, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Fz[j, :], mode="drop") else: - for i in range(5): - for j in range(5): - for k in range(5): - Jz_ext = Jz_ext.at[xpts[i, :], ypts[j, :], zpts[k, :]].add(Fz[i, j, k, :], mode="drop") - # deposit the current contributions onto the grid. - - esirkepov_x_bc = species.x_bc - esirkepov_y_bc = species.y_bc - esirkepov_z_bc = species.z_bc - # determine the boundary condition type from the particle species classes - - Jx_ext = ghost_cell_bc_esirkepov(Jx_ext, esirkepov_x_bc, esirkepov_y_bc, esirkepov_z_bc, component_axis=0) - Jy_ext = ghost_cell_bc_esirkepov(Jy_ext, esirkepov_x_bc, esirkepov_y_bc, esirkepov_z_bc, component_axis=1) - Jz_ext = ghost_cell_bc_esirkepov(Jz_ext, esirkepov_x_bc, esirkepov_y_bc, esirkepov_z_bc, component_axis=2) - # Fold the two Esirkepov ghost layers before returning to the solver's - # ordinary one-ghost current arrays. + # if only the z-axis is active, compute the 1D Esirkepov weights and deposit the currents accordingly + Wx_, Wy_, Wz_ = _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=2) + Fx = dJx * Wx_ + Fy = dJy * Wy_ + Fz = dJz * Wz_ + Jz_loc = jnp.cumsum(Fz, axis=0) + for k in range(5): + ix = xpts[2, :] + iy = ypts[2, :] + iz = zpts[k, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add(Fx[k, :], mode="drop") + tile_Jy = tile_Jy.at[ix, iy, iz].add(Fy[k, :], mode="drop") + tile_Jz = tile_Jz.at[ix, iy, iz].add(Jz_loc[k, :], mode="drop") + + return tile_Jx, tile_Jy, tile_Jz + + tx, ty, tz = jnp.meshgrid( + jnp.arange(ntx), + jnp.arange(nty), + jnp.arange(ntz), + indexing="ij", + ) + # build a meshgrid of tile indices for the x, y, and z axes to pass to the deposit function + + deposit_tiles = deposit_one_tile + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + # nested vmap to vectorize the deposit function over the tile indices for x, y, and z axes + + Jx, Jy, Jz = deposit_tiles( + particles.x, + particles.u, + particles.active, + tx, + ty, + tz, + ) + # deposit the currents for all tiles in parallel using the vectorized deposit function - Jx = Jx + eliminate_esirkepov_ghost_cells(Jx_ext) - Jy = Jy + eliminate_esirkepov_ghost_cells(Jy_ext) - Jz = Jz + eliminate_esirkepov_ghost_cells(Jz_ext) + J = fold_tiled_esirkepov_vector_ghost_cells((Jx, Jy, Jz), world, num_guard_cells=g, tile_shape=tile_shape) + # fold the deposited currents across tile boundaries, applying the appropriate boundary conditions for ghost cells + J = update_tiled_vector_ghost_cells(J, world, num_guard_cells=g, tile_shape=tile_shape) + # update the ghost cells of the folded currents to ensure consistency across tile boundaries - 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) - # refresh the ordinary one-ghost current halos before the field update + if filter == "bilinear": + J = bilinear_filter_vector(J, num_guard_cells=g) + # apply a bilinear filter to the current density tiles if specified in the filter argument + J = update_tiled_vector_ghost_cells(J, world, g, tile_shape) + # update the ghost cells of the current density tiles to reflect the contributions from neighboring tiles + elif filter == "digital": + J = digital_filter_vector(J, constants["alpha"], num_guard_cells=g) + # apply a digital filter to the current density tiles if specified in the filter argument + J = update_tiled_vector_ghost_cells(J, world, num_guard_cells=g, tile_shape=tile_shape) + # update the ghost cells of the filtered current density tiles to ensure continuity across tile boundaries - return (Jx, Jy, Jz) + return J def get_3D_esirkepov_weights( @@ -442,106 +565,4 @@ def get_3D_esirkepov_weights( ) ) - return Wx_, Wy_, Wz_ - - -def get_2D_esirkepov_weights( - x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights, N_particles, null_dim=2 -): - d_Sx = [x_weights[i] - old_x_weights[i] for i in range(len(x_weights))] - d_Sy = [y_weights[i] - old_y_weights[i] for i in range(len(y_weights))] - d_Sz = [z_weights[i] - old_z_weights[i] for i in range(len(z_weights))] - - Wx_ = jnp.zeros((len(x_weights), len(y_weights), len(z_weights), N_particles)) - Wy_ = jnp.zeros_like(Wx_) - Wz_ = jnp.zeros_like(Wx_) - - def xy_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for i in range(len(x_weights)): - for j in range(len(y_weights)): - Wx_ = Wx_.at[i, j, 2, :].set(1 / 2 * d_Sx[i] * (y_weights[j] + old_y_weights[j])) - Wy_ = Wy_.at[i, j, 2, :].set(1 / 2 * d_Sy[j] * (x_weights[i] + old_x_weights[i])) - Wz_ = Wz_.at[i, j, 2, :].set( - 1 / 3 * (x_weights[i] * y_weights[j] + old_x_weights[i] * old_y_weights[j]) - + 1 / 6 * (x_weights[i] * old_y_weights[j] + old_x_weights[i] * y_weights[j]) - ) - return Wx_, Wy_, Wz_ - - def xz_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for i in range(len(x_weights)): - for k in range(len(z_weights)): - Wx_ = Wx_.at[i, 2, k, :].set(1 / 2 * d_Sx[i] * (z_weights[k] + old_z_weights[k])) - Wy_ = Wy_.at[i, 2, k, :].set( - 1 / 3 * (x_weights[i] * z_weights[k] + old_x_weights[i] * old_z_weights[k]) - + 1 / 6 * (x_weights[i] * old_z_weights[k] + old_x_weights[i] * z_weights[k]) - ) - Wz_ = Wz_.at[i, 2, k, :].set(1 / 2 * d_Sz[k] * (x_weights[i] + old_x_weights[i])) - return Wx_, Wy_, Wz_ - - def yz_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for j in range(len(y_weights)): - for k in range(len(z_weights)): - Wx_ = Wx_.at[2, j, k, :].set( - 1 / 3 * (y_weights[j] * z_weights[k] + old_y_weights[j] * old_z_weights[k]) - + 1 / 6 * (y_weights[j] * old_z_weights[k] + old_y_weights[j] * z_weights[k]) - ) - Wy_ = Wy_.at[2, j, k, :].set(1 / 2 * d_Sy[j] * (z_weights[k] + old_z_weights[k])) - Wz_ = Wz_.at[2, j, k, :].set(1 / 2 * d_Sz[k] * (y_weights[j] + old_y_weights[j])) - return Wx_, Wy_, Wz_ - - Wx_, Wy_, Wz_ = lax.cond( - null_dim == 0, - lambda _: yz_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - lambda _: lax.cond( - null_dim == 1, - lambda _: xz_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - lambda _: xy_plane(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - operand=None, - ), - operand=None, - ) - - return Wx_, Wy_, Wz_ - - -def get_1D_esirkepov_weights( - x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights, N_particles, dim=0 -): - Wx_ = jnp.zeros((len(x_weights), len(y_weights), len(z_weights), N_particles)) - Wy_ = jnp.zeros_like(Wx_) - Wz_ = jnp.zeros_like(Wx_) - - def x_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for i in range(len(x_weights)): - Wx_ = Wx_.at[i, 2, 2, :].set((x_weights[i] - old_x_weights[i])) - Wy_ = Wy_.at[i, 2, 2, :].set((x_weights[i] + old_x_weights[i]) / 2) - Wz_ = Wz_.at[i, 2, 2, :].set((x_weights[i] + old_x_weights[i]) / 2) - return Wx_, Wy_, Wz_ - - def y_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for j in range(len(y_weights)): - Wy_ = Wy_.at[2, j, 2, :].set((y_weights[j] - old_y_weights[j])) - Wx_ = Wx_.at[2, j, 2, :].set((y_weights[j] + old_y_weights[j]) / 2) - Wz_ = Wz_.at[2, j, 2, :].set((y_weights[j] + old_y_weights[j]) / 2) - return Wx_, Wy_, Wz_ - - def z_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights): - for k in range(len(z_weights)): - Wz_ = Wz_.at[2, 2, k, :].set((z_weights[k] - old_z_weights[k])) - Wx_ = Wx_.at[2, 2, k, :].set((z_weights[k] + old_z_weights[k]) / 2) - Wy_ = Wy_.at[2, 2, k, :].set((z_weights[k] + old_z_weights[k]) / 2) - return Wx_, Wy_, Wz_ - - Wx_, Wy_, Wz_ = lax.cond( - dim == 0, - lambda _: x_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - lambda _: lax.cond( - dim == 1, - lambda _: y_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - lambda _: z_active(Wx_, Wy_, Wz_, x_weights, y_weights, z_weights, old_x_weights, old_y_weights, old_z_weights), - operand=None, - ), - operand=None, - ) - - return Wx_, Wy_, Wz_ + return Wx_, Wy_, Wz_ \ No newline at end of file diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index b813f8f8..a83f3a6a 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -1,101 +1,143 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial +from PyPIC3D.particles.particle_class import TiledParticles, SpeciesConfig from PyPIC3D.boundary_conditions.grid_and_stencil import ( - axis_has_active_cells, collapse_axis_stencil, prepare_particle_axis_stencil, ) -from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells + from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.utils import bilinear_filter +from PyPIC3D.boundary_conditions.ghost_cells import ( + fold_tiled_vector_ghost_cells, + update_tiled_vector_ghost_cells, +) +from PyPIC3D.utilities.filters import ( + bilinear_filter_vector, + digital_filter_vector, +) + +import jax +import jax.numpy as jnp -@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.""" +def _collapse_tiled_axis_stencil(points, weights, local_n, reduced_axis, g): + if reduced_axis: + collapsed_points = jnp.full((1, points.shape[1]), int(g), dtype=points.dtype) + collapsed_weights = jnp.sum(weights, axis=0, keepdims=True) + return collapsed_points, collapsed_weights + return collapse_axis_stencil(points, weights, local_n, ghost_cells=True) - if grid is None: - grid = world["grids"]["center"] + +def J_from_rhov( + particles, + species_config, + J=None, + constants=None, + world=None, + filter="bilinear", +): + """Compute tile-local direct current from centered tiled particles.""" + + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + g = int(g) + # determine the number of guard cells and the shape of each of the tiles + + tiled_grid = world["grids"]["tiled_center_grid"] + # get the grid for the tiles dx = world["dx"] dy = world["dy"] dz = world["dz"] - - Jx, Jy, Jz = J - Nx, Ny, Nz = Jx.shape - # 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() - active = species.get_active_mask().astype(x.dtype) - # 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 + # get the world resolution + + Jx_tiles, Jy_tiles, Jz_tiles = J + # unpack the current density tiles + + ntx, nty, ntz = Jx_tiles.shape[:3] + # get the number of tiles in each dimension + tile_nx, tile_ny, tile_nz = tile_shape + # unpack the tile shape + local_Nx = tile_nx + 2 * g + local_Ny = tile_ny + 2 * g + local_Nz = tile_nz + 2 * g + # piece together the total local tile shape + + shape_factor = world["shape_factor"] + # get the shape factor + + reduced_x = int(tile_nx) == 1 and int(ntx) == 1 + reduced_y = int(tile_ny) == 1 and int(nty) == 1 + reduced_z = int(tile_nz) == 1 and int(ntz) == 1 + # determine if any of the axes are dummy axes + + Jx_template = jnp.zeros_like(Jx_tiles[0, 0, 0]) + Jy_template = jnp.zeros_like(Jy_tiles[0, 0, 0]) + Jz_template = jnp.zeros_like(Jz_tiles[0, 0, 0]) + # build template tiles + + # Tile boundaries are not physical boundaries. Deposits that cross a tile + # edge should land in tile ghost cells and be exchanged by the tiled fold. + local_bc = 2 + + species_weighted_charge = species_config.charge * species_config.weight + # compute the weighted charge for each species + + def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): + # deposit the current density for a single tile, given the particle positions, velocities, and active mask + x = x_tile[..., 0].reshape(-1) + y = x_tile[..., 1].reshape(-1) + z = x_tile[..., 2].reshape(-1) + # reshape the particle positions into 1D arrays for processing + vx = u_tile[..., 0].reshape(-1) + vy = u_tile[..., 1].reshape(-1) + vz = u_tile[..., 2].reshape(-1) + # reshape the particle velocities into 1D arrays for processing + active = active_tile.reshape(-1).astype(x.dtype) + # reshape the active particle mask into a 1D array for processing + q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) + # reshape the particle charges into a 1D array for processing + dq = q / (dx * dy * dz) + # compute the charge density contribution of each particle + + x_grid = tiled_grid[0][tx, ty, tz] + y_grid = tiled_grid[1][tx, ty, tz] + z_grid = tiled_grid[2][tx, ty, tz] + # get the grid points for the current tile in each dimension x, x0, deltax_node, xpts = prepare_particle_axis_stencil( x, - grid[0], - Nx, + x_grid, + local_Nx, shape_factor, - bc_x, - wind=world["x_wind"], + local_bc, + wind=tile_nx * dx, 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], - Ny, + y_grid, + local_Ny, shape_factor, - bc_y, - wind=world["y_wind"], + local_bc, + wind=tile_ny * dy, 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], - Nz, + z_grid, + local_Nz, shape_factor, - bc_z, - wind=world["z_wind"], + local_bc, + wind=tile_nz * dz, 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. + ) + # prepare the particle positions and compute the stencil points and deltas for each axis - 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. + deltax_face = (x - x_grid[0]) - (x0 + 0.5) * dx + deltay_face = (y - y_grid[0]) - (y0 + 0.5) * dy + deltaz_face = (z - z_grid[0]) - (z0 + 0.5) * dz + # compute the deltas for the face-centered weights based on the particle positions and grid points x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( shape_factor == 1, @@ -103,116 +145,96 @@ 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, lambda _: get_first_order_weights(deltax_face, deltay_face, deltaz_face, dx, dy, dz), lambda _: get_second_order_weights(deltax_face, deltay_face, deltaz_face, dx, dy, dz), operand=None, ) - # compute the face-centered weights for deposition + # compute the weights for the node-centered and face-centered contributions based on the shape factor and deltas xpts = jnp.asarray(xpts) ypts = jnp.asarray(ypts) zpts = jnp.asarray(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) - z_weights_face = jnp.asarray(z_weights_face) - 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 - zpts_eff = zpts - x_weights_node_eff = x_weights_node - y_weights_node_eff = y_weights_node - z_weights_node_eff = z_weights_node - x_weights_face_eff = x_weights_face - y_weights_face_eff = y_weights_face - z_weights_face_eff = z_weights_face - - 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", - ) - 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 - ix = xpts_eff[i, ...] - iy = ypts_eff[j, ...] - iz = zpts_eff[k, ...] - valx = ( - (active * dq * vx) - * x_weights_face_eff[i, ...] - * y_weights_node_eff[j, ...] - * z_weights_node_eff[k, ...] - ) - valy = ( - (active * dq * vy) - * x_weights_node_eff[i, ...] - * y_weights_face_eff[j, ...] - * z_weights_node_eff[k, ...] - ) - valz = ( - (active * dq * vz) - * x_weights_node_eff[i, ...] - * y_weights_node_eff[j, ...] - * z_weights_face_eff[k, ...] - ) - 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) - Jz = fold_ghost_cells(Jz, bc_x, bc_y, bc_z) - # fold ghost cell deposits back into interior - 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) - # refresh ghost cells before any stencil-based post-processing - - def filter_func(J_, filter): - J_ = jax.lax.cond( - filter == "bilinear", - lambda J_: bilinear_filter(J_), - lambda J_: J_, - operand=J_, - ) - return J_ - - 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) + x_weights_face = jnp.asarray(x_weights_face) + y_weights_face = jnp.asarray(y_weights_face) + z_weights_face = jnp.asarray(z_weights_face) + # convert the stencil points and weights to JAX arrays for further processing + + xpts, x_weights_node = _collapse_tiled_axis_stencil(xpts, x_weights_node, local_Nx, reduced_x, g) + _, x_weights_face = _collapse_tiled_axis_stencil(xpts, x_weights_face, local_Nx, reduced_x, g) + ypts, y_weights_node = _collapse_tiled_axis_stencil(ypts, y_weights_node, local_Ny, reduced_y, g) + _, y_weights_face = _collapse_tiled_axis_stencil(ypts, y_weights_face, local_Ny, reduced_y, g) + zpts, z_weights_node = _collapse_tiled_axis_stencil(zpts, z_weights_node, local_Nz, reduced_z, g) + _, z_weights_face = _collapse_tiled_axis_stencil(zpts, z_weights_face, local_Nz, reduced_z, g) + # collapse the stencil points and weights for each axis, taking into account any reduced axes and guard cells + + tile_Jx = Jx_template + tile_Jy = Jy_template + tile_Jz = Jz_template + + for i in range(xpts.shape[0]): + for j in range(ypts.shape[0]): + for k in range(zpts.shape[0]): + ix = xpts[i] + iy = ypts[j] + iz = zpts[k] + tile_Jx = tile_Jx.at[ix, iy, iz].add( + active * dq * vx * x_weights_face[i] * y_weights_node[j] * z_weights_node[k], + mode="drop", + ) + tile_Jy = tile_Jy.at[ix, iy, iz].add( + active * dq * vy * x_weights_node[i] * y_weights_face[j] * z_weights_node[k], + mode="drop", + ) + tile_Jz = tile_Jz.at[ix, iy, iz].add( + active * dq * vz * x_weights_node[i] * y_weights_node[j] * z_weights_face[k], + mode="drop", + ) + + return tile_Jx, tile_Jy, tile_Jz + + tx, ty, tz = jnp.meshgrid( + jnp.arange(ntx), + jnp.arange(nty), + jnp.arange(ntz), + indexing="ij", + ) + # build the tile index arrays for each dimension + + deposit_tiles = deposit_one_tile + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) + # vectorize the deposit_one_tile function over the tile indices using jax.vmap + + Jx, Jy, Jz = deposit_tiles( + particles.x, + particles.u, + particles.active, + tx, + ty, + tz, + ) + # compute the current density contributions for all tiles by applying the vectorized deposit function to the particle data and tile indices + + J = fold_tiled_vector_ghost_cells((Jx, Jy, Jz), world, g, tile_shape) + # fold the ghost cells of the current density tiles to ensure continuity across tile boundaries + J = update_tiled_vector_ghost_cells(J, world, g, tile_shape) + # update the ghost cells of the current density tiles to reflect the contributions from neighboring tiles + + if filter == "bilinear": + J = bilinear_filter_vector(J, num_guard_cells=g) + # apply a bilinear filter to the current density tiles if specified in the filter argument + J = update_tiled_vector_ghost_cells(J, world, g, tile_shape) + # update the ghost cells of the current density tiles to reflect the contributions from neighboring tiles + elif filter == "digital": + J = digital_filter_vector(J, constants["alpha"], num_guard_cells=g) + # apply a digital filter to the current density tiles if specified in the filter argument + J = update_tiled_vector_ghost_cells(J, world, num_guard_cells=g, tile_shape=tile_shape) + # update the ghost cells of the filtered current density tiles to ensure continuity across tile boundaries + + return J diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 566015ae..c021e104 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -1,104 +1,429 @@ +from functools import partial + import jax from jax import jit import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import ( + BC_PERIODIC, collapse_axis_stencil, prepare_particle_axis_stencil, ) +from PyPIC3D.boundary_conditions.ghost_cells import update_tiled_ghost_cells from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells -from PyPIC3D.utils import digital_filter +from PyPIC3D.particles.particle_class import TiledParticles +from PyPIC3D.utilities.filters import digital_filter + + +def _species_scalar_to_slots(tiled_particles, species_value): + return jnp.broadcast_to( + species_value.reshape((1, 1, 1, species_value.shape[0], 1)), + tiled_particles.active.shape, + ) + + +def _species_slot_counts(tiled_particles): + occupied = tiled_particles.active | jnp.any(tiled_particles.x != 0.0, axis=-1) | jnp.any(tiled_particles.u != 0.0, axis=-1) + species_counts = jnp.sum(occupied.astype(tiled_particles.x.dtype), axis=(0, 1, 2, 4)) + return jnp.where(species_counts > 0, species_counts, 1.0) + + +def _diagnostic_mass_position(tiled_particles, world): + """ + Mirror the ordinary scalar-diagnostic particle position for tiled particles. + """ + x = tiled_particles.x[..., 0] - tiled_particles.u[..., 0] * world["dt"] / 2 + y = tiled_particles.x[..., 1] - tiled_particles.u[..., 1] * world["dt"] / 2 + z = tiled_particles.x[..., 2] - tiled_particles.u[..., 2] * world["dt"] / 2 -@jit -def compute_rho(particles, rho, world, constants): - """Compute the charge density (rho) on the vertex grid.""" + x = jax.lax.cond( + world["boundary_conditions"]["x"] == BC_PERIODIC, + lambda value: jnp.where( + value > world["x_wind"] / 2, + value - world["x_wind"], + jnp.where(value < -world["x_wind"] / 2, value + world["x_wind"], value), + ), + lambda value: value, + x, + ) + y = jax.lax.cond( + world["boundary_conditions"]["y"] == BC_PERIODIC, + lambda value: jnp.where( + value > world["y_wind"] / 2, + value - world["y_wind"], + jnp.where(value < -world["y_wind"] / 2, value + world["y_wind"], value), + ), + lambda value: value, + y, + ) + z = jax.lax.cond( + world["boundary_conditions"]["z"] == BC_PERIODIC, + lambda value: jnp.where( + value > world["z_wind"] / 2, + value - world["z_wind"], + jnp.where(value < -world["z_wind"] / 2, value + world["z_wind"], value), + ), + lambda value: value, + z, + ) + + return x, y, z + + +def _deposit_tiled_scalar_moment_to_tiles( + tiled_particles, + rho_tiles, + world, + grid, + scalar_weight, + tile_shape, + g, + position=None, + divide_by_volume=True, +): dx = world["dx"] dy = world["dy"] dz = world["dz"] - grid = world["grids"]["vertex"] - Nx, Ny, Nz = rho.shape bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] - rho = jnp.zeros_like(rho) - - for species in particles: - shape_factor = species.get_shape() - q = species.get_charge() - dq = q / dx / dy / dz - x, y, z = species.get_forward_position() - vx, vy, vz = species.get_velocity() - active = species.get_active_mask().astype(x.dtype) - dt = species.dt - # Deposit from the unwrapped half-step-back position so periodic seam - # contributions land in the ghost cells instead of being pre-wrapped. - x = x - vx * dt / 2 - y = y - vy * dt / 2 - z = z - vz * dt / 2 - - x, _, deltax, xpts = prepare_particle_axis_stencil( - x, - grid[0], - Nx, - shape_factor, - bc_x, - wind=world["x_wind"], - ghost_cells=True, - ) - y, _, deltay, ypts = prepare_particle_axis_stencil( - y, - grid[1], - Ny, - shape_factor, - bc_y, - wind=world["y_wind"], - ghost_cells=True, + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(g) + Nx = int(grid[0].shape[0]) - 2 + Ny = int(grid[1].shape[0]) - 2 + Nz = int(grid[2].shape[0]) - 2 + + if position is None: + x = tiled_particles.x[..., 0].reshape(-1) + y = tiled_particles.x[..., 1].reshape(-1) + z = tiled_particles.x[..., 2].reshape(-1) + else: + x = position[0].reshape(-1) + y = position[1].reshape(-1) + z = position[2].reshape(-1) + active = tiled_particles.active.reshape(-1).astype(x.dtype) + scalar_density = scalar_weight.reshape(-1) + if divide_by_volume: + scalar_density = scalar_density / (dx * dy * dz) + + x, _, deltax, xpts = prepare_particle_axis_stencil( + x, + grid[0], + Nx + 2, + shape_factor, + bc_x, + wind=world["x_wind"], + ghost_cells=True, + ) + y, _, deltay, ypts = prepare_particle_axis_stencil( + y, + grid[1], + Ny + 2, + shape_factor, + bc_y, + wind=world["y_wind"], + ghost_cells=True, + ) + z, _, deltaz, zpts = prepare_particle_axis_stencil( + z, + grid[2], + Nz + 2, + shape_factor, + bc_z, + wind=world["z_wind"], + ghost_cells=True, + ) + + 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, + ) + + xpts = jnp.asarray(xpts) + ypts = jnp.asarray(ypts) + zpts = jnp.asarray(zpts) + x_weights = jnp.asarray(x_weights) + y_weights = jnp.asarray(y_weights) + z_weights = jnp.asarray(z_weights) + + xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx + 2, ghost_cells=True) + ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny + 2, ghost_cells=True) + zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz + 2, ghost_cells=True) + + rho_tiles = jnp.zeros_like(rho_tiles) + + def folded_owner(points, n_cells, cells_per_tile, bc): + folded = jax.lax.cond( + bc == BC_PERIODIC, + lambda p: jnp.where(p == 0, n_cells, jnp.where(p == n_cells + 1, 1, p)), + lambda p: jnp.where(p == 0, 1, jnp.where(p == n_cells + 1, n_cells, p)), + points, ) - z, _, deltaz, zpts = prepare_particle_axis_stencil( - z, - grid[2], - Nz, - shape_factor, - bc_z, - wind=world["z_wind"], - ghost_cells=True, + sign = jax.lax.cond( + bc == BC_PERIODIC, + lambda p: jnp.ones_like(p, dtype=x.dtype), + lambda p: jnp.where((p == 0) | (p == n_cells + 1), -1.0, 1.0).astype(x.dtype), + points, ) + valid = (points >= 0) & (points <= n_cells + 1) + folded = jnp.clip(folded, 1, n_cells) + tile = (folded - 1) // cells_per_tile + local = g + (folded - 1) - tile * cells_per_tile + return tile, local, sign, valid - 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, - ) + for i in range(xpts.shape[0]): + tx, lx, sx, valid_x = folded_owner(xpts[i], Nx, tile_nx, bc_x) + for j in range(ypts.shape[0]): + ty, ly, sy, valid_y = folded_owner(ypts[j], Ny, tile_ny, bc_y) + for k in range(zpts.shape[0]): + tz, lz, sz, valid_z = folded_owner(zpts[k], Nz, tile_nz, bc_z) + valid = valid_x & valid_y & valid_z + value = ( + valid.astype(x.dtype) + * active + * sx + * sy + * sz + * scalar_density + * x_weights[i] + * y_weights[j] + * z_weights[k] + ) + rho_tiles = rho_tiles.at[tx, ty, tz, lx, ly, lz].add(value, mode="drop") + + rho_tiles = update_tiled_ghost_cells(rho_tiles, world, g, tile_shape) + + return rho_tiles + + +def _filter_tiled_scalar_with_halos(field_tiles, alpha, world, g, tile_shape): + g = int(g) + field_tiles = update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + return digital_filter(field_tiles, alpha, num_guard_cells=g) + + +def compute_rho( + particles, + species_config, + rho_tiles, + constants, + world, +): + """Compute charge density into tile-major vertex-grid scalar arrays.""" + + if not isinstance(particles, TiledParticles): + raise ValueError("Public compute_rho requires TiledParticles.") + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + grid = world["grids"]["vertex"] - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - x_weights = jnp.asarray(x_weights) - y_weights = jnp.asarray(y_weights) - z_weights = jnp.asarray(z_weights) - - xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=True) - ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=True) - zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=True) - - for i in range(xpts.shape[0]): - for j in range(ypts.shape[0]): - for k in range(zpts.shape[0]): - rho = rho.at[xpts[i], ypts[j], zpts[k]].add( - active * dq * x_weights[i] * y_weights[j] * z_weights[k], mode="drop" - ) - - rho = fold_ghost_cells(rho, bc_x, bc_y, bc_z) - # fold ghost cell deposits back into interior - rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) - # refresh ghost cells before any stencil-based post-processing + scalar_weight = _species_scalar_to_slots( + particles, + species_config.charge * species_config.weight, + ) + rho_tiles = _deposit_tiled_scalar_moment_to_tiles( + particles, + rho_tiles, + world, + grid, + scalar_weight, + tile_shape, + g, + ) alpha = constants["alpha"] - rho = digital_filter(rho, alpha) - rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) + rho_tiles = _filter_tiled_scalar_with_halos(rho_tiles, alpha, world, g, tile_shape) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world, g, tile_shape) + + return rho_tiles + + +@partial(jit, static_argnames=("tile_shape", "g")) +def compute_tiled_mass_density_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, grid=None, tile_shape=None, g=2): + """Compute mass density into tile-major vertex-grid scalar arrays.""" + if grid is None: + grid = world["grids"]["vertex"] + + scalar_weight = _species_scalar_to_slots( + tiled_particles, + species_config.mass * species_config.weight, + ) + position = _diagnostic_mass_position(tiled_particles, world) + return _deposit_tiled_scalar_moment_to_tiles( + tiled_particles, + rho_tiles, + world, + grid, + scalar_weight, + tile_shape, + g, + position=position, + ) + + +@partial(jit, static_argnames=("direction", "tile_shape", "g")) +def compute_tiled_velocity_field_from_tiled_particles( + tiled_particles, + species_config, + field_tiles, + direction, + world, + grid=None, + tile_shape=None, + g=2, +): + """Compute the legacy scalar velocity diagnostic into tile-major arrays.""" + if grid is None: + grid = world["grids"]["vertex"] + + species_counts = _species_slot_counts(tiled_particles) + species_counts = species_counts.reshape((1, 1, 1, species_counts.shape[0], 1)) + scalar_weight = tiled_particles.u[..., direction] / species_counts + position = _diagnostic_mass_position(tiled_particles, world) + + return _deposit_tiled_scalar_moment_to_tiles( + tiled_particles, + field_tiles, + world, + grid, + scalar_weight, + tile_shape, + g, + position=position, + divide_by_volume=False, + ) + + +@partial(jit, static_argnames=("direction", "tile_shape", "g")) +def compute_tiled_pressure_field_from_tiled_particles( + tiled_particles, + species_config, + field_tiles, + velocity_field_tiles, + direction, + world, + grid=None, + tile_shape=None, + g=2, +): + """Compute the legacy scalar pressure diagnostic into tile-major arrays.""" + if grid is None: + grid = world["grids"]["vertex"] + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(g) + Nx = int(grid[0].shape[0]) - 2 + Ny = int(grid[1].shape[0]) - 2 + Nz = int(grid[2].shape[0]) - 2 + + position = _diagnostic_mass_position(tiled_particles, world) + x = position[0].reshape(-1) + y = position[1].reshape(-1) + z = position[2].reshape(-1) + v = tiled_particles.u[..., direction].reshape(-1) + active = tiled_particles.active.reshape(-1).astype(x.dtype) + + x, _, deltax, xpts = prepare_particle_axis_stencil( + x, + grid[0], + Nx + 2, + shape_factor, + bc_x, + wind=world["x_wind"], + ghost_cells=True, + ) + y, _, deltay, ypts = prepare_particle_axis_stencil( + y, + grid[1], + Ny + 2, + shape_factor, + bc_y, + wind=world["y_wind"], + ghost_cells=True, + ) + z, _, deltaz, zpts = prepare_particle_axis_stencil( + z, + grid[2], + Nz + 2, + shape_factor, + bc_z, + wind=world["z_wind"], + ghost_cells=True, + ) + + 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, + ) + + xpts = jnp.asarray(xpts) + ypts = jnp.asarray(ypts) + zpts = jnp.asarray(zpts) + x_weights = jnp.asarray(x_weights) + y_weights = jnp.asarray(y_weights) + z_weights = jnp.asarray(z_weights) + + xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx + 2, ghost_cells=True) + ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny + 2, ghost_cells=True) + zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz + 2, ghost_cells=True) + + field_tiles = jnp.zeros_like(field_tiles) + + def folded_owner(points, n_cells, cells_per_tile, bc): + folded = jax.lax.cond( + bc == BC_PERIODIC, + lambda p: jnp.where(p == 0, n_cells, jnp.where(p == n_cells + 1, 1, p)), + lambda p: jnp.where(p == 0, 1, jnp.where(p == n_cells + 1, n_cells, p)), + points, + ) + sign = jax.lax.cond( + bc == BC_PERIODIC, + lambda p: jnp.ones_like(p, dtype=x.dtype), + lambda p: jnp.where((p == 0) | (p == n_cells + 1), -1.0, 1.0).astype(x.dtype), + points, + ) + valid = (points >= 0) & (points <= n_cells + 1) + folded = jnp.clip(folded, 1, n_cells) + tile = (folded - 1) // cells_per_tile + local = g + (folded - 1) - tile * cells_per_tile + return tile, local, sign, valid + + for i in range(xpts.shape[0]): + tx, lx, sx, valid_x = folded_owner(xpts[i], Nx, tile_nx, bc_x) + for j in range(ypts.shape[0]): + ty, ly, sy, valid_y = folded_owner(ypts[j], Ny, tile_ny, bc_y) + for k in range(zpts.shape[0]): + tz, lz, sz, valid_z = folded_owner(zpts[k], Nz, tile_nz, bc_z) + valid = valid_x & valid_y & valid_z + velocity_value = velocity_field_tiles[tx, ty, tz, lx, ly, lz] + vbar = v - velocity_value + value = ( + valid.astype(x.dtype) + * active + * sx + * sy + * sz + * vbar**2 + * x_weights[i] + * y_weights[j] + * z_weights[k] + ) + field_tiles = field_tiles.at[tx, ty, tz, lx, ly, lz].add(value, mode="drop") - return rho + field_tiles = update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + return field_tiles diff --git a/PyPIC3D/diagnostics/async_writer.py b/PyPIC3D/diagnostics/async_writer.py new file mode 100644 index 00000000..2d31e6c9 --- /dev/null +++ b/PyPIC3D/diagnostics/async_writer.py @@ -0,0 +1,485 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List, Mapping, Sequence, Tuple, Union +import queue +import threading +import traceback + +import jax +import numpy as np + +from PyPIC3D.diagnostics.openPMD import ( + TiledMeshLayout, + write_tiled_field_snapshot_openpmd, + write_tiled_particle_snapshot_openpmd, +) + +ArrayLike = Any +ShardIndex = Tuple[Union[slice, int], ...] +HostShard = Tuple[ShardIndex, np.ndarray] +HostShardList = List[HostShard] +FieldValue = Union[ArrayLike, Sequence[ArrayLike]] +SnapshotFieldValue = Union[HostShardList, Tuple[HostShardList, HostShardList, HostShardList]] + + +_OPENPMD_WRITE_LOCK = threading.Lock() + + +@dataclass(frozen=True) +class TiledFieldSnapshot: + """ + Host-owned snapshot of tiled field data. + + Scalar fields store a list of ``(shard_index, host_array)`` chunks. Vector + fields store one such list for each component. The first three array axes + are tile indices; the last three are tile-local grid indices including + guard cells. + """ + + step: int + time: float + fields: Mapping[str, SnapshotFieldValue] + + +@dataclass(frozen=True) +class TiledParticleSnapshot: + """ + Host-owned snapshot of tiled particle data. + + The snapshot keeps the fixed-capacity tile/species/slot layout. The writer + thread compacts active slots by species before writing openPMD records. + """ + + step: int + time: float + species_names: Tuple[str, ...] + x_shards: HostShardList + u_shards: HostShardList + active_shards: HostShardList + species_charge: np.ndarray + species_mass: np.ndarray + species_weight: np.ndarray + + +def _fields_to_tiled_output_map(fields): + """Return the tiled field components currently written by field diagnostics.""" + E, B, J, rho, phi, external_fields, *rest = fields + external_E, external_B = external_fields + return { + "E": E, + "B": B, + "J": J, + "rho": rho, + "phi": phi, + "external_E": external_E, + "external_B": external_B, + } + + +def _copy_array_to_host_shards(arr): + """ + Copy a possibly-sharded JAX array to host-owned NumPy chunks. + + Only addressable shards are copied. In ordinary single-process runs this is + the complete tiled field. Multi-process distributed output can build on this + same shard contract later. + """ + if hasattr(arr, "addressable_shards"): + out = [] + for shard in arr.addressable_shards: + host = np.array(jax.device_get(shard.data), copy=True, order="C") + out.append((tuple(shard.index), host)) + return out + + try: + arr = jax.device_get(arr) + except Exception: + pass + + host = np.array(arr, copy=True, order="C") + full_index = tuple(slice(0, n) for n in host.shape) + return [(full_index, host)] + + +def prefetch_field_map_to_host(field_map): + """ + Start asynchronous copies for JAX arrays that support it. + + ``make_tiled_field_snapshot`` still performs the real host copy before the + snapshot enters the queue. + """ + def _prefetch(value): + if hasattr(value, "copy_to_host_async"): + value.copy_to_host_async() + + for value in field_map.values(): + if isinstance(value, (tuple, list)): + for component in value: + _prefetch(component) + else: + _prefetch(value) + + +def make_tiled_field_snapshot(field_map, *, step, time): + copied = {} + + for name, value in field_map.items(): + is_vector = isinstance(value, (tuple, list)) and len(value) == 3 + if is_vector: + copied[name] = tuple(_copy_array_to_host_shards(component) for component in value) + else: + copied[name] = _copy_array_to_host_shards(value) + + return TiledFieldSnapshot( + step=int(step), + time=float(time), + fields=copied, + ) + + +def prefetch_tiled_particles_to_host(particles): + for value in (particles.x, particles.u, particles.active): + if hasattr(value, "copy_to_host_async"): + value.copy_to_host_async() + + +def _species_names_for_output(species_names, n_species): + if species_names is None: + return tuple(f"species_{species_index}" for species_index in range(n_species)) + return tuple(str(name) for name in species_names) + + +def _species_array_to_host(value): + value = jax.device_get(value) + return np.asarray(value, dtype=np.float64) + + +def make_tiled_particle_snapshot( + particles, + *, + step, + time, + species_names, + species_config, +): + names = _species_names_for_output(species_names, int(particles.active.shape[3])) + + return TiledParticleSnapshot( + step=int(step), + time=float(time), + species_names=names, + x_shards=_copy_array_to_host_shards(particles.x), + u_shards=_copy_array_to_host_shards(particles.u), + active_shards=_copy_array_to_host_shards(particles.active), + species_charge=_species_array_to_host(species_config.charge), + species_mass=_species_array_to_host(species_config.mass), + species_weight=_species_array_to_host(species_config.weight), + ) + + +class AsyncTiledOpenPMDFieldWriter: + """ + Bounded background queue writer for tiled openPMD field output. + + The simulation thread copies tile-major fields to a host snapshot and puts + that snapshot in a bounded queue. The writer thread strips guard cells and + writes tile interiors as chunks in the global openPMD mesh. + """ + + def __init__( + self, + *, + output_dir, + filename, + world, + global_shape, + tile_shape, + guard_cells=1, + active_dims=(1, 1, 1), + file_extension=".bp", + dtype=np.float64, + queue_size=2, + raise_writer_errors_on_close=True, + ): + queue_size = max(1, int(queue_size)) + + self.output_dir = output_dir + self.filename = filename + self.world = world + self.layout = TiledMeshLayout( + global_shape=tuple(int(width) for width in global_shape), + tile_shape=tuple(int(width) for width in tile_shape), + guard_cells=guard_cells, + active_dims=active_dims, + dtype=dtype, + ) + self.file_extension = file_extension + self.raise_writer_errors_on_close = bool(raise_writer_errors_on_close) + + self._queue = queue.Queue(maxsize=queue_size) + self._thread = None + self._closed = False + self._error = None + self._error_traceback = None + + def start(self): + if self._thread is not None: + return + + self._thread = threading.Thread(target=self._writer_loop, daemon=True) + self._thread.start() + + def enqueue(self, snapshot, *, block=True): + if self._closed: + raise RuntimeError("Cannot enqueue after writer.close().") + if self._error is not None: + raise RuntimeError( + "Async openPMD writer failed. Original traceback:\n" + f"{self._error_traceback}" + ) from self._error + + try: + self._queue.put(snapshot, block=block) + return True + except queue.Full: + return False + + def enqueue_fields(self, field_map, *, step, time, block=True): + snapshot = make_tiled_field_snapshot(field_map, step=step, time=time) + return self.enqueue(snapshot, block=block) + + def close(self, raise_errors=None): + if self._closed: + return + + self._closed = True + if self._thread is not None: + self._queue.put(None) + self._queue.join() + self._thread.join() + self._thread = None + + if raise_errors is None: + raise_errors = self.raise_writer_errors_on_close + if self._error is not None and raise_errors: + raise RuntimeError( + "Async openPMD writer failed. Original traceback:\n" + f"{self._error_traceback}" + ) from self._error + + def _writer_loop(self): + while True: + item = self._queue.get() + try: + if item is None: + return + + if self._error is None: + # openPMD/HDF5 native writes are serialized across field and + # particle diagnostics; host snapshot construction remains + # outside this lock. + with _OPENPMD_WRITE_LOCK: + write_tiled_field_snapshot_openpmd( + item, + output_dir=self.output_dir, + filename=self.filename, + world=self.world, + layout=self.layout, + file_extension=self.file_extension, + ) + + except BaseException as exc: + self._error = exc + self._error_traceback = traceback.format_exc() + finally: + self._queue.task_done() + + +class AsyncTiledOpenPMDParticleWriter: + """ + Bounded background queue writer for tiled openPMD particle output. + """ + + def __init__( + self, + *, + output_dir, + filename, + world, + constants, + file_extension=".bp", + dtype=np.float64, + queue_size=2, + raise_writer_errors_on_close=True, + ): + queue_size = max(1, int(queue_size)) + + self.output_dir = output_dir + self.filename = filename + self.world = world + self.constants = constants + self.file_extension = file_extension + self.dtype = dtype + self.raise_writer_errors_on_close = bool(raise_writer_errors_on_close) + + self._queue = queue.Queue(maxsize=queue_size) + self._thread = None + self._closed = False + self._error = None + self._error_traceback = None + + def start(self): + if self._thread is not None: + return + + self._thread = threading.Thread(target=self._writer_loop, daemon=True) + self._thread.start() + + def enqueue(self, snapshot, *, block=True): + if self._closed: + raise RuntimeError("Cannot enqueue after writer.close().") + if self._error is not None: + raise RuntimeError( + "Async openPMD particle writer failed. Original traceback:\n" + f"{self._error_traceback}" + ) from self._error + + try: + self._queue.put(snapshot, block=block) + return True + except queue.Full: + return False + + def enqueue_particles(self, particles, *, step, time, species_config, species_names=None, block=True): + snapshot = make_tiled_particle_snapshot( + particles, + step=step, + time=time, + species_config=species_config, + species_names=species_names, + ) + return self.enqueue(snapshot, block=block) + + def close(self, raise_errors=None): + if self._closed: + return + + self._closed = True + if self._thread is not None: + self._queue.put(None) + self._queue.join() + self._thread.join() + self._thread = None + + if raise_errors is None: + raise_errors = self.raise_writer_errors_on_close + if self._error is not None and raise_errors: + raise RuntimeError( + "Async openPMD particle writer failed. Original traceback:\n" + f"{self._error_traceback}" + ) from self._error + + def _writer_loop(self): + while True: + item = self._queue.get() + try: + if item is None: + return + + if self._error is None: + # Share the same native-write lock as field diagnostics so + # independent output streams never enter openPMD concurrently. + with _OPENPMD_WRITE_LOCK: + write_tiled_particle_snapshot_openpmd( + item, + output_dir=self.output_dir, + filename=self.filename, + world=self.world, + constants=self.constants, + file_extension=self.file_extension, + dtype=self.dtype, + ) + + except BaseException as exc: + self._error = exc + self._error_traceback = traceback.format_exc() + finally: + self._queue.task_done() + + +def create_async_tiled_openpmd_field_writer( + world, + output_dir, + *, + filename="fields", + file_extension=".h5", + queue_size=2, +): + writer = AsyncTiledOpenPMDFieldWriter( + output_dir=output_dir, + filename=filename, + world=world, + global_shape=(int(world["Nx"]), int(world["Ny"]), int(world["Nz"])), + tile_shape=tuple(int(width) for width in world["tile_shape"]), + guard_cells=int(world["guard_cells"]), + active_dims=(1, 1, 1), + file_extension=file_extension, + queue_size=queue_size, + ) + writer.start() + return writer + + +def create_async_tiled_openpmd_particle_writer( + world, + constants, + output_dir, + *, + filename="particles", + file_extension=".h5", + queue_size=2, +): + writer = AsyncTiledOpenPMDParticleWriter( + output_dir=output_dir, + filename=filename, + world=world, + constants=constants, + file_extension=file_extension, + queue_size=queue_size, + ) + writer.start() + return writer + + +def enqueue_openpmd_field_output(field_writer, fields, world, plot_t, t, *, block=True): + field_map = _fields_to_tiled_output_map(fields) + prefetch_field_map_to_host(field_map) + return field_writer.enqueue_fields( + field_map, + step=int(plot_t), + time=float(t * world["dt"]), + block=block, + ) + + +def enqueue_openpmd_particle_output( + particle_writer, + particles, + world, + plot_t, + t, + *, + species_config, + species_names=None, + block=True, +): + prefetch_tiled_particles_to_host(particles) + return particle_writer.enqueue_particles( + particles, + step=int(plot_t), + time=float(t * world["dt"]), + species_config=species_config, + species_names=species_names, + block=block, + ) diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index b41ccf09..17d5e815 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -1,15 +1,12 @@ +from PyPIC3D.deposition.rho import ( + compute_tiled_mass_density_from_tiled_particles, + compute_tiled_pressure_field_from_tiled_particles, + compute_tiled_velocity_field_from_tiled_particles, +) +from PyPIC3D.particles.particle_class import TiledParticles -from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.boundary_conditions.grid_and_stencil import collapse_axis_stencil, prepare_particle_axis_stencil -from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells -import jax -import jax.numpy as jnp -from jax import jit - - -@jit -def compute_mass_density(particles, rho, world): +def compute_mass_density(particles, rho, world, species_config=None): """ Compute the mass density (rho) for a given set of particles in a simulation world. Parameters: @@ -26,222 +23,50 @@ def compute_mass_density(particles, rho, world): Returns: ndarray: The updated charge density array. """ - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - grid = world['grids']['vertex'] - Nx, Ny, Nz = rho.shape - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the shape of the charge density array - - rho = jnp.zeros_like(rho) - # reset rho to zero - - for species in particles: - shape_factor = species.get_shape() - # get the shape factor of the species, which determines the weighting function - N_particles = species.get_number_of_particles() - mass = species.get_mass() - # get the number of particles and their mass - dm = mass / dx / dy / dz - # calculate the mass per unit volume - x, y, z = species.get_position() - # get the position of the particles in the species - active = species.get_active_mask().astype(x.dtype) - - _, _, deltax, xpts = prepare_particle_axis_stencil( - x, grid[0], Nx, shape_factor, bc_x, wind=world['x_wind'], ghost_cells=True - ) - _, _, deltay, ypts = prepare_particle_axis_stencil( - y, grid[1], Ny, shape_factor, bc_y, wind=world['y_wind'], ghost_cells=True - ) - _, _, deltaz, zpts = prepare_particle_axis_stencil( - z, grid[2], Nz, shape_factor, bc_z, wind=world['z_wind'], ghost_cells=True - ) - - 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 - ) - # get the weighting factors based on the shape factor - - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - x_weights = jnp.asarray(x_weights) - y_weights = jnp.asarray(y_weights) - z_weights = jnp.asarray(z_weights) - xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=True) - ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=True) - zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=True) - - - for i in range(3): - for j in range(3): - for k in range(3): - rho = rho.at[xpts[i], ypts[j], zpts[k]].add(active * dm * x_weights[i] * y_weights[j] * z_weights[k], mode='drop') - # distribute the mass of the particles to the grid points using the weighting factors - - rho = fold_ghost_cells(rho, bc_x, bc_y, bc_z) - rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) - return rho - -@jit -def compute_velocity_field(particles, field, direction, world): - """ - Compute the velocity field (v) for a given set of particles in a simulation world. - Parameters: - particles (list): A list of particle species, each containing methods to get the number of particles, - their positions, and their mass. - field (ndarray): The initial velocity field array to be updated. - direction (int): The direction along which to compute the velocity field (0: x, 1: y, 2: z). - world (dict): A dictionary containing the simulation world parameters, including: - - 'dx': Grid spacing in the x-direction. - - 'dy': Grid spacing in the y-direction. - - 'dz': Grid spacing in the z-direction. - - 'x_wind': Window size in the x-direction. - - 'y_wind': Window size in the y-direction. - - 'z_wind': Window size in the z-direction. - Returns: - ndarray: The updated velocity field array. - """ - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - grid = world['grids']['vertex'] - Nx, Ny, Nz = field.shape - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the shape of the velocity field array - - field = jnp.zeros_like(field) - # reset field to zero - - for species in particles: - shape_factor = species.get_shape() - # get the shape factor of the species, which determines the weighting function - N_particles = species.get_number_of_particles() - # get the number of particles - x, y, z = species.get_position() - # get the position of the particles in the species - vx, vy, vz = species.get_velocity() - # get the velocity of the particles in the species - active = species.get_active_mask().astype(x.dtype) - - dv = active * jnp.array([vx, vy, vz])[direction] / N_particles - # select the velocity component based on the direction - - _, _, deltax, xpts = prepare_particle_axis_stencil( - x, grid[0], Nx, shape_factor, bc_x, wind=world['x_wind'], ghost_cells=True - ) - _, _, deltay, ypts = prepare_particle_axis_stencil( - y, grid[1], Ny, shape_factor, bc_y, wind=world['y_wind'], ghost_cells=True - ) - _, _, deltaz, zpts = prepare_particle_axis_stencil( - z, grid[2], Nz, shape_factor, bc_z, wind=world['z_wind'], ghost_cells=True - ) - - 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 - ) - # get the weighting factors based on the shape factor - - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - x_weights = jnp.asarray(x_weights) - y_weights = jnp.asarray(y_weights) - z_weights = jnp.asarray(z_weights) - xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=True) - ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=True) - zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=True) - - for i in range(3): - for j in range(3): - for k in range(3): - field = field.at[xpts[i], ypts[j], zpts[k]].add( dv * x_weights[i] * y_weights[j] * z_weights[k], mode='drop') - # distribute the velocity of the particles to the grid points using the weighting factors - - field = fold_ghost_cells(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field - - - - -@jit -def compute_pressure_field(particles, field, velocity_field, direction, world): - - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - grid = world['grids']['vertex'] - Nx, Ny, Nz = field.shape - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the shape of the velocity field array - - field = jnp.zeros_like(field) - # reset field to zero - - for species in particles: - shape_factor = species.get_shape() - # get the shape factor of the species, which determines the weighting function - x, y, z = species.get_position() - # get the position of the particles in the species - vx, vy, vz = species.get_velocity() - # get the velocity of the particles in the species - active = species.get_active_mask().astype(x.dtype) - - - v = active * jnp.array([vx, vy, vz])[direction] - # select the velocity component based on the direction - - _, _, deltax, xpts = prepare_particle_axis_stencil( - x, grid[0], Nx, shape_factor, bc_x, wind=world['x_wind'], ghost_cells=True - ) - _, _, deltay, ypts = prepare_particle_axis_stencil( - y, grid[1], Ny, shape_factor, bc_y, wind=world['y_wind'], ghost_cells=True - ) - _, _, deltaz, zpts = prepare_particle_axis_stencil( - z, grid[2], Nz, shape_factor, bc_z, wind=world['z_wind'], ghost_cells=True - ) - - 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 - ) - # get the weighting factors based on the shape factor - - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - x_weights = jnp.asarray(x_weights) - y_weights = jnp.asarray(y_weights) - z_weights = jnp.asarray(z_weights) - xpts, x_weights = collapse_axis_stencil(xpts, x_weights, Nx, ghost_cells=True) - ypts, y_weights = collapse_axis_stencil(ypts, y_weights, Ny, ghost_cells=True) - zpts, z_weights = collapse_axis_stencil(zpts, z_weights, Nz, ghost_cells=True) - - for i in range(3): - for j in range(3): - for k in range(3): - vbar = v - velocity_field.at[xpts[i], ypts[j], zpts[k]].get() - - field = field.at[xpts[i], ypts[j], zpts[k]].add(active * vbar**2 * x_weights[i] * y_weights[j] * z_weights[k], mode='drop') - # distribute the pressure moment of the particles to the grid points using the weighting factors - - field = fold_ghost_cells(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field + if isinstance(particles, TiledParticles): + if getattr(rho, "ndim", 0) == 6: + return compute_tiled_mass_density_from_tiled_particles( + particles, + species_config, + rho, + world, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), + ) + raise ValueError("compute_mass_density requires tile-major scalar field storage.") + + raise ValueError("Public compute_mass_density requires TiledParticles.") + + +def compute_velocity_field(particles, field, direction, world, species_config=None): + if isinstance(particles, TiledParticles): + if getattr(field, "ndim", 0) == 6: + return compute_tiled_velocity_field_from_tiled_particles( + particles, + species_config, + field, + int(direction), + world, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), + ) + raise ValueError("compute_velocity_field requires tile-major scalar field storage.") + + raise ValueError("Public compute_velocity_field requires TiledParticles.") + +def compute_pressure_field(particles, field, velocity_field, direction, world, species_config=None): + if isinstance(particles, TiledParticles): + if getattr(field, "ndim", 0) == 6: + return compute_tiled_pressure_field_from_tiled_particles( + particles, + species_config, + field, + velocity_field, + int(direction), + world, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), + ) + raise ValueError("compute_pressure_field requires tile-major scalar field storage.") + + raise ValueError("Public compute_pressure_field requires TiledParticles.") diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index 00054330..1b306e8c 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -1,10 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Tuple, Union + import openpmd_api as io import jax.numpy as jnp import os import numpy as np import importlib.metadata +from PyPIC3D.diagnostics.output_adapters import fields_for_output, particles_for_output + + +@dataclass(frozen=True) +class TiledMeshLayout: + """ + Metadata needed to place tile-local interiors into a global openPMD mesh. + """ + + global_shape: Tuple[int, int, int] + tile_shape: Tuple[int, int, int] + guard_cells: Union[int, Tuple[int, int, int]] = 1 + active_dims: Tuple[int, int, int] = (1, 1, 1) + dtype: Any = np.float64 + def _ensure_openpmd_array(data, dtype=np.float64, squeeze=False): arr = np.asarray(data, dtype=dtype) if squeeze: @@ -50,15 +70,7 @@ def _configure_openpmd_mesh(mesh, world, active_dims=(1,1,1)): mesh.axis_labels = axes mesh.grid_spacing = ds mesh.grid_global_offset = offsets - - - # mesh.grid_spacing = [float(world["dx"]), float(world["dy"]), float(world["dz"])] - # mesh.grid_global_offset = [ - # -float(world["x_wind"]) / 2.0, - # -float(world["y_wind"]) / 2.0, - # -float(world["z_wind"]) / 2.0, - # ] - # mesh.axis_labels = ["x", "y", "z"] + mesh.unit_SI = 1.0 @@ -115,6 +127,170 @@ def _fields_to_interior_map(fields): return field_map +def _as_3tuple(value): + if isinstance(value, tuple): + return tuple(int(v) for v in value) + if isinstance(value, list): + return tuple(int(v) for v in value) + return (int(value), int(value), int(value)) + + +def _slice_start(index_entry): + if isinstance(index_entry, slice): + return 0 if index_entry.start is None else int(index_entry.start) + return int(index_entry) + + +def _tile_interior(tile, guard_cells): + gx, gy, gz = guard_cells + sx = slice(gx, -gx) if gx > 0 else slice(None) + sy = slice(gy, -gy) if gy > 0 else slice(None) + sz = slice(gz, -gz) if gz > 0 else slice(None) + return tile[sx, sy, sz] + + +def _iter_tile_chunks_from_host_shard(shard_index, shard_data, *, layout): + """ + Yield ``(global_mesh_offset, tile_interior)`` chunks from one host shard. + """ + if shard_data.ndim != 6: + raise ValueError( + "Expected tiled scalar field shard with shape " + "(ntx, nty, ntz, nx+2g, ny+2g, nz+2g). " + f"Got {shard_data.shape}." + ) + + tile_shape = tuple(int(width) for width in layout.tile_shape) + guard_cells = _as_3tuple(layout.guard_cells) + + tx0 = _slice_start(shard_index[0]) + ty0 = _slice_start(shard_index[1]) + tz0 = _slice_start(shard_index[2]) + ntx_local, nty_local, ntz_local = shard_data.shape[:3] + + for ltx in range(ntx_local): + for lty in range(nty_local): + for ltz in range(ntz_local): + tx = tx0 + ltx + ty = ty0 + lty + tz = tz0 + ltz + + raw_tile = shard_data[ltx, lty, ltz] + interior = _tile_interior(raw_tile, guard_cells) + if tuple(interior.shape) != tile_shape: + raise ValueError( + f"Tile interior has shape {interior.shape}, expected {tile_shape}. " + f"Check guard_cells={guard_cells} and tile_shape={tile_shape}." + ) + + offset = [ + tx * tile_shape[0], + ty * tile_shape[1], + tz * tile_shape[2], + ] + yield offset, interior + + +def _reset_scalar_mesh_record(iteration, name, *, world, layout): + mesh = iteration.meshes[name] + _configure_openpmd_mesh(mesh, world, layout.active_dims) + record = mesh[io.Mesh_Record_Component.SCALAR] + record.reset_dataset(io.Dataset(np.dtype(layout.dtype), list(layout.global_shape))) + record.unit_SI = 1.0 + return record + + +def _reset_vector_mesh_record(iteration, name, component_name, *, world, layout): + mesh = iteration.meshes[name] + _configure_openpmd_mesh(mesh, world, layout.active_dims) + record = mesh[component_name] + record.reset_dataset(io.Dataset(np.dtype(layout.dtype), list(layout.global_shape))) + record.unit_SI = 1.0 + return record + + +def write_tiled_scalar_field_chunks_to_iteration(iteration, name, host_shards, *, world, layout): + record = _reset_scalar_mesh_record(iteration, name, world=world, layout=layout) + + for shard_index, shard_data in host_shards: + for offset, tile in _iter_tile_chunks_from_host_shard( + shard_index, + shard_data, + layout=layout, + ): + tile = _ensure_openpmd_array(tile, dtype=layout.dtype) + record.store_chunk(tile, offset, list(tile.shape)) + + +def write_tiled_vector_field_chunks_to_iteration( + iteration, + name, + component_host_shards, + *, + world, + layout, + component_names=("x", "y", "z"), +): + for component_name, host_shards in zip(component_names, component_host_shards): + record = _reset_vector_mesh_record( + iteration, + name, + component_name, + world=world, + layout=layout, + ) + + for shard_index, shard_data in host_shards: + for offset, tile in _iter_tile_chunks_from_host_shard( + shard_index, + shard_data, + layout=layout, + ): + tile = _ensure_openpmd_array(tile, dtype=layout.dtype) + record.store_chunk(tile, offset, list(tile.shape)) + + +def write_tiled_field_snapshot_openpmd( + snapshot, + *, + output_dir, + filename, + world, + layout, + file_extension=".bp", +): + series = _open_openpmd_series(output_dir, filename, file_extension=file_extension) + + try: + iteration = series.iterations[int(snapshot.step)] + iteration.time = float(snapshot.time) + iteration.dt = float(world["dt"]) + iteration.time_unit_SI = 1.0 + + for name, value in snapshot.fields.items(): + is_vector = isinstance(value, tuple) and len(value) == 3 + if is_vector: + write_tiled_vector_field_chunks_to_iteration( + iteration, + name, + value, + world=world, + layout=layout, + ) + else: + write_tiled_scalar_field_chunks_to_iteration( + iteration, + name, + value, + world=world, + layout=layout, + ) + + series.flush() + finally: + series.close() + + def write_openpmd_fields_to_iteration(iteration, field_map, world, active_dims=(1,1,1)): for name, data in field_map.items(): is_vector = isinstance(data, (list, tuple)) and len(data) == 3 @@ -124,18 +300,22 @@ def write_openpmd_fields_to_iteration(iteration, field_map, world, active_dims=( _write_openpmd_scalar_mesh(iteration, name, data, world, active_dims) -def write_openpmd_particles_to_iteration(iteration, particles, constants): +def write_openpmd_particles_to_iteration(iteration, particles, constants, species_config=None, species_names=None, world=None): + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) + # Tiled particles carry inactive capacity slots; openPMD should see the + # active physical particles by species, matching the ordinary output path. + if not particles: return C = float(constants["C"]) for species in particles: - species_name = species.get_name().replace(" ", "_") + species_name = species.name.replace(" ", "_") species_group = iteration.particles[species_name] - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] gamma = 1 / jnp.sqrt(1.0 - (vx**2 + vy**2 + vz**2) / C**2) x = _ensure_openpmd_array(x, squeeze=True) @@ -149,9 +329,9 @@ def write_openpmd_particles_to_iteration(iteration, particles, constants): num_particles = x.shape[0] # number of particles in this species - particle_mass = species.get_mass() - particle_charge = species.get_charge() - weights = species.get_weight() + particle_mass = species.mass * species.weight + particle_charge = species.charge * species.weight + weights = species.weight # get the particle mass, charge, and weight for this species @@ -212,6 +392,215 @@ def write_openpmd_particles_to_iteration(iteration, particles, constants): mass.unit_SI = 1.0 +def _axis_diagnostic_position_array(x, u, dt, wind, bc): + x_diagnostic = x - u * dt / 2.0 + + if int(np.asarray(bc)) == 0: + half_wind = wind / 2.0 + x_diagnostic = np.where( + x_diagnostic > half_wind, + x_diagnostic - wind, + np.where(x_diagnostic < -half_wind, x_diagnostic + wind, x_diagnostic), + ) + + return x_diagnostic + + +def _diagnostic_position_array(x, u, world): + particle_bc = world.get("particle_boundary_conditions", {}) + dt = float(world["dt"]) + + return np.stack( + ( + _axis_diagnostic_position_array(x[:, 0], u[:, 0], dt, float(world["x_wind"]), particle_bc.get("x", 0)), + _axis_diagnostic_position_array(x[:, 1], u[:, 1], dt, float(world["y_wind"]), particle_bc.get("y", 0)), + _axis_diagnostic_position_array(x[:, 2], u[:, 2], dt, float(world["z_wind"]), particle_bc.get("z", 0)), + ), + axis=-1, + ) + + +def _count_snapshot_particles_by_species(snapshot): + counts = np.zeros(len(snapshot.species_names), dtype=np.int64) + + for active_index, active_chunk in snapshot.active_shards: + if active_chunk.ndim != 5: + raise ValueError( + "Expected active particle shard with shape " + "(ntx, nty, ntz, nspecies, max_particles_per_tile). " + f"Got {active_chunk.shape}." + ) + s0 = _slice_start(active_index[3]) + for local_s in range(active_chunk.shape[3]): + counts[s0 + local_s] += int(np.count_nonzero(active_chunk[:, :, :, local_s, :])) + + return counts + + +def _iter_snapshot_particle_chunks(snapshot, world, constants): + if len(snapshot.x_shards) != len(snapshot.u_shards) or len(snapshot.x_shards) != len(snapshot.active_shards): + raise ValueError("Particle snapshot x, u, and active shard lists must have the same length.") + + C = float(constants["C"]) + species_charge = np.asarray(snapshot.species_charge, dtype=np.float64) + species_mass = np.asarray(snapshot.species_mass, dtype=np.float64) + species_weight = np.asarray(snapshot.species_weight, dtype=np.float64) + + for (x_index, x_chunk), (_u_index, u_chunk), (_active_index, active_chunk) in zip( + snapshot.x_shards, + snapshot.u_shards, + snapshot.active_shards, + ): + if x_chunk.ndim != 6: + raise ValueError( + "Expected particle x shard with shape " + "(ntx, nty, ntz, nspecies, max_particles_per_tile, 3). " + f"Got {x_chunk.shape}." + ) + if u_chunk.shape != x_chunk.shape: + raise ValueError(f"Particle u shard shape {u_chunk.shape} must match x shape {x_chunk.shape}.") + if active_chunk.shape != x_chunk.shape[:-1]: + raise ValueError(f"Particle active shard shape {active_chunk.shape} must match x.shape[:-1] {x_chunk.shape[:-1]}.") + + s0 = _slice_start(x_index[3]) + ntx, nty, ntz, ns_local, _slots, _coords = x_chunk.shape + + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + for local_s in range(ns_local): + species_index = s0 + local_s + active = np.asarray(active_chunk[tx, ty, tz, local_s], dtype=bool) + n_active = int(np.count_nonzero(active)) + if n_active == 0: + continue + + x_live = np.array(x_chunk[tx, ty, tz, local_s][active], dtype=np.float64, copy=True, order="C") + u_live = np.array(u_chunk[tx, ty, tz, local_s][active], dtype=np.float64, copy=True, order="C") + x_diagnostic = _diagnostic_position_array(x_live, u_live, world) + + charge = np.full(n_active, float(species_charge[species_index]), dtype=np.float64) + mass = np.full(n_active, float(species_mass[species_index]), dtype=np.float64) + weight = np.full(n_active, float(species_weight[species_index]), dtype=np.float64) + + gamma = 1.0 / np.sqrt(1.0 - np.sum(u_live * u_live, axis=1) / C**2) + + yield species_index, x_diagnostic, u_live, charge, mass, weight, gamma + + +def _reset_particle_species_records(species_group, dtype, num_particles): + shape = [int(num_particles)] + dtype = np.dtype(dtype) + + for record_name in ("position", "positionOffset", "momentum"): + record = species_group[record_name] + for component in ("x", "y", "z"): + record_component = record[component] + record_component.reset_dataset(io.Dataset(dtype, shape)) + record_component.unit_SI = 1.0 + + for record_name in ("weighting", "charge", "mass"): + record_component = species_group[record_name] + record_component.reset_dataset(io.Dataset(dtype, shape)) + record_component.unit_SI = 1.0 + + +def _store_particle_record_chunk(species_group, offset, x, u, charge, mass, weight, gamma, dtype=np.float64): + num_particles = int(x.shape[0]) + if num_particles == 0: + return int(offset) + + start = [int(offset)] + extent = [num_particles] + + x = _ensure_openpmd_array(x, dtype=dtype) + u = _ensure_openpmd_array(u, dtype=dtype) + charge = _ensure_openpmd_array(charge, dtype=dtype) + mass = _ensure_openpmd_array(mass, dtype=dtype) + weight = _ensure_openpmd_array(weight, dtype=dtype) + gamma = _ensure_openpmd_array(gamma, dtype=dtype) + + for component, data in zip(("x", "y", "z"), (x[:, 0], x[:, 1], x[:, 2])): + species_group["position"][component].store_chunk(_ensure_openpmd_array(data, dtype=dtype), start, extent) + + zeros = np.zeros(num_particles, dtype=np.dtype(dtype)) + for component in ("x", "y", "z"): + species_group["positionOffset"][component].store_chunk(zeros, start, extent) + + momentum = u * (mass * weight * gamma)[:, None] + for component, data in zip(("x", "y", "z"), (momentum[:, 0], momentum[:, 1], momentum[:, 2])): + species_group["momentum"][component].store_chunk(_ensure_openpmd_array(data, dtype=dtype), start, extent) + + species_group["weighting"].store_chunk(weight, start, extent) + species_group["charge"].store_chunk(charge, start, extent) + species_group["mass"].store_chunk(mass, start, extent) + + return int(offset) + num_particles + + +def write_tiled_particle_snapshot_to_iteration(iteration, snapshot, world, constants, dtype=np.float64): + counts = _count_snapshot_particles_by_species(snapshot) + + for species_index, species_name in enumerate(snapshot.species_names): + species_group = iteration.particles[species_name.replace(" ", "_")] + _reset_particle_species_records( + species_group, + dtype=dtype, + num_particles=int(counts[species_index]), + ) + + offsets = np.zeros(len(snapshot.species_names), dtype=np.int64) + + for species_index, x, u, charge, mass, weight, gamma in _iter_snapshot_particle_chunks(snapshot, world, constants): + species_group = iteration.particles[snapshot.species_names[species_index].replace(" ", "_")] + offsets[species_index] = _store_particle_record_chunk( + species_group, + int(offsets[species_index]), + x, + u, + charge, + mass, + weight, + gamma, + dtype=dtype, + ) + + for species_name, stored, expected in zip(snapshot.species_names, offsets, counts): + if int(stored) != int(expected): + raise RuntimeError(f"Stored {stored} particles for {species_name}, expected {expected}.") + + +def write_tiled_particle_snapshot_openpmd( + snapshot, + *, + output_dir, + filename, + world, + constants, + file_extension=".bp", + dtype=np.float64, +): + series = _open_openpmd_series(output_dir, filename, file_extension=file_extension) + + try: + iteration = series.iterations[int(snapshot.step)] + iteration.time = float(snapshot.time) + iteration.dt = float(world["dt"]) + iteration.time_unit_SI = 1.0 + + write_tiled_particle_snapshot_to_iteration( + iteration, + snapshot, + world, + constants, + dtype=dtype, + ) + + series.flush() + finally: + series.close() + + def write_openpmd_fields(fields, world, output_dir, plot_t, t, filename="fields", file_extension=".bp"): """ Write all field data to an openPMD file for visualization in ParaView/VisIt. @@ -225,6 +614,7 @@ def write_openpmd_fields(fields, world, output_dir, plot_t, t, filename="fields" filename (str): Base name for the openPMD file. file_extension (str): File extension for the openPMD series (for example, ".bp"). """ + fields = fields_for_output(fields, world) field_map = _fields_to_interior_map(fields) # extract physical interior (strip ghost cells) @@ -249,7 +639,7 @@ def write_openpmd_fields(fields, world, output_dir, plot_t, t, filename="fields" # flush and close the series -def write_openpmd_particles(particles, world, constants, output_dir, plot_t, t, filename="particles", file_extension=".bp"): +def write_openpmd_particles(particles, world, constants, output_dir, plot_t, t, filename="particles", file_extension=".bp", species_config=None, species_names=None): """ Write all particle data to an openPMD file for visualization in ParaView/VisIt. @@ -271,7 +661,7 @@ def write_openpmd_particles(particles, world, constants, output_dir, plot_t, t, # set the time step iteration.time_unit_SI = 1.0 # set the time unit - write_openpmd_particles_to_iteration(iteration, particles, constants) + write_openpmd_particles_to_iteration(iteration, particles, constants, species_config=species_config, species_names=species_names, world=world) # write the particle data to the iteration series.flush() series.close() @@ -279,7 +669,7 @@ def write_openpmd_particles(particles, world, constants, output_dir, plot_t, t, -def write_openpmd_initial_particles(particles, world, constants, output_dir, filename="initial_particles.h5"): +def write_openpmd_initial_particles(particles, world, constants, output_dir, filename="initial_particles.h5", species_config=None, species_names=None): """ Write the initial particle states to separate openPMD files, one per species. @@ -290,6 +680,10 @@ def write_openpmd_initial_particles(particles, world, constants, output_dir, fil output_dir (str): Base output directory for the simulation. filename (str): Base name of the openPMD output file (species name is prepended). """ + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) + # Initial particle dumps may receive tile-major runtime storage. The + # openPMD file still contains ordinary per-species particle records. + if not particles: return @@ -305,7 +699,7 @@ def make_array_writable(arr): return arr for species in particles: - species_name = species.get_name().replace(" ", "_") + species_name = species.name.replace(" ", "_") series_filename = f"{species_name}_{filename}" series_path = os.path.join(output_path, series_filename) @@ -320,8 +714,8 @@ def make_array_writable(arr): species_group = iteration.particles[species_name] - x, y, z = species.get_forward_position() - vx, vy, vz = species.get_velocity() + x, y, z = species.x[:, 0], species.x[:, 1], species.x[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] gamma = 1 / jnp.sqrt(1.0 - (vx**2 + vy**2 + vz**2) / C**2) # compute the Lorentz factor @@ -334,8 +728,27 @@ def make_array_writable(arr): gamma = make_array_writable(gamma) num_particles = x.shape[0] - particle_mass = float(species.mass) - particle_charge = float(species.charge) + particle_mass = species.mass * species.weight + particle_charge = species.charge * species.weight + particle_weight = species.weight + + if jnp.ndim(particle_weight) == 0: + weights = np.full(num_particles, float(particle_weight), dtype=np.float64) + else: + weights = _ensure_openpmd_array(particle_weight, squeeze=True) + + if jnp.ndim(particle_mass) == 0: + masses = np.full(num_particles, float(particle_mass), dtype=np.float64) + else: + masses = _ensure_openpmd_array(particle_mass, squeeze=True) + + if jnp.ndim(particle_charge) == 0: + charges = np.full(num_particles, float(particle_charge), dtype=np.float64) + else: + charges = _ensure_openpmd_array(particle_charge, squeeze=True) + + mass_per_weight = masses / weights + charge_per_weight = charges / weights position = species_group["position"] for component, data in zip(("x", "y", "z"), (x, y, z)): @@ -357,25 +770,22 @@ def make_array_writable(arr): for component, data in zip(("x", "y", "z"), (vx, vy, vz)): record_component = momentum[component] record_component.reset_dataset(io.Dataset(data.dtype, [num_particles])) - record_component.store_chunk(data * particle_mass * gamma , [0], [num_particles]) + record_component.store_chunk(data * masses * gamma, [0], [num_particles]) record_component.unit_SI = 1.0 weighting = species_group["weighting"] - weights = np.full(num_particles, float(species.weight), dtype=np.float64) weighting.reset_dataset(io.Dataset(weights.dtype, [num_particles])) weighting.store_chunk(weights, [0], [num_particles]) weighting.unit_SI = 1.0 charge = species_group["charge"] - charges = np.full(num_particles, particle_charge, dtype=np.float64) charge.reset_dataset(io.Dataset(charges.dtype, [num_particles])) - charge.store_chunk(charges, [0], [num_particles]) + charge.store_chunk(charge_per_weight, [0], [num_particles]) charge.unit_SI = 1.0 mass = species_group["mass"] - masses = np.full(num_particles, particle_mass, dtype=np.float64) mass.reset_dataset(io.Dataset(masses.dtype, [num_particles])) - mass.store_chunk(masses, [0], [num_particles]) + mass.store_chunk(mass_per_weight, [0], [num_particles]) mass.unit_SI = 1.0 series.flush() @@ -391,6 +801,7 @@ def write_openpmd_initial_fields(fields, world, output_dir, filename="initial_fi output_dir (str): Base output directory for the simulation. filename (str): openPMD file name. """ + fields = fields_for_output(fields, world) field_map = _fields_to_interior_map(fields) # extract physical interior (strip ghost cells) diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py new file mode 100644 index 00000000..14ad39bb --- /dev/null +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -0,0 +1,205 @@ +import jax.numpy as jnp +from typing import NamedTuple + +from PyPIC3D.particles.particle_class import TiledParticles + + +class ParticleOutputRecord(NamedTuple): + name: str + species_index: int + x: jnp.ndarray + x_diagnostic: jnp.ndarray + u: jnp.ndarray + charge: jnp.ndarray + mass: jnp.ndarray + weight: jnp.ndarray + + +def _is_tiled_scalar(field): + return getattr(field, "ndim", 0) == 6 + + +def _is_tiled_vector(field): + return ( + isinstance(field, (list, tuple)) + and len(field) == 3 + and _is_tiled_scalar(field[0]) + ) + + +def _tile_shape_from_world(world): + return tuple(int(width) for width in world["tile_shape"]) + + +def _guard_depth_from_world(world): + return int(world["guard_cells"]) + + +def assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells=2): + """ + Assemble compact field tiles back into one global ghost-celled field. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + ntx, nty, ntz = field_tiles.shape[:3] + Nx = int(ntx) * tile_nx + Ny = int(nty) * tile_ny + Nz = int(ntz) * tile_nz + + field = jnp.zeros((Nx + 2, Ny + 2, Nz + 2), dtype=field_tiles.dtype) + + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + tile_with_one_guard = field_tiles[ + tx, + ty, + tz, + g - 1:g + tile_nx + 1, + g - 1:g + tile_ny + 1, + g - 1:g + tile_nz + 1, + ] + ix = tx * tile_nx + iy = ty * tile_ny + iz = tz * tile_nz + field = field.at[ix:ix + tile_nx + 2, iy:iy + tile_ny + 2, iz:iz + tile_nz + 2].set(tile_with_one_guard) + + return field + + +def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells=2): + """ + Assemble tiled vector-field components into ordinary ghost-celled arrays. + """ + + return tuple(assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells) for component in field_tiles) + + +def scalar_field_for_output(field, world): + """ + Return an ordinary ghost-celled scalar field for file formats. + + Runtime diagnostics can operate on tile-major arrays, but openPMD + mesh writers still expect one global ghost-celled array. + """ + + if not _is_tiled_scalar(field): + return field + + tile_shape = _tile_shape_from_world(world) + g = _guard_depth_from_world(world) + return assemble_tiled_scalar_field(field, world, tile_shape, num_guard_cells=g) + + +def vector_field_for_output(field, world): + """ + Return ordinary ghost-celled vector components for file formats. + """ + + if not _is_tiled_vector(field): + return field + + tile_shape = _tile_shape_from_world(world) + g = _guard_depth_from_world(world) + return assemble_tiled_vector_field(field, world, tile_shape, num_guard_cells=g) + + +def fields_for_output(fields, world): + """ + Assemble tile-major fields at the I/O boundary. + + The live solver state is left untouched. The particle-retile overflow flag + is a Python-driver diagnostic, not a physical field, so it is not included + in the returned output tuple. + """ + + E, B, J, rho, phi, external_fields, *rest = fields + external_E, external_B = external_fields + + output_fields = ( + vector_field_for_output(E, world), + vector_field_for_output(B, world), + vector_field_for_output(J, world), + scalar_field_for_output(rho, world), + scalar_field_for_output(phi, world), + ( + vector_field_for_output(external_E, world), + vector_field_for_output(external_B, world), + ), + ) + + if not rest: + return output_fields + + pml_state = rest[0] + return output_fields + (pml_state,) + + +def _axis_diagnostic_position(x, u, dt, wind, bc): + x_diagnostic = x - u * dt / 2 + + if int(jnp.asarray(bc).item()) == 0: + half_wind = wind / 2 + x_diagnostic = jnp.where( + x_diagnostic > half_wind, + x_diagnostic - wind, + jnp.where(x_diagnostic < -half_wind, x_diagnostic + wind, x_diagnostic), + ) + + return x_diagnostic + + +def _diagnostic_position(x, u, world): + if world is None: + return x + + particle_bc = world.get("particle_boundary_conditions", {}) + dt = world["dt"] + x_diagnostic = _axis_diagnostic_position(x[:, 0], u[:, 0], dt, world["x_wind"], particle_bc.get("x", 0)) + y_diagnostic = _axis_diagnostic_position(x[:, 1], u[:, 1], dt, world["y_wind"], particle_bc.get("y", 0)) + z_diagnostic = _axis_diagnostic_position(x[:, 2], u[:, 2], dt, world["z_wind"], particle_bc.get("z", 0)) + + return jnp.stack((x_diagnostic, y_diagnostic, z_diagnostic), axis=-1) + + +def particles_for_output(particles, species_config=None, species_names=None, world=None): + """ + Flatten fixed-capacity tiled particle storage for diagnostics. + """ + + if not isinstance(particles, TiledParticles): + raise TypeError("Particle output requires TiledParticles runtime storage.") + + n_species = particles.active.shape[3] + output_particles = [] + + for species_index in range(n_species): + active = particles.active[:, :, :, species_index, :].reshape(-1) + + x = particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active] + u = particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active] + n_active = int(jnp.sum(active)) + charge = jnp.full((n_active,), species_config.charge[species_index]) + mass = jnp.full((n_active,), species_config.mass[species_index]) + weight = jnp.full((n_active,), species_config.weight[species_index]) + + if species_names is None: + name = f"species_{species_index}" + else: + name = species_names[species_index] + + output_particles.append( + ParticleOutputRecord( + name=name, + species_index=species_index, + x=x, + x_diagnostic=_diagnostic_position(x, u, world), + u=u, + charge=charge, + mass=mass, + weight=weight, + ) + ) + + return output_particles diff --git a/PyPIC3D/diagnostics/plotting.py b/PyPIC3D/diagnostics/plotting.py index 1236fae8..da74a86c 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -7,7 +7,9 @@ import jax from functools import partial -def plot_positions(particles, t, x_wind, y_wind, z_wind, path): +from PyPIC3D.diagnostics.output_adapters import particles_for_output + +def plot_positions(particles, t, x_wind, y_wind, z_wind, path, species_config=None, species_names=None, world=None): """ Makes an interactive 3D plot of the positions of the particles using Plotly. @@ -23,13 +25,17 @@ def plot_positions(particles, t, x_wind, y_wind, z_wind, path): """ fig = go.Figure() + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) + # Keep tiled fixed-capacity storage out of the plotting loop. From here on + # the code sees the same species-like diagnostic objects as ordinary output. + colors = ['red', 'blue', 'green', 'purple', 'orange', 'yellow', 'cyan'] for idx, species in enumerate(particles): - x, y, z = species.get_position() + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] fig.add_trace(go.Scatter3d( x=x, y=y, z=z, mode='markers', marker=dict(size=2, color=colors[idx % len(colors)]), - name=f'Species {idx + 1}' + name=species.name )) fig.update_layout( @@ -49,7 +55,7 @@ def plot_positions(particles, t, x_wind, y_wind, z_wind, path): fig.write_html(f"{path}/data/positions/particles.{t:09}.html") -def write_particles_phase_space(particles, t, path): +def write_particles_phase_space(particles, t, path, species_config=None, species_names=None, world=None): """ Write the phase space of the particles to a file. @@ -69,10 +75,14 @@ def write_particles_phase_space(particles, t, path): os.makedirs(f"{path}/data/phase_space/z") # Create directory if it doesn't exist + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) + # Tiled storage contains fixed-capacity inactive slots; phase-space output + # writes only active particles and otherwise keeps the old species-list path. + for species in particles: - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() - name = species.get_name().replace(" ", "") + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] + name = species.name.replace(" ", "") x_phase_space = jnp.stack((x, vx), axis=-1) y_phase_space = jnp.stack((y, vy), axis=-1) @@ -83,7 +93,7 @@ def write_particles_phase_space(particles, t, path): jnp.save(f"{path}/data/phase_space/z/{name}_phase_space.{t:09}.npy", z_phase_space) # write the phase space of the particles to a file -def particles_phase_space(particles, world, t, name, path): +def particles_phase_space(particles, world, t, name, path, species_config=None, species_names=None): """ Plot the phase space of the particles. @@ -103,9 +113,13 @@ def particles_phase_space(particles, world, t, name, path): colors = ['r', 'b', 'g', 'c', 'm', 'y', 'k'] idx = 0 order = 10 + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) + # Phase-space plots are an output boundary; tile-major storage is flattened + # once here rather than leaking into the plotting loops. + for species in particles: - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] plt.scatter(x, vx, c=colors[idx], zorder=order, s=1) idx += 1 order -= 1 @@ -119,8 +133,8 @@ def particles_phase_space(particles, world, t, name, path): idx = 0 for species in particles: - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] plt.scatter(y, vy, c=colors[idx]) idx += 1 plt.xlabel("Position") @@ -132,8 +146,8 @@ def particles_phase_space(particles, world, t, name, path): idx = 0 for species in particles: - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() + x, y, z = species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2] + vx, vy, vz = species.u[:, 0], species.u[:, 1], species.u[:, 2] plt.scatter(z, vz, c=colors[idx]) idx += 1 plt.xlabel("Position") @@ -144,16 +158,14 @@ def particles_phase_space(particles, world, t, name, path): plt.close() -def plot_initial_histograms(particle_species, world, name, path): +def plot_initial_histograms(particle_record, world, name, path): """ Generates and saves histograms for the initial positions and velocities of particles in a simulation. Parameters: - particle_species (object): An object representing the particle species, - which provides methods `get_position()` and - `get_velocity()` to retrieve particle positions - (x, y, z) and velocities (vx, vy, vz). + particle_record (ParticleOutputRecord): Flattened tiled particle output + with diagnostic positions and velocities. world (dict): A dictionary containing the simulation world parameters, specifically the wind dimensions with keys 'x_wind', 'y_wind', and 'z_wind'. @@ -169,8 +181,8 @@ def plot_initial_histograms(particle_species, world, name, path): formatted as "{name}_initial__histogram.png". """ - x, y, z = particle_species.get_position() - vx, vy, vz = particle_species.get_velocity() + x, y, z = particle_record.x_diagnostic[:, 0], particle_record.x_diagnostic[:, 1], particle_record.x_diagnostic[:, 2] + vx, vy, vz = particle_record.u[:, 0], particle_record.u[:, 1], particle_record.u[:, 2] x_wind = world['x_wind'] y_wind = world['y_wind'] @@ -243,4 +255,3 @@ def write_to_file(filename, time, data): f.write(f"{time}, {data}\n") return jax.debug.callback(write_to_file, filename, time, data, ordered=True) - diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py deleted file mode 100644 index f7ff14bd..00000000 --- a/PyPIC3D/diagnostics/vtk.py +++ /dev/null @@ -1,231 +0,0 @@ -import vtk -from vtk.util import numpy_support -import numpy as np -import os -from pyevtk.hl import gridToVTK, pointsToVTK -import jax.numpy as jnp - -def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): - """ - Plot the fields in a 3D grid. - - Args: - fieldx (ndarray): Array representing the x-component of the field. - fieldy (ndarray): Array representing the y-component of the field. - fieldz (ndarray): Array representing the z-component of the field. - t (float): Time value. - name (str): Name of the field. - dx (float): Spacing between grid points in the x-direction. - dy (float): Spacing between grid points in the y-direction. - dz (float): Spacing between grid points in the z-direction. - - Returns: - None - """ - Nx = fieldx.shape[0] - Ny = fieldx.shape[1] - Nz = fieldx.shape[2] - x = np.linspace(0, Nx, Nx) * dx - y = np.linspace(0, Ny, Ny) * dy - z = np.linspace(0, Nz, Nz) * dz - - # Create directory if it doesn't exist - directory = "./plots/fields" - if not os.path.exists(directory): - os.makedirs(directory) - - gridToVTK(f"./plots/fields/{name}_{t:09}", x, y, z, \ - cellData = {f"{name}_x" : np.asarray(fieldx), \ - f"{name}_y" : np.asarray(fieldy), f"{name}_z" : np.asarray(fieldz)}) -# plot the electric fields in the vtk file format - - -def plot_vtk_particles(particles, t, path): - """ - Plot the particles in VTK format. - - Args: - particles (Particles): The particles to be plotted. - t (ndarray): The time values. - path (str): The path to save the plot. - - Returns: - None - """ - if not os.path.exists(f"{path}/data/particles"): - os.makedirs(f"{path}/data/particles") - - particle_names = [species.get_name().replace(" ", "") for species in particles] - for species in particles: - name = species.get_name().replace(" ", "") - x, y, z = map(np.asarray, species.get_position()) - vx, vy, vz = map(np.asarray, species.get_velocity()) - # Get the position and velocity of the particles - q = np.asarray( species.get_charge() * np.ones_like(vx) ) - # Get the charge of the particles - - pointsToVTK(f"{path}/data/particles/{name}.{t:09}", x, y, z, \ - {"v": (vx, vy, vz), "q": q}) - # save the particles in the vtk file format - - -def plot_field_slice_vtk(field_slices, field_names, slice, grid, t, name, path, world): - """ - Plot a slice of a field in VTK format using Python VTK library. - - Args: - field_slices (list): List of 2D field slices to be plotted. - field_names (list): List of field names corresponding to the slices. - slice (int): Slice direction (0=x-slice, 1=y-slice, 2=z-slice, 3=full 3D). - grid (tuple): The grid dimensions (x, y, z). - t (int): The time step. - name (str): The name of the field. - path (str): The path to save the plot. - world (dict): World parameters containing grid information. - - Returns: - None - """ - - x, y, z = grid - nx, ny, nz = world['Nx'], world['Ny'], world['Nz'] - dx, dy, dz = world['dx'], world['dy'], world['dz'] - - if not os.path.exists(f"{path}/data/{name}_slice"): - os.makedirs(f"{path}/data/{name}_slice") - # Create directory if it doesn't exist - - # Create VTK structured grid - structured_grid = vtk.vtkStructuredGrid() - - if slice == 0: - x = np.asarray([x[nx//2]]) - elif slice == 1: - y = np.asarray([y[ny//2]]) - elif slice == 2: - z = np.asarray([z[nz//2]]) - - structured_grid.SetDimensions(x.shape[0], y.shape[0], z.shape[0]) - # Set the dimensions of the structured grid based on the slice type - - # Efficiently create all grid points using numpy meshgrid and bulk insert - X, Y, Z = np.meshgrid(x, y, z, indexing='ij') - coords = np.column_stack((Z.ravel(), Y.ravel(), X.ravel())) - points = vtk.vtkPoints() - points.SetData(numpy_support.numpy_to_vtk(coords, deep=True)) - structured_grid.SetPoints(points) - # Create points for the structured grid based on the slice type - - for idx, (field_slice, field_name) in enumerate(zip(field_slices, field_names)): - # Ensure field_slice is 2D and handle VTK ordering - field_data = np.asarray(field_slice) - if field_data.ndim == 2: - # VTK expects data in (k,j,i) order, but our slice is typically (j,k) - field_data = field_data.T.flatten() - else: - field_data = field_data.flatten() - - vtk_array = numpy_support.numpy_to_vtk(field_data) - vtk_array.SetName(field_name) - structured_grid.GetPointData().AddArray(vtk_array) - - # Write the VTK file - writer = vtk.vtkStructuredGridWriter() - writer.SetFileName(f"{path}/data/{name}_slice/{name}_slice_{t:09}.vtk") - writer.SetInputData(structured_grid) - writer.Write() - - -def plot_vectorfield_slice_vtk(field_slices, field_names, slice, grid, t, name, path, world): - """ - Plot a slice of a field in VTK format as vector data using Python VTK library. - - Args: - field_slices (list): List of 2D field slices to be plotted. Should be [Fx, Fy, Fz] for vector fields. - field_names (list): List of field names corresponding to the slices (e.g., ['Ex', 'Ey', 'Ez']). - slice (int): Slice direction (0=x-slice, 1=y-slice, 2=z-slice, 3=full 3D). - grid (tuple): The grid dimensions (x, y, z). - t (int): The time step. - name (str): The name of the field. - path (str): The path to save the plot. - world (dict): World parameters containing grid information. - - Returns: - None - """ - - x, y, z = grid - nx, ny, nz = world['Nx'], world['Ny'], world['Nz'] - dx, dy, dz = world['dx'], world['dy'], world['dz'] - - if not os.path.exists(f"{path}/data/{name}_slice"): - os.makedirs(f"{path}/data/{name}_slice") - - # Handle slicing - if slice == 0: - x = np.asarray([x[nx//2]]) - elif slice == 1: - y = np.asarray([y[ny//2]]) - elif slice == 2: - z = np.asarray([z[nz//2]]) - - # Create VTK structured grid - structured_grid = vtk.vtkStructuredGrid() - structured_grid.SetDimensions(x.shape[0], y.shape[0], z.shape[0]) - - # Create grid points - X, Y, Z = np.meshgrid(x, y, z, indexing='ij') - coords = np.column_stack((Z.ravel(), Y.ravel(), X.ravel())) - points = vtk.vtkPoints() - points.SetData(numpy_support.numpy_to_vtk(coords, deep=True)) - structured_grid.SetPoints(points) - - # Stack field slices as vector data - # Each field_slice should be 2D, shape (len(x), len(y)) or similar - # We flatten and stack them as (N, 3) for VTK vector - for field_slice, field_name in zip(field_slices, field_names): - field_arrays = [] - for comp in field_slice: - field_data = np.asarray(comp) - # convert to np array - if field_data.ndim == 2: - field_data = field_data.T.flatten() # VTK expects Fortran order - else: - field_data = field_data.flatten() - field_arrays.append(field_data) - # Stack as (N, 3) - vector_data = np.stack(field_arrays, axis=-1) - # If shape is (N, 3), convert to VTK - vtk_vector_array = numpy_support.numpy_to_vtk(vector_data, deep=True) - vtk_vector_array.SetName(f"{field_name}_vector") - structured_grid.GetPointData().AddArray(vtk_vector_array) - - # Write the VTK file - writer = vtk.vtkStructuredGridWriter() - writer.SetFileName(f"{path}/data/{name}_slice/{name}_slice_{t:09}.vtk") - writer.SetInputData(structured_grid) - writer.Write() - - -def write_slice(field_slice, x1, x2, t, name, path, dt): - """ - Plots a slice of a field and saves it in VTK format. - Parameters: - field_slice (numpy.ndarray): The 2D slice of the field to be plotted. - x1 (numpy.ndarray): The x-coordinates of the field slice. - x2 (numpy.ndarray): The y-coordinates of the field slice. - t (int): The time step or index for the slice. - name (str): The name of the field or slice. - path (str): The directory path where the VTK file will be saved. - dt (float): The time step size (not used in the function but included in parameters). - Returns: - None - """ - - x3 = np.zeros(1) - - field_slice = jnp.expand_dims(field_slice, axis=-1) - - gridToVTK(f"{path}/data/{name}_slice/{name}_slice_{t:09}", x1, x2, x3, \ - cellData = {f"{name}" : field_slice}) - # plot the slice of the field in the vtk file format \ No newline at end of file diff --git a/PyPIC3D/errors.py b/PyPIC3D/errors.py deleted file mode 100644 index 332d6a91..00000000 --- a/PyPIC3D/errors.py +++ /dev/null @@ -1,103 +0,0 @@ -from jax import jit -import jax.numpy as jnp -from functools import partial -# import external libraries - -from PyPIC3D.solvers.pstd import spectral_laplacian, spectral_divergence -from PyPIC3D.solvers.fdtd import centered_finite_difference_laplacian, centered_finite_difference_divergence -# import functions from the PyPIC3D package - -@partial(jit, static_argnums=(4, 5)) -def compute_pe(phi, rho, constants, world, solver, bc='periodic'): - """ - Compute the relative percentage difference of the Poisson solver. - - Args: - phi (ndarray): The potential field. - rho (ndarray): The charge density. - eps (float): The permittivity. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - bc (str): The boundary condition. - - Returns: - float: The relative percentage difference of the Poisson solver. - """ - eps = constants['eps'] - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - if solver == 'spectral': - x = spectral_laplacian(phi, world) - elif solver == 'fdtd': - x = centered_finite_difference_laplacian(phi, dx, dy, dz, bc) - elif solver == 'autodiff': - return 0 - poisson_error = x + rho/eps - magnitude = jnp.mean(jnp.abs(rho/eps)) + 1e-16 - return jnp.mean(jnp.abs(poisson_error)) / magnitude - -@partial(jit, static_argnums=(4, 5)) -def compute_magnetic_divergence_error(Bx, By, Bz, world, solver, bc='periodic'): - """ - Compute the error in the divergence of the magnetic field for different solvers. - - Args: - Bx (ndarray): The x-component of the magnetic field. - By (ndarray): The y-component of the magnetic field. - Bz (ndarray): The z-component of the magnetic field. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - bc (str): The boundary condition. - - Returns: - float: The error in the divergence of the magnetic field. - """ - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - - if solver == 'spectral': - divB = spectral_divergence(Bx, By, Bz, world) - elif solver == 'fdtd': - divB = centered_finite_difference_divergence(Bx, By, Bz, dx, dy, dz, bc) - elif solver == 'autodiff': - return 0 - divergence_error = jnp.mean(jnp.abs(divB)) - return divergence_error - -@partial(jit, static_argnums=(6, 7)) -def compute_electric_divergence_error(Ex, Ey, Ez, rho, constants, world, solver, bc='periodic'): - """ - Compute the error in the divergence of the electric field using the charge density and the components of the electric field. - - Args: - Ex (ndarray): The x-component of the electric field. - Ey (ndarray): The y-component of the electric field. - Ez (ndarray): The z-component of the electric field. - rho (ndarray): The charge density. - eps (float): The permittivity. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - bc (str): The boundary condition. - - Returns: - float: The error in the divergence of the electric field. - """ - eps = constants['eps'] - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - - if solver == 'spectral': - divE = spectral_divergence(Ex, Ey, Ez, world) - elif solver == 'fdtd': - divE = centered_finite_difference_divergence(Ex, Ey, Ez, dx, dy, dz, bc) - elif solver == 'autodiff': - return 0 - - divergence_error = jnp.mean(jnp.abs(divE - rho / eps)) - return divergence_error \ No newline at end of file diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 23cc76e5..53322fac 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -1,259 +1,189 @@ # Christopher Woolford Dec 5, 2024 -# This contains the evolution loop for the 3D PIC code that calculates the electric and magnetic fields and updates the particles. +# This contains the public evolution-loop names for the 3D PIC code. -import jax.numpy as jnp -import jax -from jax import jit from functools import partial -from PyPIC3D.pusher.particle_push import ( - particle_push -) - -from PyPIC3D.solvers.first_order_yee import ( - update_E, update_B -) +from jax import jit -from PyPIC3D.solvers.electrostatic_yee import ( - calculate_electrostatic_fields -) +import jax -from PyPIC3D.solvers.vector_potential import ( - E_from_A, B_from_A, update_vector_potential +from PyPIC3D.deposition.Esirkepov import Esirkepov_current +from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.particles.particle_tile_communication import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, ) - +from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields +from PyPIC3D.solvers.first_order_yee import update_B, update_E from PyPIC3D.utils import add_external_fields -@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 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. - - Args: - particles (list): List of particle objects to be advanced. - fields (tuple): Tuple containing the field arrays (E, B, J, rho, phi, external_fields). - world (dict): Dictionary containing simulation parameters (e.g., time step 'dt'). - constants (dict): Dictionary of physical constants used in the simulation. - curl_func (callable): Function to compute the curl of a field (not used in this function). - J_func (callable): Function to compute the current density (not used in this function). - solver (callable): Function or object to solve the Poisson equation for the electric field. - - Returns: - tuple: Updated particles list and fields tuple (E, B, J, rho, phi, external_fields). - """ - - E, B, J, rho, phi, external_fields = fields - # unpack the fields - center_grid = world['grids']['center'] - vertex_grid = world['grids']['vertex'] - push_E, push_B = add_external_fields(E, B, external_fields) - # particles see evolved fields plus external-only fields - - ################ PARTICLE PUSH ######################################################################################## - for i in range(len(particles)): - - particles[i] = particle_push(particles[i], push_E, push_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 - - ############### SOLVE E FIELD ############################################################################################ - E, phi, rho = calculate_electrostatic_fields(world, particles, constants, rho, phi, solver, 'periodic') - # calculate the electric field using the Poisson equation - fields = (E, B, J, rho, phi, external_fields) - # pack the fields into a tuple +__all__ = ["time_loop_electrodynamic", "time_loop_electrostatic"] - for i in range(len(particles)): - particles[i].boundary_conditions() - # apply boundary conditions to the particles - return particles, fields +def _particle_kernel_world(world): + kernel_world = dict(world) + kernel_world.pop("current_deposition", None) + kernel_world.pop("current_filter", None) + return kernel_world -@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"): +def time_loop_electrodynamic( + particles, + species_config, + fields, + world, + constants, + 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 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, - 5) Particle boundary condition enforcement. - Parameters - ---------- - particles : Sequence[object] - Collection of particle objects. Each particle is expected to be compatible - with `particle_push(...)`, provide an `update_position()` method, and a - `boundary_conditions()` method. - fields : tuple - Tuple of field arrays/objects in the form - `(E, B, J, rho, phi, external_fields, pml_state)`. - `pml_state` is `None` when the ordinary Yee update is used. - world : dict - Simulation configuration. Must contain `world['dt']` (time step). - constants : object or dict - Physical/normalization constants required by push, deposition, and updates. - curl_func : callable - Function/operator used by field update routines to compute curls. - J_func : callable - Current deposition function with signature like - `J_func(particles, J, constants, world) -> J`. - solver : object - Reserved/placeholder for a solver interface (currently unused). - relativistic : bool, optional - If True, perform a relativistic Boris push; otherwise use the - non-relativistic variant. - Returns - ------- - particles : Sequence[object] - Updated particle collection after push, position update, and boundary - conditions. - fields : tuple - Updated fields tuple `(E, B, J, rho, phi, external_fields, pml_state)` - with new E, B, and J. `rho`, `phi`, and `external_fields` are passed - through unchanged. - Notes - ----- - - The update order is: deposit J -> update E -> update B. - - `solver` is accepted for API compatibility but is not used in this function. + Advance a tiled electrodynamic PIC system by one time step. """ + E, B, J, rho, phi, external_fields, pml_state, overflow_previous = fields + # unpack the tiled field state + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + particle_world = _particle_kernel_world(world) + # get tiled layout information from the world contract - E, B, J, rho, phi, external_fields, pml_state = fields - # unpack the fields - center_grid = world['grids']['center'] - vertex_grid = world['grids']['vertex'] push_E, push_B = add_external_fields(E, B, external_fields) # particles see evolved fields plus external-only fields - ################ PARTICLE PUSH ######################################################################################## - for i in range(len(particles)): - - particles[i] = particle_push(particles[i], push_E, push_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 - - ################ FIELD UPDATE ################################################################################################ - J = J_func(particles, J, constants, world) - # calculate the current density based on the selected method - E, pml_state = update_E(E, B, J, world, constants, curl_func, pml_state) - # update the electric field using the curl of the magnetic field - B, pml_state = update_B(E, B, world, constants, curl_func, pml_state) - # update the magnetic field using the curl of the electric field - - for i in range(len(particles)): - particles[i].boundary_conditions() - # apply boundary conditions to the particles - - fields = (E, B, J, rho, phi, external_fields, pml_state) - # pack the fields into a tuple - + particles = particle_push( + particles, + species_config, + push_E, + push_B, + particle_world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + # use the selected tiled pusher for particle velocities + + def direct_deposition_step(state): + particles, J_tiles, overflow_previous = state + particles = update_tiled_particle_positions(particles, species_config, world["dt"] / 2) + # update particle positions to the centered direct-current deposition time + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + # wrap particles and move them into their owning tiles. + overflow = overflow_previous | overflow + # keep fixed-capacity tile overflow visible to the Python driver + J_tiles = J_from_rhov( + particles, + species_config, + J_tiles, + constants, + world, + filter=world["current_filter"], + ) + # deposit current directly into tile-local Yee current arrays + particles = update_tiled_particle_positions(particles, species_config, world["dt"] / 2) + # complete the full particle position update + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + # refresh tile ownership after the full position update. + overflow = overflow_previous | overflow + return particles, J_tiles, overflow + # if the direct deposition method is selected, first refresh the particle tiles, then deposit current directly into the tiled J arrays + + def esirkepov_deposition_step(state): + particles, J_tiles, overflow_previous = state + J_tiles = Esirkepov_current(particles, species_config, J_tiles, constants, world) + # deposit current into the tiled J arrays using the Esirkepov method, which requires old and new particle positions + particles = update_tiled_particle_positions(particles, species_config, world["dt"]) + # update particle positions to the new time step + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + # refresh tile ownership after the full position update + overflow = overflow_previous | overflow + return particles, J_tiles, overflow + # if the Esirkepov deposition method is selected, first deposit current into the tiled J arrays, then refresh the particle tiles + + particles, J, overflow = jax.lax.cond( + world["current_deposition"] == "esirkepov", + esirkepov_deposition_step, + direct_deposition_step, + (particles, J, overflow_previous), + ) + # deposit current into the tiled J arrays using the selected deposition method + + E, pml_state = update_E(E, B, J, world, constants, pml_state) + # update electric field from B and the supplied current + # for no pml, the pml_state is None, and the update_E function returns None for the pml_state + + B, pml_state = update_B(E, B, world, constants, pml_state) + # update magnetic field from the newly updated electric field + # for no pml, the pml_state is None, and the update_B function returns None for the pml_state + + fields = (E, B, J, rho, phi, external_fields, pml_state, overflow) + # pack the tiled field state return particles, fields -@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"): +def time_loop_electrostatic( + particles, + species_config, + fields, + world, + constants, + 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. - - This routine: - 1) Pushes particle velocities/positions using the current E and B fields. - 2) Updates the vector potential A via the current density J. - 3) Recomputes E and B from the updated vector potential. - 4) Recomputes J from particle motion and applies particle boundary conditions. - - Parameters - ---------- - particles : Sequence - Iterable of particle objects. Each particle is expected to be compatible - with `particle_push(...)`, provide `update_position()`, and - `boundary_conditions()` methods. - fields : tuple - Field tuple in the order `(E, B, J, rho, phi, external_fields, A2, A1, A0)`, - where `external_fields` is added only for particle pushes, `A2` denotes - the newest vector potential, `A1` the previous, and `A0` the older one. - world : dict - Simulation parameters. Must at least contain `'dt'` (time step). - constants : Any - Physical/constants container passed through to lower-level routines. - curl_func : Callable - Reserved/unused in the current implementation (kept for API compatibility). - J_func : Callable - Function to compute/update current density, called as: - `J_func(particles, J, constants, world)`. - solver : Any - Reserved/unused in the current implementation (kept for API compatibility). - relativistic : bool, optional - If True, the particle pusher is invoked in relativistic mode. - - Returns - ------- - particles : Sequence - Updated particle collection after push, position update, and boundary - conditions. - fields : tuple - Updated field tuple `(E, B, J, rho, phi, external_fields, A2, A1, A0)` - after advancing A and recomputing E, B, and J. - - Notes - ----- - - The vector potential history is shifted each step (`A0 <- A1`, `A1 <- A2`), - then `A2` is computed from the current density via `update_vector_potential`. - - Electric and magnetic fields are recomputed using `E_from_A` and `B_from_A` - with `interpolation_order=2`. - - Several parameters (`curl_func`, `solver`) are - currently not used inside this function but may be part of a broader API. - """ - - E, B, J, rho, phi, external_fields, A2, A1, A0 = fields - # unpack the fields - center_grid = world['grids']['center'] - vertex_grid = world['grids']['vertex'] - push_E, push_B = add_external_fields(E, B, external_fields) - # particles see evolved fields plus external-only fields - - ################ PARTICLE PUSH ######################################################################################## - for i in range(len(particles)): - - particles[i] = particle_push(particles[i], push_E, push_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 + Advance a tiled electrostatic PIC system by one time step. - ################ FIELD UPDATE ################################################################################################ - A0 = A1 - A1 = A2 - # update the vector potential for the next iteration - A2 = update_vector_potential(J, world, constants, A1, A0) - # update the vector potential based on the current density J - - E = E_from_A(A2, A1, A0, world, center_grid, vertex_grid, interpolation_order=2) - # calculate the electric field from the vector potential using centered finite difference - B = B_from_A(A1, world, center_grid, vertex_grid, interpolation_order=2) - # calculate the magnetic field from the vector potential using centered finite difference - J = J_func(particles, J, constants, world) - # calculate the current density using the selected method - - for i in range(len(particles)): - particles[i].boundary_conditions() - # apply boundary conditions to the particles - - fields = (E, B, J, rho, phi, external_fields, A2, A1, A0) - # pack the fields into a tuple + The particle push and retile use tile-local fields. Charge density is + deposited into tiled scalar storage, assembled for the existing global + Poisson solve, then the solved potential is tiled again before computing + tile-local electrostatic E. + """ + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields + # unpack the tiled field state + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + particle_world = _particle_kernel_world(world) + # get tiled layout information from the world contract + + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) + # particles see evolved fields plus prescribed external fields + + particles = particle_push( + particles, + species_config, + push_E_tiles, + push_B_tiles, + particle_world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + # push velocities using the selected tiled particle pusher + + particles = update_tiled_particle_positions(particles, species_config, world["dt"]) + # update particle forward positions before depositing rho + + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + # keep fixed-capacity tile overflow visible to the Python driver + + E_tiles, phi_tiles, rho_tiles = calculate_tiled_electrostatic_fields( + world, + particles, + species_config, + constants, + rho_tiles, + phi_tiles, + solver, + "periodic", + ) + # solve electrostatic fields from tiled charge density + + fields = (E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow) + # pack the tiled field state return particles, fields diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index b7faa3ce..0ccdfdf5 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -2,7 +2,6 @@ from jax.experimental import mesh_utils from jax.sharding import Mesh, PartitionSpec, NamedSharding import os -import functools from functools import partial import toml import matplotlib.pyplot as plt @@ -17,49 +16,41 @@ from PyPIC3D.utils import ( courant_condition, update_parameters_from_toml, - build_yee_grid, convert_to_jax_compatible, load_external_fields_from_toml, + convert_to_jax_compatible, load_external_fields_from_toml, print_stats, particle_sanity_check, build_plasma_parameters_dict, - make_dir, compute_energy, build_collocated_grid, add_external_fields + make_dir, compute_energy, add_external_fields ) - -from PyPIC3D.solvers.pstd import ( - spectral_curl -) - -from PyPIC3D.solvers.fdtd import ( - centered_finite_difference_curl +from PyPIC3D.utilities.grids import ( + build_collocated_grid, + build_tiled_yee_grids, + build_yee_grid, ) - from PyPIC3D.diagnostics.plotting import ( plot_initial_histograms ) +from PyPIC3D.diagnostics.output_adapters import particles_for_output from PyPIC3D.diagnostics.openPMD import ( write_openpmd_initial_particles, write_openpmd_initial_fields ) -from PyPIC3D.particles.flat_particles import ( - to_flat_particles, check_flat_compat +from PyPIC3D.boundary_conditions.ghost_cells import ( + update_tiled_vector_ghost_cells, + update_tiled_vector_ghost_cells_for_pml, ) from PyPIC3D.evolve import ( - time_loop_electrodynamic, time_loop_electrostatic, time_loop_vector_potential + time_loop_electrodynamic, + time_loop_electrostatic, ) -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 from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.PML import ( - initialize_pml_state, + initialize_tiled_pml_state, load_pml_from_toml, - update_ghost_cells_for_pml, ) @@ -76,6 +67,78 @@ def _encode_field_bc(bc_name): return bc_codes[bc_name] +def _encode_particle_bc(bc_name): + """ + Encode global particle boundary condition labels into integer codes. + """ + bc_codes = { + "periodic": 0, + "reflecting": 1, + "absorbing": 2, + } + if bc_name not in bc_codes: + raise ValueError(f"Unsupported particle boundary condition: {bc_name}") + return bc_codes[bc_name] + + +def validate_field_solver(solver): + """ + Keep the active field-solver names explicit so stale configs do not silently + fall through to a different numerical update. + """ + supported_solvers = ("electrodynamic_yee", "electrostatic") + if solver not in supported_solvers: + raise ValueError( + f"Unsupported solver: {solver}. Use 'electrodynamic_yee' or 'electrostatic'." + ) + + +def _tile_shape_from_parameters(simulation_parameters): + return ( + int(simulation_parameters["particle_tile_nx"]), + int(simulation_parameters["particle_tile_ny"]), + int(simulation_parameters["particle_tile_nz"]), + ) + + +def _encode_current_calculation(current_calculation): + if current_calculation not in ("j_from_rhov", "esirkepov"): + raise ValueError("Unsupported current_calculation. Use 'j_from_rhov' or 'esirkepov'.") + if current_calculation == "esirkepov": + return "esirkepov" + return "direct" + + +def _validate_current_filter_contract(simulation_parameters): + if ( + simulation_parameters["current_calculation"] == "esirkepov" + and simulation_parameters["filter_j"] != "none" + ): + raise ValueError("Esirkepov current filtering is not supported; use filter_j='none'.") + + +def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_active): + """ + Keep the first tile-native PIC path tied to the kernels that exist today. + """ + + if simulation_parameters["current_calculation"] not in ("j_from_rhov", "esirkepov"): + raise ValueError("Yee runtime currently supports current_calculation='j_from_rhov' or 'esirkepov'") + if simulation_parameters["particle_pusher"] not in ("boris", "higuera_cary"): + raise ValueError("Yee runtime currently supports only particle_pusher='boris' or 'higuera_cary'") + if simulation_parameters["filter_j"] not in ("none", "digital", "bilinear"): + raise ValueError("Yee runtime currently supports only filter_j='none', filter_j='digital', or 'bilinear'") + tile_shape = _tile_shape_from_parameters(simulation_parameters) + grid_shape = ( + int(simulation_parameters["Nx"]), + int(simulation_parameters["Ny"]), + int(simulation_parameters["Nz"]), + ) + for cells, tile_width in zip(grid_shape, tile_shape): + if cells % tile_width != 0: + raise ValueError("Yee runtime requires the shared tile shape to divide Nx/Ny/Nz exactly") + + def default_parameters(): """ Returns a dictionary of default parameters for the simulation. @@ -96,12 +159,11 @@ def default_parameters(): "plot_errors": False, "plot_dispersion": False, 'plot_chargeconservation': False, - "plot_vtk_particles": False, - "plot_vtk_scalars" : False, - "plot_vtk_vectors" : False, "plot_openpmd_particles": False, "plot_openpmd_fields": False, "plotting_interval": 10, + "openpmd_field_queue_size": 2, + "openpmd_particle_queue_size": 2, "dump_particles": False, "dump_fields": False, } @@ -110,9 +172,10 @@ def default_parameters(): simulation_parameters = { "name": "Default Simulation", "output_dir": os.getcwd(), - "solver": "fdtd", # solver: spectral, fdtd, vector_potential, curl_curl - "fast_backend": "flat", # flat | default (flat when compatible, else fallback) - "particle_bc": "periodic", # particle boundary conditions: periodic, absorbing, reflecting + "solver": "electrodynamic_yee", # solver: electrodynamic_yee or electrostatic + "particle_x_bc": "periodic", # particle x boundary conditions: periodic, absorbing, reflecting + "particle_y_bc": "periodic", # particle y boundary conditions: periodic, absorbing, reflecting + "particle_z_bc": "periodic", # particle z boundary conditions: periodic, absorbing, reflecting # "bc": "periodic", # boundary conditions: periodic, dirichlet, neumann "x_bc": "periodic", # x boundary conditions: periodic, conducting "y_bc": "periodic", # y boundary conditions: periodic, conducting @@ -126,7 +189,6 @@ def default_parameters(): "t_wind": 1e-12, # size of the temporal window in seconds "dt": None, # time step in seconds "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 @@ -135,6 +197,11 @@ def default_parameters(): "cfl" : 1.0, # CFL condition number "ds_per_debye" : None, # number of grid spacings per debye length "shape_factor" : 1, # shape factor for the simulation (1 for 1st order, 2 for 2nd order) + "guard_cells": 2, # tile guard-cell depth shared by tiled fields and tiled currents + "particle_tile_nx": None, # number of x cells per shared field/particle tile + "particle_tile_ny": None, # number of y cells per shared field/particle tile + "particle_tile_nz": None, # number of z cells per shared field/particle tile + "particle_tile_capacity_factor": 1.0, # inactive particle slot headroom per tile "current_calculation": "j_from_rhov", # current calculation method: esirkepov, villasenor_buneman, j_from_rhov "filter_j": "bilinear", # filter for the current density: bilinear, digital, none } @@ -183,13 +250,12 @@ def initialize_simulation(toml_file): plotting_parameters (dict): Dictionary containing parameters for plotting. plasma_parameters (dict): Dictionary containing plasma parameters. M (jax.numpy.ndarray or None): Preconditioner matrix, if neural network preconditioning is used. - solver (str): Solver type, either "spectral" or "fdtd". - electrostatic (bool): Flag indicating if the simulation is electrostatic. + solver (str): Solver type, either "electrodynamic_yee" or "electrostatic". + electrostatic (bool): Derived flag indicating if the simulation is electrostatic. verbose (bool): Flag indicating if verbose output is enabled. GPUs (int): Number of GPUs to use. start (float): Start time of the simulation initialization. Nt (int): Number of time steps. - curl_func (function): Function to compute the curl of the fields. pecs (list): List of perfectly electrical conductor boundaries. lasers (list): List of laser objects initialized from the configuration file. surfaces (list): List of material surfaces initialized from the configuration file. @@ -207,18 +273,34 @@ def initialize_simulation(toml_file): x_wind, y_wind, z_wind = simulation_parameters['x_wind'], simulation_parameters['y_wind'], simulation_parameters['z_wind'] Nx, Ny, Nz = simulation_parameters['Nx'], simulation_parameters['Ny'], simulation_parameters['Nz'] t_wind = simulation_parameters['t_wind'] - electrostatic = simulation_parameters['electrostatic'] solver = simulation_parameters['solver'] + if "electrostatic" in simulation_parameters: + raise ValueError("Use solver='electrostatic' instead of the deprecated electrostatic flag.") + validate_field_solver(solver) + electrostatic = solver == "electrostatic" relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] - validate_particle_pusher(particle_pusher) + + if simulation_parameters['particle_tile_nx'] is None: + simulation_parameters["particle_tile_nx"] = int(Nx) + if simulation_parameters['particle_tile_ny'] is None: + simulation_parameters["particle_tile_ny"] = int(Ny) + if simulation_parameters['particle_tile_nz'] is None: + simulation_parameters["particle_tile_nz"] = int(Nz) + + if electrostatic: + simulation_parameters["particle_tile_nx"] = int(Nx) + simulation_parameters["particle_tile_ny"] = int(Ny) + simulation_parameters["particle_tile_nz"] = int(Nz) + # Electrostatic uses one tile covering the whole domain. The leading + # tile axes remain present, but their extent is one in each direction. + guard_cells = int(simulation_parameters["guard_cells"]) + guard_cells = max(guard_cells, 2) + _validate_current_filter_contract(simulation_parameters) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] # set the simulation parameters - # if 'ncores' in simulation_parameters: - # os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={simulation_parameters['ncores']}' - # set the number of cores to use setup_write_dir(simulation_parameters, plotting_parameters) # setup the write directory @@ -258,11 +340,20 @@ def initialize_simulation(toml_file): 'x_wind': x_wind, 'y_wind': y_wind, 'z_wind': z_wind, + 'shape_factor': simulation_parameters['shape_factor'], + 'guard_cells': guard_cells, + 'current_deposition': _encode_current_calculation(simulation_parameters['current_calculation']), + 'current_filter': simulation_parameters['filter_j'], 'boundary_conditions': { 'x': _encode_field_bc(simulation_parameters['x_bc']), 'y': _encode_field_bc(simulation_parameters['y_bc']), 'z': _encode_field_bc(simulation_parameters['z_bc']), }, + 'particle_boundary_conditions': { + 'x': _encode_particle_bc(simulation_parameters['particle_x_bc']), + 'y': _encode_particle_bc(simulation_parameters['particle_y_bc']), + 'z': _encode_particle_bc(simulation_parameters['particle_z_bc']), + }, } # set the simulation world parameters @@ -270,8 +361,9 @@ def initialize_simulation(toml_file): if toml_file is not None: raw_pml = toml_file.get("pml", []) pml_active = bool(raw_pml) - if pml_active and (solver != "fdtd" or electrostatic): - raise ValueError("PML is only supported for the fdtd electrodynamic solver") + if pml_active and electrostatic: + raise ValueError("PML is only supported for the electrodynamic_yee solver") + _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_active) world["pml"] = load_pml_from_toml(raw_pml, world, constants) world = convert_to_jax_compatible(world) @@ -279,40 +371,64 @@ def initialize_simulation(toml_file): simulation_parameters = convert_to_jax_compatible(simulation_parameters) plotting_parameters = convert_to_jax_compatible(plotting_parameters) # convert the world parameters to jax compatible format + world["guard_cells"] = guard_cells + simulation_parameters["guard_cells"] = guard_cells + # Guard-cell depth controls static tiled array shapes, so keep it as a + # Python integer for jitted tile kernels. - # if solver == "vector_potential": - # B_grid, E_grid = build_collocated_grid(world) - # # build the grid for the fields - # else: - B_grid, E_grid = build_yee_grid(world) - # build the Yee grid for the fields + if electrostatic: + B_grid, E_grid = build_collocated_grid(world) + # electrostatic E comes from a symmetric gradient and is colocated at cell centers + else: + B_grid, E_grid = build_yee_grid(world) + # build the Yee grid for the electrodynamic fields world['grids'] = { 'vertex': E_grid, 'center': B_grid, } # set the grids in the world parameters + tile_shape = _tile_shape_from_parameters(simulation_parameters) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, guard_cells) + world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + world["grids"]["tiled_center_grid"] = tiled_center_grid + simulation_parameters["tile_shape"] = tile_shape + # set the shared tiled field/particle geometry before allocating fields if not os.path.exists(f"{simulation_parameters['output_dir']}/data"): os.makedirs(f"{simulation_parameters['output_dir']}/data") # create the data directory if it doesn't exist ################################### INITIALIZE PARTICLES AND FIELDS ######################################################## - particles = load_particles_from_toml(toml_file, simulation_parameters, world, constants) - # load the particles from the configuration file - - for species in particles: - name = species.get_name() - name = name.replace(" ", "_") + particles, species_config, particle_species_names, particle_metadata = load_particles_from_toml( + toml_file, + simulation_parameters, + world, + constants, + ) + # load particles directly into tiled runtime storage + simulation_parameters["particle_species_names"] = particle_species_names + simulation_parameters["particle_species_metadata"] = particle_metadata + # keep species names and scalar metadata available for diagnostics/output + + initial_particle_records = particles_for_output( + particles, + species_config=species_config, + species_names=particle_species_names, + world=world, + ) + for particle_record in initial_particle_records: + name = particle_record.name.replace(" ", "_") # replace spaces with underscores in the name - plot_initial_histograms(species, world, path=f"{simulation_parameters['output_dir']}/data", name=name) + plot_initial_histograms(particle_record, world, path=f"{simulation_parameters['output_dir']}/data", name=name) # plot the initial histograms of the particles print_stats(world) # print the statistics of the simulation - if particles: - plasma_parameters = build_plasma_parameters_dict(world, constants, particles[0], dt) + if particle_metadata: + plasma_parameters = build_plasma_parameters_dict(world, constants, particle_metadata[0], dt) # build the plasma parameters dictionary else: plasma_parameters = {} @@ -322,27 +438,17 @@ def initialize_simulation(toml_file): # ensure the arrays for the particles are of the correct shape if plotting_parameters['dump_particles']: - write_openpmd_initial_particles(particles, world, constants, simulation_parameters['output_dir']) + write_openpmd_initial_particles( + particles, + world, + constants, + simulation_parameters['output_dir'], + species_config=species_config, + species_names=particle_species_names, + ) # write the initial particles to an openPMD file - - fast_backend = simulation_parameters.get("fast_backend", "flat") - if fast_backend == "flat": - if electrostatic or solver == "vector_potential": - print("fast_backend='flat' not supported for electrostatic/vector_potential; falling back to default") - simulation_parameters["fast_backend"] = "default" - elif not check_flat_compat(particles): - print("fast_backend='flat' incompatible with species layout; falling back to default") - simulation_parameters["fast_backend"] = "default" - else: - print("Using flat particle backend for simulation") - particles = to_flat_particles(particles) - simulation_parameters["fast_backend"] = "flat" - # check if we can use the flat particle backend for faster performance, and convert to flat particles if possible - - - - E, B, J, phi, rho = initialize_fields(Nx, Ny, Nz) + E, B, J, phi, rho = initialize_fields(world) # initialize the electric and magnetic fields external_fields = ( tuple(jax.numpy.zeros_like(comp) for comp in E), @@ -353,43 +459,39 @@ def initialize_simulation(toml_file): # load any external fields fields = [component for field in [E, B, J] for component in field] # convert the E, B, and J tuples into one big list - fields, external_fields = load_external_fields_from_toml(fields, external_fields, toml_file) + fields, external_fields = load_external_fields_from_toml(fields, external_fields, toml_file, world) # route configured fields into either evolved fields or external-only fields E, B, J = fields[:3], fields[3:6], fields[6:9] # convert the fields list back into tuples - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the boundary condition codes - if pml_active: - E = tuple(update_ghost_cells_for_pml(comp, world) for comp in E) - B = tuple(update_ghost_cells_for_pml(comp, world) for comp in B) + E = update_tiled_vector_ghost_cells_for_pml(E, world, num_guard_cells=guard_cells, tile_shape=tile_shape) + B = update_tiled_vector_ghost_cells_for_pml(B, world, num_guard_cells=guard_cells, tile_shape=tile_shape) else: - E = tuple(update_ghost_cells(comp, bc_x, bc_y, bc_z) for comp in E) - B = tuple(update_ghost_cells(comp, bc_x, bc_y, bc_z) for comp in B) + E = update_tiled_vector_ghost_cells(E, world, num_guard_cells=guard_cells, tile_shape=tile_shape) + B = update_tiled_vector_ghost_cells(B, world, num_guard_cells=guard_cells, tile_shape=tile_shape) # fill ghost cells for the initial E and B fields external_E, external_B = external_fields if pml_active: - external_E = tuple(update_ghost_cells_for_pml(comp, world) for comp in external_E) - external_B = tuple(update_ghost_cells_for_pml(comp, world) for comp in external_B) + external_E = update_tiled_vector_ghost_cells_for_pml(external_E, world, num_guard_cells=guard_cells, tile_shape=tile_shape) + external_B = update_tiled_vector_ghost_cells_for_pml(external_B, world, num_guard_cells=guard_cells, tile_shape=tile_shape) else: - external_E = tuple(update_ghost_cells(comp, bc_x, bc_y, bc_z) for comp in external_E) - external_B = tuple(update_ghost_cells(comp, bc_x, bc_y, bc_z) for comp in external_B) + external_E = update_tiled_vector_ghost_cells(external_E, world, num_guard_cells=guard_cells, tile_shape=tile_shape) + external_B = update_tiled_vector_ghost_cells(external_B, world, num_guard_cells=guard_cells, tile_shape=tile_shape) external_fields = (external_E, external_B) # fill ghost cells for external fields before they are interpolated to particles - if solver == "spectral": - curl_func = functools.partial(spectral_curl, world=world) - else: - curl_func = functools.partial(centered_finite_difference_curl, dx=dx, dy=dy, dz=dz, bc="periodic") - - ######################### COMPUTE INITIAL ENERGY ######################################################## total_E, total_B = add_external_fields(E, B, external_fields) # energy diagnostics use the same total fields that the particle pusher sees - e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants) + e_energy, b_energy, kinetic_energy = compute_energy( + particles, + total_E, + total_B, + world, + constants, + species_config=species_config, + ) # compute the initial energy of the system print(f"Initial Electric Field Energy: {e_energy:.2e} J") print(f"Initial Magnetic Field Energy: {b_energy:.2e} J") @@ -407,43 +509,33 @@ def initialize_simulation(toml_file): if electrostatic: print("Using electrostatic solver") evolve_loop = time_loop_electrostatic - - elif solver == "vector_potential": - raise NotImplementedError("Vector potential solver is not fully functional yet.") - print("Using vector potential solver") - evolve_loop = time_loop_vector_potential - else: - print(f"Using electrodynamic solver with: {solver}") + print("Using electrodynamic Yee solver") evolve_loop = time_loop_electrodynamic # set the evolve loop function based on the electrostatic flag if simulation_parameters['current_calculation'] == "esirkepov": print("Using Esirkepov current calculation method") - # raise NotImplementedError("Esirkepov current calculation method is not fully functional yet.") - J_func = Esirkepov_current elif simulation_parameters['current_calculation'] == "j_from_rhov": print(f"Using J from rhov current calculation method with filter: {simulation_parameters['filter_j']}") - J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) - if solver == "vector_potential": - A2, A1, A0 = initialize_vector_potential(J, world, constants) - # initialize the vector potential A based on the current density J - fields = (E, B, J, rho, phi, external_fields, A2, A1, A0) - # define the fields tuple for the vector potential solver - elif electrostatic: - fields = (E, B, J, rho, phi, external_fields) + external_E, external_B = external_fields + external_fields = (external_E, external_B) + print(f"Using tiled Yee storage with tile shape: {tile_shape}") + + overflow = jnp.asarray(False) + if electrostatic: + fields = (E, B, J, rho, phi, external_fields, None, overflow) # define the fields tuple for the electrostatic solver else: - if pml_active: - pml_state = initialize_pml_state(world) - else: - pml_state = None - # electrodynamic FDTD always carries the PML state slot; None means + pml_state = None + # electrodynamic Yee always carries the PML state slot; None means # ordinary, unstretched Yee derivatives. - fields = (E, B, J, rho, phi, external_fields, pml_state) - # define the fields tuple for the electrodynamic FDTD solver + if pml_active: + pml_state = initialize_tiled_pml_state(world, tile_shape) + fields = (E, B, J, rho, phi, external_fields, pml_state, overflow) + # define the fields tuple for the electrodynamic Yee solver if plotting_parameters['dump_fields']: write_openpmd_initial_fields(fields, world, simulation_parameters['output_dir'], filename="initial_fields.h5") @@ -453,53 +545,57 @@ def initialize_simulation(toml_file): if GPUs: print(f"GPUs Detected! Using GPUs for simulation\n") particles = jax.device_put(particles, jax.devices("gpu")[0]) + if species_config is not None: + species_config = jax.device_put(species_config, jax.devices("gpu")[0]) # put the particles on the GPU if GPUs are enabled return evolve_loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, \ - solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic, particle_pusher + solver, electrostatic, verbose, GPUs, Nt, relativistic, particle_pusher, species_config +def build_tiled_array(world, dtype=jnp.float64): + """ + Build one zero-filled tiled field component from the geometry in world. + """ -def initialize_fields(Nx, Ny, Nz): + tile_nx, tile_ny, tile_nz = [int(width) for width in world["tile_shape"]] + Nx = int(world["Nx"]) + Ny = int(world["Ny"]) + Nz = int(world["Nz"]) + g = int(world["guard_cells"]) + ntx = Nx // tile_nx + nty = Ny // tile_ny + ntz = Nz // tile_nz + tiled_shape = (ntx, nty, ntz, tile_nx + 2 * g, tile_ny + 2 * g, tile_nz + 2 * g) + return jnp.zeros(shape=tiled_shape, dtype=dtype) + + +def initialize_fields(world): """ Initializes the electric and magnetic field arrays, as well as the electric potential and charge density arrays. - All arrays include ghost cells and have shape (Nx+2, Ny+2, Nz+2). - The physical interior corresponds to indices [1:-1, 1:-1, 1:-1]. - - Args: - Nx (int): Number of physical grid points in the x-direction. - Ny (int): Number of physical grid points in the y-direction. - Nz (int): Number of physical grid points in the z-direction. + The field state is tiled. The tile shape and guard-cell depth are stored in + world and are shared by E, B, J, phi, and rho. Returns: - E (tuple): Electric field arrays (Ex, Ey, Ez) each with shape (Nx+2, Ny+2, Nz+2). - B (tuple): Magnetic field arrays (Bx, By, Bz) each with shape (Nx+2, Ny+2, Nz+2). - J (tuple): Current density arrays (Jx, Jy, Jz) each with shape (Nx+2, Ny+2, Nz+2). - phi (ndarray): Electric potential array with shape (Nx+2, Ny+2, Nz+2). - rho (ndarray): Charge density array with shape (Nx+2, Ny+2, Nz+2). + E, B, J, phi, rho initialized to zero in tiled storage. """ - # all fields include ghost cells: shape is (Nx+2, Ny+2, Nz+2) - ghost_shape = (Nx + 2, Ny + 2, Nz + 2) - Ex = jax.numpy.zeros(shape = ghost_shape) - Ey = jax.numpy.zeros(shape = ghost_shape) - Ez = jax.numpy.zeros(shape = ghost_shape) + Ex = build_tiled_array(world) + Ey = build_tiled_array(world) + Ez = build_tiled_array(world) # initialize the electric field arrays as 0 - - Bx = jax.numpy.zeros(shape = ghost_shape) - By = jax.numpy.zeros(shape = ghost_shape) - Bz = jax.numpy.zeros(shape = ghost_shape) + Bx = build_tiled_array(world) + By = build_tiled_array(world) + Bz = build_tiled_array(world) # initialize the magnetic field arrays as 0 - - Jx = jax.numpy.zeros(shape = ghost_shape) - Jy = jax.numpy.zeros(shape = ghost_shape) - Jz = jax.numpy.zeros(shape = ghost_shape) + Jx = build_tiled_array(world) + Jy = build_tiled_array(world) + Jz = build_tiled_array(world) # initialize the current density arrays as 0 - - phi = jax.numpy.zeros(shape = ghost_shape) - rho = jax.numpy.zeros(shape = ghost_shape) + phi = build_tiled_array(world) + rho = build_tiled_array(world) # initialize the electric potential and charge density arrays as 0 return (Ex, Ey, Ez), (Bx, By, Bz), (Jx, Jy, Jz), phi, rho diff --git a/PyPIC3D/particles/flat_particles.py b/PyPIC3D/particles/flat_particles.py deleted file mode 100644 index 40a694b5..00000000 --- a/PyPIC3D/particles/flat_particles.py +++ /dev/null @@ -1,396 +0,0 @@ -import jax -import jax.numpy as jnp - - -@jax.tree_util.register_pytree_node_class -class flat_particle_species: - def __init__( - self, - name, - N_particles, - charge, - mass, - weight, - T, - x1, - x2, - x3, - v1, - v2, - v3, - x_wind, - y_wind, - z_wind, - dx, - dy, - dz, - x_bc, - y_bc, - z_bc, - update_pos, - update_v, - update_x, - update_y, - update_z, - update_vx, - update_vy, - update_vz, - shape, - dt, - species_meta, - active_mask=None, - ): - self.name = name - self.N_particles = N_particles - self.charge = charge - self.mass = mass - self.weight = weight - self.T = T - self.x1 = x1 - self.x2 = x2 - self.x3 = x3 - self.v1 = v1 - self.v2 = v2 - self.v3 = v3 - self.x_wind = x_wind - self.y_wind = y_wind - self.z_wind = z_wind - self.dx = dx - self.dy = dy - self.dz = dz - self.x_bc = x_bc - self.y_bc = y_bc - self.z_bc = z_bc - self.update_pos = update_pos - self.update_v = update_v - self.update_x = update_x - self.update_y = update_y - self.update_z = update_z - self.update_vx = update_vx - self.update_vy = update_vy - self.update_vz = update_vz - self.shape = shape - self.dt = dt - self.species_meta = species_meta - if active_mask is None: - active_mask = jnp.ones_like(x1, dtype=bool) - self.active_mask = active_mask - - def get_name(self): - return self.name - - def get_charge(self): - return self.charge * self.weight - - def get_number_of_particles(self): - return self.N_particles - - def get_temperature(self): - return self.T - - def get_velocity(self): - return self.v1, self.v2, self.v3 - - def get_active_velocity(self): - active = self.active_mask.astype(self.v1.dtype) - return active * self.v1, active * self.v2, active * self.v3 - - def get_forward_position(self): - return self.x1, self.x2, self.x3 - - def get_active_mask(self): - return self.active_mask - - def get_position(self): - x1_back = self.x1 - self.v1 * self.dt / 2 - x2_back = self.x2 - self.v2 * self.dt / 2 - x3_back = self.x3 - self.v3 * self.dt / 2 - - half_x = self.x_wind / 2 - half_y = self.y_wind / 2 - half_z = self.z_wind / 2 - - x1_back = jnp.where( - x1_back > half_x, - x1_back - self.x_wind, - jnp.where(x1_back < -half_x, x1_back + self.x_wind, x1_back), - ) - x2_back = jnp.where( - x2_back > half_y, - x2_back - self.y_wind, - jnp.where(x2_back < -half_y, x2_back + self.y_wind, x2_back), - ) - x3_back = jnp.where( - x3_back > half_z, - x3_back - self.z_wind, - jnp.where(x3_back < -half_z, x3_back + self.z_wind, x3_back), - ) - - return x1_back, x2_back, x3_back - - def get_mass(self): - return self.mass * self.weight - - def get_weight(self): - return self.weight - - def get_shape(self): - return self.shape - - def momentum(self): - vmag = jnp.sqrt(self.v1**2 + self.v2**2 + self.v3**2) - return jnp.sum(self.active_mask.astype(vmag.dtype) * vmag * self.mass * self.weight) - - def set_velocity(self, v1, v2, v3): - if self.update_v: - active = self.active_mask - if self.update_vx: - self.v1 = jnp.where(active, v1, self.v1) - if self.update_vy: - self.v2 = jnp.where(active, v2, self.v2) - if self.update_vz: - self.v3 = jnp.where(active, v3, self.v3) - - def update_position(self): - if self.update_pos: - active = self.active_mask.astype(self.x1.dtype) - if self.update_x: - self.x1 = self.x1 + active * self.v1 * self.dt - if self.update_y: - self.x2 = self.x2 + active * self.v2 * self.dt - if self.update_z: - self.x3 = self.x3 + active * self.v3 * self.dt - - def boundary_conditions(self): - half_x = self.x_wind / 2 - half_y = self.y_wind / 2 - half_z = self.z_wind / 2 - - self.x1 = jnp.where( - self.x1 > half_x, - self.x1 - self.x_wind, - jnp.where(self.x1 < -half_x, self.x1 + self.x_wind, self.x1), - ) - self.x2 = jnp.where( - self.x2 > half_y, - self.x2 - self.y_wind, - jnp.where(self.x2 < -half_y, self.x2 + self.y_wind, self.x2), - ) - self.x3 = jnp.where( - self.x3 > half_z, - self.x3 - self.z_wind, - jnp.where(self.x3 < -half_z, self.x3 + self.z_wind, self.x3), - ) - - def tree_flatten(self): - children = (self.x1, self.x2, self.x3, self.v1, self.v2, self.v3, self.active_mask) - aux_data = ( - self.name, - self.N_particles, - self.charge, - self.mass, - self.weight, - self.T, - self.x_wind, - self.y_wind, - self.z_wind, - self.dx, - self.dy, - self.dz, - self.x_bc, - self.y_bc, - self.z_bc, - self.update_pos, - self.update_v, - self.update_x, - self.update_y, - self.update_z, - self.update_vx, - self.update_vy, - self.update_vz, - self.shape, - self.dt, - self.species_meta, - ) - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): - x1, x2, x3, v1, v2, v3, active_mask = children - ( - name, - N_particles, - charge, - mass, - weight, - T, - x_wind, - y_wind, - z_wind, - dx, - dy, - dz, - x_bc, - y_bc, - z_bc, - update_pos, - update_v, - update_x, - update_y, - update_z, - update_vx, - update_vy, - update_vz, - shape, - dt, - species_meta, - ) = aux_data - return cls( - name=name, - N_particles=N_particles, - charge=charge, - mass=mass, - weight=weight, - T=T, - x1=x1, - x2=x2, - x3=x3, - v1=v1, - v2=v2, - v3=v3, - x_wind=x_wind, - y_wind=y_wind, - z_wind=z_wind, - dx=dx, - dy=dy, - dz=dz, - x_bc=x_bc, - y_bc=y_bc, - z_bc=z_bc, - update_pos=update_pos, - update_v=update_v, - update_x=update_x, - update_y=update_y, - update_z=update_z, - update_vx=update_vx, - update_vy=update_vy, - update_vz=update_vz, - shape=shape, - dt=dt, - species_meta=species_meta, - active_mask=active_mask, - ) - - -def _normalize_attr(value): - try: - return jnp.asarray(value).item() - except Exception: - return value - - -def _same(attr_list): - norm = [_normalize_attr(v) for v in attr_list] - return len(set(norm)) == 1 - - -def check_flat_compat(particles): - if not particles: - return False - if not _same([p.get_shape() for p in particles]): - return False - if not _same([p.x_bc for p in particles]) or not _same([p.y_bc for p in particles]) or not _same([p.z_bc for p in particles]): - return False - if particles[0].x_bc != "periodic" or particles[0].y_bc != "periodic" or particles[0].z_bc != "periodic": - return False - if not _same([p.update_pos for p in particles]) or not _same([p.update_v for p in particles]): - return False - return True - - -def to_flat_particles(particles): - species_meta = [] - x_list, y_list, z_list = [], [], [] - vx_list, vy_list, vz_list = [], [], [] - active_list = [] - q_list, m_list, w_list, T_list = [], [], [], [] - - for species in particles: - x, y, z = species.get_forward_position() - vx, vy, vz = species.get_velocity() - active = species.get_active_mask() - x_list.append(x) - y_list.append(y) - z_list.append(z) - vx_list.append(vx) - vy_list.append(vy) - vz_list.append(vz) - active_list.append(active) - q_list.append(jnp.full_like(x, species.charge)) - m_list.append(jnp.full_like(x, species.mass)) - w_list.append(jnp.full_like(x, species.weight)) - T_list.append(jnp.full_like(x, species.T)) - species_meta.append( - { - "name": species.name, - "N_particles": float(species.N_particles), - "weight": float(species.weight), - "charge": float(species.charge), - "mass": float(species.mass), - "temperature": float(species.T), - "scaled mass": float(species.get_mass()), - "scaled charge": float(species.get_charge()), - "update_pos": species.update_pos, - "update_v": species.update_v, - } - ) - - x = jnp.concatenate(x_list, axis=0) - y = jnp.concatenate(y_list, axis=0) - z = jnp.concatenate(z_list, axis=0) - vx = jnp.concatenate(vx_list, axis=0) - vy = jnp.concatenate(vy_list, axis=0) - vz = jnp.concatenate(vz_list, axis=0) - active_mask = jnp.concatenate(active_list, axis=0) - charge = jnp.concatenate(q_list, axis=0) - mass = jnp.concatenate(m_list, axis=0) - weight = jnp.concatenate(w_list, axis=0) - T = jnp.concatenate(T_list, axis=0) - - first = particles[0] - flat = flat_particle_species( - name="flat_all", - N_particles=int(x.shape[0]), - charge=charge, - mass=mass, - weight=weight, - T=T, - x1=x, - x2=y, - x3=z, - v1=vx, - v2=vy, - v3=vz, - x_wind=first.x_wind, - y_wind=first.y_wind, - z_wind=first.z_wind, - dx=first.dx, - dy=first.dy, - dz=first.dz, - x_bc=first.x_bc, - y_bc=first.y_bc, - z_bc=first.z_bc, - update_pos=first.update_pos, - update_v=first.update_v, - update_x=first.update_x, - update_y=first.update_y, - update_z=first.update_z, - update_vx=first.update_vx, - update_vy=first.update_vy, - update_vz=first.update_vz, - shape=first.shape, - dt=first.dt, - species_meta=species_meta, - active_mask=active_mask, - ) - return [flat] diff --git a/PyPIC3D/particles/particle_class.py b/PyPIC3D/particles/particle_class.py new file mode 100644 index 00000000..e9a050c6 --- /dev/null +++ b/PyPIC3D/particles/particle_class.py @@ -0,0 +1,31 @@ +import jax +from typing import NamedTuple + + +class SpeciesConfig(NamedTuple): + """ + Per-species particle metadata shared by all slots in the species block. + """ + + charge: jax.Array # (species,) + mass: jax.Array # (species,) + weight: jax.Array # (species,) + update_x: jax.Array # (species, 3) + update_u: jax.Array # (species, 3) + + +class TiledParticles(NamedTuple): + """ + Tile-major particle storage on the shared tiled Yee mesh. + + The leading ``(ntx, nty, ntz)`` axes match the field-tile layout. Each + tile/species block has a fixed slot capacity chosen during initialization; + inactive slots remain present so the tiled arrays keep a static shape. + """ + + x: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile, 3) + # positions + u: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile, 3) + # velocities + active: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + # boolean array indicating whether the particle is active or not diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index bf79546f..00a1f668 100644 --- a/PyPIC3D/particles/particle_initialization.py +++ b/PyPIC3D/particles/particle_initialization.py @@ -1,9 +1,10 @@ import numpy as np import jax import jax.numpy as jnp +import math -from PyPIC3D.utils import vth_to_T, plasma_frequency, debye_length, T_to_vth -from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.utils import vth_to_T, T_to_vth +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles def grab_particle_keys(config): @@ -45,50 +46,6 @@ def load_initial_velocities(param, config, key, default, N_particles): return default -def compute_macroparticle_weight(config, particle_keys, simulation_parameters, world, constants): - x_wind = world["x_wind"] - y_wind = world["y_wind"] - z_wind = world["z_wind"] - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - kb = constants["kb"] - eps = constants["eps"] - - if simulation_parameters["ds_per_debye"]: - ds_per_debye = simulation_parameters["ds_per_debye"] - inverse_total_debye = 0 - - for toml_key in particle_keys: - N_particles = config[toml_key]["N_particles"] - charge = config[toml_key]["charge"] - mass = config[toml_key]["mass"] - if "temperature" in config[toml_key]: - T = config[toml_key]["temperature"] - elif "vth" in config[toml_key]: - T = vth_to_T(config[toml_key]["vth"], mass, kb) - - inverse_total_debye += ( - jnp.sqrt(N_particles / (x_wind * y_wind * z_wind) / (eps * kb * T)) - * jnp.abs(charge) - ) - - ds2 = 0 - for d in [dx, dy, dz]: - if d != 1: - ds2 += d**2 - - if ds2 == 0: - raise ValueError( - "Invalid configuration for 'ds_per_debye': at least one of dx, dy, dz must differ from 1." - ) - weight = 1 / (ds2) / (ds_per_debye**2) / inverse_total_debye - else: - weight = 1.0 - - return weight - - def initial_particles( N_per_cell, N_particles, @@ -121,6 +78,144 @@ def initial_particles( return x, y, z, vx, vy, vz +def _as_int(value): + return int(jnp.asarray(value).item()) + + +def _tile_axis_count(n_cells, cells_per_tile): + return int(math.ceil(_as_int(n_cells) / _as_int(cells_per_tile))) + + +def _particle_tile_indices(x1, x2, x3, world, tile_nx, tile_ny, tile_nz): + Nx = _as_int(world["Nx"]) + Ny = _as_int(world["Ny"]) + Nz = _as_int(world["Nz"]) + + x_cell = np.floor((np.asarray(x1) + float(world["x_wind"]) / 2) / float(world["dx"])).astype(int) + y_cell = np.floor((np.asarray(x2) + float(world["y_wind"]) / 2) / float(world["dy"])).astype(int) + z_cell = np.floor((np.asarray(x3) + float(world["z_wind"]) / 2) / float(world["dz"])).astype(int) + + x_cell = np.clip(x_cell, 0, Nx - 1) + y_cell = np.clip(y_cell, 0, Ny - 1) + z_cell = np.clip(z_cell, 0, Nz - 1) + + return x_cell // tile_nx, y_cell // tile_ny, z_cell // tile_nz + + +def _prepare_species_tile_data(species_arrays, world, tile_shape): + tile_nx, tile_ny, tile_nz = tile_shape + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + n_tiles = ntx * nty * ntz + + tile_counts = np.zeros((ntx, nty, ntz, len(species_arrays)), dtype=int) + species_tile_data = [] + + for species_index, (x_array, u_array, active_mask) in enumerate(species_arrays): + x_np = np.asarray(x_array) + u_np = np.asarray(u_array) + active_np = np.asarray(active_mask, dtype=bool) + + tx, ty, tz = _particle_tile_indices( + x_np[:, 0], + x_np[:, 1], + x_np[:, 2], + world, + tile_nx, + tile_ny, + tile_nz, + ) + flat_tile = (tx * nty + ty) * ntz + tz + active_indices = np.nonzero(active_np)[0] + + flat_counts = np.bincount(flat_tile[active_indices], minlength=n_tiles) + tile_counts[:, :, :, species_index] = flat_counts.reshape((ntx, nty, ntz)) + + species_tile_data.append((x_np, u_np, active_np, tx, ty, tz, flat_tile, active_indices)) + + return tile_counts, species_tile_data + + +def _pack_species_arrays_into_tiles(species_arrays, world, tile_shape, capacity_factor): + tile_nx, tile_ny, tile_nz = tile_shape + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + n_species = len(species_arrays) + + tile_counts, species_tile_data = _prepare_species_tile_data(species_arrays, world, tile_shape) + max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 + max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) + + x_tiles = np.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3), dtype=float) + u_tiles = np.zeros_like(x_tiles) + active_tiles = np.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) + + for species_index, (x_np, u_np, active_np, tx, ty, tz, flat_tile, active_indices) in enumerate(species_tile_data): + if active_indices.size == 0: + continue + + order = np.argsort(flat_tile[active_indices], kind="stable") + particle_indices = active_indices[order] + sorted_flat_tile = flat_tile[particle_indices] + + flat_counts = tile_counts[:, :, :, species_index].reshape(-1) + tile_starts = np.cumsum(flat_counts) - flat_counts + slots = np.arange(particle_indices.size) - tile_starts[sorted_flat_tile] + + x_tiles[ + tx[particle_indices], + ty[particle_indices], + tz[particle_indices], + species_index, + slots, + :, + ] = x_np[particle_indices] + u_tiles[ + tx[particle_indices], + ty[particle_indices], + tz[particle_indices], + species_index, + slots, + :, + ] = u_np[particle_indices] + active_tiles[ + tx[particle_indices], + ty[particle_indices], + tz[particle_indices], + species_index, + slots, + ] = active_np[particle_indices] + + return x_tiles, u_tiles, active_tiles, tile_counts + + +def _plasma_frequency(N_per_cell, charge, mass, weight, world, constants): + dx, dy, dz = world["dx"], world["dy"], world["dz"] + # get spatial resolution of the simulation domain + eps = constants["eps"] + # get the permittivity of free space from the constants dictionary + n = N_per_cell * weight / (dx * dy * dz) + # define the number density + sqrt_n = jnp.sqrt(n) + # take the square root of the number density + sqrt_eps = jnp.sqrt(eps) + sqrt_mass = jnp.sqrt(mass) + return sqrt_n * jnp.abs(charge) / (sqrt_eps * sqrt_mass) / (2 * jnp.pi) + + +def _debye_length(N_per_cell, charge, temperature, weight, world, constants): + eps = constants["eps"] + kb = constants["kb"] + # get the permittivity of free space and Boltzmann constant from the constants dictionary + dx, dy, dz = world["dx"], world["dy"], world["dz"] + n = N_per_cell * weight / (dx * dy * dz) + # define the number density + + return jnp.sqrt(eps * kb * temperature / (n * charge**2)) + + def load_particles_from_toml(config, simulation_parameters, world, constants): x_wind = world["x_wind"] y_wind = world["y_wind"] @@ -134,14 +229,15 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): dt = world["dt"] kb = constants["kb"] eps = constants["eps"] + # get the simulation domain dimensions, grid sizes, spatial resolutions, time step, Boltzmann constant, and permittivity of free space from the world and constants dictionaries i = 0 - particles = [] particle_keys = grab_particle_keys(config) + species_arrays = [] + species_metadata = [] - weight = compute_macroparticle_weight( - config, particle_keys, simulation_parameters, world, constants - ) + weight = 1.0 + # start with a default weight of 1.0, which will be overwritten if specified in the config for toml_key in particle_keys: key1, key2, key3 = jax.random.key(i), jax.random.key(i + 1), jax.random.key(i + 2) @@ -151,6 +247,7 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): print(f"\nInitializing particle species: {particle_name}") charge = config[toml_key]["charge"] mass = config[toml_key]["mass"] + # get the particle species name, charge, and mass from the config if "N_particles" in config[toml_key]: N_particles = config[toml_key]["N_particles"] @@ -158,6 +255,7 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): elif "N_per_cell" in config[toml_key]: N_per_cell = config[toml_key]["N_per_cell"] N_particles = int(N_per_cell * world["Nx"] * world["Ny"] * world["Nz"]) + # get the number of particles and number of particles per cell for the simulation if "temperature" in config[toml_key]: T = config[toml_key]["temperature"] @@ -168,10 +266,12 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): else: T = 1.0 vth = T_to_vth(T, mass, kb) + # from the provided configuration, determine the temperature and thermal velocity of the particle species, using either the specified temperature or thermal velocity, or defaulting to a temperature of 1.0 if neither is provided Tx = read_value("Tx", toml_key, config, T) Ty = read_value("Ty", toml_key, config, T) Tz = read_value("Tz", toml_key, config, T) + # overwrite the temperature in each direction if specified in the configuration, otherwise use the previously determined temperature xmin = read_value("xmin", toml_key, config, -x_wind / 2) xmax = read_value("xmax", toml_key, config, x_wind / 2) @@ -227,6 +327,7 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): vx = load_initial_velocities("initial_vx", config, toml_key, vx, N_particles) vy = load_initial_velocities("initial_vy", config, toml_key, vy, N_particles) vz = load_initial_velocities("initial_vz", config, toml_key, vz, N_particles) + # overwrite the initial positions and velocities of the particles if specified in the configuration, otherwise use the previously generated values if "temperature" not in config[toml_key]: T = (mass / (3 * kb * N_particles)) * ( @@ -235,21 +336,10 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): if "weight" in config[toml_key]: weight = config[toml_key]["weight"] - elif "ds_per_debye" in config[toml_key]: - ds_per_debye = config[toml_key]["ds_per_debye"] - - ds2 = 0 - for d in [dx, dy, dz]: - if d != 1: - ds2 += d**2 - - if ds2 == 0: - raise ValueError( - "Invalid configuration for 'ds_per_debye': at least one of dx, dy, dz must differ from 1." - ) - weight = ( - x_wind * y_wind * z_wind * eps * kb * T - ) / (N_particles * charge**2 * ds_per_debye**2 * ds2) + elif "number_density" in config[toml_key]: + n = config[toml_key]["number_density"] + weight = (n / N_per_cell) * (dx * dy * dz) + # convert number density to weight based on the number of particles per cell and cell volume update_pos = read_value("update_pos", toml_key, config, True) update_v = read_value("update_v", toml_key, config, True) @@ -259,44 +349,36 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): update_x = read_value("update_x", toml_key, config, True) update_y = read_value("update_y", toml_key, config, True) update_z = read_value("update_z", toml_key, config, True) + # determine whether to update the position and velocity components of the particle species based on the configuration, defaulting to True if not specified + + x_array = jnp.stack((x, y, z), axis=-1) + u_array = jnp.stack((vx, vy, vz), axis=-1) + active_mask = jnp.ones((N_particles,), dtype=bool) + + metadata = { + "name": particle_name, + "N_particles": int(N_particles), + "N_per_cell": N_per_cell, + "charge": charge, + "mass": mass, + "temperature": T, + "thermal_velocity": vth, + "weight": weight, + "x_bc": x_bc, + "y_bc": y_bc, + "z_bc": z_bc, + "update_x": (update_pos and update_x, update_pos and update_y, update_pos and update_z), + "update_u": (update_v and update_vx, update_v and update_vy, update_v and update_vz), + } + species_arrays.append((x_array, u_array, active_mask)) + species_metadata.append(metadata) + + kinetic_energy = 0.5 * weight * mass * jnp.sum(vx**2 + vy**2 + vz**2) + # calculate the total kinetic energy of the particle species + pf = _plasma_frequency(N_per_cell, charge, mass, weight, world, constants) + dl = _debye_length(N_per_cell, charge, T, weight, world, constants) + # calculate the plasma frequency and Debye length for the particle species - particle = particle_species( - name=particle_name, - N_particles=N_particles, - charge=charge, - mass=mass, - T=T, - x1=x, - x2=y, - x3=z, - v1=vx, - v2=vy, - v3=vz, - xwind=x_wind, - ywind=y_wind, - zwind=z_wind, - dx=dx, - dy=dy, - dz=dz, - weight=weight, - x_bc=x_bc, - y_bc=y_bc, - z_bc=z_bc, - update_vx=update_vx, - update_vy=update_vy, - update_vz=update_vz, - update_x=update_x, - update_y=update_y, - update_z=update_z, - update_pos=update_pos, - update_v=update_v, - shape=simulation_parameters["shape_factor"], - dt=dt, - ) - particles.append(particle) - - pf = plasma_frequency(particle, world, constants) - dl = debye_length(particle, world, constants) print(f"Number of particles: {N_particles}") print(f"Number of particles per cell: {N_per_cell}") print(f"x, y, z boundary conditions: {x_bc}, {y_bc}, {z_bc}") @@ -304,12 +386,39 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): print(f"Mass: {mass}") print(f"Temperature: {T}") print(f"Thermal Velocity: {vth}") - print(f"Particle Kinetic Energy: {particle.kinetic_energy()}") + print(f"Particle Kinetic Energy: {kinetic_energy}") print(f"Particle Species Plasma Frequency: {pf}") print(f"Time Steps Per Plasma Period: {(1 / (dt * pf))}") print(f"Particle Species Debye Length: {dl}") print(f"Particle Weight: {weight}") - print(f"Particle Species Scaled Charge: {particle.get_charge()}") - print(f"Particle Species Scaled Mass: {particle.get_mass()}") + print(f"Particle Species Scaled Charge: {charge * weight}") + print(f"Particle Species Scaled Mass: {mass * weight}") + + tile_nx, tile_ny, tile_nz = [_as_int(width) for width in world["tile_shape"]] + # determine the number of tiles in each dimension based on the tile shape and simulation domain dimensions + + capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) + # determine the maximum number of particles per tile across all species and tiles, and apply a capacity factor to allow for some buffer + + x_tiles, u_tiles, active_tiles, _tile_counts = _pack_species_arrays_into_tiles( + species_arrays, + world, + (tile_nx, tile_ny, tile_nz), + capacity_factor, + ) + + species_config = SpeciesConfig( + charge=jnp.asarray([metadata["charge"] for metadata in species_metadata]), + mass=jnp.asarray([metadata["mass"] for metadata in species_metadata]), + weight=jnp.asarray([metadata["weight"] for metadata in species_metadata]), + update_x=jnp.asarray([metadata["update_x"] for metadata in species_metadata], dtype=bool), + update_u=jnp.asarray([metadata["update_u"] for metadata in species_metadata], dtype=bool), + ) + particles = TiledParticles( + x=jnp.asarray(x_tiles), + u=jnp.asarray(u_tiles), + active=jnp.asarray(active_tiles), + ) + species_names = tuple(metadata["name"] for metadata in species_metadata) - return particles + return particles, species_config, species_names, tuple(species_metadata) diff --git a/PyPIC3D/particles/particle_tile_communication.py b/PyPIC3D/particles/particle_tile_communication.py new file mode 100644 index 00000000..87f56093 --- /dev/null +++ b/PyPIC3D/particles/particle_tile_communication.py @@ -0,0 +1,475 @@ +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import wrap_periodic_position +from PyPIC3D.particles.particle_class import TiledParticles + + +def _particle_boundary_conditions(world): + if "particle_boundary_conditions" in world: + return world["particle_boundary_conditions"] + return {"x": 0, "y": 0, "z": 0} + + +def _apply_tiled_axis_boundary(x, u, active, wind, bc): + half_wind = 0.5 * wind + periodic = bc == 0 + reflecting = bc == 1 + absorbing = bc == 2 + + periodic_x = wrap_periodic_position(x, wind) + reflected_x = jnp.where( + x > half_wind, + 2.0 * half_wind - x, + jnp.where(x < -half_wind, -2.0 * half_wind - x, x), + ) + reflected_u = jnp.where((x >= half_wind) | (x <= -half_wind), -u, u) + + x_out = jnp.where(periodic, periodic_x, jnp.where(reflecting, reflected_x, x)) + u_out = jnp.where(reflecting, reflected_u, u) + active_out = jnp.where(absorbing, active & (x <= half_wind) & (x >= -half_wind), active) + + return x_out, u_out, active_out + + +def _particle_tile_indices(x, y, z, world, tile_shape, tile_counts): + tile_nx, tile_ny, tile_nz = tile_shape + ntx, nty, ntz = tile_counts + + x_cell = jnp.floor((x + 0.5 * world["x_wind"]) / world["dx"]).astype(int) + y_cell = jnp.floor((y + 0.5 * world["y_wind"]) / world["dy"]).astype(int) + z_cell = jnp.floor((z + 0.5 * world["z_wind"]) / world["dz"]).astype(int) + + x_cell = jnp.clip(x_cell, 0, world["Nx"] - 1) + y_cell = jnp.clip(y_cell, 0, world["Ny"] - 1) + z_cell = jnp.clip(z_cell, 0, world["Nz"] - 1) + + return ( + jnp.clip(x_cell // tile_nx, 0, ntx - 1), + jnp.clip(y_cell // tile_ny, 0, nty - 1), + jnp.clip(z_cell // tile_nz, 0, ntz - 1), + ) + + +def update_tiled_particle_positions(tiled_particles, species_config, dt): + """ + Advance tile-major particle positions without changing tile ownership. + """ + + active = tiled_particles.active.astype(tiled_particles.x.dtype) + update_x = species_config.update_x.reshape((1, 1, 1, species_config.update_x.shape[0], 1, 3)) + + dx = active * tiled_particles.u[..., 0] * dt + dy = active * tiled_particles.u[..., 1] * dt + dz = active * tiled_particles.u[..., 2] * dt + + x = tiled_particles.x + x = x.at[..., 0].set(jnp.where(tiled_particles.active & update_x[..., 0], x[..., 0] + dx, x[..., 0])) + x = x.at[..., 1].set(jnp.where(tiled_particles.active & update_x[..., 1], x[..., 1] + dy, x[..., 1])) + x = x.at[..., 2].set(jnp.where(tiled_particles.active & update_x[..., 2], x[..., 2] + dz, x[..., 2])) + + return tiled_particles._replace(x=x) + + +def _movement_offsets(count): + if count == 1: + return (0,) + return (1, 0, -1) + + +def _adjacent_tile_offset(dest_tile, source_tile, tile_count): + """ + Signed adjacent offset from the source tile to the destination tile. + + The tiled particle step assumes particles move by at most one cell, so tile + ownership can only change by one neighboring tile along any active axis. + Periodic end points are represented with the physical crossing direction: + first -> last is -1, last -> first is +1. + """ + + if tile_count == 1: + return jnp.zeros_like(dest_tile) + + offset = dest_tile - source_tile + if tile_count == 2: + return offset + + offset = jnp.where(offset == tile_count - 1, -1, offset) + offset = jnp.where(offset == -(tile_count - 1), 1, offset) + + return offset + + +def _tiles_need_repack(offset_x, offset_y, offset_z, active): + moving = active & ((offset_x != 0) | (offset_y != 0) | (offset_z != 0)) + outgoing = jnp.any(moving, axis=-1) + incoming = jnp.zeros_like(outgoing) + + for ox in _movement_offsets(active.shape[0]): + for oy in _movement_offsets(active.shape[1]): + for oz in _movement_offsets(active.shape[2]): + if ox == 0 and oy == 0 and oz == 0: + continue + stream_active = moving & (offset_x == ox) & (offset_y == oy) & (offset_z == oz) + incoming_stream = jnp.any(stream_active, axis=-1) + incoming_stream = jnp.roll(incoming_stream, shift=ox, axis=0) + incoming_stream = jnp.roll(incoming_stream, shift=oy, axis=1) + incoming_stream = jnp.roll(incoming_stream, shift=oz, axis=2) + incoming = incoming | incoming_stream + + return outgoing | incoming + + +def _compact_tile_candidates( + candidate_x, + candidate_u, + candidate_active, + n_slots, + fallback_x=None, + fallback_u=None, + fallback_active=None, + repack_mask=None, +): + """ + Repack fixed-size incoming streams into the destination tile slot axis. + """ + + leading_shape = candidate_active.shape[:-1] + n_candidates = candidate_active.shape[-1] + + flat_x = candidate_x.reshape((-1, n_candidates, 3)) + flat_u = candidate_u.reshape((-1, n_candidates, 3)) + flat_active = candidate_active.reshape((-1, n_candidates)) + if fallback_x is None: + flat_fallback_x = jnp.zeros((flat_x.shape[0], n_slots, 3), dtype=flat_x.dtype) + flat_fallback_u = jnp.zeros((flat_u.shape[0], n_slots, 3), dtype=flat_u.dtype) + flat_fallback_active = jnp.zeros((flat_active.shape[0], n_slots), dtype=flat_active.dtype) + flat_repack_mask = jnp.ones((flat_active.shape[0],), dtype=bool) + else: + flat_fallback_x = fallback_x.reshape((-1, n_slots, 3)) + flat_fallback_u = fallback_u.reshape((-1, n_slots, 3)) + flat_fallback_active = fallback_active.reshape((-1, n_slots)) + flat_repack_mask = repack_mask.reshape((-1,)) + + def compact_one(local_x_in, local_u_in, local_active_in, fallback_x_in, fallback_u_in, fallback_active_in, do_repack): + rank = jnp.cumsum(local_active_in.astype(int)) - 1 + fits = local_active_in & (rank < n_slots) + safe_rank = jnp.where(fits, rank, 0) + + local_x = jnp.zeros((n_slots, 3), dtype=local_x_in.dtype) + local_u = jnp.zeros((n_slots, 3), dtype=local_u_in.dtype) + local_active_count = jnp.zeros(n_slots, dtype=int) + + valid = fits.astype(local_x_in.dtype) + local_x = local_x.at[safe_rank].add(valid[:, None] * local_x_in) + local_u = local_u.at[safe_rank].add(valid[:, None] * local_u_in) + local_active_count = local_active_count.at[safe_rank].add(fits.astype(int)) + + compacted = (local_x, local_u, local_active_count > 0, jnp.sum(local_active_in) > n_slots) + unchanged = (fallback_x_in, fallback_u_in, fallback_active_in, jnp.asarray(False)) + + return jax.lax.cond(do_repack, lambda _: compacted, lambda _: unchanged, operand=None) + + flat_new_x, flat_new_u, flat_new_active, flat_overflow = jax.vmap(compact_one)( + flat_x, + flat_u, + flat_active, + flat_fallback_x, + flat_fallback_u, + flat_fallback_active, + flat_repack_mask, + ) + + new_x = flat_new_x.reshape(leading_shape + (n_slots, 3)) + new_u = flat_new_u.reshape(leading_shape + (n_slots, 3)) + new_active = flat_new_active.reshape(leading_shape + (n_slots,)) + overflow = jnp.any(flat_overflow) + + return new_x, new_u, new_active, overflow + + +def _bounded_state_and_tile_offsets(tiled_particles, world, tile_shape): + """ + Apply physical particle boundaries and identify the adjacent tile offset. + """ + + ntx, nty, ntz, n_species, n_slots = tiled_particles.active.shape + tile_counts = (ntx, nty, ntz) + + particle_bc = _particle_boundary_conditions(world) + bounded_x = tiled_particles.x + bounded_u = tiled_particles.u + bounded_active = tiled_particles.active + + x1, u1, bounded_active = _apply_tiled_axis_boundary( + bounded_x[..., 0], + bounded_u[..., 0], + bounded_active, + world["x_wind"], + particle_bc["x"], + ) + x2, u2, bounded_active = _apply_tiled_axis_boundary( + bounded_x[..., 1], + bounded_u[..., 1], + bounded_active, + world["y_wind"], + particle_bc["y"], + ) + x3, u3, bounded_active = _apply_tiled_axis_boundary( + bounded_x[..., 2], + bounded_u[..., 2], + bounded_active, + world["z_wind"], + particle_bc["z"], + ) + + bounded_x = bounded_x.at[..., 0].set(x1) + bounded_x = bounded_x.at[..., 1].set(x2) + bounded_x = bounded_x.at[..., 2].set(x3) + bounded_u = bounded_u.at[..., 0].set(u1) + bounded_u = bounded_u.at[..., 1].set(u2) + bounded_u = bounded_u.at[..., 2].set(u3) + + dest_tx, dest_ty, dest_tz = _particle_tile_indices( + bounded_x[..., 0], + bounded_x[..., 1], + bounded_x[..., 2], + world, + tile_shape, + tile_counts, + ) + + tx = jnp.arange(ntx).reshape((ntx, 1, 1, 1, 1)) + ty = jnp.arange(nty).reshape((1, nty, 1, 1, 1)) + tz = jnp.arange(ntz).reshape((1, 1, ntz, 1, 1)) + + offset_x = _adjacent_tile_offset(dest_tx, tx, ntx) + offset_y = _adjacent_tile_offset(dest_ty, ty, nty) + offset_z = _adjacent_tile_offset(dest_tz, tz, ntz) + + return bounded_x, bounded_u, bounded_active, offset_x, offset_y, offset_z + + +def _refresh_tiled_particle_tiles_compacting(tiled_particles, world, tile_shape): + """ + Move active particles into owning tiles by compacting all neighbor streams. + """ + + bounded_x, bounded_u, bounded_active, offset_x, offset_y, offset_z = _bounded_state_and_tile_offsets( + tiled_particles, + world, + tile_shape, + ) + + ntx, nty, ntz, n_species, n_slots = bounded_active.shape + repack_mask = _tiles_need_repack(offset_x, offset_y, offset_z, bounded_active) + + x_offsets = _movement_offsets(ntx) + y_offsets = _movement_offsets(nty) + z_offsets = _movement_offsets(ntz) + + candidate_x = [] + candidate_u = [] + candidate_active = [] + + for ox in x_offsets: + for oy in y_offsets: + for oz in z_offsets: + stream_active = ( + bounded_active + & (offset_x == ox) + & (offset_y == oy) + & (offset_z == oz) + ) + stream_x = jnp.where(stream_active[..., None], bounded_x, 0.0) + stream_u = jnp.where(stream_active[..., None], bounded_u, 0.0) + + stream_x = jnp.roll(stream_x, shift=ox, axis=0) + stream_x = jnp.roll(stream_x, shift=oy, axis=1) + stream_x = jnp.roll(stream_x, shift=oz, axis=2) + stream_u = jnp.roll(stream_u, shift=ox, axis=0) + stream_u = jnp.roll(stream_u, shift=oy, axis=1) + stream_u = jnp.roll(stream_u, shift=oz, axis=2) + stream_active = jnp.roll(stream_active, shift=ox, axis=0) + stream_active = jnp.roll(stream_active, shift=oy, axis=1) + stream_active = jnp.roll(stream_active, shift=oz, axis=2) + + candidate_x.append(stream_x) + candidate_u.append(stream_u) + candidate_active.append(stream_active) + + candidate_x = jnp.concatenate(candidate_x, axis=-2) + candidate_u = jnp.concatenate(candidate_u, axis=-2) + candidate_active = jnp.concatenate(candidate_active, axis=-1) + + new_x, new_u, new_active, overflow = _compact_tile_candidates( + candidate_x, + candidate_u, + candidate_active, + n_slots, + fallback_x=bounded_x, + fallback_u=bounded_u, + fallback_active=bounded_active, + repack_mask=repack_mask, + ) + + refreshed = TiledParticles( + x=new_x, + u=new_u, + active=new_active, + ) + + return refreshed, overflow + + +def _fill_incoming_particles(stay_x, stay_u, stay_active, incoming_x, incoming_u, incoming_active): + """ + Fill inactive destination slots with incoming neighbor particles. + + The slot layout remains fixed. Tiles without incoming particles keep their + stay-particle slots untouched; tiles with incoming particles use the first + available inactive slots and report overflow when the incoming stream is + larger than the local free capacity. + """ + + leading_shape = stay_active.shape[:-1] + n_slots = stay_active.shape[-1] + n_candidates = incoming_active.shape[-1] + + flat_stay_x = stay_x.reshape((-1, n_slots, 3)) + flat_stay_u = stay_u.reshape((-1, n_slots, 3)) + flat_stay_active = stay_active.reshape((-1, n_slots)) + flat_incoming_x = incoming_x.reshape((-1, n_candidates, 3)) + flat_incoming_u = incoming_u.reshape((-1, n_candidates, 3)) + flat_incoming_active = incoming_active.reshape((-1, n_candidates)) + + slot_ids = jnp.arange(n_slots) + + def fill_one(local_x, local_u, local_active, incoming_x_in, incoming_u_in, incoming_active_in): + free = ~local_active + free_rank = jnp.cumsum(free.astype(int)) - 1 + safe_free_rank = jnp.where(free, free_rank, 0) + slot_for_rank = jnp.zeros(n_slots, dtype=slot_ids.dtype) + slot_for_rank = slot_for_rank.at[safe_free_rank].add(jnp.where(free, slot_ids, 0)) + + incoming_rank = jnp.cumsum(incoming_active_in.astype(int)) - 1 + n_free = jnp.sum(free.astype(int)) + fits = incoming_active_in & (incoming_rank < n_free) + overflow = jnp.any(incoming_active_in & (incoming_rank >= n_free)) + + safe_rank = jnp.where(fits, incoming_rank, 0) + selected_slots = slot_for_rank[safe_rank] + valid = fits.astype(local_x.dtype) + + incoming_count = jnp.zeros(n_slots, dtype=int) + local_x = local_x.at[selected_slots].add(valid[:, None] * incoming_x_in) + local_u = local_u.at[selected_slots].add(valid[:, None] * incoming_u_in) + incoming_count = incoming_count.at[selected_slots].add(fits.astype(int)) + local_active = local_active | (incoming_count > 0) + + return local_x, local_u, local_active, overflow + + flat_x, flat_u, flat_active, flat_overflow = jax.vmap(fill_one)( + flat_stay_x, + flat_stay_u, + flat_stay_active, + flat_incoming_x, + flat_incoming_u, + flat_incoming_active, + ) + + new_x = flat_x.reshape(leading_shape + (n_slots, 3)) + new_u = flat_u.reshape(leading_shape + (n_slots, 3)) + new_active = flat_active.reshape(leading_shape + (n_slots,)) + overflow = jnp.any(flat_overflow) + + return new_x, new_u, new_active, overflow + + +def _refresh_tiled_particle_tiles_sparse(tiled_particles, world, tile_shape): + """ + Move active particles into owning tiles using neighbor-only incoming streams. + """ + + bounded_x, bounded_u, bounded_active, offset_x, offset_y, offset_z = _bounded_state_and_tile_offsets( + tiled_particles, + world, + tile_shape, + ) + + ntx, nty, ntz, n_species, n_slots = bounded_active.shape + moving = bounded_active & ((offset_x != 0) | (offset_y != 0) | (offset_z != 0)) + stay_active = bounded_active & ~moving + stay_x = jnp.where(stay_active[..., None], bounded_x, 0.0) + stay_u = jnp.where(stay_active[..., None], bounded_u, 0.0) + + incoming_x = [] + incoming_u = [] + incoming_active = [] + + for ox in _movement_offsets(ntx): + for oy in _movement_offsets(nty): + for oz in _movement_offsets(ntz): + if ox == 0 and oy == 0 and oz == 0: + continue + + stream_active = ( + moving + & (offset_x == ox) + & (offset_y == oy) + & (offset_z == oz) + ) + stream_x = jnp.where(stream_active[..., None], bounded_x, 0.0) + stream_u = jnp.where(stream_active[..., None], bounded_u, 0.0) + + stream_x = jnp.roll(stream_x, shift=ox, axis=0) + stream_x = jnp.roll(stream_x, shift=oy, axis=1) + stream_x = jnp.roll(stream_x, shift=oz, axis=2) + stream_u = jnp.roll(stream_u, shift=ox, axis=0) + stream_u = jnp.roll(stream_u, shift=oy, axis=1) + stream_u = jnp.roll(stream_u, shift=oz, axis=2) + stream_active = jnp.roll(stream_active, shift=ox, axis=0) + stream_active = jnp.roll(stream_active, shift=oy, axis=1) + stream_active = jnp.roll(stream_active, shift=oz, axis=2) + + incoming_x.append(stream_x) + incoming_u.append(stream_u) + incoming_active.append(stream_active) + + if len(incoming_active) == 0: + overflow = jnp.asarray(False) + return TiledParticles(x=stay_x, u=stay_u, active=stay_active), overflow + + incoming_x = jnp.concatenate(incoming_x, axis=-2) + incoming_u = jnp.concatenate(incoming_u, axis=-2) + incoming_active = jnp.concatenate(incoming_active, axis=-1) + + new_x, new_u, new_active, overflow = _fill_incoming_particles( + stay_x, + stay_u, + stay_active, + incoming_x, + incoming_u, + incoming_active, + ) + + refreshed = TiledParticles( + x=new_x, + u=new_u, + active=new_active, + ) + + return refreshed, overflow + + +def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): + """ + Move active particles into their owning tiles while preserving static shape. + + The refresh assumes particles move by at most one cell in a timestep, so each + particle either stays in its current tile or moves to an adjacent tile. + Particles that do not fit in the destination tile capacity are dropped from + the returned active mask and reported through the overflow flag. + """ + + return _refresh_tiled_particle_tiles_sparse(tiled_particles, world, tile_shape) diff --git a/PyPIC3D/particles/species_class.py b/PyPIC3D/particles/species_class.py deleted file mode 100644 index c2d60aa7..00000000 --- a/PyPIC3D/particles/species_class.py +++ /dev/null @@ -1,363 +0,0 @@ -import jax -from jax import jit -from functools import partial -import jax.numpy as jnp -from jax.tree_util import register_pytree_node_class - - -@jit -def compute_index(x, dx, window): - """Compute grid-cell index for a position.""" - scaled_x = x + window / 2 - return jnp.floor(scaled_x / dx).astype(int) - - -@partial(jit, static_argnames=("periodic", "reflecting")) -def apply_axis_boundary_condition(x, v, wind, half_wind, periodic, reflecting): - """Apply boundary conditions to particle positions/velocities along one axis.""" - - def periodic_bc(state): - x_in, v_in = state - x_out = x_in + wind * (x_in < -half_wind) - wind * (x_in > half_wind) - return x_out, v_in - - def reflecting_bc(state): - x_in, v_in = state - v_out = jnp.where((x_in >= half_wind) | (x_in <= -half_wind), -v_in, v_in) - return x_in, v_out - - def identity_bc(state): - return state - - return jax.lax.cond( - periodic, - periodic_bc, - lambda state: jax.lax.cond(reflecting, reflecting_bc, identity_bc, state), - (x, v), - ) - - -@register_pytree_node_class -class particle_species: - """A particle species (positions/velocities + metadata) stored as a JAX pytree.""" - - def __init__( - self, - name, - N_particles, - charge, - mass, - T, - v1, - v2, - v3, - x1, - x2, - x3, - xwind, - ywind, - zwind, - dx, - dy, - dz, - weight=1, - x_bc="periodic", - y_bc="periodic", - z_bc="periodic", - update_x=True, - update_y=True, - update_z=True, - update_vx=True, - update_vy=True, - update_vz=True, - update_pos=True, - update_v=True, - shape=1, - dt=0, - active_mask=None, - ): - self.name = name - self.N_particles = N_particles - self.charge = charge - self.mass = mass - self.weight = weight - self.T = T - self.v1 = v1 - self.v2 = v2 - self.v3 = v3 - self.dx = dx - self.dy = dy - self.dz = dz - self.x_wind = xwind - self.y_wind = ywind - self.z_wind = zwind - self.half_x_wind = 0.5 * xwind - self.half_y_wind = 0.5 * ywind - self.half_z_wind = 0.5 * zwind - self.x_bc = x_bc - self.y_bc = y_bc - self.z_bc = z_bc - self.x_periodic = x_bc == "periodic" - self.x_reflecting = x_bc == "reflecting" - self.y_periodic = y_bc == "periodic" - self.y_reflecting = y_bc == "reflecting" - self.z_periodic = z_bc == "periodic" - self.z_reflecting = z_bc == "reflecting" - self.x_absorbing = x_bc == "absorbing" - self.y_absorbing = y_bc == "absorbing" - self.z_absorbing = z_bc == "absorbing" - self.update_x = update_x - self.update_y = update_y - self.update_z = update_z - self.update_vx = update_vx - self.update_vy = update_vy - self.update_vz = update_vz - self.update_pos = update_pos - self.update_v = update_v - self.shape = shape - self.dt = dt - - self.x1 = x1 - self.x2 = x2 - self.x3 = x3 - if active_mask is None: - active_mask = jnp.ones_like(x1, dtype=bool) - self.active_mask = active_mask - - def get_name(self): - return self.name - - def get_charge(self): - return self.charge * self.weight - - def get_number_of_particles(self): - return self.N_particles - - def get_temperature(self): - return self.T - - def get_velocity(self): - return self.v1, self.v2, self.v3 - - def get_active_velocity(self): - active = self.active_mask.astype(self.v1.dtype) - return active * self.v1, active * self.v2, active * self.v3 - - def get_forward_position(self): - return self.x1, self.x2, self.x3 - - def get_active_mask(self): - return self.active_mask - - def get_position(self): - x1_back = self.x1 - self.v1 * self.dt / 2 - x2_back = self.x2 - self.v2 * self.dt / 2 - x3_back = self.x3 - self.v3 * self.dt / 2 - - if self.x_bc == "periodic": - x1_back = jnp.where( - x1_back > self.x_wind / 2, - x1_back - self.x_wind, - jnp.where(x1_back < -self.x_wind / 2, x1_back + self.x_wind, x1_back), - ) - - if self.y_bc == "periodic": - x2_back = jnp.where( - x2_back > self.y_wind / 2, - x2_back - self.y_wind, - jnp.where(x2_back < -self.y_wind / 2, x2_back + self.y_wind, x2_back), - ) - - if self.z_bc == "periodic": - x3_back = jnp.where( - x3_back > self.z_wind / 2, - x3_back - self.z_wind, - jnp.where(x3_back < -self.z_wind / 2, x3_back + self.z_wind, x3_back), - ) - - return x1_back, x2_back, x3_back - - def get_mass(self): - return self.mass * self.weight - - def get_weight(self): - return self.weight - - def get_resolution(self): - return self.dx, self.dy, self.dz - - def get_shape(self): - return self.shape - - def get_index(self): - return ( - compute_index(self.x1, self.dx, self.x_wind), - compute_index(self.x2, self.dy, self.y_wind), - compute_index(self.x3, self.dz, self.z_wind), - ) - - def set_velocity(self, v1, v2, v3): - if self.update_v: - active = self.active_mask - if self.update_vx: - self.v1 = jnp.where(active, v1, self.v1) - if self.update_vy: - self.v2 = jnp.where(active, v2, self.v2) - if self.update_vz: - self.v3 = jnp.where(active, v3, self.v3) - - def set_position(self, x1, x2, x3): - self.x1 = x1 - self.x2 = x2 - self.x3 = x3 - - def set_mass(self, mass): - self.mass = mass - - def set_weight(self, weight): - self.weight = weight - - def kinetic_energy(self): - v2 = jnp.square(self.v1) + jnp.square(self.v2) + jnp.square(self.v3) - return 0.5 * self.weight * self.mass * jnp.sum(self.active_mask.astype(v2.dtype) * v2) - - def momentum(self): - vmag = jnp.sqrt(self.v1**2 + self.v2**2 + self.v3**2) - return self.mass * self.weight * jnp.sum(self.active_mask.astype(vmag.dtype) * vmag) - - def boundary_conditions(self): - x1, x2, x3 = self.x1, self.x2, self.x3 - v1, v2, v3 = self.v1, self.v2, self.v3 - - x1, v1 = apply_axis_boundary_condition( - x1, v1, self.x_wind, self.half_x_wind, self.x_periodic, self.x_reflecting - ) - x2, v2 = apply_axis_boundary_condition( - x2, v2, self.y_wind, self.half_y_wind, self.y_periodic, self.y_reflecting - ) - x3, v3 = apply_axis_boundary_condition( - x3, v3, self.z_wind, self.half_z_wind, self.z_periodic, self.z_reflecting - ) - - active_mask = self.active_mask - if self.x_absorbing: - active_mask = active_mask & (x1 <= self.half_x_wind) & (x1 >= -self.half_x_wind) - if self.y_absorbing: - active_mask = active_mask & (x2 <= self.half_y_wind) & (x2 >= -self.half_y_wind) - if self.z_absorbing: - active_mask = active_mask & (x3 <= self.half_z_wind) & (x3 >= -self.half_z_wind) - - self.x1, self.x2, self.x3 = x1, x2, x3 - self.v1, self.v2, self.v3 = v1, v2, v3 - self.active_mask = active_mask - - def update_position(self): - if self.update_pos: - active = self.active_mask.astype(self.x1.dtype) - if self.update_x: - self.x1 = self.x1 + active * self.v1 * self.dt - if self.update_y: - self.x2 = self.x2 + active * self.v2 * self.dt - if self.update_z: - self.x3 = self.x3 + active * self.v3 * self.dt - - def tree_flatten(self): - children = (self.v1, self.v2, self.v3, self.x1, self.x2, self.x3, self.active_mask) - - aux_data = ( - self.name, - self.N_particles, - self.charge, - self.mass, - self.T, - self.x_wind, - self.y_wind, - self.z_wind, - self.dx, - self.dy, - self.dz, - self.weight, - self.x_bc, - self.y_bc, - self.z_bc, - self.update_pos, - self.update_v, - self.update_x, - self.update_y, - self.update_z, - self.update_vx, - self.update_vy, - self.update_vz, - self.shape, - self.dt, - ) - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): - v1, v2, v3, x1, x2, x3, active_mask = children - - ( - name, - N_particles, - charge, - mass, - T, - x_wind, - y_wind, - z_wind, - dx, - dy, - dz, - weight, - x_bc, - y_bc, - z_bc, - update_pos, - update_v, - update_x, - update_y, - update_z, - update_vx, - update_vy, - update_vz, - shape, - dt, - ) = aux_data - - obj = cls( - name=name, - N_particles=N_particles, - charge=charge, - mass=mass, - T=T, - x1=x1, - x2=x2, - x3=x3, - v1=v1, - v2=v2, - v3=v3, - xwind=x_wind, - ywind=y_wind, - zwind=z_wind, - dx=dx, - dy=dy, - dz=dz, - weight=weight, - x_bc=x_bc, - y_bc=y_bc, - z_bc=z_bc, - update_x=update_x, - update_y=update_y, - update_z=update_z, - update_vx=update_vx, - update_vy=update_vy, - update_vz=update_vz, - update_pos=update_pos, - update_v=update_v, - shape=shape, - dt=dt, - active_mask=active_mask, - ) - - return obj diff --git a/PyPIC3D/pusher/__init__.py b/PyPIC3D/pusher/__init__.py index 6e579869..e18bee36 100644 --- a/PyPIC3D/pusher/__init__.py +++ b/PyPIC3D/pusher/__init__.py @@ -10,7 +10,5 @@ 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 index f06f4274..a3cd0fd9 100644 --- a/PyPIC3D/pusher/boris.py +++ b/PyPIC3D/pusher/boris.py @@ -6,7 +6,7 @@ from PyPIC3D.boundary_conditions.grid_and_stencil import ( BC_PERIODIC, axis_has_active_cells, - collapse_axis_stencil, + inactive_axis_index, prepare_particle_axis_stencil, ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights @@ -124,8 +124,18 @@ def relativistic_boris_single_particle(vx, vy, vz, efield_atx, efield_aty, efiel 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): +@partial(jit, static_argnames=("ghost_cells", "active_axes", "inactive_axis_indices")) +def interpolate_field_to_particles( + field, + x, + y, + z, + grid, + shape_factor, + ghost_cells=False, + active_axes=None, + inactive_axis_indices=None, +): """ Interpolate a Yee-grid field component to particle positions using PIC shape functions. @@ -143,9 +153,12 @@ def interpolate_field_to_particles(field, x, y, z, grid, shape_factor, ghost_cel 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) + if active_axes is None: + 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) + else: + x_active, y_active, z_active = active_axes _, _, deltax, xpts = prepare_particle_axis_stencil( x, x_grid, Nx, shape_factor, BC_PERIODIC, ghost_cells=ghost_cells @@ -176,25 +189,49 @@ def interpolate_field_to_particles(field, x, y, z, grid, shape_factor, ghost_cel ypts = jnp.asarray(ypts) zpts = jnp.asarray(zpts) + def collapse_axis(points, weights, axis_size, axis_active, inactive_index): + if axis_active: + return points, weights + + if inactive_index is None: + inactive_index = inactive_axis_index(axis_size, ghost_cells=ghost_cells) + + collapsed_points = jnp.full( + (1, points.shape[1]), + int(inactive_index), + dtype=points.dtype, + ) + collapsed_weights = jnp.sum(weights, axis=0, keepdims=True) + return collapsed_points, collapsed_weights + + if inactive_axis_indices is None: + inactive_axis_indices = (None, None, None) + # 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) + xpts_eff, x_weights_eff = collapse_axis( + xpts, x_weights, Nx, x_active, inactive_axis_indices[0] + ) 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) + ypts_eff, y_weights_eff = collapse_axis( + ypts, y_weights, Ny, y_active, inactive_axis_indices[1] + ) 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) + zpts_eff, z_weights_eff = collapse_axis( + zpts, z_weights, Nz, z_active, inactive_axis_indices[2] + ) def stencil_contribution(stencil_idx): i, j, k = stencil_idx diff --git a/PyPIC3D/pusher/particle_push.py b/PyPIC3D/pusher/particle_push.py index fad6e554..27174bb5 100644 --- a/PyPIC3D/pusher/particle_push.py +++ b/PyPIC3D/pusher/particle_push.py @@ -1,8 +1,7 @@ import jax -from jax import jit import jax.numpy as jnp -from functools import partial +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.pusher.boris import ( boris_single_particle, interpolate_field_to_particles, @@ -11,87 +10,86 @@ from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle -VALID_PARTICLE_PUSHERS = ("boris", "higuera_cary") - - -def validate_particle_pusher(particle_pusher): +def particle_push(particles, species_config, E_tiles, B_tiles, world, constants, relativistic=True, particle_pusher="boris"): """ - Validate the particle pusher name before entering jitted simulation code. + Push tile-major particles with the selected pusher using compact tiled Yee fields. + + Particles are assumed to live in the tile that owns their current forward + position. The configured field halos on each compact tile provide the + neighboring Yee data needed by the interpolation stencil near tile faces. """ - if particle_pusher not in VALID_PARTICLE_PUSHERS: - raise ValueError(f"particle_pusher must be one of {VALID_PARTICLE_PUSHERS}, got {particle_pusher!r}") + tile_shape = tuple(int(width) for width in world["tile_shape"]) + tile_nx, tile_ny, tile_nz = tile_shape + g = int(world["guard_cells"]) + dt = world["dt"] + shape_factor = world["shape_factor"] -@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), + tiled_center_grid = world["grids"]["tiled_center_grid"] + tiled_vertex_grid = world["grids"]["tiled_vertex_grid"] + + Ex_tiles, Ey_tiles, Ez_tiles = E_tiles + Bx_tiles, By_tiles, Bz_tiles = B_tiles + + ntx, nty, ntz = particles.x.shape[:3] + active_axes = ( + int(ntx) * int(tile_nx) > 1, + int(nty) * int(tile_ny) > 1, + int(ntz) * int(tile_nz) > 1, + ) + inactive_axis_indices = (g, g, g) + + def push_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, charge_species, mass_species, update_u_species, + Ex_tile, Ey_tile, Ez_tile, Bx_tile, By_tile, Bz_tile): + x = x_tile[..., 0].reshape(-1) + y = x_tile[..., 1].reshape(-1) + z = x_tile[..., 2].reshape(-1) + vx = u_tile[..., 0].reshape(-1) + vy = u_tile[..., 1].reshape(-1) + vz = u_tile[..., 2].reshape(-1) + q = jnp.broadcast_to(charge_species[:, jnp.newaxis], active_tile.shape).reshape(-1) + m = jnp.broadcast_to(mass_species[:, jnp.newaxis], active_tile.shape).reshape(-1) + + center_x = tiled_center_grid[0][tx, ty, tz] + center_y = tiled_center_grid[1][tx, ty, tz] + center_z = tiled_center_grid[2][tx, ty, tz] + vertex_x = tiled_vertex_grid[0][tx, ty, tz] + vertex_y = tiled_vertex_grid[1][tx, ty, tz] + vertex_z = tiled_vertex_grid[2][tx, ty, tz] + + Ex_grid = vertex_x, center_y, center_z + Ey_grid = center_x, vertex_y, center_z + Ez_grid = center_x, center_y, vertex_z + Bx_grid = center_x, vertex_y, vertex_z + By_grid = vertex_x, center_y, vertex_z + Bz_grid = vertex_x, vertex_y, center_z + + efield_atx = interpolate_field_to_particles( + Ex_tile, x, y, z, Ex_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices ) - relativistic_boris_vmap = jax.vmap( - relativistic_boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), + efield_aty = interpolate_field_to_particles( + Ey_tile, x, y, z, Ey_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices ) - higuera_cary_vmap = jax.vmap( - higuera_cary_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None), + efield_atz = interpolate_field_to_particles( + Ez_tile, x, y, z, Ez_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices + ) + + bfield_atx = interpolate_field_to_particles( + Bx_tile, x, y, z, Bx_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices + ) + bfield_aty = interpolate_field_to_particles( + By_tile, x, y, z, By_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices ) - else: + bfield_atz = interpolate_field_to_particles( + Bz_tile, x, y, z, Bz_grid, shape_factor, ghost_cells=True, + active_axes=active_axes, inactive_axis_indices=inactive_axis_indices + ) + boris_vmap = jax.vmap( boris_single_particle, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), @@ -104,23 +102,73 @@ def particle_push(particles, E, B, grid, staggered_grid, dt, constants, periodic 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 + + if particle_pusher == "boris": + new_vx, new_vy, new_vz = jax.lax.cond( + relativistic, + 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, + ) + elif particle_pusher == "higuera_cary": + new_vx, new_vy, new_vz = higuera_cary_vmap( + vx, vy, vz, + efield_atx, efield_aty, efield_atz, + bfield_atx, bfield_aty, bfield_atz, + q, m, dt, constants, + ) + else: + raise ValueError(f"Unknown particle_pusher: {particle_pusher}") + + active = active_tile.reshape(-1) + update_u1 = jnp.broadcast_to(update_u_species[:, 0, jnp.newaxis], active_tile.shape).reshape(-1) + update_u2 = jnp.broadcast_to(update_u_species[:, 1, jnp.newaxis], active_tile.shape).reshape(-1) + update_u3 = jnp.broadcast_to(update_u_species[:, 2, jnp.newaxis], active_tile.shape).reshape(-1) + + new_u = u_tile.reshape(-1, 3) + new_u = new_u.at[:, 0].set(jnp.where(active & update_u1, new_vx, vx)) + new_u = new_u.at[:, 1].set(jnp.where(active & update_u2, new_vy, vy)) + new_u = new_u.at[:, 2].set(jnp.where(active & update_u3, new_vz, vz)) + + return new_u.reshape(u_tile.shape) + + push_tiles = push_one_tile + tx = jnp.arange(ntx) + ty = jnp.arange(nty) + tz = jnp.arange(ntz) + + push_tiles = jax.vmap(push_tiles, in_axes=(None, None, 0, 0, 0, 0, None, None, None, 0, 0, 0, 0, 0, 0), out_axes=0) + push_tiles = jax.vmap(push_tiles, in_axes=(None, 0, None, 0, 0, 0, None, None, None, 0, 0, 0, 0, 0, 0), out_axes=0) + push_tiles = jax.vmap(push_tiles, in_axes=(0, None, None, 0, 0, 0, None, None, None, 0, 0, 0, 0, 0, 0), out_axes=0) + + new_u = push_tiles( + tx, ty, tz, + particles.x, + particles.u, + particles.active, + species_config.charge, + species_config.mass, + species_config.update_u, + Ex_tiles, + Ey_tiles, + Ez_tiles, + Bx_tiles, + By_tiles, + Bz_tiles, + ) + + return TiledParticles( + x=particles.x, + u=new_u, + active=particles.active, + ) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 9e3d8434..e505d499 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -1,118 +1,99 @@ import jax.numpy as jnp -from jax import jit from jax import lax -from functools import partial from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient -from PyPIC3D.solvers.pstd import spectral_gradient -from PyPIC3D.utils import digital_filter -from PyPIC3D.boundary_conditions.boundaryconditions import ( - update_ghost_cells, apply_scalar_conducting_bc -) +from PyPIC3D.utilities.filters import digital_filter +from PyPIC3D.boundary_conditions import ghost_cells -@jit -def solve_poisson_with_fft(rho, constants, world): - """ - Solve Poisson's equation in a periodic domain using FFTs. +def _active_slice(g): + return slice(g, -g) - Operates on the interior of the ghost-celled rho array and writes the - result back into the interior of phi. - Args: - rho (ndarray): Charge density field with shape (Nx+2, Ny+2, Nz+2). - constants (dict): Physical constants with key ``eps``. - world (dict): Simulation metadata with ``dx``, ``dy``, and ``dz``. +def _forward_slice(g): + return slice(g + 1, None if g == 1 else -g + 1) - Returns: - ndarray: Electrostatic potential field with shape (Nx+2, Ny+2, Nz+2). - """ - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - eps = constants['eps'] - rho_interior = rho[1:-1, 1:-1, 1:-1] - # extract the physical interior for the FFT solve +def _backward_slice(g): + return slice(g - 1, -g - 1) - rho_hat = jnp.fft.fftn(rho_interior) - nx, ny, nz = rho_hat.shape - kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi - ky = jnp.fft.fftfreq(ny, d=dy) * 2 * jnp.pi - kz = jnp.fft.fftfreq(nz, d=dz) * 2 * jnp.pi - kx, ky, kz = jnp.meshgrid(kx, ky, kz, indexing='ij') +def _as_single_tile(field): + return field[jnp.newaxis, jnp.newaxis, jnp.newaxis, :, :, :] - k_squared = kx**2 + ky**2 + kz**2 - k_squared = k_squared.at[0, 0, 0].set(1.0) - phi_hat = rho_hat / (eps * k_squared) - phi_hat = phi_hat.at[0, 0, 0].set(0.0) - phi_interior = jnp.fft.ifftn(phi_hat).real +def _refresh_single_tile_scalar(field, world, g, tile_shape, apply_conducting=False): + field_tiles = _as_single_tile(field) + if apply_conducting: + field_tiles = ghost_cells.apply_tiled_scalar_conducting_bc(field_tiles, world, num_guard_cells=g) + field_tiles = ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + return field_tiles[0, 0, 0] - phi = jnp.zeros_like(rho) - phi = phi.at[1:-1, 1:-1, 1:-1].set(phi_interior) - # write the solution back into the interior of the ghost-celled array - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) - # fill ghost cells for the potential +def _centered_finite_difference_gradient(field, dx, dy, dz): + """ + Compute the centered finite-difference gradient on a periodic scalar field. - return phi + This is the electrostatic single-tile post-processing stencil. The tiled + electrostatic path below uses the same centered difference directly on + compact tile arrays with refreshed halos. + """ + + grad_x = (jnp.roll(field, shift=-1, axis=0) - jnp.roll(field, shift=1, axis=0)) / (2.0 * dx) + grad_y = (jnp.roll(field, shift=-1, axis=1) - jnp.roll(field, shift=1, axis=1)) / (2.0 * dy) + grad_z = (jnp.roll(field, shift=-1, axis=2) - jnp.roll(field, shift=1, axis=2)) / (2.0 * dz) + + return grad_x, grad_y, grad_z -@partial(jit, static_argnames=("tol", "max_iter")) -def solve_poisson_with_conjugate_gradient(rho, phi, constants, world, tol=1e-6, max_iter=5000): +def solve_poisson_with_conjugate_gradient(rho, phi, constants, world, tol=1e-12, max_iter=5000): """ Solve Poisson's equation using matrix-free conjugate gradient. Uses ghost-cell slicing for the Laplacian stencil instead of jnp.roll. Args: - rho (ndarray): Charge density field with shape (Nx+2, Ny+2, Nz+2). - phi (ndarray): Initial guess for potential with shape (Nx+2, Ny+2, Nz+2). + rho (ndarray): Charge density field with shape (Nx+2*g, Ny+2*g, Nz+2*g). + phi (ndarray): Initial guess for potential with shape (Nx+2*g, Ny+2*g, Nz+2*g). constants (dict): Physical constants with key ``eps``. world (dict): Simulation metadata with ``dx``, ``dy``, and ``dz``. tol (float): Residual tolerance. max_iter (int): Maximum number of CG iterations. Returns: - ndarray: Electrostatic potential field with shape (Nx+2, Ny+2, Nz+2). + ndarray: Electrostatic potential field with shape (Nx+2*g, Ny+2*g, Nz+2*g). """ dx = world['dx'] dy = world['dy'] dz = world['dz'] eps = constants['eps'] - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the boundary condition codes for each axis + g = int(world["guard_cells"]) + tile_shape = tuple(int(width) for width in world["tile_shape"]) + active = _active_slice(g) + forward = _forward_slice(g) + backward = _backward_slice(g) + # get the single-tile layout from the world contract def lapl(field): - # Laplacian using ghost-cell neighbors instead of jnp.roll - dfdx2 = (field[2:, 1:-1, 1:-1] + field[:-2, 1:-1, 1:-1] - 2.0 * field[1:-1, 1:-1, 1:-1]) / (dx * dx) - dfdy2 = (field[1:-1, 2:, 1:-1] + field[1:-1, :-2, 1:-1] - 2.0 * field[1:-1, 1:-1, 1:-1]) / (dy * dy) - dfdz2 = (field[1:-1, 1:-1, 2:] + field[1:-1, 1:-1, :-2] - 2.0 * field[1:-1, 1:-1, 1:-1]) / (dz * dz) + # Laplacian using tile-local ghost-cell neighbors instead of jnp.roll. + dfdx2 = (field[forward, active, active] + field[backward, active, active] - 2.0 * field[active, active, active]) / (dx * dx) + dfdy2 = (field[active, forward, active] + field[active, backward, active] - 2.0 * field[active, active, active]) / (dy * dy) + dfdz2 = (field[active, active, forward] + field[active, active, backward] - 2.0 * field[active, active, active]) / (dz * dz) return dfdx2 + dfdy2 + dfdz2 # compute the laplacian on the interior using ghost cell neighbors def apply_bc(field): - field = apply_scalar_conducting_bc(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field - # apply boundary conditions and update ghost cells + return _refresh_single_tile_scalar(field, world, g, tile_shape, apply_conducting=True) + # apply scalar conducting boundaries and refresh ghost cells through the tiled halo path def body_fun(state): phi, r, p, k = state lapl_p = -lapl(p) - alpha = jnp.sum(r * r) / jnp.sum(p[1:-1, 1:-1, 1:-1] * lapl_p) + alpha = jnp.sum(r * r) / jnp.sum(p[active, active, active] * lapl_p) # compute the optimal step size in the direction of the residual - phi_next = phi.at[1:-1, 1:-1, 1:-1].add(alpha * p[1:-1, 1:-1, 1:-1]) + phi_next = phi.at[active, active, active].add(alpha * p[active, active, active]) # update the potential guess phi_next = apply_bc(phi_next) # apply boundary conditions to the new potential guess @@ -121,7 +102,7 @@ def body_fun(state): # compute the new residual after stepping in the direction of p beta = jnp.sum(r_next * r_next) / jnp.sum(r * r) # compute the optimal scaling for the new search direction - p_next = p.at[1:-1, 1:-1, 1:-1].set(r_next + beta * p[1:-1, 1:-1, 1:-1]) + p_next = p.at[active, active, active].set(r_next + beta * p[active, active, active]) p_next = apply_bc(p_next) # compute the new search direction @@ -137,11 +118,11 @@ def cond_fun(state): # if the residual norm is above the tolerance and we haven't exceeded max iterations, continue iterating phi = apply_bc(phi) - residual = rho[1:-1, 1:-1, 1:-1] / eps + lapl(phi) + residual = rho[active, active, active] / eps + lapl(phi) # compute the error in the Poisson equation with the current potential guess p0 = jnp.zeros_like(phi) - p0 = p0.at[1:-1, 1:-1, 1:-1].set(residual) + p0 = p0.at[active, active, active].set(residual) p0 = apply_bc(p0) # initialize the search direction with ghost cells @@ -155,19 +136,18 @@ def cond_fun(state): return apply_bc(phi) -@partial(jit, static_argnames=("solver", "bc")) def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver, bc): """ Compute electrostatic fields from charge deposition and Poisson solve. - All field arrays have shape (Nx+2, Ny+2, Nz+2) with ghost cells. + All field arrays have shape (Nx+2*g, Ny+2*g, Nz+2*g) with ghost cells. Args: world (dict): Simulation world dictionary. particles (list): Particle species list. constants (dict): Physical constants. - rho (ndarray): Charge density array with shape (Nx+2, Ny+2, Nz+2). - phi (ndarray): Potential array with shape (Nx+2, Ny+2, Nz+2). + rho (ndarray): Charge density array with shape (Nx+2*g, Ny+2*g, Nz+2*g). + phi (ndarray): Potential array with shape (Nx+2*g, Ny+2*g, Nz+2*g). solver (str): Electrostatic solver mode. bc (str): Boundary condition for finite-difference gradient. @@ -177,45 +157,105 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver dx = world['dx'] dy = world['dy'] dz = world['dz'] - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + active = _active_slice(g) - rho = compute_rho(particles, rho, world, constants) - - phi = lax.cond( - solver == "spectral", - lambda _: solve_poisson_with_fft(rho, constants, world), - lambda _: solve_poisson_with_conjugate_gradient(rho, phi, constants, world), - operand=None, - ) + phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) - phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) + phi = _refresh_single_tile_scalar(phi, world, g, tile_shape) # refresh ghost cells before any stencil-based post-processing alpha = constants['alpha'] - phi = digital_filter(phi, alpha) - phi = apply_scalar_conducting_bc(phi, bc_x, bc_y, bc_z) - phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) + phi = digital_filter(phi, alpha, num_guard_cells=g) + phi = _refresh_single_tile_scalar(phi, world, g, tile_shape, apply_conducting=True) # update ghost cells after filtering - Ex, Ey, Ez = lax.cond( - solver == "spectral", - lambda _: spectral_gradient(-1 * phi[1:-1, 1:-1, 1:-1], world), - lambda _: centered_finite_difference_gradient(-1 * phi[1:-1, 1:-1, 1:-1], dx, dy, dz, bc), - operand=None, - ) + del solver, bc + + Ex, Ey, Ez = _centered_finite_difference_gradient(-1.0 * phi[active, active, active], dx, dy, dz) # compute gradient on the interior # Place the gradient results into ghost-celled arrays Ex_full = jnp.zeros_like(phi) Ey_full = jnp.zeros_like(phi) Ez_full = jnp.zeros_like(phi) - Ex_full = Ex_full.at[1:-1, 1:-1, 1:-1].set(Ex) - Ey_full = Ey_full.at[1:-1, 1:-1, 1:-1].set(Ey) - Ez_full = Ez_full.at[1:-1, 1:-1, 1:-1].set(Ez) - Ex_full = update_ghost_cells(Ex_full, bc_x, bc_y, bc_z) - Ey_full = update_ghost_cells(Ey_full, bc_x, bc_y, bc_z) - Ez_full = update_ghost_cells(Ez_full, bc_x, bc_y, bc_z) + Ex_full = Ex_full.at[active, active, active].set(Ex) + Ey_full = Ey_full.at[active, active, active].set(Ey) + Ez_full = Ez_full.at[active, active, active].set(Ez) + Ex_full = _refresh_single_tile_scalar(Ex_full, world, g, tile_shape) + Ey_full = _refresh_single_tile_scalar(Ey_full, world, g, tile_shape) + Ez_full = _refresh_single_tile_scalar(Ez_full, world, g, tile_shape) return (Ex_full, Ey_full, Ez_full), phi, rho + + +def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): + """ + Compute ``E = -grad(phi)`` on compact scalar tiles. + + The potential halos must already contain neighboring tile/global boundary + values. This mirrors ``centered_finite_difference_gradient`` on the + assembled physical interior, then refreshes vector halos for particle + interpolation. + """ + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + g = int(g) + active = slice(g, -g) + forward = slice(g + 1, None if g == 1 else -g + 1) + backward = slice(g - 1, -g - 1) + + phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + + Ex = jnp.zeros_like(phi_tiles) + Ey = jnp.zeros_like(phi_tiles) + Ez = jnp.zeros_like(phi_tiles) + + Ex = Ex.at[:, :, :, active, active, active].set( + -1.0 * (phi_tiles[:, :, :, forward, active, active] - phi_tiles[:, :, :, backward, active, active]) / (2.0 * dx) + ) + Ey = Ey.at[:, :, :, active, active, active].set( + -1.0 * (phi_tiles[:, :, :, active, forward, active] - phi_tiles[:, :, :, active, backward, active]) / (2.0 * dy) + ) + Ez = Ez.at[:, :, :, active, active, active].set( + -1.0 * (phi_tiles[:, :, :, active, active, forward] - phi_tiles[:, :, :, active, active, backward]) / (2.0 * dz) + ) + + return ghost_cells.update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) + + +def calculate_tiled_electrostatic_fields(world, particles, species_config, constants, rho_tiles, phi_tiles, solver, bc): + """ + Compute electrostatic fields from single-tile rho deposition and a Poisson solve. + + Electrostatic runs use one tile covering the whole physical domain. The + leading tile axes remain singleton axes, so the Poisson solve acts directly + on ``rho_tiles[0, 0, 0]`` and ``phi_tiles[0, 0, 0]`` without assembling a + separate global representation. + """ + + del bc + + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + rho_tiles = compute_rho(particles, species_config, rho_tiles, constants, world) + rho = rho_tiles[0, 0, 0] + phi = phi_tiles[0, 0, 0] + + phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) + phi_tiles = phi_tiles.at[0, 0, 0].set(phi) + phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + # refresh ghost cells before filtering and tiled differentiation + + alpha = constants["alpha"] + phi_tiles = digital_filter(phi_tiles, alpha, num_guard_cells=g) + phi_tiles = ghost_cells.apply_tiled_scalar_conducting_bc(phi_tiles, world, num_guard_cells=g) + phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + # keep the same phi post-processing order as the previous electrostatic solver + + E_tiles = _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g) + + return E_tiles, phi_tiles, rho_tiles diff --git a/PyPIC3D/solvers/fdtd.py b/PyPIC3D/solvers/fdtd.py deleted file mode 100644 index 317bbb7d..00000000 --- a/PyPIC3D/solvers/fdtd.py +++ /dev/null @@ -1,210 +0,0 @@ -import jax.numpy as jnp -import jax -from jax import jit -from jax import lax -# import external libraries - -from PyPIC3D.boundary_conditions.boundaryconditions import apply_zero_boundary_condition - -def solve_poisson_sor(phi, rho, dx, dy, dz, eps, omega=1.5, tol=1e-6, max_iter=10000): - """ - Solve Poisson's equation using Successive Over-Relaxation (SOR) method. - - Args: - phi (jax.numpy.ndarray): Initial guess for the potential. - rho (jax.numpy.ndarray): Charge density. - omega (float): Relaxation factor. - tol (float): Tolerance for convergence. - max_iter (int): Maximum number of iterations. - - Returns: - jax.numpy.ndarray: Solution for the potential. - """ - phi = jnp.array(phi) - rho = jnp.array(rho) - b = dx * dy * dz / jnp.pi - - def cond_fun(val): - _, norm_diff, iter_count = val - return (norm_diff >= tol) & (iter_count < max_iter) - - def body_fun(val): - phi, _, iter_count = val - phi_old = phi.copy() - phi = phi.at[1:-1, 1:-1, 1:-1].set( - (1 - omega) * phi[1:-1, 1:-1, 1:-1] + omega / 6 * ( - phi[:-2, 1:-1, 1:-1] + phi[2:, 1:-1, 1:-1] + - phi[1:-1, :-2, 1:-1] + phi[1:-1, 2:, 1:-1] + - phi[1:-1, 1:-1, :-2] + phi[1:-1, 1:-1, 2:] - - b * rho[1:-1, 1:-1, 1:-1] / eps - ) - ) - return phi, jnp.linalg.norm(phi - phi_old), iter_count + 1 - - phi, _, _ = jax.lax.while_loop(cond_fun, body_fun, (phi, jnp.inf, 0)) - - return phi - - -#@partial(jit, static_argnums=(1, 2, 3, 4)) -def centered_finite_difference_laplacian(field, dx, dy, dz, bc): - """ - Calculates the Laplacian of a given field using centered finite difference and applies the specified boundary conditions. - - Args: - field: numpy.ndarray - The input field. - dx: float - The spacing between grid points in the x-direction. - dy: float - The spacing between grid points in the y-direction. - dz: float - The spacing between grid points in the z-direction. - boundary_condition: str - The type of boundary condition ('periodic', 'neumann', 'dirichlet'). - - Returns: - numpy.ndarray - The Laplacian of the field with the specified boundary conditions applied. - """ - - if bc == 'dirichlet': - field = apply_zero_boundary_condition(field) - # apply zero boundary condition at the edges of the field - - x_comp = (jnp.roll(field, shift=1, axis=0) + jnp.roll(field, shift=-1, axis=0) - 2*field) / (dx*dx) - y_comp = (jnp.roll(field, shift=1, axis=1) + jnp.roll(field, shift=-1, axis=1) - 2*field) / (dy*dy) - z_comp = (jnp.roll(field, shift=1, axis=2) + jnp.roll(field, shift=-1, axis=2) - 2*field) / (dz*dz) - # calculate the Laplacian of the field using centered finite difference - - if bc == 'neumann': - x_comp = apply_zero_boundary_condition(x_comp) - y_comp = apply_zero_boundary_condition(y_comp) - z_comp = apply_zero_boundary_condition(z_comp) - - return x_comp + y_comp + z_comp - -#@partial(jit, static_argnums=(3, 4, 5, 6)) -def centered_finite_difference_curl(field_x, field_y, field_z, dx, dy, dz, bc): - """ - Computes the curl of a vector field using centered finite differencing and applies the specified boundary conditions. - - Args: - field_x: numpy.ndarray - The x-component of the vector field. - field_y: numpy.ndarray - The y-component of the vector field. - field_z: numpy.ndarray - The z-component of the vector field. - dx: float - The spacing between grid points in the x-direction. - dy: float - The spacing between grid points in the y-direction. - dz: float - The spacing between grid points in the z-direction. - bc: str - The type of boundary condition ('periodic', 'neumann', 'dirichlet'). - - Returns: - tuple of numpy.ndarray - The curl components (curl_x, curl_y, curl_z) of the vector field with the specified boundary conditions applied. - """ - if bc == 'dirichlet': - field_x = apply_zero_boundary_condition(field_x) - field_y = apply_zero_boundary_condition(field_y) - field_z = apply_zero_boundary_condition(field_z) - - dFz_dy = (jnp.roll(field_z, shift=-1, axis=1) - jnp.roll(field_z, shift=1, axis=1)) / (2 * dy) - dFy_dz = (jnp.roll(field_y, shift=-1, axis=2) - jnp.roll(field_y, shift=1, axis=2)) / (2 * dz) - dFx_dz = (jnp.roll(field_x, shift=-1, axis=2) - jnp.roll(field_x, shift=1, axis=2)) / (2 * dz) - dFz_dx = (jnp.roll(field_z, shift=-1, axis=0) - jnp.roll(field_z, shift=1, axis=0)) / (2 * dx) - dFy_dx = (jnp.roll(field_y, shift=-1, axis=0) - jnp.roll(field_y, shift=1, axis=0)) / (2 * dx) - dFx_dy = (jnp.roll(field_x, shift=-1, axis=1) - jnp.roll(field_x, shift=1, axis=1)) / (2 * dy) - - curl_x = dFz_dy - dFy_dz - curl_y = dFx_dz - dFz_dx - curl_z = dFy_dx - dFx_dy - # calculate the curl of the field - - if bc == 'neumann': - curl_x = apply_zero_boundary_condition(curl_x) - curl_y = apply_zero_boundary_condition(curl_y) - curl_z = apply_zero_boundary_condition(curl_z) - - return curl_x, curl_y, curl_z - -#@partial(jit, static_argnums=(1, 2, 3, 4)) -def centered_finite_difference_gradient(field, dx, dy, dz, bc): - """ - Computes the gradient of a scalar field using centered finite differencing and applies the specified boundary conditions. - - Args: - field: numpy.ndarray - The input scalar field. - dx: float - The spacing between grid points in the x-direction. - dy: float - The spacing between grid points in the y-direction. - dz: float - The spacing between grid points in the z-direction. - bc: str - The type of boundary condition ('periodic', 'neumann', 'dirichlet'). - - Returns: - tuple of numpy.ndarray - The gradient components (grad_x, grad_y, grad_z) of the scalar field with the specified boundary conditions applied. - """ - if bc == 'dirichlet': - field = apply_zero_boundary_condition(field) - - grad_x = (jnp.roll(field, shift=-1, axis=0) - jnp.roll(field, shift=1, axis=0)) / (2 * dx) - grad_y = (jnp.roll(field, shift=-1, axis=1) - jnp.roll(field, shift=1, axis=1)) / (2 * dy) - grad_z = (jnp.roll(field, shift=-1, axis=2) - jnp.roll(field, shift=1, axis=2)) / (2 * dz) - - if bc == 'neumann': - grad_x = apply_zero_boundary_condition(grad_x) - grad_y = apply_zero_boundary_condition(grad_y) - grad_z = apply_zero_boundary_condition(grad_z) - - return grad_x, grad_y, grad_z - -#@partial(jit, static_argnums=(3, 4, 5, 6)) -def centered_finite_difference_divergence(field_x, field_y, field_z, dx, dy, dz, bc): - """ - Computes the divergence of a vector field using centered finite differencing and applies the specified boundary conditions. - - Args: - field_x: numpy.ndarray - The x-component of the vector field. - field_y: numpy.ndarray - The y-component of the vector field. - field_z: numpy.ndarray - The z-component of the vector field. - dx: float - The spacing between grid points in the x-direction. - dy: float - The spacing between grid points in the y-direction. - dz: float - The spacing between grid points in the z-direction. - bc: str - The type of boundary condition ('periodic', 'neumann', 'dirichlet'). - - Returns: - numpy.ndarray - The divergence of the vector field with the specified boundary conditions applied. - """ - if bc == 'dirichlet': - field_x = apply_zero_boundary_condition(field_x) - field_y = apply_zero_boundary_condition(field_y) - field_z = apply_zero_boundary_condition(field_z) - - div_x = (jnp.roll(field_x, shift=-1, axis=0) - jnp.roll(field_x, shift=1, axis=0)) / (2 * dx) - div_y = (jnp.roll(field_y, shift=-1, axis=1) - jnp.roll(field_y, shift=1, axis=1)) / (2 * dy) - div_z = (jnp.roll(field_z, shift=-1, axis=2) - jnp.roll(field_z, shift=1, axis=2)) / (2 * dz) - - if bc == 'neumann': - div_x = apply_zero_boundary_condition(div_x) - div_y = apply_zero_boundary_condition(div_y) - div_z = apply_zero_boundary_condition(div_z) - - return div_x + div_y + div_z \ No newline at end of file diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py index 7bc10c6c..c56ae8fb 100644 --- a/PyPIC3D/solvers/first_order_yee.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -1,210 +1,152 @@ -import jax.numpy as jnp -from jax import jit -from functools import partial -# import external libraries - -from PyPIC3D.utils import digital_filter -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells, apply_conducting_bc from PyPIC3D.boundary_conditions.PML import ( - apply_pml_to_b_curl, - apply_pml_to_e_curl, - update_ghost_cells_for_pml, + apply_tiled_pml_to_b_curl, + apply_tiled_pml_to_e_curl, ) -# import internal libraries +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.utilities.filters import digital_filter_vector -@partial(jit, static_argnames=("curl_func",)) -def update_E(E, B, J, world, constants, curl_func, pml_state=None): +def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ - Update the electric field components (Ex, Ey, Ez) based on the given parameters. - - Uses ghost cells for boundary handling. Fields have shape (Nx+2, Ny+2, Nz+2) - with 1 ghost cell on each side. The physical interior is [1:-1, 1:-1, 1:-1]. - - Args: - E (tuple): A tuple containing the electric field components (Ex, Ey, Ez). - B (tuple): A tuple containing the magnetic field components (Bx, By, Bz). - J (tuple): A tuple containing the current density components (Jx, Jy, Jz). - world (dict): A dictionary containing the world parameters such as 'dx', 'dy', 'dz', and 'dt'. - constants (dict): A dictionary containing the physical constants such as 'C' (speed of light) and 'eps' (permittivity). - curl_func (function): A function to calculate the curl of the magnetic field. - - Returns: - tuple: ``((Ex, Ey, Ez), pml_state)``. When no PML is active, - ``pml_state`` is passed through as ``None`` and the update is the - ordinary unstretched Yee update. + Update compact tiled electric fields without assembling a global field. + + The Yee curl is evaluated on each tile's physical interior after B halos + have been refreshed from neighbor tiles or field boundary conditions. """ - Ex, Ey, Ez = E - Bx, By, Bz = B - Jx, Jy, Jz = J - # unpack the E, B, and J fields - - dt = world['dt'] - dx, dy, dz = world['dx'], world['dy'], world['dz'] - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - # get the time resolution and grid spacings - - C = constants['C'] - eps = constants['eps'] - # get the necessary constants - - # forward differences using ghost cells - # B has shape (Nx+2, Ny+2, Nz+2) with valid ghost cell values - # slicing [1:-1, 2:, 1:-1] reaches into the ghost cell at the +1 boundary - dBz_dy = (Bz[1:-1, 2:, 1:-1] - Bz[1:-1, 1:-1, 1:-1]) / dy - dBy_dz = (By[1:-1, 1:-1, 2:] - By[1:-1, 1:-1, 1:-1]) / dz - dBx_dz = (Bx[1:-1, 1:-1, 2:] - Bx[1:-1, 1:-1, 1:-1]) / dz - dBx_dy = (Bx[1:-1, 2:, 1:-1] - Bx[1:-1, 1:-1, 1:-1]) / dy - dBz_dx = (Bz[2:, 1:-1, 1:-1] - Bz[1:-1, 1:-1, 1:-1]) / dx - dBy_dx = (By[2:, 1:-1, 1:-1] - By[1:-1, 1:-1, 1:-1]) / dx - # compute the forward finite differences of the magnetic field + Ex, Ey, Ez = E_tiles + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + # get the tile information + + active = slice(g, -g) + # build interior slice for active axes + forward = slice(g + 1, None if g == 1 else -g + 1) + # build forward slice used for forward differences + + if pml_state is None: + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells(B_tiles, world, g, tile_shape) + else: + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells_for_pml(B_tiles, world, g, tile_shape) + Jx, Jy, Jz = J_tiles + current = slice(g, -g) #_active_slice(g) + + dt = world["dt"] + dx, dy, dz = world["dx"], world["dy"], world["dz"] + C = constants["C"] + eps = constants["eps"] + + # Forward differences use each tile's + side guard cell. Those guards now + # contain the adjacent tile's interior value, including periodic wrap at + # the global edge. + dBz_dy = (Bz[:, :, :, active, forward, active] - Bz[:, :, :, active, active, active]) / dy + dBy_dz = (By[:, :, :, active, active, forward] - By[:, :, :, active, active, active]) / dz + dBx_dz = (Bx[:, :, :, active, active, forward] - Bx[:, :, :, active, active, active]) / dz + dBx_dy = (Bx[:, :, :, active, forward, active] - Bx[:, :, :, active, active, active]) / dy + dBz_dx = (Bz[:, :, :, forward, active, active] - Bz[:, :, :, active, active, active]) / dx + dBy_dx = (By[:, :, :, forward, active, active] - By[:, :, :, active, active, active]) / dx if pml_state is None: curl_x = dBz_dy - dBy_dz curl_y = dBx_dz - dBz_dx curl_z = dBy_dx - dBx_dy else: - # In the PML layer, each Yee derivative is interpreted as a - # coordinate-stretched derivative before the curl is assembled. - (curl_x, curl_y, curl_z), pml_state = apply_pml_to_e_curl( + (curl_x, curl_y, curl_z), pml_state = apply_tiled_pml_to_e_curl( (dBz_dy, dBy_dz, dBx_dz, dBz_dx, dBy_dx, dBx_dy), world, pml_state, ) - # calculate the curl of the magnetic field, optionally using stretched derivatives - Ex = Ex.at[1:-1, 1:-1, 1:-1].set( - Ex[1:-1, 1:-1, 1:-1] + (C**2 * curl_x - Jx[1:-1, 1:-1, 1:-1] / eps) * dt + Ex = Ex.at[:, :, :, active, active, active].set( + Ex[:, :, :, active, active, active] + + (C**2 * curl_x - Jx[:, :, :, current, current, current] / eps) * dt ) - Ey = Ey.at[1:-1, 1:-1, 1:-1].set( - Ey[1:-1, 1:-1, 1:-1] + (C**2 * curl_y - Jy[1:-1, 1:-1, 1:-1] / eps) * dt + Ey = Ey.at[:, :, :, active, active, active].set( + Ey[:, :, :, active, active, active] + + (C**2 * curl_y - Jy[:, :, :, current, current, current] / eps) * dt ) - Ez = Ez.at[1:-1, 1:-1, 1:-1].set( - Ez[1:-1, 1:-1, 1:-1] + (C**2 * curl_z - Jz[1:-1, 1:-1, 1:-1] / eps) * dt + Ez = Ez.at[:, :, :, active, active, active].set( + Ez[:, :, :, active, active, active] + + (C**2 * curl_z - Jz[:, :, :, current, current, current] / eps) * dt ) - # update the electric field from Maxwell's equations if pml_state is None: - Ex = update_ghost_cells(Ex, bc_x, bc_y, bc_z) - Ey = update_ghost_cells(Ey, bc_x, bc_y, bc_z) - Ez = update_ghost_cells(Ez, bc_x, bc_y, bc_z) + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) else: - Ex = update_ghost_cells_for_pml(Ex, world) - Ey = update_ghost_cells_for_pml(Ey, world) - Ez = update_ghost_cells_for_pml(Ez, world) - # refresh ghost cells before any stencil-based post-processing + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape) + # refresh tile halos before the digital field filter, matching the global + # ghost-cell order in the standard Yee solver. - alpha = constants['alpha'] - Ex = digital_filter(Ex, alpha) - Ey = digital_filter(Ey, alpha) - Ez = digital_filter(Ez, alpha) - # apply a digital filter to the electric field components + Ex, Ey, Ez = digital_filter_vector((Ex, Ey, Ez), constants.get("alpha", 1.0), num_guard_cells=g) - Ex, Ey, Ez = apply_conducting_bc((Ex, Ey, Ez), bc_x, bc_y, bc_z) - # enforce conducting boundary conditions on tangential E + Ex, Ey, Ez = ghost_cells.apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) if pml_state is None: - Ex = update_ghost_cells(Ex, bc_x, bc_y, bc_z) - Ey = update_ghost_cells(Ey, bc_x, bc_y, bc_z) - Ez = update_ghost_cells(Ez, bc_x, bc_y, bc_z) - else: - Ex = update_ghost_cells_for_pml(Ex, world) - Ey = update_ghost_cells_for_pml(Ey, world) - Ez = update_ghost_cells_for_pml(Ez, world) - # fill ghost cells with the updated values + return ghost_cells.update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape), None - return (Ex, Ey, Ez), pml_state + return ghost_cells.update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state -@partial(jit, static_argnames=("curl_func",)) -def update_B(E, B, world, constants, curl_func, pml_state=None): +def update_B(E_tiles, B_tiles, world, constants, pml_state=None): """ - Update the magnetic field components (Bx, By, Bz) using the curl of the electric field. - - Uses ghost cells for boundary handling. Fields have shape (Nx+2, Ny+2, Nz+2) - with 1 ghost cell on each side. The physical interior is [1:-1, 1:-1, 1:-1]. - - Args: - E (tuple): The electric field components (Ex, Ey, Ez). - B (tuple): The magnetic field components (Bx, By, Bz). - world (dict): Dictionary containing simulation parameters such as 'dx', 'dy', 'dz', and 'dt'. - constants (dict): Dictionary containing physical constants. - curl_func (function): Function to calculate the curl of the electric field. - - Returns: - tuple: ``((Bx, By, Bz), pml_state)``. When no PML is active, - ``pml_state`` is passed through as ``None`` and the update is the - ordinary unstretched Yee update. + Update compact tiled magnetic fields without assembling a global field. + + The Yee curl is evaluated on each tile's physical interior after E halos + have been refreshed from neighbor tiles or field boundary conditions. """ - dt = world['dt'] - # get the time resolution - dx, dy, dz = world['dx'], world['dy'], world['dz'] - # get the grid spacings - bc_x = world['boundary_conditions']['x'] - bc_y = world['boundary_conditions']['y'] - bc_z = world['boundary_conditions']['z'] - - Ex, Ey, Ez = E - Bx, By, Bz = B - # unpack the E and B fields - - # backward differences using ghost cells - # slicing [1:-1, :-2, 1:-1] reaches into the ghost cell at the -1 boundary - dEz_dy = (Ez[1:-1, 1:-1, 1:-1] - Ez[1:-1, :-2, 1:-1]) / dy - dEy_dz = (Ey[1:-1, 1:-1, 1:-1] - Ey[1:-1, 1:-1, :-2]) / dz - dEx_dz = (Ex[1:-1, 1:-1, 1:-1] - Ex[1:-1, 1:-1, :-2]) / dz - dEx_dy = (Ex[1:-1, 1:-1, 1:-1] - Ex[1:-1, :-2, 1:-1]) / dy - dEz_dx = (Ez[1:-1, 1:-1, 1:-1] - Ez[:-2, 1:-1, 1:-1]) / dx - dEy_dx = (Ey[1:-1, 1:-1, 1:-1] - Ey[:-2, 1:-1, 1:-1]) / dx - # compute the backward finite differences of the electric field + Bx, By, Bz = B_tiles + tile_shape = tuple(int(width) for width in world["tile_shape"]) + tile_nx, tile_ny, tile_nz = tile_shape + g = int(world["guard_cells"]) + g = int(g) + active = slice(g, -g) + # build interior slice for active axes + backward = slice(g - 1, -g - 1) + # build backward slice for active axes, used for backward differences + + if pml_state is None: + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, g, tile_shape) + else: + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells_for_pml(E_tiles, world, g, tile_shape) + dt = world["dt"] + dx, dy, dz = world["dx"], world["dy"], world["dz"] + + # Backward differences use each tile's - side guard cell. Those guards now + # contain the adjacent tile's interior value, including periodic wrap at + # the global edge. + dEz_dy = (Ez[:, :, :, active, active, active] - Ez[:, :, :, active, backward, active]) / dy + dEy_dz = (Ey[:, :, :, active, active, active] - Ey[:, :, :, active, active, backward]) / dz + dEx_dz = (Ex[:, :, :, active, active, active] - Ex[:, :, :, active, active, backward]) / dz + dEx_dy = (Ex[:, :, :, active, active, active] - Ex[:, :, :, active, backward, active]) / dy + dEz_dx = (Ez[:, :, :, active, active, active] - Ez[:, :, :, backward, active, active]) / dx + dEy_dx = (Ey[:, :, :, active, active, active] - Ey[:, :, :, backward, active, active]) / dx if pml_state is None: curl_x = dEz_dy - dEy_dz curl_y = dEx_dz - dEz_dx curl_z = dEy_dx - dEx_dy else: - # In the PML layer, the frequency-space coordinate stretch is carried by - # ADE memory terms attached to each finite-difference derivative. - (curl_x, curl_y, curl_z), pml_state = apply_pml_to_b_curl( + (curl_x, curl_y, curl_z), pml_state = apply_tiled_pml_to_b_curl( (dEz_dy, dEy_dz, dEx_dz, dEz_dx, dEy_dx, dEx_dy), world, pml_state, ) - # calculate the curl of the electric field, optionally using stretched derivatives - Bx = Bx.at[1:-1, 1:-1, 1:-1].set(Bx[1:-1, 1:-1, 1:-1] - dt * curl_x) - By = By.at[1:-1, 1:-1, 1:-1].set(By[1:-1, 1:-1, 1:-1] - dt * curl_y) - Bz = Bz.at[1:-1, 1:-1, 1:-1].set(Bz[1:-1, 1:-1, 1:-1] - dt * curl_z) - # update the magnetic field from Maxwell's equations + Bx = Bx.at[:, :, :, active, active, active].set(Bx[:, :, :, active, active, active] - dt * curl_x) + By = By.at[:, :, :, active, active, active].set(By[:, :, :, active, active, active] - dt * curl_y) + Bz = Bz.at[:, :, :, active, active, active].set(Bz[:, :, :, active, active, active] - dt * curl_z) if pml_state is None: - Bx = update_ghost_cells(Bx, bc_x, bc_y, bc_z) - By = update_ghost_cells(By, bc_x, bc_y, bc_z) - Bz = update_ghost_cells(Bz, bc_x, bc_y, bc_z) + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) else: - Bx = update_ghost_cells_for_pml(Bx, world) - By = update_ghost_cells_for_pml(By, world) - Bz = update_ghost_cells_for_pml(Bz, world) - # refresh ghost cells before any stencil-based post-processing + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape) + # refresh tile halos before the digital field filter, matching the global + # ghost-cell order in the standard Yee solver. - alpha = constants['alpha'] - Bx = digital_filter(Bx, alpha) - By = digital_filter(By, alpha) - Bz = digital_filter(Bz, alpha) - # apply a digital filter to the magnetic field components + Bx, By, Bz = digital_filter_vector((Bx, By, Bz), constants.get("alpha", 1.0), num_guard_cells=g) if pml_state is None: - Bx = update_ghost_cells(Bx, bc_x, bc_y, bc_z) - By = update_ghost_cells(By, bc_x, bc_y, bc_z) - Bz = update_ghost_cells(Bz, bc_x, bc_y, bc_z) - else: - Bx = update_ghost_cells_for_pml(Bx, world) - By = update_ghost_cells_for_pml(By, world) - Bz = update_ghost_cells_for_pml(Bz, world) - # fill ghost cells with the updated values + return ghost_cells.update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape), None - return (Bx, By, Bz), pml_state + return ghost_cells.update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape), pml_state diff --git a/PyPIC3D/solvers/pstd.py b/PyPIC3D/solvers/pstd.py deleted file mode 100644 index 69120923..00000000 --- a/PyPIC3D/solvers/pstd.py +++ /dev/null @@ -1,171 +0,0 @@ -from jax import jit -import jax -import jax.numpy as jnp -from functools import partial -# import external libraries - -def create_k_mesh(nx, ny, nz, dx, dy, dz): - """ - Create a mesh of wave numbers for FFT operations. - - Args: - nx (int): Number of grid points in the x-direction. - ny (int): Number of grid points in the y-direction. - nz (int): Number of grid points in the z-direction. - dx (float): Grid spacing in the x-direction. - dy (float): Grid spacing in the y-direction. - dz (float): Grid spacing in the z-direction. - - Returns: - tuple: Meshgrid of wave numbers (kx, ky, kz). - """ - kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi - ky = jnp.fft.fftfreq(ny, d=dy) * 2 * jnp.pi - kz = jnp.fft.fftfreq(nz, d=dz) * 2 * jnp.pi - return jnp.meshgrid(kx, ky, kz, indexing='ij') - -@jit -def spectral_poisson_solve(rho, constants, world): - """ - Solve the Poisson equation for electrostatic potential using spectral method. - - Args: - rho (ndarray): Charge density. - eps (float): Permittivity. - dx (float): Grid spacing in the x-direction. - dy (float): Grid spacing in the y-direction. - dz (float): Grid spacing in the z-direction. - - Returns: - phi (ndarray): Solution to the Poisson equation. - """ - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - eps = constants['eps'] - - krho = jnp.fft.fftn(rho) - nx, ny, nz = krho.shape - kx, ky, kz = create_k_mesh(nx, ny, nz, dx, dy, dz) - k2 = kx**2 + ky**2 + kz**2 - k2 = k2.at[0, 0, 0].set(1.0) - phi = krho / (eps*k2) - phi = phi.at[0, 0, 0].set(0) - phi = jnp.fft.ifftn(phi).real - return phi - -@jit -def spectral_divergence(xfield, yfield, zfield, world): - """ - Calculate the spectral divergence of a 3D vector field. - - This function computes the divergence of a vector field in the spectral domain - using the Fast Fourier Transform (FFT). The input fields are assumed to be - periodic in all three dimensions. - - Args: - xfield (ndarray): The x-component of the vector field. - yfield (ndarray): The y-component of the vector field. - zfield (ndarray): The z-component of the vector field. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - - Returns: - ndarray: The real part of the inverse FFT of the spectral divergence. - """ - - xfft = jnp.fft.fftn(xfield) - yfft = jnp.fft.fftn(yfield) - zfft = jnp.fft.fftn(zfield) - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - nx, ny, nz = xfft.shape - kx, ky, kz = create_k_mesh(nx, ny, nz, dx, dy, dz) - div = 1j * kx * xfft + 1j * ky * yfft + 1j * kz * zfft - return jnp.fft.ifftn(div).real - -@jit -def spectral_curl(xfield, yfield, zfield, world): - """ - Compute the curl of a 3D vector field using spectral methods. - - Args: - xfield (ndarray): The x-component of the vector field. - yfield (ndarray): The y-component of the vector field. - zfield (ndarray): The z-component of the vector field. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - - Returns: - tuple: A tuple containing the x, y, and z components of the curl of the vector field. - """ - - xfft = jnp.fft.fftn(xfield) - yfft = jnp.fft.fftn(yfield) - zfft = jnp.fft.fftn(zfield) - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - nx, ny, nz = xfft.shape - kx, ky, kz = create_k_mesh(nx, ny, nz, dx, dy, dz) - curlx = 1j * ky * zfft - 1j * kz * yfft - curly = 1j * kz * xfft - 1j * kx * zfft - curlz = 1j * kx * yfft - 1j * ky * xfft - return jnp.fft.ifftn(curlx).real, jnp.fft.ifftn(curly).real, jnp.fft.ifftn(curlz).real - -@jit -def spectral_gradient(field, world): - """ - Compute the gradient of a 3D scalar field using spectral methods. - - Args: - field (ndarray): The input scalar field. - dx (float): The grid spacing in the x-direction. - dy (float): The grid spacing in the y-direction. - dz (float): The grid spacing in the z-direction. - - Returns: - tuple: A tuple containing the x, y, and z components of the gradient of the field. - """ - - field_fft = jnp.fft.fftn(field) - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - nx, ny, nz = field_fft.shape - kx, ky, kz = create_k_mesh(nx, ny, nz, dx, dy, dz) - gradx = 1j * kx * field_fft - grady = 1j * ky * field_fft - gradz = 1j * kz * field_fft - return jnp.fft.ifftn(gradx).real, jnp.fft.ifftn(grady).real, jnp.fft.ifftn(gradz).real - -@jit -def spectral_laplacian(field, world): - """ - Calculates the Laplacian of a given field using spectral method. Assumes periodic boundary conditions. - - Args: - field: numpy.ndarray - The input field. - dx: float - The spacing between grid points in the x-direction. - dy: float - The spacing between grid points in the y-direction. - dz: float - The spacing between grid points in the z-direction. - - Returns: - numpy.ndarray - The Laplacian of the field. - """ - field_fft = jnp.fft.fftn(field) - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - nx, ny, nz = field_fft.shape - kx, ky, kz = create_k_mesh(nx, ny, nz, dx, dy, dz) - lapl = -(kx**2 + ky**2 + kz**2) * field_fft - return jnp.fft.ifftn(lapl).real \ No newline at end of file diff --git a/PyPIC3D/solvers/vector_potential.py b/PyPIC3D/solvers/vector_potential.py deleted file mode 100644 index 9dd72c11..00000000 --- a/PyPIC3D/solvers/vector_potential.py +++ /dev/null @@ -1,323 +0,0 @@ -import jax -import jax.numpy as jnp - -from PyPIC3D.solvers.fdtd import ( - centered_finite_difference_curl, - centered_finite_difference_laplacian, - centered_finite_difference_divergence, - centered_finite_difference_gradient -) -from PyPIC3D.pusher.boris import interpolate_field_to_particles -from PyPIC3D.utils import digital_filter - - -def initialize_vector_potential(J, world, constants): - """ - Initialize the vector potential A based on the current density J. - - Args: - J (tuple): Current density components (Jx, Jy, Jz). - world (dict): Simulation world parameters including 'dx', 'dy', 'dz'. - constants (dict): Physical constants including 'mu'. - - Returns: - tuple: Initial vector potential components (Ax, Ay, Az). - """ - Jx, Jy, Jz = J - - Ax = jnp.zeros_like(Jx) - Ay = jnp.zeros_like(Jy) - Az = jnp.zeros_like(Jz) - # initialize A as zero - - A = Ax, Ay, Az - - return A, A, A - -@jax.jit -def update_vector_potential(J, world, constants, A1, A0): - """ - Update the electromagnetic vector potential using an explicit second-order - finite-difference time-domain (FDTD) scheme with periodic boundary conditions. - - This routine advances the vector potential components (Ax, Ay, Az) from the - previous two time levels (A0, A1) using the wave equation form with a - current-density source and a gauge-like correction term: - A^{n+1} = 2 A^{n} - A^{n-1} + C^2 dt^2 * ( mu * J + ∇²A - ∇(∇·A) ) - A simple digital filter is applied to each updated component to reduce numerical - noise. - - Parameters - ---------- - J : tuple[array_like, array_like, array_like] - Current density components (Jx, Jy, Jz) defined on the grid. - world : dict - Simulation/grid parameters. Must contain: - - 'dx', 'dy', 'dz' : float - Grid spacings in each direction. - - 'dt' : float - Time step. - constants : dict - Physical/numerical constants. Must contain: - - 'mu' : float - Magnetic permeability used in the source term. - - 'C' : float - Wave propagation speed (e.g., speed of light). - - 'alpha' : float - Digital filter parameter passed to `digital_filter`. - A1 : tuple[array_like, array_like, array_like] - Vector potential at the current time level n: (Ax, Ay, Az). - A0 : tuple[array_like, array_like, array_like] - Vector potential at the previous time level n-1: (Ax0, Ay0, Az0). - - Returns - ------- - tuple[array_like, array_like, array_like] - Updated vector potential at the next time level n+1: (Ax_new, Ay_new, Az_new). - - Notes - ----- - - Spatial operators are computed using centered finite differences with - 'periodic' boundary conditions via: - `centered_finite_difference_laplacian`, - `centered_finite_difference_divergence`, - `centered_finite_difference_gradient`. - - The term ∇(∇·A) helps control the divergence of A depending on the chosen gauge. - - The update is explicit; stability depends on dt relative to grid spacing and C - (CFL-like constraint). - """ - - - Ax0, Ay0, Az0 = A0 - Ax, Ay, Az = A1 - Jx, Jy, Jz = J - - dx, dy, dz = world['dx'], world['dy'], world['dz'] - - mu = constants['mu'] - C = constants['C'] - dt = world['dt'] - - laplacian_Ax = centered_finite_difference_laplacian(Ax, dx, dy, dz, 'periodic') - laplacian_Ay = centered_finite_difference_laplacian(Ay, dx, dy, dz, 'periodic') - laplacian_Az = centered_finite_difference_laplacian(Az, dx, dy, dz, 'periodic') - # calculate the Laplacian of the vector potential using centered finite difference - - divA = centered_finite_difference_divergence(Ax, Ay, Az, dx, dy, dz, 'periodic') - # calculate the divergence of the vector potential using centered finite difference - - gradx, grady, gradz = centered_finite_difference_gradient(divA, dx, dy, dz, 'periodic') - # calculate the gradient of the divergence of the vector potential using centered finite difference - - Ax_new = 2 * Ax - Ax0 + C**2 * dt**2 * ( mu * Jx + laplacian_Ax - gradx ) - Ay_new = 2 * Ay - Ay0 + C**2 * dt**2 * ( mu * Jy + laplacian_Ay - grady ) - Az_new = 2 * Az - Az0 + C**2 * dt**2 * ( mu * Jz + laplacian_Az - gradz ) - # update the vector potential using centered finite difference - - alpha = constants['alpha'] - Ax_new = digital_filter(Ax_new, alpha) - Ay_new = digital_filter(Ay_new, alpha) - Az_new = digital_filter(Az_new, alpha) - # apply a digital filter to the vector potential components - - return Ax_new, Ay_new, Az_new - -@jax.jit -def E_from_A(A2, A1, A0, world, grid, staggered_grid, interpolation_order=2): - """ - Compute the electric field components from the vector potential via a leapfrog-consistent - finite difference in time and interpolate them onto Yee-staggered grids. - - This routine uses: - E(t + 1/2) = - (A(t + 1) - A(t)) / dt - where `A1` and `A0` represent the vector potential at successive time levels. The raw - finite-difference fields are computed on the base `grid` and then interpolated to the - appropriate staggered (Yee) locations for each component. - - Parameters - ---------- - A2 : tuple[array-like, array-like, array-like] - Vector potential components (Ax, Ay, Az) at a time level not used by this function - (kept for interface/time-stepping consistency). - A1 : tuple[array-like, array-like, array-like] - Vector potential components (Ax, Ay, Az) at the newer time level. - A0 : tuple[array-like, array-like, array-like] - Vector potential components (Ax, Ay, Az) at the older time level. - world : dict - Simulation state containing at least: - - 'dt' (float): time step size. - grid : tuple - Base grid specification for the vector potential and intermediate fields. The exact - structure is passed through to `interpolate_field`. - staggered_grid : tuple - Staggered grid specification (typically Yee offsets). Each component of E is placed on: - - Ex: (staggered_grid[0], grid[1], grid[2]) - - Ey: (grid[0], staggered_grid[1], grid[2]) - - Ez: (grid[0], grid[1], staggered_grid[2]) - interpolation_order : int, optional - Interpolation order forwarded to `interpolate_field` (default is 2). - - Returns - ------- - Ex_ : array-like - x-component of the electric field interpolated to its Yee-staggered grid. - Ey_ : array-like - y-component of the electric field interpolated to its Yee-staggered grid. - Ez_ : array-like - z-component of the electric field interpolated to its Yee-staggered grid. - - Notes - ----- - - The function assumes `interpolate_field(field, src_grid, dst_grid, interpolation_order=...)` - is available in scope. - - `A2` is accepted but not used in the present implementation. - """ - - - Ax2, Ay2, Az2 = A2 - Ax1, Ay1, Az1 = A1 - Ax0, Ay0, Az0 = A0 - dt = world['dt'] - - Ex = -1 * (Ax1 - Ax0) / dt - Ey = -1 * (Ay1 - Ay0) / dt - Ez = -1 * (Az1 - Az0) / dt - # calculate the electric field from the vector potential using a simple finite difference to get - # E_t+1/2 = - (A_t+1 - A_t) / dt, which is more consistent with the leapfrog time-stepping scheme - - ############## 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 - ########################################################################## - - Ex_ = interpolate_field(Ex, grid, Ex_grid, interpolation_order=interpolation_order) - Ey_ = interpolate_field(Ey, grid, Ey_grid, interpolation_order=interpolation_order) - Ez_ = interpolate_field(Ez, grid, Ez_grid, interpolation_order=interpolation_order) - # interpolate the electric field components to the Yee grid - - return Ex_, Ey_, Ez_ - -@jax.jit -def B_from_A(A, world, grid, staggered_grid, interpolation_order=2): - """ - Compute the magnetic field **B** from a vector potential **A** via a centered - finite-difference curl, then interpolate the result onto a Yee (staggered) grid. - - This routine: - 1) Computes B = ∇×A on the provided *cell-centered* grid using - `centered_finite_difference_curl(..., bc='periodic')`. - 2) Defines the appropriate staggered-grid locations for each B component. - 3) Interpolates each component from the original grid to its Yee-grid location - with `interpolate_field`. - - Parameters - ---------- - A : tuple(array_like, array_like, array_like) - Vector potential components `(Ax, Ay, Az)` defined on `grid`. - world : dict - Simulation metadata containing grid spacings with keys `'dx'`, `'dy'`, `'dz'`. - grid : tuple(array_like, array_like, array_like) - Base grid coordinate arrays `(x, y, z)` (typically cell-centered) on which - `A` is defined. - staggered_grid : tuple(array_like, array_like, array_like) - Staggered coordinate arrays `(x_s, y_s, z_s)` defining half-cell offsets - used for Yee-grid placement. - interpolation_order : int, optional - Order passed to `interpolate_field` when mapping each B component from - `grid` to its staggered target grid. Default is 2. - - Returns - ------- - Bx, By, Bz : array_like - Magnetic field components interpolated to the Yee grid: - - `Bx` on `(grid[0], staggered_grid[1], staggered_grid[2])` - - `By` on `(staggered_grid[0], grid[1], staggered_grid[2])` - - `Bz` on `(staggered_grid[0], staggered_grid[1], grid[2])` - - Notes - ----- - - Boundary conditions for the curl are currently hard-coded to `'periodic'`. - - This function assumes `centered_finite_difference_curl` and `interpolate_field` - are available in scope and accept the argument patterns used here. - """ - - Ax, Ay, Az = A - dx, dy, dz = world['dx'], world['dy'], world['dz'] - # unpack the vector potential components and grid spacing from the input parameters - - - Bx, By, Bz = centered_finite_difference_curl(Ax, Ay, Az, dx, dy, dz, 'periodic') - # calculate the magnetic field from the vector potential using centered finite difference curl - - ################## INTERPOLATION GRIDS ################################### - 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 - ########################################################################## - - Bx_ = interpolate_field(Bx, grid, Bx_grid, interpolation_order) - By_ = interpolate_field(By, grid, By_grid, interpolation_order) - Bz_ = interpolate_field(Bz, grid, Bz_grid, interpolation_order) - # interpolate the magnetic field components to the Yee grid - - return Bx_, By_, Bz_ - -def interpolate_field(field, grid, target_grid, interpolation_order=1): - """ - Interpolate a field defined on a source grid onto a target grid. - - This function creates a 3D mesh from the 1D coordinate arrays provided in - ``target_grid``, flattens the mesh coordinates, and calls - :func:`interpolate_field_to_particles` to evaluate/interpolate the field at those - points. The interpolated values are then reshaped back to the target mesh shape. - - Parameters - ---------- - field : array-like - Field values defined on ``grid``. Typically a JAX array with shape matching - the source grid (or compatible with ``interpolate_field_to_particles``). - grid : tuple - Source grid specification passed through to - :func:`interpolate_field_to_particles` (e.g., coordinate arrays and/or - metadata describing the field's native grid). - target_grid : tuple of array-like - Target coordinate arrays ``(x, y, z)``. Each element is a 1D array of - coordinates at which to construct the target mesh. - interpolation_order : int, optional - Interpolation order forwarded to :func:`interpolate_field_to_particles`. - Default is 1 (typically linear). - - Returns - ------- - jax.numpy.ndarray - Interpolated field on the target grid with shape - ``(len(x), len(y), len(z))``. - - Notes - ----- - This routine uses ``jax.numpy.meshgrid(..., indexing="ij")`` so the output axis - ordering follows ``(x, y, z)``. - """ - - - x, y, z = target_grid - # Unpack target grid coordinates - - X_target, Y_target, Z_target = jnp.meshgrid(x, y, z, indexing='ij') - # Create a meshgrid for target coordinates - - x_flat = X_target.flatten() - y_flat = Y_target.flatten() - z_flat = Z_target.flatten() - # Flatten target coordinates for interpolation - - interp_flat = interpolate_field_to_particles(field, x_flat, y_flat, z_flat, grid, interpolation_order) - # Interpolate the field values at the flattened target coordinates using the specified interpolation order - - interp_field = interp_flat.reshape(X_target.shape) - # Reshape back to the target grid shape - - return interp_field diff --git a/PyPIC3D/utilities/filters.py b/PyPIC3D/utilities/filters.py new file mode 100644 index 00000000..b944241b --- /dev/null +++ b/PyPIC3D/utilities/filters.py @@ -0,0 +1,155 @@ +import jax.numpy as jnp +from functools import partial +from jax import jit, lax, vmap + + +def apply_valid_3d_conv_last_axes(phi, kernel_3d): + """ + Apply a 3D VALID convolution over the last three axes of phi. + + phi shape: + (..., Nx + 2, Ny + 2, Nz + 2) + + kernel_3d shape: + (3, 3, 3) + + Returns: + filtered interior with shape (..., Nx, Ny, Nz) + """ + leading_shape = phi.shape[:-3] + spatial_shape = phi.shape[-3:] + + # Flatten all leading dimensions into a batch axis. + x = phi.reshape((-1, 1, *spatial_shape)) + + # lax.conv_general_dilated expects: + # lhs: (N, C, D, H, W) + # rhs: (O, I, D, H, W) + # output: (N, C, D, H, W) + kernel = kernel_3d[None, None, :, :, :] + + filtered = lax.conv_general_dilated( + x, + kernel, + window_strides=(1, 1, 1), + padding="VALID", + dimension_numbers=("NCDHW", "OIDHW", "NCDHW"), + ) + + # Remove channel dimension and restore leading dimensions. + filtered = filtered[:, 0] + + return filtered.reshape( + (*leading_shape, spatial_shape[0] - 2, spatial_shape[1] - 2, spatial_shape[2] - 2) + ) + +def _active_slice(num_guard_cells): + g = int(num_guard_cells) + return slice(g, -g) + + +def _stencil_slice(num_guard_cells): + g = int(num_guard_cells) + return slice(g - 1, None if g == 1 else -g + 1) + + +def _is_stacked_vector_field(field): + return hasattr(field, "ndim") and field.ndim >= 4 and int(field.shape[0]) == 3 + + +def _stack_vector_field(field): + if _is_stacked_vector_field(field): + return field + return jnp.stack(field, axis=0) + + +def _restore_vector_field(stacked_field, original_field): + if _is_stacked_vector_field(original_field): + return stacked_field + return stacked_field[0], stacked_field[1], stacked_field[2] + + +@partial(jit, static_argnames=("num_guard_cells",)) +def bilinear_filter(phi, num_guard_cells=1): + """ + Apply a 3D tri-linear smoothing filter to a ghost-celled field. + + The last three axes are spatial. ``num_guard_cells`` selects the physical + interior ``g:-g``; the filter reads the one-cell stencil halo around that + interior and leaves all guard cells unchanged. + """ + k1 = jnp.array([1.0, 2.0, 1.0], dtype=phi.dtype) + + kernel_3d = ( + k1[:, None, None] + * k1[None, :, None] + * k1[None, None, :] + ) / 64.0 + + stencil = _stencil_slice(num_guard_cells) + filtered = apply_valid_3d_conv_last_axes(phi[..., stencil, stencil, stencil], kernel_3d) + + active = _active_slice(num_guard_cells) + + return phi.at[..., active, active, active].set(filtered) + + +@partial(jit, static_argnames=("num_guard_cells",)) +def digital_filter(phi, alpha, num_guard_cells=1): + """ + Apply a 3D nearest-neighbor digital filter to a ghost-celled field. + + The last three axes are spatial. ``num_guard_cells`` selects the physical + interior ``g:-g``; the filter reads the one-cell stencil halo around that + interior and leaves all guard cells unchanged. + """ + neighbor_weight = (1.0 - alpha) / 6.0 + + kernel_3d = jnp.zeros((3, 3, 3), dtype=phi.dtype) + + # Center + kernel_3d = kernel_3d.at[1, 1, 1].set(alpha) + + # Face neighbors + kernel_3d = kernel_3d.at[0, 1, 1].set(neighbor_weight) + kernel_3d = kernel_3d.at[2, 1, 1].set(neighbor_weight) + + kernel_3d = kernel_3d.at[1, 0, 1].set(neighbor_weight) + kernel_3d = kernel_3d.at[1, 2, 1].set(neighbor_weight) + + kernel_3d = kernel_3d.at[1, 1, 0].set(neighbor_weight) + kernel_3d = kernel_3d.at[1, 1, 2].set(neighbor_weight) + + stencil = _stencil_slice(num_guard_cells) + filtered = apply_valid_3d_conv_last_axes(phi[..., stencil, stencil, stencil], kernel_3d) + # Apply the digital filter to the interior of phi using a 3D convolution with the specified kernel. + + active = _active_slice(num_guard_cells) + + return phi.at[..., active, active, active].set(filtered) + + +def bilinear_filter_vector(field, num_guard_cells=1): + """ + Apply the tri-linear filter component-wise to a vector field. + """ + stacked = _stack_vector_field(field) + filtered = vmap( + lambda component: bilinear_filter(component, num_guard_cells=num_guard_cells), + in_axes=0, + out_axes=0, + )(stacked) + return _restore_vector_field(filtered, field) + + +def digital_filter_vector(field, alpha, num_guard_cells=1): + """ + Apply the six-neighbor digital filter component-wise to a vector field. + """ + stacked = _stack_vector_field(field) + filtered = vmap( + lambda component: digital_filter(component, alpha, num_guard_cells=num_guard_cells), + in_axes=0, + out_axes=0, + )(stacked) + return _restore_vector_field(filtered, field) diff --git a/PyPIC3D/utilities/grids.py b/PyPIC3D/utilities/grids.py new file mode 100644 index 00000000..c053f769 --- /dev/null +++ b/PyPIC3D/utilities/grids.py @@ -0,0 +1,137 @@ +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + build_collocated_axis, + build_staggered_axis, +) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def build_collocated_grid(world): + """ + Build the collocated vertex/center grid including one ghost cell per side. + """ + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + x_wind = world["x_wind"] + y_wind = world["y_wind"] + z_wind = world["z_wind"] + Nx = world["Nx"] + Ny = world["Ny"] + Nz = world["Nz"] + + grid = ( + build_collocated_axis(-x_wind / 2, dx, Nx), + build_collocated_axis(-y_wind / 2, dy, Ny), + build_collocated_axis(-z_wind / 2, dz, Nz), + ) + + return grid, grid + + +def build_yee_grid(world): + """ + Build the Yee vertex and center grids including one ghost cell per side. + """ + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + x_wind = world["x_wind"] + y_wind = world["y_wind"] + z_wind = world["z_wind"] + Nx = world["Nx"] + Ny = world["Ny"] + Nz = world["Nz"] + + vertex_grid = ( + build_collocated_axis(-x_wind / 2, dx, Nx), + build_collocated_axis(-y_wind / 2, dy, Ny), + build_collocated_axis(-z_wind / 2, dz, Nz), + ) + center_grid = ( + build_staggered_axis(-x_wind / 2, dx, Nx), + build_staggered_axis(-y_wind / 2, dy, Ny), + build_staggered_axis(-z_wind / 2, dz, Nz), + ) + + return vertex_grid, center_grid + + +def _tile_grid_axis(global_axis_grid, world, tile_shape, tile_counts, axis_index, num_guard_cells): + tile_width = int(tile_shape[axis_index]) + tile_count = int(tile_counts[axis_index]) + g = int(num_guard_cells) + + if axis_index == 0: + d = world["dx"] + elif axis_index == 1: + d = world["dy"] + else: + d = world["dz"] + + offsets = jnp.arange(tile_width + 2 * g, dtype=global_axis_grid.dtype) + tile_indices = jnp.arange(tile_count, dtype=global_axis_grid.dtype) + axis_lines = global_axis_grid[0] + ( + offsets[jnp.newaxis, :] + tile_indices[:, jnp.newaxis] * tile_width - (g - 1) + ) * d + + axis_shape = [1, 1, 1, tile_width + 2 * g] + axis_shape[axis_index] = tile_count + tiled_shape = tuple(tile_counts) + (tile_width + 2 * g,) + + return jnp.broadcast_to(axis_lines.reshape(axis_shape), tiled_shape) + + +def _tile_grid_axes(grid, world, tile_shape, num_guard_cells=2): + """ + Build tile-local coordinate lines for a center or vertex grid. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + Nx = int(grid[0].shape[0]) - 2 + Ny = int(grid[1].shape[0]) - 2 + Nz = int(grid[2].shape[0]) - 2 + tile_counts = ( + _tile_axis_count(Nx, tile_nx), + _tile_axis_count(Ny, tile_ny), + _tile_axis_count(Nz, tile_nz), + ) + + return tuple( + _tile_grid_axis(grid[axis], world, tile_shape, tile_counts, axis, num_guard_cells) + for axis in range(3) + ) + + +def build_tiled_yee_grids(world, tile_shape=None, num_guard_cells=None): + """ + Build the tiled vertex and center grids from the untiled grids in world. + """ + + if tile_shape is None: + tile_shape = tuple(int(width) for width in world["tile_shape"]) + if num_guard_cells is None: + num_guard_cells = int(world["guard_cells"]) + + tiled_vertex_grid = _tile_grid_axes( + world["grids"]["vertex"], + world, + tile_shape, + num_guard_cells=num_guard_cells, + ) + tiled_center_grid = _tile_grid_axes( + world["grids"]["center"], + world, + tile_shape, + num_guard_cells=num_guard_cells, + ) + + return tiled_vertex_grid, tiled_center_grid diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 4b9857df..78c1e933 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -1,7 +1,6 @@ import jax import plotly import tqdm -import pyevtk from jax import jit import argparse import jax.numpy as jnp @@ -13,7 +12,6 @@ from datetime import datetime import importlib.metadata from scipy import stats -from PyPIC3D.boundary_conditions.grid_and_stencil import build_collocated_axis, build_staggered_axis # import external libraries def setup_pmd_files(file_path, name, extension=".bp"): @@ -38,73 +36,6 @@ def wrap_around(ix, size): return jnp.mod(ix, size) -@jit -def bilinear_filter(phi): - """ - Apply a 3D (tri-linear) smoothing filter to a ghost-celled 3D array using a separable - [1, 2, 1]/4 kernel in each dimension. - - The ghost cells provide the boundary information for the convolution. - The filtered result is written back into the interior of the field. - - Args: - phi (jnp.ndarray): 3D field array with shape (Nx+2, Ny+2, Nz+2) including ghost cells. - - Returns: - jnp.ndarray: Field with filtered interior and original ghost cells. - """ - k1 = jnp.array([1.0, 2.0, 1.0], dtype=phi.dtype) / 4.0 # sums to 1 - k3 = k1[:, None, None] * k1[None, :, None] * k1[None, None, :] # (3,3,3), sums to 1 - - kernel = jnp.zeros((3, 3, 3, 1, 1), dtype=phi.dtype) - kernel = kernel.at[:, :, :, 0, 0].set(k3) - - filtered = jax.lax.conv_general_dilated( - phi[jnp.newaxis, ..., jnp.newaxis], - kernel, - window_strides=(1, 1, 1), - padding="VALID", - dimension_numbers=("NDHWC", "DHWIO", "NDHWC"), - feature_group_count=1, - ) - # convolution with VALID padding on (Nx+2, Ny+2, Nz+2) produces (Nx, Ny, Nz) interior - return phi.at[1:-1, 1:-1, 1:-1].set(jnp.squeeze(filtered, axis=(0, 4))) - -@jit -def digital_filter(phi, alpha): - """ - Apply a digital filter to a ghost-celled field. - - The ghost cells provide the boundary information for the convolution. - The filtered result is written back into the interior of the field. - - Args: - phi (ndarray): Field array with shape (Nx+2, Ny+2, Nz+2) including ghost cells. - alpha (float): Filter coefficient. - - Returns: - ndarray: Field with filtered interior and original ghost cells. - """ - neighbor_weight = (1 - alpha) / 6 - kernel = jnp.zeros((3, 3, 3, 1, 1), dtype=phi.dtype) - kernel = kernel.at[1, 1, 1, 0, 0].set(alpha) - kernel = kernel.at[0, 1, 1, 0, 0].set(neighbor_weight) - kernel = kernel.at[2, 1, 1, 0, 0].set(neighbor_weight) - kernel = kernel.at[1, 0, 1, 0, 0].set(neighbor_weight) - kernel = kernel.at[1, 2, 1, 0, 0].set(neighbor_weight) - kernel = kernel.at[1, 1, 0, 0, 0].set(neighbor_weight) - kernel = kernel.at[1, 1, 2, 0, 0].set(neighbor_weight) - - filtered = jax.lax.conv_general_dilated( - phi[jnp.newaxis, ..., jnp.newaxis], - kernel, - window_strides=(1, 1, 1), - padding="VALID", - dimension_numbers=("NDHWC", "DHWIO", "NDHWC"), - ) - # convolution with VALID padding on (Nx+2, Ny+2, Nz+2) produces (Nx, Ny, Nz) interior - return phi.at[1:-1, 1:-1, 1:-1].set(jnp.squeeze(filtered, axis=(0, 4))) - def mae(x, y): """ Calculates the mean absolute error (MAE) between two arrays. @@ -172,12 +103,12 @@ def convergence_test(func): return slope -def compute_energy(particles, E, B, world, constants): +def compute_energy(particles, E, B, world, constants, species_config=None): """ Compute the total energy of the system, including electric field energy, magnetic field energy, and kinetic energy of particles. Args: - particles (list): List of particle species. + particles (TiledParticles): Tile-major particle storage. E (tuple): Electric field components (Ex, Ey, Ez). B (tuple): Magnetic field components (Bx, By, Bz). world (dict): Dictionary containing the simulation world parameters. @@ -209,8 +140,21 @@ def nd_trapezoid(arr, dxs): Ex, Ey, Ez = E Bx, By, Bz = B - # use interior slices to exclude ghost cells from the energy integral - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + # use physical interior slices to exclude ghost cells from the energy + # integral. Tiled fields have leading tile axes followed by local + # ghost-celled Yee arrays. + if Ex.ndim == 6: + g = int(world.get("guard_cells", 2)) + interior = ( + slice(None), + slice(None), + slice(None), + slice(g, -g), + slice(g, -g), + slice(g, -g), + ) + else: + interior = (slice(1, -1), slice(1, -1), slice(1, -1)) dV = dx * dy * dz # calculate the volume element E2_integral = jnp.sum(Ex[interior]**2 + Ey[interior]**2 + Ez[interior]**2) * dV @@ -219,29 +163,43 @@ def nd_trapezoid(arr, dxs): e_energy = 0.5 * constants['eps'] * E2_integral b_energy = 0.5 / constants['mu'] * B2_integral # Electric and magnetic field energy - # kinetic_energy = sum([species.kinetic_energy() for species in particles]) - C = constants['C'] # speed of light - kinetic_energy = 0.0 - for species in particles: - mass = species.get_mass() - vx, vy, vz = species.get_velocity() - v2 = vx**2 + vy**2 + vz**2 - active = species.get_active_mask().astype(v2.dtype) - gamma = 1.0 / jnp.sqrt(1 - v2 / C**2) - momentum2 = jnp.square(mass * gamma ) * v2 - # compute the squared momentum for each particle - KE = jnp.sum(active * (jnp.sqrt(momentum2 * C**2 + mass**2 * C**4) - mass * C**2)) - # compute the kinetic energy for this species - kinetic_energy += KE - # add to total kinetic energy + vx = particles.u[..., 0] + vy = particles.u[..., 1] + vz = particles.u[..., 2] + v2 = vx**2 + vy**2 + vz**2 + + active = particles.active.astype(v2.dtype) + species_mass = species_config.mass * species_config.weight + mass = jnp.broadcast_to( + species_mass.reshape((1, 1, 1, species_mass.shape[0], 1)), + particles.active.shape, + ) + gamma = 1.0 / jnp.sqrt(1 - v2 / C**2) + momentum2 = jnp.square(mass * gamma) * v2 + kinetic_energy = jnp.sum(active * (jnp.sqrt(momentum2 * C**2 + mass**2 * C**4) - mass * C**2)) # Kinetic energy of particles return e_energy, b_energy, kinetic_energy +def compute_total_momentum(particles, species_config=None): + """ + Compute the scalar momentum diagnostic for tiled particles. + """ + + vmag = jnp.sqrt(particles.u[..., 0]**2 + particles.u[..., 1]**2 + particles.u[..., 2]**2) + active = particles.active.astype(vmag.dtype) + species_mass = species_config.mass * species_config.weight + mass = jnp.broadcast_to( + species_mass.reshape((1, 1, 1, species_mass.shape[0], 1)), + particles.active.shape, + ) + return jnp.sum(active * vmag * mass) + + def add_external_fields(E, B, external_fields): """ Add prescribed external fields to the self-consistent fields. @@ -338,30 +296,18 @@ def if_verbose_print(verbose, string): def particle_sanity_check(particles): """ - Perform a sanity check on the particles to ensure consistency in their attributes. - - This function iterates over each species in the particles list and checks that the - number of particles matches the shape of their position and velocity arrays. + Perform a basic shape check for tiled particle storage. Args: - particles (list): A list of species objects, where each species object must have - the following methods: - - get_number_of_particles(): returns the number of particles (int) - - get_position(): returns a tuple of numpy arrays (x, y, z) representing - the positions of the particles - - get_velocity(): returns a tuple of numpy arrays (vx, vy, vz) representing - the velocities of the particles + particles (TiledParticles): Tile-major particle storage. Raises: - AssertionError: If the shapes of the position and velocity arrays do not match the - number of particles. + AssertionError: If the tiled position, velocity, and active arrays are inconsistent. """ - for species in particles: - N = species.get_number_of_particles() - x, y, z = species.get_position() - vx, vy, vz = species.get_velocity() - assert x.shape == y.shape == z.shape == vx.shape == vy.shape == vz.shape == (N,) + assert particles.x.shape == particles.u.shape + assert particles.x.shape[-1] == 3 + assert particles.active.shape == particles.x.shape[:-1] def print_stats(world): @@ -447,28 +393,33 @@ def build_plasma_parameters_dict(world, constants, electrons, dt): Args: world (dict): A dictionary containing the spatial resolution and wind parameters. constants (dict): A dictionary containing physical constants. - electrons (object): An object representing the electrons in the simulation. + electrons (dict): Metadata for the electron species from particle initialization. dt (float): Time step of the simulation. Returns: dict: A dictionary containing the plasma parameters. """ - me = electrons.get_mass() - Te = electrons.get_temperature() + me = electrons["mass"] + Te = electrons["temperature"] + N = electrons["N_particles"] + q = electrons["charge"] + weight = electrons["weight"] kb = constants['kb'] dx, dy, dz = world['dx'], world['dy'], world['dz'] - theoretical_freq = plasma_frequency(electrons, world, constants) - debye = debye_length(electrons, world, constants) + volume = world["x_wind"] * world["y_wind"] * world["z_wind"] + density = weight * N / volume + theoretical_freq = jnp.sqrt(density) * jnp.abs(q) / jnp.sqrt(constants["eps"] * me) + debye = jnp.sqrt(constants["eps"] * kb * Te / (density * q**2)) thermal_velocity = jnp.sqrt(3*kb*Te/me) plasma_parameters = { "Theoretical Plasma Frequency": theoretical_freq, "Debye Length": debye, "Thermal Velocity": thermal_velocity, - "Number of Electrons": electrons.get_number_of_particles(), - "Temperature of Electrons": electrons.get_temperature(), + "Number of Electrons": N, + "Temperature of Electrons": Te, "dx per debye length": debye/dx, "dy per debye length": debye/dy, "dz per debye length": debye/dz, @@ -489,91 +440,6 @@ def convert_to_jax_compatible(data): return tree_map(lambda x: jnp.array(x) if isinstance(x, (int, float, list, tuple)) else x, data) -def build_collocated_grid(world): - """ - Builds a co-allocated grid including ghost cell positions. - - The returned grid has Nx+2 points per axis (1 ghost cell on each side). - - Args: - world (dict): A dictionary containing the following keys: - - 'dx' (float): The grid spacing in the x-direction. - - 'dy' (float): The grid spacing in the y-direction. - - 'dz' (float): The grid spacing in the z-direction. - - 'x_wind' (float): The extent of the grid in the x-direction. - - 'y_wind' (float): The extent of the grid in the y-direction. - - 'z_wind' (float): The extent of the grid in the z-direction. - - Returns: - tuple: A tuple containing two elements: - - grid (tuple): A tuple of three arrays representing the grid points including ghost positions. - - grid (tuple): A duplicate of the first grid tuple. - """ - - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - x_wind = world['x_wind'] - y_wind = world['y_wind'] - z_wind = world['z_wind'] - Nx = world['Nx'] - Ny = world['Ny'] - Nz = world['Nz'] - # get the grid parameters - grid = ( - build_collocated_axis(-x_wind / 2, dx, Nx), - build_collocated_axis(-y_wind / 2, dy, Ny), - build_collocated_axis(-z_wind / 2, dz, Nz), - ) - # create the grid space with ghost cell positions - return grid, grid - -def build_yee_grid(world): - """ - Builds a Yee grid and a staggered Yee grid including ghost cell positions. - - The returned grids have Nx+2 points per axis (1 ghost cell on each side). - The physical interior corresponds to indices [1:-1]. - - Args: - world (dict): A dictionary containing the following keys: - - 'dx' (float): Grid spacing in the x-direction. - - 'dy' (float): Grid spacing in the y-direction. - - 'dz' (float): Grid spacing in the z-direction. - - 'x_wind' (float): Extent of the grid in the x-direction. - - 'y_wind' (float): Extent of the grid in the y-direction. - - 'z_wind' (float): Extent of the grid in the z-direction. - - Returns: - tuple: A tuple containing two elements: - - grid (tuple of jnp.ndarray): The Yee vertex grid with Nx+2 points including ghost positions. - - staggered_grid (tuple of jnp.ndarray): The staggered Yee grid with Nx+2 points including ghost positions. - """ - - dx = world['dx'] - dy = world['dy'] - dz = world['dz'] - x_wind = world['x_wind'] - y_wind = world['y_wind'] - z_wind = world['z_wind'] - Nx = world['Nx'] - Ny = world['Ny'] - Nz = world['Nz'] - # get the grid parameters - - grid = ( - build_collocated_axis(-x_wind / 2, dx, Nx), - build_collocated_axis(-y_wind / 2, dy, Ny), - build_collocated_axis(-z_wind / 2, dz, Nz), - ) - staggered_grid = ( - build_staggered_axis(-x_wind / 2, dx, Nx), - build_staggered_axis(-y_wind / 2, dy, Ny), - build_staggered_axis(-z_wind / 2, dz, Nz), - ) - # create the grid space with ghost cell positions on each side - return grid, staggered_grid - def precondition(NN, phi, rho, model=None): """ Precondition the Poisson equation using a neural network model. @@ -644,7 +510,39 @@ def grab_field_keys(config): field_keys.append(key) return field_keys -def load_external_fields_from_toml(fields, external_fields, config): +def _add_external_field_to_tiled_component(component, external_field, world, field_name): + """ + Add one physical field array into the active interiors of a tiled component. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in world["tile_shape"]] + g = int(world["guard_cells"]) + ntx, nty, ntz = component.shape[:3] + interior_shape = (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + if external_field.shape != interior_shape: + raise ValueError( + f"Shape mismatch for field '{field_name}': external field shape {external_field.shape} " + f"does not match expected interior shape {interior_shape}" + ) + + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = tx * tile_nx + iy = ty * tile_ny + iz = tz * tile_nz + block = external_field[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + component = component.at[ + tx, ty, tz, + g:g + tile_nx, + g:g + tile_ny, + g:g + tile_nz, + ].add(block) + + return component + + +def load_external_fields_from_toml(fields, external_fields, config, world): """ Load external fields from a TOML file. @@ -652,6 +550,7 @@ def load_external_fields_from_toml(fields, external_fields, config): fields (list): Flattened list of evolved E, B, and J field components. external_fields (tuple): External-only E and B field tuples. config (dict): Dictionary containing the configuration values. + world (dict): Simulation world with tiled field geometry. Returns: tuple: Updated evolved fields list and external field tuple. @@ -669,31 +568,39 @@ def load_external_fields_from_toml(fields, external_fields, config): external_field = jnp.load(field_path) - # External fields are (Nx, Ny, Nz), add them to the interior of the ghost-celled arrays - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) if not evolve and (field_type < 0 or field_type > 5): raise ValueError("External-only fields must be electric or magnetic field components with type 0 through 5") - destination = fields[field_type] if evolve else (external_E + external_B)[field_type] - interior_shape = destination[interior].shape - if external_field.shape != interior_shape: - raise ValueError(f"Shape mismatch for field '{field_name}': external field shape {external_field.shape} does not match expected interior shape {interior_shape}") - if evolve: # Evolved fields are part of the self-consistent Maxwell solve. # This is the original behavior and remains the default. - fields[field_type] = fields[field_type].at[interior].add(external_field) + fields[field_type] = _add_external_field_to_tiled_component( + fields[field_type], + external_field, + world, + field_name, + ) else: # External-only E/B fields are invisible to Maxwell's equations. # They are added back only for particle pushes and diagnostics. if field_type < 3: external_E = list(external_E) - external_E[field_type] = external_E[field_type].at[interior].add(external_field) + external_E[field_type] = _add_external_field_to_tiled_component( + external_E[field_type], + external_field, + world, + field_name, + ) external_E = tuple(external_E) else: external_B = list(external_B) b_index = field_type - 3 - external_B[b_index] = external_B[b_index].at[interior].add(external_field) + external_B[b_index] = _add_external_field_to_tiled_component( + external_B[b_index], + external_field, + world, + field_name, + ) external_B = tuple(external_B) print(f"Field loaded successfully: {field_name}") @@ -750,7 +657,7 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para simulation_parameters (dict): Dictionary of simulation parameters. plotting_parameters (dict): Dictionary of plotting parameters. constants (dict): Dictionary of constants. - particles (list): List of particle species. + particles (TiledParticles): Tile-major particle storage. """ output_path = simulation_parameters["output_dir"] @@ -765,22 +672,34 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para "particles": [] } - for particle in particles: - if hasattr(particle, "species_meta") and particle.species_meta: - config["particles"].extend(particle.species_meta) - continue - particle_dict = { - "name": particle.name, - "N_particles": float(particle.N_particles), - "weight": float(particle.weight), - "charge": float(particle.charge), - "mass": float(particle.mass), - "temperature": float(particle.T), - "scaled mass": float(particle.get_mass()), - "scaled charge": float(particle.get_charge()), - "update_pos": particle.update_pos, - "update_v": particle.update_v - } + n_species = particles.active.shape[3] + species_names = simulation_parameters.get("particle_species_names") + species_metadata = simulation_parameters.get("particle_species_metadata") + tile_shape = jax.tree_util.tree_map( + lambda x: x.tolist() if isinstance(x, jnp.ndarray) else x, + simulation_parameters.get("tile_shape", ()), + ) + + for species_index in range(n_species): + if species_names is None: + name = f"species_{species_index}" + else: + name = species_names[species_index] + + active_particles = int(jnp.sum(particles.active[:, :, :, species_index, :])) + if species_metadata is None: + particle_dict = {"name": name} + else: + particle_dict = dict( + jax.tree_util.tree_map( + lambda x: x.tolist() if isinstance(x, jnp.ndarray) else x, + species_metadata[species_index], + ) + ) + + particle_dict["storage"] = "tiled" + particle_dict["active_particles"] = active_particles + particle_dict["tile_shape"] = tile_shape config["particles"].append(particle_dict) config["version"] = { @@ -794,7 +713,6 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para "toml": toml.__version__, "plotly": plotly.__version__, "tqdm": tqdm.__version__, - "pyevtk": pyevtk.__version__, } config["package_versions"] = package_versions @@ -850,8 +768,20 @@ def courant_condition(courant_number, dx, dy, dz, simulation_parameters, constan C = constants['C'] - # Build dxs list with only components that are not 1 - dx_inv = list(1/d for d in (dz, dy, dx) if d != 1) + + Nx = simulation_parameters["Nx"] + Ny = simulation_parameters["Ny"] + Nz = simulation_parameters["Nz"] + # get the number of grid points in each direction + Ns = [Nx, Ny, Nz] + dxs = [dx, dy, dz] + # build a list of the spatial steps in each direction and the number of grid points in each direction + + dx_inv = [] + for d, N in zip(dxs, Ns): + if N > 1: + dx_inv.append(1/d) + # only add the inverse spatial steps for dimensions with more than one grid point dx_inv = sum(dx_inv) # sum all the inverse spatial steps @@ -860,81 +790,3 @@ def courant_condition(courant_number, dx, dy, dz, simulation_parameters, constan # compute the maximum allowed timestep return dt - - -def plasma_frequency(particle_species, world, constants): - """ - Calculate the plasma frequency. - - The plasma frequency is calculated using the properties of the electrons, - the dimensions of the world, and physical constants. - - Args: - particle_species (object): An object representing the particle species - with the following methods: - - get_charge(): returns the charge of the particle. - - get_number_of_particles(): returns the number of particles. - - get_mass(): returns the mass of the particle. - - get_temperature(): returns the temperature of the particle. - world (dict): A dictionary containing the dimensions of the world with keys: - 'x_wind', 'y_wind', and 'z_wind'. - constants (dict): A dictionary containing physical constants with key 'eps'. - - Returns: - float: The calculated plasma frequency. - """ - - - x_wind = world['x_wind'] - y_wind = world['y_wind'] - z_wind = world['z_wind'] - q = particle_species.get_charge() / particle_species.weight - N = particle_species.get_number_of_particles() - m = particle_species.get_mass() / particle_species.weight - eps = constants['eps'] - - # these values are so small that I was having issues calculating - # the plasma frequency with floating point precision so I had to - # break it down into smaller parts - sqrt_dv = jnp.sqrt( x_wind * y_wind * z_wind ) - sqrt_ne = jnp.sqrt( N * particle_species.weight) / sqrt_dv - sqrt_eps = jnp.sqrt( eps ) - sqrt_me = jnp.sqrt( m ) - pf = sqrt_ne * jnp.abs(q) / (sqrt_eps * sqrt_me) - - # pf = jnp.sqrt( N * particle_species.weight * q**2 ) / jnp.sqrt(m) / jnp.sqrt(eps) / jnp.sqrt(x_wind) - - #plasma_frequency = jnp.sqrt(number_pseudoelectrons * weight * charge_electron**2)/jnp.sqrt(mass_electron)/jnp.sqrt(epsilon_0)/jnp.sqrt(length) - - - return pf -# calculate the expected plasma frequency from analytical theory - -def debye_length(particle_species, world, constants): - """ - Calculate the Debye length of a system based on the given parameters. - - Args: - particle_species (object): An object representing the particle species - with the following methods: - - get_charge(): returns the charge of the particle. - - get_number_of_particles(): returns the number of particles. - - get_temperature(): returns the temperature of the particle. - world (dict): A dictionary containing the dimensions of the world with keys: - 'x_wind', 'y_wind', and 'z_wind'. - constants (dict): A dictionary containing physical constants with keys 'eps' and 'kb'. - - Returns: - float: Debye length of the system. - """ - - q = particle_species.get_charge() - N_particles = particle_species.get_number_of_particles() - T = particle_species.get_temperature() - eps = constants['eps'] - kb = constants['kb'] - x_wind = world['x_wind'] - y_wind = world['y_wind'] - z_wind = world['z_wind'] - n = particle_species.weight * N_particles / (x_wind * y_wind * z_wind) - return jnp.sqrt(eps * kb * T / (n * ( q / particle_species.weight )**2)) diff --git a/README.md b/README.md index 6c0fedc3..030d4ba0 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ PyPIC3D --config path/to/config.toml - Advances charged particle species with a Boris particle pusher. - Deposits current with either `j_from_rhov` or `esirkepov`. - Evolves fields with a first-order Yee electrodynamic update or an electrostatic Poisson solve. -- Writes diagnostics, VTK outputs, and optional openPMD files. +- Writes diagnostics, matplotlib phase-space outputs, and optional openPMD files. ## Installation @@ -76,7 +76,7 @@ PyPIC3D/ J.py # current deposition methods rho.py # charge deposition solvers/ # field solvers and operators - diagnostics/ # plotting, VTK, openPMD writers + diagnostics/ # plotting and openPMD writers utils.py # config, IO, helper math/utilities demos/ # runnable example configurations diff --git a/demos/convergence_testing/solver_convergence_test.py b/demos/convergence_testing/solver_convergence_test.py deleted file mode 100644 index 632243a9..00000000 --- a/demos/convergence_testing/solver_convergence_test.py +++ /dev/null @@ -1,514 +0,0 @@ -import jax.numpy as jnp -from scipy import stats - -from PyPIC3D.solvers.fdtd import ( - centered_finite_difference_laplacian, centered_finite_difference_gradient, - centered_finite_difference_divergence, centered_finite_difference_curl -) - -from PyPIC3D.solvers.pstd import ( - spectral_laplacian, spectral_gradient, - spectral_divergence, spectral_curl -) - -from PyPIC3D.utils import ( - convergence_test, mae -) - -def fdtd_laplacian_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical Laplacian of a test function - on a 3D periodic grid. - - The test function is phi(x, y, z) = x^2 + y^2 + z^2, whose Laplacian is analytically 6 everywhere. - The function constructs a cubic grid of size `nx` in each dimension, computes the numerical Laplacian - using a centered finite difference scheme with periodic boundary conditions, and compares it to the - analytical solution (6) in the interior of the domain. - - Parameters - ---------- - nx : int - Number of grid points in each spatial dimension. - - Returns - ------- - error : float - Mean squared error between the numerical and analytical Laplacian in the interior of the grid. - dx : float - Grid spacing in the x-direction (equal to y and z spacing). - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = X**2 + Y**2 + Z**2 - expected_solution = 6 * jnp.ones_like(phi) - # define a analytical phi and its expected solution - - laplacian = centered_finite_difference_laplacian(phi, dx, dy, dz, 'periodic') - # compute the numerical result - - slicer = (slice(1, -1), slice(1, -1), slice(1, -1)) - # slice for the solution comparison - - error = mae( laplacian[slicer], expected_solution[slicer] ) - # compute the mean squared error of the laplacian against the analytical solution - - return error, dx - -def fdtd_gradient_comparison(nx): - """ - Computes the mean squared error between the analytical and numerical gradients - of a 3D scalar field using centered finite differences with periodic boundary conditions. - - The scalar field is defined as: - phi = sin(2πX) + sin(2πY) + sin(2πZ) - over a symmetric cubic domain of size 0.25 in each direction, discretized into `nx` points per axis. - - Parameters - ---------- - nx : int - Number of grid points along each spatial dimension. - - Returns - ------- - error : float - The sum of mean squared errors between the numerical and analytical gradients - in the x, y, and z directions. - dx : float - The spatial grid spacing along each axis. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(2 * jnp.pi * X) + jnp.sin(2 * jnp.pi * Y) + jnp.sin(2 * jnp.pi * Z) - expected_gradx = 2 * jnp.pi * jnp.cos(2 * jnp.pi * X) - expected_grady = 2 * jnp.pi * jnp.cos(2 * jnp.pi * Y) - expected_gradz = 2 * jnp.pi * jnp.cos(2 * jnp.pi * Z) - # build analytical phi with known solution - - gradx, grady, gradz = centered_finite_difference_gradient(phi, dx, dy, dz, 'periodic') - # compute the numerical result - - x_error = mae(gradx, expected_gradx) - y_error = mae(grady, expected_grady) - z_error = mae(gradz, expected_gradz) - error = (x_error + y_error + z_error) / 3 - # compute the mean squared error of the gradient function - - return error, dx - - -def fdtd_divergence_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical divergence - of the vector field F = (X, Y, Z) on a 3D periodic grid. - - The function constructs a symmetric 3D grid of size `nx` in each dimension, - defines the vector field F = (X, Y, Z), and computes its divergence using a - centered finite difference scheme with periodic boundary conditions. The analytical - divergence of F is 3 everywhere. The function returns the mean squared error - between the computed and expected divergence, along with the grid spacing `dx`. - - Args: - nx (int): Number of grid points in each spatial dimension. - - Returns: - tuple: - error (float): Mean squared error between numerical and analytical divergence. - dx (float): Grid spacing in the x-direction. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = X - Fy = Y - Fz = Z - # F = (X, Y, Z), div F = 3 - expected_divergence = 3 * jnp.ones_like(X) - - divergence = centered_finite_difference_divergence(Fx, Fy, Fz, dx, dy, dz, 'periodic') - # compute the numerical divergence - - error = mae(divergence, expected_divergence) - # compute the mean squared error of the divergence - - return error, dx - - -def fdtd_curl_comparison(nx): - """ - Computes the mean squared error between the analytical and numerical curl of a test vector field - on a 3D periodic grid, as a function of grid resolution. - - The test vector field is F = (-Y, X, 0), whose analytical curl is (0, 0, 2). - The numerical curl is computed using a centered finite difference scheme with periodic boundaries. - - Args: - nx (int): Number of grid points in each spatial dimension. - - Returns: - tuple: - error (float): The total mean squared error between the analytical and numerical curl components. - dx (float): The grid spacing in each dimension. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = -1 * Y - Fy = X - Fz = jnp.zeros_like(Z) - # Vector field: F = (-Y, X, 0), curl F = (0, 0, 2) - - expected_curlx = jnp.zeros_like(X) - expected_curly = jnp.zeros_like(Y) - expected_curlz = 2 * jnp.ones_like(Z) - # compute the analytical curl - - curlx, curly, curlz = centered_finite_difference_curl(Fx, Fy, Fz, dx, dy, dz, 'periodic') - # compute the numerical curl - - error_x = mae(curlx, expected_curlx) - error_y = mae(curly, expected_curly) - error_z = mae(curlz, expected_curlz) - error = (error_x + error_y + error_z) / 3 - # compute the error in the numerical curl - - return error, dx - - -def pstd_laplacian_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical Laplacian of a test function - on a 3D periodic grid. - - The test function is phi(x, y, z) = x^2 + y^2 + z^2, whose Laplacian is analytically 6 everywhere. - The function constructs a cubic grid of size `nx` in each dimension, computes the numerical Laplacian - using a centered finite difference scheme with periodic boundary conditions, and compares it to the - analytical solution (6) in the interior of the domain. - - Parameters - ---------- - nx : int - Number of grid points in each spatial dimension. - - Returns - ------- - error : float - Mean squared error between the numerical and analytical Laplacian in the interior of the grid. - dx : float - Grid spacing in the x-direction (equal to y and z spacing). - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(X + Y + Z) - expected_solution = -3 * phi - # define a analytical phi and its expected solution - - laplacian = spectral_laplacian(phi, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - slicer = (slice(10, -10), slice(10, -10), slice(10, -10)) - # slice for the solution comparison - - error = mae( laplacian[slicer], expected_solution[slicer] ) - # compute the mean squared error of the laplacian against the analytical solution - - return error, dx - - -def pstd_gradient_comparison(nx): - """ - Computes the mean squared error (MSE) between the numerical and analytical gradients - of a scalar field using spectral methods. - - The function constructs a 3D meshgrid over a symmetric cubic domain, defines a scalar - field `phi = sin(X + Y + Z)`, and computes its gradient both analytically and numerically. - The numerical gradient is computed using the `spectral_gradient` function. The MSE between - the numerical and analytical gradients is calculated for each component and summed. - - Args: - nx (int): Number of grid points along each spatial dimension. - - Returns: - tuple: - error (float): The total mean squared error between the numerical and analytical gradients. - dx (float): The spatial resolution along each axis. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(X + Y + Z) - expected_gradx = jnp.cos(X + Y + Z) - expected_grady = jnp.cos(X + Y + Z) - expected_gradz = jnp.cos(X + Y + Z) - # define a analytical phi and its expected solution - - gradx, grady, gradz = spectral_gradient(phi, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error_x = mae( gradx, expected_gradx ) - error_y = mae( grady, expected_grady ) - error_z = mae( gradz, expected_gradz ) - error = (error_x + error_y + error_z)/3 - # compute the mean squared error of the gradient against the analytical solution - - return error, dx - - -def pstd_divergence_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical divergence of a vector field. - - This function constructs a 3D meshgrid over a symmetric cubic domain, defines a vector field - with components Fx, Fy, Fz = cos(X + Y + Z), and computes its divergence both analytically and - numerically (using a spectral method). It then returns the mean squared error between the two - divergence calculations, along with the spatial resolution. - - Args: - nx (int): Number of grid points along each axis. - - Returns: - tuple: - error (float): Mean squared error between the numerical and analytical divergence. - dx (float): Spatial resolution along each axis. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = jnp.cos(X + Y + Z) - Fy = jnp.cos(X + Y + Z) - Fz = jnp.cos(X + Y + Z) - expected_div = -3 * jnp.sin(X + Y + Z) - # define a analytical phi and its expected solution - - divF = spectral_divergence(Fx, Fy, Fz, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error = mae( divF, expected_div ) - # compute the mean squared error of the divergence against the analytical solution - - return error, dx - -def pstd_curl_comparison(nx): - """ - Computes the mean squared error (MSE) between the numerical and analytical curl of a vector field - defined on a 3D grid using spectral methods. - - The vector field is defined as F = (sin(Y), sin(Z), sin(X)), and its analytical curl is - (cos(Z) - cos(Y), cos(X) - cos(Z), cos(Y) - cos(X)). The function generates a symmetric 3D grid, - computes the numerical curl using the `spectral_curl` function, and returns the total MSE along - with the grid spacing. - - Args: - nx (int): Number of grid points along each spatial dimension. - - Returns: - tuple: - error (float): The sum of mean squared errors for each component of the curl. - dx (float): The spatial resolution (grid spacing) in the x-direction. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = jnp.sin(Y) - Fy = jnp.sin(Z) - Fz = jnp.sin(X) - expected_curlx = jnp.cos(Z) - jnp.cos(Y) - expected_curly = jnp.cos(X) - jnp.cos(Z) - expected_curlz = jnp.cos(Y) - jnp.cos(X) - # For curl test: F = (sin(Y), sin(Z), sin(X)), curl(F) = (cos(Z)-cos(Y), cos(X)-cos(Z), cos(Y)-cos(X)) - # define a analytical phi and its expected solution - - curlx, curly, curlz = spectral_curl(Fx, Fy, Fz, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error_x = mae( curlx, expected_curlx ) - error_y = mae( curly, expected_curly ) - error_z = mae( curlz, expected_curlz ) - error = (error_x + error_y + error_z) / 3 - # compute the mean squared error of the curl against the analytical solution - - return error, dx - - -if __name__ == "__main__": - - #################### FINITE DIFFERENCE CONVERGENCE TEST ########################## - - print("Convergence Testing of the Finite Difference Methods in PyPIC3D") - - ################### LAPLACIAN CONVERGENCE TEST ############################# - slope = convergence_test(fdtd_laplacian_comparison) - print(f"\nExpected Order of Laplacian Method: 2") - print(f"Calculated Order of Laplacian Method: {slope}") - print(f"Error in Order: {jnp.abs( 100 * (slope - 2) / 2 )} %") - ############################################################################ - - - ################## GRADIENT CONVERGENCE TEST ############################### - slope = convergence_test(fdtd_gradient_comparison) - print(f"\nExpected Order of Gradient Method: 2") - print(f"Calculated Order of Gradient Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - - ################## DIVERGENCE CONVERGENCE TEST ############################# - slope = convergence_test(fdtd_divergence_comparison) - print(f"\nExpected Order of Divergence Method: 2") - print(f"Calculated Order of Divergence Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - - ################# CURL CONVERGENCE TEST #################################### - slope = convergence_test(fdtd_curl_comparison) - print(f"\nExpected Order of Curl Method: 2") - print(f"Calculated Order of Curl Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - print("###################################################################") - ############################################################################ - - - #################### SPECTRAL METHODS CONVERGENCE TEST ########################## - print("\nConvergence Testing of the Spectral Methods in PyPIC3D") - - ################### LAPLACIAN CONVERGENCE TEST ############################# - slope = convergence_test(pstd_laplacian_comparison) - print(f"\nExpected Order of Laplacian Method: 2") - print(f"Calculated Order of Laplacian Method: {slope}") - print(f"Error in Order: {jnp.abs( 100 * (slope - 2) / 2 )} %") - ############################################################################ - - ################## GRADIENT CONVERGENCE TEST ############################### - slope = convergence_test(pstd_gradient_comparison) - print(f"\nExpected Order of Gradient Method: 2") - print(f"Calculated Order of Gradient Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - ################## DIVERGENCE CONVERGENCE TEST ############################# - slope = convergence_test(pstd_divergence_comparison) - print(f"\nExpected Order of Divergence Method: 2") - print(f"Calculated Order of Divergence Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - ################# CURL CONVERGENCE TEST #################################### - slope = convergence_test(pstd_curl_comparison) - print(f"\nExpected Order of Curl Method: 2") - print(f"Calculated Order of Curl Method: {slope}") - print(f"Error in Order: {jnp.abs(100 * (slope - 2) / 2)} %") - ############################################################################ - - ############################################################################ \ No newline at end of file diff --git a/demos/ot_vortex/orszag_tang.toml b/demos/ot_vortex/orszag_tang.toml index fa010cd5..97a1ba61 100644 --- a/demos/ot_vortex/orszag_tang.toml +++ b/demos/ot_vortex/orszag_tang.toml @@ -3,7 +3,7 @@ name = "Demo of a Orszag_TangVortex with 30,000 electrons" t_wind = 0.75e-8 bc = 'periodic' -solver = 'fdtd' +solver = 'electrodynamic_yee' Nx = 100 Ny = 100 Nz = 1 @@ -11,7 +11,6 @@ x_wind = 1 y_wind = 1 z_wind = 1 verbose = false -electrostatic = false cfl = 0.4 shape_factor = 1 diff --git a/demos/reconnection_2d/harris_current.toml b/demos/reconnection_2d/harris_current.toml index 6dcd40a3..6bf9c723 100644 --- a/demos/reconnection_2d/harris_current.toml +++ b/demos/reconnection_2d/harris_current.toml @@ -4,7 +4,7 @@ name = "Electron Positron Reconnection in Harris Current Sheet" Nt = 1500 x_bc = "periodic" z_bc = "conducting" -solver = 'fdtd' +solver = 'electrodynamic_yee' Nx = 500 Ny = 1 Nz = 500 @@ -12,14 +12,12 @@ x_wind = 0.0008 y_wind = 1 z_wind = 0.0008 verbose = false -electrostatic = false cfl = 0.1 shape_factor = 2 relativistic = true [plotting] plotting_interval = 50 -plot_vtk_particles = false plot_phasespace = true [field1] diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 359addce..19d3ad5a 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -2,28 +2,28 @@ [simulation_parameters] name = "Demo of a TS instability with 3,000 electrons" t_wind = 5.413e-8 -solver = 'fdtd' +solver = 'electrodynamic_yee' Nx = 100 Ny = 1 Nz = 1 +particle_tile_nx = 50 +particle_tile_ny = 1 +particle_tile_nz = 1 +particle_tile_capacity_factor=2.0 x_wind = 1 y_wind = 1 z_wind = 1 verbose = false -electrostatic = false cfl = 0.99 shape_factor = 2 current_calculation = "j_from_rhov" -fast_backend = "default" alpha = 1.0 relativistic = true [plotting] plotting_interval = 8 -plot_phasespace = true plot_openpmd_particles = true plot_openpmd_fields = true -dump_particles = true [particle1] name = "electron1" @@ -31,7 +31,7 @@ N_particles = 1500 charge = -1.602e-19 mass = 9.1093837e-31 vth = 14989622.9 -ds_per_debye = 8.0 +number_density = 1.0e15 initial_vx = 0.5e8 [particle2] diff --git a/demos/weibel/weibel.toml b/demos/weibel/weibel.toml index aa17799b..ac733cf2 100644 --- a/demos/weibel/weibel.toml +++ b/demos/weibel/weibel.toml @@ -3,7 +3,7 @@ name = "Demo of a 1D Weibel instability with 3,000 electrons" Nt = 6000 bc = 'periodic' -solver = 'fdtd' +solver = 'electrodynamic_yee' Nx = 201 Ny = 1 Nz = 1 @@ -11,7 +11,6 @@ x_wind = 3e-1 y_wind = 1 z_wind = 1 verbose = false -electrostatic = false cfl = 0.99 shape_factor = 2 relativistic = true @@ -27,7 +26,7 @@ name = "electron1" N_per_cell = 40 charge = -1.602e-19 mass = 9.1093837e-31 -ds_per_debye = 2.0 +number_density = 1.0e15 Tx = 296494.82871735515 Tz = 29649482.87173552 Ty = 0 diff --git a/docs/architecture.rst b/docs/architecture.rst index 9a104f11..b79fe750 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -26,7 +26,7 @@ Core Module Map - ``PyPIC3D/rho.py``: charge deposition. - ``PyPIC3D/solvers/``: field update operators and electrostatic Poisson helpers. -- ``PyPIC3D/diagnostics/``: phase-space plots, VTK output, openPMD output. +- ``PyPIC3D/diagnostics/``: phase-space plots and openPMD output. - ``PyPIC3D/utils.py``: config handling, filters, energy calculations, serialization helpers. @@ -38,18 +38,51 @@ Main runtime objects: - ``particles``: list of particle species objects. - ``fields``: - - electrodynamic/electrostatic: ``(E, B, J, rho, phi)`` - - vector potential path: ``(E, B, J, rho, phi, A2, A1, A0)`` + - electrostatic: ``(E, B, J, rho, phi, external_fields)`` + - electrodynamic: ``(E, B, J, rho, phi, external_fields, pml_state)`` - ``world``: spatial/temporal metadata, grid spacing, grid arrays, encoded field boundary conditions. - ``constants``: physical constants and filter coefficients. +Tiled State Contract +-------------------- + +The ``electrodynamic_yee`` path uses one shared tile shape for fields and particles. +The existing tile-size configuration is interpreted as this common +``tile_shape = (tile_nx, tile_ny, tile_nz)``; there are not separate field-tile +and particle-tile dimensions in the current contract. Each tile width must +divide the corresponding physical grid size exactly, so the leading tile axes +are ``(ntx, nty, ntz)`` with ``ntx * tile_nx = Nx`` and similarly for ``y`` and +``z``. + +Tiled fields store each vector component as compact per-tile arrays with +leading tile axes followed by a one-cell halo around the tile-local physical +interior: +``(ntx, nty, ntz, tile_nx + 2, tile_ny + 2, tile_nz + 2)``. The halo cells +carry neighbor-tile values or exterior field boundary conditions so tiled Yee +curls can be evaluated on the physical interior without assembling a global +field. + +Tiled particles use the same leading tile axes, followed by species and fixed +slot axes. Positions and velocities have shape +``(ntx, nty, ntz, species, max_particles_per_tile, 3)``; the active mask has +shape ``(ntx, nty, ntz, species, max_particles_per_tile)``. Species-level +metadata such as charge, mass, weight, and position/velocity update masks is +stored once in ``SpeciesConfig`` and broadcast over tile slots inside the tiled +kernels. The slot capacity is set when tiled particles are initialized. Empty +slots remain inactive so the array shape stays static during JAX updates. + +Retiling preserves this fixed-capacity layout. If later particle motion would +place more active particles in a tile/species block than its slot capacity can +hold, the tiled refresh reports overflow and the Python driver treats that as a +hard error. This avoids silently dropping active particles. + Data and Output Flow -------------------- - Diagnostics and metadata are written under ``/data``. - Text outputs include energy and momentum traces. -- Optional outputs include VTK files and openPMD files. +- Optional outputs include matplotlib phase-space arrays and openPMD files. - ``output.toml`` captures simulation stats, resolved runtime parameters, particle summaries, and package versions. diff --git a/docs/demos.rst b/docs/demos.rst index c3ee7c64..95a309b0 100644 --- a/docs/demos.rst +++ b/docs/demos.rst @@ -42,7 +42,6 @@ Convergence Scripts .. code-block:: bash python demos/convergence_testing/interpolation_convergence_test.py - python demos/convergence_testing/solver_convergence_test.py Notes ----- diff --git a/docs/index.rst b/docs/index.rst index 5a74d454..1b19e04c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -5,7 +5,7 @@ PyPIC3D Documentation PyPIC3D is a JAX-based particle-in-cell code for electrodynamic and electrostatic plasma simulation. Runs are composed with TOML, launched - from the CLI, and instrumented for diagnostics in VTK and OpenPMD. + from the CLI, and instrumented for diagnostics in matplotlib and openPMD. .. container:: hero-actions @@ -72,4 +72,3 @@ Indices and tables - :ref:`genindex` - :ref:`modindex` - :ref:`search` - diff --git a/docs/particles.rst b/docs/particles.rst index 5ad73f43..c6fcc2c8 100644 --- a/docs/particles.rst +++ b/docs/particles.rst @@ -2,7 +2,8 @@ Particle Species ================ Each ``[particleX]`` TOML section defines one particle species. PyPIC3D supports -multi-species plasmas and per-species boundary behavior. +multi-species plasmas. Particle boundary behavior is a global simulation +setting stored on ``world``. Required Fields --------------- @@ -19,13 +20,15 @@ Example .. code-block:: toml + [simulation_parameters] + shape_factor = 1 + [particle1] name = "electrons" N_particles = 30000 charge = -1.602e-19 mass = 9.1093837e-31 temperature = 293000 - shape_factor = 1 Common Optional Fields ---------------------- @@ -33,7 +36,6 @@ Common Optional Fields - Thermal setup: ``temperature`` or ``vth`` (optionally ``Tx``, ``Ty``, ``Tz``) - Weighting: ``weight`` or ``ds_per_debye`` - Spatial bounds: ``xmin/xmax``, ``ymin/ymax``, ``zmin/zmax`` -- Boundary conditions: ``x_bc``, ``y_bc``, ``z_bc`` - External initial state: ``initial_x/y/z``, ``initial_vx/vy/vz`` (``.npy``) - Update controls: ``update_pos``, ``update_v``, component-level flags @@ -52,14 +54,17 @@ location with sub-cell variation when applicable. Particle Boundary Conditions ---------------------------- -Per-species boundary options: +Global boundary options are set with +``simulation_parameters.particle_x_bc``, +``simulation_parameters.particle_y_bc``, and +``simulation_parameters.particle_z_bc``: - ``periodic`` - ``reflecting`` - ``absorbing`` Absorbing particle boundaries keep the JAX particle arrays fixed-size and mark -particles inactive with a species-local mask after they leave the domain. The +particles inactive with the particle active mask after they leave the domain. The inactive particles remain allocated for JIT shape stability, but they no longer move, push, deposit charge/current, or contribute to particle energy moments. diff --git a/docs/solvers.rst b/docs/solvers.rst index 9a521fe8..cddf6c43 100644 --- a/docs/solvers.rst +++ b/docs/solvers.rst @@ -14,7 +14,7 @@ Set: .. code-block:: toml - electrostatic = true + solver = "electrostatic" Per step, PyPIC3D: @@ -23,10 +23,10 @@ Per step, PyPIC3D: 3. Solves Poisson's equation for ``phi``. 4. Computes ``E = -grad(phi)``. -The ``solver`` key controls both the Poisson solve and gradient operator: +The electrostatic path uses the finite-difference Yee method: -- ``solver = "spectral"``: FFT Poisson + spectral gradient -- ``solver = "fdtd"`` (or any non-``"spectral"`` value): conjugate-gradient Poisson + centered finite-difference gradient +- conjugate-gradient Poisson solve for ``phi`` +- centered finite-difference gradient for ``E = -grad(phi)`` Electrodynamic mode ^^^^^^^^^^^^^^^^^^^ @@ -35,7 +35,7 @@ Set: .. code-block:: toml - electrostatic = false + solver = "electrodynamic_yee" Per step, PyPIC3D: @@ -47,12 +47,6 @@ Per step, PyPIC3D: The electrodynamic update uses first-order Yee-style kernels in ``PyPIC3D.solvers.first_order_yee``. -Vector potential mode -^^^^^^^^^^^^^^^^^^^^^ - -``solver = "vector_potential"`` currently raises ``NotImplementedError`` during -initialization. - Electrodynamic Update Equations ------------------------------- diff --git a/docs/usage.rst b/docs/usage.rst index ee64764a..56949f57 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -33,8 +33,7 @@ Minimal Working Example [simulation_parameters] name = "minimal" - solver = "fdtd" - electrostatic = false + solver = "electrodynamic_yee" relativistic = true current_calculation = "j_from_rhov" # j_from_rhov or esirkepov filter_j = "bilinear" # bilinear, digital, none @@ -50,16 +49,18 @@ Minimal Working Example x_bc = "periodic" # periodic or conducting y_bc = "periodic" z_bc = "periodic" + particle_x_bc = "periodic" # periodic, reflecting, or absorbing + particle_y_bc = "periodic" + particle_z_bc = "periodic" output_dir = "./outputs" [plotting] plotting_interval = 10 plot_phasespace = false - plot_vtk_particles = false - plot_vtk_scalars = false - plot_vtk_vectors = false plot_openpmd_particles = false plot_openpmd_fields = false + openpmd_field_queue_size = 2 + openpmd_particle_queue_size = 2 dump_particles = false dump_fields = false @@ -76,7 +77,6 @@ Minimal Working Example charge = -1.602e-19 mass = 9.1093837e-31 temperature = 1.0 - x_bc = "periodic" # periodic or reflecting Key Notes --------- @@ -90,8 +90,9 @@ Key Notes - Field boundary conditions use ``simulation_parameters.x_bc/y_bc/z_bc`` with values ``periodic`` or ``conducting``. -- Particle boundary conditions are per species (``x_bc/y_bc/z_bc``) with values - ``periodic`` or ``reflecting``. +- Particle boundary conditions are global simulation parameters: + ``particle_x_bc/particle_y_bc/particle_z_bc`` with values ``periodic``, + ``reflecting``, or ``absorbing``. - The smoothing coefficient is ``constants.alpha`` (not ``simulation_parameters.alpha``). @@ -133,5 +134,8 @@ At ``plotting_interval`` cadence, PyPIC3D writes diagnostics like: - ``kinetic_energy.txt`` - ``total_momentum.txt`` -Depending on flags, it also writes VTK and openPMD files plus -``data/output.toml`` metadata. +Depending on flags, it also writes matplotlib phase-space arrays and openPMD +files plus ``data/output.toml`` metadata. Runtime openPMD field and particle +output use bounded queues; ``openpmd_field_queue_size`` and +``openpmd_particle_queue_size`` cap pending batches so output cannot grow memory +without bound. diff --git a/pyproject.toml b/pyproject.toml index 9979253a..928f2ee3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,11 +17,14 @@ dependencies = [ "plotly", "toml", "tqdm", - "pyevtk", - "vtk", "openpmd-api", ] +[project.optional-dependencies] +test = [ + "coverage", +] + [tool.setuptools.packages.find] include = ["PyPIC3D*"] diff --git a/tests/__init__.py b/tests/__init__.py index 316da698..cc755aca 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,16 +1 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - -from . import boris_test -from . import boundaryconditions_test -from . import errors_test -from . import electrostatic_yee_test -from . import fdtd_test -from . import first_order_yee_test -from . import initialization_test -from . import particle_test -from . import pstd_test -from . import utils_test +"""Shared unittest package for PyPIC3D tests.""" diff --git a/tests/boundaryconditions_test.py b/tests/boundaryconditions_test.py deleted file mode 100644 index 0593b536..00000000 --- a/tests/boundaryconditions_test.py +++ /dev/null @@ -1,218 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.boundary_conditions.boundaryconditions import ( - apply_supergaussian_boundary_condition, apply_zero_boundary_condition, - update_ghost_cells, fold_ghost_cells, apply_conducting_bc, - apply_scalar_conducting_bc -) - -jax.config.update("jax_enable_x64", True) - -class TestBoundaryConditions(unittest.TestCase): - - def test_apply_zero_boundary_condition(self): - key = jax.random.PRNGKey(0) - field = jax.random.uniform(key, shape=(10, 10, 10)) - - result = apply_zero_boundary_condition(field) - - # Check if the boundary values are set to zero - self.assertTrue(jnp.all(result[0, :, :] == 0)) - self.assertTrue(jnp.all(result[-1, :, :] == 0)) - self.assertTrue(jnp.all(result[:, 0, :] == 0)) - self.assertTrue(jnp.all(result[:, -1, :] == 0)) - self.assertTrue(jnp.all(result[:, :, 0] == 0)) - self.assertTrue(jnp.all(result[:, :, -1] == 0)) - - -class TestGhostCells(unittest.TestCase): - """Tests for ghost cell functions used in the ghost cell boundary condition approach.""" - - def setUp(self): - # Small field with ghost cells: physical interior is 4x4x4, total 6x6x6 - self.Nx = 4 - self.Ny = 4 - self.Nz = 4 - key = jax.random.PRNGKey(42) - self.field = jax.random.uniform(key, shape=(self.Nx + 2, self.Ny + 2, self.Nz + 2)) - # set ghost cells to known values for testing - self.field = self.field.at[0, :, :].set(0.0) - self.field = self.field.at[-1, :, :].set(0.0) - self.field = self.field.at[:, 0, :].set(0.0) - self.field = self.field.at[:, -1, :].set(0.0) - self.field = self.field.at[:, :, 0].set(0.0) - self.field = self.field.at[:, :, -1].set(0.0) - # ghost cells are zero, interior has random data - - def test_update_ghost_cells_periodic(self): - """Periodic BCs: ghost cells should copy from opposite interior edge.""" - result = update_ghost_cells(self.field, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - - # x-axis: ghost[0] = interior[-2] (rightmost interior), ghost[-1] = interior[1] (leftmost interior) - self.assertTrue(jnp.allclose(result[0, :, :], result[-2, :, :])) - self.assertTrue(jnp.allclose(result[-1, :, :], result[1, :, :])) - - # y-axis: ghost[:,0,:] = interior[:,-2,:], ghost[:,-1,:] = interior[:,1,:] - self.assertTrue(jnp.allclose(result[:, 0, :], result[:, -2, :])) - self.assertTrue(jnp.allclose(result[:, -1, :], result[:, 1, :])) - - # z-axis: ghost[:,:,0] = interior[:,:,-2], ghost[:,:,-1] = interior[:,:,1] - self.assertTrue(jnp.allclose(result[:, :, 0], result[:, :, -2])) - self.assertTrue(jnp.allclose(result[:, :, -1], result[:, :, 1])) - - def test_update_ghost_cells_conducting(self): - """Conducting BCs: ghost cells should be zero.""" - # put non-zero values in ghost cells first - field = self.field.at[0, :, :].set(999.0).at[-1, :, :].set(999.0) - field = field.at[:, 0, :].set(999.0).at[:, -1, :].set(999.0) - field = field.at[:, :, 0].set(999.0).at[:, :, -1].set(999.0) - - result = update_ghost_cells(field, BC_CONDUCTING, BC_CONDUCTING, BC_CONDUCTING) - - self.assertTrue(jnp.all(result[0, :, :] == 0.0)) - self.assertTrue(jnp.all(result[-1, :, :] == 0.0)) - self.assertTrue(jnp.all(result[:, 0, :] == 0.0)) - self.assertTrue(jnp.all(result[:, -1, :] == 0.0)) - self.assertTrue(jnp.all(result[:, :, 0] == 0.0)) - self.assertTrue(jnp.all(result[:, :, -1] == 0.0)) - - def test_update_ghost_cells_mixed(self): - """Mixed BCs: periodic in x, conducting in y and z.""" - result = update_ghost_cells(self.field, BC_PERIODIC, BC_CONDUCTING, BC_CONDUCTING) - - # x-axis should be periodic - self.assertTrue(jnp.allclose(result[0, :, :], result[-2, :, :])) - self.assertTrue(jnp.allclose(result[-1, :, :], result[1, :, :])) - - # y and z ghost cells should be zero - self.assertTrue(jnp.all(result[:, 0, :] == 0.0)) - self.assertTrue(jnp.all(result[:, -1, :] == 0.0)) - self.assertTrue(jnp.all(result[:, :, 0] == 0.0)) - self.assertTrue(jnp.all(result[:, :, -1] == 0.0)) - - def test_update_ghost_cells_preserves_interior(self): - """Ghost cell update should not modify interior values.""" - interior_before = self.field[1:-1, 1:-1, 1:-1].copy() - result = update_ghost_cells(self.field, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - self.assertTrue(jnp.allclose(result[1:-1, 1:-1, 1:-1], interior_before)) - - def test_fold_ghost_cells_periodic(self): - """Periodic fold: ghost cell deposits should be added to opposite interior cells.""" - field = jnp.zeros((6, 6, 6)) - # put deposits in ghost cells - field = field.at[0, 2, 2].set(1.0) # left x ghost - field = field.at[-1, 2, 2].set(2.0) # right x ghost - field = field.at[2, 0, 2].set(3.0) # left y ghost - field = field.at[2, -1, 2].set(4.0) # right y ghost - field = field.at[2, 2, 0].set(5.0) # left z ghost - field = field.at[2, 2, -1].set(6.0) # right z ghost - # put a value in interior - field = field.at[1, 2, 2].set(10.0) - field = field.at[-2, 2, 2].set(20.0) - - result = fold_ghost_cells(field, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - - # ghost[-1] (2.0) wraps to interior[1] (10.0) -> 12.0 - self.assertAlmostEqual(float(result[1, 2, 2]), 12.0) - # ghost[0] (1.0) wraps to interior[-2] (20.0) -> 21.0 - self.assertAlmostEqual(float(result[-2, 2, 2]), 21.0) - # all ghost cells should be cleared after folding - self.assertEqual(float(result[0, 2, 2]), 0.0) - self.assertEqual(float(result[-1, 2, 2]), 0.0) - self.assertEqual(float(result[2, 0, 2]), 0.0) - self.assertEqual(float(result[2, -1, 2]), 0.0) - self.assertEqual(float(result[2, 2, 0]), 0.0) - self.assertEqual(float(result[2, 2, -1]), 0.0) - - def test_fold_ghost_cells_conducting(self): - """Conducting fold: ghost cells should be reflected.""" - field = jnp.zeros((6, 6, 6)) - field = field.at[0, 2, 2].set(1.0) - field = field.at[-1, 2, 2].set(2.0) - field = field.at[1, 2, 2].set(10.0) - field = field.at[-2, 2, 2].set(20.0) - - result = fold_ghost_cells(field, BC_CONDUCTING, BC_CONDUCTING, BC_CONDUCTING) - - # ghost[0] (1.0) reflects to same-side interior[1] (10.0) with sign flip -> 9.0 - self.assertAlmostEqual(float(result[1, 2, 2]), 9.0) - # ghost[-1] (2.0) reflects to same-side interior[-2] (20.0) with sign flip -> 18.0 - self.assertAlmostEqual(float(result[-2, 2, 2]), 18.0) - # ghost cells should be reflected - self.assertEqual(float(result[0, 2, 2]), 0.0) - self.assertEqual(float(result[-1, 2, 2]), 0.0) - - def test_apply_conducting_bc(self): - """Conducting BCs should zero tangential E at boundaries.""" - key = jax.random.PRNGKey(0) - keys = jax.random.split(key, 3) - Ex = jax.random.uniform(keys[0], shape=(6, 6, 6)) - Ey = jax.random.uniform(keys[1], shape=(6, 6, 6)) - Ez = jax.random.uniform(keys[2], shape=(6, 6, 6)) - - Ex_out, Ey_out, Ez_out = apply_conducting_bc((Ex, Ey, Ez), BC_CONDUCTING, BC_CONDUCTING, BC_CONDUCTING) - - # x-boundary conducting: Ey and Ez tangential components zeroed at x faces - self.assertTrue(jnp.all(Ey_out[1, :, :] == 0.0)) - self.assertTrue(jnp.all(Ey_out[-2, :, :] == 0.0)) - self.assertTrue(jnp.all(Ez_out[1, :, :] == 0.0)) - self.assertTrue(jnp.all(Ez_out[-2, :, :] == 0.0)) - - # y-boundary conducting: Ex and Ez tangential components zeroed at y faces - self.assertTrue(jnp.all(Ex_out[:, 1, :] == 0.0)) - self.assertTrue(jnp.all(Ex_out[:, -2, :] == 0.0)) - self.assertTrue(jnp.all(Ez_out[:, 1, :] == 0.0)) - self.assertTrue(jnp.all(Ez_out[:, -2, :] == 0.0)) - - # z-boundary conducting: Ex and Ey tangential components zeroed at z faces - self.assertTrue(jnp.all(Ex_out[:, :, 1] == 0.0)) - self.assertTrue(jnp.all(Ex_out[:, :, -2] == 0.0)) - self.assertTrue(jnp.all(Ey_out[:, :, 1] == 0.0)) - self.assertTrue(jnp.all(Ey_out[:, :, -2] == 0.0)) - - def test_apply_conducting_bc_periodic_noop(self): - """When all axes are periodic, apply_conducting_bc should not modify E.""" - key = jax.random.PRNGKey(1) - keys = jax.random.split(key, 3) - Ex = jax.random.uniform(keys[0], shape=(6, 6, 6)) - Ey = jax.random.uniform(keys[1], shape=(6, 6, 6)) - Ez = jax.random.uniform(keys[2], shape=(6, 6, 6)) - - Ex_out, Ey_out, Ez_out = apply_conducting_bc((Ex, Ey, Ez), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - - self.assertTrue(jnp.allclose(Ex_out, Ex)) - self.assertTrue(jnp.allclose(Ey_out, Ey)) - self.assertTrue(jnp.allclose(Ez_out, Ez)) - - def test_apply_scalar_conducting_bc(self): - """Scalar conducting BC should zero field at conducting boundary faces.""" - key = jax.random.PRNGKey(2) - field = jax.random.uniform(key, shape=(6, 6, 6)) - - result = apply_scalar_conducting_bc(field, BC_CONDUCTING, BC_CONDUCTING, BC_CONDUCTING) - - # field at first and last interior faces should be zero - self.assertTrue(jnp.all(result[1, :, :] == 0.0)) - self.assertTrue(jnp.all(result[-2, :, :] == 0.0)) - self.assertTrue(jnp.all(result[:, 1, :] == 0.0)) - self.assertTrue(jnp.all(result[:, -2, :] == 0.0)) - self.assertTrue(jnp.all(result[:, :, 1] == 0.0)) - self.assertTrue(jnp.all(result[:, :, -2] == 0.0)) - - def test_apply_scalar_conducting_bc_periodic_noop(self): - """When all axes are periodic, scalar conducting BC should not modify field.""" - key = jax.random.PRNGKey(3) - field = jax.random.uniform(key, shape=(6, 6, 6)) - - result = apply_scalar_conducting_bc(field, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - - self.assertTrue(jnp.allclose(result, field)) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/code_tests/__init__.py b/tests/code_tests/__init__.py new file mode 100644 index 00000000..9d74fbb6 --- /dev/null +++ b/tests/code_tests/__init__.py @@ -0,0 +1 @@ +"""Code-contract tests for PyPIC3D tiled runtime ownership.""" diff --git a/tests/code_tests/direct_deposition_test.py b/tests/code_tests/direct_deposition_test.py new file mode 100644 index 00000000..a40eab2d --- /dev/null +++ b/tests/code_tests/direct_deposition_test.py @@ -0,0 +1,966 @@ +import unittest +from pathlib import Path + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.utilities.filters import ( + digital_filter, + digital_filter_vector, +) +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_tile_communication import refresh_tiled_particle_tiles +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) +# compute the number of tiles along each axis, ensuring that the number of cells is divisible by the number of cells per tile. + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + # get the number of tiles along each axis + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + # reshape the interior of the field into tiles, and then transpose to get the correct order of axes + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + # populate the field tiles with the interior tiles, leaving the guard cells as zeros + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + # update the guard cells of the tiled field using the ghost_cells function + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + # call tile_scalar_field for each component of the vector field and return a tuple of tiled components + + +def _update_ghost_cells(field, bc_x, bc_y, bc_z): + field = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda f: f.at[0, :, :].set(f[-2, :, :]).at[-1, :, :].set(f[1, :, :]), + lambda f: f.at[0, :, :].set(0.0).at[-1, :, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda f: f.at[:, 0, :].set(f[:, -2, :]).at[:, -1, :].set(f[:, 1, :]), + lambda f: f.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda f: f.at[:, :, 0].set(f[:, :, -2]).at[:, :, -1].set(f[:, :, 1]), + lambda f: f.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0), + operand=field, + ) + return field + + +def _fold_ghost_cells(field, bc_x, bc_y, bc_z): + field = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda f: f.at[1, :, :].add(f[-1, :, :]).at[-2, :, :].add(f[0, :, :]), + lambda f: f.at[1, :, :].add(-f[0, :, :]).at[-2, :, :].add(-f[-1, :, :]), + operand=field, + ) + field = field.at[0, :, :].set(0.0).at[-1, :, :].set(0.0) + field = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda f: f.at[:, 1, :].add(f[:, -1, :]).at[:, -2, :].add(f[:, 0, :]), + lambda f: f.at[:, 1, :].add(-f[:, 0, :]).at[:, -2, :].add(-f[:, -1, :]), + operand=field, + ) + field = field.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0) + field = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda f: f.at[:, :, 1].add(f[:, :, -1]).at[:, :, -2].add(f[:, :, 0]), + lambda f: f.at[:, :, 1].add(-f[:, :, 0]).at[:, :, -2].add(-f[:, :, -1]), + operand=field, + ) + field = field.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0) + return field + + +class TestDirectDeposition(unittest.TestCase): + def _build_world(self, Nx=8, Ny=6, Nz=4, dt=0.05, boundary_conditions=None): + x_wind, y_wind, z_wind = 4.0, 3.0, 2.0 + if boundary_conditions is None: + boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} + world = { + "dx": x_wind / Nx, + "dy": y_wind / Ny, + "dz": z_wind / Nz, + "Nx": Nx, + "Ny": Ny, + "Nz": Nz, + "x_wind": x_wind, + "y_wind": y_wind, + "z_wind": z_wind, + "dt": dt, + "shape_factor": 1, + "guard_cells": 1, + "boundary_conditions": boundary_conditions, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "center": center_grid} + return world + + def _empty_J(self, world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + + def _empty_J_tiles(self, world): + tile_shape = tuple(int(width) for width in world["tile_shape"]) + tile_nx, tile_ny, tile_nz = tile_shape + g = int(world["guard_cells"]) + shape = ( + world["Nx"] // tile_nx, + world["Ny"] // tile_ny, + world["Nz"] // tile_nz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ) + return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + # create an empty tiled current density field with the appropriate shape based on the world and tile shape + + def _tile_shape(self, simulation_parameters): + return ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + # get the shape of the particle tiles from the simulation parameters + + def _world_with_tiled_grids(self, world, tile_shape): + g = int(world["guard_cells"]) + world = dict(world) + grids = dict(world["grids"]) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + grids["tiled_vertex_grid"] = tiled_vertex_grid + grids["tiled_center_grid"] = tiled_center_grid + world["grids"] = grids + return world + # create a new world dictionary that includes the tiled grids based on the given tile shape + + def _one_tile_parameters(self, world): + return { + "particle_tile_nx": world["Nx"], + "particle_tile_ny": world["Ny"], + "particle_tile_nz": world["Nz"], + } + # create simulation parameters for a single tile that covers the entire world grid + + def _species_config(self, charges, masses, weights, update_x=None, update_u=None): + n_species = len(charges) + if update_x is None: + update_x = [(True, True, True)] * n_species + if update_u is None: + update_u = [(True, True, True)] * n_species + + return SpeciesConfig( + charge=jnp.asarray(charges, dtype=float), + mass=jnp.asarray(masses, dtype=float), + weight=jnp.asarray(weights, dtype=float), + update_x=jnp.asarray(update_x, dtype=bool), + update_u=jnp.asarray(update_u, dtype=bool), + ) + # create a SpeciesConfig object with the given charges, masses, weights, and optional update flags for position and velocity + + def _empty_tiled_particles(self, world, simulation_parameters, n_species, n_slots): + tile_nx, tile_ny, tile_nz = self._tile_shape(simulation_parameters) + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + shape = (ntx, nty, ntz, n_species, n_slots, 3) + + return TiledParticles( + x=jnp.zeros(shape), + u=jnp.zeros(shape), + active=jnp.zeros(shape[:-1], dtype=bool), + ) + # create an empty TiledParticles object with the appropriate shape based on the world, simulation parameters, number of species, and number of slots + + def _set_tiled_particle(self, particles, tile, species, slot, x, u, active=True): + tx, ty, tz = tile + particles = particles._replace( + x=particles.x.at[tx, ty, tz, species, slot].set(jnp.asarray(x, dtype=float)), + u=particles.u.at[tx, ty, tz, species, slot].set(jnp.asarray(u, dtype=float)), + active=particles.active.at[tx, ty, tz, species, slot].set(active), + ) + return particles + # set the position, velocity, and active status of a specific particle in the TiledParticles object based on the given tile, species, slot, position, velocity, and active flag + + def _particles_from_slots(self, world, simulation_parameters, n_species, n_slots, slots): + particles = self._empty_tiled_particles(world, simulation_parameters, n_species, n_slots) + for tile, species, slot, position, velocity, active in slots: + particles = self._set_tiled_particle( + particles, + tile, + species, + slot, + position, + velocity, + active, + ) + return particles + # create a TiledParticles object from a list of slots, where each slot specifies the tile, species, slot index, position, velocity, and active status of a particle + + def _one_tile_particles_from_tiled(self, particles): + n_species = particles.active.shape[3] + n_slots = ( + particles.active.shape[0] + * particles.active.shape[1] + * particles.active.shape[2] + * particles.active.shape[4] + ) + return TiledParticles( + x=particles.x.transpose(3, 0, 1, 2, 4, 5).reshape(1, 1, 1, n_species, n_slots, 3), + u=particles.u.transpose(3, 0, 1, 2, 4, 5).reshape(1, 1, 1, n_species, n_slots, 3), + active=particles.active.transpose(3, 0, 1, 2, 4).reshape(1, 1, 1, n_species, n_slots), + ) + # convert a TiledParticles object into a single-tile representation by transposing and reshaping the arrays to have a single tile dimension, while preserving the species and slot dimensions + + def _centered_tiled_particles(self, particles, world, simulation_parameters): + """ + Build the tiled particle view expected by direct tiled deposition. + + ``J_from_rhov`` stores particles at the forward position and steps them + back by ``0.5*u*dt`` internally. The tiled deposition kernel is already + centered, so the test view applies that half-step before deposition and + refreshes tile ownership at the centered position. + """ + + particles = particles._replace(x=particles.x - 0.5 * particles.u * world["dt"]) + + centered_particles, overflow = refresh_tiled_particle_tiles( + particles, + world, + self._tile_shape(simulation_parameters), + ) + self.assertFalse(bool(overflow)) + # ensure that no particles have overflowed their tiles after centering + + return centered_particles + + def _assembled_tiled_current(self, particles, species_config, world, simulation_parameters, constants, filter="none"): + tile_shape = self._tile_shape(simulation_parameters) + world = self._world_with_tiled_grids(world, tile_shape) + tiled_particles = self._centered_tiled_particles(particles, world, simulation_parameters) + + J_tiles = J_from_rhov( + tiled_particles, + species_config, + self._empty_J_tiles(world), + constants, + world, + filter=filter, + ) + g = int(world["guard_cells"]) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) + # assemble the tiled current density field into a global field for comparison + + return J_tiles, J_from_tiles + + def _compare_tiled_to_one_tile(self, particles, species_config, world, simulation_parameters, filter="none", alpha=1.0): + constants = {"C": 3.0e8, "alpha": alpha} + J_tiles, J_from_tiles = self._assembled_tiled_current( + particles, species_config, world, simulation_parameters, constants, filter=filter + ) + _, J_reference = self._assembled_tiled_current( + self._one_tile_particles_from_tiled(particles), + species_config, + world, + self._one_tile_parameters(world), + constants, + filter=filter, + ) + + for tile_component in J_tiles: + self.assertEqual(tile_component.ndim, 6) + # ensure that the tiled current density components have 6 dimensions (tile_x, tile_y, tile_z, tile_nx, tile_ny, tile_nz) + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + error = jnp.max(jnp.abs(tiled_component - reference_component)) + self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=5.0e-15, atol=5.0e-15)) + # compare the assembled tiled current density components to the reference components from the single-tile deposition, ensuring they are close within a specified tolerance + + + def test_tiled_direct_deposition_matches_quadratic_with_two_guard_cells(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + world["shape_factor"] = 2 + world["guard_cells"] = 2 + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 3, + "particle_tile_nz": 2, + } + + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=2, + slots=[ + ((0, 0, 0), 0, 0, (-1.55, -1.10, -0.70), (0.18, 0.03, -0.06), True), + ((1, 0, 0), 0, 0, (-0.52, -0.55, -0.04), (-0.11, 0.17, 0.24), True), + ((1, 0, 1), 0, 0, (-0.03, -0.03, 0.03), (0.07, -0.22, 0.11), True), + ((2, 1, 1), 0, 0, (0.49, 0.02, 0.31), (-0.04, 0.19, -0.14), True), + ((2, 1, 1), 0, 1, (0.55, 0.52, 0.49), (0.21, -0.08, 0.05), True), + ((3, 1, 1), 0, 0, (1.45, 1.05, 0.72), (-0.16, 0.12, -0.19), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + # ensure the direct deposition from tiled particles matches the deposition from a single tile representation, using quadratic shape factors and two guard cells + + def test_tiled_direct_deposition_matches_quadratic_saved_style_reduced_axes(self): + world = self._build_world(Nx=20, Ny=1, Nz=1, dt=0.05) + world["shape_factor"] = 2 + world["guard_cells"] = 2 + simulation_parameters = { + "particle_tile_nx": 5, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=3, + slots=[ + ((0, 0, 0), 0, 0, (-1.95, 0.0, 0.0), (0.18, 0.03, -0.06), True), + ((0, 0, 0), 0, 1, (-1.51, 0.0, 0.0), (-0.11, 0.17, 0.24), True), + ((0, 0, 0), 0, 2, (-1.02, 0.0, 0.0), (0.07, -0.22, 0.11), True), + ((1, 0, 0), 0, 0, (-0.48, 0.0, 0.0), (-0.04, 0.19, -0.14), True), + ((2, 0, 0), 0, 0, (0.02, 0.0, 0.0), (0.21, -0.08, 0.05), True), + ((2, 0, 0), 0, 1, (0.47, 0.0, 0.0), (-0.16, 0.12, -0.19), True), + ((3, 0, 0), 0, 0, (1.04, 0.0, 0.0), (0.09, -0.15, 0.16), True), + ((3, 0, 0), 0, 1, (1.88, 0.0, 0.0), (-0.13, 0.05, -0.07), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + # test the direct deposition from tiled particles matches single tiled with reduced dimensions + + def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): + world = self._build_world(Nx=20, Ny=1, Nz=1, dt=0.05) + world["shape_factor"] = 2 + world["guard_cells"] = 2 + simulation_parameters = { + "particle_tile_nx": 5, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=3, + slots=[ + ((0, 0, 0), 0, 0, (-1.95, 0.0, 0.0), (0.18, 0.03, -0.06), True), + ((0, 0, 0), 0, 1, (-1.51, 0.0, 0.0), (-0.11, 0.17, 0.24), True), + ((0, 0, 0), 0, 2, (-1.02, 0.0, 0.0), (0.07, -0.22, 0.11), True), + ((1, 0, 0), 0, 0, (-0.48, 0.0, 0.0), (-0.04, 0.19, -0.14), True), + ((2, 0, 0), 0, 0, (0.02, 0.0, 0.0), (0.21, -0.08, 0.05), True), + ((2, 0, 0), 0, 1, (0.47, 0.0, 0.0), (-0.16, 0.12, -0.19), True), + ((3, 0, 0), 0, 0, (1.04, 0.0, 0.0), (0.09, -0.15, 0.16), True), + ((3, 0, 0), 0, 1, (1.88, 0.0, 0.0), (-0.13, 0.05, -0.07), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters, filter="bilinear") + # test the bilinear filtered direct deposition from tiled particles matches single tiled with reduced dimensions + + def _assert_tile_scalar_field_rebuilds_halos(self, world, tile_shape, num_guard_cells): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + field = jnp.arange(jnp.prod(jnp.asarray(shape)), dtype=jnp.float64).reshape(shape) + field = field.at[0, :, :].set(-1001.0) + field = field.at[-1, :, :].set(-1002.0) + field = field.at[:, 0, :].set(-1003.0) + field = field.at[:, -1, :].set(-1004.0) + field = field.at[:, :, 0].set(-1005.0) + field = field.at[:, :, -1].set(-1006.0) + + tiles = tile_scalar_field(field, world, tile_shape, num_guard_cells=num_guard_cells) + assembled = assemble_tiled_vector_field( + (tiles, tiles, tiles), + world, + tile_shape, + num_guard_cells=num_guard_cells, + )[0] + expected = _update_ghost_cells( + field, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + + self.assertTrue(jnp.allclose(assembled, expected, rtol=1.0e-15, atol=1.0e-15)) + # test that the tiled scalar field correctly rebuilds the halo regions from the interior values, comparing the assembled tiled field to the expected field with updated ghost cells + + def test_tile_scalar_field_one_guard_rebuilds_periodic_halos_from_interiors(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=1) + # test that a tiled scalar field with one guard cell correctly rebuilds periodic halo regions from the interior values + + def test_tile_scalar_field_two_guards_rebuilds_periodic_halos_from_interiors(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=2) + # test that a tiled scalar field with two guard cells correctly rebuilds periodic halo regions from the interior values + + def test_tile_scalar_field_rebuilds_conducting_halos_from_interiors(self): + world = self._build_world( + Nx=8, + Ny=6, + Nz=4, + boundary_conditions={"x": BC_CONDUCTING, "y": BC_PERIODIC, "z": BC_PERIODIC}, + ) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=1) + # test that a tiled scalar field with conducting boundary conditions correctly rebuilds halo regions from the interior values + + def test_tiled_digital_filter_matches_global_digital_filter(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + tile_shape = (2, 3, 2) + alpha = 0.6 + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + base = jnp.arange(jnp.prod(jnp.asarray(shape)), dtype=jnp.float64).reshape(shape) + base = _update_ghost_cells(base, bc_x, bc_y, bc_z) + J = ( + base / 17.0, + -0.5 * base + 0.25, + jnp.sin(base / 11.0), + ) + J = tuple(_update_ghost_cells(component, bc_x, bc_y, bc_z) for component in J) + + J_tiles = tile_vector_field(J, world, tile_shape) + filtered_tiles = digital_filter_vector(J_tiles, alpha, num_guard_cells=1) + filtered_tiles = ghost_cells.update_tiled_vector_ghost_cells(filtered_tiles, world, num_guard_cells=1, tile_shape=tile_shape) + filtered_from_tiles = assemble_tiled_vector_field(filtered_tiles, world, tile_shape) + filtered_reference = tuple( + _update_ghost_cells(digital_filter(component, alpha), bc_x, bc_y, bc_z) + for component in J + ) + + for reference_component, tiled_component in zip(filtered_reference, filtered_from_tiles): + self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=1.0e-15, atol=1.0e-15)) + # test that the digital filter applied to the tiled vector field matches the result of applying the digital filter to the global vector field, ensuring consistency between tiled and global filtering + + def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self): + world = self._build_world(Nx=4, Ny=1, Nz=1) + world = self._world_with_tiled_grids(world, (2, 1, 1)) + tiles = jnp.zeros((2, 1, 1, 4, 3, 3)) + tiles = tiles.at[0, 0, 0, -1, 1, 1].set(2.0) + tiles = tiles.at[1, 0, 0, 0, 1, 1].set(3.0) + + folded = ghost_cells.fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=world["tile_shape"]) + + self.assertEqual(float(folded[1, 0, 0, 1, 1, 1]), 2.0) + self.assertEqual(float(folded[0, 0, 0, -2, 1, 1]), 3.0) + self.assertTrue(jnp.allclose(folded[:, :, :, 0, :, :], 0.0)) + self.assertTrue(jnp.allclose(folded[:, :, :, -1, :, :], 0.0)) + # test that folding tiled ghost cells with periodic boundary conditions correctly adds current deposits to neighboring tiles, and that the guard cells are zeroed out after folding + + def test_fold_tiled_ghost_cells_two_guard_layers_adds_deposits_to_neighbors(self): + world = self._build_world(Nx=8, Ny=4, Nz=4) + num_guard_cells = 2 + world["guard_cells"] = num_guard_cells + world = self._world_with_tiled_grids(world, (4, 4, 4)) + tiles = jnp.zeros((2, 1, 1, 8, 8, 8)) + + tiles = tiles.at[1, 0, 0, 0, 2, 2].set(2.0) + tiles = tiles.at[1, 0, 0, 1, 2, 2].set(3.0) + tiles = tiles.at[0, 0, 0, -2, 2, 2].set(5.0) + tiles = tiles.at[0, 0, 0, -1, 2, 2].set(7.0) + + folded = ghost_cells.fold_tiled_ghost_cells(tiles, world, num_guard_cells, tile_shape=world["tile_shape"]) + + self.assertEqual(float(folded[0, 0, 0, 4, 2, 2]), 2.0) + self.assertEqual(float(folded[0, 0, 0, 5, 2, 2]), 3.0) + self.assertEqual(float(folded[1, 0, 0, 2, 2, 2]), 5.0) + self.assertEqual(float(folded[1, 0, 0, 3, 2, 2]), 7.0) + self.assertTrue(jnp.allclose(folded[:, :, :, :num_guard_cells, :, :], 0.0)) + self.assertTrue(jnp.allclose(folded[:, :, :, -num_guard_cells:, :, :], 0.0)) + # test that folding tiled ghost cells with two guard layers correctly adds deposits to neighboring tiles, and that the guard cells are zeroed out after folding + + def test_fold_tiled_ghost_cells_two_guard_reduced_axis_folds_to_single_active_cell(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + num_guard_cells = 2 + tile_shape = (4, 1, 1) + world["guard_cells"] = num_guard_cells + world = self._world_with_tiled_grids(world, tile_shape) + tiles = jnp.zeros((2, 1, 1, 8, 5, 5)) + + tiles = tiles.at[0, 0, 0, 2, 0, 2].set(1.0) + tiles = tiles.at[0, 0, 0, 2, 1, 2].set(2.0) + tiles = tiles.at[0, 0, 0, 2, 3, 2].set(3.0) + tiles = tiles.at[0, 0, 0, 2, 4, 2].set(4.0) + + folded = ghost_cells.fold_tiled_ghost_cells(tiles, world, num_guard_cells, tile_shape=tile_shape) + + self.assertEqual(float(folded[0, 0, 0, 2, 2, 2]), 10.0) + self.assertTrue(jnp.allclose(folded[:, :, :, :, :num_guard_cells, :], 0.0)) + self.assertTrue(jnp.allclose(folded[:, :, :, :, -num_guard_cells:, :], 0.0)) + # test that folding tiled ghost cells with two guard layers in a reduced axis configuration correctly folds all deposits into a single active cell, and that the guard cells are zeroed out after folding + + def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): + world = self._build_world( + Nx=4, + Ny=1, + Nz=1, + boundary_conditions={"x": BC_CONDUCTING, "y": BC_PERIODIC, "z": BC_PERIODIC}, + ) + world = self._world_with_tiled_grids(world, (2, 1, 1)) + tiles = jnp.zeros((2, 1, 1, 4, 3, 3)) + tiles = tiles.at[0, 0, 0, 0, 1, 1].set(4.0) + tiles = tiles.at[-1, 0, 0, -1, 1, 1].set(7.0) + tiles = tiles.at[0, 0, 0, -1, 1, 1].set(2.0) + tiles = tiles.at[1, 0, 0, 0, 1, 1].set(3.0) + + folded = ghost_cells.fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=world["tile_shape"]) + + self.assertEqual(float(folded[0, 0, 0, 1, 1, 1]), -4.0) + self.assertEqual(float(folded[-1, 0, 0, -2, 1, 1]), -7.0) + self.assertEqual(float(folded[1, 0, 0, 1, 1, 1]), 2.0) + self.assertEqual(float(folded[0, 0, 0, -2, 1, 1]), 3.0) + self.assertTrue(jnp.allclose(folded[:, :, :, 0, :, :], 0.0)) + self.assertTrue(jnp.allclose(folded[:, :, :, -1, :, :], 0.0)) + # test that folding tiled ghost cells with conducting boundary conditions correctly reflects exterior deposits back into the interior, and that the guard cells are zeroed out after folding + + def test_fold_tiled_ghost_cells_matches_global_fold_for_mixed_boundaries(self): + world = self._build_world( + Nx=4, + Ny=4, + Nz=2, + boundary_conditions={"x": BC_PERIODIC, "y": BC_CONDUCTING, "z": BC_PERIODIC}, + ) + tile_shape = (2, 2, 1) + world = self._world_with_tiled_grids(world, tile_shape) + field = jnp.zeros((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) + tiles = jnp.zeros((2, 2, 2, 4, 4, 3)) + + field = field.at[0, 2, 1].set(1.5) + tiles = tiles.at[0, 0, 0, 0, 2, 1].set(1.5) + + field = field.at[-1, 3, 2].set(-0.5) + tiles = tiles.at[-1, 1, 1, -1, 1, 1].set(-0.5) + + field = field.at[2, 0, 1].set(3.0) + tiles = tiles.at[0, 0, 0, 2, 0, 1].set(3.0) + + field = field.at[3, -1, 2].set(-4.0) + tiles = tiles.at[1, -1, 1, 1, -1, 1].set(-4.0) + + field = field.at[2, 2, 0].set(2.0) + tiles = tiles.at[0, 0, 0, 2, 2, 0].set(2.0) + + field = field.at[3, 3, -1].set(5.0) + tiles = tiles.at[1, 1, -1, 1, 1, -1].set(5.0) + + folded_tiles = ghost_cells.fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=tile_shape) + folded_tiles = ghost_cells.update_tiled_ghost_cells(folded_tiles, world, num_guard_cells=1, tile_shape=tile_shape) + folded_from_tiles = assemble_tiled_vector_field((folded_tiles, folded_tiles, folded_tiles), world, tile_shape, num_guard_cells=1)[0] + folded_reference = _update_ghost_cells( + _fold_ghost_cells( + field, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ), + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + + self.assertTrue(jnp.allclose(folded_from_tiles, folded_reference, rtol=1.0e-15, atol=1.0e-15)) + # test that folding tiled ghost cells with mixed boundary conditions matches the result of folding the global field, ensuring consistency between tiled and global folding operations + + def test_tiled_direct_deposition_returns_only_local_current_tiles(self): + world = self._build_world() + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + constants = {"C": 3.0e8, "alpha": 1.0} + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, -1.0, -0.65), (0.2, 0.0, -0.05), True), + ((1, 1, 0), 0, 0, (-0.25, -0.25, -0.15), (-0.1, 0.15, 0.25), True), + ((2, 1, 1), 0, 0, (0.65, 0.35, 0.25), (0.05, -0.2, 0.1), True), + ((3, 2, 1), 0, 0, (1.45, 1.05, 0.75), (0.3, 0.1, -0.15), True), + ], + ) + species_config = self._species_config(charges=[1.0], masses=[1.0], weights=[1.0]) + tiled_particles = self._centered_tiled_particles(particles, world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) + world = self._world_with_tiled_grids(world, tile_shape) + + J_tiles = J_from_rhov( + tiled_particles, + species_config, + self._empty_J_tiles(world), + constants, + world, + filter="none", + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) + + _, J_reference = self._assembled_tiled_current( + self._one_tile_particles_from_tiled(particles), + species_config, + world, + self._one_tile_parameters(world), + constants, + filter="none", + ) + + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=1.0e-15, atol=1.0e-15)) + # test that the direct deposition from tiled particles returns only the local current tiles, and that the assembled tiled current matches the reference current from a single-tile representation + + def test_tiled_direct_deposition_matches_J_from_rhov_for_dummy_species(self): + world = self._build_world() + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=2, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, -1.0, -0.65), (0.2, 0.0, -0.05), True), + ((0, 2, 1), 1, 0, (-1.65, 1.15, 0.35), (-0.1, 0.3, 0.1), True), + ((1, 1, 0), 0, 0, (-0.25, -0.25, -0.15), (-0.1, 0.15, 0.25), True), + ((2, 0, 0), 1, 0, (0.15, -0.75, -0.45), (0.2, -0.05, 0.05), True), + ((2, 1, 1), 0, 0, (0.65, 0.35, 0.25), (0.05, -0.2, 0.1), True), + ((3, 1, 1), 1, 0, (1.75, 0.45, 0.85), (-0.25, 0.15, -0.2), True), + ((3, 2, 1), 0, 0, (1.45, 1.05, 0.75), (0.3, 0.1, -0.15), True), + ], + ) + species_config = self._species_config( + charges=[-1.0, 2.0], + masses=[1.0, 4.0], + weights=[0.5, 0.25], + ) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + # test that the direct deposition from tiled particles with multiple species (including a dummy species) matches the deposition from a single-tile representation, ensuring consistency across species + + def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(self): + world = self._build_world() + world["guard_cells"] = 2 + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + tile_shape = self._tile_shape(simulation_parameters) + world = self._world_with_tiled_grids(world, tile_shape) + constants = {"C": 3.0e8, "alpha": 0.6} + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, -1.0, -0.65), (0.2, 0.0, -0.05), True), + ((1, 1, 0), 0, 0, (-0.25, -0.25, -0.15), (-0.1, 0.15, 0.25), True), + ((2, 1, 1), 0, 0, (0.65, 0.35, 0.25), (0.05, -0.2, 0.1), True), + ((3, 2, 1), 0, 0, (1.45, 1.05, 0.75), (0.3, 0.1, -0.15), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + tiled_particles = self._centered_tiled_particles(particles, world, simulation_parameters) + + J_tiles = J_from_rhov( + tiled_particles, + species_config, + self._empty_J_tiles(world), + constants, + world, + filter="digital", + ) + J_from_tiles = assemble_tiled_vector_field( + J_tiles, + world, + tile_shape, + num_guard_cells=int(world["guard_cells"]), + ) + _, J_reference = self._assembled_tiled_current( + self._one_tile_particles_from_tiled(particles), + species_config, + world, + self._one_tile_parameters(world), + constants, + filter="digital", + ) + + for tile_component in J_tiles: + self.assertEqual(tile_component.ndim, 6) + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=1.0e-15, atol=1.0e-15)) + # test that the public J_from_rhov function correctly dispatches tiled particles to the tile-local current deposition, and that the assembled tiled current matches the reference current from a single-tile representation + + def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): + world = self._build_world() + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + constants = {"C": 3.0e8, "alpha": 0.6} + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, -1.0, -0.65), (0.2, 0.0, -0.05), True), + ((1, 1, 0), 0, 0, (-0.25, -0.25, -0.15), (-0.1, 0.15, 0.25), True), + ((2, 1, 1), 0, 0, (0.65, 0.35, 0.25), (0.05, -0.2, 0.1), True), + ((3, 2, 1), 0, 0, (1.45, 1.05, 0.75), (0.3, 0.1, -0.15), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters, filter="digital", alpha=0.6) + # test that the direct deposition from tiled particles with a digital filter matches the deposition from a single-tile representation, ensuring consistency between tiled and global deposition with filtering + + def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 3, + "particle_tile_nz": 2, + } + constants = {"C": 3.0e8, "alpha": 1.0} + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=2, + slots=[ + ((0, 0, 0), 0, 0, (-1.55, -1.10, -0.70), (0.18, 0.03, -0.06), True), + ((1, 0, 0), 0, 0, (-0.52, -0.55, -0.04), (-0.11, 0.17, 0.24), True), + ((1, 0, 1), 0, 0, (-0.03, -0.03, 0.03), (0.07, -0.22, 0.11), True), + ((2, 1, 1), 0, 0, (0.49, 0.02, 0.31), (-0.04, 0.19, -0.14), True), + ((2, 1, 1), 0, 1, (0.55, 0.52, 0.49), (0.21, -0.08, 0.05), True), + ((3, 1, 1), 0, 0, (1.45, 1.05, 0.72), (-0.16, 0.12, -0.19), True), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters, filter="bilinear") + # test that the direct deposition from tiled particles with a bilinear filter matches the deposition from a single-tile representation, ensuring consistency between tiled and global deposition with bilinear filtering + + def test_tiled_direct_deposition_respects_active_mask(self): + world = self._build_world() + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 3, + "particle_tile_nz": 2, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, -1.0, -0.65), (0.2, 0.0, -0.05), True), + ((1, 1, 0), 0, 0, (-0.25, -0.25, -0.15), (-0.1, 0.15, 0.25), False), + ((2, 1, 1), 0, 0, (0.65, 0.35, 0.25), (0.05, -0.2, 0.1), True), + ((3, 2, 1), 0, 0, (1.45, 1.05, 0.75), (0.3, 0.1, -0.15), False), + ], + ) + species_config = self._species_config(charges=[1.0], masses=[1.0], weights=[1.0]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + # test that the direct deposition from tiled particles respects the active mask, ensuring that only active particles contribute to the current deposition, and that the assembled tiled current matches the reference current from a single-tile representation + + def test_tiled_direct_deposition_periodic_boundary_crossing(self): + world = self._build_world(Nx=10, Ny=6, Nz=4, dt=0.0) + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 1, 1), 0, 0, (-world["x_wind"] / 2 - 0.2 * world["dx"], 0.0, 0.0), (-0.25, -0.2, 0.15), True), + ((4, 1, 1), 0, 0, (world["x_wind"] / 2 + 0.1 * world["dx"], 0.0, 0.0), (0.5, 0.1, 0.0), True), + ], + ) + species_config = self._species_config(charges=[1.0], masses=[1.0], weights=[1.0]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + + def test_tiled_direct_deposition_matches_J_from_rhov_for_conducting_boundaries(self): + world = self._build_world( + Nx=8, + Ny=6, + Nz=4, + dt=0.0, + boundary_conditions={"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING}, + ) + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ( + (0, 0, 1), + 0, + 0, + (-world["x_wind"] / 2 + 0.1 * world["dx"], -world["y_wind"] / 2 + 0.1 * world["dy"], 0.0), + (0.5, 0.1, -0.15), + True, + ), + ( + (3, 1, 0), + 0, + 0, + (world["x_wind"] / 2 - 0.1 * world["dx"], 0.0, -world["z_wind"] / 2 + 0.1 * world["dz"]), + (-0.25, -0.2, 0.35), + True, + ), + ( + (2, 2, 1), + 0, + 0, + (0.0, world["y_wind"] / 2 - 0.1 * world["dy"], world["z_wind"] / 2 - 0.1 * world["dz"]), + (0.15, 0.3, -0.1), + True, + ), + ], + ) + species_config = self._species_config(charges=[1.0], masses=[1.0], weights=[1.0]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + + def test_tiled_direct_deposition_matches_J_from_rhov_for_mixed_boundaries(self): + world = self._build_world( + Nx=8, + Ny=6, + Nz=4, + dt=0.0, + boundary_conditions={"x": BC_PERIODIC, "y": BC_CONDUCTING, "z": BC_PERIODIC}, + ) + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ( + (0, 0, 1), + 0, + 0, + (-world["x_wind"] / 2 - 0.1 * world["dx"], -world["y_wind"] / 2 + 0.1 * world["dy"], 0.0), + (0.2, 0.0, -0.05), + True, + ), + ((1, 1, 0), 0, 0, (-0.5, -0.25, -world["z_wind"] / 2 - 0.1 * world["dz"]), (0.05, -0.2, 0.1), True), + ((2, 1, 1), 0, 0, (0.5, 0.25, world["z_wind"] / 2 + 0.2 * world["dz"]), (0.3, 0.1, -0.15), True), + ( + (3, 2, 1), + 0, + 0, + (world["x_wind"] / 2 + 0.2 * world["dx"], world["y_wind"] / 2 - 0.2 * world["dy"], 0.25), + (-0.1, 0.15, 0.25), + True, + ), + ], + ) + species_config = self._species_config(charges=[-1.0], masses=[1.0], weights=[0.5]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + + def test_tiled_direct_deposition_reduced_dimensions(self): + world = self._build_world(Nx=16, Ny=1, Nz=1, dt=0.02) + simulation_parameters = { + "particle_tile_nx": 4, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + particles = self._particles_from_slots( + world, + simulation_parameters, + n_species=1, + n_slots=1, + slots=[ + ((0, 0, 0), 0, 0, (-1.25, 0.0, 0.0), (0.2, 0.3, -0.05), True), + ((2, 0, 0), 0, 0, (0.15, 0.0, 0.0), (-0.1, 0.15, 0.25), True), + ((3, 0, 0), 0, 0, (1.25, 0.0, 0.0), (0.05, -0.2, 0.1), True), + ], + ) + species_config = self._species_config(charges=[1.0], masses=[1.0], weights=[1.0]) + + self._compare_tiled_to_one_tile(particles, species_config, world, simulation_parameters) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/esirkepov_test.py b/tests/code_tests/esirkepov_test.py new file mode 100644 index 00000000..e8b5948d --- /dev/null +++ b/tests/code_tests/esirkepov_test.py @@ -0,0 +1,944 @@ +import os +import tempfile +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import toml + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.deposition.Esirkepov import ( + Esirkepov_current, +) +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field, fields_for_output +from PyPIC3D.initialization import build_tiled_array, initialize_fields, initialize_simulation +from PyPIC3D.particles.particle_tile_communication import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from PyPIC3D.solvers.first_order_yee import ( + update_E, +) +from PyPIC3D.utilities.grids import build_tiled_yee_grids + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +def _update_ghost_cells(field, bc_x, bc_y, bc_z): + field = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda f: f.at[0, :, :].set(f[-2, :, :]).at[-1, :, :].set(f[1, :, :]), + lambda f: f.at[0, :, :].set(0.0).at[-1, :, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda f: f.at[:, 0, :].set(f[:, -2, :]).at[:, -1, :].set(f[:, 1, :]), + lambda f: f.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda f: f.at[:, :, 0].set(f[:, :, -2]).at[:, :, -1].set(f[:, :, 1]), + lambda f: f.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0), + operand=field, + ) + return field + + +class TestTiledEsirkepovCurrent(unittest.TestCase): + def _build_world( + self, + Nx=8, + Ny=1, + Nz=1, + dt=0.05, + shape_factor=1, + boundary_conditions=None, + particle_boundary_conditions=None, + ): + x_wind = 4.0 if Nx > 1 else 1.0 + y_wind = 4.0 if Ny > 1 else 1.0 + z_wind = 4.0 if Nz > 1 else 1.0 + if boundary_conditions is None: + boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} + if particle_boundary_conditions is None: + particle_boundary_conditions = {"x": 0, "y": 0, "z": 0} + world = { + "dx": x_wind / Nx, + "dy": y_wind / Ny, + "dz": z_wind / Nz, + "Nx": Nx, + "Ny": Ny, + "Nz": Nz, + "x_wind": x_wind, + "y_wind": y_wind, + "z_wind": z_wind, + "dt": dt, + "shape_factor": shape_factor, + "boundary_conditions": boundary_conditions, + "particle_boundary_conditions": particle_boundary_conditions, + "current_deposition": "esirkepov", + "current_filter": "none", + "guard_cells": 1, + } + center_grid = ( + jnp.linspace(-x_wind / 2 - world["dx"] / 2, x_wind / 2 + world["dx"] / 2, Nx + 2), + jnp.linspace(-y_wind / 2 - world["dy"] / 2, y_wind / 2 + world["dy"] / 2, Ny + 2), + jnp.linspace(-z_wind / 2 - world["dz"] / 2, z_wind / 2 + world["dz"] / 2, Nz + 2), + ) + world["grids"] = {"center": center_grid, "vertex": center_grid} + return world + + def _empty_J(self, world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + + def _species_config(self, charge=-1.0, mass=1.0, weight=0.5): + return SpeciesConfig( + charge=jnp.asarray([charge], dtype=float), + mass=jnp.asarray([mass], dtype=float), + weight=jnp.asarray([weight], dtype=float), + update_x=jnp.ones((1, 3), dtype=bool), + update_u=jnp.ones((1, 3), dtype=bool), + ) + + def _tile_index_for_position(self, position, world, tile_shape): + x, y, z = [float(component) for component in position] + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + + ix = int(jnp.floor((x + 0.5 * world["x_wind"]) / world["dx"])) + iy = int(jnp.floor((y + 0.5 * world["y_wind"]) / world["dy"])) + iz = int(jnp.floor((z + 0.5 * world["z_wind"]) / world["dz"])) + + ix = min(max(ix, 0), int(world["Nx"]) - 1) + iy = min(max(iy, 0), int(world["Ny"]) - 1) + iz = min(max(iz, 0), int(world["Nz"]) - 1) + + return ix // tile_nx, iy // tile_ny, iz // tile_nz + + def _empty_tiled_particles(self, world, tile_shape, n_slots): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + shape = (ntx, nty, ntz, 1, n_slots, 3) + + return TiledParticles( + x=jnp.zeros(shape), + u=jnp.zeros(shape), + active=jnp.zeros(shape[:-1], dtype=bool), + ) + + def _set_tiled_particle(self, particles, tile, slot, x, u, active=True): + tx, ty, tz = tile + return particles._replace( + x=particles.x.at[tx, ty, tz, 0, slot].set(jnp.asarray(x, dtype=float)), + u=particles.u.at[tx, ty, tz, 0, slot].set(jnp.asarray(u, dtype=float)), + active=particles.active.at[tx, ty, tz, 0, slot].set(active), + ) + + def _particles_from_arrays(self, world, tile_shape, x, u, active_mask=None): + x = jnp.asarray(x, dtype=float) + u = jnp.asarray(u, dtype=float) + active_mask = jnp.ones(x.shape[0], dtype=bool) if active_mask is None else jnp.asarray(active_mask, dtype=bool) + particles = self._empty_tiled_particles(world, tile_shape, max(1, int(x.shape[0]))) + write_counts = {} + + for particle_index in range(int(x.shape[0])): + tile = self._tile_index_for_position(x[particle_index], world, tile_shape) + slot = write_counts.get(tile, 0) + write_counts[tile] = slot + 1 + particles = self._set_tiled_particle( + particles, + tile, + slot, + x[particle_index], + u[particle_index], + bool(active_mask[particle_index]), + ) + + return particles, self._species_config() + + def _one_tile_particles_from_tiled(self, particles): + n_species = particles.active.shape[3] + n_slots = ( + particles.active.shape[0] + * particles.active.shape[1] + * particles.active.shape[2] + * particles.active.shape[4] + ) + return TiledParticles( + x=particles.x.transpose(3, 0, 1, 2, 4, 5).reshape(1, 1, 1, n_species, n_slots, 3), + u=particles.u.transpose(3, 0, 1, 2, 4, 5).reshape(1, 1, 1, n_species, n_slots, 3), + active=particles.active.transpose(3, 0, 1, 2, 4).reshape(1, 1, 1, n_species, n_slots), + ) + + def _one_dimensional_particles(self, world, x1, tile_shape): + x = jnp.stack((x1, jnp.zeros_like(x1), jnp.zeros_like(x1)), axis=-1) + u = jnp.stack( + ( + jnp.array([0.08, -0.06, 0.05], dtype=float), + jnp.zeros_like(x1), + jnp.zeros_like(x1), + ), + axis=-1, + ) + return self._particles_from_arrays(world, tile_shape, x, u) + + def _tile_shape_for_world(self, world): + return ( + 2 if world["Nx"] > 1 else 1, + 2 if world["Ny"] > 1 else 1, + 2 if world["Nz"] > 1 else 1, + ) + + def _one_tile_shape_for_world(self, world): + return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + + def _assert_tile_scalar_field_rebuilds_halos(self, world, tile_shape, num_guard_cells): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + field = jnp.arange(jnp.prod(jnp.asarray(shape)), dtype=jnp.float64).reshape(shape) + field = field.at[0, :, :].set(-1001.0) + field = field.at[-1, :, :].set(-1002.0) + field = field.at[:, 0, :].set(-1003.0) + field = field.at[:, -1, :].set(-1004.0) + field = field.at[:, :, 0].set(-1005.0) + field = field.at[:, :, -1].set(-1006.0) + + tiles = tile_scalar_field(field, world, tile_shape, num_guard_cells=num_guard_cells) + assembled = assemble_tiled_vector_field( + (tiles, tiles, tiles), + world, + tile_shape, + num_guard_cells=num_guard_cells, + )[0] + expected = _update_ghost_cells( + field, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + + self.assertTrue(jnp.allclose(assembled, expected, rtol=1.0e-15, atol=1.0e-15)) + + def test_tile_scalar_field_one_guard_rebuilds_periodic_halos_from_interiors(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=1) + + def test_tile_scalar_field_two_guards_rebuilds_periodic_halos_from_interiors(self): + world = self._build_world(Nx=8, Ny=6, Nz=4) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=2) + + def test_tile_scalar_field_rebuilds_conducting_halos_from_interiors(self): + world = self._build_world( + Nx=8, + Ny=6, + Nz=4, + boundary_conditions={"x": BC_CONDUCTING, "y": BC_PERIODIC, "z": BC_PERIODIC}, + ) + self._assert_tile_scalar_field_rebuilds_halos(world, (2, 3, 2), num_guard_cells=1) + + def test_source_has_no_legacy_particle_fixture_imports(self): + with open(__file__, "r") as source_file: + source = source_file.read() + + banned_tokens = [ + "tiled_" + "particle_" + "fixtures", + "particle_" + "species", + "to_" + "tiled_" + "particles", + ] + for token in banned_tokens: + self.assertNotIn(token, source) + + def _world_with_tiled_grids(self, world, tile_shape): + g = int(world["guard_cells"]) + world = dict(world) + grids = dict(world["grids"]) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + grids["tiled_vertex_grid"] = tiled_vertex_grid + grids["tiled_center_grid"] = tiled_center_grid + world["grids"] = grids + return world + + def _assembled_esirkepov_current(self, world, tiled_particles, species_config, constants, tile_shape): + world = self._world_with_tiled_grids(world, tile_shape) + g = int(world["guard_cells"]) + _, _, J_template, _, _ = initialize_fields(world) + J_tiles = Esirkepov_current( + tiled_particles, + species_config, + J_template, + constants, + world, + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) + + return J_tiles, J_from_tiles + + def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=None): + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + if tile_shape is None: + tile_shape = self._tile_shape_for_world(world) + + tiled_particles, species_config = self._particles_from_arrays(world, tile_shape, x_old, u) + J_tiles, J_from_tiles = self._assembled_esirkepov_current( + world, + tiled_particles, + species_config, + constants, + tile_shape, + ) + _, J_reference = self._assembled_esirkepov_current( + world, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, + constants, + self._one_tile_shape_for_world(world), + ) + + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=1.0e-12, atol=1.0e-12), + f"max diff {jnp.max(jnp.abs(tiled_component - reference_component))}", + ) + + def _basic_positions_and_velocities(self, world): + dx, dy, dz = world["dx"], world["dy"], world["dz"] + x = jnp.array( + [ + [-0.30 * world["x_wind"] if world["Nx"] > 1 else 0.0, + -0.20 * world["y_wind"] if world["Ny"] > 1 else 0.0, + -0.10 * world["z_wind"] if world["Nz"] > 1 else 0.0], + [0.05 * world["x_wind"] if world["Nx"] > 1 else 0.0, + 0.15 * world["y_wind"] if world["Ny"] > 1 else 0.0, + 0.20 * world["z_wind"] if world["Nz"] > 1 else 0.0], + [0.25 * world["x_wind"] if world["Nx"] > 1 else 0.0, + -0.30 * world["y_wind"] if world["Ny"] > 1 else 0.0, + 0.30 * world["z_wind"] if world["Nz"] > 1 else 0.0], + ] + ) + u = jnp.array( + [ + [0.35 * dx / world["dt"] if world["Nx"] > 1 else 0.0, + -0.20 * dy / world["dt"] if world["Ny"] > 1 else 0.0, + 0.25 * dz / world["dt"] if world["Nz"] > 1 else 0.0], + [-0.15 * dx / world["dt"] if world["Nx"] > 1 else 0.0, + 0.30 * dy / world["dt"] if world["Ny"] > 1 else 0.0, + -0.10 * dz / world["dt"] if world["Nz"] > 1 else 0.0], + [0.10 * dx / world["dt"] if world["Nx"] > 1 else 0.0, + 0.15 * dy / world["dt"] if world["Ny"] > 1 else 0.0, + 0.20 * dz / world["dt"] if world["Nz"] > 1 else 0.0], + ] + ) + return x, u + + def test_initialize_fields_builds_tiled_current_with_startup_guard_cells(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + world["guard_cells"] = 2 + tile_shape = (2, 1, 1) + world["tile_shape"] = tile_shape + + _, _, J_tiles, _, _ = initialize_fields(world) + + self.assertEqual(J_tiles[0].shape, (4, 1, 1, 6, 5, 5)) + self.assertTrue(jnp.allclose(J_tiles[0], 0.0)) + + def test_tiled_esirkepov_honors_one_guard_cell_startup_depth(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05, shape_factor=1) + world["guard_cells"] = 1 + x_old, u = self._basic_positions_and_velocities(world) + + self._assert_tiled_current_matches_reference(world, x_old, u, tile_shape=(2, 1, 1)) + + def test_shared_guard_current_tiles_assemble_to_one_guard_global_current(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + world["guard_cells"] = 2 + tile_shape = (2, 1, 1) + world["tile_shape"] = tile_shape + g = int(world["guard_cells"]) + _, _, J_tiles, _, _ = initialize_fields(world) + Jx, Jy, Jz = J_tiles + Jx = Jx.at[1, 0, 0, 2, 2, 2].set(3.0) + + assembled = assemble_tiled_vector_field((Jx, Jy, Jz), world, tile_shape, num_guard_cells=g) + + self.assertEqual(assembled[0].shape, (10, 3, 3)) + self.assertEqual(float(assembled[0][3, 1, 1]), 3.0) + + def test_output_adapter_uses_world_guard_cells(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + world["guard_cells"] = 2 + tile_shape = (2, 1, 1) + world["tile_shape"] = tile_shape + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + zeros = (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + g = int(world["guard_cells"]) + E_tiles = tile_vector_field(zeros, world, tile_shape, num_guard_cells=g) + B_tiles = tile_vector_field(zeros, world, tile_shape, num_guard_cells=g) + _, _, J_tiles, _, _ = initialize_fields(world) + Jx, Jy, Jz = J_tiles + Jx = Jx.at[1, 0, 0, 2, 2, 2].set(7.0) + rho = jnp.zeros(shape) + phi = jnp.zeros(shape) + fields = (E_tiles, B_tiles, (Jx, Jy, Jz), rho, phi, (E_tiles, B_tiles), None) + + output_fields = fields_for_output(fields, world) + + self.assertEqual(output_fields[2][0].shape, shape) + self.assertEqual(float(output_fields[2][0][3, 1, 1]), 7.0) + self.assertEqual(fields[2][0].shape[-3:], (6, 5, 5)) + + def test_update_E_reads_two_guard_current_interior(self): + world = self._build_world(Nx=4, Ny=1, Nz=1, dt=0.25) + constants = {"C": 1.0, "eps": 2.0} + tile_shape = (2, 1, 1) + world["guard_cells"] = 2 + g = int(world["guard_cells"]) + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + zeros = (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + E_tiles = tile_vector_field(zeros, world, tile_shape, num_guard_cells=g) + B_tiles = tile_vector_field(zeros, world, tile_shape, num_guard_cells=g) + world["tile_shape"] = tile_shape + _, _, J_tiles, _, _ = initialize_fields(world) + Jx, Jy, Jz = J_tiles + Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) + + world["tile_shape"] = tile_shape + E_after, pml_state = update_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants) + + self.assertIsNone(pml_state) + self.assertTrue(jnp.allclose(E_after[0][:, :, :, g:-g, g:-g, g:-g], -0.5)) + + def test_tiled_esirkepov_matches_global_1d_periodic_current(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + world = self._world_with_tiled_grids(world, tile_shape) + x_old = jnp.array([-1.10, -0.10, 1.05]) + tiled_particles, species_config = self._one_dimensional_particles(world, x_old, tile_shape) + + g = int(world["guard_cells"]) + _, _, J_template, _, _ = initialize_fields(world) + J_tiles = Esirkepov_current( + tiled_particles, + species_config, + J_template, + constants, + world, + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) + _, J_reference = self._assembled_esirkepov_current( + world, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, + constants, + self._one_tile_shape_for_world(world), + ) + + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=1.0e-12, atol=1.0e-12), + f"max diff {jnp.max(jnp.abs(tiled_component - reference_component))}", + ) + + def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_current(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05, shape_factor=1) + world["guard_cells"] = 2 + world = self._world_with_tiled_grids(world, (2, 1, 1)) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + x_old = jnp.array([-1.10, -0.10, 1.05]) + tiled_particles, species_config = self._one_dimensional_particles(world, x_old, world["tile_shape"]) + + _, _, J_template, _, _ = initialize_fields(world) + J_tiles = Esirkepov_current( + tiled_particles, + species_config, + J_template, + constants, + world, + ) + J_from_tiles = assemble_tiled_vector_field( + J_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + _, J_reference = self._assembled_esirkepov_current( + world, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, + constants, + self._one_tile_shape_for_world(world), + ) + + for tile_component in J_tiles: + self.assertEqual(tile_component.ndim, 6) + for reference_component, tiled_component in zip(J_reference, J_from_tiles): + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=1.0e-12, atol=1.0e-12), + f"max diff {jnp.max(jnp.abs(tiled_component - reference_component))}", + ) + + def test_public_Esirkepov_current_rejects_tiled_filter_modes(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05, shape_factor=1) + world["guard_cells"] = 2 + tile_shape = (2, 1, 1) + world = self._world_with_tiled_grids(world, tile_shape) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + tiled_particles, species_config = self._one_dimensional_particles( + world, + jnp.array([-1.10, -0.10, 1.05]), + tile_shape, + ) + + with self.assertRaisesRegex(ValueError, "Esirkepov current filtering is not supported"): + _, _, J_template, _, _ = initialize_fields(world) + Esirkepov_current( + tiled_particles, + species_config, + J_template, + constants, + world, + filter="digital", + ) + + def test_tiled_esirkepov_matches_global_current_for_dimensions_and_shapes(self): + cases = [ + (8, 1, 1), + (1, 8, 1), + (1, 1, 8), + (8, 8, 1), + (8, 1, 8), + (1, 8, 8), + (6, 6, 6), + ] + for shape_factor in (1, 2): + for Nx, Ny, Nz in cases: + with self.subTest(shape_factor=shape_factor, shape=(Nx, Ny, Nz)): + world = self._build_world(Nx=Nx, Ny=Ny, Nz=Nz, dt=0.05, shape_factor=shape_factor) + if shape_factor == 2: + world["guard_cells"] = 2 + x_old, u = self._basic_positions_and_velocities(world) + self._assert_tiled_current_matches_reference(world, x_old, u) + + def test_tiled_esirkepov_folds_internal_tile_and_periodic_boundary_crossings(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05) + dx = world["dx"] + dt = world["dt"] + x_old = jnp.array( + [ + [-1.0 * dx, 0.0, 0.0], + [0.5 * world["x_wind"] - 0.2 * dx, 0.0, 0.0], + [-0.5 * world["x_wind"] + 0.2 * dx, 0.0, 0.0], + ] + ) + u = jnp.array( + [ + [0.7 * dx / dt, 0.0, 0.0], + [0.6 * dx / dt, 0.0, 0.0], + [-0.6 * dx / dt, 0.0, 0.0], + ] + ) + + self._assert_tiled_current_matches_reference(world, x_old, u, tile_shape=(2, 1, 1)) + + def test_tiled_esirkepov_matches_global_current_with_particle_boundaries(self): + cases = [ + ( + {"x": BC_CONDUCTING, "y": BC_PERIODIC, "z": BC_PERIODIC}, + {"x": 1, "y": 0, "z": 0}, + jnp.array([[1.75, 0.0, 0.0], [-1.75, 0.0, 0.0]]), + jnp.array([[0.6, 0.0, 0.0], [-0.6, 0.0, 0.0]]), + (8, 1, 1), + ), + ( + {"x": BC_CONDUCTING, "y": BC_PERIODIC, "z": BC_CONDUCTING}, + {"x": 1, "y": 0, "z": 2}, + jnp.array([[1.75, -0.50, 1.75], [-1.75, 0.50, -1.75]]), + jnp.array([[0.6, 0.4, 0.6], [-0.6, -0.4, -0.6]]), + (8, 4, 4), + ), + ] + for field_bc, particle_bc, x_old, u_cells, shape in cases: + with self.subTest(field_bc=field_bc, particle_bc=particle_bc, shape=shape): + world = self._build_world( + Nx=shape[0], + Ny=shape[1], + Nz=shape[2], + dt=0.05, + boundary_conditions=field_bc, + particle_boundary_conditions=particle_bc, + ) + world["guard_cells"] = 2 + u = u_cells.at[:, 0].set(u_cells[:, 0] * world["dx"] / world["dt"]) + u = u.at[:, 1].set(u_cells[:, 1] * world["dy"] / world["dt"]) + u = u.at[:, 2].set(u_cells[:, 2] * world["dz"] / world["dt"]) + self._assert_tiled_current_matches_reference(world, x_old, u) + + def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, dt=0.05, shape_factor=1) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + world = self._world_with_tiled_grids(world, tile_shape) + dx = world["dx"] + x_old = jnp.array( + [ + [-1.0 * dx, 0.0, 0.0], + [0.5 * world["x_wind"] - 0.2 * dx, 0.0, 0.0], + [0.25 * world["x_wind"], 0.0, 0.0], + ] + ) + u = jnp.array( + [ + [0.7 * dx / world["dt"], 0.0, 0.0], + [0.6 * dx / world["dt"], 0.0, 0.0], + [-0.4 * dx / world["dt"], 0.0, 0.0], + ] + ) + tiled_particles, species_config = self._particles_from_arrays(world, tile_shape, x_old, u) + g = int(world["guard_cells"]) + rho_tiles = build_tiled_array(world) + + rho_old = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) + _, _, J_template, _, _ = initialize_fields(world) + J_tiles = Esirkepov_current( + tiled_particles, + species_config, + J_template, + constants, + world, + ) + new_particles = update_tiled_particle_positions(tiled_particles, species_config, world["dt"]) + new_particles, overflow = refresh_tiled_particle_tiles(new_particles, world, tile_shape) + rho_new = compute_rho(new_particles, species_config, rho_tiles, constants, world) + + self.assertFalse(bool(overflow)) + drhodt = (rho_new[:, :, :, g:-g, g:-g, g:-g] - rho_old[:, :, :, g:-g, g:-g, g:-g]) / world["dt"] + dJxdx = (J_tiles[0][:, :, :, g:-g, g:-g, g:-g] - J_tiles[0][:, :, :, g - 1:-g - 1, g:-g, g:-g]) / world["dx"] + continuity = drhodt + dJxdx + scale = jnp.maximum(1.0, jnp.max(jnp.abs(drhodt)) + jnp.max(jnp.abs(dJxdx))) + + self.assertLessEqual(float(jnp.max(jnp.abs(continuity))), float(1.0e-12 * scale)) + + def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + vx_path = os.path.join(tmpdir, "vx.npy") + np.save(x_path, np.array([-1.5, -0.5, 0.5, 1.5])) + np.save(zeros_path, np.zeros(4)) + np.save(vx_path, np.array([0.10, -0.05, 0.07, -0.02])) + + config = { + "simulation_parameters": { + "name": "tiled yee esirkepov init smoke", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 0.01, + "Nt": 1, + "shape_factor": 1, + "guard_cells": 2, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "esirkepov", + "filter_j": "none", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": {"plotting_interval": 1}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 0.5, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": vx_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + config_path = os.path.join(tmpdir, "tiled_yee_esirkepov.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + _loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) + E_tiles, _B_tiles, J_tiles, *_ = fields + + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(E_tiles[0].shape[-3:], (6, 5, 5)) + self.assertEqual(J_tiles[0].shape[-3:], (6, 5, 5)) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertNotIn("current_guard_cells", world) + self.assertEqual(world["current_deposition"], "esirkepov") + self.assertEqual(world["current_filter"], "none") + + def test_initialize_rejects_filtered_esirkepov_current(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + vx_path = os.path.join(tmpdir, "vx.npy") + np.save(x_path, np.array([-1.5, -0.5, 0.5, 1.5])) + np.save(zeros_path, np.zeros(4)) + np.save(vx_path, np.array([0.10, -0.05, 0.07, -0.02])) + + config = { + "simulation_parameters": { + "name": "reject filtered esirkepov", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 0.01, + "Nt": 1, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "esirkepov", + "filter_j": "digital", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": {"plotting_interval": 1}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 0.5, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": vx_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + with self.assertRaisesRegex(ValueError, "Esirkepov current filtering is not supported"): + initialize_simulation(config) + + def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_initial = np.array([-1.5, -0.5, 0.5, 1.5]) + vx_initial = np.array([0.10, -0.05, 0.07, -0.02]) + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + vx_path = os.path.join(tmpdir, "vx.npy") + np.save(x_path, x_initial) + np.save(zeros_path, np.zeros(4)) + np.save(vx_path, vx_initial) + + config = { + "simulation_parameters": { + "name": "tiled yee esirkepov step smoke", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 0.01, + "Nt": 1, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "esirkepov", + "filter_j": "none", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": {"plotting_interval": 1}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 0.5, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": vx_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + config_path = os.path.join(tmpdir, "tiled_yee_esirkepov_step.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, _simulation_parameters, constants, _plotting_parameters, _plasma_parameters, \ + solver, _electrostatic, _verbose, _GPUs, _Nt, relativistic, particle_pusher, species_config = initialize_simulation(toml.load(config_path)) + + particles, fields = loop( + particles, + species_config, + fields, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + active_x = np.asarray(particles.x[..., 0][particles.active]) + expected_x = np.sort(x_initial + vx_initial * float(world["dt"])) + self.assertTrue(np.allclose(np.sort(active_x), expected_x, rtol=1.0e-12, atol=1.0e-12)) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertEqual(fields[2][0].shape[-3:], (6, 5, 5)) + self.assertFalse(bool(fields[-1])) + + def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_initial = np.array([-1.5, -0.5, 0.5, 1.5]) + vx_initial = np.array([0.10, -0.05, 0.07, -0.02]) + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + vx_path = os.path.join(tmpdir, "vx.npy") + np.save(x_path, x_initial) + np.save(zeros_path, np.zeros(4)) + np.save(vx_path, vx_initial) + + config = { + "simulation_parameters": { + "name": "tiled yee esirkepov alias staging", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 0.01, + "Nt": 1, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "esirkepov", + "filter_j": "none", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": {"plotting_interval": 1}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 0.5, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": vx_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + loop, particles, fields, world, _simulation_parameters, constants, _plotting_parameters, _plasma_parameters, \ + solver, _electrostatic, _verbose, _GPUs, _Nt, relativistic, particle_pusher, species_config = initialize_simulation(config) + self.assertEqual(world["current_deposition"], "esirkepov") + particles, fields = loop( + particles, + species_config, + fields, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + reference_particles, reference_species_config = self._particles_from_arrays( + world, + self._one_tile_shape_for_world(world), + jnp.column_stack((jnp.asarray(x_initial), jnp.zeros(4), jnp.zeros(4))), + jnp.column_stack((jnp.asarray(vx_initial), jnp.zeros(4), jnp.zeros(4))), + ) + _, reference_J = self._assembled_esirkepov_current( + world, + reference_particles, + reference_species_config, + constants, + self._one_tile_shape_for_world(world), + ) + tiled_J = assemble_tiled_vector_field( + fields[2], + world, + tuple(int(width) for width in world["tile_shape"]), + num_guard_cells=int(world["guard_cells"]), + ) + + for reference_component, tiled_component in zip(reference_J, tiled_J): + self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/evolve_test.py b/tests/code_tests/evolve_test.py new file mode 100644 index 00000000..951bb8da --- /dev/null +++ b/tests/code_tests/evolve_test.py @@ -0,0 +1,206 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.evolve import ( + time_loop_electrodynamic, + time_loop_electrostatic, +) +from PyPIC3D.initialization import initialize_fields +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + +jax.config.update("jax_enable_x64", True) + + +def add_tiled_grids_to_world(world, tile_shape): + g = int(world["guard_cells"]) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + world["grids"]["tiled_center_grid"] = tiled_center_grid + # make tiled versions of the grids for the parallelized particle pushers to use + + +def one_slot_tiled_particles( + x, + u, + charge=1.0, + mass=1.0, + weight=1.0, + active=True, + update_x=(True, True, True), + update_u=(True, True, True), +): + particles = TiledParticles( + x=jnp.asarray(x, dtype=float).reshape((1, 1, 1, 1, 1, 3)), + u=jnp.asarray(u, dtype=float).reshape((1, 1, 1, 1, 1, 3)), + active=jnp.asarray([[[[[active]]]]], dtype=bool), + ) + # create a single tile of particles with one particle in it, and a single species config + species_config = SpeciesConfig( + charge=jnp.asarray([charge], dtype=float), + mass=jnp.asarray([mass], dtype=float), + weight=jnp.asarray([weight], dtype=float), + update_x=jnp.asarray([update_x], dtype=bool), + update_u=jnp.asarray([update_u], dtype=bool), + ) + return particles, species_config + + +class TestEvolveExternalFields(unittest.TestCase): + def test_public_loop_names_own_tiled_contracts(self): + self.assertEqual(time_loop_electrodynamic.__module__, "PyPIC3D.evolve") + self.assertEqual(time_loop_electrostatic.__module__, "PyPIC3D.evolve") + # test that both these methods exist in the public API and are not private methods + + def test_external_electric_field_pushes_particles_without_evolving_maxwell_fields(self): + world = { + "Nx": 3, + "Ny": 3, + "Nz": 3, + "dx": 1.0 / 3.0, + "dy": 1.0 / 3.0, + "dz": 1.0 / 3.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "current_deposition": "direct", + "current_filter": "none", + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_yee_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + # build world and grids, and set constants for the simulation + + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + world["guard_cells"] = 2 + add_tiled_grids_to_world(world, tile_shape) + E, B, J, phi, rho = initialize_fields(world) + external_E = tuple(jnp.zeros_like(comp) for comp in E) + external_B = tuple(jnp.zeros_like(comp) for comp in B) + external_E = (jnp.ones_like(external_E[0]), external_E[1], external_E[2]) + external_fields = (external_E, external_B) + # build empty fields and set an external electric field in the x-direction + + tiled_particles, species_config = one_slot_tiled_particles( + x=jnp.array([0.0, 0.0, 0.0]), + u=jnp.array([0.0, 0.0, 0.0]), + ) + # create a single particle at the origin with zero velocity + fields = ( + E, + B, + J, + rho, + phi, + external_fields, + None, + jnp.asarray(False), + ) + + tiled_particles, fields = time_loop_electrodynamic( + tiled_particles, + species_config, + fields, + world, + constants, + relativistic=False, + particle_pusher="boris", + ) # advance the simulation for one time step with the external electric field + + E_after, B_after, J_after, rho_after, phi_after, external_after, pml_state, overflow = fields + vx = tiled_particles.u[:, :, :, 0, :, 0][tiled_particles.active[:, :, :, 0, :]] + vy = tiled_particles.u[:, :, :, 0, :, 1][tiled_particles.active[:, :, :, 0, :]] + vz = tiled_particles.u[:, :, :, 0, :, 2][tiled_particles.active[:, :, :, 0, :]] + # unpack the fields and particle velocities after the time step + + self.assertIsNone(pml_state) + # ensure the PML state is None, as we are not using PML in this test + self.assertFalse(bool(overflow)) + # ensure the particle did not overflow the simulation domain + self.assertGreater(float(vx[0]), 0.0) + self.assertTrue(jnp.allclose(vy, 0.0)) + self.assertTrue(jnp.allclose(vz, 0.0)) + # ensure the particle has gained velocity in the x-direction due to the external electric field, and has not gained velocity in the y or z directions + self.assertTrue(jnp.allclose(external_after[0][0], external_fields[0][0])) + # ensure the external electric field has not changed after the time step + + def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): + world = { + "Nx": 3, + "Ny": 1, + "Nz": 1, + "dx": 1.0 / 3.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "current_deposition": "direct", + "current_filter": "none", + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_yee_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + # build world and grids, and set constants for the simulation + + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + world["guard_cells"] = 2 + add_tiled_grids_to_world(world, tile_shape) + E, B, J, phi, rho = initialize_fields(world) + external_fields = ( + tuple(jnp.zeros_like(comp) for comp in E), + tuple(jnp.zeros_like(comp) for comp in B), + ) + # build empty fields and set external fields to zero + + tiled_particles, species_config = one_slot_tiled_particles( + x=jnp.array([0.49, 0.0, 0.0]), + u=jnp.array([0.2, 0.0, 0.0]), + ) + # create a single particle at x=0.49 with a velocity in the x-direction, which will cause it to cross the absorbing boundary at x=0.5 + + fields = ( + E, + B, + J, + rho, + phi, + external_fields, + None, + jnp.asarray(False), + ) + # build the fields tuple with the initialized fields and external fields + + tiled_particles, _ = time_loop_electrodynamic( + tiled_particles, + species_config, + fields, + world, + constants, + relativistic=False, + particle_pusher="boris", + ) + # advance the simulation for one time step, which should cause the particle to cross the absorbing boundary and become inactive + + active = tiled_particles.active[:, :, :, 0, :].reshape(-1) + x = tiled_particles.x[:, :, :, 0, :, 0].reshape(-1) + # reshape the active mask and x positions of the particles to check their status after the time step + + self.assertEqual(x.shape[0], 1) + self.assertTrue(jnp.array_equal(active, jnp.array([False]))) + # ensure that the particle has become inactive after crossing the absorbing boundary + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/fluid_quantities_test.py b/tests/code_tests/fluid_quantities_test.py new file mode 100644 index 00000000..28cd277e --- /dev/null +++ b/tests/code_tests/fluid_quantities_test.py @@ -0,0 +1,383 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.diagnostics.fluid_quantities import ( + compute_mass_density, + compute_pressure_field, + compute_velocity_field, +) +from PyPIC3D.deposition.rho import ( + compute_tiled_mass_density_from_tiled_particles, + compute_tiled_pressure_field_from_tiled_particles, + compute_tiled_velocity_field_from_tiled_particles, +) +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field +from tests.initial_particles import build_tiled_particles, tiled_species +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + # get the number of tiles along each axis + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + # reshape the interior of the field into tiles, and then transpose to get the correct order of axes + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + # populate the field tiles with the interior tiles, leaving the guard cells as zeros + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + # update the guard cells of the tiled field using the ghost_cells function + +class TestTiledFluidQuantities(unittest.TestCase): + def _build_world(self, shape_factor=2): + x_wind, y_wind, z_wind = 4.0, 3.0, 2.0 + world = { + "dx": x_wind / 8, + "dy": y_wind / 6, + "dz": z_wind / 4, + "Nx": 8, + "Ny": 6, + "Nz": 4, + "x_wind": x_wind, + "y_wind": y_wind, + "z_wind": z_wind, + "dt": 0.08, + "shape_factor": shape_factor, + "guard_cells": 2, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "center": center_grid} + return world + + def _world_with_tiled_grids(self, world, tile_shape=None): + if tile_shape is None: + tile_shape = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + g = int(world["guard_cells"]) + world = dict(world) + grids = dict(world["grids"]) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + grids["tiled_vertex_grid"] = tiled_vertex_grid + grids["tiled_center_grid"] = tiled_center_grid + world["grids"] = grids + return world + + def _scalar_tiles(self, world): + return tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + def _single_tile_world(self, world): + return self._world_with_tiled_grids( + self._build_world(shape_factor=world["shape_factor"]), + tile_shape=(world["Nx"], world["Ny"], world["Nz"]), + ) + + def _assembled_mass_density(self, particles, world): + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + rho_tiles = compute_tiled_mass_density_from_tiled_particles( + tiled_particles, + species_config, + self._scalar_tiles(world), + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + return assemble_tiled_scalar_field( + rho_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + def _assembled_velocity_field(self, particles, world, direction): + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + field_tiles = compute_tiled_velocity_field_from_tiled_particles( + tiled_particles, + species_config, + self._scalar_tiles(world), + int(direction), + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + return assemble_tiled_scalar_field( + field_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + def _assembled_pressure_field(self, particles, world, direction): + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + velocity_tiles = compute_tiled_velocity_field_from_tiled_particles( + tiled_particles, + species_config, + self._scalar_tiles(world), + int(direction), + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + pressure_tiles = compute_tiled_pressure_field_from_tiled_particles( + tiled_particles, + species_config, + self._scalar_tiles(world), + velocity_tiles, + int(direction), + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + return assemble_tiled_scalar_field( + pressure_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + def _empty_scalar(self, world): + return jnp.zeros((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) + + def _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + + def _particles(self, world): + electrons = tiled_species( + name="electrons", + N_particles=5, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.75, -0.65, 0.15, 1.75, world["x_wind"] / 2 - 0.03]), + x2=jnp.array([-1.15, -0.45, 0.35, 1.05, 0.0]), + x3=jnp.array([-0.75, -0.20, 0.25, 0.80, 0.0]), + v1=jnp.array([0.2, -0.1, 0.05, 0.3, -1.4]), + v2=jnp.array([0.0, 0.15, -0.2, 0.1, 0.0]), + v3=jnp.array([-0.05, 0.25, 0.1, -0.15, 0.0]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + ions = tiled_species( + name="ions", + N_particles=3, + charge=2.0, + mass=5.0, + weight=0.25, + T=1.0, + x1=jnp.array([-1.25, -0.20, 0.75]), + x2=jnp.array([1.15, -0.75, 0.45]), + x3=jnp.array([0.35, -0.45, 0.85]), + v1=jnp.array([-0.1, 0.2, -0.25]), + v2=jnp.array([0.3, -0.05, 0.15]), + v3=jnp.array([0.1, 0.05, -0.2]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + active_mask=jnp.array([True, False, True]), + ) + return [electrons, ions] + + def test_tiled_particles_require_tile_major_mass_storage(self): + world = self._build_world(shape_factor=2) + particles = self._particles(world) + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + + with self.assertRaisesRegex(ValueError, "tile-major scalar field storage"): + compute_mass_density(tiled_particles, self._empty_scalar(world), world, species_config=species_config) + + def test_tile_major_mass_density_assembles_to_global_mass_density(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world) + particles = self._particles(world) + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + reference_world = self._single_tile_world(world) + mass_reference = self._assembled_mass_density(self._particles(reference_world), reference_world) + mass_tiles = compute_tiled_mass_density_from_tiled_particles( + tiled_particles, + species_config, + self._scalar_tiles(world), + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + mass_from_tiles = assemble_tiled_scalar_field( + mass_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + self.assertTrue( + jnp.allclose( + mass_from_tiles[1:-1, 1:-1, 1:-1], + mass_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world) + particles = self._particles(world) + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + reference_world = self._single_tile_world(world) + mass_reference = self._assembled_mass_density(self._particles(reference_world), reference_world) + rho_tiles = self._scalar_tiles(world) + mass_tiles = compute_mass_density(tiled_particles, rho_tiles, world, species_config=species_config) + mass_from_tiles = assemble_tiled_scalar_field( + mass_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + self.assertEqual(mass_tiles.ndim, 6) + self.assertTrue( + jnp.allclose( + mass_from_tiles[1:-1, 1:-1, 1:-1], + mass_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tile_major_velocity_field_assembles_to_global_velocity_field(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world) + particles = self._particles(world) + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + reference_world = self._single_tile_world(world) + velocity_reference = self._assembled_velocity_field(self._particles(reference_world), reference_world, 0) + velocity_tiles = compute_velocity_field( + tiled_particles, + self._scalar_tiles(world), + 0, + world, + species_config=species_config, + ) + velocity_from_tiles = assemble_tiled_scalar_field( + velocity_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + self.assertTrue( + jnp.allclose( + velocity_from_tiles[1:-1, 1:-1, 1:-1], + velocity_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tile_major_pressure_field_assembles_to_global_pressure_field(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world) + particles = self._particles(world) + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + reference_world = self._single_tile_world(world) + pressure_reference = self._assembled_pressure_field(self._particles(reference_world), reference_world, 0) + field_tiles = self._scalar_tiles(world) + velocity_tiles = compute_velocity_field( + tiled_particles, + field_tiles, + 0, + world, + species_config=species_config, + ) + pressure_tiles = compute_pressure_field( + tiled_particles, + field_tiles, + velocity_tiles, + 0, + world, + species_config=species_config, + ) + pressure_from_tiles = assemble_tiled_scalar_field( + pressure_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + self.assertTrue( + jnp.allclose( + pressure_from_tiles[1:-1, 1:-1, 1:-1], + pressure_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_public_fluid_quantities_reject_flat_particles(self): + world = self._build_world(shape_factor=2) + particles = self._particles(world) + field = self._empty_scalar(world) + + with self.assertRaisesRegex(ValueError, "TiledParticles"): + compute_mass_density(particles, field, world) + with self.assertRaisesRegex(ValueError, "TiledParticles"): + compute_velocity_field(particles, field, 0, world) + with self.assertRaisesRegex(ValueError, "TiledParticles"): + compute_pressure_field(particles, field, field, 0, world) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/ghost_cells_test.py b/tests/code_tests/ghost_cells_test.py new file mode 100644 index 00000000..f3c5511b --- /dev/null +++ b/tests/code_tests/ghost_cells_test.py @@ -0,0 +1,96 @@ +import unittest +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.boundary_conditions import ghost_cells + +jax.config.update("jax_enable_x64", True) + + +class TestGhostCells(unittest.TestCase): + """Tests for the tiled ghost-cell boundary condition approach.""" + + def setUp(self): + self.tile_shape = (2, 2, 2) + self.g = 1 + self.world = { + "tile_shape": self.tile_shape, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + } + + def test_update_tiled_ghost_cells_periodic_refreshes_neighbor_halos(self): + # this tests the communication of ghost cells between two tiles in a periodic domain + + field_tiles = jnp.zeros((2, 1, 1, 4, 4, 4)) + field_tiles = field_tiles.at[0, 0, 0, 1:3, 1:3, 1:3].set(1.0) + field_tiles = field_tiles.at[1, 0, 0, 1:3, 1:3, 1:3].set(2.0) + # create two tiles, one with a value of 1.0 and one with a value of 2.0 constant + # across the tile + + result = ghost_cells.update_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) + # call the update ghost cells method to communicate the ghost cells between the two tiles + + self.assertTrue(jnp.all(result[0, 0, 0, -1, 1:3, 1:3] == 2.0)) + # make sure the ghost cell on the first tile has been updated from 1.0 to 2.0 + self.assertTrue(jnp.all(result[1, 0, 0, 0, 1:3, 1:3] == 1.0)) + # make sure the ghost cell on the second tile has been updated from 2.0 to 1.0 + + def test_fold_tiled_ghost_cells_periodic_adds_to_owner_tile(self): + # this tests the folding of ghost cells back to the owner tile in a periodic domain + # this is used to confirm the current and charge deposition is correct when using ghost cells + + field_tiles = jnp.zeros((2, 1, 1, 4, 4, 4)) + # create two tiles with a shape of (4, 4, 4) and a ghost cell width of 1 + field_tiles = field_tiles.at[0, 0, 0, -1, 2, 2].set(3.0) + # set the ghost cell on the first tile at the right x boundary to a value of 3.0 + field_tiles = field_tiles.at[1, 0, 0, 0, 2, 2].set(5.0) + # set the ghost cell on the second tile at the left x boundary to a value of 5.0 + + result = ghost_cells.fold_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) + # call the fold ghost cells method to add the ghost cell values back to the owner tile + + self.assertAlmostEqual(float(result[1, 0, 0, 1, 2, 2]), 3.0) + # make sure the ghost cell value of 3.0 on the first tile has been added to the owner tile + self.assertAlmostEqual(float(result[0, 0, 0, -2, 2, 2]), 5.0) + # make sure the ghost cell value of 5.0 on the second tile has been added to the owner tile + self.assertEqual(float(result[0, 0, 0, -1, 2, 2]), 0.0) + self.assertEqual(float(result[1, 0, 0, 0, 2, 2]), 0.0) + # make sure the ghost cell values have been reset to 0.0 after folding + + def test_apply_tiled_conducting_bc_zeros_global_tangential_faces(self): + # this tests the application of conducting boundary conditions to a tiled electric field in a periodic domain + + world = { + "tile_shape": self.tile_shape, + "boundary_conditions": {"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING}, + } + # create a world with conducting boundary conditions in all directions + + E = tuple(jnp.ones((1, 1, 1, 4, 4, 4)) for _ in range(3)) + # create a tuple of three electric field components (Ex, Ey, Ez) with shape (1, 1, 1, 4, 4, 4) and all values set to 1.0 + + Ex, Ey, Ez = ghost_cells.apply_tiled_conducting_bc(E, world, self.g) + # call the apply tiled conducting boundary conditions method to zero out the tangential faces of the electric field components + + self.assertTrue(jnp.all(Ey[0, 0, 0, 1, :, :] == 0.0)) + self.assertTrue(jnp.all(Ez[0, 0, 0, 1, :, :] == 0.0)) + self.assertTrue(jnp.all(Ex[0, 0, 0, :, 1, :] == 0.0)) + self.assertTrue(jnp.all(Ez[0, 0, 0, :, 1, :] == 0.0)) + self.assertTrue(jnp.all(Ex[0, 0, 0, :, :, 1] == 0.0)) + self.assertTrue(jnp.all(Ey[0, 0, 0, :, :, 1] == 0.0)) + # make sure the tangential faces of the electric field components have been zeroed out + + def test_apply_tiled_scalar_conducting_bc_periodic_noop(self): + field = jnp.arange(4 * 4 * 4, dtype=float).reshape((1, 1, 1, 4, 4, 4)) + # create a scalar field with shape (1, 1, 1, 4, 4, 4) and values from 0 to 63 + result = ghost_cells.apply_tiled_scalar_conducting_bc(field, self.world, self.g) + # call the apply tiled scalar conducting boundary conditions method + # since the boundary conditions are periodic, this should be a no-op and the result should be equal to the input field + + self.assertTrue(jnp.allclose(result, field)) + # confirm the field has not been modified since the boundary conditions are periodic and should not affect the field + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/grid_and_stencil_test.py b/tests/code_tests/grid_and_stencil_test.py similarity index 84% rename from tests/grid_and_stencil_test.py rename to tests/code_tests/grid_and_stencil_test.py index c31ab49f..d9b4ad2d 100644 --- a/tests/grid_and_stencil_test.py +++ b/tests/code_tests/grid_and_stencil_test.py @@ -30,12 +30,16 @@ def test_wrap_periodic_position(self): wrapped = wrap_periodic_position(x, wind) expected = jnp.array([1.5, -2.0, 1.75, 2.0, -1.75, -1.5]) self.assertTrue(jnp.allclose(wrapped, expected)) + # test that the periodic wrapping works def test_axis_activity_and_inactive_index(self): self.assertTrue(axis_has_active_cells(6, ghost_cells=True)) self.assertFalse(axis_has_active_cells(3, ghost_cells=True)) self.assertEqual(inactive_axis_index(3, ghost_cells=True), 1) self.assertEqual(inactive_axis_index(1, ghost_cells=False), 0) + # test the methods that determine if the axes are active in the simulation. + # inactive axes are for quasi-dimensional simulations. + # i.e. 2D in xy has a inactive z axis. def test_axis_spacing_and_anchor_helpers(self): axis = build_collocated_axis(-1.0, 0.5, 4) @@ -46,6 +50,8 @@ def test_axis_spacing_and_anchor_helpers(self): self.assertAlmostEqual(float(uniform_axis_spacing(axis)), 0.5) self.assertTrue(jnp.array_equal(anchor, jnp.array([3, 4]))) self.assertTrue(jnp.allclose(offset, jnp.array([0.1, 0.1]))) + # test that nearest index and difference between nearest index methods + # are working properly. def test_build_axis_stencil_points_periodic_and_conducting(self): anchor = jnp.array([0, 5]) @@ -58,6 +64,7 @@ def test_build_axis_stencil_points_periodic_and_conducting(self): self.assertTrue(jnp.array_equal(periodic[:, 1], jnp.array([4, 5, 0]))) self.assertTrue(jnp.array_equal(conducting[:, 0], jnp.array([-1, 0, 1]))) self.assertTrue(jnp.array_equal(conducting[:, 1], jnp.array([4, 5, 6]))) + # test that method that builds the local interpolation and deposition stencils. def test_prepare_particle_axis_stencil_wraps_periodic_indices(self): axis = build_collocated_axis(-1.0, 0.5, 4) @@ -74,6 +81,8 @@ def test_prepare_particle_axis_stencil_wraps_periodic_indices(self): self.assertTrue(jnp.array_equal(anchor, jnp.array([5]))) self.assertTrue(jnp.allclose(offset, jnp.array([0.1]))) self.assertTrue(jnp.array_equal(points[:, 0], jnp.array([4, 5, 0]))) + # confirm that method for local interpolation and deposition stencils enforces + # periodic boundaries properly. def test_collapse_axis_stencil_for_inactive_axis(self): points = jnp.array([[0, 1], [1, 1], [2, 1]]) @@ -88,6 +97,7 @@ def test_collapse_axis_stencil_for_inactive_axis(self): self.assertTrue(jnp.array_equal(collapsed_points, jnp.array([[1, 1]]))) self.assertTrue(jnp.allclose(collapsed_weights, jnp.array([[1.0, 1.0]]))) + # test method used to simplify dummy (grid_points=1) axes. def test_build_axis_helpers_include_ghost_cells(self): collocated = build_collocated_axis(-1.0, 0.5, 4) @@ -95,7 +105,7 @@ def test_build_axis_helpers_include_ghost_cells(self): self.assertTrue(jnp.allclose(collocated, jnp.array([-1.5, -1.0, -0.5, 0.0, 0.5, 1.0]))) self.assertTrue(jnp.allclose(staggered, jnp.array([-1.25, -0.75, -0.25, 0.25, 0.75, 1.25]))) - + # test that axis construction methods include the ghost cell points if __name__ == "__main__": unittest.main() diff --git a/tests/code_tests/initialization_test.py b/tests/code_tests/initialization_test.py new file mode 100644 index 00000000..a35e26e4 --- /dev/null +++ b/tests/code_tests/initialization_test.py @@ -0,0 +1,233 @@ +import unittest +import tempfile +import os +import numpy as np +import toml +import jax +import jax.numpy as jnp +from PyPIC3D.initialization import setup_write_dir, default_parameters, initialize_simulation, validate_field_solver +from PyPIC3D.evolve import time_loop_electrodynamic, time_loop_electrostatic +from PyPIC3D.particles.particle_class import TiledParticles + +jax.config.update("jax_enable_x64", True) + +class TestInitializationFunctions(unittest.TestCase): + def setUp(self): + self.plotting_parameters, self.simulation_parameters, self.constants = default_parameters() + self.simulation_parameters['output_dir'] = 'test_output' + self.plotting_parameters['plotfields'] = False + # check the default parameters are set correctly + + def test_setup_write_dir(self): + # Should not raise + setup_write_dir(self.simulation_parameters, self.plotting_parameters) + # check that the output directory is created + + 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.assertEqual(sim["solver"], "electrodynamic_yee") + self.assertNotIn("electrostatic", sim) + self.assertNotIn("fast_backend", sim) + self.assertEqual(sim["particle_x_bc"], "periodic") + self.assertEqual(sim["particle_y_bc"], "periodic") + self.assertEqual(sim["particle_z_bc"], "periodic") + self.assertEqual(sim["guard_cells"], 2) + self.assertNotIn("plot_vtk_particles", plotting) + self.assertNotIn("plot_vtk_scalars", plotting) + self.assertNotIn("plot_vtk_vectors", plotting) + self.assertIn('eps', const) + self.assertIn('plotfields', plotting) + # check that the default parameters contain expected keys + + def test_initialize_simulation_returns_tiled_runtime_for_ordinary_electrodynamic_config(self): + with tempfile.TemporaryDirectory() as tmpdir: + zeros_path = os.path.join(tmpdir, "zeros.npy") + x_path = os.path.join(tmpdir, "x.npy") + np.save(x_path, np.array([-0.375, -0.125, 0.125, 0.375])) + np.save(zeros_path, np.zeros(4)) + config = { + "simulation_parameters": { + "name": "ordinary tiled runtime test", + "output_dir": tmpdir, + "Nx": 4, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "filter_j": "none", + }, + "plotting": {"plotting": False}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 1.0, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": zeros_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + config_path = os.path.join(tmpdir, "global_particle_bc.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) + + self.assertIs(loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(simulation_parameters["solver"], "electrodynamic_yee") + self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + self.assertIn("tiled_center_grid", world["grids"]) + self.assertIn("tiled_vertex_grid", world["grids"]) + E, B, J, rho, phi, external_fields, pml_state, overflow = fields + self.assertEqual(E[0].ndim, 6) + self.assertEqual(B[0].ndim, 6) + self.assertEqual(J[0].ndim, 6) + self.assertIsNone(pml_state) + self.assertFalse(bool(overflow)) + # dump a dummy config file to tmp directory and confirm it can be read + # in correctly + + def test_initialize_simulation_encodes_global_particle_boundary_conditions(self): + with tempfile.TemporaryDirectory() as tmpdir: + zeros_path = os.path.join(tmpdir, "zeros.npy") + np.save(zeros_path, np.zeros(1)) + config = { + "simulation_parameters": { + "name": "global particle bc test", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 1, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + "particle_x_bc": "reflecting", + "particle_y_bc": "absorbing", + "particle_z_bc": "periodic", + }, + "plotting": {"plotting": False}, + "particle1": { + "name": "electrons", + "N_particles": 1, + "charge": -1.0, + "mass": 1.0, + "temperature": 1.0, + "x_bc": "absorbing", + "initial_x": zeros_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": zeros_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + config_path = os.path.join(tmpdir, "global_particle_bc.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + _, particles, _, world, *_ = initialize_simulation(toml.load(config_path)) + + self.assertEqual(world["particle_boundary_conditions"], {"x": 1, "y": 2, "z": 0}) + self.assertIsInstance(particles, TiledParticles) + # check that the global particle boundary conditions are encoded correctly in the world dictionary + + def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): + with tempfile.TemporaryDirectory() as tmpdir: + zeros_path = os.path.join(tmpdir, "zeros.npy") + np.save(zeros_path, np.zeros(1)) + config = { + "simulation_parameters": { + "name": "electrostatic collocated grid test", + "output_dir": tmpdir, + "solver": "electrostatic", + "Nx": 4, + "Ny": 2, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + }, + "plotting": {"plotting": False}, + "particle1": { + "name": "electrons", + "N_particles": 1, + "charge": -1.0, + "mass": 1.0, + "temperature": 1.0, + "initial_x": zeros_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": zeros_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + loop, particles, fields, world, *_ = initialize_simulation(config) + + self.assertIs(loop, time_loop_electrostatic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + for vertex_axis, center_axis in zip(world["grids"]["vertex"], world["grids"]["center"]): + self.assertTrue(jnp.allclose(vertex_axis, center_axis)) + # test the initialize_simulation function with an electrostatic solver and check that it uses a collocated grid + + def test_initialize_simulation_rejects_unknown_solver(self): + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "simulation_parameters": { + "name": "unknown solver test", + "output_dir": tmpdir, + "solver": "old_solver", + "Nx": 4, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + }, + "plotting": {"plotting": False}, + } + + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + initialize_simulation(config) + # test that initialize_simulation raises an error for an unknown solver + + def test_validate_field_solver_rejects_spectral(self): + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + validate_field_solver("spectral") + + def test_validate_field_solver_accepts_only_public_runtime_modes(self): + validate_field_solver("electrodynamic_yee") + validate_field_solver("electrostatic") + + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + validate_field_solver("fdtd") + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + validate_field_solver("tiled_yee") + +if __name__ == '__main__': + unittest.main() diff --git a/tests/code_tests/openpmd_test.py b/tests/code_tests/openpmd_test.py new file mode 100644 index 00000000..9abb08ea --- /dev/null +++ b/tests/code_tests/openpmd_test.py @@ -0,0 +1,703 @@ +import unittest +import math +import threading +import time +from unittest.mock import patch + +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics import async_writer +from PyPIC3D.diagnostics import openPMD +from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +class FakeRecord: + def __init__(self): + self.shape = None + self.dataset_shape = None + self.unit_SI = None + self.data = None + self.chunks = [] + + def reset_dataset(self, dataset): + extent = getattr(dataset, "extent", None) + if extent is not None: + self.dataset_shape = tuple(extent) + + def store_chunk(self, array, offset, extent): + self.shape = tuple(extent) + self.data = jnp.asarray(array) + self.chunks.append((tuple(offset), tuple(extent), jnp.asarray(array))) + + +class FakeMesh: + def __init__(self): + self.records = {} + self.axis_labels = None + self.grid_spacing = None + self.grid_global_offset = None + + def __getitem__(self, component): + if component not in self.records: + self.records[component] = FakeRecord() + return self.records[component] + + +class FakeMeshes(dict): + def __getitem__(self, name): + if name not in self: + self[name] = FakeMesh() + return dict.__getitem__(self, name) + + +class FakeParticleRecord(dict): + def __getitem__(self, component): + if component not in self: + self[component] = FakeRecord() + return dict.__getitem__(self, component) + + +class FakeParticleSpecies(dict): + def __getitem__(self, record_name): + if record_name not in self: + if record_name in ("position", "positionOffset", "momentum"): + self[record_name] = FakeParticleRecord() + else: + self[record_name] = FakeRecord() + return dict.__getitem__(self, record_name) + + +class FakeParticles(dict): + def __getitem__(self, species_name): + if species_name not in self: + self[species_name] = FakeParticleSpecies() + return dict.__getitem__(self, species_name) + + +class FakeIteration: + def __init__(self): + self.meshes = FakeMeshes() + self.particles = FakeParticles() + self.time = None + self.dt = None + self.time_unit_SI = None + + +class FakeIterations(dict): + def __getitem__(self, iteration): + if iteration not in self: + self[iteration] = FakeIteration() + return dict.__getitem__(self, iteration) + + +class FakeSeries: + def __init__(self): + self.iterations = FakeIterations() + + def set_attribute(self, name, value): + pass + + def flush(self): + pass + + def close(self): + pass + + +def _zero_field(shape): + return tuple(jnp.zeros(shape) for _ in range(3)) + + +def _world(): + return { + "dt": 0.2, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "x_wind": 4.0, + "y_wind": 2.0, + "z_wind": 1.0, + "Nx": 4, + "Ny": 2, + "Nz": 1, + "tile_shape": (2, 1, 1), + "guard_cells": 2, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + + +def _record_data(record): + if len(record.chunks) <= 1: + return record.data + + n = sum(int(extent[0]) for _offset, extent, _data in record.chunks) + out = jnp.zeros((n,), dtype=record.chunks[0][2].dtype) + for offset, extent, data in record.chunks: + start = int(offset[0]) + stop = start + int(extent[0]) + out = out.at[start:stop].set(data) + return out + + +def _species(name, charge, mass, weight, x1): + return { + "name": name, + "charge": charge, + "mass": mass, + "weight": weight, + "x1": jnp.asarray(x1), + "u1": jnp.ones_like(x1) * 0.1, + } + + +def _particle_tile_index(x, y, z, world, tile_shape): + tile_nx, tile_ny, tile_nz = tile_shape + + x_cell = math.floor((float(x) + float(world["x_wind"]) / 2.0) / float(world["dx"])) + y_cell = math.floor((float(y) + float(world["y_wind"]) / 2.0) / float(world["dy"])) + z_cell = math.floor((float(z) + float(world["z_wind"]) / 2.0) / float(world["dz"])) + + x_cell = min(max(x_cell, 0), int(world["Nx"]) - 1) + y_cell = min(max(y_cell, 0), int(world["Ny"]) - 1) + z_cell = min(max(z_cell, 0), int(world["Nz"]) - 1) + + return x_cell // tile_nx, y_cell // tile_ny, z_cell // tile_nz + + +def _make_tiled_particles(species, world): + tile_shape = tuple(int(width) for width in world["tile_shape"]) + tile_nx, tile_ny, tile_nz = tile_shape + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + + placements = [] + tile_counts = {} + + for species_index, species_data in enumerate(species): + x1 = species_data["x1"] + u1 = species_data["u1"] + zeros = jnp.zeros_like(x1) + x = jnp.stack((x1, zeros, zeros), axis=-1) + u = jnp.stack((u1, zeros, zeros), axis=-1) + + for particle_index in range(x1.shape[0]): + tx, ty, tz = _particle_tile_index( + x[particle_index, 0], + x[particle_index, 1], + x[particle_index, 2], + world, + tile_shape, + ) + tile_key = (tx, ty, tz, species_index) + tile_counts[tile_key] = tile_counts.get(tile_key, 0) + 1 + placements.append((tile_key, x[particle_index], u[particle_index])) + + max_particles_per_tile = max(1, max(tile_counts.values())) + n_species = len(species) + + x_tiles = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) + u_tiles = jnp.zeros_like(x_tiles) + active_tiles = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) + write_counts = {} + + for tile_key, x_particle, u_particle in placements: + slot = write_counts.get(tile_key, 0) + write_counts[tile_key] = slot + 1 + + x_tiles = x_tiles.at[tile_key + (slot, slice(None))].set(x_particle) + u_tiles = u_tiles.at[tile_key + (slot, slice(None))].set(u_particle) + active_tiles = active_tiles.at[tile_key + (slot,)].set(True) + + particles = TiledParticles(x=x_tiles, u=u_tiles, active=active_tiles) + species_config = SpeciesConfig( + charge=jnp.asarray([species_data["charge"] for species_data in species]), + mass=jnp.asarray([species_data["mass"] for species_data in species]), + weight=jnp.asarray([species_data["weight"] for species_data in species]), + update_x=jnp.ones((n_species, 3), dtype=bool), + update_u=jnp.ones((n_species, 3), dtype=bool), + ) + species_names = tuple(species_data["name"] for species_data in species) + + return particles, species_config, species_names + + +class OpenPMDDiagnosticsTests(unittest.TestCase): + + def test_openpmd_field_array_preserves_thin_y_axis(self): + field_component = jnp.ones((4, 1, 6)) + + array = _ensure_openpmd_array(field_component) + + self.assertEqual(array.shape, (4, 1, 6)) + + def _expected_tiled_scalar_from_interior(self, field, world, num_guard_cells): + tile_shape = world["tile_shape"] + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + def _field_with_stale_global_ghosts(self, world, value_offset=0.0): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + field = jnp.arange(jnp.prod(jnp.asarray(shape)), dtype=jnp.float64).reshape(shape) + field = field + value_offset + field = field.at[0, :, :].set(-1000.0) + field = field.at[-1, :, :].set(-2000.0) + field = field.at[:, 0, :].set(-3000.0) + field = field.at[:, -1, :].set(-4000.0) + field = field.at[:, :, 0].set(-5000.0) + field = field.at[:, :, -1].set(-6000.0) + return field + + def test_tile_scalar_field_one_guard_rebuilds_periodic_halos_from_interiors(self): + world = _world() + field = self._field_with_stale_global_ghosts(world) + + tiles = tile_scalar_field(field, world, world["tile_shape"], num_guard_cells=1) + expected = self._expected_tiled_scalar_from_interior(field, world, num_guard_cells=1) + + self.assertTrue(jnp.allclose(tiles, expected)) + + def test_tile_scalar_field_two_guards_rebuilds_periodic_halos_from_interiors(self): + world = _world() + field = self._field_with_stale_global_ghosts(world, value_offset=10.0) + + tiles = tile_scalar_field(field, world, world["tile_shape"], num_guard_cells=2) + expected = self._expected_tiled_scalar_from_interior(field, world, num_guard_cells=2) + + self.assertTrue(jnp.allclose(tiles, expected)) + + def test_write_openpmd_fields_preserves_thin_y_mesh_metadata(self): + shape_with_ghosts = (6, 3, 8) + E = _zero_field(shape_with_ghosts) + B = _zero_field(shape_with_ghosts) + J = _zero_field(shape_with_ghosts) + rho = jnp.zeros(shape_with_ghosts) + phi = jnp.zeros(shape_with_ghosts) + external_fields = _zero_field(shape_with_ghosts), _zero_field(shape_with_ghosts) + fields = (E, B, J, rho, phi, external_fields) + world = { + "dt": 1.0, + "dx": 0.25, + "dy": 0.5, + "dz": 0.75, + "x_wind": 1.0, + "y_wind": 2.0, + "z_wind": 3.0, + } + series = FakeSeries() + + with patch.object(openPMD, "_open_openpmd_series", return_value=series): + openPMD.write_openpmd_fields(fields, world, "/tmp", plot_t=0, t=0) + + B_mesh = series.iterations[0].meshes["B"] + self.assertEqual(B_mesh.axis_labels, ["x", "y", "z"]) + self.assertEqual(B_mesh.grid_spacing, [0.25, 0.5, 0.75]) + self.assertEqual(B_mesh.grid_global_offset, [-0.5, -1.0, -1.5]) + self.assertEqual(B_mesh.records["x"].shape, (4, 1, 6)) + + def test_write_openpmd_fields_assembles_tiled_fields_before_output(self): + shape_with_ghosts = (6, 4, 4) + E = _zero_field(shape_with_ghosts) + B = _zero_field(shape_with_ghosts) + J = _zero_field(shape_with_ghosts) + rho = jnp.zeros(shape_with_ghosts) + phi = jnp.zeros(shape_with_ghosts) + external_fields = _zero_field(shape_with_ghosts), _zero_field(shape_with_ghosts) + world = { + "dt": 1.0, + "dx": 0.25, + "dy": 0.5, + "dz": 0.75, + "x_wind": 1.0, + "y_wind": 2.0, + "z_wind": 3.0, + "Nx": 4, + "Ny": 2, + "Nz": 2, + "tile_shape": (2, 1, 1), + "guard_cells": 2, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + tiled_fields = ( + tile_vector_field(E, world, world["tile_shape"]), + tile_vector_field(B, world, world["tile_shape"]), + tile_vector_field(J, world, world["tile_shape"]), + tile_scalar_field(rho, world, world["tile_shape"]), + tile_scalar_field(phi, world, world["tile_shape"]), + ( + tile_vector_field(external_fields[0], world, world["tile_shape"]), + tile_vector_field(external_fields[1], world, world["tile_shape"]), + ), + None, + ) + series = FakeSeries() + + with patch.object(openPMD, "_open_openpmd_series", return_value=series): + openPMD.write_openpmd_fields(tiled_fields, world, "/tmp", plot_t=0, t=0) + + E_mesh = series.iterations[0].meshes["E"] + rho_mesh = series.iterations[0].meshes["rho"] + self.assertEqual(E_mesh.records["x"].shape, (4, 2, 2)) + self.assertEqual(rho_mesh.records[openPMD.io.Mesh_Record_Component.SCALAR].shape, (4, 2, 2)) + + def test_tiled_field_snapshot_writes_tile_chunks_without_global_assembly(self): + world = _world() + shape_with_ghosts = (6, 4, 3) + rho_global = jnp.arange(6 * 4 * 3, dtype=jnp.float64).reshape(shape_with_ghosts) + phi_global = rho_global + 100.0 + E_global = tuple(rho_global + offset for offset in (10.0, 20.0, 30.0)) + tile_shape = world["tile_shape"] + field_map = { + "E": tile_vector_field(E_global, world, tile_shape), + "rho": tile_scalar_field(rho_global, world, tile_shape), + "phi": tile_scalar_field(phi_global, world, tile_shape), + } + snapshot = async_writer.make_tiled_field_snapshot( + field_map, + step=3, + time=0.6, + ) + layout = openPMD.TiledMeshLayout( + global_shape=(world["Nx"], world["Ny"], world["Nz"]), + tile_shape=tile_shape, + guard_cells=world["guard_cells"], + ) + series = FakeSeries() + + with patch.object(openPMD, "_open_openpmd_series", return_value=series): + openPMD.write_tiled_field_snapshot_openpmd( + snapshot, + output_dir="/tmp", + filename="fields", + world=world, + layout=layout, + file_extension=".h5", + ) + + iteration = series.iterations[3] + rho_record = iteration.meshes["rho"].records[openPMD.io.Mesh_Record_Component.SCALAR] + E_record = iteration.meshes["E"].records["x"] + + self.assertEqual(iteration.time, 0.6) + self.assertEqual(iteration.dt, world["dt"]) + self.assertEqual(iteration.meshes["rho"].axis_labels, ["x", "y", "z"]) + self.assertEqual(len(rho_record.chunks), 4) + self.assertEqual(rho_record.chunks[0][0], (0, 0, 0)) + self.assertEqual(rho_record.chunks[0][1], tile_shape) + self.assertEqual(len(E_record.chunks), 4) + self.assertEqual(E_record.chunks[0][0], (0, 0, 0)) + + def test_async_field_writer_queue_size_caps_pending_snapshots(self): + world = _world() + writer = async_writer.AsyncTiledOpenPMDFieldWriter( + output_dir="/tmp", + filename="fields", + world=world, + global_shape=(world["Nx"], world["Ny"], world["Nz"]), + tile_shape=world["tile_shape"], + guard_cells=world["guard_cells"], + queue_size=1, + ) + snapshot = async_writer.TiledFieldSnapshot(step=0, time=0.0, fields={}) + + self.assertTrue(writer.enqueue(snapshot, block=False)) + self.assertFalse(writer.enqueue(snapshot, block=False)) + writer.close(raise_errors=False) + + def test_async_field_writer_raises_worker_errors_on_close(self): + world = _world() + writer = async_writer.AsyncTiledOpenPMDFieldWriter( + output_dir="/tmp", + filename="fields", + world=world, + global_shape=(world["Nx"], world["Ny"], world["Nz"]), + tile_shape=world["tile_shape"], + guard_cells=world["guard_cells"], + queue_size=1, + ) + snapshot = async_writer.TiledFieldSnapshot(step=0, time=0.0, fields={}) + + with patch.object(async_writer, "write_tiled_field_snapshot_openpmd", side_effect=RuntimeError("disk failed")): + writer.start() + self.assertTrue(writer.enqueue(snapshot)) + with self.assertRaisesRegex(RuntimeError, "Async openPMD writer failed"): + writer.close() + + def test_tiled_particle_snapshot_writes_same_records_as_synchronous_output(self): + world = _world() + species = [ + _species("beam electrons", -1.0, 2.0, 3.0, jnp.array([-1.5, 0.5])), + _species("background ions", 1.0, 4.0, 5.0, jnp.array([1.5, -0.5])), + ] + tiled_particles, species_config, species_names = _make_tiled_particles(species, world) + constants = {"C": 10.0} + + sync_series = FakeSeries() + with patch.object(openPMD, "_open_openpmd_series", return_value=sync_series): + openPMD.write_openpmd_particles( + tiled_particles, + world, + constants, + "/tmp", + plot_t=2, + t=3, + species_config=species_config, + species_names=species_names, + ) + + snapshot = async_writer.make_tiled_particle_snapshot( + tiled_particles, + step=2, + time=3 * world["dt"], + species_names=species_names, + species_config=species_config, + ) + async_series = FakeSeries() + with patch.object(openPMD, "_open_openpmd_series", return_value=async_series): + openPMD.write_tiled_particle_snapshot_openpmd( + snapshot, + output_dir="/tmp", + filename="particles", + world=world, + constants=constants, + file_extension=".h5", + ) + + sync_electrons = sync_series.iterations[2].particles["beam_electrons"] + async_electrons = async_series.iterations[2].particles["beam_electrons"] + sync_ions = sync_series.iterations[2].particles["background_ions"] + async_ions = async_series.iterations[2].particles["background_ions"] + + self.assertEqual(async_series.iterations[2].time, 3 * world["dt"]) + self.assertEqual(_record_data(async_electrons["position"]["x"]).shape, _record_data(sync_electrons["position"]["x"]).shape) + self.assertEqual(_record_data(async_ions["position"]["x"]).shape, _record_data(sync_ions["position"]["x"]).shape) + for group_async, group_sync in ((async_electrons, sync_electrons), (async_ions, sync_ions)): + for record_name in ("position", "positionOffset", "momentum"): + for component in ("x", "y", "z"): + self.assertTrue( + jnp.allclose( + _record_data(group_async[record_name][component]), + _record_data(group_sync[record_name][component]), + ) + ) + for record_name in ("weighting", "charge", "mass"): + self.assertTrue( + jnp.allclose( + _record_data(group_async[record_name]), + _record_data(group_sync[record_name]), + ) + ) + + def test_async_particle_writer_queue_size_caps_pending_snapshots(self): + world = _world() + writer = async_writer.AsyncTiledOpenPMDParticleWriter( + output_dir="/tmp", + filename="particles", + world=world, + constants={"C": 10.0}, + queue_size=1, + ) + snapshot = async_writer.TiledParticleSnapshot( + step=0, + time=0.0, + species_names=("electrons",), + x_shards=[], + u_shards=[], + active_shards=[], + species_charge=jnp.array([-1.0]), + species_mass=jnp.array([1.0]), + species_weight=jnp.array([1.0]), + ) + + self.assertTrue(writer.enqueue(snapshot, block=False)) + self.assertFalse(writer.enqueue(snapshot, block=False)) + writer.close(raise_errors=False) + + def test_async_particle_writer_raises_worker_errors_on_close(self): + world = _world() + writer = async_writer.AsyncTiledOpenPMDParticleWriter( + output_dir="/tmp", + filename="particles", + world=world, + constants={"C": 10.0}, + queue_size=1, + ) + snapshot = async_writer.TiledParticleSnapshot( + step=0, + time=0.0, + species_names=("electrons",), + x_shards=[], + u_shards=[], + active_shards=[], + species_charge=jnp.array([-1.0]), + species_mass=jnp.array([1.0]), + species_weight=jnp.array([1.0]), + ) + + with patch.object(async_writer, "write_tiled_particle_snapshot_openpmd", side_effect=RuntimeError("disk failed")): + writer.start() + self.assertTrue(writer.enqueue(snapshot)) + with self.assertRaisesRegex(RuntimeError, "Async openPMD particle writer failed"): + writer.close() + + def test_field_and_particle_async_writers_serialize_openpmd_calls(self): + world = _world() + field_writer = async_writer.AsyncTiledOpenPMDFieldWriter( + output_dir="/tmp", + filename="fields", + world=world, + global_shape=(world["Nx"], world["Ny"], world["Nz"]), + tile_shape=world["tile_shape"], + guard_cells=world["guard_cells"], + queue_size=1, + ) + particle_writer = async_writer.AsyncTiledOpenPMDParticleWriter( + output_dir="/tmp", + filename="particles", + world=world, + constants={"C": 10.0}, + queue_size=1, + ) + field_snapshot = async_writer.TiledFieldSnapshot(step=0, time=0.0, fields={}) + particle_snapshot = async_writer.TiledParticleSnapshot( + step=0, + time=0.0, + species_names=("electrons",), + x_shards=[], + u_shards=[], + active_shards=[], + species_charge=jnp.array([-1.0]), + species_mass=jnp.array([1.0]), + species_weight=jnp.array([1.0]), + ) + + active_writes = 0 + overlaps = [] + counter_lock = threading.Lock() + field_write_started = threading.Event() + + def write_snapshot(*args, **kwargs): + nonlocal active_writes + with counter_lock: + active_writes += 1 + if active_writes > 1: + overlaps.append(active_writes) + field_write_started.set() + time.sleep(0.05) + with counter_lock: + active_writes -= 1 + + with ( + patch.object(async_writer, "write_tiled_field_snapshot_openpmd", side_effect=write_snapshot), + patch.object(async_writer, "write_tiled_particle_snapshot_openpmd", side_effect=write_snapshot), + ): + field_writer.start() + particle_writer.start() + self.assertTrue(field_writer.enqueue(field_snapshot)) + self.assertTrue(field_write_started.wait(timeout=1.0)) + self.assertTrue(particle_writer.enqueue(particle_snapshot)) + field_writer.close() + particle_writer.close() + + self.assertEqual(overlaps, []) + + def test_write_openpmd_initial_particles_flattens_tiled_particles_and_preserves_names(self): + world = _world() + species = [ + _species("beam electrons", -1.0, 2.0, 3.0, jnp.array([-1.5, 0.5])), + _species("background ions", 1.0, 4.0, 5.0, jnp.array([1.5])), + ] + tiled_particles, species_config, species_names = _make_tiled_particles(species, world) + series = FakeSeries() + + with patch.object(openPMD.io, "Series", return_value=series): + openPMD.write_openpmd_initial_particles( + tiled_particles, + world, + {"C": 10.0}, + "/tmp", + species_config=species_config, + species_names=species_names, + ) + + iteration = series.iterations[0] + electron_group = iteration.particles["beam_electrons"] + ion_group = iteration.particles["background_ions"] + self.assertEqual(electron_group["position"]["x"].shape, (2,)) + self.assertEqual(ion_group["position"]["x"].shape, (1,)) + self.assertTrue(jnp.allclose(electron_group["weighting"].data, jnp.array([3.0, 3.0]))) + self.assertTrue(jnp.allclose(ion_group["charge"].data, jnp.array([1.0]))) + expected_gamma = 1.0 / jnp.sqrt(1.0 - 0.1**2 / 10.0**2) + self.assertTrue(jnp.allclose(electron_group["momentum"]["x"].data, jnp.ones(2) * 0.1 * 6.0 * expected_gamma)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/particle_diagnostics_test.py b/tests/code_tests/particle_diagnostics_test.py new file mode 100644 index 00000000..f6fe6f53 --- /dev/null +++ b/tests/code_tests/particle_diagnostics_test.py @@ -0,0 +1,159 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.diagnostics.output_adapters import particles_for_output +from tests.initial_particles import build_tiled_particles, species_names, tiled_species + + +jax.config.update("jax_enable_x64", True) + + +class TestTiledParticleDiagnostics(unittest.TestCase): + def _world(self): + return { + "Nx": 4, + "Ny": 2, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.2, + "x_wind": 4.0, + "y_wind": 2.0, + "z_wind": 1.0, + } + + def _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + def _species(self): + world = self._world() + ions = tiled_species( + name="ions", + N_particles=4, + charge=2.0, + mass=3.0, + weight=4.0, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.array([-0.5, -0.5, 0.5, 0.5]), + x3=jnp.array([0.0, 0.0, 0.0, 0.0]), + v1=jnp.array([0.1, 0.2, 0.3, 0.4]), + v2=jnp.array([1.0, 1.1, 1.2, 1.3]), + v3=jnp.array([2.0, 2.1, 2.2, 2.3]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + active_mask=jnp.array([True, False, True, True]), + dt=world["dt"], + ) + electrons = tiled_species( + name="electrons", + N_particles=3, + charge=-1.0, + mass=0.5, + weight=8.0, + T=1.0, + x1=jnp.array([-1.25, 0.25, 1.25]), + x2=jnp.array([0.25, -0.25, 0.25]), + x3=jnp.array([0.0, 0.0, 0.0]), + v1=jnp.array([-0.1, -0.2, -0.3]), + v2=jnp.array([-1.0, -1.1, -1.2]), + v3=jnp.array([-2.0, -2.1, -2.2]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + active_mask=jnp.array([False, True, True]), + dt=world["dt"], + ) + return [ions, electrons] + + def test_flatten_tiled_particles_matches_active_original_species(self): + species_list = self._species() + tiled_particles, species_config = build_tiled_particles(species_list, self._world(), self._simulation_parameters()) + + flattened_species = particles_for_output(tiled_particles, species_config=species_config) + + self.assertEqual(len(flattened_species), len(species_list)) + for species_index, (original, flattened) in enumerate(zip(species_list, flattened_species)): + active = original["active"] + + original_x = original["x"][active] + original_u = original["u"][active] + flattened_x = flattened.x + flattened_u = flattened.u + + self.assertEqual(flattened.species_index, species_index) + self.assertTrue(jnp.allclose(flattened_x, original_x)) + self.assertTrue(jnp.allclose(flattened_u, original_u)) + self.assertTrue(jnp.allclose(flattened.charge, original["charge"] * jnp.ones(jnp.sum(active)))) + self.assertTrue(jnp.allclose(flattened.mass, original["mass"] * jnp.ones(jnp.sum(active)))) + self.assertTrue(jnp.allclose(flattened.weight, original["weight"] * jnp.ones(jnp.sum(active)))) + + def test_inactive_tiled_slots_are_not_flattened(self): + species_list = self._species() + tiled_particles, species_config = build_tiled_particles(species_list, self._world(), self._simulation_parameters()) + + flattened_species = particles_for_output(tiled_particles, species_config=species_config) + + self.assertEqual(flattened_species[0].x.shape[0], 3) + self.assertEqual(flattened_species[1].x.shape[0], 2) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].x[:, 0], -0.5))) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[1].x[:, 0], -1.25))) + + def test_absorbed_particles_do_not_appear_in_flattened_output(self): + species_list = self._species() + tiled_particles, species_config = build_tiled_particles(species_list, self._world(), self._simulation_parameters()) + + tiled_particles = tiled_particles._replace( + active=tiled_particles.active.at[1, 1, 0, 0, 1].set(False) + ) + flattened_species = particles_for_output(tiled_particles, species_config=species_config) + + self.assertEqual(flattened_species[0].x.shape[0], 2) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].x[:, 0], 1.5))) + + def test_flattened_diagnostic_positions_match_original_half_step_positions(self): + species_list = self._species() + world = self._world() + world["particle_boundary_conditions"] = { + "x": 0, + "y": 0, + "z": 0, + } + tiled_particles, species_config = build_tiled_particles(species_list, world, self._simulation_parameters()) + + flattened_species = particles_for_output(tiled_particles, species_config=species_config, world=world) + + for original, flattened in zip(species_list, flattened_species): + active = original["active"] + original_position = original["x"][active] - 0.5 * original["u"][active] * world["dt"] + flattened_position = flattened.x_diagnostic + + self.assertTrue(jnp.allclose(flattened_position, original_position)) + + def test_species_names_are_preserved_when_metadata_is_available(self): + species_list = self._species() + tiled_particles, species_config = build_tiled_particles(species_list, self._world(), self._simulation_parameters()) + names = species_names(species_list) + + flattened_species = particles_for_output(tiled_particles, species_config=species_config, species_names=names) + + self.assertEqual(flattened_species[0].name, "ions") + self.assertEqual(flattened_species[1].name, "electrons") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/particle_initialization_test.py b/tests/code_tests/particle_initialization_test.py new file mode 100644 index 00000000..b624d855 --- /dev/null +++ b/tests/code_tests/particle_initialization_test.py @@ -0,0 +1,439 @@ +import os +import inspect +import tempfile +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from PyPIC3D.particles.particle_initialization import load_particles_from_toml +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from tests.initial_particles import build_tiled_particles, tiled_species + + +jax.config.update("jax_enable_x64", True) + + +class TestTiledParticleInitialization(unittest.TestCase): + def _write_array(self, tmpdir, name, values): + path = os.path.join(tmpdir, name) + np.save(path, np.asarray(values, dtype=float)) + return path + + def test_direct_tiled_particles_preserve_inactive_slots_and_metadata(self): + world = { + "Nx": 4, + "Ny": 2, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 4.0, + "y_wind": 2.0, + "z_wind": 1.0, + } + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + species = tiled_species( + name="ions", + N_particles=3, + charge=2.0, + mass=3.0, + weight=4.0, + T=1.0, + x1=jnp.array([-1.5, 0.5, 1.5]), + x2=jnp.array([-0.5, 0.5, 0.5]), + x3=jnp.array([0.0, 0.0, 0.0]), + v1=jnp.array([0.1, 0.2, 0.3]), + v2=jnp.array([0.0, 0.0, 0.0]), + v3=jnp.array([1.0, 2.0, 3.0]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + update_y=False, + update_vx=False, + active_mask=jnp.array([True, False, True]), + dt=world["dt"], + ) + + particles, species_config = build_tiled_particles([species], world, simulation_parameters) + + self.assertIsInstance(particles, TiledParticles) + self.assertIsInstance(species_config, SpeciesConfig) + self.assertEqual(particles.x.shape, (2, 2, 1, 1, 2, 3)) + self.assertEqual(int(jnp.sum(particles.active)), 2) + self.assertFalse(hasattr(particles, "charge")) + self.assertFalse(hasattr(particles, "mass")) + self.assertFalse(hasattr(particles, "weight")) + + self.assertTrue(jnp.allclose(particles.x[0, 0, 0, 0, 0], jnp.array([-1.5, -0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 1, 0, 0, 0], jnp.array([0.5, 0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 1, 0, 0, 1], jnp.array([1.5, 0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.u[1, 1, 0, 0, 1], jnp.array([0.3, 0.0, 3.0]))) + + self.assertEqual(species_config.charge.shape, (1,)) + self.assertEqual(species_config.mass.shape, (1,)) + self.assertEqual(species_config.weight.shape, (1,)) + self.assertEqual(species_config.update_x.shape, (1, 3)) + self.assertEqual(species_config.update_u.shape, (1, 3)) + self.assertTrue(jnp.allclose(species_config.charge, jnp.array([2.0]))) + self.assertTrue(jnp.allclose(species_config.mass, jnp.array([3.0]))) + self.assertTrue(jnp.allclose(species_config.weight, jnp.array([4.0]))) + + self.assertTrue(bool(species_config.update_x[0, 0])) + self.assertFalse(bool(species_config.update_x[0, 1])) + self.assertTrue(bool(species_config.update_x[0, 2])) + self.assertFalse(bool(species_config.update_u[0, 0])) + self.assertTrue(bool(species_config.update_u[0, 1])) + self.assertTrue(bool(species_config.update_u[0, 2])) + + def test_species_metadata_is_not_slot_shaped(self): + world = { + "Nx": 4, + "Ny": 1, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + } + simulation_parameters = { + "particle_tile_nx": 1, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + species = [ + tiled_species( + name="electrons", + N_particles=2, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5]), + x2=jnp.zeros(2), + x3=jnp.zeros(2), + v1=jnp.zeros(2), + v2=jnp.zeros(2), + v3=jnp.zeros(2), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ), + tiled_species( + name="ions", + N_particles=2, + charge=2.0, + mass=5.0, + weight=0.25, + T=1.0, + x1=jnp.array([0.5, 1.5]), + x2=jnp.zeros(2), + x3=jnp.zeros(2), + v1=jnp.zeros(2), + v2=jnp.zeros(2), + v3=jnp.zeros(2), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ), + ] + + particles, species_config = build_tiled_particles(species, world, simulation_parameters) + + self.assertEqual(particles.x.shape[:4], (4, 1, 1, 2)) + self.assertEqual(species_config.charge.shape, (2,)) + self.assertEqual(species_config.mass.shape, (2,)) + self.assertEqual(species_config.weight.shape, (2,)) + self.assertEqual(species_config.update_x.shape, (2, 3)) + self.assertEqual(species_config.update_u.shape, (2, 3)) + + def test_direct_tiled_particles_can_allocate_inactive_capacity_headroom(self): + world = { + "Nx": 4, + "Ny": 1, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + } + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "particle_tile_capacity_factor": 3.0, + } + species = tiled_species( + name="ions", + N_particles=3, + charge=1.0, + mass=1.0, + weight=1.0, + T=1.0, + x1=jnp.array([-1.5, -0.5, 1.5]), + x2=jnp.zeros(3), + x3=jnp.zeros(3), + v1=jnp.zeros(3), + v2=jnp.zeros(3), + v3=jnp.zeros(3), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + particles, species_config = build_tiled_particles([species], world, simulation_parameters) + + self.assertEqual(particles.active.shape[-1], 6) + self.assertEqual(int(jnp.sum(particles.active)), 3) + self.assertEqual(species_config.charge.shape, (1,)) + + def test_load_particles_from_toml_uses_tile_axes_before_species(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = self._write_array(tmpdir, "x.npy", [-1.5, 0.5, 1.5]) + y_path = self._write_array(tmpdir, "y.npy", [-0.5, 0.5, 0.5]) + z_path = self._write_array(tmpdir, "z.npy", [0.0, 0.0, 0.0]) + vx_path = self._write_array(tmpdir, "vx.npy", [0.1, 0.2, 0.3]) + vy_path = self._write_array(tmpdir, "vy.npy", [0.0, 0.0, 0.0]) + vz_path = self._write_array(tmpdir, "vz.npy", [1.0, 2.0, 3.0]) + + world = { + "Nx": 4, + "Ny": 2, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 4.0, + "y_wind": 2.0, + "z_wind": 1.0, + "tile_shape": (2, 1, 1), + } + constants = {"kb": 1.0, "eps": 1.0} + simulation_parameters = { + "ds_per_debye": None, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + config = { + "particle1": { + "name": "electrons", + "N_particles": 3, + "charge": -1.0, + "mass": 2.0, + "weight": 4.0, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": y_path, + "initial_z": z_path, + "initial_vx": vx_path, + "initial_vy": vy_path, + "initial_vz": vz_path, + } + } + + particles, species_config, species_names, metadata = load_particles_from_toml(config, simulation_parameters, world, constants) + + self.assertIsInstance(particles, TiledParticles) + self.assertIsInstance(species_config, SpeciesConfig) + self.assertEqual(species_names, ("electrons",)) + self.assertEqual(metadata[0]["name"], "electrons") + self.assertEqual(particles.x.shape, (2, 2, 1, 1, 2, 3)) + self.assertEqual(particles.u.shape, (2, 2, 1, 1, 2, 3)) + + self.assertTrue(particles.active[0, 0, 0, 0, 0]) + self.assertTrue(particles.active[1, 1, 0, 0, 0]) + self.assertTrue(particles.active[1, 1, 0, 0, 1]) + self.assertEqual(int(jnp.sum(particles.active)), 3) + + self.assertTrue(jnp.allclose(particles.x[0, 0, 0, 0, 0], jnp.array([-1.5, -0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 1, 0, 0, 0], jnp.array([0.5, 0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 1, 0, 0, 1], jnp.array([1.5, 0.5, 0.0]))) + self.assertTrue(jnp.allclose(particles.u[1, 1, 0, 0, 1], jnp.array([0.3, 0.0, 3.0]))) + + self.assertTrue(jnp.allclose(species_config.charge, jnp.array([-1.0]))) + self.assertTrue(jnp.allclose(species_config.mass, jnp.array([2.0]))) + self.assertTrue(jnp.allclose(species_config.weight, jnp.array([4.0]))) + + def test_load_particles_from_toml_preserves_interleaved_tile_order(self): + with tempfile.TemporaryDirectory() as tmpdir: + electron_x_path = self._write_array(tmpdir, "electron_x.npy", [-1.5, 0.5, -0.5, 1.5]) + ion_x_path = self._write_array(tmpdir, "ion_x.npy", [1.5, -1.5, 0.5, -0.5]) + y_path = self._write_array(tmpdir, "y.npy", [0.0, 0.0, 0.0, 0.0]) + z_path = self._write_array(tmpdir, "z.npy", [0.0, 0.0, 0.0, 0.0]) + electron_vx_path = self._write_array(tmpdir, "electron_vx.npy", [10.0, 20.0, 30.0, 40.0]) + ion_vx_path = self._write_array(tmpdir, "ion_vx.npy", [100.0, 200.0, 300.0, 400.0]) + vy_path = self._write_array(tmpdir, "vy.npy", [0.0, 0.0, 0.0, 0.0]) + vz_path = self._write_array(tmpdir, "vz.npy", [1.0, 2.0, 3.0, 4.0]) + + world = { + "Nx": 4, + "Ny": 1, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "tile_shape": (2, 1, 1), + } + constants = {"kb": 1.0, "eps": 1.0} + simulation_parameters = { + "ds_per_debye": None, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + config = { + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 4.0, + "temperature": 1.0, + "initial_x": electron_x_path, + "initial_y": y_path, + "initial_z": z_path, + "initial_vx": electron_vx_path, + "initial_vy": vy_path, + "initial_vz": vz_path, + }, + "particle2": { + "name": "ions", + "N_particles": 4, + "charge": 1.0, + "mass": 3.0, + "weight": 5.0, + "temperature": 1.0, + "initial_x": ion_x_path, + "initial_y": y_path, + "initial_z": z_path, + "initial_vx": ion_vx_path, + "initial_vy": vy_path, + "initial_vz": vz_path, + }, + } + + particles, species_config, species_names, metadata = load_particles_from_toml(config, simulation_parameters, world, constants) + + self.assertEqual(species_names, ("electrons", "ions")) + self.assertEqual(tuple(item["name"] for item in metadata), ("electrons", "ions")) + self.assertEqual(particles.x.shape, (2, 1, 1, 2, 2, 3)) + self.assertEqual(int(jnp.sum(particles.active)), 8) + + self.assertTrue(jnp.allclose(particles.x[0, 0, 0, 0, :, 0], jnp.array([-1.5, -0.5]))) + self.assertTrue(jnp.allclose(particles.u[0, 0, 0, 0, :, 0], jnp.array([10.0, 30.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 0, 0, 0, :, 0], jnp.array([0.5, 1.5]))) + self.assertTrue(jnp.allclose(particles.u[1, 0, 0, 0, :, 0], jnp.array([20.0, 40.0]))) + + self.assertTrue(jnp.allclose(particles.x[0, 0, 0, 1, :, 0], jnp.array([-1.5, -0.5]))) + self.assertTrue(jnp.allclose(particles.u[0, 0, 0, 1, :, 0], jnp.array([200.0, 400.0]))) + self.assertTrue(jnp.allclose(particles.x[1, 0, 0, 1, :, 0], jnp.array([1.5, 0.5]))) + self.assertTrue(jnp.allclose(particles.u[1, 0, 0, 1, :, 0], jnp.array([100.0, 300.0]))) + + self.assertTrue(jnp.allclose(species_config.charge, jnp.array([-1.0, 1.0]))) + self.assertTrue(jnp.allclose(species_config.mass, jnp.array([2.0, 3.0]))) + self.assertTrue(jnp.allclose(species_config.weight, jnp.array([4.0, 5.0]))) + + def test_load_particles_from_toml_does_not_scatter_each_particle_with_jax_at(self): + source = inspect.getsource(load_particles_from_toml) + + self.assertNotIn(".at[index].set", source) + + def test_load_particles_from_toml_maps_update_flags_to_active_particles(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = self._write_array(tmpdir, "x.npy", [0.0]) + y_path = self._write_array(tmpdir, "y.npy", [0.0]) + z_path = self._write_array(tmpdir, "z.npy", [0.0]) + + world = { + "Nx": 1, + "Ny": 1, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "tile_shape": (1, 1, 1), + } + constants = {"kb": 1.0, "eps": 1.0} + simulation_parameters = { + "ds_per_debye": None, + "shape_factor": 1, + "particle_tile_nx": 1, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + config = { + "particle1": { + "name": "partly fixed", + "N_particles": 1, + "charge": 1.0, + "mass": 1.0, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": y_path, + "initial_z": z_path, + "initial_vx": 0.0, + "initial_vy": 0.0, + "initial_vz": 0.0, + "update_pos": True, + "update_x": True, + "update_y": False, + "update_z": True, + "update_v": True, + "update_vx": False, + "update_vy": True, + "update_vz": False, + } + } + + particles, species_config, species_names, metadata = load_particles_from_toml(config, simulation_parameters, world, constants) + + self.assertTrue(bool(species_config.update_x[0, 0])) + self.assertFalse(bool(species_config.update_x[0, 1])) + self.assertTrue(bool(species_config.update_x[0, 2])) + self.assertFalse(bool(species_config.update_u[0, 0])) + self.assertTrue(bool(species_config.update_u[0, 1])) + self.assertFalse(bool(species_config.update_u[0, 2])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/particle_refresh_test.py b/tests/code_tests/particle_refresh_test.py new file mode 100644 index 00000000..480ab32d --- /dev/null +++ b/tests/code_tests/particle_refresh_test.py @@ -0,0 +1,230 @@ +import unittest + +import jax +import jax.numpy as jnp + +from tests.initial_particles import build_tiled_particles, tiled_species +from PyPIC3D.particles.particle_tile_communication import ( + _adjacent_tile_offset, + _refresh_tiled_particle_tiles_compacting, + _refresh_tiled_particle_tiles_sparse, + _tiles_need_repack, + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) + + +jax.config.update("jax_enable_x64", True) + + +class TestTiledParticleRefresh(unittest.TestCase): + def _build_world(self): + return { + "Nx": 4, + "Ny": 1, + "Nz": 1, + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "dt": 1.0, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + + def _species(self, world, x1, v1, active_mask=None, update_x=True, x_bc="periodic"): + if active_mask is None: + active_mask = jnp.ones_like(jnp.asarray(x1), dtype=bool) + n_particles = len(x1) + return tiled_species( + name="moving", + N_particles=n_particles, + charge=2.0, + mass=3.0, + weight=4.0, + T=1.0, + x1=jnp.asarray(x1, dtype=float), + x2=jnp.zeros(n_particles), + x3=jnp.zeros(n_particles), + v1=jnp.asarray(v1, dtype=float), + v2=jnp.zeros(n_particles), + v3=jnp.zeros(n_particles), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + active_mask=active_mask, + update_x=update_x, + x_bc=x_bc, + ) + + def _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + def _active_rows(self, tiled_particles): + active = tiled_particles.active.reshape(-1) + x = tiled_particles.x.reshape(-1, 3)[active] + u = tiled_particles.u.reshape(-1, 3)[active] + order = jnp.argsort(x[:, 0]) + return x[order], u[order] + + def test_update_tiled_particle_positions_respects_active_and_update_flags(self): + world = self._build_world() + species = self._species( + world, + x1=[-1.5, -0.5, 0.5], + v1=[0.25, 0.5, 0.75], + active_mask=jnp.array([True, False, True]), + ) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + + moved = update_tiled_particle_positions(tiled_particles, species_config, world["dt"]) + + x, _ = self._active_rows(moved) + self.assertTrue(jnp.allclose(x[:, 0], jnp.array([-1.25, 1.25]))) + self.assertTrue(jnp.allclose(moved.x[~tiled_particles.active], tiled_particles.x[~tiled_particles.active])) + + fixed_species = self._species(world, x1=[-1.5], v1=[0.25], update_x=False) + fixed, fixed_species_config = build_tiled_particles([fixed_species], world, self._simulation_parameters()) + fixed_moved = update_tiled_particle_positions(fixed, fixed_species_config, world["dt"]) + self.assertTrue(jnp.allclose(fixed_moved.x, fixed.x)) + + def test_adjacent_tile_offset_handles_periodic_edges(self): + source = jnp.array([0, 0, 1, 1]) + dest = jnp.array([0, 1, 0, 1]) + + offset = _adjacent_tile_offset(dest, source, tile_count=2) + + self.assertTrue(jnp.array_equal(offset, jnp.array([0, 1, -1, 0]))) + + def test_tiles_need_repack_is_false_when_active_particles_stay_in_tile(self): + active = jnp.array([[[[[True, False]]], [[[True, False]]]]]) + zero_offsets = jnp.zeros_like(active, dtype=int) + + repack = _tiles_need_repack(zero_offsets, zero_offsets, zero_offsets, active) + + self.assertTrue(jnp.array_equal(repack, jnp.zeros(active.shape[:-1], dtype=bool))) + + def test_refresh_moves_particles_to_neighbor_tiles_with_static_shape(self): + world = self._build_world() + species = self._species(world, x1=[-1.5, -0.25, 0.25], v1=[0.0, 0.0, 0.0]) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved_x = tiled_particles.x.at[0, 0, 0, 0, 1, 0].set(0.25) + moved = tiled_particles._replace(x=moved_x) + + refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) + + self.assertEqual(refreshed.x.shape, tiled_particles.x.shape) + self.assertFalse(bool(overflow)) + self.assertEqual(int(jnp.sum(refreshed.active[0, 0, 0, 0])), 1) + self.assertEqual(int(jnp.sum(refreshed.active[1, 0, 0, 0])), 2) + x, u = self._active_rows(refreshed) + self.assertTrue(jnp.allclose(x[:, 0], jnp.array([-1.5, 0.25, 0.25]))) + self.assertTrue(jnp.allclose(u[:, 0], jnp.array([0.0, 0.0, 0.0]))) + + def test_refresh_wraps_periodic_particle_to_opposite_tile(self): + world = self._build_world() + species = self._species(world, x1=[1.75], v1=[0.0]) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace(x=tiled_particles.x.at[1, 0, 0, 0, 0, 0].set(2.25)) + + refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) + + self.assertFalse(bool(overflow)) + self.assertTrue(bool(refreshed.active[0, 0, 0, 0, 0])) + self.assertTrue(jnp.allclose(refreshed.x[0, 0, 0, 0, 0, 0], -1.75)) + + def test_refresh_reflects_particle_from_global_boundary_condition(self): + world = self._build_world() + world["particle_boundary_conditions"] = {"x": 1, "y": 0, "z": 0} + species = self._species(world, x1=[1.75, -1.75], v1=[0.5, -0.25], x_bc="periodic") + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace( + x=tiled_particles.x + .at[1, 0, 0, 0, 0, 0].set(2.25) + .at[0, 0, 0, 0, 0, 0].set(-2.10) + ) + + refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) + + self.assertFalse(bool(overflow)) + x, u = self._active_rows(refreshed) + self.assertTrue(jnp.allclose(x[:, 0], jnp.array([-1.90, 1.75]))) + self.assertTrue(jnp.allclose(u[:, 0], jnp.array([0.25, -0.5]))) + + def test_refresh_absorbs_particle_from_global_boundary_condition(self): + world = self._build_world() + world["particle_boundary_conditions"] = {"x": 2, "y": 0, "z": 0} + species = self._species(world, x1=[1.75, -0.25], v1=[0.5, 0.0], x_bc="periodic") + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace(x=tiled_particles.x.at[1, 0, 0, 0, 0, 0].set(2.25)) + + refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) + + self.assertFalse(bool(overflow)) + self.assertEqual(int(jnp.sum(refreshed.active)), 1) + x, u = self._active_rows(refreshed) + self.assertTrue(jnp.allclose(x[:, 0], jnp.array([-0.25]))) + self.assertTrue(jnp.allclose(u[:, 0], jnp.array([0.0]))) + + def test_refresh_reports_overflow_without_changing_shape(self): + world = self._build_world() + species = self._species(world, x1=[-1.5, 0.5, 1.5], v1=[0.0, 0.0, 0.0]) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace(x=tiled_particles.x.at[0, 0, 0, 0, 0, 0].set(0.25)) + + refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) + + self.assertEqual(refreshed.x.shape, tiled_particles.x.shape) + self.assertTrue(bool(overflow)) + self.assertEqual(int(jnp.sum(refreshed.active)), 2) + + def test_sparse_refresh_matches_compacting_refresh_for_neighbor_exchange(self): + world = self._build_world() + species = self._species( + world, + x1=[-1.5, -0.25, 0.25, 1.25], + v1=[0.0, 0.0, 0.0, 0.0], + ) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace( + x=tiled_particles.x + .at[0, 0, 0, 0, 1, 0].set(0.25) + .at[1, 0, 0, 0, 1, 0].set(-0.25) + ) + + compacted, compacted_overflow = _refresh_tiled_particle_tiles_compacting(moved, world, tile_shape=(2, 1, 1)) + sparse, sparse_overflow = _refresh_tiled_particle_tiles_sparse(moved, world, tile_shape=(2, 1, 1)) + + self.assertEqual(sparse.x.shape, moved.x.shape) + self.assertEqual(bool(sparse_overflow), bool(compacted_overflow)) + self.assertEqual(int(jnp.sum(sparse.active)), int(jnp.sum(compacted.active))) + compacted_x, compacted_u = self._active_rows(compacted) + sparse_x, sparse_u = self._active_rows(sparse) + self.assertTrue(jnp.allclose(sparse_x, compacted_x)) + self.assertTrue(jnp.allclose(sparse_u, compacted_u)) + + def test_sparse_refresh_reports_same_overflow_as_compacting_refresh(self): + world = self._build_world() + species = self._species(world, x1=[-1.5, 0.5, 1.5], v1=[0.0, 0.0, 0.0]) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) + moved = tiled_particles._replace(x=tiled_particles.x.at[0, 0, 0, 0, 0, 0].set(0.25)) + + compacted, compacted_overflow = _refresh_tiled_particle_tiles_compacting(moved, world, tile_shape=(2, 1, 1)) + sparse, sparse_overflow = _refresh_tiled_particle_tiles_sparse(moved, world, tile_shape=(2, 1, 1)) + + self.assertEqual(bool(sparse_overflow), bool(compacted_overflow)) + self.assertEqual(int(jnp.sum(sparse.active)), int(jnp.sum(compacted.active))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/pusher_test.py b/tests/code_tests/pusher_test.py new file mode 100644 index 00000000..e66ca0cd --- /dev/null +++ b/tests/code_tests/pusher_test.py @@ -0,0 +1,447 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from tests.initial_particles import build_tiled_particles, tiled_species +from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + if g != 1: + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + interior = field[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + field_tiles = field_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + def tile_at(tx, ty, tz): + start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) + size = (tile_nx + 2, tile_ny + 2, tile_nz + 2) + return jax.lax.dynamic_slice(field, start, size) + + return jnp.stack( + [ + jnp.stack( + [ + jnp.stack([tile_at(tx, ty, tz) for tz in range(ntz)], axis=0) + for ty in range(nty) + ], + axis=0, + ) + for tx in range(ntx) + ], + axis=0, + ) + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +class TestTiledParticlePusher(unittest.TestCase): + def _build_world(self, Nx=8, Ny=6, Nz=4, shape_factor=1): + world = { + "Nx": Nx, + "Ny": Ny, + "Nz": Nz, + "dx": 4.0 / Nx, + "dy": 3.0 / Ny, + "dz": 2.0 / Nz, + "dt": 0.05, + "x_wind": 4.0, + "y_wind": 3.0, + "z_wind": 2.0, + "shape_factor": shape_factor, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "center": center_grid} + return world + + def _with_tiled_grids(self, world, tile_shape, g=1): + world["tile_shape"] = tuple(int(width) for width in tile_shape) + world["guard_cells"] = int(g) + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + world["grids"]["tiled_center_grid"] = tiled_center_grid + return world + + def _copy_world_for_tile_shape(self, world, tile_shape, g): + tiled_world = dict(world) + tiled_world["grids"] = dict(world["grids"]) + return self._with_tiled_grids(tiled_world, tile_shape, g=g) + + def _simulation_parameters_for_tile_shape(self, tile_shape): + return { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + def _push_tiled_species(self, species, world, tile_shape, E, B, constants, relativistic=True, particle_pusher="boris"): + tiled_particles, species_config = build_tiled_particles( + [species], + world, + self._simulation_parameters_for_tile_shape(tile_shape), + ) + g = int(world["guard_cells"]) + return particle_push( + tiled_particles, + species_config, + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + def _deterministic_vector_field(self, world, scale): + Nx, Ny, Nz = world["Nx"], world["Ny"], world["Nz"] + ii, jj, kk = jnp.meshgrid( + jnp.arange(Nx, dtype=float), + jnp.arange(Ny, dtype=float), + jnp.arange(Nz, dtype=float), + indexing="ij", + ) + + shape = (Nx + 2, Ny + 2, Nz + 2) + Fx = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (0.2 + 0.03 * ii - 0.02 * jj + 0.04 * kk)) + Fy = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (-0.1 + 0.05 * ii + 0.01 * jj - 0.03 * kk)) + Fz = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (0.3 - 0.04 * ii + 0.02 * jj + 0.01 * kk)) + return Fx, Fy, Fz + + def _species(self, world, active_mask=None, update_vx=True, update_vy=True, update_vz=True): + if active_mask is None: + active_mask = jnp.array([True, True, True, True]) + + return tiled_species( + name="test particles", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.25, -0.25, 0.65, 1.45]), + x2=jnp.array([-1.0, -0.25, 0.35, 1.05]), + x3=jnp.array([-0.65, -0.15, 0.25, 0.75]), + v1=jnp.array([0.2, -0.1, 0.05, 0.3]), + v2=jnp.array([0.0, 0.15, -0.2, 0.1]), + v3=jnp.array([-0.05, 0.25, 0.1, -0.15]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + active_mask=active_mask, + update_vx=update_vx, + update_vy=update_vy, + update_vz=update_vz, + ) + + def _flatten_active_by_position(self, tiled_particles): + active = tiled_particles.active.reshape(-1) + x = tiled_particles.x.reshape(-1, 3)[active] + u = tiled_particles.u.reshape(-1, 3)[active] + order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def test_particle_push_matches_one_tile_boris(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + world = self._with_tiled_grids(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + + species = self._species(world) + reference = self._species(world) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, reference_tile_shape, int(world["guard_cells"])) + reference = self._push_tiled_species( + reference, + reference_world, + reference_tile_shape, + E, + B, + constants, + relativistic=False, + particle_pusher="boris", + ) + + pushed = self._push_tiled_species( + species, + world, + tile_shape, + E, + B, + constants, + relativistic=False, + particle_pusher="boris", + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + _, reference_u = self._flatten_active_by_position(reference) + + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) + + def test_particle_push_matches_one_tile_higuera_cary(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + world = self._with_tiled_grids(world, tile_shape) + E = self._deterministic_vector_field(world, scale=0.25) + B = self._deterministic_vector_field(world, scale=0.05) + + def make_species(): + return tiled_species( + name="higuera cary particles", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.95, -1.01, 0.99, 1.95]), + x2=jnp.array([-1.35, -0.01, 0.49, 1.35]), + x3=jnp.array([-0.95, -0.01, 0.49, 0.95]), + v1=jnp.array([0.02, -0.01, 0.03, -0.015]), + v2=jnp.array([0.01, 0.02, -0.015, 0.005]), + v3=jnp.array([-0.005, 0.015, 0.01, -0.02]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + species = make_species() + reference = make_species() + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, reference_tile_shape, int(world["guard_cells"])) + reference = self._push_tiled_species( + reference, + reference_world, + reference_tile_shape, + E, + B, + constants, + particle_pusher="higuera_cary", + ) + + pushed = self._push_tiled_species( + species, + world, + tile_shape, + E, + B, + constants, + particle_pusher="higuera_cary", + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + _, reference_u = self._flatten_active_by_position(reference) + + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) + + def test_particle_push_respects_active_and_update_flags(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + world = self._with_tiled_grids(world, tile_shape) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + species = self._species( + world, + active_mask=jnp.array([True, False, True, True]), + update_vx=False, + update_vy=True, + update_vz=False, + ) + tiled_particles, species_config = build_tiled_particles([species], world, simulation_parameters) + + pushed = particle_push( + tiled_particles, + species_config, + tile_vector_field(E, world, tile_shape, num_guard_cells=1), + tile_vector_field(B, world, tile_shape, num_guard_cells=1), + world, + constants, + relativistic=False, + ) + + self.assertTrue(jnp.allclose(pushed.u[..., 0], tiled_particles.u[..., 0])) + self.assertTrue(jnp.allclose(pushed.u[..., 2], tiled_particles.u[..., 2])) + self.assertTrue(jnp.allclose(pushed.u[..., 1][~tiled_particles.active], tiled_particles.u[..., 1][~tiled_particles.active])) + + def test_particle_push_matches_relativistic_one_tile_boris_on_reduced_axes(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + constants = {"C": 10.0} + tile_shape = (2, 1, 1) + world = self._with_tiled_grids(world, tile_shape) + E = self._deterministic_vector_field(world, scale=0.1) + B = self._deterministic_vector_field(world, scale=0.02) + def make_species(): + return tiled_species( + name="one dimensional", + N_particles=3, + charge=1.0, + mass=1.0, + weight=1.0, + T=1.0, + x1=jnp.array([-1.25, 0.15, 1.25]), + x2=jnp.zeros(3), + x3=jnp.zeros(3), + v1=jnp.array([0.02, -0.01, 0.03]), + v2=jnp.array([0.0, 0.01, -0.02]), + v3=jnp.array([-0.005, 0.025, 0.01]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + species = make_species() + reference = make_species() + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, reference_tile_shape, int(world["guard_cells"])) + reference = self._push_tiled_species( + reference, + reference_world, + reference_tile_shape, + E, + B, + constants, + relativistic=True, + particle_pusher="boris", + ) + + pushed = self._push_tiled_species( + species, + world, + tile_shape, + E, + B, + constants, + relativistic=True, + particle_pusher="boris", + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + _, reference_u = self._flatten_active_by_position(reference) + + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) + + def test_particle_push_matches_one_tile_boris_on_two_guard_reduced_axes(self): + world = self._build_world(Nx=8, Ny=1, Nz=1, shape_factor=2) + constants = {"C": 10.0} + tile_shape = (2, 1, 1) + g = 2 + world = self._with_tiled_grids(world, tile_shape, g=g) + E = self._deterministic_vector_field(world, scale=0.1) + B = self._deterministic_vector_field(world, scale=0.02) + + def make_species(): + return tiled_species( + name="two guard one dimensional", + N_particles=3, + charge=1.0, + mass=1.0, + weight=1.0, + T=1.0, + x1=jnp.array([-1.25, 0.15, 1.25]), + x2=jnp.zeros(3), + x3=jnp.zeros(3), + v1=jnp.array([0.02, -0.01, 0.03]), + v2=jnp.array([0.0, 0.01, -0.02]), + v3=jnp.array([-0.005, 0.025, 0.01]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + species = make_species() + reference = make_species() + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, reference_tile_shape, g) + reference = self._push_tiled_species( + reference, + reference_world, + reference_tile_shape, + E, + B, + constants, + relativistic=True, + particle_pusher="boris", + ) + + pushed = self._push_tiled_species( + species, + world, + tile_shape, + E, + B, + constants, + relativistic=True, + particle_pusher="boris", + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + _, reference_u = self._flatten_active_by_position(reference) + + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/rho_test.py b/tests/code_tests/rho_test.py new file mode 100644 index 00000000..8d3b1906 --- /dev/null +++ b/tests/code_tests/rho_test.py @@ -0,0 +1,360 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field +from tests.initial_particles import build_tiled_particles, tiled_species +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + if g != 1: + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + interior = field[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + field_tiles = field_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + def tile_at(tx, ty, tz): + start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) + size = (tile_nx + 2, tile_ny + 2, tile_nz + 2) + return jax.lax.dynamic_slice(field, start, size) + + return jnp.stack( + [ + jnp.stack( + [ + jnp.stack([tile_at(tx, ty, tz) for tz in range(ntz)], axis=0) + for ty in range(nty) + ], + axis=0, + ) + for tx in range(ntx) + ], + axis=0, + ) + + +class TestTiledRho(unittest.TestCase): + def _build_world(self, shape_factor, dt=0.08): + x_wind, y_wind, z_wind = 4.0, 3.0, 2.0 + world = { + "dx": x_wind / 8, + "dy": y_wind / 6, + "dz": z_wind / 4, + "Nx": 8, + "Ny": 6, + "Nz": 4, + "x_wind": x_wind, + "y_wind": y_wind, + "z_wind": z_wind, + "dt": dt, + "shape_factor": shape_factor, + "guard_cells": 2, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "center": center_grid} + return world + + def _world_with_tiled_grids(self, world, simulation_parameters): + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + g = int(world["guard_cells"]) + world = dict(world) + grids = dict(world["grids"]) + world["tile_shape"] = tile_shape + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + grids["tiled_vertex_grid"] = tiled_vertex_grid + grids["tiled_center_grid"] = tiled_center_grid + world["grids"] = grids + return world + + def _empty_scalar(self, world): + return jnp.zeros((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) + + def _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + + def _one_tile_parameters(self, world): + return { + "particle_tile_nx": world["Nx"], + "particle_tile_ny": world["Ny"], + "particle_tile_nz": world["Nz"], + } + + def _particles(self, world): + electrons = tiled_species( + name="electrons", + N_particles=5, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.75, -0.65, 0.15, 1.75, world["x_wind"] / 2 - 0.03]), + x2=jnp.array([-1.15, -0.45, 0.35, 1.05, 0.0]), + x3=jnp.array([-0.75, -0.20, 0.25, 0.80, 0.0]), + v1=jnp.array([0.2, -0.1, 0.05, 0.3, 1.4]), + v2=jnp.array([0.0, 0.15, -0.2, 0.1, 0.0]), + v3=jnp.array([-0.05, 0.25, 0.1, -0.15, 0.0]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + ions = tiled_species( + name="ions", + N_particles=4, + charge=2.0, + mass=4.0, + weight=0.25, + T=1.0, + x1=jnp.array([-1.25, -0.20, 0.75, 1.35]), + x2=jnp.array([1.15, -0.75, 0.45, 0.05]), + x3=jnp.array([0.35, -0.45, 0.85, -0.15]), + v1=jnp.array([-0.1, 0.2, -0.25, 0.05]), + v2=jnp.array([0.3, -0.05, 0.15, -0.2]), + v3=jnp.array([0.1, 0.05, -0.2, 0.3]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + active_mask=jnp.array([True, False, True, True]), + ) + return [electrons, ions] + + def _tiled_with_noisy_inactive_slots(self, particles, world): + tiled_particles, species_config = build_tiled_particles(particles, world, self._simulation_parameters()) + + inactive = ~tiled_particles.active + x = tiled_particles.x.at[inactive, 0].set(0.33) + x = x.at[inactive, 1].set(-0.27) + x = x.at[inactive, 2].set(0.18) + u = tiled_particles.u.at[inactive, 0].set(4.0) + + return tiled_particles._replace(x=x, u=u), species_config + + def _zero_species_velocities(self, particles): + for species in particles: + species["u"] = jnp.zeros_like(species["u"]) + return particles + + def _zero_tiled_velocities(self, tiled_particles): + return tiled_particles._replace(u=jnp.zeros_like(tiled_particles.u)) + + def _deposit_and_assemble(self, particles, world, simulation_parameters, constants): + world = self._world_with_tiled_grids(world, simulation_parameters) + tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) + rho_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + rho_tiles = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) + rho = assemble_tiled_scalar_field( + rho_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + return rho_tiles, rho + + def _compare_tiled_to_standard(self, shape_factor, alpha): + world = self._build_world(shape_factor) + constants = {"alpha": alpha} + particles = self._particles(world) + _, rho_from_tiles = self._deposit_and_assemble(particles, world, self._simulation_parameters(), constants) + _, rho_reference = self._deposit_and_assemble(particles, world, self._one_tile_parameters(world), constants) + + self.assertTrue( + jnp.allclose( + rho_from_tiles[1:-1, 1:-1, 1:-1], + rho_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tiled_rho_matches_compute_rho_for_shape_factor_1(self): + self._compare_tiled_to_standard(shape_factor=1, alpha=1.0) + + def test_tiled_rho_matches_compute_rho_for_shape_factor_2(self): + self._compare_tiled_to_standard(shape_factor=2, alpha=1.0) + + def test_tiled_rho_matches_compute_rho_after_digital_filter(self): + self._compare_tiled_to_standard(shape_factor=2, alpha=0.55) + + def test_compute_rho_uses_current_positions_not_half_step_back_positions(self): + world = self._build_world(shape_factor=2) + constants = {"alpha": 1.0} + particles = self._particles(world) + zero_velocity_particles = self._zero_species_velocities(self._particles(world)) + + _, rho_with_velocity = self._deposit_and_assemble(particles, world, self._one_tile_parameters(world), constants) + _, rho_with_zero_velocity = self._deposit_and_assemble(zero_velocity_particles, world, self._one_tile_parameters(world), constants) + + self.assertTrue( + jnp.allclose( + rho_with_velocity[1:-1, 1:-1, 1:-1], + rho_with_zero_velocity[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tiled_global_rho_uses_current_positions_not_half_step_back_positions(self): + world = self._build_world(shape_factor=2) + constants = {"alpha": 1.0} + particles = self._particles(world) + world = self._world_with_tiled_grids(world, self._simulation_parameters()) + tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) + zero_velocity_tiled_particles = self._zero_tiled_velocities(tiled_particles) + rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + + rho_with_velocity_tiles = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) + rho_with_zero_velocity_tiles = compute_rho(zero_velocity_tiled_particles, species_config, rho_tiles, constants, world) + rho_with_velocity = assemble_tiled_scalar_field(rho_with_velocity_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + rho_with_zero_velocity = assemble_tiled_scalar_field(rho_with_zero_velocity_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + + self.assertTrue( + jnp.allclose( + rho_with_velocity[1:-1, 1:-1, 1:-1], + rho_with_zero_velocity[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world, self._simulation_parameters()) + constants = {"alpha": 1.0} + particles = self._particles(world) + tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) + zero_velocity_tiled_particles = self._zero_tiled_velocities(tiled_particles) + rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"]) + + rho_tiles_with_velocity = compute_rho( + tiled_particles, + species_config, + rho_tiles, + constants, + world, + ) + rho_tiles_with_zero_velocity = compute_rho( + zero_velocity_tiled_particles, + species_config, + rho_tiles, + constants, + world, + ) + rho_with_velocity = assemble_tiled_scalar_field(rho_tiles_with_velocity, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + rho_with_zero_velocity = assemble_tiled_scalar_field(rho_tiles_with_zero_velocity, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + + self.assertTrue( + jnp.allclose( + rho_with_velocity[1:-1, 1:-1, 1:-1], + rho_with_zero_velocity[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(self): + world = self._build_world(shape_factor=2) + world = self._world_with_tiled_grids(world, self._simulation_parameters()) + constants = {"alpha": 1.0} + particles = self._particles(world) + tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) + rho_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + rho_tiles = compute_rho( + tiled_particles, + species_config, + rho_tiles, + constants, + world, + ) + rho_from_tiles = assemble_tiled_scalar_field( + rho_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + _, rho_reference = self._deposit_and_assemble(particles, self._build_world(shape_factor=2), self._one_tile_parameters(world), constants) + + self.assertTrue( + jnp.allclose( + rho_from_tiles[1:-1, 1:-1, 1:-1], + rho_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_public_compute_rho_rejects_flat_particles(self): + world = self._build_world(shape_factor=2) + constants = {"alpha": 1.0} + + with self.assertRaisesRegex(ValueError, "TiledParticles"): + compute_rho(self._particles(world), None, self._empty_scalar(world), constants, world) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/utils_test.py b/tests/code_tests/utils_test.py new file mode 100644 index 00000000..67247755 --- /dev/null +++ b/tests/code_tests/utils_test.py @@ -0,0 +1,469 @@ +import unittest +import os +import tempfile +import importlib.util +import jax +import jax.numpy as jnp +import numpy as np +import toml +import PyPIC3D +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.initialization import initialize_fields +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from tests.initial_particles import build_tiled_particles, tiled_species +from PyPIC3D.diagnostics import plotting +from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid +from PyPIC3D.utils import ( + print_stats, check_stability, + particle_sanity_check, load_external_fields_from_toml, add_external_fields, + compute_energy, dump_parameters_to_toml, +) + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + if g != 1: + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + interior = field[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + field_tiles = field_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + def tile_at(tx, ty, tz): + start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) + size = (tile_nx + 2, tile_ny + 2, tile_nz + 2) + return jax.lax.dynamic_slice(field, start, size) + + return jnp.stack( + [ + jnp.stack( + [ + jnp.stack([tile_at(tx, ty, tz) for tz in range(ntz)], axis=0) + for ty in range(nty) + ], + axis=0, + ) + for tx in range(ntx) + ], + axis=0, + ) + + +def _field_world(Nx, Ny, Nz): + return { + "Nx": Nx, + "Ny": Ny, + "Nz": Nz, + "tile_shape": (Nx, Ny, Nz), + "guard_cells": 2, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + + +def _active_interior(world): + g = int(world["guard_cells"]) + return ( + 0, + 0, + 0, + slice(g, g + int(world["Nx"])), + slice(g, g + int(world["Ny"])), + slice(g, g + int(world["Nz"])), + ) + + +def _tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +class TestUtilsFunctions(unittest.TestCase): + def setUp(self): + self.world = { + 'Nx': 4, + 'Ny': 4, + 'Nz': 4, + 'dx': 0.1, + 'dy': 0.1, + 'dz': 0.1, + 'dt': 0.01, + 'x_wind': 1.0, + 'y_wind': 1.0, + 'z_wind': 1.0 + } + self.plasma_parameters = { + "Theoretical Plasma Frequency": 1.0, + "Debye Length": 0.01, + "Thermal Velocity": 1.0, + "Number of Electrons": 10, + "dx per debye length": 2.0 + } + + def test_build_yee_grid(self): + grid, staggered = build_yee_grid(self.world) + self.assertEqual(len(grid), 3) + self.assertEqual(len(staggered), 3) + self.assertEqual(len(grid[0]), self.world['Nx'] + 2) + self.assertEqual(len(grid[1]), self.world['Ny'] + 2) + self.assertEqual(len(grid[2]), self.world['Nz'] + 2) + self.assertEqual(len(staggered[0]), self.world['Nx'] + 2) + self.assertEqual(len(staggered[1]), self.world['Ny'] + 2) + self.assertEqual(len(staggered[2]), self.world['Nz'] + 2) + # Check that the grid and staggered arrays have the expected lengths + + def test_check_stability(self): + # Should not raise + check_stability(self.plasma_parameters, 0.01) + # Check that the stability check does not raise an error + + def test_particle_sanity_check(self): + particles = TiledParticles( + x=jnp.zeros((1, 1, 1, 1, 5, 3)), + u=jnp.zeros((1, 1, 1, 1, 5, 3)), + active=jnp.ones((1, 1, 1, 1, 5), dtype=bool), + ) + # Should not raise + particle_sanity_check(particles) + # Check that the particle sanity check does not raise an error + + def test_add_external_fields_adds_components(self): + E = (jnp.ones((2, 2, 2)), jnp.ones((2, 2, 2)) * 2, jnp.ones((2, 2, 2)) * 3) + B = (jnp.ones((2, 2, 2)) * 4, jnp.ones((2, 2, 2)) * 5, jnp.ones((2, 2, 2)) * 6) + external_E = (jnp.ones((2, 2, 2)) * 10, jnp.ones((2, 2, 2)) * 20, jnp.ones((2, 2, 2)) * 30) + external_B = (jnp.ones((2, 2, 2)) * 40, jnp.ones((2, 2, 2)) * 50, jnp.ones((2, 2, 2)) * 60) + + total_E, total_B = add_external_fields(E, B, (external_E, external_B)) + + self.assertTrue(jnp.allclose(total_E[0], 11.0)) + self.assertTrue(jnp.allclose(total_E[1], 22.0)) + self.assertTrue(jnp.allclose(total_E[2], 33.0)) + self.assertTrue(jnp.allclose(total_B[0], 44.0)) + self.assertTrue(jnp.allclose(total_B[1], 55.0)) + self.assertTrue(jnp.allclose(total_B[2], 66.0)) + + def test_load_external_fields_defaults_to_evolved_fields(self): + world = _field_world(2, 2, 2) + E, B, J, phi, rho = initialize_fields(world) + fields = [component for field in [E, B, J] for component in field] + external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "ex.npy") + np.save(path, np.ones((2, 2, 2))) + config = {"field1": {"name": "Ex", "type": 0, "path": path}} + + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) + + interior = _active_interior(world) + self.assertTrue(jnp.allclose(fields[0][interior], 1.0)) + self.assertTrue(jnp.allclose(external_fields[0][0], 0.0)) + + def test_load_external_fields_evolve_true_uses_evolved_fields(self): + world = _field_world(2, 2, 2) + E, B, J, phi, rho = initialize_fields(world) + fields = [component for field in [E, B, J] for component in field] + external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "by.npy") + np.save(path, np.ones((2, 2, 2)) * 3) + config = {"field1": {"name": "By", "type": 4, "path": path, "evolve": True}} + + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) + + interior = _active_interior(world) + self.assertTrue(jnp.allclose(fields[4][interior], 3.0)) + self.assertTrue(jnp.allclose(external_fields[1][1], 0.0)) + + def test_load_external_fields_evolve_false_uses_external_fields(self): + world = _field_world(2, 2, 2) + E, B, J, phi, rho = initialize_fields(world) + fields = [component for field in [E, B, J] for component in field] + external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "bz.npy") + np.save(path, np.ones((2, 2, 2)) * 5) + config = {"field1": {"name": "external Bz", "type": 5, "path": path, "evolve": False}} + + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) + + interior = _active_interior(world) + self.assertTrue(jnp.allclose(fields[5][interior], 0.0)) + self.assertTrue(jnp.allclose(external_fields[1][2][interior], 5.0)) + + def test_load_external_fields_rejects_external_current(self): + world = _field_world(2, 2, 2) + E, B, J, phi, rho = initialize_fields(world) + fields = [component for field in [E, B, J] for component in field] + external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "jx.npy") + np.save(path, np.ones((2, 2, 2))) + config = {"field1": {"name": "external Jx", "type": 6, "path": path, "evolve": False}} + + with self.assertRaisesRegex(ValueError, "External-only fields must be electric or magnetic"): + load_external_fields_from_toml(fields, external_fields, config, world) + + def test_load_external_fields_preserves_shape_validation(self): + world = _field_world(2, 2, 2) + E, B, J, phi, rho = initialize_fields(world) + fields = [component for field in [E, B, J] for component in field] + external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "wrong.npy") + np.save(path, np.ones((3, 2, 2))) + config = {"field1": {"name": "wrong Ex", "type": 0, "path": path, "evolve": False}} + + with self.assertRaisesRegex(ValueError, "Shape mismatch"): + load_external_fields_from_toml(fields, external_fields, config, world) + + def test_energy_can_include_external_fields(self): + world = _field_world(1, 1, 1) + E, B, J, phi, rho = initialize_fields(world) + external_E = tuple(jnp.zeros_like(comp) for comp in E) + external_B = tuple(jnp.zeros_like(comp) for comp in B) + interior = _active_interior(world) + external_E = (external_E[0].at[interior].set(2.0), external_E[1], external_E[2]) + external_B = (external_B[0], external_B[1].at[interior].set(3.0), external_B[2]) + total_E, total_B = add_external_fields(E, B, (external_E, external_B)) + + world = world | {"dx": 1.0, "dy": 1.0, "dz": 1.0} + constants = {"eps": 2.0, "mu": 4.0, "C": 10.0} + particles = TiledParticles( + x=jnp.zeros((1, 1, 1, 1, 0, 3)), + u=jnp.zeros((1, 1, 1, 1, 0, 3)), + active=jnp.zeros((1, 1, 1, 1, 0), dtype=bool), + ) + species_config = SpeciesConfig( + charge=jnp.asarray([1.0]), + mass=jnp.asarray([1.0]), + weight=jnp.asarray([1.0]), + update_x=jnp.ones((1, 3), dtype=bool), + update_u=jnp.ones((1, 3), dtype=bool), + ) + e_energy, b_energy, kinetic_energy = compute_energy(particles, total_E, total_B, world, constants, species_config=species_config) + + self.assertTrue(jnp.allclose(e_energy, 4.0)) + self.assertTrue(jnp.allclose(b_energy, 1.125)) + self.assertEqual(kinetic_energy, 0.0) + + def test_compute_energy_ignores_inactive_particles(self): + world = _field_world(1, 1, 1) | { + "dx": 1.0, + "dy": 1.0, + "dz": 1.0, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + } + E, B, J, phi, rho = initialize_fields(world) + constants = {"eps": 1.0, "mu": 1.0, "C": 10.0} + tiled_particles = TiledParticles( + x=jnp.asarray([[[[[[0.6, 0.0, 0.0]]]]]], dtype=float), + u=jnp.asarray([[[[[[1.0, 0.0, 0.0]]]]]], dtype=float), + active=jnp.asarray([[[[[False]]]]], dtype=bool), + ) + species_config = SpeciesConfig( + charge=jnp.asarray([1.0]), + mass=jnp.asarray([1.0]), + weight=jnp.asarray([1.0]), + update_x=jnp.ones((1, 3), dtype=bool), + update_u=jnp.ones((1, 3), dtype=bool), + ) + # Keep a nonzero inactive velocity so this checks the kinetic-energy mask. + + _, _, kinetic_energy = compute_energy(tiled_particles, E, B, world, constants, species_config=species_config) + + self.assertTrue(jnp.allclose(kinetic_energy, 0.0)) + + def test_dump_parameters_to_toml_writes_tiled_species_summaries(self): + active = jnp.array([[[[[True, False, True], [False, True, False]]]]]) + zeros3 = jnp.zeros(active.shape + (3,)) + particles = TiledParticles( + x=zeros3, + u=zeros3, + active=active, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + os.makedirs(os.path.join(tmpdir, "data")) + simulation_parameters = { + "output_dir": tmpdir, + "particle_species_names": ("electrons", "ions"), + "tile_shape": (2, 1, 1), + } + + dump_parameters_to_toml( + {"total_time": 1.0}, + simulation_parameters, + {}, + {}, + {}, + particles, + ) + + config = toml.load(os.path.join(tmpdir, "data/output.toml")) + + self.assertEqual(config["particles"][0]["name"], "electrons") + self.assertEqual(config["particles"][0]["storage"], "tiled") + self.assertEqual(config["particles"][0]["active_particles"], 2) + self.assertEqual(config["particles"][1]["name"], "ions") + self.assertEqual(config["particles"][1]["storage"], "tiled") + self.assertEqual(config["particles"][1]["active_particles"], 1) + self.assertEqual(config["particles"][0]["tile_shape"], [2, 1, 1]) + + def test_package_does_not_export_vtk_diagnostics(self): + self.assertFalse(hasattr(PyPIC3D, "vtk")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.vtk")) + + def test_plot_positions_flattens_tiled_particles_and_preserves_species_names(self): + species = tiled_species( + name="beam electrons", + charge=-1.0, + mass=1.0, + weight=2.0, + x1=jnp.array([-0.25, 0.25]), + x2=jnp.array([0.0, 0.0]), + x3=jnp.array([0.0, 0.0]), + u1=jnp.array([0.1, 0.2]), + u2=jnp.array([0.0, 0.0]), + u3=jnp.array([0.0, 0.0]), + ) + world = { + "Nx": 4, + "Ny": 2, + "Nz": 1, + "dx": 0.25, + "dy": 0.5, + "dz": 1.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "tile_shape": (2, 1, 1), + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + tiled_particles, species_config = build_tiled_particles([species], world, simulation_parameters=simulation_parameters) + + class FakeFigure: + def __init__(self): + self.trace_names = [] + + def add_trace(self, trace): + self.trace_names.append(trace.name) + + def update_layout(self, *args, **kwargs): + pass + + def write_html(self, *args, **kwargs): + pass + + figure = FakeFigure() + with tempfile.TemporaryDirectory() as tmpdir: + with unittest.mock.patch.object(plotting.go, "Figure", return_value=figure): + plotting.plot_positions( + tiled_particles, + 0, + world["x_wind"], + world["y_wind"], + world["z_wind"], + tmpdir, + species_config=species_config, + species_names=("beam electrons",), + world=world, + ) + + self.assertEqual(figure.trace_names, ["beam electrons"]) + + def test_particles_phase_space_flattens_tiled_particles(self): + species = tiled_species( + name="electrons", + charge=-1.0, + mass=1.0, + weight=1.0, + x1=jnp.array([0.0]), + x2=jnp.array([0.0]), + x3=jnp.array([0.0]), + u1=jnp.array([0.2]), + u2=jnp.array([0.0]), + u3=jnp.array([0.0]), + ) + world = { + "Nx": 4, + "Ny": 2, + "Nz": 1, + "dx": 0.25, + "dy": 0.5, + "dz": 1.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "tile_shape": (2, 1, 1), + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + tiled_particles, species_config = build_tiled_particles([species], world, simulation_parameters=simulation_parameters) + + with tempfile.TemporaryDirectory() as tmpdir: + os.makedirs(os.path.join(tmpdir, "data/phase_space/x")) + os.makedirs(os.path.join(tmpdir, "data/phase_space/y")) + os.makedirs(os.path.join(tmpdir, "data/phase_space/z")) + with unittest.mock.patch.object(plotting.plt, "scatter") as scatter: + with unittest.mock.patch.object(plotting.plt, "savefig"): + plotting.particles_phase_space( + tiled_particles, + world, + 0, + "electrons", + tmpdir, + species_config=species_config, + species_names=("electrons",), + ) + + self.assertEqual(scatter.call_count, 3) + +if __name__ == '__main__': + unittest.main() diff --git a/tests/code_tests/yee_test.py b/tests/code_tests/yee_test.py new file mode 100644 index 00000000..5d17deb0 --- /dev/null +++ b/tests/code_tests/yee_test.py @@ -0,0 +1,635 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.solvers.first_order_yee import ( + update_B, + update_E, +) +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC + + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + + if g != 1: + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + ix = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + interior = field[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + field_tiles = field_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + + def tile_at(tx, ty, tz): + start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) + size = (tile_nx + 2, tile_ny + 2, tile_nz + 2) + return jax.lax.dynamic_slice(field, start, size) + + return jnp.stack( + [ + jnp.stack( + [ + jnp.stack([tile_at(tx, ty, tz) for tz in range(ntz)], axis=0) + for ty in range(nty) + ], + axis=0, + ) + for tx in range(ntx) + ], + axis=0, + ) + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +def _update_ghost_cells(field, bc_x, bc_y, bc_z): + field = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda f: f.at[0, :, :].set(f[-2, :, :]).at[-1, :, :].set(f[1, :, :]), + lambda f: f.at[0, :, :].set(0.0).at[-1, :, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda f: f.at[:, 0, :].set(f[:, -2, :]).at[:, -1, :].set(f[:, 1, :]), + lambda f: f.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda f: f.at[:, :, 0].set(f[:, :, -2]).at[:, :, -1].set(f[:, :, 1]), + lambda f: f.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0), + operand=field, + ) + return field + + +class TestYeeTiled(unittest.TestCase): + def _build_world(self): + world = { + "Nx": 8, + "Ny": 6, + "Nz": 4, + "dx": 0.5, + "dy": 0.5, + "dz": 0.5, + "dt": 0.05, + "x_wind": 4.0, + "y_wind": 3.0, + "z_wind": 2.0, + "shape_factor": 1, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "center": center_grid} + return world + + def _conducting_world(self): + world = self._build_world() + world["boundary_conditions"] = {"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING} + return world + + def _with_tile_metadata(self, world, tile_shape, g=2): + world["tile_shape"] = tuple(int(width) for width in tile_shape) + world["guard_cells"] = int(g) + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world, tile_shape, g) + world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + world["grids"]["tiled_center_grid"] = tiled_center_grid + return world + + def _mixed_bc_world(self): + world = self._build_world() + world["boundary_conditions"] = {"x": BC_PERIODIC, "y": BC_CONDUCTING, "z": BC_CONDUCTING} + return world + + def _fill_ghosts(self, field, world): + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + return _update_ghost_cells(field, bc_x, bc_y, bc_z) + + def _deterministic_vector_field(self, world, scale): + Nx, Ny, Nz = world["Nx"], world["Ny"], world["Nz"] + ii, jj, kk = jnp.meshgrid( + jnp.arange(Nx, dtype=float), + jnp.arange(Ny, dtype=float), + jnp.arange(Nz, dtype=float), + indexing="ij", + ) + + shape = (Nx + 2, Ny + 2, Nz + 2) + Fx = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (0.2 + 0.03 * ii - 0.02 * jj + 0.04 * kk)) + Fy = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (-0.1 + 0.05 * ii + 0.01 * jj - 0.03 * kk)) + Fz = jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (0.3 - 0.04 * ii + 0.02 * jj + 0.01 * kk)) + + return tuple(self._fill_ghosts(component, world) for component in (Fx, Fy, Fz)) + + def _copy_world_for_tile_shape(self, world, tile_shape, g=2): + reference_world = dict(world) + reference_world["boundary_conditions"] = dict(world["boundary_conditions"]) + reference_world["grids"] = dict(world["grids"]) + reference_world["tile_shape"] = tuple(int(width) for width in tile_shape) + reference_world["guard_cells"] = int(g) + return reference_world + + def _reference_update_E(self, E, B, J, world, constants): + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, tile_shape, int(world["guard_cells"])) + E_reference, pml_state = update_E( + tile_vector_field(E, reference_world, tile_shape, num_guard_cells=int(reference_world["guard_cells"])), + tile_vector_field(B, reference_world, tile_shape, num_guard_cells=int(reference_world["guard_cells"])), + tile_vector_field(J, reference_world, tile_shape, num_guard_cells=int(reference_world["guard_cells"])), + reference_world, + constants, + ) + return assemble_tiled_vector_field( + E_reference, + reference_world, + tile_shape, + num_guard_cells=int(reference_world["guard_cells"]), + ), pml_state + + def _reference_update_B(self, E, B, world, constants): + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._copy_world_for_tile_shape(world, tile_shape, int(world["guard_cells"])) + B_reference, pml_state = update_B( + tile_vector_field(E, reference_world, tile_shape, num_guard_cells=int(reference_world["guard_cells"])), + tile_vector_field(B, reference_world, tile_shape, num_guard_cells=int(reference_world["guard_cells"])), + reference_world, + constants, + ) + return assemble_tiled_vector_field( + B_reference, + reference_world, + tile_shape, + num_guard_cells=int(reference_world["guard_cells"]), + ), pml_state + + def _reference_yee_step(self, E, B, J, world, constants): + E_reference, pml_state = self._reference_update_E(E, B, J, world, constants) + B_reference, pml_state = self._reference_update_B(E_reference, B, world, constants) + return E_reference, B_reference, pml_state + + def _fill_guard_cells(self, field, world, num_guard_cells): + g = num_guard_cells + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + if bc_x == BC_PERIODIC: + field = field.at[:g, :, :].set(field[-2 * g:-g, :, :]) + field = field.at[-g:, :, :].set(field[g:2 * g, :, :]) + else: + field = field.at[:g, :, :].set(0.0) + field = field.at[-g:, :, :].set(0.0) + + if bc_y == BC_PERIODIC: + field = field.at[:, :g, :].set(field[:, -2 * g:-g, :]) + field = field.at[:, -g:, :].set(field[:, g:2 * g, :]) + else: + field = field.at[:, :g, :].set(0.0) + field = field.at[:, -g:, :].set(0.0) + + if bc_z == BC_PERIODIC: + field = field.at[:, :, :g].set(field[:, :, -2 * g:-g]) + field = field.at[:, :, -g:].set(field[:, :, g:2 * g]) + else: + field = field.at[:, :, :g].set(0.0) + field = field.at[:, :, -g:].set(0.0) + + return field + + def _tile_scalar_field_with_guard(self, field, tile_shape, num_guard_cells): + g = num_guard_cells + tile_nx, tile_ny, tile_nz = tile_shape + Nx = int(field.shape[0]) - 2 * g + Ny = int(field.shape[1]) - 2 * g + Nz = int(field.shape[2]) - 2 * g + ntx = Nx // tile_nx + nty = Ny // tile_ny + ntz = Nz // tile_nz + + tiles = [] + for tx in range(ntx): + y_tiles = [] + for ty in range(nty): + z_tiles = [] + for tz in range(ntz): + ix = tx * tile_nx + iy = ty * tile_ny + iz = tz * tile_nz + z_tiles.append( + field[ + ix:ix + tile_nx + 2 * g, + iy:iy + tile_ny + 2 * g, + iz:iz + tile_nz + 2 * g, + ] + ) + y_tiles.append(jnp.stack(z_tiles, axis=0)) + tiles.append(jnp.stack(y_tiles, axis=0)) + + return jnp.stack(tiles, axis=0) + + def test_tile_vector_field_assembles_to_original_ghost_celled_field(self): + world = self._build_world() + tile_shape = (2, 3, 2) + E = self._deterministic_vector_field(world, scale=1.0) + + E_tiles = tile_vector_field(E, world, tile_shape) + E_assembled = assemble_tiled_vector_field(E_tiles, world, tile_shape) + + for original, assembled in zip(E, E_assembled): + self.assertTrue(jnp.allclose(assembled, original, rtol=1.0e-12, atol=1.0e-12)) + + def test_tile_grid_axes_include_configured_guard_cells(self): + world = self._build_world() + tile_shape = (2, 3, 2) + g = 2 + world = self._with_tile_metadata(world, tile_shape, g) + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world) + + self.assertEqual(tiled_center_grid[0].shape, (4, 2, 2, tile_shape[0] + 2 * g)) + self.assertEqual(tiled_center_grid[1].shape, (4, 2, 2, tile_shape[1] + 2 * g)) + self.assertEqual(tiled_center_grid[2].shape, (4, 2, 2, tile_shape[2] + 2 * g)) + + for tx in range(4): + for ty in range(2): + for tz in range(2): + center_x = world["grids"]["center"][0][0] + ( + jnp.arange(tile_shape[0] + 2 * g) + tx * tile_shape[0] - (g - 1) + ) * world["dx"] + vertex_y = world["grids"]["vertex"][1][0] + ( + jnp.arange(tile_shape[1] + 2 * g) + ty * tile_shape[1] - (g - 1) + ) * world["dy"] + + self.assertTrue(jnp.allclose(tiled_center_grid[0][tx, ty, tz], center_x)) + self.assertTrue(jnp.allclose(tiled_vertex_grid[1][tx, ty, tz], vertex_y)) + + def test_update_tiled_ghost_cells_periodic_refreshes_neighbor_halos(self): + world = self._build_world() + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape, g=2) + field = self._deterministic_vector_field(world, scale=1.0)[0] + tiles = tile_vector_field((field,), world, tile_shape)[0] + + stale_tiles = tiles.at[:, :, :, 0, :, :].set(-100.0) + stale_tiles = stale_tiles.at[:, :, :, -1, :, :].set(-200.0) + stale_tiles = stale_tiles.at[:, :, :, :, 0, :].set(-300.0) + stale_tiles = stale_tiles.at[:, :, :, :, -1, :].set(-400.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) + + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + + self.assertTrue(jnp.allclose(refreshed, tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_tiled_vector_ghost_cells_periodic_refreshes_each_component(self): + world = self._build_world() + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape, g=2) + E = self._deterministic_vector_field(world, scale=1.0) + E_tiles = tile_vector_field(E, world, tile_shape) + + stale_tiles = tuple(component.at[:, :, :, 0, :, :].set(-10.0 * (i + 1)) for i, component in enumerate(E_tiles)) + refreshed = ghost_cells.update_tiled_vector_ghost_cells(stale_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + + for original_tiles, refreshed_component in zip(E_tiles, refreshed): + self.assertTrue(jnp.allclose(refreshed_component, original_tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_tiled_ghost_cells_conducting_matches_global_ghost_cells(self): + world = self._conducting_world() + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape, g=2) + field = self._deterministic_vector_field(world, scale=1.0)[0] + tiles = tile_vector_field((field,), world, tile_shape)[0] + + stale_tiles = tiles.at[:, :, :, 0, :, :].set(-100.0) + stale_tiles = stale_tiles.at[:, :, :, -1, :, :].set(-200.0) + stale_tiles = stale_tiles.at[:, :, :, :, 0, :].set(-300.0) + stale_tiles = stale_tiles.at[:, :, :, :, -1, :].set(-400.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) + + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world) + reference = _update_ghost_cells( + field, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + reference_tiles = tile_vector_field((reference,), world, tile_shape)[0] + + self.assertTrue(jnp.allclose(refreshed, reference_tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_tiled_ghost_cells_mixed_matches_global_ghost_cells(self): + world = self._mixed_bc_world() + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape, g=2) + field = self._deterministic_vector_field(world, scale=1.0)[0] + tiles = tile_vector_field((field,), world, tile_shape)[0] + + stale_tiles = tiles.at[:, :, :, 0, :, :].set(-100.0) + stale_tiles = stale_tiles.at[:, :, :, -1, :, :].set(-200.0) + stale_tiles = stale_tiles.at[:, :, :, :, 0, :].set(-300.0) + stale_tiles = stale_tiles.at[:, :, :, :, -1, :].set(-400.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) + + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world) + reference = _update_ghost_cells( + field, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + reference_tiles = tile_vector_field((reference,), world, tile_shape)[0] + + self.assertTrue(jnp.allclose(refreshed, reference_tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_tiled_ghost_cells_two_guard_layers_matches_global_refresh(self): + world = self._mixed_bc_world() + tile_shape = (2, 3, 2) + num_guard_cells = 2 + world = self._with_tile_metadata(world, tile_shape, g=num_guard_cells) + Nx, Ny, Nz = world["Nx"], world["Ny"], world["Nz"] + shape = ( + Nx + 2 * num_guard_cells, + Ny + 2 * num_guard_cells, + Nz + 2 * num_guard_cells, + ) + + field = jnp.arange(jnp.prod(jnp.asarray(shape)), dtype=jnp.float64).reshape(shape) + field = self._fill_guard_cells(field, world, num_guard_cells) + tiles = self._tile_scalar_field_with_guard(field, tile_shape, num_guard_cells) + + stale_tiles = tiles.at[:, :, :, :num_guard_cells, :, :].set(-100.0) + stale_tiles = stale_tiles.at[:, :, :, -num_guard_cells:, :, :].set(-200.0) + stale_tiles = stale_tiles.at[:, :, :, :, :num_guard_cells, :].set(-300.0) + stale_tiles = stale_tiles.at[:, :, :, :, -num_guard_cells:, :].set(-400.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, :num_guard_cells].set(-500.0) + stale_tiles = stale_tiles.at[:, :, :, :, :, -num_guard_cells:].set(-600.0) + + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world, num_guard_cells) + + self.assertTrue(jnp.allclose(refreshed, tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_tiled_vector_ghost_cells_conducting_refreshes_each_component(self): + world = self._conducting_world() + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape, g=2) + E = self._deterministic_vector_field(world, scale=1.0) + E_tiles = tile_vector_field(E, world, tile_shape) + + stale_tiles = tuple(component.at[:, :, :, 0, :, :].set(-10.0 * (i + 1)) for i, component in enumerate(E_tiles)) + refreshed = ghost_cells.update_tiled_vector_ghost_cells(stale_tiles, world) + + for original, refreshed_component in zip(E, refreshed): + reference = _update_ghost_cells( + original, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + reference_tiles = tile_vector_field((reference,), world, tile_shape)[0] + self.assertTrue(jnp.allclose(refreshed_component, reference_tiles, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_E_matches_single_tile_yee_update(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + + E_reference, _ = self._reference_update_E(E, B, J, world, constants) + E_tiled, pml_state = update_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + ) + self.assertIsNone(pml_state) + E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) + + for reference, tiled in zip(E_reference, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_E_is_the_public_tiled_update(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + E_tiles = tile_vector_field(E, world, tile_shape, num_guard_cells=2) + B_tiles = tile_vector_field(B, world, tile_shape, num_guard_cells=2) + J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=2) + + E_public, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + + self.assertIsNone(pml_state) + self.assertEqual(E_public[0].ndim, 6) + self.assertEqual(E_public[0].shape[:3], E_tiles[0].shape[:3]) + + def test_update_E_matches_single_tile_yee_update_with_conducting_boundaries(self): + world = self._conducting_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + + E_reference, _ = self._reference_update_E(E, B, J, world, constants) + E_tiled, pml_state = update_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + ) + self.assertIsNone(pml_state) + E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) + + for reference, tiled in zip(E_reference, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_B_matches_single_tile_yee_update(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + + B_reference, _ = self._reference_update_B(E, B, world, constants) + B_tiled, pml_state = update_B( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + ) + self.assertIsNone(pml_state) + B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) + + for reference, tiled in zip(B_reference, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_update_B_is_the_public_tiled_update(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + E_tiles = tile_vector_field(E, world, tile_shape, num_guard_cells=2) + B_tiles = tile_vector_field(B, world, tile_shape, num_guard_cells=2) + + B_public, pml_state = update_B(E_tiles, B_tiles, world, constants) + + self.assertIsNone(pml_state) + self.assertEqual(B_public[0].ndim, 6) + self.assertEqual(B_public[0].shape[:3], B_tiles[0].shape[:3]) + + def test_update_B_matches_single_tile_yee_update_with_conducting_boundaries(self): + world = self._conducting_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + + B_reference, _ = self._reference_update_B(E, B, world, constants) + B_tiled, pml_state = update_B( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + ) + self.assertIsNone(pml_state) + B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) + + for reference, tiled in zip(B_reference, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_electrodynamic_step_matches_single_tile_yee_sequence(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + + E_reference, B_reference, _ = self._reference_yee_step(E, B, J, world, constants) + + E_tiles = tile_vector_field(E, world, tile_shape) + B_tiles = tile_vector_field(B, world, tile_shape) + J_tiles = tile_vector_field(J, world, tile_shape) + E_tiles, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_B(E_tiles, B_tiles, world, constants, pml_state) + self.assertIsNone(pml_state) + + E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) + + for reference, tiled in zip(E_reference, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(B_reference, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_electrodynamic_step_matches_single_tile_yee_sequence_with_conducting_boundaries(self): + world = self._conducting_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + + E_reference, B_reference, _ = self._reference_yee_step(E, B, J, world, constants) + + E_tiles = tile_vector_field(E, world, tile_shape) + B_tiles = tile_vector_field(B, world, tile_shape) + J_tiles = tile_vector_field(J, world, tile_shape) + E_tiles, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_B(E_tiles, B_tiles, world, constants, pml_state) + self.assertIsNone(pml_state) + + E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) + + for reference, tiled in zip(E_reference, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(B_reference, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_evolve_updates_fields_without_pushing_particles(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world = self._with_tile_metadata(world, tile_shape) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + J = self._deterministic_vector_field(world, scale=0.05) + + E_reference, B_reference, _ = self._reference_yee_step(E, B, J, world, constants) + E_tiles = tile_vector_field(E, world, tile_shape) + B_tiles = tile_vector_field(B, world, tile_shape) + J_tiles = tile_vector_field(J, world, tile_shape) + E_tiles, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_B(E_tiles, B_tiles, world, constants, pml_state) + + E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) + + self.assertIsNone(pml_state) + for reference, tiled in zip(E_reference, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(B_reference, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for original, after in zip(tile_vector_field(J, world, tile_shape), J_tiles): + self.assertTrue(jnp.allclose(after, original, rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py deleted file mode 100644 index 7c38a36d..00000000 --- a/tests/electrostatic_yee_test.py +++ /dev/null @@ -1,205 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.solvers.electrostatic_yee import ( - solve_poisson_with_fft, - solve_poisson_with_conjugate_gradient, - calculate_electrostatic_fields, -) -from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient -from PyPIC3D.solvers.pstd import spectral_gradient -from PyPIC3D.utils import build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells - -jax.config.update("jax_enable_x64", True) - - -def apply_negative_laplacian(field, dx, dy, dz): - laplacian_x = (jnp.roll(field, shift=1, axis=0) + jnp.roll(field, shift=-1, axis=0) - 2.0 * field) / (dx * dx) - laplacian_y = (jnp.roll(field, shift=1, axis=1) + jnp.roll(field, shift=-1, axis=1) - 2.0 * field) / (dy * dy) - laplacian_z = (jnp.roll(field, shift=1, axis=2) + jnp.roll(field, shift=-1, axis=2) - 2.0 * field) / (dz * dz) - return -(laplacian_x + laplacian_y + laplacian_z) - - -class TestElectrostaticYeeMethods(unittest.TestCase): - def setUp(self): - self.Nx = 16 - self.Ny = 16 - self.Nz = 16 - self.x_wind = 2 * jnp.pi - self.y_wind = 2 * jnp.pi - self.z_wind = 2 * jnp.pi - self.dx = self.x_wind / self.Nx - self.dy = self.y_wind / self.Ny - self.dz = self.z_wind / self.Nz - - # Interior coordinate grid (no ghost cells) for analytical solutions - x = jnp.linspace(0, self.x_wind, self.Nx, endpoint=False) - y = jnp.linspace(0, self.y_wind, self.Ny, endpoint=False) - z = jnp.linspace(0, self.z_wind, self.Nz, endpoint=False) - self.X, self.Y, self.Z = jnp.meshgrid(x, y, z, indexing='ij') - - self.constants = { - "eps": 1.0, - "alpha": 1.0, - } - self.world = { - "Nx": self.Nx, - "Ny": self.Ny, - "Nz": self.Nz, - "dx": self.dx, - "dy": self.dy, - "dz": self.dz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(self.world) - self.world["grids"] = { - "vertex": vertex_grid, - "center": center_grid, - } - - # Ghost-celled initial fields: shape (Nx+2, Ny+2, Nz+2) - self.initial_rho = jnp.zeros((self.Nx + 2, self.Ny + 2, self.Nz + 2)) - self.initial_phi = jnp.zeros((self.Nx + 2, self.Ny + 2, self.Nz + 2)) - - self.particles = [ - particle_species( - name="test", - N_particles=1, - charge=1.0, - mass=1.0, - T=1.0, - v1=jnp.zeros(1), - v2=jnp.zeros(1), - v3=jnp.zeros(1), - x1=jnp.array([0.1]), - x2=jnp.array([0.2]), - x3=jnp.array([0.3]), - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - dx=self.dx, - dy=self.dy, - dz=self.dz, - weight=1.0, - x_bc="periodic", - y_bc="periodic", - z_bc="periodic", - shape=1, - dt=0.0, - ) - ] - - def test_solve_poisson_with_fft_single_mode(self): - phi_true_interior = jnp.sin(self.X + self.Y + self.Z) - rho_interior = 3.0 * self.constants["eps"] * phi_true_interior - - # Place rho into ghost-celled array - rho = jnp.zeros((self.Nx + 2, self.Ny + 2, self.Nz + 2)) - rho = rho.at[1:-1, 1:-1, 1:-1].set(rho_interior) - rho = update_ghost_cells(rho, 0, 0, 0) - - phi = solve_poisson_with_fft(rho, self.constants, self.world) - - # Compare on interior - phi_num = phi[1:-1, 1:-1, 1:-1] - phi_num = phi_num - jnp.mean(phi_num) - phi_true = phi_true_interior - jnp.mean(phi_true_interior) - - self.assertEqual(phi_num.shape, phi_true.shape) - self.assertTrue(jnp.allclose(phi_num, phi_true, atol=1e-10, rtol=1e-8)) - - def test_solve_poisson_with_conjugate_gradient_single_mode(self): - phi_true_interior = jnp.sin(self.X + self.Y + self.Z) - rhs = apply_negative_laplacian(phi_true_interior, self.dx, self.dy, self.dz) - rho_interior = rhs * self.constants["eps"] - - # Place rho into ghost-celled array - rho = jnp.zeros((self.Nx + 2, self.Ny + 2, self.Nz + 2)) - rho = rho.at[1:-1, 1:-1, 1:-1].set(rho_interior) - rho = update_ghost_cells(rho, 0, 0, 0) - - phi = solve_poisson_with_conjugate_gradient( - rho, - self.initial_phi, - self.constants, - self.world, - tol=1e-10, - max_iter=4000, - ) - - # Compare on interior - phi_num = phi[1:-1, 1:-1, 1:-1] - phi_num = phi_num - jnp.mean(phi_num) - phi_true = phi_true_interior - jnp.mean(phi_true_interior) - - self.assertEqual(phi_num.shape, phi_true.shape) - self.assertTrue(jnp.allclose(phi_num, phi_true, atol=1e-7, rtol=1e-6)) - - residual = apply_negative_laplacian(phi_num, self.dx, self.dy, self.dz) - rho_interior / self.constants["eps"] - self.assertLess(jnp.max(jnp.abs(residual)), 1e-6) - - def test_solver_mode_mapping(self): - E_spectral, phi_spectral, rho_spectral = calculate_electrostatic_fields( - self.world, - self.particles, - self.constants, - self.initial_rho, - self.initial_phi, - "spectral", - "periodic", - ) - - # The spectral solver returns ghost-celled arrays - # Verify shapes - self.assertEqual(phi_spectral.shape, (self.Nx + 2, self.Ny + 2, self.Nz + 2)) - self.assertEqual(E_spectral[0].shape, (self.Nx + 2, self.Ny + 2, self.Nz + 2)) - - # Verify spectral path: solve Poisson on interior, then gradient on interior - expected_phi = solve_poisson_with_fft(rho_spectral, self.constants, self.world) - expected_Ex, expected_Ey, expected_Ez = spectral_gradient(-1.0 * expected_phi[1:-1, 1:-1, 1:-1], self.world) - - self.assertTrue(jnp.allclose(phi_spectral[1:-1, 1:-1, 1:-1], expected_phi[1:-1, 1:-1, 1:-1], atol=1e-8, rtol=1e-8)) - self.assertTrue(jnp.allclose(E_spectral[0][1:-1, 1:-1, 1:-1], expected_Ex, atol=1e-8, rtol=1e-8)) - self.assertTrue(jnp.allclose(E_spectral[1][1:-1, 1:-1, 1:-1], expected_Ey, atol=1e-8, rtol=1e-8)) - self.assertTrue(jnp.allclose(E_spectral[2][1:-1, 1:-1, 1:-1], expected_Ez, atol=1e-8, rtol=1e-8)) - - E_fdtd, phi_fdtd, rho_fdtd = calculate_electrostatic_fields( - self.world, - self.particles, - self.constants, - self.initial_rho, - self.initial_phi, - "fdtd", - "periodic", - ) - - expected_phi_fdtd = solve_poisson_with_conjugate_gradient( - rho_fdtd, - self.initial_phi, - self.constants, - self.world, - ) - expected_Ex_fdtd, expected_Ey_fdtd, expected_Ez_fdtd = centered_finite_difference_gradient( - -1.0 * expected_phi_fdtd[1:-1, 1:-1, 1:-1], - self.dx, - self.dy, - self.dz, - "periodic", - ) - - self.assertTrue(jnp.allclose(rho_fdtd[1:-1, 1:-1, 1:-1], rho_spectral[1:-1, 1:-1, 1:-1], atol=1e-12, rtol=1e-12)) - self.assertTrue(jnp.allclose(phi_fdtd[1:-1, 1:-1, 1:-1], expected_phi_fdtd[1:-1, 1:-1, 1:-1], atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[0][1:-1, 1:-1, 1:-1], expected_Ex_fdtd, atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[1][1:-1, 1:-1, 1:-1], expected_Ey_fdtd, atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[2][1:-1, 1:-1, 1:-1], expected_Ez_fdtd, atol=1e-6, rtol=1e-6)) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/errors_test.py b/tests/errors_test.py deleted file mode 100644 index c281e548..00000000 --- a/tests/errors_test.py +++ /dev/null @@ -1,76 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - -# # Add the parent directory to the sys.path -# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) - -from PyPIC3D.errors import ( - compute_pe, compute_magnetic_divergence_error, compute_electric_divergence_error -) - -jax.config.update("jax_enable_x64", True) - -class TestErrors(unittest.TestCase): - - def setUp(self): - self.constants = {'eps': 1.0} - self.world = {'dx': 1.0, 'dy': 1.0, 'dz': 1.0} - - def test_compute_pe(self): - x = jnp.linspace(0, 1, 10) - y = jnp.linspace(0, 1, 10) - z = jnp.linspace(0, 1, 10) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - - phi = X**2 + Y**2 + Z**2 - eps = 1.0 - rho = -6 * jnp.ones_like(phi) # Laplacian of phi = -6 - dx, dy, dz = 1.0 / 9, 1.0 / 9, 1.0 / 9 - - laplacian = -6 * jnp.ones_like(phi) - - result_spectral = compute_pe(phi, rho, self.constants, self.world, solver='spectral') - result_fdtd = compute_pe(phi, rho, self.constants, self.world, solver='fdtd') - - self.assertIsInstance(result_spectral, jnp.ndarray) - self.assertIsInstance(result_fdtd, jnp.ndarray) - - - - def test_compute_magnetic_divergence_error(self): - key = jax.random.PRNGKey(0) - Bx_key, By_key, Bz_key = jax.random.split(key, 3) - Bx = jax.random.uniform(Bx_key, shape=(10, 10, 10)) - By = jax.random.uniform(By_key, shape=(10, 10, 10)) - Bz = jax.random.uniform(Bz_key, shape=(10, 10, 10)) - dx, dy, dz = 1.0, 1.0, 1.0 - - - result_spectral = compute_magnetic_divergence_error(Bx, By, Bz, self.world, solver='spectral') - result_fdtd = compute_magnetic_divergence_error(Bx, By, Bz, self.world, solver='fdtd') - - self.assertIsInstance(result_spectral, jnp.ndarray) - self.assertIsInstance(result_fdtd, jnp.ndarray) - - def test_compute_electric_divergence_error(self): - key = jax.random.PRNGKey(1) - Ex_key, Ey_key, Ez_key, rho_key = jax.random.split(key, 4) - - Ex = jax.random.uniform(Ex_key, shape=(10, 10, 10)) - Ey = jax.random.uniform(Ey_key, shape=(10, 10, 10)) - Ez = jax.random.uniform(Ez_key, shape=(10, 10, 10)) - rho = jax.random.uniform(rho_key, shape=(10, 10, 10)) - eps = 1.0 - dx, dy, dz = 1.0, 1.0, 1.0 - - - result_spectral = compute_electric_divergence_error(Ex, Ey, Ez, rho, self.constants, self.world, solver='spectral') - result_fdtd = compute_electric_divergence_error(Ex, Ey, Ez, rho, self.constants, self.world, solver='fdtd') - - self.assertIsInstance(result_spectral, jnp.ndarray) - self.assertIsInstance(result_fdtd, jnp.ndarray) -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/esirkepov_ghost_cells_test.py b/tests/esirkepov_ghost_cells_test.py deleted file mode 100644 index 66250ac3..00000000 --- a/tests/esirkepov_ghost_cells_test.py +++ /dev/null @@ -1,84 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.deposition.Esirkepov import ( - eliminate_esirkepov_ghost_cells, - enforce_bc_along_axis, -) - -jax.config.update("jax_enable_x64", True) - - -class TestExtendedEsirkepovGhostCells(unittest.TestCase): - - def test_two_periodic_ghost_layers_fold_to_opposite_interior(self): - field = jnp.zeros((8, 3, 3)) - field = field.at[0, 1, 1].set(1.0) - field = field.at[1, 1, 1].set(2.0) - field = field.at[-2, 1, 1].set(3.0) - field = field.at[-1, 1, 1].set(4.0) - field = field.at[2, 1, 1].set(10.0) - field = field.at[3, 1, 1].set(20.0) - field = field.at[-4, 1, 1].set(30.0) - field = field.at[-3, 1, 1].set(40.0) - - folded = enforce_bc_along_axis(field, axis=0, bc="periodic", component_axis=0) - - self.assertAlmostEqual(float(folded[-4, 1, 1]), 31.0) - self.assertAlmostEqual(float(folded[-3, 1, 1]), 42.0) - self.assertAlmostEqual(float(folded[2, 1, 1]), 13.0) - self.assertAlmostEqual(float(folded[3, 1, 1]), 24.0) - self.assertEqual(float(folded[0, 1, 1]), 0.0) - self.assertEqual(float(folded[1, 1, 1]), 0.0) - self.assertEqual(float(folded[-2, 1, 1]), 0.0) - self.assertEqual(float(folded[-1, 1, 1]), 0.0) - - def test_two_reflecting_ghost_layers_fold_to_same_side_with_normal_sign(self): - field = jnp.zeros((8, 3, 3)) - field = field.at[0, 1, 1].set(1.0) - field = field.at[1, 1, 1].set(2.0) - field = field.at[-2, 1, 1].set(3.0) - field = field.at[-1, 1, 1].set(4.0) - field = field.at[2, 1, 1].set(10.0) - field = field.at[3, 1, 1].set(20.0) - field = field.at[-4, 1, 1].set(30.0) - field = field.at[-3, 1, 1].set(40.0) - - folded = enforce_bc_along_axis(field, axis=0, bc="reflecting", component_axis=0) - - self.assertAlmostEqual(float(folded[2, 1, 1]), 9.0) - self.assertAlmostEqual(float(folded[3, 1, 1]), 18.0) - self.assertAlmostEqual(float(folded[-4, 1, 1]), 27.0) - self.assertAlmostEqual(float(folded[-3, 1, 1]), 36.0) - - def test_two_absorbing_ghost_layers_are_discarded(self): - field = jnp.zeros((8, 3, 3)) - field = field.at[0, 1, 1].set(1.0) - field = field.at[1, 1, 1].set(2.0) - field = field.at[-2, 1, 1].set(3.0) - field = field.at[-1, 1, 1].set(4.0) - field = field.at[2, 1, 1].set(10.0) - field = field.at[-3, 1, 1].set(40.0) - - folded = enforce_bc_along_axis(field, axis=0, bc="absorbing", component_axis=0) - - self.assertAlmostEqual(float(folded[2, 1, 1]), 10.0) - self.assertAlmostEqual(float(folded[-3, 1, 1]), 40.0) - self.assertEqual(float(folded[0, 1, 1]), 0.0) - self.assertEqual(float(folded[1, 1, 1]), 0.0) - self.assertEqual(float(folded[-2, 1, 1]), 0.0) - self.assertEqual(float(folded[-1, 1, 1]), 0.0) - - def test_eliminate_esirkepov_ghost_cells_returns_one_ghost_layer(self): - field = jnp.arange(8 * 3 * 3, dtype=float).reshape((8, 3, 3)) - - restricted = eliminate_esirkepov_ghost_cells(field) - - self.assertEqual(restricted.shape, (6, 1, 1)) - self.assertTrue(jnp.array_equal(restricted[:, 0, 0], field[1:-1, 1, 1])) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/evolve_test.py b/tests/evolve_test.py deleted file mode 100644 index 74ba461a..00000000 --- a/tests/evolve_test.py +++ /dev/null @@ -1,164 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.initialization import initialize_fields -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.utils import build_yee_grid - -jax.config.update("jax_enable_x64", True) - - -def zero_current(particles, J, constants, world): - return tuple(jnp.zeros_like(comp) for comp in J) - - -def unused_curl(Ex, Ey, Ez): - return None - - -class TestEvolveExternalFields(unittest.TestCase): - def test_external_electric_field_pushes_particles_without_evolving_maxwell_fields(self): - world = { - "Nx": 3, - "Ny": 3, - "Nz": 3, - "dx": 1.0 / 3.0, - "dy": 1.0 / 3.0, - "dz": 1.0 / 3.0, - "dt": 0.1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_yee_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} - - E, B, J, phi, rho = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) - external_E = tuple(jnp.zeros_like(comp) for comp in E) - external_B = tuple(jnp.zeros_like(comp) for comp in B) - external_E = (jnp.ones_like(external_E[0]), external_E[1], external_E[2]) - external_fields = (external_E, external_B) - - particles = [ - particle_species( - name="test", - N_particles=1, - charge=1.0, - mass=1.0, - T=0.0, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.0]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - shape=1, - dt=world["dt"], - ) - ] - fields = (E, B, J, rho, phi, external_fields, None) - - particles, fields = time_loop_electrodynamic( - particles, - fields, - world, - constants, - unused_curl, - zero_current, - solver="fdtd", - relativistic=False, - particle_pusher="boris", - ) - - E_after, B_after, J_after, rho_after, phi_after, external_after, pml_state = fields - vx, vy, vz = particles[0].get_velocity() - - self.assertIsNone(pml_state) - self.assertGreater(float(vx[0]), 0.0) - self.assertTrue(jnp.allclose(vy, 0.0)) - self.assertTrue(jnp.allclose(vz, 0.0)) - self.assertTrue(jnp.allclose(E_after[0], 0.0)) - self.assertTrue(jnp.allclose(B_after[0], 0.0)) - self.assertTrue(jnp.allclose(external_after[0][0], external_fields[0][0])) - - def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): - world = { - "Nx": 3, - "Ny": 1, - "Nz": 1, - "dx": 1.0 / 3.0, - "dy": 1.0, - "dz": 1.0, - "dt": 0.1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_yee_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} - - E, B, J, phi, rho = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) - external_fields = ( - tuple(jnp.zeros_like(comp) for comp in E), - tuple(jnp.zeros_like(comp) for comp in B), - ) - particles = [ - particle_species( - name="absorbing", - N_particles=1, - charge=1.0, - mass=1.0, - T=0.0, - weight=1.0, - v1=jnp.array([0.2]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.49]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - x_bc="absorbing", - shape=1, - dt=world["dt"], - ) - ] - fields = (E, B, J, rho, phi, external_fields, None) - - particles, _ = time_loop_electrodynamic( - particles, - fields, - world, - constants, - unused_curl, - zero_current, - solver="fdtd", - relativistic=False, - particle_pusher="boris", - ) - - x, _, _ = particles[0].get_forward_position() - - self.assertEqual(x.shape[0], 1) - self.assertTrue(jnp.array_equal(particles[0].get_active_mask(), jnp.array([False]))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/fdtd_test.py b/tests/fdtd_test.py deleted file mode 100644 index 337a4bba..00000000 --- a/tests/fdtd_test.py +++ /dev/null @@ -1,314 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp - -from PyPIC3D.solvers.fdtd import ( - centered_finite_difference_curl, centered_finite_difference_laplacian, - centered_finite_difference_gradient, centered_finite_difference_divergence -) - -from PyPIC3D.utils import convergence_test, mse - -jax.config.update("jax_enable_x64", True) - -class TestFDTDMethods(unittest.TestCase): - def setUp(self): - x = jnp.linspace(0, 1, 120) - y = jnp.linspace(0, 1, 120) - z = jnp.linspace(0, 1, 120) - self.X, self.Y, self.Z = jnp.meshgrid(x, y, z, indexing='ij') - self.Nx, self.Ny, self.Nz = 120, 120, 120 - self.dx = 1.0/119 - self.dy = 1.0/119 - self.dz = 1.0/119 - self.bc = 'periodic' - self.slicer = (slice(1, -1), slice(1, -1), slice(1, -1)) - - def test_centered_finite_difference_laplacian(self): - # Scalar field: phi = X**2 + Y**2 + Z**2, Laplacian = 2 + 2 + 2 = 6 - phi = self.X**2 + self.Y**2 + self.Z**2 - expected = 6 * jnp.ones_like(phi) - laplacian = centered_finite_difference_laplacian(phi, self.dx, self.dy, self.dz, self.bc) - self.assertEqual(laplacian.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(laplacian.dtype, phi.dtype) - self.assertTrue(jnp.allclose(laplacian[self.slicer], expected[self.slicer], rtol=1e-2, atol=1e-2)) - - def test_centered_finite_difference_gradient(self): - # Scalar field: phi = sin(2 * pi * X) + sin(2 * pi * Y) + sin(2 * pi * Z) - # Gradient: grad(phi) = (2 * pi * cos(2 * pi * X), 2 * pi * cos(2 * pi * Y), 2 * pi * cos(2 * pi * Z)) - # Expected values at interior points - phi = jnp.sin(2 * jnp.pi * self.X) + jnp.sin(2 * jnp.pi * self.Y) + jnp.sin(2 * jnp.pi * self.Z) - expected_gradx = 2 * jnp.pi * jnp.cos(2 * jnp.pi * self.X) - expected_grady = 2 * jnp.pi * jnp.cos(2 * jnp.pi * self.Y) - expected_gradz = 2 * jnp.pi * jnp.cos(2 * jnp.pi * self.Z) - gradx, grady, gradz = centered_finite_difference_gradient(phi, self.dx, self.dy, self.dz, self.bc) - self.assertEqual(gradx.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(grady.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(gradz.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(gradx.dtype, phi.dtype) - self.assertTrue(jnp.allclose(gradx[self.slicer], expected_gradx[self.slicer], rtol=1e-2, atol=1e-2)) - self.assertTrue(jnp.allclose(grady[self.slicer], expected_grady[self.slicer], rtol=1e-2, atol=1e-2)) - self.assertTrue(jnp.allclose(gradz[self.slicer], expected_gradz[self.slicer], rtol=1e-2, atol=1e-2)) - - def test_centered_finite_difference_divergence(self): - # Vector field: F = (X, Y, Z), div F = 1 + 1 + 1 = 3 - Fx = self.X - Fy = self.Y - Fz = self.Z - expected = 3 * jnp.ones_like(self.X) - divF = centered_finite_difference_divergence(Fx, Fy, Fz, self.dx, self.dy, self.dz, self.bc) - self.assertEqual(divF.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(divF.dtype, Fx.dtype) - self.assertTrue(jnp.allclose(divF[self.slicer], expected[self.slicer], rtol=1e-2, atol=1e-2)) - - def test_centered_finite_difference_curl(self): - # Vector field: F = (-Y, X, 0), curl F = (0, 0, 2) - Fx = -self.Y - Fy = self.X - Fz = jnp.zeros_like(self.X) - expected_curlx = jnp.zeros_like(self.X) - expected_curly = jnp.zeros_like(self.Y) - expected_curlz = 2 * jnp.ones_like(self.Z) - curlx, curly, curlz = centered_finite_difference_curl(Fx, Fy, Fz, self.dx, self.dy, self.dz, self.bc) - self.assertEqual(curlx.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(curly.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(curlz.shape, (self.Nx, self.Ny, self.Nz)) - self.assertEqual(curlx.dtype, Fx.dtype) - self.assertTrue(jnp.allclose(curlx[self.slicer], expected_curlx[self.slicer], rtol=1e-4, atol=1e-4)) - self.assertTrue(jnp.allclose(curly[self.slicer], expected_curly[self.slicer], rtol=1e-4, atol=1e-4)) - self.assertTrue(jnp.allclose(curlz[self.slicer], expected_curlz[self.slicer], rtol=1e-4, atol=1e-4)) - - - def test_convergence_fdtd(self): - - def fdtd_laplacian_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical Laplacian of a test function - on a 3D periodic grid. - - The test function is phi(x, y, z) = x^2 + y^2 + z^2, whose Laplacian is analytically 6 everywhere. - The function constructs a cubic grid of size `nx` in each dimension, computes the numerical Laplacian - using a centered finite difference scheme with periodic boundary conditions, and compares it to the - analytical solution (6) in the interior of the domain. - - Parameters - ---------- - nx : int - Number of grid points in each spatial dimension. - - Returns - ------- - error : float - Mean squared error between the numerical and analytical Laplacian in the interior of the grid. - dx : float - Grid spacing in the x-direction (equal to y and z spacing). - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = X**2 + Y**2 + Z**2 - expected_solution = 6 * jnp.ones_like(phi) - # define a analytical phi and its expected solution - - laplacian = centered_finite_difference_laplacian(phi, dx, dy, dz, 'periodic') - # compute the numerical result - - slicer = (slice(1, -1), slice(1, -1), slice(1, -1)) - # slice for the solution comparison - - error = mse( laplacian[slicer], expected_solution[slicer] ) - # compute the mean squared error of the laplacian against the analytical solution - - return error, dx - - def fdtd_gradient_comparison(nx): - """ - Computes the mean squared error between the analytical and numerical gradients - of a 3D scalar field using centered finite differences with periodic boundary conditions. - - The scalar field is defined as: - phi = sin(2πX) + sin(2πY) + sin(2πZ) - over a symmetric cubic domain of size 0.25 in each direction, discretized into `nx` points per axis. - - Parameters - ---------- - nx : int - Number of grid points along each spatial dimension. - - Returns - ------- - error : float - The sum of mean squared errors between the numerical and analytical gradients - in the x, y, and z directions. - dx : float - The spatial grid spacing along each axis. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(2 * jnp.pi * X) + jnp.sin(2 * jnp.pi * Y) + jnp.sin(2 * jnp.pi * Z) - expected_gradx = 2 * jnp.pi * jnp.cos(2 * jnp.pi * X) - expected_grady = 2 * jnp.pi * jnp.cos(2 * jnp.pi * Y) - expected_gradz = 2 * jnp.pi * jnp.cos(2 * jnp.pi * Z) - # build analytical phi with known solution - - gradx, grady, gradz = centered_finite_difference_gradient(phi, dx, dy, dz, 'periodic') - # compute the numerical result - - x_error = mse(gradx, expected_gradx) - y_error = mse(grady, expected_grady) - z_error = mse(gradz, expected_gradz) - error = x_error + y_error + z_error - # compute the mean squared error of the gradient function - - return error, dx - - - def fdtd_divergence_comparison(nx): - """ - Computes the mean squared error between the numerical and analytical divergence - of the vector field F = (X, Y, Z) on a 3D periodic grid. - - The function constructs a symmetric 3D grid of size `nx` in each dimension, - defines the vector field F = (X, Y, Z), and computes its divergence using a - centered finite difference scheme with periodic boundary conditions. The analytical - divergence of F is 3 everywhere. The function returns the mean squared error - between the computed and expected divergence, along with the grid spacing `dx`. - - Args: - nx (int): Number of grid points in each spatial dimension. - - Returns: - tuple: - error (float): Mean squared error between numerical and analytical divergence. - dx (float): Grid spacing in the x-direction. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = X - Fy = Y - Fz = Z - # F = (X, Y, Z), div F = 3 - expected_divergence = 3 * jnp.ones_like(X) - - divergence = centered_finite_difference_divergence(Fx, Fy, Fz, dx, dy, dz, 'periodic') - # compute the numerical divergence - - error = mse(divergence, expected_divergence) - # compute the mean squared error of the divergence - - return error, dx - - - def fdtd_curl_comparison(nx): - """ - Computes the mean squared error between the analytical and numerical curl of a test vector field - on a 3D periodic grid, as a function of grid resolution. - - The test vector field is F = (-Y, X, 0), whose analytical curl is (0, 0, 2). - The numerical curl is computed using a centered finite difference scheme with periodic boundaries. - - Args: - nx (int): Number of grid points in each spatial dimension. - - Returns: - tuple: - error (float): The total mean squared error between the analytical and numerical curl components. - dx (float): The grid spacing in each dimension. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = -1 * Y - Fy = X - Fz = jnp.zeros_like(Z) - # Vector field: F = (-Y, X, 0), curl F = (0, 0, 2) - - expected_curlx = jnp.zeros_like(X) - expected_curly = jnp.zeros_like(Y) - expected_curlz = 2 * jnp.ones_like(Z) - # compute the analytical curl - - curlx, curly, curlz = centered_finite_difference_curl(Fx, Fy, Fz, dx, dy, dz, 'periodic') - # compute the numerical curl - - error_x = mse(curlx, expected_curlx) - error_y = mse(curly, expected_curly) - error_z = mse(curlz, expected_curlz) - error = error_x + error_y + error_z + error_z - # compute the error in the numerical curl - - return error, dx - - order = convergence_test(fdtd_laplacian_comparison) - self.assertTrue(order > 2) - # compute order of fdtd laplacian - - order = convergence_test(fdtd_gradient_comparison) - self.assertTrue(order > 2) - # compute order of fdtd gradient - - order = convergence_test(fdtd_divergence_comparison) - self.assertTrue(order > 2) - # compute order of fdtd divergence - - order = convergence_test(fdtd_curl_comparison) - self.assertTrue(order > 2) - # compute order of fdtd curl - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/first_order_yee_test.py b/tests/first_order_yee_test.py deleted file mode 100644 index 74a56988..00000000 --- a/tests/first_order_yee_test.py +++ /dev/null @@ -1,271 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - - -from PyPIC3D.initialization import initialize_fields -from PyPIC3D.solvers.first_order_yee import update_E, update_B -from PyPIC3D.utils import build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells - -from PyPIC3D.solvers.fdtd import centered_finite_difference_curl - -jax.config.update("jax_enable_x64", True) - - -class TestFieldsMethods(unittest.TestCase): - - def setUp(self): - self.world = { - 'Nx': 100, - 'Ny': 100, - 'Nz': 100, - 'dx': 0.01, - 'dy': 0.01, - 'dz': 0.01, - 'dt': 0.01, - 'x_wind': 1.0, - 'y_wind': 1.0, - 'z_wind': 1.0, - 'boundary_conditions': {'x': 0, 'y': 0, 'z': 0}, - } - # Use normalized units to avoid numerical issues - self.constants = { - 'eps': 1.0, # Normalized permittivity - 'mu': 1.0, - 'C': 1.0, - 'alpha': 1.0, # Digital filter alpha value - } - self.grid, self.staggered_grid = build_yee_grid(self.world) - Nx = self.world['Nx'] - Ny = self.world['Ny'] - Nz = self.world['Nz'] - - E, B, J, phi, rho = initialize_fields(Nx, Ny, Nz) - self.Ex, self.Ey, self.Ez = E - self.Bx, self.By, self.Bz = B - self.Jx, self.Jy, self.Jz = J - self.phi = phi - self.rho = rho - # create fields with ghost cells (shape Nx+2, Ny+2, Nz+2) - - # Make coordinate grid for the physical interior - dx, dy, dz = self.world['dx'], self.world['dy'], self.world['dz'] - - x = jnp.linspace(0, 1, Nx) - y = jnp.linspace(0, 1, Ny) - z = jnp.linspace(0, 1, Nz) - self.X, self.Y, self.Z = jnp.meshgrid(x, y, z, indexing='ij') - - # Create a simpler analytical solution: phi = sin(2πx/L) * sin(2πy/L) * sin(2πz/L) - # This has zero boundary conditions (periodic compatible) - self.real_phi = jnp.sin(2 * jnp.pi * self.X) * jnp.sin(2 * jnp.pi * self.Y) * jnp.sin(2 * jnp.pi * self.Z) - - # For this phi, the Laplacian is: ∇²φ = -(2π/L)² * (sin_x*sin_y*sin_z + sin_x*sin_y*sin_z + sin_x*sin_y*sin_z) - # = -3*(2π/L)² * phi - # So rho = -∇²φ = 3*(2π/L)² * phi (with eps=1) - wave_number_squared = (2 * jnp.pi)**2 # Assuming Lx=Ly=Lz=1 - self.rho_analytical = 3 * wave_number_squared * self.real_phi - - # Analytical E field: E = -∇φ - self.real_Ex = -2 * jnp.pi * jnp.cos(2 * jnp.pi * self.X) * jnp.sin(2 * jnp.pi * self.Y) * jnp.sin(2 * jnp.pi * self.Z) - self.real_Ey = -2 * jnp.pi * jnp.sin(2 * jnp.pi * self.X) * jnp.cos(2 * jnp.pi * self.Y) * jnp.sin(2 * jnp.pi * self.Z) - self.real_Ez = -2 * jnp.pi * jnp.sin(2 * jnp.pi * self.X) * jnp.sin(2 * jnp.pi * self.Y) * jnp.cos(2 * jnp.pi * self.Z) - - def test_initialize_fields(self): - # fields now have ghost cells: shape is (Nx+2, Ny+2, Nz+2) - self.assertEqual(self.Ex.shape, (102, 102, 102)) - self.assertEqual(self.Ey.shape, (102, 102, 102)) - self.assertEqual(self.Ez.shape, (102, 102, 102)) - self.assertEqual(self.Bx.shape, (102, 102, 102)) - self.assertEqual(self.By.shape, (102, 102, 102)) - self.assertEqual(self.Bz.shape, (102, 102, 102)) - self.assertEqual(self.Jx.shape, (102, 102, 102)) - self.assertEqual(self.Jy.shape, (102, 102, 102)) - self.assertEqual(self.Jz.shape, (102, 102, 102)) - self.assertEqual(self.phi.shape, (102, 102, 102)) - self.assertEqual(self.rho.shape, (102, 102, 102)) - - def test_update_E(self): - """Test update_E against analytical electromagnetic wave solution""" - # Create analytical electromagnetic wave solution - # Use a simple plane wave: E = E0 * sin(kz - ωt), B = (E0/c) * sin(kz - ωt) - # Propagating in z-direction, E in x-direction, B in y-direction - - dx, dy, dz = self.world['dx'], self.world['dy'], self.world['dz'] - dt = self.world['dt'] - c = self.constants['C'] - eps = self.constants['eps'] - Nx, Ny, Nz = self.world['Nx'], self.world['Ny'], self.world['Nz'] - bc_x = self.world['boundary_conditions']['x'] - bc_y = self.world['boundary_conditions']['y'] - bc_z = self.world['boundary_conditions']['z'] - - # Wave parameters - wavelength = 4 * dz * 32 # 4 wavelengths across domain - k = 2 * jnp.pi / wavelength # wavenumber - omega = k * c # frequency (dispersion relation) - E0 = 1.0 # amplitude - - # Current time step (t=0) - t = 0.0 - - # Create wave fields at t=0 on interior grid - kz_phase = k * self.Z - Ex_int = E0 * jnp.sin(kz_phase - omega * t) - Ey_int = jnp.zeros_like(Ex_int) - Ez_int = jnp.zeros_like(Ex_int) - - By_int = (E0 / c) * jnp.sin(kz_phase + k * dz/2 - omega * t) - Bx_int = jnp.zeros_like(By_int) - Bz_int = jnp.zeros_like(By_int) - - # Place into ghost-celled arrays - Ex_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Ey_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Ez_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Bx_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - By_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Bz_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)) - - Ex_initial = Ex_initial.at[1:-1, 1:-1, 1:-1].set(Ex_int) - Ey_initial = Ey_initial.at[1:-1, 1:-1, 1:-1].set(Ey_int) - Ez_initial = Ez_initial.at[1:-1, 1:-1, 1:-1].set(Ez_int) - Bx_initial = Bx_initial.at[1:-1, 1:-1, 1:-1].set(Bx_int) - By_initial = By_initial.at[1:-1, 1:-1, 1:-1].set(By_int) - Bz_initial = Bz_initial.at[1:-1, 1:-1, 1:-1].set(Bz_int) - - # Fill ghost cells - Ex_initial = update_ghost_cells(Ex_initial, bc_x, bc_y, bc_z) - Ey_initial = update_ghost_cells(Ey_initial, bc_x, bc_y, bc_z) - Ez_initial = update_ghost_cells(Ez_initial, bc_x, bc_y, bc_z) - Bx_initial = update_ghost_cells(Bx_initial, bc_x, bc_y, bc_z) - By_initial = update_ghost_cells(By_initial, bc_x, bc_y, bc_z) - Bz_initial = update_ghost_cells(Bz_initial, bc_x, bc_y, bc_z) - - # Zero current density - Jx = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Jy = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Jz = jnp.zeros((Nx+2, Ny+2, Nz+2)) - - # Analytical solution at t + dt - Ex_analytical = E0 * jnp.sin(kz_phase - omega * (t + dt)) - - # Use curl function - curl_func = lambda Ex, Ey, Ez: None # curl function already defined in update_E - - # Test update_E - E_initial = (Ex_initial, Ey_initial, Ez_initial) - B_initial = (Bx_initial, By_initial, Bz_initial) - J = (Jx, Jy, Jz) - - E_computed, pml_state = update_E(E_initial, B_initial, J, self.world, self.constants, curl_func) - Ex_computed, Ey_computed, Ez_computed = E_computed - - # Check shapes (including ghost cells) - self.assertIsNone(pml_state) - self.assertEqual(Ex_computed.shape, (102, 102, 102)) - self.assertEqual(Ey_computed.shape, (102, 102, 102)) - self.assertEqual(Ez_computed.shape, (102, 102, 102)) - - # Compare with analytical solution on interior (avoid boundary effects) - interior = (slice(5, 29), slice(5, 29), slice(5, 29)) - - Ex_error = jnp.max(jnp.abs(Ex_computed[1:-1, 1:-1, 1:-1][interior] - Ex_analytical[interior])) - Ey_error = jnp.max(jnp.abs(Ey_computed[1:-1, 1:-1, 1:-1][interior])) - Ez_error = jnp.max(jnp.abs(Ez_computed[1:-1, 1:-1, 1:-1][interior])) - - # For a plane wave, Ey and Ez should remain zero - self.assertLess(Ey_error, 1e-10, "Ey should remain zero for plane wave") - self.assertLess(Ez_error, 1e-10, "Ez should remain zero for plane wave") - - # Ex should evolve correctly (tolerance accounts for discretization) - relative_error = Ex_error / E0 - self.assertLess(relative_error, 4e-3, "Ex evolution should match analytical solution within 0.4%") - - - def test_update_B(self): - """Test update_B against analytical electromagnetic wave solution""" - - dx, dy, dz = self.world['dx'], self.world['dy'], self.world['dz'] - dt = self.world['dt'] - c = self.constants['C'] - Nx, Ny, Nz = self.world['Nx'], self.world['Ny'], self.world['Nz'] - bc_x = self.world['boundary_conditions']['x'] - bc_y = self.world['boundary_conditions']['y'] - bc_z = self.world['boundary_conditions']['z'] - - # Wave parameters (same as test_update_E) - wavelength = 4 * dz * 32 - k = 2 * jnp.pi / wavelength - omega = k * c - E0 = 1.0 - - # Current time step (t=0) - t = 0.0 - - # Create wave fields at t=0 on interior grid - kz_phase = k * self.Z - Ex_int = E0 * jnp.sin(kz_phase - omega * t) - Ey_int = jnp.zeros_like(Ex_int) - Ez_int = jnp.zeros_like(Ex_int) - - By_int = (E0 / c) * jnp.sin(kz_phase + k * dz/2 - omega * t) - Bx_int = jnp.zeros_like(By_int) - Bz_int = jnp.zeros_like(By_int) - - # Place into ghost-celled arrays - Ex_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(Ex_int) - Ey_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(Ey_int) - Ez_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(Ez_int) - Bx_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(Bx_int) - By_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(By_int) - Bz_initial = jnp.zeros((Nx+2, Ny+2, Nz+2)).at[1:-1, 1:-1, 1:-1].set(Bz_int) - - - Ex_initial = update_ghost_cells(Ex_initial, bc_x, bc_y, bc_z) - Ey_initial = update_ghost_cells(Ey_initial, bc_x, bc_y, bc_z) - Ez_initial = update_ghost_cells(Ez_initial, bc_x, bc_y, bc_z) - Bx_initial = update_ghost_cells(Bx_initial, bc_x, bc_y, bc_z) - By_initial = update_ghost_cells(By_initial, bc_x, bc_y, bc_z) - Bz_initial = update_ghost_cells(Bz_initial, bc_x, bc_y, bc_z) - - # Analytical solution at t + dt - By_analytical = (E0 / c) * jnp.sin(kz_phase + k * dz/2 - omega * (t + dt)) - - # Use curl function - curl_func = lambda Ex, Ey, Ez: None # curl function already defined in update_B - - # Test update_B - E_initial = (Ex_initial, Ey_initial, Ez_initial) - B_initial = (Bx_initial, By_initial, Bz_initial) - - B_computed, pml_state = update_B(E_initial, B_initial, self.world, self.constants, curl_func) - Bx_computed, By_computed, Bz_computed = B_computed - - # Check shapes - self.assertIsNone(pml_state) - self.assertEqual(Bx_computed.shape, (102, 102, 102)) - self.assertEqual(By_computed.shape, (102, 102, 102)) - self.assertEqual(Bz_computed.shape, (102, 102, 102)) - - # Compare with analytical solution on interior (avoid boundary effects) - interior = (slice(5, 29), slice(5, 29), slice(5, 29)) - - Bx_error = jnp.max(jnp.abs(Bx_computed[1:-1, 1:-1, 1:-1][interior])) - By_error = jnp.max(jnp.abs(By_computed[1:-1, 1:-1, 1:-1][interior] - By_analytical[interior])) - Bz_error = jnp.max(jnp.abs(Bz_computed[1:-1, 1:-1, 1:-1][interior])) - - # For a plane wave, Bx and Bz should remain zero - self.assertLess(Bx_error, 1e-10, "Bx should remain zero for plane wave") - self.assertLess(Bz_error, 1e-10, "Bz should remain zero for plane wave") - - # By should evolve correctly - B0 = E0 / c - relative_error = By_error / B0 - self.assertLess(relative_error, 2e-3, "By evolution should match analytical solution within 0.2%") - -if __name__ == '__main__': - unittest.main() diff --git a/tests/initial_particles.py b/tests/initial_particles.py new file mode 100644 index 00000000..b3475d17 --- /dev/null +++ b/tests/initial_particles.py @@ -0,0 +1,204 @@ +import math + +import jax.numpy as jnp +import numpy as np + +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles + + +def tiled_species( + name, + charge, + mass, + weight=1.0, + N_particles=None, + T=None, + x1=None, + x2=None, + x3=None, + u1=None, + u2=None, + u3=None, + v1=None, + v2=None, + v3=None, + active_mask=None, + update_pos=True, + update_v=True, + update_x=True, + update_y=True, + update_z=True, + update_u=None, + update_vx=True, + update_vy=True, + update_vz=True, + **unused, +): + x1 = jnp.asarray(x1, dtype=float) + n_particles = int(x1.shape[0]) + + if u1 is None and v1 is not None: + u1 = v1 + if u2 is None and v2 is not None: + u2 = v2 + if u3 is None and v3 is not None: + u3 = v3 + + if x2 is None: + x2 = jnp.zeros(n_particles) + if x3 is None: + x3 = jnp.zeros(n_particles) + if u1 is None: + u1 = jnp.zeros(n_particles) + if u2 is None: + u2 = jnp.zeros(n_particles) + if u3 is None: + u3 = jnp.zeros(n_particles) + if active_mask is None: + active_mask = jnp.ones(n_particles, dtype=bool) + + if isinstance(update_x, (tuple, list)): + update_x_components = tuple(update_x) + else: + update_x_components = ( + bool(update_pos and update_x), + bool(update_pos and update_y), + bool(update_pos and update_z), + ) + + if update_u is None: + update_u_components = ( + bool(update_v and update_vx), + bool(update_v and update_vy), + bool(update_v and update_vz), + ) + elif isinstance(update_u, (tuple, list)): + update_u_components = tuple(update_u) + else: + update_u_components = (bool(update_u), bool(update_u), bool(update_u)) + + return { + "name": name, + "charge": charge, + "mass": mass, + "weight": weight, + "x": jnp.stack( + ( + x1, + jnp.asarray(x2, dtype=float), + jnp.asarray(x3, dtype=float), + ), + axis=-1, + ), + "u": jnp.stack( + ( + jnp.asarray(u1, dtype=float), + jnp.asarray(u2, dtype=float), + jnp.asarray(u3, dtype=float), + ), + axis=-1, + ), + "active": jnp.asarray(active_mask, dtype=bool), + "update_x": update_x_components, + "update_u": update_u_components, + } + + +def _tile_axis_count(n_cells, cells_per_tile): + return int(math.ceil(int(n_cells) / int(cells_per_tile))) + + +def _tile_indices(x, world, tile_shape): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + + x_cell = np.floor((np.asarray(x[:, 0]) + float(world["x_wind"]) / 2.0) / float(world["dx"])).astype(int) + y_cell = np.floor((np.asarray(x[:, 1]) + float(world["y_wind"]) / 2.0) / float(world["dy"])).astype(int) + z_cell = np.floor((np.asarray(x[:, 2]) + float(world["z_wind"]) / 2.0) / float(world["dz"])).astype(int) + + x_cell = np.clip(x_cell, 0, int(world["Nx"]) - 1) + y_cell = np.clip(y_cell, 0, int(world["Ny"]) - 1) + z_cell = np.clip(z_cell, 0, int(world["Nz"]) - 1) + + return x_cell // tile_nx, y_cell // tile_ny, z_cell // tile_nz + + +def tile_shape_from_parameters(simulation_parameters): + return ( + int(simulation_parameters["particle_tile_nx"]), + int(simulation_parameters["particle_tile_ny"]), + int(simulation_parameters["particle_tile_nz"]), + ) + + +def build_tiled_particles(species, world, tile_shape=None, simulation_parameters=None, capacity_factor=None): + if tile_shape is not None and isinstance(tile_shape, dict): + simulation_parameters = tile_shape + tile_shape = None + + if tile_shape is None: + tile_shape = tile_shape_from_parameters(simulation_parameters) + if capacity_factor is None: + capacity_factor = 1.0 + if simulation_parameters is not None: + capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx = _tile_axis_count(world["Nx"], tile_nx) + nty = _tile_axis_count(world["Ny"], tile_ny) + ntz = _tile_axis_count(world["Nz"], tile_nz) + n_species = len(species) + + tile_counts = np.zeros((ntx, nty, ntz, n_species), dtype=int) + particle_tile_data = [] + n_tiles = ntx * nty * ntz + + for species_index, species_data in enumerate(species): + x = np.asarray(species_data["x"], dtype=float) + u = np.asarray(species_data["u"], dtype=float) + active = np.asarray(species_data["active"], dtype=bool) + tx, ty, tz = _tile_indices(x, world, tile_shape) + flat_tile = (tx * nty + ty) * ntz + tz + particle_indices = np.arange(x.shape[0]) + + flat_counts = np.bincount(flat_tile[particle_indices], minlength=n_tiles) + tile_counts[:, :, :, species_index] = flat_counts.reshape((ntx, nty, ntz)) + particle_tile_data.append((x, u, active, tx, ty, tz, flat_tile, particle_indices)) + + max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 + max_particles_per_tile = int(math.ceil(max_particles_per_tile * float(capacity_factor))) + + x_tiles = np.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3), dtype=float) + u_tiles = np.zeros_like(x_tiles) + active_tiles = np.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) + + for species_index, (x, u, active, tx, ty, tz, flat_tile, all_particle_indices) in enumerate(particle_tile_data): + order = np.argsort(flat_tile[all_particle_indices], kind="stable") + particle_indices = all_particle_indices[order] + sorted_flat_tile = flat_tile[particle_indices] + + flat_counts = tile_counts[:, :, :, species_index].reshape(-1) + tile_starts = np.cumsum(flat_counts) - flat_counts + slots = np.arange(particle_indices.size) - tile_starts[sorted_flat_tile] + + x_tiles[tx[particle_indices], ty[particle_indices], tz[particle_indices], species_index, slots, :] = x[particle_indices] + u_tiles[tx[particle_indices], ty[particle_indices], tz[particle_indices], species_index, slots, :] = u[particle_indices] + active_tiles[tx[particle_indices], ty[particle_indices], tz[particle_indices], species_index, slots] = active[particle_indices] + + particles = TiledParticles( + x=jnp.asarray(x_tiles), + u=jnp.asarray(u_tiles), + active=jnp.asarray(active_tiles), + ) + species_config = SpeciesConfig( + charge=jnp.asarray([species_data["charge"] for species_data in species], dtype=float), + mass=jnp.asarray([species_data["mass"] for species_data in species], dtype=float), + weight=jnp.asarray([species_data["weight"] for species_data in species], dtype=float), + update_x=jnp.asarray([species_data["update_x"] for species_data in species], dtype=bool), + update_u=jnp.asarray([species_data["update_u"] for species_data in species], dtype=bool), + ) + + return particles, species_config + + +def species_names(species): + return tuple(species_data["name"] for species_data in species) diff --git a/tests/initialization_test.py b/tests/initialization_test.py deleted file mode 100644 index 4665549a..00000000 --- a/tests/initialization_test.py +++ /dev/null @@ -1,30 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -from PyPIC3D.initialization import setup_write_dir, default_parameters - -jax.config.update("jax_enable_x64", True) - -class TestInitializationFunctions(unittest.TestCase): - def setUp(self): - self.plotting_parameters, self.simulation_parameters, self.constants = default_parameters() - self.simulation_parameters['output_dir'] = 'test_output' - self.plotting_parameters['plotfields'] = False - # check the default parameters are set correctly - - def test_setup_write_dir(self): - # Should not raise - setup_write_dir(self.simulation_parameters, self.plotting_parameters) - # check that the output directory is created - - 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 - -if __name__ == '__main__': - unittest.main() diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py deleted file mode 100644 index ac3b4a37..00000000 --- a/tests/openpmd_test.py +++ /dev/null @@ -1,112 +0,0 @@ -import unittest -from unittest.mock import patch - -import jax.numpy as jnp - -from PyPIC3D.diagnostics import openPMD -from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array - - -class FakeRecord: - def __init__(self): - self.shape = None - self.unit_SI = None - - def reset_dataset(self, dataset): - pass - - def store_chunk(self, array, offset, extent): - self.shape = tuple(extent) - - -class FakeMesh: - def __init__(self): - self.records = {} - self.axis_labels = None - self.grid_spacing = None - self.grid_global_offset = None - - def __getitem__(self, component): - if component not in self.records: - self.records[component] = FakeRecord() - return self.records[component] - - -class FakeMeshes(dict): - def __getitem__(self, name): - if name not in self: - self[name] = FakeMesh() - return dict.__getitem__(self, name) - - -class FakeIteration: - def __init__(self): - self.meshes = FakeMeshes() - self.time = None - self.dt = None - self.time_unit_SI = None - - -class FakeIterations(dict): - def __getitem__(self, iteration): - if iteration not in self: - self[iteration] = FakeIteration() - return dict.__getitem__(self, iteration) - - -class FakeSeries: - def __init__(self): - self.iterations = FakeIterations() - - def flush(self): - pass - - def close(self): - pass - - -def _zero_field(shape): - return tuple(jnp.zeros(shape) for _ in range(3)) - - -class OpenPMDDiagnosticsTests(unittest.TestCase): - - def test_openpmd_field_array_preserves_thin_y_axis(self): - field_component = jnp.ones((4, 1, 6)) - - array = _ensure_openpmd_array(field_component) - - self.assertEqual(array.shape, (4, 1, 6)) - - def test_write_openpmd_fields_preserves_thin_y_mesh_metadata(self): - shape_with_ghosts = (6, 3, 8) - E = _zero_field(shape_with_ghosts) - B = _zero_field(shape_with_ghosts) - J = _zero_field(shape_with_ghosts) - rho = jnp.zeros(shape_with_ghosts) - phi = jnp.zeros(shape_with_ghosts) - external_fields = _zero_field(shape_with_ghosts), _zero_field(shape_with_ghosts) - fields = (E, B, J, rho, phi, external_fields) - world = { - "dt": 1.0, - "dx": 0.25, - "dy": 0.5, - "dz": 0.75, - "x_wind": 1.0, - "y_wind": 2.0, - "z_wind": 3.0, - } - series = FakeSeries() - - with patch.object(openPMD, "_open_openpmd_series", return_value=series): - openPMD.write_openpmd_fields(fields, world, "/tmp", plot_t=0, t=0) - - B_mesh = series.iterations[0].meshes["B"] - self.assertEqual(B_mesh.axis_labels, ["x", "y", "z"]) - self.assertEqual(B_mesh.grid_spacing, [0.25, 0.5, 0.75]) - self.assertEqual(B_mesh.grid_global_offset, [-0.5, -1.0, -1.5]) - self.assertEqual(B_mesh.records["x"].shape, (4, 1, 6)) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/particle_test.py b/tests/particle_test.py deleted file mode 100644 index 7571e0a2..00000000 --- a/tests/particle_test.py +++ /dev/null @@ -1,948 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - - -from PyPIC3D.particles.particle_initialization import ( - initial_particles, - load_particles_from_toml, - ) - -from PyPIC3D.particles.species_class import particle_species - -from PyPIC3D.deposition.J_from_rhov import J_from_rhov -from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.utils import build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells - - - -jax.config.update("jax_enable_x64", True) - -class TestParticleMethods(unittest.TestCase): - def setUp(self): - """ - Set up the initial conditions for the particle test. - - Attributes: - N_particles (int): Number of particles. - x_wind (float): Wind velocity in the x direction. - y_wind (float): Wind velocity in the y direction. - z_wind (float): Wind velocity in the z direction. - mass (float): Mass of each particle. - T (float): Temperature in Kelvin. - kb (float): Boltzmann constant. - key1 (jax.random.PRNGKey): First random key for JAX. - key2 (jax.random.PRNGKey): Second random key for JAX. - key3 (jax.random.PRNGKey): Third random key for JAX. - dt (float): Time step for the simulation. - """ - self.N_particles = 3 - self.x_wind = 10.0 - self.y_wind = 10.0 - self.z_wind = 10.0 - self.mass = 1.0 - self.T = 100.0 - self.kb = 1.38e-23 - self.key1, self.key2, self.key3 = jax.random.split(jax.random.PRNGKey(0), 3) - self.dt = 1.0 - - def test_initial_particles(self): - """ - Test the initial_particles function to ensure it generates the correct number of particles. - - This test verifies that the initial_particles function returns arrays of the correct shape - for the positions (x, y, z) and velocities (vx, vy, vz) of the particles. - - The function checks: - - The number of particles in the x, y, and z position arrays is equal to N_particles. - - The number of particles in the vx, vy, and vz velocity arrays is equal to N_particles. - - Attributes: - N_particles (int): Number of particles to generate. - x_wind (float): Wind component in the x direction. - y_wind (float): Wind component in the y direction. - z_wind (float): Wind component in the z direction. - mass (float): Mass of each particle. - T (float): Temperature. - kb (float): Boltzmann constant. - key1, key2, key3: Keys for random number generation. - """ - x, y, z, vx, vy, vz = initial_particles( - N_per_cell=0.1, - N_particles=self.N_particles, - minx=-self.x_wind/2, - maxx=self.x_wind/2, - miny=-self.y_wind/2, - maxy=self.y_wind/2, - minz=-self.z_wind/2, - maxz=self.z_wind/2, - mass=self.mass, - Tx=self.T, - Ty=self.T, - Tz=self.T, - kb=self.kb, - key1=self.key1, - key2=self.key2, - key3=self.key3, - ) - self.assertEqual(x.shape[0], self.N_particles) - self.assertEqual(y.shape[0], self.N_particles) - self.assertEqual(z.shape[0], self.N_particles) - self.assertEqual(vx.shape[0], self.N_particles) - self.assertEqual(vy.shape[0], self.N_particles) - self.assertEqual(vz.shape[0], self.N_particles) - - def test_periodic_boundary_condition(self): - """ - Test the periodic boundary condition for particle species. - This test initializes particle positions and velocities, creates a particle species - with these properties, and applies the periodic boundary condition. It then checks - if the positions of the particles are correctly wrapped around the boundaries. - The test verifies that the positions of the particles after applying the periodic - boundary condition are as expected. - Asserts: - - The x positions of the particles are [-5.0, 5.0, 0.0]. - - The y positions of the particles are [-5.0, 5.0, 0.0]. - - The z positions of the particles are [-5.0, 5.0, 0.0]. - """ - - x1 = jnp.array([5.50, -6.0, 0.0]) - x2 = jnp.array([5.50, -6.0, 0.0]) - x3 = jnp.array([5.50, -6.0, 0.0]) - - species = particle_species( - - name = "test", - N_particles = self.N_particles, - charge = 1.0, - mass = self.mass, - weight = 1.0, - T = self.T, - v1 = jnp.array([0.0, 0.0, 0.0]), - v2 = jnp.array([0.0, 0.0, 0.0]), - v3 = jnp.array([0.0, 0.0, 0.0]), - x1 = x1, - x2 = x2, - x3 = x3, - dx = 1.0, - dy = 1.0, - dz = 1.0, - xwind = self.x_wind, - ywind = self.y_wind, - zwind = self.z_wind, - ) - - x, y, z = species.get_position() - self.assertTrue(jnp.all(x == jnp.array([-4.5, 4.0, 0.0]))) - self.assertTrue(jnp.all(y == jnp.array([-4.5, 4.0, 0.0]))) - self.assertTrue(jnp.all(z == jnp.array([-4.5, 4.0, 0.0]))) - - def test_update_position(self): - """ - Test the update_position method of the particle_species class. - This test initializes a particle_species instance with predefined positions - and velocities, updates the positions using the update_position method, and - verifies that the new positions are as expected. - The test checks: - - The initial positions of the particles. - - The velocities of the particles. - - The updated positions after calling update_position. - - The correctness of the updated positions using jnp.allclose. - Asserts: - - The updated x positions are close to the expected values. - - The updated y positions are close to the expected values. - - The updated z positions are close to the expected values. - """ - x1=jnp.array([0.1, 0.2, 0.3]) - x2=jnp.array([0.1, 0.2, 0.3]) - x3=jnp.array([0.1, 0.2, 0.3]) - - species = particle_species( - name="test", - N_particles=self.N_particles, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([1.0, 2.0, 3.0]), - v2=jnp.array([1.0, 2.0, 3.0]), - v3=jnp.array([1.0, 2.0, 3.0]), - x1=x1, - x2=x2, - x3=x3, - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - dt=self.dt - ) - - species.update_position() - x, y, z = species.get_forward_position() - self.assertTrue(jnp.allclose(x, jnp.array([1.1, 2.2, 3.3]))) - self.assertTrue(jnp.allclose(y, jnp.array([1.1, 2.2, 3.3]))) - self.assertTrue(jnp.allclose(z, jnp.array([1.1, 2.2, 3.3]))) - - def test_absorbing_boundary_updates_active_mask(self): - x1 = jnp.array([0.0, self.x_wind / 2 + 0.1, -self.x_wind / 2 - 0.1]) - x2 = jnp.array([0.0, 0.0, 0.0]) - x3 = jnp.array([0.0, 0.0, 0.0]) - - species = particle_species( - name="absorbing", - N_particles=self.N_particles, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([1.0, 1.0, 1.0]), - v2=jnp.array([0.0, 0.0, 0.0]), - v3=jnp.array([0.0, 0.0, 0.0]), - x1=x1, - x2=x2, - x3=x3, - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - x_bc="absorbing", - dt=self.dt, - ) - - species.boundary_conditions() - - self.assertTrue(jnp.array_equal(species.get_active_mask(), jnp.array([True, False, False]))) - self.assertEqual(species.get_forward_position()[0].shape[0], self.N_particles) - - def test_particle_initialization_accepts_absorbing_axis_boundary(self): - world = { - "Nx": 1, - "Ny": 1, - "Nz": 1, - "dx": 1.0, - "dy": 1.0, - "dz": 1.0, - "dt": self.dt, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - } - constants = {"kb": self.kb, "eps": 1.0} - simulation_parameters = {"ds_per_debye": None, "shape_factor": 1} - config = { - "particle1": { - "name": "absorbing", - "N_particles": 1, - "charge": 1.0, - "mass": self.mass, - "temperature": self.T, - "x_bc": "absorbing", - "initial_x": 0.0, - "initial_y": 0.0, - "initial_z": 0.0, - "initial_vx": 0.0, - "initial_vy": 0.0, - "initial_vz": 0.0, - } - } - - particles = load_particles_from_toml(config, simulation_parameters, world, constants) - - self.assertEqual(particles[0].x_bc, "absorbing") - self.assertTrue(jnp.array_equal(particles[0].get_active_mask(), jnp.array([True]))) - - def test_inactive_particles_do_not_advance_or_take_new_velocity(self): - species = particle_species( - name="absorbing", - N_particles=self.N_particles, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([1.0, 2.0, 3.0]), - v2=jnp.array([0.0, 0.0, 0.0]), - v3=jnp.array([0.0, 0.0, 0.0]), - x1=jnp.array([0.0, self.x_wind / 2 + 0.1, 0.0]), - x2=jnp.array([0.0, 0.0, 0.0]), - x3=jnp.array([0.0, 0.0, 0.0]), - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - x_bc="absorbing", - dt=self.dt, - ) - species.boundary_conditions() - - species.set_velocity( - jnp.array([10.0, 20.0, 30.0]), - jnp.array([0.0, 0.0, 0.0]), - jnp.array([0.0, 0.0, 0.0]), - ) - species.update_position() - - x, _, _ = species.get_forward_position() - vx, _, _ = species.get_velocity() - - self.assertTrue(jnp.array_equal(species.get_active_mask(), jnp.array([True, False, True]))) - self.assertTrue(jnp.allclose(vx, jnp.array([10.0, 2.0, 30.0]))) - self.assertTrue(jnp.allclose(x, jnp.array([10.0, self.x_wind / 2 + 0.1, 30.0]))) - active_vx, _, _ = species.get_active_velocity() - self.assertTrue(jnp.allclose(active_vx, jnp.array([10.0, 0.0, 30.0]))) - - def test_total_KE(self): - """ - Test the total kinetic energy calculation for a particle species. - This test initializes a particle species with a given number of particles, - mass, and temperature. It sets the velocities of all particles to 1.0 in - each direction and calculates the total kinetic energy using the - `total_KE` function. The expected kinetic energy is compared to the - calculated value using `assertAlmostEqual`. - The expected kinetic energy is calculated as: - KE = 0.5 * mass * N_particles * 3 - where: - - mass is the mass of each particle - - N_particles is the number of particles - - 3 accounts for the three velocity components (v1, v2, v3) - Asserts: - The calculated kinetic energy is almost equal to the expected value. - """ - - x1=jnp.zeros(self.N_particles) - x2=jnp.zeros(self.N_particles) - x3=jnp.zeros(self.N_particles) - - species = particle_species( - name="test", - N_particles=self.N_particles, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1 = jnp.ones(self.N_particles), - v2 = jnp.ones(self.N_particles), - v3 = jnp.ones(self.N_particles), - x1=x1, - x2=x2, - x3=x3, - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - ke = species.kinetic_energy() - self.assertAlmostEqual(ke, 0.5 * self.mass * self.N_particles * 3) - - def test_particle_species(self): - """ - Test the particle_species class. - This test verifies the following functionalities of the particle_species class: - - Initialization of particle_species with given parameters. - - Retrieval of particle species name. - - Retrieval of particle charge. - - Retrieval of the number of particles. - - Retrieval of particle velocities and positions. - - Setting of particle velocities and positions. - - Setting of particle mass. - - Calculation of kinetic energy and momentum. - The test performs the following checks: - - The name of the species is correctly set and retrieved. - - The charge of the species is correctly set and retrieved. - - The number of particles is correctly set and retrieved. - - The initial velocities and positions are correctly set and retrieved. - - The velocities and positions can be updated and retrieved correctly. - - The mass can be updated and retrieved correctly. - - The kinetic energy and momentum are calculated correctly. - """ - - x1=jnp.zeros(self.N_particles) - x2=jnp.zeros(self.N_particles) - x3=jnp.zeros(self.N_particles) - - species = particle_species( - name="test", - N_particles=self.N_particles, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1 = jnp.ones(self.N_particles), - v2 = jnp.ones(self.N_particles), - v3 = jnp.ones(self.N_particles), - x1=x1, - x2=x2, - x3=x3, - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - dt=self.dt - ) - - - self.assertEqual(species.get_name(), "test") - self.assertEqual(species.get_charge(), 1.0) - self.assertEqual(species.get_number_of_particles(), self.N_particles) - self.assertTrue(jnp.all(species.get_velocity()[0] == jnp.ones(self.N_particles))) - self.assertTrue(jnp.all(species.get_forward_position()[0] == jnp.zeros(self.N_particles))) - self.assertEqual(species.get_mass(), self.mass) - species.set_velocity(jnp.zeros(self.N_particles), jnp.zeros(self.N_particles), jnp.zeros(self.N_particles)) - self.assertTrue(jnp.all(species.get_velocity()[0] == jnp.zeros(self.N_particles))) - species.set_position(jnp.ones(self.N_particles), jnp.ones(self.N_particles), jnp.ones(self.N_particles)) - self.assertTrue(jnp.all(species.get_forward_position()[0] == jnp.ones(self.N_particles))) - species.set_mass(2.0) - self.assertEqual(species.get_mass(), 2.0) - self.assertAlmostEqual(species.kinetic_energy(), 0.0) - self.assertAlmostEqual(species.momentum(), 0.0) - - def test_J_from_rhov(self): - x = jnp.array([0.0]) - y = jnp.array([0.0]) - z = jnp.array([0.0]) - vx = jnp.array([0.5]) - vy = jnp.array([0.5]) - vz = jnp.array([0.5]) - # define particle position and velocity - - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - # uniform spatial resolution in xyz - - world = { - 'dx': dx, 'dy': dy, 'dz': dz, - 'Nx': Nx, 'Ny': Ny, 'Nz': Nz, - 'x_wind': self.x_wind, 'y_wind': self.y_wind, 'z_wind': self.z_wind, - 'dt': 0.0001, - 'boundary_conditions': {'x': 0, 'y': 0, 'z': 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world['grids'] = {'vertex': vertex_grid, 'center': center_grid} - # build world with extended grids - - # ghost-celled J arrays: shape (Nx+2, Ny+2, Nz+2) - num_J = (jnp.zeros((Nx+2, Ny+2, Nz+2)), jnp.zeros((Nx+2, Ny+2, Nz+2)), jnp.zeros((Nx+2, Ny+2, Nz+2))) - - species = particle_species( - name="test", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1 = vx, - v2 = vy, - v3 = vz, - x1=x, - x2=y, - x3=z, - dx=1.0, - dy=1.0, - dz=1.0, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - constants = {'C': 3e8, 'alpha' : 1.0} - - num_J = J_from_rhov([species], num_J, constants, world) - - # Jan 12, 2025: Suppressing test for now. I have done benchmarks of the two stream and weibel - # against WarpX and have validated this method - - def test_inactive_particles_do_not_deposit_rho_or_current(self): - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "dt": 0.0, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="absorbed", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([0.5]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([self.x_wind / 2 + dx]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - x_bc="absorbing", - ) - species.boundary_conditions() - - constants = {"C": 3e8, "alpha": 1.0} - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) - J = J_from_rhov( - [species], - ( - jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), - jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), - jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), - ), - constants, - world, - filter="none", - ) - - self.assertAlmostEqual(float(jnp.sum(rho[1:-1, 1:-1, 1:-1]) * dx * dy * dz), 0.0) - self.assertAlmostEqual(float(jnp.sum(J[0][1:-1, 1:-1, 1:-1]) * dx * dy * dz), 0.0) - - def test_rho(self): - x = jnp.array([0.0]) - y = jnp.array([0.0]) - z = jnp.array([0.0]) - vx = jnp.array([0.5]) - vy = jnp.array([0.5]) - vz = jnp.array([0.5]) - # define particle position and velocity - - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - # uniform spatial resolution in xyz - - world = { - 'dx': dx, 'dy': dy, 'dz': dz, - 'Nx': Nx, 'Ny': Ny, 'Nz': Nz, - 'x_wind': self.x_wind, 'y_wind': self.y_wind, 'z_wind': self.z_wind, - 'boundary_conditions': {'x': 0, 'y': 0, 'z': 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world['grids'] = {'vertex': vertex_grid, 'center': center_grid} - # build world with extended grids - - # ghost-celled rho: shape (Nx+2, Ny+2, Nz+2) - num_rho = jnp.zeros((Nx+2, Ny+2, Nz+2)) - - # expected rho: particle at (0,0,0) with extended grid starting at -x_wind/2 - dx - # x0 = floor((0 - (-x_wind/2 - dx)) / dx) = floor((5 + 1) / 1) = 6 - # deposit at ghost-celled index [6, 6, 6] - exp_rho = jnp.zeros((Nx+2, Ny+2, Nz+2)) - exp_rho = exp_rho.at[6, 6, 6].set(1.0 / (dx*dy*dz)) - # build expected rho array with non-zero values at the correct index - - species = particle_species( - name="test", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1 = vx, - v2 = vy, - v3 = vz, - x1=x, - x2=y, - x3=z, - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - constants = {'C': 3e8, 'alpha' : 1.0} - - num_rho = compute_rho([species], num_rho, world, constants) - # compute rho - - self.assertTrue(jnp.allclose(num_rho, exp_rho)) - # check if computed rho matches expected rho - - def test_rho_periodic_boundary_exact_face(self): - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="seam", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([self.x_wind / 2]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) - interior = rho[1:-1, 1:-1, 1:-1] - cell_charge = interior * dx * dy * dz - - self.assertAlmostEqual(float(jnp.sum(cell_charge)), 1.0) - # Particle at x=x_wind/2 deposits into ghost[-1]; periodic fold wraps - # to interior[1] (the opposite-side seam point at x=-x_wind/2). - self.assertAlmostEqual(float(cell_charge[0, Ny // 2, Nz // 2]), 1.0) - - def test_rho_periodic_boundary_crossing(self): - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="crossing", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([self.x_wind / 2 + 0.1 * dx]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) - interior = rho[1:-1, 1:-1, 1:-1] * dx * dy * dz - - self.assertAlmostEqual(float(jnp.sum(interior)), 1.0) - # ghost[-1] (weight 0.9) wraps to interior[1] (opposite side); - # ghost[0] (weight 0.1) wraps to interior[-2] (opposite side). - self.assertAlmostEqual(float(interior[0, Ny // 2, Nz // 2]), 0.9) - self.assertAlmostEqual(float(interior[-1, Ny // 2, Nz // 2]), 0.1) - - def test_J_from_rhov_periodic_boundary_crossing(self): - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "dt": 0.0, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="crossing-current", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([0.5]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([self.x_wind / 2 + 0.1 * dx]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - J = J_from_rhov( - [species], - (jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), jnp.zeros((Nx + 2, Ny + 2, Nz + 2))), - {"C": 3e8, "alpha": 1.0}, - world, - filter="none", - ) - Jx_int = J[0][1:-1, 1:-1, 1:-1] - - self.assertAlmostEqual(float(jnp.sum(Jx_int) * dx * dy * dz), 0.5) - # 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 - x_wind, y_wind, z_wind = 1.0, 1.0, 1.0 - dx = x_wind / Nx - dy = y_wind / Ny - dz = z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": x_wind, - "y_wind": y_wind, - "z_wind": z_wind, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="1d", - N_particles=1, - charge=1.0, - mass=1.0, - weight=1.0, - T=0.0, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.125]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=x_wind, - ywind=y_wind, - zwind=z_wind, - ) - - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) - interior = rho[1:-1, 1:-1, 1:-1] - - self.assertEqual(interior.shape, (Nx, 1, 1)) - self.assertAlmostEqual(float(jnp.sum(interior) * dx * dy * dz), 1.0) - - def test_filtered_rho_periodic_updates_ghost_cells(self): - Nx, Ny, Nz = 10, 10, 10 - dx = self.x_wind / Nx - dy = self.y_wind / Ny - dz = self.z_wind / Nz - world = { - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": Nx, - "Ny": Ny, - "Nz": Nz, - "x_wind": self.x_wind, - "y_wind": self.y_wind, - "z_wind": self.z_wind, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - species = particle_species( - name="filtered", - N_particles=1, - charge=1.0, - mass=self.mass, - weight=1.0, - T=self.T, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.0]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - dx=dx, - dy=dy, - dz=dz, - xwind=self.x_wind, - ywind=self.y_wind, - zwind=self.z_wind, - ) - - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 0.5}) - interior = rho[1:-1, 1:-1, 1:-1] - - self.assertAlmostEqual(float(jnp.sum(interior) * dx * dy * dz), 1.0) - self.assertTrue(jnp.allclose(rho[0, :, :], rho[-2, :, :])) - self.assertTrue(jnp.allclose(rho[-1, :, :], rho[1, :, :])) - self.assertTrue(jnp.allclose(rho[:, 0, :], rho[:, -2, :])) - self.assertTrue(jnp.allclose(rho[:, -1, :], rho[:, 1, :])) - self.assertTrue(jnp.allclose(rho[:, :, 0], rho[:, :, -2])) - self.assertTrue(jnp.allclose(rho[:, :, -1], rho[:, :, 1])) - - def test_check_continuity_1D(self): - - # WORLD PARAMETERS ######## - Nx = 100 - Ny = 1 - Nz = 1 - x_wind = 1.0 - y_wind = 1.0 - z_wind = 1.0 - dx = x_wind / Nx - dy = y_wind / Ny - dz = z_wind / Nz - dt = dx / (3e8) - ########################### - - world = { - "dx": dx, "dy": dy, "dz": dz, - "Nx": Nx, "Ny": Ny, "Nz": Nz, - "x_wind": x_wind, "y_wind": y_wind, "z_wind": z_wind, - "dt": dt, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "center": center_grid} - - constants = {"C": 3e8, "alpha": 1.0} - # build constants and world parameters data structures - - vy = 0.0 - vz = 0.0 - y0 = 0.0 - z0 = 0.0 - # transverse dimensions - - vx = 0.01 # particle velocity in x - x0 = 0.0 # initial particle position in x - - species = particle_species( - name="single", - N_particles=1, - charge=1.0, - mass=1.0, - weight=1.0, - T=0.0, - v1=jnp.array([vx]), - v2=jnp.array([vy]), - v3=jnp.array([vz]), - x1=jnp.array([x0]), - x2=jnp.array([y0]), - x3=jnp.array([z0]), - dx=dx, dy=dy, dz=dz, - dt=dt, - xwind=x_wind, ywind=y_wind, zwind=z_wind, - shape=1, - ) - - # ghost-celled arrays: (Nx+2, Ny+2, Nz+2) - rho = jnp.zeros((Nx+2, Ny+2, Nz+2)) - - species.update_position() - # move particles to new position - prev_rho = compute_rho([species], rho, world, constants) - # compute rho - species.update_position() - rho = compute_rho([species], rho, world, constants) - # compute rho again - - # compute drhodt on interior - drhodt = (rho[1:-1, 1:-1, 1:-1] - prev_rho[1:-1, 1:-1, 1:-1]) / dt - # calculate backward difference for drhodt - - Jx = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Jy = jnp.zeros((Nx+2, Ny+2, Nz+2)) - Jz = jnp.zeros((Nx+2, Ny+2, Nz+2)) - J = (Jx, Jy, Jz) - J = Esirkepov_current([species], J, constants, world) - # compute J using Esirkepov method - - # update ghost cells on J for periodic divergence computation - Jx_gc = update_ghost_cells(J[0], 0, 0, 0) - - # backward difference for divergence in 1D on interior - # uses ghost cell neighbor for the left boundary - dJxdx = ( Jx_gc[1:-1, 1:-1, 1:-1] - Jx_gc[:-2, 1:-1, 1:-1] ) / dx - - continuity = drhodt + dJxdx - # rho lives on cell centers and J lives on vertices - # due to this, we need to sum over the means of the values - # to get a global measure of the continuity violation - - - self.assertLess(jnp.mean(continuity), 0.0004) - - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/physics_tests/__init__.py b/tests/physics_tests/__init__.py new file mode 100644 index 00000000..b306fdc6 --- /dev/null +++ b/tests/physics_tests/__init__.py @@ -0,0 +1 @@ +"""Numerical physics validation tests for PyPIC3D.""" diff --git a/tests/boris_test.py b/tests/physics_tests/boris_test.py similarity index 92% rename from tests/boris_test.py rename to tests/physics_tests/boris_test.py index 96c02714..b5e45de4 100644 --- a/tests/boris_test.py +++ b/tests/physics_tests/boris_test.py @@ -6,7 +6,6 @@ from functools import partial -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 @@ -88,6 +87,45 @@ def test_interpolate_field_to_particles_with_ghost_cells_on_reduced_axes(self): self.assertTrue(jnp.allclose(interpolated, 7.0)) + def test_interpolate_field_to_particles_accepts_explicit_active_axes(self): + field = jnp.zeros((7, 5, 5)) + field = field.at[:, 2, 2].set(jnp.linspace(1.0, 2.0, 7)) + field = field.at[:, 1, 2].set(100.0) + field = field.at[:, 3, 2].set(-50.0) + field = field.at[:, 2, 1].set(25.0) + field = field.at[:, 2, 3].set(-10.0) + x = jnp.array([-0.25, 0.0, 0.25]) + y = jnp.zeros_like(x) + z = jnp.zeros_like(x) + grid = ( + jnp.linspace(-0.5, 0.5, 7), + jnp.linspace(-1.0, 1.0, 5), + jnp.linspace(-1.0, 1.0, 5), + ) + + interpolated = interpolate_field_to_particles( + field, + x, + y, + z, + grid, + shape_factor=2, + ghost_cells=True, + active_axes=(True, False, False), + inactive_axis_indices=(None, 2, 2), + ) + reference = interpolate_field_to_particles( + field[:, 2:3, 2:3], + x, + y, + z, + (grid[0], grid[1][2:3], grid[2][2:3]), + shape_factor=2, + ghost_cells=False, + ) + + self.assertTrue(jnp.allclose(interpolated, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_boris_single_particle(self): @@ -205,10 +243,6 @@ def test_higuera_cary_vectorizes_with_per_particle_charge_and_mass(self): 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/physics_tests/electrostatic_yee_test.py b/tests/physics_tests/electrostatic_yee_test.py new file mode 100644 index 00000000..1b63754f --- /dev/null +++ b/tests/physics_tests/electrostatic_yee_test.py @@ -0,0 +1,239 @@ +import unittest +import os +import tempfile + +import jax +import jax.numpy as jnp +import numpy as np +import toml + +from PyPIC3D.evolve import time_loop_electrostatic +from PyPIC3D.initialization import initialize_simulation +from PyPIC3D.particles.particle_class import TiledParticles +from tests.initial_particles import tiled_species +from PyPIC3D.solvers.electrostatic_yee import ( + solve_poisson_with_conjugate_gradient, + calculate_electrostatic_fields, + _centered_finite_difference_gradient, +) +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC + +jax.config.update("jax_enable_x64", True) + + +def apply_negative_laplacian(field, dx, dy, dz): + laplacian_x = (jnp.roll(field, shift=1, axis=0) + jnp.roll(field, shift=-1, axis=0) - 2.0 * field) / (dx * dx) + laplacian_y = (jnp.roll(field, shift=1, axis=1) + jnp.roll(field, shift=-1, axis=1) - 2.0 * field) / (dy * dy) + laplacian_z = (jnp.roll(field, shift=1, axis=2) + jnp.roll(field, shift=-1, axis=2) - 2.0 * field) / (dz * dz) + return -(laplacian_x + laplacian_y + laplacian_z) + + +class TestElectrostaticYeeMethods(unittest.TestCase): + def setUp(self): + self.Nx = 16 + self.Ny = 16 + self.Nz = 16 + self.x_wind = 2 * jnp.pi + self.y_wind = 2 * jnp.pi + self.z_wind = 2 * jnp.pi + self.dx = self.x_wind / self.Nx + self.dy = self.y_wind / self.Ny + self.dz = self.z_wind / self.Nz + + # Interior coordinate grid (no ghost cells) for analytical solutions + x = jnp.linspace(0, self.x_wind, self.Nx, endpoint=False) + y = jnp.linspace(0, self.y_wind, self.Ny, endpoint=False) + z = jnp.linspace(0, self.z_wind, self.Nz, endpoint=False) + self.X, self.Y, self.Z = jnp.meshgrid(x, y, z, indexing='ij') + + self.constants = { + "eps": 1.0, + "alpha": 1.0, + } + self.world = { + "Nx": self.Nx, + "Ny": self.Ny, + "Nz": self.Nz, + "dx": self.dx, + "dy": self.dy, + "dz": self.dz, + "x_wind": self.x_wind, + "y_wind": self.y_wind, + "z_wind": self.z_wind, + "shape_factor": 1, + "guard_cells": 2, + "tile_shape": (self.Nx, self.Ny, self.Nz), + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + } + vertex_grid, center_grid = build_yee_grid(self.world) + self.world["grids"] = { + "vertex": vertex_grid, + "center": center_grid, + } + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(self.world, self.world["tile_shape"], self.world["guard_cells"]) + self.world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + self.world["grids"]["tiled_center_grid"] = tiled_center_grid + + self.g = int(self.world["guard_cells"]) + self.active = slice(self.g, -self.g) + # Single tile-local fields: shape (Nx+2*g, Ny+2*g, Nz+2*g) + self.initial_rho = jnp.zeros((self.Nx + 2 * self.g, self.Ny + 2 * self.g, self.Nz + 2 * self.g)) + self.initial_phi = jnp.zeros((self.Nx + 2 * self.g, self.Ny + 2 * self.g, self.Nz + 2 * self.g)) + + self.particles = [ + tiled_species( + name="test", + N_particles=1, + charge=1.0, + mass=1.0, + T=1.0, + v1=jnp.zeros(1), + v2=jnp.zeros(1), + v3=jnp.zeros(1), + x1=jnp.array([0.1]), + x2=jnp.array([0.2]), + x3=jnp.array([0.3]), + xwind=self.x_wind, + ywind=self.y_wind, + zwind=self.z_wind, + dx=self.dx, + dy=self.dy, + dz=self.dz, + weight=1.0, + x_bc="periodic", + y_bc="periodic", + z_bc="periodic", + dt=0.0, + ) + ] + + def test_solve_poisson_with_conjugate_gradient_single_mode(self): + phi_true_interior = jnp.sin(self.X + self.Y + self.Z) + rhs = apply_negative_laplacian(phi_true_interior, self.dx, self.dy, self.dz) + rho_interior = rhs * self.constants["eps"] + + # Place rho into ghost-celled array + rho = jnp.zeros_like(self.initial_rho) + rho = rho.at[self.active, self.active, self.active].set(rho_interior) + + phi = solve_poisson_with_conjugate_gradient( + rho, + self.initial_phi, + self.constants, + self.world, + tol=1e-10, + max_iter=4000, + ) + + # Compare on interior + phi_num = phi[self.active, self.active, self.active] + phi_num = phi_num - jnp.mean(phi_num) + phi_true = phi_true_interior - jnp.mean(phi_true_interior) + + self.assertEqual(phi_num.shape, phi_true.shape) + self.assertTrue(jnp.allclose(phi_num, phi_true, atol=1e-7, rtol=1e-6)) + + residual = apply_negative_laplacian(phi_num, self.dx, self.dy, self.dz) - rho_interior / self.constants["eps"] + self.assertLess(jnp.max(jnp.abs(residual)), 1e-6) + + def test_electrostatic_field_solve_uses_local_centered_gradient(self): + E, phi, rho = calculate_electrostatic_fields( + self.world, + self.particles, + self.constants, + self.initial_rho, + self.initial_phi, + "electrostatic", + "periodic", + ) + + expected_phi = solve_poisson_with_conjugate_gradient( + rho, + self.initial_phi, + self.constants, + self.world, + ) + expected_Ex, expected_Ey, expected_Ez = _centered_finite_difference_gradient( + -1.0 * expected_phi[self.active, self.active, self.active], + self.dx, + self.dy, + self.dz, + ) + + self.assertTrue(jnp.allclose(phi[self.active, self.active, self.active], expected_phi[self.active, self.active, self.active], atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E[0][self.active, self.active, self.active], expected_Ex, atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E[1][self.active, self.active, self.active], expected_Ey, atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E[2][self.active, self.active, self.active], expected_Ez, atol=1e-6, rtol=1e-6)) + + def test_initialize_simulation_accepts_single_tile_electrostatic(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + vx_path = os.path.join(tmpdir, "vx.npy") + np.save(x_path, np.array([-1.5, -0.5, 0.5, 1.5])) + np.save(zeros_path, np.zeros(4)) + np.save(vx_path, np.array([0.10, -0.05, 0.07, -0.02])) + + config = { + "simulation_parameters": { + "name": "electrostatic init smoke", + "output_dir": tmpdir, + "solver": "electrostatic", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 0.01, + "Nt": 1, + "shape_factor": 1, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "j_from_rhov", + "filter_j": "none", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": {"plotting_interval": 1}, + "particle1": { + "name": "electrons", + "N_particles": 4, + "charge": -1.0, + "mass": 2.0, + "weight": 0.5, + "temperature": 1.0, + "initial_x": x_path, + "initial_y": zeros_path, + "initial_z": zeros_path, + "initial_vx": vx_path, + "initial_vy": zeros_path, + "initial_vz": zeros_path, + }, + } + + loop, particles, fields, world, simulation_parameters, *_ = initialize_simulation(toml.loads(toml.dumps(config))) + + self.assertIs(loop, time_loop_electrostatic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(tuple(world["tile_shape"]), (8, 1, 1)) + self.assertEqual(tuple(simulation_parameters["tile_shape"]), (8, 1, 1)) + self.assertEqual(int(simulation_parameters["particle_tile_nx"]), 8) + self.assertEqual(int(simulation_parameters["particle_tile_ny"]), 1) + self.assertEqual(int(simulation_parameters["particle_tile_nz"]), 1) + for vertex_axis, center_axis in zip(world["grids"]["vertex"], world["grids"]["center"]): + self.assertTrue(jnp.allclose(vertex_axis, center_axis)) + for tiled_vertex_axis, tiled_center_axis in zip( + world["grids"]["tiled_vertex_grid"], + world["grids"]["tiled_center_grid"], + ): + self.assertTrue(jnp.allclose(tiled_vertex_axis, tiled_center_axis)) + self.assertEqual(fields[0][0].shape[:3], (1, 1, 1)) + self.assertEqual(fields[3].shape[:3], (1, 1, 1)) + self.assertEqual(fields[4].shape[:3], (1, 1, 1)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/physics_tests/pml_test.py b/tests/physics_tests/pml_test.py new file mode 100644 index 00000000..69a9ff27 --- /dev/null +++ b/tests/physics_tests/pml_test.py @@ -0,0 +1,572 @@ +import tempfile +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.PML import ( + PML_WALLS, + build_pml, + initialize_pml_state, + initialize_tiled_pml_state, + load_pml_from_toml, + tile_pml_profiles, +) +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field +from PyPIC3D.initialization import initialize_simulation +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles +from PyPIC3D.solvers.first_order_yee import ( + update_B, + update_E, +) +from PyPIC3D.utilities.grids import build_yee_grid +from PyPIC3D.utils import compute_energy + +jax.config.update("jax_enable_x64", True) + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Shared tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(num_guard_cells) + Nx = int(field.shape[0]) - 2 + Ny = int(field.shape[1]) - 2 + Nz = int(field.shape[2]) - 2 + ntx = _tile_axis_count(Nx, tile_nx) + nty = _tile_axis_count(Ny, tile_ny) + ntz = _tile_axis_count(Nz, tile_nz) + # get the number of tiles along each axis + + interior_tiles = field[1:-1, 1:-1, 1:-1] + interior_tiles = interior_tiles.reshape(ntx, tile_nx, nty, tile_ny, ntz, tile_nz) + interior_tiles = interior_tiles.transpose(0, 2, 4, 1, 3, 5) + # reshape the interior of the field into tiles, and then transpose to get the correct order of axes + + field_tiles = jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=field.dtype, + ) + field_tiles = field_tiles.at[:, :, :, g:-g, g:-g, g:-g].set(interior_tiles) + # populate the field tiles with the interior tiles, leaving the guard cells as zeros + + return ghost_cells.update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + # update the guard cells of the tiled field using the ghost_cells function + + +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) + + +def _update_ghost_cells(field, bc_x, bc_y, bc_z): + field = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda f: f.at[0, :, :].set(f[-2, :, :]).at[-1, :, :].set(f[1, :, :]), + lambda f: f.at[0, :, :].set(0.0).at[-1, :, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda f: f.at[:, 0, :].set(f[:, -2, :]).at[:, -1, :].set(f[:, 1, :]), + lambda f: f.at[:, 0, :].set(0.0).at[:, -1, :].set(0.0), + operand=field, + ) + field = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda f: f.at[:, :, 0].set(f[:, :, -2]).at[:, :, -1].set(f[:, :, 1]), + lambda f: f.at[:, :, 0].set(0.0).at[:, :, -1].set(0.0), + operand=field, + ) + return field + + +def _update_ghost_cells_for_pml(field, world): + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + _, pml_x, pml_y, pml_z, _ = world["pml"] + + bc_x = jnp.where((pml_x) & (bc_x == BC_PERIODIC), BC_CONDUCTING, bc_x) + bc_y = jnp.where((pml_y) & (bc_y == BC_PERIODIC), BC_CONDUCTING, bc_y) + bc_z = jnp.where((pml_z) & (bc_z == BC_PERIODIC), BC_CONDUCTING, bc_z) + + return _update_ghost_cells(field, bc_x, bc_y, bc_z) + + +def _empty_global_fields(world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + E = (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + B = (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + J = (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + return E, B, J + + +def _base_world(nx=24, ny=1, nz=1): + world = { + "Nx": nx, + "Ny": ny, + "Nz": nz, + "dx": 1.0 / nx, + "dy": 1.0, + "dz": 1.0, + "dt": 0.5 / nx, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "guard_cells": 2, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_yee_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + return world + + +def _empty_config(tmpdir, solver="electrodynamic_yee", pml=None): + sim = { + "name": "pml init test", + "output_dir": tmpdir, + "solver": solver, + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1e-10, + } + config = {"simulation_parameters": sim, "plotting": {"plotting": False}} + if pml is not None: + config["pml"] = pml + return config + + +class TestPMLConfiguration(unittest.TestCase): + def test_load_pml_from_toml_accepts_all_six_walls(self): + raw = [ + {"wall": wall, "thickness": 2, "order": 3.0, "target_reflection": 1e-8} + for wall in PML_WALLS + ] + + active, pml_x, pml_y, pml_z, profiles = load_pml_from_toml( + raw, + _base_world(nx=8, ny=8, nz=8), + {"C": 3.0}, + ) + sigma_x, sigma_y, sigma_z = profiles + + self.assertTrue(active) + self.assertTrue(pml_x) + self.assertTrue(pml_y) + self.assertTrue(pml_z) + self.assertEqual(sigma_x.shape, (10, 10, 10)) + self.assertTrue(jnp.any(sigma_x > 0.0)) + self.assertTrue(jnp.any(sigma_y > 0.0)) + self.assertTrue(jnp.any(sigma_z > 0.0)) + + def test_load_pml_from_toml_rejects_invalid_duplicate_and_oversized_walls(self): + world = _base_world(nx=8, ny=1, nz=1) + + with self.assertRaisesRegex(ValueError, "Invalid PML wall"): + load_pml_from_toml([{"wall": "x+", "thickness": 2}], world, {"C": 3.0}) + + with self.assertRaisesRegex(ValueError, "Duplicate PML wall"): + load_pml_from_toml( + [{"wall": "+x", "thickness": 2}, {"wall": "+x", "thickness": 2}], + world, + {"C": 3.0}, + ) + + with self.assertRaisesRegex(ValueError, "exceeds active cells"): + load_pml_from_toml([{"wall": "+x", "thickness": 9}], world, {"C": 3.0}) + + def test_build_pml_ramp_only_on_requested_side(self): + world = _base_world(nx=8, ny=1, nz=1) + pml_layers = (("+x", "x", 3, 2.0, 9.0),) + + sigma_x, sigma_y, sigma_z = build_pml(world, pml_layers) + + self.assertEqual(sigma_x.shape, (10, 3, 3)) + self.assertTrue(jnp.allclose(sigma_x[1:5, :, :], 0.0)) + self.assertTrue(jnp.all(sigma_x[-4:-1, :, :] > 0.0)) + self.assertTrue(float(sigma_x[-2, 1, 1]) > float(sigma_x[-4, 1, 1])) + self.assertTrue(jnp.allclose(sigma_y, 0.0)) + self.assertTrue(jnp.allclose(sigma_z, 0.0)) + + def test_build_pml_ramps_from_interior_interface_to_outer_wall(self): + world = _base_world(nx=8, ny=1, nz=1) + pml_layers = ( + ("-x", "x", 3, 2.0, 9.0), + ("+x", "x", 3, 2.0, 9.0), + ) + + sigma_x, _, _ = build_pml(world, pml_layers) + sigma_x = sigma_x[:, 1, 1] + + self.assertGreater(float(sigma_x[1]), float(sigma_x[2])) + self.assertGreater(float(sigma_x[2]), float(sigma_x[3])) + self.assertGreater(float(sigma_x[-2]), float(sigma_x[-3])) + self.assertGreater(float(sigma_x[-3]), float(sigma_x[-4])) + + +class TestPMLInitialization(unittest.TestCase): + def test_initialize_pml_state_uses_physical_interior_shape(self): + world = _base_world(nx=8, ny=4, nz=2) + + pml_state = initialize_pml_state(world) + e_memory, b_memory = pml_state + + self.assertEqual(len(e_memory), 6) + self.assertEqual(len(b_memory), 6) + for memory in e_memory: + self.assertEqual(memory.shape, (8, 4, 2)) + for memory in b_memory: + self.assertEqual(memory.shape, (8, 4, 2)) + + def test_initialize_tiled_pml_state_uses_tile_local_interior_shape(self): + world = _base_world(nx=8, ny=4, nz=2) + world["pml"] = load_pml_from_toml( + [{"wall": "+x", "thickness": 2, "sigma_max": 3.0}], + world, + {"C": 1.0}, + ) + tile_shape = (2, 2, 1) + + pml_state = initialize_tiled_pml_state(world, tile_shape) + e_memory, b_memory, tiled_profiles = pml_state + + self.assertEqual(len(e_memory), 6) + self.assertEqual(len(b_memory), 6) + for memory in e_memory: + self.assertEqual(memory.shape, (4, 2, 2, 2, 2, 1)) + for memory in b_memory: + self.assertEqual(memory.shape, (4, 2, 2, 2, 2, 1)) + for profile in tiled_profiles: + self.assertEqual(profile.shape, (4, 2, 2, 6, 6, 5)) + + def test_tile_pml_profiles_matches_global_profiles_on_tile_interiors(self): + world = _base_world(nx=8, ny=4, nz=2) + world["pml"] = load_pml_from_toml( + [ + {"wall": "-x", "thickness": 2, "sigma_max": 3.0}, + {"wall": "+x", "thickness": 2, "sigma_max": 3.0}, + ], + world, + {"C": 1.0}, + ) + tile_shape = (2, 2, 1) + + tiled_profiles = tile_pml_profiles(world, tile_shape) + global_profiles = world["pml"][-1] + assembled_profiles = assemble_tiled_vector_field(tiled_profiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) + + for assembled, reference in zip(assembled_profiles, global_profiles): + self.assertTrue( + jnp.allclose( + assembled[1:-1, 1:-1, 1:-1], + reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_initialize_simulation_rejects_pml_for_electrostatic_solver(self): + pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] + + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaisesRegex(ValueError, "PML is only supported"): + initialize_simulation(_empty_config(tmpdir, solver="electrostatic", pml=pml)) + + def test_initialize_simulation_appends_pml_state_for_electrodynamic_yee(self): + pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] + + with tempfile.TemporaryDirectory() as tmpdir: + result = initialize_simulation(_empty_config(tmpdir, pml=pml)) + + fields = result[2] + world = result[3] + active, pml_x, pml_y, pml_z, _ = world["pml"] + self.assertEqual(len(fields), 8) + self.assertIn("pml", world) + self.assertTrue(active) + self.assertTrue(pml_x) + self.assertFalse(pml_y) + self.assertFalse(pml_z) + + def test_initialize_simulation_uses_tiled_pml_state_for_electrodynamic_yee(self): + pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] + + with tempfile.TemporaryDirectory() as tmpdir: + config = _empty_config(tmpdir, solver="electrodynamic_yee", pml=pml) + config["simulation_parameters"].update( + { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "filter_j": "none", + } + ) + result = initialize_simulation(config) + + fields = result[2] + world = result[3] + pml_state = fields[6] + overflow = fields[-1] + e_memory, b_memory, tiled_profiles = pml_state + + self.assertEqual(len(fields), 8) + self.assertFalse(bool(overflow)) + self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + self.assertEqual(e_memory[0].shape, (4, 1, 1, 2, 1, 1)) + self.assertEqual(b_memory[0].shape, (4, 1, 1, 2, 1, 1)) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertEqual(tiled_profiles[0].shape, (4, 1, 1, 6, 5, 5)) + + def test_initialize_simulation_uses_none_pml_state_without_pml_for_electrodynamic_yee(self): + with tempfile.TemporaryDirectory() as tmpdir: + result = initialize_simulation(_empty_config(tmpdir)) + + fields = result[2] + world = result[3] + active, pml_x, pml_y, pml_z, _ = world["pml"] + + self.assertEqual(len(fields), 8) + self.assertIsNone(fields[6]) + self.assertFalse(bool(fields[-1])) + self.assertFalse(active) + self.assertFalse(pml_x) + self.assertFalse(pml_y) + self.assertFalse(pml_z) + + +class TestPMLFDTDBehavior(unittest.TestCase): + def test_no_pml_state_returns_field_and_none_state(self): + world = _base_world(nx=4, ny=4, nz=4) + constants = {"C": 2.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 2, 2) + world["tile_shape"] = tile_shape + E, B, J = _empty_global_fields(world) + B = (B[0], B[1], B[2].at[1:-1, 1:-1, 1:-1].set(1.0)) + + E_tiles = tile_vector_field(E, world, tile_shape) + B_tiles = tile_vector_field(B, world, tile_shape) + J_tiles = tile_vector_field(J, world, tile_shape) + E_after, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + B_after, pml_state = update_B(E_after, B_tiles, world, constants, pml_state) + + self.assertEqual(len(E_after), 3) + self.assertEqual(len(B_after), 3) + self.assertIsNone(pml_state) + + def test_tiled_pml_matches_single_tile_pml_for_one_yee_step(self): + world = _base_world(nx=8, ny=4, nz=2) + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + world["pml"] = load_pml_from_toml( + [ + {"wall": "-x", "thickness": 2, "sigma_max": 4.0}, + {"wall": "+x", "thickness": 2, "sigma_max": 4.0}, + ], + world, + constants, + ) + tile_shape = (2, 2, 1) + E, B, J = _empty_global_fields(world) + + x = world["grids"]["vertex"][0][1:-1] + y = world["grids"]["vertex"][1][1:-1] + z = world["grids"]["vertex"][2][1:-1] + X, Y, Z = jnp.meshgrid(x, y, z, indexing="ij") + Ex, Ey, Ez = E + Bx, By, Bz = B + Jx, Jy, Jz = J + Ey = Ey.at[1:-1, 1:-1, 1:-1].set(jnp.sin(2.0 * jnp.pi * X) + 0.1 * Y) + Ez = Ez.at[1:-1, 1:-1, 1:-1].set(0.2 * X - 0.3 * Z) + By = By.at[1:-1, 1:-1, 1:-1].set(0.4 * X + 0.2 * Z) + Bz = Bz.at[1:-1, 1:-1, 1:-1].set(jnp.cos(2.0 * jnp.pi * X) - 0.1 * Y) + Jx = Jx.at[1:-1, 1:-1, 1:-1].set(0.05 * X) + Jy = Jy.at[1:-1, 1:-1, 1:-1].set(-0.02 * Y) + Jz = Jz.at[1:-1, 1:-1, 1:-1].set(0.03 * Z) + E = (Ex, Ey, Ez) + B = (Bx, By, Bz) + J = (Jx, Jy, Jz) + E = tuple(_update_ghost_cells_for_pml(component, world) for component in E) + B = tuple(_update_ghost_cells_for_pml(component, world) for component in B) + + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + world["tile_shape"] = reference_tile_shape + reference_pml_state = initialize_tiled_pml_state(world, reference_tile_shape) + reference_E, reference_pml_state = update_E( + tile_vector_field(E, world, reference_tile_shape), + tile_vector_field(B, world, reference_tile_shape), + tile_vector_field(J, world, reference_tile_shape), + world, + constants, + reference_pml_state, + ) + reference_B, reference_pml_state = update_B( + reference_E, + tile_vector_field(B, world, reference_tile_shape), + world, + constants, + reference_pml_state, + ) + E_reference = assemble_tiled_vector_field( + reference_E, + world, + reference_tile_shape, + num_guard_cells=int(world["guard_cells"]), + ) + B_reference = assemble_tiled_vector_field( + reference_B, + world, + reference_tile_shape, + num_guard_cells=int(world["guard_cells"]), + ) + + world["tile_shape"] = tile_shape + tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) + E_tiles, tiled_pml_state = update_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + tiled_pml_state, + ) + B_tiles, tiled_pml_state = update_B( + E_tiles, + tile_vector_field(B, world, tile_shape), + world, + constants, + tiled_pml_state, + ) + + E_tiled = assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) + B_tiled = assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) + + interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + for reference, tiled in zip(E_reference, E_tiled): + self.assertTrue(jnp.allclose(tiled[interior], reference[interior], rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(B_reference, B_tiled): + self.assertTrue(jnp.allclose(tiled[interior], reference[interior], rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): + world = _base_world(nx=40, ny=1, nz=1) + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + world["pml"] = load_pml_from_toml( + [ + {"wall": "-x", "thickness": 8, "order": 3.0, "sigma_max": 60.0}, + {"wall": "+x", "thickness": 8, "order": 3.0, "sigma_max": 60.0}, + ], + world, + constants, + ) + tile_shape = (4, 1, 1) + world["tile_shape"] = tile_shape + tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) + E, B, J = _empty_global_fields(world) + + x = world["grids"]["vertex"][0][1:-1] + pulse = jnp.exp(-((x + 0.30) / 0.04) ** 2) + Ex, Ey, Ez = E + Bx, By, Bz = B + Ey = Ey.at[1:-1, 1, 1].set(pulse) + Bz = Bz.at[1:-1, 1, 1].set(pulse) + E_tiles = tile_vector_field((Ex, Ey, Ez), world, tile_shape) + B_tiles = tile_vector_field((Bx, By, Bz), world, tile_shape) + J_tiles = tile_vector_field(J, world, tile_shape) + + empty_particles = TiledParticles( + x=jnp.zeros((1, 1, 1, 0, 0, 3)), + u=jnp.zeros((1, 1, 1, 0, 0, 3)), + active=jnp.zeros((1, 1, 1, 0, 0), dtype=bool), + ) + empty_species_config = SpeciesConfig( + charge=jnp.zeros((0,)), + mass=jnp.zeros((0,)), + weight=jnp.zeros((0,)), + update_x=jnp.zeros((0, 3), dtype=bool), + update_u=jnp.zeros((0, 3), dtype=bool), + ) + + initial_energy = sum( + compute_energy( + empty_particles, + E_tiles, + B_tiles, + world, + constants, + species_config=empty_species_config, + )[:2] + ) + def step(E_tiles, B_tiles, tiled_pml_state): + E_tiles, tiled_pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants, tiled_pml_state) + B_tiles, tiled_pml_state = update_B(E_tiles, B_tiles, world, constants, tiled_pml_state) + return E_tiles, B_tiles, tiled_pml_state + + step = jax.jit(step) + for _ in range(60): + E_tiles, B_tiles, tiled_pml_state = step(E_tiles, B_tiles, tiled_pml_state) + + final_energy = sum( + compute_energy( + empty_particles, + E_tiles, + B_tiles, + world, + constants, + species_config=empty_species_config, + )[:2] + ) + self.assertTrue(jnp.isfinite(final_energy)) + self.assertLess(float(final_energy), 0.65 * float(initial_energy)) + + def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): + world = _base_world(nx=8, ny=1, nz=1) + constants = {"C": 1.0, "eps": 2.0, "mu": 1.0, "alpha": 1.0} + world["pml"] = load_pml_from_toml( + [{"wall": "+x", "thickness": 2, "order": 3.0, "sigma_max": 4.0}], + world, + constants, + ) + tile_shape = (2, 1, 1) + world["tile_shape"] = tile_shape + E, B, J = _empty_global_fields(world) + Ex, Ey, Ez = E + Bx, By, Bz = B + Ey = Ey.at[1:-1, 1, 1].set(jnp.linspace(0.0, 0.3, world["Nx"])) + Bz = Bz.at[1:-1, 1, 1].set(jnp.linspace(0.1, -0.2, world["Nx"])) + E_tiles = tile_vector_field((Ex, Ey, Ez), world, tile_shape) + B_tiles = tile_vector_field((Bx, By, Bz), world, tile_shape) + + g = int(world["guard_cells"]) + J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + Jx, Jy, Jz = J_tiles + Jx = Jx.at[:, :, :, 1:-1, 1:-1, 1:-1].set(0.25) + J_tiles = (Jx, Jy, Jz) + + pml_state = initialize_tiled_pml_state(world, tile_shape) + E_after, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state) + B_after, pml_state = update_B(E_after, B_tiles, world, constants, pml_state) + + for component in E_after + B_after: + self.assertTrue(jnp.all(jnp.isfinite(component))) + for memory in pml_state[0] + pml_state[1]: + self.assertTrue(jnp.all(jnp.isfinite(memory))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/physics_tests/pusher_schmitz_test.py b/tests/physics_tests/pusher_schmitz_test.py new file mode 100644 index 00000000..bd0b3576 --- /dev/null +++ b/tests/physics_tests/pusher_schmitz_test.py @@ -0,0 +1,203 @@ +import math +import unittest + +import jax + +from PyPIC3D.pusher.boris import relativistic_boris_single_particle +from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle + + +jax.config.update("jax_enable_x64", True) + + +class TestSchmitzParticlePusherBenchmarks(unittest.TestCase): + """ + Focused tests from Schmitz 2026, Section 3, for pushers supported here. + + The paper compares several relativistic pushers. PyPIC3D currently exposes + relativistic Boris and Higuera-Cary, so the checks below use those two + update maps and preserve the paper's normalised parameters where possible. + """ + + constants = {"C": 1.0} + q = 1.0 + m = 1.0 + + def _gamma(self, v): + v2 = sum(component * component for component in v) + return 1.0 / math.sqrt(1.0 - v2) + + def _velocity_from_u(self, u): + gamma = math.sqrt(1.0 + sum(component * component for component in u)) + return tuple(component / gamma for component in u) + + def _step(self, pusher, v, E, B, dt): + vx, vy, vz = pusher( + v[0], + v[1], + v[2], + E[0], + E[1], + E[2], + B[0], + B[1], + B[2], + self.q, + self.m, + dt, + self.constants, + ) + return float(vx), float(vy), float(vz) + + def _advance(self, pusher, v, E, B, dt, steps): + for _ in range(steps): + v = self._step(pusher, v, E, B, dt) + return v + + def test_constant_magnetic_field_preserves_energy_for_boris_and_higuera_cary(self): + # Section 3.1: E = 0, B = B_z zhat, gamma0 = 1.001 and 10, + # with ten steps of dt omega_c / (2 pi) = 0.1. + B = (0.0, 0.0, 1.0) + E = (0.0, 0.0, 0.0) + + for gamma0 in (1.001, 10.0): + u0 = (math.sqrt(gamma0**2 - 1.0), 0.0, 0.0) + v0 = self._velocity_from_u(u0) + dt = 0.2 * math.pi * gamma0 + + for pusher in (relativistic_boris_single_particle, higuera_cary_single_particle): + v = self._advance(pusher, v0, E, B, dt, 10) + + self.assertAlmostEqual(self._gamma(v), gamma0, delta=2.0e-12) + + def test_force_free_crossed_fields_expose_boris_relativistic_weakness(self): + # Section 3.2: E = -v0 x B, gamma0 = 10. Higuera-Cary is designed + # to keep the force-free trajectory exact; Boris is retained as a + # bounded known-weakness comparison rather than forced to pass HC's + # near-machine-precision criterion. + gamma0 = 10.0 + u0 = math.sqrt(gamma0**2 - 1.0) + v0 = u0 / gamma0 + E = (0.0, v0, 0.0) + B = (0.0, 0.0, 1.0) + dt = 0.2 * math.pi * gamma0 + + higuera_cary_v = (v0, 0.0, 0.0) + boris_v = (v0, 0.0, 0.0) + higuera_cary_angle = 0.0 + higuera_cary_energy_error = 0.0 + boris_angle = 0.0 + boris_energy_error = 0.0 + + for _ in range(20): + higuera_cary_v = self._step(higuera_cary_single_particle, higuera_cary_v, E, B, dt) + boris_v = self._step(relativistic_boris_single_particle, boris_v, E, B, dt) + + higuera_cary_angle = max(higuera_cary_angle, abs(math.atan2(higuera_cary_v[1], higuera_cary_v[0]))) + higuera_cary_energy_error = max( + higuera_cary_energy_error, + abs((self._gamma(higuera_cary_v) - gamma0) / (gamma0 - 1.0)), + ) + boris_angle = max(boris_angle, abs(math.atan2(boris_v[1], boris_v[0]))) + boris_energy_error = max( + boris_energy_error, + abs((self._gamma(boris_v) - gamma0) / (gamma0 - 1.0)), + ) + + self.assertLess(higuera_cary_angle, 1.0e-13) + self.assertLess(higuera_cary_energy_error, 1.0e-12) + self.assertGreater(boris_angle, 1.0e-3) + self.assertLess(boris_angle, 2.0e-1) + self.assertGreater(boris_energy_error, 1.0e-2) + + def _boosted_gyration_exact_velocity(self, lab_time, gamma_particle=10.0, gamma_frame=5.0): + beta_frame = math.sqrt(1.0 - 1.0 / gamma_frame**2) + u_particle = math.sqrt(gamma_particle**2 - 1.0) + v_particle = u_particle / gamma_particle + omega = 1.0 / gamma_particle + radius = v_particle / omega + + def x_moving_frame(t): + return radius * (1.0 - math.cos(omega * t)) + + def vx_moving_frame(t): + return v_particle * math.sin(omega * t) + + def vy_moving_frame(t): + return v_particle * math.cos(omega * t) + + moving_time = lab_time / gamma_frame + for _ in range(8): + residual = gamma_frame * ( + moving_time + beta_frame * x_moving_frame(moving_time) + ) - lab_time + derivative = gamma_frame * ( + 1.0 + beta_frame * vx_moving_frame(moving_time) + ) + moving_time -= residual / derivative + + vx_m = vx_moving_frame(moving_time) + vy_m = vy_moving_frame(moving_time) + denominator = 1.0 + beta_frame * vx_m + return ( + (vx_m + beta_frame) / denominator, + vy_m / (gamma_frame * denominator), + 0.0, + ) + + def test_lorentz_boosted_crossed_field_gyration_matches_exact_velocity(self): + # Section 3.3: gamma_P = 10 and gamma_M = 5. Fields are the Lorentz + # transform of a pure magnetic field in the moving frame, so the exact + # lab-frame velocity is obtained by transforming the circular orbit. + gamma_particle = 10.0 + gamma_frame = 5.0 + beta_frame = math.sqrt(1.0 - 1.0 / gamma_frame**2) + v_particle = math.sqrt(1.0 - 1.0 / gamma_particle**2) + E = (0.0, beta_frame * gamma_frame, 0.0) + B = (0.0, 0.0, gamma_frame) + v0 = (beta_frame, v_particle / gamma_frame, 0.0) + gyration_period = 2.0 * math.pi * gamma_particle * gamma_frame + dt = 1.0e-3 * gyration_period + steps = 20 + + exact_v = self._boosted_gyration_exact_velocity(steps * dt, gamma_particle, gamma_frame) + higuera_cary_v = self._advance(higuera_cary_single_particle, v0, E, B, dt, steps) + boris_v = self._advance(relativistic_boris_single_particle, v0, E, B, dt, steps) + + higuera_cary_error = math.sqrt(sum((higuera_cary_v[i] - exact_v[i]) ** 2 for i in range(3))) + boris_error = math.sqrt(sum((boris_v[i] - exact_v[i]) ** 2 for i in range(3))) + + self.assertLess(higuera_cary_error, 1.0e-6) + self.assertLess(boris_error, 1.0e-3) + + def test_oscillating_parallel_electric_field_returns_to_initial_energy(self): + # Section 3.7: B = B0 zhat and E = E0 zhat cos(omega0 t), with + # gamma_perp = 1.1, omega0 = omega_perp / 2, and + # E0 = 10 m omega0 c / q. A midpoint field sample preserves the + # oscillatory impulse over five electric-field periods. + gamma_perp = 1.1 + u_perp = math.sqrt(gamma_perp**2 - 1.0) + v0 = self._velocity_from_u((u_perp, 0.0, 0.0)) + B0 = 1.0 + omega_perp = B0 / gamma_perp + omega0 = 0.5 * omega_perp + E0 = 10.0 * omega0 + electric_period = 2.0 * math.pi / omega0 + steps_per_period = 100 + dt = electric_period / steps_per_period + steps = 5 * steps_per_period + + for pusher in (relativistic_boris_single_particle, higuera_cary_single_particle): + v = v0 + for step in range(steps): + t_mid = (step + 0.5) * dt + E = (0.0, 0.0, E0 * math.cos(omega0 * t_mid)) + B = (0.0, 0.0, B0) + v = self._step(pusher, v, E, B, dt) + + self.assertAlmostEqual(v[2], 0.0, delta=1.0e-10) + self.assertAlmostEqual(self._gamma(v), gamma_perp, delta=1.0e-10) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/physics_tests/yee_convergence_test.py b/tests/physics_tests/yee_convergence_test.py new file mode 100644 index 00000000..73510ccf --- /dev/null +++ b/tests/physics_tests/yee_convergence_test.py @@ -0,0 +1,92 @@ +import math +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.initialization import initialize_fields +from PyPIC3D.solvers.first_order_yee import update_B +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +class TestYeeConvergence(unittest.TestCase): + def _world(self, Nx, tile_shape): + x_wind = 2.0 * math.pi + world = { + "Nx": Nx, + "Ny": 1, + "Nz": 1, + "x_wind": x_wind, + "y_wind": 1.0, + "z_wind": 1.0, + "dx": x_wind / Nx, + "dy": 1.0, + "dz": 1.0, + "dt": 1.0e-3, + "tile_shape": tile_shape, + "guard_cells": 2, + "boundary_conditions": { + "x": BC_PERIODIC, + "y": BC_PERIODIC, + "z": BC_PERIODIC, + }, + "field_boundary_conditions": (BC_PERIODIC, BC_PERIODIC, BC_PERIODIC), + "grids": {}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"]["vertex"] = vertex_grid + world["grids"]["center"] = center_grid + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids(world) + world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid + world["grids"]["tiled_center_grid"] = tiled_center_grid + return world + + def call_updateB(self, Nx, tile_shape): + world = self._world(Nx, tile_shape) + E, B, _J, _phi, _rho = initialize_fields(world) + Ex, Ey, Ez = E + g = int(world["guard_cells"]) + active = slice(g, -g) + + x_center = world["grids"]["tiled_center_grid"][0][:, :, :, active] + Ez_values = jnp.sin(x_center)[:, :, :, :, jnp.newaxis, jnp.newaxis] + Ez = Ez.at[:, :, :, active, active, active].set(Ez_values) + + _Bx, By, _Bz = update_B((Ex, Ey, Ez), B, world, {"alpha": 1.0})[0] + + x_vertex = world["grids"]["tiled_vertex_grid"][0][:, :, :, active] + exact_By = world["dt"] * jnp.cos(x_vertex) + diff = By[:, :, :, active, active, active] - exact_By[:, :, :, :, jnp.newaxis, jnp.newaxis] + + return float(jnp.sqrt(jnp.mean(diff**2))) + + def test_public_tiled_update_B_converges_for_one_tile_and_multi_tile_layouts(self): + # A no-particle periodic E_z = sin(x) mode isolates the Yee curl + # dE_z/dx used in the B_y update. Errors are evaluated on tile-local + # B_y interiors before any global field assembly. + layouts = ( + lambda Nx: (Nx, 1, 1), + lambda _Nx: (8, 1, 1), + lambda _Nx: (16, 1, 1), + ) + + for layout in layouts: + errors = [] + for Nx in (32, 64, 128): + errors.append(self.call_updateB(Nx, layout(Nx))) + # step through one timestep of the update B method + + first_order = math.log(errors[0] / errors[1], 2.0) + second_order = math.log(errors[1] / errors[2], 2.0) + # calculate the observed order of convergence for the first and second refinement steps + + self.assertGreater(first_order, 1.8) + self.assertGreater(second_order, 1.8) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pml_test.py b/tests/pml_test.py deleted file mode 100644 index 3c85e122..00000000 --- a/tests/pml_test.py +++ /dev/null @@ -1,238 +0,0 @@ -import tempfile -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.boundary_conditions.PML import ( - PML_WALLS, - build_pml, - initialize_pml_state, - load_pml_from_toml, -) -from PyPIC3D.initialization import initialize_fields, initialize_simulation -from PyPIC3D.solvers.first_order_yee import update_B, update_E -from PyPIC3D.utils import build_yee_grid, compute_energy - -jax.config.update("jax_enable_x64", True) - - -def _base_world(nx=24, ny=1, nz=1): - world = { - "Nx": nx, - "Ny": ny, - "Nz": nz, - "dx": 1.0 / nx, - "dy": 1.0, - "dz": 1.0, - "dt": 0.5 / nx, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_yee_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - return world - - -def _empty_config(tmpdir, solver="fdtd", electrostatic=False, pml=None): - sim = { - "name": "pml init test", - "output_dir": tmpdir, - "solver": solver, - "electrostatic": electrostatic, - "Nx": 8, - "Ny": 1, - "Nz": 1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "Nt": 1, - "dt": 1e-10, - "fast_backend": "default", - } - config = {"simulation_parameters": sim, "plotting": {"plotting": False}} - if pml is not None: - config["pml"] = pml - return config - - -class TestPMLConfiguration(unittest.TestCase): - def test_load_pml_from_toml_accepts_all_six_walls(self): - raw = [ - {"wall": wall, "thickness": 2, "order": 3.0, "target_reflection": 1e-8} - for wall in PML_WALLS - ] - - active, pml_x, pml_y, pml_z, profiles = load_pml_from_toml( - raw, - _base_world(nx=8, ny=8, nz=8), - {"C": 3.0}, - ) - sigma_x, sigma_y, sigma_z = profiles - - self.assertTrue(active) - self.assertTrue(pml_x) - self.assertTrue(pml_y) - self.assertTrue(pml_z) - self.assertEqual(sigma_x.shape, (10, 10, 10)) - self.assertTrue(jnp.any(sigma_x > 0.0)) - self.assertTrue(jnp.any(sigma_y > 0.0)) - self.assertTrue(jnp.any(sigma_z > 0.0)) - - def test_load_pml_from_toml_rejects_invalid_duplicate_and_oversized_walls(self): - world = _base_world(nx=8, ny=1, nz=1) - - with self.assertRaisesRegex(ValueError, "Invalid PML wall"): - load_pml_from_toml([{"wall": "x+", "thickness": 2}], world, {"C": 3.0}) - - with self.assertRaisesRegex(ValueError, "Duplicate PML wall"): - load_pml_from_toml( - [{"wall": "+x", "thickness": 2}, {"wall": "+x", "thickness": 2}], - world, - {"C": 3.0}, - ) - - with self.assertRaisesRegex(ValueError, "exceeds active cells"): - load_pml_from_toml([{"wall": "+x", "thickness": 9}], world, {"C": 3.0}) - - def test_build_pml_ramp_only_on_requested_side(self): - world = _base_world(nx=8, ny=1, nz=1) - pml_layers = (("+x", "x", 3, 2.0, 9.0),) - - sigma_x, sigma_y, sigma_z = build_pml(world, pml_layers) - - self.assertEqual(sigma_x.shape, (10, 3, 3)) - self.assertTrue(jnp.allclose(sigma_x[1:5, :, :], 0.0)) - self.assertTrue(jnp.all(sigma_x[-4:-1, :, :] > 0.0)) - self.assertTrue(float(sigma_x[-2, 1, 1]) > float(sigma_x[-4, 1, 1])) - self.assertTrue(jnp.allclose(sigma_y, 0.0)) - self.assertTrue(jnp.allclose(sigma_z, 0.0)) - - def test_build_pml_ramps_from_interior_interface_to_outer_wall(self): - world = _base_world(nx=8, ny=1, nz=1) - pml_layers = ( - ("-x", "x", 3, 2.0, 9.0), - ("+x", "x", 3, 2.0, 9.0), - ) - - sigma_x, _, _ = build_pml(world, pml_layers) - sigma_x = sigma_x[:, 1, 1] - - self.assertGreater(float(sigma_x[1]), float(sigma_x[2])) - self.assertGreater(float(sigma_x[2]), float(sigma_x[3])) - self.assertGreater(float(sigma_x[-2]), float(sigma_x[-3])) - self.assertGreater(float(sigma_x[-3]), float(sigma_x[-4])) - - -class TestPMLInitialization(unittest.TestCase): - def test_initialize_pml_state_uses_physical_interior_shape(self): - world = _base_world(nx=8, ny=4, nz=2) - - pml_state = initialize_pml_state(world) - e_memory, b_memory = pml_state - - self.assertEqual(len(e_memory), 6) - self.assertEqual(len(b_memory), 6) - for memory in e_memory: - self.assertEqual(memory.shape, (8, 4, 2)) - for memory in b_memory: - self.assertEqual(memory.shape, (8, 4, 2)) - - def test_initialize_simulation_rejects_pml_for_non_fdtd_solvers(self): - pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] - - with tempfile.TemporaryDirectory() as tmpdir: - with self.assertRaisesRegex(ValueError, "PML is only supported"): - initialize_simulation(_empty_config(tmpdir, solver="spectral", pml=pml)) - - with tempfile.TemporaryDirectory() as tmpdir: - with self.assertRaisesRegex(ValueError, "PML is only supported"): - initialize_simulation(_empty_config(tmpdir, solver="vector_potential", pml=pml)) - - with tempfile.TemporaryDirectory() as tmpdir: - with self.assertRaisesRegex(ValueError, "PML is only supported"): - initialize_simulation(_empty_config(tmpdir, electrostatic=True, pml=pml)) - - def test_initialize_simulation_appends_pml_state_for_fdtd(self): - pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] - - with tempfile.TemporaryDirectory() as tmpdir: - result = initialize_simulation(_empty_config(tmpdir, pml=pml)) - - fields = result[2] - world = result[3] - active, pml_x, pml_y, pml_z, _ = world["pml"] - self.assertEqual(len(fields), 7) - self.assertIn("pml", world) - self.assertTrue(active) - self.assertTrue(pml_x) - self.assertFalse(pml_y) - self.assertFalse(pml_z) - - def test_initialize_simulation_uses_none_pml_state_without_pml_for_fdtd(self): - with tempfile.TemporaryDirectory() as tmpdir: - result = initialize_simulation(_empty_config(tmpdir)) - - fields = result[2] - world = result[3] - active, pml_x, pml_y, pml_z, _ = world["pml"] - - self.assertEqual(len(fields), 7) - self.assertIsNone(fields[-1]) - self.assertFalse(active) - self.assertFalse(pml_x) - self.assertFalse(pml_y) - self.assertFalse(pml_z) - - -class TestPMLFDTDBehavior(unittest.TestCase): - def test_no_pml_state_returns_field_and_none_state(self): - world = _base_world(nx=4, ny=4, nz=4) - constants = {"C": 2.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) - B = (B[0], B[1], B[2].at[1:-1, 1:-1, 1:-1].set(1.0)) - - E_after, pml_state = update_E(E, B, J, world, constants, lambda *args: None) - B_after, pml_state = update_B(E_after, B, world, constants, lambda *args: None, pml_state) - - self.assertEqual(len(E_after), 3) - self.assertEqual(len(B_after), 3) - self.assertIsNone(pml_state) - - def test_pml_absorbs_field_energy_in_particle_free_1d_wave(self): - world = _base_world(nx=80, ny=1, nz=1) - constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - world["pml"] = load_pml_from_toml( - [ - {"wall": "-x", "thickness": 12, "order": 3.0, "sigma_max": 40.0}, - {"wall": "+x", "thickness": 12, "order": 3.0, "sigma_max": 40.0}, - ], - world, - constants, - ) - pml_state = initialize_pml_state(world) - E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) - - x = world["grids"]["vertex"][0][1:-1] - pulse = jnp.exp(-((x + 0.25) / 0.06) ** 2) - Ex, Ey, Ez = E - Bx, By, Bz = B - Ey = Ey.at[1:-1, 1, 1].set(pulse) - Bz = Bz.at[1:-1, 1, 1].set(pulse) - E = (Ex, Ey, Ez) - B = (Bx, By, Bz) - - initial_energy = sum(compute_energy([], E, B, world, constants)[:2]) - for _ in range(180): - E, pml_state = update_E(E, B, J, world, constants, lambda *args: None, pml_state) - B, pml_state = update_B(E, B, world, constants, lambda *args: None, pml_state) - - final_energy = sum(compute_energy([], E, B, world, constants)[:2]) - self.assertTrue(jnp.isfinite(final_energy)) - self.assertLess(float(final_energy), 0.35 * float(initial_energy)) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/pstd_test.py b/tests/pstd_test.py deleted file mode 100644 index 85a6780c..00000000 --- a/tests/pstd_test.py +++ /dev/null @@ -1,329 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -from PyPIC3D.solvers.pstd import spectral_gradient, spectral_poisson_solve, spectral_curl, spectral_laplacian, spectral_divergence -from PyPIC3D.utils import mse, convergence_test, mae - -jax.config.update("jax_enable_x64", True) - -class TestSpectralMethods(unittest.TestCase): - def setUp(self): - # Use a single Fourier mode for all tests - self.Nx = 16 - self.Ny = 16 - self.Nz = 16 - self.Lx = 2 * jnp.pi - self.Ly = 2 * jnp.pi - self.Lz = 2 * jnp.pi - self.dx = self.Lx / self.Nx - self.dy = self.Ly / self.Ny - self.dz = self.Lz / self.Nz - x = jnp.linspace(0, self.Lx, self.Nx, endpoint=False) - y = jnp.linspace(0, self.Ly, self.Ny, endpoint=False) - z = jnp.linspace(0, self.Lz, self.Nz, endpoint=False) - self.X, self.Y, self.Z = jnp.meshgrid(x, y, z, indexing='ij') - # build a grid of points in 3D space - self.phi = jnp.sin(self.X + self.Y + self.Z) - # use a simple function for phi = sin(X + Y + Z) - self.constants = {'eps': 1.0} - self.world = {'dx': self.dx, 'dy': self.dy, 'dz': self.dz} - # build constants and world parameters - self.gradx = jnp.cos(self.X + self.Y + self.Z) - self.grady = jnp.cos(self.X + self.Y + self.Z) - self.gradz = jnp.cos(self.X + self.Y + self.Z) - # Analytical derivatives - self.laplacian = -3 * self.phi - # Laplacian: -3 * sin(X + Y + Z) - self.rho = 3 * self.phi # For Poisson: Laplacian(phi) = -rho, so rho = 3*phi - self.Ex = self.gradx - self.Ey = self.grady - self.Ez = self.gradz - # For divergence test: E = grad(phi) - self.Fx = jnp.sin(self.Y) - self.Fy = jnp.sin(self.Z) - self.Fz = jnp.sin(self.X) - self.curlx = jnp.cos(self.Z) - jnp.cos(self.Y) - self.curly = jnp.cos(self.X) - jnp.cos(self.Z) - self.curlz = jnp.cos(self.Y) - jnp.cos(self.X) - # For curl test: F = (sin(Y), sin(Z), sin(X)), curl(F) = (cos(Z)-cos(Y), cos(X)-cos(Z), cos(Y)-cos(X)) - - def test_spectral_curl(self): - # Use a gradient field, whose curl should be zero - Fx = self.gradx - Fy = self.grady - Fz = self.gradz - curlx, curly, curlz = spectral_curl(Fx, Fy, Fz, self.world) - maxerrx = jnp.max(jnp.abs(curlx)) - maxerry = jnp.max(jnp.abs(curly)) - maxerrz = jnp.max(jnp.abs(curlz)) - # print('Curl (grad field) max abs error:', float(maxerrx), float(maxerry), float(maxerrz)) - self.assertTrue(jnp.allclose(curlx, 0, atol=1e-10, rtol=1e-8)) - self.assertTrue(jnp.allclose(curly, 0, atol=1e-10, rtol=1e-8)) - self.assertTrue(jnp.allclose(curlz, 0, atol=1e-10, rtol=1e-8)) - - def test_spectral_poisson_solve(self): - phi_num = spectral_poisson_solve(self.rho, self.constants, self.world) - phi_num = phi_num - jnp.mean(phi_num) - phi_true = self.phi - jnp.mean(self.phi) - maxerr = jnp.max(jnp.abs(phi_num - phi_true)) - # calculate the maximum absolute error - self.assertEqual(phi_num.shape, self.phi.shape) - self.assertTrue(jnp.allclose(phi_num, phi_true, atol=1e-10, rtol=1e-8)) - - def test_spectral_divergence(self): - divE = spectral_divergence(self.Ex, self.Ey, self.Ez, self.world) - maxerr = jnp.max(jnp.abs(divE - self.laplacian)) - # calculate the maximum absolute error - self.assertEqual(divE.shape, self.phi.shape) - self.assertTrue(jnp.allclose(divE, self.laplacian, atol=1e-10, rtol=1e-8)) - - def test_spectral_laplacian(self): - laplacian_num = spectral_laplacian(self.phi, self.world) - maxerr = jnp.max(jnp.abs(laplacian_num - self.laplacian)) - # calculate the maximum absolute error - self.assertEqual(laplacian_num.shape, self.phi.shape) - self.assertTrue(jnp.allclose(laplacian_num, self.laplacian, atol=1e-10, rtol=1e-8)) - - def test_spectral_gradient(self): - gradx, grady, gradz = spectral_gradient(self.phi, self.world) - maxerrx = jnp.max(jnp.abs(gradx - self.gradx)) - maxerry = jnp.max(jnp.abs(grady - self.grady)) - maxerrz = jnp.max(jnp.abs(gradz - self.gradz)) - # calculate the maximum absolute error - self.assertEqual(gradx.shape, self.phi.shape) - self.assertEqual(grady.shape, self.phi.shape) - self.assertEqual(gradz.shape, self.phi.shape) - self.assertTrue(jnp.allclose(gradx, self.gradx, atol=1e-10, rtol=1e-8)) - self.assertTrue(jnp.allclose(grady, self.grady, atol=1e-10, rtol=1e-8)) - self.assertTrue(jnp.allclose(gradz, self.gradz, atol=1e-10, rtol=1e-8)) - - - def test_convergence_pstd(self): - - def pstd_laplacian_comparison(nx): - """ - Computes the mean absolute error (MAE) between the numerical and analytical Laplacian of a test function - on a 3D periodic grid. - - The test function is phi(x, y, z) = x^2 + y^2 + z^2, whose Laplacian is analytically 6 everywhere. - The function constructs a cubic grid of size `nx` in each dimension, computes the numerical Laplacian - using a centered finite difference scheme with periodic boundary conditions, and compares it to the - analytical solution (6) in the interior of the domain. - - Parameters - ---------- - nx : int - Number of grid points in each spatial dimension. - - Returns - ------- - error : float - Mean squared error between the numerical and analytical Laplacian in the interior of the grid. - dx : float - Grid spacing in the x-direction (equal to y and z spacing). - """ - x_wind = 1.0 - y_wind = 1.0 - z_wind = 1.0 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(X + Y + Z) - expected_solution = -3 * phi - # define a analytical phi and its expected solution - - laplacian = spectral_laplacian(phi, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - slicer = (slice(10, -10), slice(10, -10), slice(10, -10)) - # slice for the solution comparison - - error = mae( laplacian[slicer], expected_solution[slicer] ) - # compute the mean absolute error of the laplacian against the analytical solution - - return error, dx - - - def pstd_gradient_comparison(nx): - """ - Computes the mean absolute error (MAE) between the numerical and analytical gradients - of a scalar field using spectral methods. - - The function constructs a 3D meshgrid over a symmetric cubic domain, defines a scalar - field `phi = sin(X + Y + Z)`, and computes its gradient both analytically and numerically. - The numerical gradient is computed using the `spectral_gradient` function. The MAE between - the numerical and analytical gradients is calculated for each component and summed. - - Args: - nx (int): Number of grid points along each spatial dimension. - - Returns: - tuple: - error (float): The total mean squared error between the numerical and analytical gradients. - dx (float): The spatial resolution along each axis. - """ - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - phi = jnp.sin(X + Y + Z) - expected_gradx = jnp.cos(X + Y + Z) - expected_grady = jnp.cos(X + Y + Z) - expected_gradz = jnp.cos(X + Y + Z) - # define a analytical phi and its expected solution - - gradx, grady, gradz = spectral_gradient(phi, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error_x = mae( gradx, expected_gradx ) - error_y = mae( grady, expected_grady ) - error_z = mae( gradz, expected_gradz ) - error = (error_x + error_y + error_z) / 3 - # compute the mean absolute error of the gradient against the analytical solution - - return error, dx - - - def pstd_divergence_comparison(nx): - """ - Computes the mean absolute error between the numerical and analytical divergence of a vector field. - - This function constructs a 3D meshgrid over a symmetric cubic domain, defines a vector field - with components Fx, Fy, Fz = cos(X + Y + Z), and computes its divergence both analytically and - numerically (using a spectral method). It then returns the mean absolute error between the two - divergence calculations, along with the spatial resolution. - - Args: - nx (int): Number of grid points along each axis. - - Returns: - tuple: - error (float): Mean squared error between the numerical and analytical divergence. - dx (float): Spatial resolution along each axis. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = jnp.cos(X + Y + Z) - Fy = jnp.cos(X + Y + Z) - Fz = jnp.cos(X + Y + Z) - expected_div = -3 * jnp.sin(X + Y + Z) - # define a analytical phi and its expected solution - - divF = spectral_divergence(Fx, Fy, Fz, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error = mae( divF, expected_div ) - # compute the mean absolute error of the divergence against the analytical solution - - return error, dx - - def pstd_curl_comparison(nx): - """ - Computes the mean absolute error (MAE) between the numerical and analytical curl of a vector field - defined on a 3D grid using spectral methods. - - The vector field is defined as F = (sin(Y), sin(Z), sin(X)), and its analytical curl is - (cos(Z) - cos(Y), cos(X) - cos(Z), cos(Y) - cos(X)). The function generates a symmetric 3D grid, - computes the numerical curl using the `spectral_curl` function, and returns the total MAE along - with the grid spacing. - - Args: - nx (int): Number of grid points along each spatial dimension. - - Returns: - tuple: - error (float): The sum of mean squared errors for each component of the curl. - dx (float): The spatial resolution (grid spacing) in the x-direction. - """ - - x_wind = 0.25 - y_wind = 0.25 - z_wind = 0.25 - # symmetric world volume - - x = jnp.linspace(-x_wind/2, x_wind/2, nx) - y = jnp.linspace(-y_wind/2, y_wind/2, nx) - z = jnp.linspace(-z_wind/2, z_wind/2, nx) - X, Y, Z = jnp.meshgrid(x, y, z, indexing='ij') - # build x, y, z meshgrid - - dx = x_wind / nx - dy = y_wind / nx - dz = z_wind / nx - # spatial resolution - - Fx = jnp.sin(Y) - Fy = jnp.sin(Z) - Fz = jnp.sin(X) - expected_curlx = jnp.cos(Z) - jnp.cos(Y) - expected_curly = jnp.cos(X) - jnp.cos(Z) - expected_curlz = jnp.cos(Y) - jnp.cos(X) - # For curl test: F = (sin(Y), sin(Z), sin(X)), curl(F) = (cos(Z)-cos(Y), cos(X)-cos(Z), cos(Y)-cos(X)) - # define a analytical phi and its expected solution - - curlx, curly, curlz = spectral_curl(Fx, Fy, Fz, world={'dx': dx, 'dy': dy, 'dz': dz}) - # compute the numerical result - - error_x = mae( curlx, expected_curlx ) - error_y = mae( curly, expected_curly ) - error_z = mae( curlz, expected_curlz ) - error = (error_x + error_y + error_z) / 3 - # compute the mean absolute error of the curl against the analytical solution - - return error, dx - - order = convergence_test(pstd_laplacian_comparison) - self.assertTrue(order > 1.9) - # compute order of pstd laplacian - - order = convergence_test(pstd_gradient_comparison) - self.assertTrue(order > 1.9) - # compute order of pstd gradient - - order = convergence_test(pstd_divergence_comparison) - self.assertTrue(order > 1.9) - # compute order of pstd divergence - - order = convergence_test(pstd_curl_comparison) - self.assertTrue(order > 1.9) - # compute order of pstd curl - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/utils_test.py b/tests/utils_test.py deleted file mode 100644 index 3b4bde56..00000000 --- a/tests/utils_test.py +++ /dev/null @@ -1,211 +0,0 @@ -import unittest -import os -import tempfile -import jax -import jax.numpy as jnp -import numpy as np -from PyPIC3D.initialization import initialize_fields -from PyPIC3D.utils import ( - print_stats, build_yee_grid, build_collocated_grid, check_stability, - particle_sanity_check, load_external_fields_from_toml, add_external_fields, - compute_energy, -) -from PyPIC3D.particles.species_class import particle_species - -jax.config.update("jax_enable_x64", True) - -class TestUtilsFunctions(unittest.TestCase): - def setUp(self): - self.world = { - 'Nx': 4, - 'Ny': 4, - 'Nz': 4, - 'dx': 0.1, - 'dy': 0.1, - 'dz': 0.1, - 'dt': 0.01, - 'x_wind': 1.0, - 'y_wind': 1.0, - 'z_wind': 1.0 - } - self.plasma_parameters = { - "Theoretical Plasma Frequency": 1.0, - "Debye Length": 0.01, - "Thermal Velocity": 1.0, - "Number of Electrons": 10, - "dx per debye length": 2.0 - } - - def test_build_yee_grid(self): - grid, staggered = build_yee_grid(self.world) - self.assertEqual(len(grid), 3) - self.assertEqual(len(staggered), 3) - self.assertEqual(len(grid[0]), self.world['Nx'] + 2) - self.assertEqual(len(grid[1]), self.world['Ny'] + 2) - self.assertEqual(len(grid[2]), self.world['Nz'] + 2) - self.assertEqual(len(staggered[0]), self.world['Nx'] + 2) - self.assertEqual(len(staggered[1]), self.world['Ny'] + 2) - self.assertEqual(len(staggered[2]), self.world['Nz'] + 2) - # Check that the grid and staggered arrays have the expected lengths - - def test_check_stability(self): - # Should not raise - check_stability(self.plasma_parameters, 0.01) - # Check that the stability check does not raise an error - - def test_particle_sanity_check(self): - class DummyParticles: - def __iter__(self): - return iter([self]) - def get_position(self): - N = 5 - return (jnp.zeros(N), jnp.zeros(N), jnp.zeros(N)) - def get_velocity(self): - N = 5 - return (jnp.zeros(N), jnp.zeros(N), jnp.zeros(N)) - def get_number_of_particles(self): - return 5 - particles = DummyParticles() - # Should not raise - particle_sanity_check(particles) - # Check that the particle sanity check does not raise an error - - def test_add_external_fields_adds_components(self): - E = (jnp.ones((2, 2, 2)), jnp.ones((2, 2, 2)) * 2, jnp.ones((2, 2, 2)) * 3) - B = (jnp.ones((2, 2, 2)) * 4, jnp.ones((2, 2, 2)) * 5, jnp.ones((2, 2, 2)) * 6) - external_E = (jnp.ones((2, 2, 2)) * 10, jnp.ones((2, 2, 2)) * 20, jnp.ones((2, 2, 2)) * 30) - external_B = (jnp.ones((2, 2, 2)) * 40, jnp.ones((2, 2, 2)) * 50, jnp.ones((2, 2, 2)) * 60) - - total_E, total_B = add_external_fields(E, B, (external_E, external_B)) - - self.assertTrue(jnp.allclose(total_E[0], 11.0)) - self.assertTrue(jnp.allclose(total_E[1], 22.0)) - self.assertTrue(jnp.allclose(total_E[2], 33.0)) - self.assertTrue(jnp.allclose(total_B[0], 44.0)) - self.assertTrue(jnp.allclose(total_B[1], 55.0)) - self.assertTrue(jnp.allclose(total_B[2], 66.0)) - - def test_load_external_fields_defaults_to_evolved_fields(self): - E, B, J, phi, rho = initialize_fields(2, 2, 2) - fields = [component for field in [E, B, J] for component in field] - external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) - - with tempfile.TemporaryDirectory() as tmpdir: - path = os.path.join(tmpdir, "ex.npy") - np.save(path, np.ones((2, 2, 2))) - config = {"field1": {"name": "Ex", "type": 0, "path": path}} - - fields, external_fields = load_external_fields_from_toml(fields, external_fields, config) - - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - self.assertTrue(jnp.allclose(fields[0][interior], 1.0)) - self.assertTrue(jnp.allclose(external_fields[0][0], 0.0)) - - def test_load_external_fields_evolve_true_uses_evolved_fields(self): - E, B, J, phi, rho = initialize_fields(2, 2, 2) - fields = [component for field in [E, B, J] for component in field] - external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) - - with tempfile.TemporaryDirectory() as tmpdir: - path = os.path.join(tmpdir, "by.npy") - np.save(path, np.ones((2, 2, 2)) * 3) - config = {"field1": {"name": "By", "type": 4, "path": path, "evolve": True}} - - fields, external_fields = load_external_fields_from_toml(fields, external_fields, config) - - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - self.assertTrue(jnp.allclose(fields[4][interior], 3.0)) - self.assertTrue(jnp.allclose(external_fields[1][1], 0.0)) - - def test_load_external_fields_evolve_false_uses_external_fields(self): - E, B, J, phi, rho = initialize_fields(2, 2, 2) - fields = [component for field in [E, B, J] for component in field] - external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) - - with tempfile.TemporaryDirectory() as tmpdir: - path = os.path.join(tmpdir, "bz.npy") - np.save(path, np.ones((2, 2, 2)) * 5) - config = {"field1": {"name": "external Bz", "type": 5, "path": path, "evolve": False}} - - fields, external_fields = load_external_fields_from_toml(fields, external_fields, config) - - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - self.assertTrue(jnp.allclose(fields[5][interior], 0.0)) - self.assertTrue(jnp.allclose(external_fields[1][2][interior], 5.0)) - - def test_load_external_fields_rejects_external_current(self): - E, B, J, phi, rho = initialize_fields(2, 2, 2) - fields = [component for field in [E, B, J] for component in field] - external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) - - with tempfile.TemporaryDirectory() as tmpdir: - path = os.path.join(tmpdir, "jx.npy") - np.save(path, np.ones((2, 2, 2))) - config = {"field1": {"name": "external Jx", "type": 6, "path": path, "evolve": False}} - - with self.assertRaisesRegex(ValueError, "External-only fields must be electric or magnetic"): - load_external_fields_from_toml(fields, external_fields, config) - - def test_load_external_fields_preserves_shape_validation(self): - E, B, J, phi, rho = initialize_fields(2, 2, 2) - fields = [component for field in [E, B, J] for component in field] - external_fields = (tuple(jnp.zeros_like(comp) for comp in E), tuple(jnp.zeros_like(comp) for comp in B)) - - with tempfile.TemporaryDirectory() as tmpdir: - path = os.path.join(tmpdir, "wrong.npy") - np.save(path, np.ones((3, 2, 2))) - config = {"field1": {"name": "wrong Ex", "type": 0, "path": path, "evolve": False}} - - with self.assertRaisesRegex(ValueError, "Shape mismatch"): - load_external_fields_from_toml(fields, external_fields, config) - - def test_energy_can_include_external_fields(self): - E, B, J, phi, rho = initialize_fields(1, 1, 1) - external_E = tuple(jnp.zeros_like(comp) for comp in E) - external_B = tuple(jnp.zeros_like(comp) for comp in B) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - external_E = (external_E[0].at[interior].set(2.0), external_E[1], external_E[2]) - external_B = (external_B[0], external_B[1].at[interior].set(3.0), external_B[2]) - total_E, total_B = add_external_fields(E, B, (external_E, external_B)) - - world = {"dx": 1.0, "dy": 1.0, "dz": 1.0, "Nx": 1, "Ny": 1, "Nz": 1} - constants = {"eps": 2.0, "mu": 4.0, "C": 10.0} - e_energy, b_energy, kinetic_energy = compute_energy([], total_E, total_B, world, constants) - - self.assertTrue(jnp.allclose(e_energy, 4.0)) - self.assertTrue(jnp.allclose(b_energy, 1.125)) - self.assertEqual(kinetic_energy, 0.0) - - def test_compute_energy_ignores_inactive_particles(self): - E, B, J, phi, rho = initialize_fields(1, 1, 1) - world = {"dx": 1.0, "dy": 1.0, "dz": 1.0, "Nx": 1, "Ny": 1, "Nz": 1} - constants = {"eps": 1.0, "mu": 1.0, "C": 10.0} - species = particle_species( - name="absorbed", - N_particles=1, - charge=1.0, - mass=1.0, - weight=1.0, - T=0.0, - v1=jnp.array([1.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.6]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=1.0, - dy=1.0, - dz=1.0, - x_bc="absorbing", - ) - species.boundary_conditions() - - _, _, kinetic_energy = compute_energy([species], E, B, world, constants) - - self.assertTrue(jnp.allclose(kinetic_energy, 0.0)) - -if __name__ == '__main__': - unittest.main()