From 72e83a2879ee357887ed3a8ba92790aeadbcea6e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 25 Jun 2026 11:37:56 -0500 Subject: [PATCH 01/83] removing old vector potential solver to clean up base before particle tiling refactor --- PyPIC3D/__init__.py | 1 - PyPIC3D/evolve.py | 100 --------- PyPIC3D/initialization.py | 38 ++-- PyPIC3D/solvers/vector_potential.py | 323 ---------------------------- docs/architecture.rst | 4 +- docs/solvers.rst | 8 +- tests/initialization_test.py | 26 ++- tests/pml_test.py | 4 - 8 files changed, 45 insertions(+), 459 deletions(-) delete mode 100644 PyPIC3D/solvers/vector_potential.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index 9e66c0ef..e325f2a7 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -13,7 +13,6 @@ 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 diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 23cc76e5..cd12e0ef 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -18,10 +18,6 @@ calculate_electrostatic_fields ) -from PyPIC3D.solvers.vector_potential import ( - E_from_A, B_from_A, update_vector_potential -) - from PyPIC3D.utils import add_external_fields @partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) @@ -160,100 +156,4 @@ def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_f # pack the fields into a tuple - 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"): - """ - 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 - - ################ 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 - - return particles, fields diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index b7faa3ce..930f66b5 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -19,7 +19,7 @@ update_parameters_from_toml, build_yee_grid, 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 ( @@ -45,14 +45,13 @@ 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 @@ -76,6 +75,16 @@ def _encode_field_bc(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 = ("fdtd", "spectral") + if solver not in supported_solvers: + raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd' or 'spectral'.") + + def default_parameters(): """ Returns a dictionary of default parameters for the simulation. @@ -110,7 +119,7 @@ def default_parameters(): simulation_parameters = { "name": "Default Simulation", "output_dir": os.getcwd(), - "solver": "fdtd", # solver: spectral, fdtd, vector_potential, curl_curl + "solver": "fdtd", # solver: spectral, fdtd "fast_backend": "flat", # flat | default (flat when compatible, else fallback) "particle_bc": "periodic", # particle boundary conditions: periodic, absorbing, reflecting # "bc": "periodic", # boundary conditions: periodic, dirichlet, neumann @@ -209,6 +218,7 @@ def initialize_simulation(toml_file): t_wind = simulation_parameters['t_wind'] electrostatic = simulation_parameters['electrostatic'] solver = simulation_parameters['solver'] + validate_field_solver(solver) relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] validate_particle_pusher(particle_pusher) @@ -280,10 +290,6 @@ def initialize_simulation(toml_file): plotting_parameters = convert_to_jax_compatible(plotting_parameters) # convert the world parameters to jax compatible format - # 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 @@ -328,8 +334,8 @@ def initialize_simulation(toml_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") + if electrostatic: + print("fast_backend='flat' not supported for electrostatic; 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") @@ -408,11 +414,6 @@ def initialize_simulation(toml_file): 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}") evolve_loop = time_loop_electrodynamic @@ -427,12 +428,7 @@ def initialize_simulation(toml_file): 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: + if electrostatic: fields = (E, B, J, rho, phi, external_fields) # define the fields tuple for the electrostatic solver else: 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/docs/architecture.rst b/docs/architecture.rst index 9a104f11..60948150 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -38,8 +38,8 @@ 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. diff --git a/docs/solvers.rst b/docs/solvers.rst index 9a521fe8..764dfc16 100644 --- a/docs/solvers.rst +++ b/docs/solvers.rst @@ -26,7 +26,7 @@ Per step, PyPIC3D: The ``solver`` key controls both the Poisson solve and gradient operator: - ``solver = "spectral"``: FFT Poisson + spectral gradient -- ``solver = "fdtd"`` (or any non-``"spectral"`` value): conjugate-gradient Poisson + centered finite-difference gradient +- ``solver = "fdtd"``: conjugate-gradient Poisson + centered finite-difference gradient Electrodynamic mode ^^^^^^^^^^^^^^^^^^^ @@ -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/tests/initialization_test.py b/tests/initialization_test.py index 4665549a..7fe60a62 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -1,7 +1,8 @@ import unittest +import tempfile import jax import jax.numpy as jnp -from PyPIC3D.initialization import setup_write_dir, default_parameters +from PyPIC3D.initialization import setup_write_dir, default_parameters, initialize_simulation jax.config.update("jax_enable_x64", True) @@ -26,5 +27,28 @@ def test_default_parameters(self): self.assertIn('plotfields', plotting) # check that the default parameters contain expected keys + 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, + "fast_backend": "default", + }, + "plotting": {"plotting": False}, + } + + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + initialize_simulation(config) + if __name__ == '__main__': unittest.main() diff --git a/tests/pml_test.py b/tests/pml_test.py index 3c85e122..b48a4f02 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -147,10 +147,6 @@ def test_initialize_simulation_rejects_pml_for_non_fdtd_solvers(self): 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)) From 37a398dd6997c562f5118a2c851c379d4dbb8f02 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 25 Jun 2026 18:27:38 -0500 Subject: [PATCH 02/83] starting to implement the tiled particles --- PyPIC3D/deposition/Esirkepov.py | 2 +- PyPIC3D/deposition/J_from_rhov.py | 4 +- PyPIC3D/deposition/direct_deposition_tiled.py | 201 ++++++++++++ PyPIC3D/deposition/rho.py | 2 +- PyPIC3D/diagnostics/fluid_quantities.py | 9 +- PyPIC3D/electrodynamic_tiled.py | 36 +++ PyPIC3D/evolve.py | 4 +- PyPIC3D/initialization.py | 4 + PyPIC3D/particles/flat_particles.py | 11 - PyPIC3D/particles/particle_initialization.py | 1 - PyPIC3D/particles/species_class.py | 8 - .../tiled_particle_initialization.py | 135 ++++++++ PyPIC3D/particles/tiled_particle_refresh.py | 228 ++++++++++++++ PyPIC3D/particles/tiled_particles.py | 23 ++ PyPIC3D/pusher/particle_push.py | 9 +- PyPIC3D/pusher/tiled_pusher.py | 152 +++++++++ PyPIC3D/solvers/yee_tiled.py | 278 +++++++++++++++++ docs/particles.rst | 4 +- tests/direct_deposition_tiled_test.py | 293 ++++++++++++++++++ tests/electrostatic_yee_test.py | 2 +- tests/evolve_test.py | 4 +- tests/particle_test.py | 16 +- tests/tiled_particle_initialization_test.py | 218 +++++++++++++ tests/tiled_particle_refresh_test.py | 145 +++++++++ tests/tiled_pusher_test.py | 235 ++++++++++++++ tests/yee_tiled_test.py | 233 ++++++++++++++ 26 files changed, 2215 insertions(+), 42 deletions(-) create mode 100644 PyPIC3D/deposition/direct_deposition_tiled.py create mode 100644 PyPIC3D/electrodynamic_tiled.py create mode 100644 PyPIC3D/particles/tiled_particle_initialization.py create mode 100644 PyPIC3D/particles/tiled_particle_refresh.py create mode 100644 PyPIC3D/particles/tiled_particles.py create mode 100644 PyPIC3D/pusher/tiled_pusher.py create mode 100644 PyPIC3D/solvers/yee_tiled.py create mode 100644 tests/direct_deposition_tiled_test.py create mode 100644 tests/tiled_particle_initialization_test.py create mode 100644 tests/tiled_particle_refresh_test.py create mode 100644 tests/tiled_pusher_test.py create mode 100644 tests/yee_tiled_test.py diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 3a812488..5a38faf2 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -177,6 +177,7 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] Jx = Jx.at[:, :, :].set(0) Jy = Jy.at[:, :, :].set(0) @@ -211,7 +212,6 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): 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 diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index b813f8f8..879f609d 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -30,6 +30,7 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] # determine the boundary conditions for each axis to handle ghost cells correctly Jx = Jx.at[:, :, :].set(0) @@ -38,10 +39,9 @@ def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): # 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 + # get the particle information needed for deposition x, y, z = species.get_forward_position() vx, vy, vz = species.get_velocity() diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py new file mode 100644 index 00000000..ed38409d --- /dev/null +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -0,0 +1,201 @@ +import jax +from jax import jit +import jax.numpy as jnp +from functools import partial + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + collapse_axis_stencil, + prepare_particle_axis_stencil, +) +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.solvers.yee_tiled import ( + fold_tiled_vector_ghost_cells_periodic, + update_tiled_vector_ghost_cells_periodic, +) +from PyPIC3D.utils import bilinear_filter + + +@partial(jit, static_argnames=("filter",)) +def direct_J_from_tiled_particles(tiled_particles, J_tiles, constants, world, grid=None, filter="bilinear"): + """ + Compute direct current deposition from tile-major particle storage. + + This follows the same time centering and Yee-component stencil as the + global ``J_from_rhov`` path, but current stays in compact tile-local + ghost-celled arrays throughout the deposition. + """ + + del constants + + if grid is None: + grid = world["grids"]["center"] + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + dt = world["dt"] + + Jx_tiles, Jy_tiles, Jz_tiles = J_tiles + ntx, nty, ntz = Jx_tiles.shape[:3] + tile_nx = Jx_tiles.shape[3] - 2 + tile_ny = Jx_tiles.shape[4] - 2 + tile_nz = Jx_tiles.shape[5] - 2 + Nx = ntx * tile_nx + 2 + Ny = nty * tile_ny + 2 + Nz = ntz * tile_nz + 2 + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] + + 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]) + + def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, ty, tz): + 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) + active = active_tile.reshape(-1).astype(x.dtype) + dq = (charge_tile * weight_tile).reshape(-1) / (dx * dy * dz) + + # Match J_from_rhov: the current is deposited from the half-step-back + # particle position, while the velocity is the time-centered velocity. + x = x - vx * dt / 2 + y = y - vy * dt / 2 + z = z - vz * dt / 2 + + x, x0, deltax_node, xpts = prepare_particle_axis_stencil( + x, + grid[0], + Nx, + shape_factor, + bc_x, + wind=world["x_wind"], + ghost_cells=True, + ) + y, y0, deltay_node, ypts = prepare_particle_axis_stencil( + y, + grid[1], + Ny, + shape_factor, + bc_y, + wind=world["y_wind"], + ghost_cells=True, + ) + z, z0, deltaz_node, zpts = prepare_particle_axis_stencil( + z, + grid[2], + Nz, + shape_factor, + bc_z, + wind=world["z_wind"], + ghost_cells=True, + ) + + 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 + + x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( + shape_factor == 1, + lambda _: get_first_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), + lambda _: get_second_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), + operand=None, + ) + 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, + ) + + xpts = jnp.asarray(xpts) + ypts = jnp.asarray(ypts) + zpts = jnp.asarray(zpts) + x_weights_node = jnp.asarray(x_weights_node) + y_weights_node = jnp.asarray(y_weights_node) + z_weights_node = jnp.asarray(z_weights_node) + x_weights_face = jnp.asarray(x_weights_face) + y_weights_face = jnp.asarray(y_weights_face) + z_weights_face = jnp.asarray(z_weights_face) + + 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) + + xpts = xpts - tx * tile_nx + ypts = ypts - ty * tile_ny + zpts = zpts - tz * tile_nz + # Convert global ghost-celled indices to tile-local ghost-celled indices. + + 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", + ) + + deposit_tiles = deposit_one_tile + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + + Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( + tiled_particles.x, + tiled_particles.u, + tiled_particles.charge, + tiled_particles.weight, + tiled_particles.active, + tx, + ty, + tz, + ) + + J_tiles = fold_tiled_vector_ghost_cells_periodic((Jx_tiles, Jy_tiles, Jz_tiles)) + J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + + def filter_tile(tile): + apply_filter = jax.vmap(jax.vmap(jax.vmap(bilinear_filter))) + return jax.lax.cond( + filter == "bilinear", + apply_filter, + lambda tile: tile, + operand=tile, + ) + + J_tiles = tuple(filter_tile(component) for component in J_tiles) + J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + + return J_tiles diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 566015ae..a59090c6 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -22,11 +22,11 @@ def compute_rho(particles, rho, world, constants): 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() diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index b41ccf09..5905bfda 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -34,14 +34,13 @@ def compute_mass_density(particles, rho, world): bc_x = world['boundary_conditions']['x'] bc_y = world['boundary_conditions']['y'] bc_z = world['boundary_conditions']['z'] + shape_factor = world['shape_factor'] # 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 @@ -117,14 +116,13 @@ def compute_velocity_field(particles, field, direction, world): bc_x = world['boundary_conditions']['x'] bc_y = world['boundary_conditions']['y'] bc_z = world['boundary_conditions']['z'] + shape_factor = world['shape_factor'] # 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() @@ -188,14 +186,13 @@ def compute_pressure_field(particles, field, velocity_field, direction, world): bc_x = world['boundary_conditions']['x'] bc_y = world['boundary_conditions']['y'] bc_z = world['boundary_conditions']['z'] + shape_factor = world['shape_factor'] # 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() diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py new file mode 100644 index 00000000..5e3f5dea --- /dev/null +++ b/PyPIC3D/electrodynamic_tiled.py @@ -0,0 +1,36 @@ +from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E + + +def time_loop_electrodynamic_tiled( + particles, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape, + relativistic=True, + particle_pusher="boris", +): + """ + Advance tiled electrodynamic fields by one Maxwell step without pushing particles. + """ + + del J_func, solver, relativistic, particle_pusher + + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields + # unpack the tiled field state. This first tiled-field pass updates + # Maxwell fields only; particles and particle boundary conditions are + # intentionally not touched. + + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + # update electric field from B and the supplied tiled current + + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + # update magnetic field from the newly updated electric field + + fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) + # pack the tiled field state + + return particles, fields diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index cd12e0ef..2a86aa09 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -54,7 +54,7 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], push_E, push_B, vertex_grid, center_grid, world['dt'], constants, relativistic=relativistic, particle_pusher=particle_pusher) + particles[i] = particle_push(particles[i], push_E, push_B, vertex_grid, center_grid, world, constants, relativistic=relativistic, particle_pusher=particle_pusher) # use the selected particle pusher for particle velocities particles[i].update_position() @@ -134,7 +134,7 @@ def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_f ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], push_E, push_B, center_grid, vertex_grid, world['dt'], constants, relativistic=relativistic, particle_pusher=particle_pusher) + particles[i] = particle_push(particles[i], push_E, push_B, center_grid, vertex_grid, world, constants, relativistic=relativistic, particle_pusher=particle_pusher) # use the selected particle pusher for particle velocities particles[i].update_position() diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 930f66b5..bcdc48d5 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -144,6 +144,9 @@ 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) + "particle_tile_nx": 1, # number of x cells per particle tile + "particle_tile_ny": 1, # number of y cells per particle tile + "particle_tile_nz": 1, # number of z cells per particle 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 } @@ -268,6 +271,7 @@ def initialize_simulation(toml_file): 'x_wind': x_wind, 'y_wind': y_wind, 'z_wind': z_wind, + 'shape_factor': simulation_parameters['shape_factor'], 'boundary_conditions': { 'x': _encode_field_bc(simulation_parameters['x_bc']), 'y': _encode_field_bc(simulation_parameters['y_bc']), diff --git a/PyPIC3D/particles/flat_particles.py b/PyPIC3D/particles/flat_particles.py index 40a694b5..1ce9671f 100644 --- a/PyPIC3D/particles/flat_particles.py +++ b/PyPIC3D/particles/flat_particles.py @@ -35,7 +35,6 @@ def __init__( update_vx, update_vy, update_vz, - shape, dt, species_meta, active_mask=None, @@ -69,7 +68,6 @@ def __init__( 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: @@ -134,9 +132,6 @@ def get_mass(self): 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) @@ -208,7 +203,6 @@ def tree_flatten(self): self.update_vx, self.update_vy, self.update_vz, - self.shape, self.dt, self.species_meta, ) @@ -241,7 +235,6 @@ def tree_unflatten(cls, aux_data, children): update_vx, update_vy, update_vz, - shape, dt, species_meta, ) = aux_data @@ -275,7 +268,6 @@ def tree_unflatten(cls, aux_data, children): update_vx=update_vx, update_vy=update_vy, update_vz=update_vz, - shape=shape, dt=dt, species_meta=species_meta, active_mask=active_mask, @@ -297,8 +289,6 @@ def _same(attr_list): 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": @@ -388,7 +378,6 @@ def to_flat_particles(particles): 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, diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index bf79546f..e6fc58fe 100644 --- a/PyPIC3D/particles/particle_initialization.py +++ b/PyPIC3D/particles/particle_initialization.py @@ -290,7 +290,6 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): update_z=update_z, update_pos=update_pos, update_v=update_v, - shape=simulation_parameters["shape_factor"], dt=dt, ) particles.append(particle) diff --git a/PyPIC3D/particles/species_class.py b/PyPIC3D/particles/species_class.py index c2d60aa7..b7b92cff 100644 --- a/PyPIC3D/particles/species_class.py +++ b/PyPIC3D/particles/species_class.py @@ -72,7 +72,6 @@ def __init__( update_vz=True, update_pos=True, update_v=True, - shape=1, dt=0, active_mask=None, ): @@ -114,7 +113,6 @@ def __init__( self.update_vz = update_vz self.update_pos = update_pos self.update_v = update_v - self.shape = shape self.dt = dt self.x1 = x1 @@ -186,9 +184,6 @@ def get_weight(self): 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), @@ -288,7 +283,6 @@ def tree_flatten(self): self.update_vx, self.update_vy, self.update_vz, - self.shape, self.dt, ) return children, aux_data @@ -321,7 +315,6 @@ def tree_unflatten(cls, aux_data, children): update_vx, update_vy, update_vz, - shape, dt, ) = aux_data @@ -355,7 +348,6 @@ def tree_unflatten(cls, aux_data, children): update_vz=update_vz, update_pos=update_pos, update_v=update_v, - shape=shape, dt=dt, active_mask=active_mask, ) diff --git a/PyPIC3D/particles/tiled_particle_initialization.py b/PyPIC3D/particles/tiled_particle_initialization.py new file mode 100644 index 00000000..1e7a99cf --- /dev/null +++ b/PyPIC3D/particles/tiled_particle_initialization.py @@ -0,0 +1,135 @@ +import math + +import jax.numpy as jnp +import numpy as np + +from PyPIC3D.particles.particle_initialization import load_particles_from_toml +from PyPIC3D.particles.tiled_particles import TiledParticles + + +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 to_tiled_particles(particles, world, simulation_parameters): + """ + Pack an existing particle species list into tile-major arrays. + + This routine only changes storage: particles are grouped by configurable + spatial tiles with leading axes ``(ntx, nty, ntz, species, slot, ...)``. + """ + + tile_nx = _as_int(simulation_parameters.get("particle_tile_nx", 1)) + tile_ny = _as_int(simulation_parameters.get("particle_tile_ny", 1)) + tile_nz = _as_int(simulation_parameters.get("particle_tile_nz", 1)) + + 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(particles) + + tile_counts = np.zeros((ntx, nty, ntz, n_species), dtype=int) + particle_tile_indices = [] + + for species_index, species in enumerate(particles): + x1, x2, x3 = species.get_forward_position() + tx, ty, tz = _particle_tile_indices(x1, x2, x3, world, tile_nx, tile_ny, tile_nz) + particle_tile_indices.append((tx, ty, tz)) + + for p in range(_as_int(species.N_particles)): + tile_counts[tx[p], ty[p], tz[p], species_index] += 1 + + max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 + + x = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) + u = jnp.zeros_like(x) + charge = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile)) + mass = jnp.zeros_like(charge) + weight = jnp.zeros_like(charge) + active = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) + update_x1 = jnp.zeros_like(active) + update_x2 = jnp.zeros_like(active) + update_x3 = jnp.zeros_like(active) + update_u1 = jnp.zeros_like(active) + update_u2 = jnp.zeros_like(active) + update_u3 = jnp.zeros_like(active) + + next_slot = np.zeros_like(tile_counts) + + for species_index, species in enumerate(particles): + x1, x2, x3 = species.get_forward_position() + u1, u2, u3 = species.get_velocity() + tx, ty, tz = particle_tile_indices[species_index] + + for p in range(_as_int(species.N_particles)): + tile_index = (tx[p], ty[p], tz[p], species_index) + slot = next_slot[tile_index] + next_slot[tile_index] += 1 + index = tile_index + (slot,) + + x = x.at[index + (0,)].set(x1[p]) + x = x.at[index + (1,)].set(x2[p]) + x = x.at[index + (2,)].set(x3[p]) + u = u.at[index + (0,)].set(u1[p]) + u = u.at[index + (1,)].set(u2[p]) + u = u.at[index + (2,)].set(u3[p]) + + charge = charge.at[index].set(species.charge) + mass = mass.at[index].set(species.mass) + weight = weight.at[index].set(species.weight) + active = active.at[index].set(species.active_mask[p]) + + update_x1 = update_x1.at[index].set(species.update_pos and species.update_x) + update_x2 = update_x2.at[index].set(species.update_pos and species.update_y) + update_x3 = update_x3.at[index].set(species.update_pos and species.update_z) + update_u1 = update_u1.at[index].set(species.update_v and species.update_vx) + update_u2 = update_u2.at[index].set(species.update_v and species.update_vy) + update_u3 = update_u3.at[index].set(species.update_v and species.update_vz) + + return TiledParticles( + x=x, + u=u, + charge=charge, + mass=mass, + weight=weight, + active=active, + update_x1=update_x1, + update_x2=update_x2, + update_x3=update_x3, + update_u1=update_u1, + update_u2=update_u2, + update_u3=update_u3, + ) + + +def load_tiled_particles_from_toml(config, simulation_parameters, world, constants): + """ + Initialize particles from TOML input and pack them into tile-major arrays. + + The particle values are produced by the existing particle initializer. This + wrapper only changes storage by calling ``to_tiled_particles`` on the + initialized species list. + """ + + particles = load_particles_from_toml(config, simulation_parameters, world, constants) + return to_tiled_particles(particles, world, simulation_parameters) diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py new file mode 100644 index 00000000..eb7c2c4f --- /dev/null +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -0,0 +1,228 @@ +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import wrap_periodic_position +from PyPIC3D.particles.tiled_particles import TiledParticles + + +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, world): + """ + Advance tile-major particle positions without changing tile ownership. + """ + + dt = world["dt"] + active = tiled_particles.active.astype(tiled_particles.x.dtype) + + 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 & tiled_particles.update_x1, x[..., 0] + dx, x[..., 0])) + x = x.at[..., 1].set(jnp.where(tiled_particles.active & tiled_particles.update_x2, x[..., 1] + dy, x[..., 1])) + x = x.at[..., 2].set(jnp.where(tiled_particles.active & tiled_particles.update_x3, x[..., 2] + dz, x[..., 2])) + + return tiled_particles._replace(x=x) + + +def _neighbor_offsets(count): + offsets = [] + seen = set() + for offset in (-1, 0, 1): + key = offset % count + if key not in seen: + offsets.append(offset) + seen.add(key) + return tuple(offsets) + + +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 + destination tile only reads candidate particles from its neighboring tiles. + Particles that do not fit in the destination tile capacity are dropped from + the returned active mask and reported through the overflow flag. + """ + + ntx, nty, ntz, n_species, n_slots = tiled_particles.active.shape + tile_counts = (ntx, nty, ntz) + + wrapped_x = tiled_particles.x + wrapped_x = wrapped_x.at[..., 0].set(wrap_periodic_position(wrapped_x[..., 0], world["x_wind"])) + wrapped_x = wrapped_x.at[..., 1].set(wrap_periodic_position(wrapped_x[..., 1], world["y_wind"])) + wrapped_x = wrapped_x.at[..., 2].set(wrap_periodic_position(wrapped_x[..., 2], world["z_wind"])) + + dest_tx, dest_ty, dest_tz = _particle_tile_indices( + wrapped_x[..., 0], + wrapped_x[..., 1], + wrapped_x[..., 2], + world, + tile_shape, + tile_counts, + ) + + new_x = jnp.zeros_like(tiled_particles.x) + new_u = jnp.zeros_like(tiled_particles.u) + new_charge = jnp.zeros_like(tiled_particles.charge) + new_mass = jnp.zeros_like(tiled_particles.mass) + new_weight = jnp.zeros_like(tiled_particles.weight) + new_active = jnp.zeros_like(tiled_particles.active) + new_update_x1 = jnp.zeros_like(tiled_particles.update_x1) + new_update_x2 = jnp.zeros_like(tiled_particles.update_x2) + new_update_x3 = jnp.zeros_like(tiled_particles.update_x3) + new_update_u1 = jnp.zeros_like(tiled_particles.update_u1) + new_update_u2 = jnp.zeros_like(tiled_particles.update_u2) + new_update_u3 = jnp.zeros_like(tiled_particles.update_u3) + overflow = jnp.asarray(False) + + x_offsets = _neighbor_offsets(ntx) + y_offsets = _neighbor_offsets(nty) + z_offsets = _neighbor_offsets(ntz) + + for tx in range(ntx): + for ty in range(nty): + for tz in range(ntz): + for species in range(n_species): + candidate_x = [] + candidate_u = [] + candidate_charge = [] + candidate_mass = [] + candidate_weight = [] + candidate_active = [] + candidate_update_x1 = [] + candidate_update_x2 = [] + candidate_update_x3 = [] + candidate_update_u1 = [] + candidate_update_u2 = [] + candidate_update_u3 = [] + candidate_dest_tx = [] + candidate_dest_ty = [] + candidate_dest_tz = [] + + for ox in x_offsets: + sx = (tx + ox) % ntx + for oy in y_offsets: + sy = (ty + oy) % nty + for oz in z_offsets: + sz = (tz + oz) % ntz + index = (sx, sy, sz, species) + candidate_x.append(wrapped_x[index]) + candidate_u.append(tiled_particles.u[index]) + candidate_charge.append(tiled_particles.charge[index]) + candidate_mass.append(tiled_particles.mass[index]) + candidate_weight.append(tiled_particles.weight[index]) + candidate_active.append(tiled_particles.active[index]) + candidate_update_x1.append(tiled_particles.update_x1[index]) + candidate_update_x2.append(tiled_particles.update_x2[index]) + candidate_update_x3.append(tiled_particles.update_x3[index]) + candidate_update_u1.append(tiled_particles.update_u1[index]) + candidate_update_u2.append(tiled_particles.update_u2[index]) + candidate_update_u3.append(tiled_particles.update_u3[index]) + candidate_dest_tx.append(dest_tx[index]) + candidate_dest_ty.append(dest_ty[index]) + candidate_dest_tz.append(dest_tz[index]) + + candidate_x = jnp.concatenate(candidate_x, axis=0) + candidate_u = jnp.concatenate(candidate_u, axis=0) + candidate_charge = jnp.concatenate(candidate_charge, axis=0) + candidate_mass = jnp.concatenate(candidate_mass, axis=0) + candidate_weight = jnp.concatenate(candidate_weight, axis=0) + candidate_active = jnp.concatenate(candidate_active, axis=0) + candidate_update_x1 = jnp.concatenate(candidate_update_x1, axis=0) + candidate_update_x2 = jnp.concatenate(candidate_update_x2, axis=0) + candidate_update_x3 = jnp.concatenate(candidate_update_x3, axis=0) + candidate_update_u1 = jnp.concatenate(candidate_update_u1, axis=0) + candidate_update_u2 = jnp.concatenate(candidate_update_u2, axis=0) + candidate_update_u3 = jnp.concatenate(candidate_update_u3, axis=0) + candidate_dest_tx = jnp.concatenate(candidate_dest_tx, axis=0) + candidate_dest_ty = jnp.concatenate(candidate_dest_ty, axis=0) + candidate_dest_tz = jnp.concatenate(candidate_dest_tz, axis=0) + + keep = ( + candidate_active + & (candidate_dest_tx == tx) + & (candidate_dest_ty == ty) + & (candidate_dest_tz == tz) + ) + rank = jnp.cumsum(keep.astype(int)) - 1 + fits = keep & (rank < n_slots) + safe_rank = jnp.where(fits, rank, 0) + + local_x = jnp.zeros_like(new_x[tx, ty, tz, species]) + local_u = jnp.zeros_like(new_u[tx, ty, tz, species]) + local_charge = jnp.zeros_like(new_charge[tx, ty, tz, species]) + local_mass = jnp.zeros_like(new_mass[tx, ty, tz, species]) + local_weight = jnp.zeros_like(new_weight[tx, ty, tz, species]) + local_active_count = jnp.zeros(n_slots, dtype=int) + local_update_x1_count = jnp.zeros(n_slots, dtype=int) + local_update_x2_count = jnp.zeros(n_slots, dtype=int) + local_update_x3_count = jnp.zeros(n_slots, dtype=int) + local_update_u1_count = jnp.zeros(n_slots, dtype=int) + local_update_u2_count = jnp.zeros(n_slots, dtype=int) + local_update_u3_count = jnp.zeros(n_slots, dtype=int) + + valid = fits.astype(candidate_x.dtype) + local_x = local_x.at[safe_rank].add(valid[:, None] * candidate_x) + local_u = local_u.at[safe_rank].add(valid[:, None] * candidate_u) + local_charge = local_charge.at[safe_rank].add(valid * candidate_charge) + local_mass = local_mass.at[safe_rank].add(valid * candidate_mass) + local_weight = local_weight.at[safe_rank].add(valid * candidate_weight) + local_active_count = local_active_count.at[safe_rank].add(fits.astype(int)) + local_update_x1_count = local_update_x1_count.at[safe_rank].add((fits & candidate_update_x1).astype(int)) + local_update_x2_count = local_update_x2_count.at[safe_rank].add((fits & candidate_update_x2).astype(int)) + local_update_x3_count = local_update_x3_count.at[safe_rank].add((fits & candidate_update_x3).astype(int)) + local_update_u1_count = local_update_u1_count.at[safe_rank].add((fits & candidate_update_u1).astype(int)) + local_update_u2_count = local_update_u2_count.at[safe_rank].add((fits & candidate_update_u2).astype(int)) + local_update_u3_count = local_update_u3_count.at[safe_rank].add((fits & candidate_update_u3).astype(int)) + + index = (tx, ty, tz, species) + new_x = new_x.at[index].set(local_x) + new_u = new_u.at[index].set(local_u) + new_charge = new_charge.at[index].set(local_charge) + new_mass = new_mass.at[index].set(local_mass) + new_weight = new_weight.at[index].set(local_weight) + new_active = new_active.at[index].set(local_active_count > 0) + new_update_x1 = new_update_x1.at[index].set(local_update_x1_count > 0) + new_update_x2 = new_update_x2.at[index].set(local_update_x2_count > 0) + new_update_x3 = new_update_x3.at[index].set(local_update_x3_count > 0) + new_update_u1 = new_update_u1.at[index].set(local_update_u1_count > 0) + new_update_u2 = new_update_u2.at[index].set(local_update_u2_count > 0) + new_update_u3 = new_update_u3.at[index].set(local_update_u3_count > 0) + overflow = overflow | (jnp.sum(keep) > n_slots) + + refreshed = TiledParticles( + x=new_x, + u=new_u, + charge=new_charge, + mass=new_mass, + weight=new_weight, + active=new_active, + update_x1=new_update_x1, + update_x2=new_update_x2, + update_x3=new_update_x3, + update_u1=new_update_u1, + update_u2=new_update_u2, + update_u3=new_update_u3, + ) + + return refreshed, overflow diff --git a/PyPIC3D/particles/tiled_particles.py b/PyPIC3D/particles/tiled_particles.py new file mode 100644 index 00000000..3e7d26f0 --- /dev/null +++ b/PyPIC3D/particles/tiled_particles.py @@ -0,0 +1,23 @@ +import jax +import jax.numpy as jnp +from typing import NamedTuple + +class TiledParticles(NamedTuple): + 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 + charge: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + mass: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + # charge and mass of the particles + weight: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + # macroparticle weight of the particles + active: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + # boolean array indicating whether the particle is active or not + update_x1: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + update_x2: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + update_x3: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + update_u1: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + update_u2: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + update_u3: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) + # boolean array indicating whether the particle's position or velocity should be updated or not \ No newline at end of file diff --git a/PyPIC3D/pusher/particle_push.py b/PyPIC3D/pusher/particle_push.py index fad6e554..7739facb 100644 --- a/PyPIC3D/pusher/particle_push.py +++ b/PyPIC3D/pusher/particle_push.py @@ -23,7 +23,7 @@ def validate_particle_pusher(particle_pusher): @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"): +def particle_push(particles, E, B, grid, staggered_grid, world, constants, periodic=True, relativistic=True, particle_pusher="boris"): """ Updates particle velocities using the selected particle pusher. @@ -33,7 +33,7 @@ def particle_push(particles, E, B, grid, staggered_grid, dt, constants, periodic 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. + world (dict): Simulation parameters, including ``dt`` and ``shape_factor``. constants (dict): Dictionary containing physical constants. periodic (bool): Kept for compatibility with the existing call signature. relativistic (bool): Selects relativistic or non-relativistic Boris. @@ -48,8 +48,9 @@ def particle_push(particles, E, B, grid, staggered_grid, dt, constants, periodic 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 + dt = world["dt"] + shape_factor = world["shape_factor"] + # interpolation order is a global simulation property ################## INTERPOLATION GRIDS ################################### Ex_grid = staggered_grid[0], grid[1], grid[2] diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py new file mode 100644 index 00000000..fc5b6e08 --- /dev/null +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -0,0 +1,152 @@ +import jax +from jax import jit +import jax.numpy as jnp +from functools import partial + +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.pusher.boris import ( + boris_single_particle, + interpolate_field_to_particles, + relativistic_boris_single_particle, +) + + +def _tile_axis(axis, tile_index, cells_per_tile): + start = tile_index * cells_per_tile + return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) + + +@partial(jit, static_argnames=("tile_shape", "relativistic")) +def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, relativistic=True): + """ + Push tile-major particles with Boris using compact tiled Yee fields. + + Particles are assumed to live in the tile that owns their current forward + position. The one-cell field halos on each compact tile provide the + neighboring Yee data needed by the interpolation stencil near tile faces. + """ + + tile_nx, tile_ny, tile_nz = tile_shape + dt = world["dt"] + shape_factor = world["shape_factor"] + + center_grid = world["grids"]["center"] + vertex_grid = world["grids"]["vertex"] + + Ex_tiles, Ey_tiles, Ez_tiles = E_tiles + Bx_tiles, By_tiles, Bz_tiles = B_tiles + + ntx, nty, ntz = tiled_particles.x.shape[:3] + + def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_tile, update_u1, update_u2, update_u3, + 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 = charge_tile.reshape(-1) + m = mass_tile.reshape(-1) + + center_x = _tile_axis(center_grid[0], tx, tile_nx) + center_y = _tile_axis(center_grid[1], ty, tile_ny) + center_z = _tile_axis(center_grid[2], tz, tile_nz) + vertex_x = _tile_axis(vertex_grid[0], tx, tile_nx) + vertex_y = _tile_axis(vertex_grid[1], ty, tile_ny) + vertex_z = _tile_axis(vertex_grid[2], tz, tile_nz) + + 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) + efield_aty = interpolate_field_to_particles(Ey_tile, x, y, z, Ey_grid, shape_factor, ghost_cells=True) + efield_atz = interpolate_field_to_particles(Ez_tile, x, y, z, Ez_grid, shape_factor, ghost_cells=True) + + bfield_atx = interpolate_field_to_particles(Bx_tile, x, y, z, Bx_grid, shape_factor, ghost_cells=True) + bfield_aty = interpolate_field_to_particles(By_tile, x, y, z, By_grid, shape_factor, ghost_cells=True) + bfield_atz = interpolate_field_to_particles(Bz_tile, x, y, z, Bz_grid, shape_factor, ghost_cells=True) + + boris_vmap = jax.vmap( + boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + relativistic_boris_vmap = jax.vmap( + relativistic_boris_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + + 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, + ) + + active = active_tile.reshape(-1) + update_u1 = update_u1.reshape(-1) + update_u2 = update_u2.reshape(-1) + update_u3 = update_u3.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, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + push_tiles = jax.vmap(push_tiles, in_axes=(None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + push_tiles = jax.vmap(push_tiles, in_axes=(0, None, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + + new_u = push_tiles( + tx, ty, tz, + tiled_particles.x, + tiled_particles.u, + tiled_particles.charge, + tiled_particles.mass, + tiled_particles.active, + tiled_particles.update_u1, + tiled_particles.update_u2, + tiled_particles.update_u3, + Ex_tiles, + Ey_tiles, + Ez_tiles, + Bx_tiles, + By_tiles, + Bz_tiles, + ) + + return TiledParticles( + x=tiled_particles.x, + u=new_u, + charge=tiled_particles.charge, + mass=tiled_particles.mass, + weight=tiled_particles.weight, + active=tiled_particles.active, + update_x1=tiled_particles.update_x1, + update_x2=tiled_particles.update_x2, + update_x3=tiled_particles.update_x3, + update_u1=tiled_particles.update_u1, + update_u2=tiled_particles.update_u2, + update_u3=tiled_particles.update_u3, + ) diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py new file mode 100644 index 00000000..db9a9ab1 --- /dev/null +++ b/PyPIC3D/solvers/yee_tiled.py @@ -0,0 +1,278 @@ +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells + + +def _tile_axis_count(n_cells, cells_per_tile): + if int(n_cells) % int(cells_per_tile) != 0: + raise ValueError("Field tile sizes must divide the physical grid dimensions exactly.") + return int(n_cells) // int(cells_per_tile) + + +def _tile_scalar_field(field, tile_shape): + """ + Split a ghost-celled field into compact field tiles with one-cell halos. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + 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) + + 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): + """ + Split ``(Fx, Fy, Fz)`` into compact ghost-celled tiles. + + Each scalar component gets leading tile axes ``(ntx, nty, ntz)`` followed by + the tile-local ghost-celled field shape + ``(tile_nx + 2, tile_ny + 2, tile_nz + 2)``. + """ + + del world + return tuple(_tile_scalar_field(component, tile_shape) for component in field) + + +def _assemble_scalar_tiles(field_tiles, world, tile_shape): + """ + Assemble compact field tiles back into one global ghost-celled field. + """ + + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + ntx, nty, ntz = field_tiles.shape[:3] + Nx = int(world["Nx"]) + Ny = int(world["Ny"]) + Nz = int(world["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): + interior = field_tiles[tx, ty, tz, 1:-1, 1:-1, 1:-1] + ix = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + field = field.at[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz].set(interior) + + 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 assemble_tiled_vector_field(field_tiles, world, tile_shape): + """ + Assemble tiled vector-field components into ordinary ghost-celled arrays. + """ + + return tuple(_assemble_scalar_tiles(component, world, tile_shape) for component in field_tiles) + + +def update_tiled_ghost_cells_periodic(field_tiles): + """ + Refresh one-cell tile halos by periodic exchange with neighbor tiles. + + ``field_tiles`` has shape + ``(ntx, nty, ntz, tile_nx + 2, tile_ny + 2, tile_nz + 2)``. The last + three axes are the tile-local ghost-celled field. The first three axes + identify the tile in the global tiling. Periodic global boundaries are + handled by wrapping along the corresponding tile axis. + """ + + # x-halos: previous tile right interior -> left ghost, next tile left + # interior -> right ghost. + field_tiles = field_tiles.at[:, :, :, 0, :, :].set( + jnp.roll(field_tiles[:, :, :, -2, :, :], shift=1, axis=0) + ) + field_tiles = field_tiles.at[:, :, :, -1, :, :].set( + jnp.roll(field_tiles[:, :, :, 1, :, :], shift=-1, axis=0) + ) + + # y-halos use the x-refreshed field so tile-edge/corner guard cells match + # the sequential ghost-cell convention used by the global arrays. + field_tiles = field_tiles.at[:, :, :, :, 0, :].set( + jnp.roll(field_tiles[:, :, :, :, -2, :], shift=1, axis=1) + ) + field_tiles = field_tiles.at[:, :, :, :, -1, :].set( + jnp.roll(field_tiles[:, :, :, :, 1, :], shift=-1, axis=1) + ) + + # z-halos complete the edge and corner values after x/y have been refreshed. + field_tiles = field_tiles.at[:, :, :, :, :, 0].set( + jnp.roll(field_tiles[:, :, :, :, :, -2], shift=1, axis=2) + ) + field_tiles = field_tiles.at[:, :, :, :, :, -1].set( + jnp.roll(field_tiles[:, :, :, :, :, 1], shift=-1, axis=2) + ) + + return field_tiles + + +def update_tiled_vector_ghost_cells_periodic(field_tiles): + """ + Refresh periodic tile halos for each component of a vector field. + """ + + return tuple(update_tiled_ghost_cells_periodic(component) for component in field_tiles) + + +def fold_tiled_ghost_cells_periodic(field_tiles): + """ + Add tile-ghost deposits into neighboring tile interiors, then clear ghosts. + + This is the current-deposition analogue of ``fold_ghost_cells`` for a + periodic tiled layout. Unlike field halo refresh, this operation is + additive: deposits in ghost cells belong to the adjacent tile's physical + interior. + """ + + field_tiles = field_tiles.at[:, :, :, -2, :, :].add( + jnp.roll(field_tiles[:, :, :, 0, :, :], shift=-1, axis=0) + ) + field_tiles = field_tiles.at[:, :, :, 1, :, :].add( + jnp.roll(field_tiles[:, :, :, -1, :, :], shift=1, axis=0) + ) + field_tiles = field_tiles.at[:, :, :, 0, :, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, -1, :, :].set(0.0) + + field_tiles = field_tiles.at[:, :, :, :, -2, :].add( + jnp.roll(field_tiles[:, :, :, :, 0, :], shift=-1, axis=1) + ) + field_tiles = field_tiles.at[:, :, :, :, 1, :].add( + jnp.roll(field_tiles[:, :, :, :, -1, :], shift=1, axis=1) + ) + field_tiles = field_tiles.at[:, :, :, :, 0, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, -1, :].set(0.0) + + field_tiles = field_tiles.at[:, :, :, :, :, -2].add( + jnp.roll(field_tiles[:, :, :, :, :, 0], shift=-1, axis=2) + ) + field_tiles = field_tiles.at[:, :, :, :, :, 1].add( + jnp.roll(field_tiles[:, :, :, :, :, -1], shift=1, axis=2) + ) + field_tiles = field_tiles.at[:, :, :, :, :, 0].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :, -1].set(0.0) + + return field_tiles + + +def fold_tiled_vector_ghost_cells_periodic(field_tiles): + """ + Fold periodic tile-ghost deposits for each vector component. + """ + + return tuple(fold_tiled_ghost_cells_periodic(component) for component in field_tiles) + + +def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape): + """ + Update compact tiled electric fields without assembling a global field. + + This tile-native path is periodic-only and intentionally skips the digital + filter used by the global Yee update. The Yee curl is evaluated on each + tile's physical interior after B halos have been refreshed from neighbor + tiles. + """ + + del curl_func, tile_shape + + Ex, Ey, Ez = E_tiles + Bx, By, Bz = update_tiled_vector_ghost_cells_periodic(B_tiles) + Jx, Jy, Jz = J_tiles + + 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[:, :, :, 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 + + curl_x = dBz_dy - dBy_dz + curl_y = dBx_dz - dBz_dx + curl_z = dBy_dx - dBx_dy + + 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 + ) + 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 + ) + 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 + ) + + return update_tiled_vector_ghost_cells_periodic((Ex, Ey, Ez)) + + +def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): + """ + Update compact tiled magnetic fields without assembling a global field. + + This tile-native path is periodic-only and intentionally skips the digital + filter used by the global Yee update. The Yee curl is evaluated on each + tile's physical interior after E halos have been refreshed from neighbor + tiles. + """ + + del constants, curl_func, tile_shape + + Ex, Ey, Ez = update_tiled_vector_ghost_cells_periodic(E_tiles) + Bx, By, Bz = B_tiles + + 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[:, :, :, 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 + + curl_x = dEz_dy - dEy_dz + curl_y = dEx_dz - dEz_dx + curl_z = dEy_dx - dEx_dy + + 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) + + return update_tiled_vector_ghost_cells_periodic((Bx, By, Bz)) diff --git a/docs/particles.rst b/docs/particles.rst index 5ad73f43..da9e5683 100644 --- a/docs/particles.rst +++ b/docs/particles.rst @@ -19,13 +19,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 ---------------------- diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py new file mode 100644 index 00000000..d1b621a7 --- /dev/null +++ b/tests/direct_deposition_tiled_test.py @@ -0,0 +1,293 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, fold_tiled_ghost_cells_periodic +from PyPIC3D.utils import build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +class TestDirectDepositionTiled(unittest.TestCase): + def _build_world(self, Nx=8, Ny=6, Nz=4, dt=0.05): + x_wind, y_wind, z_wind = 4.0, 3.0, 2.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": 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 _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, simulation_parameters): + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + tile_nx, tile_ny, tile_nz = tile_shape + shape = ( + world["Nx"] // tile_nx, + world["Ny"] // tile_ny, + world["Nz"] // tile_nz, + tile_nx + 2, + tile_ny + 2, + tile_nz + 2, + ) + return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) + + def _compare_tiled_to_standard(self, particles, world, simulation_parameters): + constants = {"C": 3.0e8, "alpha": 1.0} + tiled_particles = to_tiled_particles(particles, world, simulation_parameters) + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + + J_reference = J_from_rhov(particles, self._empty_J(world), constants, world, filter="none") + J_tiles = direct_J_from_tiled_particles( + tiled_particles, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="none", + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + 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)) + + def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self): + 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 = fold_tiled_ghost_cells_periodic(tiles) + + 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)) + + 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} + species = particle_species( + name="local current tiles", + N_particles=4, + charge=1.0, + mass=1.0, + weight=1.0, + 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"], + ) + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + + J_tiles = direct_J_from_tiled_particles( + tiled_particles, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="none", + ) + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="none") + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + 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)) + + 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, + } + + electrons = particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=1.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"], + ) + ions = particle_species( + name="ions", + N_particles=3, + charge=2.0, + mass=4.0, + weight=0.25, + T=1.0, + x1=jnp.array([-1.65, 0.15, 1.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"], + ) + + self._compare_tiled_to_standard([electrons, ions], world, simulation_parameters) + + 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, + } + species = particle_species( + name="partly active", + N_particles=4, + charge=1.0, + mass=1.0, + weight=1.0, + 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=jnp.array([True, False, True, False]), + ) + + self._compare_tiled_to_standard([species], world, simulation_parameters) + + 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, + } + species = particle_species( + name="crossing", + N_particles=2, + charge=1.0, + mass=1.0, + weight=1.0, + T=1.0, + x1=jnp.array([world["x_wind"] / 2 + 0.1 * world["dx"], -world["x_wind"] / 2 - 0.2 * world["dx"]]), + x2=jnp.array([0.0, 0.0]), + x3=jnp.array([0.0, 0.0]), + v1=jnp.array([0.5, -0.25]), + v2=jnp.array([0.1, -0.2]), + v3=jnp.array([0.0, 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"], + ) + + self._compare_tiled_to_standard([species], 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, + } + species = particle_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.array([0.0, 0.0, 0.0]), + x3=jnp.array([0.0, 0.0, 0.0]), + v1=jnp.array([0.2, -0.1, 0.05]), + v2=jnp.array([0.3, 0.15, -0.2]), + v3=jnp.array([-0.05, 0.25, 0.1]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + self._compare_tiled_to_standard([species], world, simulation_parameters) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index 7c38a36d..d688320c 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -56,6 +56,7 @@ def setUp(self): "x_wind": self.x_wind, "y_wind": self.y_wind, "z_wind": self.z_wind, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(self.world) @@ -91,7 +92,6 @@ def setUp(self): x_bc="periodic", y_bc="periodic", z_bc="periodic", - shape=1, dt=0.0, ) ] diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 74ba461a..a59f132d 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -32,6 +32,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field "x_wind": 1.0, "y_wind": 1.0, "z_wind": 1.0, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } center_grid, vertex_grid = build_yee_grid(world) @@ -63,7 +64,6 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field dx=world["dx"], dy=world["dy"], dz=world["dz"], - shape=1, dt=world["dt"], ) ] @@ -104,6 +104,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): "x_wind": 1.0, "y_wind": 1.0, "z_wind": 1.0, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } center_grid, vertex_grid = build_yee_grid(world) @@ -136,7 +137,6 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): dy=world["dy"], dz=world["dz"], x_bc="absorbing", - shape=1, dt=world["dt"], ) ] diff --git a/tests/particle_test.py b/tests/particle_test.py index 7571e0a2..1a5dbab8 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -234,9 +234,10 @@ def test_particle_initialization_accepts_absorbing_axis_boundary(self): "x_wind": self.x_wind, "y_wind": self.y_wind, "z_wind": self.z_wind, + "shape_factor": 2, } constants = {"kb": self.kb, "eps": 1.0} - simulation_parameters = {"ds_per_debye": None, "shape_factor": 1} + simulation_parameters = {"ds_per_debye": None, "shape_factor": 2} config = { "particle1": { "name": "absorbing", @@ -258,6 +259,9 @@ def test_particle_initialization_accepts_absorbing_axis_boundary(self): self.assertEqual(particles[0].x_bc, "absorbing") self.assertTrue(jnp.array_equal(particles[0].get_active_mask(), jnp.array([True]))) + self.assertEqual(world["shape_factor"], 2) + self.assertFalse(hasattr(particles[0], "shape")) + self.assertFalse(hasattr(particles[0], "get_shape")) def test_inactive_particles_do_not_advance_or_take_new_velocity(self): species = particle_species( @@ -430,6 +434,7 @@ def test_J_from_rhov(self): 'Nx': Nx, 'Ny': Ny, 'Nz': Nz, 'x_wind': self.x_wind, 'y_wind': self.y_wind, 'z_wind': self.z_wind, 'dt': 0.0001, + 'shape_factor': 1, 'boundary_conditions': {'x': 0, 'y': 0, 'z': 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -483,6 +488,7 @@ def test_inactive_particles_do_not_deposit_rho_or_current(self): "y_wind": self.y_wind, "z_wind": self.z_wind, "dt": 0.0, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -547,6 +553,7 @@ def test_rho(self): '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, + 'shape_factor': 1, 'boundary_conditions': {'x': 0, 'y': 0, 'z': 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -607,6 +614,7 @@ def test_rho_periodic_boundary_exact_face(self): "x_wind": self.x_wind, "y_wind": self.y_wind, "z_wind": self.z_wind, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -657,6 +665,7 @@ def test_rho_periodic_boundary_crossing(self): "x_wind": self.x_wind, "y_wind": self.y_wind, "z_wind": self.z_wind, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -708,6 +717,7 @@ def test_J_from_rhov_periodic_boundary_crossing(self): "y_wind": self.y_wind, "z_wind": self.z_wind, "dt": 0.0, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -765,6 +775,7 @@ def test_rho_reduced_dimensions_periodic(self): "x_wind": x_wind, "y_wind": y_wind, "z_wind": z_wind, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -812,6 +823,7 @@ def test_filtered_rho_periodic_updates_ghost_cells(self): "x_wind": self.x_wind, "y_wind": self.y_wind, "z_wind": self.z_wind, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -869,6 +881,7 @@ def test_check_continuity_1D(self): "Nx": Nx, "Ny": Ny, "Nz": Nz, "x_wind": x_wind, "y_wind": y_wind, "z_wind": z_wind, "dt": dt, + "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } vertex_grid, center_grid = build_yee_grid(world) @@ -902,7 +915,6 @@ def test_check_continuity_1D(self): 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) diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py new file mode 100644 index 00000000..35f56644 --- /dev/null +++ b/tests/tiled_particle_initialization_test.py @@ -0,0 +1,218 @@ +import os +import tempfile +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from PyPIC3D.particles.tiled_particle_initialization import load_tiled_particles_from_toml +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.species_class import particle_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_to_tiled_particles_packs_existing_particle_species_list(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 = particle_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 = to_tiled_particles([species], world, simulation_parameters) + + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(particles.x.shape, (2, 2, 1, 1, 2, 3)) + self.assertEqual(int(jnp.sum(particles.active)), 2) + + 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(particles.charge[particles.x[..., 0] != 0.0], 2.0)) + self.assertTrue(jnp.allclose(particles.mass[particles.x[..., 0] != 0.0], 3.0)) + self.assertTrue(jnp.allclose(particles.weight[particles.x[..., 0] != 0.0], 4.0)) + + self.assertTrue(bool(particles.update_x1[1, 1, 0, 0, 1])) + self.assertFalse(bool(particles.update_x2[1, 1, 0, 0, 1])) + self.assertTrue(bool(particles.update_x3[1, 1, 0, 0, 1])) + self.assertFalse(bool(particles.update_u1[1, 1, 0, 0, 1])) + self.assertTrue(bool(particles.update_u2[1, 1, 0, 0, 1])) + self.assertTrue(bool(particles.update_u3[1, 1, 0, 0, 1])) + + def test_load_tiled_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, + } + 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 = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(particles.x.shape, (2, 2, 1, 1, 2, 3)) + self.assertEqual(particles.u.shape, (2, 2, 1, 1, 2, 3)) + self.assertEqual(particles.charge.shape, (2, 2, 1, 1, 2)) + + 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(particles.charge[particles.active], -1.0)) + self.assertTrue(jnp.allclose(particles.mass[particles.active], 2.0)) + self.assertTrue(jnp.allclose(particles.weight[particles.active], 4.0)) + + def test_load_tiled_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, + } + 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 = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + + self.assertTrue(bool(particles.update_x1[0, 0, 0, 0, 0])) + self.assertFalse(bool(particles.update_x2[0, 0, 0, 0, 0])) + self.assertTrue(bool(particles.update_x3[0, 0, 0, 0, 0])) + self.assertFalse(bool(particles.update_u1[0, 0, 0, 0, 0])) + self.assertTrue(bool(particles.update_u2[0, 0, 0, 0, 0])) + self.assertFalse(bool(particles.update_u3[0, 0, 0, 0, 0])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py new file mode 100644 index 00000000..4221bdd0 --- /dev/null +++ b/tests/tiled_particle_refresh_test.py @@ -0,0 +1,145 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particle_refresh import ( + 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}, + } + + def _species(self, world, x1, v1, active_mask=None, update_x=True): + if active_mask is None: + active_mask = jnp.ones_like(jnp.asarray(x1), dtype=bool) + n_particles = len(x1) + return particle_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, + ) + + 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] + charge = tiled_particles.charge.reshape(-1)[active] + mass = tiled_particles.mass.reshape(-1)[active] + weight = tiled_particles.weight.reshape(-1)[active] + order = jnp.argsort(x[:, 0]) + return x[order], u[order], charge[order], mass[order], weight[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 = to_tiled_particles([species], world, self._simulation_parameters()) + + moved = update_tiled_particle_positions(tiled_particles, world) + + 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 = to_tiled_particles([fixed_species], world, self._simulation_parameters()) + fixed_moved = update_tiled_particle_positions(fixed, world) + self.assertTrue(jnp.allclose(fixed_moved.x, fixed.x)) + + 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 = to_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, charge, mass, weight = 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]))) + self.assertTrue(jnp.allclose(charge, 2.0)) + self.assertTrue(jnp.allclose(mass, 3.0)) + self.assertTrue(jnp.allclose(weight, 4.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 = to_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_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 = to_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) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py new file mode 100644 index 00000000..84810eef --- /dev/null +++ b/tests/tiled_pusher_test.py @@ -0,0 +1,235 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.pusher.tiled_pusher import tiled_particle_push +from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.utils import build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +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 _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 particle_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_tiled_particle_push_matches_flat_boris(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + 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) + reference = self._species(world) + reference = particle_push( + reference, + E, + B, + world["grids"]["center"], + world["grids"]["vertex"], + world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + pushed = tiled_particle_push( + tiled_particles, + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + tile_shape, + relativistic=False, + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + flat_u = jnp.stack(reference.get_velocity(), axis=1) + flat_x = jnp.stack(reference.get_forward_position(), axis=1) + order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + + self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_particle_push_respects_active_and_update_flags(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + 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 = to_tiled_particles([species], world, simulation_parameters) + + pushed = tiled_particle_push( + tiled_particles, + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + tile_shape, + 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_tiled_particle_push_matches_relativistic_flat_boris_on_reduced_axes(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + constants = {"C": 10.0} + tile_shape = (2, 1, 1) + 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=0.1) + B = self._deterministic_vector_field(world, scale=0.02) + def make_species(): + return particle_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 = particle_push( + reference, + E, + B, + world["grids"]["center"], + world["grids"]["vertex"], + world, + constants, + relativistic=True, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + pushed = tiled_particle_push( + tiled_particles, + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + tile_shape, + relativistic=True, + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + flat_u = jnp.stack(reference.get_velocity(), axis=1) + flat_x = jnp.stack(reference.get_forward_position(), axis=1) + order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + + self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py new file mode 100644 index 00000000..c983c55a --- /dev/null +++ b/tests/yee_tiled_test.py @@ -0,0 +1,233 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.initialization import initialize_fields +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_vector_field, + tile_vector_field, + update_tiled_ghost_cells_periodic, + update_tiled_vector_ghost_cells_periodic, + update_tiled_B, + update_tiled_E, +) +from PyPIC3D.utils import build_yee_grid +from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells + + +jax.config.update("jax_enable_x64", True) + + +def unused_curl(Ex, Ey, Ez): + return None + + +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 _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 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_update_tiled_ghost_cells_periodic_refreshes_neighbor_halos(self): + world = self._build_world() + tile_shape = (2, 3, 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 = update_tiled_ghost_cells_periodic(stale_tiles) + + 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) + 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 = update_tiled_vector_ghost_cells_periodic(stale_tiles) + + 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_E_matches_standard_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) + 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, _ = update_E(E, B, J, world, constants, unused_curl) + E_tiled = update_tiled_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + unused_curl, + tile_shape, + ) + 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_tiled_B_matches_standard_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) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + + B_reference, _ = update_B(E, B, world, constants, unused_curl) + B_tiled = update_tiled_B( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + unused_curl, + tile_shape, + ) + 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_standard_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) + 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, _ = update_E(E, B, J, world, constants, unused_curl) + B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + + 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape) + + 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) + 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, _ = update_E(E, B, J, world, constants, unused_curl) + B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + particles = object() + external_fields = ( + tuple(jnp.zeros_like(component) for component in E), + tuple(jnp.zeros_like(component) for component in B), + ) + fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + jnp.zeros_like(E[0]), + jnp.zeros_like(E[0]), + external_fields, + None, + ) + + particles_after, fields_after = time_loop_electrodynamic_tiled( + particles, + fields, + world, + constants, + unused_curl, + J_func=None, + solver="fdtd", + tile_shape=tile_shape, + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, rho, phi, external_after, pml_state = fields_after + 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.assertIs(particles_after, particles) + self.assertIsNone(pml_state) + self.assertIs(external_after, external_fields) + 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(fields[2], J_tiles): + self.assertTrue(jnp.allclose(after, original, rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() From 3f98629565ad34d2887d0de4683d37856fa10b70 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 25 Jun 2026 19:16:01 -0500 Subject: [PATCH 03/83] migrating to tiled particles and fields --- PyPIC3D/__main__.py | 55 ++- PyPIC3D/deposition/J_from_rhov.py | 12 +- PyPIC3D/deposition/direct_deposition_tiled.py | 41 +- PyPIC3D/electrodynamic_tiled.py | 74 +++- PyPIC3D/initialization.py | 83 +++- PyPIC3D/solvers/yee_tiled.py | 127 +++++- PyPIC3D/utils.py | 106 +++-- tests/direct_deposition_tiled_test.py | 92 +++- tests/tiled_yee_integration_test.py | 407 ++++++++++++++++++ tests/yee_tiled_test.py | 149 +++++++ 10 files changed, 1058 insertions(+), 88 deletions(-) create mode 100644 tests/tiled_yee_integration_test.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index f313f02a..8200c1e6 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -29,7 +29,7 @@ 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 ( @@ -42,11 +42,41 @@ from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field + # Importing functions from the PyPIC3D package ############################################################################################################ +def _fields_for_output(fields, world, simulation_parameters): + if simulation_parameters["solver"] != "tiled_yee": + return fields + + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, *rest = fields + tile_shape = tuple(int(width) for width in world["tile_shape"]) + external_E, external_B = external_fields + + E = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B = assemble_tiled_vector_field(B_tiles, world, tile_shape) + J = assemble_tiled_vector_field(J_tiles, world, tile_shape) + external_fields = ( + assemble_tiled_vector_field(external_E, world, tile_shape), + assemble_tiled_vector_field(external_B, world, tile_shape), + ) + + return (E, B, J, rho, phi, external_fields, pml_state) + + +def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): + if simulation_parameters["solver"] != "tiled_yee" or len(fields) < 8: + return + + overflow = fields[-1] + if bool(jax.device_get(overflow)): + raise RuntimeError("tiled_yee particle tile capacity overflowed during periodic retile") + + def run_PyPIC3D(config_file): ##################################### INITIALIZE SIMULATION ################################################ @@ -62,6 +92,12 @@ def run_PyPIC3D(config_file): scalar_field_names = ["rho", "mass_density"] vector_field_names = ["E", "B", "J"] + tiled_run = simulation_parameters["solver"] == "tiled_yee" + + if tiled_run and plotting_parameters["plot_openpmd_particles"]: + raise ValueError("plot_openpmd_particles is not supported for tiled_yee particle storage") + if tiled_run and (plotting_parameters["plot_phasespace"] or plotting_parameters["plot_vtk_particles"]): + raise ValueError("particle phase-space and VTK particle output are not supported for tiled_yee particle storage") E, B, J, rho, phi, external_fields, *rest = fields # unpack the fields @@ -100,7 +136,7 @@ def run_PyPIC3D(config_file): 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 = compute_total_momentum(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 @@ -115,6 +151,8 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_scalars']: + if tiled_run: + raise ValueError("plot_vtk_scalars is not supported for tiled_yee particle storage") rho = compute_rho(particles, rho, world, constants) # calculate the charge density based on the particle positions mass_density = compute_mass_density(particles, rho, world) @@ -128,6 +166,9 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_vectors']: + output_fields = _fields_for_output(fields, world, simulation_parameters) + E, B, J, rho, phi, external_fields, *rest = output_fields + # assemble tiled fields before VTK output 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]], @@ -145,11 +186,12 @@ def run_PyPIC3D(config_file): # 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_openpmd_fields(_fields_for_output(fields, world, simulation_parameters), 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 + if not tiled_run: + fields = (E, B, J, rho, phi, external_fields, *rest) + # repack the fields for non-tiled diagnostics that updated rho particles, fields = jit_loop( particles, @@ -163,6 +205,9 @@ def run_PyPIC3D(config_file): 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. return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index 879f609d..fa80636f 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -10,7 +10,7 @@ ) 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.utils import bilinear_filter, digital_filter @partial(jit, static_argnames=("filter",)) @@ -198,12 +198,10 @@ def idx_and_dJ_values(idx): # 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_, - ) + if filter == "bilinear": + J_ = bilinear_filter(J_) + elif filter == "digital": + J_ = digital_filter(J_, constants["alpha"]) return J_ Jx = filter_func(Jx, filter) diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index ed38409d..946e829c 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -12,7 +12,29 @@ fold_tiled_vector_ghost_cells_periodic, update_tiled_vector_ghost_cells_periodic, ) -from PyPIC3D.utils import bilinear_filter +from PyPIC3D.utils import bilinear_filter, digital_filter + + +def digital_filter_tiled_current_density(J_tiles, alpha): + """ + Apply the global digital current filter to compact periodic current tiles. + + Tile halos are refreshed before the stencil so every tile sees the same + neighbor values that the assembled ghost-celled current would see. The + halos are refreshed again after filtering so downstream Yee updates read + filtered guard-cell values. + """ + + J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + + def filter_component(component_tiles): + filter_tiles = jax.vmap(jax.vmap(jax.vmap(lambda tile: digital_filter(tile, alpha)))) + return filter_tiles(component_tiles) + + J_tiles = tuple(filter_component(component) for component in J_tiles) + J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + + return J_tiles @partial(jit, static_argnames=("filter",)) @@ -25,8 +47,6 @@ def direct_J_from_tiled_particles(tiled_particles, J_tiles, constants, world, gr ghost-celled arrays throughout the deposition. """ - del constants - if grid is None: grid = world["grids"]["center"] @@ -186,16 +206,11 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, J_tiles = fold_tiled_vector_ghost_cells_periodic((Jx_tiles, Jy_tiles, Jz_tiles)) J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) - def filter_tile(tile): + if filter == "bilinear": apply_filter = jax.vmap(jax.vmap(jax.vmap(bilinear_filter))) - return jax.lax.cond( - filter == "bilinear", - apply_filter, - lambda tile: tile, - operand=tile, - ) - - J_tiles = tuple(filter_tile(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + J_tiles = tuple(apply_filter(component) for component in J_tiles) + J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + elif filter == "digital": + J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"]) return J_tiles diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 5e3f5dea..5bc0c36d 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -1,4 +1,13 @@ +import jax.numpy as jnp + +from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.particles.tiled_particle_refresh import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) +from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E +from PyPIC3D.utils import add_external_fields def time_loop_electrodynamic_tiled( @@ -9,20 +18,69 @@ def time_loop_electrodynamic_tiled( curl_func, J_func, solver, - tile_shape, + tile_shape=None, relativistic=True, particle_pusher="boris", ): """ - Advance tiled electrodynamic fields by one Maxwell step without pushing particles. + Advance a periodic tiled electrodynamic PIC system by one time step. """ - del J_func, solver, relativistic, particle_pusher + del solver, particle_pusher + + if len(fields) == 7: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields + overflow_previous = jnp.asarray(False) + else: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow_previous = fields + # unpack the tile-major field state. ``overflow_previous`` carries the + # particle-retile overflow diagnostic out to the Python driver. + + if tile_shape is None: + tile_shape = tuple(int(width) - 2 for width in E_tiles[0].shape[3:]) + + if not hasattr(particles, "active"): + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) + return particles, fields + # keep the original field-only helper behavior for tests and standalone + # Maxwell equivalence checks that pass a non-particle sentinel. + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) + # particles see evolved fields plus external-only fields, as in the + # standard electrodynamic path. + + particles = tiled_particle_push( + particles, + push_E_tiles, + push_B_tiles, + world, + constants, + tile_shape, + relativistic=relativistic, + ) + # use the tiled Boris pusher for particle velocities + + particles = update_tiled_particle_positions(particles, world) + # update particle forward positions before current deposition + + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + # wrap periodic particles and move them into their owning tiles. Overflow + # means the fixed tile capacity was exceeded and the Python driver must + # reject the step rather than silently dropping particles. - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields - # unpack the tiled field state. This first tiled-field pass updates - # Maxwell fields only; particles and particle boundary conditions are - # intentionally not touched. + if J_func is None: + J_tiles = direct_J_from_tiled_particles( + particles, + J_tiles, + constants, + world, + filter="none", + ) + else: + J_tiles = J_func(particles, J_tiles, constants, world) + # deposit current directly into tile-local Yee current arrays E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) # update electric field from B and the supplied tiled current @@ -30,7 +88,7 @@ def time_loop_electrodynamic_tiled( B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) # update magnetic field from the newly updated electric field - fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) + fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) # pack the tiled field state return particles, fields diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index bcdc48d5..d4ab3d60 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -43,15 +43,28 @@ to_flat_particles, check_flat_compat ) +from PyPIC3D.particles.tiled_particle_initialization import ( + to_tiled_particles +) + +from PyPIC3D.solvers.yee_tiled import ( + tile_vector_field +) + from PyPIC3D.evolve import ( time_loop_electrodynamic, time_loop_electrostatic ) +from PyPIC3D.electrodynamic_tiled import ( + time_loop_electrodynamic_tiled +) + 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.deposition.direct_deposition_tiled import direct_J_from_tiled_particles from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -80,9 +93,43 @@ 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 = ("fdtd", "spectral") + supported_solvers = ("fdtd", "spectral", "tiled_yee") if solver not in supported_solvers: - raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd' or 'spectral'.") + raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd', 'spectral', or 'tiled_yee'.") + + +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 _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_active): + """ + Keep the first tile-native PIC path tied to the kernels that exist today. + """ + + if electrostatic: + raise ValueError("tiled_yee is only supported for electrodynamic simulations") + if pml_active: + raise ValueError("tiled_yee is periodic-only and does not support PML") + if simulation_parameters["current_calculation"] != "j_from_rhov": + raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") + if simulation_parameters["particle_pusher"] != "boris": + raise ValueError("tiled_yee currently supports only particle_pusher='boris'") + if simulation_parameters["filter_j"] not in ("none", "digital"): + raise ValueError("tiled_yee currently supports only filter_j='none' or filter_j='digital'") + 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("tiled_yee requires particle_tile_nx/ny/nz to divide Nx/Ny/Nz exactly") def default_parameters(): @@ -286,6 +333,8 @@ def initialize_simulation(toml_file): 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 solver == "tiled_yee": + _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) @@ -337,7 +386,10 @@ def initialize_simulation(toml_file): fast_backend = simulation_parameters.get("fast_backend", "flat") - if fast_backend == "flat": + if fast_backend == "flat" and solver == "tiled_yee": + print("fast_backend='flat' not used for tiled_yee; using tiled particle storage") + simulation_parameters["fast_backend"] = "tiled" + elif fast_backend == "flat": if electrostatic: print("fast_backend='flat' not supported for electrostatic; falling back to default") simulation_parameters["fast_backend"] = "default" @@ -420,7 +472,10 @@ def initialize_simulation(toml_file): else: print(f"Using electrodynamic solver with: {solver}") - evolve_loop = time_loop_electrodynamic + if solver == "tiled_yee": + evolve_loop = time_loop_electrodynamic_tiled + else: + evolve_loop = time_loop_electrodynamic # set the evolve loop function based on the electrostatic flag if simulation_parameters['current_calculation'] == "esirkepov": @@ -429,7 +484,10 @@ def initialize_simulation(toml_file): 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 == "tiled_yee": + J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) + else: + J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) if electrostatic: @@ -442,6 +500,21 @@ def initialize_simulation(toml_file): pml_state = None # electrodynamic FDTD always carries the PML state slot; None means # ordinary, unstretched Yee derivatives. + if solver == "tiled_yee": + tile_shape = _tile_shape_from_parameters(simulation_parameters) + world["tile_shape"] = tile_shape + particles = to_tiled_particles(particles, world, simulation_parameters) + E = tile_vector_field(E, world, tile_shape) + B = tile_vector_field(B, world, tile_shape) + J = tile_vector_field(J, world, tile_shape) + external_E, external_B = external_fields + external_fields = ( + tile_vector_field(external_E, world, tile_shape), + tile_vector_field(external_B, world, tile_shape), + ) + simulation_parameters["fast_backend"] = "tiled" + simulation_parameters["tile_shape"] = tile_shape + print(f"Using tiled Yee storage with tile shape: {tile_shape}") fields = (E, B, J, rho, phi, external_fields, pml_state) # define the fields tuple for the electrodynamic FDTD solver diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index db9a9ab1..1475a2cb 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -2,6 +2,7 @@ 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 def _tile_axis_count(n_cells, cells_per_tile): @@ -92,25 +93,35 @@ def assemble_tiled_vector_field(field_tiles, world, tile_shape): return tuple(_assemble_scalar_tiles(component, world, tile_shape) for component in field_tiles) -def update_tiled_ghost_cells_periodic(field_tiles): +def update_tiled_ghost_cells(field_tiles, world): """ - Refresh one-cell tile halos by periodic exchange with neighbor tiles. + Refresh one-cell tile halos using the field boundary conditions. ``field_tiles`` has shape ``(ntx, nty, ntz, tile_nx + 2, tile_ny + 2, tile_nz + 2)``. The last three axes are the tile-local ghost-celled field. The first three axes - identify the tile in the global tiling. Periodic global boundaries are - handled by wrapping along the corresponding tile axis. + identify the tile in the global tiling. """ - # x-halos: previous tile right interior -> left ghost, next tile left - # interior -> right ghost. + 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 + # update_ghost_cells on the assembled field. field_tiles = field_tiles.at[:, :, :, 0, :, :].set( jnp.roll(field_tiles[:, :, :, -2, :, :], shift=1, axis=0) ) field_tiles = field_tiles.at[:, :, :, -1, :, :].set( jnp.roll(field_tiles[:, :, :, 1, :, :], shift=-1, axis=0) ) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles.at[0, :, :, 0, :, :].set(0.0).at[-1, :, :, -1, :, :].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 the global arrays. @@ -120,6 +131,12 @@ def update_tiled_ghost_cells_periodic(field_tiles): field_tiles = field_tiles.at[:, :, :, :, -1, :].set( jnp.roll(field_tiles[:, :, :, :, 1, :], shift=-1, axis=1) ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles.at[:, 0, :, :, 0, :].set(0.0).at[:, -1, :, :, -1, :].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) # z-halos complete the edge and corner values after x/y have been refreshed. field_tiles = field_tiles.at[:, :, :, :, :, 0].set( @@ -128,10 +145,33 @@ def update_tiled_ghost_cells_periodic(field_tiles): field_tiles = field_tiles.at[:, :, :, :, :, -1].set( jnp.roll(field_tiles[:, :, :, :, :, 1], shift=-1, axis=2) ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles.at[:, :, 0, :, :, 0].set(0.0).at[:, :, -1, :, :, -1].set(0.0), + lambda tiles: tiles, + operand=field_tiles, + ) return field_tiles +def update_tiled_ghost_cells_periodic(field_tiles): + """ + Refresh one-cell tile halos with all-periodic boundary conditions. + """ + + world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} + return update_tiled_ghost_cells(field_tiles, world) + + +def update_tiled_vector_ghost_cells(field_tiles, world): + """ + Refresh tile halos for each component of a vector field. + """ + + return tuple(update_tiled_ghost_cells(component, world) for component in field_tiles) + + def update_tiled_vector_ghost_cells_periodic(field_tiles): """ Refresh periodic tile halos for each component of a vector field. @@ -140,6 +180,59 @@ def update_tiled_vector_ghost_cells_periodic(field_tiles): return tuple(update_tiled_ghost_cells_periodic(component) for component in field_tiles) +def apply_tiled_conducting_bc(E_tiles, world): + """ + 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 + + Ey = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda f: f.at[0, :, :, 1, :, :].set(0.0).at[-1, :, :, -2, :, :].set(0.0), + lambda f: f, + operand=Ey, + ) + Ez = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda f: f.at[0, :, :, 1, :, :].set(0.0).at[-1, :, :, -2, :, :].set(0.0), + lambda f: f, + operand=Ez, + ) + + Ex = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda f: f.at[:, 0, :, :, 1, :].set(0.0).at[:, -1, :, :, -2, :].set(0.0), + lambda f: f, + operand=Ex, + ) + Ez = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda f: f.at[:, 0, :, :, 1, :].set(0.0).at[:, -1, :, :, -2, :].set(0.0), + lambda f: f, + operand=Ez, + ) + + Ex = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda f: f.at[:, :, 0, :, :, 1].set(0.0).at[:, :, -1, :, :, -2].set(0.0), + lambda f: f, + operand=Ex, + ) + Ey = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda f: f.at[:, :, 0, :, :, 1].set(0.0).at[:, :, -1, :, :, -2].set(0.0), + lambda f: f, + operand=Ey, + ) + + return Ex, Ey, Ez + + def fold_tiled_ghost_cells_periodic(field_tiles): """ Add tile-ghost deposits into neighboring tile interiors, then clear ghosts. @@ -192,16 +285,14 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ """ Update compact tiled electric fields without assembling a global field. - This tile-native path is periodic-only and intentionally skips the digital - filter used by the global Yee update. The Yee curl is evaluated on each - tile's physical interior after B halos have been refreshed from neighbor - tiles. + The Yee curl is evaluated on each tile's physical interior after B halos + have been refreshed from neighbor tiles or field boundary conditions. """ del curl_func, tile_shape Ex, Ey, Ez = E_tiles - Bx, By, Bz = update_tiled_vector_ghost_cells_periodic(B_tiles) + Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world) Jx, Jy, Jz = J_tiles dt = world["dt"] @@ -236,22 +327,22 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ + (C**2 * curl_z - Jz[:, :, :, 1:-1, 1:-1, 1:-1] / eps) * dt ) - return update_tiled_vector_ghost_cells_periodic((Ex, Ey, Ez)) + Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world) + + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): """ Update compact tiled magnetic fields without assembling a global field. - This tile-native path is periodic-only and intentionally skips the digital - filter used by the global Yee update. The Yee curl is evaluated on each - tile's physical interior after E halos have been refreshed from neighbor - tiles. + The Yee curl is evaluated on each tile's physical interior after E halos + have been refreshed from neighbor tiles or field boundary conditions. """ del constants, curl_func, tile_shape - Ex, Ey, Ez = update_tiled_vector_ghost_cells_periodic(E_tiles) + Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world) Bx, By, Bz = B_tiles dt = world["dt"] @@ -275,4 +366,4 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): 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) - return update_tiled_vector_ghost_cells_periodic((Bx, By, Bz)) + return update_tiled_vector_ghost_cells((Bx, By, Bz), world) diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 4b9857df..3126de19 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -209,8 +209,20 @@ 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: + interior = ( + slice(None), + slice(None), + slice(None), + slice(1, -1), + slice(1, -1), + slice(1, -1), + ) + 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 @@ -224,24 +236,48 @@ def nd_trapezoid(arr, dxs): C = constants['C'] # speed of light kinetic_energy = 0.0 - for species in particles: - mass = species.get_mass() - vx, vy, vz = species.get_velocity() + if hasattr(particles, "active") and hasattr(particles, "u"): + vx = particles.u[..., 0] + vy = particles.u[..., 1] + vz = particles.u[..., 2] v2 = vx**2 + vy**2 + vz**2 - active = species.get_active_mask().astype(v2.dtype) + active = particles.active.astype(v2.dtype) + mass = particles.mass * particles.weight 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 + momentum2 = jnp.square(mass * gamma) * v2 + kinetic_energy = jnp.sum(active * (jnp.sqrt(momentum2 * C**2 + mass**2 * C**4) - mass * C**2)) + else: + 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 # Kinetic energy of particles return e_energy, b_energy, kinetic_energy +def compute_total_momentum(particles): + """ + Compute the scalar momentum diagnostic for either species-list or tiled particles. + """ + + if hasattr(particles, "active") and hasattr(particles, "u"): + vmag = jnp.sqrt(particles.u[..., 0]**2 + particles.u[..., 1]**2 + particles.u[..., 2]**2) + active = particles.active.astype(vmag.dtype) + return jnp.sum(active * vmag * particles.mass * particles.weight) + + return sum(particle_species.momentum() for particle_species in particles) + + def add_external_fields(E, B, external_fields): """ Add prescribed external fields to the self-consistent fields. @@ -765,23 +801,35 @@ 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 - } - config["particles"].append(particle_dict) + if hasattr(particles, "active") and hasattr(particles, "u"): + config["particles"].append( + { + "storage": "tiled", + "active_particles": int(jnp.sum(particles.active)), + "tile_shape": jax.tree_util.tree_map( + lambda x: x.tolist() if isinstance(x, jnp.ndarray) else x, + simulation_parameters.get("tile_shape", ()), + ), + } + ) + else: + 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 + } + config["particles"].append(particle_dict) config["version"] = { "PyPIC3D_version": importlib.metadata.version('PyPIC3D'), diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index d1b621a7..88b2726a 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -3,12 +3,20 @@ import jax import jax.numpy as jnp +from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.deposition.J_from_rhov import J_from_rhov -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.direct_deposition_tiled import ( + digital_filter_tiled_current_density, + direct_J_from_tiled_particles, +) from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, fold_tiled_ghost_cells_periodic -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_vector_field, + fold_tiled_ghost_cells_periodic, + tile_vector_field, +) +from PyPIC3D.utils import build_yee_grid, digital_filter jax.config.update("jax_enable_x64", True) @@ -80,6 +88,35 @@ def _compare_tiled_to_standard(self, particles, world, simulation_parameters): 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)) + 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_tiled_current_density(J_tiles, alpha) + 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-12, atol=1.0e-12)) + def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self): tiles = jnp.zeros((2, 1, 1, 4, 3, 3)) tiles = tiles.at[0, 0, 0, -1, 1, 1].set(2.0) @@ -194,6 +231,55 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_dummy_species(self): self._compare_tiled_to_standard([electrons, ions], world, simulation_parameters) + 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} + species = particle_species( + name="digital filtered current", + N_particles=4, + charge=-1.0, + mass=1.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"], + ) + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + + J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") + J_tiles = direct_J_from_tiled_particles( + tiled_particles, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="digital", + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + 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)) + def test_tiled_direct_deposition_respects_active_mask(self): world = self._build_world() simulation_parameters = { diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py new file mode 100644 index 00000000..9a83e459 --- /dev/null +++ b/tests/tiled_yee_integration_test.py @@ -0,0 +1,407 @@ +import functools +import os +import tempfile +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import toml + +from PyPIC3D.evolve import time_loop_electrodynamic +from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.initialization import initialize_simulation, validate_field_solver +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, tile_vector_field +from PyPIC3D.utils import build_yee_grid, compute_energy + + +jax.config.update("jax_enable_x64", True) + + +def unused_curl(Ex, Ey, Ez): + return None + + +class TestTiledYeeIntegration(unittest.TestCase): + def _build_world(self): + world = { + "Nx": 8, + "Ny": 1, + "Nz": 1, + "dx": 0.5, + "dy": 1.0, + "dz": 1.0, + "dt": 0.01, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.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 _empty_fields(self, world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + zero = jnp.zeros(shape) + E = (zero, zero, zero) + B = (zero, zero, zero) + J = (zero, zero, zero) + external_fields = (E, B) + return E, B, J, zero, zero, external_fields, None + + def _species(self, world): + return particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + 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.lexsort((x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def test_validate_field_solver_accepts_tiled_yee(self): + validate_field_solver("tiled_yee") + + def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(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 init smoke", + "output_dir": tmpdir, + "solver": "tiled_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": "j_from_rhov", + "filter_j": "none", + "fast_backend": "default", + "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.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) + + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + + def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(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 digital filter init smoke", + "output_dir": tmpdir, + "solver": "tiled_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": "j_from_rhov", + "filter_j": "digital", + "fast_backend": "default", + "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_digital.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["filter_j"], "digital") + + def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(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 conducting init smoke", + "output_dir": tmpdir, + "solver": "tiled_yee", + "Nx": 8, + "Ny": 1, + "Nz": 1, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "x_bc": "conducting", + "y_bc": "conducting", + "z_bc": "conducting", + "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", + "fast_backend": "default", + "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_conducting.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) + + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(world["boundary_conditions"], {"x": 1, "y": 1, "z": 1}) + + def test_tiled_yee_step_matches_standard_periodic_yee_step(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + reference_particles, reference_fields = time_loop_electrodynamic( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="none"), + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + tiled_particles, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=None, + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) + reference_E, reference_B, reference_J, *_ = reference_fields + + flat_x = jnp.stack(reference_particles[0].get_forward_position(), axis=1) + flat_u = jnp.stack(reference_particles[0].get_velocity(), axis=1) + flat_order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + tiled_x, tiled_u = self._active_rows(tiled_particles) + + self.assertTrue(jnp.allclose(tiled_x, flat_x[flat_order], rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(tiled_u, flat_u[flat_order], rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_E, E_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_B, B_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape)): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + reference_energy = compute_energy(reference_particles, reference_E, reference_B, world, constants) + tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants) + for reference, tiled in zip(reference_energy, tiled_energy): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_yee_step_uses_digital_filtered_current(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} + tile_shape = (2, 1, 1) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + _, reference_fields = time_loop_electrodynamic( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="digital"), + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + _, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=functools.partial(direct_J_from_tiled_particles, filter="digital"), + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + _, _, reference_J, *_ = reference_fields + _, _, J_tiles, *_ = tiled_fields + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + for reference, tiled in zip(reference_J, J_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index c983c55a..ada896c1 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -9,13 +9,16 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, tile_vector_field, + update_tiled_ghost_cells, update_tiled_ghost_cells_periodic, + update_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells_periodic, update_tiled_B, update_tiled_E, ) from PyPIC3D.utils import build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC jax.config.update("jax_enable_x64", True) @@ -45,6 +48,16 @@ def _build_world(self): 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 _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"] @@ -107,6 +120,73 @@ def test_update_tiled_vector_ghost_cells_periodic_refreshes_each_component(self) 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) + 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 = 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) + 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 = 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_vector_ghost_cells_conducting_refreshes_each_component(self): + world = self._conducting_world() + tile_shape = (2, 3, 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 = 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_tiled_E_matches_standard_yee_update(self): world = self._build_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} @@ -130,6 +210,29 @@ def test_update_tiled_E_matches_standard_yee_update(self): 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_tiled_E_matches_standard_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) + 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, _ = update_E(E, B, J, world, constants, unused_curl) + E_tiled = update_tiled_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + unused_curl, + tile_shape, + ) + 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_tiled_B_matches_standard_yee_update(self): world = self._build_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} @@ -151,6 +254,27 @@ def test_update_tiled_B_matches_standard_yee_update(self): 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_tiled_B_matches_standard_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) + E = self._deterministic_vector_field(world, scale=1.0) + B = self._deterministic_vector_field(world, scale=0.2) + + B_reference, _ = update_B(E, B, world, constants, unused_curl) + B_tiled = update_tiled_B( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + unused_curl, + tile_shape, + ) + 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_standard_yee_sequence(self): world = self._build_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} @@ -176,6 +300,31 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): 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_standard_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) + 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, _ = update_E(E, B, J, world, constants, unused_curl) + B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + + 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape) + + 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} From 3445d8e066a0dc5990da8860f972ee7a5d8aeb60 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 25 Jun 2026 19:54:18 -0500 Subject: [PATCH 04/83] implementing tiled pml --- PyPIC3D/boundary_conditions/PML.py | 166 +++++++++++++++++++++++++ PyPIC3D/electrodynamic_tiled.py | 24 +++- PyPIC3D/initialization.py | 16 +-- PyPIC3D/solvers/yee_tiled.py | 80 ++++++++++-- tests/pml_test.py | 193 +++++++++++++++++++++++++++++ 5 files changed, 455 insertions(+), 24 deletions(-) diff --git a/PyPIC3D/boundary_conditions/PML.py b/PyPIC3D/boundary_conditions/PML.py index 5f3572ab..9a4e724f 100644 --- a/PyPIC3D/boundary_conditions/PML.py +++ b/PyPIC3D/boundary_conditions/PML.py @@ -1,5 +1,6 @@ import math +import jax import jax.numpy as jnp from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -128,6 +129,81 @@ 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): + """ + 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) + + 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 keep + the same one-cell tile halos as the field tiles, so the derivative stretch + uses ``[..., 1:-1, 1:-1, 1:-1]`` next to the Yee derivative arrays. + """ + + _, _, _, _, profiles = world["pml"] + return tuple(_tile_scalar_profile(profile, tile_shape) 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,6 +325,96 @@ def apply_pml_to_b_curl(derivatives, world, pml_state): return (curl_x, curl_y, curl_z), (e_memory, b_memory) +def apply_tiled_pml_to_e_curl(derivatives, world, pml_state): + """ + Stretch tile-local B derivatives before assembling Ampere's-law curls. + """ + + 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 + sigma_x = sigma_x[:, :, :, 1:-1, 1:-1, 1:-1] + sigma_y = sigma_y[:, :, :, 1:-1, 1:-1, 1:-1] + sigma_z = sigma_z[:, :, :, 1:-1, 1:-1, 1:-1] + 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. + """ + + 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 + sigma_x = sigma_x[:, :, :, 1:-1, 1:-1, 1:-1] + sigma_y = sigma_y[:, :, :, 1:-1, 1:-1, 1:-1] + sigma_z = sigma_z[:, :, :, 1:-1, 1:-1, 1:-1] + 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 (curl_x, curl_y, curl_z), (e_memory, b_memory, tiled_profiles) + + def update_ghost_cells_for_pml(field, world): """ Fill ghost cells without periodically wrapping across a PML wall. diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 5bc0c36d..fe4f81af 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -40,8 +40,16 @@ def time_loop_electrodynamic_tiled( tile_shape = tuple(int(width) - 2 for width in E_tiles[0].shape[3:]) if not hasattr(particles, "active"): - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + if pml_state is None: + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + else: + E_tiles, pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state + ) + B_tiles, pml_state = update_tiled_B( + E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state + ) fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) return particles, fields # keep the original field-only helper behavior for tests and standalone @@ -82,10 +90,18 @@ def time_loop_electrodynamic_tiled( J_tiles = J_func(particles, J_tiles, constants, world) # deposit current directly into tile-local Yee current arrays - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + if pml_state is None: + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + else: + E_tiles, pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state + ) # update electric field from B and the supplied tiled current - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + if pml_state is None: + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + else: + B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state) # update magnetic field from the newly updated electric field fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index d4ab3d60..2186b969 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -70,6 +70,7 @@ 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, ) @@ -113,8 +114,6 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ if electrostatic: raise ValueError("tiled_yee is only supported for electrodynamic simulations") - if pml_active: - raise ValueError("tiled_yee is periodic-only and does not support PML") if simulation_parameters["current_calculation"] != "j_from_rhov": raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") if simulation_parameters["particle_pusher"] != "boris": @@ -331,8 +330,8 @@ 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 (solver not in ("fdtd", "tiled_yee") or electrostatic): + raise ValueError("PML is only supported for the fdtd or tiled_yee electrodynamic solvers") if solver == "tiled_yee": _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_active) world["pml"] = load_pml_from_toml(raw_pml, world, constants) @@ -494,15 +493,14 @@ def initialize_simulation(toml_file): fields = (E, B, J, rho, phi, external_fields) # define the fields tuple for the electrostatic solver else: - if pml_active: - pml_state = initialize_pml_state(world) - else: - pml_state = None + pml_state = None # electrodynamic FDTD always carries the PML state slot; None means # ordinary, unstretched Yee derivatives. if solver == "tiled_yee": tile_shape = _tile_shape_from_parameters(simulation_parameters) world["tile_shape"] = tile_shape + if pml_active: + pml_state = initialize_tiled_pml_state(world, tile_shape) particles = to_tiled_particles(particles, world, simulation_parameters) E = tile_vector_field(E, world, tile_shape) B = tile_vector_field(B, world, tile_shape) @@ -515,6 +513,8 @@ def initialize_simulation(toml_file): simulation_parameters["fast_backend"] = "tiled" simulation_parameters["tile_shape"] = tile_shape print(f"Using tiled Yee storage with tile shape: {tile_shape}") + elif pml_active: + pml_state = initialize_pml_state(world) fields = (E, B, J, rho, phi, external_fields, pml_state) # define the fields tuple for the electrodynamic FDTD solver diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 1475a2cb..dbc17ffa 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -1,6 +1,10 @@ import jax import jax.numpy as jnp +from PyPIC3D.boundary_conditions.PML import ( + apply_tiled_pml_to_b_curl, + apply_tiled_pml_to_e_curl, +) from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC @@ -180,6 +184,32 @@ def update_tiled_vector_ghost_cells_periodic(field_tiles): return tuple(update_tiled_ghost_cells_periodic(component) for component in field_tiles) +def update_tiled_ghost_cells_for_pml(field_tiles, world): + """ + 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}} + return update_tiled_ghost_cells(field_tiles, pml_world) + + +def update_tiled_vector_ghost_cells_for_pml(field_tiles, world): + """ + Refresh tile halos for a vector field with PML-active exterior walls. + """ + + return tuple(update_tiled_ghost_cells_for_pml(component, world) for component in field_tiles) + + def apply_tiled_conducting_bc(E_tiles, world): """ Zero tangential electric-field components on global conducting faces. @@ -281,7 +311,7 @@ def fold_tiled_vector_ghost_cells_periodic(field_tiles): return tuple(fold_tiled_ghost_cells_periodic(component) for component in field_tiles) -def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape): +def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -292,7 +322,10 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ del curl_func, tile_shape Ex, Ey, Ez = E_tiles - Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world) + if pml_state is None: + Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world) + else: + Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml(B_tiles, world) Jx, Jy, Jz = J_tiles dt = world["dt"] @@ -310,9 +343,16 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ 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 - curl_x = dBz_dy - dBy_dz - curl_y = dBx_dz - dBz_dx - curl_z = dBy_dx - dBx_dy + if pml_state is None: + curl_x = dBz_dy - dBy_dz + curl_y = dBx_dz - dBz_dx + curl_z = dBy_dx - dBx_dy + else: + (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, + ) Ex = Ex.at[:, :, :, 1:-1, 1:-1, 1:-1].set( Ex[:, :, :, 1:-1, 1:-1, 1:-1] @@ -329,10 +369,13 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world) - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + if pml_state is None: + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world), pml_state -def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): + +def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state=None): """ Update compact tiled magnetic fields without assembling a global field. @@ -342,7 +385,10 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): del constants, curl_func, tile_shape - Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world) + if pml_state is None: + Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world) + else: + Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml(E_tiles, world) Bx, By, Bz = B_tiles dt = world["dt"] @@ -358,12 +404,22 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape): 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 - curl_x = dEz_dy - dEy_dz - curl_y = dEx_dz - dEz_dx - curl_z = dEy_dx - dEx_dy + if pml_state is None: + curl_x = dEz_dy - dEy_dz + curl_y = dEx_dz - dEz_dx + curl_z = dEy_dx - dEx_dy + else: + (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, + ) 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) - return update_tiled_vector_ghost_cells((Bx, By, Bz), world) + if pml_state is None: + return update_tiled_vector_ghost_cells((Bx, By, Bz), world) + + return update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world), pml_state diff --git a/tests/pml_test.py b/tests/pml_test.py index b48a4f02..4cbce570 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -8,10 +8,19 @@ PML_WALLS, build_pml, initialize_pml_state, + initialize_tiled_pml_state, load_pml_from_toml, + tile_pml_profiles, + update_ghost_cells_for_pml, ) from PyPIC3D.initialization import initialize_fields, initialize_simulation from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_vector_field, + tile_vector_field, + update_tiled_B, + update_tiled_E, +) from PyPIC3D.utils import build_yee_grid, compute_energy jax.config.update("jax_enable_x64", True) @@ -140,6 +149,53 @@ def test_initialize_pml_state_uses_physical_interior_shape(self): 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, 4, 4, 3)) + + 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) + + 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_non_fdtd_solvers(self): pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] @@ -167,6 +223,32 @@ def test_initialize_simulation_appends_pml_state_for_fdtd(self): self.assertFalse(pml_y) self.assertFalse(pml_z) + def test_initialize_simulation_uses_tiled_pml_state_for_tiled_yee(self): + pml = [{"wall": "+x", "thickness": 2, "sigma_max": 1.0}] + + with tempfile.TemporaryDirectory() as tmpdir: + config = _empty_config(tmpdir, solver="tiled_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[-1] + e_memory, b_memory, tiled_profiles = pml_state + + self.assertEqual(len(fields), 7) + 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(tiled_profiles[0].shape, (4, 1, 1, 4, 3, 3)) + def test_initialize_simulation_uses_none_pml_state_without_pml_for_fdtd(self): with tempfile.TemporaryDirectory() as tmpdir: result = initialize_simulation(_empty_config(tmpdir)) @@ -229,6 +311,117 @@ def test_pml_absorbs_field_energy_in_particle_free_1d_wave(self): self.assertTrue(jnp.isfinite(final_energy)) self.assertLess(float(final_energy), 0.35 * float(initial_energy)) + def test_tiled_pml_matches_global_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, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + + 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) + + pml_state = initialize_pml_state(world) + E_reference, pml_state = update_E(E, B, J, world, constants, lambda *args: None, pml_state) + B_reference, pml_state = update_B(E_reference, B, world, constants, lambda *args: None, pml_state) + + tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) + E_tiles, tiled_pml_state = update_tiled_E( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + world, + constants, + lambda *args: None, + tile_shape, + tiled_pml_state, + ) + B_tiles, tiled_pml_state = update_tiled_B( + E_tiles, + tile_vector_field(B, world, tile_shape), + world, + constants, + lambda *args: None, + tile_shape, + tiled_pml_state, + ) + + E_tiled = assemble_tiled_vector_field(E_tiles, world, tile_shape) + B_tiled = assemble_tiled_vector_field(B_tiles, world, tile_shape) + + 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) + tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) + E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + + 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) + + initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) + def step(E_tiles, B_tiles, tiled_pml_state): + E_tiles, tiled_pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, tiled_pml_state + ) + B_tiles, tiled_pml_state = update_tiled_B( + E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, 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([], E_tiles, B_tiles, world, constants)[:2]) + self.assertTrue(jnp.isfinite(final_energy)) + self.assertLess(float(final_energy), 0.65 * float(initial_energy)) + if __name__ == "__main__": unittest.main() From bd958e395e88750b6e599fec65522652113340a0 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 25 Jun 2026 20:21:23 -0500 Subject: [PATCH 05/83] expanding tiled particles boundary condition --- PyPIC3D/deposition/Esirkepov.py | 58 ++++++++++++++-- PyPIC3D/evolve.py | 4 +- PyPIC3D/initialization.py | 23 ++++++- PyPIC3D/particles/flat_particles.py | 56 ++++++++++----- PyPIC3D/particles/species_class.py | 57 ++++++++-------- PyPIC3D/particles/tiled_particle_refresh.py | 76 ++++++++++++++++++--- docs/particles.rst | 11 +-- docs/usage.rst | 9 ++- tests/evolve_test.py | 4 +- tests/initialization_test.py | 54 +++++++++++++++ tests/particle_test.py | 49 +++++++++++-- tests/tiled_particle_refresh_test.py | 37 +++++++++- tests/utils_test.py | 4 +- 13 files changed, 361 insertions(+), 81 deletions(-) diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 5a38faf2..45b53760 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -126,6 +126,46 @@ def enforce_bc_along_axis(field, axis, bc, component_axis): return field + +def enforce_particle_bc_code_along_axis(field, axis, bc, component_axis): + """ + Fold two Esirkepov ghost layers using global particle boundary-condition codes. + """ + + field_axis = jnp.moveaxis(field, axis, 0) + sign = -1.0 if axis == component_axis else 1.0 + + def periodic_bc(field_axis): + 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, :, :]) + return field_axis + + def reflecting_bc(field_axis): + 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, :, :]) + return field_axis + + def absorbing_bc(field_axis): + return field_axis + + field_axis = lax.switch(bc, (periodic_bc, reflecting_bc, absorbing_bc), field_axis) + field = jnp.moveaxis(field_axis, 0, axis) + field = clear_ghost_cells(field, axis) + + return field + + +def ghost_cell_particle_bc_codes_esirkepov(field, bc_x, bc_y, bc_z, component_axis): + field = enforce_particle_bc_code_along_axis(field, axis=0, bc=bc_x, component_axis=component_axis) + field = enforce_particle_bc_code_along_axis(field, axis=1, bc=bc_y, component_axis=component_axis) + field = enforce_particle_bc_code_along_axis(field, axis=2, bc=bc_z, component_axis=component_axis) + return field + + def ghost_cell_bc_esirkepov(field, bc_x, bc_y, bc_z, component_axis): """ Apply boundary conditions to ghost cells using Esirkepov method. @@ -385,14 +425,18 @@ def Esirkepov_current(particles, J, constants, world, grid=None, filter=None): 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 + particle_bc = world.get("particle_boundary_conditions", {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}) + # determine the particle boundary condition type from the global world state - 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) + Jx_ext = ghost_cell_particle_bc_codes_esirkepov( + Jx_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=0 + ) + Jy_ext = ghost_cell_particle_bc_codes_esirkepov( + Jy_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=1 + ) + Jz_ext = ghost_cell_particle_bc_codes_esirkepov( + Jz_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=2 + ) # Fold the two Esirkepov ghost layers before returning to the solver's # ordinary one-ghost current arrays. diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 2a86aa09..49e2fe6e 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -68,7 +68,7 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu # pack the fields into a tuple for i in range(len(particles)): - particles[i].boundary_conditions() + particles[i].boundary_conditions(world) # apply boundary conditions to the particles return particles, fields @@ -149,7 +149,7 @@ def time_loop_electrodynamic(particles, fields, world, constants, curl_func, J_f # update the magnetic field using the curl of the electric field for i in range(len(particles)): - particles[i].boundary_conditions() + particles[i].boundary_conditions(world) # apply boundary conditions to the particles fields = (E, B, J, rho, phi, external_fields, pml_state) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 2186b969..0894ad25 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -89,6 +89,20 @@ 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 @@ -167,7 +181,9 @@ def default_parameters(): "output_dir": os.getcwd(), "solver": "fdtd", # solver: spectral, fdtd "fast_backend": "flat", # flat | default (flat when compatible, else fallback) - "particle_bc": "periodic", # particle boundary conditions: periodic, absorbing, reflecting + "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 @@ -323,6 +339,11 @@ def initialize_simulation(toml_file): '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 diff --git a/PyPIC3D/particles/flat_particles.py b/PyPIC3D/particles/flat_particles.py index 1ce9671f..fe34813b 100644 --- a/PyPIC3D/particles/flat_particles.py +++ b/PyPIC3D/particles/flat_particles.py @@ -2,6 +2,25 @@ import jax.numpy as jnp +@jax.jit +def _apply_axis_boundary_condition(x, v, wind, half_wind, bc): + periodic = bc == 0 + reflecting = bc == 1 + + periodic_x = x + wind * (x < -half_wind) - wind * (x > half_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_v = jnp.where((x >= half_wind) | (x <= -half_wind), -v, v) + + x_out = jnp.where(periodic, periodic_x, jnp.where(reflecting, reflected_x, x)) + v_out = jnp.where(reflecting, reflected_v, v) + + return x_out, v_out + + @jax.tree_util.register_pytree_node_class class flat_particle_species: def __init__( @@ -156,27 +175,34 @@ def update_position(self): if self.update_z: self.x3 = self.x3 + active * self.v3 * self.dt - def boundary_conditions(self): + def boundary_conditions(self, world=None): half_x = self.x_wind / 2 half_y = self.y_wind / 2 half_z = self.z_wind / 2 + if world is None or "particle_boundary_conditions" not in world: + particle_bc = {"x": 0, "y": 0, "z": 0} + else: + particle_bc = world["particle_boundary_conditions"] - 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.x1, self.v1 = _apply_axis_boundary_condition( + self.x1, self.v1, self.x_wind, half_x, particle_bc["x"] ) - 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.x2, self.v2 = _apply_axis_boundary_condition( + self.x2, self.v2, self.y_wind, half_y, particle_bc["y"] ) - 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), + self.x3, self.v3 = _apply_axis_boundary_condition( + self.x3, self.v3, self.z_wind, half_z, particle_bc["z"] ) + active_mask = self.active_mask + x_inside = (self.x1 <= half_x) & (self.x1 >= -half_x) + y_inside = (self.x2 <= half_y) & (self.x2 >= -half_y) + z_inside = (self.x3 <= half_z) & (self.x3 >= -half_z) + active_mask = jnp.where(particle_bc["x"] == 2, active_mask & x_inside, active_mask) + active_mask = jnp.where(particle_bc["y"] == 2, active_mask & y_inside, active_mask) + active_mask = jnp.where(particle_bc["z"] == 2, active_mask & z_inside, active_mask) + self.active_mask = active_mask + def tree_flatten(self): children = (self.x1, self.x2, self.x3, self.v1, self.v2, self.v3, self.active_mask) aux_data = ( @@ -289,10 +315,6 @@ def _same(attr_list): def check_flat_compat(particles): if not 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 diff --git a/PyPIC3D/particles/species_class.py b/PyPIC3D/particles/species_class.py index b7b92cff..484c6f87 100644 --- a/PyPIC3D/particles/species_class.py +++ b/PyPIC3D/particles/species_class.py @@ -1,6 +1,5 @@ import jax from jax import jit -from functools import partial import jax.numpy as jnp from jax.tree_util import register_pytree_node_class @@ -12,29 +11,25 @@ def compute_index(x, dx, window): 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): +@jit +def apply_axis_boundary_condition(x, v, wind, half_wind, bc): """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 + periodic = bc == 0 + reflecting = bc == 1 - 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 + periodic_x = x + wind * (x < -half_wind) - wind * (x > half_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_v = jnp.where((x >= half_wind) | (x <= -half_wind), -v, v) - def identity_bc(state): - return state + x_out = jnp.where(periodic, periodic_x, jnp.where(reflecting, reflected_x, x)) + v_out = jnp.where(reflecting, reflected_v, v) - return jax.lax.cond( - periodic, - periodic_bc, - lambda state: jax.lax.cond(reflecting, reflecting_bc, identity_bc, state), - (x, v), - ) + return x_out, v_out @register_pytree_node_class @@ -220,27 +215,31 @@ 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): + def boundary_conditions(self, world=None): x1, x2, x3 = self.x1, self.x2, self.x3 v1, v2, v3 = self.v1, self.v2, self.v3 + if world is None or "particle_boundary_conditions" not in world: + particle_bc = {"x": 0, "y": 0, "z": 0} + else: + particle_bc = world["particle_boundary_conditions"] x1, v1 = apply_axis_boundary_condition( - x1, v1, self.x_wind, self.half_x_wind, self.x_periodic, self.x_reflecting + x1, v1, self.x_wind, self.half_x_wind, particle_bc["x"] ) x2, v2 = apply_axis_boundary_condition( - x2, v2, self.y_wind, self.half_y_wind, self.y_periodic, self.y_reflecting + x2, v2, self.y_wind, self.half_y_wind, particle_bc["y"] ) x3, v3 = apply_axis_boundary_condition( - x3, v3, self.z_wind, self.half_z_wind, self.z_periodic, self.z_reflecting + x3, v3, self.z_wind, self.half_z_wind, particle_bc["z"] ) 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) + x_inside = (x1 <= self.half_x_wind) & (x1 >= -self.half_x_wind) + y_inside = (x2 <= self.half_y_wind) & (x2 >= -self.half_y_wind) + z_inside = (x3 <= self.half_z_wind) & (x3 >= -self.half_z_wind) + active_mask = jnp.where(particle_bc["x"] == 2, active_mask & x_inside, active_mask) + active_mask = jnp.where(particle_bc["y"] == 2, active_mask & y_inside, active_mask) + active_mask = jnp.where(particle_bc["z"] == 2, active_mask & z_inside, active_mask) self.x1, self.x2, self.x3 = x1, x2, x3 self.v1, self.v2, self.v3 = v1, v2, v3 diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index eb7c2c4f..c544a264 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -4,6 +4,33 @@ from PyPIC3D.particles.tiled_particles 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 @@ -67,15 +94,44 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): ntx, nty, ntz, n_species, n_slots = tiled_particles.active.shape tile_counts = (ntx, nty, ntz) - wrapped_x = tiled_particles.x - wrapped_x = wrapped_x.at[..., 0].set(wrap_periodic_position(wrapped_x[..., 0], world["x_wind"])) - wrapped_x = wrapped_x.at[..., 1].set(wrap_periodic_position(wrapped_x[..., 1], world["y_wind"])) - wrapped_x = wrapped_x.at[..., 2].set(wrap_periodic_position(wrapped_x[..., 2], world["z_wind"])) + 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( - wrapped_x[..., 0], - wrapped_x[..., 1], - wrapped_x[..., 2], + bounded_x[..., 0], + bounded_x[..., 1], + bounded_x[..., 2], world, tile_shape, tile_counts, @@ -126,12 +182,12 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): for oz in z_offsets: sz = (tz + oz) % ntz index = (sx, sy, sz, species) - candidate_x.append(wrapped_x[index]) - candidate_u.append(tiled_particles.u[index]) + candidate_x.append(bounded_x[index]) + candidate_u.append(bounded_u[index]) candidate_charge.append(tiled_particles.charge[index]) candidate_mass.append(tiled_particles.mass[index]) candidate_weight.append(tiled_particles.weight[index]) - candidate_active.append(tiled_particles.active[index]) + candidate_active.append(bounded_active[index]) candidate_update_x1.append(tiled_particles.update_x1[index]) candidate_update_x2.append(tiled_particles.update_x2[index]) candidate_update_x3.append(tiled_particles.update_x3[index]) diff --git a/docs/particles.rst b/docs/particles.rst index da9e5683..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 --------------- @@ -35,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 @@ -54,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/usage.rst b/docs/usage.rst index ee64764a..fadb696a 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -50,6 +50,9 @@ 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] @@ -76,7 +79,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 +92,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``). diff --git a/tests/evolve_test.py b/tests/evolve_test.py index a59f132d..2e8ef63d 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -34,6 +34,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field "z_wind": 1.0, "shape_factor": 1, "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} @@ -106,6 +107,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): "z_wind": 1.0, "shape_factor": 1, "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} @@ -136,7 +138,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): dx=world["dx"], dy=world["dy"], dz=world["dz"], - x_bc="absorbing", + x_bc="periodic", dt=world["dt"], ) ] diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 7fe60a62..6a23b4c9 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -1,5 +1,8 @@ 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 @@ -23,10 +26,61 @@ def test_default_parameters(self): self.assertIn('Nx', sim) self.assertIn('particle_pusher', sim) self.assertEqual(sim['particle_pusher'], 'boris') + self.assertEqual(sim["particle_x_bc"], "periodic") + self.assertEqual(sim["particle_y_bc"], "periodic") + self.assertEqual(sim["particle_z_bc"], "periodic") self.assertIn('eps', const) self.assertIn('plotfields', plotting) # check that the default parameters contain expected keys + 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": "fdtd", + "Nx": 1, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + "fast_backend": "default", + "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.assertEqual(particles[0].x_bc, "absorbing") + def test_initialize_simulation_rejects_unknown_solver(self): with tempfile.TemporaryDirectory() as tmpdir: config = { diff --git a/tests/particle_test.py b/tests/particle_test.py index 1a5dbab8..a2653280 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -189,7 +189,48 @@ def test_update_position(self): 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_reflecting_boundary_mirrors_position_from_global_world_bc(self): + world = { + "particle_boundary_conditions": {"x": 1, "y": 0, "z": 0}, + } + x1 = jnp.array([self.x_wind / 2 + 0.1, -self.x_wind / 2 - 0.2, 0.0]) + x2 = jnp.array([0.0, 0.0, 0.0]) + x3 = jnp.array([0.0, 0.0, 0.0]) + + species = particle_species( + name="reflecting", + 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=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="periodic", + dt=self.dt, + ) + + species.boundary_conditions(world) + + x, _, _ = species.get_forward_position() + vx, _, _ = species.get_velocity() + self.assertTrue(jnp.allclose(x, jnp.array([self.x_wind / 2 - 0.1, -self.x_wind / 2 + 0.2, 0.0]))) + self.assertTrue(jnp.allclose(vx, jnp.array([-1.0, 2.0, 3.0]))) + def test_absorbing_boundary_updates_active_mask(self): + world = { + "particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}, + } 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]) @@ -213,11 +254,11 @@ def test_absorbing_boundary_updates_active_mask(self): xwind=self.x_wind, ywind=self.y_wind, zwind=self.z_wind, - x_bc="absorbing", + x_bc="periodic", dt=self.dt, ) - species.boundary_conditions() + species.boundary_conditions(world) 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) @@ -286,7 +327,7 @@ def test_inactive_particles_do_not_advance_or_take_new_velocity(self): x_bc="absorbing", dt=self.dt, ) - species.boundary_conditions() + species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) species.set_velocity( jnp.array([10.0, 20.0, 30.0]), @@ -515,7 +556,7 @@ def test_inactive_particles_do_not_deposit_rho_or_current(self): zwind=self.z_wind, x_bc="absorbing", ) - species.boundary_conditions() + species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) constants = {"C": 3e8, "alpha": 1.0} rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 4221bdd0..7928c2f0 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -28,9 +28,10 @@ def _build_world(self): "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): + 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) @@ -56,6 +57,7 @@ def _species(self, world, x1, v1, active_mask=None, update_x=True): dt=world["dt"], active_mask=active_mask, update_x=update_x, + x_bc=x_bc, ) def _simulation_parameters(self): @@ -128,6 +130,39 @@ def test_refresh_wraps_periodic_particle_to_opposite_tile(self): 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 = to_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 = to_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]) diff --git a/tests/utils_test.py b/tests/utils_test.py index 3b4bde56..ac50ab7e 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -199,9 +199,9 @@ def test_compute_energy_ignores_inactive_particles(self): dx=1.0, dy=1.0, dz=1.0, - x_bc="absorbing", + x_bc="periodic", ) - species.boundary_conditions() + species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) _, _, kinetic_energy = compute_energy([species], E, B, world, constants) From f7cd42349ab1a1a60ebb42b58debd303e9d79e42 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 05:46:43 -0500 Subject: [PATCH 06/83] completing tasks 1 and 2 of migration plan --- PyPIC3D/__main__.py | 12 +- PyPIC3D/diagnostics/openPMD.py | 12 +- PyPIC3D/diagnostics/plotting.py | 9 +- PyPIC3D/diagnostics/vtk.py | 10 +- PyPIC3D/initialization.py | 10 +- .../particles/tiled_particle_diagnostics.py | 117 +++++++++++++ PyPIC3D/particles/tiled_particles.py | 10 +- PyPIC3D/solvers/yee_tiled.py | 4 +- docs/architecture.rst | 32 ++++ tests/tiled_particle_diagnostics_test.py | 160 ++++++++++++++++++ 10 files changed, 354 insertions(+), 22 deletions(-) create mode 100644 PyPIC3D/particles/tiled_particle_diagnostics.py create mode 100644 tests/tiled_particle_diagnostics_test.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 8200c1e6..4fcc1f1b 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -93,11 +93,7 @@ def run_PyPIC3D(config_file): scalar_field_names = ["rho", "mass_density"] vector_field_names = ["E", "B", "J"] tiled_run = simulation_parameters["solver"] == "tiled_yee" - - if tiled_run and plotting_parameters["plot_openpmd_particles"]: - raise ValueError("plot_openpmd_particles is not supported for tiled_yee particle storage") - if tiled_run and (plotting_parameters["plot_phasespace"] or plotting_parameters["plot_vtk_particles"]): - raise ValueError("particle phase-space and VTK particle output are not supported for tiled_yee particle storage") + particle_species_names = simulation_parameters.get("particle_species_names") E, B, J, rho, phi, external_fields, *rest = fields # unpack the fields @@ -146,7 +142,7 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_phasespace']: - write_particles_phase_space(particles, t, output_dir) + write_particles_phase_space(particles, t, output_dir, species_names=particle_species_names, world=world) @@ -178,11 +174,11 @@ def run_PyPIC3D(config_file): # Plot the vector fields in VTK format if plotting_parameters['plot_vtk_particles']: - plot_vtk_particles(particles, plot_num, output_dir) + plot_vtk_particles(particles, plot_num, output_dir, species_names=particle_species_names, world=world) # 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_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5", species_names=particle_species_names) # Write the particles in openPMD format if plotting_parameters['plot_openpmd_fields']: diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index 00054330..b6f24983 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -5,6 +5,8 @@ import numpy as np import importlib.metadata +from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species + def _ensure_openpmd_array(data, dtype=np.float64, squeeze=False): arr = np.asarray(data, dtype=dtype) if squeeze: @@ -124,7 +126,11 @@ 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_names=None, world=None): + particles = flatten_tiled_particles_by_species(particles, 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 @@ -249,7 +255,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_names=None): """ Write all particle data to an openPMD file for visualization in ParaView/VisIt. @@ -271,7 +277,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_names=species_names, world=world) # write the particle data to the iteration series.flush() series.close() diff --git a/PyPIC3D/diagnostics/plotting.py b/PyPIC3D/diagnostics/plotting.py index 1236fae8..e2ae747a 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -7,6 +7,8 @@ import jax from functools import partial +from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species + def plot_positions(particles, t, x_wind, y_wind, z_wind, path): """ Makes an interactive 3D plot of the positions of the particles using Plotly. @@ -49,7 +51,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_names=None, world=None): """ Write the phase space of the particles to a file. @@ -69,6 +71,10 @@ def write_particles_phase_space(particles, t, path): os.makedirs(f"{path}/data/phase_space/z") # Create directory if it doesn't exist + particles = flatten_tiled_particles_by_species(particles, 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() @@ -243,4 +249,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 index f7ff14bd..014d6e01 100644 --- a/PyPIC3D/diagnostics/vtk.py +++ b/PyPIC3D/diagnostics/vtk.py @@ -5,6 +5,8 @@ from pyevtk.hl import gridToVTK, pointsToVTK import jax.numpy as jnp +from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species + def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): """ Plot the fields in a 3D grid. @@ -40,7 +42,7 @@ def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): # plot the electric fields in the vtk file format -def plot_vtk_particles(particles, t, path): +def plot_vtk_particles(particles, t, path, species_names=None, world=None): """ Plot the particles in VTK format. @@ -55,6 +57,10 @@ def plot_vtk_particles(particles, t, path): if not os.path.exists(f"{path}/data/particles"): os.makedirs(f"{path}/data/particles") + particles = flatten_tiled_particles_by_species(particles, species_names=species_names, world=world) + # Tiled particle arrays are fixed-capacity by tile; VTK particle output only + # receives the active diagnostic particles. + particle_names = [species.get_name().replace(" ", "") for species in particles] for species in particles: name = species.get_name().replace(" ", "") @@ -228,4 +234,4 @@ def write_slice(field_slice, x1, x2, t, name, path, dt): 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 + # plot the slice of the field in the vtk file format diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 0894ad25..dcbdadc0 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -142,7 +142,7 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ ) for cells, tile_width in zip(grid_shape, tile_shape): if cells % tile_width != 0: - raise ValueError("tiled_yee requires particle_tile_nx/ny/nz to divide Nx/Ny/Nz exactly") + raise ValueError("tiled_yee requires the shared tile shape to divide Nx/Ny/Nz exactly") def default_parameters(): @@ -206,9 +206,9 @@ 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) - "particle_tile_nx": 1, # number of x cells per particle tile - "particle_tile_ny": 1, # number of y cells per particle tile - "particle_tile_nz": 1, # number of z cells per particle tile + "particle_tile_nx": 1, # number of x cells per shared field/particle tile + "particle_tile_ny": 1, # number of y cells per shared field/particle tile + "particle_tile_nz": 1, # number of z cells per shared field/particle 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 } @@ -379,6 +379,8 @@ def initialize_simulation(toml_file): particles = load_particles_from_toml(toml_file, simulation_parameters, world, constants) # load the particles from the configuration file + simulation_parameters["particle_species_names"] = tuple(species.get_name() for species in particles) + # keep species names available after tiled storage drops per-species objects for species in particles: name = species.get_name() diff --git a/PyPIC3D/particles/tiled_particle_diagnostics.py b/PyPIC3D/particles/tiled_particle_diagnostics.py new file mode 100644 index 00000000..d5e8e391 --- /dev/null +++ b/PyPIC3D/particles/tiled_particle_diagnostics.py @@ -0,0 +1,117 @@ +from typing import NamedTuple + +import jax.numpy as jnp + +from PyPIC3D.particles.tiled_particles import TiledParticles + + +class DiagnosticParticleSpecies(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 get_name(self): + return self.name + + def get_number_of_particles(self): + return self.x.shape[0] + + def get_forward_position(self): + return self.x[:, 0], self.x[:, 1], self.x[:, 2] + + def get_position(self): + return self.x_diagnostic[:, 0], self.x_diagnostic[:, 1], self.x_diagnostic[:, 2] + + def get_velocity(self): + return self.u[:, 0], self.u[:, 1], self.u[:, 2] + + def get_active_mask(self): + return jnp.ones(self.x.shape[0], dtype=bool) + + def get_charge(self): + return self.charge * self.weight + + def get_mass(self): + return self.mass * self.weight + + def get_weight(self): + return self.weight + + +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 flatten_tiled_particles_by_species(tiled_particles, species_names=None, world=None): + """ + Flatten active tile-major particles into diagnostic species arrays. + + This is an output adapter only: it removes inactive slots introduced by the + fixed-capacity tiled storage and preserves the stored particle position, + velocity, charge, mass, weight, and species index. + """ + + if not isinstance(tiled_particles, TiledParticles): + return tiled_particles + + n_species = tiled_particles.active.shape[3] + flattened_species = [] + + for species_index in range(n_species): + active = tiled_particles.active[:, :, :, species_index, :].reshape(-1) + + x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active] + u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active] + charge = tiled_particles.charge[:, :, :, species_index, :].reshape(-1)[active] + mass = tiled_particles.mass[:, :, :, species_index, :].reshape(-1)[active] + weight = tiled_particles.weight[:, :, :, species_index, :].reshape(-1)[active] + + x_diagnostic = _diagnostic_position(x, u, world) + + if species_names is None: + name = f"species_{species_index}" + else: + name = species_names[species_index] + + flattened_species.append( + DiagnosticParticleSpecies( + name=name, + species_index=species_index, + x=x, + x_diagnostic=x_diagnostic, + u=u, + charge=charge, + mass=mass, + weight=weight, + ) + ) + + return flattened_species diff --git a/PyPIC3D/particles/tiled_particles.py b/PyPIC3D/particles/tiled_particles.py index 3e7d26f0..03d7593a 100644 --- a/PyPIC3D/particles/tiled_particles.py +++ b/PyPIC3D/particles/tiled_particles.py @@ -3,6 +3,14 @@ from typing import NamedTuple 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) @@ -20,4 +28,4 @@ class TiledParticles(NamedTuple): update_u1: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) update_u2: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) update_u3: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - # boolean array indicating whether the particle's position or velocity should be updated or not \ No newline at end of file + # boolean array indicating whether the particle's position or velocity should be updated or not diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index dbc17ffa..90b13d46 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -11,13 +11,13 @@ def _tile_axis_count(n_cells, cells_per_tile): if int(n_cells) % int(cells_per_tile) != 0: - raise ValueError("Field tile sizes must divide the physical grid dimensions exactly.") + 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, tile_shape): """ - Split a ghost-celled field into compact field tiles with one-cell halos. + Split a ghost-celled field into compact tiles using the shared tile shape. """ tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] diff --git a/docs/architecture.rst b/docs/architecture.rst index 60948150..9962d4ee 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -45,6 +45,38 @@ Main runtime objects: boundary conditions. - ``constants``: physical constants and filter coefficients. +Tiled State Contract +-------------------- + +The ``tiled_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)``; scalar particle data +such as charge, mass, weight, active flags, and update masks have shape +``(ntx, nty, ntz, species, max_particles_per_tile)``. 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 -------------------- diff --git a/tests/tiled_particle_diagnostics_test.py b/tests/tiled_particle_diagnostics_test.py new file mode 100644 index 00000000..8368ad26 --- /dev/null +++ b/tests/tiled_particle_diagnostics_test.py @@ -0,0 +1,160 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles + + +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 = particle_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 = particle_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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + + flattened_species = flatten_tiled_particles_by_species(tiled_particles) + + self.assertEqual(len(flattened_species), len(species_list)) + for species_index, (original, flattened) in enumerate(zip(species_list, flattened_species)): + active = original.get_active_mask() + + original_x = jnp.stack(original.get_forward_position(), axis=-1)[active] + original_u = jnp.stack(original.get_velocity(), axis=-1)[active] + flattened_x = jnp.stack(flattened.get_forward_position(), axis=-1) + flattened_u = jnp.stack(flattened.get_velocity(), axis=-1) + + 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + + flattened_species = flatten_tiled_particles_by_species(tiled_particles) + + self.assertEqual(flattened_species[0].get_number_of_particles(), 3) + self.assertEqual(flattened_species[1].get_number_of_particles(), 2) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].get_forward_position()[0], -0.5))) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[1].get_forward_position()[0], -1.25))) + + def test_absorbed_particles_do_not_appear_in_flattened_output(self): + species_list = self._species() + tiled_particles = to_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 = flatten_tiled_particles_by_species(tiled_particles) + + self.assertEqual(flattened_species[0].get_number_of_particles(), 2) + self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].get_forward_position()[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 = to_tiled_particles(species_list, world, self._simulation_parameters()) + + flattened_species = flatten_tiled_particles_by_species(tiled_particles, world=world) + + for original, flattened in zip(species_list, flattened_species): + active = original.get_active_mask() + original_position = jnp.stack(original.get_position(), axis=-1)[active] + flattened_position = jnp.stack(flattened.get_position(), axis=-1) + + 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + species_names = tuple(species.get_name() for species in species_list) + + flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_names=species_names) + + self.assertEqual(flattened_species[0].get_name(), "ions") + self.assertEqual(flattened_species[1].get_name(), "electrons") + + +if __name__ == "__main__": + unittest.main() From f5749e64f8eec93273cba6993dc7f816227292cf Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 06:02:38 -0500 Subject: [PATCH 07/83] implmenting task 4 --- PyPIC3D/deposition/direct_deposition_tiled.py | 20 +- PyPIC3D/solvers/yee_tiled.py | 82 ++++++-- tests/direct_deposition_tiled_test.py | 186 +++++++++++++++++- tests/tiled_yee_integration_test.py | 70 ++++++- 4 files changed, 329 insertions(+), 29 deletions(-) diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index 946e829c..d40489b5 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -9,15 +9,15 @@ ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( - fold_tiled_vector_ghost_cells_periodic, - update_tiled_vector_ghost_cells_periodic, + fold_tiled_vector_ghost_cells, + update_tiled_vector_ghost_cells, ) from PyPIC3D.utils import bilinear_filter, digital_filter -def digital_filter_tiled_current_density(J_tiles, alpha): +def digital_filter_tiled_current_density(J_tiles, alpha, world): """ - Apply the global digital current filter to compact periodic current tiles. + Apply the global digital current filter to compact current tiles. Tile halos are refreshed before the stencil so every tile sees the same neighbor values that the assembled ghost-celled current would see. The @@ -25,14 +25,14 @@ def digital_filter_tiled_current_density(J_tiles, alpha): filtered guard-cell values. """ - J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) def filter_component(component_tiles): filter_tiles = jax.vmap(jax.vmap(jax.vmap(lambda tile: digital_filter(tile, alpha)))) return filter_tiles(component_tiles) J_tiles = tuple(filter_component(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) return J_tiles @@ -203,14 +203,14 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, tz, ) - J_tiles = fold_tiled_vector_ghost_cells_periodic((Jx_tiles, Jy_tiles, Jz_tiles)) - J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) if filter == "bilinear": apply_filter = jax.vmap(jax.vmap(jax.vmap(bilinear_filter))) J_tiles = tuple(apply_filter(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells_periodic(J_tiles) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) elif filter == "digital": - J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"]) + J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"], world) return J_tiles diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 90b13d46..d79b4052 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -263,39 +263,80 @@ def apply_tiled_conducting_bc(E_tiles, world): return Ex, Ey, Ez -def fold_tiled_ghost_cells_periodic(field_tiles): +def fold_tiled_ghost_cells(field_tiles, world): """ - Add tile-ghost deposits into neighboring tile interiors, then clear ghosts. + Add tile-ghost deposits into owning interiors, then clear ghosts. This is the current-deposition analogue of ``fold_ghost_cells`` for a - periodic tiled layout. Unlike field halo refresh, this operation is - additive: deposits in ghost cells belong to the adjacent tile's physical - interior. + tiled layout. 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 sign convention + used by the global ghost-cell fold. """ + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + x_lower_ghost = field_tiles[:, :, :, 0, :, :] + x_upper_ghost = field_tiles[:, :, :, -1, :, :] field_tiles = field_tiles.at[:, :, :, -2, :, :].add( - jnp.roll(field_tiles[:, :, :, 0, :, :], shift=-1, axis=0) + jnp.roll(x_lower_ghost, shift=-1, axis=0) ) field_tiles = field_tiles.at[:, :, :, 1, :, :].add( - jnp.roll(field_tiles[:, :, :, -1, :, :], shift=1, axis=0) + jnp.roll(x_upper_ghost, shift=1, axis=0) + ) + field_tiles = jax.lax.cond( + bc_x == BC_CONDUCTING, + lambda tiles: tiles + .at[-1, :, :, -2, :, :].add(-x_lower_ghost[0, :, :, :, :]) + .at[0, :, :, 1, :, :].add(-x_upper_ghost[-1, :, :, :, :]) + .at[0, :, :, 1, :, :].add(-x_lower_ghost[0, :, :, :, :]) + .at[-1, :, :, -2, :, :].add(-x_upper_ghost[-1, :, :, :, :]), + lambda tiles: tiles, + operand=field_tiles, ) field_tiles = field_tiles.at[:, :, :, 0, :, :].set(0.0) field_tiles = field_tiles.at[:, :, :, -1, :, :].set(0.0) + y_lower_ghost = field_tiles[:, :, :, :, 0, :] + y_upper_ghost = field_tiles[:, :, :, :, -1, :] field_tiles = field_tiles.at[:, :, :, :, -2, :].add( - jnp.roll(field_tiles[:, :, :, :, 0, :], shift=-1, axis=1) + jnp.roll(y_lower_ghost, shift=-1, axis=1) ) field_tiles = field_tiles.at[:, :, :, :, 1, :].add( - jnp.roll(field_tiles[:, :, :, :, -1, :], shift=1, axis=1) + jnp.roll(y_upper_ghost, shift=1, axis=1) + ) + field_tiles = jax.lax.cond( + bc_y == BC_CONDUCTING, + lambda tiles: tiles + .at[:, -1, :, :, -2, :].add(-y_lower_ghost[:, 0, :, :, :]) + .at[:, 0, :, :, 1, :].add(-y_upper_ghost[:, -1, :, :, :]) + .at[:, 0, :, :, 1, :].add(-y_lower_ghost[:, 0, :, :, :]) + .at[:, -1, :, :, -2, :].add(-y_upper_ghost[:, -1, :, :, :]), + lambda tiles: tiles, + operand=field_tiles, ) field_tiles = field_tiles.at[:, :, :, :, 0, :].set(0.0) field_tiles = field_tiles.at[:, :, :, :, -1, :].set(0.0) + z_lower_ghost = field_tiles[:, :, :, :, :, 0] + z_upper_ghost = field_tiles[:, :, :, :, :, -1] field_tiles = field_tiles.at[:, :, :, :, :, -2].add( - jnp.roll(field_tiles[:, :, :, :, :, 0], shift=-1, axis=2) + jnp.roll(z_lower_ghost, shift=-1, axis=2) ) field_tiles = field_tiles.at[:, :, :, :, :, 1].add( - jnp.roll(field_tiles[:, :, :, :, :, -1], shift=1, axis=2) + jnp.roll(z_upper_ghost, shift=1, axis=2) + ) + field_tiles = jax.lax.cond( + bc_z == BC_CONDUCTING, + lambda tiles: tiles + .at[:, :, -1, :, :, -2].add(-z_lower_ghost[:, :, 0, :, :]) + .at[:, :, 0, :, :, 1].add(-z_upper_ghost[:, :, -1, :, :]) + .at[:, :, 0, :, :, 1].add(-z_lower_ghost[:, :, 0, :, :]) + .at[:, :, -1, :, :, -2].add(-z_upper_ghost[:, :, -1, :, :]), + lambda tiles: tiles, + operand=field_tiles, ) field_tiles = field_tiles.at[:, :, :, :, :, 0].set(0.0) field_tiles = field_tiles.at[:, :, :, :, :, -1].set(0.0) @@ -303,9 +344,26 @@ def fold_tiled_ghost_cells_periodic(field_tiles): return field_tiles +def fold_tiled_ghost_cells_periodic(field_tiles): + """ + Fold all-periodic tile-ghost deposits for one scalar component. + """ + + world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} + return fold_tiled_ghost_cells(field_tiles, world) + + +def fold_tiled_vector_ghost_cells(field_tiles, world): + """ + Fold tile-ghost deposits for each vector component. + """ + + return tuple(fold_tiled_ghost_cells(component, world) for component in field_tiles) + + def fold_tiled_vector_ghost_cells_periodic(field_tiles): """ - Fold periodic tile-ghost deposits for each vector component. + Fold all-periodic tile-ghost deposits for each vector component. """ return tuple(fold_tiled_ghost_cells_periodic(component) for component in field_tiles) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 88b2726a..ab3f486c 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -3,7 +3,8 @@ import jax import jax.numpy as jnp -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells +from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_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.deposition.direct_deposition_tiled import ( digital_filter_tiled_current_density, @@ -13,6 +14,7 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, + fold_tiled_ghost_cells, fold_tiled_ghost_cells_periodic, tile_vector_field, ) @@ -23,8 +25,10 @@ class TestDirectDepositionTiled(unittest.TestCase): - def _build_world(self, Nx=8, Ny=6, Nz=4, dt=0.05): + 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, @@ -37,7 +41,7 @@ def _build_world(self, Nx=8, Ny=6, Nz=4, dt=0.05): "z_wind": z_wind, "dt": dt, "shape_factor": 1, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) world["grids"] = {"vertex": vertex_grid, "center": center_grid} @@ -107,7 +111,7 @@ def test_tiled_digital_filter_matches_global_digital_filter(self): 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_tiled_current_density(J_tiles, alpha) + filtered_tiles = digital_filter_tiled_current_density(J_tiles, alpha, world) 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) @@ -118,17 +122,88 @@ def test_tiled_digital_filter_matches_global_digital_filter(self): self.assertTrue(jnp.allclose(tiled_component, reference_component, rtol=1.0e-12, atol=1.0e-12)) def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self): + world = self._build_world(Nx=4, Ny=1, Nz=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 = fold_tiled_ghost_cells_periodic(tiles) + folded = fold_tiled_ghost_cells(tiles, world) 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)) + legacy_folded = fold_tiled_ghost_cells_periodic(tiles) + self.assertTrue(jnp.allclose(folded, legacy_folded, rtol=1.0e-12, atol=1.0e-12)) + + 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}, + ) + 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 = fold_tiled_ghost_cells(tiles, world) + + 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)) + + 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) + 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 = fold_tiled_ghost_cells(tiles, world) + folded_from_tiles = assemble_tiled_vector_field((folded_tiles, folded_tiles, folded_tiles), world, tile_shape)[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-12, atol=1.0e-12)) + def test_tiled_direct_deposition_returns_only_local_current_tiles(self): world = self._build_world() simulation_parameters = { @@ -343,6 +418,107 @@ def test_tiled_direct_deposition_periodic_boundary_crossing(self): self._compare_tiled_to_standard([species], 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, + } + species = particle_species( + name="conducting wall deposits", + N_particles=3, + charge=1.0, + mass=1.0, + weight=1.0, + T=1.0, + x1=jnp.array([ + -world["x_wind"] / 2 + 0.1 * world["dx"], + world["x_wind"] / 2 - 0.1 * world["dx"], + 0.0, + ]), + x2=jnp.array([ + -world["y_wind"] / 2 + 0.1 * world["dy"], + 0.0, + world["y_wind"] / 2 - 0.1 * world["dy"], + ]), + x3=jnp.array([ + 0.0, + -world["z_wind"] / 2 + 0.1 * world["dz"], + world["z_wind"] / 2 - 0.1 * world["dz"], + ]), + v1=jnp.array([0.5, -0.25, 0.15]), + v2=jnp.array([0.1, -0.2, 0.3]), + v3=jnp.array([-0.15, 0.35, -0.1]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + self._compare_tiled_to_standard([species], 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, + } + species = particle_species( + name="mixed wall deposits", + N_particles=4, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([ + -world["x_wind"] / 2 - 0.1 * world["dx"], + world["x_wind"] / 2 + 0.2 * world["dx"], + -0.5, + 0.5, + ]), + x2=jnp.array([ + -world["y_wind"] / 2 + 0.1 * world["dy"], + world["y_wind"] / 2 - 0.2 * world["dy"], + -0.25, + 0.25, + ]), + x3=jnp.array([ + 0.0, + 0.25, + -world["z_wind"] / 2 - 0.1 * world["dz"], + world["z_wind"] / 2 + 0.2 * world["dz"], + ]), + 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"], + ) + + self._compare_tiled_to_standard([species], 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 = { diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 9a83e459..76a2ed24 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -18,6 +18,7 @@ from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, tile_vector_field from PyPIC3D.utils import build_yee_grid, compute_energy +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC jax.config.update("jax_enable_x64", True) @@ -28,7 +29,9 @@ def unused_curl(Ex, Ey, Ez): class TestTiledYeeIntegration(unittest.TestCase): - def _build_world(self): + def _build_world(self, boundary_conditions=None): + if boundary_conditions is None: + boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} world = { "Nx": 8, "Ny": 1, @@ -41,7 +44,7 @@ def _build_world(self): "y_wind": 1.0, "z_wind": 1.0, "shape_factor": 1, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) world["grids"] = {"vertex": vertex_grid, "center": center_grid} @@ -344,6 +347,69 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): for reference, tiled in zip(reference_energy, tiled_energy): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_yee_step_matches_standard_conducting_yee_step(self): + world = self._build_world( + boundary_conditions={"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING} + ) + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + _, reference_fields = time_loop_electrodynamic( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="none"), + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + _, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=None, + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + reference_E, reference_B, reference_J, *_ = reference_fields + + for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, tile_shape)): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, tile_shape)): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape)): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_yee_step_uses_digital_filtered_current(self): world = self._build_world() constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} From d926da3209ea94e75b83b032e6c6d22c54b2ee94 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 06:33:00 -0500 Subject: [PATCH 08/83] implementing tasks 3 and 5 --- PyPIC3D/__main__.py | 6 +- PyPIC3D/deposition/rho_tiled.py | 188 ++++++++++++++++++++++++ PyPIC3D/diagnostics/fluid_quantities.py | 7 +- PyPIC3D/initialization.py | 4 +- tests/direct_deposition_tiled_test.py | 49 ++++++ tests/tiled_fluid_quantities_test.py | 127 ++++++++++++++++ tests/tiled_rho_test.py | 141 ++++++++++++++++++ tests/tiled_yee_integration_test.py | 117 +++++++++++++++ 8 files changed, 634 insertions(+), 5 deletions(-) create mode 100644 PyPIC3D/deposition/rho_tiled.py create mode 100644 tests/tiled_fluid_quantities_test.py create mode 100644 tests/tiled_rho_test.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 4fcc1f1b..3e25237a 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -41,6 +41,7 @@ ) from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field @@ -148,8 +149,9 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_scalars']: if tiled_run: - raise ValueError("plot_vtk_scalars is not supported for tiled_yee particle storage") - rho = compute_rho(particles, rho, world, constants) + rho = compute_rho_from_tiled_particles(particles, rho, world, constants) + else: + 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 diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py new file mode 100644 index 00000000..b3994684 --- /dev/null +++ b/PyPIC3D/deposition/rho_tiled.py @@ -0,0 +1,188 @@ +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.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 digital_filter + + +def _deposit_tiled_scalar_moment( + tiled_particles, + rho, + world, + grid, + scalar_weight, + half_step_back, + wrap_periodic_position=False, +): + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + dt = world["dt"] + 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"] + + x = tiled_particles.x[..., 0].reshape(-1) + y = tiled_particles.x[..., 1].reshape(-1) + z = tiled_particles.x[..., 2].reshape(-1) + vx = tiled_particles.u[..., 0].reshape(-1) + vy = tiled_particles.u[..., 1].reshape(-1) + vz = tiled_particles.u[..., 2].reshape(-1) + active = tiled_particles.active.reshape(-1).astype(x.dtype) + scalar_density = scalar_weight.reshape(-1) / (dx * dy * dz) + + if half_step_back: + x = x - vx * dt / 2 + y = y - vy * dt / 2 + z = z - vz * dt / 2 + # Rho deposits stay unwrapped so seam-crossing charge lands in ghost + # cells before folding. Scalar diagnostics that mirror + # species.get_position() wrap periodic half-step positions first. + if wrap_periodic_position: + x = jax.lax.cond( + bc_x == BC_PERIODIC, + lambda x_value: jnp.where( + x_value > world["x_wind"] / 2, + x_value - world["x_wind"], + jnp.where(x_value < -world["x_wind"] / 2, x_value + world["x_wind"], x_value), + ), + lambda x_value: x_value, + x, + ) + y = jax.lax.cond( + bc_y == BC_PERIODIC, + lambda y_value: jnp.where( + y_value > world["y_wind"] / 2, + y_value - world["y_wind"], + jnp.where(y_value < -world["y_wind"] / 2, y_value + world["y_wind"], y_value), + ), + lambda y_value: y_value, + y, + ) + z = jax.lax.cond( + bc_z == BC_PERIODIC, + lambda z_value: jnp.where( + z_value > world["z_wind"] / 2, + z_value - world["z_wind"], + jnp.where(z_value < -world["z_wind"] / 2, z_value + world["z_wind"], z_value), + ), + lambda z_value: z_value, + z, + ) + + 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, + ) + z, _, 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, + ) + + 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) + + rho = jnp.zeros_like(rho) + + 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 * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], + mode="drop", + ) + + 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_rho_from_tiled_particles(tiled_particles, rho, world, constants, grid=None): + """Compute charge density on the vertex grid from tile-major particles.""" + if grid is None: + grid = world["grids"]["vertex"] + + scalar_weight = tiled_particles.charge * tiled_particles.weight + rho = _deposit_tiled_scalar_moment( + tiled_particles, + rho, + world, + grid, + scalar_weight, + half_step_back=True, + wrap_periodic_position=False, + ) + + alpha = constants["alpha"] + rho = digital_filter(rho, alpha) + rho = update_ghost_cells( + rho, + world["boundary_conditions"]["x"], + world["boundary_conditions"]["y"], + world["boundary_conditions"]["z"], + ) + + return rho + + +@jit +def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid=None): + """Compute mass density on the vertex grid from tile-major particles.""" + if grid is None: + grid = world["grids"]["vertex"] + + scalar_weight = tiled_particles.mass * tiled_particles.weight + return _deposit_tiled_scalar_moment( + tiled_particles, + rho, + world, + grid, + scalar_weight, + half_step_back=True, + wrap_periodic_position=True, + ) diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index 5905bfda..b07ec22a 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -2,6 +2,8 @@ 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 +from PyPIC3D.deposition.rho_tiled import compute_mass_density_from_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles import jax import jax.numpy as jnp @@ -26,6 +28,9 @@ def compute_mass_density(particles, rho, world): Returns: ndarray: The updated charge density array. """ + if isinstance(particles, TiledParticles): + return compute_mass_density_from_tiled_particles(particles, rho, world) + dx = world['dx'] dy = world['dy'] dz = world['dz'] @@ -47,7 +52,7 @@ def compute_mass_density(particles, rho, world): 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 + # get the particle position used by the existing scalar diagnostic active = species.get_active_mask().astype(x.dtype) _, _, deltax, xpts = prepare_particle_axis_stencil( diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index dcbdadc0..ed2316d6 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -132,8 +132,8 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") if simulation_parameters["particle_pusher"] != "boris": raise ValueError("tiled_yee currently supports only particle_pusher='boris'") - if simulation_parameters["filter_j"] not in ("none", "digital"): - raise ValueError("tiled_yee currently supports only filter_j='none' or filter_j='digital'") + if simulation_parameters["filter_j"] not in ("none", "digital", "bilinear"): + raise ValueError("tiled_yee currently supports only filter_j='none', filter_j='digital', or filter_j='bilinear'") tile_shape = _tile_shape_from_parameters(simulation_parameters) grid_shape = ( int(simulation_parameters["Nx"]), diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index ab3f486c..dbac1d8f 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -355,6 +355,55 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): 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)) + 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} + species = particle_species( + name="bilinear filtered current", + N_particles=6, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.55, -0.52, -0.03, 0.49, 0.55, 1.45]), + x2=jnp.array([-1.10, -0.55, -0.03, 0.02, 0.52, 1.05]), + x3=jnp.array([-0.70, -0.04, 0.03, 0.31, 0.49, 0.72]), + v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16]), + v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12]), + v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19]), + 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_particles = to_tiled_particles([species], world, simulation_parameters) + tile_shape = ( + simulation_parameters["particle_tile_nx"], + simulation_parameters["particle_tile_ny"], + simulation_parameters["particle_tile_nz"], + ) + + J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") + J_tiles = direct_J_from_tiled_particles( + tiled_particles, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="bilinear", + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + 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)) + def test_tiled_direct_deposition_respects_active_mask(self): world = self._build_world() simulation_parameters = { diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py new file mode 100644 index 00000000..099fc60f --- /dev/null +++ b/tests/tiled_fluid_quantities_test.py @@ -0,0 +1,127 @@ +import inspect +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D import __main__ as pypic_main +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.diagnostics.fluid_quantities import compute_mass_density +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.utils import build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +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, + "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 _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 = particle_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 = particle_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_mass_density_matches_compute_mass_density_on_interior(self): + world = self._build_world(shape_factor=2) + particles = self._particles(world) + tiled_particles = to_tiled_particles(particles, world, self._simulation_parameters()) + + inactive = ~tiled_particles.active + x = tiled_particles.x.at[inactive, 0].set(0.10) + x = x.at[inactive, 1].set(0.20) + x = x.at[inactive, 2].set(-0.30) + mass = tiled_particles.mass.at[inactive].set(1000.0) + weight = tiled_particles.weight.at[inactive].set(1000.0) + tiled_particles = tiled_particles._replace(x=x, mass=mass, weight=weight) + + mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) + mass_tiled = compute_mass_density(tiled_particles, self._empty_scalar(world), world) + + self.assertTrue( + jnp.allclose( + mass_tiled[1:-1, 1:-1, 1:-1], + mass_reference[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tiled_scalar_vtk_path_no_longer_rejects_tiled_yee(self): + run_source = inspect.getsource(pypic_main.run_PyPIC3D) + + self.assertNotIn("plot_vtk_scalars is not supported for tiled_yee", run_source) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py new file mode 100644 index 00000000..01525fc1 --- /dev/null +++ b/tests/tiled_rho_test.py @@ -0,0 +1,141 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.utils import build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +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, + "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 _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 = particle_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 = particle_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 = to_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) + charge = tiled_particles.charge.at[inactive].set(1000.0) + weight = tiled_particles.weight.at[inactive].set(1000.0) + + return tiled_particles._replace(x=x, u=u, charge=charge, weight=weight) + + def _compare_tiled_to_standard(self, shape_factor, alpha): + world = self._build_world(shape_factor) + constants = {"alpha": alpha} + particles = self._particles(world) + tiled_particles = self._tiled_with_noisy_inactive_slots(particles, world) + + rho_reference = compute_rho(particles, self._empty_scalar(world), world, constants) + rho_tiled = compute_rho_from_tiled_particles( + tiled_particles, + self._empty_scalar(world), + world, + constants, + ) + + self.assertTrue( + jnp.allclose( + rho_tiled[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) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 76a2ed24..fb76d2a0 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -210,6 +210,65 @@ def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(self): self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(simulation_parameters["filter_j"], "digital") + def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(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 bilinear filter init smoke", + "output_dir": tmpdir, + "solver": "tiled_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": "j_from_rhov", + "filter_j": "bilinear", + "fast_backend": "default", + "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_bilinear.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["filter_j"], "bilinear") + def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(self): with tempfile.TemporaryDirectory() as tmpdir: x_path = os.path.join(tmpdir, "x.npy") @@ -468,6 +527,64 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): for reference, tiled in zip(reference_J, J_from_tiles): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_yee_step_uses_bilinear_filtered_current(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + _, reference_fields = time_loop_electrodynamic( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="bilinear"), + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + _, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + _, _, reference_J, *_ = reference_fields + _, _, J_tiles, *_ = tiled_fields + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + + for reference, tiled in zip(reference_J, J_from_tiles): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + if __name__ == "__main__": unittest.main() From abb72584f9dc1381609dd342eb95f5b0b0d334a2 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 10:02:41 -0500 Subject: [PATCH 09/83] implement task 7 --- PyPIC3D/electrodynamic_tiled.py | 5 ++- PyPIC3D/initialization.py | 4 +- PyPIC3D/pusher/tiled_pusher.py | 45 +++++++++++++------- tests/tiled_pusher_test.py | 66 +++++++++++++++++++++++++++++ tests/tiled_yee_integration_test.py | 59 ++++++++++++++++++++++++++ 5 files changed, 160 insertions(+), 19 deletions(-) diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index fe4f81af..2fff8f8f 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -26,7 +26,7 @@ def time_loop_electrodynamic_tiled( Advance a periodic tiled electrodynamic PIC system by one time step. """ - del solver, particle_pusher + del solver if len(fields) == 7: E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields @@ -66,8 +66,9 @@ def time_loop_electrodynamic_tiled( constants, tile_shape, relativistic=relativistic, + particle_pusher=particle_pusher, ) - # use the tiled Boris pusher for particle velocities + # use the selected tiled pusher for particle velocities particles = update_tiled_particle_positions(particles, world) # update particle forward positions before current deposition diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index ed2316d6..f2a926b4 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -130,8 +130,8 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ raise ValueError("tiled_yee is only supported for electrodynamic simulations") if simulation_parameters["current_calculation"] != "j_from_rhov": raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") - if simulation_parameters["particle_pusher"] != "boris": - raise ValueError("tiled_yee currently supports only particle_pusher='boris'") + if simulation_parameters["particle_pusher"] not in ("boris", "higuera_cary"): + raise ValueError("tiled_yee currently supports only particle_pusher='boris' or 'higuera_cary'") if simulation_parameters["filter_j"] not in ("none", "digital", "bilinear"): raise ValueError("tiled_yee currently supports only filter_j='none', filter_j='digital', or filter_j='bilinear'") tile_shape = _tile_shape_from_parameters(simulation_parameters) diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index fc5b6e08..e646c553 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -9,6 +9,7 @@ interpolate_field_to_particles, relativistic_boris_single_particle, ) +from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle def _tile_axis(axis, tile_index, cells_per_tile): @@ -16,10 +17,10 @@ def _tile_axis(axis, tile_index, cells_per_tile): return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) -@partial(jit, static_argnames=("tile_shape", "relativistic")) -def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, relativistic=True): +@partial(jit, static_argnames=("tile_shape", "relativistic", "particle_pusher")) +def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, relativistic=True, particle_pusher="boris"): """ - Push tile-major particles with Boris using compact tiled Yee fields. + 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 one-cell field halos on each compact tile provide the @@ -79,23 +80,37 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til relativistic_boris_single_particle, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), ) + higuera_cary_vmap = jax.vmap( + higuera_cary_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) - new_vx, new_vy, new_vz = jax.lax.cond( - relativistic, - lambda _: relativistic_boris_vmap( + 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, - ), - lambda _: boris_vmap( - vx, vy, vz, - efield_atx, efield_aty, efield_atz, - bfield_atx, bfield_aty, bfield_atz, - q, m, dt, constants, - ), - operand=None, - ) + ) + else: + raise ValueError(f"Unknown particle_pusher: {particle_pusher}") active = active_tile.reshape(-1) update_u1 = update_u1.reshape(-1) diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 84810eef..f2f4520b 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -130,6 +130,72 @@ def test_tiled_particle_push_matches_flat_boris(self): self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_particle_push_matches_flat_higuera_cary(self): + world = self._build_world() + constants = {"C": 10.0} + tile_shape = (2, 3, 2) + 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=0.25) + B = self._deterministic_vector_field(world, scale=0.05) + + def make_species(): + return particle_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 = particle_push( + reference, + E, + B, + world["grids"]["center"], + world["grids"]["vertex"], + world, + constants, + particle_pusher="higuera_cary", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + pushed = tiled_particle_push( + tiled_particles, + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + world, + constants, + tile_shape, + particle_pusher="higuera_cary", + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + flat_u = jnp.stack(reference.get_velocity(), axis=1) + flat_x = jnp.stack(reference.get_forward_position(), axis=1) + order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + + self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_particle_push_respects_active_and_update_flags(self): world = self._build_world() constants = {"C": 10.0} diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index fb76d2a0..df8b72c1 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -151,6 +151,65 @@ def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(sel self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(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 higuera cary init smoke", + "output_dir": tmpdir, + "solver": "tiled_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": "j_from_rhov", + "filter_j": "none", + "fast_backend": "default", + "particle_pusher": "higuera_cary", + "relativistic": True, + }, + "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_higuera_cary.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["particle_pusher"], "higuera_cary") + def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(self): with tempfile.TemporaryDirectory() as tmpdir: x_path = os.path.join(tmpdir, "x.npy") From 17771f4da6e3c039b0068efe03880d1451f32cfe Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 10:38:54 -0500 Subject: [PATCH 10/83] implemented task 8 and removed spectral solvers for leaner codebase --- PyPIC3D/__init__.py | 1 - PyPIC3D/__main__.py | 12 +- PyPIC3D/deposition/rho_tiled.py | 168 +++++++++ PyPIC3D/electrostatic_tiled.py | 85 +++++ PyPIC3D/errors.py | 24 +- PyPIC3D/initialization.py | 64 ++-- PyPIC3D/solvers/electrostatic_yee.py | 143 ++++---- PyPIC3D/solvers/pstd.py | 171 --------- PyPIC3D/solvers/yee_tiled.py | 25 +- .../solver_convergence_test.py | 246 ------------- docs/solvers.rst | 6 +- tests/__init__.py | 1 - tests/electrostatic_yee_test.py | 46 --- tests/errors_test.py | 8 +- tests/initialization_test.py | 6 +- tests/pml_test.py | 4 - tests/pstd_test.py | 329 ------------------ tests/tiled_electrostatic_test.py | 281 +++++++++++++++ 18 files changed, 689 insertions(+), 931 deletions(-) create mode 100644 PyPIC3D/electrostatic_tiled.py delete mode 100644 PyPIC3D/solvers/pstd.py delete mode 100644 tests/pstd_test.py create mode 100644 tests/tiled_electrostatic_test.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index e325f2a7..eb950691 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -13,7 +13,6 @@ from . import pusher from . import evolve -from .solvers import pstd from .solvers import fdtd from .solvers import first_order_yee from .solvers import electrostatic_yee diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 3e25237a..5a39cc83 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -43,7 +43,7 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, assemble_tiled_vector_field # Importing functions from the PyPIC3D package @@ -54,13 +54,21 @@ def _fields_for_output(fields, world, simulation_parameters): if simulation_parameters["solver"] != "tiled_yee": return fields - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, *rest = fields + if len(fields) == 6: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields = fields + pml_state = None + else: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, *rest = fields tile_shape = tuple(int(width) for width in world["tile_shape"]) external_E, external_B = external_fields E = assemble_tiled_vector_field(E_tiles, world, tile_shape) B = assemble_tiled_vector_field(B_tiles, world, tile_shape) J = assemble_tiled_vector_field(J_tiles, world, tile_shape) + if getattr(rho, "ndim", 0) == 6: + rho = assemble_tiled_scalar_field(rho, world, tile_shape) + if getattr(phi, "ndim", 0) == 6: + phi = assemble_tiled_scalar_field(phi, world, tile_shape) external_fields = ( assemble_tiled_vector_field(external_E, world, tile_shape), assemble_tiled_vector_field(external_B, world, tile_shape), diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index b3994684..c45ecd98 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -9,6 +9,10 @@ ) 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.solvers.yee_tiled import ( + fold_tiled_ghost_cells, + update_tiled_ghost_cells, +) from PyPIC3D.utils import digital_filter @@ -141,6 +145,147 @@ def _deposit_tiled_scalar_moment( return rho +def _tile_axis(axis, tile_index, cells_per_tile): + start = tile_index * cells_per_tile + return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) + + +def _deposit_tiled_scalar_moment_to_tiles( + tiled_particles, + rho_tiles, + world, + grid, + scalar_weight, + half_step_back, +): + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + dt = world["dt"] + 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 = rho_tiles.shape[3] - 2 + tile_ny = rho_tiles.shape[4] - 2 + tile_nz = rho_tiles.shape[5] - 2 + ntx, nty, ntz = rho_tiles.shape[:3] + + def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile): + local_grid = ( + _tile_axis(grid[0], tx, tile_nx), + _tile_axis(grid[1], ty, tile_ny), + _tile_axis(grid[2], tz, tile_nz), + ) + + 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) + active = active_tile.reshape(-1).astype(x.dtype) + scalar_density = scalar_weight_tile.reshape(-1) / (dx * dy * dz) + + if half_step_back: + x = x - vx * dt / 2 + y = y - vy * dt / 2 + z = z - vz * dt / 2 + + x, _, deltax, xpts = prepare_particle_axis_stencil( + x, + local_grid[0], + tile_nx + 2, + shape_factor, + bc_x, + wind=world["x_wind"], + ghost_cells=True, + ) + y, _, deltay, ypts = prepare_particle_axis_stencil( + y, + local_grid[1], + tile_ny + 2, + shape_factor, + bc_y, + wind=world["y_wind"], + ghost_cells=True, + ) + z, _, deltaz, zpts = prepare_particle_axis_stencil( + z, + local_grid[2], + tile_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, tile_nx + 2, ghost_cells=True) + ypts, y_weights = collapse_axis_stencil(ypts, y_weights, tile_ny + 2, ghost_cells=True) + zpts, z_weights = collapse_axis_stencil(zpts, z_weights, tile_nz + 2, ghost_cells=True) + + rho_tile = jnp.zeros((tile_nx + 2, tile_ny + 2, tile_nz + 2), dtype=rho_tiles.dtype) + + for i in range(xpts.shape[0]): + for j in range(ypts.shape[0]): + for k in range(zpts.shape[0]): + rho_tile = rho_tile.at[xpts[i], ypts[j], zpts[k]].add( + active * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], + mode="drop", + ) + + return rho_tile + + tx = jnp.arange(ntx) + ty = jnp.arange(nty) + tz = jnp.arange(ntz) + deposit_tiles = deposit_one_tile + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(None, None, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(None, 0, None, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, None, None, 0, 0, 0, 0), out_axes=0) + + rho_tiles = deposit_tiles( + tx, + ty, + tz, + tiled_particles.x, + tiled_particles.u, + tiled_particles.active, + scalar_weight, + ) + + rho_tiles = fold_tiled_ghost_cells(rho_tiles, world) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world) + + return rho_tiles + + +def _digital_filter_tiled_scalar(field_tiles, alpha, world): + field_tiles = update_tiled_ghost_cells(field_tiles, world) + + filter_tiles = digital_filter + filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) + filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) + filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) + + return filter_tiles(field_tiles, alpha) + + @jit def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, grid=None): """Compute charge density on the vertex grid from tile-major particles.""" @@ -170,6 +315,29 @@ def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, gri return rho +@jit +def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, grid=None): + """Compute charge density into tile-major vertex-grid scalar arrays.""" + if grid is None: + grid = world["grids"]["vertex"] + + scalar_weight = tiled_particles.charge * tiled_particles.weight + rho_tiles = _deposit_tiled_scalar_moment_to_tiles( + tiled_particles, + rho_tiles, + world, + grid, + scalar_weight, + half_step_back=True, + ) + + alpha = constants["alpha"] + rho_tiles = _digital_filter_tiled_scalar(rho_tiles, alpha, world) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world) + + return rho_tiles + + @jit def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid=None): """Compute mass density on the vertex grid from tile-major particles.""" diff --git a/PyPIC3D/electrostatic_tiled.py b/PyPIC3D/electrostatic_tiled.py new file mode 100644 index 00000000..e8bef6dc --- /dev/null +++ b/PyPIC3D/electrostatic_tiled.py @@ -0,0 +1,85 @@ +import jax.numpy as jnp + +from PyPIC3D.particles.tiled_particle_refresh import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) +from PyPIC3D.pusher.tiled_pusher import tiled_particle_push +from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields +from PyPIC3D.utils import add_external_fields + + +def time_loop_electrostatic_tiled( + particles, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=None, + relativistic=True, + particle_pusher="boris", +): + """ + Advance a tiled electrostatic PIC system by one time step. + + The particle push and retile use tile-local Yee 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. + """ + + del curl_func, J_func + + if len(fields) == 6: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields = fields + pml_state = None + overflow_previous = jnp.asarray(False) + elif len(fields) == 7: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state = fields + overflow_previous = jnp.asarray(False) + else: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields + + if tile_shape is None: + tile_shape = tuple(int(width) - 2 for width in E_tiles[0].shape[3:]) + + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) + # particles see evolved fields plus prescribed external fields + + particles = tiled_particle_push( + particles, + push_E_tiles, + push_B_tiles, + world, + constants, + tile_shape, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + # push velocities using the selected tiled particle pusher + + particles = update_tiled_particle_positions(particles, world) + # 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, + constants, + rho_tiles, + phi_tiles, + solver, + "periodic", + tile_shape, + ) + # 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/errors.py b/PyPIC3D/errors.py index 332d6a91..6f9d5689 100644 --- a/PyPIC3D/errors.py +++ b/PyPIC3D/errors.py @@ -3,7 +3,6 @@ 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 @@ -28,12 +27,9 @@ def compute_pe(phi, rho, constants, world, solver, bc='periodic'): 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': + if solver == 'autodiff': return 0 + x = centered_finite_difference_laplacian(phi, dx, dy, dz, bc) poisson_error = x + rho/eps magnitude = jnp.mean(jnp.abs(rho/eps)) + 1e-16 return jnp.mean(jnp.abs(poisson_error)) / magnitude @@ -59,12 +55,9 @@ def compute_magnetic_divergence_error(Bx, By, Bz, world, solver, bc='periodic'): 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': + if solver == 'autodiff': return 0 + divB = centered_finite_difference_divergence(Bx, By, Bz, dx, dy, dz, bc) divergence_error = jnp.mean(jnp.abs(divB)) return divergence_error @@ -92,12 +85,9 @@ def compute_electric_divergence_error(Ex, Ey, Ez, rho, constants, world, solver, 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': + if solver == 'autodiff': return 0 + divE = centered_finite_difference_divergence(Ex, Ey, Ez, dx, dy, dz, bc) divergence_error = jnp.mean(jnp.abs(divE - rho / eps)) - return divergence_error \ No newline at end of file + return divergence_error diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index f2a926b4..24d539c3 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -22,10 +22,6 @@ make_dir, compute_energy, add_external_fields ) -from PyPIC3D.solvers.pstd import ( - spectral_curl -) - from PyPIC3D.solvers.fdtd import ( centered_finite_difference_curl ) @@ -48,7 +44,8 @@ ) from PyPIC3D.solvers.yee_tiled import ( - tile_vector_field + tile_scalar_field, + tile_vector_field, ) @@ -60,6 +57,10 @@ time_loop_electrodynamic_tiled ) +from PyPIC3D.electrostatic_tiled import ( + time_loop_electrostatic_tiled +) + from PyPIC3D.pusher.particle_push import validate_particle_pusher from PyPIC3D.deposition.Esirkepov import Esirkepov_current @@ -108,9 +109,9 @@ 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 = ("fdtd", "spectral", "tiled_yee") + supported_solvers = ("fdtd", "tiled_yee") if solver not in supported_solvers: - raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd', 'spectral', or 'tiled_yee'.") + raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd' or 'tiled_yee'.") def _tile_shape_from_parameters(simulation_parameters): @@ -126,8 +127,6 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ Keep the first tile-native PIC path tied to the kernels that exist today. """ - if electrostatic: - raise ValueError("tiled_yee is only supported for electrodynamic simulations") if simulation_parameters["current_calculation"] != "j_from_rhov": raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") if simulation_parameters["particle_pusher"] not in ("boris", "higuera_cary"): @@ -179,7 +178,7 @@ def default_parameters(): simulation_parameters = { "name": "Default Simulation", "output_dir": os.getcwd(), - "solver": "fdtd", # solver: spectral, fdtd + "solver": "fdtd", # solver: fdtd or tiled_yee "fast_backend": "flat", # flat | default (flat when compatible, else fallback) "particle_x_bc": "periodic", # particle x boundary conditions: periodic, absorbing, reflecting "particle_y_bc": "periodic", # particle y boundary conditions: periodic, absorbing, reflecting @@ -257,7 +256,7 @@ 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". + solver (str): Solver type, either "fdtd" or "tiled_yee". electrostatic (bool): Flag indicating if the simulation is electrostatic. verbose (bool): Flag indicating if verbose output is enabled. GPUs (int): Number of GPUs to use. @@ -464,10 +463,7 @@ def initialize_simulation(toml_file): 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") + curl_func = functools.partial(centered_finite_difference_curl, dx=dx, dy=dy, dz=dz, bc="periodic") ######################### COMPUTE INITIAL ENERGY ######################################################## @@ -490,7 +486,10 @@ def initialize_simulation(toml_file): if electrostatic: print("Using electrostatic solver") - evolve_loop = time_loop_electrostatic + if solver == "tiled_yee": + evolve_loop = time_loop_electrostatic_tiled + else: + evolve_loop = time_loop_electrostatic else: print(f"Using electrodynamic solver with: {solver}") @@ -512,7 +511,26 @@ def initialize_simulation(toml_file): J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) + if solver == "tiled_yee": + tile_shape = _tile_shape_from_parameters(simulation_parameters) + world["tile_shape"] = tile_shape + particles = to_tiled_particles(particles, world, simulation_parameters) + E = tile_vector_field(E, world, tile_shape) + B = tile_vector_field(B, world, tile_shape) + J = tile_vector_field(J, world, tile_shape) + external_E, external_B = external_fields + external_fields = ( + tile_vector_field(external_E, world, tile_shape), + tile_vector_field(external_B, world, tile_shape), + ) + simulation_parameters["fast_backend"] = "tiled" + simulation_parameters["tile_shape"] = tile_shape + print(f"Using tiled Yee storage with tile shape: {tile_shape}") + if electrostatic: + if solver == "tiled_yee": + rho = tile_scalar_field(rho, world, tile_shape) + phi = tile_scalar_field(phi, world, tile_shape) fields = (E, B, J, rho, phi, external_fields) # define the fields tuple for the electrostatic solver else: @@ -520,22 +538,8 @@ def initialize_simulation(toml_file): # electrodynamic FDTD always carries the PML state slot; None means # ordinary, unstretched Yee derivatives. if solver == "tiled_yee": - tile_shape = _tile_shape_from_parameters(simulation_parameters) - world["tile_shape"] = tile_shape if pml_active: pml_state = initialize_tiled_pml_state(world, tile_shape) - particles = to_tiled_particles(particles, world, simulation_parameters) - E = tile_vector_field(E, world, tile_shape) - B = tile_vector_field(B, world, tile_shape) - J = tile_vector_field(J, world, tile_shape) - external_E, external_B = external_fields - external_fields = ( - tile_vector_field(external_E, world, tile_shape), - tile_vector_field(external_B, world, tile_shape), - ) - simulation_parameters["fast_backend"] = "tiled" - simulation_parameters["tile_shape"] = tile_shape - print(f"Using tiled Yee storage with tile shape: {tile_shape}") elif pml_active: pml_state = initialize_pml_state(world) fields = (E, B, J, rho, phi, external_fields, pml_state) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 9e3d8434..9398b7cd 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -4,64 +4,18 @@ from functools import partial from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles 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 ) - - -@jit -def solve_poisson_with_fft(rho, constants, world): - """ - Solve Poisson's equation in a periodic domain using FFTs. - - 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``. - - 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 - - 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') - - 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 - - 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 - - return phi +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_scalar_field, + tile_scalar_field, + update_tiled_ghost_cells, + update_tiled_vector_ghost_cells, +) @partial(jit, static_argnames=("tol", "max_iter")) @@ -183,12 +137,7 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver 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) # refresh ghost cells before any stencil-based post-processing @@ -199,12 +148,7 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) # 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, - ) + Ex, Ey, Ez = centered_finite_difference_gradient(-1 * phi[1:-1, 1:-1, 1:-1], dx, dy, dz, bc) # compute gradient on the interior # Place the gradient results into ghost-celled arrays @@ -219,3 +163,72 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver Ez_full = update_ghost_cells(Ez_full, bc_x, bc_y, bc_z) return (Ex_full, Ey_full, Ez_full), phi, rho + + +def _centered_tiled_electrostatic_gradient(phi_tiles, world): + """ + 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"] + + phi_tiles = update_tiled_ghost_cells(phi_tiles, world) + + Ex = jnp.zeros_like(phi_tiles) + Ey = jnp.zeros_like(phi_tiles) + Ez = jnp.zeros_like(phi_tiles) + + Ex = Ex.at[:, :, :, 1:-1, 1:-1, 1:-1].set( + -1.0 * (phi_tiles[:, :, :, 2:, 1:-1, 1:-1] - phi_tiles[:, :, :, :-2, 1:-1, 1:-1]) / (2.0 * dx) + ) + Ey = Ey.at[:, :, :, 1:-1, 1:-1, 1:-1].set( + -1.0 * (phi_tiles[:, :, :, 1:-1, 2:, 1:-1] - phi_tiles[:, :, :, 1:-1, :-2, 1:-1]) / (2.0 * dy) + ) + Ez = Ez.at[:, :, :, 1:-1, 1:-1, 1:-1].set( + -1.0 * (phi_tiles[:, :, :, 1:-1, 1:-1, 2:] - phi_tiles[:, :, :, 1:-1, 1:-1, :-2]) / (2.0 * dz) + ) + + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + + +def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, phi_tiles, solver, bc, tile_shape): + """ + Compute electrostatic fields from tiled rho deposition and a global Poisson solve. + + Rho is deposited into tile-major scalar storage, assembled for the existing + Poisson solver, and the solved global potential is tiled again before the + electric field is differentiated on tile-local halos. + """ + + del bc + + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + + rho_tiles = compute_tiled_rho_from_tiled_particles(particles, rho_tiles, world, constants) + rho = assemble_tiled_scalar_field(rho_tiles, world, tile_shape) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape) + + phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) + + phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) + # refresh ghost cells before filtering and tiled differentiation + + 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) + # keep the same global phi post-processing order as the untiled solver + + phi_tiles = tile_scalar_field(phi, world, tile_shape) + E_tiles = _centered_tiled_electrostatic_gradient(phi_tiles, world) + + return E_tiles, phi_tiles, rho_tiles 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/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index d79b4052..9309d35a 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -15,11 +15,13 @@ def _tile_axis_count(n_cells, cells_per_tile): return int(n_cells) // int(cells_per_tile) -def _tile_scalar_field(field, tile_shape): +def tile_scalar_field(field, world, tile_shape): """ Split a ghost-celled field into compact tiles using the shared tile shape. """ + del world + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] Nx = int(field.shape[0]) - 2 Ny = int(field.shape[1]) - 2 @@ -48,6 +50,10 @@ def tile_at(tx, ty, tz): ) +def _tile_scalar_field(field, tile_shape): + return tile_scalar_field(field, None, tile_shape) + + def tile_vector_field(field, world, tile_shape): """ Split ``(Fx, Fy, Fz)`` into compact ghost-celled tiles. @@ -57,20 +63,19 @@ def tile_vector_field(field, world, tile_shape): ``(tile_nx + 2, tile_ny + 2, tile_nz + 2)``. """ - del world - return tuple(_tile_scalar_field(component, tile_shape) for component in field) + return tuple(tile_scalar_field(component, world, tile_shape) for component in field) -def _assemble_scalar_tiles(field_tiles, world, tile_shape): +def assemble_tiled_scalar_field(field_tiles, world, tile_shape): """ Assemble compact field tiles back into one global ghost-celled field. """ tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] ntx, nty, ntz = field_tiles.shape[:3] - Nx = int(world["Nx"]) - Ny = int(world["Ny"]) - Nz = int(world["Nz"]) + 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) @@ -94,7 +99,11 @@ def assemble_tiled_vector_field(field_tiles, world, tile_shape): Assemble tiled vector-field components into ordinary ghost-celled arrays. """ - return tuple(_assemble_scalar_tiles(component, world, tile_shape) for component in field_tiles) + return tuple(assemble_tiled_scalar_field(component, world, tile_shape) for component in field_tiles) + + +def _assemble_scalar_tiles(field_tiles, world, tile_shape): + return assemble_tiled_scalar_field(field_tiles, world, tile_shape) def update_tiled_ghost_cells(field_tiles, world): diff --git a/demos/convergence_testing/solver_convergence_test.py b/demos/convergence_testing/solver_convergence_test.py index 632243a9..a3d55203 100644 --- a/demos/convergence_testing/solver_convergence_test.py +++ b/demos/convergence_testing/solver_convergence_test.py @@ -6,11 +6,6 @@ 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 ) @@ -229,216 +224,6 @@ def fdtd_curl_comparison(nx): 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 ########################## @@ -480,35 +265,4 @@ def pstd_curl_comparison(nx): ############################################################################ - #################### 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/docs/solvers.rst b/docs/solvers.rst index 764dfc16..1de32a71 100644 --- a/docs/solvers.rst +++ b/docs/solvers.rst @@ -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"``: conjugate-gradient Poisson + centered finite-difference gradient +- conjugate-gradient Poisson solve for ``phi`` +- centered finite-difference gradient for ``E = -grad(phi)`` Electrodynamic mode ^^^^^^^^^^^^^^^^^^^ diff --git a/tests/__init__.py b/tests/__init__.py index 316da698..fa27bf25 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -12,5 +12,4 @@ from . import first_order_yee_test from . import initialization_test from . import particle_test -from . import pstd_test from . import utils_test diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index d688320c..d7111ce7 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -5,12 +5,10 @@ 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 @@ -96,25 +94,6 @@ def setUp(self): ) ] - 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) @@ -146,30 +125,6 @@ def test_solve_poisson_with_conjugate_gradient_single_mode(self): 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, @@ -194,7 +149,6 @@ def test_solver_mode_mapping(self): "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)) diff --git a/tests/errors_test.py b/tests/errors_test.py index c281e548..db20162f 100644 --- a/tests/errors_test.py +++ b/tests/errors_test.py @@ -32,10 +32,8 @@ def test_compute_pe(self): 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) @@ -49,10 +47,8 @@ def test_compute_magnetic_divergence_error(self): 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): @@ -67,10 +63,8 @@ def test_compute_electric_divergence_error(self): 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 + unittest.main() diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 6a23b4c9..ec33302f 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -5,7 +5,7 @@ import toml import jax import jax.numpy as jnp -from PyPIC3D.initialization import setup_write_dir, default_parameters, initialize_simulation +from PyPIC3D.initialization import setup_write_dir, default_parameters, initialize_simulation, validate_field_solver jax.config.update("jax_enable_x64", True) @@ -104,5 +104,9 @@ def test_initialize_simulation_rejects_unknown_solver(self): with self.assertRaisesRegex(ValueError, "Unsupported solver"): initialize_simulation(config) + def test_validate_field_solver_rejects_spectral(self): + with self.assertRaisesRegex(ValueError, "Unsupported solver"): + validate_field_solver("spectral") + if __name__ == '__main__': unittest.main() diff --git a/tests/pml_test.py b/tests/pml_test.py index 4cbce570..6b682c37 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -199,10 +199,6 @@ def test_tile_pml_profiles_matches_global_profiles_on_tile_interiors(self): 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, electrostatic=True, pml=pml)) 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/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py new file mode 100644 index 00000000..5683f7e8 --- /dev/null +++ b/tests/tiled_electrostatic_test.py @@ -0,0 +1,281 @@ +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_PERIODIC +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles +from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled +from PyPIC3D.evolve import time_loop_electrostatic +from PyPIC3D.initialization import initialize_simulation +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_scalar_field, + assemble_tiled_vector_field, + tile_scalar_field, + tile_vector_field, +) +from PyPIC3D.utils import build_yee_grid + + +jax.config.update("jax_enable_x64", True) + + +def unused_curl(Ex, Ey, Ez): + return None + + +class TestTiledElectrostatic(unittest.TestCase): + def _build_world(self, shape_factor=1): + world = { + "Nx": 8, + "Ny": 1, + "Nz": 1, + "dx": 0.5, + "dy": 1.0, + "dz": 1.0, + "dt": 0.01, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": shape_factor, + "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} + world["tile_shape"] = (2, 1, 1) + return world + + def _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + def _empty_fields(self, world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + zero = jnp.zeros(shape) + E = (zero, zero, zero) + B = (zero, zero, zero) + J = (zero, zero, zero) + external_fields = (E, B) + return E, B, J, zero, zero, external_fields + + def _species(self, world): + return particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + def _neutral_species(self, world): + electrons = self._species(world) + ions = particle_species( + name="ions", + N_particles=4, + charge=1.0, + mass=4.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.25, -0.25, 0.75, 1.75]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([-0.03, 0.04, -0.02, 0.01]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + return [electrons, ions] + + def test_scalar_tiles_round_trip_global_rho_and_phi(self): + world = self._build_world() + tile_shape = world["tile_shape"] + scalar = jnp.arange((world["Nx"] + 2) * (world["Ny"] + 2) * (world["Nz"] + 2), dtype=float) + scalar = scalar.reshape((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) + + scalar_tiles = tile_scalar_field(scalar, world, tile_shape) + scalar_from_tiles = assemble_tiled_scalar_field(scalar_tiles, world, tile_shape) + + self.assertTrue(jnp.allclose(scalar_from_tiles[1:-1, 1:-1, 1:-1], scalar[1:-1, 1:-1, 1:-1])) + + def test_tiled_rho_assembles_to_global_rho(self): + world = self._build_world(shape_factor=1) + constants = {"alpha": 1.0} + species = self._species(world) + tiled_particles = to_tiled_particles([species], world, self._simulation_parameters()) + _, _, _, rho, _, _ = self._empty_fields(world) + rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) + + reference_rho = compute_rho([self._species(world)], rho, world, constants) + tiled_rho = compute_tiled_rho_from_tiled_particles( + tiled_particles, + rho_tiles, + world, + constants, + ) + assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"]) + + self.assertTrue( + jnp.allclose( + assembled_rho[1:-1, 1:-1, 1:-1], + reference_rho[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): + world = self._build_world(shape_factor=1) + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + E, B, J, rho, phi, external_fields = self._empty_fields(world) + + reference_particles, reference_fields = time_loop_electrostatic( + self._neutral_species(world), + (E, B, J, rho, phi, external_fields), + world, + constants, + unused_curl, + J_func=None, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles(self._neutral_species(world), world, self._simulation_parameters()) + 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, + ) + + tiled_particles, tiled_fields = time_loop_electrostatic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=None, + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields + + for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, world["tile_shape"])): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-8, atol=1.0e-8)) + for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, world["tile_shape"])): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, world["tile_shape"])): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"]) + phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, world["tile_shape"]) + self.assertTrue(jnp.allclose(rho_from_tiles[1:-1, 1:-1, 1:-1], reference_rho[1:-1, 1:-1, 1:-1], rtol=1.0e-12, atol=1.0e-12)) + tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] + reference_phi_interior = reference_phi[1:-1, 1:-1, 1:-1] + tiled_phi_interior = tiled_phi_interior - jnp.mean(tiled_phi_interior) + reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) + self.assertTrue(jnp.allclose(tiled_phi_interior, reference_phi_interior, rtol=1.0e-8, atol=1.0e-8)) + + def test_initialize_simulation_accepts_tiled_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": "tiled electrostatic init smoke", + "output_dir": tmpdir, + "solver": "tiled_yee", + "electrostatic": True, + "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", + "fast_backend": "default", + "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_tiled) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + self.assertEqual(tuple(simulation_parameters["tile_shape"]), (2, 1, 1)) + self.assertEqual(fields[0][0].shape[:3], (4, 1, 1)) + self.assertEqual(fields[3].shape[:3], (4, 1, 1)) + self.assertEqual(fields[4].shape[:3], (4, 1, 1)) + + +if __name__ == "__main__": + unittest.main() From c14ea293f76d277d9b9f5e73c3616d01a4a07bd3 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 11:11:29 -0500 Subject: [PATCH 11/83] implemented task 9 --- PyPIC3D/__main__.py | 62 +++++++---------- PyPIC3D/deposition/rho_tiled.py | 33 +++++++++ PyPIC3D/diagnostics/openPMD.py | 6 +- PyPIC3D/diagnostics/output_adapters.py | 94 ++++++++++++++++++++++++++ PyPIC3D/diagnostics/plotting.py | 4 +- PyPIC3D/diagnostics/vtk.py | 4 +- PyPIC3D/utils.py | 30 +++++--- tests/openpmd_test.py | 45 ++++++++++++ tests/tiled_fluid_quantities_test.py | 26 +++++++ tests/utils_test.py | 50 +++++++++++++- 10 files changed, 299 insertions(+), 55 deletions(-) create mode 100644 PyPIC3D/diagnostics/output_adapters.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 5a39cc83..57767361 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -41,42 +41,18 @@ ) from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles - -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, assemble_tiled_vector_field +from PyPIC3D.deposition.rho_tiled import ( + compute_rho_from_tiled_particles, + compute_tiled_mass_density_from_tiled_particles, + compute_tiled_rho_from_tiled_particles, +) +from PyPIC3D.diagnostics.output_adapters import fields_for_output, scalar_field_for_output # Importing functions from the PyPIC3D package ############################################################################################################ -def _fields_for_output(fields, world, simulation_parameters): - if simulation_parameters["solver"] != "tiled_yee": - return fields - - if len(fields) == 6: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields = fields - pml_state = None - else: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, *rest = fields - tile_shape = tuple(int(width) for width in world["tile_shape"]) - external_E, external_B = external_fields - - E = assemble_tiled_vector_field(E_tiles, world, tile_shape) - B = assemble_tiled_vector_field(B_tiles, world, tile_shape) - J = assemble_tiled_vector_field(J_tiles, world, tile_shape) - if getattr(rho, "ndim", 0) == 6: - rho = assemble_tiled_scalar_field(rho, world, tile_shape) - if getattr(phi, "ndim", 0) == 6: - phi = assemble_tiled_scalar_field(phi, world, tile_shape) - external_fields = ( - assemble_tiled_vector_field(external_E, world, tile_shape), - assemble_tiled_vector_field(external_B, world, tile_shape), - ) - - return (E, B, J, rho, phi, external_fields, pml_state) - - def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): if simulation_parameters["solver"] != "tiled_yee" or len(fields) < 8: return @@ -156,23 +132,31 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_scalars']: - if tiled_run: - rho = compute_rho_from_tiled_particles(particles, rho, world, constants) + if getattr(rho, "ndim", 0) == 6: + rho = compute_tiled_rho_from_tiled_particles(particles, rho, world, constants) + mass_density = compute_tiled_mass_density_from_tiled_particles(particles, rho, world) + rho_output = scalar_field_for_output(rho, world) + mass_density_output = scalar_field_for_output(mass_density, world) else: - 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 + if tiled_run: + rho = compute_rho_from_tiled_particles(particles, rho, world, constants) + else: + 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 + rho_output = rho + mass_density_output = mass_density 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]] + fields_mag = [rho_output[1:-1, y_mid, 1:-1], mass_density_output[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']: - output_fields = _fields_for_output(fields, world, simulation_parameters) + output_fields = fields_for_output(fields, world) E, B, J, rho, phi, external_fields, *rest = output_fields # assemble tiled fields before VTK output y_mid = world['Ny']//2 + 1 @@ -192,7 +176,7 @@ def run_PyPIC3D(config_file): # Write the particles in openPMD format if plotting_parameters['plot_openpmd_fields']: - write_openpmd_fields(_fields_for_output(fields, world, simulation_parameters), world, os.path.join(output_dir, "data"), plot_num, t, "fields", ".h5") + write_openpmd_fields(fields, world, os.path.join(output_dir, "data"), plot_num, t, "fields", ".h5") # Write the fields in openPMD format if not tiled_run: diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index c45ecd98..1264887a 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -11,6 +11,7 @@ from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( fold_tiled_ghost_cells, + tile_scalar_field, update_tiled_ghost_cells, ) from PyPIC3D.utils import digital_filter @@ -354,3 +355,35 @@ def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid= half_step_back=True, wrap_periodic_position=True, ) + + +@jit +def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world, grid=None): + """Compute mass density into tile-major vertex-grid scalar arrays.""" + if grid is None: + grid = world["grids"]["vertex"] + + scalar_weight = tiled_particles.mass * tiled_particles.weight + tile_shape = tuple(int(width) - 2 for width in rho_tiles.shape[3:]) + Nx = int(rho_tiles.shape[0]) * tile_shape[0] + Ny = int(rho_tiles.shape[1]) * tile_shape[1] + Nz = int(rho_tiles.shape[2]) * tile_shape[2] + rho = jnp.zeros( + ( + Nx + 2, + Ny + 2, + Nz + 2, + ), + dtype=rho_tiles.dtype, + ) + rho = _deposit_tiled_scalar_moment( + tiled_particles, + rho, + world, + grid, + scalar_weight, + half_step_back=True, + wrap_periodic_position=True, + ) + + return tile_scalar_field(rho, world, tile_shape) diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index b6f24983..99725d35 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -5,7 +5,7 @@ import numpy as np import importlib.metadata -from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.diagnostics.output_adapters import fields_for_output, particles_for_output def _ensure_openpmd_array(data, dtype=np.float64, squeeze=False): arr = np.asarray(data, dtype=dtype) @@ -127,7 +127,7 @@ def write_openpmd_fields_to_iteration(iteration, field_map, world, active_dims=( def write_openpmd_particles_to_iteration(iteration, particles, constants, species_names=None, world=None): - particles = flatten_tiled_particles_by_species(particles, species_names=species_names, world=world) + particles = particles_for_output(particles, 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. @@ -231,6 +231,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) @@ -397,6 +398,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..1f23be73 --- /dev/null +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -0,0 +1,94 @@ +from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, assemble_tiled_vector_field + + +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_field(field): + return tuple(int(width) - 2 for width in field.shape[3:]) + + +def _tile_shape_from_world_or_field(world, field): + if "tile_shape" in world: + return tuple(int(width) for width in world["tile_shape"]) + return _tile_shape_from_field(field) + + +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 and VTK + mesh writers still expect one global ghost-celled array. + """ + + if not _is_tiled_scalar(field): + return field + + tile_shape = _tile_shape_from_world_or_field(world, field) + return assemble_tiled_scalar_field(field, world, tile_shape) + + +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_or_field(world, field[0]) + return assemble_tiled_vector_field(field, world, tile_shape) + + +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 particles_for_output(particles, species_names=None, world=None): + """ + Flatten fixed-capacity tiled particle storage for diagnostics. + """ + + return flatten_tiled_particles_by_species( + particles, + species_names=species_names, + world=world, + ) diff --git a/PyPIC3D/diagnostics/plotting.py b/PyPIC3D/diagnostics/plotting.py index e2ae747a..9e283128 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -7,7 +7,7 @@ import jax from functools import partial -from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.diagnostics.output_adapters import particles_for_output def plot_positions(particles, t, x_wind, y_wind, z_wind, path): """ @@ -71,7 +71,7 @@ def write_particles_phase_space(particles, t, path, species_names=None, world=No os.makedirs(f"{path}/data/phase_space/z") # Create directory if it doesn't exist - particles = flatten_tiled_particles_by_species(particles, species_names=species_names, world=world) + particles = particles_for_output(particles, 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. diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py index 014d6e01..92374427 100644 --- a/PyPIC3D/diagnostics/vtk.py +++ b/PyPIC3D/diagnostics/vtk.py @@ -5,7 +5,7 @@ from pyevtk.hl import gridToVTK, pointsToVTK import jax.numpy as jnp -from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.diagnostics.output_adapters import particles_for_output def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): """ @@ -57,7 +57,7 @@ def plot_vtk_particles(particles, t, path, species_names=None, world=None): if not os.path.exists(f"{path}/data/particles"): os.makedirs(f"{path}/data/particles") - particles = flatten_tiled_particles_by_species(particles, species_names=species_names, world=world) + particles = particles_for_output(particles, species_names=species_names, world=world) # Tiled particle arrays are fixed-capacity by tile; VTK particle output only # receives the active diagnostic particles. diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 3126de19..ba3a603f 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -802,16 +802,28 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para } if hasattr(particles, "active") and hasattr(particles, "u"): - config["particles"].append( - { - "storage": "tiled", - "active_particles": int(jnp.sum(particles.active)), - "tile_shape": jax.tree_util.tree_map( - lambda x: x.tolist() if isinstance(x, jnp.ndarray) else x, - simulation_parameters.get("tile_shape", ()), - ), - } + n_species = particles.active.shape[3] + species_names = simulation_parameters.get("particle_species_names") + 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, :])) + config["particles"].append( + { + "name": name, + "storage": "tiled", + "active_particles": active_particles, + "tile_shape": tile_shape, + } + ) else: for particle in particles: if hasattr(particle, "species_meta") and particle.species_meta: diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index ac3b4a37..f6b4d028 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -5,6 +5,7 @@ from PyPIC3D.diagnostics import openPMD from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array +from PyPIC3D.solvers.yee_tiled import tile_scalar_field, tile_vector_field class FakeRecord: @@ -107,6 +108,50 @@ def test_write_openpmd_fields_preserves_thin_y_mesh_metadata(self): 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), + "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)) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 099fc60f..6b869677 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -7,8 +7,10 @@ from PyPIC3D import __main__ as pypic_main from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC from PyPIC3D.diagnostics.fluid_quantities import compute_mass_density +from PyPIC3D.deposition.rho_tiled import compute_tiled_mass_density_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field from PyPIC3D.utils import build_yee_grid @@ -117,6 +119,30 @@ def test_tiled_mass_density_matches_compute_mass_density_on_interior(self): ) ) + def test_tile_major_mass_density_assembles_to_global_mass_density(self): + world = self._build_world(shape_factor=2) + world["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + particles = self._particles(world) + tiled_particles = to_tiled_particles(particles, world, self._simulation_parameters()) + rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"]) + + mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) + mass_tiles = compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world) + mass_from_tiles = assemble_tiled_scalar_field(mass_tiles, world, world["tile_shape"]) + + 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_tiled_scalar_vtk_path_no_longer_rejects_tiled_yee(self): run_source = inspect.getsource(pypic_main.run_PyPIC3D) diff --git a/tests/utils_test.py b/tests/utils_test.py index ac50ab7e..a5f57f1f 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -4,11 +4,13 @@ import jax import jax.numpy as jnp import numpy as np +import toml from PyPIC3D.initialization import initialize_fields +from PyPIC3D.particles.tiled_particles import TiledParticles 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, + compute_energy, dump_parameters_to_toml, ) from PyPIC3D.particles.species_class import particle_species @@ -207,5 +209,51 @@ def test_compute_energy_ignores_inactive_particles(self): 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,)) + ones = jnp.ones(active.shape) + particles = TiledParticles( + x=zeros3, + u=zeros3, + charge=ones, + mass=ones, + weight=ones, + active=active, + update_x1=active, + update_x2=active, + update_x3=active, + update_u1=active, + update_u2=active, + update_u3=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]) + if __name__ == '__main__': unittest.main() From 15738fe4622e2708052107bea3042a73cabd155a Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 14:09:06 -0500 Subject: [PATCH 12/83] attempting to fix direct deposition --- PyPIC3D/deposition/direct_deposition_tiled.py | 103 ++++++++++-------- PyPIC3D/electrodynamic_tiled.py | 16 ++- PyPIC3D/particles/tiled_particle_refresh.py | 3 +- 3 files changed, 73 insertions(+), 49 deletions(-) diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index d40489b5..eff76716 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -37,14 +37,36 @@ def filter_component(component_tiles): return J_tiles +def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d): + """ + Build the ghost-celled coordinate line for one compact tile. + + The first local grid point is the global ghost-cell origin shifted by the + tile's active-cell offset. With this convention, local active indices are + 1..tile_n and local ghost indices are 0 and tile_n + 1. + """ + + offsets = jnp.arange(local_n, dtype=global_axis_grid.dtype) + return global_axis_grid[0] + (offsets + tile_index * tile_n) * d + + @partial(jit, static_argnames=("filter",)) -def direct_J_from_tiled_particles(tiled_particles, J_tiles, constants, world, grid=None, filter="bilinear"): +def direct_J_from_tiled_particles( + tiled_particles, + J_tiles, + constants, + world, + grid=None, + filter="bilinear", +): """ - Compute direct current deposition from tile-major particle storage. + Compute direct current deposition using tile-local stencils only. - This follows the same time centering and Yee-component stencil as the - global ``J_from_rhov`` path, but current stays in compact tile-local - ghost-celled arrays throughout the deposition. + ``tiled_particles.x`` is assumed to already be centered at the current + deposition time and refreshed into the tiles that own those centered + positions. In other words, callers that store particles at the forward + position should pass a temporary, refreshed deposition view with + ``x_dep = x_forward - 0.5 * u * dt``. """ if grid is None: @@ -53,25 +75,25 @@ def direct_J_from_tiled_particles(tiled_particles, J_tiles, constants, world, gr dx = world["dx"] dy = world["dy"] dz = world["dz"] - dt = world["dt"] Jx_tiles, Jy_tiles, Jz_tiles = J_tiles ntx, nty, ntz = Jx_tiles.shape[:3] tile_nx = Jx_tiles.shape[3] - 2 tile_ny = Jx_tiles.shape[4] - 2 tile_nz = Jx_tiles.shape[5] - 2 - Nx = ntx * tile_nx + 2 - Ny = nty * tile_ny + 2 - Nz = ntz * tile_nz + 2 - bc_x = world["boundary_conditions"]["x"] - bc_y = world["boundary_conditions"]["y"] - bc_z = world["boundary_conditions"]["z"] + local_Nx = tile_nx + 2 + local_Ny = tile_ny + 2 + local_Nz = tile_nz + 2 shape_factor = world["shape_factor"] 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]) + # 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 + def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, ty, tz): x = x_tile[..., 0].reshape(-1) y = x_tile[..., 1].reshape(-1) @@ -82,43 +104,41 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, active = active_tile.reshape(-1).astype(x.dtype) dq = (charge_tile * weight_tile).reshape(-1) / (dx * dy * dz) - # Match J_from_rhov: the current is deposited from the half-step-back - # particle position, while the velocity is the time-centered velocity. - x = x - vx * dt / 2 - y = y - vy * dt / 2 - z = z - vz * dt / 2 + x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx) + y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy) + z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz) 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, ) 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, ) 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, ) - 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 + 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 x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( shape_factor == 1, @@ -143,17 +163,12 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, y_weights_face = jnp.asarray(y_weights_face) z_weights_face = jnp.asarray(z_weights_face) - 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) - - xpts = xpts - tx * tile_nx - ypts = ypts - ty * tile_ny - zpts = zpts - tz * tile_nz - # Convert global ghost-celled indices to tile-local ghost-celled indices. + xpts, x_weights_node = collapse_axis_stencil(xpts, x_weights_node, local_Nx, ghost_cells=True) + _, x_weights_face = collapse_axis_stencil(xpts, x_weights_face, local_Nx, ghost_cells=True) + ypts, y_weights_node = collapse_axis_stencil(ypts, y_weights_node, local_Ny, ghost_cells=True) + _, y_weights_face = collapse_axis_stencil(ypts, y_weights_face, local_Ny, ghost_cells=True) + zpts, z_weights_node = collapse_axis_stencil(zpts, z_weights_node, local_Nz, ghost_cells=True) + _, z_weights_face = collapse_axis_stencil(zpts, z_weights_face, local_Nz, ghost_cells=True) tile_Jx = Jx_template tile_Jy = Jy_template @@ -213,4 +228,4 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, elif filter == "digital": J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"], world) - return J_tiles + return J_tiles \ No newline at end of file diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 2fff8f8f..474c0126 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -70,9 +70,8 @@ def time_loop_electrodynamic_tiled( ) # use the selected tiled pusher for particle velocities - particles = update_tiled_particle_positions(particles, world) - # update particle forward positions before current deposition - + particles = update_tiled_particle_positions(particles, world["dt"]/2) + # update particle forward positions before current deposition half a time step particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow # wrap periodic particles and move them into their owning tiles. Overflow @@ -91,6 +90,17 @@ def time_loop_electrodynamic_tiled( J_tiles = J_func(particles, J_tiles, constants, world) # deposit current directly into tile-local Yee current arrays + + particles = update_tiled_particle_positions(particles, world["dt"]/2) + # update particle forward positions before current deposition half a time step + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + # wrap periodic particles and move them into their owning tiles. Overflow + # means the fixed tile capacity was exceeded and the Python driver must + # reject the step rather than silently dropping particles. + + + if pml_state is None: E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) else: diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index c544a264..88a91b74 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -50,12 +50,11 @@ def _particle_tile_indices(x, y, z, world, tile_shape, tile_counts): ) -def update_tiled_particle_positions(tiled_particles, world): +def update_tiled_particle_positions(tiled_particles, dt): """ Advance tile-major particle positions without changing tile ownership. """ - dt = world["dt"] active = tiled_particles.active.astype(tiled_particles.x.dtype) dx = active * tiled_particles.u[..., 0] * dt From b0c0980593d5d16a33736336e9e2f9d18d332b6b Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 15:25:15 -0500 Subject: [PATCH 13/83] testing tiled particle deposition --- tests/direct_deposition_tiled_test.py | 60 +++++++++++++++++---------- 1 file changed, 38 insertions(+), 22 deletions(-) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index dbac1d8f..679be9a5 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -12,6 +12,7 @@ ) from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particle_refresh import refresh_tiled_particle_tiles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, fold_tiled_ghost_cells, @@ -68,15 +69,42 @@ def _empty_J_tiles(self, world, simulation_parameters): ) return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) - def _compare_tiled_to_standard(self, particles, world, simulation_parameters): - constants = {"C": 3.0e8, "alpha": 1.0} - tiled_particles = to_tiled_particles(particles, world, simulation_parameters) - tile_shape = ( + def _tile_shape(self, simulation_parameters): + return ( simulation_parameters["particle_tile_nx"], simulation_parameters["particle_tile_ny"], simulation_parameters["particle_tile_nz"], ) + 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. + """ + + tiled_particles = to_tiled_particles(particles, world, simulation_parameters) + tiled_particles = tiled_particles._replace( + x=tiled_particles.x - 0.5 * tiled_particles.u * world["dt"] + ) + + centered_particles, overflow = refresh_tiled_particle_tiles( + tiled_particles, + world, + self._tile_shape(simulation_parameters), + ) + self.assertFalse(bool(overflow)) + + return centered_particles + + def _compare_tiled_to_standard(self, particles, world, simulation_parameters): + constants = {"C": 3.0e8, "alpha": 1.0} + tiled_particles = self._centered_tiled_particles(particles, world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) + J_reference = J_from_rhov(particles, self._empty_J(world), constants, world, filter="none") J_tiles = direct_J_from_tiled_particles( tiled_particles, @@ -233,7 +261,7 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) J_tiles = direct_J_from_tiled_particles( tiled_particles, @@ -242,11 +270,7 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): world, filter="none", ) - tile_shape = ( - simulation_parameters["particle_tile_nx"], - simulation_parameters["particle_tile_ny"], - simulation_parameters["particle_tile_nz"], - ) + tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="none") J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) @@ -335,12 +359,8 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) - tile_shape = ( - simulation_parameters["particle_tile_nx"], - simulation_parameters["particle_tile_ny"], - simulation_parameters["particle_tile_nz"], - ) + tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") J_tiles = direct_J_from_tiled_particles( @@ -384,12 +404,8 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) - tile_shape = ( - simulation_parameters["particle_tile_nx"], - simulation_parameters["particle_tile_ny"], - simulation_parameters["particle_tile_nz"], - ) + tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") J_tiles = direct_J_from_tiled_particles( From 3284814a443d3fae11e823eab07e2ea88805f5ab Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 15:52:58 -0500 Subject: [PATCH 14/83] running long two stream test --- tests/tiled_yee_integration_test.py | 263 ++++++++++++++++++++++++++++ 1 file changed, 263 insertions(+) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index df8b72c1..c0871e02 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -23,6 +23,9 @@ jax.config.update("jax_enable_x64", True) +ROUND_OFF_RTOL = 1.0e-11 +ROUND_OFF_ATOL = 1.0e-11 + def unused_curl(Ex, Ey, Ez): return None @@ -89,6 +92,210 @@ def _active_rows(self, tiled_particles): order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] + def _long_two_stream_state(self): + world = { + "Nx": 16, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.002, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (4, 1, 1) + Nt = 1500 + + x0 = jnp.linspace( + -0.5 * world["x_wind"], + 0.5 * world["x_wind"], + world["Nx"], + endpoint=False, + ) + 0.5 * world["dx"] + x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) + zeros = jnp.zeros_like(x0) + beam_speed = 0.08 + + species_kwargs = { + "N_particles": world["Nx"], + "mass": 1.0, + "weight": 1.0, + "T": 0.0, + "x1": x0, + "x2": zeros, + "x3": zeros, + "v2": zeros, + "v3": zeros, + "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 = [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + reference_particles = [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + reference_fields = (E, B, J, rho, phi, external_fields, pml_state) + + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + tiled_particles = to_tiled_particles(particles, world, simulation_parameters) + tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + step=0, + ) + + return reference_particles, reference_fields, tiled_particles, tiled_fields, world, constants, tile_shape, Nt + + def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): + current_slots = tiled_particles.active.shape[-1] + if current_slots >= min_slots: + return tiled_particles + + pad_width = [(0, 0)] * tiled_particles.x.ndim + pad_width[-2] = (0, min_slots - current_slots) + x = jnp.pad(tiled_particles.x, pad_width) + u = jnp.pad(tiled_particles.u, pad_width) + + scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim + scalar_pad_width[-1] = (0, min_slots - current_slots) + charge = jnp.pad(tiled_particles.charge, scalar_pad_width) + mass = jnp.pad(tiled_particles.mass, scalar_pad_width) + weight = jnp.pad(tiled_particles.weight, scalar_pad_width) + active = jnp.pad(tiled_particles.active, scalar_pad_width) + update_x1 = jnp.pad(tiled_particles.update_x1, scalar_pad_width) + update_x2 = jnp.pad(tiled_particles.update_x2, scalar_pad_width) + update_x3 = jnp.pad(tiled_particles.update_x3, scalar_pad_width) + update_u1 = jnp.pad(tiled_particles.update_u1, scalar_pad_width) + update_u2 = jnp.pad(tiled_particles.update_u2, scalar_pad_width) + update_u3 = jnp.pad(tiled_particles.update_u3, scalar_pad_width) + + return TiledParticles( + x=x, + u=u, + charge=charge, + mass=mass, + weight=weight, + active=active, + update_x1=update_x1, + update_x2=update_x2, + update_x3=update_x3, + update_u1=update_u1, + update_u2=update_u2, + update_u3=update_u3, + ) + + def _standard_species_rows(self, particles, species_index): + species = particles[species_index] + active = species.get_active_mask() + x = jnp.stack(species.get_forward_position(), axis=1)[active] + u = jnp.stack(species.get_velocity(), axis=1)[active] + order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _tiled_species_rows(self, tiled_particles, species_index): + active = tiled_particles.active[:, :, :, species_index, :] + x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): + for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): + diff = jnp.max(jnp.abs(reference_component - tiled_component)) + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: {name}{component_name} max abs diff {diff}", + ) + + def _assert_tiled_state_matches_standard( + self, + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + step, + ): + if len(tiled_fields) > 7: + self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") + + for species_index in range(len(reference_particles)): + reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) + self.assertTrue( + jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} position mismatch", + ) + self.assertTrue( + jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} velocity mismatch", + ) + + reference_E, reference_B, reference_J, *_ = reference_fields + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + self._assert_vector_fields_close( + reference_E, + assemble_tiled_vector_field(E_tiles, world, tile_shape), + step, + "E", + ) + self._assert_vector_fields_close( + reference_B, + assemble_tiled_vector_field(B_tiles, world, tile_shape), + step, + "B", + ) + self._assert_vector_fields_close( + reference_J, + assemble_tiled_vector_field(J_tiles, world, tile_shape), + step, + "J", + ) + def test_validate_field_solver_accepts_tiled_yee(self): validate_field_solver("tiled_yee") @@ -644,6 +851,62 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): for reference, tiled in zip(reference_J, J_from_tiles): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_YEE") == "1", + "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", + ) + def test_long_two_stream_tiled_yee_matches_standard_every_step(self): + ( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + constants, + tile_shape, + Nt, + ) = self._long_two_stream_state() + + standard_J = functools.partial(J_from_rhov, filter="none") + tiled_loop = jax.jit( + time_loop_electrodynamic_tiled, + static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher"), + ) + + for t in range(Nt): + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_fields, + world, + constants, + unused_curl, + J_func=standard_J, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=None, + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + t + 1, + ) + if __name__ == "__main__": unittest.main() From e59343411eb29df9e7fc3bf30a8f13f9dddd9dc2 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 16:31:35 -0500 Subject: [PATCH 15/83] getting rid of halfstep back logic bug in electrostatic yee --- PyPIC3D/deposition/rho.py | 7 -- PyPIC3D/deposition/rho_tiled.py | 122 +++++++++++++++----------------- PyPIC3D/electrostatic_tiled.py | 2 +- 3 files changed, 58 insertions(+), 73 deletions(-) diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index a59090c6..13e5eed2 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -30,14 +30,7 @@ def compute_rho(particles, rho, world, constants): 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, diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index 1264887a..f6cad0b5 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -23,67 +23,29 @@ def _deposit_tiled_scalar_moment( world, grid, scalar_weight, - half_step_back, - wrap_periodic_position=False, + position=None, ): dx = world["dx"] dy = world["dy"] dz = world["dz"] - dt = world["dt"] 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"] - x = tiled_particles.x[..., 0].reshape(-1) - y = tiled_particles.x[..., 1].reshape(-1) - z = tiled_particles.x[..., 2].reshape(-1) - vx = tiled_particles.u[..., 0].reshape(-1) - vy = tiled_particles.u[..., 1].reshape(-1) - vz = tiled_particles.u[..., 2].reshape(-1) + 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, y, z = position + x = x.reshape(-1) + y = y.reshape(-1) + z = z.reshape(-1) active = tiled_particles.active.reshape(-1).astype(x.dtype) scalar_density = scalar_weight.reshape(-1) / (dx * dy * dz) - if half_step_back: - x = x - vx * dt / 2 - y = y - vy * dt / 2 - z = z - vz * dt / 2 - # Rho deposits stay unwrapped so seam-crossing charge lands in ghost - # cells before folding. Scalar diagnostics that mirror - # species.get_position() wrap periodic half-step positions first. - if wrap_periodic_position: - x = jax.lax.cond( - bc_x == BC_PERIODIC, - lambda x_value: jnp.where( - x_value > world["x_wind"] / 2, - x_value - world["x_wind"], - jnp.where(x_value < -world["x_wind"] / 2, x_value + world["x_wind"], x_value), - ), - lambda x_value: x_value, - x, - ) - y = jax.lax.cond( - bc_y == BC_PERIODIC, - lambda y_value: jnp.where( - y_value > world["y_wind"] / 2, - y_value - world["y_wind"], - jnp.where(y_value < -world["y_wind"] / 2, y_value + world["y_wind"], y_value), - ), - lambda y_value: y_value, - y, - ) - z = jax.lax.cond( - bc_z == BC_PERIODIC, - lambda z_value: jnp.where( - z_value > world["z_wind"] / 2, - z_value - world["z_wind"], - jnp.where(z_value < -world["z_wind"] / 2, z_value + world["z_wind"], z_value), - ), - lambda z_value: z_value, - z, - ) - x, _, deltax, xpts = prepare_particle_axis_stencil( x, grid[0], @@ -146,6 +108,49 @@ def _deposit_tiled_scalar_moment( return rho +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 + + 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 _tile_axis(axis, tile_index, cells_per_tile): start = tile_index * cells_per_tile return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) @@ -157,12 +162,10 @@ def _deposit_tiled_scalar_moment_to_tiles( world, grid, scalar_weight, - half_step_back, ): dx = world["dx"] dy = world["dy"] dz = world["dz"] - dt = world["dt"] bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] @@ -183,17 +186,9 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_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) active = active_tile.reshape(-1).astype(x.dtype) scalar_density = scalar_weight_tile.reshape(-1) / (dx * dy * dz) - if half_step_back: - x = x - vx * dt / 2 - y = y - vy * dt / 2 - z = z - vz * dt / 2 - x, _, deltax, xpts = prepare_particle_axis_stencil( x, local_grid[0], @@ -300,8 +295,6 @@ def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, gri world, grid, scalar_weight, - half_step_back=True, - wrap_periodic_position=False, ) alpha = constants["alpha"] @@ -329,7 +322,6 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, co world, grid, scalar_weight, - half_step_back=True, ) alpha = constants["alpha"] @@ -346,14 +338,14 @@ def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid= grid = world["grids"]["vertex"] scalar_weight = tiled_particles.mass * tiled_particles.weight + position = _diagnostic_mass_position(tiled_particles, world) return _deposit_tiled_scalar_moment( tiled_particles, rho, world, grid, scalar_weight, - half_step_back=True, - wrap_periodic_position=True, + position=position, ) @@ -364,6 +356,7 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, grid = world["grids"]["vertex"] scalar_weight = tiled_particles.mass * tiled_particles.weight + position = _diagnostic_mass_position(tiled_particles, world) tile_shape = tuple(int(width) - 2 for width in rho_tiles.shape[3:]) Nx = int(rho_tiles.shape[0]) * tile_shape[0] Ny = int(rho_tiles.shape[1]) * tile_shape[1] @@ -382,8 +375,7 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world, grid, scalar_weight, - half_step_back=True, - wrap_periodic_position=True, + position=position, ) return tile_scalar_field(rho, world, tile_shape) diff --git a/PyPIC3D/electrostatic_tiled.py b/PyPIC3D/electrostatic_tiled.py index e8bef6dc..dee48805 100644 --- a/PyPIC3D/electrostatic_tiled.py +++ b/PyPIC3D/electrostatic_tiled.py @@ -60,7 +60,7 @@ def time_loop_electrostatic_tiled( ) # push velocities using the selected tiled particle pusher - particles = update_tiled_particle_positions(particles, world) + particles = update_tiled_particle_positions(particles, world["dt"]) # update particle forward positions before depositing rho particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) From 861db43e4afc73076ec3155b35444ff606b1a504 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 26 Jun 2026 17:08:54 -0500 Subject: [PATCH 16/83] validated electrostatic tiled solver --- PyPIC3D/evolve.py | 9 +- PyPIC3D/initialization.py | 10 +- PyPIC3D/solvers/electrostatic_yee.py | 2 +- tests/evolve_test.py | 193 ++++++++++++++++++- tests/initialization_test.py | 43 +++++ tests/tiled_electrostatic_test.py | 278 ++++++++++++++++++++++++++- tests/tiled_particle_refresh_test.py | 4 +- tests/tiled_rho_test.py | 97 +++++++++- 8 files changed, 621 insertions(+), 15 deletions(-) diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 49e2fe6e..eeb2f871 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -54,12 +54,15 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu ################ PARTICLE PUSH ######################################################################################## for i in range(len(particles)): - particles[i] = particle_push(particles[i], push_E, push_B, vertex_grid, center_grid, world, constants, relativistic=relativistic, particle_pusher=particle_pusher) + particles[i] = particle_push(particles[i], push_E, push_B, center_grid, vertex_grid, world, constants, relativistic=relativistic, particle_pusher=particle_pusher) # use the selected particle pusher for particle velocities particles[i].update_position() # update the particle positions + particles[i].boundary_conditions(world) + # electrostatic rho is deposited from the current in-domain 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 @@ -67,10 +70,6 @@ def time_loop_electrostatic(particles, fields, world, constants, curl_func, J_fu fields = (E, B, J, rho, phi, external_fields) # pack the fields into a tuple - for i in range(len(particles)): - particles[i].boundary_conditions(world) - # apply boundary conditions to the particles - return particles, fields diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 24d539c3..5e47301f 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -16,7 +16,7 @@ from PyPIC3D.utils import ( courant_condition, - update_parameters_from_toml, + update_parameters_from_toml, build_collocated_grid, build_yee_grid, convert_to_jax_compatible, load_external_fields_from_toml, print_stats, particle_sanity_check, build_plasma_parameters_dict, make_dir, compute_energy, add_external_fields @@ -362,8 +362,12 @@ def initialize_simulation(toml_file): plotting_parameters = convert_to_jax_compatible(plotting_parameters) # convert the world parameters to jax compatible format - 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, diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 9398b7cd..c18dcec0 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -19,7 +19,7 @@ @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. diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 2e8ef63d..031a4820 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -3,10 +3,12 @@ import jax import jax.numpy as jnp -from PyPIC3D.evolve import time_loop_electrodynamic +from PyPIC3D.evolve import time_loop_electrodynamic, time_loop_electrostatic +from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.utils import build_collocated_grid, build_yee_grid jax.config.update("jax_enable_x64", True) @@ -20,6 +22,193 @@ def unused_curl(Ex, Ey, Ez): class TestEvolveExternalFields(unittest.TestCase): + def _electrostatic_grid_order_species(self, world, charge): + return particle_species( + name=f"charge {charge}", + N_particles=1, + charge=charge, + mass=1.0, + T=0.0, + weight=1.0, + v1=jnp.array([0.0]), + v2=jnp.array([0.0]), + v3=jnp.array([0.0]), + x1=jnp.array([0.07]), + 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"], + dt=world["dt"], + ) + + def test_electrostatic_push_uses_same_grid_order_as_electrodynamic(self): + world = { + "Nx": 4, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} + + E, B, J, phi, rho = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + Ex = jnp.broadcast_to(vertex_grid[0][:, None, None], E[0].shape) + E = (Ex, E[1], E[2]) + external_fields = ( + tuple(jnp.zeros_like(comp) for comp in E), + tuple(jnp.zeros_like(comp) for comp in B), + ) + + pushed_particles = [ + self._electrostatic_grid_order_species(world, 1.0), + self._electrostatic_grid_order_species(world, -1.0), + ] + reference_particles = [ + self._electrostatic_grid_order_species(world, 1.0), + self._electrostatic_grid_order_species(world, -1.0), + ] + + pushed_particles, _ = time_loop_electrostatic( + pushed_particles, + (E, B, J, rho, phi, external_fields), + world, + constants, + unused_curl, + J_func=None, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + for i in range(len(reference_particles)): + reference_particles[i] = particle_push( + reference_particles[i], + E, + B, + center_grid, + vertex_grid, + world, + constants, + relativistic=False, + particle_pusher="boris", + ) + reference_particles[i].update_position() + reference_particles[i].boundary_conditions(world) + + for pushed, reference in zip(pushed_particles, reference_particles): + pushed_vx, pushed_vy, pushed_vz = pushed.get_velocity() + reference_vx, reference_vy, reference_vz = reference.get_velocity() + self.assertTrue(jnp.allclose(pushed_vx, reference_vx, rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(pushed_vy, reference_vy, rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(pushed_vz, reference_vz, rtol=1.0e-12, atol=1.0e-12)) + + def test_electrostatic_rho_uses_boundary_applied_particle_positions(self): + world = { + "Nx": 16, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.002, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_collocated_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.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="right_crossing_charge", + N_particles=1, + charge=1.0, + mass=1.0, + T=0.0, + weight=1.0, + v1=jnp.array([0.10]), + v2=jnp.array([0.0]), + v3=jnp.array([0.0]), + x1=jnp.array([1.99995]), + 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"], + dt=world["dt"], + ), + particle_species( + name="neutralizing_charge", + N_particles=1, + charge=-1.0, + mass=1.0, + T=0.0, + weight=1.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"], + dt=world["dt"], + ), + ] + + particles, fields = time_loop_electrostatic( + particles, + (E, B, J, rho, phi, external_fields), + world, + constants, + unused_curl, + J_func=None, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + _, _, _, rho_after, _, _ = fields + recomputed_rho = compute_rho(particles, jnp.zeros_like(rho_after), world, constants) + + self.assertTrue( + jnp.allclose( + rho_after[1:-1, 1:-1, 1:-1], + recomputed_rho[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + def test_external_electric_field_pushes_particles_without_evolving_maxwell_fields(self): world = { "Nx": 3, diff --git a/tests/initialization_test.py b/tests/initialization_test.py index ec33302f..2bc64d1c 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -6,6 +6,7 @@ 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_electrostatic jax.config.update("jax_enable_x64", True) @@ -81,6 +82,48 @@ def test_initialize_simulation_encodes_global_particle_boundary_conditions(self) self.assertEqual(world["particle_boundary_conditions"], {"x": 1, "y": 2, "z": 0}) self.assertEqual(particles[0].x_bc, "absorbing") + 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": "fdtd", + "electrostatic": True, + "Nx": 4, + "Ny": 2, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "Nt": 1, + "dt": 1.0e-10, + "fast_backend": "default", + }, + "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, _, _, world, *_ = initialize_simulation(config) + + self.assertIs(loop, time_loop_electrostatic) + for vertex_axis, center_axis in zip(world["grids"]["vertex"], world["grids"]["center"]): + self.assertTrue(jnp.allclose(vertex_axis, center_axis)) + def test_initialize_simulation_rejects_unknown_solver(self): with tempfile.TemporaryDirectory() as tmpdir: config = { diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 5683f7e8..fd96ee03 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -22,11 +22,14 @@ tile_scalar_field, tile_vector_field, ) -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.utils import build_collocated_grid, build_yee_grid jax.config.update("jax_enable_x64", True) +ROUND_OFF_RTOL = 1.0e-11 +ROUND_OFF_ATOL = 1.0e-11 + def unused_curl(Ex, Ey, Ez): return None @@ -117,6 +120,221 @@ def _neutral_species(self, world): ) return [electrons, ions] + def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): + current_slots = tiled_particles.active.shape[-1] + if current_slots >= min_slots: + return tiled_particles + + pad_width = [(0, 0)] * tiled_particles.x.ndim + pad_width[-2] = (0, min_slots - current_slots) + x = jnp.pad(tiled_particles.x, pad_width) + u = jnp.pad(tiled_particles.u, pad_width) + + scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim + scalar_pad_width[-1] = (0, min_slots - current_slots) + + return TiledParticles( + x=x, + u=u, + charge=jnp.pad(tiled_particles.charge, scalar_pad_width), + mass=jnp.pad(tiled_particles.mass, scalar_pad_width), + weight=jnp.pad(tiled_particles.weight, scalar_pad_width), + active=jnp.pad(tiled_particles.active, scalar_pad_width), + update_x1=jnp.pad(tiled_particles.update_x1, scalar_pad_width), + update_x2=jnp.pad(tiled_particles.update_x2, scalar_pad_width), + update_x3=jnp.pad(tiled_particles.update_x3, scalar_pad_width), + update_u1=jnp.pad(tiled_particles.update_u1, scalar_pad_width), + update_u2=jnp.pad(tiled_particles.update_u2, scalar_pad_width), + update_u3=jnp.pad(tiled_particles.update_u3, scalar_pad_width), + ) + + def _long_two_stream_species(self, world): + x0 = jnp.linspace( + -0.5 * world["x_wind"], + 0.5 * world["x_wind"], + world["Nx"], + endpoint=False, + ) + 0.5 * world["dx"] + x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) + zeros = jnp.zeros_like(x0) + beam_speed = 0.08 + + species_kwargs = { + "N_particles": world["Nx"], + "mass": 1.0, + "weight": 1.0, + "T": 0.0, + "x1": x0, + "x2": zeros, + "x3": zeros, + "v2": zeros, + "v3": zeros, + "xwind": world["x_wind"], + "ywind": world["y_wind"], + "zwind": world["z_wind"], + "dx": world["dx"], + "dy": world["dy"], + "dz": world["dz"], + "dt": world["dt"], + } + return [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + + def _long_two_stream_state(self): + world = { + "Nx": 16, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.002, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_collocated_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (4, 1, 1) + world["tile_shape"] = tile_shape + Nt = 1500 + + E, B, J, rho, phi, external_fields = self._empty_fields(world) + reference_fields = (E, B, J, rho, phi, external_fields) + reference_particles = self._long_two_stream_species(world) + + tiled_particles = to_tiled_particles( + self._long_two_stream_species(world), + world, + { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + }, + ) + tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + tile_scalar_field(rho, world, tile_shape), + tile_scalar_field(phi, world, tile_shape), + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + + self._assert_long_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + step=0, + ) + + return reference_particles, reference_fields, tiled_particles, tiled_fields, world, constants, tile_shape, Nt + + def _standard_species_rows(self, particles, species_index): + species = particles[species_index] + active = species.get_active_mask() + x = jnp.stack(species.get_forward_position(), axis=1)[active] + u = jnp.stack(species.get_velocity(), axis=1)[active] + order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _tiled_species_rows(self, tiled_particles, species_index): + active = tiled_particles.active[:, :, :, species_index, :] + x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): + for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): + diff = jnp.max(jnp.abs(reference_component - tiled_component)) + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: {name}{component_name} max abs diff {diff}", + ) + + def _assert_long_tiled_state_matches_standard( + self, + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + step, + ): + if len(tiled_fields) > 7: + self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") + + for species_index in range(len(reference_particles)): + reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) + self.assertTrue( + jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} position mismatch", + ) + self.assertTrue( + jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} velocity mismatch", + ) + + reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + self._assert_vector_fields_close( + reference_E, + assemble_tiled_vector_field(E_tiles, world, tile_shape), + step, + "E", + ) + self._assert_vector_fields_close( + reference_B, + assemble_tiled_vector_field(B_tiles, world, tile_shape), + step, + "B", + ) + self._assert_vector_fields_close( + reference_J, + assemble_tiled_vector_field(J_tiles, world, tile_shape), + step, + "J", + ) + + rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, tile_shape) + phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, tile_shape) + self.assertTrue( + jnp.allclose( + rho_from_tiles[1:-1, 1:-1, 1:-1], + reference_rho[1:-1, 1:-1, 1:-1], + rtol=ROUND_OFF_RTOL, + atol=ROUND_OFF_ATOL, + ), + f"step {step}: rho max abs diff {jnp.max(jnp.abs(rho_from_tiles - reference_rho))}", + ) + tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] + reference_phi_interior = reference_phi[1:-1, 1:-1, 1:-1] + tiled_phi_interior = tiled_phi_interior - jnp.mean(tiled_phi_interior) + reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) + self.assertTrue( + jnp.allclose(tiled_phi_interior, reference_phi_interior, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: phi max abs diff {jnp.max(jnp.abs(tiled_phi_interior - reference_phi_interior))}", + ) + def test_scalar_tiles_round_trip_global_rho_and_phi(self): world = self._build_world() tile_shape = world["tile_shape"] @@ -216,6 +434,62 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) self.assertTrue(jnp.allclose(tiled_phi_interior, reference_phi_interior, rtol=1.0e-8, atol=1.0e-8)) + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", + "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step tiled electrostatic comparison.", + ) + def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self): + ( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + constants, + tile_shape, + Nt, + ) = self._long_two_stream_state() + + tiled_loop = jax.jit( + time_loop_electrostatic_tiled, + static_argnames=("curl_func", "J_func", "solver", "tile_shape", "relativistic", "particle_pusher"), + ) + + for step in range(1, Nt + 1): + reference_particles, reference_fields = time_loop_electrostatic( + reference_particles, + reference_fields, + world, + constants, + unused_curl, + J_func=None, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=None, + solver="tiled_yee", + tile_shape=tile_shape, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_long_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + step, + ) + def test_initialize_simulation_accepts_tiled_electrostatic(self): with tempfile.TemporaryDirectory() as tmpdir: x_path = os.path.join(tmpdir, "x.npy") @@ -272,6 +546,8 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): self.assertIsInstance(particles, TiledParticles) self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) self.assertEqual(tuple(simulation_parameters["tile_shape"]), (2, 1, 1)) + for vertex_axis, center_axis in zip(world["grids"]["vertex"], world["grids"]["center"]): + self.assertTrue(jnp.allclose(vertex_axis, center_axis)) self.assertEqual(fields[0][0].shape[:3], (4, 1, 1)) self.assertEqual(fields[3].shape[:3], (4, 1, 1)) self.assertEqual(fields[4].shape[:3], (4, 1, 1)) diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 7928c2f0..04240521 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -87,7 +87,7 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): ) tiled_particles = to_tiled_particles([species], world, self._simulation_parameters()) - moved = update_tiled_particle_positions(tiled_particles, world) + moved = update_tiled_particle_positions(tiled_particles, world["dt"]) x, _, _, _, _ = self._active_rows(moved) self.assertTrue(jnp.allclose(x[:, 0], jnp.array([-1.25, 1.25]))) @@ -95,7 +95,7 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): fixed_species = self._species(world, x1=[-1.5], v1=[0.25], update_x=False) fixed = to_tiled_particles([fixed_species], world, self._simulation_parameters()) - fixed_moved = update_tiled_particle_positions(fixed, world) + fixed_moved = update_tiled_particle_positions(fixed, world["dt"]) self.assertTrue(jnp.allclose(fixed_moved.x, fixed.x)) def test_refresh_moves_particles_to_neighbor_tiles_with_static_shape(self): diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 01525fc1..69fc4722 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -5,9 +5,10 @@ from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles +from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field from PyPIC3D.utils import build_yee_grid @@ -104,6 +105,16 @@ def _tiled_with_noisy_inactive_slots(self, particles, world): return tiled_particles._replace(x=x, u=u, charge=charge, weight=weight) + def _zero_species_velocities(self, particles): + for species in particles: + species.v1 = jnp.zeros_like(species.v1) + species.v2 = jnp.zeros_like(species.v2) + species.v3 = jnp.zeros_like(species.v3) + return particles + + def _zero_tiled_velocities(self, tiled_particles): + return tiled_particles._replace(u=jnp.zeros_like(tiled_particles.u)) + def _compare_tiled_to_standard(self, shape_factor, alpha): world = self._build_world(shape_factor) constants = {"alpha": alpha} @@ -136,6 +147,90 @@ def test_tiled_rho_matches_compute_rho_for_shape_factor_2(self): 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 = compute_rho(particles, self._empty_scalar(world), world, constants) + rho_with_zero_velocity = compute_rho(zero_velocity_particles, self._empty_scalar(world), 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) + tiled_particles = self._tiled_with_noisy_inactive_slots(particles, world) + zero_velocity_tiled_particles = self._zero_tiled_velocities(tiled_particles) + + rho_with_velocity = compute_rho_from_tiled_particles( + tiled_particles, + self._empty_scalar(world), + world, + constants, + ) + rho_with_zero_velocity = compute_rho_from_tiled_particles( + zero_velocity_tiled_particles, + self._empty_scalar(world), + 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_tile_major_rho_uses_current_positions_not_half_step_back_positions(self): + world = self._build_world(shape_factor=2) + world["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + constants = {"alpha": 1.0} + particles = self._particles(world) + tiled_particles = 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_tiled_rho_from_tiled_particles( + tiled_particles, + rho_tiles, + world, + constants, + ) + rho_tiles_with_zero_velocity = compute_tiled_rho_from_tiled_particles( + zero_velocity_tiled_particles, + rho_tiles, + world, + constants, + ) + rho_with_velocity = assemble_tiled_scalar_field(rho_tiles_with_velocity, world, world["tile_shape"]) + rho_with_zero_velocity = assemble_tiled_scalar_field(rho_tiles_with_zero_velocity, world, world["tile_shape"]) + + 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, + ) + ) + if __name__ == "__main__": unittest.main() From 4a8dbddc20f8f6c51a9ed84ea3177509d0dc2d68 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 27 Jun 2026 06:11:56 -0500 Subject: [PATCH 17/83] generalizing guard cell handling to incorporate larger numbers of gaurd cells --- PyPIC3D/solvers/yee_tiled.py | 137 +++++++++++++------------- tests/direct_deposition_tiled_test.py | 19 ++++ tests/yee_tiled_test.py | 86 ++++++++++++++++ 3 files changed, 175 insertions(+), 67 deletions(-) diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 9309d35a..20757633 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -106,16 +106,18 @@ def _assemble_scalar_tiles(field_tiles, world, tile_shape): return assemble_tiled_scalar_field(field_tiles, world, tile_shape) -def update_tiled_ghost_cells(field_tiles, world): +def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): """ - Refresh one-cell tile halos using the field boundary conditions. + Refresh tile halos using the field boundary conditions. ``field_tiles`` has shape - ``(ntx, nty, ntz, tile_nx + 2, tile_ny + 2, tile_nz + 2)``. The last - three axes are the tile-local ghost-celled field. The first three axes + ``(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) bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] @@ -123,44 +125,44 @@ def update_tiled_ghost_cells(field_tiles, world): # x-halos: internal tile boundaries exchange neighboring interiors. At a # conducting global wall, the exterior ghost face is zero, matching # update_ghost_cells on the assembled field. - field_tiles = field_tiles.at[:, :, :, 0, :, :].set( - jnp.roll(field_tiles[:, :, :, -2, :, :], shift=1, axis=0) + field_tiles = field_tiles.at[:, :, :, :g, :, :].set( + jnp.roll(field_tiles[:, :, :, -2 * g:-g, :, :], shift=1, axis=0) ) - field_tiles = field_tiles.at[:, :, :, -1, :, :].set( - jnp.roll(field_tiles[:, :, :, 1, :, :], 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, :, :, 0, :, :].set(0.0).at[-1, :, :, -1, :, :].set(0.0), + 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 the global arrays. - field_tiles = field_tiles.at[:, :, :, :, 0, :].set( - jnp.roll(field_tiles[:, :, :, :, -2, :], shift=1, axis=1) + field_tiles = field_tiles.at[:, :, :, :, :g, :].set( + jnp.roll(field_tiles[:, :, :, :, -2 * g:-g, :], shift=1, axis=1) ) - field_tiles = field_tiles.at[:, :, :, :, -1, :].set( - jnp.roll(field_tiles[:, :, :, :, 1, :], 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, :, :, 0, :].set(0.0).at[:, -1, :, :, -1, :].set(0.0), + 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. - field_tiles = field_tiles.at[:, :, :, :, :, 0].set( - jnp.roll(field_tiles[:, :, :, :, :, -2], shift=1, axis=2) + field_tiles = field_tiles.at[:, :, :, :, :, :g].set( + jnp.roll(field_tiles[:, :, :, :, :, -2 * g:-g], shift=1, axis=2) ) - field_tiles = field_tiles.at[:, :, :, :, :, -1].set( - jnp.roll(field_tiles[:, :, :, :, :, 1], 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, :, :, 0].set(0.0).at[:, :, -1, :, :, -1].set(0.0), + lambda tiles: tiles.at[:, :, 0, :, :, :g].set(0.0).at[:, :, -1, :, :, -g:].set(0.0), lambda tiles: tiles, operand=field_tiles, ) @@ -168,32 +170,32 @@ def update_tiled_ghost_cells(field_tiles, world): return field_tiles -def update_tiled_ghost_cells_periodic(field_tiles): +def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): """ - Refresh one-cell tile halos with all-periodic boundary conditions. + Refresh tile halos with all-periodic boundary conditions. """ world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} - return update_tiled_ghost_cells(field_tiles, world) + return update_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def update_tiled_vector_ghost_cells(field_tiles, world): +def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1): """ Refresh tile halos for each component of a vector field. """ - return tuple(update_tiled_ghost_cells(component, world) for component in field_tiles) + return tuple(update_tiled_ghost_cells(component, world, num_guard_cells) for component in field_tiles) -def update_tiled_vector_ghost_cells_periodic(field_tiles): +def update_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): """ Refresh periodic tile halos for each component of a vector field. """ - return tuple(update_tiled_ghost_cells_periodic(component) for component in field_tiles) + return tuple(update_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) -def update_tiled_ghost_cells_for_pml(field_tiles, world): +def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1): """ Refresh tile halos without wrapping across PML-active global walls. """ @@ -208,15 +210,15 @@ def update_tiled_ghost_cells_for_pml(field_tiles, world): 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}} - return update_tiled_ghost_cells(field_tiles, pml_world) + return update_tiled_ghost_cells(field_tiles, pml_world, num_guard_cells) -def update_tiled_vector_ghost_cells_for_pml(field_tiles, world): +def update_tiled_vector_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1): """ Refresh tile halos for a vector field with PML-active exterior walls. """ - return tuple(update_tiled_ghost_cells_for_pml(component, world) for component in field_tiles) + return tuple(update_tiled_ghost_cells_for_pml(component, world, num_guard_cells) for component in field_tiles) def apply_tiled_conducting_bc(E_tiles, world): @@ -272,7 +274,7 @@ def apply_tiled_conducting_bc(E_tiles, world): return Ex, Ey, Ez -def fold_tiled_ghost_cells(field_tiles, world): +def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): """ Add tile-ghost deposits into owning interiors, then clear ghosts. @@ -283,99 +285,100 @@ def fold_tiled_ghost_cells(field_tiles, world): used by the global ghost-cell fold. """ + g = int(num_guard_cells) bc_x = world["boundary_conditions"]["x"] bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] - x_lower_ghost = field_tiles[:, :, :, 0, :, :] - x_upper_ghost = field_tiles[:, :, :, -1, :, :] - field_tiles = field_tiles.at[:, :, :, -2, :, :].add( + 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[:, :, :, 1, :, :].add( + 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, :, :].add(-x_lower_ghost[0, :, :, :, :]) - .at[0, :, :, 1, :, :].add(-x_upper_ghost[-1, :, :, :, :]) - .at[0, :, :, 1, :, :].add(-x_lower_ghost[0, :, :, :, :]) - .at[-1, :, :, -2, :, :].add(-x_upper_ghost[-1, :, :, :, :]), + .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[:, :, :, 0, :, :].set(0.0) - field_tiles = field_tiles.at[:, :, :, -1, :, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :g, :, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, -g:, :, :].set(0.0) - y_lower_ghost = field_tiles[:, :, :, :, 0, :] - y_upper_ghost = field_tiles[:, :, :, :, -1, :] - field_tiles = field_tiles.at[:, :, :, :, -2, :].add( + 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[:, :, :, :, 1, :].add( + 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, :].add(-y_lower_ghost[:, 0, :, :, :]) - .at[:, 0, :, :, 1, :].add(-y_upper_ghost[:, -1, :, :, :]) - .at[:, 0, :, :, 1, :].add(-y_lower_ghost[:, 0, :, :, :]) - .at[:, -1, :, :, -2, :].add(-y_upper_ghost[:, -1, :, :, :]), + .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[:, :, :, :, 0, :].set(0.0) - field_tiles = field_tiles.at[:, :, :, :, -1, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :g, :].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, -g:, :].set(0.0) - z_lower_ghost = field_tiles[:, :, :, :, :, 0] - z_upper_ghost = field_tiles[:, :, :, :, :, -1] - field_tiles = field_tiles.at[:, :, :, :, :, -2].add( + 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[:, :, :, :, :, 1].add( + 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].add(-z_lower_ghost[:, :, 0, :, :]) - .at[:, :, 0, :, :, 1].add(-z_upper_ghost[:, :, -1, :, :]) - .at[:, :, 0, :, :, 1].add(-z_lower_ghost[:, :, 0, :, :]) - .at[:, :, -1, :, :, -2].add(-z_upper_ghost[:, :, -1, :, :]), + .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[:, :, :, :, :, 0].set(0.0) - field_tiles = field_tiles.at[:, :, :, :, :, -1].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :, :g].set(0.0) + field_tiles = field_tiles.at[:, :, :, :, :, -g:].set(0.0) return field_tiles -def fold_tiled_ghost_cells_periodic(field_tiles): +def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): """ Fold all-periodic tile-ghost deposits for one scalar component. """ world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} - return fold_tiled_ghost_cells(field_tiles, world) + return fold_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def fold_tiled_vector_ghost_cells(field_tiles, world): +def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1): """ Fold tile-ghost deposits for each vector component. """ - return tuple(fold_tiled_ghost_cells(component, world) for component in field_tiles) + return tuple(fold_tiled_ghost_cells(component, world, num_guard_cells) for component in field_tiles) -def fold_tiled_vector_ghost_cells_periodic(field_tiles): +def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): """ Fold all-periodic tile-ghost deposits for each vector component. """ - return tuple(fold_tiled_ghost_cells_periodic(component) for component in field_tiles) + return tuple(fold_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state=None): diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 679be9a5..6e284419 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -165,6 +165,25 @@ def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self legacy_folded = fold_tiled_ghost_cells_periodic(tiles) self.assertTrue(jnp.allclose(folded, legacy_folded, rtol=1.0e-12, atol=1.0e-12)) + 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 + 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells) + + 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)) + def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): world = self._build_world( Nx=4, diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index ada896c1..eaf42fe6 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -80,6 +80,66 @@ def _deterministic_vector_field(self, world, scale): return tuple(self._fill_ghosts(component, world) for component in (Fx, Fy, Fz)) + 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) @@ -168,6 +228,32 @@ def test_update_tiled_ghost_cells_mixed_matches_global_ghost_cells(self): 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 + 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 = 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) From 892afb5fcc78d7e39db530f0a8434e00ddcee519 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 27 Jun 2026 06:48:52 -0500 Subject: [PATCH 18/83] implementing a tiled esirkepov current --- PyPIC3D/deposition/esirkepov_tiled.py | 240 +++++++++++++++++++++ PyPIC3D/electrodynamic_tiled.py | 62 +++--- PyPIC3D/initialization.py | 33 ++- PyPIC3D/solvers/yee_tiled.py | 73 +++++-- tests/tiled_esirkepov_test.py | 288 ++++++++++++++++++++++++++ 5 files changed, 651 insertions(+), 45 deletions(-) create mode 100644 PyPIC3D/deposition/esirkepov_tiled.py create mode 100644 tests/tiled_esirkepov_test.py diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py new file mode 100644 index 00000000..7c003a07 --- /dev/null +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -0,0 +1,240 @@ +import jax +from jax import jit +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + compute_particle_anchor, + particle_axis_offset, +) +from PyPIC3D.deposition.Esirkepov import ( + collapse_redundant_axis, + get_1D_esirkepov_weights, + get_2D_esirkepov_weights, + get_3D_esirkepov_weights, + shift_old_stencil, +) +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.solvers.yee_tiled import ( + fold_tiled_vector_ghost_cells, + update_tiled_vector_ghost_cells, +) + + +def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_cells): + """ + Build one local tile axis with ``num_guard_cells`` guard points per side. + """ + + offsets = jnp.arange(local_n, dtype=global_axis_grid.dtype) + return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d + + +@jit +def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=None): + """ + Deposit Esirkepov current into tile-local two-guard current buffers. + + ``tiled_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 boundary conditions and retile + ownership are applied by the caller after deposition. + """ + + if grid is None: + grid = world["grids"]["center"] + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + dt = world["dt"] + shape_factor = world["shape_factor"] + + Jx_tiles, Jy_tiles, Jz_tiles = J_tiles + ntx, nty, ntz = Jx_tiles.shape[:3] + g = 2 + tile_nx = Jx_tiles.shape[3] - 2 * g + tile_ny = Jx_tiles.shape[4] - 2 * g + tile_nz = Jx_tiles.shape[5] - 2 * g + local_Nx = tile_nx + 2 * g + local_Ny = tile_ny + 2 * g + local_Nz = tile_nz + 2 * g + + x_active = ntx * tile_nx > 1 + y_active = nty * tile_ny > 1 + z_active = ntz * tile_nz > 1 + + 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]) + + def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, update_x1, update_x2, update_x3, tx, ty, tz): + old_x = x_tile[..., 0].reshape(-1) + old_y = x_tile[..., 1].reshape(-1) + old_z = x_tile[..., 2].reshape(-1) + vx = u_tile[..., 0].reshape(-1) + vy = u_tile[..., 1].reshape(-1) + vz = u_tile[..., 2].reshape(-1) + active = active_tile.reshape(-1).astype(old_x.dtype) + q = (charge_tile * weight_tile).reshape(-1) + N_particles = old_x.shape[0] + + x = old_x + jnp.where(update_x1.reshape(-1), vx * dt, 0.0) + y = old_y + jnp.where(update_x2.reshape(-1), vy * dt, 0.0) + z = old_z + jnp.where(update_x3.reshape(-1), vz * dt, 0.0) + + x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) + y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) + z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz, g) + + 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) + 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) + + deltax = particle_axis_offset(x, x0, x_grid) + deltay = particle_axis_offset(y, y0, y_grid) + deltaz = particle_axis_offset(z, z0, z_grid) + 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) + + shift_x = x0 - old_x0 + shift_y = y0 - old_y0 + shift_z = z0 - old_z0 + + offsets = jnp.asarray([-2, -1, 0, 1, 2], dtype=x0.dtype) + xpts = x0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + ypts = y0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + zpts = z0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + + xw, yw, zw = 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, + ) + oxw, oyw, ozw = jax.lax.cond( + shape_factor == 1, + lambda _: get_first_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), + lambda _: get_second_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), + operand=None, + ) + + 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] + + oxw = shift_old_stencil(oxw, shift_x) + oyw = shift_old_stencil(oyw, shift_y) + ozw = shift_old_stencil(ozw, shift_z) + + 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) + + if x_active and y_active and z_active: + 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 not x_active: + null_dim = 0 + elif not y_active: + null_dim = 1 + else: + null_dim = 2 + 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): + 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): + Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=1) + else: + Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=2) + + dJx = jax.lax.cond( + x_active, + 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), + lambda _: active * q * vy / (dx * dy * dz), + operand=None, + ) + dJz = jax.lax.cond( + z_active, + lambda _: active * (-(q / (dx * dy)) / dt), + lambda _: active * q * vz / (dx * dy * dz), + operand=None, + ) + + 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) + + tile_Jx = Jx_template + tile_Jy = Jy_template + tile_Jz = Jz_template + + for i in range(5): + for j in range(5): + for k in range(5): + ix = xpts[i, :] + iy = ypts[j, :] + iz = zpts[k, :] + tile_Jx = tile_Jx.at[ix, iy, iz].add( + jax.lax.cond(x_active, lambda _: Jx_loc[i, j, k, :], lambda _: Fx[i, j, k, :], operand=None), + mode="drop", + ) + tile_Jy = tile_Jy.at[ix, iy, iz].add( + jax.lax.cond(y_active, lambda _: Jy_loc[i, j, k, :], lambda _: Fy[i, j, k, :], operand=None), + mode="drop", + ) + tile_Jz = tile_Jz.at[ix, iy, iz].add( + jax.lax.cond(z_active, lambda _: Jz_loc[i, j, k, :], lambda _: Fz[i, j, k, :], operand=None), + 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", + ) + + deposit_tiles = deposit_one_tile + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + + Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( + tiled_particles.x, + tiled_particles.u, + tiled_particles.charge, + tiled_particles.weight, + tiled_particles.active, + tiled_particles.update_x1, + tiled_particles.update_x2, + tiled_particles.update_x3, + tx, + ty, + tz, + ) + + J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=2) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells=2) + + return J_tiles diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 474c0126..9e8e36bc 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -1,6 +1,7 @@ import jax.numpy as jnp from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current from PyPIC3D.particles.tiled_particle_refresh import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, @@ -70,35 +71,40 @@ def time_loop_electrodynamic_tiled( ) # use the selected tiled pusher for particle velocities - particles = update_tiled_particle_positions(particles, world["dt"]/2) - # update particle forward positions before current deposition half a time step - particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) - overflow = overflow_previous | overflow - # wrap periodic particles and move them into their owning tiles. Overflow - # means the fixed tile capacity was exceeded and the Python driver must - # reject the step rather than silently dropping particles. - - if J_func is None: - J_tiles = direct_J_from_tiled_particles( - particles, - J_tiles, - constants, - world, - filter="none", - ) - else: - J_tiles = J_func(particles, J_tiles, constants, world) - # deposit current directly into tile-local Yee current arrays - - - particles = update_tiled_particle_positions(particles, world["dt"]/2) - # update particle forward positions before current deposition half a time step - particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) - overflow = overflow_previous | overflow - # wrap periodic particles and move them into their owning tiles. Overflow - # means the fixed tile capacity was exceeded and the Python driver must - # reject the step rather than silently dropping particles. + use_esirkepov = J_func is tiled_esirkepov_current + # Esirkepov needs old and new particle positions. The deposition kernel + # predicts the new positions locally, then the actual particle state is + # advanced and retiled after the current has been computed. + if use_esirkepov: + J_tiles = J_func(particles, J_tiles, constants, world) + particles = update_tiled_particle_positions(particles, world["dt"]) + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + else: + particles = update_tiled_particle_positions(particles, world["dt"]/2) + # update particle positions to the centered direct-current deposition time + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + # wrap periodic particles and move them into their owning tiles. + + if J_func is None: + J_tiles = direct_J_from_tiled_particles( + particles, + J_tiles, + constants, + world, + filter="none", + ) + else: + J_tiles = J_func(particles, J_tiles, constants, world) + # deposit current directly into tile-local Yee current arrays + + particles = update_tiled_particle_positions(particles, world["dt"]/2) + # complete the full particle position update + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + # refresh tile ownership after the full position update. if pml_state is None: diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 5e47301f..5c714adf 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -44,6 +44,7 @@ ) from PyPIC3D.solvers.yee_tiled import ( + empty_tiled_vector_field, tile_scalar_field, tile_vector_field, ) @@ -66,6 +67,7 @@ from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -122,13 +124,27 @@ def _tile_shape_from_parameters(simulation_parameters): ) +CURRENT_J_FROM_RHOV = 0 +CURRENT_ESIRKEPOV = 1 + + +def _encode_current_calculation(current_calculation): + current_codes = { + "j_from_rhov": CURRENT_J_FROM_RHOV, + "esirkepov": CURRENT_ESIRKEPOV, + } + if current_calculation not in current_codes: + raise ValueError("Unsupported current_calculation. Use 'j_from_rhov' or 'esirkepov'.") + return current_codes[current_calculation] + + 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"] != "j_from_rhov": - raise ValueError("tiled_yee currently supports only current_calculation='j_from_rhov'") + if simulation_parameters["current_calculation"] not in ("j_from_rhov", "esirkepov"): + raise ValueError("tiled_yee currently supports current_calculation='j_from_rhov' or 'esirkepov'") if simulation_parameters["particle_pusher"] not in ("boris", "higuera_cary"): raise ValueError("tiled_yee currently supports only particle_pusher='boris' or 'higuera_cary'") if simulation_parameters["filter_j"] not in ("none", "digital", "bilinear"): @@ -333,6 +349,8 @@ def initialize_simulation(toml_file): 'y_wind': y_wind, 'z_wind': z_wind, 'shape_factor': simulation_parameters['shape_factor'], + 'current_calculation': _encode_current_calculation(simulation_parameters['current_calculation']), + 'current_guard_cells': 2 if simulation_parameters['current_calculation'] == "esirkepov" else 1, 'boundary_conditions': { 'x': _encode_field_bc(simulation_parameters['x_bc']), 'y': _encode_field_bc(simulation_parameters['y_bc']), @@ -505,8 +523,10 @@ def initialize_simulation(toml_file): 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 + if solver == "tiled_yee": + J_func = tiled_esirkepov_current + else: + 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']}") if solver == "tiled_yee": @@ -521,7 +541,10 @@ def initialize_simulation(toml_file): particles = to_tiled_particles(particles, world, simulation_parameters) E = tile_vector_field(E, world, tile_shape) B = tile_vector_field(B, world, tile_shape) - J = tile_vector_field(J, world, tile_shape) + if simulation_parameters["current_calculation"] == "esirkepov": + J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2, dtype=E[0].dtype) + else: + J = tile_vector_field(J, world, tile_shape) external_E, external_B = external_fields external_fields = ( tile_vector_field(external_E, world, tile_shape), diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 20757633..7fe5e34d 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -15,14 +15,47 @@ def _tile_axis_count(n_cells, cells_per_tile): return int(n_cells) // int(cells_per_tile) -def tile_scalar_field(field, world, tile_shape): +def empty_tiled_scalar_field(world, tile_shape, num_guard_cells=1, dtype=None): """ - Split a ghost-celled field into compact tiles using the shared tile shape. + Allocate an empty tile-major scalar field with ``num_guard_cells`` guards. + """ + + 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) + g = int(num_guard_cells) + if dtype is None: + dtype = jnp.float64 + + return jnp.zeros( + ( + ntx, + nty, + ntz, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, + ), + dtype=dtype, + ) + + +def empty_tiled_vector_field(world, tile_shape, num_guard_cells=1, dtype=None): + """ + Allocate empty tile-major vector-field components. """ - del world + return tuple(empty_tiled_scalar_field(world, tile_shape, num_guard_cells, dtype) for _ in range(3)) + + +def tile_scalar_field(field, world, tile_shape, num_guard_cells=1): + """ + Split a ghost-celled field into compact tiles using the shared 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 @@ -30,6 +63,18 @@ def tile_scalar_field(field, world, tile_shape): nty = _tile_axis_count(Ny, tile_ny) ntz = _tile_axis_count(Nz, tile_nz) + if g != 1: + field_tiles = empty_tiled_scalar_field(world, tile_shape, g, 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 update_tiled_ghost_cells(field_tiles, world, g) + 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) @@ -54,7 +99,7 @@ def _tile_scalar_field(field, tile_shape): return tile_scalar_field(field, None, tile_shape) -def tile_vector_field(field, world, tile_shape): +def tile_vector_field(field, world, tile_shape, num_guard_cells=1): """ Split ``(Fx, Fy, Fz)`` into compact ghost-celled tiles. @@ -63,15 +108,16 @@ def tile_vector_field(field, world, tile_shape): ``(tile_nx + 2, tile_ny + 2, tile_nz + 2)``. """ - return tuple(tile_scalar_field(component, world, tile_shape) for component in field) + return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) -def assemble_tiled_scalar_field(field_tiles, world, tile_shape): +def assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells=1): """ 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 @@ -82,7 +128,7 @@ def assemble_tiled_scalar_field(field_tiles, world, tile_shape): for tx in range(ntx): for ty in range(nty): for tz in range(ntz): - interior = field_tiles[tx, ty, tz, 1:-1, 1:-1, 1:-1] + interior = field_tiles[tx, ty, tz, g:-g, g:-g, g:-g] ix = 1 + tx * tile_nx iy = 1 + ty * tile_ny iz = 1 + tz * tile_nz @@ -94,12 +140,12 @@ def assemble_tiled_scalar_field(field_tiles, world, tile_shape): return update_ghost_cells(field, bc_x, bc_y, bc_z) -def assemble_tiled_vector_field(field_tiles, world, tile_shape): +def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells=1): """ Assemble tiled vector-field components into ordinary ghost-celled arrays. """ - return tuple(assemble_tiled_scalar_field(component, world, tile_shape) for component in field_tiles) + return tuple(assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells) for component in field_tiles) def _assemble_scalar_tiles(field_tiles, world, tile_shape): @@ -397,6 +443,9 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ else: Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml(B_tiles, world) Jx, Jy, Jz = J_tiles + gx = (Jx.shape[3] - (Ex.shape[3] - 2)) // 2 + gy = (Jx.shape[4] - (Ex.shape[4] - 2)) // 2 + gz = (Jx.shape[5] - (Ex.shape[5] - 2)) // 2 dt = world["dt"] dx, dy, dz = world["dx"], world["dy"], world["dz"] @@ -426,15 +475,15 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ 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 + + (C**2 * curl_x - Jx[:, :, :, gx:-gx, gy:-gy, gz:-gz] / 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 + + (C**2 * curl_y - Jy[:, :, :, gx:-gx, gy:-gy, gz:-gz] / 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 + + (C**2 * curl_z - Jz[:, :, :, gx:-gx, gy:-gy, gz:-gz] / eps) * dt ) Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py new file mode 100644 index 00000000..8d7cfbb5 --- /dev/null +++ b/tests/tiled_esirkepov_test.py @@ -0,0 +1,288 @@ +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_PERIODIC +from PyPIC3D.deposition.Esirkepov import Esirkepov_current +from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current +from PyPIC3D.initialization import CURRENT_ESIRKEPOV, initialize_simulation +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_vector_field, + empty_tiled_vector_field, + tile_vector_field, + update_tiled_E, +) + +jax.config.update("jax_enable_x64", True) + + +class TestTiledEsirkepovCurrent(unittest.TestCase): + def _build_world(self, Nx=8, Ny=1, Nz=1, dt=0.05): + x_wind, y_wind, z_wind = 4.0, 1.0, 1.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": 1, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "current_calculation": CURRENT_ESIRKEPOV, + "current_guard_cells": 2, + } + 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(self, world, x1): + return particle_species( + name="electrons", + N_particles=3, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=x1, + x2=jnp.zeros_like(x1), + x3=jnp.zeros_like(x1), + v1=jnp.array([0.08, -0.06, 0.05]), + v2=jnp.zeros_like(x1), + v3=jnp.zeros_like(x1), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + def test_empty_tiled_vector_field_can_allocate_two_current_guard_cells(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + tile_shape = (2, 1, 1) + + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + + self.assertEqual(J_tiles[0].shape, (4, 1, 1, 6, 5, 5)) + self.assertTrue(jnp.allclose(J_tiles[0], 0.0)) + + def test_two_guard_current_tiles_assemble_to_one_guard_global_current(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + tile_shape = (2, 1, 1) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + 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=2) + + self.assertEqual(assembled[0].shape, (10, 3, 3)) + self.assertEqual(float(assembled[0][3, 1, 1]), 3.0) + + def test_update_tiled_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) + 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) + B_tiles = tile_vector_field(zeros, world, tile_shape) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + Jx, Jy, Jz = J_tiles + Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) + + E_after = update_tiled_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants, None, tile_shape) + + self.assertTrue(jnp.allclose(E_after[0][:, :, :, 1:-1, 1:-1, 1:-1], -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) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + x_old = jnp.array([-1.10, -0.10, 1.05]) + old_species = self._species(world, x_old) + x_new = x_old + old_species.v1 * world["dt"] + new_species = self._species(world, x_new) + + tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) + J_tiles = tiled_esirkepov_current( + tiled_particles, + empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + constants, + world, + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=2) + + 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_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": "tiled_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", + "fast_backend": "default", + "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:], (4, 3, 3)) + self.assertEqual(J_tiles[0].shape[-3:], (6, 5, 5)) + self.assertEqual(int(world["current_guard_cells"]), 2) + self.assertEqual(int(world["current_calculation"]), CURRENT_ESIRKEPOV) + + 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": "tiled_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", + "fast_backend": "default", + "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, curl_func, J_func, relativistic, particle_pusher = initialize_simulation(toml.load(config_path)) + + particles, fields = loop( + particles, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + 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(fields[2][0].shape[-3:], (6, 5, 5)) + self.assertFalse(bool(fields[-1])) + + +if __name__ == "__main__": + unittest.main() From 1c07cb9c0acb834d11ed85e8fac58a99f2b50526 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 28 Jun 2026 07:05:39 -0500 Subject: [PATCH 19/83] working on finishing task 0 --- PyPIC3D/deposition/current_methods.py | 12 + PyPIC3D/deposition/esirkepov_tiled.py | 57 ++- PyPIC3D/diagnostics/output_adapters.py | 12 +- PyPIC3D/electrodynamic_tiled.py | 54 ++- PyPIC3D/initialization.py | 25 +- .../tiled_particle_initialization.py | 2 + PyPIC3D/solvers/yee_tiled.py | 50 ++- tests/pml_test.py | 47 +++ tests/tiled_esirkepov_integration_test.py | 91 ++++ tests/tiled_esirkepov_test.py | 398 +++++++++++++++++- tests/tiled_particle_initialization_test.py | 46 ++ tests/tiled_yee_integration_test.py | 74 ++++ 12 files changed, 828 insertions(+), 40 deletions(-) create mode 100644 PyPIC3D/deposition/current_methods.py create mode 100644 tests/tiled_esirkepov_integration_test.py diff --git a/PyPIC3D/deposition/current_methods.py b/PyPIC3D/deposition/current_methods.py new file mode 100644 index 00000000..a130c532 --- /dev/null +++ b/PyPIC3D/deposition/current_methods.py @@ -0,0 +1,12 @@ +CURRENT_J_FROM_RHOV = 0 +CURRENT_ESIRKEPOV = 1 + + +def encode_current_calculation(current_calculation): + current_codes = { + "j_from_rhov": CURRENT_J_FROM_RHOV, + "esirkepov": CURRENT_ESIRKEPOV, + } + if current_calculation not in current_codes: + raise ValueError("Unsupported current_calculation. Use 'j_from_rhov' or 'esirkepov'.") + return current_codes[current_calculation] diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 7c003a07..5d99f9b3 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -15,7 +15,6 @@ ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( - fold_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells, ) @@ -29,6 +28,60 @@ def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_ return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d +def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2): + """ + Fold two-guard Esirkepov current deposits using particle boundary codes. + + Internal tile faces are pure ownership transfers. At global particle + boundaries, periodic deposits wrap, reflecting deposits mirror into the + same-side physical layer with the Esirkepov component sign, and absorbing + deposits are discarded. + """ + + g = int(num_guard_cells) + particle_bc = world.get("particle_boundary_conditions", {"x": 0, "y": 0, "z": 0}) + + def fold_axis(field_tiles, axis, bc): + 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 + + 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"]) + field_tiles = fold_axis(field_tiles, 1, particle_bc["y"]) + field_tiles = fold_axis(field_tiles, 2, particle_bc["z"]) + + return field_tiles + + +def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=2): + return tuple( + fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells) + for component_axis, component in enumerate(field_tiles) + ) + + @jit def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=None): """ @@ -234,7 +287,7 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, upda tz, ) - J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=2) + J_tiles = fold_tiled_esirkepov_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=2) J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells=2) return J_tiles diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index 1f23be73..740f6c42 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -24,6 +24,15 @@ def _tile_shape_from_world_or_field(world, field): return _tile_shape_from_field(field) +def _guard_depth_from_field(field, tile_shape): + guard_depths = tuple((int(local_width) - int(tile_width)) // 2 for local_width, tile_width in zip(field.shape[3:], tile_shape)) + if any(int(local_width) != int(tile_width) + 2 * guard_depth for local_width, tile_width, guard_depth in zip(field.shape[3:], tile_shape, guard_depths)): + raise ValueError("Tiled field local shape is incompatible with world['tile_shape'].") + if guard_depths[0] != guard_depths[1] or guard_depths[1] != guard_depths[2]: + raise ValueError("Output assembly requires the same guard depth on each local field axis.") + return guard_depths[0] + + def scalar_field_for_output(field, world): """ Return an ordinary ghost-celled scalar field for file formats. @@ -48,7 +57,8 @@ def vector_field_for_output(field, world): return field tile_shape = _tile_shape_from_world_or_field(world, field[0]) - return assemble_tiled_vector_field(field, world, tile_shape) + num_guard_cells = _guard_depth_from_field(field[0], tile_shape) + return assemble_tiled_vector_field(field, world, tile_shape, num_guard_cells=num_guard_cells) def fields_for_output(fields, world): diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 9e8e36bc..7ae6bde4 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -1,7 +1,8 @@ import jax.numpy as jnp +import jax from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles -from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current +from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV from PyPIC3D.particles.tiled_particle_refresh import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, @@ -71,33 +72,38 @@ def time_loop_electrodynamic_tiled( ) # use the selected tiled pusher for particle velocities - use_esirkepov = J_func is tiled_esirkepov_current - # Esirkepov needs old and new particle positions. The deposition kernel - # predicts the new positions locally, then the actual particle state is - # advanced and retiled after the current has been computed. - - if use_esirkepov: - J_tiles = J_func(particles, J_tiles, constants, world) + if J_func is None: + def current_deposition(particles, J_tiles, constants, world): + return direct_J_from_tiled_particles( + particles, + J_tiles, + constants, + world, + filter="none", + ) + else: + current_deposition = J_func + + def esirkepov_step(state): + particles, J_tiles, overflow_previous = state + # Esirkepov needs old and new particle positions. The deposition kernel + # predicts the new positions locally, then the actual particle state is + # advanced and retiled after the current has been computed. + J_tiles = current_deposition(particles, J_tiles, constants, world) particles = update_tiled_particle_positions(particles, world["dt"]) particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow - else: + return particles, J_tiles, overflow + + def direct_current_step(state): + particles, J_tiles, overflow_previous = state particles = update_tiled_particle_positions(particles, world["dt"]/2) # update particle positions to the centered direct-current deposition time particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow # wrap periodic particles and move them into their owning tiles. - if J_func is None: - J_tiles = direct_J_from_tiled_particles( - particles, - J_tiles, - constants, - world, - filter="none", - ) - else: - J_tiles = J_func(particles, J_tiles, constants, world) + J_tiles = current_deposition(particles, J_tiles, constants, world) # deposit current directly into tile-local Yee current arrays particles = update_tiled_particle_positions(particles, world["dt"]/2) @@ -105,7 +111,15 @@ def time_loop_electrodynamic_tiled( particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow # refresh tile ownership after the full position update. - + return particles, J_tiles, overflow + + current_calculation = world.get("current_calculation", CURRENT_J_FROM_RHOV) + particles, J_tiles, overflow = jax.lax.cond( + current_calculation == CURRENT_ESIRKEPOV, + esirkepov_step, + direct_current_step, + (particles, J_tiles, overflow_previous), + ) if pml_state is None: E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 5c714adf..7507865d 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -66,6 +66,11 @@ from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.current_methods import ( + CURRENT_ESIRKEPOV, + CURRENT_J_FROM_RHOV, + encode_current_calculation, +) from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current @@ -124,18 +129,16 @@ def _tile_shape_from_parameters(simulation_parameters): ) -CURRENT_J_FROM_RHOV = 0 -CURRENT_ESIRKEPOV = 1 +def _encode_current_calculation(current_calculation): + return encode_current_calculation(current_calculation) -def _encode_current_calculation(current_calculation): - current_codes = { - "j_from_rhov": CURRENT_J_FROM_RHOV, - "esirkepov": CURRENT_ESIRKEPOV, - } - if current_calculation not in current_codes: - raise ValueError("Unsupported current_calculation. Use 'j_from_rhov' or 'esirkepov'.") - return current_codes[current_calculation] +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): @@ -224,6 +227,7 @@ def default_parameters(): "particle_tile_nx": 1, # number of x cells per shared field/particle tile "particle_tile_ny": 1, # number of y cells per shared field/particle tile "particle_tile_nz": 1, # 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 } @@ -302,6 +306,7 @@ def initialize_simulation(toml_file): relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] validate_particle_pusher(particle_pusher) + _validate_current_filter_contract(simulation_parameters) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] # set the simulation parameters diff --git a/PyPIC3D/particles/tiled_particle_initialization.py b/PyPIC3D/particles/tiled_particle_initialization.py index 1e7a99cf..4dca373f 100644 --- a/PyPIC3D/particles/tiled_particle_initialization.py +++ b/PyPIC3D/particles/tiled_particle_initialization.py @@ -60,6 +60,8 @@ def to_tiled_particles(particles, world, simulation_parameters): tile_counts[tx[p], ty[p], tz[p], species_index] += 1 max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 + capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) + max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) x = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) u = jnp.zeros_like(x) diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 7fe5e34d..a370c0b2 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -320,6 +320,36 @@ def apply_tiled_conducting_bc(E_tiles, world): return Ex, Ey, Ez +def digital_filter_tiled_scalar(field_tiles, alpha): + """ + Apply the standard six-neighbor digital filter to each tile interior. + + Tile halos play the same role as global ghost cells in ``digital_filter``: + the filtered values replace only the physical tile interiors. + """ + + neighbor_weight = (1 - alpha) / 6 + filtered = ( + alpha * field_tiles[:, :, :, 1:-1, 1:-1, 1:-1] + + neighbor_weight * field_tiles[:, :, :, :-2, 1:-1, 1:-1] + + neighbor_weight * field_tiles[:, :, :, 2:, 1:-1, 1:-1] + + neighbor_weight * field_tiles[:, :, :, 1:-1, :-2, 1:-1] + + neighbor_weight * field_tiles[:, :, :, 1:-1, 2:, 1:-1] + + neighbor_weight * field_tiles[:, :, :, 1:-1, 1:-1, :-2] + + neighbor_weight * field_tiles[:, :, :, 1:-1, 1:-1, 2:] + ) + + return field_tiles.at[:, :, :, 1:-1, 1:-1, 1:-1].set(filtered) + + +def digital_filter_tiled_vector(field_tiles, alpha): + """ + Apply the standard field filter component-wise in tile-major storage. + """ + + return tuple(digital_filter_tiled_scalar(component, alpha) for component in field_tiles) + + def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): """ Add tile-ghost deposits into owning interiors, then clear ghosts. @@ -486,6 +516,15 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ + (C**2 * curl_z - Jz[:, :, :, gx:-gx, gy:-gy, gz:-gz] / eps) * dt ) + if pml_state is None: + Ex, Ey, Ez = update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + else: + Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world) + # refresh tile halos before the digital field filter, matching the global + # ghost-cell order in the standard Yee solver. + + Ex, Ey, Ez = digital_filter_tiled_vector((Ex, Ey, Ez), constants.get("alpha", 1.0)) + Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world) if pml_state is None: @@ -502,7 +541,7 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pm have been refreshed from neighbor tiles or field boundary conditions. """ - del constants, curl_func, tile_shape + del curl_func, tile_shape if pml_state is None: Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world) @@ -538,6 +577,15 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pm 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) + if pml_state is None: + Bx, By, Bz = update_tiled_vector_ghost_cells((Bx, By, Bz), world) + else: + Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world) + # refresh tile halos before the digital field filter, matching the global + # ghost-cell order in the standard Yee solver. + + Bx, By, Bz = digital_filter_tiled_vector((Bx, By, Bz), constants.get("alpha", 1.0)) + if pml_state is None: return update_tiled_vector_ghost_cells((Bx, By, Bz), world) diff --git a/tests/pml_test.py b/tests/pml_test.py index 6b682c37..77d738ef 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -17,6 +17,7 @@ from PyPIC3D.solvers.first_order_yee import update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, + empty_tiled_vector_field, tile_vector_field, update_tiled_B, update_tiled_E, @@ -418,6 +419,52 @@ def step(E_tiles, B_tiles, tiled_pml_state): self.assertTrue(jnp.isfinite(final_energy)) self.assertLess(float(final_energy), 0.65 * float(initial_energy)) + def test_tiled_pml_step_reads_two_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) + E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + 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) + + J_one_guard = tile_vector_field(J, world, tile_shape) + Jx, Jy, Jz = J_one_guard + Jx = Jx.at[:, :, :, 1:-1, 1:-1, 1:-1].set(0.25) + J_one_guard = (Jx, Jy, Jz) + J_two_guard = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + Jx2, Jy2, Jz2 = J_two_guard + Jx2 = Jx2.at[:, :, :, 2:-2, 2:-2, 2:-2].set(0.25) + J_two_guard = (Jx2, Jy2, Jz2) + + pml_one = initialize_tiled_pml_state(world, tile_shape) + pml_two = initialize_tiled_pml_state(world, tile_shape) + E_one, pml_one = update_tiled_E( + E_tiles, B_tiles, J_one_guard, world, constants, lambda *args: None, tile_shape, pml_one + ) + B_one, pml_one = update_tiled_B( + E_one, B_tiles, world, constants, lambda *args: None, tile_shape, pml_one + ) + E_two, pml_two = update_tiled_E( + E_tiles, B_tiles, J_two_guard, world, constants, lambda *args: None, tile_shape, pml_two + ) + B_two, pml_two = update_tiled_B( + E_two, B_tiles, world, constants, lambda *args: None, tile_shape, pml_two + ) + + for one_component, two_component in zip(E_one + B_one, E_two + B_two): + self.assertTrue(jnp.allclose(one_component, two_component, rtol=1.0e-12, atol=1.0e-12)) + for one_state, two_state in zip(pml_one[0] + pml_one[1], pml_two[0] + pml_two[1]): + self.assertTrue(jnp.allclose(one_state, two_state, rtol=1.0e-12, atol=1.0e-12)) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_esirkepov_integration_test.py b/tests/tiled_esirkepov_integration_test.py new file mode 100644 index 00000000..3602b159 --- /dev/null +++ b/tests/tiled_esirkepov_integration_test.py @@ -0,0 +1,91 @@ +import os +import tempfile +import unittest + +import jax +import numpy as np + +from PyPIC3D.__main__ import run_PyPIC3D + + +jax.config.update("jax_enable_x64", True) + + +class TiledEsirkepovIntegrationTest(unittest.TestCase): + def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(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 esirkepov output smoke", + "output_dir": tmpdir, + "solver": "tiled_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", + "fast_backend": "default", + "particle_pusher": "boris", + "relativistic": False, + }, + "plotting": { + "plotting_interval": 1, + "plot_vtk_vectors": True, + "plot_vtk_particles": True, + "plot_openpmd_fields": True, + "plot_openpmd_particles": True, + "plot_phasespace": False, + }, + "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, + }, + } + + _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world = run_PyPIC3D(config) + _E, _B, J, *_rest = fields + + self.assertEqual(J[0].shape[-3:], (6, 5, 5)) + self.assertEqual(tuple(int(width) for width in world["tile_shape"]), (2, 1, 1)) + for filename in ( + "total_energy.txt", + "electric_field_energy.txt", + "magnetic_field_energy.txt", + "kinetic_energy.txt", + "total_momentum.txt", + ): + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", filename))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "fields.h5"))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles.h5"))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice", "vector_field_slice_000000000.vtk"))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 8d7cfbb5..4f958ee2 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -1,21 +1,29 @@ import os import tempfile import unittest +import functools import jax import jax.numpy as jnp import numpy as np import toml -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current +from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles +from PyPIC3D.diagnostics.output_adapters import fields_for_output from PyPIC3D.initialization import CURRENT_ESIRKEPOV, initialize_simulation +from PyPIC3D.particles.tiled_particle_refresh import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, + empty_tiled_scalar_field, empty_tiled_vector_field, tile_vector_field, update_tiled_E, @@ -25,8 +33,23 @@ class TestTiledEsirkepovCurrent(unittest.TestCase): - def _build_world(self, Nx=8, Ny=1, Nz=1, dt=0.05): - x_wind, y_wind, z_wind = 4.0, 1.0, 1.0 + 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, @@ -38,9 +61,9 @@ def _build_world(self, Nx=8, Ny=1, Nz=1, dt=0.05): "y_wind": y_wind, "z_wind": z_wind, "dt": dt, - "shape_factor": 1, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "shape_factor": shape_factor, + "boundary_conditions": boundary_conditions, + "particle_boundary_conditions": particle_boundary_conditions, "current_calculation": CURRENT_ESIRKEPOV, "current_guard_cells": 2, } @@ -79,6 +102,96 @@ def _species(self, world, x1): dt=world["dt"], ) + def _species_from_arrays(self, world, x, u, weight=0.5, active_mask=None): + return particle_species( + name="electrons", + N_particles=x.shape[0], + charge=-1.0, + mass=1.0, + weight=weight, + T=1.0, + x1=x[:, 0], + x2=x[:, 1], + x3=x[:, 2], + v1=u[:, 0], + v2=u[:, 1], + v3=u[:, 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=active_mask, + ) + + 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 _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) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + old_species = self._species_from_arrays(world, x_old, u) + new_species = self._species_from_arrays(world, x_old + u * world["dt"], u) + tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + + J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) + J_tiles = tiled_esirkepov_current( + tiled_particles, + empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + constants, + world, + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=2) + + 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_empty_tiled_vector_field_can_allocate_two_current_guard_cells(self): world = self._build_world(Nx=8, Ny=1, Nz=1) tile_shape = (2, 1, 1) @@ -100,6 +213,27 @@ def test_two_guard_current_tiles_assemble_to_one_guard_global_current(self): self.assertEqual(assembled[0].shape, (10, 3, 3)) self.assertEqual(float(assembled[0][3, 1, 1]), 3.0) + def test_output_adapter_infers_two_guard_current_depth(self): + world = self._build_world(Nx=8, Ny=1, Nz=1) + 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)) + E_tiles = tile_vector_field(zeros, world, tile_shape) + B_tiles = tile_vector_field(zeros, world, tile_shape) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + 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_tiled_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} @@ -146,6 +280,123 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): f"max diff {jnp.max(jnp.abs(tiled_component - reference_component))}", ) + 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) + 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, + ) + 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) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + 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], + ] + ) + old_species = self._species_from_arrays(world, x_old, u) + tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + rho_tiles = empty_tiled_scalar_field(world, tile_shape) + + rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants) + J_tiles = tiled_esirkepov_current( + tiled_particles, + empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + constants, + world, + ) + new_particles = update_tiled_particle_positions(tiled_particles, world["dt"]) + new_particles, overflow = refresh_tiled_particle_tiles(new_particles, world, tile_shape) + rho_new = compute_tiled_rho_from_tiled_particles(new_particles, rho_tiles, world, constants) + + self.assertFalse(bool(overflow)) + drhodt = (rho_new[:, :, :, 1:-1, 1:-1, 1:-1] - rho_old[:, :, :, 1:-1, 1:-1, 1:-1]) / world["dt"] + dJxdx = (J_tiles[0][:, :, :, 2:-2, 2:-2, 2:-2] - J_tiles[0][:, :, :, 1:-3, 2:-2, 2:-2]) / 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") @@ -207,6 +458,58 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): self.assertEqual(int(world["current_guard_cells"]), 2) self.assertEqual(int(world["current_calculation"]), CURRENT_ESIRKEPOV) + 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": "tiled_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", + "fast_backend": "default", + "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]) @@ -283,6 +586,89 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): 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": "tiled_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", + "fast_backend": "default", + "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, curl_func, _J_func, relativistic, particle_pusher = initialize_simulation(config) + alias_J_func = functools.partial(tiled_esirkepov_current) + particles, fields = loop( + particles, + fields, + world, + constants, + curl_func, + alias_J_func, + solver, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + old_species = self._species_from_arrays( + 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))), + ) + new_species = self._species_from_arrays( + world, + jnp.column_stack((jnp.asarray(x_initial + vx_initial * world["dt"]), jnp.zeros(4), jnp.zeros(4))), + jnp.column_stack((jnp.asarray(vx_initial), jnp.zeros(4), jnp.zeros(4))), + ) + reference_J = Esirkepov_current([new_species], self._empty_J(world), constants, world) + tiled_J = assemble_tiled_vector_field(fields[2], world, tuple(int(width) for width in world["tile_shape"]), num_guard_cells=2) + + 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/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py index 35f56644..9ff776a1 100644 --- a/tests/tiled_particle_initialization_test.py +++ b/tests/tiled_particle_initialization_test.py @@ -87,6 +87,52 @@ def test_to_tiled_particles_packs_existing_particle_species_list(self): self.assertTrue(bool(particles.update_u2[1, 1, 0, 0, 1])) self.assertTrue(bool(particles.update_u3[1, 1, 0, 0, 1])) + def test_to_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 = particle_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 = to_tiled_particles([species], world, simulation_parameters) + + self.assertEqual(particles.active.shape[-1], 6) + self.assertEqual(int(jnp.sum(particles.active)), 3) + def test_load_tiled_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]) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index c0871e02..08784ae3 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -19,6 +19,7 @@ from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, tile_vector_field from PyPIC3D.utils import build_yee_grid, compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells jax.config.update("jax_enable_x64", True) @@ -299,6 +300,79 @@ def _assert_tiled_state_matches_standard( def test_validate_field_solver_accepts_tiled_yee(self): validate_field_solver("tiled_yee") + def test_tiled_yee_field_update_applies_standard_field_filter(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} + tile_shape = (2, 1, 1) + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + + x = jnp.arange(world["Nx"], dtype=jnp.float64)[:, None, None] + seed = jnp.zeros(shape) + Ex = update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ey = update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ez = update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bx = update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + By = update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bz = update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + E = (Ex, Ey, Ez) + B = (Bx, By, Bz) + J = (seed, seed, seed) + rho = seed + phi = seed + external_fields = ((seed, seed, seed), (seed, seed, seed)) + + _, reference_fields = time_loop_electrodynamic( + [], + (E, B, J, rho, phi, external_fields, None), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_fields = ( + tile_vector_field(E, world, tile_shape), + tile_vector_field(B, world, tile_shape), + tile_vector_field(J, world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape), + tile_vector_field(external_fields[1], world, tile_shape), + ), + None, + ) + _, tiled_fields = time_loop_electrodynamic_tiled( + [], + tiled_fields, + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J, + solver="tiled_yee", + relativistic=False, + particle_pusher="boris", + ) + + reference_E, reference_B, *_ = reference_fields + E_tiles, B_tiles, *_ = tiled_fields + self._assert_vector_fields_close( + reference_E, + assemble_tiled_vector_field(E_tiles, world, tile_shape), + 1, + "E", + ) + self._assert_vector_fields_close( + reference_B, + assemble_tiled_vector_field(B_tiles, world, tile_shape), + 1, + "B", + ) + def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(self): with tempfile.TemporaryDirectory() as tmpdir: x_path = os.path.join(tmpdir, "x.npy") From 32dd054f77918701a6367de3f3ac5e0dbb019e18 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 30 Jun 2026 10:41:07 -0500 Subject: [PATCH 20/83] generalizing tile code to work with arbitrary numbers of guard cells --- PyPIC3D/__main__.py | 56 ++- PyPIC3D/deposition/direct_deposition_tiled.py | 142 ++++-- PyPIC3D/deposition/esirkepov_tiled.py | 62 ++- PyPIC3D/deposition/rho_tiled.py | 86 ++-- PyPIC3D/diagnostics/output_adapters.py | 30 +- PyPIC3D/electrodynamic_tiled.py | 31 +- PyPIC3D/electrostatic_tiled.py | 8 +- PyPIC3D/initialization.py | 27 +- PyPIC3D/pusher/tiled_pusher.py | 26 +- PyPIC3D/solvers/electrostatic_yee.py | 35 +- PyPIC3D/solvers/yee_tiled.py | 463 +++++++++++------- tests/direct_deposition_tiled_test.py | 161 +++++- tests/openpmd_test.py | 1 + tests/pml_test.py | 46 +- tests/tiled_electrostatic_test.py | 7 +- tests/tiled_esirkepov_test.py | 93 ++-- tests/tiled_fluid_quantities_test.py | 8 +- tests/tiled_pusher_test.py | 4 + tests/tiled_rho_test.py | 4 + tests/tiled_yee_integration_test.py | 217 +++++++- tests/yee_tiled_test.py | 13 +- 21 files changed, 1087 insertions(+), 433 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 57767361..0af810ef 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -68,8 +68,6 @@ def run_PyPIC3D(config_file): loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic, particle_pusher = initialize_simulation(config_file) # initialize the simulation - jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher')) - dt = world['dt'] output_dir = simulation_parameters['output_dir'] vertex_grid = tuple(g[1:-1] for g in world['grids']['vertex']) @@ -78,7 +76,16 @@ def run_PyPIC3D(config_file): scalar_field_names = ["rho", "mass_density"] vector_field_names = ["E", "B", "J"] tiled_run = simulation_parameters["solver"] == "tiled_yee" + tile_shape = tuple(int(width) for width in world["tile_shape"]) if tiled_run else None + g = int(world["guard_cells"]) if tiled_run else None particle_species_names = simulation_parameters.get("particle_species_names") + if tiled_run: + jit_loop = jax.jit( + loop, + static_argnames=('curl_func', 'J_func', 'solver', 'tile_shape', 'g', 'relativistic', 'particle_pusher'), + ) + else: + jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher')) E, B, J, rho, phi, external_fields, *rest = fields # unpack the fields @@ -133,8 +140,12 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_scalars']: if getattr(rho, "ndim", 0) == 6: - rho = compute_tiled_rho_from_tiled_particles(particles, rho, world, constants) - mass_density = compute_tiled_mass_density_from_tiled_particles(particles, rho, world) + rho = compute_tiled_rho_from_tiled_particles( + particles, rho, world, constants, tile_shape=tile_shape, g=g + ) + mass_density = compute_tiled_mass_density_from_tiled_particles( + particles, rho, world, tile_shape=tile_shape, g=g + ) rho_output = scalar_field_for_output(rho, world) mass_density_output = scalar_field_for_output(mass_density, world) else: @@ -183,17 +194,32 @@ def run_PyPIC3D(config_file): fields = (E, B, J, rho, phi, external_fields, *rest) # repack the fields for non-tiled diagnostics that updated rho - particles, fields = jit_loop( - particles, - fields, - world, - constants, - curl_func, - J_func, - solver, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) + if tiled_run: + particles, fields = jit_loop( + particles, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=tile_shape, + g=g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + else: + 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 _raise_if_tiled_particles_overflowed(fields, simulation_parameters) # fixed tile capacity overflow would silently drop particles; stop as diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index eff76716..49d79a57 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -9,13 +9,74 @@ ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( + digital_filter_tiled_vector, fold_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells, ) -from PyPIC3D.utils import bilinear_filter, digital_filter -def digital_filter_tiled_current_density(J_tiles, alpha, world): +def bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=1, tile_shape=None): + """ + Apply the global bilinear current filter to compact current tiles. + + The filter is the same separable [1, 2, 1]/4 stencil used by + ``bilinear_filter`` on assembled fields. Only the physical tile cells are + overwritten; guard cells are refreshed before and after the stencil. + """ + + g = int(num_guard_cells) + active = slice(g, -g) + backward = slice(g - 1, -g - 1) + forward = slice(g + 1, None if g == 1 else -g + 1) + + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) + + def filter_component(component): + center = component[:, :, :, active, active, active] + filtered = ( + 8.0 * center + + 4.0 * ( + component[:, :, :, backward, active, active] + + component[:, :, :, forward, active, active] + + component[:, :, :, active, backward, active] + + component[:, :, :, active, forward, active] + + component[:, :, :, active, active, backward] + + component[:, :, :, active, active, forward] + ) + + 2.0 * ( + component[:, :, :, backward, backward, active] + + component[:, :, :, backward, forward, active] + + component[:, :, :, forward, backward, active] + + component[:, :, :, forward, forward, active] + + component[:, :, :, backward, active, backward] + + component[:, :, :, backward, active, forward] + + component[:, :, :, forward, active, backward] + + component[:, :, :, forward, active, forward] + + component[:, :, :, active, backward, backward] + + component[:, :, :, active, backward, forward] + + component[:, :, :, active, forward, backward] + + component[:, :, :, active, forward, forward] + ) + + ( + component[:, :, :, backward, backward, backward] + + component[:, :, :, backward, backward, forward] + + component[:, :, :, backward, forward, backward] + + component[:, :, :, backward, forward, forward] + + component[:, :, :, forward, backward, backward] + + component[:, :, :, forward, backward, forward] + + component[:, :, :, forward, forward, backward] + + component[:, :, :, forward, forward, forward] + ) + ) / 64.0 + return component.at[:, :, :, active, active, active].set(filtered) + + J_tiles = tuple(filter_component(component) for component in J_tiles) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) + + return J_tiles + + +def digital_filter_tiled_current_density(J_tiles, alpha, world, num_guard_cells=1, tile_shape=None): """ Apply the global digital current filter to compact current tiles. @@ -25,32 +86,37 @@ def digital_filter_tiled_current_density(J_tiles, alpha, world): filtered guard-cell values. """ - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells, tile_shape) - def filter_component(component_tiles): - filter_tiles = jax.vmap(jax.vmap(jax.vmap(lambda tile: digital_filter(tile, alpha)))) - return filter_tiles(component_tiles) - - J_tiles = tuple(filter_component(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) + J_tiles = digital_filter_tiled_vector(J_tiles, alpha, num_guard_cells) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells, tile_shape) return J_tiles -def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d): +def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_cells): """ Build the ghost-celled coordinate line for one compact tile. - The first local grid point is the global ghost-cell origin shifted by the - tile's active-cell offset. With this convention, local active indices are - 1..tile_n and local ghost indices are 0 and tile_n + 1. + The first local grid point is shifted by the configured tile guard depth. + With one guard cell this is the old compact-tile convention; with two guard + cells the extra point gives quadratic shape factors their full local + stencil at tile faces. """ offsets = jnp.arange(local_n, dtype=global_axis_grid.dtype) - return global_axis_grid[0] + (offsets + tile_index * tile_n) * d + return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d -@partial(jit, static_argnames=("filter",)) +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) + + +@partial(jit, static_argnames=("filter", "tile_shape", "g")) def direct_J_from_tiled_particles( tiled_particles, J_tiles, @@ -58,6 +124,8 @@ def direct_J_from_tiled_particles( world, grid=None, filter="bilinear", + tile_shape=None, + g=1, ): """ Compute direct current deposition using tile-local stencils only. @@ -78,13 +146,15 @@ def direct_J_from_tiled_particles( Jx_tiles, Jy_tiles, Jz_tiles = J_tiles ntx, nty, ntz = Jx_tiles.shape[:3] - tile_nx = Jx_tiles.shape[3] - 2 - tile_ny = Jx_tiles.shape[4] - 2 - tile_nz = Jx_tiles.shape[5] - 2 - local_Nx = tile_nx + 2 - local_Ny = tile_ny + 2 - local_Nz = tile_nz + 2 + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(g) + local_Nx = tile_nx + 2 * g + local_Ny = tile_ny + 2 * g + local_Nz = tile_nz + 2 * g shape_factor = world["shape_factor"] + reduced_x = int(ntx) == 1 and tile_nx == 1 + reduced_y = int(nty) == 1 and tile_ny == 1 + reduced_z = int(ntz) == 1 and tile_nz == 1 Jx_template = jnp.zeros_like(Jx_tiles[0, 0, 0]) Jy_template = jnp.zeros_like(Jy_tiles[0, 0, 0]) @@ -104,9 +174,9 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, active = active_tile.reshape(-1).astype(x.dtype) dq = (charge_tile * weight_tile).reshape(-1) / (dx * dy * dz) - x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx) - y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy) - z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz) + x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) + y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) + z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz, g) x, x0, deltax_node, xpts = prepare_particle_axis_stencil( x, @@ -163,12 +233,12 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, y_weights_face = jnp.asarray(y_weights_face) z_weights_face = jnp.asarray(z_weights_face) - xpts, x_weights_node = collapse_axis_stencil(xpts, x_weights_node, local_Nx, ghost_cells=True) - _, x_weights_face = collapse_axis_stencil(xpts, x_weights_face, local_Nx, ghost_cells=True) - ypts, y_weights_node = collapse_axis_stencil(ypts, y_weights_node, local_Ny, ghost_cells=True) - _, y_weights_face = collapse_axis_stencil(ypts, y_weights_face, local_Ny, ghost_cells=True) - zpts, z_weights_node = collapse_axis_stencil(zpts, z_weights_node, local_Nz, ghost_cells=True) - _, z_weights_face = collapse_axis_stencil(zpts, z_weights_face, local_Nz, ghost_cells=True) + 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) tile_Jx = Jx_template tile_Jy = Jy_template @@ -218,14 +288,12 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, tz, ) - J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) + J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, g, tile_shape) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) if filter == "bilinear": - apply_filter = jax.vmap(jax.vmap(jax.vmap(bilinear_filter))) - J_tiles = tuple(apply_filter(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world) + J_tiles = bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=g, tile_shape=tile_shape) elif filter == "digital": - J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"], world) + J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"], world, num_guard_cells=g, tile_shape=tile_shape) - return J_tiles \ No newline at end of file + return J_tiles diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 5d99f9b3..9a856bb6 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -1,6 +1,7 @@ import jax from jax import jit import jax.numpy as jnp +from functools import partial from PyPIC3D.boundary_conditions.grid_and_stencil import ( compute_particle_anchor, @@ -28,9 +29,9 @@ def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_ return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d -def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2): +def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=1, tile_shape=None): """ - Fold two-guard Esirkepov current deposits using particle boundary codes. + Fold Esirkepov current deposits using the configured guard depth. Internal tile faces are pure ownership transfers. At global particle boundaries, periodic deposits wrap, reflecting deposits mirror into the @@ -39,14 +40,43 @@ def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_gua """ g = int(num_guard_cells) + if tile_shape is None: + tile_shape = tuple(int(width) - 2 * g for width in field_tiles.shape[3:]) + 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}) - def fold_axis(field_tiles, axis, bc): + 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 + if reduced_axis: + ghost_sum = ( + jnp.sum(lower_ghost, axis=3, keepdims=True) + + jnp.sum(upper_ghost, axis=3, keepdims=True) + ) + + def periodic_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(ghost_sum) + + def reflecting_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(sign * ghost_sum) + + 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, :, :, :, :, :]) @@ -68,24 +98,24 @@ def absorbing_boundary(tiles): 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"]) - field_tiles = fold_axis(field_tiles, 1, particle_bc["y"]) - field_tiles = fold_axis(field_tiles, 2, particle_bc["z"]) + 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 fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=2): +def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): return tuple( - fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells) + fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells, tile_shape) for component_axis, component in enumerate(field_tiles) ) -@jit -def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=None): +@partial(jit, static_argnames=("tile_shape", "g")) +def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=None, tile_shape=None, g=1): """ - Deposit Esirkepov current into tile-local two-guard current buffers. + Deposit Esirkepov current into tile-local current buffers. ``tiled_particles.x`` is the old particle position after the velocity push. The future position used in Esirkepov's charge-conserving difference is @@ -104,10 +134,8 @@ def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=Non Jx_tiles, Jy_tiles, Jz_tiles = J_tiles ntx, nty, ntz = Jx_tiles.shape[:3] - g = 2 - tile_nx = Jx_tiles.shape[3] - 2 * g - tile_ny = Jx_tiles.shape[4] - 2 * g - tile_nz = Jx_tiles.shape[5] - 2 * g + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(g) local_Nx = tile_nx + 2 * g local_Ny = tile_ny + 2 * g local_Nz = tile_nz + 2 * g @@ -287,7 +315,7 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, upda tz, ) - J_tiles = fold_tiled_esirkepov_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=2) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells=2) + J_tiles = fold_tiled_esirkepov_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=g, tile_shape=tile_shape) + J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells=g, tile_shape=tile_shape) return J_tiles diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index f6cad0b5..69ff36e0 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -1,6 +1,7 @@ import jax from jax import jit import jax.numpy as jnp +from functools import partial from PyPIC3D.boundary_conditions.grid_and_stencil import ( BC_PERIODIC, @@ -10,6 +11,7 @@ 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.solvers.yee_tiled import ( + digital_filter_tiled_scalar, fold_tiled_ghost_cells, tile_scalar_field, update_tiled_ghost_cells, @@ -151,9 +153,18 @@ def _diagnostic_mass_position(tiled_particles, world): return x, y, z -def _tile_axis(axis, tile_index, cells_per_tile): - start = tile_index * cells_per_tile - return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) +def _tile_axis(axis, tile_index, cells_per_tile, num_guard_cells, d): + local_n = cells_per_tile + 2 * num_guard_cells + offsets = jnp.arange(local_n, dtype=axis.dtype) + return axis[0] + (offsets + tile_index * cells_per_tile - (num_guard_cells - 1)) * d + + +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) def _deposit_tiled_scalar_moment_to_tiles( @@ -162,6 +173,8 @@ def _deposit_tiled_scalar_moment_to_tiles( world, grid, scalar_weight, + tile_shape, + g, ): dx = world["dx"] dy = world["dy"] @@ -171,16 +184,21 @@ def _deposit_tiled_scalar_moment_to_tiles( bc_z = world["boundary_conditions"]["z"] shape_factor = world["shape_factor"] - tile_nx = rho_tiles.shape[3] - 2 - tile_ny = rho_tiles.shape[4] - 2 - tile_nz = rho_tiles.shape[5] - 2 + tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] + g = int(g) + local_Nx = tile_nx + 2 * g + local_Ny = tile_ny + 2 * g + local_Nz = tile_nz + 2 * g ntx, nty, ntz = rho_tiles.shape[:3] + reduced_x = int(ntx) == 1 and tile_nx == 1 + reduced_y = int(nty) == 1 and tile_ny == 1 + reduced_z = int(ntz) == 1 and tile_nz == 1 def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile): local_grid = ( - _tile_axis(grid[0], tx, tile_nx), - _tile_axis(grid[1], ty, tile_ny), - _tile_axis(grid[2], tz, tile_nz), + _tile_axis(grid[0], tx, tile_nx, g, dx), + _tile_axis(grid[1], ty, tile_ny, g, dy), + _tile_axis(grid[2], tz, tile_nz, g, dz), ) x = x_tile[..., 0].reshape(-1) @@ -192,7 +210,7 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile x, _, deltax, xpts = prepare_particle_axis_stencil( x, local_grid[0], - tile_nx + 2, + local_Nx, shape_factor, bc_x, wind=world["x_wind"], @@ -201,7 +219,7 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile y, _, deltay, ypts = prepare_particle_axis_stencil( y, local_grid[1], - tile_ny + 2, + local_Ny, shape_factor, bc_y, wind=world["y_wind"], @@ -210,7 +228,7 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile z, _, deltaz, zpts = prepare_particle_axis_stencil( z, local_grid[2], - tile_nz + 2, + local_Nz, shape_factor, bc_z, wind=world["z_wind"], @@ -231,11 +249,11 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile y_weights = jnp.asarray(y_weights) z_weights = jnp.asarray(z_weights) - xpts, x_weights = collapse_axis_stencil(xpts, x_weights, tile_nx + 2, ghost_cells=True) - ypts, y_weights = collapse_axis_stencil(ypts, y_weights, tile_ny + 2, ghost_cells=True) - zpts, z_weights = collapse_axis_stencil(zpts, z_weights, tile_nz + 2, ghost_cells=True) + xpts, x_weights = _collapse_tiled_axis_stencil(xpts, x_weights, local_Nx, reduced_x, g) + ypts, y_weights = _collapse_tiled_axis_stencil(ypts, y_weights, local_Ny, reduced_y, g) + zpts, z_weights = _collapse_tiled_axis_stencil(zpts, z_weights, local_Nz, reduced_z, g) - rho_tile = jnp.zeros((tile_nx + 2, tile_ny + 2, tile_nz + 2), dtype=rho_tiles.dtype) + rho_tile = jnp.zeros((local_Nx, local_Ny, local_Nz), dtype=rho_tiles.dtype) for i in range(xpts.shape[0]): for j in range(ypts.shape[0]): @@ -265,21 +283,16 @@ def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile scalar_weight, ) - rho_tiles = fold_tiled_ghost_cells(rho_tiles, world) - rho_tiles = update_tiled_ghost_cells(rho_tiles, world) + rho_tiles = fold_tiled_ghost_cells(rho_tiles, world, g, tile_shape) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world, g, tile_shape) return rho_tiles -def _digital_filter_tiled_scalar(field_tiles, alpha, world): - field_tiles = update_tiled_ghost_cells(field_tiles, world) - - filter_tiles = digital_filter - filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) - filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) - filter_tiles = jax.vmap(filter_tiles, in_axes=(0, None), out_axes=0) - - return filter_tiles(field_tiles, alpha) +def _digital_filter_tiled_scalar(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_tiled_scalar(field_tiles, alpha, g) @jit @@ -309,8 +322,8 @@ def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, gri return rho -@jit -def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, grid=None): +@partial(jit, static_argnames=("tile_shape", "g")) +def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, grid=None, tile_shape=None, g=1): """Compute charge density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] @@ -322,11 +335,13 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, co world, grid, scalar_weight, + tile_shape, + g, ) alpha = constants["alpha"] - rho_tiles = _digital_filter_tiled_scalar(rho_tiles, alpha, world) - rho_tiles = update_tiled_ghost_cells(rho_tiles, world) + rho_tiles = _digital_filter_tiled_scalar(rho_tiles, alpha, world, g, tile_shape) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world, g, tile_shape) return rho_tiles @@ -349,15 +364,16 @@ def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid= ) -@jit -def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world, grid=None): +@partial(jit, static_argnames=("tile_shape", "g")) +def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world, grid=None, tile_shape=None, g=1): """Compute mass density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] scalar_weight = tiled_particles.mass * tiled_particles.weight position = _diagnostic_mass_position(tiled_particles, world) - tile_shape = tuple(int(width) - 2 for width in rho_tiles.shape[3:]) + g = int(g) + tile_shape = tuple(int(width) for width in tile_shape) Nx = int(rho_tiles.shape[0]) * tile_shape[0] Ny = int(rho_tiles.shape[1]) * tile_shape[1] Nz = int(rho_tiles.shape[2]) * tile_shape[2] @@ -378,4 +394,4 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, position=position, ) - return tile_scalar_field(rho, world, tile_shape) + return tile_scalar_field(rho, world, tile_shape, num_guard_cells=g) diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index 740f6c42..c30c2dc5 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -14,23 +14,12 @@ def _is_tiled_vector(field): ) -def _tile_shape_from_field(field): - return tuple(int(width) - 2 for width in field.shape[3:]) +def _tile_shape_from_world(world): + return tuple(int(width) for width in world["tile_shape"]) -def _tile_shape_from_world_or_field(world, field): - if "tile_shape" in world: - return tuple(int(width) for width in world["tile_shape"]) - return _tile_shape_from_field(field) - - -def _guard_depth_from_field(field, tile_shape): - guard_depths = tuple((int(local_width) - int(tile_width)) // 2 for local_width, tile_width in zip(field.shape[3:], tile_shape)) - if any(int(local_width) != int(tile_width) + 2 * guard_depth for local_width, tile_width, guard_depth in zip(field.shape[3:], tile_shape, guard_depths)): - raise ValueError("Tiled field local shape is incompatible with world['tile_shape'].") - if guard_depths[0] != guard_depths[1] or guard_depths[1] != guard_depths[2]: - raise ValueError("Output assembly requires the same guard depth on each local field axis.") - return guard_depths[0] +def _guard_depth_from_world(world): + return int(world["guard_cells"]) def scalar_field_for_output(field, world): @@ -44,8 +33,9 @@ def scalar_field_for_output(field, world): if not _is_tiled_scalar(field): return field - tile_shape = _tile_shape_from_world_or_field(world, field) - return assemble_tiled_scalar_field(field, world, tile_shape) + 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): @@ -56,9 +46,9 @@ def vector_field_for_output(field, world): if not _is_tiled_vector(field): return field - tile_shape = _tile_shape_from_world_or_field(world, field[0]) - num_guard_cells = _guard_depth_from_field(field[0], tile_shape) - return assemble_tiled_vector_field(field, world, tile_shape, num_guard_cells=num_guard_cells) + 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): diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 7ae6bde4..1f407e41 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -21,6 +21,7 @@ def time_loop_electrodynamic_tiled( J_func, solver, tile_shape=None, + g=None, relativistic=True, particle_pusher="boris", ): @@ -38,19 +39,20 @@ def time_loop_electrodynamic_tiled( # unpack the tile-major field state. ``overflow_previous`` carries the # particle-retile overflow diagnostic out to the Python driver. - if tile_shape is None: - tile_shape = tuple(int(width) - 2 for width in E_tiles[0].shape[3:]) + if tile_shape is None or g is None: + raise ValueError("tiled electrodynamic updates require explicit tile_shape and g.") + g = int(g) if not hasattr(particles, "active"): if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) else: E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state ) B_tiles, pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state + E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state ) fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) return particles, fields @@ -67,19 +69,22 @@ def time_loop_electrodynamic_tiled( world, constants, tile_shape, + g, relativistic=relativistic, particle_pusher=particle_pusher, ) # use the selected tiled pusher for particle velocities if J_func is None: - def current_deposition(particles, J_tiles, constants, world): + def current_deposition(particles, J_tiles, constants, world, tile_shape=None, g=None): return direct_J_from_tiled_particles( particles, J_tiles, constants, world, filter="none", + tile_shape=tile_shape, + g=g, ) else: current_deposition = J_func @@ -89,7 +94,7 @@ def esirkepov_step(state): # Esirkepov needs old and new particle positions. The deposition kernel # predicts the new positions locally, then the actual particle state is # advanced and retiled after the current has been computed. - J_tiles = current_deposition(particles, J_tiles, constants, world) + J_tiles = current_deposition(particles, J_tiles, constants, world, tile_shape=tile_shape, g=g) particles = update_tiled_particle_positions(particles, world["dt"]) particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow @@ -103,7 +108,7 @@ def direct_current_step(state): overflow = overflow_previous | overflow # wrap periodic particles and move them into their owning tiles. - J_tiles = current_deposition(particles, J_tiles, constants, world) + J_tiles = current_deposition(particles, J_tiles, constants, world, tile_shape=tile_shape, g=g) # deposit current directly into tile-local Yee current arrays particles = update_tiled_particle_positions(particles, world["dt"]/2) @@ -122,17 +127,17 @@ def direct_current_step(state): ) if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) else: E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state ) # update electric field from B and the supplied tiled current if pml_state is None: - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) else: - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state) + B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state) # update magnetic field from the newly updated electric field fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) diff --git a/PyPIC3D/electrostatic_tiled.py b/PyPIC3D/electrostatic_tiled.py index dee48805..acd00820 100644 --- a/PyPIC3D/electrostatic_tiled.py +++ b/PyPIC3D/electrostatic_tiled.py @@ -18,6 +18,7 @@ def time_loop_electrostatic_tiled( J_func, solver, tile_shape=None, + g=None, relativistic=True, particle_pusher="boris", ): @@ -42,8 +43,9 @@ def time_loop_electrostatic_tiled( else: E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields - if tile_shape is None: - tile_shape = tuple(int(width) - 2 for width in E_tiles[0].shape[3:]) + if tile_shape is None or g is None: + raise ValueError("tiled electrostatic updates require explicit tile_shape and g.") + g = int(g) push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) # particles see evolved fields plus prescribed external fields @@ -55,6 +57,7 @@ def time_loop_electrostatic_tiled( world, constants, tile_shape, + g, relativistic=relativistic, particle_pusher=particle_pusher, ) @@ -76,6 +79,7 @@ def time_loop_electrostatic_tiled( solver, "periodic", tile_shape, + g, ) # solve electrostatic fields from tiled charge density diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 7507865d..78acf964 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -224,6 +224,7 @@ 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": 1, # tile guard-cell depth shared by tiled fields and tiled currents "particle_tile_nx": 1, # number of x cells per shared field/particle tile "particle_tile_ny": 1, # number of y cells per shared field/particle tile "particle_tile_nz": 1, # number of z cells per shared field/particle tile @@ -306,6 +307,7 @@ def initialize_simulation(toml_file): relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] validate_particle_pusher(particle_pusher) + guard_cells = int(simulation_parameters["guard_cells"]) _validate_current_filter_contract(simulation_parameters) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] @@ -354,8 +356,8 @@ def initialize_simulation(toml_file): 'y_wind': y_wind, 'z_wind': z_wind, 'shape_factor': simulation_parameters['shape_factor'], + 'guard_cells': guard_cells, 'current_calculation': _encode_current_calculation(simulation_parameters['current_calculation']), - 'current_guard_cells': 2 if simulation_parameters['current_calculation'] == "esirkepov" else 1, 'boundary_conditions': { 'x': _encode_field_bc(simulation_parameters['x_bc']), 'y': _encode_field_bc(simulation_parameters['y_bc']), @@ -384,6 +386,10 @@ 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 electrostatic: B_grid, E_grid = build_collocated_grid(world) @@ -543,17 +549,20 @@ def initialize_simulation(toml_file): if solver == "tiled_yee": tile_shape = _tile_shape_from_parameters(simulation_parameters) world["tile_shape"] = tile_shape + guard_cells = int(world["guard_cells"]) particles = to_tiled_particles(particles, world, simulation_parameters) - E = tile_vector_field(E, world, tile_shape) - B = tile_vector_field(B, world, tile_shape) + E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) + B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) if simulation_parameters["current_calculation"] == "esirkepov": - J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2, dtype=E[0].dtype) + J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) + J_func = tiled_esirkepov_current else: - J = tile_vector_field(J, world, tile_shape) + J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) + J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) external_E, external_B = external_fields external_fields = ( - tile_vector_field(external_E, world, tile_shape), - tile_vector_field(external_B, world, tile_shape), + tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), + tile_vector_field(external_B, world, tile_shape, num_guard_cells=guard_cells), ) simulation_parameters["fast_backend"] = "tiled" simulation_parameters["tile_shape"] = tile_shape @@ -561,8 +570,8 @@ def initialize_simulation(toml_file): if electrostatic: if solver == "tiled_yee": - rho = tile_scalar_field(rho, world, tile_shape) - phi = tile_scalar_field(phi, world, tile_shape) + rho = tile_scalar_field(rho, world, tile_shape, num_guard_cells=guard_cells) + phi = tile_scalar_field(phi, world, tile_shape, num_guard_cells=guard_cells) fields = (E, B, J, rho, phi, external_fields) # define the fields tuple for the electrostatic solver else: diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index e646c553..611efce5 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -12,18 +12,19 @@ from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle -def _tile_axis(axis, tile_index, cells_per_tile): - start = tile_index * cells_per_tile - return jax.lax.dynamic_slice(axis, (start,), (cells_per_tile + 2,)) +def _tile_axis(axis, tile_index, cells_per_tile, num_guard_cells, d): + local_n = cells_per_tile + 2 * num_guard_cells + offsets = jnp.arange(local_n, dtype=axis.dtype) + return axis[0] + (offsets + tile_index * cells_per_tile - (num_guard_cells - 1)) * d -@partial(jit, static_argnames=("tile_shape", "relativistic", "particle_pusher")) -def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, relativistic=True, particle_pusher="boris"): +@partial(jit, static_argnames=("tile_shape", "g", "relativistic", "particle_pusher")) +def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, g, relativistic=True, particle_pusher="boris"): """ 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 one-cell field halos on each compact tile provide the + position. The configured field halos on each compact tile provide the neighboring Yee data needed by the interpolation stencil near tile faces. """ @@ -36,6 +37,7 @@ def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, til Ex_tiles, Ey_tiles, Ez_tiles = E_tiles Bx_tiles, By_tiles, Bz_tiles = B_tiles + g = int(g) ntx, nty, ntz = tiled_particles.x.shape[:3] @@ -50,12 +52,12 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til q = charge_tile.reshape(-1) m = mass_tile.reshape(-1) - center_x = _tile_axis(center_grid[0], tx, tile_nx) - center_y = _tile_axis(center_grid[1], ty, tile_ny) - center_z = _tile_axis(center_grid[2], tz, tile_nz) - vertex_x = _tile_axis(vertex_grid[0], tx, tile_nx) - vertex_y = _tile_axis(vertex_grid[1], ty, tile_ny) - vertex_z = _tile_axis(vertex_grid[2], tz, tile_nz) + center_x = _tile_axis(center_grid[0], tx, tile_nx, g, world["dx"]) + center_y = _tile_axis(center_grid[1], ty, tile_ny, g, world["dy"]) + center_z = _tile_axis(center_grid[2], tz, tile_nz, g, world["dz"]) + vertex_x = _tile_axis(vertex_grid[0], tx, tile_nx, g, world["dx"]) + vertex_y = _tile_axis(vertex_grid[1], ty, tile_ny, g, world["dy"]) + vertex_z = _tile_axis(vertex_grid[2], tz, tile_nz, g, world["dz"]) Ex_grid = vertex_x, center_y, center_z Ey_grid = center_x, vertex_y, center_z diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index c18dcec0..9a75c785 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -165,7 +165,7 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver return (Ex_full, Ey_full, Ez_full), phi, rho -def _centered_tiled_electrostatic_gradient(phi_tiles, world): +def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): """ Compute ``E = -grad(phi)`` on compact scalar tiles. @@ -178,27 +178,31 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world): 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 = update_tiled_ghost_cells(phi_tiles, world) + phi_tiles = 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[:, :, :, 1:-1, 1:-1, 1:-1].set( - -1.0 * (phi_tiles[:, :, :, 2:, 1:-1, 1:-1] - phi_tiles[:, :, :, :-2, 1:-1, 1:-1]) / (2.0 * dx) + 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[:, :, :, 1:-1, 1:-1, 1:-1].set( - -1.0 * (phi_tiles[:, :, :, 1:-1, 2:, 1:-1] - phi_tiles[:, :, :, 1:-1, :-2, 1:-1]) / (2.0 * dy) + 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[:, :, :, 1:-1, 1:-1, 1:-1].set( - -1.0 * (phi_tiles[:, :, :, 1:-1, 1:-1, 2:] - phi_tiles[:, :, :, 1:-1, 1:-1, :-2]) / (2.0 * dz) + Ez = Ez.at[:, :, :, active, active, active].set( + -1.0 * (phi_tiles[:, :, :, active, active, forward] - phi_tiles[:, :, :, active, active, backward]) / (2.0 * dz) ) - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) -def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, phi_tiles, solver, bc, tile_shape): +def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, phi_tiles, solver, bc, tile_shape, g): """ Compute electrostatic fields from tiled rho deposition and a global Poisson solve. @@ -213,9 +217,10 @@ def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] - rho_tiles = compute_tiled_rho_from_tiled_particles(particles, rho_tiles, world, constants) - rho = assemble_tiled_scalar_field(rho_tiles, world, tile_shape) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape) + g = int(g) + rho_tiles = compute_tiled_rho_from_tiled_particles(particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + rho = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) @@ -228,7 +233,7 @@ def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) # keep the same global phi post-processing order as the untiled solver - phi_tiles = tile_scalar_field(phi, world, tile_shape) - E_tiles = _centered_tiled_electrostatic_gradient(phi_tiles, world) + phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) + E_tiles = _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g) return E_tiles, phi_tiles, rho_tiles diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index a370c0b2..f857e5c1 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -15,6 +15,30 @@ def _tile_axis_count(n_cells, cells_per_tile): return int(n_cells) // int(cells_per_tile) +def _active_slice(g): + return slice(g, -g) + + +def _forward_slice(g): + return slice(g + 1, None if g == 1 else -g + 1) + + +def _backward_slice(g): + return slice(g - 1, -g - 1) + + +def _reduced_tiled_axes(field_tiles, tile_shape, g): + if tile_shape is None: + tile_shape = tuple(int(width) - 2 * int(g) for width in field_tiles.shape[3:]) + 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 empty_tiled_scalar_field(world, tile_shape, num_guard_cells=1, dtype=None): """ Allocate an empty tile-major scalar field with ``num_guard_cells`` guards. @@ -73,7 +97,7 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=1): 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 update_tiled_ghost_cells(field_tiles, world, g) + return 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) @@ -152,7 +176,7 @@ def _assemble_scalar_tiles(field_tiles, world, tile_shape): return assemble_tiled_scalar_field(field_tiles, world, tile_shape) -def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): +def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Refresh tile halos using the field boundary conditions. @@ -164,6 +188,7 @@ def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): """ g = int(num_guard_cells) + 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"] @@ -171,47 +196,83 @@ def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): # x-halos: internal tile boundaries exchange neighboring interiors. At a # conducting global wall, the exterior ghost face is zero, matching # update_ghost_cells on the assembled field. - 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, - ) + 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 the global arrays. - 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, - ) + 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. - 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, - ) + 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 @@ -225,12 +286,12 @@ def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): return update_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1): +def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Refresh tile halos for each component of a vector field. """ - return tuple(update_tiled_ghost_cells(component, world, num_guard_cells) for component in field_tiles) + return tuple(update_tiled_ghost_cells(component, world, num_guard_cells, tile_shape) for component in field_tiles) def update_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): @@ -241,7 +302,7 @@ def update_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): return tuple(update_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) -def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1): +def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Refresh tile halos without wrapping across PML-active global walls. """ @@ -256,18 +317,18 @@ def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1): 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}} - return update_tiled_ghost_cells(field_tiles, pml_world, num_guard_cells) + 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=1): +def update_tiled_vector_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Refresh tile halos for a vector field with PML-active exterior walls. """ - return tuple(update_tiled_ghost_cells_for_pml(component, world, num_guard_cells) for component in field_tiles) + return tuple(update_tiled_ghost_cells_for_pml(component, world, num_guard_cells, tile_shape) for component in field_tiles) -def apply_tiled_conducting_bc(E_tiles, world): +def apply_tiled_conducting_bc(E_tiles, world, num_guard_cells=1): """ Zero tangential electric-field components on global conducting faces. """ @@ -277,42 +338,45 @@ def apply_tiled_conducting_bc(E_tiles, world): 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, :, :, 1, :, :].set(0.0).at[-1, :, :, -2, :, :].set(0.0), + 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, :, :, 1, :, :].set(0.0).at[-1, :, :, -2, :, :].set(0.0), + 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, :, :, 1, :].set(0.0).at[:, -1, :, :, -2, :].set(0.0), + 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, :, :, 1, :].set(0.0).at[:, -1, :, :, -2, :].set(0.0), + 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, :, :, 1].set(0.0).at[:, :, -1, :, :, -2].set(0.0), + 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, :, :, 1].set(0.0).at[:, :, -1, :, :, -2].set(0.0), + lambda f: f.at[:, :, 0, :, :, lower].set(0.0).at[:, :, -1, :, :, upper].set(0.0), lambda f: f, operand=Ey, ) @@ -320,7 +384,7 @@ def apply_tiled_conducting_bc(E_tiles, world): return Ex, Ey, Ez -def digital_filter_tiled_scalar(field_tiles, alpha): +def digital_filter_tiled_scalar(field_tiles, alpha, num_guard_cells=1): """ Apply the standard six-neighbor digital filter to each tile interior. @@ -329,28 +393,33 @@ def digital_filter_tiled_scalar(field_tiles, alpha): """ neighbor_weight = (1 - alpha) / 6 + g = int(num_guard_cells) + active = _active_slice(g) + backward = _backward_slice(g) + forward = _forward_slice(g) + filtered = ( - alpha * field_tiles[:, :, :, 1:-1, 1:-1, 1:-1] - + neighbor_weight * field_tiles[:, :, :, :-2, 1:-1, 1:-1] - + neighbor_weight * field_tiles[:, :, :, 2:, 1:-1, 1:-1] - + neighbor_weight * field_tiles[:, :, :, 1:-1, :-2, 1:-1] - + neighbor_weight * field_tiles[:, :, :, 1:-1, 2:, 1:-1] - + neighbor_weight * field_tiles[:, :, :, 1:-1, 1:-1, :-2] - + neighbor_weight * field_tiles[:, :, :, 1:-1, 1:-1, 2:] + alpha * field_tiles[:, :, :, active, active, active] + + neighbor_weight * field_tiles[:, :, :, backward, active, active] + + neighbor_weight * field_tiles[:, :, :, forward, active, active] + + neighbor_weight * field_tiles[:, :, :, active, backward, active] + + neighbor_weight * field_tiles[:, :, :, active, forward, active] + + neighbor_weight * field_tiles[:, :, :, active, active, backward] + + neighbor_weight * field_tiles[:, :, :, active, active, forward] ) - return field_tiles.at[:, :, :, 1:-1, 1:-1, 1:-1].set(filtered) + return field_tiles.at[:, :, :, active, active, active].set(filtered) -def digital_filter_tiled_vector(field_tiles, alpha): +def digital_filter_tiled_vector(field_tiles, alpha, num_guard_cells=1): """ Apply the standard field filter component-wise in tile-major storage. """ - return tuple(digital_filter_tiled_scalar(component, alpha) for component in field_tiles) + return tuple(digital_filter_tiled_scalar(component, alpha, num_guard_cells) for component in field_tiles) -def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): +def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Add tile-ghost deposits into owning interiors, then clear ghosts. @@ -362,72 +431,115 @@ def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1): """ g = int(num_guard_cells) + 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_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_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) - 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_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) - 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) + 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 @@ -441,12 +553,12 @@ def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): return fold_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1): +def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): """ Fold tile-ghost deposits for each vector component. """ - return tuple(fold_tiled_ghost_cells(component, world, num_guard_cells) for component in field_tiles) + return tuple(fold_tiled_ghost_cells(component, world, num_guard_cells, tile_shape) for component in field_tiles) def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): @@ -457,7 +569,7 @@ def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): return tuple(fold_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) -def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, pml_state=None): +def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -465,17 +577,19 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ have been refreshed from neighbor tiles or field boundary conditions. """ - del curl_func, tile_shape + del curl_func Ex, Ey, Ez = E_tiles + tile_nx, tile_ny, tile_nz = tile_shape + g = int(g) + active = _active_slice(g) + forward = _forward_slice(g) if pml_state is None: - Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world) + Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world, g, tile_shape) else: - Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml(B_tiles, world) + Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml(B_tiles, world, g, tile_shape) Jx, Jy, Jz = J_tiles - gx = (Jx.shape[3] - (Ex.shape[3] - 2)) // 2 - gy = (Jx.shape[4] - (Ex.shape[4] - 2)) // 2 - gz = (Jx.shape[5] - (Ex.shape[5] - 2)) // 2 + current = _active_slice(g) dt = world["dt"] dx, dy, dz = world["dx"], world["dy"], world["dz"] @@ -485,12 +599,12 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ # 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[:, :, :, 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 + 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 @@ -503,37 +617,37 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ pml_state, ) - Ex = Ex.at[:, :, :, 1:-1, 1:-1, 1:-1].set( - Ex[:, :, :, 1:-1, 1:-1, 1:-1] - + (C**2 * curl_x - Jx[:, :, :, gx:-gx, gy:-gy, gz:-gz] / 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[:, :, :, gx:-gx, gy:-gy, gz:-gz] / 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[:, :, :, gx:-gx, gy:-gy, gz:-gz] / eps) * dt + Ez = Ez.at[:, :, :, active, active, active].set( + Ez[:, :, :, active, active, active] + + (C**2 * curl_z - Jz[:, :, :, current, current, current] / eps) * dt ) if pml_state is None: - Ex, Ey, Ez = update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + Ex, Ey, Ez = update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) else: - Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world) + Ex, Ey, Ez = 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. - Ex, Ey, Ez = digital_filter_tiled_vector((Ex, Ey, Ez), constants.get("alpha", 1.0)) + Ex, Ey, Ez = digital_filter_tiled_vector((Ex, Ey, Ez), constants.get("alpha", 1.0), g) - Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world) + Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) if pml_state is None: - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world) + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) - return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world), pml_state + return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state -def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pml_state=None): +def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state=None): """ Update compact tiled magnetic fields without assembling a global field. @@ -541,26 +655,29 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pm have been refreshed from neighbor tiles or field boundary conditions. """ - del curl_func, tile_shape + del curl_func + Bx, By, Bz = B_tiles + tile_nx, tile_ny, tile_nz = tile_shape + g = int(g) + active = _active_slice(g) + backward = _backward_slice(g) if pml_state is None: - Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world) + Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world, g, tile_shape) else: - Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml(E_tiles, world) - Bx, By, Bz = B_tiles - + Ex, Ey, Ez = 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[:, :, :, 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 + 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 @@ -573,20 +690,20 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, pm pml_state, ) - 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) + 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, By, Bz = update_tiled_vector_ghost_cells((Bx, By, Bz), world) + Bx, By, Bz = update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) else: - Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world) + Bx, By, Bz = 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. - Bx, By, Bz = digital_filter_tiled_vector((Bx, By, Bz), constants.get("alpha", 1.0)) + Bx, By, Bz = digital_filter_tiled_vector((Bx, By, Bz), constants.get("alpha", 1.0), g) if pml_state is None: - return update_tiled_vector_ghost_cells((Bx, By, Bz), world) + return update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) - return update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world), pml_state + return update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape), pml_state diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 6e284419..89de8196 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -42,6 +42,7 @@ def _build_world(self, Nx=8, Ny=6, Nz=4, dt=0.05, boundary_conditions=None): "z_wind": z_wind, "dt": dt, "shape_factor": 1, + "guard_cells": 1, "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) @@ -59,13 +60,14 @@ def _empty_J_tiles(self, world, simulation_parameters): simulation_parameters["particle_tile_nz"], ) 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, - tile_ny + 2, - tile_nz + 2, + tile_nx + 2 * g, + tile_ny + 2 * g, + tile_nz + 2 * g, ) return (jnp.zeros(shape), jnp.zeros(shape), jnp.zeros(shape)) @@ -112,14 +114,134 @@ def _compare_tiled_to_standard(self, particles, world, simulation_parameters): constants, world, filter="none", + tile_shape=tile_shape, + g=int(world["guard_cells"]), ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + g = int(world["guard_cells"]) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) 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)) + 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, + } + + species = particle_species( + name="quadratic current", + N_particles=6, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.55, -0.52, -0.03, 0.49, 0.55, 1.45]), + x2=jnp.array([-1.10, -0.55, -0.03, 0.02, 0.52, 1.05]), + x3=jnp.array([-0.70, -0.04, 0.03, 0.31, 0.49, 0.72]), + v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16]), + v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12]), + v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + self._compare_tiled_to_standard([species], world, simulation_parameters) + + 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, + } + + species = particle_species( + name="quadratic reduced current", + N_particles=8, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.95, -1.51, -1.02, -0.48, 0.02, 0.47, 1.04, 1.88]), + x2=jnp.zeros(8), + x3=jnp.zeros(8), + v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16, 0.09, -0.13]), + v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12, -0.15, 0.05]), + v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19, 0.16, -0.07]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + self._compare_tiled_to_standard([species], world, simulation_parameters) + + 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, + } + constants = {"C": 3.0e8, "alpha": 1.0} + species = particle_species( + name="bilinear quadratic reduced current", + N_particles=8, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.95, -1.51, -1.02, -0.48, 0.02, 0.47, 1.04, 1.88]), + x2=jnp.zeros(8), + x3=jnp.zeros(8), + v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16, 0.09, -0.13]), + v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12, -0.15, 0.05]), + v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19, 0.16, -0.07]), + 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_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) + + J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") + J_tiles = direct_J_from_tiled_particles( + tiled_particles, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="bilinear", + tile_shape=tile_shape, + g=int(world["guard_cells"]), + ) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) + + 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)) + 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) @@ -184,6 +306,23 @@ def test_fold_tiled_ghost_cells_two_guard_layers_adds_deposits_to_neighbors(self self.assertTrue(jnp.allclose(folded[:, :, :, :num_guard_cells, :, :], 0.0)) self.assertTrue(jnp.allclose(folded[:, :, :, -num_guard_cells:, :, :], 0.0)) + 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) + 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 = 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)) + def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): world = self._build_world( Nx=4, @@ -281,6 +420,7 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): dt=world["dt"], ) tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tile_shape = self._tile_shape(simulation_parameters) J_tiles = direct_J_from_tiled_particles( tiled_particles, @@ -288,10 +428,11 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): constants, world, filter="none", + tile_shape=tile_shape, + g=int(world["guard_cells"]), ) - tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="none") - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) 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)) @@ -388,8 +529,10 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): constants, world, filter="digital", + tile_shape=tile_shape, + g=int(world["guard_cells"]), ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) 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)) @@ -433,8 +576,10 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): constants, world, filter="bilinear", + tile_shape=tile_shape, + g=int(world["guard_cells"]), ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) 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)) diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index f6b4d028..b7bbc3d9 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -128,6 +128,7 @@ def test_write_openpmd_fields_assembles_tiled_fields_before_output(self): "Ny": 2, "Nz": 2, "tile_shape": (2, 1, 1), + "guard_cells": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } tiled_fields = ( diff --git a/tests/pml_test.py b/tests/pml_test.py index 77d738ef..b380db75 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -355,6 +355,7 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): constants, lambda *args: None, tile_shape, + 1, tiled_pml_state, ) B_tiles, tiled_pml_state = update_tiled_B( @@ -364,6 +365,7 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): constants, lambda *args: None, tile_shape, + 1, tiled_pml_state, ) @@ -404,10 +406,10 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) def step(E_tiles, B_tiles, tiled_pml_state): E_tiles, tiled_pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, tiled_pml_state + E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, 1, tiled_pml_state ) B_tiles, tiled_pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, tiled_pml_state + E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, 1, tiled_pml_state ) return E_tiles, B_tiles, tiled_pml_state @@ -419,7 +421,7 @@ def step(E_tiles, B_tiles, tiled_pml_state): self.assertTrue(jnp.isfinite(final_energy)) self.assertLess(float(final_energy), 0.65 * float(initial_energy)) - def test_tiled_pml_step_reads_two_guard_current_without_changing_timing(self): + 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( @@ -436,34 +438,24 @@ def test_tiled_pml_step_reads_two_guard_current_without_changing_timing(self): E_tiles = tile_vector_field((Ex, Ey, Ez), world, tile_shape) B_tiles = tile_vector_field((Bx, By, Bz), world, tile_shape) - J_one_guard = tile_vector_field(J, world, tile_shape) - Jx, Jy, Jz = J_one_guard + g = 1 + 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_one_guard = (Jx, Jy, Jz) - J_two_guard = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) - Jx2, Jy2, Jz2 = J_two_guard - Jx2 = Jx2.at[:, :, :, 2:-2, 2:-2, 2:-2].set(0.25) - J_two_guard = (Jx2, Jy2, Jz2) - - pml_one = initialize_tiled_pml_state(world, tile_shape) - pml_two = initialize_tiled_pml_state(world, tile_shape) - E_one, pml_one = update_tiled_E( - E_tiles, B_tiles, J_one_guard, world, constants, lambda *args: None, tile_shape, pml_one - ) - B_one, pml_one = update_tiled_B( - E_one, B_tiles, world, constants, lambda *args: None, tile_shape, pml_one - ) - E_two, pml_two = update_tiled_E( - E_tiles, B_tiles, J_two_guard, world, constants, lambda *args: None, tile_shape, pml_two + J_tiles = (Jx, Jy, Jz) + + pml_state = initialize_tiled_pml_state(world, tile_shape) + E_after, pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, g, pml_state ) - B_two, pml_two = update_tiled_B( - E_two, B_tiles, world, constants, lambda *args: None, tile_shape, pml_two + B_after, pml_state = update_tiled_B( + E_after, B_tiles, world, constants, lambda *args: None, tile_shape, g, pml_state ) - for one_component, two_component in zip(E_one + B_one, E_two + B_two): - self.assertTrue(jnp.allclose(one_component, two_component, rtol=1.0e-12, atol=1.0e-12)) - for one_state, two_state in zip(pml_one[0] + pml_one[1], pml_two[0] + pml_two[1]): - self.assertTrue(jnp.allclose(one_state, two_state, rtol=1.0e-12, atol=1.0e-12)) + 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__": diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index fd96ee03..ce356dad 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -360,6 +360,8 @@ def test_tiled_rho_assembles_to_global_rho(self): rho_tiles, world, constants, + tile_shape=world["tile_shape"], + g=1, ) assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"]) @@ -411,6 +413,8 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): unused_curl, J_func=None, solver="tiled_yee", + tile_shape=world["tile_shape"], + g=1, relativistic=False, particle_pusher="boris", ) @@ -452,7 +456,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self tiled_loop = jax.jit( time_loop_electrostatic_tiled, - static_argnames=("curl_func", "J_func", "solver", "tile_shape", "relativistic", "particle_pusher"), + static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) for step in range(1, Nt + 1): @@ -476,6 +480,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self J_func=None, solver="tiled_yee", tile_shape=tile_shape, + g=1, relativistic=False, particle_pusher="boris", ) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 4f958ee2..ef3318c4 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -65,7 +65,7 @@ def _build_world( "boundary_conditions": boundary_conditions, "particle_boundary_conditions": particle_boundary_conditions, "current_calculation": CURRENT_ESIRKEPOV, - "current_guard_cells": 2, + "guard_cells": 1, } center_grid = ( jnp.linspace(-x_wind / 2 - world["dx"] / 2, x_wind / 2 + world["dx"] / 2, Nx + 2), @@ -142,6 +142,7 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } + g = int(world["guard_cells"]) old_species = self._species_from_arrays(world, x_old, u) new_species = self._species_from_arrays(world, x_old + u * world["dt"], u) @@ -150,11 +151,13 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) J_tiles = tiled_esirkepov_current( tiled_particles, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, + tile_shape=tile_shape, + g=g, ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=2) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) for reference_component, tiled_component in zip(J_reference, J_from_tiles): self.assertTrue( @@ -192,36 +195,48 @@ def _basic_positions_and_velocities(self, world): ) return x, u - def test_empty_tiled_vector_field_can_allocate_two_current_guard_cells(self): + def test_empty_tiled_vector_field_uses_startup_guard_cells(self): world = self._build_world(Nx=8, Ny=1, Nz=1) + world["guard_cells"] = 2 tile_shape = (2, 1, 1) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=int(world["guard_cells"])) self.assertEqual(J_tiles[0].shape, (4, 1, 1, 6, 5, 5)) self.assertTrue(jnp.allclose(J_tiles[0], 0.0)) - def test_two_guard_current_tiles_assemble_to_one_guard_global_current(self): + 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) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + g = int(world["guard_cells"]) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) 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=2) + 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_infers_two_guard_current_depth(self): + 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)) - E_tiles = tile_vector_field(zeros, world, tile_shape) - B_tiles = tile_vector_field(zeros, world, tile_shape) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + 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 = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) Jx, Jy, Jz = J_tiles Jx = Jx.at[1, 0, 0, 2, 2, 2].set(7.0) rho = jnp.zeros(shape) @@ -238,17 +253,19 @@ def test_update_tiled_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) - B_tiles = tile_vector_field(zeros, world, tile_shape) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=2) + 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 = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) Jx, Jy, Jz = J_tiles Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) - E_after = update_tiled_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants, None, tile_shape) + E_after = update_tiled_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants, None, tile_shape, g) - self.assertTrue(jnp.allclose(E_after[0][:, :, :, 1:-1, 1:-1, 1:-1], -0.5)) + 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) @@ -266,13 +283,16 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) + g = int(world["guard_cells"]) J_tiles = tiled_esirkepov_current( tiled_particles, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, + tile_shape=tile_shape, + g=g, ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=2) + J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) for reference_component, tiled_component in zip(J_reference, J_from_tiles): self.assertTrue( @@ -294,6 +314,8 @@ def test_tiled_esirkepov_matches_global_current_for_dimensions_and_shapes(self): 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) @@ -345,6 +367,7 @@ def test_tiled_esirkepov_matches_global_current_with_particle_boundaries(self): 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"]) @@ -376,22 +399,25 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): ) old_species = self._species_from_arrays(world, x_old, u) tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) - rho_tiles = empty_tiled_scalar_field(world, tile_shape) + g = int(world["guard_cells"]) + rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) - rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants) + rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) J_tiles = tiled_esirkepov_current( tiled_particles, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=2), + empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, + tile_shape=tile_shape, + g=g, ) new_particles = update_tiled_particle_positions(tiled_particles, world["dt"]) new_particles, overflow = refresh_tiled_particle_tiles(new_particles, world, tile_shape) - rho_new = compute_tiled_rho_from_tiled_particles(new_particles, rho_tiles, world, constants) + rho_new = compute_tiled_rho_from_tiled_particles(new_particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) self.assertFalse(bool(overflow)) - drhodt = (rho_new[:, :, :, 1:-1, 1:-1, 1:-1] - rho_old[:, :, :, 1:-1, 1:-1, 1:-1]) / world["dt"] - dJxdx = (J_tiles[0][:, :, :, 2:-2, 2:-2, 2:-2] - J_tiles[0][:, :, :, 1:-3, 2:-2, 2:-2]) / world["dx"] + 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))) @@ -420,6 +446,7 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): "dt": 0.01, "Nt": 1, "shape_factor": 1, + "guard_cells": 2, "particle_tile_nx": 2, "particle_tile_ny": 1, "particle_tile_nz": 1, @@ -453,9 +480,10 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): E_tiles, _B_tiles, J_tiles, *_ = fields self.assertIsInstance(particles, TiledParticles) - self.assertEqual(E_tiles[0].shape[-3:], (4, 3, 3)) + self.assertEqual(E_tiles[0].shape[-3:], (6, 5, 5)) self.assertEqual(J_tiles[0].shape[-3:], (6, 5, 5)) - self.assertEqual(int(world["current_guard_cells"]), 2) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertNotIn("current_guard_cells", world) self.assertEqual(int(world["current_calculation"]), CURRENT_ESIRKEPOV) def test_initialize_rejects_filtered_esirkepov_current(self): @@ -576,6 +604,7 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): J_func, solver, tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), relativistic=relativistic, particle_pusher=particle_pusher, ) @@ -583,7 +612,7 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): 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(fields[2][0].shape[-3:], (6, 5, 5)) + self.assertEqual(fields[2][0].shape[-3:], (4, 3, 3)) self.assertFalse(bool(fields[-1])) def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(self): @@ -649,6 +678,7 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s alias_J_func, solver, tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), relativistic=relativistic, particle_pusher=particle_pusher, ) @@ -664,7 +694,12 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s jnp.column_stack((jnp.asarray(vx_initial), jnp.zeros(4), jnp.zeros(4))), ) reference_J = Esirkepov_current([new_species], self._empty_J(world), constants, world) - tiled_J = assemble_tiled_vector_field(fields[2], world, tuple(int(width) for width in world["tile_shape"]), num_guard_cells=2) + 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)) diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 6b869677..18750816 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -131,7 +131,13 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"]) mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) - mass_tiles = compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world) + mass_tiles = compute_tiled_mass_density_from_tiled_particles( + tiled_particles, + rho_tiles, + world, + tile_shape=world["tile_shape"], + g=1, + ) mass_from_tiles = assemble_tiled_scalar_field(mass_tiles, world, world["tile_shape"]) self.assertTrue( diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index f2f4520b..5fbaad3a 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -120,6 +120,7 @@ def test_tiled_particle_push_matches_flat_boris(self): world, constants, tile_shape, + 1, relativistic=False, ) @@ -186,6 +187,7 @@ def make_species(): world, constants, tile_shape, + 1, particle_pusher="higuera_cary", ) @@ -223,6 +225,7 @@ def test_tiled_particle_push_respects_active_and_update_flags(self): world, constants, tile_shape, + 1, relativistic=False, ) @@ -286,6 +289,7 @@ def make_species(): world, constants, tile_shape, + 1, relativistic=True, ) diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 69fc4722..ab9c70ed 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -212,12 +212,16 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self rho_tiles, world, constants, + tile_shape=world["tile_shape"], + g=1, ) rho_tiles_with_zero_velocity = compute_tiled_rho_from_tiled_particles( zero_velocity_tiled_particles, rho_tiles, world, constants, + tile_shape=world["tile_shape"], + g=1, ) rho_with_velocity = assemble_tiled_scalar_field(rho_tiles_with_velocity, world, world["tile_shape"]) rho_with_zero_velocity = assemble_tiled_scalar_field(rho_tiles_with_zero_velocity, world, world["tile_shape"]) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 08784ae3..ccfa127a 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -9,14 +9,17 @@ import toml from PyPIC3D.evolve import time_loop_electrodynamic +from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, tile_vector_field +from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, empty_tiled_vector_field, tile_vector_field from PyPIC3D.utils import build_yee_grid, compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -48,6 +51,7 @@ def _build_world(self, boundary_conditions=None): "y_wind": 1.0, "z_wind": 1.0, "shape_factor": 1, + "guard_cells": 1, "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) @@ -93,7 +97,7 @@ def _active_rows(self, tiled_particles): order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def _long_two_stream_state(self): + def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): world = { "Nx": 16, "Ny": 1, @@ -105,7 +109,9 @@ def _long_two_stream_state(self): "x_wind": 4.0, "y_wind": 1.0, "z_wind": 1.0, - "shape_factor": 1, + "shape_factor": 2, + "guard_cells": 2, + "current_calculation": current_calculation, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, } @@ -165,15 +171,21 @@ def _long_two_stream_state(self): } tiled_particles = to_tiled_particles(particles, world, simulation_parameters) tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) + g = int(world["guard_cells"]) + if current_calculation == CURRENT_ESIRKEPOV: + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) + else: + J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + J_tiles, rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), ), None, ) @@ -278,21 +290,22 @@ def _assert_tiled_state_matches_standard( reference_E, reference_B, reference_J, *_ = reference_fields E_tiles, B_tiles, J_tiles, *_ = tiled_fields + num_guard_cells = int(world.get("guard_cells", 1)) self._assert_vector_fields_close( reference_E, - assemble_tiled_vector_field(E_tiles, world, tile_shape), + assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "E", ) self._assert_vector_fields_close( reference_B, - assemble_tiled_vector_field(B_tiles, world, tile_shape), + assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "B", ) self._assert_vector_fields_close( reference_J, - assemble_tiled_vector_field(J_tiles, world, tile_shape), + assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "J", ) @@ -352,8 +365,10 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J, + J_func=lambda particles, J, constants, world, tile_shape=None, g=None: J, solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -718,6 +733,8 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): unused_curl, J_func=None, solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -795,6 +812,8 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): unused_curl, J_func=None, solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -856,6 +875,8 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): unused_curl, J_func=functools.partial(direct_J_from_tiled_particles, filter="digital"), solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -914,6 +935,8 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): unused_curl, J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -925,11 +948,116 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): for reference, tiled in zip(reference_J, J_from_tiles): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_standard(self): + world = self._build_world() + world["shape_factor"] = 2 + world["guard_cells"] = 2 + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + + species = particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.array([0.03, -0.04, 0.02, -0.01]), + v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + reference_species = particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.array([0.03, -0.04, 0.02, -0.01]), + v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + _, reference_fields = time_loop_electrodynamic( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + world, + constants, + unused_curl, + J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="bilinear"), + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + + tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])), + ), + None, + ) + _, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), + solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + [reference_species], + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + 1, + ) + @unittest.skipUnless( os.environ.get("RUN_SLOW_TILED_YEE") == "1", "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", ) - def test_long_two_stream_tiled_yee_matches_standard_every_step(self): + def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_standard_every_step(self): ( reference_particles, reference_fields, @@ -939,12 +1067,12 @@ def test_long_two_stream_tiled_yee_matches_standard_every_step(self): constants, tile_shape, Nt, - ) = self._long_two_stream_state() + ) = self._long_two_stream_state(CURRENT_J_FROM_RHOV) standard_J = functools.partial(J_from_rhov, filter="none") tiled_loop = jax.jit( time_loop_electrodynamic_tiled, - static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher"), + static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) for t in range(Nt): @@ -967,6 +1095,65 @@ def test_long_two_stream_tiled_yee_matches_standard_every_step(self): unused_curl, J_func=None, solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + t + 1, + ) + + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_YEE") == "1", + "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", + ) + def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standard_every_step(self): + ( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + constants, + tile_shape, + Nt, + ) = self._long_two_stream_state(CURRENT_ESIRKEPOV) + + tiled_loop = jax.jit( + time_loop_electrodynamic_tiled, + static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), + ) + + for t in range(Nt): + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_fields, + world, + constants, + unused_curl, + J_func=Esirkepov_current, + solver="fdtd", + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + tiled_fields, + world, + constants, + unused_curl, + J_func=tiled_esirkepov_current, + solver="tiled_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index eaf42fe6..4f9e7921 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -290,6 +290,7 @@ def test_update_tiled_E_matches_standard_yee_update(self): constants, unused_curl, tile_shape, + 1, ) E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) @@ -313,6 +314,7 @@ def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(s constants, unused_curl, tile_shape, + 1, ) E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) @@ -334,6 +336,7 @@ def test_update_tiled_B_matches_standard_yee_update(self): constants, unused_curl, tile_shape, + 1, ) B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) @@ -355,6 +358,7 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s constants, unused_curl, tile_shape, + 1, ) B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) @@ -375,8 +379,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 1) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 1) E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) @@ -400,8 +404,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 1) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 1) E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) @@ -445,6 +449,7 @@ def test_tiled_evolve_updates_fields_without_pushing_particles(self): J_func=None, solver="fdtd", tile_shape=tile_shape, + g=1, relativistic=False, particle_pusher="boris", ) From dbec8ce001ab446641354bb410d48eac8a7fdcfe Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 30 Jun 2026 13:22:03 -0500 Subject: [PATCH 21/83] using species metadata structure in addition to tiled particles --- PyPIC3D/__main__.py | 25 ++-- PyPIC3D/deposition/direct_deposition_tiled.py | 16 ++- PyPIC3D/deposition/esirkepov_tiled.py | 28 ++-- PyPIC3D/deposition/rho_tiled.py | 35 +++-- PyPIC3D/diagnostics/fluid_quantities.py | 4 +- PyPIC3D/diagnostics/openPMD.py | 8 +- PyPIC3D/diagnostics/output_adapters.py | 3 +- PyPIC3D/diagnostics/plotting.py | 4 +- PyPIC3D/diagnostics/vtk.py | 4 +- PyPIC3D/electrodynamic_tiled.py | 15 +- PyPIC3D/electrostatic_tiled.py | 5 +- PyPIC3D/initialization.py | 7 +- .../particles/tiled_particle_diagnostics.py | 9 +- .../tiled_particle_initialization.py | 61 ++++---- PyPIC3D/particles/tiled_particle_refresh.py | 81 +---------- PyPIC3D/particles/tiled_particles.py | 26 ++-- PyPIC3D/pusher/tiled_pusher.py | 37 ++--- PyPIC3D/solvers/electrostatic_yee.py | 4 +- PyPIC3D/utils.py | 17 ++- docs/architecture.rst | 11 +- tests/direct_deposition_tiled_test.py | 19 ++- tests/tiled_electrostatic_test.py | 17 +-- tests/tiled_esirkepov_test.py | 21 +-- tests/tiled_fluid_quantities_test.py | 11 +- tests/tiled_particle_diagnostics_test.py | 20 +-- tests/tiled_particle_initialization_test.py | 130 ++++++++++++++---- tests/tiled_particle_refresh_test.py | 34 ++--- tests/tiled_pusher_test.py | 12 +- tests/tiled_rho_test.py | 17 ++- tests/tiled_yee_integration_test.py | 38 ++--- tests/utils_test.py | 10 -- tests/yee_tiled_test.py | 1 + 32 files changed, 378 insertions(+), 352 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 0af810ef..cbf2a304 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -65,7 +65,7 @@ def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): 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, curl_func, J_func, relativistic, particle_pusher, species_config = initialize_simulation(config_file) # initialize the simulation dt = world['dt'] @@ -91,7 +91,7 @@ def run_PyPIC3D(config_file): # 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 @@ -116,7 +116,7 @@ def run_PyPIC3D(config_file): 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 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)) @@ -124,7 +124,7 @@ def run_PyPIC3D(config_file): 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) + 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 @@ -134,27 +134,27 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_phasespace']: - write_particles_phase_space(particles, t, output_dir, species_names=particle_species_names, world=world) + write_particles_phase_space(particles, t, output_dir, species_config=species_config, species_names=particle_species_names, world=world) if plotting_parameters['plot_vtk_scalars']: if getattr(rho, "ndim", 0) == 6: rho = compute_tiled_rho_from_tiled_particles( - particles, rho, world, constants, tile_shape=tile_shape, g=g + particles, species_config, rho, world, constants, tile_shape=tile_shape, g=g ) mass_density = compute_tiled_mass_density_from_tiled_particles( - particles, rho, world, tile_shape=tile_shape, g=g + particles, species_config, rho, world, tile_shape=tile_shape, g=g ) rho_output = scalar_field_for_output(rho, world) mass_density_output = scalar_field_for_output(mass_density, world) else: if tiled_run: - rho = compute_rho_from_tiled_particles(particles, rho, world, constants) + rho = compute_rho_from_tiled_particles(particles, species_config, rho, world, constants) else: rho = compute_rho(particles, rho, world, constants) # calculate the charge density based on the particle positions - mass_density = compute_mass_density(particles, rho, world) + mass_density = compute_mass_density(particles, rho, world, species_config=species_config) # calculate the mass density based on the particle positions rho_output = rho mass_density_output = mass_density @@ -179,11 +179,11 @@ def run_PyPIC3D(config_file): # Plot the vector fields in VTK format if plotting_parameters['plot_vtk_particles']: - plot_vtk_particles(particles, plot_num, output_dir, species_names=particle_species_names, world=world) + plot_vtk_particles(particles, plot_num, output_dir, species_config=species_config, species_names=particle_species_names, world=world) # 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", species_names=particle_species_names) + write_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5", species_config=species_config, species_names=particle_species_names) # Write the particles in openPMD format if plotting_parameters['plot_openpmd_fields']: @@ -197,6 +197,7 @@ def run_PyPIC3D(config_file): if tiled_run: particles, fields = jit_loop( particles, + species_config, fields, world, constants, @@ -256,7 +257,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/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index 49d79a57..fff93360 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -119,6 +119,7 @@ def _collapse_tiled_axis_stencil(points, weights, local_n, reduced_axis, g): @partial(jit, static_argnames=("filter", "tile_shape", "g")) def direct_J_from_tiled_particles( tiled_particles, + species_config, J_tiles, constants, world, @@ -164,7 +165,9 @@ def direct_J_from_tiled_particles( # edge should land in tile ghost cells and be exchanged by the tiled fold. local_bc = 2 - def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, ty, tz): + species_weighted_charge = species_config.charge * species_config.weight + + def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): x = x_tile[..., 0].reshape(-1) y = x_tile[..., 1].reshape(-1) z = x_tile[..., 2].reshape(-1) @@ -172,7 +175,8 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, vy = u_tile[..., 1].reshape(-1) vz = u_tile[..., 2].reshape(-1) active = active_tile.reshape(-1).astype(x.dtype) - dq = (charge_tile * weight_tile).reshape(-1) / (dx * dy * dz) + q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) + dq = q / (dx * dy * dz) x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) @@ -273,15 +277,13 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, tx, ) deposit_tiles = deposit_one_tile - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 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) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( tiled_particles.x, tiled_particles.u, - tiled_particles.charge, - tiled_particles.weight, tiled_particles.active, tx, ty, diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 9a856bb6..8ca27a7b 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -113,7 +113,7 @@ def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells= @partial(jit, static_argnames=("tile_shape", "g")) -def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=None, tile_shape=None, g=1): +def tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=1): """ Deposit Esirkepov current into tile-local current buffers. @@ -148,7 +148,9 @@ def tiled_esirkepov_current(tiled_particles, J_tiles, constants, world, grid=Non Jy_template = jnp.zeros_like(Jy_tiles[0, 0, 0]) Jz_template = jnp.zeros_like(Jz_tiles[0, 0, 0]) - def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, update_x1, update_x2, update_x3, tx, ty, tz): + species_weighted_charge = species_config.charge * species_config.weight + + 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) @@ -156,12 +158,15 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, upda vy = u_tile[..., 1].reshape(-1) vz = u_tile[..., 2].reshape(-1) active = active_tile.reshape(-1).astype(old_x.dtype) - q = (charge_tile * weight_tile).reshape(-1) + q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) N_particles = old_x.shape[0] + 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) - x = old_x + jnp.where(update_x1.reshape(-1), vx * dt, 0.0) - y = old_y + jnp.where(update_x2.reshape(-1), vy * dt, 0.0) - z = old_z + jnp.where(update_x3.reshape(-1), vz * dt, 0.0) + 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) x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) @@ -297,19 +302,14 @@ def deposit_one_tile(x_tile, u_tile, charge_tile, weight_tile, active_tile, upda ) deposit_tiles = deposit_one_tile - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 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) + deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, 0, 0, 0, 0, 0), out_axes=0) Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( tiled_particles.x, tiled_particles.u, - tiled_particles.charge, - tiled_particles.weight, tiled_particles.active, - tiled_particles.update_x1, - tiled_particles.update_x2, - tiled_particles.update_x3, tx, ty, tz, diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index 69ff36e0..2d67e7fb 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -19,6 +19,13 @@ from PyPIC3D.utils 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 _deposit_tiled_scalar_moment( tiled_particles, rho, @@ -296,12 +303,15 @@ def _digital_filter_tiled_scalar(field_tiles, alpha, world, g, tile_shape): @jit -def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, grid=None): +def compute_rho_from_tiled_particles(tiled_particles, species_config, rho, world, constants, grid=None): """Compute charge density on the vertex grid from tile-major particles.""" if grid is None: grid = world["grids"]["vertex"] - scalar_weight = tiled_particles.charge * tiled_particles.weight + scalar_weight = _species_scalar_to_slots( + tiled_particles, + species_config.charge * species_config.weight, + ) rho = _deposit_tiled_scalar_moment( tiled_particles, rho, @@ -323,12 +333,15 @@ def compute_rho_from_tiled_particles(tiled_particles, rho, world, constants, gri @partial(jit, static_argnames=("tile_shape", "g")) -def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, grid=None, tile_shape=None, g=1): +def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, grid=None, tile_shape=None, g=1): """Compute charge density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] - scalar_weight = tiled_particles.charge * tiled_particles.weight + scalar_weight = _species_scalar_to_slots( + tiled_particles, + species_config.charge * species_config.weight, + ) rho_tiles = _deposit_tiled_scalar_moment_to_tiles( tiled_particles, rho_tiles, @@ -347,12 +360,15 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, co @jit -def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid=None): +def compute_mass_density_from_tiled_particles(tiled_particles, species_config, rho, world, grid=None): """Compute mass density on the vertex grid from tile-major particles.""" if grid is None: grid = world["grids"]["vertex"] - scalar_weight = tiled_particles.mass * tiled_particles.weight + 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( tiled_particles, @@ -365,12 +381,15 @@ def compute_mass_density_from_tiled_particles(tiled_particles, rho, world, grid= @partial(jit, static_argnames=("tile_shape", "g")) -def compute_tiled_mass_density_from_tiled_particles(tiled_particles, rho_tiles, world, grid=None, tile_shape=None, g=1): +def compute_tiled_mass_density_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, grid=None, tile_shape=None, g=1): """Compute mass density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] - scalar_weight = tiled_particles.mass * tiled_particles.weight + scalar_weight = _species_scalar_to_slots( + tiled_particles, + species_config.mass * species_config.weight, + ) position = _diagnostic_mass_position(tiled_particles, world) g = int(g) tile_shape = tuple(int(width) for width in tile_shape) diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index b07ec22a..06d48fc3 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -11,7 +11,7 @@ @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: @@ -29,7 +29,7 @@ def compute_mass_density(particles, rho, world): ndarray: The updated charge density array. """ if isinstance(particles, TiledParticles): - return compute_mass_density_from_tiled_particles(particles, rho, world) + return compute_mass_density_from_tiled_particles(particles, species_config, rho, world) dx = world['dx'] dy = world['dy'] diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index 99725d35..aeb57c9d 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -126,8 +126,8 @@ 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, species_names=None, world=None): - particles = particles_for_output(particles, species_names=species_names, world=world) +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. @@ -256,7 +256,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", species_names=None): +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. @@ -278,7 +278,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, species_names=species_names, world=world) + 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() diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index c30c2dc5..3d15f313 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -82,13 +82,14 @@ def fields_for_output(fields, world): return output_fields + (pml_state,) -def particles_for_output(particles, species_names=None, world=None): +def particles_for_output(particles, species_config=None, species_names=None, world=None): """ Flatten fixed-capacity tiled particle storage for diagnostics. """ return flatten_tiled_particles_by_species( particles, + species_config=species_config, species_names=species_names, world=world, ) diff --git a/PyPIC3D/diagnostics/plotting.py b/PyPIC3D/diagnostics/plotting.py index 9e283128..cae3d2d7 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -51,7 +51,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, species_names=None, world=None): +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. @@ -71,7 +71,7 @@ def write_particles_phase_space(particles, t, path, species_names=None, world=No os.makedirs(f"{path}/data/phase_space/z") # Create directory if it doesn't exist - particles = particles_for_output(particles, species_names=species_names, world=world) + 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. diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py index 92374427..ba8e378c 100644 --- a/PyPIC3D/diagnostics/vtk.py +++ b/PyPIC3D/diagnostics/vtk.py @@ -42,7 +42,7 @@ def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): # plot the electric fields in the vtk file format -def plot_vtk_particles(particles, t, path, species_names=None, world=None): +def plot_vtk_particles(particles, t, path, species_config=None, species_names=None, world=None): """ Plot the particles in VTK format. @@ -57,7 +57,7 @@ def plot_vtk_particles(particles, t, path, species_names=None, world=None): if not os.path.exists(f"{path}/data/particles"): os.makedirs(f"{path}/data/particles") - particles = particles_for_output(particles, species_names=species_names, world=world) + particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) # Tiled particle arrays are fixed-capacity by tile; VTK particle output only # receives the active diagnostic particles. diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py index 1f407e41..20e7c81e 100644 --- a/PyPIC3D/electrodynamic_tiled.py +++ b/PyPIC3D/electrodynamic_tiled.py @@ -14,6 +14,7 @@ def time_loop_electrodynamic_tiled( particles, + species_config, fields, world, constants, @@ -64,6 +65,7 @@ def time_loop_electrodynamic_tiled( particles = tiled_particle_push( particles, + species_config, push_E_tiles, push_B_tiles, world, @@ -76,9 +78,10 @@ def time_loop_electrodynamic_tiled( # use the selected tiled pusher for particle velocities if J_func is None: - def current_deposition(particles, J_tiles, constants, world, tile_shape=None, g=None): + def current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=None, g=None): return direct_J_from_tiled_particles( particles, + species_config, J_tiles, constants, world, @@ -94,24 +97,24 @@ def esirkepov_step(state): # Esirkepov needs old and new particle positions. The deposition kernel # predicts the new positions locally, then the actual particle state is # advanced and retiled after the current has been computed. - J_tiles = current_deposition(particles, J_tiles, constants, world, tile_shape=tile_shape, g=g) - particles = update_tiled_particle_positions(particles, world["dt"]) + J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + particles = update_tiled_particle_positions(particles, species_config, world["dt"]) particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow return particles, J_tiles, overflow def direct_current_step(state): particles, J_tiles, overflow_previous = state - particles = update_tiled_particle_positions(particles, world["dt"]/2) + 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) overflow = overflow_previous | overflow # wrap periodic particles and move them into their owning tiles. - J_tiles = current_deposition(particles, J_tiles, constants, world, tile_shape=tile_shape, g=g) + J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) # deposit current directly into tile-local Yee current arrays - particles = update_tiled_particle_positions(particles, world["dt"]/2) + 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) overflow = overflow_previous | overflow diff --git a/PyPIC3D/electrostatic_tiled.py b/PyPIC3D/electrostatic_tiled.py index acd00820..918d1f85 100644 --- a/PyPIC3D/electrostatic_tiled.py +++ b/PyPIC3D/electrostatic_tiled.py @@ -11,6 +11,7 @@ def time_loop_electrostatic_tiled( particles, + species_config, fields, world, constants, @@ -52,6 +53,7 @@ def time_loop_electrostatic_tiled( particles = tiled_particle_push( particles, + species_config, push_E_tiles, push_B_tiles, world, @@ -63,7 +65,7 @@ def time_loop_electrostatic_tiled( ) # push velocities using the selected tiled particle pusher - particles = update_tiled_particle_positions(particles, world["dt"]) + 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) @@ -73,6 +75,7 @@ def time_loop_electrostatic_tiled( E_tiles, phi_tiles, rho_tiles = calculate_tiled_electrostatic_fields( world, particles, + species_config, constants, rho_tiles, phi_tiles, diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 78acf964..13131438 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -546,11 +546,12 @@ def initialize_simulation(toml_file): J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) + species_config = None if solver == "tiled_yee": tile_shape = _tile_shape_from_parameters(simulation_parameters) world["tile_shape"] = tile_shape guard_cells = int(world["guard_cells"]) - particles = to_tiled_particles(particles, world, simulation_parameters) + particles, species_config = to_tiled_particles(particles, world, simulation_parameters) E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) if simulation_parameters["current_calculation"] == "esirkepov": @@ -594,11 +595,13 @@ 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, curl_func, J_func, relativistic, particle_pusher, species_config diff --git a/PyPIC3D/particles/tiled_particle_diagnostics.py b/PyPIC3D/particles/tiled_particle_diagnostics.py index d5e8e391..13c5ef9a 100644 --- a/PyPIC3D/particles/tiled_particle_diagnostics.py +++ b/PyPIC3D/particles/tiled_particle_diagnostics.py @@ -70,7 +70,7 @@ def _diagnostic_position(x, u, world): return jnp.stack((x_diagnostic, y_diagnostic, z_diagnostic), axis=-1) -def flatten_tiled_particles_by_species(tiled_particles, species_names=None, world=None): +def flatten_tiled_particles_by_species(tiled_particles, species_config=None, species_names=None, world=None): """ Flatten active tile-major particles into diagnostic species arrays. @@ -90,9 +90,10 @@ def flatten_tiled_particles_by_species(tiled_particles, species_names=None, worl x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active] u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active] - charge = tiled_particles.charge[:, :, :, species_index, :].reshape(-1)[active] - mass = tiled_particles.mass[:, :, :, species_index, :].reshape(-1)[active] - weight = tiled_particles.weight[:, :, :, species_index, :].reshape(-1)[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]) x_diagnostic = _diagnostic_position(x, u, world) diff --git a/PyPIC3D/particles/tiled_particle_initialization.py b/PyPIC3D/particles/tiled_particle_initialization.py index 4dca373f..ef542b45 100644 --- a/PyPIC3D/particles/tiled_particle_initialization.py +++ b/PyPIC3D/particles/tiled_particle_initialization.py @@ -4,7 +4,7 @@ import numpy as np from PyPIC3D.particles.particle_initialization import load_particles_from_toml -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles def _as_int(value): @@ -65,16 +65,33 @@ def to_tiled_particles(particles, world, simulation_parameters): x = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) u = jnp.zeros_like(x) - charge = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile)) - mass = jnp.zeros_like(charge) - weight = jnp.zeros_like(charge) active = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) - update_x1 = jnp.zeros_like(active) - update_x2 = jnp.zeros_like(active) - update_x3 = jnp.zeros_like(active) - update_u1 = jnp.zeros_like(active) - update_u2 = jnp.zeros_like(active) - update_u3 = jnp.zeros_like(active) + + charge = jnp.asarray([species.charge for species in particles]) + mass = jnp.asarray([species.mass for species in particles]) + weight = jnp.asarray([species.weight for species in particles]) + update_x = jnp.asarray( + [ + [ + species.update_pos and species.update_x, + species.update_pos and species.update_y, + species.update_pos and species.update_z, + ] + for species in particles + ], + dtype=bool, + ) + update_u = jnp.asarray( + [ + [ + species.update_v and species.update_vx, + species.update_v and species.update_vy, + species.update_v and species.update_vz, + ] + for species in particles + ], + dtype=bool, + ) next_slot = np.zeros_like(tile_counts) @@ -96,32 +113,16 @@ def to_tiled_particles(particles, world, simulation_parameters): u = u.at[index + (1,)].set(u2[p]) u = u.at[index + (2,)].set(u3[p]) - charge = charge.at[index].set(species.charge) - mass = mass.at[index].set(species.mass) - weight = weight.at[index].set(species.weight) active = active.at[index].set(species.active_mask[p]) - update_x1 = update_x1.at[index].set(species.update_pos and species.update_x) - update_x2 = update_x2.at[index].set(species.update_pos and species.update_y) - update_x3 = update_x3.at[index].set(species.update_pos and species.update_z) - update_u1 = update_u1.at[index].set(species.update_v and species.update_vx) - update_u2 = update_u2.at[index].set(species.update_v and species.update_vy) - update_u3 = update_u3.at[index].set(species.update_v and species.update_vz) - - return TiledParticles( - x=x, - u=u, + species_config = SpeciesConfig( charge=charge, mass=mass, weight=weight, - active=active, - update_x1=update_x1, - update_x2=update_x2, - update_x3=update_x3, - update_u1=update_u1, - update_u2=update_u2, - update_u3=update_u3, + update_x=update_x, + update_u=update_u, ) + return TiledParticles(x=x, u=u, active=active), species_config def load_tiled_particles_from_toml(config, simulation_parameters, world, constants): diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index 88a91b74..a0859e6f 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -50,21 +50,22 @@ def _particle_tile_indices(x, y, z, world, tile_shape, tile_counts): ) -def update_tiled_particle_positions(tiled_particles, dt): +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 & tiled_particles.update_x1, x[..., 0] + dx, x[..., 0])) - x = x.at[..., 1].set(jnp.where(tiled_particles.active & tiled_particles.update_x2, x[..., 1] + dy, x[..., 1])) - x = x.at[..., 2].set(jnp.where(tiled_particles.active & tiled_particles.update_x3, x[..., 2] + dz, x[..., 2])) + 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) @@ -138,16 +139,7 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): new_x = jnp.zeros_like(tiled_particles.x) new_u = jnp.zeros_like(tiled_particles.u) - new_charge = jnp.zeros_like(tiled_particles.charge) - new_mass = jnp.zeros_like(tiled_particles.mass) - new_weight = jnp.zeros_like(tiled_particles.weight) new_active = jnp.zeros_like(tiled_particles.active) - new_update_x1 = jnp.zeros_like(tiled_particles.update_x1) - new_update_x2 = jnp.zeros_like(tiled_particles.update_x2) - new_update_x3 = jnp.zeros_like(tiled_particles.update_x3) - new_update_u1 = jnp.zeros_like(tiled_particles.update_u1) - new_update_u2 = jnp.zeros_like(tiled_particles.update_u2) - new_update_u3 = jnp.zeros_like(tiled_particles.update_u3) overflow = jnp.asarray(False) x_offsets = _neighbor_offsets(ntx) @@ -160,16 +152,7 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): for species in range(n_species): candidate_x = [] candidate_u = [] - candidate_charge = [] - candidate_mass = [] - candidate_weight = [] candidate_active = [] - candidate_update_x1 = [] - candidate_update_x2 = [] - candidate_update_x3 = [] - candidate_update_u1 = [] - candidate_update_u2 = [] - candidate_update_u3 = [] candidate_dest_tx = [] candidate_dest_ty = [] candidate_dest_tz = [] @@ -183,32 +166,14 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): index = (sx, sy, sz, species) candidate_x.append(bounded_x[index]) candidate_u.append(bounded_u[index]) - candidate_charge.append(tiled_particles.charge[index]) - candidate_mass.append(tiled_particles.mass[index]) - candidate_weight.append(tiled_particles.weight[index]) candidate_active.append(bounded_active[index]) - candidate_update_x1.append(tiled_particles.update_x1[index]) - candidate_update_x2.append(tiled_particles.update_x2[index]) - candidate_update_x3.append(tiled_particles.update_x3[index]) - candidate_update_u1.append(tiled_particles.update_u1[index]) - candidate_update_u2.append(tiled_particles.update_u2[index]) - candidate_update_u3.append(tiled_particles.update_u3[index]) candidate_dest_tx.append(dest_tx[index]) candidate_dest_ty.append(dest_ty[index]) candidate_dest_tz.append(dest_tz[index]) candidate_x = jnp.concatenate(candidate_x, axis=0) candidate_u = jnp.concatenate(candidate_u, axis=0) - candidate_charge = jnp.concatenate(candidate_charge, axis=0) - candidate_mass = jnp.concatenate(candidate_mass, axis=0) - candidate_weight = jnp.concatenate(candidate_weight, axis=0) candidate_active = jnp.concatenate(candidate_active, axis=0) - candidate_update_x1 = jnp.concatenate(candidate_update_x1, axis=0) - candidate_update_x2 = jnp.concatenate(candidate_update_x2, axis=0) - candidate_update_x3 = jnp.concatenate(candidate_update_x3, axis=0) - candidate_update_u1 = jnp.concatenate(candidate_update_u1, axis=0) - candidate_update_u2 = jnp.concatenate(candidate_update_u2, axis=0) - candidate_update_u3 = jnp.concatenate(candidate_update_u3, axis=0) candidate_dest_tx = jnp.concatenate(candidate_dest_tx, axis=0) candidate_dest_ty = jnp.concatenate(candidate_dest_ty, axis=0) candidate_dest_tz = jnp.concatenate(candidate_dest_tz, axis=0) @@ -225,59 +190,23 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): local_x = jnp.zeros_like(new_x[tx, ty, tz, species]) local_u = jnp.zeros_like(new_u[tx, ty, tz, species]) - local_charge = jnp.zeros_like(new_charge[tx, ty, tz, species]) - local_mass = jnp.zeros_like(new_mass[tx, ty, tz, species]) - local_weight = jnp.zeros_like(new_weight[tx, ty, tz, species]) local_active_count = jnp.zeros(n_slots, dtype=int) - local_update_x1_count = jnp.zeros(n_slots, dtype=int) - local_update_x2_count = jnp.zeros(n_slots, dtype=int) - local_update_x3_count = jnp.zeros(n_slots, dtype=int) - local_update_u1_count = jnp.zeros(n_slots, dtype=int) - local_update_u2_count = jnp.zeros(n_slots, dtype=int) - local_update_u3_count = jnp.zeros(n_slots, dtype=int) valid = fits.astype(candidate_x.dtype) local_x = local_x.at[safe_rank].add(valid[:, None] * candidate_x) local_u = local_u.at[safe_rank].add(valid[:, None] * candidate_u) - local_charge = local_charge.at[safe_rank].add(valid * candidate_charge) - local_mass = local_mass.at[safe_rank].add(valid * candidate_mass) - local_weight = local_weight.at[safe_rank].add(valid * candidate_weight) local_active_count = local_active_count.at[safe_rank].add(fits.astype(int)) - local_update_x1_count = local_update_x1_count.at[safe_rank].add((fits & candidate_update_x1).astype(int)) - local_update_x2_count = local_update_x2_count.at[safe_rank].add((fits & candidate_update_x2).astype(int)) - local_update_x3_count = local_update_x3_count.at[safe_rank].add((fits & candidate_update_x3).astype(int)) - local_update_u1_count = local_update_u1_count.at[safe_rank].add((fits & candidate_update_u1).astype(int)) - local_update_u2_count = local_update_u2_count.at[safe_rank].add((fits & candidate_update_u2).astype(int)) - local_update_u3_count = local_update_u3_count.at[safe_rank].add((fits & candidate_update_u3).astype(int)) index = (tx, ty, tz, species) new_x = new_x.at[index].set(local_x) new_u = new_u.at[index].set(local_u) - new_charge = new_charge.at[index].set(local_charge) - new_mass = new_mass.at[index].set(local_mass) - new_weight = new_weight.at[index].set(local_weight) new_active = new_active.at[index].set(local_active_count > 0) - new_update_x1 = new_update_x1.at[index].set(local_update_x1_count > 0) - new_update_x2 = new_update_x2.at[index].set(local_update_x2_count > 0) - new_update_x3 = new_update_x3.at[index].set(local_update_x3_count > 0) - new_update_u1 = new_update_u1.at[index].set(local_update_u1_count > 0) - new_update_u2 = new_update_u2.at[index].set(local_update_u2_count > 0) - new_update_u3 = new_update_u3.at[index].set(local_update_u3_count > 0) overflow = overflow | (jnp.sum(keep) > n_slots) refreshed = TiledParticles( x=new_x, u=new_u, - charge=new_charge, - mass=new_mass, - weight=new_weight, active=new_active, - update_x1=new_update_x1, - update_x2=new_update_x2, - update_x3=new_update_x3, - update_u1=new_update_u1, - update_u2=new_update_u2, - update_u3=new_update_u3, ) return refreshed, overflow diff --git a/PyPIC3D/particles/tiled_particles.py b/PyPIC3D/particles/tiled_particles.py index 03d7593a..e9a050c6 100644 --- a/PyPIC3D/particles/tiled_particles.py +++ b/PyPIC3D/particles/tiled_particles.py @@ -1,7 +1,19 @@ import jax -import jax.numpy as jnp 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. @@ -15,17 +27,5 @@ class TiledParticles(NamedTuple): # positions u: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile, 3) # velocities - charge: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - mass: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - # charge and mass of the particles - weight: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - # macroparticle weight of the particles active: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) # boolean array indicating whether the particle is active or not - update_x1: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - update_x2: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - update_x3: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - update_u1: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - update_u2: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - update_u3: jax.Array # (ntx, nty, ntz, species, max_particles_per_tile) - # boolean array indicating whether the particle's position or velocity should be updated or not diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index 611efce5..1018bffa 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -19,7 +19,7 @@ def _tile_axis(axis, tile_index, cells_per_tile, num_guard_cells, d): @partial(jit, static_argnames=("tile_shape", "g", "relativistic", "particle_pusher")) -def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, tile_shape, g, relativistic=True, particle_pusher="boris"): +def tiled_particle_push(tiled_particles, species_config, E_tiles, B_tiles, world, constants, tile_shape, g, relativistic=True, particle_pusher="boris"): """ Push tile-major particles with the selected pusher using compact tiled Yee fields. @@ -41,7 +41,7 @@ def tiled_particle_push(tiled_particles, E_tiles, B_tiles, world, constants, til ntx, nty, ntz = tiled_particles.x.shape[:3] - def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_tile, update_u1, update_u2, update_u3, + 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) @@ -49,8 +49,8 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til vx = u_tile[..., 0].reshape(-1) vy = u_tile[..., 1].reshape(-1) vz = u_tile[..., 2].reshape(-1) - q = charge_tile.reshape(-1) - m = mass_tile.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 = _tile_axis(center_grid[0], tx, tile_nx, g, world["dx"]) center_y = _tile_axis(center_grid[1], ty, tile_ny, g, world["dy"]) @@ -115,9 +115,9 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til raise ValueError(f"Unknown particle_pusher: {particle_pusher}") active = active_tile.reshape(-1) - update_u1 = update_u1.reshape(-1) - update_u2 = update_u2.reshape(-1) - update_u3 = update_u3.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)) @@ -131,20 +131,18 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til ty = jnp.arange(nty) tz = jnp.arange(ntz) - push_tiles = jax.vmap(push_tiles, in_axes=(None, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - push_tiles = jax.vmap(push_tiles, in_axes=(None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) - push_tiles = jax.vmap(push_tiles, in_axes=(0, None, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), out_axes=0) + 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, tiled_particles.x, tiled_particles.u, - tiled_particles.charge, - tiled_particles.mass, tiled_particles.active, - tiled_particles.update_u1, - tiled_particles.update_u2, - tiled_particles.update_u3, + species_config.charge, + species_config.mass, + species_config.update_u, Ex_tiles, Ey_tiles, Ez_tiles, @@ -156,14 +154,5 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, charge_tile, mass_tile, active_til return TiledParticles( x=tiled_particles.x, u=new_u, - charge=tiled_particles.charge, - mass=tiled_particles.mass, - weight=tiled_particles.weight, active=tiled_particles.active, - update_x1=tiled_particles.update_x1, - update_x2=tiled_particles.update_x2, - update_x3=tiled_particles.update_x3, - update_u1=tiled_particles.update_u1, - update_u2=tiled_particles.update_u2, - update_u3=tiled_particles.update_u3, ) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 9a75c785..c5256ea9 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -202,7 +202,7 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) -def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, phi_tiles, solver, bc, tile_shape, g): +def calculate_tiled_electrostatic_fields(world, particles, species_config, constants, rho_tiles, phi_tiles, solver, bc, tile_shape, g): """ Compute electrostatic fields from tiled rho deposition and a global Poisson solve. @@ -218,7 +218,7 @@ def calculate_tiled_electrostatic_fields(world, particles, constants, rho_tiles, bc_z = world["boundary_conditions"]["z"] g = int(g) - rho_tiles = compute_tiled_rho_from_tiled_particles(particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + rho_tiles = compute_tiled_rho_from_tiled_particles(particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) rho = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index ba3a603f..4950d1d2 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -172,7 +172,7 @@ 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. @@ -242,7 +242,11 @@ def nd_trapezoid(arr, dxs): vz = particles.u[..., 2] v2 = vx**2 + vy**2 + vz**2 active = particles.active.astype(v2.dtype) - mass = particles.mass * particles.weight + 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)) @@ -265,7 +269,7 @@ def nd_trapezoid(arr, dxs): return e_energy, b_energy, kinetic_energy -def compute_total_momentum(particles): +def compute_total_momentum(particles, species_config=None): """ Compute the scalar momentum diagnostic for either species-list or tiled particles. """ @@ -273,7 +277,12 @@ def compute_total_momentum(particles): if hasattr(particles, "active") and hasattr(particles, "u"): vmag = jnp.sqrt(particles.u[..., 0]**2 + particles.u[..., 1]**2 + particles.u[..., 2]**2) active = particles.active.astype(vmag.dtype) - return jnp.sum(active * vmag * particles.mass * particles.weight) + 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) return sum(particle_species.momentum() for particle_species in particles) diff --git a/docs/architecture.rst b/docs/architecture.rst index 9962d4ee..54424bd2 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -66,11 +66,12 @@ 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)``; scalar particle data -such as charge, mass, weight, active flags, and update masks have shape -``(ntx, nty, ntz, species, max_particles_per_tile)``. The slot capacity is set -when tiled particles are initialized. Empty slots remain inactive so the array -shape stays static during JAX updates. +``(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 diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 89de8196..fdfbb4f3 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -88,7 +88,7 @@ def _centered_tiled_particles(self, particles, world, simulation_parameters): refreshes tile ownership at the centered position. """ - tiled_particles = to_tiled_particles(particles, world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) tiled_particles = tiled_particles._replace( x=tiled_particles.x - 0.5 * tiled_particles.u * world["dt"] ) @@ -100,16 +100,17 @@ def _centered_tiled_particles(self, particles, world, simulation_parameters): ) self.assertFalse(bool(overflow)) - return centered_particles + return centered_particles, species_config def _compare_tiled_to_standard(self, particles, world, simulation_parameters): constants = {"C": 3.0e8, "alpha": 1.0} - tiled_particles = self._centered_tiled_particles(particles, world, simulation_parameters) + tiled_particles, species_config = self._centered_tiled_particles(particles, world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov(particles, self._empty_J(world), constants, world, filter="none") J_tiles = direct_J_from_tiled_particles( tiled_particles, + species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, @@ -224,12 +225,13 @@ def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") J_tiles = direct_J_from_tiled_particles( tiled_particles, + species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, @@ -419,11 +421,12 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) J_tiles = direct_J_from_tiled_particles( tiled_particles, + species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, @@ -519,12 +522,13 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") J_tiles = direct_J_from_tiled_particles( tiled_particles, + species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, @@ -566,12 +570,13 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles = self._centered_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") J_tiles = direct_J_from_tiled_particles( tiled_particles, + species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index ce356dad..d2682e9a 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -136,16 +136,7 @@ def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): return TiledParticles( x=x, u=u, - charge=jnp.pad(tiled_particles.charge, scalar_pad_width), - mass=jnp.pad(tiled_particles.mass, scalar_pad_width), - weight=jnp.pad(tiled_particles.weight, scalar_pad_width), active=jnp.pad(tiled_particles.active, scalar_pad_width), - update_x1=jnp.pad(tiled_particles.update_x1, scalar_pad_width), - update_x2=jnp.pad(tiled_particles.update_x2, scalar_pad_width), - update_x3=jnp.pad(tiled_particles.update_x3, scalar_pad_width), - update_u1=jnp.pad(tiled_particles.update_u1, scalar_pad_width), - update_u2=jnp.pad(tiled_particles.update_u2, scalar_pad_width), - update_u3=jnp.pad(tiled_particles.update_u3, scalar_pad_width), ) def _long_two_stream_species(self, world): @@ -211,7 +202,7 @@ def _long_two_stream_state(self): reference_fields = (E, B, J, rho, phi, external_fields) reference_particles = self._long_two_stream_species(world) - tiled_particles = to_tiled_particles( + tiled_particles, species_config = to_tiled_particles( self._long_two_stream_species(world), world, { @@ -350,13 +341,14 @@ def test_tiled_rho_assembles_to_global_rho(self): world = self._build_world(shape_factor=1) constants = {"alpha": 1.0} species = self._species(world) - tiled_particles = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) _, _, _, rho, _, _ = self._empty_fields(world) rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) reference_rho = compute_rho([self._species(world)], rho, world, constants) tiled_rho = compute_tiled_rho_from_tiled_particles( tiled_particles, + species_config, rho_tiles, world, constants, @@ -391,7 +383,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles(self._neutral_species(world), world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(self._neutral_species(world), world, self._simulation_parameters()) tiled_fields = ( tile_vector_field(E, world, world["tile_shape"]), tile_vector_field(B, world, world["tile_shape"]), @@ -407,6 +399,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): tiled_particles, tiled_fields = time_loop_electrostatic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index ef3318c4..0b89019c 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -146,11 +146,12 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No old_species = self._species_from_arrays(world, x_old, u) new_species = self._species_from_arrays(world, x_old + u * world["dt"], u) - tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) J_tiles = tiled_esirkepov_current( tiled_particles, + species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, @@ -281,11 +282,12 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): x_new = x_old + old_species.v1 * world["dt"] new_species = self._species(world, x_new) - tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) g = int(world["guard_cells"]) J_tiles = tiled_esirkepov_current( tiled_particles, + species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, @@ -398,22 +400,23 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): ] ) old_species = self._species_from_arrays(world, x_old, u) - tiled_particles = to_tiled_particles([old_species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) g = int(world["guard_cells"]) rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) - rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) J_tiles = tiled_esirkepov_current( tiled_particles, + species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, tile_shape=tile_shape, g=g, ) - new_particles = update_tiled_particle_positions(tiled_particles, world["dt"]) + 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_tiled_rho_from_tiled_particles(new_particles, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + rho_new = compute_tiled_rho_from_tiled_particles(new_particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) self.assertFalse(bool(overflow)) drhodt = (rho_new[:, :, :, g:-g, g:-g, g:-g] - rho_old[:, :, :, g:-g, g:-g, g:-g]) / world["dt"] @@ -593,10 +596,11 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): toml.dump(config, f) 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(toml.load(config_path)) + solver, _electrostatic, _verbose, _GPUs, _Nt, curl_func, J_func, relativistic, particle_pusher, species_config = initialize_simulation(toml.load(config_path)) particles, fields = loop( particles, + species_config, fields, world, constants, @@ -667,10 +671,11 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s } 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) + solver, _electrostatic, _verbose, _GPUs, _Nt, curl_func, _J_func, relativistic, particle_pusher, species_config = initialize_simulation(config) alias_J_func = functools.partial(tiled_esirkepov_current) particles, fields = loop( particles, + species_config, fields, world, constants, diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 18750816..4d8e8a92 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -97,18 +97,16 @@ def _particles(self, world): def test_tiled_mass_density_matches_compute_mass_density_on_interior(self): world = self._build_world(shape_factor=2) particles = self._particles(world) - tiled_particles = to_tiled_particles(particles, world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) inactive = ~tiled_particles.active x = tiled_particles.x.at[inactive, 0].set(0.10) x = x.at[inactive, 1].set(0.20) x = x.at[inactive, 2].set(-0.30) - mass = tiled_particles.mass.at[inactive].set(1000.0) - weight = tiled_particles.weight.at[inactive].set(1000.0) - tiled_particles = tiled_particles._replace(x=x, mass=mass, weight=weight) + tiled_particles = tiled_particles._replace(x=x) mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) - mass_tiled = compute_mass_density(tiled_particles, self._empty_scalar(world), world) + mass_tiled = compute_mass_density(tiled_particles, self._empty_scalar(world), world, species_config=species_config) self.assertTrue( jnp.allclose( @@ -127,12 +125,13 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): self._simulation_parameters()["particle_tile_nz"], ) particles = self._particles(world) - tiled_particles = to_tiled_particles(particles, world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"]) mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) mass_tiles = compute_tiled_mass_density_from_tiled_particles( tiled_particles, + species_config, rho_tiles, world, tile_shape=world["tile_shape"], diff --git a/tests/tiled_particle_diagnostics_test.py b/tests/tiled_particle_diagnostics_test.py index 8368ad26..9522423c 100644 --- a/tests/tiled_particle_diagnostics_test.py +++ b/tests/tiled_particle_diagnostics_test.py @@ -83,9 +83,9 @@ def _species(self): def test_flatten_tiled_particles_matches_active_original_species(self): species_list = self._species() - tiled_particles = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles) + flattened_species = flatten_tiled_particles_by_species(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)): @@ -105,9 +105,9 @@ def test_flatten_tiled_particles_matches_active_original_species(self): def test_inactive_tiled_slots_are_not_flattened(self): species_list = self._species() - tiled_particles = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles) + flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config) self.assertEqual(flattened_species[0].get_number_of_particles(), 3) self.assertEqual(flattened_species[1].get_number_of_particles(), 2) @@ -116,12 +116,12 @@ def test_inactive_tiled_slots_are_not_flattened(self): def test_absorbed_particles_do_not_appear_in_flattened_output(self): species_list = self._species() - tiled_particles = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + tiled_particles, species_config = to_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 = flatten_tiled_particles_by_species(tiled_particles) + flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config) self.assertEqual(flattened_species[0].get_number_of_particles(), 2) self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].get_forward_position()[0], 1.5))) @@ -134,9 +134,9 @@ def test_flattened_diagnostic_positions_match_original_half_step_positions(self) "y": 0, "z": 0, } - tiled_particles = to_tiled_particles(species_list, world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(species_list, world, self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, world=world) + flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config, world=world) for original, flattened in zip(species_list, flattened_species): active = original.get_active_mask() @@ -147,10 +147,10 @@ def test_flattened_diagnostic_positions_match_original_half_step_positions(self) def test_species_names_are_preserved_when_metadata_is_available(self): species_list = self._species() - tiled_particles = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) species_names = tuple(species.get_name() for species in species_list) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_names=species_names) + flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config, species_names=species_names) self.assertEqual(flattened_species[0].get_name(), "ions") self.assertEqual(flattened_species[1].get_name(), "electrons") diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py index 9ff776a1..e294c467 100644 --- a/tests/tiled_particle_initialization_test.py +++ b/tests/tiled_particle_initialization_test.py @@ -8,7 +8,7 @@ from PyPIC3D.particles.tiled_particle_initialization import load_tiled_particles_from_toml from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles from PyPIC3D.particles.species_class import particle_species @@ -65,27 +65,108 @@ def test_to_tiled_particles_packs_existing_particle_species_list(self): dt=world["dt"], ) - particles = to_tiled_particles([species], world, simulation_parameters) + particles, species_config = to_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.assertTrue(jnp.allclose(particles.charge[particles.x[..., 0] != 0.0], 2.0)) - self.assertTrue(jnp.allclose(particles.mass[particles.x[..., 0] != 0.0], 3.0)) - self.assertTrue(jnp.allclose(particles.weight[particles.x[..., 0] != 0.0], 4.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(particles.update_x1[1, 1, 0, 0, 1])) - self.assertFalse(bool(particles.update_x2[1, 1, 0, 0, 1])) - self.assertTrue(bool(particles.update_x3[1, 1, 0, 0, 1])) - self.assertFalse(bool(particles.update_u1[1, 1, 0, 0, 1])) - self.assertTrue(bool(particles.update_u2[1, 1, 0, 0, 1])) - self.assertTrue(bool(particles.update_u3[1, 1, 0, 0, 1])) + 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 = [ + particle_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"], + ), + particle_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 = to_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_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): world = { @@ -128,10 +209,11 @@ def test_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): dt=world["dt"], ) - particles = to_tiled_particles([species], world, simulation_parameters) + particles, species_config = to_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_tiled_particles_from_toml_uses_tile_axes_before_species(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -179,12 +261,12 @@ def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): } } - particles = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + particles, species_config = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) self.assertIsInstance(particles, TiledParticles) + self.assertIsInstance(species_config, SpeciesConfig) self.assertEqual(particles.x.shape, (2, 2, 1, 1, 2, 3)) self.assertEqual(particles.u.shape, (2, 2, 1, 1, 2, 3)) - self.assertEqual(particles.charge.shape, (2, 2, 1, 1, 2)) self.assertTrue(particles.active[0, 0, 0, 0, 0]) self.assertTrue(particles.active[1, 1, 0, 0, 0]) @@ -196,9 +278,9 @@ def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): 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(particles.charge[particles.active], -1.0)) - self.assertTrue(jnp.allclose(particles.mass[particles.active], 2.0)) - self.assertTrue(jnp.allclose(particles.weight[particles.active], 4.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_tiled_particles_from_toml_maps_update_flags_to_active_particles(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -250,14 +332,14 @@ def test_load_tiled_particles_from_toml_maps_update_flags_to_active_particles(se } } - particles = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + particles, species_config = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) - self.assertTrue(bool(particles.update_x1[0, 0, 0, 0, 0])) - self.assertFalse(bool(particles.update_x2[0, 0, 0, 0, 0])) - self.assertTrue(bool(particles.update_x3[0, 0, 0, 0, 0])) - self.assertFalse(bool(particles.update_u1[0, 0, 0, 0, 0])) - self.assertTrue(bool(particles.update_u2[0, 0, 0, 0, 0])) - self.assertFalse(bool(particles.update_u3[0, 0, 0, 0, 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.assertFalse(bool(species_config.update_u[0, 2])) if __name__ == "__main__": diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 04240521..41779e6c 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -71,11 +71,8 @@ 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] - charge = tiled_particles.charge.reshape(-1)[active] - mass = tiled_particles.mass.reshape(-1)[active] - weight = tiled_particles.weight.reshape(-1)[active] order = jnp.argsort(x[:, 0]) - return x[order], u[order], charge[order], mass[order], weight[order] + return x[order], u[order] def test_update_tiled_particle_positions_respects_active_and_update_flags(self): world = self._build_world() @@ -85,23 +82,23 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): v1=[0.25, 0.5, 0.75], active_mask=jnp.array([True, False, True]), ) - tiled_particles = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) - moved = update_tiled_particle_positions(tiled_particles, world["dt"]) + moved = update_tiled_particle_positions(tiled_particles, species_config, world["dt"]) - x, _, _, _, _ = self._active_rows(moved) + 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 = to_tiled_particles([fixed_species], world, self._simulation_parameters()) - fixed_moved = update_tiled_particle_positions(fixed, world["dt"]) + fixed, fixed_species_config = to_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_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 = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_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) @@ -111,17 +108,14 @@ def test_refresh_moves_particles_to_neighbor_tiles_with_static_shape(self): 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, charge, mass, weight = self._active_rows(refreshed) + 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]))) - self.assertTrue(jnp.allclose(charge, 2.0)) - self.assertTrue(jnp.allclose(mass, 3.0)) - self.assertTrue(jnp.allclose(weight, 4.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 = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_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)) @@ -134,7 +128,7 @@ 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 = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_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) @@ -144,7 +138,7 @@ def test_refresh_reflects_particle_from_global_boundary_condition(self): refreshed, overflow = refresh_tiled_particle_tiles(moved, world, tile_shape=(2, 1, 1)) self.assertFalse(bool(overflow)) - x, u, _, _, _ = self._active_rows(refreshed) + 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]))) @@ -152,21 +146,21 @@ 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 = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_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) + 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 = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = to_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)) diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 5fbaad3a..fc132535 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -112,9 +112,10 @@ def test_tiled_particle_push_matches_flat_boris(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) pushed = tiled_particle_push( tiled_particles, + species_config, tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, @@ -179,9 +180,10 @@ def make_species(): particle_pusher="higuera_cary", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) pushed = tiled_particle_push( tiled_particles, + species_config, tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, @@ -216,10 +218,11 @@ def test_tiled_particle_push_respects_active_and_update_flags(self): update_vy=True, update_vz=False, ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) pushed = tiled_particle_push( tiled_particles, + species_config, tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, @@ -281,9 +284,10 @@ def make_species(): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) pushed = tiled_particle_push( tiled_particles, + species_config, tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index ab9c70ed..5aafaf6f 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -93,17 +93,15 @@ def _particles(self, world): return [electrons, ions] def _tiled_with_noisy_inactive_slots(self, particles, world): - tiled_particles = to_tiled_particles(particles, world, self._simulation_parameters()) + tiled_particles, species_config = to_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) - charge = tiled_particles.charge.at[inactive].set(1000.0) - weight = tiled_particles.weight.at[inactive].set(1000.0) - return tiled_particles._replace(x=x, u=u, charge=charge, weight=weight) + return tiled_particles._replace(x=x, u=u), species_config def _zero_species_velocities(self, particles): for species in particles: @@ -119,11 +117,12 @@ def _compare_tiled_to_standard(self, shape_factor, alpha): world = self._build_world(shape_factor) constants = {"alpha": alpha} particles = self._particles(world) - tiled_particles = self._tiled_with_noisy_inactive_slots(particles, world) + tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) rho_reference = compute_rho(particles, self._empty_scalar(world), world, constants) rho_tiled = compute_rho_from_tiled_particles( tiled_particles, + species_config, self._empty_scalar(world), world, constants, @@ -169,17 +168,19 @@ def test_tiled_global_rho_uses_current_positions_not_half_step_back_positions(se world = self._build_world(shape_factor=2) constants = {"alpha": 1.0} particles = self._particles(world) - tiled_particles = self._tiled_with_noisy_inactive_slots(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_with_velocity = compute_rho_from_tiled_particles( tiled_particles, + species_config, self._empty_scalar(world), world, constants, ) rho_with_zero_velocity = compute_rho_from_tiled_particles( zero_velocity_tiled_particles, + species_config, self._empty_scalar(world), world, constants, @@ -203,12 +204,13 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self ) constants = {"alpha": 1.0} particles = self._particles(world) - tiled_particles = self._tiled_with_noisy_inactive_slots(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_tiled_rho_from_tiled_particles( tiled_particles, + species_config, rho_tiles, world, constants, @@ -217,6 +219,7 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self ) rho_tiles_with_zero_velocity = compute_tiled_rho_from_tiled_particles( zero_velocity_tiled_particles, + species_config, rho_tiles, world, constants, diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index ccfa127a..7435b862 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -169,7 +169,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } - tiled_particles = to_tiled_particles(particles, world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) g = int(world["guard_cells"]) if current_calculation == CURRENT_ESIRKEPOV: @@ -214,30 +214,12 @@ def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim scalar_pad_width[-1] = (0, min_slots - current_slots) - charge = jnp.pad(tiled_particles.charge, scalar_pad_width) - mass = jnp.pad(tiled_particles.mass, scalar_pad_width) - weight = jnp.pad(tiled_particles.weight, scalar_pad_width) active = jnp.pad(tiled_particles.active, scalar_pad_width) - update_x1 = jnp.pad(tiled_particles.update_x1, scalar_pad_width) - update_x2 = jnp.pad(tiled_particles.update_x2, scalar_pad_width) - update_x3 = jnp.pad(tiled_particles.update_x3, scalar_pad_width) - update_u1 = jnp.pad(tiled_particles.update_u1, scalar_pad_width) - update_u2 = jnp.pad(tiled_particles.update_u2, scalar_pad_width) - update_u3 = jnp.pad(tiled_particles.update_u3, scalar_pad_width) return TiledParticles( x=x, u=u, - charge=charge, - mass=mass, - weight=weight, active=active, - update_x1=update_x1, - update_x2=update_x2, - update_x3=update_x3, - update_u1=update_u1, - update_u2=update_u2, - update_u3=update_u3, ) def _standard_species_rows(self, particles, species_index): @@ -361,6 +343,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): ) _, tiled_fields = time_loop_electrodynamic_tiled( [], + None, tiled_fields, world, constants, @@ -712,7 +695,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) tiled_fields = ( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -727,6 +710,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): ) tiled_particles, tiled_fields = time_loop_electrodynamic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, @@ -759,7 +743,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) reference_energy = compute_energy(reference_particles, reference_E, reference_B, world, constants) - tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants) + tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants, species_config=species_config) for reference, tiled in zip(reference_energy, tiled_energy): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) @@ -791,7 +775,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) tiled_fields = ( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -806,6 +790,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): ) _, tiled_fields = time_loop_electrodynamic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, @@ -854,7 +839,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) tiled_fields = ( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -869,6 +854,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): ) _, tiled_fields = time_loop_electrodynamic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, @@ -914,7 +900,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) tiled_fields = ( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -929,6 +915,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): ) _, tiled_fields = time_loop_electrodynamic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, @@ -1016,7 +1003,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand particle_pusher="boris", ) - tiled_particles = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) tiled_fields = ( tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), @@ -1031,6 +1018,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ) _, tiled_fields = time_loop_electrodynamic_tiled( tiled_particles, + species_config, tiled_fields, world, constants, diff --git a/tests/utils_test.py b/tests/utils_test.py index a5f57f1f..4c18b303 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -212,20 +212,10 @@ def test_compute_energy_ignores_inactive_particles(self): 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,)) - ones = jnp.ones(active.shape) particles = TiledParticles( x=zeros3, u=zeros3, - charge=ones, - mass=ones, - weight=ones, active=active, - update_x1=active, - update_x2=active, - update_x3=active, - update_u1=active, - update_u2=active, - update_u3=active, ) with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 4f9e7921..a347c061 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -442,6 +442,7 @@ def test_tiled_evolve_updates_fields_without_pushing_particles(self): particles_after, fields_after = time_loop_electrodynamic_tiled( particles, + None, fields, world, constants, From 4ec6aaf53a17af432f5ecf9098c5e2d78311b36e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 30 Jun 2026 14:03:21 -0500 Subject: [PATCH 22/83] using a tiled grid structure instead of computing tiled axes in real time --- PyPIC3D/deposition/direct_deposition_tiled.py | 28 ++++---- PyPIC3D/deposition/esirkepov_tiled.py | 23 ++++--- PyPIC3D/deposition/rho_tiled.py | 64 +++++++++--------- PyPIC3D/initialization.py | 13 ++++ PyPIC3D/pusher/tiled_pusher.py | 33 ++++++---- PyPIC3D/solvers/yee_tiled.py | 65 +++++++++++++++++++ tests/tiled_electrostatic_test.py | 9 +++ tests/yee_tiled_test.py | 25 +++++++ 8 files changed, 189 insertions(+), 71 deletions(-) diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index fff93360..3b051d69 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -11,6 +11,7 @@ from PyPIC3D.solvers.yee_tiled import ( digital_filter_tiled_vector, fold_tiled_vector_ghost_cells, + tiled_grid_axes_from_world, update_tiled_vector_ghost_cells, ) @@ -94,20 +95,6 @@ def digital_filter_tiled_current_density(J_tiles, alpha, world, num_guard_cells= return J_tiles -def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_cells): - """ - Build the ghost-celled coordinate line for one compact tile. - - The first local grid point is shifted by the configured tile guard depth. - With one guard cell this is the old compact-tile convention; with two guard - cells the extra point gives quadratic shape factors their full local - stencil at tile faces. - """ - - offsets = jnp.arange(local_n, dtype=global_axis_grid.dtype) - return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d - - 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) @@ -140,6 +127,13 @@ def direct_J_from_tiled_particles( if grid is None: grid = world["grids"]["center"] + tiled_grid = tiled_grid_axes_from_world( + world, + grid, + "tiled_center_grid", + tile_shape, + g, + ) dx = world["dx"] dy = world["dy"] @@ -178,9 +172,9 @@ def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) dq = q / (dx * dy * dz) - x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) - y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) - z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz, g) + x_grid = tiled_grid[0][tx, ty, tz] + y_grid = tiled_grid[1][tx, ty, tz] + z_grid = tiled_grid[2][tx, ty, tz] x, x0, deltax_node, xpts = prepare_particle_axis_stencil( x, diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 8ca27a7b..bd6d8494 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -16,19 +16,11 @@ ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( + tiled_grid_axes_from_world, update_tiled_vector_ghost_cells, ) -def _tile_axis_grid(global_axis_grid, tile_index, tile_n, local_n, d, num_guard_cells): - """ - Build one local tile axis with ``num_guard_cells`` guard points per side. - """ - - offsets = jnp.arange(local_n, dtype=global_axis_grid.dtype) - return global_axis_grid[0] + (offsets + tile_index * tile_n - (num_guard_cells - 1)) * d - - def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=1, tile_shape=None): """ Fold Esirkepov current deposits using the configured guard depth. @@ -125,6 +117,13 @@ def tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, if grid is None: grid = world["grids"]["center"] + tiled_grid = tiled_grid_axes_from_world( + world, + grid, + "tiled_center_grid", + tile_shape, + g, + ) dx = world["dx"] dy = world["dy"] @@ -168,9 +167,9 @@ def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): y = old_y + jnp.where(update_x2, vy * dt, 0.0) z = old_z + jnp.where(update_x3, vz * dt, 0.0) - x_grid = _tile_axis_grid(grid[0], tx, tile_nx, local_Nx, dx, g) - y_grid = _tile_axis_grid(grid[1], ty, tile_ny, local_Ny, dy, g) - z_grid = _tile_axis_grid(grid[2], tz, tile_nz, local_Nz, dz, g) + x_grid = tiled_grid[0][tx, ty, tz] + y_grid = tiled_grid[1][tx, ty, tz] + z_grid = tiled_grid[2][tx, ty, tz] x0 = compute_particle_anchor(x, x_grid, shape_factor) y0 = compute_particle_anchor(y, y_grid, shape_factor) diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index 2d67e7fb..9a8962cf 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -13,7 +13,7 @@ from PyPIC3D.solvers.yee_tiled import ( digital_filter_tiled_scalar, fold_tiled_ghost_cells, - tile_scalar_field, + tiled_grid_axes_from_world, update_tiled_ghost_cells, ) from PyPIC3D.utils import digital_filter @@ -160,12 +160,6 @@ def _diagnostic_mass_position(tiled_particles, world): return x, y, z -def _tile_axis(axis, tile_index, cells_per_tile, num_guard_cells, d): - local_n = cells_per_tile + 2 * num_guard_cells - offsets = jnp.arange(local_n, dtype=axis.dtype) - return axis[0] + (offsets + tile_index * cells_per_tile - (num_guard_cells - 1)) * d - - 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) @@ -182,6 +176,8 @@ def _deposit_tiled_scalar_moment_to_tiles( scalar_weight, tile_shape, g, + tiled_grid, + position=None, ): dx = world["dx"] dy = world["dy"] @@ -203,14 +199,19 @@ def _deposit_tiled_scalar_moment_to_tiles( def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile): local_grid = ( - _tile_axis(grid[0], tx, tile_nx, g, dx), - _tile_axis(grid[1], ty, tile_ny, g, dy), - _tile_axis(grid[2], tz, tile_nz, g, dz), + tiled_grid[0][tx, ty, tz], + tiled_grid[1][tx, ty, tz], + tiled_grid[2][tx, ty, tz], ) - x = x_tile[..., 0].reshape(-1) - y = x_tile[..., 1].reshape(-1) - z = x_tile[..., 2].reshape(-1) + if position is None: + x = x_tile[..., 0].reshape(-1) + y = x_tile[..., 1].reshape(-1) + z = x_tile[..., 2].reshape(-1) + else: + x = position[0][tx, ty, tz].reshape(-1) + y = position[1][tx, ty, tz].reshape(-1) + z = position[2][tx, ty, tz].reshape(-1) active = active_tile.reshape(-1).astype(x.dtype) scalar_density = scalar_weight_tile.reshape(-1) / (dx * dy * dz) @@ -337,6 +338,13 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_ """Compute charge density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] + tiled_grid = tiled_grid_axes_from_world( + world, + grid, + "tiled_vertex_grid", + tile_shape, + g, + ) scalar_weight = _species_scalar_to_slots( tiled_particles, @@ -350,6 +358,7 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_ scalar_weight, tile_shape, g, + tiled_grid, ) alpha = constants["alpha"] @@ -385,32 +394,27 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, species_con """Compute mass density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] + tiled_grid = tiled_grid_axes_from_world( + world, + grid, + "tiled_vertex_grid", + tile_shape, + g, + ) scalar_weight = _species_scalar_to_slots( tiled_particles, species_config.mass * species_config.weight, ) position = _diagnostic_mass_position(tiled_particles, world) - g = int(g) - tile_shape = tuple(int(width) for width in tile_shape) - Nx = int(rho_tiles.shape[0]) * tile_shape[0] - Ny = int(rho_tiles.shape[1]) * tile_shape[1] - Nz = int(rho_tiles.shape[2]) * tile_shape[2] - rho = jnp.zeros( - ( - Nx + 2, - Ny + 2, - Nz + 2, - ), - dtype=rho_tiles.dtype, - ) - rho = _deposit_tiled_scalar_moment( + return _deposit_tiled_scalar_moment_to_tiles( tiled_particles, - rho, + rho_tiles, world, grid, scalar_weight, + tile_shape, + g, + tiled_grid, position=position, ) - - return tile_scalar_field(rho, world, tile_shape, num_guard_cells=g) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 13131438..de5701cb 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -45,6 +45,7 @@ from PyPIC3D.solvers.yee_tiled import ( empty_tiled_vector_field, + tile_grid_axes, tile_scalar_field, tile_vector_field, ) @@ -551,6 +552,18 @@ def initialize_simulation(toml_file): tile_shape = _tile_shape_from_parameters(simulation_parameters) world["tile_shape"] = tile_shape guard_cells = int(world["guard_cells"]) + world["grids"]["tiled_center_grid"] = tile_grid_axes( + world["grids"]["center"], + world, + tile_shape, + num_guard_cells=guard_cells, + ) + world["grids"]["tiled_vertex_grid"] = tile_grid_axes( + world["grids"]["vertex"], + world, + tile_shape, + num_guard_cells=guard_cells, + ) particles, species_config = to_tiled_particles(particles, world, simulation_parameters) E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index 1018bffa..dddbb0dc 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -10,12 +10,7 @@ relativistic_boris_single_particle, ) from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle - - -def _tile_axis(axis, tile_index, cells_per_tile, num_guard_cells, d): - local_n = cells_per_tile + 2 * num_guard_cells - offsets = jnp.arange(local_n, dtype=axis.dtype) - return axis[0] + (offsets + tile_index * cells_per_tile - (num_guard_cells - 1)) * d +from PyPIC3D.solvers.yee_tiled import tiled_grid_axes_from_world @partial(jit, static_argnames=("tile_shape", "g", "relativistic", "particle_pusher")) @@ -34,6 +29,20 @@ def tiled_particle_push(tiled_particles, species_config, E_tiles, B_tiles, world center_grid = world["grids"]["center"] vertex_grid = world["grids"]["vertex"] + tiled_center_grid = tiled_grid_axes_from_world( + world, + center_grid, + "tiled_center_grid", + tile_shape, + g, + ) + tiled_vertex_grid = tiled_grid_axes_from_world( + world, + vertex_grid, + "tiled_vertex_grid", + tile_shape, + g, + ) Ex_tiles, Ey_tiles, Ez_tiles = E_tiles Bx_tiles, By_tiles, Bz_tiles = B_tiles @@ -52,12 +61,12 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, charge_species, mass_ 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 = _tile_axis(center_grid[0], tx, tile_nx, g, world["dx"]) - center_y = _tile_axis(center_grid[1], ty, tile_ny, g, world["dy"]) - center_z = _tile_axis(center_grid[2], tz, tile_nz, g, world["dz"]) - vertex_x = _tile_axis(vertex_grid[0], tx, tile_nx, g, world["dx"]) - vertex_y = _tile_axis(vertex_grid[1], ty, tile_ny, g, world["dy"]) - vertex_z = _tile_axis(vertex_grid[2], tz, tile_nz, g, world["dz"]) + 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 diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index f857e5c1..2924a68d 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -73,6 +73,71 @@ def empty_tiled_vector_field(world, tile_shape, num_guard_cells=1, dtype=None): return tuple(empty_tiled_scalar_field(world, tile_shape, num_guard_cells, dtype) for _ in range(3)) +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=1): + """ + Build tile-local coordinate lines for a center or vertex grid. + + Each returned axis has leading tile indices followed by the local + ghost-celled coordinate line for that axis. These arrays carry the same + coordinate convention previously rebuilt inside tiled interpolation and + deposition kernels. + """ + + 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 tiled_grid_axes_from_world(world, grid, tiled_grid_key, tile_shape, num_guard_cells=1): + """ + Read precomputed tile-local grid axes, with a construction fallback. + + Initialization stores the tiled grids in ``world['grids']`` for production + tiled runs. The fallback keeps focused tests and older call sites that + construct small worlds by hand on the same numerical convention. + """ + + grids = world["grids"] + if tiled_grid_key in grids: + return grids[tiled_grid_key] + + return tile_grid_axes(grid, world, tile_shape, num_guard_cells) + + def tile_scalar_field(field, world, tile_shape, num_guard_cells=1): """ Split a ghost-celled field into compact tiles using the shared tile shape. diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index d2682e9a..962480ff 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -546,6 +546,15 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): self.assertEqual(tuple(simulation_parameters["tile_shape"]), (2, 1, 1)) for vertex_axis, center_axis in zip(world["grids"]["vertex"], world["grids"]["center"]): self.assertTrue(jnp.allclose(vertex_axis, center_axis)) + self.assertIn("tiled_center_grid", world["grids"]) + self.assertIn("tiled_vertex_grid", world["grids"]) + self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (4, 1, 1, 4)) + self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (4, 1, 1, 3)) + 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], (4, 1, 1)) self.assertEqual(fields[3].shape[:3], (4, 1, 1)) self.assertEqual(fields[4].shape[:3], (4, 1, 1)) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index a347c061..1c9a2547 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -8,6 +8,7 @@ from PyPIC3D.solvers.first_order_yee import update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, + tile_grid_axes, tile_vector_field, update_tiled_ghost_cells, update_tiled_ghost_cells_periodic, @@ -151,6 +152,30 @@ def test_tile_vector_field_assembles_to_original_ghost_celled_field(self): 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 + tiled_center_grid = tile_grid_axes(world["grids"]["center"], world, tile_shape, num_guard_cells=g) + tiled_vertex_grid = tile_grid_axes(world["grids"]["vertex"], world, tile_shape, num_guard_cells=g) + + 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) From c91989dbe3e58c77816b6a322c24b26e32df1b8e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 30 Jun 2026 15:31:52 -0500 Subject: [PATCH 23/83] refactoring and defaulting to 2 guard cells unless specified --- PyPIC3D/boundary_conditions/PML.py | 40 +- PyPIC3D/deposition/direct_deposition_tiled.py | 6 +- PyPIC3D/deposition/esirkepov_tiled.py | 6 +- PyPIC3D/deposition/rho_tiled.py | 4 +- PyPIC3D/diagnostics/tiled_benchmarking.py | 358 ++++++++++++++++++ PyPIC3D/initialization.py | 4 +- PyPIC3D/solvers/yee_tiled.py | 128 +++++-- PyPIC3D/utils.py | 7 +- tests/direct_deposition_tiled_test.py | 10 +- tests/initialization_test.py | 1 + tests/openpmd_test.py | 2 +- tests/pml_test.py | 22 +- tests/tiled_electrostatic_test.py | 24 +- tests/tiled_esirkepov_test.py | 3 +- tests/tiled_refactor_contract_test.py | 103 +++++ tests/tiled_rho_test.py | 9 +- tests/tiled_yee_integration_test.py | 2 +- tests/yee_tiled_test.py | 16 +- 18 files changed, 652 insertions(+), 93 deletions(-) create mode 100644 PyPIC3D/diagnostics/tiled_benchmarking.py create mode 100644 tests/tiled_refactor_contract_test.py diff --git a/PyPIC3D/boundary_conditions/PML.py b/PyPIC3D/boundary_conditions/PML.py index 9a4e724f..bfa336ef 100644 --- a/PyPIC3D/boundary_conditions/PML.py +++ b/PyPIC3D/boundary_conditions/PML.py @@ -135,7 +135,7 @@ def _tile_axis_count(n_cells, cells_per_tile): return int(n_cells) // int(cells_per_tile) -def _tile_scalar_profile(profile, tile_shape): +def _tile_scalar_profile(profile, tile_shape, num_guard_cells=2): """ Split one ghost-celled PML profile into compact tile-local profile arrays. """ @@ -147,6 +147,22 @@ def _tile_scalar_profile(profile, tile_shape): 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 = 1 + tx * tile_nx + iy = 1 + ty * tile_ny + iz = 1 + tz * tile_nz + interior = profile[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] + profile_tiles = profile_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + return profile_tiles def tile_at(tx, ty, tz): start = (tx * tile_nx, ty * tile_ny, tz * tile_nz) @@ -172,13 +188,13 @@ 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 keep - the same one-cell tile halos as the field tiles, so the derivative stretch - uses ``[..., 1:-1, 1:-1, 1:-1]`` next to the Yee derivative arrays. + 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"] - return tuple(_tile_scalar_profile(profile, tile_shape) for profile in profiles) + 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): @@ -342,9 +358,10 @@ def apply_tiled_pml_to_e_curl(derivatives, world, pml_state): ) = e_memory sigma_x, sigma_y, sigma_z = tiled_profiles - sigma_x = sigma_x[:, :, :, 1:-1, 1:-1, 1:-1] - sigma_y = sigma_y[:, :, :, 1:-1, 1:-1, 1:-1] - sigma_z = sigma_z[:, :, :, 1:-1, 1:-1, 1:-1] + 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) @@ -387,9 +404,10 @@ def apply_tiled_pml_to_b_curl(derivatives, world, pml_state): ) = b_memory sigma_x, sigma_y, sigma_z = tiled_profiles - sigma_x = sigma_x[:, :, :, 1:-1, 1:-1, 1:-1] - sigma_y = sigma_y[:, :, :, 1:-1, 1:-1, 1:-1] - sigma_z = sigma_z[:, :, :, 1:-1, 1:-1, 1:-1] + 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) diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index 3b051d69..60ee54ee 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -16,7 +16,7 @@ ) -def bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=1, tile_shape=None): +def bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=2, tile_shape=None): """ Apply the global bilinear current filter to compact current tiles. @@ -77,7 +77,7 @@ def filter_component(component): return J_tiles -def digital_filter_tiled_current_density(J_tiles, alpha, world, num_guard_cells=1, tile_shape=None): +def digital_filter_tiled_current_density(J_tiles, alpha, world, num_guard_cells=2, tile_shape=None): """ Apply the global digital current filter to compact current tiles. @@ -113,7 +113,7 @@ def direct_J_from_tiled_particles( grid=None, filter="bilinear", tile_shape=None, - g=1, + g=2, ): """ Compute direct current deposition using tile-local stencils only. diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index bd6d8494..4476f5c7 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -21,7 +21,7 @@ ) -def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=1, tile_shape=None): +def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2, tile_shape=None): """ Fold Esirkepov current deposits using the configured guard depth. @@ -97,7 +97,7 @@ def absorbing_boundary(tiles): return field_tiles -def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): +def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): return tuple( fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells, tile_shape) for component_axis, component in enumerate(field_tiles) @@ -105,7 +105,7 @@ def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells= @partial(jit, static_argnames=("tile_shape", "g")) -def tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=1): +def tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=2): """ Deposit Esirkepov current into tile-local current buffers. diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index 9a8962cf..c5c3e1ac 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -334,7 +334,7 @@ def compute_rho_from_tiled_particles(tiled_particles, species_config, rho, world @partial(jit, static_argnames=("tile_shape", "g")) -def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, grid=None, tile_shape=None, g=1): +def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, grid=None, tile_shape=None, g=2): """Compute charge density into tile-major vertex-grid scalar arrays.""" if grid is None: grid = world["grids"]["vertex"] @@ -390,7 +390,7 @@ def compute_mass_density_from_tiled_particles(tiled_particles, species_config, r @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=1): +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"] diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py new file mode 100644 index 00000000..95f0aa23 --- /dev/null +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -0,0 +1,358 @@ +from dataclasses import dataclass +from functools import partial + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.deposition.current_methods import CURRENT_J_FROM_RHOV +from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particle_refresh import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) +from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.pusher.tiled_pusher import tiled_particle_push +from PyPIC3D.solvers.yee_tiled import ( + empty_tiled_scalar_field, + empty_tiled_vector_field, + tile_grid_axes, + tile_vector_field, + update_tiled_B, + update_tiled_E, +) +from PyPIC3D.utils import add_external_fields, build_yee_grid, compute_energy + + +@dataclass +class StageSpec: + name: str + function: object + args: tuple + kwargs: dict + static_argnames: tuple = () + + +@dataclass +class SyntheticTiledYeeState: + particles: TiledParticles + species_config: object + fields: tuple + world: dict + constants: dict + tile_shape: tuple + g: int + + +def unused_curl(Ex, Ey, Ez): + return None + + +def _pad_tiled_particle_capacity(tiled_particles, slots_per_tile): + current_slots = tiled_particles.active.shape[-1] + if current_slots >= slots_per_tile: + return tiled_particles + + pad_slots = int(slots_per_tile) - int(current_slots) + vector_pad = [(0, 0)] * tiled_particles.x.ndim + vector_pad[-2] = (0, pad_slots) + scalar_pad = [(0, 0)] * tiled_particles.active.ndim + scalar_pad[-1] = (0, pad_slots) + + return TiledParticles( + x=jnp.pad(tiled_particles.x, vector_pad), + u=jnp.pad(tiled_particles.u, vector_pad), + active=jnp.pad(tiled_particles.active, scalar_pad), + ) + + +def build_synthetic_tiled_yee_state( + nx=16, + ny=16, + nz=16, + particles_per_species=2048, + species_count=2, + shape=1, + seed=0, + tile_shape=(8, 8, 8), + slots_per_tile=None, + g=2, + dt=None, +): + """ + Build a deterministic tiled electrodynamic Yee benchmark state. + """ + + x_wind = 1.0 + y_wind = 1.0 + z_wind = 1.0 + dx = x_wind / nx + dy = y_wind / ny + dz = z_wind / nz + if dt is None: + dt = 0.05 * min(dx, dy, dz) + + world = { + "dt": dt, + "Nt": 1, + "dx": dx, + "dy": dy, + "dz": dz, + "Nx": nx, + "Ny": ny, + "Nz": nz, + "x_wind": x_wind, + "y_wind": y_wind, + "z_wind": z_wind, + "shape_factor": shape, + "current_calculation": CURRENT_J_FROM_RHOV, + "tile_shape": tile_shape, + "guard_cells": g, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + world["grids"]["tiled_center_grid"] = tile_grid_axes(center_grid, world, tile_shape, num_guard_cells=g) + world["grids"]["tiled_vertex_grid"] = tile_grid_axes(vertex_grid, world, tile_shape, num_guard_cells=g) + + constants = { + "eps": 1.0, + "mu": 1.0, + "C": 1.0, + "kb": 1.0, + "alpha": 1.0, + } + + key = jax.random.PRNGKey(seed) + particles = [] + for species_index in range(species_count): + key, x_key, y_key, z_key, vx_key, vy_key, vz_key = jax.random.split(key, 7) + x1 = jax.random.uniform( + x_key, (particles_per_species,), minval=-0.45 * x_wind, maxval=0.45 * x_wind + ) + x2 = jax.random.uniform( + y_key, (particles_per_species,), minval=-0.45 * y_wind, maxval=0.45 * y_wind + ) + x3 = jax.random.uniform( + z_key, (particles_per_species,), minval=-0.45 * z_wind, maxval=0.45 * z_wind + ) + v1 = 0.01 * jax.random.normal(vx_key, (particles_per_species,)) + v2 = 0.01 * jax.random.normal(vy_key, (particles_per_species,)) + v3 = 0.01 * jax.random.normal(vz_key, (particles_per_species,)) + + charge_sign = 1.0 if species_index % 2 == 0 else -1.0 + particles.append( + particle_species( + name=f"benchmark_species_{species_index}", + N_particles=particles_per_species, + charge=charge_sign, + mass=1.0, + T=0.0, + v1=v1, + v2=v2, + v3=v3, + x1=x1, + x2=x2, + x3=x3, + xwind=x_wind, + ywind=y_wind, + zwind=z_wind, + dx=dx, + dy=dy, + dz=dz, + dt=dt, + ) + ) + + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + "particle_tile_capacity_factor": 1.0, + } + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + if slots_per_tile is not None: + tiled_particles = _pad_tiled_particle_capacity(tiled_particles, slots_per_tile) + + global_shape = (nx + 2, ny + 2, nz + 2) + E = tuple(jnp.zeros(global_shape) for _ in range(3)) + B = tuple(jnp.zeros(global_shape) for _ in range(3)) + E_tiles = tile_vector_field(E, world, tile_shape, num_guard_cells=g) + B_tiles = tile_vector_field(B, world, tile_shape, num_guard_cells=g) + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) + rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) + phi_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) + external_fields = ( + tuple(jnp.zeros_like(component) for component in E_tiles), + tuple(jnp.zeros_like(component) for component in B_tiles), + ) + fields = (E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, None) + + return SyntheticTiledYeeState( + particles=tiled_particles, + species_config=species_config, + fields=fields, + world=world, + constants=constants, + tile_shape=tile_shape, + g=g, + ) + + +def _tiled_particle_push_stage(case, relativistic=True, particle_pusher="boris"): + E_tiles, B_tiles, *_ = case.fields + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, case.fields[5]) + with jax.named_scope("tiled_particle_push"): + return tiled_particle_push( + case.particles, + case.species_config, + push_E_tiles, + push_B_tiles, + case.world, + case.constants, + case.tile_shape, + case.g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + +def _tiled_particle_retile_stage(case): + with jax.named_scope("tiled_particle_retile"): + moved = update_tiled_particle_positions(case.particles, case.species_config, case.world["dt"]) + return refresh_tiled_particle_tiles(moved, case.world, case.tile_shape) + + +def _tiled_current_deposition_stage(case, particles, J_func): + _, _, J_tiles, *_ = case.fields + with jax.named_scope("tiled_current_deposition"): + return J_func( + particles, + case.species_config, + J_tiles, + case.constants, + case.world, + tile_shape=case.tile_shape, + g=case.g, + ) + + +def _tiled_field_update_stage(case, J_tiles): + E_tiles, B_tiles, *_ = case.fields + with jax.named_scope("tiled_field_update"): + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, case.world, case.constants, unused_curl, case.tile_shape, case.g) + B_tiles = update_tiled_B(E_tiles, B_tiles, case.world, case.constants, unused_curl, case.tile_shape, case.g) + return E_tiles, B_tiles + + +def _tiled_diagnostics_stage(case): + E_tiles, B_tiles, *_ = case.fields + with jax.named_scope("tiled_diagnostics"): + return compute_energy( + case.particles, + E_tiles, + B_tiles, + case.world, + case.constants, + species_config=case.species_config, + ) + + +def _tiled_output_bridge_stage(case): + E_tiles, B_tiles, *_ = case.fields + with jax.named_scope("tiled_output_bridge"): + E = vector_field_for_output(E_tiles, case.world) + B = vector_field_for_output(B_tiles, case.world) + particles = particles_for_output(case.particles, case.species_config) + return E, B, particles + + +def benchmark_tiled_pic_step(case, J_func, relativistic=True, particle_pusher="boris"): + with jax.named_scope("tiled_pic_step"): + return time_loop_electrodynamic_tiled( + case.particles, + case.species_config, + case.fields, + case.world, + case.constants, + unused_curl, + J_func, + "tiled_yee", + tile_shape=case.tile_shape, + g=case.g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + +def build_tiled_stage_specs(case, J_func=None, relativistic=True, particle_pusher="boris"): + """ + Return independently timed tiled Yee stage specs. + + These specs identify the expensive tile-management stages before larger + kernel rewrites: push, retile, deposition, field update, diagnostics, and + the current global output bridge. + """ + + if J_func is None: + J_func = partial(direct_J_from_tiled_particles, filter="none") + + pushed_particles = _tiled_particle_push_stage(case, relativistic=relativistic, particle_pusher=particle_pusher) + deposited_J = _tiled_current_deposition_stage(case, pushed_particles, J_func) + + return { + "tiled_pic_step": StageSpec( + "tiled_pic_step", + benchmark_tiled_pic_step, + (case, J_func), + {"relativistic": relativistic, "particle_pusher": particle_pusher}, + ("J_func", "relativistic", "particle_pusher"), + ), + "tiled_particle_push": StageSpec( + "tiled_particle_push", + _tiled_particle_push_stage, + (case,), + {"relativistic": relativistic, "particle_pusher": particle_pusher}, + ("relativistic", "particle_pusher"), + ), + "tiled_particle_retile": StageSpec( + "tiled_particle_retile", + _tiled_particle_retile_stage, + (case,), + {}, + (), + ), + "tiled_current_deposition": StageSpec( + "tiled_current_deposition", + _tiled_current_deposition_stage, + (case, pushed_particles, J_func), + {}, + ("J_func",), + ), + "tiled_field_update": StageSpec( + "tiled_field_update", + _tiled_field_update_stage, + (case, deposited_J), + {}, + (), + ), + "tiled_diagnostics": StageSpec( + "tiled_diagnostics", + _tiled_diagnostics_stage, + (case,), + {}, + (), + ), + "tiled_output_bridge": StageSpec( + "tiled_output_bridge", + _tiled_output_bridge_stage, + (case,), + {}, + (), + ), + } diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index de5701cb..6f03df4e 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -225,7 +225,7 @@ 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": 1, # tile guard-cell depth shared by tiled fields and tiled currents + "guard_cells": 2, # tile guard-cell depth shared by tiled fields and tiled currents "particle_tile_nx": 1, # number of x cells per shared field/particle tile "particle_tile_ny": 1, # number of y cells per shared field/particle tile "particle_tile_nz": 1, # number of z cells per shared field/particle tile @@ -309,6 +309,8 @@ def initialize_simulation(toml_file): particle_pusher = simulation_parameters['particle_pusher'] validate_particle_pusher(particle_pusher) guard_cells = int(simulation_parameters["guard_cells"]) + if solver == "tiled_yee": + guard_cells = max(guard_cells, 2) _validate_current_filter_contract(simulation_parameters) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 2924a68d..e8325f5b 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -39,7 +39,7 @@ def _reduced_tiled_axes(field_tiles, tile_shape, g): ) -def empty_tiled_scalar_field(world, tile_shape, num_guard_cells=1, dtype=None): +def empty_tiled_scalar_field(world, tile_shape, num_guard_cells=2, dtype=None): """ Allocate an empty tile-major scalar field with ``num_guard_cells`` guards. """ @@ -65,7 +65,7 @@ def empty_tiled_scalar_field(world, tile_shape, num_guard_cells=1, dtype=None): ) -def empty_tiled_vector_field(world, tile_shape, num_guard_cells=1, dtype=None): +def empty_tiled_vector_field(world, tile_shape, num_guard_cells=2, dtype=None): """ Allocate empty tile-major vector-field components. """ @@ -73,6 +73,40 @@ def empty_tiled_vector_field(world, tile_shape, num_guard_cells=1, dtype=None): return tuple(empty_tiled_scalar_field(world, tile_shape, num_guard_cells, dtype) for _ in range(3)) +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): + """ + Stack tiled vector components onto one leading component axis. + + The public tiled field state still uses the existing ``(Fx, Fy, Fz)`` + tuples. This packed view gives halo/fold/filter routines one vector-field + surface, which is the shape later shard-local exchange should operate on. + """ + + 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): + """ + Return a tiled vector field as the existing ``(Fx, Fy, Fz)`` tuple. + """ + + 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 _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]) @@ -96,7 +130,7 @@ def _tile_grid_axis(global_axis_grid, world, tile_shape, tile_counts, axis_index return jnp.broadcast_to(axis_lines.reshape(axis_shape), tiled_shape) -def tile_grid_axes(grid, world, tile_shape, num_guard_cells=1): +def tile_grid_axes(grid, world, tile_shape, num_guard_cells=2): """ Build tile-local coordinate lines for a center or vertex grid. @@ -122,7 +156,7 @@ def tile_grid_axes(grid, world, tile_shape, num_guard_cells=1): ) -def tiled_grid_axes_from_world(world, grid, tiled_grid_key, tile_shape, num_guard_cells=1): +def tiled_grid_axes_from_world(world, grid, tiled_grid_key, tile_shape, num_guard_cells=2): """ Read precomputed tile-local grid axes, with a construction fallback. @@ -138,7 +172,7 @@ def tiled_grid_axes_from_world(world, grid, tiled_grid_key, tile_shape, num_guar return tile_grid_axes(grid, world, tile_shape, num_guard_cells) -def tile_scalar_field(field, world, tile_shape, num_guard_cells=1): +def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): """ Split a ghost-celled field into compact tiles using the shared tile shape. """ @@ -188,7 +222,7 @@ def _tile_scalar_field(field, tile_shape): return tile_scalar_field(field, None, tile_shape) -def tile_vector_field(field, world, tile_shape, num_guard_cells=1): +def tile_vector_field(field, world, tile_shape, num_guard_cells=2): """ Split ``(Fx, Fy, Fz)`` into compact ghost-celled tiles. @@ -200,7 +234,7 @@ def tile_vector_field(field, world, tile_shape, num_guard_cells=1): return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) -def assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells=1): +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. """ @@ -229,7 +263,7 @@ def assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells= return update_ghost_cells(field, bc_x, bc_y, bc_z) -def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells=1): +def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells=2): """ Assemble tiled vector-field components into ordinary ghost-celled arrays. """ @@ -241,7 +275,7 @@ def _assemble_scalar_tiles(field_tiles, world, tile_shape): return assemble_tiled_scalar_field(field_tiles, world, tile_shape) -def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): +def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): """ Refresh tile halos using the field boundary conditions. @@ -342,7 +376,7 @@ def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=N return field_tiles -def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): +def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=2): """ Refresh tile halos with all-periodic boundary conditions. """ @@ -351,23 +385,35 @@ def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): return update_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): +def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): """ - Refresh tile halos for each component of a vector field. + Refresh tile halos for a tiled vector field. """ - return tuple(update_tiled_ghost_cells(component, world, num_guard_cells, tile_shape) for component in field_tiles) + 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_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): +def update_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): """ Refresh periodic tile halos for each component of a vector field. """ - return tuple(update_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) + stacked_tiles = stack_tiled_vector_field(field_tiles) + refreshed = jax.vmap( + lambda component: update_tiled_ghost_cells_periodic(component, num_guard_cells), + 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=1, tile_shape=None): +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. """ @@ -385,15 +431,21 @@ def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=1, tile 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=1, tile_shape=None): +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. """ - return tuple(update_tiled_ghost_cells_for_pml(component, world, num_guard_cells, tile_shape) for component in field_tiles) + 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=1): +def apply_tiled_conducting_bc(E_tiles, world, num_guard_cells=2): """ Zero tangential electric-field components on global conducting faces. """ @@ -449,7 +501,7 @@ def apply_tiled_conducting_bc(E_tiles, world, num_guard_cells=1): return Ex, Ey, Ez -def digital_filter_tiled_scalar(field_tiles, alpha, num_guard_cells=1): +def digital_filter_tiled_scalar(field_tiles, alpha, num_guard_cells=2): """ Apply the standard six-neighbor digital filter to each tile interior. @@ -476,15 +528,21 @@ def digital_filter_tiled_scalar(field_tiles, alpha, num_guard_cells=1): return field_tiles.at[:, :, :, active, active, active].set(filtered) -def digital_filter_tiled_vector(field_tiles, alpha, num_guard_cells=1): +def digital_filter_tiled_vector(field_tiles, alpha, num_guard_cells=2): """ Apply the standard field filter component-wise in tile-major storage. """ - return tuple(digital_filter_tiled_scalar(component, alpha, num_guard_cells) for component in field_tiles) + stacked_tiles = stack_tiled_vector_field(field_tiles) + filtered = jax.vmap( + lambda component: digital_filter_tiled_scalar(component, alpha, num_guard_cells), + in_axes=0, + out_axes=0, + )(stacked_tiles) + return _restore_tiled_vector_layout(filtered, field_tiles) -def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): +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. @@ -609,7 +667,7 @@ def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=Non return field_tiles -def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): +def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=2): """ Fold all-periodic tile-ghost deposits for one scalar component. """ @@ -618,20 +676,32 @@ def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=1): return fold_tiled_ghost_cells(field_tiles, world, num_guard_cells) -def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=1, tile_shape=None): +def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): """ - Fold tile-ghost deposits for each vector component. + Fold tile-ghost deposits for a tiled vector field. """ - return tuple(fold_tiled_ghost_cells(component, world, num_guard_cells, tile_shape) for component in field_tiles) + 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) -def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=1): +def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): """ Fold all-periodic tile-ghost deposits for each vector component. """ - return tuple(fold_tiled_ghost_cells_periodic(component, num_guard_cells) for component in field_tiles) + stacked_tiles = stack_tiled_vector_field(field_tiles) + folded = jax.vmap( + lambda component: fold_tiled_ghost_cells_periodic(component, num_guard_cells), + in_axes=0, + out_axes=0, + )(stacked_tiles) + return _restore_tiled_vector_layout(folded, field_tiles) def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state=None): diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 4950d1d2..6db5dd8a 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -213,13 +213,14 @@ def nd_trapezoid(arr, dxs): # 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(1, -1), - slice(1, -1), - slice(1, -1), + slice(g, -g), + slice(g, -g), + slice(g, -g), ) else: interior = (slice(1, -1), slice(1, -1), slice(1, -1)) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index fdfbb4f3..56fc7172 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -279,14 +279,14 @@ def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self 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 = fold_tiled_ghost_cells(tiles, world) + folded = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) 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)) - legacy_folded = fold_tiled_ghost_cells_periodic(tiles) + legacy_folded = fold_tiled_ghost_cells_periodic(tiles, num_guard_cells=1) self.assertTrue(jnp.allclose(folded, legacy_folded, rtol=1.0e-12, atol=1.0e-12)) def test_fold_tiled_ghost_cells_two_guard_layers_adds_deposits_to_neighbors(self): @@ -338,7 +338,7 @@ def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): 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 = fold_tiled_ghost_cells(tiles, world) + folded = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) self.assertEqual(float(folded[0, 0, 0, 1, 1, 1]), -4.0) self.assertEqual(float(folded[-1, 0, 0, -2, 1, 1]), -7.0) @@ -376,8 +376,8 @@ def test_fold_tiled_ghost_cells_matches_global_fold_for_mixed_boundaries(self): field = field.at[3, 3, -1].set(5.0) tiles = tiles.at[1, 1, -1, 1, 1, -1].set(5.0) - folded_tiles = fold_tiled_ghost_cells(tiles, world) - folded_from_tiles = assemble_tiled_vector_field((folded_tiles, folded_tiles, folded_tiles), world, tile_shape)[0] + folded_tiles = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) + 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, diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 2bc64d1c..a30719a3 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -30,6 +30,7 @@ def test_default_parameters(self): 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.assertIn('eps', const) self.assertIn('plotfields', plotting) # check that the default parameters contain expected keys diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index b7bbc3d9..fbc70d7e 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -128,7 +128,7 @@ def test_write_openpmd_fields_assembles_tiled_fields_before_output(self): "Ny": 2, "Nz": 2, "tile_shape": (2, 1, 1), - "guard_cells": 1, + "guard_cells": 2, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } tiled_fields = ( diff --git a/tests/pml_test.py b/tests/pml_test.py index b380db75..d5a96424 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -39,6 +39,7 @@ def _base_world(nx=24, ny=1, nz=1): "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) @@ -169,7 +170,7 @@ def test_initialize_tiled_pml_state_uses_tile_local_interior_shape(self): 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, 4, 4, 3)) + 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) @@ -185,7 +186,7 @@ def test_tile_pml_profiles_matches_global_profiles_on_tile_interiors(self): tiled_profiles = tile_pml_profiles(world, tile_shape) global_profiles = world["pml"][-1] - assembled_profiles = assemble_tiled_vector_field(tiled_profiles, world, tile_shape) + 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( @@ -244,7 +245,8 @@ def test_initialize_simulation_uses_tiled_pml_state_for_tiled_yee(self): 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(tiled_profiles[0].shape, (4, 1, 1, 4, 3, 3)) + 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_fdtd(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -355,7 +357,7 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): constants, lambda *args: None, tile_shape, - 1, + int(world["guard_cells"]), tiled_pml_state, ) B_tiles, tiled_pml_state = update_tiled_B( @@ -365,12 +367,12 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): constants, lambda *args: None, tile_shape, - 1, + int(world["guard_cells"]), tiled_pml_state, ) - E_tiled = assemble_tiled_vector_field(E_tiles, world, tile_shape) - B_tiled = assemble_tiled_vector_field(B_tiles, world, tile_shape) + 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): @@ -406,10 +408,10 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) def step(E_tiles, B_tiles, tiled_pml_state): E_tiles, tiled_pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, 1, tiled_pml_state + E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, int(world["guard_cells"]), tiled_pml_state ) B_tiles, tiled_pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, 1, tiled_pml_state + E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, int(world["guard_cells"]), tiled_pml_state ) return E_tiles, B_tiles, tiled_pml_state @@ -438,7 +440,7 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): E_tiles = tile_vector_field((Ex, Ey, Ez), world, tile_shape) B_tiles = tile_vector_field((Bx, By, Bz), world, tile_shape) - g = 1 + 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) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 962480ff..4a954ea5 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -49,6 +49,7 @@ def _build_world(self, shape_factor=1): "y_wind": 1.0, "z_wind": 1.0, "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) @@ -353,9 +354,9 @@ def test_tiled_rho_assembles_to_global_rho(self): world, constants, tile_shape=world["tile_shape"], - g=1, + g=int(world["guard_cells"]), ) - assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"]) + assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) self.assertTrue( jnp.allclose( @@ -407,7 +408,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): J_func=None, solver="tiled_yee", tile_shape=world["tile_shape"], - g=1, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -415,15 +416,15 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields - for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, world["tile_shape"])): + for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-8, atol=1.0e-8)) - for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, world["tile_shape"])): + for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, world["tile_shape"])): + for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"]) - phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, world["tile_shape"]) + rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) self.assertTrue(jnp.allclose(rho_from_tiles[1:-1, 1:-1, 1:-1], reference_rho[1:-1, 1:-1, 1:-1], rtol=1.0e-12, atol=1.0e-12)) tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] reference_phi_interior = reference_phi[1:-1, 1:-1, 1:-1] @@ -473,7 +474,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self J_func=None, solver="tiled_yee", tile_shape=tile_shape, - g=1, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -548,8 +549,9 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): self.assertTrue(jnp.allclose(vertex_axis, center_axis)) self.assertIn("tiled_center_grid", world["grids"]) self.assertIn("tiled_vertex_grid", world["grids"]) - self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (4, 1, 1, 4)) - self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (4, 1, 1, 3)) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (4, 1, 1, 6)) + self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (4, 1, 1, 5)) for tiled_vertex_axis, tiled_center_axis in zip( world["grids"]["tiled_vertex_grid"], world["grids"]["tiled_center_grid"], diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 0b89019c..bc9c6176 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -616,7 +616,8 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): 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(fields[2][0].shape[-3:], (4, 3, 3)) + 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): diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py new file mode 100644 index 00000000..c82085c7 --- /dev/null +++ b/tests/tiled_refactor_contract_test.py @@ -0,0 +1,103 @@ +import inspect +import unittest + +import jax.numpy as jnp + +from PyPIC3D import electrodynamic_tiled +from PyPIC3D.solvers import electrostatic_yee +from PyPIC3D.solvers import yee_tiled +from PyPIC3D.utils import build_yee_grid + + +class TestTiledRefactorContracts(unittest.TestCase): + def _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 _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) + return jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (ii + 2.0 * jj - kk)) + + def test_tiled_vector_field_stack_roundtrips_and_refreshes_halos(self): + world = self._world() + tile_shape = (2, 3, 2) + E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) + E_tiles = yee_tiled.tile_vector_field(E, world, tile_shape) + + stacked = yee_tiled.stack_tiled_vector_field(E_tiles) + self.assertEqual(stacked.shape[0], 3) + + unstacked = yee_tiled.unstack_tiled_vector_field(stacked) + for original, recovered in zip(E_tiles, unstacked): + self.assertTrue(jnp.allclose(original, recovered, rtol=1.0e-12, atol=1.0e-12)) + + refreshed = yee_tiled.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=1, tile_shape=tile_shape) + refreshed_from_stack = yee_tiled.unstack_tiled_vector_field( + yee_tiled.update_tiled_vector_ghost_cells(stacked, world, num_guard_cells=1, tile_shape=tile_shape) + ) + for reference, candidate in zip(refreshed, refreshed_from_stack): + self.assertTrue(jnp.allclose(reference, candidate, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_benchmark_stage_specs_expose_tiled_hot_paths(self): + from PyPIC3D.diagnostics import tiled_benchmarking + + case = tiled_benchmarking.build_synthetic_tiled_yee_state( + nx=4, + ny=4, + nz=4, + particles_per_species=8, + species_count=1, + tile_shape=(2, 2, 2), + slots_per_tile=8, + ) + stage_specs = tiled_benchmarking.build_tiled_stage_specs(case) + + expected = { + "tiled_pic_step", + "tiled_particle_push", + "tiled_particle_retile", + "tiled_current_deposition", + "tiled_field_update", + "tiled_diagnostics", + "tiled_output_bridge", + } + self.assertTrue(expected.issubset(stage_specs)) + self.assertEqual(stage_specs["tiled_pic_step"].static_argnames, ("J_func", "relativistic", "particle_pusher")) + self.assertEqual(stage_specs["tiled_current_deposition"].static_argnames, ("J_func",)) + + def test_electrodynamic_tiled_hot_step_does_not_assemble_global_fields(self): + source = inspect.getsource(electrodynamic_tiled.time_loop_electrodynamic_tiled) + self.assertNotIn("assemble_tiled_scalar_field", source) + self.assertNotIn("assemble_tiled_vector_field", source) + self.assertNotIn("fields_for_output", source) + + def test_electrostatic_global_poisson_bridge_stays_explicit(self): + source = inspect.getsource(electrostatic_yee.calculate_tiled_electrostatic_fields) + self.assertIn("assemble_tiled_scalar_field", source) + self.assertIn("tile_scalar_field", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 5aafaf6f..903f76ec 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -30,6 +30,7 @@ def _build_world(self, shape_factor, dt=0.08): "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) @@ -215,7 +216,7 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self world, constants, tile_shape=world["tile_shape"], - g=1, + g=int(world["guard_cells"]), ) rho_tiles_with_zero_velocity = compute_tiled_rho_from_tiled_particles( zero_velocity_tiled_particles, @@ -224,10 +225,10 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self world, constants, tile_shape=world["tile_shape"], - g=1, + g=int(world["guard_cells"]), ) - rho_with_velocity = assemble_tiled_scalar_field(rho_tiles_with_velocity, world, world["tile_shape"]) - rho_with_zero_velocity = assemble_tiled_scalar_field(rho_tiles_with_zero_velocity, world, world["tile_shape"]) + 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( diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 7435b862..d473e1f3 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -51,7 +51,7 @@ def _build_world(self, boundary_conditions=None): "y_wind": 1.0, "z_wind": 1.0, "shape_factor": 1, - "guard_cells": 1, + "guard_cells": 2, "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 1c9a2547..88e9945a 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -315,7 +315,7 @@ def test_update_tiled_E_matches_standard_yee_update(self): constants, unused_curl, tile_shape, - 1, + 2, ) E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) @@ -339,7 +339,7 @@ def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(s constants, unused_curl, tile_shape, - 1, + 2, ) E_from_tiles = assemble_tiled_vector_field(E_tiled, world, tile_shape) @@ -361,7 +361,7 @@ def test_update_tiled_B_matches_standard_yee_update(self): constants, unused_curl, tile_shape, - 1, + 2, ) B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) @@ -383,7 +383,7 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s constants, unused_curl, tile_shape, - 1, + 2, ) B_from_tiles = assemble_tiled_vector_field(B_tiled, world, tile_shape) @@ -404,8 +404,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 1) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 1) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) @@ -429,8 +429,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 1) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 1) + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) From 1cc74d8caae940953c0e1c62d07cfe1d2319a441 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 04:24:52 -0500 Subject: [PATCH 24/83] refactoring tiled particle refresh to take into account single cell particle movement --- PyPIC3D/particles/tiled_particle_refresh.py | 198 ++++++++++++-------- PyPIC3D/pusher/boris.py | 55 +++++- PyPIC3D/pusher/tiled_pusher.py | 36 +++- tests/boris_test.py | 39 ++++ tests/tiled_particle_refresh_test.py | 9 + tests/tiled_pusher_test.py | 86 ++++++++- 6 files changed, 325 insertions(+), 98 deletions(-) diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index a0859e6f..a119eeca 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -1,3 +1,4 @@ +import jax import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import wrap_periodic_position @@ -70,15 +71,75 @@ def update_tiled_particle_positions(tiled_particles, species_config, dt): return tiled_particles._replace(x=x) -def _neighbor_offsets(count): - offsets = [] - seen = set() - for offset in (-1, 0, 1): - key = offset % count - if key not in seen: - offsets.append(offset) - seen.add(key) - return tuple(offsets) +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 _compact_tile_candidates(candidate_x, candidate_u, candidate_active, n_slots): + """ + 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)) + + def compact_one(local_x_in, local_u_in, local_active_in): + 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)) + + return local_x, local_u, local_active_count > 0, jnp.sum(local_active_in) > n_slots + + flat_new_x, flat_new_u, flat_new_active, flat_overflow = jax.vmap(compact_one)( + flat_x, + flat_u, + flat_active, + ) + + 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 refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): @@ -86,7 +147,7 @@ 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 - destination tile only reads candidate particles from its neighboring tiles. + 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. """ @@ -137,71 +198,58 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): tile_counts, ) - new_x = jnp.zeros_like(tiled_particles.x) - new_u = jnp.zeros_like(tiled_particles.u) - new_active = jnp.zeros_like(tiled_particles.active) - overflow = jnp.asarray(False) - - x_offsets = _neighbor_offsets(ntx) - y_offsets = _neighbor_offsets(nty) - z_offsets = _neighbor_offsets(ntz) - - for tx in range(ntx): - for ty in range(nty): - for tz in range(ntz): - for species in range(n_species): - candidate_x = [] - candidate_u = [] - candidate_active = [] - candidate_dest_tx = [] - candidate_dest_ty = [] - candidate_dest_tz = [] - - for ox in x_offsets: - sx = (tx + ox) % ntx - for oy in y_offsets: - sy = (ty + oy) % nty - for oz in z_offsets: - sz = (tz + oz) % ntz - index = (sx, sy, sz, species) - candidate_x.append(bounded_x[index]) - candidate_u.append(bounded_u[index]) - candidate_active.append(bounded_active[index]) - candidate_dest_tx.append(dest_tx[index]) - candidate_dest_ty.append(dest_ty[index]) - candidate_dest_tz.append(dest_tz[index]) - - candidate_x = jnp.concatenate(candidate_x, axis=0) - candidate_u = jnp.concatenate(candidate_u, axis=0) - candidate_active = jnp.concatenate(candidate_active, axis=0) - candidate_dest_tx = jnp.concatenate(candidate_dest_tx, axis=0) - candidate_dest_ty = jnp.concatenate(candidate_dest_ty, axis=0) - candidate_dest_tz = jnp.concatenate(candidate_dest_tz, axis=0) - - keep = ( - candidate_active - & (candidate_dest_tx == tx) - & (candidate_dest_ty == ty) - & (candidate_dest_tz == tz) - ) - rank = jnp.cumsum(keep.astype(int)) - 1 - fits = keep & (rank < n_slots) - safe_rank = jnp.where(fits, rank, 0) - - local_x = jnp.zeros_like(new_x[tx, ty, tz, species]) - local_u = jnp.zeros_like(new_u[tx, ty, tz, species]) - local_active_count = jnp.zeros(n_slots, dtype=int) - - valid = fits.astype(candidate_x.dtype) - local_x = local_x.at[safe_rank].add(valid[:, None] * candidate_x) - local_u = local_u.at[safe_rank].add(valid[:, None] * candidate_u) - local_active_count = local_active_count.at[safe_rank].add(fits.astype(int)) - - index = (tx, ty, tz, species) - new_x = new_x.at[index].set(local_x) - new_u = new_u.at[index].set(local_u) - new_active = new_active.at[index].set(local_active_count > 0) - overflow = overflow | (jnp.sum(keep) > n_slots) + 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) + + 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, + ) refreshed = TiledParticles( x=new_x, 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/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index dddbb0dc..b531efd6 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -49,6 +49,12 @@ def tiled_particle_push(tiled_particles, species_config, E_tiles, B_tiles, world g = int(g) ntx, nty, ntz = tiled_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): @@ -75,13 +81,31 @@ def push_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, charge_species, mass_ 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) - efield_aty = interpolate_field_to_particles(Ey_tile, x, y, z, Ey_grid, shape_factor, ghost_cells=True) - efield_atz = interpolate_field_to_particles(Ez_tile, x, y, z, Ez_grid, shape_factor, ghost_cells=True) + 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 + ) + 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 + ) + 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) - bfield_aty = interpolate_field_to_particles(By_tile, x, y, z, By_grid, shape_factor, ghost_cells=True) - bfield_atz = interpolate_field_to_particles(Bz_tile, x, y, z, Bz_grid, shape_factor, ghost_cells=True) + 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 + ) + 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, diff --git a/tests/boris_test.py b/tests/boris_test.py index 96c02714..e94b1628 100644 --- a/tests/boris_test.py +++ b/tests/boris_test.py @@ -88,6 +88,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): diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 41779e6c..f4497dbb 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -6,6 +6,7 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh import ( + _adjacent_tile_offset, refresh_tiled_particle_tiles, update_tiled_particle_positions, ) @@ -95,6 +96,14 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): 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_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]) diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index fc132535..ed7fd515 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -116,8 +116,8 @@ def test_tiled_particle_push_matches_flat_boris(self): pushed = tiled_particle_push( tiled_particles, species_config, - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), + tile_vector_field(E, world, tile_shape, num_guard_cells=1), + tile_vector_field(B, world, tile_shape, num_guard_cells=1), world, constants, tile_shape, @@ -184,8 +184,8 @@ def make_species(): pushed = tiled_particle_push( tiled_particles, species_config, - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), + tile_vector_field(E, world, tile_shape, num_guard_cells=1), + tile_vector_field(B, world, tile_shape, num_guard_cells=1), world, constants, tile_shape, @@ -223,8 +223,8 @@ def test_tiled_particle_push_respects_active_and_update_flags(self): pushed = tiled_particle_push( tiled_particles, species_config, - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), + tile_vector_field(E, world, tile_shape, num_guard_cells=1), + tile_vector_field(B, world, tile_shape, num_guard_cells=1), world, constants, tile_shape, @@ -288,8 +288,8 @@ def make_species(): pushed = tiled_particle_push( tiled_particles, species_config, - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), + tile_vector_field(E, world, tile_shape, num_guard_cells=1), + tile_vector_field(B, world, tile_shape, num_guard_cells=1), world, constants, tile_shape, @@ -304,6 +304,76 @@ def make_species(): self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + def test_tiled_particle_push_matches_flat_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 + 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=0.1) + B = self._deterministic_vector_field(world, scale=0.02) + + def make_species(): + return particle_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 = particle_push( + reference, + E, + B, + world["grids"]["center"], + world["grids"]["vertex"], + world, + constants, + relativistic=True, + particle_pusher="boris", + ) + + tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + pushed = tiled_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, + tile_shape, + g, + relativistic=True, + ) + + _, tiled_u = self._flatten_active_by_position(pushed) + flat_u = jnp.stack(reference.get_velocity(), axis=1) + flat_x = jnp.stack(reference.get_forward_position(), axis=1) + order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + + self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + if __name__ == "__main__": unittest.main() From c199719b583eebe7ebb0fd2e7753170197d65872 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 04:56:17 -0500 Subject: [PATCH 25/83] refactoring esirkepov tiled deposition --- PyPIC3D/deposition/esirkepov_tiled.py | 16 +++++++++++++--- tests/tiled_esirkepov_test.py | 6 +++++- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 4476f5c7..1f91b8a0 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -21,6 +21,12 @@ ) +def _active_stencil_indices(axis_active): + if axis_active: + return (0, 1, 2, 3, 4) + return (2,) + + def fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2, tile_shape=None): """ Fold Esirkepov current deposits using the configured guard depth. @@ -272,9 +278,13 @@ def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): tile_Jy = Jy_template tile_Jz = Jz_template - for i in range(5): - for j in range(5): - for k in range(5): + x_stencil_indices = _active_stencil_indices(x_active) + y_stencil_indices = _active_stencil_indices(y_active) + z_stencil_indices = _active_stencil_indices(z_active) + + for i in x_stencil_indices: + for j in y_stencil_indices: + for k in z_stencil_indices: ix = xpts[i, :] iy = ypts[j, :] iz = zpts[k, :] diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index bc9c6176..35a2f7cc 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -10,7 +10,7 @@ from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current +from PyPIC3D.deposition.esirkepov_tiled import _active_stencil_indices, tiled_esirkepov_current from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.diagnostics.output_adapters import fields_for_output from PyPIC3D.initialization import CURRENT_ESIRKEPOV, initialize_simulation @@ -268,6 +268,10 @@ def test_update_tiled_E_reads_two_guard_current_interior(self): self.assertTrue(jnp.allclose(E_after[0][:, :, :, g:-g, g:-g, g:-g], -0.5)) + def test_reduced_axis_esirkepov_scatter_uses_only_collapsed_stencil_index(self): + self.assertEqual(_active_stencil_indices(True), (0, 1, 2, 3, 4)) + self.assertEqual(_active_stencil_indices(False), (2,)) + 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} From 1e549a17bd6969f01931a8fde82cbafa422643a3 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 06:24:33 -0500 Subject: [PATCH 26/83] refactoring esirkepov quasi-dimensional weight computation and the stay particle retile --- PyPIC3D/deposition/esirkepov_tiled.py | 230 +++++++++++++++----- PyPIC3D/particles/tiled_particle_refresh.py | 59 ++++- tests/tiled_esirkepov_test.py | 39 +++- tests/tiled_particle_refresh_test.py | 9 + 4 files changed, 281 insertions(+), 56 deletions(-) diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 1f91b8a0..3c7c6c77 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -9,8 +9,6 @@ ) from PyPIC3D.deposition.Esirkepov import ( collapse_redundant_axis, - get_1D_esirkepov_weights, - get_2D_esirkepov_weights, get_3D_esirkepov_weights, shift_old_stencil, ) @@ -27,6 +25,82 @@ def _active_stencil_indices(axis_active): return (2,) +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 fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2, tile_shape=None): """ Fold Esirkepov current deposits using the configured guard depth. @@ -229,25 +303,6 @@ def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): 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) - if x_active and y_active and z_active: - 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 not x_active: - null_dim = 0 - elif not y_active: - null_dim = 1 - else: - null_dim = 2 - 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): - 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): - Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=1) - else: - Wx_, Wy_, Wz_ = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles, dim=2) - dJx = jax.lax.cond( x_active, lambda _: active * (-(q / (dy * dz)) / dt), @@ -267,39 +322,114 @@ def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): operand=None, ) - 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) - tile_Jx = Jx_template tile_Jy = Jy_template tile_Jz = Jz_template - x_stencil_indices = _active_stencil_indices(x_active) - y_stencil_indices = _active_stencil_indices(y_active) - z_stencil_indices = _active_stencil_indices(z_active) - - for i in x_stencil_indices: - for j in y_stencil_indices: - for k in z_stencil_indices: - ix = xpts[i, :] - iy = ypts[j, :] - iz = zpts[k, :] - tile_Jx = tile_Jx.at[ix, iy, iz].add( - jax.lax.cond(x_active, lambda _: Jx_loc[i, j, k, :], lambda _: Fx[i, j, k, :], operand=None), - mode="drop", - ) - tile_Jy = tile_Jy.at[ix, iy, iz].add( - jax.lax.cond(y_active, lambda _: Jy_loc[i, j, k, :], lambda _: Fy[i, j, k, :], operand=None), - mode="drop", - ) - tile_Jz = tile_Jz.at[ix, iy, iz].add( - jax.lax.cond(z_active, lambda _: Jz_loc[i, j, k, :], lambda _: Fz[i, j, k, :], operand=None), - mode="drop", - ) + if x_active and y_active and z_active: + 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) + + for i in range(5): + for j in range(5): + for k 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 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): + 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): + 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): + 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): + 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): + 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: + 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 diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index a119eeca..0cd089e9 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -100,7 +100,36 @@ def _adjacent_tile_offset(dest_tile, source_tile, tile_count): return offset -def _compact_tile_candidates(candidate_x, candidate_u, candidate_active, n_slots): +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. """ @@ -111,8 +140,18 @@ def _compact_tile_candidates(candidate_x, candidate_u, candidate_active, n_slots 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)) - - def compact_one(local_x_in, local_u_in, local_active_in): + 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) @@ -126,12 +165,19 @@ def compact_one(local_x_in, local_u_in, local_active_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)) - return local_x, local_u, local_active_count > 0, jnp.sum(local_active_in) > n_slots + 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)) @@ -205,6 +251,7 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): 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) + repack_mask = _tiles_need_repack(offset_x, offset_y, offset_z, bounded_active) x_offsets = _movement_offsets(ntx) y_offsets = _movement_offsets(nty) @@ -249,6 +296,10 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): candidate_u, candidate_active, n_slots, + fallback_x=bounded_x, + fallback_u=bounded_u, + fallback_active=bounded_active, + repack_mask=repack_mask, ) refreshed = TiledParticles( diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 35a2f7cc..94416dc1 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -9,8 +9,13 @@ import toml from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.esirkepov_tiled import _active_stencil_indices, tiled_esirkepov_current +from PyPIC3D.deposition.Esirkepov import Esirkepov_current, get_1D_esirkepov_weights, get_2D_esirkepov_weights +from PyPIC3D.deposition.esirkepov_tiled import ( + _active_stencil_indices, + _compact_1d_esirkepov_weights, + _compact_2d_esirkepov_weights, + tiled_esirkepov_current, +) from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.diagnostics.output_adapters import fields_for_output from PyPIC3D.initialization import CURRENT_ESIRKEPOV, initialize_simulation @@ -272,6 +277,36 @@ def test_reduced_axis_esirkepov_scatter_uses_only_collapsed_stencil_index(self): self.assertEqual(_active_stencil_indices(True), (0, 1, 2, 3, 4)) self.assertEqual(_active_stencil_indices(False), (2,)) + def test_compact_1d_esirkepov_weights_match_full_center_line(self): + xw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.3]), jnp.array([0.6, 0.5]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] + oxw = [jnp.array([0.0, 0.0]), jnp.array([0.1, 0.2]), jnp.array([0.7, 0.6]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] + yw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] + zw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] + oyw = yw + ozw = zw + + Wx_line, Wy_line, Wz_line = _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=0) + Wx_full, Wy_full, Wz_full = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles=2, dim=0) + + self.assertTrue(jnp.allclose(Wx_line, Wx_full[:, 2, 2, :])) + self.assertTrue(jnp.allclose(Wy_line, Wy_full[:, 2, 2, :])) + self.assertTrue(jnp.allclose(Wz_line, Wz_full[:, 2, 2, :])) + + def test_compact_2d_esirkepov_weights_match_full_active_plane(self): + xw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.3]), jnp.array([0.6, 0.5]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] + oxw = [jnp.array([0.0, 0.0]), jnp.array([0.1, 0.2]), jnp.array([0.7, 0.6]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] + yw = [jnp.array([0.0, 0.0]), jnp.array([0.3, 0.2]), jnp.array([0.4, 0.5]), jnp.array([0.3, 0.3]), jnp.array([0.0, 0.0])] + oyw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.1]), jnp.array([0.5, 0.6]), jnp.array([0.3, 0.3]), jnp.array([0.0, 0.0])] + zw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] + ozw = zw + + Wx_plane, Wy_plane, Wz_plane = _compact_2d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, null_dim=2) + Wx_full, Wy_full, Wz_full = get_2D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles=2, null_dim=2) + + self.assertTrue(jnp.allclose(Wx_plane, Wx_full[:, :, 2, :])) + self.assertTrue(jnp.allclose(Wy_plane, Wy_full[:, :, 2, :])) + self.assertTrue(jnp.allclose(Wz_plane, Wz_full[:, :, 2, :])) + 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} diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index f4497dbb..3ef93de0 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -7,6 +7,7 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh import ( _adjacent_tile_offset, + _tiles_need_repack, refresh_tiled_particle_tiles, update_tiled_particle_positions, ) @@ -104,6 +105,14 @@ def test_adjacent_tile_offset_handles_periodic_edges(self): 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]) From 61ae6306d67a23a11a3240a16fdb3ac8badd1e64 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 07:38:55 -0500 Subject: [PATCH 27/83] refactoring tiled particles --- PyPIC3D/particles/tiled_particle_refresh.py | 178 +++++++++++++++++++- tests/tiled_particle_refresh_test.py | 39 +++++ 2 files changed, 210 insertions(+), 7 deletions(-) diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/tiled_particle_refresh.py index 0cd089e9..2d9c906a 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/tiled_particle_refresh.py @@ -188,14 +188,9 @@ def compact_one(local_x_in, local_u_in, local_active_in, fallback_x_in, fallback return new_x, new_u, new_active, overflow -def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): +def _bounded_state_and_tile_offsets(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. + Apply physical particle boundaries and identify the adjacent tile offset. """ ntx, nty, ntz, n_species, n_slots = tiled_particles.active.shape @@ -251,6 +246,22 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): 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) @@ -309,3 +320,156 @@ def refresh_tiled_particle_tiles(tiled_particles, world, tile_shape): ) 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/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 3ef93de0..491f3ccb 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -7,6 +7,8 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh 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, @@ -187,6 +189,43 @@ def test_refresh_reports_overflow_without_changing_shape(self): 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 = to_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 = to_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() From 33f024d7f36299598fe019e3768a29c137a33cc8 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 10:54:59 -0500 Subject: [PATCH 28/83] finished tasks 0 and 1 --- PyPIC3D/__main__.py | 62 +++----- PyPIC3D/diagnostics/tiled_benchmarking.py | 2 +- PyPIC3D/initialization.py | 172 ++++++++-------------- demos/ot_vortex/orszag_tang.toml | 1 - demos/reconnection_2d/harris_current.toml | 1 - demos/two_stream/two_stream.toml | 2 - demos/weibel/weibel.toml | 1 - docs/architecture.rst | 2 +- docs/solvers.rst | 4 +- docs/usage.rst | 3 +- tests/initialization_test.py | 88 +++++++++-- tests/pml_test.py | 16 +- tests/tiled_electrostatic_test.py | 8 +- tests/tiled_esirkepov_integration_test.py | 3 +- tests/tiled_esirkepov_test.py | 12 +- tests/tiled_fluid_quantities_test.py | 8 +- tests/tiled_yee_integration_test.py | 33 ++--- 17 files changed, 200 insertions(+), 218 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index cbf2a304..38c0ed96 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -54,12 +54,12 @@ def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): - if simulation_parameters["solver"] != "tiled_yee" or len(fields) < 8: + if len(fields) < 8: return overflow = fields[-1] if bool(jax.device_get(overflow)): - raise RuntimeError("tiled_yee particle tile capacity overflowed during periodic retile") + raise RuntimeError("tiled particle tile capacity overflowed during periodic retile") def run_PyPIC3D(config_file): @@ -75,17 +75,14 @@ def run_PyPIC3D(config_file): scalar_field_names = ["rho", "mass_density"] vector_field_names = ["E", "B", "J"] - tiled_run = simulation_parameters["solver"] == "tiled_yee" - tile_shape = tuple(int(width) for width in world["tile_shape"]) if tiled_run else None - g = int(world["guard_cells"]) if tiled_run else None + tiled_run = True + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) particle_species_names = simulation_parameters.get("particle_species_names") - if tiled_run: - jit_loop = jax.jit( - loop, - static_argnames=('curl_func', 'J_func', 'solver', 'tile_shape', 'g', 'relativistic', 'particle_pusher'), - ) - else: - jit_loop = jax.jit(loop, static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher')) + jit_loop = jax.jit( + loop, + static_argnames=('curl_func', 'J_func', 'solver', 'tile_shape', 'g', 'relativistic', 'particle_pusher'), + ) E, B, J, rho, phi, external_fields, *rest = fields # unpack the fields @@ -194,33 +191,20 @@ def run_PyPIC3D(config_file): fields = (E, B, J, rho, phi, external_fields, *rest) # repack the fields for non-tiled diagnostics that updated rho - if tiled_run: - particles, fields = jit_loop( - particles, - species_config, - fields, - world, - constants, - curl_func, - J_func, - solver, - tile_shape=tile_shape, - g=g, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - else: - particles, fields = jit_loop( - particles, - fields, - world, - constants, - curl_func, - J_func, - solver, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) + particles, fields = jit_loop( + particles, + species_config, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=tile_shape, + g=g, + 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 diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index 95f0aa23..dd31e9b9 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -282,7 +282,7 @@ def benchmark_tiled_pic_step(case, J_func, relativistic=True, particle_pusher="b case.constants, unused_curl, J_func, - "tiled_yee", + "electrodynamic_yee", tile_shape=case.tile_shape, g=case.g, relativistic=relativistic, diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 6f03df4e..c5fd24fa 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -35,10 +35,6 @@ write_openpmd_initial_particles, write_openpmd_initial_fields ) -from PyPIC3D.particles.flat_particles import ( - to_flat_particles, check_flat_compat -) - from PyPIC3D.particles.tiled_particle_initialization import ( to_tiled_particles ) @@ -51,10 +47,6 @@ ) -from PyPIC3D.evolve import ( - time_loop_electrodynamic, time_loop_electrostatic -) - from PyPIC3D.electrodynamic_tiled import ( time_loop_electrodynamic_tiled ) @@ -65,11 +57,7 @@ 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.deposition.current_methods import ( - CURRENT_ESIRKEPOV, - CURRENT_J_FROM_RHOV, encode_current_calculation, ) from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles @@ -78,7 +66,6 @@ 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, @@ -117,9 +104,11 @@ 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 = ("fdtd", "tiled_yee") + supported_solvers = ("electrodynamic_yee", "electrostatic") if solver not in supported_solvers: - raise ValueError(f"Unsupported solver: {solver}. Use 'fdtd' or 'tiled_yee'.") + raise ValueError( + f"Unsupported solver: {solver}. Use 'electrodynamic_yee' or 'electrostatic'." + ) def _tile_shape_from_parameters(simulation_parameters): @@ -148,11 +137,11 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ """ if simulation_parameters["current_calculation"] not in ("j_from_rhov", "esirkepov"): - raise ValueError("tiled_yee currently supports current_calculation='j_from_rhov' or '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("tiled_yee currently supports only particle_pusher='boris' or '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("tiled_yee currently supports only filter_j='none', filter_j='digital', or filter_j='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"]), @@ -161,7 +150,7 @@ def _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_ ) for cells, tile_width in zip(grid_shape, tile_shape): if cells % tile_width != 0: - raise ValueError("tiled_yee requires the shared tile shape to divide Nx/Ny/Nz exactly") + raise ValueError("Yee runtime requires the shared tile shape to divide Nx/Ny/Nz exactly") def default_parameters(): @@ -198,8 +187,7 @@ def default_parameters(): simulation_parameters = { "name": "Default Simulation", "output_dir": os.getcwd(), - "solver": "fdtd", # solver: fdtd or tiled_yee - "fast_backend": "flat", # flat | default (flat when compatible, else fallback) + "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 @@ -216,7 +204,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 @@ -278,8 +265,8 @@ 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 "fdtd" or "tiled_yee". - 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. @@ -302,15 +289,16 @@ 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) guard_cells = int(simulation_parameters["guard_cells"]) - if solver == "tiled_yee": - guard_cells = max(guard_cells, 2) + guard_cells = max(guard_cells, 2) _validate_current_filter_contract(simulation_parameters) verbose = simulation_parameters['verbose'] GPUs = simulation_parameters['GPUs'] @@ -378,10 +366,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 not in ("fdtd", "tiled_yee") or electrostatic): - raise ValueError("PML is only supported for the fdtd or tiled_yee electrodynamic solvers") - if solver == "tiled_yee": - _validate_tiled_yee_configuration(simulation_parameters, electrostatic, pml_active) + 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) @@ -441,26 +428,6 @@ def initialize_simulation(toml_file): write_openpmd_initial_particles(particles, world, constants, simulation_parameters['output_dir']) # write the initial particles to an openPMD file - - fast_backend = simulation_parameters.get("fast_backend", "flat") - if fast_backend == "flat" and solver == "tiled_yee": - print("fast_backend='flat' not used for tiled_yee; using tiled particle storage") - simulation_parameters["fast_backend"] = "tiled" - elif fast_backend == "flat": - if electrostatic: - print("fast_backend='flat' not supported for electrostatic; 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) # initialize the electric and magnetic fields external_fields = ( @@ -522,85 +489,66 @@ def initialize_simulation(toml_file): if electrostatic: print("Using electrostatic solver") - if solver == "tiled_yee": - evolve_loop = time_loop_electrostatic_tiled - else: - evolve_loop = time_loop_electrostatic - + evolve_loop = time_loop_electrostatic_tiled else: - print(f"Using electrodynamic solver with: {solver}") - if solver == "tiled_yee": - evolve_loop = time_loop_electrodynamic_tiled - else: - evolve_loop = time_loop_electrodynamic + print("Using electrodynamic Yee solver") + evolve_loop = time_loop_electrodynamic_tiled # set the evolve loop function based on the electrostatic flag if simulation_parameters['current_calculation'] == "esirkepov": print("Using Esirkepov current calculation method") - if solver == "tiled_yee": - J_func = tiled_esirkepov_current - else: - J_func = Esirkepov_current + J_func = tiled_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']}") - if solver == "tiled_yee": - J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) - else: - J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) + J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) species_config = None - if solver == "tiled_yee": - tile_shape = _tile_shape_from_parameters(simulation_parameters) - world["tile_shape"] = tile_shape - guard_cells = int(world["guard_cells"]) - world["grids"]["tiled_center_grid"] = tile_grid_axes( - world["grids"]["center"], - world, - tile_shape, - num_guard_cells=guard_cells, - ) - world["grids"]["tiled_vertex_grid"] = tile_grid_axes( - world["grids"]["vertex"], - world, - tile_shape, - num_guard_cells=guard_cells, - ) - particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) - B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) - if simulation_parameters["current_calculation"] == "esirkepov": - J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) - J_func = tiled_esirkepov_current - else: - J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) - J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) - external_E, external_B = external_fields - external_fields = ( - tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), - tile_vector_field(external_B, world, tile_shape, num_guard_cells=guard_cells), - ) - simulation_parameters["fast_backend"] = "tiled" - simulation_parameters["tile_shape"] = tile_shape - print(f"Using tiled Yee storage with tile shape: {tile_shape}") + tile_shape = _tile_shape_from_parameters(simulation_parameters) + world["tile_shape"] = tile_shape + guard_cells = int(world["guard_cells"]) + world["grids"]["tiled_center_grid"] = tile_grid_axes( + world["grids"]["center"], + world, + tile_shape, + num_guard_cells=guard_cells, + ) + world["grids"]["tiled_vertex_grid"] = tile_grid_axes( + world["grids"]["vertex"], + world, + tile_shape, + num_guard_cells=guard_cells, + ) + particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) + B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) + if simulation_parameters["current_calculation"] == "esirkepov": + J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) + J_func = tiled_esirkepov_current + else: + J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) + J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) + external_E, external_B = external_fields + external_fields = ( + tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), + tile_vector_field(external_B, world, tile_shape, num_guard_cells=guard_cells), + ) + simulation_parameters["tile_shape"] = tile_shape + print(f"Using tiled Yee storage with tile shape: {tile_shape}") if electrostatic: - if solver == "tiled_yee": - rho = tile_scalar_field(rho, world, tile_shape, num_guard_cells=guard_cells) - phi = tile_scalar_field(phi, world, tile_shape, num_guard_cells=guard_cells) + rho = tile_scalar_field(rho, world, tile_shape, num_guard_cells=guard_cells) + phi = tile_scalar_field(phi, world, tile_shape, num_guard_cells=guard_cells) fields = (E, B, J, rho, phi, external_fields) # define the fields tuple for the electrostatic solver else: pml_state = None - # electrodynamic FDTD always carries the PML state slot; None means + # electrodynamic Yee always carries the PML state slot; None means # ordinary, unstretched Yee derivatives. - if solver == "tiled_yee": - if pml_active: - pml_state = initialize_tiled_pml_state(world, tile_shape) - elif pml_active: - pml_state = initialize_pml_state(world) + if pml_active: + pml_state = initialize_tiled_pml_state(world, tile_shape) fields = (E, B, J, rho, phi, external_fields, pml_state) - # define the fields tuple for the electrodynamic FDTD solver + # 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") diff --git a/demos/ot_vortex/orszag_tang.toml b/demos/ot_vortex/orszag_tang.toml index fa010cd5..e919d48f 100644 --- a/demos/ot_vortex/orszag_tang.toml +++ b/demos/ot_vortex/orszag_tang.toml @@ -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..d5d21e63 100644 --- a/demos/reconnection_2d/harris_current.toml +++ b/demos/reconnection_2d/harris_current.toml @@ -12,7 +12,6 @@ x_wind = 0.0008 y_wind = 1 z_wind = 0.0008 verbose = false -electrostatic = false cfl = 0.1 shape_factor = 2 relativistic = true diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 359addce..5cac4a46 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -10,11 +10,9 @@ 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 diff --git a/demos/weibel/weibel.toml b/demos/weibel/weibel.toml index aa17799b..dd3e3a3a 100644 --- a/demos/weibel/weibel.toml +++ b/demos/weibel/weibel.toml @@ -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 diff --git a/docs/architecture.rst b/docs/architecture.rst index 54424bd2..fd196694 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -48,7 +48,7 @@ Main runtime objects: Tiled State Contract -------------------- -The ``tiled_yee`` path uses one shared tile shape for fields and particles. +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 diff --git a/docs/solvers.rst b/docs/solvers.rst index 1de32a71..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: @@ -35,7 +35,7 @@ Set: .. code-block:: toml - electrostatic = false + solver = "electrodynamic_yee" Per step, PyPIC3D: diff --git a/docs/usage.rst b/docs/usage.rst index fadb696a..a84ce052 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 diff --git a/tests/initialization_test.py b/tests/initialization_test.py index a30719a3..0541717d 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -6,7 +6,9 @@ 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_electrostatic +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled +from PyPIC3D.particles.tiled_particles import TiledParticles jax.config.update("jax_enable_x64", True) @@ -27,6 +29,9 @@ def test_default_parameters(self): 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") @@ -35,6 +40,62 @@ def test_default_parameters(self): 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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + 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 = fields + self.assertEqual(E[0].ndim, 6) + self.assertEqual(B[0].ndim, 6) + self.assertEqual(J[0].ndim, 6) + def test_initialize_simulation_encodes_global_particle_boundary_conditions(self): with tempfile.TemporaryDirectory() as tmpdir: zeros_path = os.path.join(tmpdir, "zeros.npy") @@ -43,7 +104,7 @@ def test_initialize_simulation_encodes_global_particle_boundary_conditions(self) "simulation_parameters": { "name": "global particle bc test", "output_dir": tmpdir, - "solver": "fdtd", + "solver": "electrodynamic_yee", "Nx": 1, "Ny": 1, "Nz": 1, @@ -52,7 +113,6 @@ def test_initialize_simulation_encodes_global_particle_boundary_conditions(self) "z_wind": 1.0, "Nt": 1, "dt": 1.0e-10, - "fast_backend": "default", "particle_x_bc": "reflecting", "particle_y_bc": "absorbing", "particle_z_bc": "periodic", @@ -81,7 +141,7 @@ def test_initialize_simulation_encodes_global_particle_boundary_conditions(self) _, particles, _, world, *_ = initialize_simulation(toml.load(config_path)) self.assertEqual(world["particle_boundary_conditions"], {"x": 1, "y": 2, "z": 0}) - self.assertEqual(particles[0].x_bc, "absorbing") + self.assertIsInstance(particles, TiledParticles) def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -91,8 +151,7 @@ def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): "simulation_parameters": { "name": "electrostatic collocated grid test", "output_dir": tmpdir, - "solver": "fdtd", - "electrostatic": True, + "solver": "electrostatic", "Nx": 4, "Ny": 2, "Nz": 1, @@ -101,7 +160,6 @@ def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): "z_wind": 1.0, "Nt": 1, "dt": 1.0e-10, - "fast_backend": "default", }, "plotting": {"plotting": False}, "particle1": { @@ -119,9 +177,11 @@ def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): }, } - loop, _, _, world, *_ = initialize_simulation(config) + loop, particles, fields, world, *_ = initialize_simulation(config) - self.assertIs(loop, time_loop_electrostatic) + self.assertIs(loop, time_loop_electrostatic_tiled) + 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)) @@ -140,7 +200,6 @@ def test_initialize_simulation_rejects_unknown_solver(self): "z_wind": 1.0, "Nt": 1, "dt": 1.0e-10, - "fast_backend": "default", }, "plotting": {"plotting": False}, } @@ -152,5 +211,14 @@ 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/pml_test.py b/tests/pml_test.py index d5a96424..561f4b2e 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -47,12 +47,11 @@ def _base_world(nx=24, ny=1, nz=1): return world -def _empty_config(tmpdir, solver="fdtd", electrostatic=False, pml=None): +def _empty_config(tmpdir, solver="electrodynamic_yee", pml=None): sim = { "name": "pml init test", "output_dir": tmpdir, "solver": solver, - "electrostatic": electrostatic, "Nx": 8, "Ny": 1, "Nz": 1, @@ -61,7 +60,6 @@ def _empty_config(tmpdir, solver="fdtd", electrostatic=False, pml=None): "z_wind": 1.0, "Nt": 1, "dt": 1e-10, - "fast_backend": "default", } config = {"simulation_parameters": sim, "plotting": {"plotting": False}} if pml is not None: @@ -198,14 +196,14 @@ def test_tile_pml_profiles_matches_global_profiles_on_tile_interiors(self): ) ) - def test_initialize_simulation_rejects_pml_for_non_fdtd_solvers(self): + 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, electrostatic=True, pml=pml)) + initialize_simulation(_empty_config(tmpdir, solver="electrostatic", pml=pml)) - def test_initialize_simulation_appends_pml_state_for_fdtd(self): + 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: @@ -221,11 +219,11 @@ def test_initialize_simulation_appends_pml_state_for_fdtd(self): self.assertFalse(pml_y) self.assertFalse(pml_z) - def test_initialize_simulation_uses_tiled_pml_state_for_tiled_yee(self): + 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="tiled_yee", pml=pml) + config = _empty_config(tmpdir, solver="electrodynamic_yee", pml=pml) config["simulation_parameters"].update( { "particle_tile_nx": 2, @@ -248,7 +246,7 @@ def test_initialize_simulation_uses_tiled_pml_state_for_tiled_yee(self): 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_fdtd(self): + 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)) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 4a954ea5..79c627ad 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -406,7 +406,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): constants, unused_curl, J_func=None, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=world["tile_shape"], g=int(world["guard_cells"]), relativistic=False, @@ -472,7 +472,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self constants, unused_curl, J_func=None, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -502,8 +502,7 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): "simulation_parameters": { "name": "tiled electrostatic init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", - "electrostatic": True, + "solver": "electrostatic", "Nx": 8, "Ny": 1, "Nz": 1, @@ -518,7 +517,6 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, diff --git a/tests/tiled_esirkepov_integration_test.py b/tests/tiled_esirkepov_integration_test.py index 3602b159..982721f1 100644 --- a/tests/tiled_esirkepov_integration_test.py +++ b/tests/tiled_esirkepov_integration_test.py @@ -25,7 +25,7 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): "simulation_parameters": { "name": "tiled esirkepov output smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -40,7 +40,6 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): "particle_tile_nz": 1, "current_calculation": "esirkepov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 94416dc1..fa4e5770 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -478,7 +478,7 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): "simulation_parameters": { "name": "tiled yee esirkepov init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -494,7 +494,6 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): "particle_tile_nz": 1, "current_calculation": "esirkepov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -541,7 +540,7 @@ def test_initialize_rejects_filtered_esirkepov_current(self): "simulation_parameters": { "name": "reject filtered esirkepov", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -556,7 +555,6 @@ def test_initialize_rejects_filtered_esirkepov_current(self): "particle_tile_nz": 1, "current_calculation": "esirkepov", "filter_j": "digital", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -595,7 +593,7 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): "simulation_parameters": { "name": "tiled yee esirkepov step smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -610,7 +608,6 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): "particle_tile_nz": 1, "current_calculation": "esirkepov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -674,7 +671,7 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s "simulation_parameters": { "name": "tiled yee esirkepov alias staging", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -689,7 +686,6 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s "particle_tile_nz": 1, "current_calculation": "esirkepov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 4d8e8a92..9fbfb257 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -137,7 +137,9 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): tile_shape=world["tile_shape"], g=1, ) - mass_from_tiles = assemble_tiled_scalar_field(mass_tiles, world, world["tile_shape"]) + mass_from_tiles = assemble_tiled_scalar_field( + mass_tiles, world, world["tile_shape"], num_guard_cells=1 + ) self.assertTrue( jnp.allclose( @@ -148,10 +150,10 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): ) ) - def test_tiled_scalar_vtk_path_no_longer_rejects_tiled_yee(self): + def test_tiled_scalar_vtk_path_no_longer_rejects_electrodynamic_yee(self): run_source = inspect.getsource(pypic_main.run_PyPIC3D) - self.assertNotIn("plot_vtk_scalars is not supported for tiled_yee", run_source) + self.assertNotIn("plot_vtk_scalars is not supported for electrodynamic_yee", run_source) if __name__ == "__main__": diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index d473e1f3..1e4ec259 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -293,7 +293,7 @@ def _assert_tiled_state_matches_standard( ) def test_validate_field_solver_accepts_tiled_yee(self): - validate_field_solver("tiled_yee") + validate_field_solver("electrodynamic_yee") def test_tiled_yee_field_update_applies_standard_field_filter(self): world = self._build_world() @@ -349,7 +349,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): constants, unused_curl, J_func=lambda particles, J, constants, world, tile_shape=None, g=None: J, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -384,7 +384,7 @@ def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(sel "simulation_parameters": { "name": "tiled yee init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -399,7 +399,6 @@ def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(sel "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -443,7 +442,7 @@ def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(self): "simulation_parameters": { "name": "tiled yee higuera cary init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -458,7 +457,6 @@ def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(self): "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "higuera_cary", "relativistic": True, }, @@ -502,7 +500,7 @@ def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(self): "simulation_parameters": { "name": "tiled yee digital filter init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -517,7 +515,6 @@ def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(self): "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "digital", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -561,7 +558,7 @@ def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(self): "simulation_parameters": { "name": "tiled yee bilinear filter init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -576,7 +573,6 @@ def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(self): "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "bilinear", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -620,7 +616,7 @@ def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(sel "simulation_parameters": { "name": "tiled yee conducting init smoke", "output_dir": tmpdir, - "solver": "tiled_yee", + "solver": "electrodynamic_yee", "Nx": 8, "Ny": 1, "Nz": 1, @@ -638,7 +634,6 @@ def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(sel "particle_tile_nz": 1, "current_calculation": "j_from_rhov", "filter_j": "none", - "fast_backend": "default", "particle_pusher": "boris", "relativistic": False, }, @@ -716,7 +711,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): constants, unused_curl, J_func=None, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -796,7 +791,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): constants, unused_curl, J_func=None, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -860,7 +855,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): constants, unused_curl, J_func=functools.partial(direct_J_from_tiled_particles, filter="digital"), - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -921,7 +916,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): constants, unused_curl, J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -1024,7 +1019,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand constants, unused_curl, J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -1082,7 +1077,7 @@ def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_s constants, unused_curl, J_func=None, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, @@ -1139,7 +1134,7 @@ def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standa constants, unused_curl, J_func=tiled_esirkepov_current, - solver="tiled_yee", + solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, From cf761c65640d40b16d646088e452d8d06dd9c546 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 13:15:34 -0500 Subject: [PATCH 29/83] migrating codebase to tiled fields and particles for sharding/parallelization --- PyPIC3D/__main__.py | 23 +- PyPIC3D/deposition/Esirkepov.py | 45 ++- PyPIC3D/deposition/J_from_rhov.py | 48 ++- PyPIC3D/deposition/rho.py | 26 +- PyPIC3D/deposition/rho_tiled.py | 502 +++++++++++++++++++----- PyPIC3D/diagnostics/fluid_quantities.py | 76 +++- PyPIC3D/diagnostics/openPMD.py | 38 +- PyPIC3D/diagnostics/plotting.py | 14 +- PyPIC3D/diagnostics/vtk.py | 8 +- PyPIC3D/evolve.py | 176 +++------ PyPIC3D/initialization.py | 25 +- PyPIC3D/solvers/first_order_yee.py | 43 +- tests/direct_deposition_tiled_test.py | 56 +++ tests/evolve_test.py | 24 +- tests/initialization_test.py | 7 +- tests/openpmd_test.py | 115 ++++++ tests/tiled_electrostatic_test.py | 6 +- tests/tiled_esirkepov_test.py | 69 +++- tests/tiled_fluid_quantities_test.py | 151 ++++++- tests/tiled_rho_test.py | 41 ++ tests/tiled_yee_integration_test.py | 18 +- tests/utils_test.py | 170 ++++++++ tests/yee_tiled_test.py | 38 ++ 23 files changed, 1411 insertions(+), 308 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 38c0ed96..0e638a39 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -41,11 +41,6 @@ ) from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.deposition.rho_tiled import ( - compute_rho_from_tiled_particles, - compute_tiled_mass_density_from_tiled_particles, - compute_tiled_rho_from_tiled_particles, -) from PyPIC3D.diagnostics.output_adapters import fields_for_output, scalar_field_for_output @@ -75,7 +70,6 @@ def run_PyPIC3D(config_file): scalar_field_names = ["rho", "mass_density"] vector_field_names = ["E", "B", "J"] - tiled_run = True tile_shape = tuple(int(width) for width in world["tile_shape"]) g = int(world["guard_cells"]) particle_species_names = simulation_parameters.get("particle_species_names") @@ -137,19 +131,12 @@ def run_PyPIC3D(config_file): if plotting_parameters['plot_vtk_scalars']: if getattr(rho, "ndim", 0) == 6: - rho = compute_tiled_rho_from_tiled_particles( - particles, species_config, rho, world, constants, tile_shape=tile_shape, g=g - ) - mass_density = compute_tiled_mass_density_from_tiled_particles( - particles, species_config, rho, world, tile_shape=tile_shape, g=g - ) + rho = compute_rho(particles, rho, world, constants, species_config=species_config) + mass_density = compute_mass_density(particles, rho, world, species_config=species_config) rho_output = scalar_field_for_output(rho, world) mass_density_output = scalar_field_for_output(mass_density, world) else: - if tiled_run: - rho = compute_rho_from_tiled_particles(particles, species_config, rho, world, constants) - else: - rho = compute_rho(particles, rho, world, constants) + rho = compute_rho(particles, rho, world, constants, species_config=species_config) # calculate the charge density based on the particle positions mass_density = compute_mass_density(particles, rho, world, species_config=species_config) # calculate the mass density based on the particle positions @@ -187,10 +174,6 @@ def run_PyPIC3D(config_file): write_openpmd_fields(fields, world, os.path.join(output_dir, "data"), plot_num, t, "fields", ".h5") # Write the fields in openPMD format - if not tiled_run: - fields = (E, B, J, rho, phi, external_fields, *rest) - # repack the fields for non-tiled diagnostics that updated rho - particles, fields = jit_loop( particles, species_config, diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 45b53760..4a3c9d6a 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -12,6 +12,7 @@ ) from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.particles.tiled_particles import TiledParticles def shift_old_stencil(old_w_list, shift): @@ -206,8 +207,50 @@ def eliminate_esirkepov_ghost_cells(field): # 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): +def Esirkepov_current( + particles, + J, + constants, + world, + grid=None, + filter=None, + species_config=None, + tile_shape=None, + g=None, +): """Esirkepov current deposition supporting 1D/2D/3D via inactive dims.""" + if isinstance(particles, TiledParticles): + if filter not in (None, "none"): + raise ValueError("Esirkepov current filtering is not supported; use filter='none'.") + + if species_config is None: + species_config = J + J = constants + constants = world + world = grid + grid = None + # Tiled Esirkepov stores old particle positions. The tiled kernel + # predicts the new positions locally and leaves particle retile staging + # to the caller, matching the Task 0 old/new position contract. + + if tile_shape is None: + tile_shape = tuple(int(width) for width in world["tile_shape"]) + if g is None: + g = int(world["guard_cells"]) + + from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current + + return tiled_esirkepov_current( + particles, + species_config, + J, + constants, + world, + grid=grid, + tile_shape=tile_shape, + g=int(g), + ) + if grid is None: grid = world["grids"]["center"] diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index fa80636f..b2c21c01 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -9,12 +9,58 @@ prepare_particle_axis_stencil, ) from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells +from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.utils import bilinear_filter, digital_filter +def J_from_rhov( + particles, + J, + constants, + world, + grid=None, + filter="bilinear", + species_config=None, + tile_shape=None, + g=None, +): + """Compute current density from velocity-weighted particle charge.""" + + if isinstance(particles, TiledParticles): + if species_config is None: + species_config = J + J = constants + constants = world + world = grid + grid = None + # Tiled runtime callers pass ``species_config`` as the second positional + # argument so the public current name can be used without global current + # assembly in the deposition path. + + if tile_shape is None: + tile_shape = tuple(int(width) for width in world["tile_shape"]) + if g is None: + g = int(world["guard_cells"]) + + return direct_J_from_tiled_particles( + particles, + species_config, + J, + constants, + world, + grid=grid, + filter=filter, + tile_shape=tile_shape, + g=int(g), + ) + + return _J_from_rhov_flat(particles, J, constants, world, grid=grid, filter=filter) + + @partial(jit, static_argnames=("filter",)) -def J_from_rhov(particles, J, constants, world, grid=None, filter="bilinear"): +def _J_from_rhov_flat(particles, J, constants, world, grid=None, filter="bilinear"): """Compute current density (Jx,Jy,Jz) by depositing particle velocities.""" if grid is None: diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 13e5eed2..204019d8 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -8,11 +8,35 @@ ) 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.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles +from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.utils import digital_filter +def compute_rho(particles, rho, world, constants, species_config=None): + """Compute charge density for flat or tile-major particle storage.""" + if isinstance(particles, TiledParticles): + if species_config is None: + species_config = world["species_config"] + + if getattr(rho, "ndim", 0) == 6: + return compute_tiled_rho_from_tiled_particles( + particles, + species_config, + rho, + world, + constants, + tile_shape=tuple(int(width) for width in world["tile_shape"]), + g=int(world["guard_cells"]), + ) + + return compute_rho_from_tiled_particles(particles, species_config, rho, world, constants) + + return _compute_rho_flat(particles, rho, world, constants) + + @jit -def compute_rho(particles, rho, world, constants): +def _compute_rho_flat(particles, rho, world, constants): """Compute the charge density (rho) on the vertex grid.""" dx = world["dx"] dy = world["dy"] diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index c5c3e1ac..67f27107 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -12,7 +12,6 @@ from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.solvers.yee_tiled import ( digital_filter_tiled_scalar, - fold_tiled_ghost_cells, tiled_grid_axes_from_world, update_tiled_ghost_cells, ) @@ -26,6 +25,12 @@ def _species_scalar_to_slots(tiled_particles, species_value): ) +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 _deposit_tiled_scalar_moment( tiled_particles, rho, @@ -33,6 +38,7 @@ def _deposit_tiled_scalar_moment( grid, scalar_weight, position=None, + divide_by_volume=True, ): dx = world["dx"] dy = world["dy"] @@ -53,7 +59,9 @@ def _deposit_tiled_scalar_moment( y = y.reshape(-1) z = z.reshape(-1) active = tiled_particles.active.reshape(-1).astype(x.dtype) - scalar_density = scalar_weight.reshape(-1) / (dx * dy * dz) + 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, @@ -160,14 +168,6 @@ def _diagnostic_mass_position(tiled_particles, world): return x, y, z -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) - - def _deposit_tiled_scalar_moment_to_tiles( tiled_particles, rho_tiles, @@ -178,6 +178,7 @@ def _deposit_tiled_scalar_moment_to_tiles( g, tiled_grid, position=None, + divide_by_volume=True, ): dx = world["dx"] dy = world["dy"] @@ -189,109 +190,110 @@ def _deposit_tiled_scalar_moment_to_tiles( tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] g = int(g) - local_Nx = tile_nx + 2 * g - local_Ny = tile_ny + 2 * g - local_Nz = tile_nz + 2 * g - ntx, nty, ntz = rho_tiles.shape[:3] - reduced_x = int(ntx) == 1 and tile_nx == 1 - reduced_y = int(nty) == 1 and tile_ny == 1 - reduced_z = int(ntz) == 1 and tile_nz == 1 - - def deposit_one_tile(tx, ty, tz, x_tile, u_tile, active_tile, scalar_weight_tile): - local_grid = ( - tiled_grid[0][tx, ty, tz], - tiled_grid[1][tx, ty, tz], - tiled_grid[2][tx, ty, tz], - ) + 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 = x_tile[..., 0].reshape(-1) - y = x_tile[..., 1].reshape(-1) - z = x_tile[..., 2].reshape(-1) - else: - x = position[0][tx, ty, tz].reshape(-1) - y = position[1][tx, ty, tz].reshape(-1) - z = position[2][tx, ty, tz].reshape(-1) - active = active_tile.reshape(-1).astype(x.dtype) - scalar_density = scalar_weight_tile.reshape(-1) / (dx * dy * dz) - - x, _, deltax, xpts = prepare_particle_axis_stencil( - x, - local_grid[0], - local_Nx, - shape_factor, - bc_x, - wind=world["x_wind"], - ghost_cells=True, - ) - y, _, deltay, ypts = prepare_particle_axis_stencil( - y, - local_grid[1], - local_Ny, - shape_factor, - bc_y, - wind=world["y_wind"], - ghost_cells=True, - ) - z, _, deltaz, zpts = prepare_particle_axis_stencil( - z, - local_grid[2], - local_Nz, - shape_factor, - bc_z, - wind=world["z_wind"], - ghost_cells=True, - ) + 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_weights, y_weights, z_weights = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(deltax, deltay, deltaz, dx, dy, dz), - lambda _: get_second_order_weights(deltax, deltay, deltaz, dx, dy, dz), - operand=None, - ) + x, _, 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, + ) - 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_tiled_axis_stencil(xpts, x_weights, local_Nx, reduced_x, g) - ypts, y_weights = _collapse_tiled_axis_stencil(ypts, y_weights, local_Ny, reduced_y, g) - zpts, z_weights = _collapse_tiled_axis_stencil(zpts, z_weights, local_Nz, reduced_z, g) - - rho_tile = jnp.zeros((local_Nx, local_Ny, local_Nz), dtype=rho_tiles.dtype) - - for i in range(xpts.shape[0]): - for j in range(ypts.shape[0]): - for k in range(zpts.shape[0]): - rho_tile = rho_tile.at[xpts[i], ypts[j], zpts[k]].add( - active * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], - mode="drop", - ) - - return rho_tile - - tx = jnp.arange(ntx) - ty = jnp.arange(nty) - tz = jnp.arange(ntz) - deposit_tiles = deposit_one_tile - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(None, None, 0, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(None, 0, None, 0, 0, 0, 0), out_axes=0) - deposit_tiles = jax.vmap(deposit_tiles, in_axes=(0, None, None, 0, 0, 0, 0), out_axes=0) - - rho_tiles = deposit_tiles( - tx, - ty, - tz, - tiled_particles.x, - tiled_particles.u, - tiled_particles.active, - scalar_weight, + 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, ) - rho_tiles = fold_tiled_ghost_cells(rho_tiles, world, g, tile_shape) + 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, + ) + 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 + 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 @@ -418,3 +420,287 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, species_con tiled_grid, position=position, ) + + +@partial(jit, static_argnames=("direction",)) +def compute_velocity_field_from_tiled_particles(tiled_particles, species_config, field, direction, world, grid=None): + """Compute the legacy scalar velocity diagnostic from tile-major particles.""" + 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( + tiled_particles, + field, + world, + grid, + scalar_weight, + position=position, + divide_by_volume=False, + ) + + +@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"] + tiled_grid = tiled_grid_axes_from_world( + world, + grid, + "tiled_vertex_grid", + tile_shape, + g, + ) + + 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, + tiled_grid, + position=position, + divide_by_volume=False, + ) + + +@partial(jit, static_argnames=("direction",)) +def compute_pressure_field_from_tiled_particles( + tiled_particles, + species_config, + field, + velocity_field, + direction, + world, + grid=None, +): + """Compute the legacy scalar pressure diagnostic from tile-major particles.""" + if grid is None: + grid = world["grids"]["vertex"] + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + Nx, Ny, Nz = field.shape + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] + + x, y, z = _diagnostic_mass_position(tiled_particles, world) + x = x.reshape(-1) + y = y.reshape(-1) + z = z.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, + 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, + ) + z, _, 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, + ) + + 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) + + field = jnp.zeros_like(field) + + for i in range(xpts.shape[0]): + for j in range(ypts.shape[0]): + for k in range(zpts.shape[0]): + 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", + ) + + field = fold_ghost_cells(field, bc_x, bc_y, bc_z) + field = update_ghost_cells(field, bc_x, bc_y, bc_z) + return field + + +@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") + + field_tiles = update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + return field_tiles diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index 06d48fc3..e3fafea8 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -2,15 +2,22 @@ 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 -from PyPIC3D.deposition.rho_tiled import compute_mass_density_from_tiled_particles +from PyPIC3D.deposition.rho_tiled import ( + compute_mass_density_from_tiled_particles, + compute_pressure_field_from_tiled_particles, + compute_tiled_mass_density_from_tiled_particles, + compute_tiled_pressure_field_from_tiled_particles, + compute_tiled_velocity_field_from_tiled_particles, + compute_velocity_field_from_tiled_particles, +) from PyPIC3D.particles.tiled_particles import TiledParticles import jax import jax.numpy as jnp from jax import jit +from functools import partial -@jit 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. @@ -29,8 +36,22 @@ def compute_mass_density(particles, rho, world, species_config=None): ndarray: The updated charge density array. """ 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"]), + ) return compute_mass_density_from_tiled_particles(particles, species_config, rho, world) + return _compute_mass_density_flat(particles, rho, world) + + +@jit +def _compute_mass_density_flat(particles, rho, world): dx = world['dx'] dy = world['dy'] dz = world['dz'] @@ -94,8 +115,26 @@ def compute_mass_density(particles, rho, world, species_config=None): rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) return rho -@jit -def compute_velocity_field(particles, field, direction, world): + +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"]), + ) + return compute_velocity_field_from_tiled_particles(particles, species_config, field, int(direction), world) + + return _compute_velocity_field_flat(particles, field, direction, world) + + +@partial(jit, static_argnames=("direction",)) +def _compute_velocity_field_flat(particles, field, direction, world): """ Compute the velocity field (v) for a given set of particles in a simulation world. Parameters: @@ -180,8 +219,33 @@ def compute_velocity_field(particles, field, direction, world): -@jit -def compute_pressure_field(particles, field, velocity_field, direction, world): +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"]), + ) + return compute_pressure_field_from_tiled_particles( + particles, + species_config, + field, + velocity_field, + int(direction), + world, + ) + + return _compute_pressure_field_flat(particles, field, velocity_field, direction, world) + + +@partial(jit, static_argnames=("direction",)) +def _compute_pressure_field_flat(particles, field, velocity_field, direction, world): dx = world['dx'] dy = world['dy'] diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index aeb57c9d..1548ec57 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -286,7 +286,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. @@ -297,6 +297,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 @@ -341,8 +345,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.get_mass() + particle_charge = species.get_charge() + particle_weight = species.get_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)): @@ -364,25 +387,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() diff --git a/PyPIC3D/diagnostics/plotting.py b/PyPIC3D/diagnostics/plotting.py index cae3d2d7..6a37cd43 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -9,7 +9,7 @@ from PyPIC3D.diagnostics.output_adapters import particles_for_output -def plot_positions(particles, t, x_wind, y_wind, z_wind, path): +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. @@ -25,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() 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.get_name() )) fig.update_layout( @@ -89,7 +93,7 @@ def write_particles_phase_space(particles, t, path, species_config=None, species 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. @@ -109,6 +113,10 @@ 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() diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py index ba8e378c..b4ec4385 100644 --- a/PyPIC3D/diagnostics/vtk.py +++ b/PyPIC3D/diagnostics/vtk.py @@ -5,9 +5,9 @@ from pyevtk.hl import gridToVTK, pointsToVTK import jax.numpy as jnp -from PyPIC3D.diagnostics.output_adapters import particles_for_output +from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output -def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): +def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz, world=None): """ Plot the fields in a 3D grid. @@ -24,6 +24,10 @@ def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz): Returns: None """ + if world is not None: + fieldx, fieldy, fieldz = vector_field_for_output((fieldx, fieldy, fieldz), world) + # Tiled field storage is assembled only here, at the file-output boundary. + Nx = fieldx.shape[0] Ny = fieldx.shape[1] Nz = fieldx.shape[2] diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index eeb2f871..1d603511 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -1,158 +1,104 @@ # 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 jax import jit -from PyPIC3D.solvers.first_order_yee import ( - update_E, update_B -) +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled as time_loop_electrodynamic +from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled as time_loop_electrostatic +from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.solvers.electrostatic_yee import calculate_electrostatic_fields +from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.utils import add_external_fields -from PyPIC3D.solvers.electrostatic_yee import ( - calculate_electrostatic_fields -) -from PyPIC3D.utils import add_external_fields +__all__ = ["time_loop_electrodynamic", "time_loop_electrostatic"] + @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"): +def _time_loop_electrostatic_global_reference( + 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). + Old global electrostatic loop retained as a reference path for tests. """ E, B, J, rho, phi, external_fields = fields - # unpack the fields - center_grid = world['grids']['center'] - vertex_grid = world['grids']['vertex'] + 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, constants, relativistic=relativistic, particle_pusher=particle_pusher) - # use the selected particle pusher for particle velocities - + particles[i] = particle_push( + particles[i], + push_E, + push_B, + center_grid, + vertex_grid, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) particles[i].update_position() - # update the particle positions - particles[i].boundary_conditions(world) - # electrostatic rho is deposited from the current in-domain 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 + E, phi, rho = calculate_electrostatic_fields(world, particles, constants, rho, phi, solver, "periodic") fields = (E, B, J, rho, phi, external_fields) - # pack the fields into a tuple - return particles, fields @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_global_reference( + particles, + fields, + world, + constants, + curl_func, + J_func, + solver, + relativistic=True, + particle_pusher="boris", +): """ - Advance an electrodynamic Particle-In-Cell (PIC) system by one time step. - This routine performs, in order: - 1) Particle push using the 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. + Old global electrodynamic loop retained as a reference path for tests. """ - E, B, J, rho, phi, external_fields, pml_state = fields - # unpack the fields - center_grid = world['grids']['center'] - vertex_grid = world['grids']['vertex'] + 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, constants, relativistic=relativistic, particle_pusher=particle_pusher) - # use the selected particle pusher for particle velocities - + particles[i] = particle_push( + particles[i], + push_E, + push_B, + center_grid, + vertex_grid, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) 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(world) - # apply boundary conditions to the particles fields = (E, B, J, rho, phi, external_fields, pml_state) - # pack the fields into a tuple - - return particles, fields diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index c5fd24fa..ff59b3f4 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -47,12 +47,9 @@ ) -from PyPIC3D.electrodynamic_tiled import ( - time_loop_electrodynamic_tiled -) - -from PyPIC3D.electrostatic_tiled import ( - time_loop_electrostatic_tiled +from PyPIC3D.evolve import ( + time_loop_electrodynamic, + time_loop_electrostatic, ) from PyPIC3D.pusher.particle_push import validate_particle_pusher @@ -60,8 +57,8 @@ from PyPIC3D.deposition.current_methods import ( encode_current_calculation, ) -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles -from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current +from PyPIC3D.deposition.Esirkepov import Esirkepov_current +from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -489,18 +486,18 @@ def initialize_simulation(toml_file): if electrostatic: print("Using electrostatic solver") - evolve_loop = time_loop_electrostatic_tiled + evolve_loop = time_loop_electrostatic else: print("Using electrodynamic Yee solver") - evolve_loop = time_loop_electrodynamic_tiled + 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") - J_func = tiled_esirkepov_current + 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(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) + J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) species_config = None @@ -524,10 +521,10 @@ def initialize_simulation(toml_file): B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) if simulation_parameters["current_calculation"] == "esirkepov": J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) - J_func = tiled_esirkepov_current + J_func = Esirkepov_current else: J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) - J_func = functools.partial(direct_J_from_tiled_particles, filter=simulation_parameters['filter_j']) + J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) external_E, external_B = external_fields external_fields = ( tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py index 7bc10c6c..3729c453 100644 --- a/PyPIC3D/solvers/first_order_yee.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -10,11 +10,12 @@ apply_pml_to_e_curl, update_ghost_cells_for_pml, ) +from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E # import internal libraries @partial(jit, static_argnames=("curl_func",)) -def update_E(E, B, J, world, constants, curl_func, pml_state=None): +def _update_E_global(E, B, J, world, constants, curl_func, pml_state=None): """ Update the electric field components (Ex, Ey, Ez) based on the given parameters. @@ -120,7 +121,7 @@ def update_E(E, B, J, world, constants, curl_func, pml_state=None): @partial(jit, static_argnames=("curl_func",)) -def update_B(E, B, world, constants, curl_func, pml_state=None): +def _update_B_global(E, B, world, constants, curl_func, pml_state=None): """ Update the magnetic field components (Bx, By, Bz) using the curl of the electric field. @@ -208,3 +209,41 @@ def update_B(E, B, world, constants, curl_func, pml_state=None): # fill ghost cells with the updated values return (Bx, By, Bz), pml_state + + +def update_E(E, B, J, world, constants, curl_func, pml_state=None): + """ + Update electric fields through the public Yee helper. + + Tiled field/current arrays are routed to the tile-local Yee update using + the explicit ``world['tile_shape']`` and ``world['guard_cells']`` contract. + Ordinary global ghost-celled arrays keep the historical first-order Yee + update in ``_update_E_global``. + """ + + if E[0].ndim == 6: + tile_shape = tuple(world["tile_shape"]) + g = int(world["guard_cells"]) + if pml_state is None: + return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g), None + return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g, pml_state) + + return _update_E_global(E, B, J, world, constants, curl_func, pml_state) + + +def update_B(E, B, world, constants, curl_func, pml_state=None): + """ + Update magnetic fields through the public Yee helper. + + Tiled arrays follow the tile-local Yee path and ordinary global arrays + follow the original global ghost-cell implementation. + """ + + if E[0].ndim == 6: + tile_shape = tuple(world["tile_shape"]) + g = int(world["guard_cells"]) + if pml_state is None: + return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g), None + return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g, pml_state) + + return _update_B_global(E, B, world, constants, curl_func, pml_state) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 56fc7172..d623cee1 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -493,6 +493,62 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_dummy_species(self): self._compare_tiled_to_standard([electrons, ions], world, simulation_parameters) + def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(self): + world = self._build_world() + world["guard_cells"] = 2 + world["tile_shape"] = (2, 2, 2) + simulation_parameters = { + "particle_tile_nx": 2, + "particle_tile_ny": 2, + "particle_tile_nz": 2, + } + constants = {"C": 3.0e8, "alpha": 0.6} + species = particle_species( + name="public tiled direct current", + N_particles=4, + charge=-1.0, + mass=1.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"], + ) + tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) + + J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") + J_tiles = J_from_rhov( + tiled_particles, + species_config, + self._empty_J_tiles(world, simulation_parameters), + constants, + world, + filter="digital", + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + J_from_tiles = assemble_tiled_vector_field( + J_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + 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)) + def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): world = self._build_world() simulation_parameters = { diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 031a4820..38917fef 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -3,7 +3,14 @@ import jax import jax.numpy as jnp -from PyPIC3D.evolve import time_loop_electrodynamic, time_loop_electrostatic +from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled +from PyPIC3D.evolve import ( + _time_loop_electrodynamic_global_reference, + _time_loop_electrostatic_global_reference, + time_loop_electrodynamic, + time_loop_electrostatic, +) from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species @@ -22,6 +29,13 @@ def unused_curl(Ex, Ey, Ez): class TestEvolveExternalFields(unittest.TestCase): + def test_public_loop_names_use_tiled_contracts(self): + electrodynamic_loop = time_loop_electrodynamic.func if hasattr(time_loop_electrodynamic, "func") else time_loop_electrodynamic + electrostatic_loop = time_loop_electrostatic.func if hasattr(time_loop_electrostatic, "func") else time_loop_electrostatic + + self.assertIs(electrodynamic_loop, time_loop_electrodynamic_tiled) + self.assertIs(electrostatic_loop, time_loop_electrostatic_tiled) + def _electrostatic_grid_order_species(self, world, charge): return particle_species( name=f"charge {charge}", @@ -82,7 +96,7 @@ def test_electrostatic_push_uses_same_grid_order_as_electrodynamic(self): self._electrostatic_grid_order_species(world, -1.0), ] - pushed_particles, _ = time_loop_electrostatic( + pushed_particles, _ = _time_loop_electrostatic_global_reference( pushed_particles, (E, B, J, rho, phi, external_fields), world, @@ -186,7 +200,7 @@ def test_electrostatic_rho_uses_boundary_applied_particle_positions(self): ), ] - particles, fields = time_loop_electrostatic( + particles, fields = _time_loop_electrostatic_global_reference( particles, (E, B, J, rho, phi, external_fields), world, @@ -259,7 +273,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field ] fields = (E, B, J, rho, phi, external_fields, None) - particles, fields = time_loop_electrodynamic( + particles, fields = _time_loop_electrodynamic_global_reference( particles, fields, world, @@ -333,7 +347,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): ] fields = (E, B, J, rho, phi, external_fields, None) - particles, _ = time_loop_electrodynamic( + particles, _ = _time_loop_electrodynamic_global_reference( particles, fields, world, diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 0541717d..c538f6e1 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -6,8 +6,7 @@ import jax import jax.numpy as jnp from PyPIC3D.initialization import setup_write_dir, default_parameters, initialize_simulation, validate_field_solver -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled -from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled +from PyPIC3D.evolve import time_loop_electrodynamic, time_loop_electrostatic from PyPIC3D.particles.tiled_particles import TiledParticles jax.config.update("jax_enable_x64", True) @@ -85,7 +84,7 @@ def test_initialize_simulation_returns_tiled_runtime_for_ordinary_electrodynamic loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(simulation_parameters["solver"], "electrodynamic_yee") self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) @@ -179,7 +178,7 @@ def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): loop, particles, fields, world, *_ = initialize_simulation(config) - self.assertIs(loop, time_loop_electrostatic_tiled) + 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"]): diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index fbc70d7e..9e679996 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -6,18 +6,22 @@ from PyPIC3D.diagnostics import openPMD from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array from PyPIC3D.solvers.yee_tiled import tile_scalar_field, tile_vector_field +from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles class FakeRecord: def __init__(self): self.shape = None self.unit_SI = None + self.data = None def reset_dataset(self, dataset): pass def store_chunk(self, array, offset, extent): self.shape = tuple(extent) + self.data = jnp.asarray(array) class FakeMesh: @@ -40,9 +44,34 @@ def __getitem__(self, name): 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 @@ -59,6 +88,9 @@ class FakeSeries: def __init__(self): self.iterations = FakeIterations() + def set_attribute(self, name, value): + pass + def flush(self): pass @@ -70,6 +102,59 @@ 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 _simulation_parameters(): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + +def _species(name, charge, mass, weight, x1): + world = _world() + return particle_species( + name=name, + N_particles=x1.shape[0], + charge=charge, + mass=mass, + weight=weight, + T=0.0, + x1=x1, + x2=jnp.zeros_like(x1), + x3=jnp.zeros_like(x1), + v1=jnp.ones_like(x1) * 0.1, + v2=jnp.zeros_like(x1), + v3=jnp.zeros_like(x1), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + active_mask=jnp.ones(x1.shape[0], dtype=bool), + dt=world["dt"], + ) + + class OpenPMDDiagnosticsTests(unittest.TestCase): def test_openpmd_field_array_preserves_thin_y_axis(self): @@ -153,6 +238,36 @@ def test_write_openpmd_fields_assembles_tiled_fields_before_output(self): 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_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 = to_tiled_particles(species, world, _simulation_parameters()) + species_names = tuple(s.get_name() for s in species) + 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/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 79c627ad..f18338bb 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -11,7 +11,7 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled -from PyPIC3D.evolve import time_loop_electrostatic +from PyPIC3D.evolve import _time_loop_electrostatic_global_reference from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -372,7 +372,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} E, B, J, rho, phi, external_fields = self._empty_fields(world) - reference_particles, reference_fields = time_loop_electrostatic( + reference_particles, reference_fields = _time_loop_electrostatic_global_reference( self._neutral_species(world), (E, B, J, rho, phi, external_fields), world, @@ -454,7 +454,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self ) for step in range(1, Nt + 1): - reference_particles, reference_fields = time_loop_electrostatic( + reference_particles, reference_fields = _time_loop_electrostatic_global_reference( reference_particles, reference_fields, world, diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index fa4e5770..3be8f8ba 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -16,9 +16,10 @@ _compact_2d_esirkepov_weights, tiled_esirkepov_current, ) +from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.diagnostics.output_adapters import fields_for_output -from PyPIC3D.initialization import CURRENT_ESIRKEPOV, initialize_simulation +from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.tiled_particle_refresh import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, @@ -341,6 +342,72 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): 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["tile_shape"] = (2, 1, 1) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + simulation_parameters = { + "particle_tile_nx": world["tile_shape"][0], + "particle_tile_ny": world["tile_shape"][1], + "particle_tile_nz": world["tile_shape"][2], + } + x_old = jnp.array([-1.10, -0.10, 1.05]) + old_species = self._species(world, x_old) + x_new = x_old + old_species.v1 * world["dt"] + new_species = self._species(world, x_new) + + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) + J_tiles = Esirkepov_current( + tiled_particles, + species_config, + empty_tiled_vector_field(world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])), + constants, + world, + tile_shape=world["tile_shape"], + g=int(world["guard_cells"]), + ) + J_from_tiles = assemble_tiled_vector_field( + J_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + 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) + constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + old_species = self._species(world, jnp.array([-1.10, -0.10, 1.05])) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + + with self.assertRaisesRegex(ValueError, "Esirkepov current filtering is not supported"): + Esirkepov_current( + tiled_particles, + species_config, + empty_tiled_vector_field(world, tile_shape, num_guard_cells=int(world["guard_cells"])), + constants, + world, + filter="digital", + tile_shape=tile_shape, + g=int(world["guard_cells"]), + ) + def test_tiled_esirkepov_matches_global_current_for_dimensions_and_shapes(self): cases = [ (8, 1, 1), diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 9fbfb257..6f51ee84 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -6,7 +6,11 @@ from PyPIC3D import __main__ as pypic_main from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.diagnostics.fluid_quantities import compute_mass_density +from PyPIC3D.diagnostics.fluid_quantities import ( + compute_mass_density, + compute_pressure_field, + compute_velocity_field, +) from PyPIC3D.deposition.rho_tiled import compute_tiled_mass_density_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -32,6 +36,7 @@ def _build_world(self, shape_factor=2): "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) @@ -126,7 +131,12 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): ) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) - rho_tiles = tile_scalar_field(self._empty_scalar(world), world, world["tile_shape"]) + rho_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) mass_tiles = compute_tiled_mass_density_from_tiled_particles( @@ -135,10 +145,13 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): rho_tiles, world, tile_shape=world["tile_shape"], - g=1, + g=int(world["guard_cells"]), ) mass_from_tiles = assemble_tiled_scalar_field( - mass_tiles, world, world["tile_shape"], num_guard_cells=1 + mass_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), ) self.assertTrue( @@ -150,6 +163,136 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): ) ) + def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(self): + world = self._build_world(shape_factor=2) + world["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + particles = self._particles(world) + tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + rho_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + mass_reference = compute_mass_density(particles, self._empty_scalar(world), 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["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + particles = self._particles(world) + tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + field_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + velocity_reference = compute_velocity_field(particles, self._empty_scalar(world), 0, world) + velocity_tiles = compute_velocity_field( + tiled_particles, + field_tiles, + 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["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + particles = self._particles(world) + tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + field_tiles = tile_scalar_field( + self._empty_scalar(world), + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + velocity_reference = compute_velocity_field(particles, self._empty_scalar(world), 0, world) + pressure_reference = compute_pressure_field( + particles, + self._empty_scalar(world), + velocity_reference, + 0, + 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_tiled_scalar_vtk_path_no_longer_rejects_electrodynamic_yee(self): run_source = inspect.getsource(pypic_main.run_PyPIC3D) diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 903f76ec..7332f656 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -239,6 +239,47 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self ) ) + def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(self): + world = self._build_world(shape_factor=2) + world["tile_shape"] = ( + self._simulation_parameters()["particle_tile_nx"], + self._simulation_parameters()["particle_tile_ny"], + self._simulation_parameters()["particle_tile_nz"], + ) + 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_reference = compute_rho(particles, self._empty_scalar(world), world, constants) + rho_tiles = compute_rho( + tiled_particles, + rho_tiles, + world, + constants, + species_config=species_config, + ) + rho_from_tiles = assemble_tiled_scalar_field( + rho_tiles, + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + + 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, + ) + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 1e4ec259..d612c1ca 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -8,7 +8,7 @@ import numpy as np import toml -from PyPIC3D.evolve import time_loop_electrodynamic +from PyPIC3D.evolve import _time_loop_electrodynamic_global_reference from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV @@ -317,7 +317,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): phi = seed external_fields = ((seed, seed, seed), (seed, seed, seed)) - _, reference_fields = time_loop_electrodynamic( + _, reference_fields = _time_loop_electrodynamic_global_reference( [], (E, B, J, rho, phi, external_fields, None), world, @@ -678,7 +678,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, @@ -758,7 +758,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - _, reference_fields = time_loop_electrodynamic( + _, reference_fields = _time_loop_electrodynamic_global_reference( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, @@ -822,7 +822,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - _, reference_fields = time_loop_electrodynamic( + _, reference_fields = _time_loop_electrodynamic_global_reference( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, @@ -883,7 +883,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - _, reference_fields = time_loop_electrodynamic( + _, reference_fields = _time_loop_electrodynamic_global_reference( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, @@ -986,7 +986,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - _, reference_fields = time_loop_electrodynamic( + _, reference_fields = _time_loop_electrodynamic_global_reference( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, @@ -1059,7 +1059,7 @@ def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_s ) for t in range(Nt): - reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( reference_particles, reference_fields, world, @@ -1116,7 +1116,7 @@ def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standa ) for t in range(Nt): - reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( reference_particles, reference_fields, world, diff --git a/tests/utils_test.py b/tests/utils_test.py index 4c18b303..e16d62bf 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -7,6 +7,9 @@ import toml from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics +from PyPIC3D.solvers.yee_tiled import tile_vector_field 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, @@ -245,5 +248,172 @@ def test_dump_parameters_to_toml_writes_tiled_species_summaries(self): self.assertEqual(config["particles"][1]["active_particles"], 1) self.assertEqual(config["particles"][0]["tile_shape"], [2, 1, 1]) + def test_plot_positions_flattens_tiled_particles_and_preserves_species_names(self): + species = particle_species( + name="beam electrons", + N_particles=2, + charge=-1.0, + mass=1.0, + weight=2.0, + T=0.0, + x1=jnp.array([-0.25, 0.25]), + x2=jnp.array([0.0, 0.0]), + x3=jnp.array([0.0, 0.0]), + v1=jnp.array([0.1, 0.2]), + v2=jnp.array([0.0, 0.0]), + v3=jnp.array([0.0, 0.0]), + xwind=1.0, + ywind=1.0, + zwind=1.0, + dx=0.25, + dy=0.5, + dz=1.0, + dt=0.1, + ) + 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, + "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 = to_tiled_particles([species], world, 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 = particle_species( + name="electrons", + N_particles=1, + charge=-1.0, + mass=1.0, + weight=1.0, + T=0.0, + x1=jnp.array([0.0]), + x2=jnp.array([0.0]), + x3=jnp.array([0.0]), + v1=jnp.array([0.2]), + v2=jnp.array([0.0]), + v3=jnp.array([0.0]), + xwind=1.0, + ywind=1.0, + zwind=1.0, + dx=0.25, + dy=0.5, + dz=1.0, + dt=0.1, + ) + 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, + "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 = to_tiled_particles([species], world, 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) + + def test_vtk_plot_fields_assembles_tiled_components_at_output_boundary(self): + world = { + "Nx": 4, + "Ny": 2, + "Nz": 2, + "dx": 0.25, + "dy": 0.5, + "dz": 0.75, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.5, + "tile_shape": (2, 1, 1), + "guard_cells": 2, + "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + field = tuple(jnp.zeros((8, 6, 6)) + component for component in (1.0, 2.0, 3.0)) + tiled_field = tile_vector_field(field, world, world["tile_shape"], num_guard_cells=world["guard_cells"]) + + with unittest.mock.patch.object(vtk_diagnostics, "gridToVTK") as grid_to_vtk: + vtk_diagnostics.plot_fields( + tiled_field[0], + tiled_field[1], + tiled_field[2], + 0, + "E", + world["dx"], + world["dy"], + world["dz"], + world=world, + ) + + cell_data = grid_to_vtk.call_args.kwargs["cellData"] + self.assertEqual(cell_data["E_x"].shape, (6, 4, 4)) + self.assertTrue(jnp.allclose(cell_data["E_y"], 2.0)) + if __name__ == '__main__': unittest.main() diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 88e9945a..b1353ca2 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -322,6 +322,26 @@ def test_update_tiled_E_matches_standard_yee_update(self): 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_first_order_yee_public_update_E_routes_tiled_arrays(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world["tile_shape"] = tile_shape + world["guard_cells"] = 2 + 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, unused_curl) + E_reference = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) + + self.assertIsNone(pml_state) + for public, reference in zip(E_public, E_reference): + self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(self): world = self._conducting_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} @@ -368,6 +388,24 @@ def test_update_tiled_B_matches_standard_yee_update(self): 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_first_order_yee_public_update_B_routes_tiled_arrays(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 3, 2) + world["tile_shape"] = tile_shape + world["guard_cells"] = 2 + 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, unused_curl) + B_reference = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) + + self.assertIsNone(pml_state) + for public, reference in zip(B_public, B_reference): + self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(self): world = self._conducting_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} From 2d6484af1ed3af4874a59e252b5ea27e51de246d Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 15:27:22 -0500 Subject: [PATCH 30/83] finishing up tasks 2-5, 7, and 8 --- PyPIC3D/deposition/J_from_rhov.py | 2 +- PyPIC3D/deposition/rho.py | 2 +- PyPIC3D/diagnostics/fluid_quantities.py | 6 +- PyPIC3D/evolve.py | 6 +- PyPIC3D/solvers/electrostatic_yee.py | 4 +- PyPIC3D/solvers/first_order_yee.py | 12 ++-- tests/direct_deposition_tiled_test.py | 41 +++++++++-- tests/evolve_test.py | 4 +- tests/first_order_yee_test.py | 6 +- tests/particle_test.py | 26 +++---- tests/pml_test.py | 14 ++-- tests/tiled_electrostatic_test.py | 4 +- tests/tiled_fluid_quantities_test.py | 27 ++++++-- tests/tiled_rho_test.py | 17 +++-- tests/tiled_yee_integration_test.py | 90 ++++++++++++++++++------- tests/yee_tiled_test.py | 34 +++++++--- 16 files changed, 197 insertions(+), 98 deletions(-) diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index b2c21c01..65318f0c 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -56,7 +56,7 @@ def J_from_rhov( g=int(g), ) - return _J_from_rhov_flat(particles, J, constants, world, grid=grid, filter=filter) + raise ValueError("Public J_from_rhov requires TiledParticles; use _J_from_rhov_flat for flat reference tests.") @partial(jit, static_argnames=("filter",)) diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 204019d8..fdd7f8fd 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -32,7 +32,7 @@ def compute_rho(particles, rho, world, constants, species_config=None): return compute_rho_from_tiled_particles(particles, species_config, rho, world, constants) - return _compute_rho_flat(particles, rho, world, constants) + raise ValueError("Public compute_rho requires TiledParticles; use _compute_rho_flat for flat reference tests.") @jit diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index e3fafea8..c5be4eb9 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -47,7 +47,7 @@ def compute_mass_density(particles, rho, world, species_config=None): ) return compute_mass_density_from_tiled_particles(particles, species_config, rho, world) - return _compute_mass_density_flat(particles, rho, world) + raise ValueError("Public compute_mass_density requires TiledParticles; use _compute_mass_density_flat for flat reference tests.") @jit @@ -130,7 +130,7 @@ def compute_velocity_field(particles, field, direction, world, species_config=No ) return compute_velocity_field_from_tiled_particles(particles, species_config, field, int(direction), world) - return _compute_velocity_field_flat(particles, field, direction, world) + raise ValueError("Public compute_velocity_field requires TiledParticles; use _compute_velocity_field_flat for flat reference tests.") @partial(jit, static_argnames=("direction",)) @@ -241,7 +241,7 @@ def compute_pressure_field(particles, field, velocity_field, direction, world, s world, ) - return _compute_pressure_field_flat(particles, field, velocity_field, direction, world) + raise ValueError("Public compute_pressure_field requires TiledParticles; use _compute_pressure_field_flat for flat reference tests.") @partial(jit, static_argnames=("direction",)) diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 1d603511..45761dea 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -9,7 +9,7 @@ from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled as time_loop_electrostatic from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.solvers.electrostatic_yee import calculate_electrostatic_fields -from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global from PyPIC3D.utils import add_external_fields @@ -94,8 +94,8 @@ def _time_loop_electrodynamic_global_reference( particles[i].update_position() J = J_func(particles, J, constants, world) - E, pml_state = update_E(E, B, J, world, constants, curl_func, pml_state) - B, pml_state = update_B(E, B, world, constants, curl_func, pml_state) + E, pml_state = _update_E_global(E, B, J, world, constants, curl_func, pml_state) + B, pml_state = _update_B_global(E, B, world, constants, curl_func, pml_state) for i in range(len(particles)): particles[i].boundary_conditions(world) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index c5256ea9..b9213d8c 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -3,7 +3,7 @@ from jax import lax from functools import partial -from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient from PyPIC3D.utils import digital_filter @@ -135,7 +135,7 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver bc_y = world['boundary_conditions']['y'] bc_z = world['boundary_conditions']['z'] - rho = compute_rho(particles, rho, world, constants) + rho = _compute_rho_flat(particles, rho, world, constants) phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py index 3729c453..d1876083 100644 --- a/PyPIC3D/solvers/first_order_yee.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -217,8 +217,8 @@ def update_E(E, B, J, world, constants, curl_func, pml_state=None): Tiled field/current arrays are routed to the tile-local Yee update using the explicit ``world['tile_shape']`` and ``world['guard_cells']`` contract. - Ordinary global ghost-celled arrays keep the historical first-order Yee - update in ``_update_E_global``. + Flat global ghost-celled arrays are private reference inputs; tests that + need them should call ``_update_E_global`` directly. """ if E[0].ndim == 6: @@ -228,15 +228,15 @@ def update_E(E, B, J, world, constants, curl_func, pml_state=None): return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g), None return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g, pml_state) - return _update_E_global(E, B, J, world, constants, curl_func, pml_state) + raise ValueError("Public update_E requires tiled field arrays; use _update_E_global for flat reference tests.") def update_B(E, B, world, constants, curl_func, pml_state=None): """ Update magnetic fields through the public Yee helper. - Tiled arrays follow the tile-local Yee path and ordinary global arrays - follow the original global ghost-cell implementation. + Flat global ghost-celled arrays are private reference inputs; tests that + need them should call ``_update_B_global`` directly. """ if E[0].ndim == 6: @@ -246,4 +246,4 @@ def update_B(E, B, world, constants, curl_func, pml_state=None): return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g), None return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g, pml_state) - return _update_B_global(E, B, world, constants, curl_func, pml_state) + raise ValueError("Public update_B requires tiled field arrays; use _update_B_global for flat reference tests.") diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index d623cee1..93e41ccd 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -5,7 +5,7 @@ from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_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.deposition.J_from_rhov import J_from_rhov, _J_from_rhov_flat from PyPIC3D.deposition.direct_deposition_tiled import ( digital_filter_tiled_current_density, direct_J_from_tiled_particles, @@ -107,7 +107,7 @@ def _compare_tiled_to_standard(self, particles, world, simulation_parameters): tiled_particles, species_config = self._centered_tiled_particles(particles, world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_reference = J_from_rhov(particles, self._empty_J(world), constants, world, filter="none") + J_reference = _J_from_rhov_flat(particles, self._empty_J(world), constants, world, filter="none") J_tiles = direct_J_from_tiled_particles( tiled_particles, species_config, @@ -126,6 +126,33 @@ def _compare_tiled_to_standard(self, particles, world, simulation_parameters): 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)) + def test_public_J_from_rhov_rejects_flat_particles(self): + world = self._build_world() + constants = {"C": 3.0e8, "alpha": 1.0} + species = particle_species( + name="flat direct current", + N_particles=1, + charge=-1.0, + mass=1.0, + T=1.0, + x1=jnp.array([0.0]), + x2=jnp.array([0.0]), + x3=jnp.array([0.0]), + v1=jnp.array([0.0]), + v2=jnp.array([0.0]), + v3=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"], + dt=world["dt"], + ) + + with self.assertRaisesRegex(ValueError, "TiledParticles"): + J_from_rhov([species], self._empty_J(world), constants, world, filter="none") + 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 @@ -228,7 +255,7 @@ def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") + J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") J_tiles = direct_J_from_tiled_particles( tiled_particles, species_config, @@ -434,7 +461,7 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): tile_shape=tile_shape, g=int(world["guard_cells"]), ) - J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="none") + J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="none") J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) for reference_component, tiled_component in zip(J_reference, J_from_tiles): @@ -526,7 +553,7 @@ def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(sel ) tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) - J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") + J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="digital") J_tiles = J_from_rhov( tiled_particles, species_config, @@ -581,7 +608,7 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="digital") + J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="digital") J_tiles = direct_J_from_tiled_particles( tiled_particles, species_config, @@ -629,7 +656,7 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_reference = J_from_rhov([species], self._empty_J(world), constants, world, filter="bilinear") + J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") J_tiles = direct_J_from_tiled_particles( tiled_particles, species_config, diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 38917fef..44780d5c 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -11,7 +11,7 @@ time_loop_electrodynamic, time_loop_electrostatic, ) -from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species from PyPIC3D.pusher.particle_push import particle_push @@ -212,7 +212,7 @@ def test_electrostatic_rho_uses_boundary_applied_particle_positions(self): particle_pusher="boris", ) _, _, _, rho_after, _, _ = fields - recomputed_rho = compute_rho(particles, jnp.zeros_like(rho_after), world, constants) + recomputed_rho = _compute_rho_flat(particles, jnp.zeros_like(rho_after), world, constants) self.assertTrue( jnp.allclose( diff --git a/tests/first_order_yee_test.py b/tests/first_order_yee_test.py index 74a56988..a7449588 100644 --- a/tests/first_order_yee_test.py +++ b/tests/first_order_yee_test.py @@ -6,7 +6,7 @@ from PyPIC3D.initialization import initialize_fields -from PyPIC3D.solvers.first_order_yee import update_E, update_B +from PyPIC3D.solvers.first_order_yee import _update_E_global, _update_B_global from PyPIC3D.utils import build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -161,7 +161,7 @@ def test_update_E(self): 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) + E_computed, pml_state = _update_E_global(E_initial, B_initial, J, self.world, self.constants, curl_func) Ex_computed, Ey_computed, Ez_computed = E_computed # Check shapes (including ghost cells) @@ -242,7 +242,7 @@ def test_update_B(self): 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) + B_computed, pml_state = _update_B_global(E_initial, B_initial, self.world, self.constants, curl_func) Bx_computed, By_computed, Bz_computed = B_computed # Check shapes diff --git a/tests/particle_test.py b/tests/particle_test.py index a2653280..40068909 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -12,9 +12,9 @@ from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.J_from_rhov import _J_from_rhov_flat from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.utils import build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -508,7 +508,7 @@ def test_J_from_rhov(self): constants = {'C': 3e8, 'alpha' : 1.0} - num_J = J_from_rhov([species], num_J, constants, world) + num_J = _J_from_rhov_flat([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 @@ -559,8 +559,8 @@ def test_inactive_particles_do_not_deposit_rho_or_current(self): species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) constants = {"C": 3e8, "alpha": 1.0} - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) - J = J_from_rhov( + rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) + J = _J_from_rhov_flat( [species], ( jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), @@ -634,7 +634,7 @@ def test_rho(self): constants = {'C': 3e8, 'alpha' : 1.0} - num_rho = compute_rho([species], num_rho, world, constants) + num_rho = _compute_rho_flat([species], num_rho, world, constants) # compute rho self.assertTrue(jnp.allclose(num_rho, exp_rho)) @@ -682,7 +682,7 @@ def test_rho_periodic_boundary_exact_face(self): zwind=self.z_wind, ) - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_flat([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 @@ -733,7 +733,7 @@ def test_rho_periodic_boundary_crossing(self): zwind=self.z_wind, ) - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_flat([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) @@ -785,7 +785,7 @@ def test_J_from_rhov_periodic_boundary_crossing(self): zwind=self.z_wind, ) - J = J_from_rhov( + J = _J_from_rhov_flat( [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}, @@ -843,7 +843,7 @@ def test_rho_reduced_dimensions_periodic(self): zwind=z_wind, ) - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_flat([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)) @@ -891,7 +891,7 @@ def test_filtered_rho_periodic_updates_ghost_cells(self): zwind=self.z_wind, ) - rho = compute_rho([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 0.5}) + rho = _compute_rho_flat([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) @@ -963,10 +963,10 @@ def test_check_continuity_1D(self): species.update_position() # move particles to new position - prev_rho = compute_rho([species], rho, world, constants) + prev_rho = _compute_rho_flat([species], rho, world, constants) # compute rho species.update_position() - rho = compute_rho([species], rho, world, constants) + rho = _compute_rho_flat([species], rho, world, constants) # compute rho again # compute drhodt on interior diff --git a/tests/pml_test.py b/tests/pml_test.py index 561f4b2e..b3e41fc7 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -14,7 +14,7 @@ update_ghost_cells_for_pml, ) from PyPIC3D.initialization import initialize_fields, initialize_simulation -from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, empty_tiled_vector_field, @@ -269,8 +269,8 @@ def test_no_pml_state_returns_field_and_none_state(self): 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) + E_after, pml_state = _update_E_global(E, B, J, world, constants, lambda *args: None) + B_after, pml_state = _update_B_global(E_after, B, world, constants, lambda *args: None, pml_state) self.assertEqual(len(E_after), 3) self.assertEqual(len(B_after), 3) @@ -301,8 +301,8 @@ def test_pml_absorbs_field_energy_in_particle_free_1d_wave(self): 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) + E, pml_state = _update_E_global(E, B, J, world, constants, lambda *args: None, pml_state) + B, pml_state = _update_B_global(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)) @@ -343,8 +343,8 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): B = tuple(update_ghost_cells_for_pml(component, world) for component in B) pml_state = initialize_pml_state(world) - E_reference, pml_state = update_E(E, B, J, world, constants, lambda *args: None, pml_state) - B_reference, pml_state = update_B(E_reference, B, world, constants, lambda *args: None, pml_state) + E_reference, pml_state = _update_E_global(E, B, J, world, constants, lambda *args: None, pml_state) + B_reference, pml_state = _update_B_global(E_reference, B, world, constants, lambda *args: None, pml_state) tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) E_tiles, tiled_pml_state = update_tiled_E( diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index f18338bb..2c8e18d8 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -8,7 +8,7 @@ import toml from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled from PyPIC3D.evolve import _time_loop_electrostatic_global_reference @@ -346,7 +346,7 @@ def test_tiled_rho_assembles_to_global_rho(self): _, _, _, rho, _, _ = self._empty_fields(world) rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) - reference_rho = compute_rho([self._species(world)], rho, world, constants) + reference_rho = _compute_rho_flat([self._species(world)], rho, world, constants) tiled_rho = compute_tiled_rho_from_tiled_particles( tiled_particles, species_config, diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 6f51ee84..9c941328 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -10,6 +10,9 @@ compute_mass_density, compute_pressure_field, compute_velocity_field, + _compute_mass_density_flat, + _compute_pressure_field_flat, + _compute_velocity_field_flat, ) from PyPIC3D.deposition.rho_tiled import compute_tiled_mass_density_from_tiled_particles from PyPIC3D.particles.species_class import particle_species @@ -110,7 +113,7 @@ def test_tiled_mass_density_matches_compute_mass_density_on_interior(self): x = x.at[inactive, 2].set(-0.30) tiled_particles = tiled_particles._replace(x=x) - mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) + mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) mass_tiled = compute_mass_density(tiled_particles, self._empty_scalar(world), world, species_config=species_config) self.assertTrue( @@ -138,7 +141,7 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): num_guard_cells=int(world["guard_cells"]), ) - mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) + mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) mass_tiles = compute_tiled_mass_density_from_tiled_particles( tiled_particles, species_config, @@ -179,7 +182,7 @@ def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(s num_guard_cells=int(world["guard_cells"]), ) - mass_reference = compute_mass_density(particles, self._empty_scalar(world), world) + mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) mass_tiles = compute_mass_density(tiled_particles, rho_tiles, world, species_config=species_config) mass_from_tiles = assemble_tiled_scalar_field( mass_tiles, @@ -214,7 +217,7 @@ def test_tile_major_velocity_field_assembles_to_global_velocity_field(self): num_guard_cells=int(world["guard_cells"]), ) - velocity_reference = compute_velocity_field(particles, self._empty_scalar(world), 0, world) + velocity_reference = _compute_velocity_field_flat(particles, self._empty_scalar(world), 0, world) velocity_tiles = compute_velocity_field( tiled_particles, field_tiles, @@ -254,8 +257,8 @@ def test_tile_major_pressure_field_assembles_to_global_pressure_field(self): num_guard_cells=int(world["guard_cells"]), ) - velocity_reference = compute_velocity_field(particles, self._empty_scalar(world), 0, world) - pressure_reference = compute_pressure_field( + velocity_reference = _compute_velocity_field_flat(particles, self._empty_scalar(world), 0, world) + pressure_reference = _compute_pressure_field_flat( particles, self._empty_scalar(world), velocity_reference, @@ -293,6 +296,18 @@ def test_tile_major_pressure_field_assembles_to_global_pressure_field(self): ) ) + 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) + def test_tiled_scalar_vtk_path_no_longer_rejects_electrodynamic_yee(self): run_source = inspect.getsource(pypic_main.run_PyPIC3D) diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 7332f656..049c3c4f 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -4,7 +4,7 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.deposition.rho import compute_rho, _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -120,7 +120,7 @@ def _compare_tiled_to_standard(self, shape_factor, alpha): particles = self._particles(world) tiled_particles, species_config = self._tiled_with_noisy_inactive_slots(particles, world) - rho_reference = compute_rho(particles, self._empty_scalar(world), world, constants) + rho_reference = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) rho_tiled = compute_rho_from_tiled_particles( tiled_particles, species_config, @@ -153,8 +153,8 @@ def test_compute_rho_uses_current_positions_not_half_step_back_positions(self): particles = self._particles(world) zero_velocity_particles = self._zero_species_velocities(self._particles(world)) - rho_with_velocity = compute_rho(particles, self._empty_scalar(world), world, constants) - rho_with_zero_velocity = compute_rho(zero_velocity_particles, self._empty_scalar(world), world, constants) + rho_with_velocity = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) + rho_with_zero_velocity = _compute_rho_flat(zero_velocity_particles, self._empty_scalar(world), world, constants) self.assertTrue( jnp.allclose( @@ -256,7 +256,7 @@ def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(sel num_guard_cells=int(world["guard_cells"]), ) - rho_reference = compute_rho(particles, self._empty_scalar(world), world, constants) + rho_reference = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) rho_tiles = compute_rho( tiled_particles, rho_tiles, @@ -280,6 +280,13 @@ def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(sel ) ) + 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), self._empty_scalar(world), world, constants) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index d612c1ca..5282f3dc 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -10,7 +10,7 @@ from PyPIC3D.evolve import _time_loop_electrodynamic_global_reference from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.deposition.J_from_rhov import _J_from_rhov_flat from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current @@ -200,7 +200,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): step=0, ) - return reference_particles, reference_fields, tiled_particles, tiled_fields, world, constants, tile_shape, Nt + return reference_particles, reference_fields, tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): current_slots = tiled_particles.active.shape[-1] @@ -292,6 +292,26 @@ def _assert_tiled_state_matches_standard( "J", ) + def _max_tiled_state_difference(self, reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape): + differences = {} + for species_index in range(len(reference_particles)): + reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) + differences[f"x{species_index}"] = float(jnp.max(jnp.abs(tiled_x - reference_x))) + differences[f"u{species_index}"] = float(jnp.max(jnp.abs(tiled_u - reference_u))) + + reference_E, reference_B, reference_J, *_ = reference_fields + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + num_guard_cells = int(world.get("guard_cells", 1)) + for name, reference_field, tiled_field in ( + ("E", reference_E, assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), + ("B", reference_B, assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), + ("J", reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), + ): + differences[name] = max(float(jnp.max(jnp.abs(tiled - reference))) for reference, tiled in zip(reference_field, tiled_field)) + + return differences + def test_validate_field_solver_accepts_tiled_yee(self): validate_field_solver("electrodynamic_yee") @@ -684,7 +704,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="none"), + J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="none"), solver="fdtd", relativistic=False, particle_pusher="boris", @@ -764,7 +784,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="none"), + J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="none"), solver="fdtd", relativistic=False, particle_pusher="boris", @@ -828,7 +848,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="digital"), + J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="digital"), solver="fdtd", relativistic=False, particle_pusher="boris", @@ -889,7 +909,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="bilinear"), + J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="bilinear"), solver="fdtd", relativistic=False, particle_pusher="boris", @@ -992,7 +1012,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J_from_rhov(particles, J, constants, world, filter="bilinear"), + J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="bilinear"), solver="fdtd", relativistic=False, particle_pusher="boris", @@ -1040,11 +1060,12 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand os.environ.get("RUN_SLOW_TILED_YEE") == "1", "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", ) - def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_standard_every_step(self): + def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): ( reference_particles, reference_fields, tiled_particles, + species_config, tiled_fields, world, constants, @@ -1052,7 +1073,7 @@ def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_s Nt, ) = self._long_two_stream_state(CURRENT_J_FROM_RHOV) - standard_J = functools.partial(J_from_rhov, filter="none") + standard_J = functools.partial(_J_from_rhov_flat, filter="none") tiled_loop = jax.jit( time_loop_electrodynamic_tiled, static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), @@ -1072,6 +1093,7 @@ def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_s ) tiled_particles, tiled_fields = tiled_loop( tiled_particles, + species_config, tiled_fields, world, constants, @@ -1084,25 +1106,33 @@ def test_long_two_stream_tiled_yee_direct_current_quadratic_two_guards_matches_s particle_pusher="boris", ) - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - tile_shape, - t + 1, + differences = self._max_tiled_state_difference( + reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape ) + if t == 0: + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + t + 1, + ) + if any(value > ROUND_OFF_ATOL for value in differences.values()): + print(f"direct-current roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") + break @unittest.skipUnless( os.environ.get("RUN_SLOW_TILED_YEE") == "1", "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", ) - def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standard_every_step(self): + def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): ( reference_particles, reference_fields, tiled_particles, + species_config, tiled_fields, world, constants, @@ -1129,6 +1159,7 @@ def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standa ) tiled_particles, tiled_fields = tiled_loop( tiled_particles, + species_config, tiled_fields, world, constants, @@ -1141,15 +1172,22 @@ def test_long_two_stream_tiled_yee_esirkepov_quadratic_two_guards_matches_standa particle_pusher="boris", ) - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - tile_shape, - t + 1, + differences = self._max_tiled_state_difference( + reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape ) + if t == 0: + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + tile_shape, + t + 1, + ) + if any(value > ROUND_OFF_ATOL for value in differences.values()): + print(f"Esirkepov roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") + break if __name__ == "__main__": diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index b1353ca2..6aa7ac35 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -5,7 +5,7 @@ from PyPIC3D.initialization import initialize_fields from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled -from PyPIC3D.solvers.first_order_yee import update_B, update_E +from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global, update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, tile_grid_axes, @@ -306,7 +306,7 @@ def test_update_tiled_E_matches_standard_yee_update(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = update_E(E, B, J, world, constants, unused_curl) + E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) E_tiled = update_tiled_E( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -342,6 +342,18 @@ def test_first_order_yee_public_update_E_routes_tiled_arrays(self): for public, reference in zip(E_public, E_reference): self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) + def test_first_order_yee_public_update_rejects_global_arrays(self): + world = self._build_world() + constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + 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) + + with self.assertRaisesRegex(ValueError, "tiled field arrays"): + update_E(E, B, J, world, constants, unused_curl) + with self.assertRaisesRegex(ValueError, "tiled field arrays"): + update_B(E, B, world, constants, unused_curl) + def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(self): world = self._conducting_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} @@ -350,7 +362,7 @@ def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(s B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = update_E(E, B, J, world, constants, unused_curl) + E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) E_tiled = update_tiled_E( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -373,7 +385,7 @@ def test_update_tiled_B_matches_standard_yee_update(self): E = self._deterministic_vector_field(world, scale=1.0) B = self._deterministic_vector_field(world, scale=0.2) - B_reference, _ = update_B(E, B, world, constants, unused_curl) + B_reference, _ = _update_B_global(E, B, world, constants, unused_curl) B_tiled = update_tiled_B( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -413,7 +425,7 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s E = self._deterministic_vector_field(world, scale=1.0) B = self._deterministic_vector_field(world, scale=0.2) - B_reference, _ = update_B(E, B, world, constants, unused_curl) + B_reference, _ = _update_B_global(E, B, world, constants, unused_curl) B_tiled = update_tiled_B( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), @@ -436,8 +448,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = update_E(E, B, J, world, constants, unused_curl) - B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) + B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) E_tiles = tile_vector_field(E, world, tile_shape) B_tiles = tile_vector_field(B, world, tile_shape) @@ -461,8 +473,8 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = update_E(E, B, J, world, constants, unused_curl) - B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) + B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) E_tiles = tile_vector_field(E, world, tile_shape) B_tiles = tile_vector_field(B, world, tile_shape) @@ -486,8 +498,8 @@ def test_tiled_evolve_updates_fields_without_pushing_particles(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = update_E(E, B, J, world, constants, unused_curl) - B_reference, _ = update_B(E_reference, B, world, constants, unused_curl) + E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) + B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) particles = object() external_fields = ( tuple(jnp.zeros_like(component) for component in E), From 57da22a6c6da47360b10c4b7f862c151f4775686 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 16:20:27 -0500 Subject: [PATCH 31/83] finishing migration towards complete tiling --- PyPIC3D/deposition/Esirkepov.py | 4 +- PyPIC3D/deposition/J_from_rhov.py | 4 +- PyPIC3D/deposition/direct_deposition_tiled.py | 2 +- PyPIC3D/deposition/esirkepov_tiled.py | 2 +- PyPIC3D/diagnostics/tiled_benchmarking.py | 8 +- PyPIC3D/electrodynamic_tiled.py | 149 ------- PyPIC3D/electrostatic_tiled.py | 92 ---- PyPIC3D/evolve.py | 223 +++++++++- PyPIC3D/particles/flat_particles.py | 407 ------------------ tests/direct_deposition_tiled_test.py | 11 +- tests/evolve_test.py | 11 +- tests/tiled_electrostatic_test.py | 8 +- tests/tiled_esirkepov_test.py | 9 +- tests/tiled_refactor_contract_test.py | 6 +- tests/tiled_yee_integration_test.py | 40 +- tests/yee_tiled_test.py | 4 +- 16 files changed, 271 insertions(+), 709 deletions(-) delete mode 100644 PyPIC3D/electrodynamic_tiled.py delete mode 100644 PyPIC3D/electrostatic_tiled.py delete mode 100644 PyPIC3D/particles/flat_particles.py diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 4a3c9d6a..2972df0c 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -238,9 +238,9 @@ def Esirkepov_current( if g is None: g = int(world["guard_cells"]) - from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current + from PyPIC3D.deposition.esirkepov_tiled import _tiled_esirkepov_current - return tiled_esirkepov_current( + return _tiled_esirkepov_current( particles, species_config, J, diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index 65318f0c..18808d74 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -9,7 +9,7 @@ prepare_particle_axis_stencil, ) from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.direct_deposition_tiled import _direct_J_from_tiled_particles from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.utils import bilinear_filter, digital_filter @@ -44,7 +44,7 @@ def J_from_rhov( if g is None: g = int(world["guard_cells"]) - return direct_J_from_tiled_particles( + return _direct_J_from_tiled_particles( particles, species_config, J, diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py index 60ee54ee..69de8f21 100644 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ b/PyPIC3D/deposition/direct_deposition_tiled.py @@ -104,7 +104,7 @@ def _collapse_tiled_axis_stencil(points, weights, local_n, reduced_axis, g): @partial(jit, static_argnames=("filter", "tile_shape", "g")) -def direct_J_from_tiled_particles( +def _direct_J_from_tiled_particles( tiled_particles, species_config, J_tiles, diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py index 3c7c6c77..4e61e32f 100644 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ b/PyPIC3D/deposition/esirkepov_tiled.py @@ -185,7 +185,7 @@ def fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells= @partial(jit, static_argnames=("tile_shape", "g")) -def tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=2): +def _tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=2): """ Deposit Esirkepov current into tile-local current buffers. diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index dd31e9b9..2d5cb3c5 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -6,9 +6,9 @@ from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC from PyPIC3D.deposition.current_methods import CURRENT_J_FROM_RHOV -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles +from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh import ( @@ -274,7 +274,7 @@ def _tiled_output_bridge_stage(case): def benchmark_tiled_pic_step(case, J_func, relativistic=True, particle_pusher="boris"): with jax.named_scope("tiled_pic_step"): - return time_loop_electrodynamic_tiled( + return time_loop_electrodynamic( case.particles, case.species_config, case.fields, @@ -300,7 +300,7 @@ def build_tiled_stage_specs(case, J_func=None, relativistic=True, particle_pushe """ if J_func is None: - J_func = partial(direct_J_from_tiled_particles, filter="none") + J_func = partial(J_from_rhov, filter="none") pushed_particles = _tiled_particle_push_stage(case, relativistic=relativistic, particle_pusher=particle_pusher) deposited_J = _tiled_current_deposition_stage(case, pushed_particles, J_func) diff --git a/PyPIC3D/electrodynamic_tiled.py b/PyPIC3D/electrodynamic_tiled.py deleted file mode 100644 index 20e7c81e..00000000 --- a/PyPIC3D/electrodynamic_tiled.py +++ /dev/null @@ -1,149 +0,0 @@ -import jax.numpy as jnp -import jax - -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles -from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV -from PyPIC3D.particles.tiled_particle_refresh import ( - refresh_tiled_particle_tiles, - update_tiled_particle_positions, -) -from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E -from PyPIC3D.utils import add_external_fields - - -def time_loop_electrodynamic_tiled( - particles, - species_config, - fields, - world, - constants, - curl_func, - J_func, - solver, - tile_shape=None, - g=None, - relativistic=True, - particle_pusher="boris", -): - """ - Advance a periodic tiled electrodynamic PIC system by one time step. - """ - - del solver - - if len(fields) == 7: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields - overflow_previous = jnp.asarray(False) - else: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow_previous = fields - # unpack the tile-major field state. ``overflow_previous`` carries the - # particle-retile overflow diagnostic out to the Python driver. - - if tile_shape is None or g is None: - raise ValueError("tiled electrodynamic updates require explicit tile_shape and g.") - g = int(g) - - if not hasattr(particles, "active"): - if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) - else: - E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - B_tiles, pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) - return particles, fields - # keep the original field-only helper behavior for tests and standalone - # Maxwell equivalence checks that pass a non-particle sentinel. - push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) - # particles see evolved fields plus external-only fields, as in the - # standard electrodynamic path. - - particles = tiled_particle_push( - particles, - species_config, - push_E_tiles, - push_B_tiles, - world, - constants, - tile_shape, - g, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - # use the selected tiled pusher for particle velocities - - if J_func is None: - def current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=None, g=None): - return direct_J_from_tiled_particles( - particles, - species_config, - J_tiles, - constants, - world, - filter="none", - tile_shape=tile_shape, - g=g, - ) - else: - current_deposition = J_func - - def esirkepov_step(state): - particles, J_tiles, overflow_previous = state - # Esirkepov needs old and new particle positions. The deposition kernel - # predicts the new positions locally, then the actual particle state is - # advanced and retiled after the current has been computed. - J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) - particles = update_tiled_particle_positions(particles, species_config, world["dt"]) - particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) - overflow = overflow_previous | overflow - return particles, J_tiles, overflow - - def direct_current_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) - overflow = overflow_previous | overflow - # wrap periodic particles and move them into their owning tiles. - - J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) - # 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) - overflow = overflow_previous | overflow - # refresh tile ownership after the full position update. - return particles, J_tiles, overflow - - current_calculation = world.get("current_calculation", CURRENT_J_FROM_RHOV) - particles, J_tiles, overflow = jax.lax.cond( - current_calculation == CURRENT_ESIRKEPOV, - esirkepov_step, - direct_current_step, - (particles, J_tiles, overflow_previous), - ) - - if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) - else: - E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - # update electric field from B and the supplied tiled current - - if pml_state is None: - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) - else: - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state) - # update magnetic field from the newly updated electric field - - fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) - # pack the tiled field state - - return particles, fields diff --git a/PyPIC3D/electrostatic_tiled.py b/PyPIC3D/electrostatic_tiled.py deleted file mode 100644 index 918d1f85..00000000 --- a/PyPIC3D/electrostatic_tiled.py +++ /dev/null @@ -1,92 +0,0 @@ -import jax.numpy as jnp - -from PyPIC3D.particles.tiled_particle_refresh import ( - refresh_tiled_particle_tiles, - update_tiled_particle_positions, -) -from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields -from PyPIC3D.utils import add_external_fields - - -def time_loop_electrostatic_tiled( - particles, - species_config, - fields, - world, - constants, - curl_func, - J_func, - solver, - tile_shape=None, - g=None, - relativistic=True, - particle_pusher="boris", -): - """ - Advance a tiled electrostatic PIC system by one time step. - - The particle push and retile use tile-local Yee 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. - """ - - del curl_func, J_func - - if len(fields) == 6: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields = fields - pml_state = None - overflow_previous = jnp.asarray(False) - elif len(fields) == 7: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state = fields - overflow_previous = jnp.asarray(False) - else: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields - - if tile_shape is None or g is None: - raise ValueError("tiled electrostatic updates require explicit tile_shape and g.") - g = int(g) - - push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) - # particles see evolved fields plus prescribed external fields - - particles = tiled_particle_push( - particles, - species_config, - push_E_tiles, - push_B_tiles, - world, - constants, - tile_shape, - g, - 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", - tile_shape, - g, - ) - # 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/evolve.py b/PyPIC3D/evolve.py index 45761dea..74902964 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -3,19 +3,238 @@ from functools import partial +import jax +import jax.numpy as jnp from jax import jit -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled as time_loop_electrodynamic -from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled as time_loop_electrostatic +from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV +from PyPIC3D.deposition.J_from_rhov import J_from_rhov +from PyPIC3D.particles.tiled_particle_refresh import ( + refresh_tiled_particle_tiles, + update_tiled_particle_positions, +) from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.solvers.electrostatic_yee import calculate_electrostatic_fields +from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global +from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E from PyPIC3D.utils import add_external_fields __all__ = ["time_loop_electrodynamic", "time_loop_electrostatic"] +def time_loop_electrodynamic( + particles, + species_config, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=None, + g=None, + relativistic=True, + particle_pusher="boris", +): + """ + Advance a tiled electrodynamic PIC system by one time step. + """ + + del solver + + if len(fields) == 7: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields + overflow_previous = jnp.asarray(False) + else: + E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow_previous = fields + # unpack the tile-major field state. ``overflow_previous`` carries the + # particle-retile overflow diagnostic out to the Python driver. + + if tile_shape is None or g is None: + raise ValueError("tiled electrodynamic updates require explicit tile_shape and g.") + g = int(g) + + if not hasattr(particles, "active"): + if pml_state is None: + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) + else: + E_tiles, pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state + ) + B_tiles, pml_state = update_tiled_B( + E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state + ) + fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) + return particles, fields + # keep field-only helper behavior for standalone Maxwell equivalence checks. + + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) + # particles see evolved fields plus external-only fields, as in the standard + # electrodynamic path. + + particles = tiled_particle_push( + particles, + species_config, + push_E_tiles, + push_B_tiles, + world, + constants, + tile_shape, + g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + # use the selected tiled pusher for particle velocities + + if J_func is None: + current_deposition = partial(J_from_rhov, filter="none") + else: + current_deposition = J_func + + def esirkepov_step(state): + particles, J_tiles, overflow_previous = state + # Esirkepov needs old and new particle positions. The deposition kernel + # predicts the new positions locally, then the actual particle state is + # advanced and retiled after the current has been computed. + J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + particles = update_tiled_particle_positions(particles, species_config, world["dt"]) + particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) + overflow = overflow_previous | overflow + return particles, J_tiles, overflow + + def direct_current_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) + overflow = overflow_previous | overflow + # wrap periodic particles and move them into their owning tiles. + + J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + # 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) + overflow = overflow_previous | overflow + # refresh tile ownership after the full position update. + return particles, J_tiles, overflow + + current_calculation = world.get("current_calculation", CURRENT_J_FROM_RHOV) + particles, J_tiles, overflow = jax.lax.cond( + current_calculation == CURRENT_ESIRKEPOV, + esirkepov_step, + direct_current_step, + (particles, J_tiles, overflow_previous), + ) + + if pml_state is None: + E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) + else: + E_tiles, pml_state = update_tiled_E( + E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state + ) + # update electric field from B and the supplied tiled current + + if pml_state is None: + B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) + else: + B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state) + # update magnetic field from the newly updated electric field + + fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) + # pack the tiled field state + + return particles, fields + + +def time_loop_electrostatic( + particles, + species_config, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=None, + g=None, + relativistic=True, + particle_pusher="boris", +): + """ + Advance a tiled electrostatic PIC system by one time step. + + 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. + """ + + del curl_func, J_func + + if len(fields) == 6: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields = fields + pml_state = None + overflow_previous = jnp.asarray(False) + elif len(fields) == 7: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state = fields + overflow_previous = jnp.asarray(False) + else: + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields + + if tile_shape is None or g is None: + raise ValueError("tiled electrostatic updates require explicit tile_shape and g.") + g = int(g) + + push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) + # particles see evolved fields plus prescribed external fields + + particles = tiled_particle_push( + particles, + species_config, + push_E_tiles, + push_B_tiles, + world, + constants, + tile_shape, + g, + 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", + tile_shape, + g, + ) + # 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 + + @partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) def _time_loop_electrostatic_global_reference( particles, diff --git a/PyPIC3D/particles/flat_particles.py b/PyPIC3D/particles/flat_particles.py deleted file mode 100644 index fe34813b..00000000 --- a/PyPIC3D/particles/flat_particles.py +++ /dev/null @@ -1,407 +0,0 @@ -import jax -import jax.numpy as jnp - - -@jax.jit -def _apply_axis_boundary_condition(x, v, wind, half_wind, bc): - periodic = bc == 0 - reflecting = bc == 1 - - periodic_x = x + wind * (x < -half_wind) - wind * (x > half_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_v = jnp.where((x >= half_wind) | (x <= -half_wind), -v, v) - - x_out = jnp.where(periodic, periodic_x, jnp.where(reflecting, reflected_x, x)) - v_out = jnp.where(reflecting, reflected_v, v) - - return x_out, v_out - - -@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, - 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.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 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, world=None): - half_x = self.x_wind / 2 - half_y = self.y_wind / 2 - half_z = self.z_wind / 2 - if world is None or "particle_boundary_conditions" not in world: - particle_bc = {"x": 0, "y": 0, "z": 0} - else: - particle_bc = world["particle_boundary_conditions"] - - self.x1, self.v1 = _apply_axis_boundary_condition( - self.x1, self.v1, self.x_wind, half_x, particle_bc["x"] - ) - self.x2, self.v2 = _apply_axis_boundary_condition( - self.x2, self.v2, self.y_wind, half_y, particle_bc["y"] - ) - self.x3, self.v3 = _apply_axis_boundary_condition( - self.x3, self.v3, self.z_wind, half_z, particle_bc["z"] - ) - - active_mask = self.active_mask - x_inside = (self.x1 <= half_x) & (self.x1 >= -half_x) - y_inside = (self.x2 <= half_y) & (self.x2 >= -half_y) - z_inside = (self.x3 <= half_z) & (self.x3 >= -half_z) - active_mask = jnp.where(particle_bc["x"] == 2, active_mask & x_inside, active_mask) - active_mask = jnp.where(particle_bc["y"] == 2, active_mask & y_inside, active_mask) - active_mask = jnp.where(particle_bc["z"] == 2, active_mask & z_inside, active_mask) - self.active_mask = active_mask - - 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.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, - 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, - 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.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, - dt=first.dt, - species_meta=species_meta, - active_mask=active_mask, - ) - return [flat] diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 93e41ccd..9b072880 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -8,7 +8,6 @@ from PyPIC3D.deposition.J_from_rhov import J_from_rhov, _J_from_rhov_flat from PyPIC3D.deposition.direct_deposition_tiled import ( digital_filter_tiled_current_density, - direct_J_from_tiled_particles, ) from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -108,7 +107,7 @@ def _compare_tiled_to_standard(self, particles, world, simulation_parameters): tile_shape = self._tile_shape(simulation_parameters) J_reference = _J_from_rhov_flat(particles, self._empty_J(world), constants, world, filter="none") - J_tiles = direct_J_from_tiled_particles( + J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), @@ -256,7 +255,7 @@ def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): tile_shape = self._tile_shape(simulation_parameters) J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") - J_tiles = direct_J_from_tiled_particles( + J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), @@ -451,7 +450,7 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_tiles = direct_J_from_tiled_particles( + J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), @@ -609,7 +608,7 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): tile_shape = self._tile_shape(simulation_parameters) J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="digital") - J_tiles = direct_J_from_tiled_particles( + J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), @@ -657,7 +656,7 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): tile_shape = self._tile_shape(simulation_parameters) J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") - J_tiles = direct_J_from_tiled_particles( + J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 44780d5c..e742f8dd 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -3,8 +3,6 @@ import jax import jax.numpy as jnp -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled -from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled from PyPIC3D.evolve import ( _time_loop_electrodynamic_global_reference, _time_loop_electrostatic_global_reference, @@ -29,12 +27,9 @@ def unused_curl(Ex, Ey, Ez): class TestEvolveExternalFields(unittest.TestCase): - def test_public_loop_names_use_tiled_contracts(self): - electrodynamic_loop = time_loop_electrodynamic.func if hasattr(time_loop_electrodynamic, "func") else time_loop_electrodynamic - electrostatic_loop = time_loop_electrostatic.func if hasattr(time_loop_electrostatic, "func") else time_loop_electrostatic - - self.assertIs(electrodynamic_loop, time_loop_electrodynamic_tiled) - self.assertIs(electrostatic_loop, time_loop_electrostatic_tiled) + 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") def _electrostatic_grid_order_species(self, world, charge): return particle_species( diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 2c8e18d8..52b98c8d 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -10,7 +10,7 @@ from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles -from PyPIC3D.electrostatic_tiled import time_loop_electrostatic_tiled +from PyPIC3D.evolve import time_loop_electrostatic from PyPIC3D.evolve import _time_loop_electrostatic_global_reference from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.species_class import particle_species @@ -398,7 +398,7 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): None, ) - tiled_particles, tiled_fields = time_loop_electrostatic_tiled( + tiled_particles, tiled_fields = time_loop_electrostatic( tiled_particles, species_config, tiled_fields, @@ -449,7 +449,7 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self ) = self._long_two_stream_state() tiled_loop = jax.jit( - time_loop_electrostatic_tiled, + time_loop_electrostatic, static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) @@ -539,7 +539,7 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): loop, particles, fields, world, simulation_parameters, *_ = initialize_simulation(toml.loads(toml.dumps(config))) - self.assertIs(loop, time_loop_electrostatic_tiled) + self.assertIs(loop, time_loop_electrostatic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) self.assertEqual(tuple(simulation_parameters["tile_shape"]), (2, 1, 1)) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 3be8f8ba..60e945f3 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -14,7 +14,6 @@ _active_stencil_indices, _compact_1d_esirkepov_weights, _compact_2d_esirkepov_weights, - tiled_esirkepov_current, ) from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles @@ -155,7 +154,7 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) - J_tiles = tiled_esirkepov_current( + J_tiles = Esirkepov_current( tiled_particles, species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), @@ -325,7 +324,7 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) g = int(world["guard_cells"]) - J_tiles = tiled_esirkepov_current( + J_tiles = Esirkepov_current( tiled_particles, species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), @@ -511,7 +510,7 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) - J_tiles = tiled_esirkepov_current( + J_tiles = Esirkepov_current( tiled_particles, species_config, empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), @@ -775,7 +774,7 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s loop, particles, fields, world, _simulation_parameters, constants, _plotting_parameters, _plasma_parameters, \ solver, _electrostatic, _verbose, _GPUs, _Nt, curl_func, _J_func, relativistic, particle_pusher, species_config = initialize_simulation(config) - alias_J_func = functools.partial(tiled_esirkepov_current) + alias_J_func = functools.partial(Esirkepov_current) particles, fields = loop( particles, species_config, diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index c82085c7..e36fb0fd 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -3,7 +3,7 @@ import jax.numpy as jnp -from PyPIC3D import electrodynamic_tiled +from PyPIC3D import evolve from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled from PyPIC3D.utils import build_yee_grid @@ -87,8 +87,8 @@ def test_tiled_benchmark_stage_specs_expose_tiled_hot_paths(self): self.assertEqual(stage_specs["tiled_pic_step"].static_argnames, ("J_func", "relativistic", "particle_pusher")) self.assertEqual(stage_specs["tiled_current_deposition"].static_argnames, ("J_func",)) - def test_electrodynamic_tiled_hot_step_does_not_assemble_global_fields(self): - source = inspect.getsource(electrodynamic_tiled.time_loop_electrodynamic_tiled) + def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): + source = inspect.getsource(evolve.time_loop_electrodynamic) self.assertNotIn("assemble_tiled_scalar_field", source) self.assertNotIn("assemble_tiled_vector_field", source) self.assertNotIn("fields_for_output", source) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 5282f3dc..70dc4120 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -10,11 +10,9 @@ from PyPIC3D.evolve import _time_loop_electrodynamic_global_reference from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.J_from_rhov import _J_from_rhov_flat +from PyPIC3D.deposition.J_from_rhov import J_from_rhov, _J_from_rhov_flat from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV -from PyPIC3D.deposition.direct_deposition_tiled import direct_J_from_tiled_particles -from PyPIC3D.deposition.esirkepov_tiled import tiled_esirkepov_current -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -361,7 +359,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic_tiled( + _, tiled_fields = time_loop_electrodynamic( [], None, tiled_fields, @@ -444,7 +442,7 @@ def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(sel loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) @@ -502,7 +500,7 @@ def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(self): loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(simulation_parameters["particle_pusher"], "higuera_cary") @@ -560,7 +558,7 @@ def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(self): loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(simulation_parameters["filter_j"], "digital") @@ -618,7 +616,7 @@ def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(self): loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(simulation_parameters["filter_j"], "bilinear") @@ -679,7 +677,7 @@ def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(sel loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic_tiled) + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) self.assertIsInstance(particles, TiledParticles) self.assertEqual(fields[0][0].ndim, 6) self.assertEqual(world["boundary_conditions"], {"x": 1, "y": 1, "z": 1}) @@ -723,7 +721,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): ), None, ) - tiled_particles, tiled_fields = time_loop_electrodynamic_tiled( + tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -803,7 +801,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic_tiled( + _, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -867,14 +865,14 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic_tiled( + _, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, world, constants, unused_curl, - J_func=functools.partial(direct_J_from_tiled_particles, filter="digital"), + J_func=functools.partial(J_from_rhov, filter="digital"), solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), @@ -928,14 +926,14 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic_tiled( + _, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, world, constants, unused_curl, - J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), + J_func=functools.partial(J_from_rhov, filter="bilinear"), solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), @@ -1031,14 +1029,14 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ), None, ) - _, tiled_fields = time_loop_electrodynamic_tiled( + _, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, world, constants, unused_curl, - J_func=functools.partial(direct_J_from_tiled_particles, filter="bilinear"), + J_func=functools.partial(J_from_rhov, filter="bilinear"), solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), @@ -1075,7 +1073,7 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): standard_J = functools.partial(_J_from_rhov_flat, filter="none") tiled_loop = jax.jit( - time_loop_electrodynamic_tiled, + time_loop_electrodynamic, static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) @@ -1141,7 +1139,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): ) = self._long_two_stream_state(CURRENT_ESIRKEPOV) tiled_loop = jax.jit( - time_loop_electrodynamic_tiled, + time_loop_electrodynamic, static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) @@ -1164,7 +1162,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): world, constants, unused_curl, - J_func=tiled_esirkepov_current, + J_func=Esirkepov_current, solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 6aa7ac35..35d1f203 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -4,7 +4,7 @@ import jax.numpy as jnp from PyPIC3D.initialization import initialize_fields -from PyPIC3D.electrodynamic_tiled import time_loop_electrodynamic_tiled +from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global, update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, @@ -515,7 +515,7 @@ def test_tiled_evolve_updates_fields_without_pushing_particles(self): None, ) - particles_after, fields_after = time_loop_electrodynamic_tiled( + particles_after, fields_after = time_loop_electrodynamic( particles, None, fields, From 687f4df78666f17def7d9480434b1193f64b4e56 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 17:09:46 -0500 Subject: [PATCH 32/83] refactoring tiled particle logic --- PyPIC3D/deposition/Esirkepov.py | 253 +--------------- PyPIC3D/deposition/J_from_rhov.py | 218 +------------- PyPIC3D/evolve.py | 47 --- tests/direct_deposition_tiled_test.py | 109 +++---- tests/evolve_test.py | 78 ++++- tests/particle_test.py | 228 +------------- tests/tiled_esirkepov_test.py | 55 ++-- tests/tiled_refactor_contract_test.py | 5 + tests/tiled_yee_integration_test.py | 408 ++++++++++++++++++-------- 9 files changed, 428 insertions(+), 973 deletions(-) diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 2972df0c..15465419 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -1,17 +1,7 @@ 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.particles.tiled_particles import TiledParticles @@ -251,248 +241,7 @@ def Esirkepov_current( g=int(g), ) - if grid is None: - grid = world["grids"]["center"] - - 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"] - shape_factor = world["shape_factor"] - - 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) - 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 - - shift_x = x0 - old_x0 - shift_y = y0 - old_y0 - shift_z = z0 - old_z0 - # compute how many grid points the particles have shifted in each direction. - - offsets = jnp.asarray([-2, -1, 0, 1, 2], dtype=x0.dtype) - # 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. - - xw, yw, zw = 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, - ) - oxw, oyw, ozw = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), - lambda _: get_second_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), - operand=None, - ) - # compute the particle shape weights at the current and old time steps - - tmp = jnp.zeros_like(xw[0]) - - 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. - - 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, - ) - - 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 - - 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), - 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), - 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), - operand=None, - ) - # compute the current contribution prefactors for each active dimension according to Esirkepov's formula - - Fx = dJx * Wx_ - Fy = dJy * Wy_ - Fz = dJz * Wz_ - - Jx_loc = jnp.zeros_like(Fx) - Jy_loc = jnp.zeros_like(Fy) - Jz_loc = jnp.zeros_like(Fz) - - 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): - 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): - 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: - 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") - 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. - - particle_bc = world.get("particle_boundary_conditions", {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}) - # determine the particle boundary condition type from the global world state - - Jx_ext = ghost_cell_particle_bc_codes_esirkepov( - Jx_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=0 - ) - Jy_ext = ghost_cell_particle_bc_codes_esirkepov( - Jy_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=1 - ) - Jz_ext = ghost_cell_particle_bc_codes_esirkepov( - Jz_ext, particle_bc["x"], particle_bc["y"], particle_bc["z"], component_axis=2 - ) - # Fold the two Esirkepov ghost layers before returning to the solver's - # ordinary one-ghost current arrays. - - Jx = Jx + eliminate_esirkepov_ghost_cells(Jx_ext) - Jy = Jy + eliminate_esirkepov_ghost_cells(Jy_ext) - Jz = Jz + eliminate_esirkepov_ghost_cells(Jz_ext) - - 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 - - return (Jx, Jy, Jz) + raise ValueError("Public Esirkepov_current requires TiledParticles.") def get_3D_esirkepov_weights( diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index 18808d74..5c332e68 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -1,18 +1,5 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -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.direct_deposition_tiled import _direct_J_from_tiled_particles -from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.utils import bilinear_filter, digital_filter def J_from_rhov( @@ -56,207 +43,4 @@ def J_from_rhov( g=int(g), ) - raise ValueError("Public J_from_rhov requires TiledParticles; use _J_from_rhov_flat for flat reference tests.") - - -@partial(jit, static_argnames=("filter",)) -def _J_from_rhov_flat(particles, J, constants, world, grid=None, filter="bilinear"): - """Compute current density (Jx,Jy,Jz) by depositing particle velocities.""" - - if grid is None: - grid = world["grids"]["center"] - - 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"] - shape_factor = world["shape_factor"] - # 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: - charge = species.get_charge() - dq = charge / (dx * dy * dz) - # get the particle information needed for deposition - - 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 - - x, x0, deltax_node, xpts = prepare_particle_axis_stencil( - x, - grid[0], - Nx, - shape_factor, - bc_x, - wind=world["x_wind"], - ghost_cells=True, - ) - # prepare the particle axis stencil for the x-axis, which includes determining the grid points that - # contribute to the deposition and the corresponding weights based on the particle's position and - # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. - - y, y0, deltay_node, ypts = prepare_particle_axis_stencil( - y, - grid[1], - Ny, - shape_factor, - bc_y, - wind=world["y_wind"], - ghost_cells=True, - ) - # prepare the particle axis stencil for the y-axis, which includes determining the grid points that - # contribute to the deposition and the corresponding weights based on the particle's position and - # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. - - z, z0, deltaz_node, zpts = prepare_particle_axis_stencil( - z, - grid[2], - Nz, - shape_factor, - bc_z, - wind=world["z_wind"], - ghost_cells=True, - ) - # prepare the particle axis stencil for the z-axis, which includes determining the grid points that - # contribute to the deposition and the corresponding weights based on the particle's position and - # shape factor. This is done for all three axes (x, y, z) to handle 3D deposition correctly. - - deltax_face = (x - grid[0][0]) - (x0 + 0.5) * dx - deltay_face = (y - grid[1][0]) - (y0 + 0.5) * dy - deltaz_face = (z - grid[2][0]) - (z0 + 0.5) * dz - # compute the distances from the particles to the face-centered grid points. - - x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), - 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 - - 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): - if filter == "bilinear": - J_ = bilinear_filter(J_) - elif filter == "digital": - J_ = digital_filter(J_, constants["alpha"]) - 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) + raise ValueError("Public J_from_rhov requires TiledParticles.") diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 74902964..15c9e7d8 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -17,7 +17,6 @@ from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.solvers.electrostatic_yee import calculate_electrostatic_fields from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields -from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E from PyPIC3D.utils import add_external_fields @@ -275,49 +274,3 @@ def _time_loop_electrostatic_global_reference( fields = (E, B, J, rho, phi, external_fields) return particles, fields - - -@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) -def _time_loop_electrodynamic_global_reference( - particles, - fields, - world, - constants, - curl_func, - J_func, - solver, - relativistic=True, - particle_pusher="boris", -): - """ - Old global electrodynamic loop retained as a reference path for tests. - """ - - E, B, J, rho, phi, external_fields, pml_state = fields - center_grid = world["grids"]["center"] - vertex_grid = world["grids"]["vertex"] - push_E, push_B = add_external_fields(E, B, external_fields) - - for i in range(len(particles)): - particles[i] = particle_push( - particles[i], - push_E, - push_B, - center_grid, - vertex_grid, - world, - constants, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - particles[i].update_position() - - J = J_func(particles, J, constants, world) - E, pml_state = _update_E_global(E, B, J, world, constants, curl_func, pml_state) - B, pml_state = _update_B_global(E, B, world, constants, curl_func, pml_state) - - for i in range(len(particles)): - particles[i].boundary_conditions(world) - - fields = (E, B, J, rho, phi, external_fields, pml_state) - return particles, fields diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 9b072880..80fa9d49 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -5,7 +5,7 @@ from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.deposition.J_from_rhov import J_from_rhov, _J_from_rhov_flat +from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.deposition.direct_deposition_tiled import ( digital_filter_tiled_current_density, ) @@ -77,6 +77,13 @@ def _tile_shape(self, simulation_parameters): simulation_parameters["particle_tile_nz"], ) + def _one_tile_parameters(self, world): + return { + "particle_tile_nx": world["Nx"], + "particle_tile_ny": world["Ny"], + "particle_tile_nz": world["Nz"], + } + def _centered_tiled_particles(self, particles, world, simulation_parameters): """ Build the tiled particle view expected by direct tiled deposition. @@ -101,25 +108,34 @@ def _centered_tiled_particles(self, particles, world, simulation_parameters): return centered_particles, species_config - def _compare_tiled_to_standard(self, particles, world, simulation_parameters): - constants = {"C": 3.0e8, "alpha": 1.0} + def _assembled_tiled_current(self, particles, world, simulation_parameters, constants, filter="none"): tiled_particles, species_config = self._centered_tiled_particles(particles, world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) - J_reference = _J_from_rhov_flat(particles, self._empty_J(world), constants, world, filter="none") J_tiles = J_from_rhov( tiled_particles, species_config, self._empty_J_tiles(world, simulation_parameters), constants, world, - filter="none", + filter=filter, tile_shape=tile_shape, g=int(world["guard_cells"]), ) g = int(world["guard_cells"]) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) + return J_tiles, J_from_tiles + + def _compare_tiled_to_one_tile(self, particles, world, simulation_parameters, filter="none", alpha=1.0): + constants = {"C": 3.0e8, "alpha": alpha} + J_tiles, J_from_tiles = self._assembled_tiled_current( + particles, world, simulation_parameters, constants, filter=filter + ) + _, J_reference = self._assembled_tiled_current( + particles, world, self._one_tile_parameters(world), constants, filter=filter + ) + for tile_component in J_tiles: self.assertEqual(tile_component.ndim, 6) for reference_component, tiled_component in zip(J_reference, J_from_tiles): @@ -184,7 +200,7 @@ def test_tiled_direct_deposition_matches_quadratic_with_two_guard_cells(self): dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) 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) @@ -218,7 +234,7 @@ def test_tiled_direct_deposition_matches_quadratic_saved_style_reduced_axes(self dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): world = self._build_world(Nx=20, Ny=1, Nz=1, dt=0.05) @@ -229,7 +245,6 @@ def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): "particle_tile_ny": 1, "particle_tile_nz": 1, } - constants = {"C": 3.0e8, "alpha": 1.0} species = particle_species( name="bilinear quadratic reduced current", N_particles=8, @@ -251,24 +266,7 @@ def test_tiled_direct_deposition_bilinear_matches_quadratic_reduced_axes(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) - tile_shape = self._tile_shape(simulation_parameters) - - J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") - J_tiles = J_from_rhov( - tiled_particles, - species_config, - self._empty_J_tiles(world, simulation_parameters), - constants, - world, - filter="bilinear", - tile_shape=tile_shape, - g=int(world["guard_cells"]), - ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) - - 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)) + self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="bilinear") def test_tiled_digital_filter_matches_global_digital_filter(self): world = self._build_world(Nx=8, Ny=6, Nz=4) @@ -460,9 +458,12 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): tile_shape=tile_shape, g=int(world["guard_cells"]), ) - J_reference = _J_from_rhov_flat([species], self._empty_J(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( + [species], 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-12, atol=1.0e-12)) @@ -517,7 +518,7 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_dummy_species(self): dt=world["dt"], ) - self._compare_tiled_to_standard([electrons, ions], world, simulation_parameters) + self._compare_tiled_to_one_tile([electrons, ions], world, simulation_parameters) def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(self): world = self._build_world() @@ -552,7 +553,6 @@ def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(sel ) tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) - J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="digital") J_tiles = J_from_rhov( tiled_particles, species_config, @@ -569,6 +569,9 @@ def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(sel world["tile_shape"], num_guard_cells=int(world["guard_cells"]), ) + _, J_reference = self._assembled_tiled_current( + [species], world, self._one_tile_parameters(world), constants, filter="digital" + ) for tile_component in J_tiles: self.assertEqual(tile_component.ndim, 6) @@ -604,24 +607,7 @@ def test_tiled_direct_deposition_digital_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) - tile_shape = self._tile_shape(simulation_parameters) - - J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="digital") - J_tiles = J_from_rhov( - tiled_particles, - species_config, - self._empty_J_tiles(world, simulation_parameters), - constants, - world, - filter="digital", - tile_shape=tile_shape, - g=int(world["guard_cells"]), - ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) - - 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)) + self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="digital", alpha=0.6) def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): world = self._build_world(Nx=8, Ny=6, Nz=4) @@ -652,24 +638,7 @@ def test_tiled_direct_deposition_bilinear_filter_matches_J_from_rhov(self): dz=world["dz"], dt=world["dt"], ) - tiled_particles, species_config = self._centered_tiled_particles([species], world, simulation_parameters) - tile_shape = self._tile_shape(simulation_parameters) - - J_reference = _J_from_rhov_flat([species], self._empty_J(world), constants, world, filter="bilinear") - J_tiles = J_from_rhov( - tiled_particles, - species_config, - self._empty_J_tiles(world, simulation_parameters), - constants, - world, - filter="bilinear", - tile_shape=tile_shape, - g=int(world["guard_cells"]), - ) - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) - - 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)) + self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="bilinear") def test_tiled_direct_deposition_respects_active_mask(self): world = self._build_world() @@ -701,7 +670,7 @@ def test_tiled_direct_deposition_respects_active_mask(self): active_mask=jnp.array([True, False, True, False]), ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) def test_tiled_direct_deposition_periodic_boundary_crossing(self): world = self._build_world(Nx=10, Ny=6, Nz=4, dt=0.0) @@ -732,7 +701,7 @@ def test_tiled_direct_deposition_periodic_boundary_crossing(self): dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) def test_tiled_direct_deposition_matches_J_from_rhov_for_conducting_boundaries(self): world = self._build_world( @@ -781,7 +750,7 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_conducting_boundaries(s dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) def test_tiled_direct_deposition_matches_J_from_rhov_for_mixed_boundaries(self): world = self._build_world( @@ -833,7 +802,7 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_mixed_boundaries(self): dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) def test_tiled_direct_deposition_reduced_dimensions(self): world = self._build_world(Nx=16, Ny=1, Nz=1, dt=0.02) @@ -864,7 +833,7 @@ def test_tiled_direct_deposition_reduced_dimensions(self): dt=world["dt"], ) - self._compare_tiled_to_standard([species], world, simulation_parameters) + self._compare_tiled_to_one_tile([species], world, simulation_parameters) if __name__ == "__main__": diff --git a/tests/evolve_test.py b/tests/evolve_test.py index e742f8dd..2e015c14 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -4,7 +4,6 @@ import jax.numpy as jnp from PyPIC3D.evolve import ( - _time_loop_electrodynamic_global_reference, _time_loop_electrostatic_global_reference, time_loop_electrodynamic, time_loop_electrostatic, @@ -12,13 +11,15 @@ from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push +from PyPIC3D.solvers.yee_tiled import tile_vector_field from PyPIC3D.utils import build_collocated_grid, build_yee_grid jax.config.update("jax_enable_x64", True) -def zero_current(particles, J, constants, world): +def zero_current(particles, species_config, J, constants, world, tile_shape=None, g=None): return tuple(jnp.zeros_like(comp) for comp in J) @@ -266,30 +267,54 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field dt=world["dt"], ) ] - fields = (E, B, J, rho, phi, external_fields, None) + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + world["guard_cells"] = 2 + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + external_E_tiles = tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])) + external_B_tiles = tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])) + fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + rho, + phi, + (external_E_tiles, external_B_tiles), + None, + ) - particles, fields = _time_loop_electrodynamic_global_reference( - particles, + tiled_particles, fields = time_loop_electrodynamic( + tiled_particles, + species_config, fields, world, constants, unused_curl, zero_current, - solver="fdtd", + solver="electrodynamic_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), 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() + 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, :]] self.assertIsNone(pml_state) + self.assertFalse(bool(overflow)) 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])) + self.assertTrue(jnp.allclose(external_after[0][0], external_E_tiles[0])) def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): world = { @@ -340,24 +365,47 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): dt=world["dt"], ) ] - fields = (E, B, J, rho, phi, external_fields, None) + tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + world["guard_cells"] = 2 + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])), + ), + None, + ) - particles, _ = _time_loop_electrodynamic_global_reference( - particles, + tiled_particles, _ = time_loop_electrodynamic( + tiled_particles, + species_config, fields, world, constants, unused_curl, zero_current, - solver="fdtd", + solver="electrodynamic_yee", + tile_shape=tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) - x, _, _ = particles[0].get_forward_position() + active = tiled_particles.active[:, :, :, 0, :].reshape(-1) + x = tiled_particles.x[:, :, :, 0, :, 0].reshape(-1) self.assertEqual(x.shape[0], 1) - self.assertTrue(jnp.array_equal(particles[0].get_active_mask(), jnp.array([False]))) + self.assertTrue(jnp.array_equal(active, jnp.array([False]))) if __name__ == "__main__": diff --git a/tests/particle_test.py b/tests/particle_test.py index 40068909..34aded4d 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -12,11 +12,8 @@ from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.deposition.J_from_rhov import _J_from_rhov_flat -from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.utils import build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -455,65 +452,7 @@ def test_particle_species(self): 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, - '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} - # 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_flat([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): + def test_inactive_particles_do_not_deposit_rho(self): Nx, Ny, Nz = 10, 10, 10 dx = self.x_wind / Nx dy = self.y_wind / Ny @@ -560,20 +499,8 @@ def test_inactive_particles_do_not_deposit_rho_or_current(self): constants = {"C": 3e8, "alpha": 1.0} rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) - J = _J_from_rhov_flat( - [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]) @@ -742,64 +669,6 @@ def test_rho_periodic_boundary_crossing(self): 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, - "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} - - 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_flat( - [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 @@ -902,100 +771,5 @@ def test_filtered_rho_periodic_updates_ghost_cells(self): 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, - "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} - - 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, - ) - - # 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_flat([species], rho, world, constants) - # compute rho - species.update_position() - rho = _compute_rho_flat([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/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 60e945f3..1d762974 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -138,22 +138,17 @@ def _tile_shape_for_world(self, world): 2 if world["Nz"] > 1 else 1, ) - 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) + def _one_tile_shape_for_world(self, world): + return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + + def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape): simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } g = int(world["guard_cells"]) - - old_species = self._species_from_arrays(world, x_old, u) - new_species = self._species_from_arrays(world, x_old + u * world["dt"], u) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - - J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -165,6 +160,22 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No ) 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) + + old_species = self._species_from_arrays(world, x_old, u) + J_tiles, J_from_tiles = self._assembled_esirkepov_current(world, old_species, constants, tile_shape) + _, J_reference = self._assembled_esirkepov_current( + world, + old_species, + 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), @@ -318,11 +329,8 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): } x_old = jnp.array([-1.10, -0.10, 1.05]) old_species = self._species(world, x_old) - x_new = x_old + old_species.v1 * world["dt"] - new_species = self._species(world, x_new) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) g = int(world["guard_cells"]) J_tiles = Esirkepov_current( tiled_particles, @@ -334,6 +342,12 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): g=g, ) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) + _, J_reference = self._assembled_esirkepov_current( + world, + old_species, + constants, + self._one_tile_shape_for_world(world), + ) for reference_component, tiled_component in zip(J_reference, J_from_tiles): self.assertTrue( @@ -353,11 +367,8 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre } x_old = jnp.array([-1.10, -0.10, 1.05]) old_species = self._species(world, x_old) - x_new = x_old + old_species.v1 * world["dt"] - new_species = self._species(world, x_new) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - J_reference = Esirkepov_current([new_species], self._empty_J(world), constants, world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -373,6 +384,12 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre world["tile_shape"], num_guard_cells=int(world["guard_cells"]), ) + _, J_reference = self._assembled_esirkepov_current( + world, + old_species, + constants, + self._one_tile_shape_for_world(world), + ) for tile_component in J_tiles: self.assertEqual(tile_component.ndim, 6) @@ -795,12 +812,12 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s 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))), ) - new_species = self._species_from_arrays( + _, reference_J = self._assembled_esirkepov_current( world, - jnp.column_stack((jnp.asarray(x_initial + vx_initial * world["dt"]), jnp.zeros(4), jnp.zeros(4))), - jnp.column_stack((jnp.asarray(vx_initial), jnp.zeros(4), jnp.zeros(4))), + old_species, + constants, + self._one_tile_shape_for_world(world), ) - reference_J = Esirkepov_current([new_species], self._empty_J(world), constants, world) tiled_J = assemble_tiled_vector_field( fields[2], world, diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index e36fb0fd..81c9b433 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -4,6 +4,7 @@ import jax.numpy as jnp from PyPIC3D import evolve +from PyPIC3D.deposition import J_from_rhov from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled from PyPIC3D.utils import build_yee_grid @@ -98,6 +99,10 @@ def test_electrostatic_global_poisson_bridge_stays_explicit(self): self.assertIn("assemble_tiled_scalar_field", source) self.assertIn("tile_scalar_field", source) + def test_flat_current_and_electrodynamic_references_are_removed(self): + self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) + self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 70dc4120..e18a79f4 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -8,9 +8,8 @@ import numpy as np import toml -from PyPIC3D.evolve import _time_loop_electrodynamic_global_reference from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.J_from_rhov import J_from_rhov, _J_from_rhov_flat +from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.initialization import initialize_simulation, validate_field_solver @@ -118,6 +117,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (4, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) Nt = 1500 x0 = jnp.linspace( @@ -153,15 +153,13 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), ] - reference_particles = [ + reference_species = [ particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), ] E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_fields = (E, B, J, rho, phi, external_fields, pml_state) - simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -187,6 +185,14 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): ), None, ) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + reference_species, + (E, B, J, rho, phi, external_fields, pml_state), + world, + reference_tile_shape, + current_calculation, + ) + reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) self._assert_tiled_state_matches_standard( reference_particles, @@ -194,11 +200,24 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): tiled_particles, tiled_fields, world, + reference_tile_shape, tile_shape, step=0, ) - return reference_particles, reference_fields, tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt + return ( + reference_particles, + reference_species_config, + reference_fields, + reference_tile_shape, + tiled_particles, + species_config, + tiled_fields, + world, + constants, + tile_shape, + Nt, + ) def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): current_slots = tiled_particles.active.shape[-1] @@ -220,14 +239,6 @@ def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): active=active, ) - def _standard_species_rows(self, particles, species_index): - species = particles[species_index] - active = species.get_active_mask() - x = jnp.stack(species.get_forward_position(), axis=1)[active] - u = jnp.stack(species.get_velocity(), axis=1)[active] - order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - def _tiled_species_rows(self, tiled_particles, species_index): active = tiled_particles.active[:, :, :, species_index, :] x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] @@ -235,6 +246,33 @@ def _tiled_species_rows(self, tiled_particles, species_index): order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] + def _build_tiled_state(self, particles, fields, world, tile_shape, current_calculation=CURRENT_J_FROM_RHOV): + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + E, B, J, rho, phi, external_fields, pml_state = fields + g = int(world["guard_cells"]) + if current_calculation == CURRENT_ESIRKEPOV: + J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) + else: + J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + tiled_fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + J_tiles, + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), + ), + pml_state, + ) + return tiled_particles, species_config, tiled_fields + def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): diff = jnp.max(jnp.abs(reference_component - tiled_component)) @@ -250,14 +288,16 @@ def _assert_tiled_state_matches_standard( tiled_particles, tiled_fields, world, + reference_tile_shape, tile_shape, step, ): if len(tiled_fields) > 7: self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - for species_index in range(len(reference_particles)): - reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + species_count = reference_particles.active.shape[3] + for species_index in range(species_count): + reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) self.assertTrue( jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), @@ -272,28 +312,38 @@ def _assert_tiled_state_matches_standard( E_tiles, B_tiles, J_tiles, *_ = tiled_fields num_guard_cells = int(world.get("guard_cells", 1)) self._assert_vector_fields_close( - reference_E, + assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "E", ) self._assert_vector_fields_close( - reference_B, + assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "B", ) self._assert_vector_fields_close( - reference_J, + assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), step, "J", ) - def _max_tiled_state_difference(self, reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape): + def _max_tiled_state_difference( + self, + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + world, + tile_shape, + ): differences = {} - for species_index in range(len(reference_particles)): - reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + species_count = reference_particles.active.shape[3] + for species_index in range(species_count): + reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) differences[f"x{species_index}"] = float(jnp.max(jnp.abs(tiled_x - reference_x))) differences[f"u{species_index}"] = float(jnp.max(jnp.abs(tiled_u - reference_u))) @@ -302,9 +352,21 @@ def _max_tiled_state_difference(self, reference_particles, reference_fields, til E_tiles, B_tiles, J_tiles, *_ = tiled_fields num_guard_cells = int(world.get("guard_cells", 1)) for name, reference_field, tiled_field in ( - ("E", reference_E, assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), - ("B", reference_B, assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), - ("J", reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells)), + ( + "E", + assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), + ( + "B", + assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), + ( + "J", + assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), ): differences[name] = max(float(jnp.max(jnp.abs(tiled - reference))) for reference, tiled in zip(reference_field, tiled_field)) @@ -334,15 +396,30 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): rho = seed phi = seed external_fields = ((seed, seed, seed), (seed, seed, seed)) - - _, reference_fields = _time_loop_electrodynamic_global_reference( + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_fields = ( + tile_vector_field(E, world, reference_tile_shape), + tile_vector_field(B, world, reference_tile_shape), + tile_vector_field(J, world, reference_tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], world, reference_tile_shape), + tile_vector_field(external_fields[1], world, reference_tile_shape), + ), + None, + ) + _, reference_fields = time_loop_electrodynamic( [], - (E, B, J, rho, phi, external_fields, None), + None, + reference_fields, world, constants, unused_curl, - J_func=lambda particles, J, constants, world: J, - solver="fdtd", + J_func=lambda particles, J, constants, world, tile_shape=None, g=None: J, + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -359,7 +436,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic( + tiled_particles, tiled_fields = time_loop_electrodynamic( [], None, tiled_fields, @@ -377,13 +454,13 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): reference_E, reference_B, *_ = reference_fields E_tiles, B_tiles, *_ = tiled_fields self._assert_vector_fields_close( - reference_E, + assemble_tiled_vector_field(reference_E, world, reference_tile_shape), assemble_tiled_vector_field(E_tiles, world, tile_shape), 1, "E", ) self._assert_vector_fields_close( - reference_B, + assemble_tiled_vector_field(reference_B, world, reference_tile_shape), assemble_tiled_vector_field(B_tiles, world, tile_shape), 1, "B", @@ -686,6 +763,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): world = self._build_world() constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -695,17 +773,11 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): species = self._species(world) reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, - constants, - unused_curl, - J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="none"), - solver="fdtd", - relativistic=False, - particle_pusher="boris", + reference_tile_shape, ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) @@ -721,6 +793,20 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): ), None, ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + world, + constants, + unused_curl, + J_func=None, + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, @@ -736,26 +822,26 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): particle_pusher="boris", ) - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - E_from_tiles = assemble_tiled_vector_field(E_tiles, world, tile_shape) - B_from_tiles = assemble_tiled_vector_field(B_tiles, world, tile_shape) - reference_E, reference_B, reference_J, *_ = reference_fields - - flat_x = jnp.stack(reference_particles[0].get_forward_position(), axis=1) - flat_u = jnp.stack(reference_particles[0].get_velocity(), axis=1) - flat_order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) - tiled_x, tiled_u = self._active_rows(tiled_particles) - - self.assertTrue(jnp.allclose(tiled_x, flat_x[flat_order], rtol=1.0e-12, atol=1.0e-12)) - self.assertTrue(jnp.allclose(tiled_u, flat_u[flat_order], rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_E, E_from_tiles): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_B, B_from_tiles): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape)): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - - reference_energy = compute_energy(reference_particles, reference_E, reference_B, world, constants) + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + reference_E, reference_B, *_ = reference_fields + E_tiles, B_tiles, *_ = tiled_fields + reference_energy = compute_energy( + reference_particles, + reference_E, + reference_B, + world, + constants, + species_config=reference_species_config, + ) tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants, species_config=species_config) for reference, tiled in zip(reference_energy, tiled_energy): self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) @@ -766,6 +852,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): ) constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -775,17 +862,11 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): species = self._species(world) reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - _, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, - constants, - unused_curl, - J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="none"), - solver="fdtd", - relativistic=False, - particle_pusher="boris", + reference_tile_shape, ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) @@ -801,7 +882,21 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic( + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + world, + constants, + unused_curl, + J_func=None, + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -816,20 +911,22 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): particle_pusher="boris", ) - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - reference_E, reference_B, reference_J, *_ = reference_fields - - for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, tile_shape)): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, tile_shape)): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, tile_shape)): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) def test_tiled_yee_step_uses_digital_filtered_current(self): world = self._build_world() constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -839,17 +936,11 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): species = self._species(world) reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - _, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, - constants, - unused_curl, - J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="digital"), - solver="fdtd", - relativistic=False, - particle_pusher="boris", + reference_tile_shape, ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) @@ -865,7 +956,21 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic( + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + world, + constants, + unused_curl, + J_func=functools.partial(J_from_rhov, filter="digital"), + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -880,17 +985,22 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): particle_pusher="boris", ) - _, _, reference_J, *_ = reference_fields - _, _, J_tiles, *_ = tiled_fields - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) - - for reference, tiled in zip(reference_J, J_from_tiles): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) def test_tiled_yee_step_uses_bilinear_filtered_current(self): world = self._build_world() constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -900,17 +1010,11 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): species = self._species(world) reference_species = self._species(world) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - _, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, - constants, - unused_curl, - J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="bilinear"), - solver="fdtd", - relativistic=False, - particle_pusher="boris", + reference_tile_shape, ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) @@ -926,7 +1030,21 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): ), None, ) - _, tiled_fields = time_loop_electrodynamic( + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + world, + constants, + unused_curl, + J_func=functools.partial(J_from_rhov, filter="bilinear"), + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -941,12 +1059,16 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): particle_pusher="boris", ) - _, _, reference_J, *_ = reference_fields - _, _, J_tiles, *_ = tiled_fields - J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape) - - for reference, tiled in zip(reference_J, J_from_tiles): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_standard(self): world = self._build_world() @@ -954,6 +1076,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand world["guard_cells"] = 2 constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -1004,16 +1127,11 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ) E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - _, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), world, - constants, - unused_curl, - J_func=lambda particles, J, constants, world: _J_from_rhov_flat(particles, J, constants, world, filter="bilinear"), - solver="fdtd", - relativistic=False, - particle_pusher="boris", + reference_tile_shape, ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) @@ -1029,7 +1147,21 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ), None, ) - _, tiled_fields = time_loop_electrodynamic( + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + world, + constants, + unused_curl, + J_func=functools.partial(J_from_rhov, filter="bilinear"), + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( tiled_particles, species_config, tiled_fields, @@ -1045,11 +1177,12 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand ) self._assert_tiled_state_matches_standard( - [reference_species], + reference_particles, reference_fields, tiled_particles, tiled_fields, world, + reference_tile_shape, tile_shape, 1, ) @@ -1061,7 +1194,9 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): ( reference_particles, + reference_species_config, reference_fields, + reference_tile_shape, tiled_particles, species_config, tiled_fields, @@ -1071,21 +1206,23 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): Nt, ) = self._long_two_stream_state(CURRENT_J_FROM_RHOV) - standard_J = functools.partial(_J_from_rhov_flat, filter="none") tiled_loop = jax.jit( time_loop_electrodynamic, static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), ) for t in range(Nt): - reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_fields = tiled_loop( reference_particles, + reference_species_config, reference_fields, world, constants, unused_curl, - J_func=standard_J, - solver="fdtd", + J_func=None, + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1105,7 +1242,13 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): ) differences = self._max_tiled_state_difference( - reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + world, + tile_shape, ) if t == 0: self._assert_tiled_state_matches_standard( @@ -1114,6 +1257,7 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): tiled_particles, tiled_fields, world, + reference_tile_shape, tile_shape, t + 1, ) @@ -1128,7 +1272,9 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): ( reference_particles, + reference_species_config, reference_fields, + reference_tile_shape, tiled_particles, species_config, tiled_fields, @@ -1144,14 +1290,17 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): ) for t in range(Nt): - reference_particles, reference_fields = _time_loop_electrodynamic_global_reference( + reference_particles, reference_fields = tiled_loop( reference_particles, + reference_species_config, reference_fields, world, constants, unused_curl, J_func=Esirkepov_current, - solver="fdtd", + solver="electrodynamic_yee", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1171,7 +1320,13 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): ) differences = self._max_tiled_state_difference( - reference_particles, reference_fields, tiled_particles, tiled_fields, world, tile_shape + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + world, + tile_shape, ) if t == 0: self._assert_tiled_state_matches_standard( @@ -1180,6 +1335,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): tiled_particles, tiled_fields, world, + reference_tile_shape, tile_shape, t + 1, ) From 8c6d402297b685f996a5a193e735e30da2b57ea7 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 1 Jul 2026 17:26:35 -0500 Subject: [PATCH 33/83] remove global tile reference methods --- PyPIC3D/evolve.py | 44 ----- tests/evolve_test.py | 192 +--------------------- tests/tiled_electrostatic_test.py | 226 ++++++++++++++++---------- tests/tiled_refactor_contract_test.py | 1 + 4 files changed, 139 insertions(+), 324 deletions(-) diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 15c9e7d8..85f11df8 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -13,9 +13,7 @@ refresh_tiled_particle_tiles, update_tiled_particle_positions, ) -from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.electrostatic_yee import calculate_electrostatic_fields from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E from PyPIC3D.utils import add_external_fields @@ -232,45 +230,3 @@ def time_loop_electrostatic( # pack the tiled field state return particles, fields - - -@partial(jit, static_argnames=("curl_func", "J_func", "solver", "relativistic", "particle_pusher")) -def _time_loop_electrostatic_global_reference( - particles, - fields, - world, - constants, - curl_func, - J_func, - solver, - relativistic=True, - particle_pusher="boris", -): - """ - Old global electrostatic loop retained as a reference path for tests. - """ - - E, B, J, rho, phi, external_fields = fields - center_grid = world["grids"]["center"] - vertex_grid = world["grids"]["vertex"] - push_E, push_B = add_external_fields(E, B, external_fields) - - for i in range(len(particles)): - particles[i] = particle_push( - particles[i], - push_E, - push_B, - center_grid, - vertex_grid, - world, - constants, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - particles[i].update_position() - particles[i].boundary_conditions(world) - - E, phi, rho = calculate_electrostatic_fields(world, particles, constants, rho, phi, solver, "periodic") - - fields = (E, B, J, rho, phi, external_fields) - return particles, fields diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 2e015c14..e076d142 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -4,17 +4,14 @@ import jax.numpy as jnp from PyPIC3D.evolve import ( - _time_loop_electrostatic_global_reference, time_loop_electrodynamic, time_loop_electrostatic, ) -from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.solvers.yee_tiled import tile_vector_field -from PyPIC3D.utils import build_collocated_grid, build_yee_grid +from PyPIC3D.utils import build_yee_grid jax.config.update("jax_enable_x64", True) @@ -32,193 +29,6 @@ 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") - def _electrostatic_grid_order_species(self, world, charge): - return particle_species( - name=f"charge {charge}", - N_particles=1, - charge=charge, - mass=1.0, - T=0.0, - weight=1.0, - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([0.07]), - 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"], - dt=world["dt"], - ) - - def test_electrostatic_push_uses_same_grid_order_as_electrodynamic(self): - world = { - "Nx": 4, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} - - E, B, J, phi, rho = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) - Ex = jnp.broadcast_to(vertex_grid[0][:, None, None], E[0].shape) - E = (Ex, E[1], E[2]) - external_fields = ( - tuple(jnp.zeros_like(comp) for comp in E), - tuple(jnp.zeros_like(comp) for comp in B), - ) - - pushed_particles = [ - self._electrostatic_grid_order_species(world, 1.0), - self._electrostatic_grid_order_species(world, -1.0), - ] - reference_particles = [ - self._electrostatic_grid_order_species(world, 1.0), - self._electrostatic_grid_order_species(world, -1.0), - ] - - pushed_particles, _ = _time_loop_electrostatic_global_reference( - pushed_particles, - (E, B, J, rho, phi, external_fields), - world, - constants, - unused_curl, - J_func=None, - solver="fdtd", - relativistic=False, - particle_pusher="boris", - ) - - for i in range(len(reference_particles)): - reference_particles[i] = particle_push( - reference_particles[i], - E, - B, - center_grid, - vertex_grid, - world, - constants, - relativistic=False, - particle_pusher="boris", - ) - reference_particles[i].update_position() - reference_particles[i].boundary_conditions(world) - - for pushed, reference in zip(pushed_particles, reference_particles): - pushed_vx, pushed_vy, pushed_vz = pushed.get_velocity() - reference_vx, reference_vy, reference_vz = reference.get_velocity() - self.assertTrue(jnp.allclose(pushed_vx, reference_vx, rtol=1.0e-12, atol=1.0e-12)) - self.assertTrue(jnp.allclose(pushed_vy, reference_vy, rtol=1.0e-12, atol=1.0e-12)) - self.assertTrue(jnp.allclose(pushed_vz, reference_vz, rtol=1.0e-12, atol=1.0e-12)) - - def test_electrostatic_rho_uses_boundary_applied_particle_positions(self): - world = { - "Nx": 16, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.002, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_collocated_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.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="right_crossing_charge", - N_particles=1, - charge=1.0, - mass=1.0, - T=0.0, - weight=1.0, - v1=jnp.array([0.10]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - x1=jnp.array([1.99995]), - 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"], - dt=world["dt"], - ), - particle_species( - name="neutralizing_charge", - N_particles=1, - charge=-1.0, - mass=1.0, - T=0.0, - weight=1.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"], - dt=world["dt"], - ), - ] - - particles, fields = _time_loop_electrostatic_global_reference( - particles, - (E, B, J, rho, phi, external_fields), - world, - constants, - unused_curl, - J_func=None, - solver="fdtd", - relativistic=False, - particle_pusher="boris", - ) - _, _, _, rho_after, _, _ = fields - recomputed_rho = _compute_rho_flat(particles, jnp.zeros_like(rho_after), world, constants) - - self.assertTrue( - jnp.allclose( - rho_after[1:-1, 1:-1, 1:-1], - recomputed_rho[1:-1, 1:-1, 1:-1], - rtol=1.0e-12, - atol=1.0e-12, - ) - ) - def test_external_electric_field_pushes_particles_without_evolving_maxwell_fields(self): world = { "Nx": 3, diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 52b98c8d..024b09ea 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -11,7 +11,6 @@ from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.evolve import time_loop_electrostatic -from PyPIC3D.evolve import _time_loop_electrostatic_global_reference from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -64,6 +63,16 @@ def _simulation_parameters(self): "particle_tile_nz": 1, } + 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 _one_tile_shape(self, world): + return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + def _empty_fields(self, world): shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) zero = jnp.zeros(shape) @@ -140,6 +149,32 @@ def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): active=jnp.pad(tiled_particles.active, scalar_pad_width), ) + def _build_tiled_state(self, particles, fields, world, tile_shape, min_slots=None): + tiled_particles, species_config = to_tiled_particles( + particles, + world, + self._simulation_parameters_for_tile_shape(tile_shape), + ) + if min_slots is not None: + tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots) + + E, B, J, rho, phi, external_fields = fields + g = int(world["guard_cells"]) + tiled_fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + tile_vector_field(J, world, tile_shape, num_guard_cells=g), + tile_scalar_field(rho, world, tile_shape, num_guard_cells=g), + tile_scalar_field(phi, world, tile_shape, num_guard_cells=g), + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), + ), + None, + ) + + return tiled_particles, species_config, tiled_fields + def _long_two_stream_species(self, world): x0 = jnp.linspace( -0.5 * world["x_wind"], @@ -196,39 +231,30 @@ def _long_two_stream_state(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (4, 1, 1) + reference_tile_shape = self._one_tile_shape(world) world["tile_shape"] = tile_shape Nt = 1500 E, B, J, rho, phi, external_fields = self._empty_fields(world) - reference_fields = (E, B, J, rho, phi, external_fields) - reference_particles = self._long_two_stream_species(world) - - tiled_particles, species_config = to_tiled_particles( + fields = (E, B, J, rho, phi, external_fields) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( self._long_two_stream_species(world), + fields, world, - { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - }, + reference_tile_shape, ) - tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) - tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), - tile_scalar_field(rho, world, tile_shape), - tile_scalar_field(phi, world, tile_shape), - ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), - ), - None, + tiled_particles, species_config, tiled_fields = self._build_tiled_state( + self._long_two_stream_species(world), + fields, + world, + tile_shape, + min_slots=12, ) - self._assert_long_tiled_state_matches_standard( + self._assert_tiled_state_matches_reference( reference_particles, reference_fields, + reference_tile_shape, tiled_particles, tiled_fields, world, @@ -236,15 +262,19 @@ def _long_two_stream_state(self): step=0, ) - return reference_particles, reference_fields, tiled_particles, tiled_fields, world, constants, tile_shape, Nt - - def _standard_species_rows(self, particles, species_index): - species = particles[species_index] - active = species.get_active_mask() - x = jnp.stack(species.get_forward_position(), axis=1)[active] - u = jnp.stack(species.get_velocity(), axis=1)[active] - order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] + return ( + reference_particles, + reference_species_config, + reference_fields, + reference_tile_shape, + tiled_particles, + species_config, + tiled_fields, + world, + constants, + tile_shape, + Nt, + ) def _tiled_species_rows(self, tiled_particles, species_index): active = tiled_particles.active[:, :, :, species_index, :] @@ -261,10 +291,11 @@ def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): f"step {step}: {name}{component_name} max abs diff {diff}", ) - def _assert_long_tiled_state_matches_standard( + def _assert_tiled_state_matches_reference( self, reference_particles, reference_fields, + reference_tile_shape, tiled_particles, tiled_fields, world, @@ -274,8 +305,9 @@ def _assert_long_tiled_state_matches_standard( if len(tiled_fields) > 7: self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - for species_index in range(len(reference_particles)): - reference_x, reference_u = self._standard_species_rows(reference_particles, species_index) + species_count = reference_particles.active.shape[3] + for species_index in range(species_count): + reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) self.assertTrue( jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), @@ -288,38 +320,51 @@ def _assert_long_tiled_state_matches_standard( reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + g = int(world["guard_cells"]) self._assert_vector_fields_close( - reference_E, - assemble_tiled_vector_field(E_tiles, world, tile_shape), + assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=g), + assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=g), step, "E", ) self._assert_vector_fields_close( - reference_B, - assemble_tiled_vector_field(B_tiles, world, tile_shape), + assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=g), + assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=g), step, "B", ) self._assert_vector_fields_close( - reference_J, - assemble_tiled_vector_field(J_tiles, world, tile_shape), + assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=g), + assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g), step, "J", ) - rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, tile_shape) - phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, tile_shape) + reference_rho_from_tiles = assemble_tiled_scalar_field( + reference_rho, + world, + reference_tile_shape, + num_guard_cells=g, + ) + reference_phi_from_tiles = assemble_tiled_scalar_field( + reference_phi, + world, + reference_tile_shape, + num_guard_cells=g, + ) + rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) + phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) self.assertTrue( jnp.allclose( rho_from_tiles[1:-1, 1:-1, 1:-1], - reference_rho[1:-1, 1:-1, 1:-1], + reference_rho_from_tiles[1:-1, 1:-1, 1:-1], rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL, ), - f"step {step}: rho max abs diff {jnp.max(jnp.abs(rho_from_tiles - reference_rho))}", + f"step {step}: rho max abs diff {jnp.max(jnp.abs(rho_from_tiles - reference_rho_from_tiles))}", ) tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] - reference_phi_interior = reference_phi[1:-1, 1:-1, 1:-1] + reference_phi_interior = reference_phi_from_tiles[1:-1, 1:-1, 1:-1] tiled_phi_interior = tiled_phi_interior - jnp.mean(tiled_phi_interior) reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) self.assertTrue( @@ -367,37 +412,40 @@ def test_tiled_rho_assembles_to_global_rho(self): ) ) - def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): + def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): world = self._build_world(shape_factor=1) constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} E, B, J, rho, phi, external_fields = self._empty_fields(world) + fields = (E, B, J, rho, phi, external_fields) + reference_tile_shape = self._one_tile_shape(world) - reference_particles, reference_fields = _time_loop_electrostatic_global_reference( + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( self._neutral_species(world), - (E, B, J, rho, phi, external_fields), + fields, + world, + reference_tile_shape, + ) + tiled_particles, species_config, tiled_fields = self._build_tiled_state( + self._neutral_species(world), + fields, + world, + world["tile_shape"], + ) + + reference_particles, reference_fields = time_loop_electrostatic( + reference_particles, + reference_species_config, + reference_fields, world, constants, unused_curl, J_func=None, - solver="fdtd", + solver="electrostatic", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) - - tiled_particles, species_config = to_tiled_particles(self._neutral_species(world), world, self._simulation_parameters()) - 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, - ) - tiled_particles, tiled_fields = time_loop_electrostatic( tiled_particles, species_config, @@ -406,41 +454,36 @@ def test_tiled_electrostatic_step_matches_global_step_after_assembly(self): constants, unused_curl, J_func=None, - solver="electrodynamic_yee", + solver="electrostatic", tile_shape=world["tile_shape"], g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields - - for reference, tiled in zip(reference_E, assemble_tiled_vector_field(E_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-8, atol=1.0e-8)) - for reference, tiled in zip(reference_B, assemble_tiled_vector_field(B_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - for reference, tiled in zip(reference_J, assemble_tiled_vector_field(J_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"]))): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - - rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - self.assertTrue(jnp.allclose(rho_from_tiles[1:-1, 1:-1, 1:-1], reference_rho[1:-1, 1:-1, 1:-1], rtol=1.0e-12, atol=1.0e-12)) - tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] - reference_phi_interior = reference_phi[1:-1, 1:-1, 1:-1] - tiled_phi_interior = tiled_phi_interior - jnp.mean(tiled_phi_interior) - reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) - self.assertTrue(jnp.allclose(tiled_phi_interior, reference_phi_interior, rtol=1.0e-8, atol=1.0e-8)) + self._assert_tiled_state_matches_reference( + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + world, + world["tile_shape"], + step=1, + ) @unittest.skipUnless( os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step tiled electrostatic comparison.", ) - def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self): + def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(self): ( reference_particles, + reference_species_config, reference_fields, + reference_tile_shape, tiled_particles, + species_config, tiled_fields, world, constants, @@ -454,34 +497,39 @@ def test_long_two_stream_tiled_electrostatic_matches_global_step_every_step(self ) for step in range(1, Nt + 1): - reference_particles, reference_fields = _time_loop_electrostatic_global_reference( + reference_particles, reference_fields = tiled_loop( reference_particles, + reference_species_config, reference_fields, world, constants, unused_curl, J_func=None, - solver="fdtd", + solver="electrostatic", + tile_shape=reference_tile_shape, + g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) tiled_particles, tiled_fields = tiled_loop( tiled_particles, + species_config, tiled_fields, world, constants, unused_curl, J_func=None, - solver="electrodynamic_yee", + solver="electrostatic", tile_shape=tile_shape, g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) - self._assert_long_tiled_state_matches_standard( + self._assert_tiled_state_matches_reference( reference_particles, reference_fields, + reference_tile_shape, tiled_particles, tiled_fields, world, diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 81c9b433..cd23e8a2 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -102,6 +102,7 @@ def test_electrostatic_global_poisson_bridge_stays_explicit(self): def test_flat_current_and_electrodynamic_references_are_removed(self): self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) + self.assertFalse(hasattr(evolve, "_time_loop_electrostatic_global_reference")) if __name__ == "__main__": From d1e3e46cf439f850be6a2c353af7882c6ae3941e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 3 Jul 2026 08:33:22 -0500 Subject: [PATCH 34/83] refactoring direct deposition current --- PyPIC3D/__main__.py | 34 +- PyPIC3D/deposition/J_from_rhov.py | 270 +++++++++++++--- PyPIC3D/deposition/direct_deposition_tiled.py | 295 ------------------ PyPIC3D/deposition/rho.py | 2 +- PyPIC3D/deposition/rho_tiled.py | 9 +- PyPIC3D/diagnostics/tiled_benchmarking.py | 2 - PyPIC3D/evolve.py | 5 +- PyPIC3D/solvers/electrostatic_yee.py | 2 +- PyPIC3D/solvers/first_order_yee.py | 2 +- PyPIC3D/solvers/yee_tiled.py | 46 +-- PyPIC3D/utilities/filters.py | 155 +++++++++ PyPIC3D/utils.py | 67 ---- tests/direct_deposition_tiled_test.py | 61 ++-- tests/evolve_test.py | 2 +- tests/tiled_refactor_contract_test.py | 6 + tests/tiled_yee_integration_test.py | 232 ++++++++------ 16 files changed, 617 insertions(+), 573 deletions(-) delete mode 100644 PyPIC3D/deposition/direct_deposition_tiled.py create mode 100644 PyPIC3D/utilities/filters.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 0e638a39..49cfdc7d 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -73,9 +73,36 @@ def run_PyPIC3D(config_file): tile_shape = tuple(int(width) for width in world["tile_shape"]) g = int(world["guard_cells"]) particle_species_names = simulation_parameters.get("particle_species_names") + + def loop_with_static_world( + particles, + species_config, + fields, + constants, + curl_func, + J_func, + solver, + relativistic=True, + particle_pusher="boris", + ): + return loop( + particles, + species_config, + fields, + world, + constants, + curl_func, + J_func, + solver, + tile_shape=tile_shape, + g=g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + jit_loop = jax.jit( - loop, - static_argnames=('curl_func', 'J_func', 'solver', 'tile_shape', 'g', 'relativistic', 'particle_pusher'), + loop_with_static_world, + static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher'), ) E, B, J, rho, phi, external_fields, *rest = fields @@ -178,13 +205,10 @@ def run_PyPIC3D(config_file): particles, species_config, fields, - world, constants, curl_func, J_func, solver, - tile_shape=tile_shape, - g=g, relativistic=relativistic, particle_pusher=particle_pusher, ) diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index 5c332e68..00ed27d7 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -1,46 +1,240 @@ -from PyPIC3D.deposition.direct_deposition_tiled import _direct_J_from_tiled_particles -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.tiled_particles import TiledParticles, SpeciesConfig + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + collapse_axis_stencil, + prepare_particle_axis_stencil, +) + +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.solvers.yee_tiled 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 + +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) def J_from_rhov( particles, - J, - constants, - world, - grid=None, + species_config, + J=None, + constants=None, + world=None, filter="bilinear", - species_config=None, - tile_shape=None, - g=None, ): - """Compute current density from velocity-weighted particle charge.""" - - if isinstance(particles, TiledParticles): - if species_config is None: - species_config = J - J = constants - constants = world - world = grid - grid = None - # Tiled runtime callers pass ``species_config`` as the second positional - # argument so the public current name can be used without global current - # assembly in the deposition path. - - if tile_shape is None: - tile_shape = tuple(int(width) for width in world["tile_shape"]) - if g is None: - g = int(world["guard_cells"]) - - return _direct_J_from_tiled_particles( - particles, - species_config, - J, - constants, - world, - grid=grid, - filter=filter, - tile_shape=tile_shape, - g=int(g), + """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"] + # 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, + x_grid, + local_Nx, + shape_factor, + local_bc, + wind=tile_nx * dx, + ghost_cells=True, + ) + y, y0, deltay_node, ypts = prepare_particle_axis_stencil( + y, + y_grid, + local_Ny, + shape_factor, + local_bc, + wind=tile_ny * dy, + ghost_cells=True, + ) + z, z0, deltaz_node, zpts = prepare_particle_axis_stencil( + z, + z_grid, + local_Nz, + shape_factor, + local_bc, + wind=tile_nz * dz, + ghost_cells=True, ) + # prepare the particle positions and compute the stencil points and deltas for each axis + + 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, + lambda _: get_first_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), + lambda _: get_second_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), + operand=None, + ) + 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 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) + x_weights_node = jnp.asarray(x_weights_node) + y_weights_node = jnp.asarray(y_weights_node) + z_weights_node = jnp.asarray(z_weights_node) + 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 - raise ValueError("Public J_from_rhov requires TiledParticles.") + return J diff --git a/PyPIC3D/deposition/direct_deposition_tiled.py b/PyPIC3D/deposition/direct_deposition_tiled.py deleted file mode 100644 index 69de8f21..00000000 --- a/PyPIC3D/deposition/direct_deposition_tiled.py +++ /dev/null @@ -1,295 +0,0 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -from PyPIC3D.boundary_conditions.grid_and_stencil import ( - collapse_axis_stencil, - prepare_particle_axis_stencil, -) -from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.solvers.yee_tiled import ( - digital_filter_tiled_vector, - fold_tiled_vector_ghost_cells, - tiled_grid_axes_from_world, - update_tiled_vector_ghost_cells, -) - - -def bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=2, tile_shape=None): - """ - Apply the global bilinear current filter to compact current tiles. - - The filter is the same separable [1, 2, 1]/4 stencil used by - ``bilinear_filter`` on assembled fields. Only the physical tile cells are - overwritten; guard cells are refreshed before and after the stencil. - """ - - g = int(num_guard_cells) - active = slice(g, -g) - backward = slice(g - 1, -g - 1) - forward = slice(g + 1, None if g == 1 else -g + 1) - - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) - - def filter_component(component): - center = component[:, :, :, active, active, active] - filtered = ( - 8.0 * center - + 4.0 * ( - component[:, :, :, backward, active, active] - + component[:, :, :, forward, active, active] - + component[:, :, :, active, backward, active] - + component[:, :, :, active, forward, active] - + component[:, :, :, active, active, backward] - + component[:, :, :, active, active, forward] - ) - + 2.0 * ( - component[:, :, :, backward, backward, active] - + component[:, :, :, backward, forward, active] - + component[:, :, :, forward, backward, active] - + component[:, :, :, forward, forward, active] - + component[:, :, :, backward, active, backward] - + component[:, :, :, backward, active, forward] - + component[:, :, :, forward, active, backward] - + component[:, :, :, forward, active, forward] - + component[:, :, :, active, backward, backward] - + component[:, :, :, active, backward, forward] - + component[:, :, :, active, forward, backward] - + component[:, :, :, active, forward, forward] - ) - + ( - component[:, :, :, backward, backward, backward] - + component[:, :, :, backward, backward, forward] - + component[:, :, :, backward, forward, backward] - + component[:, :, :, backward, forward, forward] - + component[:, :, :, forward, backward, backward] - + component[:, :, :, forward, backward, forward] - + component[:, :, :, forward, forward, backward] - + component[:, :, :, forward, forward, forward] - ) - ) / 64.0 - return component.at[:, :, :, active, active, active].set(filtered) - - J_tiles = tuple(filter_component(component) for component in J_tiles) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) - - return J_tiles - - -def digital_filter_tiled_current_density(J_tiles, alpha, world, num_guard_cells=2, tile_shape=None): - """ - Apply the global digital current filter to compact current tiles. - - Tile halos are refreshed before the stencil so every tile sees the same - neighbor values that the assembled ghost-celled current would see. The - halos are refreshed again after filtering so downstream Yee updates read - filtered guard-cell values. - """ - - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells, tile_shape) - - J_tiles = digital_filter_tiled_vector(J_tiles, alpha, num_guard_cells) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells, tile_shape) - - return J_tiles - - -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) - - -@partial(jit, static_argnames=("filter", "tile_shape", "g")) -def _direct_J_from_tiled_particles( - tiled_particles, - species_config, - J_tiles, - constants, - world, - grid=None, - filter="bilinear", - tile_shape=None, - g=2, -): - """ - Compute direct current deposition using tile-local stencils only. - - ``tiled_particles.x`` is assumed to already be centered at the current - deposition time and refreshed into the tiles that own those centered - positions. In other words, callers that store particles at the forward - position should pass a temporary, refreshed deposition view with - ``x_dep = x_forward - 0.5 * u * dt``. - """ - - if grid is None: - grid = world["grids"]["center"] - tiled_grid = tiled_grid_axes_from_world( - world, - grid, - "tiled_center_grid", - tile_shape, - g, - ) - - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - - Jx_tiles, Jy_tiles, Jz_tiles = J_tiles - ntx, nty, ntz = Jx_tiles.shape[:3] - tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] - g = int(g) - local_Nx = tile_nx + 2 * g - local_Ny = tile_ny + 2 * g - local_Nz = tile_nz + 2 * g - shape_factor = world["shape_factor"] - reduced_x = int(ntx) == 1 and tile_nx == 1 - reduced_y = int(nty) == 1 and tile_ny == 1 - reduced_z = int(ntz) == 1 and tile_nz == 1 - - 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]) - - # 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 - - def deposit_one_tile(x_tile, u_tile, active_tile, tx, ty, tz): - 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) - active = active_tile.reshape(-1).astype(x.dtype) - q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) - dq = q / (dx * dy * dz) - - x_grid = tiled_grid[0][tx, ty, tz] - y_grid = tiled_grid[1][tx, ty, tz] - z_grid = tiled_grid[2][tx, ty, tz] - - x, x0, deltax_node, xpts = prepare_particle_axis_stencil( - x, - x_grid, - local_Nx, - shape_factor, - local_bc, - wind=tile_nx * dx, - ghost_cells=True, - ) - y, y0, deltay_node, ypts = prepare_particle_axis_stencil( - y, - y_grid, - local_Ny, - shape_factor, - local_bc, - wind=tile_ny * dy, - ghost_cells=True, - ) - z, z0, deltaz_node, zpts = prepare_particle_axis_stencil( - z, - z_grid, - local_Nz, - shape_factor, - local_bc, - wind=tile_nz * dz, - ghost_cells=True, - ) - - 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 - - x_weights_node, y_weights_node, z_weights_node = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), - lambda _: get_second_order_weights(deltax_node, deltay_node, deltaz_node, dx, dy, dz), - operand=None, - ) - 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, - ) - - xpts = jnp.asarray(xpts) - ypts = jnp.asarray(ypts) - zpts = jnp.asarray(zpts) - x_weights_node = jnp.asarray(x_weights_node) - y_weights_node = jnp.asarray(y_weights_node) - z_weights_node = jnp.asarray(z_weights_node) - x_weights_face = jnp.asarray(x_weights_face) - y_weights_face = jnp.asarray(y_weights_face) - z_weights_face = jnp.asarray(z_weights_face) - - 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) - - 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", - ) - - 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) - - Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( - tiled_particles.x, - tiled_particles.u, - tiled_particles.active, - tx, - ty, - tz, - ) - - J_tiles = fold_tiled_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, g, tile_shape) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, g, tile_shape) - - if filter == "bilinear": - J_tiles = bilinear_filter_tiled_current_density(J_tiles, world, num_guard_cells=g, tile_shape=tile_shape) - elif filter == "digital": - J_tiles = digital_filter_tiled_current_density(J_tiles, constants["alpha"], world, num_guard_cells=g, tile_shape=tile_shape) - - return J_tiles diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index fdd7f8fd..f5fbd913 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -10,7 +10,7 @@ from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.utils import digital_filter +from PyPIC3D.utilities.filters import digital_filter def compute_rho(particles, rho, world, constants, species_config=None): diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py index 67f27107..85673a31 100644 --- a/PyPIC3D/deposition/rho_tiled.py +++ b/PyPIC3D/deposition/rho_tiled.py @@ -11,11 +11,10 @@ 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.solvers.yee_tiled import ( - digital_filter_tiled_scalar, tiled_grid_axes_from_world, update_tiled_ghost_cells, ) -from PyPIC3D.utils import digital_filter +from PyPIC3D.utilities.filters import digital_filter def _species_scalar_to_slots(tiled_particles, species_value): @@ -299,10 +298,10 @@ def folded_owner(points, n_cells, cells_per_tile, bc): return rho_tiles -def _digital_filter_tiled_scalar(field_tiles, alpha, world, g, tile_shape): +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_tiled_scalar(field_tiles, alpha, g) + return digital_filter(field_tiles, alpha, num_guard_cells=g) @jit @@ -364,7 +363,7 @@ def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_ ) alpha = constants["alpha"] - rho_tiles = _digital_filter_tiled_scalar(rho_tiles, alpha, world, g, tile_shape) + 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 diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index 2d5cb3c5..6f9c46ac 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -237,8 +237,6 @@ def _tiled_current_deposition_stage(case, particles, J_func): J_tiles, case.constants, case.world, - tile_shape=case.tile_shape, - g=case.g, ) diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 85f11df8..2a1f88e2 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -8,6 +8,7 @@ from jax import jit from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV +from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.particles.tiled_particle_refresh import ( refresh_tiled_particle_tiles, @@ -97,7 +98,7 @@ def esirkepov_step(state): # Esirkepov needs old and new particle positions. The deposition kernel # predicts the new positions locally, then the actual particle state is # advanced and retiled after the current has been computed. - J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + J_tiles = Esirkepov_current(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) particles = update_tiled_particle_positions(particles, species_config, world["dt"]) particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow @@ -111,7 +112,7 @@ def direct_current_step(state): overflow = overflow_previous | overflow # wrap periodic particles and move them into their owning tiles. - J_tiles = current_deposition(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + J_tiles = current_deposition(particles, species_config, J_tiles, constants, world) # deposit current directly into tile-local Yee current arrays particles = update_tiled_particle_positions(particles, species_config, world["dt"] / 2) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index b9213d8c..1b94833f 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -6,7 +6,7 @@ from PyPIC3D.deposition.rho import _compute_rho_flat from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient -from PyPIC3D.utils import digital_filter +from PyPIC3D.utilities.filters import digital_filter from PyPIC3D.boundary_conditions.boundaryconditions import ( update_ghost_cells, apply_scalar_conducting_bc ) diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py index d1876083..ee622c0d 100644 --- a/PyPIC3D/solvers/first_order_yee.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -3,7 +3,7 @@ from functools import partial # import external libraries -from PyPIC3D.utils import digital_filter +from PyPIC3D.utilities.filters 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, diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index e8325f5b..2f548092 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -7,6 +7,7 @@ ) from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.utilities.filters import digital_filter_vector def _tile_axis_count(n_cells, cells_per_tile): @@ -501,47 +502,6 @@ def apply_tiled_conducting_bc(E_tiles, world, num_guard_cells=2): return Ex, Ey, Ez -def digital_filter_tiled_scalar(field_tiles, alpha, num_guard_cells=2): - """ - Apply the standard six-neighbor digital filter to each tile interior. - - Tile halos play the same role as global ghost cells in ``digital_filter``: - the filtered values replace only the physical tile interiors. - """ - - neighbor_weight = (1 - alpha) / 6 - g = int(num_guard_cells) - active = _active_slice(g) - backward = _backward_slice(g) - forward = _forward_slice(g) - - filtered = ( - alpha * field_tiles[:, :, :, active, active, active] - + neighbor_weight * field_tiles[:, :, :, backward, active, active] - + neighbor_weight * field_tiles[:, :, :, forward, active, active] - + neighbor_weight * field_tiles[:, :, :, active, backward, active] - + neighbor_weight * field_tiles[:, :, :, active, forward, active] - + neighbor_weight * field_tiles[:, :, :, active, active, backward] - + neighbor_weight * field_tiles[:, :, :, active, active, forward] - ) - - return field_tiles.at[:, :, :, active, active, active].set(filtered) - - -def digital_filter_tiled_vector(field_tiles, alpha, num_guard_cells=2): - """ - Apply the standard field filter component-wise in tile-major storage. - """ - - stacked_tiles = stack_tiled_vector_field(field_tiles) - filtered = jax.vmap( - lambda component: digital_filter_tiled_scalar(component, alpha, num_guard_cells), - in_axes=0, - out_axes=0, - )(stacked_tiles) - return _restore_tiled_vector_layout(filtered, 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. @@ -772,7 +732,7 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ # refresh tile halos before the digital field filter, matching the global # ghost-cell order in the standard Yee solver. - Ex, Ey, Ez = digital_filter_tiled_vector((Ex, Ey, Ez), constants.get("alpha", 1.0), g) + Ex, Ey, Ez = digital_filter_vector((Ex, Ey, Ez), constants.get("alpha", 1.0), num_guard_cells=g) Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) @@ -836,7 +796,7 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, # refresh tile halos before the digital field filter, matching the global # ghost-cell order in the standard Yee solver. - Bx, By, Bz = digital_filter_tiled_vector((Bx, By, Bz), constants.get("alpha", 1.0), g) + Bx, By, Bz = digital_filter_vector((Bx, By, Bz), constants.get("alpha", 1.0), num_guard_cells=g) if pml_state is None: return update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) 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/utils.py b/PyPIC3D/utils.py index 6db5dd8a..e3a96475 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -38,73 +38,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. diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 80fa9d49..a0da534c 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -6,8 +6,9 @@ from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_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.deposition.direct_deposition_tiled import ( - digital_filter_tiled_current_density, +from PyPIC3D.utilities.filters import ( + digital_filter, + digital_filter_vector, ) from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles @@ -16,9 +17,11 @@ assemble_tiled_vector_field, fold_tiled_ghost_cells, fold_tiled_ghost_cells_periodic, + tile_grid_axes, tile_vector_field, + update_tiled_vector_ghost_cells, ) -from PyPIC3D.utils import build_yee_grid, digital_filter +from PyPIC3D.utils import build_yee_grid jax.config.update("jax_enable_x64", True) @@ -52,12 +55,8 @@ 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, simulation_parameters): - tile_shape = ( - simulation_parameters["particle_tile_nx"], - simulation_parameters["particle_tile_ny"], - simulation_parameters["particle_tile_nz"], - ) + 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 = ( @@ -77,6 +76,26 @@ def _tile_shape(self, simulation_parameters): simulation_parameters["particle_tile_nz"], ) + 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 + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids + return world + def _one_tile_parameters(self, world): return { "particle_tile_nx": world["Nx"], @@ -109,18 +128,17 @@ def _centered_tiled_particles(self, particles, world, simulation_parameters): return centered_particles, species_config def _assembled_tiled_current(self, particles, world, simulation_parameters, constants, filter="none"): - tiled_particles, species_config = self._centered_tiled_particles(particles, world, simulation_parameters) tile_shape = self._tile_shape(simulation_parameters) + world = self._world_with_tiled_grids(world, tile_shape) + tiled_particles, species_config = self._centered_tiled_particles(particles, world, simulation_parameters) J_tiles = J_from_rhov( tiled_particles, species_config, - self._empty_J_tiles(world, simulation_parameters), + self._empty_J_tiles(world), constants, world, filter=filter, - tile_shape=tile_shape, - g=int(world["guard_cells"]), ) g = int(world["guard_cells"]) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) @@ -287,7 +305,8 @@ def test_tiled_digital_filter_matches_global_digital_filter(self): 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_tiled_current_density(J_tiles, alpha, world) + filtered_tiles = digital_filter_vector(J_tiles, alpha, num_guard_cells=1) + filtered_tiles = 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) @@ -447,16 +466,15 @@ def test_tiled_direct_deposition_returns_only_local_current_tiles(self): ) tiled_particles, species_config = self._centered_tiled_particles([species], 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, simulation_parameters), + self._empty_J_tiles(world), constants, world, filter="none", - tile_shape=tile_shape, - g=int(world["guard_cells"]), ) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=int(world["guard_cells"])) @@ -523,12 +541,13 @@ def test_tiled_direct_deposition_matches_J_from_rhov_for_dummy_species(self): def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(self): world = self._build_world() world["guard_cells"] = 2 - world["tile_shape"] = (2, 2, 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} species = particle_species( name="public tiled direct current", @@ -556,17 +575,15 @@ def test_public_J_from_rhov_dispatches_tiled_particles_to_tile_local_current(sel J_tiles = J_from_rhov( tiled_particles, species_config, - self._empty_J_tiles(world, simulation_parameters), + self._empty_J_tiles(world), constants, world, filter="digital", - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), ) J_from_tiles = assemble_tiled_vector_field( J_tiles, world, - world["tile_shape"], + tile_shape, num_guard_cells=int(world["guard_cells"]), ) _, J_reference = self._assembled_tiled_current( diff --git a/tests/evolve_test.py b/tests/evolve_test.py index e076d142..649f0c56 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -16,7 +16,7 @@ jax.config.update("jax_enable_x64", True) -def zero_current(particles, species_config, J, constants, world, tile_shape=None, g=None): +def zero_current(particles, species_config, J, constants, world): return tuple(jnp.zeros_like(comp) for comp in J) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index cd23e8a2..c3b31712 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -101,9 +101,15 @@ def test_electrostatic_global_poisson_bridge_stays_explicit(self): def test_flat_current_and_electrodynamic_references_are_removed(self): self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) + self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_tiled")) self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) self.assertFalse(hasattr(evolve, "_time_loop_electrostatic_global_reference")) + def test_direct_current_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(J_from_rhov.J_from_rhov) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index e18a79f4..350011ae 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -16,7 +16,7 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, empty_tiled_vector_field, tile_vector_field +from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, empty_tiled_vector_field, tile_grid_axes, tile_vector_field from PyPIC3D.utils import build_yee_grid, compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -55,6 +55,26 @@ def _build_world(self, boundary_conditions=None): world["grids"] = {"vertex": vertex_grid, "center": center_grid} return world + 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 + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids + return world + def _empty_fields(self, world): shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) zero = jnp.zeros(shape) @@ -118,6 +138,8 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (4, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) Nt = 1500 x0 = jnp.linspace( @@ -165,30 +187,30 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) - g = int(world["guard_cells"]) + g = int(tiled_world["guard_cells"]) if current_calculation == CURRENT_ESIRKEPOV: - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) + J_tiles = empty_tiled_vector_field(tiled_world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) else: - J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=g), - tile_vector_field(B, world, tile_shape, num_guard_cells=g), + tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=g), + tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=g), J_tiles, rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=g), ), None, ) reference_particles, reference_species_config, reference_fields = self._build_tiled_state( reference_species, (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, current_calculation, ) @@ -199,7 +221,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): reference_fields, tiled_particles, tiled_fields, - world, + tiled_world, reference_tile_shape, tile_shape, step=0, @@ -210,10 +232,11 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): reference_species_config, reference_fields, reference_tile_shape, + reference_world, tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, tile_shape, Nt, @@ -273,6 +296,39 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, current_calcu ) return tiled_particles, species_config, tiled_fields + def _jitted_electrodynamic_loop_with_static_world(self, world, tile_shape, J_func): + g = int(world["guard_cells"]) + + def loop_with_static_world( + particles, + species_config, + fields, + constants, + curl_func, + solver, + relativistic=True, + particle_pusher="boris", + ): + return time_loop_electrodynamic( + particles, + species_config, + fields, + world, + constants, + curl_func, + J_func=J_func, + solver=solver, + tile_shape=tile_shape, + g=g, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + return jax.jit( + loop_with_static_world, + static_argnames=("curl_func", "solver", "relativistic", "particle_pusher"), + ) + def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): diff = jnp.max(jnp.abs(reference_component - tiled_component)) @@ -416,7 +472,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world, tile_shape=None, g=None: J, + J_func=lambda particles, species_config, J, constants, world: J, solver="electrodynamic_yee", tile_shape=reference_tile_shape, g=int(world["guard_cells"]), @@ -443,7 +499,7 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): world, constants, unused_curl, - J_func=lambda particles, J, constants, world, tile_shape=None, g=None: J, + J_func=lambda particles, species_config, J, constants, world: J, solver="electrodynamic_yee", tile_shape=tile_shape, g=int(world["guard_cells"]), @@ -764,6 +820,8 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -776,20 +834,20 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, ) @@ -797,7 +855,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=None, @@ -811,7 +869,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=None, @@ -853,6 +911,8 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -865,20 +925,20 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, ) @@ -886,7 +946,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=None, @@ -900,7 +960,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=None, @@ -927,6 +987,8 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -939,20 +1001,20 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, ) @@ -960,7 +1022,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="digital"), @@ -974,7 +1036,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="digital"), @@ -1001,6 +1063,8 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -1013,20 +1077,20 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, ) @@ -1034,7 +1098,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="bilinear"), @@ -1048,7 +1112,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="bilinear"), @@ -1077,6 +1141,8 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -1130,20 +1196,20 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand reference_particles, reference_species_config, reference_fields = self._build_tiled_state( [reference_species], (E, B, J, rho, phi, external_fields, pml_state), - world, + reference_world, reference_tile_shape, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), ), None, ) @@ -1151,7 +1217,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="bilinear"), @@ -1165,7 +1231,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=functools.partial(J_from_rhov, filter="bilinear"), @@ -1197,32 +1263,27 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): reference_species_config, reference_fields, reference_tile_shape, + reference_world, tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, tile_shape, Nt, ) = self._long_two_stream_state(CURRENT_J_FROM_RHOV) - tiled_loop = jax.jit( - time_loop_electrodynamic, - static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), - ) + reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world, reference_tile_shape, None) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world, tile_shape, None) for t in range(Nt): - reference_particles, reference_fields = tiled_loop( + reference_particles, reference_fields = reference_loop( reference_particles, reference_species_config, reference_fields, - world, constants, unused_curl, - J_func=None, solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1230,13 +1291,9 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): tiled_particles, species_config, tiled_fields, - world, constants, unused_curl, - J_func=None, solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1247,7 +1304,7 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): reference_tile_shape, tiled_particles, tiled_fields, - world, + tiled_world, tile_shape, ) if t == 0: @@ -1256,7 +1313,7 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): reference_fields, tiled_particles, tiled_fields, - world, + tiled_world, reference_tile_shape, tile_shape, t + 1, @@ -1275,32 +1332,31 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): reference_species_config, reference_fields, reference_tile_shape, + reference_world, tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, tile_shape, Nt, ) = self._long_two_stream_state(CURRENT_ESIRKEPOV) - tiled_loop = jax.jit( - time_loop_electrodynamic, - static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), + reference_loop = self._jitted_electrodynamic_loop_with_static_world( + reference_world, + reference_tile_shape, + Esirkepov_current, ) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world, tile_shape, Esirkepov_current) for t in range(Nt): - reference_particles, reference_fields = tiled_loop( + reference_particles, reference_fields = reference_loop( reference_particles, reference_species_config, reference_fields, - world, constants, unused_curl, - J_func=Esirkepov_current, solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1308,13 +1364,9 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): tiled_particles, species_config, tiled_fields, - world, constants, unused_curl, - J_func=Esirkepov_current, solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1325,7 +1377,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): reference_tile_shape, tiled_particles, tiled_fields, - world, + tiled_world, tile_shape, ) if t == 0: @@ -1334,7 +1386,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): reference_fields, tiled_particles, tiled_fields, - world, + tiled_world, reference_tile_shape, tile_shape, t + 1, From 53f9cb05dc9f91747e8aa6b97ea88b708a374d7d Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 3 Jul 2026 09:13:02 -0500 Subject: [PATCH 35/83] refactoring the esirkepov current density --- PyPIC3D/deposition/Esirkepov.py | 530 ++++++++++++++++++++++++-- PyPIC3D/deposition/esirkepov_tiled.py | 460 ---------------------- PyPIC3D/evolve.py | 2 +- 3 files changed, 492 insertions(+), 500 deletions(-) delete mode 100644 PyPIC3D/deposition/esirkepov_tiled.py diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 15465419..65006182 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -2,7 +2,14 @@ import jax.numpy as jnp from jax import lax -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + compute_particle_anchor, + particle_axis_offset, +) +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights +from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.solvers.yee_tiled import update_tiled_vector_ghost_cells +from PyPIC3D.utilities.filters import bilinear_filter_vector, digital_filter_vector def shift_old_stencil(old_w_list, shift): @@ -197,51 +204,496 @@ def eliminate_esirkepov_ghost_cells(field): # remove one layer on each side, returning the solver's ordinary one-ghost current array. return field[tuple(slices)] + +def _active_stencil_indices(axis_active): + if axis_active: + return (0, 1, 2, 3, 4) + return (2,) + + +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 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. + """ + + 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}) + + 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 + + if reduced_axis: + ghost_sum = ( + jnp.sum(lower_ghost, axis=3, keepdims=True) + + jnp.sum(upper_ghost, axis=3, keepdims=True) + ) + + def periodic_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(ghost_sum) + + def reflecting_boundary(tiles): + return tiles.at[:, :, :, g:g + 1, :, :].add(sign * ghost_sum) + + 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 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, + particles: TiledParticles, + species_config: SpeciesConfig, J, constants, world, - grid=None, filter=None, - species_config=None, - tile_shape=None, - g=None, ): - """Esirkepov current deposition supporting 1D/2D/3D via inactive dims.""" - if isinstance(particles, TiledParticles): - if filter not in (None, "none"): - raise ValueError("Esirkepov current filtering is not supported; use filter='none'.") - - if species_config is None: - species_config = J - J = constants - constants = world - world = grid - grid = None - # Tiled Esirkepov stores old particle positions. The tiled kernel - # predicts the new positions locally and leaves particle retile staging - # to the caller, matching the Task 0 old/new position contract. - - if tile_shape is None: - tile_shape = tuple(int(width) for width in world["tile_shape"]) - if g is None: - g = int(world["guard_cells"]) - - from PyPIC3D.deposition.esirkepov_tiled import _tiled_esirkepov_current - - return _tiled_esirkepov_current( - particles, - species_config, - J, - constants, - world, - grid=grid, - tile_shape=tile_shape, - g=int(g), - ) + """ + 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. + """ + - raise ValueError("Public Esirkepov_current requires TiledParticles.") + 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 + # 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 + # 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) + xpts = x0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + ypts = y0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + zpts = z0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] + # compute the 5-point stencil indices for the new anchor points in each axis + + xw, yw, zw = 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, + ) + oxw, oyw, ozw = jax.lax.cond( + shape_factor == 1, + lambda _: get_first_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), + lambda _: get_second_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), + operand=None, + ) + # 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] + # 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 to align with the new anchor points based on the computed shifts + + 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), + lambda _: active * q * vx / (dx * dy * dz), + operand=None, + ) + dJy = jax.lax.cond( + y_active, + 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), + lambda _: active * q * vz / (dx * dy * dz), + operand=None, + ) + # 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. + + tile_Jx = Jx_template + tile_Jy = Jy_template + tile_Jz = Jz_template + # initialize the local tile current arrays to zero + + 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) + + for i in range(5): + for j in range(5): + for k 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): + 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): + 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): + 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: + # 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 + + 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 + + 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 def get_3D_esirkepov_weights( diff --git a/PyPIC3D/deposition/esirkepov_tiled.py b/PyPIC3D/deposition/esirkepov_tiled.py deleted file mode 100644 index 4e61e32f..00000000 --- a/PyPIC3D/deposition/esirkepov_tiled.py +++ /dev/null @@ -1,460 +0,0 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -from PyPIC3D.boundary_conditions.grid_and_stencil import ( - compute_particle_anchor, - particle_axis_offset, -) -from PyPIC3D.deposition.Esirkepov import ( - collapse_redundant_axis, - get_3D_esirkepov_weights, - shift_old_stencil, -) -from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.solvers.yee_tiled import ( - tiled_grid_axes_from_world, - update_tiled_vector_ghost_cells, -) - - -def _active_stencil_indices(axis_active): - if axis_active: - return (0, 1, 2, 3, 4) - return (2,) - - -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 fold_tiled_esirkepov_ghost_cells(field_tiles, world, component_axis, num_guard_cells=2, tile_shape=None): - """ - Fold Esirkepov current deposits using the configured guard depth. - - Internal tile faces are pure ownership transfers. At global particle - boundaries, periodic deposits wrap, reflecting deposits mirror into the - same-side physical layer with the Esirkepov component sign, and absorbing - deposits are discarded. - """ - - g = int(num_guard_cells) - if tile_shape is None: - tile_shape = tuple(int(width) - 2 * g for width in field_tiles.shape[3:]) - 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}) - - 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 - - if reduced_axis: - ghost_sum = ( - jnp.sum(lower_ghost, axis=3, keepdims=True) - + jnp.sum(upper_ghost, axis=3, keepdims=True) - ) - - def periodic_boundary(tiles): - return tiles.at[:, :, :, g:g + 1, :, :].add(ghost_sum) - - def reflecting_boundary(tiles): - return tiles.at[:, :, :, g:g + 1, :, :].add(sign * ghost_sum) - - 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 fold_tiled_esirkepov_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): - return tuple( - fold_tiled_esirkepov_ghost_cells(component, world, component_axis, num_guard_cells, tile_shape) - for component_axis, component in enumerate(field_tiles) - ) - - -@partial(jit, static_argnames=("tile_shape", "g")) -def _tiled_esirkepov_current(tiled_particles, species_config, J_tiles, constants, world, grid=None, tile_shape=None, g=2): - """ - Deposit Esirkepov current into tile-local current buffers. - - ``tiled_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 boundary conditions and retile - ownership are applied by the caller after deposition. - """ - - if grid is None: - grid = world["grids"]["center"] - tiled_grid = tiled_grid_axes_from_world( - world, - grid, - "tiled_center_grid", - tile_shape, - g, - ) - - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - dt = world["dt"] - shape_factor = world["shape_factor"] - - Jx_tiles, Jy_tiles, Jz_tiles = J_tiles - ntx, nty, ntz = Jx_tiles.shape[:3] - tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] - g = int(g) - local_Nx = tile_nx + 2 * g - local_Ny = tile_ny + 2 * g - local_Nz = tile_nz + 2 * g - - x_active = ntx * tile_nx > 1 - y_active = nty * tile_ny > 1 - z_active = ntz * tile_nz > 1 - - 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]) - - species_weighted_charge = species_config.charge * species_config.weight - - 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) - vx = u_tile[..., 0].reshape(-1) - vy = u_tile[..., 1].reshape(-1) - vz = u_tile[..., 2].reshape(-1) - active = active_tile.reshape(-1).astype(old_x.dtype) - q = jnp.broadcast_to(species_weighted_charge[:, jnp.newaxis], active_tile.shape).reshape(-1) - N_particles = old_x.shape[0] - 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) - - 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) - - x_grid = tiled_grid[0][tx, ty, tz] - y_grid = tiled_grid[1][tx, ty, tz] - z_grid = tiled_grid[2][tx, ty, tz] - - 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) - 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) - - deltax = particle_axis_offset(x, x0, x_grid) - deltay = particle_axis_offset(y, y0, y_grid) - deltaz = particle_axis_offset(z, z0, z_grid) - 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) - - shift_x = x0 - old_x0 - shift_y = y0 - old_y0 - shift_z = z0 - old_z0 - - offsets = jnp.asarray([-2, -1, 0, 1, 2], dtype=x0.dtype) - xpts = x0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] - ypts = y0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] - zpts = z0[jnp.newaxis, ...] + offsets[:, jnp.newaxis] - - xw, yw, zw = 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, - ) - oxw, oyw, ozw = jax.lax.cond( - shape_factor == 1, - lambda _: get_first_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), - lambda _: get_second_order_weights(old_deltax, old_deltay, old_deltaz, dx, dy, dz), - operand=None, - ) - - 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] - - oxw = shift_old_stencil(oxw, shift_x) - oyw = shift_old_stencil(oyw, shift_y) - ozw = shift_old_stencil(ozw, shift_z) - - 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) - - dJx = jax.lax.cond( - x_active, - 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), - lambda _: active * q * vy / (dx * dy * dz), - operand=None, - ) - dJz = jax.lax.cond( - z_active, - lambda _: active * (-(q / (dx * dy)) / dt), - lambda _: active * q * vz / (dx * dy * dz), - operand=None, - ) - - tile_Jx = Jx_template - tile_Jy = Jy_template - tile_Jz = Jz_template - - if x_active and y_active and z_active: - 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) - - for i in range(5): - for j in range(5): - for k 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 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): - 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): - 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): - 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): - 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): - 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: - 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", - ) - - 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) - - Jx_tiles, Jy_tiles, Jz_tiles = deposit_tiles( - tiled_particles.x, - tiled_particles.u, - tiled_particles.active, - tx, - ty, - tz, - ) - - J_tiles = fold_tiled_esirkepov_vector_ghost_cells((Jx_tiles, Jy_tiles, Jz_tiles), world, num_guard_cells=g, tile_shape=tile_shape) - J_tiles = update_tiled_vector_ghost_cells(J_tiles, world, num_guard_cells=g, tile_shape=tile_shape) - - return J_tiles diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 2a1f88e2..da7ed53d 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -98,7 +98,7 @@ def esirkepov_step(state): # Esirkepov needs old and new particle positions. The deposition kernel # predicts the new positions locally, then the actual particle state is # advanced and retiled after the current has been computed. - J_tiles = Esirkepov_current(particles, species_config, J_tiles, constants, world, tile_shape=tile_shape, g=g) + J_tiles = Esirkepov_current(particles, species_config, J_tiles, constants, world) particles = update_tiled_particle_positions(particles, species_config, world["dt"]) particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) overflow = overflow_previous | overflow From 1562cb37fdf9501d540461cfef44319d2e954a79 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 3 Jul 2026 09:13:21 -0500 Subject: [PATCH 36/83] refactoring unit tests for esirkepov refactor --- tests/direct_deposition_tiled_test.py | 7 +++-- tests/evolve_test.py | 21 ++++++++++++- tests/tiled_esirkepov_test.py | 43 +++++++++++++++++++-------- tests/tiled_refactor_contract_test.py | 10 +++++++ 4 files changed, 65 insertions(+), 16 deletions(-) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index a0da534c..5dbcc02e 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -162,6 +162,8 @@ def _compare_tiled_to_one_tile(self, particles, world, simulation_parameters, fi def test_public_J_from_rhov_rejects_flat_particles(self): world = self._build_world() constants = {"C": 3.0e8, "alpha": 1.0} + simulation_parameters = self._one_tile_parameters(world) + world = self._world_with_tiled_grids(world, self._tile_shape(simulation_parameters)) species = particle_species( name="flat direct current", N_particles=1, @@ -182,9 +184,10 @@ def test_public_J_from_rhov_rejects_flat_particles(self): dz=world["dz"], dt=world["dt"], ) + _, species_config = to_tiled_particles([species], world, simulation_parameters) - with self.assertRaisesRegex(ValueError, "TiledParticles"): - J_from_rhov([species], self._empty_J(world), constants, world, filter="none") + with self.assertRaises(AttributeError): + J_from_rhov([species], species_config, self._empty_J_tiles(world), constants, world, filter="none") def test_tiled_direct_deposition_matches_quadratic_with_two_guard_cells(self): world = self._build_world(Nx=8, Ny=6, Nz=4) diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 649f0c56..9d4b9e29 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -10,7 +10,7 @@ from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.solvers.yee_tiled import tile_grid_axes, tile_vector_field from PyPIC3D.utils import build_yee_grid jax.config.update("jax_enable_x64", True) @@ -24,6 +24,23 @@ def unused_curl(Ex, Ey, Ez): return None +def add_tiled_grids_to_world(world, tile_shape): + g = int(world["guard_cells"]) + world["tile_shape"] = tile_shape + world["grids"]["tiled_center_grid"] = tile_grid_axes( + world["grids"]["center"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"]["tiled_vertex_grid"] = tile_grid_axes( + world["grids"]["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + + class TestEvolveExternalFields(unittest.TestCase): def test_public_loop_names_own_tiled_contracts(self): self.assertEqual(time_loop_electrodynamic.__module__, "PyPIC3D.evolve") @@ -84,6 +101,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field "particle_tile_nz": tile_shape[2], } world["guard_cells"] = 2 + add_tiled_grids_to_world(world, tile_shape) tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) external_E_tiles = tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])) external_B_tiles = tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])) @@ -182,6 +200,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): "particle_tile_nz": tile_shape[2], } world["guard_cells"] = 2 + add_tiled_grids_to_world(world, tile_shape) tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) fields = ( tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 1d762974..81555c47 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -9,11 +9,13 @@ import toml from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.deposition.Esirkepov import Esirkepov_current, get_1D_esirkepov_weights, get_2D_esirkepov_weights -from PyPIC3D.deposition.esirkepov_tiled import ( +from PyPIC3D.deposition.Esirkepov import ( _active_stencil_indices, _compact_1d_esirkepov_weights, _compact_2d_esirkepov_weights, + Esirkepov_current, + get_1D_esirkepov_weights, + get_2D_esirkepov_weights, ) from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles @@ -30,6 +32,7 @@ assemble_tiled_vector_field, empty_tiled_scalar_field, empty_tiled_vector_field, + tile_grid_axes, tile_vector_field, update_tiled_E, ) @@ -141,7 +144,28 @@ def _tile_shape_for_world(self, world): def _one_tile_shape_for_world(self, world): return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + 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 + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids + return world + def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape): + world = self._world_with_tiled_grids(world, tile_shape) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -155,8 +179,6 @@ def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, - tile_shape=tile_shape, - g=g, ) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) @@ -322,6 +344,7 @@ 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) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -338,8 +361,6 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, - tile_shape=tile_shape, - g=g, ) J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) _, J_reference = self._assembled_esirkepov_current( @@ -358,7 +379,7 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): 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["tile_shape"] = (2, 1, 1) + world = self._world_with_tiled_grids(world, (2, 1, 1)) constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} simulation_parameters = { "particle_tile_nx": world["tile_shape"][0], @@ -375,8 +396,6 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre empty_tiled_vector_field(world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])), constants, world, - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), ) J_from_tiles = assemble_tiled_vector_field( J_tiles, @@ -403,6 +422,7 @@ 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} simulation_parameters = { "particle_tile_nx": tile_shape[0], @@ -420,8 +440,6 @@ def test_public_Esirkepov_current_rejects_tiled_filter_modes(self): constants, world, filter="digital", - tile_shape=tile_shape, - g=int(world["guard_cells"]), ) def test_tiled_esirkepov_matches_global_current_for_dimensions_and_shapes(self): @@ -501,6 +519,7 @@ 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) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -533,8 +552,6 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), constants, world, - tile_shape=tile_shape, - g=g, ) new_particles = update_tiled_particle_positions(tiled_particles, species_config, world["dt"]) new_particles, overflow = refresh_tiled_particle_tiles(new_particles, world, tile_shape) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index c3b31712..21900bd2 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -1,9 +1,11 @@ import inspect +import importlib.util import unittest import jax.numpy as jnp from PyPIC3D import evolve +from PyPIC3D.deposition import Esirkepov from PyPIC3D.deposition import J_from_rhov from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled @@ -110,6 +112,14 @@ def test_direct_current_api_reads_tile_metadata_from_world(self): self.assertNotIn("tile_shape", signature.parameters) self.assertNotIn("g", signature.parameters) + def test_esirkepov_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(Esirkepov.Esirkepov_current) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + + def test_tiled_esirkepov_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.esirkepov_tiled")) + if __name__ == "__main__": unittest.main() From 1f41275c08fcdc75f44ad6d38a87ef746c0dcdef Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Fri, 3 Jul 2026 09:59:39 -0500 Subject: [PATCH 37/83] updating tiled charge density calculation --- PyPIC3D/__main__.py | 9 +- PyPIC3D/deposition/Esirkepov.py | 5 +- PyPIC3D/deposition/current_methods.py | 12 - PyPIC3D/deposition/rho.py | 704 ++++++++++++++++++--- PyPIC3D/deposition/rho_tiled.py | 705 ---------------------- PyPIC3D/diagnostics/fluid_quantities.py | 2 +- PyPIC3D/diagnostics/tiled_benchmarking.py | 3 +- PyPIC3D/evolve.py | 10 +- PyPIC3D/initialization.py | 9 +- PyPIC3D/solvers/electrostatic_yee.py | 12 +- tests/particle_test.py | 47 +- tests/tiled_electrostatic_test.py | 80 ++- tests/tiled_esirkepov_test.py | 11 +- tests/tiled_fluid_quantities_test.py | 53 +- tests/tiled_refactor_contract_test.py | 13 + tests/tiled_rho_test.py | 125 ++-- tests/tiled_yee_integration_test.py | 17 +- 17 files changed, 873 insertions(+), 944 deletions(-) delete mode 100644 PyPIC3D/deposition/current_methods.py delete mode 100644 PyPIC3D/deposition/rho_tiled.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 49cfdc7d..ca4edb9a 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -158,17 +158,12 @@ def loop_with_static_world( if plotting_parameters['plot_vtk_scalars']: if getattr(rho, "ndim", 0) == 6: - rho = compute_rho(particles, rho, world, constants, species_config=species_config) + rho = compute_rho(particles, species_config, rho, constants, world) mass_density = compute_mass_density(particles, rho, world, species_config=species_config) rho_output = scalar_field_for_output(rho, world) mass_density_output = scalar_field_for_output(mass_density, world) else: - rho = compute_rho(particles, rho, world, constants, species_config=species_config) - # calculate the charge density based on the particle positions - mass_density = compute_mass_density(particles, rho, world, species_config=species_config) - # calculate the mass density based on the particle positions - rho_output = rho - mass_density_output = mass_density + raise ValueError("Runtime scalar diagnostics require tiled rho storage.") y_mid = world['Ny']//2 + 1 # midplane index shifted by 1 for ghost cells diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 65006182..0fbce74f 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -9,7 +9,6 @@ from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles from PyPIC3D.solvers.yee_tiled import update_tiled_vector_ghost_cells -from PyPIC3D.utilities.filters import bilinear_filter_vector, digital_filter_vector def shift_old_stencil(old_w_list, shift): @@ -384,6 +383,10 @@ def Esirkepov_current( 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 diff --git a/PyPIC3D/deposition/current_methods.py b/PyPIC3D/deposition/current_methods.py deleted file mode 100644 index a130c532..00000000 --- a/PyPIC3D/deposition/current_methods.py +++ /dev/null @@ -1,12 +0,0 @@ -CURRENT_J_FROM_RHOV = 0 -CURRENT_ESIRKEPOV = 1 - - -def encode_current_calculation(current_calculation): - current_codes = { - "j_from_rhov": CURRENT_J_FROM_RHOV, - "esirkepov": CURRENT_ESIRKEPOV, - } - if current_calculation not in current_codes: - raise ValueError("Unsupported current_calculation. Use 'j_from_rhov' or 'esirkepov'.") - return current_codes[current_calculation] diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index f5fbd913..9bcac765 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -1,121 +1,659 @@ +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.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.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles +from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.solvers.yee_tiled import update_tiled_ghost_cells from PyPIC3D.utilities.filters import digital_filter -def compute_rho(particles, rho, world, constants, species_config=None): - """Compute charge density for flat or tile-major particle storage.""" - if isinstance(particles, TiledParticles): - if species_config is None: - species_config = world["species_config"] - - if getattr(rho, "ndim", 0) == 6: - return compute_tiled_rho_from_tiled_particles( - particles, - species_config, - rho, - world, - constants, - tile_shape=tuple(int(width) for width in world["tile_shape"]), - g=int(world["guard_cells"]), - ) +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, + ) - return compute_rho_from_tiled_particles(particles, species_config, rho, world, constants) - raise ValueError("Public compute_rho requires TiledParticles; use _compute_rho_flat for flat reference tests.") +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) -@jit -def _compute_rho_flat(particles, rho, world, constants): - """Compute the charge density (rho) on the vertex grid.""" +def _deposit_tiled_scalar_moment( + tiled_particles, + rho, + world, + grid, + scalar_weight, + 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"] + 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, y, z = position + x = x.reshape(-1) + y = y.reshape(-1) + z = z.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, + 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, + ) + z, _, 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, + ) + + 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) + rho = jnp.zeros_like(rho) - for species in particles: - q = species.get_charge() - dq = q / dx / dy / dz - x, y, z = species.get_forward_position() - active = species.get_active_mask().astype(x.dtype) - - 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, + 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 * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], + mode="drop", + ) + + rho = fold_ghost_cells(rho, bc_x, bc_y, bc_z) + rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) + + return rho + + +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 + + 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"] + 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 + + 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") - 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) + rho_tiles = update_tiled_ghost_cells(rho_tiles, world, g, tile_shape) - 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) + return rho_tiles - 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 +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"] + + 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 + return rho_tiles + + +@jit +def compute_mass_density_from_tiled_particles(tiled_particles, species_config, rho, world, grid=None): + """Compute mass density on the vertex grid from tile-major particles.""" + 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( + tiled_particles, + rho, + world, + grid, + scalar_weight, + position=position, + ) + + +@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",)) +def compute_velocity_field_from_tiled_particles(tiled_particles, species_config, field, direction, world, grid=None): + """Compute the legacy scalar velocity diagnostic from tile-major particles.""" + 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( + tiled_particles, + field, + world, + grid, + scalar_weight, + position=position, + divide_by_volume=False, + ) + + +@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",)) +def compute_pressure_field_from_tiled_particles( + tiled_particles, + species_config, + field, + velocity_field, + direction, + world, + grid=None, +): + """Compute the legacy scalar pressure diagnostic from tile-major particles.""" + if grid is None: + grid = world["grids"]["vertex"] + + dx = world["dx"] + dy = world["dy"] + dz = world["dz"] + Nx, Ny, Nz = field.shape + bc_x = world["boundary_conditions"]["x"] + bc_y = world["boundary_conditions"]["y"] + bc_z = world["boundary_conditions"]["z"] + shape_factor = world["shape_factor"] + + x, y, z = _diagnostic_mass_position(tiled_particles, world) + x = x.reshape(-1) + y = y.reshape(-1) + z = z.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, + 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, + ) + z, _, 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, + ) + + 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) + + field = jnp.zeros_like(field) + + for i in range(xpts.shape[0]): + for j in range(ypts.shape[0]): + for k in range(zpts.shape[0]): + 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", + ) + + field = fold_ghost_cells(field, bc_x, bc_y, bc_z) + field = update_ghost_cells(field, bc_x, bc_y, bc_z) + return field + + +@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") + + field_tiles = update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + return field_tiles diff --git a/PyPIC3D/deposition/rho_tiled.py b/PyPIC3D/deposition/rho_tiled.py deleted file mode 100644 index 85673a31..00000000 --- a/PyPIC3D/deposition/rho_tiled.py +++ /dev/null @@ -1,705 +0,0 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -from PyPIC3D.boundary_conditions.grid_and_stencil import ( - BC_PERIODIC, - 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.solvers.yee_tiled import ( - tiled_grid_axes_from_world, - update_tiled_ghost_cells, -) -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 _deposit_tiled_scalar_moment( - tiled_particles, - rho, - world, - grid, - scalar_weight, - position=None, - divide_by_volume=True, -): - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - 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"] - - 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, y, z = position - x = x.reshape(-1) - y = y.reshape(-1) - z = z.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, - 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, - ) - z, _, 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, - ) - - 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) - - rho = jnp.zeros_like(rho) - - 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 * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], - mode="drop", - ) - - rho = fold_ghost_cells(rho, bc_x, bc_y, bc_z) - rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) - - return rho - - -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 - - 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, - tiled_grid, - position=None, - divide_by_volume=True, -): - 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 - - 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, - ) - 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 - 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) - - -@jit -def compute_rho_from_tiled_particles(tiled_particles, species_config, rho, world, constants, grid=None): - """Compute charge density on the vertex grid from tile-major particles.""" - if grid is None: - grid = world["grids"]["vertex"] - - scalar_weight = _species_scalar_to_slots( - tiled_particles, - species_config.charge * species_config.weight, - ) - rho = _deposit_tiled_scalar_moment( - tiled_particles, - rho, - world, - grid, - scalar_weight, - ) - - alpha = constants["alpha"] - rho = digital_filter(rho, alpha) - rho = update_ghost_cells( - rho, - world["boundary_conditions"]["x"], - world["boundary_conditions"]["y"], - world["boundary_conditions"]["z"], - ) - - return rho - - -@partial(jit, static_argnames=("tile_shape", "g")) -def compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, grid=None, tile_shape=None, g=2): - """Compute charge density into tile-major vertex-grid scalar arrays.""" - if grid is None: - grid = world["grids"]["vertex"] - tiled_grid = tiled_grid_axes_from_world( - world, - grid, - "tiled_vertex_grid", - tile_shape, - g, - ) - - scalar_weight = _species_scalar_to_slots( - tiled_particles, - species_config.charge * species_config.weight, - ) - rho_tiles = _deposit_tiled_scalar_moment_to_tiles( - tiled_particles, - rho_tiles, - world, - grid, - scalar_weight, - tile_shape, - g, - tiled_grid, - ) - - alpha = constants["alpha"] - 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 - - -@jit -def compute_mass_density_from_tiled_particles(tiled_particles, species_config, rho, world, grid=None): - """Compute mass density on the vertex grid from tile-major particles.""" - 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( - tiled_particles, - rho, - world, - grid, - scalar_weight, - position=position, - ) - - -@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"] - tiled_grid = tiled_grid_axes_from_world( - world, - grid, - "tiled_vertex_grid", - tile_shape, - g, - ) - - 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, - tiled_grid, - position=position, - ) - - -@partial(jit, static_argnames=("direction",)) -def compute_velocity_field_from_tiled_particles(tiled_particles, species_config, field, direction, world, grid=None): - """Compute the legacy scalar velocity diagnostic from tile-major particles.""" - 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( - tiled_particles, - field, - world, - grid, - scalar_weight, - position=position, - divide_by_volume=False, - ) - - -@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"] - tiled_grid = tiled_grid_axes_from_world( - world, - grid, - "tiled_vertex_grid", - tile_shape, - g, - ) - - 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, - tiled_grid, - position=position, - divide_by_volume=False, - ) - - -@partial(jit, static_argnames=("direction",)) -def compute_pressure_field_from_tiled_particles( - tiled_particles, - species_config, - field, - velocity_field, - direction, - world, - grid=None, -): - """Compute the legacy scalar pressure diagnostic from tile-major particles.""" - if grid is None: - grid = world["grids"]["vertex"] - - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - Nx, Ny, Nz = field.shape - bc_x = world["boundary_conditions"]["x"] - bc_y = world["boundary_conditions"]["y"] - bc_z = world["boundary_conditions"]["z"] - shape_factor = world["shape_factor"] - - x, y, z = _diagnostic_mass_position(tiled_particles, world) - x = x.reshape(-1) - y = y.reshape(-1) - z = z.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, - 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, - ) - z, _, 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, - ) - - 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) - - field = jnp.zeros_like(field) - - for i in range(xpts.shape[0]): - for j in range(ypts.shape[0]): - for k in range(zpts.shape[0]): - 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", - ) - - field = fold_ghost_cells(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field - - -@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") - - field_tiles = update_tiled_ghost_cells(field_tiles, world, g, tile_shape) - return field_tiles diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index c5be4eb9..a141be1c 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -2,7 +2,7 @@ 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 -from PyPIC3D.deposition.rho_tiled import ( +from PyPIC3D.deposition.rho import ( compute_mass_density_from_tiled_particles, compute_pressure_field_from_tiled_particles, compute_tiled_mass_density_from_tiled_particles, diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index 6f9c46ac..3cafe326 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -5,7 +5,6 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.current_methods import CURRENT_J_FROM_RHOV from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output from PyPIC3D.evolve import time_loop_electrodynamic @@ -109,7 +108,7 @@ def build_synthetic_tiled_yee_state( "y_wind": y_wind, "z_wind": z_wind, "shape_factor": shape, - "current_calculation": CURRENT_J_FROM_RHOV, + "use_esirkepov_current": False, "tile_shape": tile_shape, "guard_cells": g, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index da7ed53d..342fabe5 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -3,11 +3,11 @@ from functools import partial -import jax import jax.numpy as jnp from jax import jit -from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV +import jax + from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.particles.tiled_particle_refresh import ( @@ -122,9 +122,9 @@ def direct_current_step(state): # refresh tile ownership after the full position update. return particles, J_tiles, overflow - current_calculation = world.get("current_calculation", CURRENT_J_FROM_RHOV) + use_esirkepov_current = world.get("use_esirkepov_current", False) particles, J_tiles, overflow = jax.lax.cond( - current_calculation == CURRENT_ESIRKEPOV, + use_esirkepov_current, esirkepov_step, direct_current_step, (particles, J_tiles, overflow_previous), @@ -222,8 +222,6 @@ def time_loop_electrostatic( phi_tiles, solver, "periodic", - tile_shape, - g, ) # solve electrostatic fields from tiled charge density diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index ff59b3f4..dd2fcc96 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -54,9 +54,6 @@ from PyPIC3D.pusher.particle_push import validate_particle_pusher -from PyPIC3D.deposition.current_methods import ( - encode_current_calculation, -) from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov @@ -117,7 +114,9 @@ def _tile_shape_from_parameters(simulation_parameters): def _encode_current_calculation(current_calculation): - return 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'.") + return current_calculation == "esirkepov" def _validate_current_filter_contract(simulation_parameters): @@ -345,7 +344,7 @@ def initialize_simulation(toml_file): 'z_wind': z_wind, 'shape_factor': simulation_parameters['shape_factor'], 'guard_cells': guard_cells, - 'current_calculation': _encode_current_calculation(simulation_parameters['current_calculation']), + 'use_esirkepov_current': _encode_current_calculation(simulation_parameters['current_calculation']), 'boundary_conditions': { 'x': _encode_field_bc(simulation_parameters['x_bc']), 'y': _encode_field_bc(simulation_parameters['y_bc']), diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 1b94833f..9636686e 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -3,8 +3,7 @@ from jax import lax from functools import partial -from PyPIC3D.deposition.rho import _compute_rho_flat -from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles +from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient from PyPIC3D.utilities.filters import digital_filter from PyPIC3D.boundary_conditions.boundaryconditions import ( @@ -135,8 +134,6 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver bc_y = world['boundary_conditions']['y'] bc_z = world['boundary_conditions']['z'] - rho = _compute_rho_flat(particles, rho, world, constants) - phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) @@ -202,7 +199,7 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): return 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, tile_shape, g): +def calculate_tiled_electrostatic_fields(world, particles, species_config, constants, rho_tiles, phi_tiles, solver, bc): """ Compute electrostatic fields from tiled rho deposition and a global Poisson solve. @@ -217,8 +214,9 @@ def calculate_tiled_electrostatic_fields(world, particles, species_config, const bc_y = world["boundary_conditions"]["y"] bc_z = world["boundary_conditions"]["z"] - g = int(g) - rho_tiles = compute_tiled_rho_from_tiled_particles(particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + 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 = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) diff --git a/tests/particle_test.py b/tests/particle_test.py index 34aded4d..ebbf1355 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -12,13 +12,46 @@ from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.deposition.rho import _compute_rho_flat +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field from PyPIC3D.utils import build_yee_grid jax.config.update("jax_enable_x64", True) + +def _compute_rho_one_tile(particles, rho, world, constants): + world = dict(world) + world["guard_cells"] = int(world.get("guard_cells", 2)) + world["tile_shape"] = (world["Nx"], world["Ny"], world["Nz"]) + grids = dict(world["grids"]) + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + world["tile_shape"], + num_guard_cells=int(world["guard_cells"]), + ) + world["grids"] = grids + + simulation_parameters = { + "particle_tile_nx": world["Nx"], + "particle_tile_ny": world["Ny"], + "particle_tile_nz": world["Nz"], + } + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + rho_tiles = tile_scalar_field(rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + rho_tiles = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) + return assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + + class TestParticleMethods(unittest.TestCase): def setUp(self): """ @@ -498,7 +531,7 @@ def test_inactive_particles_do_not_deposit_rho(self): species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) constants = {"C": 3e8, "alpha": 1.0} - rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) + rho = _compute_rho_one_tile([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) self.assertAlmostEqual(float(jnp.sum(rho[1:-1, 1:-1, 1:-1]) * dx * dy * dz), 0.0) @@ -561,7 +594,7 @@ def test_rho(self): constants = {'C': 3e8, 'alpha' : 1.0} - num_rho = _compute_rho_flat([species], num_rho, world, constants) + num_rho = _compute_rho_one_tile([species], num_rho, world, constants) # compute rho self.assertTrue(jnp.allclose(num_rho, exp_rho)) @@ -609,7 +642,7 @@ def test_rho_periodic_boundary_exact_face(self): zwind=self.z_wind, ) - rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_one_tile([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 @@ -660,7 +693,7 @@ def test_rho_periodic_boundary_crossing(self): zwind=self.z_wind, ) - rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_one_tile([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) @@ -712,7 +745,7 @@ def test_rho_reduced_dimensions_periodic(self): zwind=z_wind, ) - rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 1.0}) + rho = _compute_rho_one_tile([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)) @@ -760,7 +793,7 @@ def test_filtered_rho_periodic_updates_ghost_cells(self): zwind=self.z_wind, ) - rho = _compute_rho_flat([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, {"C": 3e8, "alpha": 0.5}) + rho = _compute_rho_one_tile([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) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 024b09ea..e191202c 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -8,8 +8,7 @@ import toml from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.rho import _compute_rho_flat -from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles +from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.evolve import time_loop_electrostatic from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.species_class import particle_species @@ -18,6 +17,7 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_scalar_field, assemble_tiled_vector_field, + tile_grid_axes, tile_scalar_field, tile_vector_field, ) @@ -54,6 +54,27 @@ def _build_world(self, shape_factor=1): vertex_grid, center_grid = build_yee_grid(world) world["grids"] = {"vertex": vertex_grid, "center": center_grid} world["tile_shape"] = (2, 1, 1) + world = self._world_with_tiled_grids(world, world["tile_shape"]) + return world + + 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 + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids return world def _simulation_parameters(self): @@ -223,6 +244,7 @@ def _long_two_stream_state(self): "y_wind": 1.0, "z_wind": 1.0, "shape_factor": 1, + "guard_cells": 2, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, } @@ -391,18 +413,36 @@ def test_tiled_rho_assembles_to_global_rho(self): _, _, _, rho, _, _ = self._empty_fields(world) rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) - reference_rho = _compute_rho_flat([self._species(world)], rho, world, constants) - tiled_rho = compute_tiled_rho_from_tiled_particles( + tiled_rho = compute_rho( tiled_particles, species_config, rho_tiles, - world, constants, - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), + world, ) assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + reference_world = self._world_with_tiled_grids(world, self._one_tile_shape(world)) + reference_particles, reference_species_config = to_tiled_particles( + [self._species(reference_world)], + reference_world, + self._simulation_parameters_for_tile_shape(reference_world["tile_shape"]), + ) + reference_rho_tiles = tile_scalar_field(rho, reference_world, reference_world["tile_shape"], num_guard_cells=int(reference_world["guard_cells"])) + reference_rho_tiles = compute_rho( + reference_particles, + reference_species_config, + reference_rho_tiles, + constants, + reference_world, + ) + reference_rho = assemble_tiled_scalar_field( + reference_rho_tiles, + reference_world, + reference_world["tile_shape"], + num_guard_cells=int(reference_world["guard_cells"]), + ) + self.assertTrue( jnp.allclose( assembled_rho[1:-1, 1:-1, 1:-1], @@ -418,31 +458,33 @@ def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): E, B, J, rho, phi, external_fields = self._empty_fields(world) fields = (E, B, J, rho, phi, external_fields) reference_tile_shape = self._one_tile_shape(world) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, world["tile_shape"]) reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - self._neutral_species(world), + self._neutral_species(reference_world), fields, - world, + reference_world, reference_tile_shape, ) tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._neutral_species(world), + self._neutral_species(tiled_world), fields, - world, - world["tile_shape"], + tiled_world, + tiled_world["tile_shape"], ) reference_particles, reference_fields = time_loop_electrostatic( reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, unused_curl, J_func=None, solver="electrostatic", tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), + g=int(reference_world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -450,13 +492,13 @@ def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, unused_curl, J_func=None, solver="electrostatic", - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), + tile_shape=tiled_world["tile_shape"], + g=int(tiled_world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -467,8 +509,8 @@ def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): reference_tile_shape, tiled_particles, tiled_fields, - world, - world["tile_shape"], + tiled_world, + tiled_world["tile_shape"], step=1, ) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 81555c47..b3bc411b 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -17,8 +17,7 @@ get_1D_esirkepov_weights, get_2D_esirkepov_weights, ) -from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV -from PyPIC3D.deposition.rho_tiled import compute_tiled_rho_from_tiled_particles +from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.diagnostics.output_adapters import fields_for_output from PyPIC3D.initialization import initialize_simulation from PyPIC3D.particles.tiled_particle_refresh import ( @@ -72,7 +71,7 @@ def _build_world( "shape_factor": shape_factor, "boundary_conditions": boundary_conditions, "particle_boundary_conditions": particle_boundary_conditions, - "current_calculation": CURRENT_ESIRKEPOV, + "use_esirkepov_current": True, "guard_cells": 1, } center_grid = ( @@ -545,7 +544,7 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): g = int(world["guard_cells"]) rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) - rho_old = compute_tiled_rho_from_tiled_particles(tiled_particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + rho_old = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -555,7 +554,7 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): ) 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_tiled_rho_from_tiled_particles(new_particles, species_config, rho_tiles, world, constants, tile_shape=tile_shape, g=g) + 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"] @@ -625,7 +624,7 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): 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(int(world["current_calculation"]), CURRENT_ESIRKEPOV) + self.assertTrue(bool(world["use_esirkepov_current"])) def test_initialize_rejects_filtered_esirkepov_current(self): with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 9c941328..8d590ff2 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -14,10 +14,10 @@ _compute_pressure_field_flat, _compute_velocity_field_flat, ) -from PyPIC3D.deposition.rho_tiled import compute_tiled_mass_density_from_tiled_particles +from PyPIC3D.deposition.rho import compute_tiled_mass_density_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field from PyPIC3D.utils import build_yee_grid @@ -46,6 +46,31 @@ def _build_world(self, shape_factor=2): world["grids"] = {"vertex": vertex_grid, "center": center_grid} return world + def _world_with_tiled_grids(self, world): + 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 + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids + return world + def _empty_scalar(self, world): return jnp.zeros((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) @@ -127,11 +152,7 @@ def test_tiled_mass_density_matches_compute_mass_density_on_interior(self): def test_tile_major_mass_density_assembles_to_global_mass_density(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) rho_tiles = tile_scalar_field( @@ -168,11 +189,7 @@ def test_tile_major_mass_density_assembles_to_global_mass_density(self): def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) rho_tiles = tile_scalar_field( @@ -203,11 +220,7 @@ def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(s def test_tile_major_velocity_field_assembles_to_global_velocity_field(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) field_tiles = tile_scalar_field( @@ -243,11 +256,7 @@ def test_tile_major_velocity_field_assembles_to_global_velocity_field(self): def test_tile_major_pressure_field_assembles_to_global_pressure_field(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) field_tiles = tile_scalar_field( diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 21900bd2..d56b9139 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -7,6 +7,7 @@ from PyPIC3D import evolve from PyPIC3D.deposition import Esirkepov from PyPIC3D.deposition import J_from_rhov +from PyPIC3D.deposition import rho from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled from PyPIC3D.utils import build_yee_grid @@ -120,6 +121,18 @@ def test_esirkepov_api_reads_tile_metadata_from_world(self): def test_tiled_esirkepov_module_is_removed(self): self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.esirkepov_tiled")) + def test_rho_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(rho.compute_rho) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + + def test_flat_rho_and_tiled_rho_module_are_removed(self): + self.assertFalse(hasattr(rho, "_compute_rho_flat")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.rho_tiled")) + + def test_current_methods_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.current_methods")) + if __name__ == "__main__": unittest.main() diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 049c3c4f..cf0978fb 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -4,11 +4,10 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.rho import compute_rho, _compute_rho_flat -from PyPIC3D.deposition.rho_tiled import compute_rho_from_tiled_particles, compute_tiled_rho_from_tiled_particles +from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field from PyPIC3D.utils import build_yee_grid @@ -37,6 +36,31 @@ def _build_world(self, shape_factor, dt=0.08): 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 + grids["tiled_vertex_grid"] = tile_grid_axes( + grids["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + grids["tiled_center_grid"] = tile_grid_axes( + grids["center"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"] = grids + return world + def _empty_scalar(self, world): return jnp.zeros((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) @@ -47,6 +71,13 @@ def _simulation_parameters(self): "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 = particle_species( name="electrons", @@ -114,24 +145,34 @@ def _zero_species_velocities(self, particles): def _zero_tiled_velocities(self, tiled_particles): return tiled_particles._replace(u=jnp.zeros_like(tiled_particles.u)) - def _compare_tiled_to_standard(self, shape_factor, alpha): - world = self._build_world(shape_factor) - constants = {"alpha": alpha} - particles = self._particles(world) + 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_reference = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) - rho_tiled = compute_rho_from_tiled_particles( - tiled_particles, - species_config, + rho_tiles = tile_scalar_field( self._empty_scalar(world), world, - constants, + 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_tiled[1:-1, 1:-1, 1:-1], + 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, @@ -153,8 +194,8 @@ def test_compute_rho_uses_current_positions_not_half_step_back_positions(self): particles = self._particles(world) zero_velocity_particles = self._zero_species_velocities(self._particles(world)) - rho_with_velocity = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) - rho_with_zero_velocity = _compute_rho_flat(zero_velocity_particles, self._empty_scalar(world), world, constants) + _, 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( @@ -169,23 +210,15 @@ def test_tiled_global_rho_uses_current_positions_not_half_step_back_positions(se 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 = compute_rho_from_tiled_particles( - tiled_particles, - species_config, - self._empty_scalar(world), - world, - constants, - ) - rho_with_zero_velocity = compute_rho_from_tiled_particles( - zero_velocity_tiled_particles, - species_config, - self._empty_scalar(world), - world, - constants, - ) + 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( @@ -198,34 +231,26 @@ def test_tiled_global_rho_uses_current_positions_not_half_step_back_positions(se def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + 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_tiled_rho_from_tiled_particles( + rho_tiles_with_velocity = compute_rho( tiled_particles, species_config, rho_tiles, - world, constants, - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), + world, ) - rho_tiles_with_zero_velocity = compute_tiled_rho_from_tiled_particles( + rho_tiles_with_zero_velocity = compute_rho( zero_velocity_tiled_particles, species_config, rho_tiles, - world, constants, - tile_shape=world["tile_shape"], - g=int(world["guard_cells"]), + 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"])) @@ -241,11 +266,7 @@ def test_tile_major_rho_uses_current_positions_not_half_step_back_positions(self def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(self): world = self._build_world(shape_factor=2) - world["tile_shape"] = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + 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) @@ -256,13 +277,12 @@ def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(sel num_guard_cells=int(world["guard_cells"]), ) - rho_reference = _compute_rho_flat(particles, self._empty_scalar(world), world, constants) rho_tiles = compute_rho( tiled_particles, + species_config, rho_tiles, - world, constants, - species_config=species_config, + world, ) rho_from_tiles = assemble_tiled_scalar_field( rho_tiles, @@ -270,6 +290,7 @@ def test_compute_rho_dispatches_to_tile_major_deposition_for_tiled_particles(sel 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( @@ -285,7 +306,7 @@ def test_public_compute_rho_rejects_flat_particles(self): constants = {"alpha": 1.0} with self.assertRaisesRegex(ValueError, "TiledParticles"): - compute_rho(self._particles(world), self._empty_scalar(world), world, constants) + compute_rho(self._particles(world), None, self._empty_scalar(world), constants, world) if __name__ == "__main__": diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 350011ae..cdc9e114 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -10,7 +10,6 @@ from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov -from PyPIC3D.deposition.current_methods import CURRENT_ESIRKEPOV, CURRENT_J_FROM_RHOV from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species @@ -114,7 +113,7 @@ def _active_rows(self, tiled_particles): order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): + def _long_two_stream_state(self, use_esirkepov_current=False): world = { "Nx": 16, "Ny": 1, @@ -128,7 +127,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): "z_wind": 1.0, "shape_factor": 2, "guard_cells": 2, - "current_calculation": current_calculation, + "use_esirkepov_current": use_esirkepov_current, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, } @@ -190,7 +189,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) g = int(tiled_world["guard_cells"]) - if current_calculation == CURRENT_ESIRKEPOV: + if use_esirkepov_current: J_tiles = empty_tiled_vector_field(tiled_world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) @@ -212,7 +211,7 @@ def _long_two_stream_state(self, current_calculation=CURRENT_J_FROM_RHOV): (E, B, J, rho, phi, external_fields, pml_state), reference_world, reference_tile_shape, - current_calculation, + use_esirkepov_current, ) reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) @@ -269,7 +268,7 @@ def _tiled_species_rows(self, tiled_particles, species_index): order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def _build_tiled_state(self, particles, fields, world, tile_shape, current_calculation=CURRENT_J_FROM_RHOV): + def _build_tiled_state(self, particles, fields, world, tile_shape, use_esirkepov_current=False): simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -278,7 +277,7 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, current_calcu tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) E, B, J, rho, phi, external_fields, pml_state = fields g = int(world["guard_cells"]) - if current_calculation == CURRENT_ESIRKEPOV: + if use_esirkepov_current: J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) @@ -1271,7 +1270,7 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): constants, tile_shape, Nt, - ) = self._long_two_stream_state(CURRENT_J_FROM_RHOV) + ) = self._long_two_stream_state(False) reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world, reference_tile_shape, None) tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world, tile_shape, None) @@ -1340,7 +1339,7 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): constants, tile_shape, Nt, - ) = self._long_two_stream_state(CURRENT_ESIRKEPOV) + ) = self._long_two_stream_state(True) reference_loop = self._jitted_electrodynamic_loop_with_static_world( reference_world, From 49a2549194dd4a83e5d037460895c17b0778c012 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 06:00:07 -0500 Subject: [PATCH 38/83] refactoring evolve loop --- PyPIC3D/__main__.py | 24 +-- PyPIC3D/diagnostics/tiled_benchmarking.py | 50 +++--- PyPIC3D/evolve.py | 162 +++++++---------- PyPIC3D/initialization.py | 31 +--- PyPIC3D/solvers/first_order_yee.py | 16 +- PyPIC3D/solvers/yee_tiled.py | 16 +- tests/evolve_test.py | 26 +-- tests/initialization_test.py | 4 +- tests/pml_test.py | 38 ++-- tests/tiled_electrostatic_test.py | 48 ++--- tests/tiled_esirkepov_test.py | 27 ++- tests/tiled_refactor_contract_test.py | 19 +- tests/tiled_yee_integration_test.py | 202 ++++++++-------------- tests/yee_tiled_test.py | 106 +++++------- 14 files changed, 301 insertions(+), 468 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index ca4edb9a..ad69c15a 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -60,7 +60,7 @@ def _raise_if_tiled_particles_overflowed(fields, simulation_parameters): 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, species_config = 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 dt = world['dt'] @@ -79,30 +79,34 @@ def loop_with_static_world( species_config, fields, constants, - curl_func, - J_func, 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, - curl_func, - J_func, - solver, - tile_shape=tile_shape, - g=g, relativistic=relativistic, particle_pusher=particle_pusher, ) jit_loop = jax.jit( loop_with_static_world, - static_argnames=('curl_func', 'J_func', 'solver', 'relativistic', 'particle_pusher'), + static_argnames=('solver', 'relativistic', 'particle_pusher'), ) E, B, J, rho, phi, external_fields, *rest = fields @@ -201,8 +205,6 @@ def loop_with_static_world( species_config, fields, constants, - curl_func, - J_func, solver, relativistic=relativistic, particle_pusher=particle_pusher, diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index 3cafe326..93bb4b40 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -1,10 +1,9 @@ from dataclasses import dataclass -from functools import partial - import jax import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC +from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output from PyPIC3D.evolve import time_loop_electrodynamic @@ -51,6 +50,13 @@ def unused_curl(Ex, Ey, Ez): return None +def _particle_kernel_world(world): + kernel_world = dict(world) + kernel_world.pop("current_deposition", None) + kernel_world.pop("current_filter", None) + return kernel_world + + def _pad_tiled_particle_capacity(tiled_particles, slots_per_tile): current_slots = tiled_particles.active.shape[-1] if current_slots >= slots_per_tile: @@ -108,7 +114,8 @@ def build_synthetic_tiled_yee_state( "y_wind": y_wind, "z_wind": z_wind, "shape_factor": shape, - "use_esirkepov_current": False, + "current_deposition": "direct", + "current_filter": "none", "tile_shape": tile_shape, "guard_cells": g, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, @@ -190,7 +197,7 @@ def build_synthetic_tiled_yee_state( tuple(jnp.zeros_like(component) for component in E_tiles), tuple(jnp.zeros_like(component) for component in B_tiles), ) - fields = (E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, None) + fields = (E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, None, jnp.asarray(False)) return SyntheticTiledYeeState( particles=tiled_particles, @@ -212,7 +219,7 @@ def _tiled_particle_push_stage(case, relativistic=True, particle_pusher="boris") case.species_config, push_E_tiles, push_B_tiles, - case.world, + _particle_kernel_world(case.world), case.constants, case.tile_shape, case.g, @@ -227,23 +234,26 @@ def _tiled_particle_retile_stage(case): return refresh_tiled_particle_tiles(moved, case.world, case.tile_shape) -def _tiled_current_deposition_stage(case, particles, J_func): +def _tiled_current_deposition_stage(case, particles): _, _, J_tiles, *_ = case.fields with jax.named_scope("tiled_current_deposition"): - return J_func( + if case.world["current_deposition"] == "esirkepov": + return Esirkepov_current(particles, case.species_config, J_tiles, case.constants, case.world) + return J_from_rhov( particles, case.species_config, J_tiles, case.constants, case.world, + filter=case.world["current_filter"], ) def _tiled_field_update_stage(case, J_tiles): E_tiles, B_tiles, *_ = case.fields with jax.named_scope("tiled_field_update"): - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, case.world, case.constants, unused_curl, case.tile_shape, case.g) - B_tiles = update_tiled_B(E_tiles, B_tiles, case.world, case.constants, unused_curl, case.tile_shape, case.g) + E_tiles, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, case.world, case.constants) + B_tiles, _pml_state = update_tiled_B(E_tiles, B_tiles, case.world, case.constants, pml_state) return E_tiles, B_tiles @@ -269,7 +279,7 @@ def _tiled_output_bridge_stage(case): return E, B, particles -def benchmark_tiled_pic_step(case, J_func, relativistic=True, particle_pusher="boris"): +def benchmark_tiled_pic_step(case, relativistic=True, particle_pusher="boris"): with jax.named_scope("tiled_pic_step"): return time_loop_electrodynamic( case.particles, @@ -277,17 +287,12 @@ def benchmark_tiled_pic_step(case, J_func, relativistic=True, particle_pusher="b case.fields, case.world, case.constants, - unused_curl, - J_func, - "electrodynamic_yee", - tile_shape=case.tile_shape, - g=case.g, relativistic=relativistic, particle_pusher=particle_pusher, ) -def build_tiled_stage_specs(case, J_func=None, relativistic=True, particle_pusher="boris"): +def build_tiled_stage_specs(case, relativistic=True, particle_pusher="boris"): """ Return independently timed tiled Yee stage specs. @@ -296,19 +301,16 @@ def build_tiled_stage_specs(case, J_func=None, relativistic=True, particle_pushe the current global output bridge. """ - if J_func is None: - J_func = partial(J_from_rhov, filter="none") - pushed_particles = _tiled_particle_push_stage(case, relativistic=relativistic, particle_pusher=particle_pusher) - deposited_J = _tiled_current_deposition_stage(case, pushed_particles, J_func) + deposited_J = _tiled_current_deposition_stage(case, pushed_particles) return { "tiled_pic_step": StageSpec( "tiled_pic_step", benchmark_tiled_pic_step, - (case, J_func), + (case,), {"relativistic": relativistic, "particle_pusher": particle_pusher}, - ("J_func", "relativistic", "particle_pusher"), + ("relativistic", "particle_pusher"), ), "tiled_particle_push": StageSpec( "tiled_particle_push", @@ -327,9 +329,9 @@ def build_tiled_stage_specs(case, J_func=None, relativistic=True, particle_pushe "tiled_current_deposition": StageSpec( "tiled_current_deposition", _tiled_current_deposition_stage, - (case, pushed_particles, J_func), + (case, pushed_particles), {}, - ("J_func",), + (), ), "tiled_field_update": StageSpec( "tiled_field_update", diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 342fabe5..79eaf49a 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -3,7 +3,6 @@ from functools import partial -import jax.numpy as jnp from jax import jit import jax @@ -23,17 +22,19 @@ __all__ = ["time_loop_electrodynamic", "time_loop_electrostatic"] +def _particle_kernel_world(world): + kernel_world = dict(world) + kernel_world.pop("current_deposition", None) + kernel_world.pop("current_filter", None) + return kernel_world + + def time_loop_electrodynamic( particles, species_config, fields, world, constants, - curl_func, - J_func, - solver, - tile_shape=None, - g=None, relativistic=True, particle_pusher="boris", ): @@ -41,45 +42,23 @@ def time_loop_electrodynamic( Advance a tiled electrodynamic PIC system by one time step. """ - del solver - - if len(fields) == 7: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state = fields - overflow_previous = jnp.asarray(False) - else: - E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow_previous = fields - # unpack the tile-major field state. ``overflow_previous`` carries the - # particle-retile overflow diagnostic out to the Python driver. - - if tile_shape is None or g is None: - raise ValueError("tiled electrodynamic updates require explicit tile_shape and g.") - g = int(g) - - if not hasattr(particles, "active"): - if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) - else: - E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - B_tiles, pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state) - return particles, fields - # keep field-only helper behavior for standalone Maxwell equivalence checks. + E, B, J, rho, phi, external_fields, pml_state, overflow_previous = fields + # unpack the tiled field state - push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, external_fields) - # particles see evolved fields plus external-only fields, as in the standard - # electrodynamic path. + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + particle_world = _particle_kernel_world(world) + # get tiled layout information from the world contract + + push_E, push_B = add_external_fields(E, B, external_fields) + # particles see evolved fields plus external-only fields particles = tiled_particle_push( particles, species_config, - push_E_tiles, - push_B_tiles, - world, + push_E, + push_B, + particle_world, constants, tile_shape, g, @@ -88,63 +67,60 @@ def time_loop_electrodynamic( ) # use the selected tiled pusher for particle velocities - if J_func is None: - current_deposition = partial(J_from_rhov, filter="none") - else: - current_deposition = J_func - - def esirkepov_step(state): - particles, J_tiles, overflow_previous = state - # Esirkepov needs old and new particle positions. The deposition kernel - # predicts the new positions locally, then the actual particle state is - # advanced and retiled after the current has been computed. - J_tiles = Esirkepov_current(particles, species_config, J_tiles, constants, world) - particles = update_tiled_particle_positions(particles, species_config, world["dt"]) - particles, overflow = refresh_tiled_particle_tiles(particles, world, tile_shape) - overflow = overflow_previous | overflow - return particles, J_tiles, overflow - - def direct_current_step(state): + 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 - # wrap periodic particles and move them into their owning tiles. - - J_tiles = current_deposition(particles, species_config, J_tiles, constants, world) + # 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) - overflow = overflow_previous | overflow # 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 - use_esirkepov_current = world.get("use_esirkepov_current", False) - particles, J_tiles, overflow = jax.lax.cond( - use_esirkepov_current, - esirkepov_step, - direct_current_step, - (particles, J_tiles, overflow_previous), + 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 - if pml_state is None: - E_tiles = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g) - else: - E_tiles, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state - ) - # update electric field from B and the supplied tiled current + E, pml_state = update_tiled_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_tiled_E function returns None for the pml_state - if pml_state is None: - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g) - else: - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state) + B, pml_state = update_tiled_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_tiled_B function returns None for the pml_state - fields = (E_tiles, B_tiles, J_tiles, rho, phi, external_fields, pml_state, overflow) + fields = (E, B, J, rho, phi, external_fields, pml_state, overflow) # pack the tiled field state return particles, fields @@ -156,11 +132,7 @@ def time_loop_electrostatic( fields, world, constants, - curl_func, - J_func, solver, - tile_shape=None, - g=None, relativistic=True, particle_pusher="boris", ): @@ -173,21 +145,13 @@ def time_loop_electrostatic( tile-local electrostatic E. """ - del curl_func, J_func - - if len(fields) == 6: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields = fields - pml_state = None - overflow_previous = jnp.asarray(False) - elif len(fields) == 7: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state = fields - overflow_previous = jnp.asarray(False) - else: - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, pml_state, overflow_previous = fields + # unpack the tiled field state - if tile_shape is None or g is None: - raise ValueError("tiled electrostatic updates require explicit tile_shape and g.") - g = int(g) + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + 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 @@ -197,7 +161,7 @@ def time_loop_electrostatic( species_config, push_E_tiles, push_B_tiles, - world, + particle_world, constants, tile_shape, g, diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index dd2fcc96..93da958c 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 @@ -22,11 +21,6 @@ make_dir, compute_energy, add_external_fields ) -from PyPIC3D.solvers.fdtd import ( - centered_finite_difference_curl -) - - from PyPIC3D.diagnostics.plotting import ( plot_initial_histograms ) @@ -54,9 +48,6 @@ 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.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 ( @@ -116,7 +107,9 @@ def _tile_shape_from_parameters(simulation_parameters): 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'.") - return current_calculation == "esirkepov" + if current_calculation == "esirkepov": + return "esirkepov" + return "direct" def _validate_current_filter_contract(simulation_parameters): @@ -267,7 +260,6 @@ def initialize_simulation(toml_file): 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. @@ -344,7 +336,8 @@ def initialize_simulation(toml_file): 'z_wind': z_wind, 'shape_factor': simulation_parameters['shape_factor'], 'guard_cells': guard_cells, - 'use_esirkepov_current': _encode_current_calculation(simulation_parameters['current_calculation']), + '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']), @@ -462,9 +455,6 @@ def initialize_simulation(toml_file): external_fields = (external_E, external_B) # fill ghost cells for external fields before they are interpolated to particles - 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 @@ -493,10 +483,8 @@ def initialize_simulation(toml_file): if simulation_parameters['current_calculation'] == "esirkepov": print("Using Esirkepov current calculation method") - 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']) species_config = None @@ -520,10 +508,8 @@ def initialize_simulation(toml_file): B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) if simulation_parameters["current_calculation"] == "esirkepov": J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) - J_func = Esirkepov_current else: J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) - J_func = functools.partial(J_from_rhov, filter=simulation_parameters['filter_j']) external_E, external_B = external_fields external_fields = ( tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), @@ -532,10 +518,11 @@ def initialize_simulation(toml_file): simulation_parameters["tile_shape"] = tile_shape print(f"Using tiled Yee storage with tile shape: {tile_shape}") + overflow = jnp.asarray(False) if electrostatic: rho = tile_scalar_field(rho, world, tile_shape, num_guard_cells=guard_cells) phi = tile_scalar_field(phi, world, tile_shape, num_guard_cells=guard_cells) - fields = (E, B, J, rho, phi, external_fields) + fields = (E, B, J, rho, phi, external_fields, None, overflow) # define the fields tuple for the electrostatic solver else: pml_state = None @@ -543,7 +530,7 @@ def initialize_simulation(toml_file): # ordinary, unstretched Yee derivatives. if pml_active: pml_state = initialize_tiled_pml_state(world, tile_shape) - fields = (E, B, J, rho, phi, external_fields, pml_state) + 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']: @@ -560,7 +547,7 @@ def initialize_simulation(toml_file): return evolve_loop, particles, fields, world, simulation_parameters, constants, plotting_parameters, plasma_parameters, \ - solver, electrostatic, verbose, GPUs, Nt, curl_func, J_func, relativistic, particle_pusher, species_config + solver, electrostatic, verbose, GPUs, Nt, relativistic, particle_pusher, species_config diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py index ee622c0d..97866f17 100644 --- a/PyPIC3D/solvers/first_order_yee.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -211,7 +211,7 @@ def _update_B_global(E, B, world, constants, curl_func, pml_state=None): return (Bx, By, Bz), pml_state -def update_E(E, B, J, world, constants, curl_func, pml_state=None): +def update_E(E, B, J, world, constants, pml_state=None): """ Update electric fields through the public Yee helper. @@ -222,16 +222,12 @@ def update_E(E, B, J, world, constants, curl_func, pml_state=None): """ if E[0].ndim == 6: - tile_shape = tuple(world["tile_shape"]) - g = int(world["guard_cells"]) - if pml_state is None: - return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g), None - return update_tiled_E(E, B, J, world, constants, curl_func, tile_shape, g, pml_state) + return update_tiled_E(E, B, J, world, constants, pml_state) raise ValueError("Public update_E requires tiled field arrays; use _update_E_global for flat reference tests.") -def update_B(E, B, world, constants, curl_func, pml_state=None): +def update_B(E, B, world, constants, pml_state=None): """ Update magnetic fields through the public Yee helper. @@ -240,10 +236,6 @@ def update_B(E, B, world, constants, curl_func, pml_state=None): """ if E[0].ndim == 6: - tile_shape = tuple(world["tile_shape"]) - g = int(world["guard_cells"]) - if pml_state is None: - return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g), None - return update_tiled_B(E, B, world, constants, curl_func, tile_shape, g, pml_state) + return update_tiled_B(E, B, world, constants, pml_state) raise ValueError("Public update_B requires tiled field arrays; use _update_B_global for flat reference tests.") diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 2f548092..043805d0 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -664,7 +664,7 @@ def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): return _restore_tiled_vector_layout(folded, field_tiles) -def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_shape, g, pml_state=None): +def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -672,10 +672,10 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ have been refreshed from neighbor tiles or field boundary conditions. """ - del curl_func - Ex, Ey, Ez = E_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 = _active_slice(g) forward = _forward_slice(g) @@ -737,12 +737,12 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, curl_func, tile_ Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) if pml_state is None: - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) + return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape), None return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state -def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, pml_state=None): +def update_tiled_B(E_tiles, B_tiles, world, constants, pml_state=None): """ Update compact tiled magnetic fields without assembling a global field. @@ -750,10 +750,10 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, have been refreshed from neighbor tiles or field boundary conditions. """ - del curl_func - 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 = _active_slice(g) backward = _backward_slice(g) @@ -799,6 +799,6 @@ def update_tiled_B(E_tiles, B_tiles, world, constants, curl_func, tile_shape, g, Bx, By, Bz = digital_filter_vector((Bx, By, Bz), constants.get("alpha", 1.0), num_guard_cells=g) if pml_state is None: - return update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) + return update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape), None return update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape), pml_state diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 9d4b9e29..ca0a5ba0 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -16,14 +16,6 @@ jax.config.update("jax_enable_x64", True) -def zero_current(particles, species_config, J, constants, world): - return tuple(jnp.zeros_like(comp) for comp in J) - - -def unused_curl(Ex, Ey, Ez): - return None - - def add_tiled_grids_to_world(world, tile_shape): g = int(world["guard_cells"]) world["tile_shape"] = tile_shape @@ -59,6 +51,8 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field "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}, } @@ -113,6 +107,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field phi, (external_E_tiles, external_B_tiles), None, + jnp.asarray(False), ) tiled_particles, fields = time_loop_electrodynamic( @@ -121,11 +116,6 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field fields, world, constants, - unused_curl, - zero_current, - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -140,8 +130,6 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field 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_E_tiles[0])) def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): @@ -157,6 +145,8 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): "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}, } @@ -213,6 +203,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])), ), None, + jnp.asarray(False), ) tiled_particles, _ = time_loop_electrodynamic( @@ -221,11 +212,6 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): fields, world, constants, - unused_curl, - zero_current, - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) diff --git a/tests/initialization_test.py b/tests/initialization_test.py index c538f6e1..5f64be6e 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -90,10 +90,12 @@ def test_initialize_simulation_returns_tiled_runtime_for_ordinary_electrodynamic 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 = fields + 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)) def test_initialize_simulation_encodes_global_particle_boundary_conditions(self): with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/pml_test.py b/tests/pml_test.py index b3e41fc7..98c73279 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -212,7 +212,7 @@ def test_initialize_simulation_appends_pml_state_for_electrodynamic_yee(self): fields = result[2] world = result[3] active, pml_x, pml_y, pml_z, _ = world["pml"] - self.assertEqual(len(fields), 7) + self.assertEqual(len(fields), 8) self.assertIn("pml", world) self.assertTrue(active) self.assertTrue(pml_x) @@ -236,10 +236,12 @@ def test_initialize_simulation_uses_tiled_pml_state_for_electrodynamic_yee(self) fields = result[2] world = result[3] - pml_state = fields[-1] + pml_state = fields[6] + overflow = fields[-1] e_memory, b_memory, tiled_profiles = pml_state - self.assertEqual(len(fields), 7) + 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)) @@ -254,8 +256,9 @@ def test_initialize_simulation_uses_none_pml_state_without_pml_for_electrodynami world = result[3] active, pml_x, pml_y, pml_z, _ = world["pml"] - self.assertEqual(len(fields), 7) - self.assertIsNone(fields[-1]) + 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) @@ -346,6 +349,7 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): E_reference, pml_state = _update_E_global(E, B, J, world, constants, lambda *args: None, pml_state) B_reference, pml_state = _update_B_global(E_reference, B, world, constants, lambda *args: None, pml_state) + world["tile_shape"] = tile_shape tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) E_tiles, tiled_pml_state = update_tiled_E( tile_vector_field(E, world, tile_shape), @@ -353,9 +357,6 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): tile_vector_field(J, world, tile_shape), world, constants, - lambda *args: None, - tile_shape, - int(world["guard_cells"]), tiled_pml_state, ) B_tiles, tiled_pml_state = update_tiled_B( @@ -363,9 +364,6 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): tile_vector_field(B, world, tile_shape), world, constants, - lambda *args: None, - tile_shape, - int(world["guard_cells"]), tiled_pml_state, ) @@ -390,6 +388,7 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): constants, ) tile_shape = (4, 1, 1) + world["tile_shape"] = tile_shape tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) @@ -405,12 +404,8 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) def step(E_tiles, B_tiles, tiled_pml_state): - E_tiles, tiled_pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, int(world["guard_cells"]), tiled_pml_state - ) - B_tiles, tiled_pml_state = update_tiled_B( - E_tiles, B_tiles, world, constants, lambda *args: None, tile_shape, int(world["guard_cells"]), tiled_pml_state - ) + E_tiles, tiled_pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, tiled_pml_state) + B_tiles, tiled_pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, tiled_pml_state) return E_tiles, B_tiles, tiled_pml_state step = jax.jit(step) @@ -430,6 +425,7 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): constants, ) tile_shape = (2, 1, 1) + world["tile_shape"] = tile_shape E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) Ex, Ey, Ez = E Bx, By, Bz = B @@ -445,12 +441,8 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): J_tiles = (Jx, Jy, Jz) pml_state = initialize_tiled_pml_state(world, tile_shape) - E_after, pml_state = update_tiled_E( - E_tiles, B_tiles, J_tiles, world, constants, lambda *args: None, tile_shape, g, pml_state - ) - B_after, pml_state = update_tiled_B( - E_after, B_tiles, world, constants, lambda *args: None, tile_shape, g, pml_state - ) + E_after, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state) + B_after, pml_state = update_tiled_B(E_after, B_tiles, world, constants, pml_state) for component in E_after + B_after: self.assertTrue(jnp.all(jnp.isfinite(component))) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index e191202c..86083aaa 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -30,10 +30,6 @@ ROUND_OFF_ATOL = 1.0e-11 -def unused_curl(Ex, Ey, Ez): - return None - - class TestTiledElectrostatic(unittest.TestCase): def _build_world(self, shape_factor=1): world = { @@ -192,6 +188,7 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, min_slots=Non tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), ), None, + jnp.asarray(False), ) return tiled_particles, species_config, tiled_fields @@ -254,21 +251,22 @@ def _long_two_stream_state(self): constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (4, 1, 1) reference_tile_shape = self._one_tile_shape(world) - world["tile_shape"] = tile_shape + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) Nt = 1500 E, B, J, rho, phi, external_fields = self._empty_fields(world) fields = (E, B, J, rho, phi, external_fields) reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - self._long_two_stream_species(world), + self._long_two_stream_species(reference_world), fields, - world, + reference_world, reference_tile_shape, ) tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._long_two_stream_species(world), + self._long_two_stream_species(tiled_world), fields, - world, + tiled_world, tile_shape, min_slots=12, ) @@ -279,7 +277,7 @@ def _long_two_stream_state(self): reference_tile_shape, tiled_particles, tiled_fields, - world, + tiled_world, tile_shape, step=0, ) @@ -292,7 +290,8 @@ def _long_two_stream_state(self): tiled_particles, species_config, tiled_fields, - world, + reference_world, + tiled_world, constants, tile_shape, Nt, @@ -480,11 +479,7 @@ def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): reference_fields, reference_world, constants, - unused_curl, - J_func=None, solver="electrostatic", - tile_shape=reference_tile_shape, - g=int(reference_world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -494,11 +489,7 @@ def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): tiled_fields, tiled_world, constants, - unused_curl, - J_func=None, solver="electrostatic", - tile_shape=tiled_world["tile_shape"], - g=int(tiled_world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -527,7 +518,8 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se tiled_particles, species_config, tiled_fields, - world, + reference_world, + tiled_world, constants, tile_shape, Nt, @@ -535,7 +527,7 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se tiled_loop = jax.jit( time_loop_electrostatic, - static_argnames=("curl_func", "J_func", "solver", "tile_shape", "g", "relativistic", "particle_pusher"), + static_argnames=("solver", "relativistic", "particle_pusher"), ) for step in range(1, Nt + 1): @@ -543,13 +535,9 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se reference_particles, reference_species_config, reference_fields, - world, + reference_world, constants, - unused_curl, - J_func=None, solver="electrostatic", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -557,13 +545,9 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se tiled_particles, species_config, tiled_fields, - world, + tiled_world, constants, - unused_curl, - J_func=None, solver="electrostatic", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -574,7 +558,7 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se reference_tile_shape, tiled_particles, tiled_fields, - world, + tiled_world, tile_shape, step, ) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index b3bc411b..6e3140d8 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -1,7 +1,6 @@ import os import tempfile import unittest -import functools import jax import jax.numpy as jnp @@ -71,7 +70,8 @@ def _build_world( "shape_factor": shape_factor, "boundary_conditions": boundary_conditions, "particle_boundary_conditions": particle_boundary_conditions, - "use_esirkepov_current": True, + "current_deposition": "esirkepov", + "current_filter": "none", "guard_cells": 1, } center_grid = ( @@ -301,8 +301,10 @@ def test_update_tiled_E_reads_two_guard_current_interior(self): Jx, Jy, Jz = J_tiles Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) - E_after = update_tiled_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants, None, tile_shape, g) + world["tile_shape"] = tile_shape + E_after, pml_state = update_tiled_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_reduced_axis_esirkepov_scatter_uses_only_collapsed_stencil_index(self): @@ -624,7 +626,8 @@ def test_initialize_tiled_yee_esirkepov_uses_two_guard_current_tiles(self): self.assertEqual(J_tiles[0].shape[-3:], (6, 5, 5)) self.assertEqual(int(world["guard_cells"]), 2) self.assertNotIn("current_guard_cells", world) - self.assertTrue(bool(world["use_esirkepov_current"])) + 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: @@ -731,7 +734,7 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): toml.dump(config, f) loop, particles, fields, world, _simulation_parameters, constants, _plotting_parameters, _plasma_parameters, \ - solver, _electrostatic, _verbose, _GPUs, _Nt, curl_func, J_func, relativistic, particle_pusher, species_config = initialize_simulation(toml.load(config_path)) + solver, _electrostatic, _verbose, _GPUs, _Nt, relativistic, particle_pusher, species_config = initialize_simulation(toml.load(config_path)) particles, fields = loop( particles, @@ -739,11 +742,6 @@ def test_tiled_yee_esirkepov_loop_advances_particles_after_deposition(self): fields, world, constants, - curl_func, - J_func, - solver, - tile_shape=tuple(int(width) for width in world["tile_shape"]), - g=int(world["guard_cells"]), relativistic=relativistic, particle_pusher=particle_pusher, ) @@ -806,19 +804,14 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s } loop, particles, fields, world, _simulation_parameters, constants, _plotting_parameters, _plasma_parameters, \ - solver, _electrostatic, _verbose, _GPUs, _Nt, curl_func, _J_func, relativistic, particle_pusher, species_config = initialize_simulation(config) - alias_J_func = functools.partial(Esirkepov_current) + 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, - curl_func, - alias_J_func, - solver, - tile_shape=tuple(int(width) for width in world["tile_shape"]), - g=int(world["guard_cells"]), relativistic=relativistic, particle_pusher=particle_pusher, ) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index d56b9139..2a9b2da0 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -88,8 +88,23 @@ def test_tiled_benchmark_stage_specs_expose_tiled_hot_paths(self): "tiled_output_bridge", } self.assertTrue(expected.issubset(stage_specs)) - self.assertEqual(stage_specs["tiled_pic_step"].static_argnames, ("J_func", "relativistic", "particle_pusher")) - self.assertEqual(stage_specs["tiled_current_deposition"].static_argnames, ("J_func",)) + self.assertEqual(stage_specs["tiled_pic_step"].static_argnames, ("relativistic", "particle_pusher")) + self.assertEqual(stage_specs["tiled_current_deposition"].static_argnames, ()) + + def test_evolve_loops_read_runtime_controls_from_world(self): + electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) + electrostatic_signature = inspect.signature(evolve.time_loop_electrostatic) + for name in ("curl_func", "J_func", "tile_shape", "g", "current_deposition"): + self.assertNotIn(name, electrodynamic_signature.parameters) + for name in ("curl_func", "J_func", "tile_shape", "g"): + self.assertNotIn(name, electrostatic_signature.parameters) + + def test_tiled_yee_updates_read_tile_metadata_from_world(self): + update_E_signature = inspect.signature(yee_tiled.update_tiled_E) + update_B_signature = inspect.signature(yee_tiled.update_tiled_B) + for name in ("curl_func", "tile_shape", "g"): + self.assertNotIn(name, update_E_signature.parameters) + self.assertNotIn(name, update_B_signature.parameters) def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): source = inspect.getsource(evolve.time_loop_electrodynamic) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index cdc9e114..1448453d 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -1,4 +1,3 @@ -import functools import os import tempfile import unittest @@ -8,14 +7,19 @@ import numpy as np import toml -from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.J_from_rhov import J_from_rhov from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_vector_field, empty_tiled_vector_field, tile_grid_axes, tile_vector_field +from PyPIC3D.solvers.yee_tiled import ( + assemble_tiled_vector_field, + empty_tiled_vector_field, + tile_grid_axes, + tile_vector_field, + update_tiled_B, + update_tiled_E, +) from PyPIC3D.utils import build_yee_grid, compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -27,10 +31,6 @@ ROUND_OFF_ATOL = 1.0e-11 -def unused_curl(Ex, Ey, Ez): - return None - - class TestTiledYeeIntegration(unittest.TestCase): def _build_world(self, boundary_conditions=None): if boundary_conditions is None: @@ -48,6 +48,8 @@ def _build_world(self, boundary_conditions=None): "z_wind": 1.0, "shape_factor": 1, "guard_cells": 2, + "current_deposition": "direct", + "current_filter": "none", "boundary_conditions": boundary_conditions, } vertex_grid, center_grid = build_yee_grid(world) @@ -113,7 +115,7 @@ def _active_rows(self, tiled_particles): order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def _long_two_stream_state(self, use_esirkepov_current=False): + def _long_two_stream_state(self, current_deposition="direct", current_filter="none"): world = { "Nx": 16, "Ny": 1, @@ -127,7 +129,8 @@ def _long_two_stream_state(self, use_esirkepov_current=False): "z_wind": 1.0, "shape_factor": 2, "guard_cells": 2, - "use_esirkepov_current": use_esirkepov_current, + "current_deposition": current_deposition, + "current_filter": current_filter, "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, } @@ -189,7 +192,7 @@ def _long_two_stream_state(self, use_esirkepov_current=False): tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) g = int(tiled_world["guard_cells"]) - if use_esirkepov_current: + if current_deposition == "esirkepov": J_tiles = empty_tiled_vector_field(tiled_world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) @@ -205,13 +208,14 @@ def _long_two_stream_state(self, use_esirkepov_current=False): tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=g), ), None, + jnp.asarray(False), ) reference_particles, reference_species_config, reference_fields = self._build_tiled_state( reference_species, (E, B, J, rho, phi, external_fields, pml_state), reference_world, reference_tile_shape, - use_esirkepov_current, + current_deposition, ) reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) @@ -268,7 +272,7 @@ def _tiled_species_rows(self, tiled_particles, species_index): order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def _build_tiled_state(self, particles, fields, world, tile_shape, use_esirkepov_current=False): + def _build_tiled_state(self, particles, fields, world, tile_shape, current_deposition="direct"): simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -277,7 +281,7 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, use_esirkepov tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) E, B, J, rho, phi, external_fields, pml_state = fields g = int(world["guard_cells"]) - if use_esirkepov_current: + if current_deposition == "esirkepov": J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) @@ -292,19 +296,16 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, use_esirkepov tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), ), pml_state, + jnp.asarray(False), ) return tiled_particles, species_config, tiled_fields - def _jitted_electrodynamic_loop_with_static_world(self, world, tile_shape, J_func): - g = int(world["guard_cells"]) - + def _jitted_electrodynamic_loop_with_static_world(self, world): def loop_with_static_world( particles, species_config, fields, constants, - curl_func, - solver, relativistic=True, particle_pusher="boris", ): @@ -314,18 +315,13 @@ def loop_with_static_world( fields, world, constants, - curl_func, - J_func=J_func, - solver=solver, - tile_shape=tile_shape, - g=g, relativistic=relativistic, particle_pusher=particle_pusher, ) return jax.jit( loop_with_static_world, - static_argnames=("curl_func", "solver", "relativistic", "particle_pusher"), + static_argnames=("relativistic", "particle_pusher"), ) def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): @@ -452,71 +448,65 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): phi = seed external_fields = ((seed, seed, seed), (seed, seed, seed)) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) reference_fields = ( - tile_vector_field(E, world, reference_tile_shape), - tile_vector_field(B, world, reference_tile_shape), - tile_vector_field(J, world, reference_tile_shape), + tile_vector_field(E, reference_world, reference_tile_shape), + tile_vector_field(B, reference_world, reference_tile_shape), + tile_vector_field(J, reference_world, reference_tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, reference_tile_shape), - tile_vector_field(external_fields[1], world, reference_tile_shape), + tile_vector_field(external_fields[0], reference_world, reference_tile_shape), + tile_vector_field(external_fields[1], reference_world, reference_tile_shape), ), None, + jnp.asarray(False), ) - _, reference_fields = time_loop_electrodynamic( - [], - None, - reference_fields, - world, + reference_E, reference_pml_state = update_tiled_E( + reference_fields[0], + reference_fields[1], + reference_fields[2], + reference_world, constants, - unused_curl, - J_func=lambda particles, species_config, J, constants, world: J, - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), - relativistic=False, - particle_pusher="boris", ) + reference_B, reference_pml_state = update_tiled_B( + reference_E, + reference_fields[1], + reference_world, + constants, + reference_pml_state, + ) + reference_fields = (reference_E, reference_B, *reference_fields[2:]) tiled_fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), rho, phi, ( - tile_vector_field(external_fields[0], world, tile_shape), - tile_vector_field(external_fields[1], world, tile_shape), + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, + jnp.asarray(False), ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - [], - None, - tiled_fields, - world, - constants, - unused_curl, - J_func=lambda particles, species_config, J, constants, world: J, - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), - relativistic=False, - particle_pusher="boris", - ) + tiled_E, tiled_pml_state = update_tiled_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) + tiled_B, tiled_pml_state = update_tiled_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) + tiled_fields = (tiled_E, tiled_B, *tiled_fields[2:]) reference_E, reference_B, *_ = reference_fields E_tiles, B_tiles, *_ = tiled_fields self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, world, reference_tile_shape), - assemble_tiled_vector_field(E_tiles, world, tile_shape), + assemble_tiled_vector_field(reference_E, reference_world, reference_tile_shape), + assemble_tiled_vector_field(E_tiles, tiled_world, tile_shape), 1, "E", ) self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, world, reference_tile_shape), - assemble_tiled_vector_field(B_tiles, world, tile_shape), + assemble_tiled_vector_field(reference_B, reference_world, reference_tile_shape), + assemble_tiled_vector_field(B_tiles, tiled_world, tile_shape), 1, "B", ) @@ -849,6 +839,7 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, + jnp.asarray(False), ) reference_particles, reference_fields = time_loop_electrodynamic( reference_particles, @@ -856,11 +847,6 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): reference_fields, reference_world, constants, - unused_curl, - J_func=None, - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -870,11 +856,6 @@ def test_tiled_yee_step_matches_standard_periodic_yee_step(self): tiled_fields, tiled_world, constants, - unused_curl, - J_func=None, - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -940,6 +921,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, + jnp.asarray(False), ) reference_particles, reference_fields = time_loop_electrodynamic( reference_particles, @@ -947,11 +929,6 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): reference_fields, reference_world, constants, - unused_curl, - J_func=None, - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -961,11 +938,6 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): tiled_fields, tiled_world, constants, - unused_curl, - J_func=None, - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -983,6 +955,7 @@ def test_tiled_yee_step_matches_standard_conducting_yee_step(self): def test_tiled_yee_step_uses_digital_filtered_current(self): world = self._build_world() + world["current_filter"] = "digital" constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) @@ -1016,6 +989,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, + jnp.asarray(False), ) reference_particles, reference_fields = time_loop_electrodynamic( reference_particles, @@ -1023,11 +997,6 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): reference_fields, reference_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="digital"), - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1037,11 +1006,6 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): tiled_fields, tiled_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="digital"), - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1059,6 +1023,7 @@ def test_tiled_yee_step_uses_digital_filtered_current(self): def test_tiled_yee_step_uses_bilinear_filtered_current(self): world = self._build_world() + world["current_filter"] = "bilinear" constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) @@ -1092,6 +1057,7 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): tile_vector_field(external_fields[1], tiled_world, tile_shape), ), None, + jnp.asarray(False), ) reference_particles, reference_fields = time_loop_electrodynamic( reference_particles, @@ -1099,11 +1065,6 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): reference_fields, reference_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="bilinear"), - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1113,11 +1074,6 @@ def test_tiled_yee_step_uses_bilinear_filtered_current(self): tiled_fields, tiled_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="bilinear"), - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1137,6 +1093,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand world = self._build_world() world["shape_factor"] = 2 world["guard_cells"] = 2 + world["current_filter"] = "bilinear" constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) @@ -1211,6 +1168,7 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), ), None, + jnp.asarray(False), ) reference_particles, reference_fields = time_loop_electrodynamic( reference_particles, @@ -1218,11 +1176,6 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand reference_fields, reference_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="bilinear"), - solver="electrodynamic_yee", - tile_shape=reference_tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1232,11 +1185,6 @@ def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_stand tiled_fields, tiled_world, constants, - unused_curl, - J_func=functools.partial(J_from_rhov, filter="bilinear"), - solver="electrodynamic_yee", - tile_shape=tile_shape, - g=int(world["guard_cells"]), relativistic=False, particle_pusher="boris", ) @@ -1270,10 +1218,10 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): constants, tile_shape, Nt, - ) = self._long_two_stream_state(False) + ) = self._long_two_stream_state("direct", "none") - reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world, reference_tile_shape, None) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world, tile_shape, None) + reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) for t in range(Nt): reference_particles, reference_fields = reference_loop( @@ -1281,8 +1229,6 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): reference_species_config, reference_fields, constants, - unused_curl, - solver="electrodynamic_yee", relativistic=False, particle_pusher="boris", ) @@ -1291,8 +1237,6 @@ def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): species_config, tiled_fields, constants, - unused_curl, - solver="electrodynamic_yee", relativistic=False, particle_pusher="boris", ) @@ -1339,14 +1283,10 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): constants, tile_shape, Nt, - ) = self._long_two_stream_state(True) + ) = self._long_two_stream_state("esirkepov", "none") - reference_loop = self._jitted_electrodynamic_loop_with_static_world( - reference_world, - reference_tile_shape, - Esirkepov_current, - ) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world, tile_shape, Esirkepov_current) + reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) for t in range(Nt): reference_particles, reference_fields = reference_loop( @@ -1354,8 +1294,6 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): reference_species_config, reference_fields, constants, - unused_curl, - solver="electrodynamic_yee", relativistic=False, particle_pusher="boris", ) @@ -1364,8 +1302,6 @@ def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): species_config, tiled_fields, constants, - unused_curl, - solver="electrodynamic_yee", relativistic=False, particle_pusher="boris", ) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 35d1f203..782854a9 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -4,7 +4,6 @@ import jax.numpy as jnp from PyPIC3D.initialization import initialize_fields -from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global, update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, @@ -54,6 +53,11 @@ def _conducting_world(self): 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) + return world + def _mixed_bc_world(self): world = self._build_world() world["boundary_conditions"] = {"x": BC_PERIODIC, "y": BC_CONDUCTING, "z": BC_CONDUCTING} @@ -302,21 +306,20 @@ def test_update_tiled_E_matches_standard_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, _ = _update_E_global(E, B, J, world, constants, unused_curl) - E_tiled = update_tiled_E( + E_tiled, pml_state = update_tiled_E( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), tile_vector_field(J, world, tile_shape), world, constants, - unused_curl, - tile_shape, - 2, ) + 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): @@ -326,8 +329,7 @@ def test_first_order_yee_public_update_E_routes_tiled_arrays(self): world = self._build_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 3, 2) - world["tile_shape"] = tile_shape - world["guard_cells"] = 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) @@ -335,10 +337,11 @@ def test_first_order_yee_public_update_E_routes_tiled_arrays(self): 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, unused_curl) - E_reference = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) + E_public, pml_state = update_E(E_tiles, B_tiles, J_tiles, world, constants) + E_reference, reference_pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) self.assertIsNone(pml_state) + self.assertIsNone(reference_pml_state) for public, reference in zip(E_public, E_reference): self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) @@ -350,29 +353,28 @@ def test_first_order_yee_public_update_rejects_global_arrays(self): J = self._deterministic_vector_field(world, scale=0.05) with self.assertRaisesRegex(ValueError, "tiled field arrays"): - update_E(E, B, J, world, constants, unused_curl) + update_E(E, B, J, world, constants) with self.assertRaisesRegex(ValueError, "tiled field arrays"): - update_B(E, B, world, constants, unused_curl) + update_B(E, B, world, constants) def test_update_tiled_E_matches_standard_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, _ = _update_E_global(E, B, J, world, constants, unused_curl) - E_tiled = update_tiled_E( + E_tiled, pml_state = update_tiled_E( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), tile_vector_field(J, world, tile_shape), world, constants, - unused_curl, - tile_shape, - 2, ) + 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): @@ -382,19 +384,18 @@ def test_update_tiled_B_matches_standard_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, _ = _update_B_global(E, B, world, constants, unused_curl) - B_tiled = update_tiled_B( + B_tiled, pml_state = update_tiled_B( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, constants, - unused_curl, - tile_shape, - 2, ) + 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): @@ -404,17 +405,17 @@ def test_first_order_yee_public_update_B_routes_tiled_arrays(self): world = self._build_world() constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} tile_shape = (2, 3, 2) - world["tile_shape"] = tile_shape - world["guard_cells"] = 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, unused_curl) - B_reference = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) + B_public, pml_state = update_B(E_tiles, B_tiles, world, constants) + B_reference, reference_pml_state = update_tiled_B(E_tiles, B_tiles, world, constants) self.assertIsNone(pml_state) + self.assertIsNone(reference_pml_state) for public, reference in zip(B_public, B_reference): self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) @@ -422,19 +423,18 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s 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, _ = _update_B_global(E, B, world, constants, unused_curl) - B_tiled = update_tiled_B( + B_tiled, pml_state = update_tiled_B( tile_vector_field(E, world, tile_shape), tile_vector_field(B, world, tile_shape), world, constants, - unused_curl, - tile_shape, - 2, ) + 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): @@ -444,6 +444,7 @@ def test_tiled_electrodynamic_step_matches_standard_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) @@ -454,8 +455,9 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) + E_tiles, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_tiled_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) @@ -469,6 +471,7 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting 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) @@ -479,8 +482,9 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting 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 = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, unused_curl, tile_shape, 2) - B_tiles = update_tiled_B(E_tiles, B_tiles, world, constants, unused_curl, tile_shape, 2) + E_tiles, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_tiled_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) @@ -494,54 +498,28 @@ 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, _ = _update_E_global(E, B, J, world, constants, unused_curl) B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) - particles = object() - external_fields = ( - tuple(jnp.zeros_like(component) for component in E), - tuple(jnp.zeros_like(component) for component in B), - ) - fields = ( - tile_vector_field(E, world, tile_shape), - tile_vector_field(B, world, tile_shape), - tile_vector_field(J, world, tile_shape), - jnp.zeros_like(E[0]), - jnp.zeros_like(E[0]), - external_fields, - None, - ) - - particles_after, fields_after = time_loop_electrodynamic( - particles, - None, - fields, - world, - constants, - unused_curl, - J_func=None, - solver="fdtd", - tile_shape=tile_shape, - g=1, - relativistic=False, - particle_pusher="boris", - ) + 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_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) + B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, pml_state) - E_tiles, B_tiles, J_tiles, rho, phi, external_after, pml_state = fields_after 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.assertIs(particles_after, particles) self.assertIsNone(pml_state) - self.assertIs(external_after, external_fields) 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(fields[2], J_tiles): + 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)) From 333ab7cc023fe20cce58a3d44fdd4097c1c4370a Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 06:29:05 -0500 Subject: [PATCH 39/83] starting to refactor yee evolution --- PyPIC3D/__init__.py | 1 - PyPIC3D/diagnostics/tiled_benchmarking.py | 8 +- PyPIC3D/evolve.py | 10 +- PyPIC3D/solvers/first_order_yee.py | 241 ------------------- PyPIC3D/solvers/yee_tiled.py | 4 +- tests/__init__.py | 1 - tests/first_order_yee_test.py | 271 ---------------------- tests/pml_test.py | 93 ++++---- tests/tiled_esirkepov_test.py | 6 +- tests/tiled_refactor_contract_test.py | 6 +- tests/tiled_yee_integration_test.py | 12 +- tests/yee_tiled_test.py | 132 ++++++----- 12 files changed, 147 insertions(+), 638 deletions(-) delete mode 100644 PyPIC3D/solvers/first_order_yee.py delete mode 100644 tests/first_order_yee_test.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index eb950691..a041a169 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -14,7 +14,6 @@ from . import evolve from .solvers import fdtd -from .solvers import first_order_yee from .solvers import electrostatic_yee from .particles import particle_initialization diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py index 93bb4b40..8089bf4e 100644 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ b/PyPIC3D/diagnostics/tiled_benchmarking.py @@ -20,8 +20,8 @@ empty_tiled_vector_field, tile_grid_axes, tile_vector_field, - update_tiled_B, - update_tiled_E, + update_B, + update_E, ) from PyPIC3D.utils import add_external_fields, build_yee_grid, compute_energy @@ -252,8 +252,8 @@ def _tiled_current_deposition_stage(case, particles): def _tiled_field_update_stage(case, J_tiles): E_tiles, B_tiles, *_ = case.fields with jax.named_scope("tiled_field_update"): - E_tiles, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, case.world, case.constants) - B_tiles, _pml_state = update_tiled_B(E_tiles, B_tiles, case.world, case.constants, pml_state) + E_tiles, pml_state = update_E(E_tiles, B_tiles, J_tiles, case.world, case.constants) + B_tiles, _pml_state = update_B(E_tiles, B_tiles, case.world, case.constants, pml_state) return E_tiles, B_tiles diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 79eaf49a..97b4c0e9 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -15,7 +15,7 @@ ) from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields -from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E +from PyPIC3D.solvers.yee_tiled import update_B, update_E from PyPIC3D.utils import add_external_fields @@ -112,13 +112,13 @@ def esirkepov_deposition_step(state): ) # deposit current into the tiled J arrays using the selected deposition method - E, pml_state = update_tiled_E(E, B, J, world, constants, pml_state) + 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_tiled_E function returns None for the pml_state + # for no pml, the pml_state is None, and the update_E function returns None for the pml_state - B, pml_state = update_tiled_B(E, B, world, constants, 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_tiled_B function returns None for the pml_state + # 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 diff --git a/PyPIC3D/solvers/first_order_yee.py b/PyPIC3D/solvers/first_order_yee.py deleted file mode 100644 index 97866f17..00000000 --- a/PyPIC3D/solvers/first_order_yee.py +++ /dev/null @@ -1,241 +0,0 @@ -import jax.numpy as jnp -from jax import jit -from functools import partial -# import external libraries - -from PyPIC3D.utilities.filters 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, -) -from PyPIC3D.solvers.yee_tiled import update_tiled_B, update_tiled_E -# import internal libraries - - -@partial(jit, static_argnames=("curl_func",)) -def _update_E_global(E, B, J, world, constants, curl_func, 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. - """ - - 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 - - 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( - (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 - ) - 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 - ) - 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 - ) - # 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) - 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 - - 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 = apply_conducting_bc((Ex, Ey, Ez), bc_x, bc_y, bc_z) - # enforce conducting boundary conditions on tangential E - - 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 (Ex, Ey, Ez), pml_state - - -@partial(jit, static_argnames=("curl_func",)) -def _update_B_global(E, B, world, constants, curl_func, 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. - """ - - 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 - - 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( - (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 - - 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) - # refresh ghost cells before any stencil-based post-processing - - 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 - - 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 (Bx, By, Bz), pml_state - - -def update_E(E, B, J, world, constants, pml_state=None): - """ - Update electric fields through the public Yee helper. - - Tiled field/current arrays are routed to the tile-local Yee update using - the explicit ``world['tile_shape']`` and ``world['guard_cells']`` contract. - Flat global ghost-celled arrays are private reference inputs; tests that - need them should call ``_update_E_global`` directly. - """ - - if E[0].ndim == 6: - return update_tiled_E(E, B, J, world, constants, pml_state) - - raise ValueError("Public update_E requires tiled field arrays; use _update_E_global for flat reference tests.") - - -def update_B(E, B, world, constants, pml_state=None): - """ - Update magnetic fields through the public Yee helper. - - Flat global ghost-celled arrays are private reference inputs; tests that - need them should call ``_update_B_global`` directly. - """ - - if E[0].ndim == 6: - return update_tiled_B(E, B, world, constants, pml_state) - - raise ValueError("Public update_B requires tiled field arrays; use _update_B_global for flat reference tests.") diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 043805d0..0d2294e7 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -664,7 +664,7 @@ def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): return _restore_tiled_vector_layout(folded, field_tiles) -def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): +def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -742,7 +742,7 @@ def update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state -def update_tiled_B(E_tiles, B_tiles, world, constants, pml_state=None): +def update_B(E_tiles, B_tiles, world, constants, pml_state=None): """ Update compact tiled magnetic fields without assembling a global field. diff --git a/tests/__init__.py b/tests/__init__.py index fa27bf25..881a2c49 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -9,7 +9,6 @@ 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 utils_test diff --git a/tests/first_order_yee_test.py b/tests/first_order_yee_test.py deleted file mode 100644 index a7449588..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_global, _update_B_global -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_global(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_global(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/pml_test.py b/tests/pml_test.py index 98c73279..77b61a04 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -14,13 +14,12 @@ update_ghost_cells_for_pml, ) from PyPIC3D.initialization import initialize_fields, initialize_simulation -from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, empty_tiled_vector_field, tile_vector_field, - update_tiled_B, - update_tiled_E, + update_B, + update_E, ) from PyPIC3D.utils import build_yee_grid, compute_energy @@ -269,49 +268,22 @@ 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, _, _ = 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_global(E, B, J, world, constants, lambda *args: None) - B_after, pml_state = _update_B_global(E_after, B, world, constants, lambda *args: None, pml_state) + 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_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_global(E, B, J, world, constants, lambda *args: None, pml_state) - B, pml_state = _update_B_global(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)) - - def test_tiled_pml_matches_global_pml_for_one_yee_step(self): + 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( @@ -345,13 +317,40 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): 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) - pml_state = initialize_pml_state(world) - E_reference, pml_state = _update_E_global(E, B, J, world, constants, lambda *args: None, pml_state) - B_reference, pml_state = _update_B_global(E_reference, B, world, constants, lambda *args: None, pml_state) + 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_tiled_E( + 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), @@ -359,7 +358,7 @@ def test_tiled_pml_matches_global_pml_for_one_yee_step(self): constants, tiled_pml_state, ) - B_tiles, tiled_pml_state = update_tiled_B( + B_tiles, tiled_pml_state = update_B( E_tiles, tile_vector_field(B, world, tile_shape), world, @@ -404,8 +403,8 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) def step(E_tiles, B_tiles, tiled_pml_state): - E_tiles, tiled_pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, tiled_pml_state) - B_tiles, tiled_pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, 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) @@ -441,8 +440,8 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): J_tiles = (Jx, Jy, Jz) pml_state = initialize_tiled_pml_state(world, tile_shape) - E_after, pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state) - B_after, pml_state = update_tiled_B(E_after, B_tiles, world, constants, pml_state) + 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))) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 6e3140d8..25944d1c 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -32,7 +32,7 @@ empty_tiled_vector_field, tile_grid_axes, tile_vector_field, - update_tiled_E, + update_E, ) jax.config.update("jax_enable_x64", True) @@ -287,7 +287,7 @@ def test_output_adapter_uses_world_guard_cells(self): 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_tiled_E_reads_two_guard_current_interior(self): + 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) @@ -302,7 +302,7 @@ def test_update_tiled_E_reads_two_guard_current_interior(self): Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) world["tile_shape"] = tile_shape - E_after, pml_state = update_tiled_E(E_tiles, B_tiles, (Jx, Jy, Jz), world, constants) + 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)) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 2a9b2da0..59d41a47 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -100,11 +100,13 @@ def test_evolve_loops_read_runtime_controls_from_world(self): self.assertNotIn(name, electrostatic_signature.parameters) def test_tiled_yee_updates_read_tile_metadata_from_world(self): - update_E_signature = inspect.signature(yee_tiled.update_tiled_E) - update_B_signature = inspect.signature(yee_tiled.update_tiled_B) + update_E_signature = inspect.signature(yee_tiled.update_E) + update_B_signature = inspect.signature(yee_tiled.update_B) for name in ("curl_func", "tile_shape", "g"): self.assertNotIn(name, update_E_signature.parameters) self.assertNotIn(name, update_B_signature.parameters) + self.assertFalse(hasattr(yee_tiled, "update_tiled_E")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): source = inspect.getsource(evolve.time_loop_electrodynamic) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 1448453d..50a0700b 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -17,8 +17,8 @@ empty_tiled_vector_field, tile_grid_axes, tile_vector_field, - update_tiled_B, - update_tiled_E, + update_B, + update_E, ) from PyPIC3D.utils import build_yee_grid, compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC @@ -463,14 +463,14 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): None, jnp.asarray(False), ) - reference_E, reference_pml_state = update_tiled_E( + reference_E, reference_pml_state = update_E( reference_fields[0], reference_fields[1], reference_fields[2], reference_world, constants, ) - reference_B, reference_pml_state = update_tiled_B( + reference_B, reference_pml_state = update_B( reference_E, reference_fields[1], reference_world, @@ -492,8 +492,8 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): None, jnp.asarray(False), ) - tiled_E, tiled_pml_state = update_tiled_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) - tiled_B, tiled_pml_state = update_tiled_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) + tiled_E, tiled_pml_state = update_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) + tiled_B, tiled_pml_state = update_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) tiled_fields = (tiled_E, tiled_B, *tiled_fields[2:]) reference_E, reference_B, *_ = reference_fields diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 782854a9..e9379869 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -4,7 +4,6 @@ import jax.numpy as jnp from PyPIC3D.initialization import initialize_fields -from PyPIC3D.solvers.first_order_yee import _update_B_global, _update_E_global, update_B, update_E from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, tile_grid_axes, @@ -13,8 +12,8 @@ update_tiled_ghost_cells_periodic, update_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells_periodic, - update_tiled_B, - update_tiled_E, + update_B, + update_E, ) from PyPIC3D.utils import build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -24,10 +23,6 @@ jax.config.update("jax_enable_x64", True) -def unused_curl(Ex, Ey, Ez): - return None - - class TestYeeTiled(unittest.TestCase): def _build_world(self): world = { @@ -85,6 +80,52 @@ def _deterministic_vector_field(self, world, scale): 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"] @@ -302,7 +343,7 @@ def test_update_tiled_vector_ghost_cells_conducting_refreshes_each_component(sel 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_tiled_E_matches_standard_yee_update(self): + 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) @@ -311,8 +352,8 @@ def test_update_tiled_E_matches_standard_yee_update(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) - E_tiled, pml_state = update_tiled_E( + 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), @@ -325,7 +366,7 @@ def test_update_tiled_E_matches_standard_yee_update(self): 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_first_order_yee_public_update_E_routes_tiled_arrays(self): + 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) @@ -338,26 +379,12 @@ def test_first_order_yee_public_update_E_routes_tiled_arrays(self): 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) - E_reference, reference_pml_state = update_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) self.assertIsNone(pml_state) - self.assertIsNone(reference_pml_state) - for public, reference in zip(E_public, E_reference): - self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) - - def test_first_order_yee_public_update_rejects_global_arrays(self): - world = self._build_world() - constants = {"C": 1.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - 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) - - with self.assertRaisesRegex(ValueError, "tiled field arrays"): - update_E(E, B, J, world, constants) - with self.assertRaisesRegex(ValueError, "tiled field arrays"): - update_B(E, B, world, constants) + self.assertEqual(E_public[0].ndim, 6) + self.assertEqual(E_public[0].shape[:3], E_tiles[0].shape[:3]) - def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(self): + 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) @@ -366,8 +393,8 @@ def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(s B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) - E_tiled, pml_state = update_tiled_E( + 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), @@ -380,7 +407,7 @@ def test_update_tiled_E_matches_standard_yee_update_with_conducting_boundaries(s 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_tiled_B_matches_standard_yee_update(self): + 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) @@ -388,8 +415,8 @@ def test_update_tiled_B_matches_standard_yee_update(self): E = self._deterministic_vector_field(world, scale=1.0) B = self._deterministic_vector_field(world, scale=0.2) - B_reference, _ = _update_B_global(E, B, world, constants, unused_curl) - B_tiled, pml_state = update_tiled_B( + 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, @@ -401,7 +428,7 @@ def test_update_tiled_B_matches_standard_yee_update(self): 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_first_order_yee_public_update_B_routes_tiled_arrays(self): + 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) @@ -412,14 +439,12 @@ def test_first_order_yee_public_update_B_routes_tiled_arrays(self): 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) - B_reference, reference_pml_state = update_tiled_B(E_tiles, B_tiles, world, constants) self.assertIsNone(pml_state) - self.assertIsNone(reference_pml_state) - for public, reference in zip(B_public, B_reference): - self.assertTrue(jnp.allclose(public, reference, rtol=1.0e-12, atol=1.0e-12)) + self.assertEqual(B_public[0].ndim, 6) + self.assertEqual(B_public[0].shape[:3], B_tiles[0].shape[:3]) - def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(self): + 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) @@ -427,8 +452,8 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s E = self._deterministic_vector_field(world, scale=1.0) B = self._deterministic_vector_field(world, scale=0.2) - B_reference, _ = _update_B_global(E, B, world, constants, unused_curl) - B_tiled, pml_state = update_tiled_B( + 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, @@ -440,7 +465,7 @@ def test_update_tiled_B_matches_standard_yee_update_with_conducting_boundaries(s 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_standard_yee_sequence(self): + 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) @@ -449,14 +474,13 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) - B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) + 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_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, pml_state) + 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) @@ -467,7 +491,7 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence(self): 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_standard_yee_sequence_with_conducting_boundaries(self): + 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) @@ -476,14 +500,13 @@ def test_tiled_electrodynamic_step_matches_standard_yee_sequence_with_conducting B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) - B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) + 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_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, pml_state) + 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) @@ -503,13 +526,12 @@ def test_tiled_evolve_updates_fields_without_pushing_particles(self): B = self._deterministic_vector_field(world, scale=0.2) J = self._deterministic_vector_field(world, scale=0.05) - E_reference, _ = _update_E_global(E, B, J, world, constants, unused_curl) - B_reference, _ = _update_B_global(E_reference, B, world, constants, unused_curl) + 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_tiled_E(E_tiles, B_tiles, J_tiles, world, constants) - B_tiles, pml_state = update_tiled_B(E_tiles, B_tiles, world, constants, pml_state) + 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) From c248e3a6e225bc8fb4b8362e03fcf145ac872194 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 12:12:34 -0500 Subject: [PATCH 40/83] remove computation of tiled grid redundant --- PyPIC3D/pusher/tiled_pusher.py | 19 ++----------------- PyPIC3D/solvers/yee_tiled.py | 16 ---------------- tests/tiled_pusher_test.py | 24 +++++++++++++++++++++++- tests/tiled_refactor_contract_test.py | 10 ++++++++++ 4 files changed, 35 insertions(+), 34 deletions(-) diff --git a/PyPIC3D/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py index b531efd6..f7e09c22 100644 --- a/PyPIC3D/pusher/tiled_pusher.py +++ b/PyPIC3D/pusher/tiled_pusher.py @@ -10,7 +10,6 @@ relativistic_boris_single_particle, ) from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle -from PyPIC3D.solvers.yee_tiled import tiled_grid_axes_from_world @partial(jit, static_argnames=("tile_shape", "g", "relativistic", "particle_pusher")) @@ -27,22 +26,8 @@ def tiled_particle_push(tiled_particles, species_config, E_tiles, B_tiles, world dt = world["dt"] shape_factor = world["shape_factor"] - center_grid = world["grids"]["center"] - vertex_grid = world["grids"]["vertex"] - tiled_center_grid = tiled_grid_axes_from_world( - world, - center_grid, - "tiled_center_grid", - tile_shape, - g, - ) - tiled_vertex_grid = tiled_grid_axes_from_world( - world, - vertex_grid, - "tiled_vertex_grid", - tile_shape, - g, - ) + 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 diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 0d2294e7..d5c1b7c8 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -157,22 +157,6 @@ def tile_grid_axes(grid, world, tile_shape, num_guard_cells=2): ) -def tiled_grid_axes_from_world(world, grid, tiled_grid_key, tile_shape, num_guard_cells=2): - """ - Read precomputed tile-local grid axes, with a construction fallback. - - Initialization stores the tiled grids in ``world['grids']`` for production - tiled runs. The fallback keeps focused tests and older call sites that - construct small worlds by hand on the same numerical convention. - """ - - grids = world["grids"] - if tiled_grid_key in grids: - return grids[tiled_grid_key] - - return tile_grid_axes(grid, world, tile_shape, num_guard_cells) - - def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): """ Split a ghost-celled field into compact tiles using the shared tile shape. diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index ed7fd515..89f7123b 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -7,7 +7,7 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.solvers.yee_tiled import tile_grid_axes, tile_vector_field from PyPIC3D.utils import build_yee_grid @@ -34,6 +34,23 @@ def _build_world(self, Nx=8, Ny=6, Nz=4, shape_factor=1): 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) + world["grids"]["tiled_center_grid"] = tile_grid_axes( + world["grids"]["center"], + world, + tile_shape, + num_guard_cells=g, + ) + world["grids"]["tiled_vertex_grid"] = tile_grid_axes( + world["grids"]["vertex"], + world, + tile_shape, + num_guard_cells=g, + ) + return world + def _deterministic_vector_field(self, world, scale): Nx, Ny, Nz = world["Nx"], world["Ny"], world["Nz"] ii, jj, kk = jnp.meshgrid( @@ -90,6 +107,7 @@ def test_tiled_particle_push_matches_flat_boris(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], @@ -136,6 +154,7 @@ def test_tiled_particle_push_matches_flat_higuera_cary(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], @@ -204,6 +223,7 @@ def test_tiled_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], @@ -240,6 +260,7 @@ def test_tiled_particle_push_matches_relativistic_flat_boris_on_reduced_axes(sel 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) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], @@ -309,6 +330,7 @@ def test_tiled_particle_push_matches_flat_boris_on_two_guard_reduced_axes(self): constants = {"C": 10.0} tile_shape = (2, 1, 1) g = 2 + world = self._with_tiled_grids(world, tile_shape, g=g) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 59d41a47..66408d66 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -107,6 +107,16 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertNotIn(name, update_B_signature.parameters) self.assertFalse(hasattr(yee_tiled, "update_tiled_E")) self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) + self.assertFalse(hasattr(yee_tiled, "tiled_grid_axes_from_world")) + + def test_tiled_pusher_reads_tiled_grids_from_world(self): + from PyPIC3D.pusher import tiled_pusher + + source = inspect.getsource(tiled_pusher.tiled_particle_push) + self.assertIn('world["grids"]["tiled_center_grid"]', source) + self.assertIn('world["grids"]["tiled_vertex_grid"]', source) + self.assertNotIn("tiled_grid_axes_from_world", source) + self.assertNotIn("tile_grid_axes", source) def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): source = inspect.getsource(evolve.time_loop_electrodynamic) From 4a966399cbb3da8623f0012e75086af462505003 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 14:23:27 -0500 Subject: [PATCH 41/83] Refactoring grid computation --- PyPIC3D/diagnostics/tiled_benchmarking.py | 357 ---------------------- PyPIC3D/initialization.py | 25 +- PyPIC3D/solvers/yee_tiled.py | 49 --- PyPIC3D/utilities/grids.py | 137 +++++++++ PyPIC3D/utils.py | 86 ------ tests/direct_deposition_tiled_test.py | 3 +- tests/electrostatic_yee_test.py | 2 +- tests/evolve_test.py | 4 +- tests/particle_test.py | 4 +- tests/pml_test.py | 3 +- tests/tiled_electrostatic_test.py | 3 +- tests/tiled_esirkepov_test.py | 2 +- tests/tiled_fluid_quantities_test.py | 4 +- tests/tiled_pusher_test.py | 4 +- tests/tiled_refactor_contract_test.py | 44 +-- tests/tiled_rho_test.py | 4 +- tests/tiled_yee_integration_test.py | 4 +- tests/utils_test.py | 3 +- tests/yee_tiled_test.py | 3 +- 19 files changed, 184 insertions(+), 557 deletions(-) delete mode 100644 PyPIC3D/diagnostics/tiled_benchmarking.py create mode 100644 PyPIC3D/utilities/grids.py diff --git a/PyPIC3D/diagnostics/tiled_benchmarking.py b/PyPIC3D/diagnostics/tiled_benchmarking.py deleted file mode 100644 index 8089bf4e..00000000 --- a/PyPIC3D/diagnostics/tiled_benchmarking.py +++ /dev/null @@ -1,357 +0,0 @@ -from dataclasses import dataclass -import jax -import jax.numpy as jnp - -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC -from PyPIC3D.deposition.Esirkepov import Esirkepov_current -from PyPIC3D.deposition.J_from_rhov import J_from_rhov -from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output -from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.particles.tiled_particle_refresh import ( - refresh_tiled_particle_tiles, - update_tiled_particle_positions, -) -from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import ( - empty_tiled_scalar_field, - empty_tiled_vector_field, - tile_grid_axes, - tile_vector_field, - update_B, - update_E, -) -from PyPIC3D.utils import add_external_fields, build_yee_grid, compute_energy - - -@dataclass -class StageSpec: - name: str - function: object - args: tuple - kwargs: dict - static_argnames: tuple = () - - -@dataclass -class SyntheticTiledYeeState: - particles: TiledParticles - species_config: object - fields: tuple - world: dict - constants: dict - tile_shape: tuple - g: int - - -def unused_curl(Ex, Ey, Ez): - return None - - -def _particle_kernel_world(world): - kernel_world = dict(world) - kernel_world.pop("current_deposition", None) - kernel_world.pop("current_filter", None) - return kernel_world - - -def _pad_tiled_particle_capacity(tiled_particles, slots_per_tile): - current_slots = tiled_particles.active.shape[-1] - if current_slots >= slots_per_tile: - return tiled_particles - - pad_slots = int(slots_per_tile) - int(current_slots) - vector_pad = [(0, 0)] * tiled_particles.x.ndim - vector_pad[-2] = (0, pad_slots) - scalar_pad = [(0, 0)] * tiled_particles.active.ndim - scalar_pad[-1] = (0, pad_slots) - - return TiledParticles( - x=jnp.pad(tiled_particles.x, vector_pad), - u=jnp.pad(tiled_particles.u, vector_pad), - active=jnp.pad(tiled_particles.active, scalar_pad), - ) - - -def build_synthetic_tiled_yee_state( - nx=16, - ny=16, - nz=16, - particles_per_species=2048, - species_count=2, - shape=1, - seed=0, - tile_shape=(8, 8, 8), - slots_per_tile=None, - g=2, - dt=None, -): - """ - Build a deterministic tiled electrodynamic Yee benchmark state. - """ - - x_wind = 1.0 - y_wind = 1.0 - z_wind = 1.0 - dx = x_wind / nx - dy = y_wind / ny - dz = z_wind / nz - if dt is None: - dt = 0.05 * min(dx, dy, dz) - - world = { - "dt": dt, - "Nt": 1, - "dx": dx, - "dy": dy, - "dz": dz, - "Nx": nx, - "Ny": ny, - "Nz": nz, - "x_wind": x_wind, - "y_wind": y_wind, - "z_wind": z_wind, - "shape_factor": shape, - "current_deposition": "direct", - "current_filter": "none", - "tile_shape": tile_shape, - "guard_cells": g, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - world["grids"]["tiled_center_grid"] = tile_grid_axes(center_grid, world, tile_shape, num_guard_cells=g) - world["grids"]["tiled_vertex_grid"] = tile_grid_axes(vertex_grid, world, tile_shape, num_guard_cells=g) - - constants = { - "eps": 1.0, - "mu": 1.0, - "C": 1.0, - "kb": 1.0, - "alpha": 1.0, - } - - key = jax.random.PRNGKey(seed) - particles = [] - for species_index in range(species_count): - key, x_key, y_key, z_key, vx_key, vy_key, vz_key = jax.random.split(key, 7) - x1 = jax.random.uniform( - x_key, (particles_per_species,), minval=-0.45 * x_wind, maxval=0.45 * x_wind - ) - x2 = jax.random.uniform( - y_key, (particles_per_species,), minval=-0.45 * y_wind, maxval=0.45 * y_wind - ) - x3 = jax.random.uniform( - z_key, (particles_per_species,), minval=-0.45 * z_wind, maxval=0.45 * z_wind - ) - v1 = 0.01 * jax.random.normal(vx_key, (particles_per_species,)) - v2 = 0.01 * jax.random.normal(vy_key, (particles_per_species,)) - v3 = 0.01 * jax.random.normal(vz_key, (particles_per_species,)) - - charge_sign = 1.0 if species_index % 2 == 0 else -1.0 - particles.append( - particle_species( - name=f"benchmark_species_{species_index}", - N_particles=particles_per_species, - charge=charge_sign, - mass=1.0, - T=0.0, - v1=v1, - v2=v2, - v3=v3, - x1=x1, - x2=x2, - x3=x3, - xwind=x_wind, - ywind=y_wind, - zwind=z_wind, - dx=dx, - dy=dy, - dz=dz, - dt=dt, - ) - ) - - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - "particle_tile_capacity_factor": 1.0, - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - if slots_per_tile is not None: - tiled_particles = _pad_tiled_particle_capacity(tiled_particles, slots_per_tile) - - global_shape = (nx + 2, ny + 2, nz + 2) - E = tuple(jnp.zeros(global_shape) for _ in range(3)) - B = tuple(jnp.zeros(global_shape) for _ in range(3)) - E_tiles = tile_vector_field(E, world, tile_shape, num_guard_cells=g) - B_tiles = tile_vector_field(B, world, tile_shape, num_guard_cells=g) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) - rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) - phi_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) - external_fields = ( - tuple(jnp.zeros_like(component) for component in E_tiles), - tuple(jnp.zeros_like(component) for component in B_tiles), - ) - fields = (E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, external_fields, None, jnp.asarray(False)) - - return SyntheticTiledYeeState( - particles=tiled_particles, - species_config=species_config, - fields=fields, - world=world, - constants=constants, - tile_shape=tile_shape, - g=g, - ) - - -def _tiled_particle_push_stage(case, relativistic=True, particle_pusher="boris"): - E_tiles, B_tiles, *_ = case.fields - push_E_tiles, push_B_tiles = add_external_fields(E_tiles, B_tiles, case.fields[5]) - with jax.named_scope("tiled_particle_push"): - return tiled_particle_push( - case.particles, - case.species_config, - push_E_tiles, - push_B_tiles, - _particle_kernel_world(case.world), - case.constants, - case.tile_shape, - case.g, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - - -def _tiled_particle_retile_stage(case): - with jax.named_scope("tiled_particle_retile"): - moved = update_tiled_particle_positions(case.particles, case.species_config, case.world["dt"]) - return refresh_tiled_particle_tiles(moved, case.world, case.tile_shape) - - -def _tiled_current_deposition_stage(case, particles): - _, _, J_tiles, *_ = case.fields - with jax.named_scope("tiled_current_deposition"): - if case.world["current_deposition"] == "esirkepov": - return Esirkepov_current(particles, case.species_config, J_tiles, case.constants, case.world) - return J_from_rhov( - particles, - case.species_config, - J_tiles, - case.constants, - case.world, - filter=case.world["current_filter"], - ) - - -def _tiled_field_update_stage(case, J_tiles): - E_tiles, B_tiles, *_ = case.fields - with jax.named_scope("tiled_field_update"): - E_tiles, pml_state = update_E(E_tiles, B_tiles, J_tiles, case.world, case.constants) - B_tiles, _pml_state = update_B(E_tiles, B_tiles, case.world, case.constants, pml_state) - return E_tiles, B_tiles - - -def _tiled_diagnostics_stage(case): - E_tiles, B_tiles, *_ = case.fields - with jax.named_scope("tiled_diagnostics"): - return compute_energy( - case.particles, - E_tiles, - B_tiles, - case.world, - case.constants, - species_config=case.species_config, - ) - - -def _tiled_output_bridge_stage(case): - E_tiles, B_tiles, *_ = case.fields - with jax.named_scope("tiled_output_bridge"): - E = vector_field_for_output(E_tiles, case.world) - B = vector_field_for_output(B_tiles, case.world) - particles = particles_for_output(case.particles, case.species_config) - return E, B, particles - - -def benchmark_tiled_pic_step(case, relativistic=True, particle_pusher="boris"): - with jax.named_scope("tiled_pic_step"): - return time_loop_electrodynamic( - case.particles, - case.species_config, - case.fields, - case.world, - case.constants, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - - -def build_tiled_stage_specs(case, relativistic=True, particle_pusher="boris"): - """ - Return independently timed tiled Yee stage specs. - - These specs identify the expensive tile-management stages before larger - kernel rewrites: push, retile, deposition, field update, diagnostics, and - the current global output bridge. - """ - - pushed_particles = _tiled_particle_push_stage(case, relativistic=relativistic, particle_pusher=particle_pusher) - deposited_J = _tiled_current_deposition_stage(case, pushed_particles) - - return { - "tiled_pic_step": StageSpec( - "tiled_pic_step", - benchmark_tiled_pic_step, - (case,), - {"relativistic": relativistic, "particle_pusher": particle_pusher}, - ("relativistic", "particle_pusher"), - ), - "tiled_particle_push": StageSpec( - "tiled_particle_push", - _tiled_particle_push_stage, - (case,), - {"relativistic": relativistic, "particle_pusher": particle_pusher}, - ("relativistic", "particle_pusher"), - ), - "tiled_particle_retile": StageSpec( - "tiled_particle_retile", - _tiled_particle_retile_stage, - (case,), - {}, - (), - ), - "tiled_current_deposition": StageSpec( - "tiled_current_deposition", - _tiled_current_deposition_stage, - (case, pushed_particles), - {}, - (), - ), - "tiled_field_update": StageSpec( - "tiled_field_update", - _tiled_field_update_stage, - (case, deposited_J), - {}, - (), - ), - "tiled_diagnostics": StageSpec( - "tiled_diagnostics", - _tiled_diagnostics_stage, - (case,), - {}, - (), - ), - "tiled_output_bridge": StageSpec( - "tiled_output_bridge", - _tiled_output_bridge_stage, - (case,), - {}, - (), - ), - } diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 93da958c..6a0d7d69 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -15,11 +15,16 @@ from PyPIC3D.utils import ( courant_condition, - update_parameters_from_toml, build_collocated_grid, - build_yee_grid, convert_to_jax_compatible, load_external_fields_from_toml, + update_parameters_from_toml, + convert_to_jax_compatible, load_external_fields_from_toml, print_stats, particle_sanity_check, build_plasma_parameters_dict, make_dir, compute_energy, add_external_fields ) +from PyPIC3D.utilities.grids import ( + build_collocated_grid, + build_tiled_yee_grids, + build_yee_grid, +) from PyPIC3D.diagnostics.plotting import ( plot_initial_histograms @@ -35,7 +40,6 @@ from PyPIC3D.solvers.yee_tiled import ( empty_tiled_vector_field, - tile_grid_axes, tile_scalar_field, tile_vector_field, ) @@ -491,18 +495,9 @@ def initialize_simulation(toml_file): tile_shape = _tile_shape_from_parameters(simulation_parameters) world["tile_shape"] = tile_shape guard_cells = int(world["guard_cells"]) - world["grids"]["tiled_center_grid"] = tile_grid_axes( - world["grids"]["center"], - world, - tile_shape, - num_guard_cells=guard_cells, - ) - world["grids"]["tiled_vertex_grid"] = tile_grid_axes( - world["grids"]["vertex"], - world, - tile_shape, - num_guard_cells=guard_cells, - ) + 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 particles, species_config = to_tiled_particles(particles, world, simulation_parameters) E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index d5c1b7c8..022e72bc 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -108,55 +108,6 @@ def _restore_tiled_vector_layout(stacked_tiles, original_tiles): return unstack_tiled_vector_field(stacked_tiles) -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. - - Each returned axis has leading tile indices followed by the local - ghost-celled coordinate line for that axis. These arrays carry the same - coordinate convention previously rebuilt inside tiled interpolation and - deposition kernels. - """ - - 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 tile_scalar_field(field, world, tile_shape, num_guard_cells=2): """ Split a ghost-celled field into compact tiles using the shared tile shape. diff --git a/PyPIC3D/utilities/grids.py b/PyPIC3D/utilities/grids.py new file mode 100644 index 00000000..8c621497 --- /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 e3a96475..2a69cdca 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -13,7 +13,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"): @@ -468,91 +467,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. diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 5dbcc02e..d2710e26 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -17,11 +17,10 @@ assemble_tiled_vector_field, fold_tiled_ghost_cells, fold_tiled_ghost_cells_periodic, - tile_grid_axes, tile_vector_field, update_tiled_vector_ghost_cells, ) -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index d7111ce7..25169860 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -9,7 +9,7 @@ calculate_electrostatic_fields, ) from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.utilities.grids import build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells jax.config.update("jax_enable_x64", True) diff --git a/tests/evolve_test.py b/tests/evolve_test.py index ca0a5ba0..db20fb23 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -10,8 +10,8 @@ from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import tile_grid_axes, tile_vector_field -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/particle_test.py b/tests/particle_test.py index ebbf1355..01af9959 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -14,8 +14,8 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes diff --git a/tests/pml_test.py b/tests/pml_test.py index 77b61a04..c15c3bde 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -21,7 +21,8 @@ update_B, update_E, ) -from PyPIC3D.utils import build_yee_grid, compute_energy +from PyPIC3D.utilities.grids import build_yee_grid +from PyPIC3D.utils import compute_energy jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 86083aaa..a2330068 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -17,11 +17,10 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_scalar_field, assemble_tiled_vector_field, - tile_grid_axes, tile_scalar_field, tile_vector_field, ) -from PyPIC3D.utils import build_collocated_grid, build_yee_grid +from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 25944d1c..4c778aca 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -30,10 +30,10 @@ assemble_tiled_vector_field, empty_tiled_scalar_field, empty_tiled_vector_field, - tile_grid_axes, tile_vector_field, update_E, ) +from PyPIC3D.utilities.grids import tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 8d590ff2..4ab81220 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -17,8 +17,8 @@ from PyPIC3D.deposition.rho import compute_tiled_mass_density_from_tiled_particles from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 89f7123b..7814db76 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -7,8 +7,8 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import tile_grid_axes, tile_vector_field -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 66408d66..219686cb 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -10,7 +10,7 @@ from PyPIC3D.deposition import rho from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.utilities.grids import build_yee_grid class TestTiledRefactorContracts(unittest.TestCase): @@ -64,33 +64,6 @@ def test_tiled_vector_field_stack_roundtrips_and_refreshes_halos(self): for reference, candidate in zip(refreshed, refreshed_from_stack): self.assertTrue(jnp.allclose(reference, candidate, rtol=1.0e-12, atol=1.0e-12)) - def test_tiled_benchmark_stage_specs_expose_tiled_hot_paths(self): - from PyPIC3D.diagnostics import tiled_benchmarking - - case = tiled_benchmarking.build_synthetic_tiled_yee_state( - nx=4, - ny=4, - nz=4, - particles_per_species=8, - species_count=1, - tile_shape=(2, 2, 2), - slots_per_tile=8, - ) - stage_specs = tiled_benchmarking.build_tiled_stage_specs(case) - - expected = { - "tiled_pic_step", - "tiled_particle_push", - "tiled_particle_retile", - "tiled_current_deposition", - "tiled_field_update", - "tiled_diagnostics", - "tiled_output_bridge", - } - self.assertTrue(expected.issubset(stage_specs)) - self.assertEqual(stage_specs["tiled_pic_step"].static_argnames, ("relativistic", "particle_pusher")) - self.assertEqual(stage_specs["tiled_current_deposition"].static_argnames, ()) - def test_evolve_loops_read_runtime_controls_from_world(self): electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) electrostatic_signature = inspect.signature(evolve.time_loop_electrostatic) @@ -108,6 +81,21 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertFalse(hasattr(yee_tiled, "update_tiled_E")) self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) self.assertFalse(hasattr(yee_tiled, "tiled_grid_axes_from_world")) + self.assertFalse(hasattr(yee_tiled, "tile_grid_axes")) + + def test_grid_builders_live_in_utilities(self): + from PyPIC3D import utils + from PyPIC3D.utilities import grids + + self.assertFalse(hasattr(utils, "build_yee_grid")) + self.assertFalse(hasattr(utils, "build_collocated_grid")) + self.assertTrue(hasattr(grids, "build_yee_grid")) + self.assertTrue(hasattr(grids, "build_collocated_grid")) + self.assertTrue(hasattr(grids, "tile_grid_axes")) + self.assertTrue(hasattr(grids, "build_tiled_yee_grids")) + + def test_tiled_benchmarking_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) def test_tiled_pusher_reads_tiled_grids_from_world(self): from PyPIC3D.pusher import tiled_pusher diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index cf0978fb..241bf560 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -7,8 +7,8 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_grid_axes, tile_scalar_field -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes jax.config.update("jax_enable_x64", True) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 50a0700b..27721b38 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -15,12 +15,12 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, empty_tiled_vector_field, - tile_grid_axes, tile_vector_field, update_B, update_E, ) -from PyPIC3D.utils import build_yee_grid, compute_energy +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utils import compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells diff --git a/tests/utils_test.py b/tests/utils_test.py index e16d62bf..c7ae2979 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -10,8 +10,9 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid from PyPIC3D.utils import ( - print_stats, build_yee_grid, build_collocated_grid, check_stability, + print_stats, check_stability, particle_sanity_check, load_external_fields_from_toml, add_external_fields, compute_energy, dump_parameters_to_toml, ) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index e9379869..61827ba4 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -6,7 +6,6 @@ from PyPIC3D.initialization import initialize_fields from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - tile_grid_axes, tile_vector_field, update_tiled_ghost_cells, update_tiled_ghost_cells_periodic, @@ -15,7 +14,7 @@ update_B, update_E, ) -from PyPIC3D.utils import build_yee_grid +from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC From 73702b492aa613bffc24ba2db9d467ece5f78817 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 15:24:43 -0500 Subject: [PATCH 42/83] refactoring code --- PyPIC3D/initialization.py | 75 ++++++++++++--------------- PyPIC3D/solvers/yee_tiled.py | 47 +++++------------ tests/evolve_test.py | 4 +- tests/pml_test.py | 9 ++-- tests/tiled_esirkepov_test.py | 34 +++++++----- tests/tiled_refactor_contract_test.py | 30 +++++++++++ tests/tiled_yee_integration_test.py | 15 ++---- tests/utils_test.py | 21 +++++--- 8 files changed, 120 insertions(+), 115 deletions(-) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 6a0d7d69..7c0f9890 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -39,7 +39,6 @@ ) from PyPIC3D.solvers.yee_tiled import ( - empty_tiled_vector_field, tile_scalar_field, tile_vector_field, ) @@ -296,9 +295,6 @@ def initialize_simulation(toml_file): 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 @@ -421,7 +417,7 @@ def initialize_simulation(toml_file): write_openpmd_initial_particles(particles, world, constants, simulation_parameters['output_dir']) # write the initial particles to an openPMD file - 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), @@ -499,12 +495,11 @@ def initialize_simulation(toml_file): world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid world["grids"]["tiled_center_grid"] = tiled_center_grid particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - E = tile_vector_field(E, world, tile_shape, num_guard_cells=guard_cells) - B = tile_vector_field(B, world, tile_shape, num_guard_cells=guard_cells) - if simulation_parameters["current_calculation"] == "esirkepov": - J = empty_tiled_vector_field(world, tile_shape, num_guard_cells=guard_cells, dtype=E[0].dtype) - else: - J = tile_vector_field(J, world, tile_shape, num_guard_cells=guard_cells) + # convert the particles to tiled storage and get the species configuration + + + + external_E, external_B = external_fields external_fields = ( tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), @@ -545,46 +540,44 @@ def initialize_simulation(toml_file): solver, electrostatic, verbose, GPUs, Nt, relativistic, particle_pusher, species_config - -def initialize_fields(Nx, Ny, Nz): +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. + Untiled arrays include one ghost cell on each side and have shape + (Nx+2, Ny+2, Nz+2). Tiled arrays use the startup tile shape and guard + depth stored in world. 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 either global or 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) + 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 + # get the grid and tile dimensions + + tiled_shape = (ntx, nty, ntz, tile_nx + 2 * g, tile_ny + 2 * g, tile_nz + 2 * g) + + Ex = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Ey = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Ez = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) # 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 = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + By = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Bz = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) # 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 = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Jy = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Jz = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) # initialize the current density arrays as 0 - - phi = jax.numpy.zeros(shape = ghost_shape) - rho = jax.numpy.zeros(shape = ghost_shape) + phi = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + rho = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) # 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/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 022e72bc..040956a2 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -40,40 +40,6 @@ def _reduced_tiled_axes(field_tiles, tile_shape, g): ) -def empty_tiled_scalar_field(world, tile_shape, num_guard_cells=2, dtype=None): - """ - Allocate an empty tile-major scalar field with ``num_guard_cells`` guards. - """ - - 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) - g = int(num_guard_cells) - if dtype is None: - dtype = jnp.float64 - - return jnp.zeros( - ( - ntx, - nty, - ntz, - tile_nx + 2 * g, - tile_ny + 2 * g, - tile_nz + 2 * g, - ), - dtype=dtype, - ) - - -def empty_tiled_vector_field(world, tile_shape, num_guard_cells=2, dtype=None): - """ - Allocate empty tile-major vector-field components. - """ - - return tuple(empty_tiled_scalar_field(world, tile_shape, num_guard_cells, dtype) for _ in range(3)) - - 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 @@ -123,7 +89,17 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): ntz = _tile_axis_count(Nz, tile_nz) if g != 1: - field_tiles = empty_tiled_scalar_field(world, tile_shape, g, field.dtype) + 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): @@ -326,6 +302,7 @@ def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_ 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), diff --git a/tests/evolve_test.py b/tests/evolve_test.py index db20fb23..8f571b05 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -60,7 +60,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field 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"]) + 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]) @@ -154,7 +154,7 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): 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"]) + 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), diff --git a/tests/pml_test.py b/tests/pml_test.py index c15c3bde..356d3432 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -16,7 +16,6 @@ from PyPIC3D.initialization import initialize_fields, initialize_simulation from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - empty_tiled_vector_field, tile_vector_field, update_B, update_E, @@ -271,7 +270,7 @@ def test_no_pml_state_returns_field_and_none_state(self): 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, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + E, B, J, _, _ = initialize_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) @@ -296,7 +295,7 @@ def test_tiled_pml_matches_single_tile_pml_for_one_yee_step(self): constants, ) tile_shape = (2, 2, 1) - E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + E, B, J, _, _ = initialize_fields(world) x = world["grids"]["vertex"][0][1:-1] y = world["grids"]["vertex"][1][1:-1] @@ -390,7 +389,7 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): tile_shape = (4, 1, 1) world["tile_shape"] = tile_shape tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) - E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + E, B, J, _, _ = initialize_fields(world) x = world["grids"]["vertex"][0][1:-1] pulse = jnp.exp(-((x + 0.30) / 0.04) ** 2) @@ -426,7 +425,7 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): ) tile_shape = (2, 1, 1) world["tile_shape"] = tile_shape - E, B, J, _, _ = initialize_fields(world["Nx"], world["Ny"], world["Nz"]) + E, B, J, _, _ = initialize_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"])) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 4c778aca..d960c40e 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -18,7 +18,7 @@ ) from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.diagnostics.output_adapters import fields_for_output -from PyPIC3D.initialization import initialize_simulation +from PyPIC3D.initialization import build_tiled_array, initialize_fields, initialize_simulation from PyPIC3D.particles.tiled_particle_refresh import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, @@ -28,8 +28,6 @@ from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - empty_tiled_scalar_field, - empty_tiled_vector_field, tile_vector_field, update_E, ) @@ -172,10 +170,11 @@ def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape } g = int(world["guard_cells"]) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + _, _, J_template, _, _ = initialize_fields(world, tiled=True) J_tiles = Esirkepov_current( tiled_particles, species_config, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), + J_template, constants, world, ) @@ -233,12 +232,13 @@ def _basic_positions_and_velocities(self, world): ) return x, u - def test_empty_tiled_vector_field_uses_startup_guard_cells(self): + 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 = empty_tiled_vector_field(world, tile_shape, num_guard_cells=int(world["guard_cells"])) + _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) self.assertEqual(J_tiles[0].shape, (4, 1, 1, 6, 5, 5)) self.assertTrue(jnp.allclose(J_tiles[0], 0.0)) @@ -254,8 +254,9 @@ 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 = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) + _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) Jx, Jy, Jz = J_tiles Jx = Jx.at[1, 0, 0, 2, 2, 2].set(3.0) @@ -274,7 +275,7 @@ def test_output_adapter_uses_world_guard_cells(self): 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 = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) + _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) Jx, Jy, Jz = J_tiles Jx = Jx.at[1, 0, 0, 2, 2, 2].set(7.0) rho = jnp.zeros(shape) @@ -297,7 +298,8 @@ def test_update_E_reads_two_guard_current_interior(self): 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) - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g) + world["tile_shape"] = tile_shape + _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) Jx, Jy, Jz = J_tiles Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) @@ -356,10 +358,11 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) g = int(world["guard_cells"]) + _, _, J_template, _, _ = initialize_fields(world, tiled=True) J_tiles = Esirkepov_current( tiled_particles, species_config, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), + J_template, constants, world, ) @@ -391,10 +394,11 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre old_species = self._species(world, x_old) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + _, _, J_template, _, _ = initialize_fields(world, tiled=True) J_tiles = Esirkepov_current( tiled_particles, species_config, - empty_tiled_vector_field(world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])), + J_template, constants, world, ) @@ -434,10 +438,11 @@ def test_public_Esirkepov_current_rejects_tiled_filter_modes(self): tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) with self.assertRaisesRegex(ValueError, "Esirkepov current filtering is not supported"): + _, _, J_template, _, _ = initialize_fields(world, tiled=True) Esirkepov_current( tiled_particles, species_config, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=int(world["guard_cells"])), + J_template, constants, world, filter="digital", @@ -544,13 +549,14 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): old_species = self._species_from_arrays(world, x_old, u) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) g = int(world["guard_cells"]) - rho_tiles = empty_tiled_scalar_field(world, tile_shape, num_guard_cells=g) + rho_tiles = build_tiled_array(world) rho_old = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) + _, _, J_template, _, _ = initialize_fields(world, tiled=True) J_tiles = Esirkepov_current( tiled_particles, species_config, - empty_tiled_vector_field(world, tile_shape, num_guard_cells=g), + J_template, constants, world, ) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 219686cb..3ef8b2c2 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -8,6 +8,7 @@ from PyPIC3D.deposition import Esirkepov from PyPIC3D.deposition import J_from_rhov from PyPIC3D.deposition import rho +from PyPIC3D.initialization import build_tiled_array, initialize_fields from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled from PyPIC3D.utilities.grids import build_yee_grid @@ -82,6 +83,35 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) self.assertFalse(hasattr(yee_tiled, "tiled_grid_axes_from_world")) self.assertFalse(hasattr(yee_tiled, "tile_grid_axes")) + self.assertFalse(hasattr(yee_tiled, "empty_tiled_scalar_field")) + self.assertFalse(hasattr(yee_tiled, "empty_tiled_vector_field")) + + def test_tiled_array_initialization_lives_in_initialization(self): + world = self._world() + world["tile_shape"] = (2, 3, 2) + world["guard_cells"] = 2 + + array = build_tiled_array(world, dtype=jnp.float32) + self.assertEqual(array.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(array.dtype, jnp.float32) + self.assertTrue(jnp.all(array == 0.0)) + + def test_initialize_fields_builds_tiled_field_state_from_world(self): + world = self._world() + world["tile_shape"] = (2, 3, 2) + world["guard_cells"] = 2 + + E, B, J, phi, rho_tiles = initialize_fields(world, tiled=True, dtype=jnp.float32) + for vector_field in (E, B, J): + self.assertEqual(len(vector_field), 3) + for component in vector_field: + self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(component.dtype, jnp.float32) + self.assertTrue(jnp.all(component == 0.0)) + for scalar_field in (phi, rho_tiles): + self.assertEqual(scalar_field.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(scalar_field.dtype, jnp.float32) + self.assertTrue(jnp.all(scalar_field == 0.0)) def test_grid_builders_live_in_utilities(self): from PyPIC3D import utils diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 27721b38..c71650ed 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -8,13 +8,12 @@ import toml from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.initialization import initialize_simulation, validate_field_solver +from PyPIC3D.initialization import initialize_fields, initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - empty_tiled_vector_field, tile_vector_field, update_B, update_E, @@ -77,13 +76,9 @@ def _world_with_tiled_grids(self, world, tile_shape): return world def _empty_fields(self, world): - shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - zero = jnp.zeros(shape) - E = (zero, zero, zero) - B = (zero, zero, zero) - J = (zero, zero, zero) + E, B, J, phi, rho = initialize_fields(world) external_fields = (E, B) - return E, B, J, zero, zero, external_fields, None + return E, B, J, rho, phi, external_fields, None def _species(self, world): return particle_species( @@ -193,7 +188,7 @@ def _long_two_stream_state(self, current_deposition="direct", current_filter="no tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) g = int(tiled_world["guard_cells"]) if current_deposition == "esirkepov": - J_tiles = empty_tiled_vector_field(tiled_world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) + _, _, J_tiles, _, _ = initialize_fields(tiled_world, tiled=True, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) @@ -282,7 +277,7 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, current_depos E, B, J, rho, phi, external_fields, pml_state = fields g = int(world["guard_cells"]) if current_deposition == "esirkepov": - J_tiles = empty_tiled_vector_field(world, tile_shape, num_guard_cells=g, dtype=E[0].dtype) + _, _, J_tiles, _, _ = initialize_fields(world, tiled=True, dtype=E[0].dtype) else: J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) tiled_fields = ( diff --git a/tests/utils_test.py b/tests/utils_test.py index c7ae2979..db4e43e2 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -20,6 +20,10 @@ jax.config.update("jax_enable_x64", True) +def _field_world(Nx, Ny, Nz): + return {"Nx": Nx, "Ny": Ny, "Nz": Nz} + + class TestUtilsFunctions(unittest.TestCase): def setUp(self): self.world = { @@ -92,7 +96,7 @@ def test_add_external_fields_adds_components(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -108,7 +112,7 @@ def test_load_external_fields_defaults_to_evolved_fields(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -124,7 +128,7 @@ def test_load_external_fields_evolve_true_uses_evolved_fields(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -140,7 +144,7 @@ def test_load_external_fields_evolve_false_uses_external_fields(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -153,7 +157,7 @@ def test_load_external_fields_rejects_external_current(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -166,7 +170,7 @@ def test_load_external_fields_preserves_shape_validation(self): 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) + E, B, J, phi, rho = initialize_fields(_field_world(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)) @@ -183,7 +187,7 @@ def test_energy_can_include_external_fields(self): self.assertEqual(kinetic_energy, 0.0) def test_compute_energy_ignores_inactive_particles(self): - E, B, J, phi, rho = initialize_fields(1, 1, 1) + E, B, J, phi, rho = initialize_fields(_field_world(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( @@ -396,7 +400,8 @@ def test_vtk_plot_fields_assembles_tiled_components_at_output_boundary(self): "guard_cells": 2, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, } - field = tuple(jnp.zeros((8, 6, 6)) + component for component in (1.0, 2.0, 3.0)) + field_shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + field = tuple(jnp.zeros(field_shape) + component for component in (1.0, 2.0, 3.0)) tiled_field = tile_vector_field(field, world, world["tile_shape"], num_guard_cells=world["guard_cells"]) with unittest.mock.patch.object(vtk_diagnostics, "gridToVTK") as grid_to_vtk: From 51777d0429845b4002bb29cdc92d412cead1bf63 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 16:03:09 -0500 Subject: [PATCH 43/83] refactoring tiling --- PyPIC3D/initialization.py | 101 ++++++++++++-------------- PyPIC3D/solvers/yee_tiled.py | 89 ++++------------------- PyPIC3D/utilities/grids.py | 6 +- PyPIC3D/utils.py | 63 +++++++++++++--- tests/direct_deposition_tiled_test.py | 42 +++++------ tests/evolve_test.py | 53 +++++--------- tests/openpmd_test.py | 6 +- tests/particle_test.py | 15 ++-- tests/pml_test.py | 24 ++++-- tests/tiled_electrostatic_test.py | 22 ++---- tests/tiled_esirkepov_test.py | 41 +++++------ tests/tiled_fluid_quantities_test.py | 17 +---- tests/tiled_pusher_test.py | 23 +++--- tests/tiled_refactor_contract_test.py | 43 ++++++----- tests/tiled_rho_test.py | 17 +---- tests/tiled_yee_integration_test.py | 38 +++++----- tests/utils_test.py | 71 ++++++++++++------ tests/yee_tiled_test.py | 28 ++++--- 18 files changed, 336 insertions(+), 363 deletions(-) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 7c0f9890..8a9f653d 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -39,8 +39,8 @@ ) from PyPIC3D.solvers.yee_tiled import ( - tile_scalar_field, - tile_vector_field, + update_tiled_vector_ghost_cells, + update_tiled_vector_ghost_cells_for_pml, ) @@ -52,11 +52,9 @@ from PyPIC3D.pusher.particle_push import validate_particle_pusher 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_tiled_pml_state, load_pml_from_toml, - update_ghost_cells_for_pml, ) @@ -382,6 +380,13 @@ def initialize_simulation(toml_file): '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") @@ -428,30 +433,25 @@ 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 @@ -488,30 +488,15 @@ def initialize_simulation(toml_file): species_config = None - tile_shape = _tile_shape_from_parameters(simulation_parameters) - world["tile_shape"] = tile_shape - guard_cells = int(world["guard_cells"]) - 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 particles, species_config = to_tiled_particles(particles, world, simulation_parameters) # convert the particles to tiled storage and get the species configuration - - - external_E, external_B = external_fields - external_fields = ( - tile_vector_field(external_E, world, tile_shape, num_guard_cells=guard_cells), - tile_vector_field(external_B, world, tile_shape, num_guard_cells=guard_cells), - ) - simulation_parameters["tile_shape"] = tile_shape + external_fields = (external_E, external_B) print(f"Using tiled Yee storage with tile shape: {tile_shape}") overflow = jnp.asarray(False) if electrostatic: - rho = tile_scalar_field(rho, world, tile_shape, num_guard_cells=guard_cells) - phi = tile_scalar_field(phi, world, tile_shape, num_guard_cells=guard_cells) fields = (E, B, J, rho, phi, external_fields, None, overflow) # define the fields tuple for the electrostatic solver else: @@ -540,16 +525,9 @@ def initialize_simulation(toml_file): solver, electrostatic, verbose, GPUs, Nt, relativistic, particle_pusher, species_config -def initialize_fields(world): +def build_tiled_array(world, dtype=jnp.float64): """ - Initializes the electric and magnetic field arrays, as well as the electric potential and charge density arrays. - - Untiled arrays include one ghost cell on each side and have shape - (Nx+2, Ny+2, Nz+2). Tiled arrays use the startup tile shape and guard - depth stored in world. - - Returns: - E, B, J, phi, rho initialized to zero in either global or tiled storage. + Build one zero-filled tiled field component from the geometry in world. """ tile_nx, tile_ny, tile_nz = [int(width) for width in world["tile_shape"]] @@ -560,24 +538,35 @@ def initialize_fields(world): ntx = Nx // tile_nx nty = Ny // tile_ny ntz = Nz // tile_nz - # get the grid and tile dimensions - 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. + + 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, B, J, phi, rho initialized to zero in tiled storage. + """ - Ex = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - Ey = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - Ez = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Ex = build_tiled_array(world) + Ey = build_tiled_array(world) + Ez = build_tiled_array(world) # initialize the electric field arrays as 0 - Bx = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - By = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - Bz = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Bx = build_tiled_array(world) + By = build_tiled_array(world) + Bz = build_tiled_array(world) # initialize the magnetic field arrays as 0 - Jx = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - Jy = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - Jz = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + Jx = build_tiled_array(world) + Jy = build_tiled_array(world) + Jz = build_tiled_array(world) # initialize the current density arrays as 0 - phi = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) - rho = jnp.zeros(shape=tiled_shape, dtype=jnp.float64) + 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/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 040956a2..8d8e73c2 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -29,8 +29,6 @@ def _backward_slice(g): def _reduced_tiled_axes(field_tiles, tile_shape, g): - if tile_shape is None: - tile_shape = tuple(int(width) - 2 * int(g) for width in field_tiles.shape[3:]) tile_nx, tile_ny, tile_nz = [int(width) for width in tile_shape] ntx, nty, ntz = field_tiles.shape[:3] return ( @@ -44,7 +42,7 @@ 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): +def _stack_tiled_vector_field(field_tiles): """ Stack tiled vector components onto one leading component axis. @@ -58,7 +56,7 @@ def stack_tiled_vector_field(field_tiles): return jnp.stack(field_tiles, axis=0) -def unstack_tiled_vector_field(field_tiles): +def _unstack_tiled_vector_field(field_tiles): """ Return a tiled vector field as the existing ``(Fx, Fy, Fz)`` tuple. """ @@ -71,7 +69,7 @@ def unstack_tiled_vector_field(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) + return _unstack_tiled_vector_field(stacked_tiles) def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): @@ -130,22 +128,6 @@ def tile_at(tx, ty, tz): ) -def _tile_scalar_field(field, tile_shape): - return tile_scalar_field(field, None, tile_shape) - - -def tile_vector_field(field, world, tile_shape, num_guard_cells=2): - """ - Split ``(Fx, Fy, Fz)`` into compact ghost-celled tiles. - - Each scalar component gets leading tile axes ``(ntx, nty, ntz)`` followed by - the tile-local ghost-celled field shape - ``(tile_nx + 2, tile_ny + 2, tile_nz + 2)``. - """ - - return tuple(tile_scalar_field(component, world, tile_shape, num_guard_cells) for component in field) - - 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. @@ -183,10 +165,6 @@ def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells= return tuple(assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells) for component in field_tiles) -def _assemble_scalar_tiles(field_tiles, world, tile_shape): - return assemble_tiled_scalar_field(field_tiles, world, tile_shape) - - def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): """ Refresh tile halos using the field boundary conditions. @@ -199,6 +177,8 @@ def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=N """ 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"] @@ -288,22 +268,13 @@ def update_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=N return field_tiles -def update_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=2): - """ - Refresh tile halos with all-periodic boundary conditions. - """ - - world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} - return update_tiled_ghost_cells(field_tiles, world, num_guard_cells) - - 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) + 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, @@ -312,20 +283,6 @@ def update_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_ return _restore_tiled_vector_layout(refreshed, field_tiles) -def update_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): - """ - Refresh periodic tile halos for each component of a vector field. - """ - - stacked_tiles = stack_tiled_vector_field(field_tiles) - refreshed = jax.vmap( - lambda component: update_tiled_ghost_cells_periodic(component, num_guard_cells), - 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. @@ -340,7 +297,10 @@ def update_tiled_ghost_cells_for_pml(field_tiles, world, num_guard_cells=2, tile 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}} + 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) @@ -349,7 +309,7 @@ def update_tiled_vector_ghost_cells_for_pml(field_tiles, world, num_guard_cells= Refresh tile halos for a vector field with PML-active exterior walls. """ - stacked_tiles = stack_tiled_vector_field(field_tiles) + 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, @@ -426,6 +386,8 @@ def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=Non """ 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"] @@ -539,21 +501,12 @@ def fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=Non return field_tiles -def fold_tiled_ghost_cells_periodic(field_tiles, num_guard_cells=2): - """ - Fold all-periodic tile-ghost deposits for one scalar component. - """ - - world = {"boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}} - return fold_tiled_ghost_cells(field_tiles, world, num_guard_cells) - - 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) + 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, @@ -562,20 +515,6 @@ def fold_tiled_vector_ghost_cells(field_tiles, world, num_guard_cells=2, tile_sh return _restore_tiled_vector_layout(folded, field_tiles) -def fold_tiled_vector_ghost_cells_periodic(field_tiles, num_guard_cells=2): - """ - Fold all-periodic tile-ghost deposits for each vector component. - """ - - stacked_tiles = stack_tiled_vector_field(field_tiles) - folded = jax.vmap( - lambda component: fold_tiled_ghost_cells_periodic(component, num_guard_cells), - in_axes=0, - out_axes=0, - )(stacked_tiles) - return _restore_tiled_vector_layout(folded, field_tiles) - - def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ Update compact tiled electric fields without assembling a global field. diff --git a/PyPIC3D/utilities/grids.py b/PyPIC3D/utilities/grids.py index 8c621497..c053f769 100644 --- a/PyPIC3D/utilities/grids.py +++ b/PyPIC3D/utilities/grids.py @@ -90,7 +90,7 @@ def _tile_grid_axis(global_axis_grid, world, tile_shape, tile_counts, axis_index return jnp.broadcast_to(axis_lines.reshape(axis_shape), tiled_shape) -def tile_grid_axes(grid, world, tile_shape, num_guard_cells=2): +def _tile_grid_axes(grid, world, tile_shape, num_guard_cells=2): """ Build tile-local coordinate lines for a center or vertex grid. """ @@ -121,13 +121,13 @@ def build_tiled_yee_grids(world, tile_shape=None, num_guard_cells=None): if num_guard_cells is None: num_guard_cells = int(world["guard_cells"]) - tiled_vertex_grid = tile_grid_axes( + tiled_vertex_grid = _tile_grid_axes( world["grids"]["vertex"], world, tile_shape, num_guard_cells=num_guard_cells, ) - tiled_center_grid = tile_grid_axes( + tiled_center_grid = _tile_grid_axes( world["grids"]["center"], world, tile_shape, diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 2a69cdca..a180107a 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -537,7 +537,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. @@ -545,6 +577,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. @@ -562,31 +595,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}") diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index d2710e26..358d1916 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -16,16 +16,19 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, fold_tiled_ghost_cells, - fold_tiled_ghost_cells_periodic, - tile_vector_field, + tile_scalar_field, update_tiled_vector_ghost_cells, ) -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) +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 TestDirectDepositionTiled(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 @@ -80,18 +83,9 @@ def _world_with_tiled_grids(self, world, tile_shape): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) + 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 @@ -320,23 +314,23 @@ def test_tiled_digital_filter_matches_global_digital_filter(self): 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) + folded = 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)) - legacy_folded = fold_tiled_ghost_cells_periodic(tiles, num_guard_cells=1) - self.assertTrue(jnp.allclose(folded, legacy_folded, rtol=1.0e-12, atol=1.0e-12)) - 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) @@ -344,7 +338,7 @@ def test_fold_tiled_ghost_cells_two_guard_layers_adds_deposits_to_neighbors(self 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells) + folded = 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) @@ -357,6 +351,8 @@ def test_fold_tiled_ghost_cells_two_guard_reduced_axis_folds_to_single_active_ce 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) @@ -377,13 +373,14 @@ def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) + folded = 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) @@ -400,6 +397,7 @@ def test_fold_tiled_ghost_cells_matches_global_fold_for_mixed_boundaries(self): 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)) @@ -421,7 +419,7 @@ def test_fold_tiled_ghost_cells_matches_global_fold_for_mixed_boundaries(self): field = field.at[3, 3, -1].set(5.0) tiles = tiles.at[1, 1, -1, 1, 1, -1].set(5.0) - folded_tiles = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1) + folded_tiles = fold_tiled_ghost_cells(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( diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 8f571b05..8d6b5ca9 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -10,8 +10,7 @@ from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import tile_vector_field -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) @@ -19,18 +18,9 @@ def add_tiled_grids_to_world(world, tile_shape): g = int(world["guard_cells"]) world["tile_shape"] = tile_shape - world["grids"]["tiled_center_grid"] = tile_grid_axes( - world["grids"]["center"], - world, - tile_shape, - num_guard_cells=g, - ) - world["grids"]["tiled_vertex_grid"] = tile_grid_axes( - world["grids"]["vertex"], - world, - tile_shape, - num_guard_cells=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 class TestEvolveExternalFields(unittest.TestCase): @@ -60,6 +50,9 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field world["grids"] = {"center": center_grid, "vertex": vertex_grid} constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} + 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) @@ -88,24 +81,19 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field dt=world["dt"], ) ] - tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } - world["guard_cells"] = 2 - add_tiled_grids_to_world(world, tile_shape) tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - external_E_tiles = tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])) - external_B_tiles = tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])) fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + E, + B, + J, rho, phi, - (external_E_tiles, external_B_tiles), + external_fields, None, jnp.asarray(False), ) @@ -130,7 +118,7 @@ def test_external_electric_field_pushes_particles_without_evolving_maxwell_field 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(external_after[0][0], external_E_tiles[0])) + self.assertTrue(jnp.allclose(external_after[0][0], external_fields[0][0])) def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): world = { @@ -154,6 +142,9 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): world["grids"] = {"center": center_grid, "vertex": vertex_grid} constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} + 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), @@ -183,25 +174,19 @@ def test_absorbing_particle_mask_survives_jitted_electrodynamic_step(self): dt=world["dt"], ) ] - tile_shape = (world["Nx"], world["Ny"], world["Nz"]) simulation_parameters = { "particle_tile_nx": tile_shape[0], "particle_tile_ny": tile_shape[1], "particle_tile_nz": tile_shape[2], } - world["guard_cells"] = 2 - add_tiled_grids_to_world(world, tile_shape) tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(B, world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(J, world, tile_shape, num_guard_cells=int(world["guard_cells"])), + E, + B, + J, rho, phi, - ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=int(world["guard_cells"])), - ), + external_fields, None, jnp.asarray(False), ) diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index 9e679996..9c245f15 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -5,11 +5,15 @@ from PyPIC3D.diagnostics import openPMD from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array -from PyPIC3D.solvers.yee_tiled import tile_scalar_field, tile_vector_field +from PyPIC3D.solvers.yee_tiled import tile_scalar_field from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +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 diff --git a/tests/particle_test.py b/tests/particle_test.py index 01af9959..d70205bf 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -15,7 +15,7 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -27,18 +27,13 @@ def _compute_rho_one_tile(particles, rho, world, constants): world["guard_cells"] = int(world.get("guard_cells", 2)) world["tile_shape"] = (world["Nx"], world["Ny"], world["Nz"]) grids = dict(world["grids"]) - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], + tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids( world, world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), - ) - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), + int(world["guard_cells"]), ) + grids["tiled_vertex_grid"] = tiled_vertex_grid + grids["tiled_center_grid"] = tiled_center_grid world["grids"] = grids simulation_parameters = { diff --git a/tests/pml_test.py b/tests/pml_test.py index 356d3432..37e6dafe 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -13,10 +13,10 @@ tile_pml_profiles, update_ghost_cells_for_pml, ) -from PyPIC3D.initialization import initialize_fields, initialize_simulation +from PyPIC3D.initialization import initialize_simulation from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - tile_vector_field, + tile_scalar_field, update_B, update_E, ) @@ -26,6 +26,18 @@ jax.config.update("jax_enable_x64", True) +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 _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, @@ -270,7 +282,7 @@ def test_no_pml_state_returns_field_and_none_state(self): 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, _, _ = initialize_fields(world) + 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) @@ -295,7 +307,7 @@ def test_tiled_pml_matches_single_tile_pml_for_one_yee_step(self): constants, ) tile_shape = (2, 2, 1) - E, B, J, _, _ = initialize_fields(world) + E, B, J = _empty_global_fields(world) x = world["grids"]["vertex"][0][1:-1] y = world["grids"]["vertex"][1][1:-1] @@ -389,7 +401,7 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): tile_shape = (4, 1, 1) world["tile_shape"] = tile_shape tiled_pml_state = initialize_tiled_pml_state(world, tile_shape) - E, B, J, _, _ = initialize_fields(world) + E, B, J = _empty_global_fields(world) x = world["grids"]["vertex"][0][1:-1] pulse = jnp.exp(-((x + 0.30) / 0.04) ** 2) @@ -425,7 +437,7 @@ def test_tiled_pml_step_uses_shared_guard_current_without_changing_timing(self): ) tile_shape = (2, 1, 1) world["tile_shape"] = tile_shape - E, B, J, _, _ = initialize_fields(world) + 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"])) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index a2330068..c854b5c6 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -18,9 +18,8 @@ assemble_tiled_scalar_field, assemble_tiled_vector_field, tile_scalar_field, - tile_vector_field, ) -from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_collocated_grid, build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) @@ -29,6 +28,10 @@ ROUND_OFF_ATOL = 1.0e-11 +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 TestTiledElectrostatic(unittest.TestCase): def _build_world(self, shape_factor=1): world = { @@ -57,18 +60,9 @@ def _world_with_tiled_grids(self, world, tile_shape): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) + 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 diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index d960c40e..7934a983 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -28,14 +28,18 @@ from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - tile_vector_field, + tile_scalar_field, update_E, ) -from PyPIC3D.utilities.grids import tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids jax.config.update("jax_enable_x64", True) +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 TestTiledEsirkepovCurrent(unittest.TestCase): def _build_world( self, @@ -146,18 +150,9 @@ def _world_with_tiled_grids(self, world, tile_shape): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) + 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 @@ -170,7 +165,7 @@ def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape } g = int(world["guard_cells"]) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - _, _, J_template, _, _ = initialize_fields(world, tiled=True) + _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -238,7 +233,7 @@ def test_initialize_fields_builds_tiled_current_with_startup_guard_cells(self): tile_shape = (2, 1, 1) world["tile_shape"] = tile_shape - _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) + _, _, 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)) @@ -256,7 +251,7 @@ def test_shared_guard_current_tiles_assemble_to_one_guard_global_current(self): tile_shape = (2, 1, 1) world["tile_shape"] = tile_shape g = int(world["guard_cells"]) - _, _, J_tiles, _, _ = initialize_fields(world, tiled=True) + _, _, J_tiles, _, _ = initialize_fields(world) Jx, Jy, Jz = J_tiles Jx = Jx.at[1, 0, 0, 2, 2, 2].set(3.0) @@ -275,7 +270,7 @@ def test_output_adapter_uses_world_guard_cells(self): 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, tiled=True) + _, _, 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) @@ -299,7 +294,7 @@ def test_update_E_reads_two_guard_current_interior(self): 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, tiled=True) + _, _, J_tiles, _, _ = initialize_fields(world) Jx, Jy, Jz = J_tiles Jx = Jx.at[:, :, :, 2:-2, 2:-2, 2:-2].set(4.0) @@ -358,7 +353,7 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) g = int(world["guard_cells"]) - _, _, J_template, _, _ = initialize_fields(world, tiled=True) + _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -394,7 +389,7 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre old_species = self._species(world, x_old) tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - _, _, J_template, _, _ = initialize_fields(world, tiled=True) + _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, species_config, @@ -438,7 +433,7 @@ def test_public_Esirkepov_current_rejects_tiled_filter_modes(self): tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) with self.assertRaisesRegex(ValueError, "Esirkepov current filtering is not supported"): - _, _, J_template, _, _ = initialize_fields(world, tiled=True) + _, _, J_template, _, _ = initialize_fields(world) Esirkepov_current( tiled_particles, species_config, @@ -552,7 +547,7 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): rho_tiles = build_tiled_array(world) rho_old = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) - _, _, J_template, _, _ = initialize_fields(world, tiled=True) + _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, species_config, diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 4ab81220..f4a76173 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -18,7 +18,7 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) @@ -56,18 +56,9 @@ def _world_with_tiled_grids(self, world): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) + 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 diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 7814db76..3e58247b 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -7,13 +7,17 @@ from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import tile_vector_field -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.solvers.yee_tiled import tile_scalar_field +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) +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 = { @@ -37,18 +41,9 @@ def _build_world(self, Nx=8, Ny=6, Nz=4, shape_factor=1): 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) - world["grids"]["tiled_center_grid"] = tile_grid_axes( - world["grids"]["center"], - world, - tile_shape, - num_guard_cells=g, - ) - world["grids"]["tiled_vertex_grid"] = tile_grid_axes( - world["grids"]["vertex"], - world, - tile_shape, - num_guard_cells=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 _deterministic_vector_field(self, world, scale): diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 3ef8b2c2..ea93bc6f 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -11,7 +11,7 @@ from PyPIC3D.initialization import build_tiled_array, initialize_fields from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled -from PyPIC3D.utilities.grids import build_yee_grid +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid class TestTiledRefactorContracts(unittest.TestCase): @@ -29,9 +29,14 @@ def _world(self): "z_wind": 2.0, "shape_factor": 1, "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "tile_shape": (2, 3, 2), + "guard_cells": 2, } vertex_grid, center_grid = build_yee_grid(world) world["grids"] = {"vertex": vertex_grid, "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 _field(self, world, scale): @@ -45,25 +50,16 @@ def _field(self, world, scale): shape = (Nx + 2, Ny + 2, Nz + 2) return jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (ii + 2.0 * jj - kk)) - def test_tiled_vector_field_stack_roundtrips_and_refreshes_halos(self): + def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): world = self._world() tile_shape = (2, 3, 2) E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) - E_tiles = yee_tiled.tile_vector_field(E, world, tile_shape) + E_tiles = tuple(yee_tiled.tile_scalar_field(component, world, tile_shape) for component in E) - stacked = yee_tiled.stack_tiled_vector_field(E_tiles) - self.assertEqual(stacked.shape[0], 3) - - unstacked = yee_tiled.unstack_tiled_vector_field(stacked) - for original, recovered in zip(E_tiles, unstacked): - self.assertTrue(jnp.allclose(original, recovered, rtol=1.0e-12, atol=1.0e-12)) - - refreshed = yee_tiled.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=1, tile_shape=tile_shape) - refreshed_from_stack = yee_tiled.unstack_tiled_vector_field( - yee_tiled.update_tiled_vector_ghost_cells(stacked, world, num_guard_cells=1, tile_shape=tile_shape) - ) - for reference, candidate in zip(refreshed, refreshed_from_stack): - self.assertTrue(jnp.allclose(reference, candidate, rtol=1.0e-12, atol=1.0e-12)) + refreshed = yee_tiled.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + self.assertEqual(len(refreshed), 3) + for component in refreshed: + self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) def test_evolve_loops_read_runtime_controls_from_world(self): electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) @@ -83,6 +79,13 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) self.assertFalse(hasattr(yee_tiled, "tiled_grid_axes_from_world")) self.assertFalse(hasattr(yee_tiled, "tile_grid_axes")) + self.assertFalse(hasattr(yee_tiled, "tile_vector_field")) + self.assertFalse(hasattr(yee_tiled, "stack_tiled_vector_field")) + self.assertFalse(hasattr(yee_tiled, "unstack_tiled_vector_field")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells_periodic")) + self.assertFalse(hasattr(yee_tiled, "fold_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(yee_tiled, "fold_tiled_vector_ghost_cells_periodic")) self.assertFalse(hasattr(yee_tiled, "empty_tiled_scalar_field")) self.assertFalse(hasattr(yee_tiled, "empty_tiled_vector_field")) @@ -101,16 +104,16 @@ def test_initialize_fields_builds_tiled_field_state_from_world(self): world["tile_shape"] = (2, 3, 2) world["guard_cells"] = 2 - E, B, J, phi, rho_tiles = initialize_fields(world, tiled=True, dtype=jnp.float32) + E, B, J, phi, rho_tiles = initialize_fields(world) for vector_field in (E, B, J): self.assertEqual(len(vector_field), 3) for component in vector_field: self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(component.dtype, jnp.float32) + self.assertEqual(component.dtype, jnp.float64) self.assertTrue(jnp.all(component == 0.0)) for scalar_field in (phi, rho_tiles): self.assertEqual(scalar_field.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(scalar_field.dtype, jnp.float32) + self.assertEqual(scalar_field.dtype, jnp.float64) self.assertTrue(jnp.all(scalar_field == 0.0)) def test_grid_builders_live_in_utilities(self): @@ -121,7 +124,7 @@ def test_grid_builders_live_in_utilities(self): self.assertFalse(hasattr(utils, "build_collocated_grid")) self.assertTrue(hasattr(grids, "build_yee_grid")) self.assertTrue(hasattr(grids, "build_collocated_grid")) - self.assertTrue(hasattr(grids, "tile_grid_axes")) + self.assertFalse(hasattr(grids, "tile_grid_axes")) self.assertTrue(hasattr(grids, "build_tiled_yee_grids")) def test_tiled_benchmarking_module_is_removed(self): diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 241bf560..331433f8 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -8,7 +8,7 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) @@ -46,18 +46,9 @@ def _world_with_tiled_grids(self, world, simulation_parameters): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) + 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 diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index c71650ed..d0afce9f 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -8,17 +8,17 @@ import toml from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.initialization import initialize_fields, initialize_simulation, validate_field_solver +from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - tile_vector_field, + tile_scalar_field, update_B, update_E, ) -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid from PyPIC3D.utils import compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells @@ -30,6 +30,10 @@ ROUND_OFF_ATOL = 1.0e-11 +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 TestTiledYeeIntegration(unittest.TestCase): def _build_world(self, boundary_conditions=None): if boundary_conditions is None: @@ -60,23 +64,19 @@ def _world_with_tiled_grids(self, world, tile_shape): world = dict(world) grids = dict(world["grids"]) world["tile_shape"] = tile_shape - grids["tiled_center_grid"] = tile_grid_axes( - grids["center"], - world, - tile_shape, - num_guard_cells=g, - ) - grids["tiled_vertex_grid"] = tile_grid_axes( - grids["vertex"], - world, - tile_shape, - num_guard_cells=g, - ) + 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_fields(self, world): - E, B, J, phi, rho = initialize_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)) + phi = jnp.zeros(shape) + rho = jnp.zeros(shape) external_fields = (E, B) return E, B, J, rho, phi, external_fields, None @@ -188,7 +188,8 @@ def _long_two_stream_state(self, current_deposition="direct", current_filter="no tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) g = int(tiled_world["guard_cells"]) if current_deposition == "esirkepov": - _, _, J_tiles, _, _ = initialize_fields(tiled_world, tiled=True, dtype=E[0].dtype) + J_template = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) + J_tiles = tuple(jnp.zeros_like(component) for component in J_template) else: J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) @@ -277,7 +278,8 @@ def _build_tiled_state(self, particles, fields, world, tile_shape, current_depos E, B, J, rho, phi, external_fields, pml_state = fields g = int(world["guard_cells"]) if current_deposition == "esirkepov": - _, _, J_tiles, _, _ = initialize_fields(world, tiled=True, dtype=E[0].dtype) + J_template = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + J_tiles = tuple(jnp.zeros_like(component) for component in J_template) else: J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) tiled_fields = ( diff --git a/tests/utils_test.py b/tests/utils_test.py index db4e43e2..e74e71ad 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -9,7 +9,7 @@ from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics -from PyPIC3D.solvers.yee_tiled import tile_vector_field +from PyPIC3D.solvers.yee_tiled import tile_scalar_field from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid from PyPIC3D.utils import ( print_stats, check_stability, @@ -21,7 +21,30 @@ jax.config.update("jax_enable_x64", True) def _field_world(Nx, Ny, Nz): - return {"Nx": Nx, "Ny": Ny, "Nz": 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): @@ -96,7 +119,8 @@ def test_add_external_fields_adds_components(self): 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(_field_world(2, 2, 2)) + 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)) @@ -105,14 +129,15 @@ def test_load_external_fields_defaults_to_evolved_fields(self): 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) + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + 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): - E, B, J, phi, rho = initialize_fields(_field_world(2, 2, 2)) + 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)) @@ -121,14 +146,15 @@ def test_load_external_fields_evolve_true_uses_evolved_fields(self): 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) + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + 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): - E, B, J, phi, rho = initialize_fields(_field_world(2, 2, 2)) + 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)) @@ -137,14 +163,15 @@ def test_load_external_fields_evolve_false_uses_external_fields(self): 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) + fields, external_fields = load_external_fields_from_toml(fields, external_fields, config, world) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + 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): - E, B, J, phi, rho = initialize_fields(_field_world(2, 2, 2)) + 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)) @@ -154,10 +181,11 @@ def test_load_external_fields_rejects_external_current(self): 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) + load_external_fields_from_toml(fields, external_fields, config, world) def test_load_external_fields_preserves_shape_validation(self): - E, B, J, phi, rho = initialize_fields(_field_world(2, 2, 2)) + 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)) @@ -167,18 +195,19 @@ def test_load_external_fields_preserves_shape_validation(self): 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) + load_external_fields_from_toml(fields, external_fields, config, world) def test_energy_can_include_external_fields(self): - E, B, J, phi, rho = initialize_fields(_field_world(1, 1, 1)) + 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 = (slice(1, -1), slice(1, -1), slice(1, -1)) + 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 = {"dx": 1.0, "dy": 1.0, "dz": 1.0, "Nx": 1, "Ny": 1, "Nz": 1} + world = world | {"dx": 1.0, "dy": 1.0, "dz": 1.0} constants = {"eps": 2.0, "mu": 4.0, "C": 10.0} e_energy, b_energy, kinetic_energy = compute_energy([], total_E, total_B, world, constants) @@ -187,8 +216,8 @@ def test_energy_can_include_external_fields(self): self.assertEqual(kinetic_energy, 0.0) def test_compute_energy_ignores_inactive_particles(self): - E, B, J, phi, rho = initialize_fields(_field_world(1, 1, 1)) - world = {"dx": 1.0, "dy": 1.0, "dz": 1.0, "Nx": 1, "Ny": 1, "Nz": 1} + world = _field_world(1, 1, 1) | {"dx": 1.0, "dy": 1.0, "dz": 1.0} + E, B, J, phi, rho = initialize_fields(world) constants = {"eps": 1.0, "mu": 1.0, "C": 10.0} species = particle_species( name="absorbed", @@ -402,7 +431,7 @@ def test_vtk_plot_fields_assembles_tiled_components_at_output_boundary(self): } field_shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) field = tuple(jnp.zeros(field_shape) + component for component in (1.0, 2.0, 3.0)) - tiled_field = tile_vector_field(field, world, world["tile_shape"], num_guard_cells=world["guard_cells"]) + tiled_field = _tile_vector_field(field, world, world["tile_shape"], num_guard_cells=world["guard_cells"]) with unittest.mock.patch.object(vtk_diagnostics, "gridToVTK") as grid_to_vtk: vtk_diagnostics.plot_fields( diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 61827ba4..8cbfa2a7 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -3,18 +3,15 @@ import jax import jax.numpy as jnp -from PyPIC3D.initialization import initialize_fields from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - tile_vector_field, + tile_scalar_field, update_tiled_ghost_cells, - update_tiled_ghost_cells_periodic, update_tiled_vector_ghost_cells, - update_tiled_vector_ghost_cells_periodic, update_B, update_E, ) -from PyPIC3D.utilities.grids import build_yee_grid, tile_grid_axes +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC @@ -22,6 +19,10 @@ jax.config.update("jax_enable_x64", True) +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 TestYeeTiled(unittest.TestCase): def _build_world(self): world = { @@ -50,6 +51,9 @@ def _conducting_world(self): 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): @@ -200,8 +204,8 @@ def test_tile_grid_axes_include_configured_guard_cells(self): world = self._build_world() tile_shape = (2, 3, 2) g = 2 - tiled_center_grid = tile_grid_axes(world["grids"]["center"], world, tile_shape, num_guard_cells=g) - tiled_vertex_grid = tile_grid_axes(world["grids"]["vertex"], world, tile_shape, num_guard_cells=g) + 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)) @@ -223,6 +227,7 @@ def test_tile_grid_axes_include_configured_guard_cells(self): 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] @@ -233,18 +238,19 @@ def test_update_tiled_ghost_cells_periodic_refreshes_neighbor_halos(self): stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) - refreshed = update_tiled_ghost_cells_periodic(stale_tiles) + refreshed = 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 = update_tiled_vector_ghost_cells_periodic(stale_tiles) + refreshed = 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)) @@ -252,6 +258,7 @@ def test_update_tiled_vector_ghost_cells_periodic_refreshes_each_component(self) 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] @@ -276,6 +283,7 @@ def test_update_tiled_ghost_cells_conducting_matches_global_ghost_cells(self): 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] @@ -301,6 +309,7 @@ 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, @@ -326,6 +335,7 @@ def test_update_tiled_ghost_cells_two_guard_layers_matches_global_refresh(self): 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) From 55ae2c2903cd5eb7f9bc05fa21ba96fdeff654ba Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sat, 4 Jul 2026 16:52:25 -0500 Subject: [PATCH 44/83] refactoring ghost cell logic --- PyPIC3D/boundary_conditions/PML.py | 38 +- PyPIC3D/boundary_conditions/__init__.py | 15 +- .../boundary_conditions/boundaryconditions.py | 232 +-------- PyPIC3D/boundary_conditions/ghost_cells.py | 418 ++++++++++++++++ PyPIC3D/deposition/Esirkepov.py | 2 +- PyPIC3D/deposition/J_from_rhov.py | 2 +- PyPIC3D/deposition/rho.py | 232 +-------- PyPIC3D/diagnostics/fluid_quantities.py | 253 +--------- PyPIC3D/initialization.py | 2 +- PyPIC3D/solvers/electrostatic_yee.py | 75 +-- PyPIC3D/solvers/yee_tiled.py | 449 ++---------------- tests/boundaryconditions_test.py | 232 ++------- tests/direct_deposition_tiled_test.py | 74 ++- tests/electrostatic_yee_test.py | 35 +- tests/pml_test.py | 41 +- tests/tiled_fluid_quantities_test.py | 173 ++++--- tests/tiled_refactor_contract_test.py | 31 +- tests/tiled_yee_integration_test.py | 35 +- tests/yee_tiled_test.py | 46 +- 19 files changed, 897 insertions(+), 1488 deletions(-) create mode 100644 PyPIC3D/boundary_conditions/ghost_cells.py diff --git a/PyPIC3D/boundary_conditions/PML.py b/PyPIC3D/boundary_conditions/PML.py index bfa336ef..87c7875c 100644 --- a/PyPIC3D/boundary_conditions/PML.py +++ b/PyPIC3D/boundary_conditions/PML.py @@ -3,7 +3,6 @@ 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 @@ -157,11 +156,18 @@ def _tile_scalar_profile(profile, tile_shape, num_guard_cells=2): 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 = profile[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz] - profile_tiles = profile_tiles.at[tx, ty, tz, g:-g, g:-g, g:-g].set(interior) + 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): @@ -431,23 +437,3 @@ def apply_tiled_pml_to_b_curl(derivatives, world, pml_state): ) return (curl_x, curl_y, curl_z), (e_memory, b_memory, tiled_profiles) - - -def update_ghost_cells_for_pml(field, world): - """ - 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. - """ - 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) 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 0fbce74f..b1bf2096 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -7,8 +7,8 @@ particle_axis_offset, ) 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.tiled_particles import SpeciesConfig, TiledParticles -from PyPIC3D.solvers.yee_tiled import update_tiled_vector_ghost_cells def shift_old_stencil(old_w_list, shift): diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index 00ed27d7..b5760eab 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -6,7 +6,7 @@ ) from PyPIC3D.deposition.shapes import get_first_order_weights, get_second_order_weights -from PyPIC3D.solvers.yee_tiled import ( +from PyPIC3D.boundary_conditions.ghost_cells import ( fold_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells, ) diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 9bcac765..5745b2da 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -9,10 +9,9 @@ collapse_axis_stencil, prepare_particle_axis_stencil, ) -from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells +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.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import update_tiled_ghost_cells from PyPIC3D.utilities.filters import digital_filter @@ -29,100 +28,6 @@ def _species_slot_counts(tiled_particles): return jnp.where(species_counts > 0, species_counts, 1.0) -def _deposit_tiled_scalar_moment( - tiled_particles, - rho, - world, - grid, - scalar_weight, - position=None, - divide_by_volume=True, -): - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - 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"] - - 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, y, z = position - x = x.reshape(-1) - y = y.reshape(-1) - z = z.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, - 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, - ) - z, _, 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, - ) - - 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) - - rho = jnp.zeros_like(rho) - - 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 * scalar_density * x_weights[i] * y_weights[j] * z_weights[k], - mode="drop", - ) - - rho = fold_ghost_cells(rho, bc_x, bc_y, bc_z) - rho = update_ghost_cells(rho, bc_x, bc_y, bc_z) - - return rho - - def _diagnostic_mass_position(tiled_particles, world): """ Mirror the ordinary scalar-diagnostic particle position for tiled particles. @@ -339,27 +244,6 @@ def compute_rho( return rho_tiles -@jit -def compute_mass_density_from_tiled_particles(tiled_particles, species_config, rho, world, grid=None): - """Compute mass density on the vertex grid from tile-major particles.""" - 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( - tiled_particles, - rho, - world, - grid, - scalar_weight, - position=position, - ) - - @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.""" @@ -383,28 +267,6 @@ def compute_tiled_mass_density_from_tiled_particles(tiled_particles, species_con ) -@partial(jit, static_argnames=("direction",)) -def compute_velocity_field_from_tiled_particles(tiled_particles, species_config, field, direction, world, grid=None): - """Compute the legacy scalar velocity diagnostic from tile-major particles.""" - 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( - tiled_particles, - field, - world, - grid, - scalar_weight, - position=position, - divide_by_volume=False, - ) - - @partial(jit, static_argnames=("direction", "tile_shape", "g")) def compute_tiled_velocity_field_from_tiled_particles( tiled_particles, @@ -438,98 +300,6 @@ def compute_tiled_velocity_field_from_tiled_particles( ) -@partial(jit, static_argnames=("direction",)) -def compute_pressure_field_from_tiled_particles( - tiled_particles, - species_config, - field, - velocity_field, - direction, - world, - grid=None, -): - """Compute the legacy scalar pressure diagnostic from tile-major particles.""" - if grid is None: - grid = world["grids"]["vertex"] - - dx = world["dx"] - dy = world["dy"] - dz = world["dz"] - Nx, Ny, Nz = field.shape - bc_x = world["boundary_conditions"]["x"] - bc_y = world["boundary_conditions"]["y"] - bc_z = world["boundary_conditions"]["z"] - shape_factor = world["shape_factor"] - - x, y, z = _diagnostic_mass_position(tiled_particles, world) - x = x.reshape(-1) - y = y.reshape(-1) - z = z.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, - 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, - ) - z, _, 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, - ) - - 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) - - field = jnp.zeros_like(field) - - for i in range(xpts.shape[0]): - for j in range(ypts.shape[0]): - for k in range(zpts.shape[0]): - 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", - ) - - field = fold_ghost_cells(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field - - @partial(jit, static_argnames=("direction", "tile_shape", "g")) def compute_tiled_pressure_field_from_tiled_particles( tiled_particles, diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index a141be1c..fcca6905 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -1,22 +1,10 @@ - -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 from PyPIC3D.deposition.rho import ( - compute_mass_density_from_tiled_particles, - compute_pressure_field_from_tiled_particles, compute_tiled_mass_density_from_tiled_particles, compute_tiled_pressure_field_from_tiled_particles, compute_tiled_velocity_field_from_tiled_particles, - compute_velocity_field_from_tiled_particles, ) from PyPIC3D.particles.tiled_particles import TiledParticles -import jax -import jax.numpy as jnp -from jax import jit -from functools import partial - def compute_mass_density(particles, rho, world, species_config=None): """ @@ -45,75 +33,9 @@ def compute_mass_density(particles, rho, world, species_config=None): tile_shape=tuple(int(width) for width in world["tile_shape"]), g=int(world["guard_cells"]), ) - return compute_mass_density_from_tiled_particles(particles, species_config, rho, world) - - raise ValueError("Public compute_mass_density requires TiledParticles; use _compute_mass_density_flat for flat reference tests.") - - -@jit -def _compute_mass_density_flat(particles, rho, world): - 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'] - # get the shape of the charge density array - - rho = jnp.zeros_like(rho) - # reset rho to zero - - for species in particles: - 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 particle position used by the existing scalar diagnostic - active = species.get_active_mask().astype(x.dtype) + raise ValueError("compute_mass_density requires tile-major scalar field storage.") - _, _, 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 + raise ValueError("Public compute_mass_density requires TiledParticles.") def compute_velocity_field(particles, field, direction, world, species_config=None): @@ -128,96 +50,9 @@ def compute_velocity_field(particles, field, direction, world, species_config=No tile_shape=tuple(int(width) for width in world["tile_shape"]), g=int(world["guard_cells"]), ) - return compute_velocity_field_from_tiled_particles(particles, species_config, field, int(direction), world) - - raise ValueError("Public compute_velocity_field requires TiledParticles; use _compute_velocity_field_flat for flat reference tests.") - - -@partial(jit, static_argnames=("direction",)) -def _compute_velocity_field_flat(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'] - shape_factor = world['shape_factor'] - # get the shape of the velocity field array - - field = jnp.zeros_like(field) - # reset field to zero - - for species in particles: - 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 - - + 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): @@ -232,82 +67,6 @@ def compute_pressure_field(particles, field, velocity_field, direction, world, s tile_shape=tuple(int(width) for width in world["tile_shape"]), g=int(world["guard_cells"]), ) - return compute_pressure_field_from_tiled_particles( - particles, - species_config, - field, - velocity_field, - int(direction), - world, - ) - - raise ValueError("Public compute_pressure_field requires TiledParticles; use _compute_pressure_field_flat for flat reference tests.") - - -@partial(jit, static_argnames=("direction",)) -def _compute_pressure_field_flat(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'] - shape_factor = world['shape_factor'] - # get the shape of the velocity field array - - field = jnp.zeros_like(field) - # reset field to zero - - for species in 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) - - - 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 + raise ValueError("compute_pressure_field requires tile-major scalar field storage.") - field = fold_ghost_cells(field, bc_x, bc_y, bc_z) - field = update_ghost_cells(field, bc_x, bc_y, bc_z) - return field + raise ValueError("Public compute_pressure_field requires TiledParticles.") diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 8a9f653d..4afd8e9c 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -38,7 +38,7 @@ to_tiled_particles ) -from PyPIC3D.solvers.yee_tiled import ( +from PyPIC3D.boundary_conditions.ghost_cells import ( update_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells_for_pml, ) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index 9636686e..f5fe333b 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -1,23 +1,16 @@ 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.utilities.filters import digital_filter -from PyPIC3D.boundary_conditions.boundaryconditions import ( - update_ghost_cells, apply_scalar_conducting_bc -) +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_scalar_field, tile_scalar_field, - update_tiled_ghost_cells, - update_tiled_vector_ghost_cells, ) -@partial(jit, static_argnames=("tol", "max_iter")) 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. @@ -40,10 +33,9 @@ def solve_poisson_with_conjugate_gradient(rho, phi, constants, world, tol=1e-12, 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 + tile_shape = tuple(int(width) for width in world["tile_shape"]) + g = int(world["guard_cells"]) + # get the tiled layout from the world contract def lapl(field): # Laplacian using ghost-cell neighbors instead of jnp.roll @@ -54,10 +46,11 @@ def lapl(field): # 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 + field_tiles = tile_scalar_field(field, world, tile_shape, num_guard_cells=g) + 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 assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells=g) + # apply scalar conducting boundaries and refresh ghost cells through the tiled halo path def body_fun(state): phi, r, p, k = state @@ -108,7 +101,6 @@ 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. @@ -130,19 +122,22 @@ 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"]) phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) - phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) + phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) + phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) # 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_tiles = tile_scalar_field(phi, world, tile_shape, 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) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) # update ghost cells after filtering Ex, Ey, Ez = centered_finite_difference_gradient(-1 * phi[1:-1, 1:-1, 1:-1], dx, dy, dz, bc) @@ -155,9 +150,19 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver 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) + E_tiles = tuple( + ghost_cells.update_tiled_ghost_cells( + tile_scalar_field(component, world, tile_shape, num_guard_cells=g), + world, + g, + tile_shape, + ) + for component in (Ex_full, Ey_full, Ez_full) + ) + Ex_full, Ey_full, Ez_full = tuple( + assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells=g) + for component in E_tiles + ) return (Ex_full, Ey_full, Ez_full), phi, rho @@ -180,7 +185,7 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): forward = slice(g + 1, None if g == 1 else -g + 1) backward = slice(g - 1, -g - 1) - phi_tiles = update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + 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) @@ -196,7 +201,7 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): -1.0 * (phi_tiles[:, :, :, active, active, forward] - phi_tiles[:, :, :, active, active, backward]) / (2.0 * dz) ) - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) + 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): @@ -210,10 +215,6 @@ def calculate_tiled_electrostatic_fields(world, particles, species_config, const del bc - 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"]) rho_tiles = compute_rho(particles, species_config, rho_tiles, constants, world) @@ -222,13 +223,17 @@ def calculate_tiled_electrostatic_fields(world, particles, species_config, const phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) - phi = update_ghost_cells(phi, bc_x, bc_y, bc_z) + phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) + phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) # refresh ghost cells before filtering and tiled differentiation 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_tiles = tile_scalar_field(phi, world, tile_shape, 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) + phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) # keep the same global phi post-processing order as the untiled solver phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/yee_tiled.py index 8d8e73c2..3f19ab99 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/yee_tiled.py @@ -5,8 +5,7 @@ apply_tiled_pml_to_b_curl, apply_tiled_pml_to_e_curl, ) -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.utilities.filters import digital_filter_vector @@ -28,50 +27,6 @@ def _backward_slice(g): return slice(g - 1, -g - 1) -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): - """ - Stack tiled vector components onto one leading component axis. - - The public tiled field state still uses the existing ``(Fx, Fy, Fz)`` - tuples. This packed view gives halo/fold/filter routines one vector-field - surface, which is the shape later shard-local exchange should operate on. - """ - - 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): - """ - Return a tiled vector field as the existing ``(Fx, Fy, Fz)`` tuple. - """ - - 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 tile_scalar_field(field, world, tile_shape, num_guard_cells=2): """ Split a ghost-celled field into compact tiles using the shared tile shape. @@ -106,7 +61,7 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): 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 update_tiled_ghost_cells(field_tiles, world, g, tile_shape) + 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) @@ -145,16 +100,20 @@ def assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells= for tx in range(ntx): for ty in range(nty): for tz in range(ntz): - interior = field_tiles[tx, ty, tz, g:-g, g:-g, g:-g] - ix = 1 + tx * tile_nx - iy = 1 + ty * tile_ny - iz = 1 + tz * tile_nz - field = field.at[ix:ix + tile_nx, iy:iy + tile_ny, iz:iz + tile_nz].set(interior) - - 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) + 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): @@ -165,356 +124,6 @@ def assemble_tiled_vector_field(field_tiles, world, tile_shape, num_guard_cells= return tuple(assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells) for component in field_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 - # update_ghost_cells on the assembled field. - 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 the global 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 fold_tiled_ghost_cells(field_tiles, world, num_guard_cells=2, tile_shape=None): - """ - Add tile-ghost deposits into owning interiors, then clear ghosts. - - This is the current-deposition analogue of ``fold_ghost_cells`` for a - tiled layout. 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 sign convention - used by the 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) - - def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -531,9 +140,9 @@ def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): active = _active_slice(g) forward = _forward_slice(g) if pml_state is None: - Bx, By, Bz = update_tiled_vector_ghost_cells(B_tiles, world, g, tile_shape) + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells(B_tiles, world, g, tile_shape) else: - Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml(B_tiles, world, g, tile_shape) + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells_for_pml(B_tiles, world, g, tile_shape) Jx, Jy, Jz = J_tiles current = _active_slice(g) @@ -577,20 +186,20 @@ def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): ) if pml_state is None: - Ex, Ey, Ez = update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape) else: - Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape) + 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. Ex, Ey, Ez = digital_filter_vector((Ex, Ey, Ez), constants.get("alpha", 1.0), num_guard_cells=g) - Ex, Ey, Ez = apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) + Ex, Ey, Ez = ghost_cells.apply_tiled_conducting_bc((Ex, Ey, Ez), world, num_guard_cells=g) if pml_state is None: - return update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape), None + return ghost_cells.update_tiled_vector_ghost_cells((Ex, Ey, Ez), world, g, tile_shape), None - return update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state + return ghost_cells.update_tiled_vector_ghost_cells_for_pml((Ex, Ey, Ez), world, g, tile_shape), pml_state def update_B(E_tiles, B_tiles, world, constants, pml_state=None): @@ -609,9 +218,9 @@ def update_B(E_tiles, B_tiles, world, constants, pml_state=None): active = _active_slice(g) backward = _backward_slice(g) if pml_state is None: - Ex, Ey, Ez = update_tiled_vector_ghost_cells(E_tiles, world, g, tile_shape) + Ex, Ey, Ez = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, g, tile_shape) else: - Ex, Ey, Ez = update_tiled_vector_ghost_cells_for_pml(E_tiles, world, g, tile_shape) + 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"] @@ -641,15 +250,15 @@ def update_B(E_tiles, B_tiles, world, constants, pml_state=None): Bz = Bz.at[:, :, :, active, active, active].set(Bz[:, :, :, active, active, active] - dt * curl_z) if pml_state is None: - Bx, By, Bz = update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) + Bx, By, Bz = ghost_cells.update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape) else: - Bx, By, Bz = update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape) + 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. Bx, By, Bz = digital_filter_vector((Bx, By, Bz), constants.get("alpha", 1.0), num_guard_cells=g) if pml_state is None: - return update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape), None + return ghost_cells.update_tiled_vector_ghost_cells((Bx, By, Bz), world, g, tile_shape), None - return update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape), pml_state + return ghost_cells.update_tiled_vector_ghost_cells_for_pml((Bx, By, Bz), world, g, tile_shape), pml_state diff --git a/tests/boundaryconditions_test.py b/tests/boundaryconditions_test.py index 0593b536..d0b4b435 100644 --- a/tests/boundaryconditions_test.py +++ b/tests/boundaryconditions_test.py @@ -1,15 +1,12 @@ 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 ) +from PyPIC3D.boundary_conditions import ghost_cells jax.config.update("jax_enable_x64", True) @@ -31,185 +28,58 @@ def test_apply_zero_boundary_condition(self): class TestGhostCells(unittest.TestCase): - """Tests for ghost cell functions used in the ghost cell boundary condition approach.""" + """Tests for the tiled 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.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): + 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) + + result = ghost_cells.update_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) + + self.assertTrue(jnp.all(result[0, 0, 0, -1, 1:3, 1:3] == 2.0)) + self.assertTrue(jnp.all(result[1, 0, 0, 0, 1:3, 1:3] == 1.0)) + + def test_fold_tiled_ghost_cells_periodic_adds_to_owner_tile(self): + field_tiles = jnp.zeros((2, 1, 1, 4, 4, 4)) + field_tiles = field_tiles.at[0, 0, 0, -1, 2, 2].set(3.0) + field_tiles = field_tiles.at[1, 0, 0, 0, 2, 2].set(5.0) + + result = ghost_cells.fold_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) + + self.assertAlmostEqual(float(result[1, 0, 0, 1, 2, 2]), 3.0) + self.assertAlmostEqual(float(result[0, 0, 0, -2, 2, 2]), 5.0) + self.assertEqual(float(result[0, 0, 0, -1, 2, 2]), 0.0) + self.assertEqual(float(result[1, 0, 0, 0, 2, 2]), 0.0) + + def test_apply_tiled_conducting_bc_zeros_global_tangential_faces(self): + world = { + "tile_shape": self.tile_shape, + "boundary_conditions": {"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING}, + } + E = tuple(jnp.ones((1, 1, 1, 4, 4, 4)) for _ in range(3)) + + Ex, Ey, Ez = ghost_cells.apply_tiled_conducting_bc(E, world, self.g) + + 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)) + + 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)) + + result = ghost_cells.apply_tiled_scalar_conducting_bc(field, self.world, self.g) self.assertTrue(jnp.allclose(result, field)) diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 358d1916..9ecc2c0a 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -3,7 +3,7 @@ import jax import jax.numpy as jnp -from PyPIC3D.boundary_conditions.boundaryconditions import fold_ghost_cells, update_ghost_cells +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 ( @@ -15,9 +15,7 @@ from PyPIC3D.particles.tiled_particle_refresh import refresh_tiled_particle_tiles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, - fold_tiled_ghost_cells, tile_scalar_field, - update_tiled_vector_ghost_cells, ) from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -29,6 +27,53 @@ 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 _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 TestDirectDepositionTiled(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 @@ -292,20 +337,20 @@ def test_tiled_digital_filter_matches_global_digital_filter(self): 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) + 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 = 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 = update_tiled_vector_ghost_cells(filtered_tiles, world, num_guard_cells=1, tile_shape=tile_shape) + 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) + _update_ghost_cells(digital_filter(component, alpha), bc_x, bc_y, bc_z) for component in J ) @@ -319,7 +364,7 @@ def test_fold_tiled_ghost_cells_periodic_adds_current_deposits_to_neighbors(self 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=world["tile_shape"]) + 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) @@ -338,7 +383,7 @@ def test_fold_tiled_ghost_cells_two_guard_layers_adds_deposits_to_neighbors(self 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells, tile_shape=world["tile_shape"]) + 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) @@ -360,7 +405,7 @@ def test_fold_tiled_ghost_cells_two_guard_reduced_axis_folds_to_single_active_ce 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells, tile_shape=tile_shape) + 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)) @@ -380,7 +425,7 @@ def test_fold_tiled_ghost_cells_conducting_reflects_exterior_deposits(self): 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 = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=world["tile_shape"]) + 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) @@ -419,10 +464,11 @@ def test_fold_tiled_ghost_cells_matches_global_fold_for_mixed_boundaries(self): field = field.at[3, 3, -1].set(5.0) tiles = tiles.at[1, 1, -1, 1, 1, -1].set(5.0) - folded_tiles = fold_tiled_ghost_cells(tiles, world, num_guard_cells=1, tile_shape=tile_shape) + 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( + folded_reference = _update_ghost_cells( + _fold_ghost_cells( field, world["boundary_conditions"]["x"], world["boundary_conditions"]["y"], diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index 25169860..e2ef661e 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -9,8 +9,8 @@ calculate_electrostatic_fields, ) from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient -from PyPIC3D.utilities.grids import build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells +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) @@ -22,6 +22,28 @@ def apply_negative_laplacian(field, dx, dy, dz): return -(laplacian_x + laplacian_y + laplacian_z) +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 TestElectrostaticYeeMethods(unittest.TestCase): def setUp(self): self.Nx = 16 @@ -55,13 +77,18 @@ def setUp(self): "y_wind": self.y_wind, "z_wind": self.z_wind, "shape_factor": 1, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, + "guard_cells": 2, + "tile_shape": (8, 8, 8), + "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 # 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)) @@ -102,7 +129,7 @@ def test_solve_poisson_with_conjugate_gradient_single_mode(self): # 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) + rho = _update_ghost_cells(rho, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) phi = solve_poisson_with_conjugate_gradient( rho, diff --git a/tests/pml_test.py b/tests/pml_test.py index 37e6dafe..461a0768 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -11,8 +11,8 @@ initialize_tiled_pml_state, load_pml_from_toml, tile_pml_profiles, - update_ghost_cells_for_pml, ) +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.initialization import initialize_simulation from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, @@ -30,6 +30,41 @@ 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)) @@ -326,8 +361,8 @@ def test_tiled_pml_matches_single_tile_pml_for_one_yee_step(self): 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) + 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 diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index f4a76173..44dd50b7 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -10,11 +10,12 @@ compute_mass_density, compute_pressure_field, compute_velocity_field, - _compute_mass_density_flat, - _compute_pressure_field_flat, - _compute_velocity_field_flat, ) -from PyPIC3D.deposition.rho import compute_tiled_mass_density_from_tiled_particles +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.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field @@ -46,12 +47,13 @@ def _build_world(self, shape_factor=2): world["grids"] = {"vertex": vertex_grid, "center": center_grid} return world - def _world_with_tiled_grids(self, world): - tile_shape = ( - self._simulation_parameters()["particle_tile_nx"], - self._simulation_parameters()["particle_tile_ny"], - self._simulation_parameters()["particle_tile_nz"], - ) + 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"]) @@ -62,6 +64,83 @@ def _world_with_tiled_grids(self, world): 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 = to_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 = to_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 = to_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)) @@ -118,46 +197,25 @@ def _particles(self, world): ) return [electrons, ions] - def test_tiled_mass_density_matches_compute_mass_density_on_interior(self): + 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 = to_tiled_particles(particles, world, self._simulation_parameters()) - inactive = ~tiled_particles.active - x = tiled_particles.x.at[inactive, 0].set(0.10) - x = x.at[inactive, 1].set(0.20) - x = x.at[inactive, 2].set(-0.30) - tiled_particles = tiled_particles._replace(x=x) - - mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) - mass_tiled = compute_mass_density(tiled_particles, self._empty_scalar(world), world, species_config=species_config) - - self.assertTrue( - jnp.allclose( - mass_tiled[1:-1, 1:-1, 1:-1], - mass_reference[1:-1, 1:-1, 1:-1], - rtol=1.0e-12, - atol=1.0e-12, - ) - ) + 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 = to_tiled_particles(particles, world, self._simulation_parameters()) - rho_tiles = tile_scalar_field( - self._empty_scalar(world), - world, - world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), - ) - - mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) + 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, - rho_tiles, + self._scalar_tiles(world), world, tile_shape=world["tile_shape"], g=int(world["guard_cells"]), @@ -183,14 +241,9 @@ def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(s world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) - rho_tiles = tile_scalar_field( - self._empty_scalar(world), - world, - world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), - ) - - mass_reference = _compute_mass_density_flat(particles, self._empty_scalar(world), world) + 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, @@ -214,17 +267,11 @@ def test_tile_major_velocity_field_assembles_to_global_velocity_field(self): world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) - field_tiles = tile_scalar_field( - self._empty_scalar(world), - world, - world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), - ) - - velocity_reference = _compute_velocity_field_flat(particles, self._empty_scalar(world), 0, world) + 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, - field_tiles, + self._scalar_tiles(world), 0, world, species_config=species_config, @@ -250,21 +297,9 @@ def test_tile_major_pressure_field_assembles_to_global_pressure_field(self): world = self._world_with_tiled_grids(world) particles = self._particles(world) tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) - field_tiles = tile_scalar_field( - self._empty_scalar(world), - world, - world["tile_shape"], - num_guard_cells=int(world["guard_cells"]), - ) - - velocity_reference = _compute_velocity_field_flat(particles, self._empty_scalar(world), 0, world) - pressure_reference = _compute_pressure_field_flat( - particles, - self._empty_scalar(world), - velocity_reference, - 0, - world, - ) + 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, diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index ea93bc6f..6eae7cd6 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -8,6 +8,8 @@ from PyPIC3D.deposition import Esirkepov from PyPIC3D.deposition import J_from_rhov from PyPIC3D.deposition import rho +from PyPIC3D.boundary_conditions import boundaryconditions +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.initialization import build_tiled_array, initialize_fields from PyPIC3D.solvers import electrostatic_yee from PyPIC3D.solvers import yee_tiled @@ -56,7 +58,7 @@ def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) E_tiles = tuple(yee_tiled.tile_scalar_field(component, world, tile_shape) for component in E) - refreshed = yee_tiled.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + refreshed = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) self.assertEqual(len(refreshed), 3) for component in refreshed: self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) @@ -82,6 +84,13 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertFalse(hasattr(yee_tiled, "tile_vector_field")) self.assertFalse(hasattr(yee_tiled, "stack_tiled_vector_field")) self.assertFalse(hasattr(yee_tiled, "unstack_tiled_vector_field")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells_for_pml")) + self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells_for_pml")) + self.assertFalse(hasattr(yee_tiled, "fold_tiled_ghost_cells")) + self.assertFalse(hasattr(yee_tiled, "fold_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(yee_tiled, "apply_tiled_conducting_bc")) self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells_periodic")) self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells_periodic")) self.assertFalse(hasattr(yee_tiled, "fold_tiled_ghost_cells_periodic")) @@ -89,6 +98,26 @@ def test_tiled_yee_updates_read_tile_metadata_from_world(self): self.assertFalse(hasattr(yee_tiled, "empty_tiled_scalar_field")) self.assertFalse(hasattr(yee_tiled, "empty_tiled_vector_field")) + def test_ghost_cell_logic_lives_in_boundary_conditions_ghost_cells(self): + for name in ( + "update_tiled_ghost_cells", + "update_tiled_vector_ghost_cells", + "update_tiled_ghost_cells_for_pml", + "update_tiled_vector_ghost_cells_for_pml", + "fold_tiled_ghost_cells", + "fold_tiled_vector_ghost_cells", + "apply_tiled_conducting_bc", + ): + self.assertTrue(hasattr(ghost_cells, name)) + + for name in ( + "update_ghost_cells", + "fold_ghost_cells", + "apply_conducting_bc", + "apply_scalar_conducting_bc", + ): + self.assertFalse(hasattr(boundaryconditions, name)) + def test_tiled_array_initialization_lives_in_initialization(self): world = self._world() world["tile_shape"] = (2, 3, 2) diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index d0afce9f..86c60d94 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -21,7 +21,6 @@ from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid from PyPIC3D.utils import compute_energy from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells jax.config.update("jax_enable_x64", True) @@ -34,6 +33,28 @@ 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 TestTiledYeeIntegration(unittest.TestCase): def _build_world(self, boundary_conditions=None): if boundary_conditions is None: @@ -432,12 +453,12 @@ def test_tiled_yee_field_update_applies_standard_field_filter(self): x = jnp.arange(world["Nx"], dtype=jnp.float64)[:, None, None] seed = jnp.zeros(shape) - Ex = update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ey = update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ez = update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bx = update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - By = update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bz = update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ex = _update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ey = _update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ez = _update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bx = _update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + By = _update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bz = _update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) E = (Ex, Ey, Ez) B = (Bx, By, Bz) J = (seed, seed, seed) diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index 8cbfa2a7..cfe027be 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -6,13 +6,11 @@ from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_vector_field, tile_scalar_field, - update_tiled_ghost_cells, - update_tiled_vector_ghost_cells, update_B, update_E, ) +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid -from PyPIC3D.boundary_conditions.boundaryconditions import update_ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC @@ -23,6 +21,28 @@ 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 = { @@ -65,7 +85,7 @@ 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) + 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"] @@ -238,7 +258,7 @@ def test_update_tiled_ghost_cells_periodic_refreshes_neighbor_halos(self): stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) - refreshed = update_tiled_ghost_cells(stale_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + 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)) @@ -250,7 +270,7 @@ def test_update_tiled_vector_ghost_cells_periodic_refreshes_each_component(self) 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 = update_tiled_vector_ghost_cells(stale_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + 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)) @@ -269,8 +289,8 @@ def test_update_tiled_ghost_cells_conducting_matches_global_ghost_cells(self): stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) - refreshed = update_tiled_ghost_cells(stale_tiles, world) - reference = update_ghost_cells( + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world) + reference = _update_ghost_cells( field, world["boundary_conditions"]["x"], world["boundary_conditions"]["y"], @@ -294,8 +314,8 @@ def test_update_tiled_ghost_cells_mixed_matches_global_ghost_cells(self): stale_tiles = stale_tiles.at[:, :, :, :, :, 0].set(-500.0) stale_tiles = stale_tiles.at[:, :, :, :, :, -1].set(-600.0) - refreshed = update_tiled_ghost_cells(stale_tiles, world) - reference = update_ghost_cells( + refreshed = ghost_cells.update_tiled_ghost_cells(stale_tiles, world) + reference = _update_ghost_cells( field, world["boundary_conditions"]["x"], world["boundary_conditions"]["y"], @@ -328,7 +348,7 @@ def test_update_tiled_ghost_cells_two_guard_layers_matches_global_refresh(self): stale_tiles = stale_tiles.at[:, :, :, :, :, :num_guard_cells].set(-500.0) stale_tiles = stale_tiles.at[:, :, :, :, :, -num_guard_cells:].set(-600.0) - refreshed = update_tiled_ghost_cells(stale_tiles, world, num_guard_cells) + 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)) @@ -340,10 +360,10 @@ def test_update_tiled_vector_ghost_cells_conducting_refreshes_each_component(sel 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 = update_tiled_vector_ghost_cells(stale_tiles, world) + refreshed = ghost_cells.update_tiled_vector_ghost_cells(stale_tiles, world) for original, refreshed_component in zip(E, refreshed): - reference = update_ghost_cells( + reference = _update_ghost_cells( original, world["boundary_conditions"]["x"], world["boundary_conditions"]["y"], From 4ff9d767a6c47a61ade1bd2042a154ae5834d634 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 03:47:48 -0500 Subject: [PATCH 45/83] moving electrostatic yee to a global tile --- PyPIC3D/initialization.py | 6 + PyPIC3D/solvers/electrostatic_yee.py | 130 ++++++------- tests/electrostatic_yee_test.py | 49 ++--- tests/tiled_electrostatic_test.py | 262 +++++--------------------- tests/tiled_refactor_contract_test.py | 7 +- 5 files changed, 133 insertions(+), 321 deletions(-) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 4afd8e9c..52196961 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -286,6 +286,12 @@ def initialize_simulation(toml_file): relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] validate_particle_pusher(particle_pusher) + 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) diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index f5fe333b..b03bf882 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -5,10 +5,30 @@ from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient from PyPIC3D.utilities.filters import digital_filter from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_scalar_field, - tile_scalar_field, -) + + +def _active_slice(g): + return slice(g, -g) + + +def _forward_slice(g): + return slice(g + 1, None if g == 1 else -g + 1) + + +def _backward_slice(g): + return slice(g - 1, -g - 1) + + +def _as_single_tile(field): + return field[jnp.newaxis, jnp.newaxis, jnp.newaxis, :, :, :] + + +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] def solve_poisson_with_conjugate_gradient(rho, phi, constants, world, tol=1e-12, max_iter=5000): @@ -18,47 +38,47 @@ def solve_poisson_with_conjugate_gradient(rho, phi, constants, world, tol=1e-12, 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'] - tile_shape = tuple(int(width) for width in world["tile_shape"]) g = int(world["guard_cells"]) - # get the tiled layout from the world contract + 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_tiles = tile_scalar_field(field, world, tile_shape, num_guard_cells=g) - 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 assemble_tiled_scalar_field(field_tiles, world, tile_shape, num_guard_cells=g) + 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 @@ -67,7 +87,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 @@ -83,11 +103,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 @@ -105,14 +125,14 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver """ 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. @@ -124,45 +144,31 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver dz = world['dz'] tile_shape = tuple(int(width) for width in world["tile_shape"]) g = int(world["guard_cells"]) + active = _active_slice(g) phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) - phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) - phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) + 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_tiles = tile_scalar_field(phi, world, tile_shape, 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) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) + 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 = centered_finite_difference_gradient(-1 * phi[1:-1, 1:-1, 1:-1], dx, dy, dz, bc) + Ex, Ey, Ez = centered_finite_difference_gradient(-1 * phi[active, active, active], dx, dy, dz, bc) # 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) - E_tiles = tuple( - ghost_cells.update_tiled_ghost_cells( - tile_scalar_field(component, world, tile_shape, num_guard_cells=g), - world, - g, - tile_shape, - ) - for component in (Ex_full, Ey_full, Ez_full) - ) - Ex_full, Ey_full, Ez_full = tuple( - assemble_tiled_scalar_field(component, world, tile_shape, num_guard_cells=g) - for component in E_tiles - ) + 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 @@ -206,11 +212,12 @@ def _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g): def calculate_tiled_electrostatic_fields(world, particles, species_config, constants, rho_tiles, phi_tiles, solver, bc): """ - Compute electrostatic fields from tiled rho deposition and a global Poisson solve. + Compute electrostatic fields from single-tile rho deposition and a Poisson solve. - Rho is deposited into tile-major scalar storage, assembled for the existing - Poisson solver, and the solved global potential is tiled again before the - electric field is differentiated on tile-local halos. + 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 @@ -218,25 +225,20 @@ def calculate_tiled_electrostatic_fields(world, particles, species_config, const 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 = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) + rho = rho_tiles[0, 0, 0] + phi = phi_tiles[0, 0, 0] phi = solve_poisson_with_conjugate_gradient(rho, phi, constants, world) - - phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) + phi_tiles = phi_tiles.at[0, 0, 0].set(phi) phi_tiles = ghost_cells.update_tiled_ghost_cells(phi_tiles, world, g, tile_shape) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) # refresh ghost cells before filtering and tiled differentiation alpha = constants["alpha"] - phi = digital_filter(phi, alpha) - phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) + 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) - phi = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) - # keep the same global phi post-processing order as the untiled solver + # keep the same phi post-processing order as the previous electrostatic solver - phi_tiles = tile_scalar_field(phi, world, tile_shape, num_guard_cells=g) E_tiles = _centered_tiled_electrostatic_gradient(phi_tiles, world, tile_shape, g) return E_tiles, phi_tiles, rho_tiles diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index e2ef661e..8b981723 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -22,28 +22,6 @@ def apply_negative_laplacian(field, dx, dy, dz): return -(laplacian_x + laplacian_y + laplacian_z) -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 TestElectrostaticYeeMethods(unittest.TestCase): def setUp(self): self.Nx = 16 @@ -78,7 +56,7 @@ def setUp(self): "z_wind": self.z_wind, "shape_factor": 1, "guard_cells": 2, - "tile_shape": (8, 8, 8), + "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) @@ -90,9 +68,11 @@ def setUp(self): self.world["grids"]["tiled_vertex_grid"] = tiled_vertex_grid self.world["grids"]["tiled_center_grid"] = tiled_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.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 = [ particle_species( @@ -127,9 +107,8 @@ def test_solve_poisson_with_conjugate_gradient_single_mode(self): 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, BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + 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, @@ -141,7 +120,7 @@ def test_solve_poisson_with_conjugate_gradient_single_mode(self): ) # Compare on interior - phi_num = phi[1:-1, 1:-1, 1:-1] + 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) @@ -169,17 +148,17 @@ def test_solver_mode_mapping(self): 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], + -1.0 * expected_phi_fdtd[self.active, self.active, self.active], self.dx, self.dy, self.dz, "periodic", ) - 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)) + self.assertTrue(jnp.allclose(phi_fdtd[self.active, self.active, self.active], expected_phi_fdtd[self.active, self.active, self.active], atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E_fdtd[0][self.active, self.active, self.active], expected_Ex_fdtd, atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E_fdtd[1][self.active, self.active, self.active], expected_Ey_fdtd, atol=1e-6, rtol=1e-6)) + self.assertTrue(jnp.allclose(E_fdtd[2][self.active, self.active, self.active], expected_Ez_fdtd, atol=1e-6, rtol=1e-6)) if __name__ == "__main__": diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index c854b5c6..443e17c5 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -16,7 +16,6 @@ from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.yee_tiled import ( assemble_tiled_scalar_field, - assemble_tiled_vector_field, tile_scalar_field, ) from PyPIC3D.utilities.grids import build_collocated_grid, build_tiled_yee_grids, build_yee_grid @@ -24,10 +23,6 @@ jax.config.update("jax_enable_x64", True) -ROUND_OFF_RTOL = 1.0e-11 -ROUND_OFF_ATOL = 1.0e-11 - - 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) @@ -242,149 +237,20 @@ def _long_two_stream_state(self): world["grids"] = {"center": center_grid, "vertex": vertex_grid} constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (4, 1, 1) - reference_tile_shape = self._one_tile_shape(world) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, tile_shape) + tile_shape = self._one_tile_shape(world) + world = self._world_with_tiled_grids(world, tile_shape) Nt = 1500 E, B, J, rho, phi, external_fields = self._empty_fields(world) fields = (E, B, J, rho, phi, external_fields) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - self._long_two_stream_species(reference_world), - fields, - reference_world, - reference_tile_shape, - ) tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._long_two_stream_species(tiled_world), + self._long_two_stream_species(world), fields, - tiled_world, - tile_shape, - min_slots=12, - ) - - self._assert_tiled_state_matches_reference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - step=0, - ) - - return ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - tiled_particles, - species_config, - tiled_fields, - reference_world, - tiled_world, - constants, + world, tile_shape, - Nt, ) - def _tiled_species_rows(self, tiled_particles, species_index): - active = tiled_particles.active[:, :, :, species_index, :] - x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - - def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): - for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): - diff = jnp.max(jnp.abs(reference_component - tiled_component)) - self.assertTrue( - jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: {name}{component_name} max abs diff {diff}", - ) - - def _assert_tiled_state_matches_reference( - self, - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - world, - tile_shape, - step, - ): - if len(tiled_fields) > 7: - self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - - species_count = reference_particles.active.shape[3] - for species_index in range(species_count): - reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) - tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) - self.assertTrue( - jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} position mismatch", - ) - self.assertTrue( - jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} velocity mismatch", - ) - - reference_E, reference_B, reference_J, reference_rho, reference_phi, *_ = reference_fields - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - g = int(world["guard_cells"]) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=g), - assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=g), - step, - "E", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=g), - assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=g), - step, - "B", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=g), - assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g), - step, - "J", - ) - - reference_rho_from_tiles = assemble_tiled_scalar_field( - reference_rho, - world, - reference_tile_shape, - num_guard_cells=g, - ) - reference_phi_from_tiles = assemble_tiled_scalar_field( - reference_phi, - world, - reference_tile_shape, - num_guard_cells=g, - ) - rho_from_tiles = assemble_tiled_scalar_field(rho_tiles, world, tile_shape, num_guard_cells=g) - phi_from_tiles = assemble_tiled_scalar_field(phi_tiles, world, tile_shape, num_guard_cells=g) - self.assertTrue( - jnp.allclose( - rho_from_tiles[1:-1, 1:-1, 1:-1], - reference_rho_from_tiles[1:-1, 1:-1, 1:-1], - rtol=ROUND_OFF_RTOL, - atol=ROUND_OFF_ATOL, - ), - f"step {step}: rho max abs diff {jnp.max(jnp.abs(rho_from_tiles - reference_rho_from_tiles))}", - ) - tiled_phi_interior = phi_from_tiles[1:-1, 1:-1, 1:-1] - reference_phi_interior = reference_phi_from_tiles[1:-1, 1:-1, 1:-1] - tiled_phi_interior = tiled_phi_interior - jnp.mean(tiled_phi_interior) - reference_phi_interior = reference_phi_interior - jnp.mean(reference_phi_interior) - self.assertTrue( - jnp.allclose(tiled_phi_interior, reference_phi_interior, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: phi max abs diff {jnp.max(jnp.abs(tiled_phi_interior - reference_phi_interior))}", - ) + return tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt def test_scalar_tiles_round_trip_global_rho_and_phi(self): world = self._build_world() @@ -444,79 +310,46 @@ def test_tiled_rho_assembles_to_global_rho(self): ) ) - def test_tiled_electrostatic_step_matches_one_tile_step_after_assembly(self): + def test_tiled_electrostatic_step_uses_single_tile_state(self): world = self._build_world(shape_factor=1) constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} E, B, J, rho, phi, external_fields = self._empty_fields(world) fields = (E, B, J, rho, phi, external_fields) - reference_tile_shape = self._one_tile_shape(world) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, world["tile_shape"]) - - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - self._neutral_species(reference_world), - fields, - reference_world, - reference_tile_shape, - ) + tile_shape = self._one_tile_shape(world) + world = self._world_with_tiled_grids(world, tile_shape) tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._neutral_species(tiled_world), + self._neutral_species(world), fields, - tiled_world, - tiled_world["tile_shape"], - ) - - reference_particles, reference_fields = time_loop_electrostatic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", + world, + tile_shape, ) tiled_particles, tiled_fields = time_loop_electrostatic( tiled_particles, species_config, tiled_fields, - tiled_world, + world, constants, solver="electrostatic", relativistic=False, particle_pusher="boris", ) - self._assert_tiled_state_matches_reference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tiled_world["tile_shape"], - step=1, - ) + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(B_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(J_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(rho_tiles.shape[:3], (1, 1, 1)) + self.assertEqual(phi_tiles.shape[:3], (1, 1, 1)) + self.assertFalse(bool(tiled_fields[-1])) + self.assertTrue(jnp.all(jnp.isfinite(rho_tiles))) + self.assertTrue(jnp.all(jnp.isfinite(phi_tiles))) @unittest.skipUnless( os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", - "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step tiled electrostatic comparison.", + "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step single-tile electrostatic check.", ) - def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(self): - ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - tiled_particles, - species_config, - tiled_fields, - reference_world, - tiled_world, - constants, - tile_shape, - Nt, - ) = self._long_two_stream_state() + def test_long_two_stream_single_tile_electrostatic_remains_finite(self): + tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt = self._long_two_stream_state() tiled_loop = jax.jit( time_loop_electrostatic, @@ -524,37 +357,25 @@ def test_long_two_stream_tiled_electrostatic_matches_one_tile_step_every_step(se ) for step in range(1, Nt + 1): - reference_particles, reference_fields = tiled_loop( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", - ) tiled_particles, tiled_fields = tiled_loop( tiled_particles, species_config, tiled_fields, - tiled_world, + world, constants, solver="electrostatic", relativistic=False, particle_pusher="boris", ) - self._assert_tiled_state_matches_reference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - step, - ) + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + self.assertEqual(tuple(world["tile_shape"]), tuple(tile_shape)) + self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) + self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") + for component in E_tiles + B_tiles + J_tiles: + self.assertTrue(jnp.all(jnp.isfinite(component)), f"step {step}: nonfinite vector field") + self.assertTrue(jnp.all(jnp.isfinite(rho_tiles)), f"step {step}: nonfinite rho") + self.assertTrue(jnp.all(jnp.isfinite(phi_tiles)), f"step {step}: nonfinite phi") def test_initialize_simulation_accepts_tiled_electrostatic(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -608,23 +429,26 @@ def test_initialize_simulation_accepts_tiled_electrostatic(self): self.assertIs(loop, time_loop_electrostatic) self.assertIsInstance(particles, TiledParticles) - self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) - self.assertEqual(tuple(simulation_parameters["tile_shape"]), (2, 1, 1)) + 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)) self.assertIn("tiled_center_grid", world["grids"]) self.assertIn("tiled_vertex_grid", world["grids"]) self.assertEqual(int(world["guard_cells"]), 2) - self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (4, 1, 1, 6)) - self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (4, 1, 1, 5)) + self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (1, 1, 1, 12)) + self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (1, 1, 1, 5)) 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], (4, 1, 1)) - self.assertEqual(fields[3].shape[:3], (4, 1, 1)) - self.assertEqual(fields[4].shape[:3], (4, 1, 1)) + 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__": diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 6eae7cd6..07e80ae4 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -174,10 +174,11 @@ def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): self.assertNotIn("assemble_tiled_vector_field", source) self.assertNotIn("fields_for_output", source) - def test_electrostatic_global_poisson_bridge_stays_explicit(self): + def test_electrostatic_solver_uses_single_tile_arrays_directly(self): source = inspect.getsource(electrostatic_yee.calculate_tiled_electrostatic_fields) - self.assertIn("assemble_tiled_scalar_field", source) - self.assertIn("tile_scalar_field", source) + self.assertNotIn("assemble_tiled_scalar_field", source) + self.assertNotIn("tile_scalar_field", source) + self.assertIn("[0, 0, 0]", source) def test_flat_current_and_electrodynamic_references_are_removed(self): self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) From 6fc0a95ec161214f1efe35ab2f781f061cdafa53 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 04:22:07 -0500 Subject: [PATCH 46/83] refactoring yee methods --- PyPIC3D/diagnostics/output_adapters.py | 44 +++++- PyPIC3D/evolve.py | 2 +- .../{yee_tiled.py => first_order_yee.py} | 138 ++---------------- tests/direct_deposition_tiled_test.py | 63 +++++++- tests/openpmd_test.py | 61 +++++++- tests/particle_test.py | 61 +++++++- tests/pml_test.py | 64 +++++++- tests/tiled_electrostatic_test.py | 65 ++++++++- tests/tiled_esirkepov_test.py | 65 ++++++++- tests/tiled_fluid_quantities_test.py | 61 +++++++- tests/tiled_pusher_test.py | 60 +++++++- tests/tiled_refactor_contract_test.py | 89 ++++++++--- tests/tiled_rho_test.py | 61 +++++++- tests/tiled_yee_integration_test.py | 64 +++++++- tests/utils_test.py | 61 +++++++- tests/yee_tiled_test.py | 63 +++++++- 16 files changed, 844 insertions(+), 178 deletions(-) rename PyPIC3D/solvers/{yee_tiled.py => first_order_yee.py} (62%) diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index 3d15f313..7ac9f372 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -1,5 +1,6 @@ +import jax.numpy as jnp + from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, assemble_tiled_vector_field def _is_tiled_scalar(field): @@ -22,6 +23,47 @@ 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. diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 97b4c0e9..75f3f8ce 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -15,7 +15,7 @@ ) from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.solvers.electrostatic_yee import calculate_tiled_electrostatic_fields -from PyPIC3D.solvers.yee_tiled import update_B, update_E +from PyPIC3D.solvers.first_order_yee import update_B, update_E from PyPIC3D.utils import add_external_fields diff --git a/PyPIC3D/solvers/yee_tiled.py b/PyPIC3D/solvers/first_order_yee.py similarity index 62% rename from PyPIC3D/solvers/yee_tiled.py rename to PyPIC3D/solvers/first_order_yee.py index 3f19ab99..c56ae8fb 100644 --- a/PyPIC3D/solvers/yee_tiled.py +++ b/PyPIC3D/solvers/first_order_yee.py @@ -1,6 +1,3 @@ -import jax -import jax.numpy as jnp - from PyPIC3D.boundary_conditions.PML import ( apply_tiled_pml_to_b_curl, apply_tiled_pml_to_e_curl, @@ -9,121 +6,6 @@ from PyPIC3D.utilities.filters import digital_filter_vector -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 _active_slice(g): - return slice(g, -g) - - -def _forward_slice(g): - return slice(g + 1, None if g == 1 else -g + 1) - - -def _backward_slice(g): - return slice(g - 1, -g - 1) - - -def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): - """ - Split a ghost-celled field into compact tiles using the shared 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) - - 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 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 update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): """ Update compact tiled electric fields without assembling a global field. @@ -134,17 +16,20 @@ def update_E(E_tiles, B_tiles, J_tiles, world, constants, pml_state=None): Ex, Ey, Ez = E_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 = _active_slice(g) - forward = _forward_slice(g) + # 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 = _active_slice(g) + current = slice(g, -g) #_active_slice(g) dt = world["dt"] dx, dy, dz = world["dx"], world["dy"], world["dz"] @@ -215,8 +100,11 @@ def update_B(E_tiles, B_tiles, world, constants, pml_state=None): tile_nx, tile_ny, tile_nz = tile_shape g = int(world["guard_cells"]) g = int(g) - active = _active_slice(g) - backward = _backward_slice(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: diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 9ecc2c0a..f4f2426c 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -13,16 +13,71 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh import refresh_tiled_particle_tiles -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_vector_field, - tile_scalar_field, -) +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) + + +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) diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index 9c245f15..bff59e59 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -1,15 +1,74 @@ import unittest from unittest.mock import patch +import jax import jax.numpy as jnp +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics import openPMD from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array -from PyPIC3D.solvers.yee_tiled import tile_scalar_field from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +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) diff --git a/tests/particle_test.py b/tests/particle_test.py index d70205bf..7346411e 100644 --- a/tests/particle_test.py +++ b/tests/particle_test.py @@ -5,6 +5,7 @@ import os +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.particles.particle_initialization import ( initial_particles, load_particles_from_toml, @@ -13,8 +14,8 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_field from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -22,6 +23,64 @@ 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 _compute_rho_one_tile(particles, rho, world, constants): world = dict(world) world["guard_cells"] = int(world.get("guard_cells", 2)) diff --git a/tests/pml_test.py b/tests/pml_test.py index 461a0768..6d66e7f9 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -12,11 +12,11 @@ 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.solvers.yee_tiled import ( - assemble_tiled_vector_field, - tile_scalar_field, +from PyPIC3D.solvers.first_order_yee import ( update_B, update_E, ) @@ -26,6 +26,64 @@ 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) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 443e17c5..caed8fdb 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -14,15 +14,72 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_scalar_field, - tile_scalar_field, -) +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field from PyPIC3D.utilities.grids import build_collocated_grid, 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) diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 7934a983..4460d2d3 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -17,7 +17,8 @@ get_2D_esirkepov_weights, ) from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.diagnostics.output_adapters import fields_for_output +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.tiled_particle_refresh import ( refresh_tiled_particle_tiles, @@ -26,9 +27,7 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_vector_field, - tile_scalar_field, +from PyPIC3D.solvers.first_order_yee import ( update_E, ) from PyPIC3D.utilities.grids import build_tiled_yee_grids @@ -36,6 +35,64 @@ 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) diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 44dd50b7..0380d042 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -5,6 +5,7 @@ import jax.numpy as jnp from PyPIC3D import __main__ as pypic_main +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, @@ -16,15 +17,73 @@ 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 PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_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) + + +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 TestTiledFluidQuantities(unittest.TestCase): def _build_world(self, shape_factor=2): x_wind, y_wind, z_wind = 4.0, 3.0, 2.0 diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 3e58247b..97cb0297 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -3,17 +3,75 @@ import jax import jax.numpy as jnp +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.pusher.tiled_pusher import tiled_particle_push -from PyPIC3D.solvers.yee_tiled import tile_scalar_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) + + +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) diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 07e80ae4..b6e85dc1 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -10,12 +10,43 @@ from PyPIC3D.deposition import rho from PyPIC3D.boundary_conditions import boundaryconditions from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics import output_adapters from PyPIC3D.initialization import build_tiled_array, initialize_fields from PyPIC3D.solvers import electrostatic_yee -from PyPIC3D.solvers import yee_tiled +from PyPIC3D.solvers import first_order_yee from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid +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) + ntx = int(world["Nx"]) // tile_nx + nty = int(world["Ny"]) // tile_ny + ntz = int(world["Nz"]) // tile_nz + + 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) + + class TestTiledRefactorContracts(unittest.TestCase): def _world(self): world = { @@ -56,7 +87,7 @@ def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): world = self._world() tile_shape = (2, 3, 2) E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) - E_tiles = tuple(yee_tiled.tile_scalar_field(component, world, tile_shape) for component in E) + E_tiles = tuple(tile_scalar_field(component, world, tile_shape) for component in E) refreshed = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) self.assertEqual(len(refreshed), 3) @@ -72,31 +103,41 @@ def test_evolve_loops_read_runtime_controls_from_world(self): self.assertNotIn(name, electrostatic_signature.parameters) def test_tiled_yee_updates_read_tile_metadata_from_world(self): - update_E_signature = inspect.signature(yee_tiled.update_E) - update_B_signature = inspect.signature(yee_tiled.update_B) + update_E_signature = inspect.signature(first_order_yee.update_E) + update_B_signature = inspect.signature(first_order_yee.update_B) for name in ("curl_func", "tile_shape", "g"): self.assertNotIn(name, update_E_signature.parameters) self.assertNotIn(name, update_B_signature.parameters) - self.assertFalse(hasattr(yee_tiled, "update_tiled_E")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_B")) - self.assertFalse(hasattr(yee_tiled, "tiled_grid_axes_from_world")) - self.assertFalse(hasattr(yee_tiled, "tile_grid_axes")) - self.assertFalse(hasattr(yee_tiled, "tile_vector_field")) - self.assertFalse(hasattr(yee_tiled, "stack_tiled_vector_field")) - self.assertFalse(hasattr(yee_tiled, "unstack_tiled_vector_field")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells_for_pml")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells_for_pml")) - self.assertFalse(hasattr(yee_tiled, "fold_tiled_ghost_cells")) - self.assertFalse(hasattr(yee_tiled, "fold_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(yee_tiled, "apply_tiled_conducting_bc")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(yee_tiled, "update_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(yee_tiled, "fold_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(yee_tiled, "fold_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(yee_tiled, "empty_tiled_scalar_field")) - self.assertFalse(hasattr(yee_tiled, "empty_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_E")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_B")) + self.assertFalse(hasattr(first_order_yee, "tiled_grid_axes_from_world")) + self.assertFalse(hasattr(first_order_yee, "tile_grid_axes")) + self.assertFalse(hasattr(first_order_yee, "tile_vector_field")) + self.assertFalse(hasattr(first_order_yee, "stack_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "unstack_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_for_pml")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_for_pml")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "apply_tiled_conducting_bc")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "empty_tiled_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "empty_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "tile_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "assemble_tiled_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "assemble_tiled_vector_field")) + + def test_global_assembly_lives_at_output_boundary(self): + self.assertTrue(hasattr(output_adapters, "assemble_tiled_scalar_field")) + self.assertTrue(hasattr(output_adapters, "assemble_tiled_vector_field")) + + source = inspect.getsource(output_adapters) + self.assertNotIn("from PyPIC3D.solvers.yee_tiled import", source) def test_ghost_cell_logic_lives_in_boundary_conditions_ghost_cells(self): for name in ( diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index 331433f8..f10c74d1 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -3,17 +3,76 @@ 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 PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -from PyPIC3D.solvers.yee_tiled import assemble_tiled_scalar_field, tile_scalar_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) + + +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 diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 86c60d94..405f3fa5 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -12,9 +12,9 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_vector_field, - tile_scalar_field, +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field +from PyPIC3D.solvers.first_order_yee import ( update_B, update_E, ) @@ -29,6 +29,64 @@ ROUND_OFF_ATOL = 1.0e-11 +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) diff --git a/tests/utils_test.py b/tests/utils_test.py index e74e71ad..7ce46ba6 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -5,11 +5,11 @@ import jax.numpy as jnp import numpy as np import toml +from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics -from PyPIC3D.solvers.yee_tiled import tile_scalar_field from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid from PyPIC3D.utils import ( print_stats, check_stability, @@ -20,6 +20,65 @@ 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, diff --git a/tests/yee_tiled_test.py b/tests/yee_tiled_test.py index cfe027be..5d17deb0 100644 --- a/tests/yee_tiled_test.py +++ b/tests/yee_tiled_test.py @@ -3,12 +3,11 @@ import jax import jax.numpy as jnp -from PyPIC3D.solvers.yee_tiled import ( - assemble_tiled_vector_field, - tile_scalar_field, +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 @@ -17,6 +16,64 @@ 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) From 2b3bb4c352a4c2db444c95f76c51d67ff3578ee5 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 05:36:06 -0500 Subject: [PATCH 47/83] migrating tiled pusher and refactoring --- PyPIC3D/evolve.py | 12 +- PyPIC3D/pusher/particle_push.py | 228 ++++++++++++++++---------- PyPIC3D/pusher/tiled_pusher.py | 176 -------------------- tests/tiled_pusher_test.py | 180 ++++++++++---------- tests/tiled_refactor_contract_test.py | 14 +- 5 files changed, 244 insertions(+), 366 deletions(-) delete mode 100644 PyPIC3D/pusher/tiled_pusher.py diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 75f3f8ce..4f9baffe 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -13,7 +13,7 @@ refresh_tiled_particle_tiles, update_tiled_particle_positions, ) -from PyPIC3D.pusher.tiled_pusher import tiled_particle_push +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 @@ -46,22 +46,19 @@ def time_loop_electrodynamic( # unpack the tiled field state tile_shape = tuple(int(width) for width in world["tile_shape"]) - g = int(world["guard_cells"]) particle_world = _particle_kernel_world(world) # get tiled layout information from the world contract push_E, push_B = add_external_fields(E, B, external_fields) # particles see evolved fields plus external-only fields - particles = tiled_particle_push( + particles = particle_push( particles, species_config, push_E, push_B, particle_world, constants, - tile_shape, - g, relativistic=relativistic, particle_pusher=particle_pusher, ) @@ -149,22 +146,19 @@ def time_loop_electrostatic( # unpack the tiled field state tile_shape = tuple(int(width) for width in world["tile_shape"]) - g = int(world["guard_cells"]) 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 = tiled_particle_push( + particles = particle_push( particles, species_config, push_E_tiles, push_B_tiles, particle_world, constants, - tile_shape, - g, relativistic=relativistic, particle_pusher=particle_pusher, ) diff --git a/PyPIC3D/pusher/particle_push.py b/PyPIC3D/pusher/particle_push.py index 7739facb..0a794a5b 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.tiled_particles import TiledParticles from PyPIC3D.pusher.boris import ( boris_single_particle, interpolate_field_to_particles, @@ -22,77 +21,86 @@ def validate_particle_pusher(particle_pusher): raise ValueError(f"particle_pusher must be one of {VALID_PARTICLE_PUSHERS}, got {particle_pusher!r}") -@partial(jit, static_argnames=("periodic", "relativistic", "particle_pusher")) -def particle_push(particles, E, B, grid, staggered_grid, world, constants, periodic=True, relativistic=True, particle_pusher="boris"): +def particle_push(particles, species_config, E_tiles, B_tiles, world, constants, 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. - world (dict): Simulation parameters, including ``dt`` and ``shape_factor``. - 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. + 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. """ - 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 + 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"] - # interpolation order is a global simulation property - - ################## 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 ) - else: + + 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 + ) + 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), @@ -105,23 +113,73 @@ def particle_push(particles, E, B, grid, staggered_grid, world, constants, perio 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/pusher/tiled_pusher.py b/PyPIC3D/pusher/tiled_pusher.py deleted file mode 100644 index f7e09c22..00000000 --- a/PyPIC3D/pusher/tiled_pusher.py +++ /dev/null @@ -1,176 +0,0 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import partial - -from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.pusher.boris import ( - boris_single_particle, - interpolate_field_to_particles, - relativistic_boris_single_particle, -) -from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle - - -@partial(jit, static_argnames=("tile_shape", "g", "relativistic", "particle_pusher")) -def tiled_particle_push(tiled_particles, species_config, E_tiles, B_tiles, world, constants, tile_shape, g, relativistic=True, particle_pusher="boris"): - """ - 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. - """ - - tile_nx, tile_ny, tile_nz = tile_shape - dt = world["dt"] - shape_factor = world["shape_factor"] - - 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 - g = int(g) - - ntx, nty, ntz = tiled_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 - ) - 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 - ) - 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 - ) - 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), - ) - relativistic_boris_vmap = jax.vmap( - relativistic_boris_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), - ) - higuera_cary_vmap = jax.vmap( - higuera_cary_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), - ) - - 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, - tiled_particles.x, - tiled_particles.u, - tiled_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=tiled_particles.x, - u=new_u, - active=tiled_particles.active, - ) diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index 97cb0297..e2c20146 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -7,7 +7,6 @@ from PyPIC3D.particles.species_class import particle_species from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push -from PyPIC3D.pusher.tiled_pusher import tiled_particle_push from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -104,6 +103,36 @@ def _with_tiled_grids(self, world, tile_shape, g=1): 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 = to_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( @@ -156,63 +185,50 @@ def _flatten_active_by_position(self, tiled_particles): order = jnp.lexsort((x[:, 2], x[:, 1], x[:, 0])) return x[order], u[order] - def test_tiled_particle_push_matches_flat_boris(self): + 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) - 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) reference = self._species(world) - reference = particle_push( + 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, - world["grids"]["center"], - world["grids"]["vertex"], - world, constants, relativistic=False, particle_pusher="boris", ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) - pushed = tiled_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), + pushed = self._push_tiled_species( + species, world, - constants, tile_shape, - 1, + E, + B, + constants, relativistic=False, + particle_pusher="boris", ) _, tiled_u = self._flatten_active_by_position(pushed) - flat_u = jnp.stack(reference.get_velocity(), axis=1) - flat_x = jnp.stack(reference.get_forward_position(), axis=1) - order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + _, reference_u = self._flatten_active_by_position(reference) - self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) - def test_tiled_particle_push_matches_flat_higuera_cary(self): + 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) - 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=0.25) B = self._deterministic_vector_field(world, scale=0.05) @@ -241,38 +257,34 @@ def make_species(): species = make_species() reference = make_species() - reference = particle_push( + 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, - world["grids"]["center"], - world["grids"]["vertex"], - world, constants, particle_pusher="higuera_cary", ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) - pushed = tiled_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), + pushed = self._push_tiled_species( + species, world, - constants, tile_shape, - 1, + E, + B, + constants, particle_pusher="higuera_cary", ) _, tiled_u = self._flatten_active_by_position(pushed) - flat_u = jnp.stack(reference.get_velocity(), axis=1) - flat_x = jnp.stack(reference.get_forward_position(), axis=1) - order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + _, reference_u = self._flatten_active_by_position(reference) - self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) - def test_tiled_particle_push_respects_active_and_update_flags(self): + def test_particle_push_respects_active_and_update_flags(self): world = self._build_world() constants = {"C": 10.0} tile_shape = (2, 3, 2) @@ -293,15 +305,13 @@ def test_tiled_particle_push_respects_active_and_update_flags(self): ) tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) - pushed = tiled_particle_push( + 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, - tile_shape, - 1, relativistic=False, ) @@ -309,16 +319,11 @@ def test_tiled_particle_push_respects_active_and_update_flags(self): 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_tiled_particle_push_matches_relativistic_flat_boris_on_reduced_axes(self): + 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) - 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=0.1) B = self._deterministic_vector_field(world, scale=0.02) def make_species(): @@ -346,49 +351,41 @@ def make_species(): species = make_species() reference = make_species() - reference = particle_push( + 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, - world["grids"]["center"], - world["grids"]["vertex"], - world, constants, relativistic=True, particle_pusher="boris", ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) - pushed = tiled_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), + pushed = self._push_tiled_species( + species, world, - constants, tile_shape, - 1, + E, + B, + constants, relativistic=True, + particle_pusher="boris", ) _, tiled_u = self._flatten_active_by_position(pushed) - flat_u = jnp.stack(reference.get_velocity(), axis=1) - flat_x = jnp.stack(reference.get_forward_position(), axis=1) - order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + _, reference_u = self._flatten_active_by_position(reference) - self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) - def test_tiled_particle_push_matches_flat_boris_on_two_guard_reduced_axes(self): + 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) - 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=0.1) B = self._deterministic_vector_field(world, scale=0.02) @@ -417,37 +414,34 @@ def make_species(): species = make_species() reference = make_species() - reference = particle_push( + 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, - world["grids"]["center"], - world["grids"]["vertex"], - world, constants, relativistic=True, particle_pusher="boris", ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) - pushed = tiled_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), + pushed = self._push_tiled_species( + species, world, - constants, tile_shape, - g, + E, + B, + constants, relativistic=True, + particle_pusher="boris", ) _, tiled_u = self._flatten_active_by_position(pushed) - flat_u = jnp.stack(reference.get_velocity(), axis=1) - flat_x = jnp.stack(reference.get_forward_position(), axis=1) - order = jnp.lexsort((flat_x[:, 2], flat_x[:, 1], flat_x[:, 0])) + _, reference_u = self._flatten_active_by_position(reference) - self.assertTrue(jnp.allclose(tiled_u, flat_u[order], rtol=1.0e-12, atol=1.0e-12)) + self.assertTrue(jnp.allclose(tiled_u, reference_u, rtol=1.0e-12, atol=1.0e-12)) if __name__ == "__main__": diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index b6e85dc1..38933e38 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -1,4 +1,5 @@ import inspect +import importlib import importlib.util import unittest @@ -200,14 +201,21 @@ def test_grid_builders_live_in_utilities(self): def test_tiled_benchmarking_module_is_removed(self): self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) - def test_tiled_pusher_reads_tiled_grids_from_world(self): - from PyPIC3D.pusher import tiled_pusher + def test_particle_push_owns_tiled_pusher_contract(self): + particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") - source = inspect.getsource(tiled_pusher.tiled_particle_push) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.pusher.tiled_pusher")) + signature = inspect.signature(particle_push.particle_push) + for name in ("grid", "staggered_grid", "tile_shape", "g"): + self.assertNotIn(name, signature.parameters) + + source = inspect.getsource(particle_push.particle_push) self.assertIn('world["grids"]["tiled_center_grid"]', source) self.assertIn('world["grids"]["tiled_vertex_grid"]', source) self.assertNotIn("tiled_grid_axes_from_world", source) self.assertNotIn("tile_grid_axes", source) + self.assertNotIn("get_forward_position", source) + self.assertNotIn("set_velocity", source) def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): source = inspect.getsource(evolve.time_loop_electrodynamic) From cb07e79d5ea0f4b3a5efcf60f050045b66bda015 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 16:02:26 -0500 Subject: [PATCH 48/83] refactoring --- PyPIC3D/initialization.py | 3 +-- PyPIC3D/pusher/__init__.py | 2 -- PyPIC3D/pusher/particle_push.py | 11 ----------- tests/boris_test.py | 5 ----- 4 files changed, 1 insertion(+), 20 deletions(-) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 52196961..0046eb20 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -49,7 +49,6 @@ time_loop_electrostatic, ) -from PyPIC3D.pusher.particle_push import validate_particle_pusher from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC from PyPIC3D.boundary_conditions.PML import ( @@ -285,7 +284,7 @@ def initialize_simulation(toml_file): electrostatic = solver == "electrostatic" relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['particle_pusher'] - validate_particle_pusher(particle_pusher) + if electrostatic: simulation_parameters["particle_tile_nx"] = int(Nx) simulation_parameters["particle_tile_ny"] = int(Ny) 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/particle_push.py b/PyPIC3D/pusher/particle_push.py index 0a794a5b..0bdb09d6 100644 --- a/PyPIC3D/pusher/particle_push.py +++ b/PyPIC3D/pusher/particle_push.py @@ -10,17 +10,6 @@ from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle -VALID_PARTICLE_PUSHERS = ("boris", "higuera_cary") - - -def validate_particle_pusher(particle_pusher): - """ - Validate the particle pusher name before entering jitted simulation code. - """ - if particle_pusher not in VALID_PARTICLE_PUSHERS: - raise ValueError(f"particle_pusher must be one of {VALID_PARTICLE_PUSHERS}, got {particle_pusher!r}") - - def particle_push(particles, species_config, E_tiles, B_tiles, world, constants, relativistic=True, particle_pusher="boris"): """ Push tile-major particles with the selected pusher using compact tiled Yee fields. diff --git a/tests/boris_test.py b/tests/boris_test.py index e94b1628..b5e45de4 100644 --- a/tests/boris_test.py +++ b/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 @@ -244,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): From 7304f7488e6ac507128d883679da62d0bd644570 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 16:22:10 -0500 Subject: [PATCH 49/83] refactoring and removing deprecated code --- PyPIC3D/__init__.py | 3 - PyPIC3D/boris.py | 6 - PyPIC3D/errors.py | 93 ------ PyPIC3D/solvers/electrostatic_yee.py | 21 +- PyPIC3D/solvers/fdtd.py | 210 ------------ .../solver_convergence_test.py | 268 --------------- demos/ot_vortex/orszag_tang.toml | 2 +- demos/reconnection_2d/harris_current.toml | 2 +- demos/two_stream/two_stream.toml | 2 +- demos/weibel/weibel.toml | 2 +- docs/demos.rst | 1 - tests/__init__.py | 2 - tests/electrostatic_yee_test.py | 25 +- tests/errors_test.py | 70 ---- tests/fdtd_test.py | 314 ------------------ tests/tiled_refactor_contract_test.py | 6 + 16 files changed, 41 insertions(+), 986 deletions(-) delete mode 100644 PyPIC3D/boris.py delete mode 100644 PyPIC3D/errors.py delete mode 100644 PyPIC3D/solvers/fdtd.py delete mode 100644 demos/convergence_testing/solver_convergence_test.py delete mode 100644 tests/errors_test.py delete mode 100644 tests/fdtd_test.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index a041a169..b00bde95 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -6,14 +6,11 @@ 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 fdtd from .solvers import electrostatic_yee from .particles import particle_initialization 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/errors.py b/PyPIC3D/errors.py deleted file mode 100644 index 6f9d5689..00000000 --- a/PyPIC3D/errors.py +++ /dev/null @@ -1,93 +0,0 @@ -from jax import jit -import jax.numpy as jnp -from functools import partial -# import external libraries - -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 == 'autodiff': - return 0 - x = centered_finite_difference_laplacian(phi, dx, dy, dz, bc) - 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 == 'autodiff': - return 0 - divB = centered_finite_difference_divergence(Bx, By, Bz, dx, dy, dz, bc) - 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 == 'autodiff': - return 0 - divE = centered_finite_difference_divergence(Ex, Ey, Ez, dx, dy, dz, bc) - - divergence_error = jnp.mean(jnp.abs(divE - rho / eps)) - return divergence_error diff --git a/PyPIC3D/solvers/electrostatic_yee.py b/PyPIC3D/solvers/electrostatic_yee.py index b03bf882..e505d499 100644 --- a/PyPIC3D/solvers/electrostatic_yee.py +++ b/PyPIC3D/solvers/electrostatic_yee.py @@ -2,7 +2,6 @@ from jax import lax from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.solvers.fdtd import centered_finite_difference_gradient from PyPIC3D.utilities.filters import digital_filter from PyPIC3D.boundary_conditions import ghost_cells @@ -31,6 +30,22 @@ def _refresh_single_tile_scalar(field, world, g, tile_shape, apply_conducting=Fa return field_tiles[0, 0, 0] +def _centered_finite_difference_gradient(field, dx, dy, dz): + """ + Compute the centered finite-difference gradient on a periodic scalar field. + + 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 + + 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. @@ -156,7 +171,9 @@ def calculate_electrostatic_fields(world, particles, constants, rho, phi, solver phi = _refresh_single_tile_scalar(phi, world, g, tile_shape, apply_conducting=True) # update ghost cells after filtering - Ex, Ey, Ez = centered_finite_difference_gradient(-1 * phi[active, active, active], dx, dy, dz, bc) + 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 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/demos/convergence_testing/solver_convergence_test.py b/demos/convergence_testing/solver_convergence_test.py deleted file mode 100644 index a3d55203..00000000 --- a/demos/convergence_testing/solver_convergence_test.py +++ /dev/null @@ -1,268 +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.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 - - -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("###################################################################") - ############################################################################ - - - ############################################################################ diff --git a/demos/ot_vortex/orszag_tang.toml b/demos/ot_vortex/orszag_tang.toml index e919d48f..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 diff --git a/demos/reconnection_2d/harris_current.toml b/demos/reconnection_2d/harris_current.toml index d5d21e63..8c7e76a8 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 diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 5cac4a46..3f013a43 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -2,7 +2,7 @@ [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 diff --git a/demos/weibel/weibel.toml b/demos/weibel/weibel.toml index dd3e3a3a..5e185e8e 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 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/tests/__init__.py b/tests/__init__.py index 881a2c49..65f04b5b 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -6,9 +6,7 @@ from . import boris_test from . import boundaryconditions_test -from . import errors_test from . import electrostatic_yee_test -from . import fdtd_test from . import initialization_test from . import particle_test from . import utils_test diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index 8b981723..38b12df0 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -7,8 +7,8 @@ from PyPIC3D.solvers.electrostatic_yee import ( solve_poisson_with_conjugate_gradient, calculate_electrostatic_fields, + _centered_finite_difference_gradient, ) -from PyPIC3D.solvers.fdtd import 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 @@ -130,35 +130,34 @@ def test_solve_poisson_with_conjugate_gradient_single_mode(self): 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_fdtd, phi_fdtd, rho_fdtd = calculate_electrostatic_fields( + 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, - "fdtd", + "electrostatic", "periodic", ) - expected_phi_fdtd = solve_poisson_with_conjugate_gradient( - rho_fdtd, + expected_phi = solve_poisson_with_conjugate_gradient( + rho, 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[self.active, self.active, self.active], + 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, - "periodic", ) - self.assertTrue(jnp.allclose(phi_fdtd[self.active, self.active, self.active], expected_phi_fdtd[self.active, self.active, self.active], atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[0][self.active, self.active, self.active], expected_Ex_fdtd, atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[1][self.active, self.active, self.active], expected_Ey_fdtd, atol=1e-6, rtol=1e-6)) - self.assertTrue(jnp.allclose(E_fdtd[2][self.active, self.active, self.active], expected_Ez_fdtd, atol=1e-6, rtol=1e-6)) + 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)) if __name__ == "__main__": diff --git a/tests/errors_test.py b/tests/errors_test.py deleted file mode 100644 index db20162f..00000000 --- a/tests/errors_test.py +++ /dev/null @@ -1,70 +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_fdtd = compute_pe(phi, rho, self.constants, self.world, solver='fdtd') - - 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_fdtd = compute_magnetic_divergence_error(Bx, By, Bz, self.world, solver='fdtd') - - 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_fdtd = compute_electric_divergence_error(Ex, Ey, Ez, rho, self.constants, self.world, solver='fdtd') - - self.assertIsInstance(result_fdtd, jnp.ndarray) -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/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 38933e38..02a48129 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -201,6 +201,12 @@ def test_grid_builders_live_in_utilities(self): def test_tiled_benchmarking_module_is_removed(self): self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) + def test_deprecated_fdtd_errors_and_root_boris_modules_are_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.solvers.fdtd")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.errors")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.boris")) + self.assertIsNotNone(importlib.util.find_spec("PyPIC3D.pusher.boris")) + def test_particle_push_owns_tiled_pusher_contract(self): particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") From 753decb36956340b955e9e8abd905c41e9934d0a Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 18:07:40 -0500 Subject: [PATCH 50/83] completely removing old particle species class --- PyPIC3D/__init__.py | 1 - PyPIC3D/diagnostics/openPMD.py | 24 +- PyPIC3D/diagnostics/output_adapters.py | 82 +- PyPIC3D/diagnostics/plotting.py | 34 +- PyPIC3D/diagnostics/vtk.py | 9 +- PyPIC3D/initialization.py | 60 +- PyPIC3D/particles/particle_initialization.py | 181 +++- PyPIC3D/particles/species_class.py | 354 ------- .../particles/tiled_particle_diagnostics.py | 118 --- .../tiled_particle_initialization.py | 138 --- PyPIC3D/utils.py | 252 ++--- tests/__init__.py | 1 - tests/direct_deposition_tiled_test.py | 4 +- tests/electrostatic_yee_test.py | 2 +- tests/evolve_test.py | 4 +- tests/openpmd_test.py | 4 +- tests/particle_test.py | 862 ------------------ tests/pml_test.py | 36 +- tests/tiled_electrostatic_test.py | 4 +- tests/tiled_esirkepov_test.py | 4 +- tests/tiled_fluid_quantities_test.py | 4 +- tests/tiled_particle_diagnostics_test.py | 38 +- tests/tiled_particle_fixtures.py | 280 ++++++ tests/tiled_particle_initialization_test.py | 18 +- tests/tiled_particle_refresh_test.py | 4 +- tests/tiled_pusher_test.py | 4 +- tests/tiled_refactor_contract_test.py | 10 + tests/tiled_rho_test.py | 4 +- tests/tiled_yee_integration_test.py | 4 +- tests/utils_test.py | 52 +- 30 files changed, 760 insertions(+), 1832 deletions(-) delete mode 100644 PyPIC3D/particles/species_class.py delete mode 100644 PyPIC3D/particles/tiled_particle_diagnostics.py delete mode 100644 PyPIC3D/particles/tiled_particle_initialization.py delete mode 100644 tests/particle_test.py create mode 100644 tests/tiled_particle_fixtures.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index b00bde95..d597a907 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -14,7 +14,6 @@ from .solvers import electrostatic_yee from .particles import particle_initialization -from .particles import species_class from .deposition import shapes from .deposition import Esirkepov diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index 1548ec57..950485b1 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -137,11 +137,11 @@ def write_openpmd_particles_to_iteration(iteration, particles, constants, specie 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) @@ -155,9 +155,9 @@ def write_openpmd_particles_to_iteration(iteration, particles, constants, specie 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 @@ -316,7 +316,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) @@ -331,8 +331,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 @@ -345,9 +345,9 @@ def make_array_writable(arr): gamma = make_array_writable(gamma) num_particles = x.shape[0] - particle_mass = species.get_mass() - particle_charge = species.get_charge() - particle_weight = species.get_weight() + 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) diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index 7ac9f372..dc395db5 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -1,6 +1,18 @@ import jax.numpy as jnp +from typing import NamedTuple -from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species +from PyPIC3D.particles.tiled_particles 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): @@ -124,14 +136,70 @@ def fields_for_output(fields, world): 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. """ - return flatten_tiled_particles_by_species( - particles, - species_config=species_config, - species_names=species_names, - world=world, - ) + 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 6a37cd43..da74a86c 100644 --- a/PyPIC3D/diagnostics/plotting.py +++ b/PyPIC3D/diagnostics/plotting.py @@ -31,11 +31,11 @@ def plot_positions(particles, t, x_wind, y_wind, z_wind, path, species_config=No 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=species.get_name() + name=species.name )) fig.update_layout( @@ -80,9 +80,9 @@ def write_particles_phase_space(particles, t, path, species_config=None, species # 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) @@ -118,8 +118,8 @@ def particles_phase_space(particles, world, t, name, path, species_config=None, # 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 @@ -133,8 +133,8 @@ def particles_phase_space(particles, world, t, name, path, species_config=None, 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") @@ -146,8 +146,8 @@ def particles_phase_space(particles, world, t, name, path, species_config=None, 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") @@ -158,16 +158,14 @@ def particles_phase_space(particles, world, t, name, path, species_config=None, 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'. @@ -183,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'] diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py index b4ec4385..ae0928dc 100644 --- a/PyPIC3D/diagnostics/vtk.py +++ b/PyPIC3D/diagnostics/vtk.py @@ -65,13 +65,12 @@ def plot_vtk_particles(particles, t, path, species_config=None, species_names=No # Tiled particle arrays are fixed-capacity by tile; VTK particle output only # receives the active diagnostic 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()) + name = species.name.replace(" ", "") + x, y, z = map(np.asarray, (species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2])) + vx, vy, vz = map(np.asarray, (species.u[:, 0], species.u[:, 1], species.u[:, 2])) # Get the position and velocity of the particles - q = np.asarray( species.get_charge() * np.ones_like(vx) ) + q = np.asarray(species.charge * species.weight * np.ones_like(vx)) # Get the charge of the particles pointsToVTK(f"{path}/data/particles/{name}.{t:09}", x, y, z, \ diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 0046eb20..7b562155 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -29,15 +29,12 @@ 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.tiled_particle_initialization import ( - to_tiled_particles -) - from PyPIC3D.boundary_conditions.ghost_cells import ( update_tiled_vector_ghost_cells, update_tiled_vector_ghost_cells_for_pml, @@ -398,23 +395,34 @@ def initialize_simulation(toml_file): # 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 - simulation_parameters["particle_species_names"] = tuple(species.get_name() for species in particles) - # keep species names available after tiled storage drops per-species objects - - 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 = {} @@ -424,7 +432,14 @@ 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 E, B, J, phi, rho = initialize_fields(world) @@ -463,7 +478,14 @@ def initialize_simulation(toml_file): ######################### 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") @@ -492,10 +514,6 @@ def initialize_simulation(toml_file): print(f"Using J from rhov current calculation method with filter: {simulation_parameters['filter_j']}") - species_config = None - particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - # convert the particles to tiled storage and get the species configuration - external_E, external_B = external_fields external_fields = (external_E, external_B) print(f"Using tiled Yee storage with tile shape: {tile_shape}") diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index e6fc58fe..f4b224d5 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.tiled_particles import SpeciesConfig, TiledParticles def grab_particle_keys(config): @@ -121,6 +122,54 @@ 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 _scaled_plasma_frequency(N_particles, charge, mass, weight, world, constants): + x_wind = world["x_wind"] + y_wind = world["y_wind"] + z_wind = world["z_wind"] + eps = constants["eps"] + + sqrt_dv = jnp.sqrt(x_wind * y_wind * z_wind) + sqrt_ne = jnp.sqrt(N_particles * weight) / sqrt_dv + sqrt_eps = jnp.sqrt(eps) + sqrt_mass = jnp.sqrt(mass) + return sqrt_ne * jnp.abs(charge) / (sqrt_eps * sqrt_mass) + + +def _scaled_debye_length(N_particles, charge, temperature, weight, world, constants): + eps = constants["eps"] + kb = constants["kb"] + x_wind = world["x_wind"] + y_wind = world["y_wind"] + z_wind = world["z_wind"] + + density = weight * N_particles / (x_wind * y_wind * z_wind) + return jnp.sqrt(eps * kb * temperature / (density * charge**2)) + + def load_particles_from_toml(config, simulation_parameters, world, constants): x_wind = world["x_wind"] y_wind = world["y_wind"] @@ -136,8 +185,9 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): eps = constants["eps"] i = 0 - particles = [] particle_keys = grab_particle_keys(config) + species_arrays = [] + species_metadata = [] weight = compute_macroparticle_weight( config, particle_keys, simulation_parameters, world, constants @@ -260,42 +310,32 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): update_y = read_value("update_y", toml_key, config, True) update_z = read_value("update_z", toml_key, config, True) - 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, - dt=dt, - ) - particles.append(particle) + 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) + pf = _scaled_plasma_frequency(N_particles, charge, mass, weight, world, constants) + dl = _scaled_debye_length(N_particles, charge, T, weight, world, constants) - 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}") @@ -303,12 +343,69 @@ 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"]] + + 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 = np.zeros((ntx, nty, ntz, n_species), dtype=int) + particle_tile_indices = [] + + for species_index, (x_array, _u_array, active_mask) in enumerate(species_arrays): + tx, ty, tz = _particle_tile_indices(x_array[:, 0], x_array[:, 1], x_array[:, 2], world, tile_nx, tile_ny, tile_nz) + particle_tile_indices.append((tx, ty, tz)) + + for p in range(int(active_mask.shape[0])): + if bool(active_mask[p]): + tile_counts[tx[p], ty[p], tz[p], species_index] += 1 + + max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 + capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) + max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) + + 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) + + next_slot = np.zeros_like(tile_counts) + + for species_index, (x_array, u_array, active_mask) in enumerate(species_arrays): + tx, ty, tz = particle_tile_indices[species_index] + + for p in range(int(active_mask.shape[0])): + if not bool(active_mask[p]): + continue + tile_index = (tx[p], ty[p], tz[p], species_index) + slot = next_slot[tile_index] + next_slot[tile_index] += 1 + index = tile_index + (slot,) + + x_tiles = x_tiles.at[index].set(x_array[p]) + u_tiles = u_tiles.at[index].set(u_array[p]) + active_tiles = active_tiles.at[index].set(True) + + 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=x_tiles, + u=u_tiles, + active=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/species_class.py b/PyPIC3D/particles/species_class.py deleted file mode 100644 index 484c6f87..00000000 --- a/PyPIC3D/particles/species_class.py +++ /dev/null @@ -1,354 +0,0 @@ -import jax -from jax import jit -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) - - -@jit -def apply_axis_boundary_condition(x, v, wind, half_wind, bc): - """Apply boundary conditions to particle positions/velocities along one axis.""" - - periodic = bc == 0 - reflecting = bc == 1 - - periodic_x = x + wind * (x < -half_wind) - wind * (x > half_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_v = jnp.where((x >= half_wind) | (x <= -half_wind), -v, v) - - x_out = jnp.where(periodic, periodic_x, jnp.where(reflecting, reflected_x, x)) - v_out = jnp.where(reflecting, reflected_v, v) - - return x_out, v_out - - -@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, - 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.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_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, world=None): - x1, x2, x3 = self.x1, self.x2, self.x3 - v1, v2, v3 = self.v1, self.v2, self.v3 - if world is None or "particle_boundary_conditions" not in world: - particle_bc = {"x": 0, "y": 0, "z": 0} - else: - particle_bc = world["particle_boundary_conditions"] - - x1, v1 = apply_axis_boundary_condition( - x1, v1, self.x_wind, self.half_x_wind, particle_bc["x"] - ) - x2, v2 = apply_axis_boundary_condition( - x2, v2, self.y_wind, self.half_y_wind, particle_bc["y"] - ) - x3, v3 = apply_axis_boundary_condition( - x3, v3, self.z_wind, self.half_z_wind, particle_bc["z"] - ) - - active_mask = self.active_mask - x_inside = (x1 <= self.half_x_wind) & (x1 >= -self.half_x_wind) - y_inside = (x2 <= self.half_y_wind) & (x2 >= -self.half_y_wind) - z_inside = (x3 <= self.half_z_wind) & (x3 >= -self.half_z_wind) - active_mask = jnp.where(particle_bc["x"] == 2, active_mask & x_inside, active_mask) - active_mask = jnp.where(particle_bc["y"] == 2, active_mask & y_inside, active_mask) - active_mask = jnp.where(particle_bc["z"] == 2, active_mask & z_inside, active_mask) - - 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.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, - 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, - dt=dt, - active_mask=active_mask, - ) - - return obj diff --git a/PyPIC3D/particles/tiled_particle_diagnostics.py b/PyPIC3D/particles/tiled_particle_diagnostics.py deleted file mode 100644 index 13c5ef9a..00000000 --- a/PyPIC3D/particles/tiled_particle_diagnostics.py +++ /dev/null @@ -1,118 +0,0 @@ -from typing import NamedTuple - -import jax.numpy as jnp - -from PyPIC3D.particles.tiled_particles import TiledParticles - - -class DiagnosticParticleSpecies(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 get_name(self): - return self.name - - def get_number_of_particles(self): - return self.x.shape[0] - - def get_forward_position(self): - return self.x[:, 0], self.x[:, 1], self.x[:, 2] - - def get_position(self): - return self.x_diagnostic[:, 0], self.x_diagnostic[:, 1], self.x_diagnostic[:, 2] - - def get_velocity(self): - return self.u[:, 0], self.u[:, 1], self.u[:, 2] - - def get_active_mask(self): - return jnp.ones(self.x.shape[0], dtype=bool) - - def get_charge(self): - return self.charge * self.weight - - def get_mass(self): - return self.mass * self.weight - - def get_weight(self): - return self.weight - - -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 flatten_tiled_particles_by_species(tiled_particles, species_config=None, species_names=None, world=None): - """ - Flatten active tile-major particles into diagnostic species arrays. - - This is an output adapter only: it removes inactive slots introduced by the - fixed-capacity tiled storage and preserves the stored particle position, - velocity, charge, mass, weight, and species index. - """ - - if not isinstance(tiled_particles, TiledParticles): - return tiled_particles - - n_species = tiled_particles.active.shape[3] - flattened_species = [] - - for species_index in range(n_species): - active = tiled_particles.active[:, :, :, species_index, :].reshape(-1) - - x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active] - u = tiled_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]) - - x_diagnostic = _diagnostic_position(x, u, world) - - if species_names is None: - name = f"species_{species_index}" - else: - name = species_names[species_index] - - flattened_species.append( - DiagnosticParticleSpecies( - name=name, - species_index=species_index, - x=x, - x_diagnostic=x_diagnostic, - u=u, - charge=charge, - mass=mass, - weight=weight, - ) - ) - - return flattened_species diff --git a/PyPIC3D/particles/tiled_particle_initialization.py b/PyPIC3D/particles/tiled_particle_initialization.py deleted file mode 100644 index ef542b45..00000000 --- a/PyPIC3D/particles/tiled_particle_initialization.py +++ /dev/null @@ -1,138 +0,0 @@ -import math - -import jax.numpy as jnp -import numpy as np - -from PyPIC3D.particles.particle_initialization import load_particles_from_toml -from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles - - -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 to_tiled_particles(particles, world, simulation_parameters): - """ - Pack an existing particle species list into tile-major arrays. - - This routine only changes storage: particles are grouped by configurable - spatial tiles with leading axes ``(ntx, nty, ntz, species, slot, ...)``. - """ - - tile_nx = _as_int(simulation_parameters.get("particle_tile_nx", 1)) - tile_ny = _as_int(simulation_parameters.get("particle_tile_ny", 1)) - tile_nz = _as_int(simulation_parameters.get("particle_tile_nz", 1)) - - 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(particles) - - tile_counts = np.zeros((ntx, nty, ntz, n_species), dtype=int) - particle_tile_indices = [] - - for species_index, species in enumerate(particles): - x1, x2, x3 = species.get_forward_position() - tx, ty, tz = _particle_tile_indices(x1, x2, x3, world, tile_nx, tile_ny, tile_nz) - particle_tile_indices.append((tx, ty, tz)) - - for p in range(_as_int(species.N_particles)): - tile_counts[tx[p], ty[p], tz[p], species_index] += 1 - - max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 - capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) - max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) - - x = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) - u = jnp.zeros_like(x) - active = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile), dtype=bool) - - charge = jnp.asarray([species.charge for species in particles]) - mass = jnp.asarray([species.mass for species in particles]) - weight = jnp.asarray([species.weight for species in particles]) - update_x = jnp.asarray( - [ - [ - species.update_pos and species.update_x, - species.update_pos and species.update_y, - species.update_pos and species.update_z, - ] - for species in particles - ], - dtype=bool, - ) - update_u = jnp.asarray( - [ - [ - species.update_v and species.update_vx, - species.update_v and species.update_vy, - species.update_v and species.update_vz, - ] - for species in particles - ], - dtype=bool, - ) - - next_slot = np.zeros_like(tile_counts) - - for species_index, species in enumerate(particles): - x1, x2, x3 = species.get_forward_position() - u1, u2, u3 = species.get_velocity() - tx, ty, tz = particle_tile_indices[species_index] - - for p in range(_as_int(species.N_particles)): - tile_index = (tx[p], ty[p], tz[p], species_index) - slot = next_slot[tile_index] - next_slot[tile_index] += 1 - index = tile_index + (slot,) - - x = x.at[index + (0,)].set(x1[p]) - x = x.at[index + (1,)].set(x2[p]) - x = x.at[index + (2,)].set(x3[p]) - u = u.at[index + (0,)].set(u1[p]) - u = u.at[index + (1,)].set(u2[p]) - u = u.at[index + (2,)].set(u3[p]) - - active = active.at[index].set(species.active_mask[p]) - - species_config = SpeciesConfig( - charge=charge, - mass=mass, - weight=weight, - update_x=update_x, - update_u=update_u, - ) - return TiledParticles(x=x, u=u, active=active), species_config - - -def load_tiled_particles_from_toml(config, simulation_parameters, world, constants): - """ - Initialize particles from TOML input and pack them into tile-major arrays. - - The particle values are produced by the existing particle initializer. This - wrapper only changes storage by calling ``to_tiled_particles`` on the - initialized species list. - """ - - particles = load_particles_from_toml(config, simulation_parameters, world, constants) - return to_tiled_particles(particles, world, simulation_parameters) diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index a180107a..1677c8d6 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -109,7 +109,7 @@ 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. @@ -164,38 +164,22 @@ 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 - if hasattr(particles, "active") and hasattr(particles, "u"): - 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)) - else: - 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 @@ -204,20 +188,17 @@ def nd_trapezoid(arr, dxs): def compute_total_momentum(particles, species_config=None): """ - Compute the scalar momentum diagnostic for either species-list or tiled particles. + Compute the scalar momentum diagnostic for tiled particles. """ - if hasattr(particles, "active") and hasattr(particles, "u"): - 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) - - return sum(particle_species.momentum() for particle_species in 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): @@ -316,30 +297,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): @@ -425,28 +394,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, @@ -684,7 +658,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"] @@ -699,47 +673,30 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para "particles": [] } - if hasattr(particles, "active") and hasattr(particles, "u"): - n_species = particles.active.shape[3] - species_names = simulation_parameters.get("particle_species_names") - tile_shape = jax.tree_util.tree_map( - lambda x: x.tolist() if isinstance(x, jnp.ndarray) else x, - simulation_parameters.get("tile_shape", ()), - ) + 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, :])) - config["particles"].append( - { - "name": name, - "storage": "tiled", - "active_particles": active_particles, - "tile_shape": tile_shape, - } - ) - else: - 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 - } - config["particles"].append(particle_dict) + 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(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"] = { "PyPIC3D_version": importlib.metadata.version('PyPIC3D'), @@ -819,80 +776,3 @@ def courant_condition(courant_number, dx, dy, dz, simulation_parameters, constan 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/tests/__init__.py b/tests/__init__.py index 65f04b5b..6e6e82c3 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -8,5 +8,4 @@ from . import boundaryconditions_test from . import electrostatic_yee_test from . import initialization_test -from . import particle_test from . import utils_test diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index f4f2426c..278e2bd9 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -10,8 +10,8 @@ digital_filter, digital_filter_vector, ) -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh 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 diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py index 38b12df0..3879d8ca 100644 --- a/tests/electrostatic_yee_test.py +++ b/tests/electrostatic_yee_test.py @@ -3,7 +3,7 @@ import jax import jax.numpy as jnp -from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.solvers.electrostatic_yee import ( solve_poisson_with_conjugate_gradient, calculate_electrostatic_fields, diff --git a/tests/evolve_test.py b/tests/evolve_test.py index 8d6b5ca9..c922221e 100644 --- a/tests/evolve_test.py +++ b/tests/evolve_test.py @@ -8,8 +8,8 @@ time_loop_electrostatic, ) from PyPIC3D.initialization import initialize_fields -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid jax.config.update("jax_enable_x64", True) diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index bff59e59..a96bc317 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -7,8 +7,8 @@ from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics import openPMD from PyPIC3D.diagnostics.openPMD import _ensure_openpmd_array -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles def _tile_axis_count(n_cells, cells_per_tile): diff --git a/tests/particle_test.py b/tests/particle_test.py deleted file mode 100644 index 7346411e..00000000 --- a/tests/particle_test.py +++ /dev/null @@ -1,862 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - - -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.particles.particle_initialization import ( - initial_particles, - load_particles_from_toml, - ) - -from PyPIC3D.particles.species_class import particle_species - -from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles -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 _compute_rho_one_tile(particles, rho, world, constants): - world = dict(world) - world["guard_cells"] = int(world.get("guard_cells", 2)) - world["tile_shape"] = (world["Nx"], world["Ny"], world["Nz"]) - grids = dict(world["grids"]) - tiled_vertex_grid, tiled_center_grid = build_tiled_yee_grids( - world, - world["tile_shape"], - int(world["guard_cells"]), - ) - grids["tiled_vertex_grid"] = tiled_vertex_grid - grids["tiled_center_grid"] = tiled_center_grid - world["grids"] = grids - - simulation_parameters = { - "particle_tile_nx": world["Nx"], - "particle_tile_ny": world["Ny"], - "particle_tile_nz": world["Nz"], - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - rho_tiles = tile_scalar_field(rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - rho_tiles = compute_rho(tiled_particles, species_config, rho_tiles, constants, world) - return assemble_tiled_scalar_field(rho_tiles, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - - -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_reflecting_boundary_mirrors_position_from_global_world_bc(self): - world = { - "particle_boundary_conditions": {"x": 1, "y": 0, "z": 0}, - } - x1 = jnp.array([self.x_wind / 2 + 0.1, -self.x_wind / 2 - 0.2, 0.0]) - x2 = jnp.array([0.0, 0.0, 0.0]) - x3 = jnp.array([0.0, 0.0, 0.0]) - - species = particle_species( - name="reflecting", - 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=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="periodic", - dt=self.dt, - ) - - species.boundary_conditions(world) - - x, _, _ = species.get_forward_position() - vx, _, _ = species.get_velocity() - self.assertTrue(jnp.allclose(x, jnp.array([self.x_wind / 2 - 0.1, -self.x_wind / 2 + 0.2, 0.0]))) - self.assertTrue(jnp.allclose(vx, jnp.array([-1.0, 2.0, 3.0]))) - - def test_absorbing_boundary_updates_active_mask(self): - world = { - "particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}, - } - 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="periodic", - dt=self.dt, - ) - - species.boundary_conditions(world) - - 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, - "shape_factor": 2, - } - constants = {"kb": self.kb, "eps": 1.0} - simulation_parameters = {"ds_per_debye": None, "shape_factor": 2} - 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]))) - self.assertEqual(world["shape_factor"], 2) - self.assertFalse(hasattr(particles[0], "shape")) - self.assertFalse(hasattr(particles[0], "get_shape")) - - 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({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) - - 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_inactive_particles_do_not_deposit_rho(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, - "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} - - 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({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) - - constants = {"C": 3e8, "alpha": 1.0} - rho = _compute_rho_one_tile([species], jnp.zeros((Nx + 2, Ny + 2, Nz + 2)), world, constants) - - self.assertAlmostEqual(float(jnp.sum(rho[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, - '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} - # 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_one_tile([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, - "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} - - 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_one_tile([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, - "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} - - 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_one_tile([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_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, - "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} - - 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_one_tile([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, - "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} - - 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_one_tile([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])) - -if __name__ == '__main__': - unittest.main() diff --git a/tests/pml_test.py b/tests/pml_test.py index 6d66e7f9..7ef855ea 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -16,6 +16,7 @@ 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.tiled_particles import SpeciesConfig, TiledParticles from PyPIC3D.solvers.first_order_yee import ( update_B, update_E, @@ -506,7 +507,29 @@ def test_tiled_pml_absorbs_field_energy_in_particle_free_1d_wave(self): B_tiles = tile_vector_field((Bx, By, Bz), world, tile_shape) J_tiles = tile_vector_field(J, world, tile_shape) - initial_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) + 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) @@ -516,7 +539,16 @@ def step(E_tiles, B_tiles, tiled_pml_state): for _ in range(60): E_tiles, B_tiles, tiled_pml_state = step(E_tiles, B_tiles, tiled_pml_state) - final_energy = sum(compute_energy([], E_tiles, B_tiles, world, constants)[:2]) + 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)) diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index caed8fdb..93e8645d 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -11,8 +11,8 @@ from PyPIC3D.deposition.rho import compute_rho from PyPIC3D.evolve import time_loop_electrostatic from PyPIC3D.initialization import initialize_simulation -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index 4460d2d3..e11dbd4a 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -24,8 +24,8 @@ refresh_tiled_particle_tiles, update_tiled_particle_positions, ) -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.solvers.first_order_yee import ( update_E, diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 0380d042..3f2ba87b 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -18,8 +18,8 @@ compute_tiled_velocity_field_from_tiled_particles, ) from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid diff --git a/tests/tiled_particle_diagnostics_test.py b/tests/tiled_particle_diagnostics_test.py index 9522423c..69ce6682 100644 --- a/tests/tiled_particle_diagnostics_test.py +++ b/tests/tiled_particle_diagnostics_test.py @@ -3,9 +3,9 @@ import jax import jax.numpy as jnp -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_diagnostics import flatten_tiled_particles_by_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.diagnostics.output_adapters import particles_for_output +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles jax.config.update("jax_enable_x64", True) @@ -85,7 +85,7 @@ def test_flatten_tiled_particles_matches_active_original_species(self): species_list = self._species() tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config) + 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)): @@ -93,8 +93,8 @@ def test_flatten_tiled_particles_matches_active_original_species(self): original_x = jnp.stack(original.get_forward_position(), axis=-1)[active] original_u = jnp.stack(original.get_velocity(), axis=-1)[active] - flattened_x = jnp.stack(flattened.get_forward_position(), axis=-1) - flattened_u = jnp.stack(flattened.get_velocity(), axis=-1) + flattened_x = flattened.x + flattened_u = flattened.u self.assertEqual(flattened.species_index, species_index) self.assertTrue(jnp.allclose(flattened_x, original_x)) @@ -107,12 +107,12 @@ def test_inactive_tiled_slots_are_not_flattened(self): species_list = self._species() tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config) + flattened_species = particles_for_output(tiled_particles, species_config=species_config) - self.assertEqual(flattened_species[0].get_number_of_particles(), 3) - self.assertEqual(flattened_species[1].get_number_of_particles(), 2) - self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].get_forward_position()[0], -0.5))) - self.assertFalse(jnp.any(jnp.isclose(flattened_species[1].get_forward_position()[0], -1.25))) + 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() @@ -121,10 +121,10 @@ def test_absorbed_particles_do_not_appear_in_flattened_output(self): tiled_particles = tiled_particles._replace( active=tiled_particles.active.at[1, 1, 0, 0, 1].set(False) ) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config) + flattened_species = particles_for_output(tiled_particles, species_config=species_config) - self.assertEqual(flattened_species[0].get_number_of_particles(), 2) - self.assertFalse(jnp.any(jnp.isclose(flattened_species[0].get_forward_position()[0], 1.5))) + 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() @@ -136,12 +136,12 @@ def test_flattened_diagnostic_positions_match_original_half_step_positions(self) } tiled_particles, species_config = to_tiled_particles(species_list, world, self._simulation_parameters()) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config, world=world) + flattened_species = particles_for_output(tiled_particles, species_config=species_config, world=world) for original, flattened in zip(species_list, flattened_species): active = original.get_active_mask() original_position = jnp.stack(original.get_position(), axis=-1)[active] - flattened_position = jnp.stack(flattened.get_position(), axis=-1) + flattened_position = flattened.x_diagnostic self.assertTrue(jnp.allclose(flattened_position, original_position)) @@ -150,10 +150,10 @@ def test_species_names_are_preserved_when_metadata_is_available(self): tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) species_names = tuple(species.get_name() for species in species_list) - flattened_species = flatten_tiled_particles_by_species(tiled_particles, species_config=species_config, species_names=species_names) + flattened_species = particles_for_output(tiled_particles, species_config=species_config, species_names=species_names) - self.assertEqual(flattened_species[0].get_name(), "ions") - self.assertEqual(flattened_species[1].get_name(), "electrons") + self.assertEqual(flattened_species[0].name, "ions") + self.assertEqual(flattened_species[1].name, "electrons") if __name__ == "__main__": diff --git a/tests/tiled_particle_fixtures.py b/tests/tiled_particle_fixtures.py new file mode 100644 index 00000000..79961fb0 --- /dev/null +++ b/tests/tiled_particle_fixtures.py @@ -0,0 +1,280 @@ +import math + +import jax.numpy as jnp +import numpy as np + +from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles + + +class particle_species: + """ + Test-only particle constructor for building tiled particle fixtures. + + Production code no longer has a flat species class. Several numerical tests + still need a compact way to name positions, velocities, masks, and per-species + constants before packing them into TiledParticles. + """ + + def __init__( + self, + name, + N_particles, + charge, + mass, + weight=1.0, + T=0.0, + x1=None, + x2=None, + x3=None, + v1=None, + v2=None, + v3=None, + xwind=1.0, + ywind=1.0, + zwind=1.0, + dx=1.0, + dy=1.0, + dz=1.0, + dt=1.0, + update_pos=True, + update_x=True, + update_y=True, + update_z=True, + update_v=True, + update_vx=True, + update_vy=True, + update_vz=True, + active_mask=None, + x_bc="periodic", + y_bc="periodic", + z_bc="periodic", + ): + self.name = name + self.N_particles = int(N_particles) + self.charge = charge + self.mass = mass + self.weight = weight + self.T = T + self.xwind = xwind + self.ywind = ywind + self.zwind = zwind + self.dx = dx + self.dy = dy + self.dz = dz + self.dt = dt + self.x_bc = x_bc + self.y_bc = y_bc + self.z_bc = z_bc + + self.update_pos = update_pos + self.update_x = update_pos and update_x + self.update_y = update_pos and update_y + self.update_z = update_pos and update_z + self.update_v = update_v + self.update_vx = update_v and update_vx + self.update_vy = update_v and update_vy + self.update_vz = update_v and update_vz + + self.x1 = jnp.asarray(x1 if x1 is not None else jnp.zeros(self.N_particles)) + self.x2 = jnp.asarray(x2 if x2 is not None else jnp.zeros(self.N_particles)) + self.x3 = jnp.asarray(x3 if x3 is not None else jnp.zeros(self.N_particles)) + self.v1 = jnp.asarray(v1 if v1 is not None else jnp.zeros(self.N_particles)) + self.v2 = jnp.asarray(v2 if v2 is not None else jnp.zeros(self.N_particles)) + self.v3 = jnp.asarray(v3 if v3 is not None else jnp.zeros(self.N_particles)) + + if active_mask is None: + self.active_mask = jnp.ones(self.N_particles, dtype=bool) + else: + self.active_mask = jnp.asarray(active_mask, dtype=bool) + + def get_name(self): + return self.name + + def get_charge(self): + return self.charge + + def get_mass(self): + return self.mass + + def get_weight(self): + return self.weight + + def get_temperature(self): + return self.T + + def get_number_of_particles(self): + return self.N_particles + + def get_velocity(self): + return self.v1, self.v2, self.v3 + + def get_active_mask(self): + return self.active_mask + + def get_forward_position(self): + return self.x1, self.x2, self.x3 + + def _half_step_axis(self, x, v, wind, bc): + x_diagnostic = x - v * self.dt / 2 + if bc == "periodic": + 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 get_position(self): + return ( + self._half_step_axis(self.x1, self.v1, self.xwind, self.x_bc), + self._half_step_axis(self.x2, self.v2, self.ywind, self.y_bc), + self._half_step_axis(self.x3, self.v3, self.zwind, self.z_bc), + ) + + def set_position(self, x1, x2, x3): + self.x1 = jnp.asarray(x1) + self.x2 = jnp.asarray(x2) + self.x3 = jnp.asarray(x3) + + def set_velocity(self, v1, v2, v3): + self.v1 = jnp.asarray(v1) + self.v2 = jnp.asarray(v2) + self.v3 = jnp.asarray(v3) + + def set_mass(self, mass): + self.mass = mass + + def set_weight(self, weight): + self.weight = weight + + def update_position(self): + self.x1 = jnp.where(self.update_x, self.x1 + self.v1 * self.dt, self.x1) + self.x2 = jnp.where(self.update_y, self.x2 + self.v2 * self.dt, self.x2) + self.x3 = jnp.where(self.update_z, self.x3 + self.v3 * self.dt, self.x3) + + def _apply_axis_bc(self, x, wind, bc): + half_wind = wind / 2 + if int(bc) == 0: + x = jnp.where(x > half_wind, x - wind, x) + x = jnp.where(x < -half_wind, x + wind, x) + return x, self.active_mask + if int(bc) == 2: + active = self.active_mask & (x >= -half_wind) & (x <= half_wind) + return x, active + return x, self.active_mask + + def boundary_conditions(self, world): + particle_bc = world.get("particle_boundary_conditions", {"x": 0, "y": 0, "z": 0}) + self.x1, self.active_mask = self._apply_axis_bc(self.x1, self.xwind, particle_bc.get("x", 0)) + self.x2, self.active_mask = self._apply_axis_bc(self.x2, self.ywind, particle_bc.get("y", 0)) + self.x3, self.active_mask = self._apply_axis_bc(self.x3, self.zwind, particle_bc.get("z", 0)) + + def kinetic_energy(self): + v2 = self.v1**2 + self.v2**2 + self.v3**2 + return jnp.sum(self.active_mask.astype(v2.dtype) * 0.5 * self.mass * self.weight * v2) + + def momentum(self): + vmag = jnp.sqrt(self.v1**2 + self.v2**2 + self.v3**2) + return jnp.sum(self.active_mask.astype(vmag.dtype) * self.mass * self.weight * vmag) + + +def _tile_shape(world, simulation_parameters): + if "tile_shape" in world: + return tuple(int(width) for width in world["tile_shape"]) + return ( + int(simulation_parameters.get("particle_tile_nx", world["Nx"])), + int(simulation_parameters.get("particle_tile_ny", world["Ny"])), + int(simulation_parameters.get("particle_tile_nz", world["Nz"])), + ) + + +def _tile_axis_count(n_cells, cells_per_tile): + return int(math.ceil(int(n_cells) / int(cells_per_tile))) + + +def _particle_tile_indices(x1, x2, x3, world, tile_shape): + tile_nx, tile_ny, tile_nz = tile_shape + Nx = int(world["Nx"]) + Ny = int(world["Ny"]) + Nz = 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 to_tiled_particles(particles, world, simulation_parameters): + tile_shape = _tile_shape(world, simulation_parameters) + 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) + + tile_counts = np.zeros((ntx, nty, ntz, len(particles)), dtype=int) + species_data = [] + + for species_index, species in enumerate(particles): + x1, x2, x3 = species.get_forward_position() + v1, v2, v3 = species.get_velocity() + active = species.get_active_mask() + tx, ty, tz = _particle_tile_indices(x1, x2, x3, world, tile_shape) + + for particle_index in range(species.get_number_of_particles()): + tile_counts[tx[particle_index], ty[particle_index], tz[particle_index], species_index] += 1 + + species_data.append((x1, x2, x3, v1, v2, v3, active, tx, ty, tz)) + + capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) + max_particles_per_tile = int(math.ceil(max(1, tile_counts.max()) * capacity_factor)) + + x_tiles = np.zeros((ntx, nty, ntz, len(particles), max_particles_per_tile, 3), dtype=float) + u_tiles = np.zeros_like(x_tiles) + active_tiles = np.zeros((ntx, nty, ntz, len(particles), max_particles_per_tile), dtype=bool) + write_counts = np.zeros_like(tile_counts) + + for species_index, (x1, x2, x3, v1, v2, v3, active, tx, ty, tz) in enumerate(species_data): + for particle_index in range(particles[species_index].get_number_of_particles()): + tile = (tx[particle_index], ty[particle_index], tz[particle_index], species_index) + slot = write_counts[tile] + write_counts[tile] += 1 + x_tiles[tile + (slot, slice(None))] = [ + float(x1[particle_index]), + float(x2[particle_index]), + float(x3[particle_index]), + ] + u_tiles[tile + (slot, slice(None))] = [ + float(v1[particle_index]), + float(v2[particle_index]), + float(v3[particle_index]), + ] + active_tiles[tile + (slot,)] = bool(active[particle_index]) + + species_config = SpeciesConfig( + charge=jnp.asarray([species.charge for species in particles]), + mass=jnp.asarray([species.mass for species in particles]), + weight=jnp.asarray([species.weight for species in particles]), + update_x=jnp.asarray( + [[species.update_x, species.update_y, species.update_z] for species in particles], + dtype=bool, + ), + update_u=jnp.asarray( + [[species.update_vx, species.update_vy, species.update_vz] for species in particles], + dtype=bool, + ), + ) + + return ( + TiledParticles( + x=jnp.asarray(x_tiles), + u=jnp.asarray(u_tiles), + active=jnp.asarray(active_tiles), + ), + species_config, + ) diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py index e294c467..f7bc7487 100644 --- a/tests/tiled_particle_initialization_test.py +++ b/tests/tiled_particle_initialization_test.py @@ -6,10 +6,10 @@ import jax.numpy as jnp import numpy as np -from PyPIC3D.particles.tiled_particle_initialization import load_tiled_particles_from_toml -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.particle_initialization import load_particles_from_toml +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles -from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import particle_species jax.config.update("jax_enable_x64", True) @@ -215,7 +215,7 @@ def test_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): self.assertEqual(int(jnp.sum(particles.active)), 3) self.assertEqual(species_config.charge.shape, (1,)) - def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): + 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]) @@ -235,6 +235,7 @@ def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): "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 = { @@ -261,10 +262,12 @@ def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): } } - particles, species_config = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + 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)) @@ -282,7 +285,7 @@ def test_load_tiled_particles_from_toml_uses_tile_axes_before_species(self): 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_tiled_particles_from_toml_maps_update_flags_to_active_particles(self): + 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]) @@ -299,6 +302,7 @@ def test_load_tiled_particles_from_toml_maps_update_flags_to_active_particles(se "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 = { @@ -332,7 +336,7 @@ def test_load_tiled_particles_from_toml_maps_update_flags_to_active_particles(se } } - particles, species_config = load_tiled_particles_from_toml(config, simulation_parameters, world, constants) + 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])) diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index 491f3ccb..f30bab6d 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -3,8 +3,8 @@ import jax import jax.numpy as jnp -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particle_refresh import ( _adjacent_tile_offset, _refresh_tiled_particle_tiles_compacting, diff --git a/tests/tiled_pusher_test.py b/tests/tiled_pusher_test.py index e2c20146..4328276a 100644 --- a/tests/tiled_pusher_test.py +++ b/tests/tiled_pusher_test.py @@ -4,8 +4,8 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.pusher.particle_push import particle_push from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid diff --git a/tests/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py index 02a48129..b39758b9 100644 --- a/tests/tiled_refactor_contract_test.py +++ b/tests/tiled_refactor_contract_test.py @@ -207,6 +207,16 @@ def test_deprecated_fdtd_errors_and_root_boris_modules_are_removed(self): self.assertIsNone(importlib.util.find_spec("PyPIC3D.boris")) self.assertIsNotNone(importlib.util.find_spec("PyPIC3D.pusher.boris")) + def test_old_particle_species_surfaces_are_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.species_class")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.tiled_particle_diagnostics")) + + from PyPIC3D.particles import particle_initialization + + self.assertFalse(hasattr(particle_initialization, "to_tiled_particles")) + source = inspect.getsource(particle_initialization.load_particles_from_toml) + self.assertNotIn("particle_species", source) + def test_particle_push_owns_tiled_pusher_contract(self): particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py index f10c74d1..410e88f9 100644 --- a/tests/tiled_rho_test.py +++ b/tests/tiled_rho_test.py @@ -7,8 +7,8 @@ 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 PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 405f3fa5..3707a7ae 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -9,8 +9,8 @@ from PyPIC3D.evolve import time_loop_electrodynamic from PyPIC3D.initialization import initialize_simulation, validate_field_solver -from PyPIC3D.particles.species_class import particle_species -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.tiled_particles import TiledParticles from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field diff --git a/tests/utils_test.py b/tests/utils_test.py index 7ce46ba6..1aa10bca 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -7,8 +7,8 @@ import toml from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.initialization import initialize_fields -from PyPIC3D.particles.tiled_particles import TiledParticles -from PyPIC3D.particles.tiled_particle_initialization import to_tiled_particles +from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid from PyPIC3D.utils import ( @@ -16,7 +16,7 @@ particle_sanity_check, load_external_fields_from_toml, add_external_fields, compute_energy, dump_parameters_to_toml, ) -from PyPIC3D.particles.species_class import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import particle_species jax.config.update("jax_enable_x64", True) @@ -146,18 +146,11 @@ def test_check_stability(self): # 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() + 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 @@ -268,14 +261,33 @@ def test_energy_can_include_external_fields(self): world = world | {"dx": 1.0, "dy": 1.0, "dz": 1.0} constants = {"eps": 2.0, "mu": 4.0, "C": 10.0} - e_energy, b_energy, kinetic_energy = compute_energy([], total_E, total_B, world, constants) + 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} + 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} species = particle_species( @@ -301,7 +313,9 @@ def test_compute_energy_ignores_inactive_particles(self): ) species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) - _, _, kinetic_energy = compute_energy([species], E, B, world, constants) + tiled_particles, species_config = to_tiled_particles([species], world, {"particle_tile_nx": 1, "particle_tile_ny": 1, "particle_tile_nz": 1}) + + _, _, kinetic_energy = compute_energy(tiled_particles, E, B, world, constants, species_config=species_config) self.assertTrue(jnp.allclose(kinetic_energy, 0.0)) @@ -374,6 +388,7 @@ def test_plot_positions_flattens_tiled_particles_and_preserves_species_names(sel "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 = { @@ -446,6 +461,7 @@ def test_particles_phase_space_flattens_tiled_particles(self): "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 = { From 7ea2232e0fd3fbf146c602e78a5b9c84a10b27f6 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 18:21:32 -0500 Subject: [PATCH 51/83] refactoring particle initialization --- PyPIC3D/particles/particle_initialization.py | 112 ++++++------------- 1 file changed, 34 insertions(+), 78 deletions(-) diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index f4b224d5..07be0a07 100644 --- a/PyPIC3D/particles/particle_initialization.py +++ b/PyPIC3D/particles/particle_initialization.py @@ -46,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, @@ -146,28 +102,29 @@ def _particle_tile_indices(x1, x2, x3, world, tile_nx, tile_ny, tile_nz): return x_cell // tile_nx, y_cell // tile_ny, z_cell // tile_nz -def _scaled_plasma_frequency(N_particles, charge, mass, weight, world, constants): - x_wind = world["x_wind"] - y_wind = world["y_wind"] - z_wind = world["z_wind"] +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"] - - sqrt_dv = jnp.sqrt(x_wind * y_wind * z_wind) - sqrt_ne = jnp.sqrt(N_particles * weight) / sqrt_dv + # 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_ne * jnp.abs(charge) / (sqrt_eps * sqrt_mass) + return sqrt_n * jnp.abs(charge) / (sqrt_eps * sqrt_mass) / (2 * jnp.pi) -def _scaled_debye_length(N_particles, charge, temperature, weight, world, constants): +def _debye_length(N_per_cell, charge, temperature, weight, world, constants): eps = constants["eps"] kb = constants["kb"] - x_wind = world["x_wind"] - y_wind = world["y_wind"] - z_wind = world["z_wind"] + # 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 - density = weight * N_particles / (x_wind * y_wind * z_wind) - return jnp.sqrt(eps * kb * temperature / (density * charge**2)) + return jnp.sqrt(eps * kb * temperature / (n * charge**2)) def load_particles_from_toml(config, simulation_parameters, world, constants): @@ -183,15 +140,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 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) @@ -201,6 +158,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"] @@ -208,6 +166,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"] @@ -218,10 +177,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) @@ -277,6 +238,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)) * ( @@ -285,21 +247,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) @@ -309,6 +260,7 @@ 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) @@ -333,8 +285,10 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): species_metadata.append(metadata) kinetic_energy = 0.5 * weight * mass * jnp.sum(vx**2 + vy**2 + vz**2) - pf = _scaled_plasma_frequency(N_particles, charge, mass, weight, world, constants) - dl = _scaled_debye_length(N_particles, charge, T, weight, world, constants) + # 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 print(f"Number of particles: {N_particles}") print(f"Number of particles per cell: {N_per_cell}") @@ -352,6 +306,7 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): 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 ntx = _tile_axis_count(world["Nx"], tile_nx) nty = _tile_axis_count(world["Ny"], tile_ny) @@ -372,6 +327,7 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) + # 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 = jnp.zeros((ntx, nty, ntz, n_species, max_particles_per_tile, 3)) u_tiles = jnp.zeros_like(x_tiles) From 6a339b6b2d7e50c9dcb767d6142cd60f337b850e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Sun, 5 Jul 2026 18:27:37 -0500 Subject: [PATCH 52/83] refactoring default initializations --- PyPIC3D/initialization.py | 13 ++++++++++--- demos/two_stream/two_stream.toml | 2 +- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 7b562155..c47da485 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -199,9 +199,9 @@ def default_parameters(): "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": 1, # number of x cells per shared field/particle tile - "particle_tile_ny": 1, # number of y cells per shared field/particle tile - "particle_tile_nz": 1, # number of z cells per shared field/particle tile + "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 @@ -281,6 +281,13 @@ def initialize_simulation(toml_file): electrostatic = solver == "electrostatic" relativistic = simulation_parameters['relativistic'] particle_pusher = simulation_parameters['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) diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 3f013a43..a6d39bf0 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -29,7 +29,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] From 5b7ed383ecf21b0e9c896ad4b25aec65836f281c Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 05:03:26 -0500 Subject: [PATCH 53/83] fixing bug --- PyPIC3D/__main__.py | 4 +- PyPIC3D/utils.py | 8 ++- tests/tiled_esirkepov_integration_test.py | 64 ++++++++++++++++++++++- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index ad69c15a..09ca4713 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -215,7 +215,7 @@ def loop_with_static_world( # soon as it is detected. - return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world + return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world, species_config def main(): ###################### JAX SETTINGS ######################################################################## @@ -234,7 +234,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() diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 1677c8d6..fe9c6868 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -691,7 +691,12 @@ def dump_parameters_to_toml(simulation_stats, simulation_parameters, plasma_para if species_metadata is None: particle_dict = {"name": name} else: - particle_dict = dict(species_metadata[species_index]) + 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 @@ -775,4 +780,3 @@ def courant_condition(courant_number, dx, dy, dz, simulation_parameters, constan # compute the maximum allowed timestep return dt - diff --git a/tests/tiled_esirkepov_integration_test.py b/tests/tiled_esirkepov_integration_test.py index 982721f1..6a2a7843 100644 --- a/tests/tiled_esirkepov_integration_test.py +++ b/tests/tiled_esirkepov_integration_test.py @@ -1,10 +1,15 @@ import os import tempfile import unittest +from contextlib import redirect_stdout +from io import StringIO +from unittest.mock import patch import jax import numpy as np +import toml +from PyPIC3D import __main__ as pypic_main from PyPIC3D.__main__ import run_PyPIC3D @@ -67,7 +72,7 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): }, } - _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world = run_PyPIC3D(config) + _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world, _species_config = run_PyPIC3D(config) _E, _B, J, *_rest = fields self.assertEqual(J[0].shape[-3:], (6, 5, 5)) @@ -85,6 +90,63 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice", "vector_field_slice_000000000.vtk"))) self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) + def test_main_prints_final_energy_for_tiled_runtime(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + config_path = os.path.join(tmpdir, "final_energy.toml") + np.save(x_path, np.array([0.0])) + np.save(zeros_path, np.zeros(1)) + + config = { + "simulation_parameters": { + "name": "final energy smoke", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 4, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 1.0e-10, + "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": 1, + "charge": -1.0, + "mass": 1.0, + "weight": 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, + }, + } + with open(config_path, "w") as f: + toml.dump(config, f) + + with patch("sys.argv", ["PyPIC3D", "--config", config_path]): + stdout = StringIO() + with redirect_stdout(stdout): + pypic_main.main() + + self.assertIn("Final Kinetic Energy", stdout.getvalue()) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "output.toml"))) + if __name__ == "__main__": unittest.main() From 18c1f8afd21587b7c5924444a76984a11969d357 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 05:08:36 -0500 Subject: [PATCH 54/83] refactoring weibel instability --- demos/weibel/weibel.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/weibel/weibel.toml b/demos/weibel/weibel.toml index 5e185e8e..ac733cf2 100644 --- a/demos/weibel/weibel.toml +++ b/demos/weibel/weibel.toml @@ -26,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 From d6895f5309ad661bbbfad65259f202dfc6037524 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 05:10:34 -0500 Subject: [PATCH 55/83] refactoring code --- PyPIC3D/deposition/Esirkepov.py | 2 +- PyPIC3D/deposition/J_from_rhov.py | 2 +- PyPIC3D/deposition/rho.py | 2 +- PyPIC3D/diagnostics/fluid_quantities.py | 2 +- PyPIC3D/diagnostics/output_adapters.py | 2 +- PyPIC3D/evolve.py | 2 +- PyPIC3D/particles/{tiled_particles.py => particle_class.py} | 0 PyPIC3D/particles/particle_initialization.py | 2 +- ...led_particle_refresh.py => particle_tile_communication.py} | 2 +- PyPIC3D/pusher/particle_push.py | 2 +- tests/direct_deposition_tiled_test.py | 2 +- tests/initialization_test.py | 2 +- tests/pml_test.py | 2 +- tests/tiled_electrostatic_test.py | 2 +- tests/tiled_esirkepov_test.py | 4 ++-- tests/tiled_particle_fixtures.py | 2 +- tests/tiled_particle_initialization_test.py | 2 +- tests/tiled_particle_refresh_test.py | 2 +- tests/tiled_yee_integration_test.py | 2 +- tests/utils_test.py | 2 +- 20 files changed, 20 insertions(+), 20 deletions(-) rename PyPIC3D/particles/{tiled_particles.py => particle_class.py} (100%) rename PyPIC3D/particles/{tiled_particle_refresh.py => particle_tile_communication.py} (99%) diff --git a/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index b1bf2096..8c5bdc92 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -8,7 +8,7 @@ ) 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.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles def shift_old_stencil(old_w_list, shift): diff --git a/PyPIC3D/deposition/J_from_rhov.py b/PyPIC3D/deposition/J_from_rhov.py index b5760eab..a83f3a6a 100644 --- a/PyPIC3D/deposition/J_from_rhov.py +++ b/PyPIC3D/deposition/J_from_rhov.py @@ -1,4 +1,4 @@ -from PyPIC3D.particles.tiled_particles import TiledParticles, SpeciesConfig +from PyPIC3D.particles.particle_class import TiledParticles, SpeciesConfig from PyPIC3D.boundary_conditions.grid_and_stencil import ( collapse_axis_stencil, diff --git a/PyPIC3D/deposition/rho.py b/PyPIC3D/deposition/rho.py index 5745b2da..c021e104 100644 --- a/PyPIC3D/deposition/rho.py +++ b/PyPIC3D/deposition/rho.py @@ -11,7 +11,7 @@ ) 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.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.utilities.filters import digital_filter diff --git a/PyPIC3D/diagnostics/fluid_quantities.py b/PyPIC3D/diagnostics/fluid_quantities.py index fcca6905..17d5e815 100644 --- a/PyPIC3D/diagnostics/fluid_quantities.py +++ b/PyPIC3D/diagnostics/fluid_quantities.py @@ -3,7 +3,7 @@ compute_tiled_pressure_field_from_tiled_particles, compute_tiled_velocity_field_from_tiled_particles, ) -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles def compute_mass_density(particles, rho, world, species_config=None): diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index dc395db5..6591f960 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -1,7 +1,7 @@ import jax.numpy as jnp from typing import NamedTuple -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles class ParticleOutputRecord(NamedTuple): diff --git a/PyPIC3D/evolve.py b/PyPIC3D/evolve.py index 4f9baffe..53322fac 100644 --- a/PyPIC3D/evolve.py +++ b/PyPIC3D/evolve.py @@ -9,7 +9,7 @@ from PyPIC3D.deposition.Esirkepov import Esirkepov_current from PyPIC3D.deposition.J_from_rhov import J_from_rhov -from PyPIC3D.particles.tiled_particle_refresh import ( +from PyPIC3D.particles.particle_tile_communication import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, ) diff --git a/PyPIC3D/particles/tiled_particles.py b/PyPIC3D/particles/particle_class.py similarity index 100% rename from PyPIC3D/particles/tiled_particles.py rename to PyPIC3D/particles/particle_class.py diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index 07be0a07..c3f429e9 100644 --- a/PyPIC3D/particles/particle_initialization.py +++ b/PyPIC3D/particles/particle_initialization.py @@ -4,7 +4,7 @@ import math from PyPIC3D.utils import vth_to_T, T_to_vth -from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles def grab_particle_keys(config): diff --git a/PyPIC3D/particles/tiled_particle_refresh.py b/PyPIC3D/particles/particle_tile_communication.py similarity index 99% rename from PyPIC3D/particles/tiled_particle_refresh.py rename to PyPIC3D/particles/particle_tile_communication.py index 2d9c906a..87f56093 100644 --- a/PyPIC3D/particles/tiled_particle_refresh.py +++ b/PyPIC3D/particles/particle_tile_communication.py @@ -2,7 +2,7 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions.grid_and_stencil import wrap_periodic_position -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles def _particle_boundary_conditions(world): diff --git a/PyPIC3D/pusher/particle_push.py b/PyPIC3D/pusher/particle_push.py index 0bdb09d6..27174bb5 100644 --- a/PyPIC3D/pusher/particle_push.py +++ b/PyPIC3D/pusher/particle_push.py @@ -1,7 +1,7 @@ import jax import jax.numpy as jnp -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.pusher.boris import ( boris_single_particle, interpolate_field_to_particles, diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py index 278e2bd9..3379ea3b 100644 --- a/tests/direct_deposition_tiled_test.py +++ b/tests/direct_deposition_tiled_test.py @@ -12,7 +12,7 @@ ) from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particle_refresh import refresh_tiled_particle_tiles +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 diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 5f64be6e..6cf0047c 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -7,7 +7,7 @@ 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.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles jax.config.update("jax_enable_x64", True) diff --git a/tests/pml_test.py b/tests/pml_test.py index 7ef855ea..de4b6859 100644 --- a/tests/pml_test.py +++ b/tests/pml_test.py @@ -16,7 +16,7 @@ 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.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles from PyPIC3D.solvers.first_order_yee import ( update_B, update_E, diff --git a/tests/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py index 93e8645d..68eb6863 100644 --- a/tests/tiled_electrostatic_test.py +++ b/tests/tiled_electrostatic_test.py @@ -13,7 +13,7 @@ from PyPIC3D.initialization import initialize_simulation from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field from PyPIC3D.utilities.grids import build_collocated_grid, build_tiled_yee_grids, build_yee_grid diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py index e11dbd4a..a75555f7 100644 --- a/tests/tiled_esirkepov_test.py +++ b/tests/tiled_esirkepov_test.py @@ -20,13 +20,13 @@ 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.tiled_particle_refresh import ( +from PyPIC3D.particles.particle_tile_communication import ( refresh_tiled_particle_tiles, update_tiled_particle_positions, ) from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.solvers.first_order_yee import ( update_E, ) diff --git a/tests/tiled_particle_fixtures.py b/tests/tiled_particle_fixtures.py index 79961fb0..d8f05cf6 100644 --- a/tests/tiled_particle_fixtures.py +++ b/tests/tiled_particle_fixtures.py @@ -3,7 +3,7 @@ import jax.numpy as jnp import numpy as np -from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles class particle_species: diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py index f7bc7487..7cfa18e4 100644 --- a/tests/tiled_particle_initialization_test.py +++ b/tests/tiled_particle_initialization_test.py @@ -8,7 +8,7 @@ from PyPIC3D.particles.particle_initialization import load_particles_from_toml from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles from PyPIC3D.tests.tiled_particle_fixtures import particle_species diff --git a/tests/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py index f30bab6d..9046b045 100644 --- a/tests/tiled_particle_refresh_test.py +++ b/tests/tiled_particle_refresh_test.py @@ -5,7 +5,7 @@ from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particle_refresh import ( +from PyPIC3D.particles.particle_tile_communication import ( _adjacent_tile_offset, _refresh_tiled_particle_tiles_compacting, _refresh_tiled_particle_tiles_sparse, diff --git a/tests/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py index 3707a7ae..2c544984 100644 --- a/tests/tiled_yee_integration_test.py +++ b/tests/tiled_yee_integration_test.py @@ -11,7 +11,7 @@ from PyPIC3D.initialization import initialize_simulation, validate_field_solver from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.tiled_particles import TiledParticles +from PyPIC3D.particles.particle_class import TiledParticles from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field from PyPIC3D.solvers.first_order_yee import ( diff --git a/tests/utils_test.py b/tests/utils_test.py index 1aa10bca..845fb46f 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -7,7 +7,7 @@ import toml from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.initialization import initialize_fields -from PyPIC3D.particles.tiled_particles import SpeciesConfig, TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid From 573ede0bbc0515b7eb67bc5c0c2a3f8a5d239014 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 05:31:10 -0500 Subject: [PATCH 56/83] removing deprecated vtk plotting --- PyPIC3D/__init__.py | 1 - PyPIC3D/__main__.py | 51 ----- PyPIC3D/diagnostics/output_adapters.py | 2 +- PyPIC3D/diagnostics/vtk.py | 240 ---------------------- PyPIC3D/initialization.py | 3 - PyPIC3D/utils.py | 2 - README.md | 4 +- demos/reconnection_2d/harris_current.toml | 1 - docs/architecture.rst | 4 +- docs/index.rst | 3 +- docs/usage.rst | 7 +- pyproject.toml | 2 - tests/initialization_test.py | 3 + tests/tiled_esirkepov_integration_test.py | 11 +- tests/tiled_fluid_quantities_test.py | 8 - tests/utils_test.py | 44 +--- 16 files changed, 23 insertions(+), 363 deletions(-) delete mode 100644 PyPIC3D/diagnostics/vtk.py diff --git a/PyPIC3D/__init__.py b/PyPIC3D/__init__.py index d597a907..2abab663 100644 --- a/PyPIC3D/__init__.py +++ b/PyPIC3D/__init__.py @@ -23,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 09ca4713..21387944 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 @@ -23,10 +22,6 @@ write_openpmd_particles, write_openpmd_fields ) -from PyPIC3D.diagnostics.vtk import ( - plot_field_slice_vtk, plot_vectorfield_slice_vtk, plot_vtk_particles -) - from PyPIC3D.utils import ( dump_parameters_to_toml, load_config_file, compute_energy, setup_pmd_files, add_external_fields, compute_total_momentum @@ -36,13 +31,6 @@ initialize_simulation ) -from PyPIC3D.diagnostics.fluid_quantities import ( - compute_mass_density -) - -from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.diagnostics.output_adapters import fields_for_output, scalar_field_for_output - # Importing functions from the PyPIC3D package ############################################################################################################ @@ -65,13 +53,6 @@ def run_PyPIC3D(config_file): 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) - - scalar_field_names = ["rho", "mass_density"] - vector_field_names = ["E", "B", "J"] - tile_shape = tuple(int(width) for width in world["tile_shape"]) - g = int(world["guard_cells"]) particle_species_names = simulation_parameters.get("particle_species_names") def loop_with_static_world( @@ -160,38 +141,6 @@ def loop_with_static_world( - if plotting_parameters['plot_vtk_scalars']: - if getattr(rho, "ndim", 0) == 6: - rho = compute_rho(particles, species_config, rho, constants, world) - mass_density = compute_mass_density(particles, rho, world, species_config=species_config) - rho_output = scalar_field_for_output(rho, world) - mass_density_output = scalar_field_for_output(mass_density, world) - else: - raise ValueError("Runtime scalar diagnostics require tiled rho storage.") - - y_mid = world['Ny']//2 + 1 - # midplane index shifted by 1 for ghost cells - fields_mag = [rho_output[1:-1, y_mid, 1:-1], mass_density_output[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']: - output_fields = fields_for_output(fields, world) - E, B, J, rho, phi, external_fields, *rest = output_fields - # assemble tiled fields before VTK output - 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, species_config=species_config, species_names=particle_species_names, world=world) - # 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", species_config=species_config, species_names=particle_species_names) # Write the particles in openPMD format diff --git a/PyPIC3D/diagnostics/output_adapters.py b/PyPIC3D/diagnostics/output_adapters.py index 6591f960..14ad39bb 100644 --- a/PyPIC3D/diagnostics/output_adapters.py +++ b/PyPIC3D/diagnostics/output_adapters.py @@ -80,7 +80,7 @@ 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 and VTK + Runtime diagnostics can operate on tile-major arrays, but openPMD mesh writers still expect one global ghost-celled array. """ diff --git a/PyPIC3D/diagnostics/vtk.py b/PyPIC3D/diagnostics/vtk.py deleted file mode 100644 index ae0928dc..00000000 --- a/PyPIC3D/diagnostics/vtk.py +++ /dev/null @@ -1,240 +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 - -from PyPIC3D.diagnostics.output_adapters import particles_for_output, vector_field_for_output - -def plot_fields(fieldx, fieldy, fieldz, t, name, dx, dy, dz, world=None): - """ - 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 - """ - if world is not None: - fieldx, fieldy, fieldz = vector_field_for_output((fieldx, fieldy, fieldz), world) - # Tiled field storage is assembled only here, at the file-output boundary. - - 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, species_config=None, species_names=None, world=None): - """ - 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") - - particles = particles_for_output(particles, species_config=species_config, species_names=species_names, world=world) - # Tiled particle arrays are fixed-capacity by tile; VTK particle output only - # receives the active diagnostic particles. - - for species in particles: - name = species.name.replace(" ", "") - x, y, z = map(np.asarray, (species.x_diagnostic[:, 0], species.x_diagnostic[:, 1], species.x_diagnostic[:, 2])) - vx, vy, vz = map(np.asarray, (species.u[:, 0], species.u[:, 1], species.u[:, 2])) - # Get the position and velocity of the particles - q = np.asarray(species.charge * species.weight * 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 diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index c47da485..3b037943 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -159,9 +159,6 @@ 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, diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index fe9c6868..0bbbfa72 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 @@ -714,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 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/reconnection_2d/harris_current.toml b/demos/reconnection_2d/harris_current.toml index 8c7e76a8..6bf9c723 100644 --- a/demos/reconnection_2d/harris_current.toml +++ b/demos/reconnection_2d/harris_current.toml @@ -18,7 +18,6 @@ relativistic = true [plotting] plotting_interval = 50 -plot_vtk_particles = false plot_phasespace = true [field1] diff --git a/docs/architecture.rst b/docs/architecture.rst index fd196694..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. @@ -83,6 +83,6 @@ 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/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/usage.rst b/docs/usage.rst index a84ce052..8e362203 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -57,9 +57,6 @@ Minimal Working Example [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 dump_particles = false @@ -135,5 +132,5 @@ 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. diff --git a/pyproject.toml b/pyproject.toml index 9979253a..61c1b095 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,8 +17,6 @@ dependencies = [ "plotly", "toml", "tqdm", - "pyevtk", - "vtk", "openpmd-api", ] diff --git a/tests/initialization_test.py b/tests/initialization_test.py index 6cf0047c..5f7f0f4c 100644 --- a/tests/initialization_test.py +++ b/tests/initialization_test.py @@ -35,6 +35,9 @@ def test_default_parameters(self): 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 diff --git a/tests/tiled_esirkepov_integration_test.py b/tests/tiled_esirkepov_integration_test.py index 6a2a7843..95b7e95e 100644 --- a/tests/tiled_esirkepov_integration_test.py +++ b/tests/tiled_esirkepov_integration_test.py @@ -17,7 +17,7 @@ class TiledEsirkepovIntegrationTest(unittest.TestCase): - def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): + def test_short_output_enabled_tiled_esirkepov_run_writes_openpmd_without_vtk(self): with tempfile.TemporaryDirectory() as tmpdir: x_path = os.path.join(tmpdir, "x.npy") zeros_path = os.path.join(tmpdir, "zeros.npy") @@ -50,11 +50,9 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): }, "plotting": { "plotting_interval": 1, - "plot_vtk_vectors": True, - "plot_vtk_particles": True, "plot_openpmd_fields": True, "plot_openpmd_particles": True, - "plot_phasespace": False, + "plot_phasespace": True, }, "particle1": { "name": "electrons", @@ -87,8 +85,9 @@ def test_short_output_enabled_tiled_esirkepov_run_keeps_live_state_tiled(self): self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", filename))) self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "fields.h5"))) self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles.h5"))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice", "vector_field_slice_000000000.vtk"))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "phase_space", "x", "electrons_phase_space.000000000.npy"))) def test_main_prints_final_energy_for_tiled_runtime(self): with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py index 3f2ba87b..83b1a1fa 100644 --- a/tests/tiled_fluid_quantities_test.py +++ b/tests/tiled_fluid_quantities_test.py @@ -1,10 +1,8 @@ -import inspect import unittest import jax import jax.numpy as jnp -from PyPIC3D import __main__ as pypic_main from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.boundary_conditions.grid_and_stencil import BC_PERIODIC from PyPIC3D.diagnostics.fluid_quantities import ( @@ -402,11 +400,5 @@ def test_public_fluid_quantities_reject_flat_particles(self): with self.assertRaisesRegex(ValueError, "TiledParticles"): compute_pressure_field(particles, field, field, 0, world) - def test_tiled_scalar_vtk_path_no_longer_rejects_electrodynamic_yee(self): - run_source = inspect.getsource(pypic_main.run_PyPIC3D) - - self.assertNotIn("plot_vtk_scalars is not supported for electrodynamic_yee", run_source) - - if __name__ == "__main__": unittest.main() diff --git a/tests/utils_test.py b/tests/utils_test.py index 845fb46f..060415da 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -1,15 +1,17 @@ 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 PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.diagnostics import plotting, vtk as vtk_diagnostics +from PyPIC3D.diagnostics import plotting from PyPIC3D.utilities.grids import build_collocated_grid, build_yee_grid from PyPIC3D.utils import ( print_stats, check_stability, @@ -355,6 +357,10 @@ def test_dump_parameters_to_toml_writes_tiled_species_summaries(self): 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 = particle_species( name="beam electrons", @@ -489,41 +495,5 @@ def test_particles_phase_space_flattens_tiled_particles(self): self.assertEqual(scatter.call_count, 3) - def test_vtk_plot_fields_assembles_tiled_components_at_output_boundary(self): - world = { - "Nx": 4, - "Ny": 2, - "Nz": 2, - "dx": 0.25, - "dy": 0.5, - "dz": 0.75, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.5, - "tile_shape": (2, 1, 1), - "guard_cells": 2, - "boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - field_shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - field = tuple(jnp.zeros(field_shape) + component for component in (1.0, 2.0, 3.0)) - tiled_field = _tile_vector_field(field, world, world["tile_shape"], num_guard_cells=world["guard_cells"]) - - with unittest.mock.patch.object(vtk_diagnostics, "gridToVTK") as grid_to_vtk: - vtk_diagnostics.plot_fields( - tiled_field[0], - tiled_field[1], - tiled_field[2], - 0, - "E", - world["dx"], - world["dy"], - world["dz"], - world=world, - ) - - cell_data = grid_to_vtk.call_args.kwargs["cellData"] - self.assertEqual(cell_data["E_x"].shape, (6, 4, 4)) - self.assertTrue(jnp.allclose(cell_data["E_y"], 2.0)) - if __name__ == '__main__': unittest.main() From 1f7720f6ff312456ff83f5ecb19a9969155896a7 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 05:56:38 -0500 Subject: [PATCH 57/83] migrating towards tiled field writer for openPMD --- PyPIC3D/__main__.py | 112 +++++---- PyPIC3D/diagnostics/openPMD.py | 433 +++++++++++++++++++++++++++++++++ PyPIC3D/initialization.py | 1 + docs/usage.rst | 5 +- tests/openpmd_test.py | 90 ++++++- 5 files changed, 593 insertions(+), 48 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 21387944..c17eadd8 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -19,7 +19,9 @@ ) from PyPIC3D.diagnostics.openPMD import ( - write_openpmd_particles, write_openpmd_fields + write_openpmd_particles, + create_async_tiled_openpmd_field_writer, + enqueue_openpmd_field_output, ) from PyPIC3D.utils import ( @@ -98,7 +100,16 @@ def loop_with_static_world( # 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") + field_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") # setup the openPMD files if needed @@ -106,62 +117,71 @@ def loop_with_static_world( ###################################################### SIMULATION LOOP ##################################### - for t in tqdm(range(Nt)): + loop_error = None + try: + for t in tqdm(range(Nt)): - # plot the data - if t % plotting_parameters['plotting_interval'] == 0: + # plot the data + if t % plotting_parameters['plotting_interval'] == 0: - plot_num = t // plotting_parameters['plotting_interval'] - # determine the plot number + plot_num = t // plotting_parameters['plotting_interval'] + # determine the plot number - E, B, J, rho, phi, external_fields, *rest = fields - # unpack the fields + 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 + 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 - # for species in particles: - # write_data(f"{output_dir}/data/{species.name}_kinetic_energy.txt", t * dt, species.kinetic_energy()) + # 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, species_config=species_config, species_names=particle_species_names, world=world) + # if plotting_parameters['plot_phasespace']: + # write_particles_phase_space(particles, t, output_dir, species_config=species_config, species_names=particle_species_names, world=world) - if plotting_parameters['plot_openpmd_particles']: - write_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5", species_config=species_config, species_names=particle_species_names) - # Write the particles in openPMD format + if plotting_parameters['plot_openpmd_particles']: + write_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5", species_config=species_config, species_names=particle_species_names) + # 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 + 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. + 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: + if field_writer is not None: + field_writer.close(raise_errors=loop_error is None) return Nt, plotting_parameters, simulation_parameters, plasma_parameters, constants, particles, fields, world, species_config diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index 950485b1..edfa738a 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -1,5 +1,14 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union +import queue +import threading +import traceback + import openpmd_api as io +import jax import jax.numpy as jnp import os import numpy as np @@ -7,6 +16,42 @@ from PyPIC3D.diagnostics.output_adapters import fields_for_output, particles_for_output +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]] + + +@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 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: @@ -117,6 +162,394 @@ def _fields_to_interior_map(fields): return field_map +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 _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 _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 _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() + + +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: + 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() + + +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 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 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 diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 3b037943..3580447d 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -162,6 +162,7 @@ def default_parameters(): "plot_openpmd_particles": False, "plot_openpmd_fields": False, "plotting_interval": 10, + "openpmd_field_queue_size": 2, "dump_particles": False, "dump_fields": False, } diff --git a/docs/usage.rst b/docs/usage.rst index 8e362203..211ab52f 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -59,6 +59,7 @@ Minimal Working Example plot_phasespace = false plot_openpmd_particles = false plot_openpmd_fields = false + openpmd_field_queue_size = 2 dump_particles = false dump_fields = false @@ -133,4 +134,6 @@ At ``plotting_interval`` cadence, PyPIC3D writes diagnostics like: - ``total_momentum.txt`` Depending on flags, it also writes matplotlib phase-space arrays and openPMD -files plus ``data/output.toml`` metadata. +files plus ``data/output.toml`` metadata. Runtime openPMD field output uses a +bounded queue; ``openpmd_field_queue_size`` caps the number of pending field +batches so output cannot grow memory without bound. diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index a96bc317..af6e2988 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -76,15 +76,20 @@ def tile_vector_field(field, world, tile_shape, num_guard_cells=2): 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): - pass + 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: @@ -301,6 +306,89 @@ def test_write_openpmd_fields_assembles_tiled_fields_before_output(self): 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 = openPMD.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 = openPMD.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 = openPMD.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 = openPMD.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 = openPMD.TiledFieldSnapshot(step=0, time=0.0, fields={}) + + with patch.object(openPMD, "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_write_openpmd_initial_particles_flattens_tiled_particles_and_preserves_names(self): world = _world() species = [ From 2ccb9c94a064dbe150b2bea23e7b44878cca6151 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 06:50:54 -0500 Subject: [PATCH 58/83] refactoring async_writer --- PyPIC3D/__main__.py | 11 +- PyPIC3D/diagnostics/async_writer.py | 262 +++++++++++++++++++++++++++ PyPIC3D/diagnostics/openPMD.py | 263 +--------------------------- tests/openpmd_test.py | 13 +- 4 files changed, 278 insertions(+), 271 deletions(-) create mode 100644 PyPIC3D/diagnostics/async_writer.py diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index c17eadd8..93246252 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -18,12 +18,15 @@ write_particles_phase_space, write_data ) -from PyPIC3D.diagnostics.openPMD import ( - write_openpmd_particles, +from PyPIC3D.diagnostics.async_writer import ( create_async_tiled_openpmd_field_writer, enqueue_openpmd_field_output, ) +from PyPIC3D.diagnostics.openPMD import ( + write_openpmd_particles, +) + from PyPIC3D.utils import ( dump_parameters_to_toml, load_config_file, compute_energy, setup_pmd_files, add_external_fields, compute_total_momentum @@ -149,8 +152,8 @@ def loop_with_static_world( # 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, species_config=species_config, species_names=particle_species_names, world=world) + if plotting_parameters['plot_phasespace']: + write_particles_phase_space(particles, t, output_dir, species_config=species_config, species_names=particle_species_names, world=world) diff --git a/PyPIC3D/diagnostics/async_writer.py b/PyPIC3D/diagnostics/async_writer.py new file mode 100644 index 00000000..76c9fe63 --- /dev/null +++ b/PyPIC3D/diagnostics/async_writer.py @@ -0,0 +1,262 @@ +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, +) + +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]] + + +@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] + + +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, + ) + + +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: + 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() + + +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 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, + ) diff --git a/PyPIC3D/diagnostics/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index edfa738a..d362d6ae 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -2,13 +2,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union -import queue -import threading -import traceback +from typing import Any, Tuple, Union import openpmd_api as io -import jax import jax.numpy as jnp import os import numpy as np @@ -16,29 +12,6 @@ from PyPIC3D.diagnostics.output_adapters import fields_for_output, particles_for_output -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]] - - -@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 TiledMeshLayout: @@ -97,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 @@ -162,21 +127,6 @@ def _fields_to_interior_map(fields): return field_map -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 _as_3tuple(value): if isinstance(value, tuple): return tuple(int(v) for v in value) @@ -191,67 +141,6 @@ def _slice_start(index_entry): return int(index_entry) -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 _tile_interior(tile, guard_cells): gx, gy, gz = guard_cells sx = slice(gx, -gx) if gx > 0 else slice(None) @@ -402,154 +291,6 @@ def write_tiled_field_snapshot_openpmd( series.close() -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: - 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() - - -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 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 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 diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index af6e2988..31427614 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -5,6 +5,7 @@ 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.tests.tiled_particle_fixtures import particle_species @@ -318,7 +319,7 @@ def test_tiled_field_snapshot_writes_tile_chunks_without_global_assembly(self): "rho": tile_scalar_field(rho_global, world, tile_shape), "phi": tile_scalar_field(phi_global, world, tile_shape), } - snapshot = openPMD.make_tiled_field_snapshot( + snapshot = async_writer.make_tiled_field_snapshot( field_map, step=3, time=0.6, @@ -355,7 +356,7 @@ def test_tiled_field_snapshot_writes_tile_chunks_without_global_assembly(self): def test_async_field_writer_queue_size_caps_pending_snapshots(self): world = _world() - writer = openPMD.AsyncTiledOpenPMDFieldWriter( + writer = async_writer.AsyncTiledOpenPMDFieldWriter( output_dir="/tmp", filename="fields", world=world, @@ -364,7 +365,7 @@ def test_async_field_writer_queue_size_caps_pending_snapshots(self): guard_cells=world["guard_cells"], queue_size=1, ) - snapshot = openPMD.TiledFieldSnapshot(step=0, time=0.0, fields={}) + snapshot = async_writer.TiledFieldSnapshot(step=0, time=0.0, fields={}) self.assertTrue(writer.enqueue(snapshot, block=False)) self.assertFalse(writer.enqueue(snapshot, block=False)) @@ -372,7 +373,7 @@ def test_async_field_writer_queue_size_caps_pending_snapshots(self): def test_async_field_writer_raises_worker_errors_on_close(self): world = _world() - writer = openPMD.AsyncTiledOpenPMDFieldWriter( + writer = async_writer.AsyncTiledOpenPMDFieldWriter( output_dir="/tmp", filename="fields", world=world, @@ -381,9 +382,9 @@ def test_async_field_writer_raises_worker_errors_on_close(self): guard_cells=world["guard_cells"], queue_size=1, ) - snapshot = openPMD.TiledFieldSnapshot(step=0, time=0.0, fields={}) + snapshot = async_writer.TiledFieldSnapshot(step=0, time=0.0, fields={}) - with patch.object(openPMD, "write_tiled_field_snapshot_openpmd", side_effect=RuntimeError("disk failed")): + 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"): From c99352ff70f7a444324dc6935f5d5cc3a913ca49 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 07:13:58 -0500 Subject: [PATCH 59/83] now using asnyc computation for particles --- PyPIC3D/__main__.py | 45 ++++-- PyPIC3D/diagnostics/async_writer.py | 213 ++++++++++++++++++++++++++++ PyPIC3D/diagnostics/openPMD.py | 209 +++++++++++++++++++++++++++ PyPIC3D/initialization.py | 1 + docs/usage.rst | 8 +- tests/openpmd_test.py | 131 +++++++++++++++++ 6 files changed, 594 insertions(+), 13 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 93246252..7183ee40 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -20,11 +20,9 @@ from PyPIC3D.diagnostics.async_writer import ( create_async_tiled_openpmd_field_writer, + create_async_tiled_openpmd_particle_writer, enqueue_openpmd_field_output, -) - -from PyPIC3D.diagnostics.openPMD import ( - write_openpmd_particles, + enqueue_openpmd_particle_output, ) from PyPIC3D.utils import ( @@ -104,6 +102,7 @@ def loop_with_static_world( initial_energy = e_energy + b_energy + kinetic_energy 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( @@ -113,7 +112,16 @@ def loop_with_static_world( 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") + 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 ############################################################################################################ @@ -157,9 +165,17 @@ def loop_with_static_world( - if plotting_parameters['plot_openpmd_particles']: - write_openpmd_particles(particles, world, constants, os.path.join(output_dir, "data"), plot_num, t, "particles", ".h5", species_config=species_config, species_names=particle_species_names) - # Write the particles in openPMD format + 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) @@ -183,8 +199,17 @@ def loop_with_static_world( loop_error = exc raise finally: - if field_writer is not None: - field_writer.close(raise_errors=loop_error is None) + 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 diff --git a/PyPIC3D/diagnostics/async_writer.py b/PyPIC3D/diagnostics/async_writer.py index 76c9fe63..7e13c967 100644 --- a/PyPIC3D/diagnostics/async_writer.py +++ b/PyPIC3D/diagnostics/async_writer.py @@ -12,6 +12,7 @@ from PyPIC3D.diagnostics.openPMD import ( TiledMeshLayout, write_tiled_field_snapshot_openpmd, + write_tiled_particle_snapshot_openpmd, ) ArrayLike = Any @@ -38,6 +39,26 @@ class TiledFieldSnapshot: 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 @@ -114,6 +135,46 @@ def make_tiled_field_snapshot(field_map, *, step, time): ) +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. @@ -228,6 +289,115 @@ def _writer_loop(self): 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: + 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, @@ -251,6 +421,27 @@ def create_async_tiled_openpmd_field_writer( 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) @@ -260,3 +451,25 @@ def enqueue_openpmd_field_output(field_writer, fields, world, plot_t, t, *, bloc 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/openPMD.py b/PyPIC3D/diagnostics/openPMD.py index d362d6ae..1b306e8c 100644 --- a/PyPIC3D/diagnostics/openPMD.py +++ b/PyPIC3D/diagnostics/openPMD.py @@ -392,6 +392,215 @@ def write_openpmd_particles_to_iteration(iteration, particles, constants, specie 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. diff --git a/PyPIC3D/initialization.py b/PyPIC3D/initialization.py index 3580447d..0ccdfdf5 100644 --- a/PyPIC3D/initialization.py +++ b/PyPIC3D/initialization.py @@ -163,6 +163,7 @@ def default_parameters(): "plot_openpmd_fields": False, "plotting_interval": 10, "openpmd_field_queue_size": 2, + "openpmd_particle_queue_size": 2, "dump_particles": False, "dump_fields": False, } diff --git a/docs/usage.rst b/docs/usage.rst index 211ab52f..56949f57 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -60,6 +60,7 @@ Minimal Working Example plot_openpmd_particles = false plot_openpmd_fields = false openpmd_field_queue_size = 2 + openpmd_particle_queue_size = 2 dump_particles = false dump_fields = false @@ -134,6 +135,7 @@ At ``plotting_interval`` cadence, PyPIC3D writes diagnostics like: - ``total_momentum.txt`` Depending on flags, it also writes matplotlib phase-space arrays and openPMD -files plus ``data/output.toml`` metadata. Runtime openPMD field output uses a -bounded queue; ``openpmd_field_queue_size`` caps the number of pending field -batches so output cannot grow memory without bound. +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/tests/openpmd_test.py b/tests/openpmd_test.py index 31427614..64f208b1 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -198,6 +198,19 @@ def _simulation_parameters(): } +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): world = _world() return particle_species( @@ -390,6 +403,124 @@ def test_async_field_writer_raises_worker_errors_on_close(self): 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 = to_tiled_particles(species, world, _simulation_parameters()) + species_names = tuple(s.get_name() for s in species) + 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_write_openpmd_initial_particles_flattens_tiled_particles_and_preserves_names(self): world = _world() species = [ From c7c418c42bb7d412d78537109198906f540bf206 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Mon, 6 Jul 2026 08:12:21 -0500 Subject: [PATCH 60/83] serializing async writer for particle and fields openPMD tiled. --- PyPIC3D/__main__.py | 9 ----- PyPIC3D/diagnostics/async_writer.py | 44 ++++++++++++-------- demos/two_stream/two_stream.toml | 2 - tests/openpmd_test.py | 63 +++++++++++++++++++++++++++++ 4 files changed, 90 insertions(+), 28 deletions(-) diff --git a/PyPIC3D/__main__.py b/PyPIC3D/__main__.py index 7183ee40..773f3fe1 100644 --- a/PyPIC3D/__main__.py +++ b/PyPIC3D/__main__.py @@ -156,15 +156,6 @@ def loop_with_static_world( 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, species_config=species_config, species_names=particle_species_names, world=world) - - - if particle_writer is not None: enqueue_openpmd_particle_output( particle_writer, diff --git a/PyPIC3D/diagnostics/async_writer.py b/PyPIC3D/diagnostics/async_writer.py index 7e13c967..2d31e6c9 100644 --- a/PyPIC3D/diagnostics/async_writer.py +++ b/PyPIC3D/diagnostics/async_writer.py @@ -23,6 +23,9 @@ SnapshotFieldValue = Union[HostShardList, Tuple[HostShardList, HostShardList, HostShardList]] +_OPENPMD_WRITE_LOCK = threading.Lock() + + @dataclass(frozen=True) class TiledFieldSnapshot: """ @@ -273,14 +276,18 @@ def _writer_loop(self): return if self._error is None: - 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, - ) + # 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 @@ -381,15 +388,18 @@ def _writer_loop(self): return if self._error is None: - 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, - ) + # 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 diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index a6d39bf0..29cf92c9 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -18,10 +18,8 @@ relativistic = true [plotting] plotting_interval = 8 -plot_phasespace = true plot_openpmd_particles = true plot_openpmd_fields = true -dump_particles = true [particle1] name = "electron1" diff --git a/tests/openpmd_test.py b/tests/openpmd_test.py index 64f208b1..da2788fd 100644 --- a/tests/openpmd_test.py +++ b/tests/openpmd_test.py @@ -1,4 +1,6 @@ import unittest +import threading +import time from unittest.mock import patch import jax @@ -521,6 +523,67 @@ def test_async_particle_writer_raises_worker_errors_on_close(self): 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 = [ From 344a173b07becf61d91892ae5c5495848a31c4ab Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 7 Jul 2026 05:12:39 -0500 Subject: [PATCH 61/83] fixing tiled particle initialization --- PyPIC3D/particles/particle_initialization.py | 139 +++++++++++++------ tests/tiled_particle_initialization_test.py | 93 ++++++++++++- 2 files changed, 189 insertions(+), 43 deletions(-) diff --git a/PyPIC3D/particles/particle_initialization.py b/PyPIC3D/particles/particle_initialization.py index c3f429e9..00a1f668 100644 --- a/PyPIC3D/particles/particle_initialization.py +++ b/PyPIC3D/particles/particle_initialization.py @@ -102,6 +102,95 @@ def _particle_tile_indices(x1, x2, x3, world, tile_nx, tile_ny, tile_nz): 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 @@ -308,47 +397,15 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): 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 - 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 = np.zeros((ntx, nty, ntz, n_species), dtype=int) - particle_tile_indices = [] - - for species_index, (x_array, _u_array, active_mask) in enumerate(species_arrays): - tx, ty, tz = _particle_tile_indices(x_array[:, 0], x_array[:, 1], x_array[:, 2], world, tile_nx, tile_ny, tile_nz) - particle_tile_indices.append((tx, ty, tz)) - - for p in range(int(active_mask.shape[0])): - if bool(active_mask[p]): - tile_counts[tx[p], ty[p], tz[p], species_index] += 1 - - max_particles_per_tile = int(np.max(tile_counts)) if tile_counts.size else 0 capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) - max_particles_per_tile = int(math.ceil(max_particles_per_tile * capacity_factor)) # 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 = 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) - - next_slot = np.zeros_like(tile_counts) - - for species_index, (x_array, u_array, active_mask) in enumerate(species_arrays): - tx, ty, tz = particle_tile_indices[species_index] - - for p in range(int(active_mask.shape[0])): - if not bool(active_mask[p]): - continue - tile_index = (tx[p], ty[p], tz[p], species_index) - slot = next_slot[tile_index] - next_slot[tile_index] += 1 - index = tile_index + (slot,) - - x_tiles = x_tiles.at[index].set(x_array[p]) - u_tiles = u_tiles.at[index].set(u_array[p]) - active_tiles = active_tiles.at[index].set(True) + 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]), @@ -358,9 +415,9 @@ def load_particles_from_toml(config, simulation_parameters, world, constants): update_u=jnp.asarray([metadata["update_u"] for metadata in species_metadata], dtype=bool), ) particles = TiledParticles( - x=x_tiles, - u=u_tiles, - active=active_tiles, + 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) diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py index 7cfa18e4..fbe1097e 100644 --- a/tests/tiled_particle_initialization_test.py +++ b/tests/tiled_particle_initialization_test.py @@ -1,4 +1,5 @@ import os +import inspect import tempfile import unittest @@ -7,9 +8,8 @@ import numpy as np from PyPIC3D.particles.particle_initialization import load_particles_from_toml -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles -from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from tiled_particle_fixtures import particle_species, to_tiled_particles jax.config.update("jax_enable_x64", True) @@ -285,6 +285,95 @@ def test_load_particles_from_toml_uses_tile_axes_before_species(self): 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]) From 57eb992ab63deea41938b2f7a516e0624cba6760 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 8 Jul 2026 06:08:09 -0500 Subject: [PATCH 62/83] refactoring ghost cells test --- tests/code_tests/ghost_cells_test.py | 96 ++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tests/code_tests/ghost_cells_test.py 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() From 59a364892a2b459e58fc755eb9f5bee69ae79c18 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 8 Jul 2026 08:38:15 -0500 Subject: [PATCH 63/83] refactoring evolve test --- tests/code_tests/evolve_test.py | 206 ++++++++++++++++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 tests/code_tests/evolve_test.py 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() From b18bd85756b4e3a32f76388fad4c4c28c5f50683 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 8 Jul 2026 08:39:00 -0500 Subject: [PATCH 64/83] removing deprecated tests --- tests/__init__.py | 12 +- tests/boris_test.py | 452 ------ tests/boundaryconditions_test.py | 88 -- tests/direct_deposition_tiled_test.py | 958 ------------- tests/electrostatic_yee_test.py | 164 --- tests/esirkepov_ghost_cells_test.py | 84 -- tests/evolve_test.py | 212 --- tests/grid_and_stencil_test.py | 101 -- tests/initialization_test.py | 228 --- tests/openpmd_test.py | 619 -------- tests/pml_test.py | 590 -------- tests/tiled_electrostatic_test.py | 512 ------- tests/tiled_esirkepov_integration_test.py | 151 -- tests/tiled_esirkepov_test.py | 900 ------------ tests/tiled_fluid_quantities_test.py | 404 ------ tests/tiled_particle_diagnostics_test.py | 160 --- tests/tiled_particle_initialization_test.py | 439 ------ tests/tiled_particle_refresh_test.py | 231 --- tests/tiled_pusher_test.py | 448 ------ tests/tiled_refactor_contract_test.py | 281 ---- tests/tiled_rho_test.py | 363 ----- tests/tiled_yee_integration_test.py | 1411 ------------------- tests/utils_test.py | 499 ------- tests/yee_tiled_test.py | 635 --------- 24 files changed, 1 insertion(+), 9941 deletions(-) delete mode 100644 tests/boris_test.py delete mode 100644 tests/boundaryconditions_test.py delete mode 100644 tests/direct_deposition_tiled_test.py delete mode 100644 tests/electrostatic_yee_test.py delete mode 100644 tests/esirkepov_ghost_cells_test.py delete mode 100644 tests/evolve_test.py delete mode 100644 tests/grid_and_stencil_test.py delete mode 100644 tests/initialization_test.py delete mode 100644 tests/openpmd_test.py delete mode 100644 tests/pml_test.py delete mode 100644 tests/tiled_electrostatic_test.py delete mode 100644 tests/tiled_esirkepov_integration_test.py delete mode 100644 tests/tiled_esirkepov_test.py delete mode 100644 tests/tiled_fluid_quantities_test.py delete mode 100644 tests/tiled_particle_diagnostics_test.py delete mode 100644 tests/tiled_particle_initialization_test.py delete mode 100644 tests/tiled_particle_refresh_test.py delete mode 100644 tests/tiled_pusher_test.py delete mode 100644 tests/tiled_refactor_contract_test.py delete mode 100644 tests/tiled_rho_test.py delete mode 100644 tests/tiled_yee_integration_test.py delete mode 100644 tests/utils_test.py delete mode 100644 tests/yee_tiled_test.py diff --git a/tests/__init__.py b/tests/__init__.py index 6e6e82c3..cc755aca 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,11 +1 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os - -from . import boris_test -from . import boundaryconditions_test -from . import electrostatic_yee_test -from . import initialization_test -from . import utils_test +"""Shared unittest package for PyPIC3D tests.""" diff --git a/tests/boris_test.py b/tests/boris_test.py deleted file mode 100644 index b5e45de4..00000000 --- a/tests/boris_test.py +++ /dev/null @@ -1,452 +0,0 @@ -import unittest -import jax -import jax.numpy as jnp -import sys -import os -from functools import partial - - -from PyPIC3D.pusher.boris import boris_single_particle, interpolate_field_to_particles -from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle -from PyPIC3D.utils import mae, convergence_test - -jax.config.update("jax_enable_x64", True) - -class TestBorisMethods(unittest.TestCase): - - def setUp(self): - self.q = 1.0 - self.m = 1.0 - self.x = jnp.array([0.0]) - self.y = jnp.array([0.0]) - self.z = jnp.array([0.0]) - # single particle at the origin - - self.vx = jnp.array([1.0]) - self.vy = jnp.array([0.0]) - self.vz = jnp.array([0.0]) - # single particle moving in the x direction - - self.Ex = jnp.zeros( (5, 5, 5) ) - self.Ey = jnp.zeros( (5, 5, 5) ) - self.Ez = jnp.zeros( (5, 5, 5) ) - self.Bx = jnp.zeros( (5, 5, 5) ) - self.By = jnp.zeros( (5, 5, 5) ) - self.Bz = jnp.ones( (5, 5, 5) ) - # grid of 5x5x5 with a uniform magnetic field in the z direction - - self.grid = jnp.arange(-1/2, 1/2, (1/5)), jnp.arange(-1/2, 1/2, (1/5)), jnp.arange(-1/2, 1/2, (1/5)) - self.staggered_grid = jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)), jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)), jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)) - self.dt = 0.1 - # grid and staggered grid for a 5x5x5 grid with a spacing of 1/5 and a timestep of 0.1 - - def test_boris(self): - - efield_atx = interpolate_field_to_particles(self.Ex, self.x, self.y, self.z, self.grid, shape_factor=1) - efield_aty = interpolate_field_to_particles(self.Ey, self.x, self.y, self.z, self.grid, shape_factor=1) - efield_atz = interpolate_field_to_particles(self.Ez, self.x, self.y, self.z, self.grid, shape_factor=1) - # calculate the electric field at the particle positions - bfield_atx = interpolate_field_to_particles(self.Bx, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) - bfield_aty = interpolate_field_to_particles(self.By, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) - bfield_atz = interpolate_field_to_particles(self.Bz, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) - - boris_vmap = jax.vmap(boris_single_particle, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None)) - newvx, newvy, newvz = boris_vmap(self.vx, self.vy, self.vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, self.q, self.m, self.dt, None) - - self.assertIsInstance(newvx, jnp.ndarray) - self.assertIsInstance(newvy, jnp.ndarray) - self.assertIsInstance(newvz, jnp.ndarray) - # make sure the velocities are jax arrays - self.assertTrue( jnp.allclose(jnp.abs(newvx), 1.0, rtol = 1e-2)) - # make sure the x velocity is unchanged - self.assertLess( newvy, 0.1 ) - # make sure the y velocity is unchanged - self.assertLess( newvz, 0.1 ) - # make sure the z velocity is 0.0 from the magnetic field - - def test_interpolate_field_to_particles_with_ghost_cells_on_reduced_axes(self): - field = jnp.full((6, 3, 3), 7.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, 6), - jnp.array([-1.0, 0.0, 1.0]), - jnp.array([-1.0, 0.0, 1.0]), - ) - - interpolated = interpolate_field_to_particles( - field, - x, - y, - z, - grid, - shape_factor=2, - ghost_cells=True, - ) - - 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): - vx, vy, vz = 1.0, 0.0, 0.0 - x, y, z = 0.0, 0.0, 0.0 - - q = 1.0 - m = 1.0 - E = jnp.array([0.0, 0.0, 0.0]) - B = jnp.array([0.0, 1.0, 0.0]) - dt = 0.1 - n_steps = 5000 - - xs = [] - ys = [] - zs = [] - - for i in range(n_steps): - vx, vy, vz = boris_single_particle(vx, vy, vz, E[0], E[1], E[2], B[0], B[1], B[2], q, m, dt, None) - x += vx * dt - y += vy * dt - z += vz * dt - - xs.append(x) - ys.append(y) - zs.append(z) - - - def measure_xz_radius(xs, zs): - """ - Measure the radius of the XZ cut by calculating the distance of each point - in the XZ plane from the origin and returning the average and maximum radius. - - Parameters: - xs (list): List of x-coordinates. - zs (list): List of z-coordinates. - - Returns: - tuple: A tuple containing the average radius and maximum radius. - """ - xs = jnp.array(xs) - zs = jnp.array(zs) - # Calculate the distance of each point from the origin in the XZ plane - radii = [jnp.sqrt(x**2 + z**2) for x, z in zip(xs, zs)] - - # Compute the average and maximum radius - avg_radius = jnp.mean(jnp.asarray(radii)) - - return avg_radius - - avg_radius = measure_xz_radius(xs, zs) - #print(f"Average radius: {avg_radius}") - self.assertTrue( jnp.isclose(avg_radius, 1.28, atol=0.5) ) - - def test_higuera_cary_zero_field_preserves_velocity(self): - vx, vy, vz = higuera_cary_single_particle( - 0.2, -0.1, 0.05, - 0.0, 0.0, 0.0, - 0.0, 0.0, 0.0, - 1.0, 1.0, 0.1, - {"C": 10.0}, - ) - # zero fields should not change the particle velocity - self.assertTrue(jnp.allclose(jnp.array([vx, vy, vz]), jnp.array([0.2, -0.1, 0.05]))) - - def test_higuera_cary_magnetic_field_preserves_speed(self): - vx0, vy0, vz0 = 0.2, 0.0, 0.0 - vx, vy, vz = higuera_cary_single_particle( - vx0, vy0, vz0, - 0.0, 0.0, 0.0, - 0.0, 0.0, 1.0, - 1.0, 1.0, 0.1, - {"C": 10.0}, - ) - # magnetic rotation should conserve speed - old_speed = jnp.sqrt(vx0**2 + vy0**2 + vz0**2) - new_speed = jnp.sqrt(vx**2 + vy**2 + vz**2) - self.assertTrue(jnp.allclose(new_speed, old_speed, rtol=1e-12, atol=1e-12)) - - def test_higuera_cary_electric_field_accelerates_along_field(self): - vx, vy, vz = higuera_cary_single_particle( - 0.0, 0.0, 0.0, - 1.0, 0.0, 0.0, - 0.0, 0.0, 0.0, - 1.0, 1.0, 0.1, - {"C": 10.0}, - ) - # electric field in x should accelerate the particle in x only - self.assertGreater(vx, 0.0) - self.assertTrue(jnp.allclose(jnp.array([vy, vz]), jnp.array([0.0, 0.0]))) - - def test_higuera_cary_vectorizes_with_per_particle_charge_and_mass(self): - higuera_cary_vmap = jax.vmap( - higuera_cary_single_particle, - in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), - ) - newvx, newvy, newvz = higuera_cary_vmap( - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([1.0, 1.0]), - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([0.0, 0.0]), - jnp.array([1.0, 2.0]), - jnp.array([1.0, 1.0]), - 0.1, - {"C": 10.0}, - ) - # per-particle charge and mass should work through the vectorized path - self.assertEqual(newvx.shape, (2,)) - self.assertGreater(newvx[1], newvx[0]) - self.assertTrue(jnp.allclose(newvy, jnp.zeros(2))) - self.assertTrue(jnp.allclose(newvz, jnp.zeros(2))) - - - def test_trilinear_convergence(self): - - # Run convergence tests for each interpolation method - ################### BASIC TRIGONOMETRIC TEST ############################# - slope = convergence_test(partial(interpolation_wave_test, shape_factor=1)) - # measure the convergence rate of the trilinear interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - ################### POLYNOMIAL TEST ####################################### - slope = convergence_test(partial(interpolation_polynomial_test, shape_factor=1)) - # measure the convergence rate of the trilinear interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - ################### HIGH-FREQUENCY TEST ################################### - slope = convergence_test(partial(interpolation_oscillatory_test, shape_factor=1)) - # measure the convergence rate of the trilinear interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - def test_quadratic_convergence(self): - - ################### BASIC TRIGONOMETRIC TEST ############################### - slope = convergence_test(partial(interpolation_wave_test, shape_factor=2)) - # measure the convergence rate of the second-order shape interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - ################### POLYNOMIAL TEST ######################################## - slope = convergence_test(partial(interpolation_polynomial_test, shape_factor=2)) - # measure the convergence rate of the second-order shape interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - ################### HIGH-FREQUENCY TEST #################################### - slope = convergence_test(partial(interpolation_oscillatory_test, shape_factor=2)) - # measure the convergence rate of the second-order shape interpolation - self.assertTrue(slope > 1.9) - # assert that the order of the error is at least 2nd order - ############################################################################ - - -def interpolation_wave_test(nx, shape_factor): - # Define a symmetric domain - x_wind = 2.0 * jnp.pi - y_wind = 2.0 * jnp.pi - z_wind = 2.0 * jnp.pi - - # Create uniform grid - x_grid = jnp.linspace(0, x_wind, nx) - y_grid = jnp.linspace(0, y_wind, nx) - z_grid = jnp.linspace(0, z_wind, nx) - - dx = x_wind / (nx - 1) - dy = y_wind / (nx - 1) - dz = z_wind / (nx - 1) - - # Create meshgrid for field evaluation - X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') - - # Define analytical test function: smooth trigonometric function - # f(x,y,z) = sin(x) * cos(y) * sin(z) - analytical_field = jnp.sin(X) * jnp.cos(Y) * jnp.sin(Z) - - grid = (x_grid, y_grid, z_grid) - - # Create test points that are offset from grid points - # This tests the interpolation accuracy between grid points - n_test = nx - x_test = jnp.linspace(dx/3, x_wind - dx/3, n_test) - y_test = jnp.linspace(dy/3, y_wind - dy/3, n_test) - z_test = jnp.linspace(dz/3, z_wind - dz/3, n_test) - - X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') - - # Flatten for vectorized interpolation - x_test_flat = X_test.flatten() - y_test_flat = Y_test.flatten() - z_test_flat = Z_test.flatten() - - # Compute analytical values at test points - analytical_values = jnp.sin(x_test_flat) * jnp.cos(y_test_flat) * jnp.sin(z_test_flat) - - # Interpolate values at test points - interpolated_values = interpolate_field_to_particles( - analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor - ) - - # Compute mean squared error - error = mae(interpolated_values, analytical_values) - - return error, dx - - -def interpolation_polynomial_test(nx, shape_factor): - - # Define domain - x_wind = 1.0 - y_wind = 1.0 - z_wind = 1.0 - - # Create uniform grid - x_grid = jnp.linspace(0, x_wind, nx) - y_grid = jnp.linspace(0, y_wind, nx) - z_grid = jnp.linspace(0, z_wind, nx) - - dx = x_wind / (nx - 1) - # create grid spacing - - # Create meshgrid - X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') - - # Use a polynomial test function: f(x,y,z) = x³y + y³z + z³x - # This has continuous derivatives and tests interpolation of curved surfaces - analytical_field = X**3 * Y + Y**3 * Z + Z**3 * X - - grid = (x_grid, y_grid, z_grid) - - # Test at mid-points between grid points (maximum interpolation error) - n_test = max(4, nx) # Ensure we have enough test points - x_test = jnp.linspace(dx/2, x_wind - dx/2, n_test) - y_test = jnp.linspace(dx/2, y_wind - dx/2, n_test) - z_test = jnp.linspace(dx/2, z_wind - dx/2, n_test) - - X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') - - # Flatten - x_test_flat = X_test.flatten() - y_test_flat = Y_test.flatten() - z_test_flat = Z_test.flatten() - - # Analytical values - analytical_values = (x_test_flat**3 * y_test_flat + - y_test_flat**3 * z_test_flat + - z_test_flat**3 * x_test_flat) - - # Interpolated values - interpolated_values = interpolate_field_to_particles( - analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor - ) - - # Compute error - error = mae(interpolated_values, analytical_values) - - return error, dx - - -def interpolation_oscillatory_test(nx, shape_factor): - - # Define domain - x_wind = jnp.pi - y_wind = jnp.pi - z_wind = jnp.pi - - # Create grid - x_grid = jnp.linspace(0, x_wind, nx) - y_grid = jnp.linspace(0, y_wind, nx) - z_grid = jnp.linspace(0, z_wind, nx) - - dx = x_wind / (nx - 1) - # create grid spacing - - # Create meshgrid - X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') - - # High-frequency test function: f(x,y,z) = cos(2x) * sin(3y) * cos(4z) - analytical_field = jnp.cos(2*X) * jnp.sin(3*Y) * jnp.cos(4*Z) - - grid = (x_grid, y_grid, z_grid) - - # Test points slightly offset from grid - n_test = max(6, nx) - x_test = jnp.linspace(dx/4, x_wind - dx/4, n_test) - y_test = jnp.linspace(dx/4, y_wind - dx/4, n_test) - z_test = jnp.linspace(dx/4, z_wind - dx/4, n_test) - - X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') - - # Flatten - x_test_flat = X_test.flatten() - y_test_flat = Y_test.flatten() - z_test_flat = Z_test.flatten() - - # Analytical values - analytical_values = (jnp.cos(2*x_test_flat) * - jnp.sin(3*y_test_flat) * - jnp.cos(4*z_test_flat)) - - # Interpolated values - interpolated_values = interpolate_field_to_particles( - analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor - ) - - # Compute error - error = mae(interpolated_values, analytical_values) - - return error, dx - -if __name__ == '__main__': - unittest.main() diff --git a/tests/boundaryconditions_test.py b/tests/boundaryconditions_test.py deleted file mode 100644 index d0b4b435..00000000 --- a/tests/boundaryconditions_test.py +++ /dev/null @@ -1,88 +0,0 @@ -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.boundaryconditions import ( - apply_supergaussian_boundary_condition, apply_zero_boundary_condition, -) -from PyPIC3D.boundary_conditions import ghost_cells - -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 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): - 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) - - result = ghost_cells.update_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) - - self.assertTrue(jnp.all(result[0, 0, 0, -1, 1:3, 1:3] == 2.0)) - self.assertTrue(jnp.all(result[1, 0, 0, 0, 1:3, 1:3] == 1.0)) - - def test_fold_tiled_ghost_cells_periodic_adds_to_owner_tile(self): - field_tiles = jnp.zeros((2, 1, 1, 4, 4, 4)) - field_tiles = field_tiles.at[0, 0, 0, -1, 2, 2].set(3.0) - field_tiles = field_tiles.at[1, 0, 0, 0, 2, 2].set(5.0) - - result = ghost_cells.fold_tiled_ghost_cells(field_tiles, self.world, self.g, self.tile_shape) - - self.assertAlmostEqual(float(result[1, 0, 0, 1, 2, 2]), 3.0) - self.assertAlmostEqual(float(result[0, 0, 0, -2, 2, 2]), 5.0) - self.assertEqual(float(result[0, 0, 0, -1, 2, 2]), 0.0) - self.assertEqual(float(result[1, 0, 0, 0, 2, 2]), 0.0) - - def test_apply_tiled_conducting_bc_zeros_global_tangential_faces(self): - world = { - "tile_shape": self.tile_shape, - "boundary_conditions": {"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING}, - } - E = tuple(jnp.ones((1, 1, 1, 4, 4, 4)) for _ in range(3)) - - Ex, Ey, Ez = ghost_cells.apply_tiled_conducting_bc(E, world, self.g) - - 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)) - - 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)) - - result = ghost_cells.apply_tiled_scalar_conducting_bc(field, self.world, self.g) - - self.assertTrue(jnp.allclose(result, field)) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/direct_deposition_tiled_test.py b/tests/direct_deposition_tiled_test.py deleted file mode 100644 index 3379ea3b..00000000 --- a/tests/direct_deposition_tiled_test.py +++ /dev/null @@ -1,958 +0,0 @@ -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_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.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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) - - -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 - - -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 TestDirectDepositionTiled(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)) - - def _tile_shape(self, simulation_parameters): - return ( - simulation_parameters["particle_tile_nx"], - simulation_parameters["particle_tile_ny"], - simulation_parameters["particle_tile_nz"], - ) - - 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 _one_tile_parameters(self, world): - return { - "particle_tile_nx": world["Nx"], - "particle_tile_ny": world["Ny"], - "particle_tile_nz": world["Nz"], - } - - 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. - """ - - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - tiled_particles = tiled_particles._replace( - x=tiled_particles.x - 0.5 * tiled_particles.u * world["dt"] - ) - - centered_particles, overflow = refresh_tiled_particle_tiles( - tiled_particles, - world, - self._tile_shape(simulation_parameters), - ) - self.assertFalse(bool(overflow)) - - return centered_particles, species_config - - def _assembled_tiled_current(self, particles, world, simulation_parameters, constants, filter="none"): - tile_shape = self._tile_shape(simulation_parameters) - world = self._world_with_tiled_grids(world, tile_shape) - tiled_particles, species_config = 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) - - return J_tiles, J_from_tiles - - def _compare_tiled_to_one_tile(self, particles, world, simulation_parameters, filter="none", alpha=1.0): - constants = {"C": 3.0e8, "alpha": alpha} - J_tiles, J_from_tiles = self._assembled_tiled_current( - particles, world, simulation_parameters, constants, filter=filter - ) - _, J_reference = self._assembled_tiled_current( - particles, world, self._one_tile_parameters(world), constants, filter=filter - ) - - 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)) - - def test_public_J_from_rhov_rejects_flat_particles(self): - world = self._build_world() - constants = {"C": 3.0e8, "alpha": 1.0} - simulation_parameters = self._one_tile_parameters(world) - world = self._world_with_tiled_grids(world, self._tile_shape(simulation_parameters)) - species = particle_species( - name="flat direct current", - N_particles=1, - charge=-1.0, - mass=1.0, - T=1.0, - x1=jnp.array([0.0]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - v1=jnp.array([0.0]), - v2=jnp.array([0.0]), - v3=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"], - dt=world["dt"], - ) - _, species_config = to_tiled_particles([species], world, simulation_parameters) - - with self.assertRaises(AttributeError): - J_from_rhov([species], species_config, self._empty_J_tiles(world), constants, world, filter="none") - - 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, - } - - species = particle_species( - name="quadratic current", - N_particles=6, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.55, -0.52, -0.03, 0.49, 0.55, 1.45]), - x2=jnp.array([-1.10, -0.55, -0.03, 0.02, 0.52, 1.05]), - x3=jnp.array([-0.70, -0.04, 0.03, 0.31, 0.49, 0.72]), - v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16]), - v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12]), - v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - self._compare_tiled_to_one_tile([species], world, simulation_parameters) - - 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, - } - - species = particle_species( - name="quadratic reduced current", - N_particles=8, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.95, -1.51, -1.02, -0.48, 0.02, 0.47, 1.04, 1.88]), - x2=jnp.zeros(8), - x3=jnp.zeros(8), - v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16, 0.09, -0.13]), - v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12, -0.15, 0.05]), - v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19, 0.16, -0.07]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - self._compare_tiled_to_one_tile([species], world, simulation_parameters) - - 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, - } - species = particle_species( - name="bilinear quadratic reduced current", - N_particles=8, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.95, -1.51, -1.02, -0.48, 0.02, 0.47, 1.04, 1.88]), - x2=jnp.zeros(8), - x3=jnp.zeros(8), - v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16, 0.09, -0.13]), - v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12, -0.15, 0.05]), - v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19, 0.16, -0.07]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="bilinear") - - 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-12, atol=1.0e-12)) - - 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)) - - 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)) - - 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)) - - 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)) - - 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-12, atol=1.0e-12)) - - 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} - species = particle_species( - name="local current tiles", - N_particles=4, - charge=1.0, - mass=1.0, - weight=1.0, - 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"], - ) - tiled_particles, species_config = self._centered_tiled_particles([species], 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( - [species], 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-12, atol=1.0e-12)) - - 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, - } - - electrons = particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=1.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"], - ) - ions = particle_species( - name="ions", - N_particles=3, - charge=2.0, - mass=4.0, - weight=0.25, - T=1.0, - x1=jnp.array([-1.65, 0.15, 1.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"], - ) - - self._compare_tiled_to_one_tile([electrons, ions], world, simulation_parameters) - - 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} - species = particle_species( - name="public tiled direct current", - N_particles=4, - charge=-1.0, - mass=1.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"], - ) - tiled_particles, species_config = self._centered_tiled_particles([species], 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( - [species], 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-12, atol=1.0e-12)) - - 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} - species = particle_species( - name="digital filtered current", - N_particles=4, - charge=-1.0, - mass=1.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"], - ) - self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="digital", alpha=0.6) - - 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} - species = particle_species( - name="bilinear filtered current", - N_particles=6, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.55, -0.52, -0.03, 0.49, 0.55, 1.45]), - x2=jnp.array([-1.10, -0.55, -0.03, 0.02, 0.52, 1.05]), - x3=jnp.array([-0.70, -0.04, 0.03, 0.31, 0.49, 0.72]), - v1=jnp.array([0.18, -0.11, 0.07, -0.04, 0.21, -0.16]), - v2=jnp.array([0.03, 0.17, -0.22, 0.19, -0.08, 0.12]), - v3=jnp.array([-0.06, 0.24, 0.11, -0.14, 0.05, -0.19]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - self._compare_tiled_to_one_tile([species], world, simulation_parameters, filter="bilinear") - - 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, - } - species = particle_species( - name="partly active", - N_particles=4, - charge=1.0, - mass=1.0, - weight=1.0, - 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=jnp.array([True, False, True, False]), - ) - - self._compare_tiled_to_one_tile([species], world, simulation_parameters) - - 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, - } - species = particle_species( - name="crossing", - N_particles=2, - charge=1.0, - mass=1.0, - weight=1.0, - T=1.0, - x1=jnp.array([world["x_wind"] / 2 + 0.1 * world["dx"], -world["x_wind"] / 2 - 0.2 * world["dx"]]), - x2=jnp.array([0.0, 0.0]), - x3=jnp.array([0.0, 0.0]), - v1=jnp.array([0.5, -0.25]), - v2=jnp.array([0.1, -0.2]), - v3=jnp.array([0.0, 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"], - ) - - self._compare_tiled_to_one_tile([species], 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, - } - species = particle_species( - name="conducting wall deposits", - N_particles=3, - charge=1.0, - mass=1.0, - weight=1.0, - T=1.0, - x1=jnp.array([ - -world["x_wind"] / 2 + 0.1 * world["dx"], - world["x_wind"] / 2 - 0.1 * world["dx"], - 0.0, - ]), - x2=jnp.array([ - -world["y_wind"] / 2 + 0.1 * world["dy"], - 0.0, - world["y_wind"] / 2 - 0.1 * world["dy"], - ]), - x3=jnp.array([ - 0.0, - -world["z_wind"] / 2 + 0.1 * world["dz"], - world["z_wind"] / 2 - 0.1 * world["dz"], - ]), - v1=jnp.array([0.5, -0.25, 0.15]), - v2=jnp.array([0.1, -0.2, 0.3]), - v3=jnp.array([-0.15, 0.35, -0.1]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - self._compare_tiled_to_one_tile([species], 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, - } - species = particle_species( - name="mixed wall deposits", - N_particles=4, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=jnp.array([ - -world["x_wind"] / 2 - 0.1 * world["dx"], - world["x_wind"] / 2 + 0.2 * world["dx"], - -0.5, - 0.5, - ]), - x2=jnp.array([ - -world["y_wind"] / 2 + 0.1 * world["dy"], - world["y_wind"] / 2 - 0.2 * world["dy"], - -0.25, - 0.25, - ]), - x3=jnp.array([ - 0.0, - 0.25, - -world["z_wind"] / 2 - 0.1 * world["dz"], - world["z_wind"] / 2 + 0.2 * world["dz"], - ]), - 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"], - ) - - self._compare_tiled_to_one_tile([species], 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, - } - species = particle_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.array([0.0, 0.0, 0.0]), - x3=jnp.array([0.0, 0.0, 0.0]), - v1=jnp.array([0.2, -0.1, 0.05]), - v2=jnp.array([0.3, 0.15, -0.2]), - v3=jnp.array([-0.05, 0.25, 0.1]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - self._compare_tiled_to_one_tile([species], world, simulation_parameters) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/electrostatic_yee_test.py b/tests/electrostatic_yee_test.py deleted file mode 100644 index 3879d8ca..00000000 --- a/tests/electrostatic_yee_test.py +++ /dev/null @@ -1,164 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.tests.tiled_particle_fixtures import particle_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 = [ - 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", - 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)) - - -if __name__ == "__main__": - unittest.main() 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 c922221e..00000000 --- a/tests/evolve_test.py +++ /dev/null @@ -1,212 +0,0 @@ -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.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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 - - -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") - - 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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} - - 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) - - 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"], - dt=world["dt"], - ) - ] - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - 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", - ) - - 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, :]] - - self.assertIsNone(pml_state) - self.assertFalse(bool(overflow)) - 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(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, - "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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.0} - - 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), - ) - 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="periodic", - dt=world["dt"], - ) - ] - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - fields = ( - E, - B, - J, - rho, - phi, - external_fields, - None, - jnp.asarray(False), - ) - - tiled_particles, _ = time_loop_electrodynamic( - tiled_particles, - species_config, - fields, - world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - active = tiled_particles.active[:, :, :, 0, :].reshape(-1) - x = tiled_particles.x[:, :, :, 0, :, 0].reshape(-1) - - self.assertEqual(x.shape[0], 1) - self.assertTrue(jnp.array_equal(active, jnp.array([False]))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/grid_and_stencil_test.py b/tests/grid_and_stencil_test.py deleted file mode 100644 index c31ab49f..00000000 --- a/tests/grid_and_stencil_test.py +++ /dev/null @@ -1,101 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.boundary_conditions.grid_and_stencil import ( - BC_CONDUCTING, - BC_PERIODIC, - axis_has_active_cells, - build_axis_stencil_points, - build_collocated_axis, - build_staggered_axis, - collapse_axis_stencil, - compute_particle_anchor, - inactive_axis_index, - particle_axis_offset, - prepare_particle_axis_stencil, - uniform_axis_spacing, - wrap_periodic_position, -) - -jax.config.update("jax_enable_x64", True) - - -class TestGhostCellHelpers(unittest.TestCase): - - def test_wrap_periodic_position(self): - wind = 4.0 - x = jnp.array([-2.5, -2.0, 1.75, 2.0, 2.25, 6.5]) - 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)) - - 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) - - def test_axis_spacing_and_anchor_helpers(self): - axis = build_collocated_axis(-1.0, 0.5, 4) - position = jnp.array([0.1, 0.6]) - anchor = compute_particle_anchor(position, axis, shape_factor=1) - offset = particle_axis_offset(position, anchor, axis) - - 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]))) - - def test_build_axis_stencil_points_periodic_and_conducting(self): - anchor = jnp.array([0, 5]) - offsets = jnp.array([-1, 0, 1]) - - periodic = build_axis_stencil_points(anchor, 6, BC_PERIODIC, offsets) - conducting = build_axis_stencil_points(anchor, 6, BC_CONDUCTING, offsets) - - self.assertTrue(jnp.array_equal(periodic[:, 0], jnp.array([5, 0, 1]))) - 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]))) - - def test_prepare_particle_axis_stencil_wraps_periodic_indices(self): - axis = build_collocated_axis(-1.0, 0.5, 4) - position = jnp.array([1.1]) - _, anchor, offset, points = prepare_particle_axis_stencil( - position, - axis, - axis_size=6, - shape_factor=1, - bc=BC_PERIODIC, - ghost_cells=True, - ) - - 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]))) - - def test_collapse_axis_stencil_for_inactive_axis(self): - points = jnp.array([[0, 1], [1, 1], [2, 1]]) - weights = jnp.array([[0.2, 0.1], [0.5, 0.6], [0.3, 0.3]]) - - collapsed_points, collapsed_weights = collapse_axis_stencil( - points, - weights, - axis_size=3, - ghost_cells=True, - ) - - self.assertTrue(jnp.array_equal(collapsed_points, jnp.array([[1, 1]]))) - self.assertTrue(jnp.allclose(collapsed_weights, jnp.array([[1.0, 1.0]]))) - - def test_build_axis_helpers_include_ghost_cells(self): - collocated = build_collocated_axis(-1.0, 0.5, 4) - staggered = build_staggered_axis(-1.0, 0.5, 4) - - 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]))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/initialization_test.py b/tests/initialization_test.py deleted file mode 100644 index 5f7f0f4c..00000000 --- a/tests/initialization_test.py +++ /dev/null @@ -1,228 +0,0 @@ -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.func if hasattr(loop, "func") else 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)) - - 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) - - 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)) - - 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) - - 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/openpmd_test.py b/tests/openpmd_test.py deleted file mode 100644 index da2788fd..00000000 --- a/tests/openpmd_test.py +++ /dev/null @@ -1,619 +0,0 @@ -import unittest -import threading -import time -from unittest.mock import patch - -import jax -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.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles - - -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 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 _simulation_parameters(): - return { - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - } - - -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): - world = _world() - return particle_species( - name=name, - N_particles=x1.shape[0], - charge=charge, - mass=mass, - weight=weight, - T=0.0, - x1=x1, - x2=jnp.zeros_like(x1), - x3=jnp.zeros_like(x1), - v1=jnp.ones_like(x1) * 0.1, - v2=jnp.zeros_like(x1), - v3=jnp.zeros_like(x1), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - active_mask=jnp.ones(x1.shape[0], dtype=bool), - dt=world["dt"], - ) - - -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)) - - 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 = to_tiled_particles(species, world, _simulation_parameters()) - species_names = tuple(s.get_name() for s in species) - 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 = to_tiled_particles(species, world, _simulation_parameters()) - species_names = tuple(s.get_name() for s in species) - 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/pml_test.py b/tests/pml_test.py deleted file mode 100644 index de4b6859..00000000 --- a/tests/pml_test.py +++ /dev/null @@ -1,590 +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, - 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) - - 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 - - -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/tiled_electrostatic_test.py b/tests/tiled_electrostatic_test.py deleted file mode 100644 index 68eb6863..00000000 --- a/tests/tiled_electrostatic_test.py +++ /dev/null @@ -1,512 +0,0 @@ -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_PERIODIC -from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.evolve import time_loop_electrostatic -from PyPIC3D.initialization import initialize_simulation -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import TiledParticles -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field -from PyPIC3D.utilities.grids import build_collocated_grid, 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 TestTiledElectrostatic(unittest.TestCase): - def _build_world(self, shape_factor=1): - world = { - "Nx": 8, - "Ny": 1, - "Nz": 1, - "dx": 0.5, - "dy": 1.0, - "dz": 1.0, - "dt": 0.01, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "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} - world["tile_shape"] = (2, 1, 1) - world = self._world_with_tiled_grids(world, world["tile_shape"]) - return world - - 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 _simulation_parameters(self): - return { - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - } - - 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 _one_tile_shape(self, world): - return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) - - def _empty_fields(self, world): - shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - zero = jnp.zeros(shape) - E = (zero, zero, zero) - B = (zero, zero, zero) - J = (zero, zero, zero) - external_fields = (E, B) - return E, B, J, zero, zero, external_fields - - def _species(self, world): - return particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - def _neutral_species(self, world): - electrons = self._species(world) - ions = particle_species( - name="ions", - N_particles=4, - charge=1.0, - mass=4.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.25, -0.25, 0.75, 1.75]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([-0.03, 0.04, -0.02, 0.01]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - return [electrons, ions] - - def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): - current_slots = tiled_particles.active.shape[-1] - if current_slots >= min_slots: - return tiled_particles - - pad_width = [(0, 0)] * tiled_particles.x.ndim - pad_width[-2] = (0, min_slots - current_slots) - x = jnp.pad(tiled_particles.x, pad_width) - u = jnp.pad(tiled_particles.u, pad_width) - - scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim - scalar_pad_width[-1] = (0, min_slots - current_slots) - - return TiledParticles( - x=x, - u=u, - active=jnp.pad(tiled_particles.active, scalar_pad_width), - ) - - def _build_tiled_state(self, particles, fields, world, tile_shape, min_slots=None): - tiled_particles, species_config = to_tiled_particles( - particles, - world, - self._simulation_parameters_for_tile_shape(tile_shape), - ) - if min_slots is not None: - tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots) - - E, B, J, rho, phi, external_fields = fields - g = int(world["guard_cells"]) - tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=g), - tile_vector_field(B, world, tile_shape, num_guard_cells=g), - tile_vector_field(J, world, tile_shape, num_guard_cells=g), - tile_scalar_field(rho, world, tile_shape, num_guard_cells=g), - tile_scalar_field(phi, world, tile_shape, num_guard_cells=g), - ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), - ), - None, - jnp.asarray(False), - ) - - return tiled_particles, species_config, tiled_fields - - def _long_two_stream_species(self, world): - x0 = jnp.linspace( - -0.5 * world["x_wind"], - 0.5 * world["x_wind"], - world["Nx"], - endpoint=False, - ) + 0.5 * world["dx"] - x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) - zeros = jnp.zeros_like(x0) - beam_speed = 0.08 - - species_kwargs = { - "N_particles": world["Nx"], - "mass": 1.0, - "weight": 1.0, - "T": 0.0, - "x1": x0, - "x2": zeros, - "x3": zeros, - "v2": zeros, - "v3": zeros, - "xwind": world["x_wind"], - "ywind": world["y_wind"], - "zwind": world["z_wind"], - "dx": world["dx"], - "dy": world["dy"], - "dz": world["dz"], - "dt": world["dt"], - } - return [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - - def _long_two_stream_state(self): - world = { - "Nx": 16, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.002, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "guard_cells": 2, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_collocated_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = self._one_tile_shape(world) - world = self._world_with_tiled_grids(world, tile_shape) - Nt = 1500 - - E, B, J, rho, phi, external_fields = self._empty_fields(world) - fields = (E, B, J, rho, phi, external_fields) - tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._long_two_stream_species(world), - fields, - world, - tile_shape, - ) - - return tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt - - def test_scalar_tiles_round_trip_global_rho_and_phi(self): - world = self._build_world() - tile_shape = world["tile_shape"] - scalar = jnp.arange((world["Nx"] + 2) * (world["Ny"] + 2) * (world["Nz"] + 2), dtype=float) - scalar = scalar.reshape((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) - - scalar_tiles = tile_scalar_field(scalar, world, tile_shape) - scalar_from_tiles = assemble_tiled_scalar_field(scalar_tiles, world, tile_shape) - - self.assertTrue(jnp.allclose(scalar_from_tiles[1:-1, 1:-1, 1:-1], scalar[1:-1, 1:-1, 1:-1])) - - def test_tiled_rho_assembles_to_global_rho(self): - world = self._build_world(shape_factor=1) - constants = {"alpha": 1.0} - species = self._species(world) - tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) - _, _, _, rho, _, _ = self._empty_fields(world) - rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) - - tiled_rho = compute_rho( - tiled_particles, - species_config, - rho_tiles, - constants, - world, - ) - assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - - reference_world = self._world_with_tiled_grids(world, self._one_tile_shape(world)) - reference_particles, reference_species_config = to_tiled_particles( - [self._species(reference_world)], - reference_world, - self._simulation_parameters_for_tile_shape(reference_world["tile_shape"]), - ) - reference_rho_tiles = tile_scalar_field(rho, reference_world, reference_world["tile_shape"], num_guard_cells=int(reference_world["guard_cells"])) - reference_rho_tiles = compute_rho( - reference_particles, - reference_species_config, - reference_rho_tiles, - constants, - reference_world, - ) - reference_rho = assemble_tiled_scalar_field( - reference_rho_tiles, - reference_world, - reference_world["tile_shape"], - num_guard_cells=int(reference_world["guard_cells"]), - ) - - self.assertTrue( - jnp.allclose( - assembled_rho[1:-1, 1:-1, 1:-1], - reference_rho[1:-1, 1:-1, 1:-1], - rtol=1.0e-12, - atol=1.0e-12, - ) - ) - - def test_tiled_electrostatic_step_uses_single_tile_state(self): - world = self._build_world(shape_factor=1) - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - E, B, J, rho, phi, external_fields = self._empty_fields(world) - fields = (E, B, J, rho, phi, external_fields) - tile_shape = self._one_tile_shape(world) - world = self._world_with_tiled_grids(world, tile_shape) - tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._neutral_species(world), - fields, - world, - tile_shape, - ) - tiled_particles, tiled_fields = time_loop_electrostatic( - tiled_particles, - species_config, - tiled_fields, - world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", - ) - - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(B_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(J_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(rho_tiles.shape[:3], (1, 1, 1)) - self.assertEqual(phi_tiles.shape[:3], (1, 1, 1)) - self.assertFalse(bool(tiled_fields[-1])) - self.assertTrue(jnp.all(jnp.isfinite(rho_tiles))) - self.assertTrue(jnp.all(jnp.isfinite(phi_tiles))) - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", - "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step single-tile electrostatic check.", - ) - def test_long_two_stream_single_tile_electrostatic_remains_finite(self): - tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt = self._long_two_stream_state() - - tiled_loop = jax.jit( - time_loop_electrostatic, - static_argnames=("solver", "relativistic", "particle_pusher"), - ) - - for step in range(1, Nt + 1): - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", - ) - - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - self.assertEqual(tuple(world["tile_shape"]), tuple(tile_shape)) - self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) - self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - for component in E_tiles + B_tiles + J_tiles: - self.assertTrue(jnp.all(jnp.isfinite(component)), f"step {step}: nonfinite vector field") - self.assertTrue(jnp.all(jnp.isfinite(rho_tiles)), f"step {step}: nonfinite rho") - self.assertTrue(jnp.all(jnp.isfinite(phi_tiles)), f"step {step}: nonfinite phi") - - def test_initialize_simulation_accepts_tiled_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": "tiled 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)) - self.assertIn("tiled_center_grid", world["grids"]) - self.assertIn("tiled_vertex_grid", world["grids"]) - self.assertEqual(int(world["guard_cells"]), 2) - self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (1, 1, 1, 12)) - self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (1, 1, 1, 5)) - 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/tiled_esirkepov_integration_test.py b/tests/tiled_esirkepov_integration_test.py deleted file mode 100644 index 95b7e95e..00000000 --- a/tests/tiled_esirkepov_integration_test.py +++ /dev/null @@ -1,151 +0,0 @@ -import os -import tempfile -import unittest -from contextlib import redirect_stdout -from io import StringIO -from unittest.mock import patch - -import jax -import numpy as np -import toml - -from PyPIC3D import __main__ as pypic_main -from PyPIC3D.__main__ import run_PyPIC3D - - -jax.config.update("jax_enable_x64", True) - - -class TiledEsirkepovIntegrationTest(unittest.TestCase): - def test_short_output_enabled_tiled_esirkepov_run_writes_openpmd_without_vtk(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 esirkepov output 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, - "plot_openpmd_fields": True, - "plot_openpmd_particles": True, - "plot_phasespace": True, - }, - "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, - }, - } - - _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world, _species_config = run_PyPIC3D(config) - _E, _B, J, *_rest = fields - - self.assertEqual(J[0].shape[-3:], (6, 5, 5)) - self.assertEqual(tuple(int(width) for width in world["tile_shape"]), (2, 1, 1)) - for filename in ( - "total_energy.txt", - "electric_field_energy.txt", - "magnetic_field_energy.txt", - "kinetic_energy.txt", - "total_momentum.txt", - ): - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", filename))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "fields.h5"))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles.h5"))) - self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice"))) - self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "phase_space", "x", "electrons_phase_space.000000000.npy"))) - - def test_main_prints_final_energy_for_tiled_runtime(self): - with tempfile.TemporaryDirectory() as tmpdir: - x_path = os.path.join(tmpdir, "x.npy") - zeros_path = os.path.join(tmpdir, "zeros.npy") - config_path = os.path.join(tmpdir, "final_energy.toml") - np.save(x_path, np.array([0.0])) - np.save(zeros_path, np.zeros(1)) - - config = { - "simulation_parameters": { - "name": "final energy smoke", - "output_dir": tmpdir, - "solver": "electrodynamic_yee", - "Nx": 4, - "Ny": 1, - "Nz": 1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "dt": 1.0e-10, - "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": 1, - "charge": -1.0, - "mass": 1.0, - "weight": 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, - }, - } - with open(config_path, "w") as f: - toml.dump(config, f) - - with patch("sys.argv", ["PyPIC3D", "--config", config_path]): - stdout = StringIO() - with redirect_stdout(stdout): - pypic_main.main() - - self.assertIn("Final Kinetic Energy", stdout.getvalue()) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "output.toml"))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/tiled_esirkepov_test.py b/tests/tiled_esirkepov_test.py deleted file mode 100644 index a75555f7..00000000 --- a/tests/tiled_esirkepov_test.py +++ /dev/null @@ -1,900 +0,0 @@ -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 ( - _active_stencil_indices, - _compact_1d_esirkepov_weights, - _compact_2d_esirkepov_weights, - Esirkepov_current, - get_1D_esirkepov_weights, - get_2D_esirkepov_weights, -) -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.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import 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) - - 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 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(self, world, x1): - return particle_species( - name="electrons", - N_particles=3, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=x1, - x2=jnp.zeros_like(x1), - x3=jnp.zeros_like(x1), - v1=jnp.array([0.08, -0.06, 0.05]), - v2=jnp.zeros_like(x1), - v3=jnp.zeros_like(x1), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - def _species_from_arrays(self, world, x, u, weight=0.5, active_mask=None): - return particle_species( - name="electrons", - N_particles=x.shape[0], - charge=-1.0, - mass=1.0, - weight=weight, - T=1.0, - x1=x[:, 0], - x2=x[:, 1], - x3=x[:, 2], - v1=u[:, 0], - v2=u[:, 1], - v3=u[:, 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=active_mask, - ) - - 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 _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, old_species, constants, tile_shape): - world = self._world_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], - } - g = int(world["guard_cells"]) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - _, _, 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) - - old_species = self._species_from_arrays(world, x_old, u) - J_tiles, J_from_tiles = self._assembled_esirkepov_current(world, old_species, constants, tile_shape) - _, J_reference = self._assembled_esirkepov_current( - world, - old_species, - 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_reduced_axis_esirkepov_scatter_uses_only_collapsed_stencil_index(self): - self.assertEqual(_active_stencil_indices(True), (0, 1, 2, 3, 4)) - self.assertEqual(_active_stencil_indices(False), (2,)) - - def test_compact_1d_esirkepov_weights_match_full_center_line(self): - xw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.3]), jnp.array([0.6, 0.5]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] - oxw = [jnp.array([0.0, 0.0]), jnp.array([0.1, 0.2]), jnp.array([0.7, 0.6]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] - yw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] - zw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] - oyw = yw - ozw = zw - - Wx_line, Wy_line, Wz_line = _compact_1d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, dim=0) - Wx_full, Wy_full, Wz_full = get_1D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles=2, dim=0) - - self.assertTrue(jnp.allclose(Wx_line, Wx_full[:, 2, 2, :])) - self.assertTrue(jnp.allclose(Wy_line, Wy_full[:, 2, 2, :])) - self.assertTrue(jnp.allclose(Wz_line, Wz_full[:, 2, 2, :])) - - def test_compact_2d_esirkepov_weights_match_full_active_plane(self): - xw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.3]), jnp.array([0.6, 0.5]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] - oxw = [jnp.array([0.0, 0.0]), jnp.array([0.1, 0.2]), jnp.array([0.7, 0.6]), jnp.array([0.2, 0.2]), jnp.array([0.0, 0.0])] - yw = [jnp.array([0.0, 0.0]), jnp.array([0.3, 0.2]), jnp.array([0.4, 0.5]), jnp.array([0.3, 0.3]), jnp.array([0.0, 0.0])] - oyw = [jnp.array([0.0, 0.0]), jnp.array([0.2, 0.1]), jnp.array([0.5, 0.6]), jnp.array([0.3, 0.3]), jnp.array([0.0, 0.0])] - zw = [jnp.zeros(2), jnp.zeros(2), jnp.ones(2), jnp.zeros(2), jnp.zeros(2)] - ozw = zw - - Wx_plane, Wy_plane, Wz_plane = _compact_2d_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, null_dim=2) - Wx_full, Wy_full, Wz_full = get_2D_esirkepov_weights(xw, yw, zw, oxw, oyw, ozw, N_particles=2, null_dim=2) - - self.assertTrue(jnp.allclose(Wx_plane, Wx_full[:, :, 2, :])) - self.assertTrue(jnp.allclose(Wy_plane, Wy_full[:, :, 2, :])) - self.assertTrue(jnp.allclose(Wz_plane, Wz_full[:, :, 2, :])) - - 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) - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - x_old = jnp.array([-1.10, -0.10, 1.05]) - old_species = self._species(world, x_old) - - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - 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, - old_species, - 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} - simulation_parameters = { - "particle_tile_nx": world["tile_shape"][0], - "particle_tile_ny": world["tile_shape"][1], - "particle_tile_nz": world["tile_shape"][2], - } - x_old = jnp.array([-1.10, -0.10, 1.05]) - old_species = self._species(world, x_old) - - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - _, _, 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, - old_species, - 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} - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - old_species = self._species(world, jnp.array([-1.10, -0.10, 1.05])) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - - 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) - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - 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], - ] - ) - old_species = self._species_from_arrays(world, x_old, u) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) - 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, - ) - - old_species = self._species_from_arrays( - 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, - old_species, - 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/tiled_fluid_quantities_test.py b/tests/tiled_fluid_quantities_test.py deleted file mode 100644 index 83b1a1fa..00000000 --- a/tests/tiled_fluid_quantities_test.py +++ /dev/null @@ -1,404 +0,0 @@ -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 PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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 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 = to_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 = to_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 = to_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 = particle_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 = particle_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 = to_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 = to_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 = to_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 = to_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 = to_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/tiled_particle_diagnostics_test.py b/tests/tiled_particle_diagnostics_test.py deleted file mode 100644 index 69ce6682..00000000 --- a/tests/tiled_particle_diagnostics_test.py +++ /dev/null @@ -1,160 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.diagnostics.output_adapters import particles_for_output -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles - - -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 = particle_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 = particle_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 = to_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.get_active_mask() - - original_x = jnp.stack(original.get_forward_position(), axis=-1)[active] - original_u = jnp.stack(original.get_velocity(), axis=-1)[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 = to_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 = to_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 = to_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.get_active_mask() - original_position = jnp.stack(original.get_position(), axis=-1)[active] - 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - species_names = tuple(species.get_name() for species in species_list) - - flattened_species = particles_for_output(tiled_particles, species_config=species_config, species_names=species_names) - - self.assertEqual(flattened_species[0].name, "ions") - self.assertEqual(flattened_species[1].name, "electrons") - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/tiled_particle_initialization_test.py b/tests/tiled_particle_initialization_test.py deleted file mode 100644 index fbe1097e..00000000 --- a/tests/tiled_particle_initialization_test.py +++ /dev/null @@ -1,439 +0,0 @@ -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 tiled_particle_fixtures import particle_species, to_tiled_particles - - -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_to_tiled_particles_packs_existing_particle_species_list(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 = particle_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 = to_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 = [ - particle_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"], - ), - particle_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 = to_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_to_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 = particle_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 = to_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/tiled_particle_refresh_test.py b/tests/tiled_particle_refresh_test.py deleted file mode 100644 index 9046b045..00000000 --- a/tests/tiled_particle_refresh_test.py +++ /dev/null @@ -1,231 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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 particle_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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/tiled_pusher_test.py b/tests/tiled_pusher_test.py deleted file mode 100644 index 4328276a..00000000 --- a/tests/tiled_pusher_test.py +++ /dev/null @@ -1,448 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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 = to_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 particle_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 particle_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 = to_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 particle_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 particle_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/tiled_refactor_contract_test.py b/tests/tiled_refactor_contract_test.py deleted file mode 100644 index b39758b9..00000000 --- a/tests/tiled_refactor_contract_test.py +++ /dev/null @@ -1,281 +0,0 @@ -import inspect -import importlib -import importlib.util -import unittest - -import jax.numpy as jnp - -from PyPIC3D import evolve -from PyPIC3D.deposition import Esirkepov -from PyPIC3D.deposition import J_from_rhov -from PyPIC3D.deposition import rho -from PyPIC3D.boundary_conditions import boundaryconditions -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics import output_adapters -from PyPIC3D.initialization import build_tiled_array, initialize_fields -from PyPIC3D.solvers import electrostatic_yee -from PyPIC3D.solvers import first_order_yee -from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid - - -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) - ntx = int(world["Nx"]) // tile_nx - nty = int(world["Ny"]) // tile_ny - ntz = int(world["Nz"]) // tile_nz - - 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) - - -class TestTiledRefactorContracts(unittest.TestCase): - def _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}, - "tile_shape": (2, 3, 2), - "guard_cells": 2, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "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 _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) - return jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (ii + 2.0 * jj - kk)) - - def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): - world = self._world() - tile_shape = (2, 3, 2) - E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) - E_tiles = tuple(tile_scalar_field(component, world, tile_shape) for component in E) - - refreshed = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) - self.assertEqual(len(refreshed), 3) - for component in refreshed: - self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) - - def test_evolve_loops_read_runtime_controls_from_world(self): - electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) - electrostatic_signature = inspect.signature(evolve.time_loop_electrostatic) - for name in ("curl_func", "J_func", "tile_shape", "g", "current_deposition"): - self.assertNotIn(name, electrodynamic_signature.parameters) - for name in ("curl_func", "J_func", "tile_shape", "g"): - self.assertNotIn(name, electrostatic_signature.parameters) - - def test_tiled_yee_updates_read_tile_metadata_from_world(self): - update_E_signature = inspect.signature(first_order_yee.update_E) - update_B_signature = inspect.signature(first_order_yee.update_B) - for name in ("curl_func", "tile_shape", "g"): - self.assertNotIn(name, update_E_signature.parameters) - self.assertNotIn(name, update_B_signature.parameters) - self.assertFalse(hasattr(first_order_yee, "update_tiled_E")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_B")) - self.assertFalse(hasattr(first_order_yee, "tiled_grid_axes_from_world")) - self.assertFalse(hasattr(first_order_yee, "tile_grid_axes")) - self.assertFalse(hasattr(first_order_yee, "tile_vector_field")) - self.assertFalse(hasattr(first_order_yee, "stack_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "unstack_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_for_pml")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_for_pml")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "apply_tiled_conducting_bc")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "empty_tiled_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "empty_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "tile_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "assemble_tiled_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "assemble_tiled_vector_field")) - - def test_global_assembly_lives_at_output_boundary(self): - self.assertTrue(hasattr(output_adapters, "assemble_tiled_scalar_field")) - self.assertTrue(hasattr(output_adapters, "assemble_tiled_vector_field")) - - source = inspect.getsource(output_adapters) - self.assertNotIn("from PyPIC3D.solvers.yee_tiled import", source) - - def test_ghost_cell_logic_lives_in_boundary_conditions_ghost_cells(self): - for name in ( - "update_tiled_ghost_cells", - "update_tiled_vector_ghost_cells", - "update_tiled_ghost_cells_for_pml", - "update_tiled_vector_ghost_cells_for_pml", - "fold_tiled_ghost_cells", - "fold_tiled_vector_ghost_cells", - "apply_tiled_conducting_bc", - ): - self.assertTrue(hasattr(ghost_cells, name)) - - for name in ( - "update_ghost_cells", - "fold_ghost_cells", - "apply_conducting_bc", - "apply_scalar_conducting_bc", - ): - self.assertFalse(hasattr(boundaryconditions, name)) - - def test_tiled_array_initialization_lives_in_initialization(self): - world = self._world() - world["tile_shape"] = (2, 3, 2) - world["guard_cells"] = 2 - - array = build_tiled_array(world, dtype=jnp.float32) - self.assertEqual(array.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(array.dtype, jnp.float32) - self.assertTrue(jnp.all(array == 0.0)) - - def test_initialize_fields_builds_tiled_field_state_from_world(self): - world = self._world() - world["tile_shape"] = (2, 3, 2) - world["guard_cells"] = 2 - - E, B, J, phi, rho_tiles = initialize_fields(world) - for vector_field in (E, B, J): - self.assertEqual(len(vector_field), 3) - for component in vector_field: - self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(component.dtype, jnp.float64) - self.assertTrue(jnp.all(component == 0.0)) - for scalar_field in (phi, rho_tiles): - self.assertEqual(scalar_field.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(scalar_field.dtype, jnp.float64) - self.assertTrue(jnp.all(scalar_field == 0.0)) - - def test_grid_builders_live_in_utilities(self): - from PyPIC3D import utils - from PyPIC3D.utilities import grids - - self.assertFalse(hasattr(utils, "build_yee_grid")) - self.assertFalse(hasattr(utils, "build_collocated_grid")) - self.assertTrue(hasattr(grids, "build_yee_grid")) - self.assertTrue(hasattr(grids, "build_collocated_grid")) - self.assertFalse(hasattr(grids, "tile_grid_axes")) - self.assertTrue(hasattr(grids, "build_tiled_yee_grids")) - - def test_tiled_benchmarking_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) - - def test_deprecated_fdtd_errors_and_root_boris_modules_are_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.solvers.fdtd")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.errors")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.boris")) - self.assertIsNotNone(importlib.util.find_spec("PyPIC3D.pusher.boris")) - - def test_old_particle_species_surfaces_are_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.species_class")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.tiled_particle_diagnostics")) - - from PyPIC3D.particles import particle_initialization - - self.assertFalse(hasattr(particle_initialization, "to_tiled_particles")) - source = inspect.getsource(particle_initialization.load_particles_from_toml) - self.assertNotIn("particle_species", source) - - def test_particle_push_owns_tiled_pusher_contract(self): - particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") - - self.assertIsNone(importlib.util.find_spec("PyPIC3D.pusher.tiled_pusher")) - signature = inspect.signature(particle_push.particle_push) - for name in ("grid", "staggered_grid", "tile_shape", "g"): - self.assertNotIn(name, signature.parameters) - - source = inspect.getsource(particle_push.particle_push) - self.assertIn('world["grids"]["tiled_center_grid"]', source) - self.assertIn('world["grids"]["tiled_vertex_grid"]', source) - self.assertNotIn("tiled_grid_axes_from_world", source) - self.assertNotIn("tile_grid_axes", source) - self.assertNotIn("get_forward_position", source) - self.assertNotIn("set_velocity", source) - - def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): - source = inspect.getsource(evolve.time_loop_electrodynamic) - self.assertNotIn("assemble_tiled_scalar_field", source) - self.assertNotIn("assemble_tiled_vector_field", source) - self.assertNotIn("fields_for_output", source) - - def test_electrostatic_solver_uses_single_tile_arrays_directly(self): - source = inspect.getsource(electrostatic_yee.calculate_tiled_electrostatic_fields) - self.assertNotIn("assemble_tiled_scalar_field", source) - self.assertNotIn("tile_scalar_field", source) - self.assertIn("[0, 0, 0]", source) - - def test_flat_current_and_electrodynamic_references_are_removed(self): - self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) - self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_tiled")) - self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) - self.assertFalse(hasattr(evolve, "_time_loop_electrostatic_global_reference")) - - def test_direct_current_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(J_from_rhov.J_from_rhov) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_esirkepov_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(Esirkepov.Esirkepov_current) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_tiled_esirkepov_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.esirkepov_tiled")) - - def test_rho_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(rho.compute_rho) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_flat_rho_and_tiled_rho_module_are_removed(self): - self.assertFalse(hasattr(rho, "_compute_rho_flat")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.rho_tiled")) - - def test_current_methods_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.current_methods")) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/tiled_rho_test.py b/tests/tiled_rho_test.py deleted file mode 100644 index 410e88f9..00000000 --- a/tests/tiled_rho_test.py +++ /dev/null @@ -1,363 +0,0 @@ -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 PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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 = particle_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 = particle_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 = to_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.v1 = jnp.zeros_like(species.v1) - species.v2 = jnp.zeros_like(species.v2) - species.v3 = jnp.zeros_like(species.v3) - 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/tiled_yee_integration_test.py b/tests/tiled_yee_integration_test.py deleted file mode 100644 index 2c544984..00000000 --- a/tests/tiled_yee_integration_test.py +++ /dev/null @@ -1,1411 +0,0 @@ -import os -import tempfile -import unittest - -import jax -import jax.numpy as jnp -import numpy as np -import toml - -from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.initialization import initialize_simulation, validate_field_solver -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import TiledParticles -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field -from PyPIC3D.solvers.first_order_yee import ( - update_B, - update_E, -) -from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid -from PyPIC3D.utils import compute_energy -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC - - -jax.config.update("jax_enable_x64", True) - -ROUND_OFF_RTOL = 1.0e-11 -ROUND_OFF_ATOL = 1.0e-11 - - -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 TestTiledYeeIntegration(unittest.TestCase): - def _build_world(self, boundary_conditions=None): - if boundary_conditions is None: - boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} - world = { - "Nx": 8, - "Ny": 1, - "Nz": 1, - "dx": 0.5, - "dy": 1.0, - "dz": 1.0, - "dt": 0.01, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "guard_cells": 2, - "current_deposition": "direct", - "current_filter": "none", - "boundary_conditions": boundary_conditions, - } - 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): - 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_fields(self, 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)) - phi = jnp.zeros(shape) - rho = jnp.zeros(shape) - external_fields = (E, B) - return E, B, J, rho, phi, external_fields, None - - def _species(self, world): - return particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - 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.lexsort((x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - - def _long_two_stream_state(self, current_deposition="direct", current_filter="none"): - world = { - "Nx": 16, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.002, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 2, - "guard_cells": 2, - "current_deposition": current_deposition, - "current_filter": current_filter, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (4, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, tile_shape) - Nt = 1500 - - x0 = jnp.linspace( - -0.5 * world["x_wind"], - 0.5 * world["x_wind"], - world["Nx"], - endpoint=False, - ) + 0.5 * world["dx"] - x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) - zeros = jnp.zeros_like(x0) - beam_speed = 0.08 - - species_kwargs = { - "N_particles": world["Nx"], - "mass": 1.0, - "weight": 1.0, - "T": 0.0, - "x1": x0, - "x2": zeros, - "x3": zeros, - "v2": zeros, - "v3": zeros, - "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 = [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - reference_species = [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) - tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) - g = int(tiled_world["guard_cells"]) - if current_deposition == "esirkepov": - J_template = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) - J_tiles = tuple(jnp.zeros_like(component) for component in J_template) - else: - J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) - - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=g), - tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=g), - J_tiles, - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=g), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - reference_species, - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - current_deposition, - ) - reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - step=0, - ) - - return ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) - - def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): - current_slots = tiled_particles.active.shape[-1] - if current_slots >= min_slots: - return tiled_particles - - pad_width = [(0, 0)] * tiled_particles.x.ndim - pad_width[-2] = (0, min_slots - current_slots) - x = jnp.pad(tiled_particles.x, pad_width) - u = jnp.pad(tiled_particles.u, pad_width) - - scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim - scalar_pad_width[-1] = (0, min_slots - current_slots) - active = jnp.pad(tiled_particles.active, scalar_pad_width) - - return TiledParticles( - x=x, - u=u, - active=active, - ) - - def _tiled_species_rows(self, tiled_particles, species_index): - active = tiled_particles.active[:, :, :, species_index, :] - x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - - def _build_tiled_state(self, particles, fields, world, tile_shape, current_deposition="direct"): - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - E, B, J, rho, phi, external_fields, pml_state = fields - g = int(world["guard_cells"]) - if current_deposition == "esirkepov": - J_template = tile_vector_field(J, world, tile_shape, num_guard_cells=g) - J_tiles = tuple(jnp.zeros_like(component) for component in J_template) - else: - J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) - tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=g), - tile_vector_field(B, world, tile_shape, num_guard_cells=g), - J_tiles, - rho, - phi, - ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), - ), - pml_state, - jnp.asarray(False), - ) - return tiled_particles, species_config, tiled_fields - - def _jitted_electrodynamic_loop_with_static_world(self, world): - def loop_with_static_world( - particles, - species_config, - fields, - constants, - relativistic=True, - particle_pusher="boris", - ): - return time_loop_electrodynamic( - particles, - species_config, - fields, - world, - constants, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - - return jax.jit( - loop_with_static_world, - static_argnames=("relativistic", "particle_pusher"), - ) - - def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): - for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): - diff = jnp.max(jnp.abs(reference_component - tiled_component)) - self.assertTrue( - jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: {name}{component_name} max abs diff {diff}", - ) - - def _assert_tiled_state_matches_standard( - self, - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - step, - ): - if len(tiled_fields) > 7: - self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - - species_count = reference_particles.active.shape[3] - for species_index in range(species_count): - reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) - tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) - self.assertTrue( - jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} position mismatch", - ) - self.assertTrue( - jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} velocity mismatch", - ) - - reference_E, reference_B, reference_J, *_ = reference_fields - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - num_guard_cells = int(world.get("guard_cells", 1)) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "E", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "B", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "J", - ) - - def _max_tiled_state_difference( - self, - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - world, - tile_shape, - ): - differences = {} - species_count = reference_particles.active.shape[3] - for species_index in range(species_count): - reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) - tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) - differences[f"x{species_index}"] = float(jnp.max(jnp.abs(tiled_x - reference_x))) - differences[f"u{species_index}"] = float(jnp.max(jnp.abs(tiled_u - reference_u))) - - reference_E, reference_B, reference_J, *_ = reference_fields - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - num_guard_cells = int(world.get("guard_cells", 1)) - for name, reference_field, tiled_field in ( - ( - "E", - assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ( - "B", - assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ( - "J", - assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ): - differences[name] = max(float(jnp.max(jnp.abs(tiled - reference))) for reference, tiled in zip(reference_field, tiled_field)) - - return differences - - def test_validate_field_solver_accepts_tiled_yee(self): - validate_field_solver("electrodynamic_yee") - - def test_tiled_yee_field_update_applies_standard_field_filter(self): - world = self._build_world() - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} - tile_shape = (2, 1, 1) - shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - - x = jnp.arange(world["Nx"], dtype=jnp.float64)[:, None, None] - seed = jnp.zeros(shape) - Ex = _update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ey = _update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ez = _update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bx = _update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - By = _update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bz = _update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - E = (Ex, Ey, Ez) - B = (Bx, By, Bz) - J = (seed, seed, seed) - rho = seed - phi = seed - external_fields = ((seed, seed, seed), (seed, seed, seed)) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, tile_shape) - reference_fields = ( - tile_vector_field(E, reference_world, reference_tile_shape), - tile_vector_field(B, reference_world, reference_tile_shape), - tile_vector_field(J, reference_world, reference_tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], reference_world, reference_tile_shape), - tile_vector_field(external_fields[1], reference_world, reference_tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_E, reference_pml_state = update_E( - reference_fields[0], - reference_fields[1], - reference_fields[2], - reference_world, - constants, - ) - reference_B, reference_pml_state = update_B( - reference_E, - reference_fields[1], - reference_world, - constants, - reference_pml_state, - ) - reference_fields = (reference_E, reference_B, *reference_fields[2:]) - - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - tiled_E, tiled_pml_state = update_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) - tiled_B, tiled_pml_state = update_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) - tiled_fields = (tiled_E, tiled_B, *tiled_fields[2:]) - - reference_E, reference_B, *_ = reference_fields - E_tiles, B_tiles, *_ = tiled_fields - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, reference_world, reference_tile_shape), - assemble_tiled_vector_field(E_tiles, tiled_world, tile_shape), - 1, - "E", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, reference_world, reference_tile_shape), - assemble_tiled_vector_field(B_tiles, tiled_world, tile_shape), - 1, - "B", - ) - - def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(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 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, - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee.toml") - with open(config_path, "w") as f: - toml.dump(config, f) - - loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) - - def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(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 higuera cary 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "filter_j": "none", - "particle_pusher": "higuera_cary", - "relativistic": True, - }, - "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_higuera_cary.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["particle_pusher"], "higuera_cary") - - def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(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 digital filter 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee_digital.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["filter_j"], "digital") - - def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(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 bilinear filter 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "filter_j": "bilinear", - "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_bilinear.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["filter_j"], "bilinear") - - def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(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 conducting 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, - "x_bc": "conducting", - "y_bc": "conducting", - "z_bc": "conducting", - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee_conducting.toml") - with open(config_path, "w") as f: - toml.dump(config, f) - - loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(world["boundary_conditions"], {"x": 1, "y": 1, "z": 1}) - - def test_tiled_yee_step_matches_standard_periodic_yee_step(self): - world = self._build_world() - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - reference_E, reference_B, *_ = reference_fields - E_tiles, B_tiles, *_ = tiled_fields - reference_energy = compute_energy( - reference_particles, - reference_E, - reference_B, - world, - constants, - species_config=reference_species_config, - ) - tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants, species_config=species_config) - for reference, tiled in zip(reference_energy, tiled_energy): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - - def test_tiled_yee_step_matches_standard_conducting_yee_step(self): - world = self._build_world( - boundary_conditions={"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING} - ) - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_uses_digital_filtered_current(self): - world = self._build_world() - world["current_filter"] = "digital" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_uses_bilinear_filtered_current(self): - world = self._build_world() - world["current_filter"] = "bilinear" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_standard(self): - world = self._build_world() - world["shape_factor"] = 2 - world["guard_cells"] = 2 - world["current_filter"] = "bilinear" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.array([0.03, -0.04, 0.02, -0.01]), - v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - reference_species = particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.array([0.03, -0.04, 0.02, -0.01]), - v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_YEE") == "1", - "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", - ) - def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): - ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) = self._long_two_stream_state("direct", "none") - - reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) - - for t in range(Nt): - reference_particles, reference_fields = reference_loop( - reference_particles, - reference_species_config, - reference_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - - differences = self._max_tiled_state_difference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - ) - if t == 0: - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - t + 1, - ) - if any(value > ROUND_OFF_ATOL for value in differences.values()): - print(f"direct-current roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") - break - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_YEE") == "1", - "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", - ) - def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): - ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) = self._long_two_stream_state("esirkepov", "none") - - reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) - - for t in range(Nt): - reference_particles, reference_fields = reference_loop( - reference_particles, - reference_species_config, - reference_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - - differences = self._max_tiled_state_difference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - ) - if t == 0: - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - t + 1, - ) - if any(value > ROUND_OFF_ATOL for value in differences.values()): - print(f"Esirkepov roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") - break - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/utils_test.py b/tests/utils_test.py deleted file mode 100644 index 060415da..00000000 --- a/tests/utils_test.py +++ /dev/null @@ -1,499 +0,0 @@ -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 PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -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, -) -from PyPIC3D.tests.tiled_particle_fixtures import particle_species - -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} - 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="periodic", - ) - species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) - - tiled_particles, species_config = to_tiled_particles([species], world, {"particle_tile_nx": 1, "particle_tile_ny": 1, "particle_tile_nz": 1}) - - _, _, 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 = particle_species( - name="beam electrons", - N_particles=2, - charge=-1.0, - mass=1.0, - weight=2.0, - T=0.0, - x1=jnp.array([-0.25, 0.25]), - x2=jnp.array([0.0, 0.0]), - x3=jnp.array([0.0, 0.0]), - v1=jnp.array([0.1, 0.2]), - v2=jnp.array([0.0, 0.0]), - v3=jnp.array([0.0, 0.0]), - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=0.25, - dy=0.5, - dz=1.0, - dt=0.1, - ) - 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 = to_tiled_particles([species], world, 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 = particle_species( - name="electrons", - N_particles=1, - charge=-1.0, - mass=1.0, - weight=1.0, - T=0.0, - x1=jnp.array([0.0]), - x2=jnp.array([0.0]), - x3=jnp.array([0.0]), - v1=jnp.array([0.2]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=0.25, - dy=0.5, - dz=1.0, - dt=0.1, - ) - 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 = to_tiled_particles([species], world, 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/yee_tiled_test.py b/tests/yee_tiled_test.py deleted file mode 100644 index 5d17deb0..00000000 --- a/tests/yee_tiled_test.py +++ /dev/null @@ -1,635 +0,0 @@ -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() From e400bb21621f25a10ce68310e0eb79db68d998b9 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 8 Jul 2026 14:14:40 -0500 Subject: [PATCH 65/83] refactoring direct deposition unit-tests --- tests/code_tests/direct_deposition_test.py | 966 +++++++++++++++++++++ 1 file changed, 966 insertions(+) create mode 100644 tests/code_tests/direct_deposition_test.py 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() From b1044c52e31bc80e013274aad16eb224f16d0ee6 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Wed, 8 Jul 2026 15:46:26 -0500 Subject: [PATCH 66/83] fixing issue with cfl condition --- PyPIC3D/utils.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/PyPIC3D/utils.py b/PyPIC3D/utils.py index 0bbbfa72..78c1e933 100644 --- a/PyPIC3D/utils.py +++ b/PyPIC3D/utils.py @@ -768,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 From d51397f47b68fca1cbfaf9fa0d909562e0eba4dc Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 9 Jul 2026 16:18:24 -0500 Subject: [PATCH 67/83] refactoring and restructuring unit tests --- .coveragerc | 11 + PyPIC3D/deposition/Esirkepov.py | 272 +--- pyproject.toml | 5 + tests/code_tests/__init__.py | 1 + tests/code_tests/grid_and_stencil_test.py | 101 ++ tests/code_tests/initialization_test.py | 228 +++ tests/code_tests/openpmd_test.py | 619 ++++++++ .../tiled_esirkepov_integration_test.py | 151 ++ tests/code_tests/tiled_esirkepov_test.py | 861 ++++++++++ .../code_tests/tiled_fluid_quantities_test.py | 404 +++++ .../tiled_particle_diagnostics_test.py | 160 ++ .../tiled_particle_initialization_test.py | 439 +++++ .../code_tests/tiled_particle_refresh_test.py | 231 +++ tests/code_tests/tiled_pusher_test.py | 448 ++++++ .../tiled_refactor_contract_test.py | 281 ++++ tests/code_tests/tiled_rho_test.py | 363 +++++ tests/code_tests/utils_test.py | 499 ++++++ tests/code_tests/yee_tiled_test.py | 635 ++++++++ tests/physics_tests/__init__.py | 1 + tests/physics_tests/boris_test.py | 452 ++++++ tests/physics_tests/electrostatic_yee_test.py | 164 ++ tests/physics_tests/pml_test.py | 590 +++++++ tests/physics_tests/pusher_schmitz_test.py | 203 +++ .../physics_tests/tiled_electrostatic_test.py | 512 ++++++ .../tiled_yee_integration_test.py | 1411 +++++++++++++++++ tests/physics_tests/yee_convergence_test.py | 90 ++ 26 files changed, 8861 insertions(+), 271 deletions(-) create mode 100644 .coveragerc create mode 100644 tests/code_tests/__init__.py create mode 100644 tests/code_tests/grid_and_stencil_test.py create mode 100644 tests/code_tests/initialization_test.py create mode 100644 tests/code_tests/openpmd_test.py create mode 100644 tests/code_tests/tiled_esirkepov_integration_test.py create mode 100644 tests/code_tests/tiled_esirkepov_test.py create mode 100644 tests/code_tests/tiled_fluid_quantities_test.py create mode 100644 tests/code_tests/tiled_particle_diagnostics_test.py create mode 100644 tests/code_tests/tiled_particle_initialization_test.py create mode 100644 tests/code_tests/tiled_particle_refresh_test.py create mode 100644 tests/code_tests/tiled_pusher_test.py create mode 100644 tests/code_tests/tiled_refactor_contract_test.py create mode 100644 tests/code_tests/tiled_rho_test.py create mode 100644 tests/code_tests/utils_test.py create mode 100644 tests/code_tests/yee_tiled_test.py create mode 100644 tests/physics_tests/__init__.py create mode 100644 tests/physics_tests/boris_test.py create mode 100644 tests/physics_tests/electrostatic_yee_test.py create mode 100644 tests/physics_tests/pml_test.py create mode 100644 tests/physics_tests/pusher_schmitz_test.py create mode 100644 tests/physics_tests/tiled_electrostatic_test.py create mode 100644 tests/physics_tests/tiled_yee_integration_test.py create mode 100644 tests/physics_tests/yee_convergence_test.py 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/PyPIC3D/deposition/Esirkepov.py b/PyPIC3D/deposition/Esirkepov.py index 8c5bdc92..bf17d5e3 100644 --- a/PyPIC3D/deposition/Esirkepov.py +++ b/PyPIC3D/deposition/Esirkepov.py @@ -42,174 +42,6 @@ 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 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. - """ - - - field_axis = jnp.moveaxis(field, axis, 0) - # move the axis to be folded to the front for easier indexing. - - if bc == "reflecting": - 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. - - - 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 - - return field - - -def enforce_particle_bc_code_along_axis(field, axis, bc, component_axis): - """ - Fold two Esirkepov ghost layers using global particle boundary-condition codes. - """ - - field_axis = jnp.moveaxis(field, axis, 0) - sign = -1.0 if axis == component_axis else 1.0 - - def periodic_bc(field_axis): - 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, :, :]) - return field_axis - - def reflecting_bc(field_axis): - 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, :, :]) - return field_axis - - def absorbing_bc(field_axis): - return field_axis - - field_axis = lax.switch(bc, (periodic_bc, reflecting_bc, absorbing_bc), field_axis) - field = jnp.moveaxis(field_axis, 0, axis) - field = clear_ghost_cells(field, axis) - - return field - - -def ghost_cell_particle_bc_codes_esirkepov(field, bc_x, bc_y, bc_z, component_axis): - field = enforce_particle_bc_code_along_axis(field, axis=0, bc=bc_x, component_axis=component_axis) - field = enforce_particle_bc_code_along_axis(field, axis=1, bc=bc_y, component_axis=component_axis) - field = enforce_particle_bc_code_along_axis(field, axis=2, bc=bc_z, component_axis=component_axis) - return field - - -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 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 _active_stencil_indices(axis_active): - if axis_active: - return (0, 1, 2, 3, 4) - return (2,) - - 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) @@ -733,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/pyproject.toml b/pyproject.toml index 61c1b095..928f2ee3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,11 @@ dependencies = [ "openpmd-api", ] +[project.optional-dependencies] +test = [ + "coverage", +] + [tool.setuptools.packages.find] include = ["PyPIC3D*"] 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/grid_and_stencil_test.py b/tests/code_tests/grid_and_stencil_test.py new file mode 100644 index 00000000..c31ab49f --- /dev/null +++ b/tests/code_tests/grid_and_stencil_test.py @@ -0,0 +1,101 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions.grid_and_stencil import ( + BC_CONDUCTING, + BC_PERIODIC, + axis_has_active_cells, + build_axis_stencil_points, + build_collocated_axis, + build_staggered_axis, + collapse_axis_stencil, + compute_particle_anchor, + inactive_axis_index, + particle_axis_offset, + prepare_particle_axis_stencil, + uniform_axis_spacing, + wrap_periodic_position, +) + +jax.config.update("jax_enable_x64", True) + + +class TestGhostCellHelpers(unittest.TestCase): + + def test_wrap_periodic_position(self): + wind = 4.0 + x = jnp.array([-2.5, -2.0, 1.75, 2.0, 2.25, 6.5]) + 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)) + + 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) + + def test_axis_spacing_and_anchor_helpers(self): + axis = build_collocated_axis(-1.0, 0.5, 4) + position = jnp.array([0.1, 0.6]) + anchor = compute_particle_anchor(position, axis, shape_factor=1) + offset = particle_axis_offset(position, anchor, axis) + + 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]))) + + def test_build_axis_stencil_points_periodic_and_conducting(self): + anchor = jnp.array([0, 5]) + offsets = jnp.array([-1, 0, 1]) + + periodic = build_axis_stencil_points(anchor, 6, BC_PERIODIC, offsets) + conducting = build_axis_stencil_points(anchor, 6, BC_CONDUCTING, offsets) + + self.assertTrue(jnp.array_equal(periodic[:, 0], jnp.array([5, 0, 1]))) + 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]))) + + def test_prepare_particle_axis_stencil_wraps_periodic_indices(self): + axis = build_collocated_axis(-1.0, 0.5, 4) + position = jnp.array([1.1]) + _, anchor, offset, points = prepare_particle_axis_stencil( + position, + axis, + axis_size=6, + shape_factor=1, + bc=BC_PERIODIC, + ghost_cells=True, + ) + + 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]))) + + def test_collapse_axis_stencil_for_inactive_axis(self): + points = jnp.array([[0, 1], [1, 1], [2, 1]]) + weights = jnp.array([[0.2, 0.1], [0.5, 0.6], [0.3, 0.3]]) + + collapsed_points, collapsed_weights = collapse_axis_stencil( + points, + weights, + axis_size=3, + ghost_cells=True, + ) + + self.assertTrue(jnp.array_equal(collapsed_points, jnp.array([[1, 1]]))) + self.assertTrue(jnp.allclose(collapsed_weights, jnp.array([[1.0, 1.0]]))) + + def test_build_axis_helpers_include_ghost_cells(self): + collocated = build_collocated_axis(-1.0, 0.5, 4) + staggered = build_staggered_axis(-1.0, 0.5, 4) + + 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]))) + + +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..5f7f0f4c --- /dev/null +++ b/tests/code_tests/initialization_test.py @@ -0,0 +1,228 @@ +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.func if hasattr(loop, "func") else 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)) + + 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) + + 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)) + + 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) + + 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..da2788fd --- /dev/null +++ b/tests/code_tests/openpmd_test.py @@ -0,0 +1,619 @@ +import unittest +import threading +import time +from unittest.mock import patch + +import jax +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.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles + + +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 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 _simulation_parameters(): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + +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): + world = _world() + return particle_species( + name=name, + N_particles=x1.shape[0], + charge=charge, + mass=mass, + weight=weight, + T=0.0, + x1=x1, + x2=jnp.zeros_like(x1), + x3=jnp.zeros_like(x1), + v1=jnp.ones_like(x1) * 0.1, + v2=jnp.zeros_like(x1), + v3=jnp.zeros_like(x1), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + active_mask=jnp.ones(x1.shape[0], dtype=bool), + dt=world["dt"], + ) + + +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)) + + 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 = to_tiled_particles(species, world, _simulation_parameters()) + species_names = tuple(s.get_name() for s in species) + 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 = to_tiled_particles(species, world, _simulation_parameters()) + species_names = tuple(s.get_name() for s in species) + 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/tiled_esirkepov_integration_test.py b/tests/code_tests/tiled_esirkepov_integration_test.py new file mode 100644 index 00000000..e5aec31a --- /dev/null +++ b/tests/code_tests/tiled_esirkepov_integration_test.py @@ -0,0 +1,151 @@ +import os +import tempfile +import unittest +from contextlib import redirect_stdout +from io import StringIO +from unittest.mock import patch + +import jax +import numpy as np +import toml + +from PyPIC3D import __main__ as pypic_main +from PyPIC3D.__main__ import run_PyPIC3D + + +jax.config.update("jax_enable_x64", True) + + +class TiledEsirkepovIntegrationTest(unittest.TestCase): + def test_short_output_enabled_tiled_esirkepov_run_writes_openpmd_without_vtk(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 esirkepov output 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, + "plot_openpmd_fields": True, + "plot_openpmd_particles": True, + "plot_phasespace": True, + }, + "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, + }, + } + + _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world, _species_config = run_PyPIC3D(config) + _E, _B, J, *_rest = fields + + self.assertEqual(J[0].shape[-3:], (6, 5, 5)) + self.assertEqual(tuple(int(width) for width in world["tile_shape"]), (2, 1, 1)) + for filename in ( + "total_energy.txt", + "electric_field_energy.txt", + "magnetic_field_energy.txt", + "kinetic_energy.txt", + "total_momentum.txt", + ): + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", filename))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "fields.h5"))) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles.h5"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) + self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "phase_space"))) + + def test_main_prints_final_energy_for_tiled_runtime(self): + with tempfile.TemporaryDirectory() as tmpdir: + x_path = os.path.join(tmpdir, "x.npy") + zeros_path = os.path.join(tmpdir, "zeros.npy") + config_path = os.path.join(tmpdir, "final_energy.toml") + np.save(x_path, np.array([0.0])) + np.save(zeros_path, np.zeros(1)) + + config = { + "simulation_parameters": { + "name": "final energy smoke", + "output_dir": tmpdir, + "solver": "electrodynamic_yee", + "Nx": 4, + "Ny": 1, + "Nz": 1, + "x_wind": 1.0, + "y_wind": 1.0, + "z_wind": 1.0, + "dt": 1.0e-10, + "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": 1, + "charge": -1.0, + "mass": 1.0, + "weight": 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, + }, + } + with open(config_path, "w") as f: + toml.dump(config, f) + + with patch("sys.argv", ["PyPIC3D", "--config", config_path]): + stdout = StringIO() + with redirect_stdout(stdout): + pypic_main.main() + + self.assertIn("Final Kinetic Energy", stdout.getvalue()) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "output.toml"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/tiled_esirkepov_test.py b/tests/code_tests/tiled_esirkepov_test.py new file mode 100644 index 00000000..6cec6dce --- /dev/null +++ b/tests/code_tests/tiled_esirkepov_test.py @@ -0,0 +1,861 @@ +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 tests.tiled_particle_fixtures import particle_species +from tests.tiled_particle_fixtures import to_tiled_particles +from PyPIC3D.particles.particle_class import 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) + + 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 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(self, world, x1): + return particle_species( + name="electrons", + N_particles=3, + charge=-1.0, + mass=1.0, + weight=0.5, + T=1.0, + x1=x1, + x2=jnp.zeros_like(x1), + x3=jnp.zeros_like(x1), + v1=jnp.array([0.08, -0.06, 0.05]), + v2=jnp.zeros_like(x1), + v3=jnp.zeros_like(x1), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + def _species_from_arrays(self, world, x, u, weight=0.5, active_mask=None): + return particle_species( + name="electrons", + N_particles=x.shape[0], + charge=-1.0, + mass=1.0, + weight=weight, + T=1.0, + x1=x[:, 0], + x2=x[:, 1], + x3=x[:, 2], + v1=u[:, 0], + v2=u[:, 1], + v3=u[:, 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=active_mask, + ) + + 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 _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, old_species, constants, tile_shape): + world = self._world_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], + } + g = int(world["guard_cells"]) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + _, _, 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) + + old_species = self._species_from_arrays(world, x_old, u) + J_tiles, J_from_tiles = self._assembled_esirkepov_current(world, old_species, constants, tile_shape) + _, J_reference = self._assembled_esirkepov_current( + world, + old_species, + 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) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + x_old = jnp.array([-1.10, -0.10, 1.05]) + old_species = self._species(world, x_old) + + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + 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, + old_species, + 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} + simulation_parameters = { + "particle_tile_nx": world["tile_shape"][0], + "particle_tile_ny": world["tile_shape"][1], + "particle_tile_nz": world["tile_shape"][2], + } + x_old = jnp.array([-1.10, -0.10, 1.05]) + old_species = self._species(world, x_old) + + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + _, _, 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, + old_species, + 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} + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + old_species = self._species(world, jnp.array([-1.10, -0.10, 1.05])) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + + 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) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + 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], + ] + ) + old_species = self._species_from_arrays(world, x_old, u) + tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + 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, + ) + + old_species = self._species_from_arrays( + 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, + old_species, + 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/tiled_fluid_quantities_test.py b/tests/code_tests/tiled_fluid_quantities_test.py new file mode 100644 index 00000000..83b1a1fa --- /dev/null +++ b/tests/code_tests/tiled_fluid_quantities_test.py @@ -0,0 +1,404 @@ +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 PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 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 = to_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 = to_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 = to_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 = particle_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 = particle_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 = to_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 = to_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 = to_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 = to_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 = to_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/tiled_particle_diagnostics_test.py b/tests/code_tests/tiled_particle_diagnostics_test.py new file mode 100644 index 00000000..69ce6682 --- /dev/null +++ b/tests/code_tests/tiled_particle_diagnostics_test.py @@ -0,0 +1,160 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.diagnostics.output_adapters import particles_for_output +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles + + +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 = particle_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 = particle_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 = to_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.get_active_mask() + + original_x = jnp.stack(original.get_forward_position(), axis=-1)[active] + original_u = jnp.stack(original.get_velocity(), axis=-1)[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 = to_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 = to_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 = to_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.get_active_mask() + original_position = jnp.stack(original.get_position(), axis=-1)[active] + 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + species_names = tuple(species.get_name() for species in species_list) + + flattened_species = particles_for_output(tiled_particles, species_config=species_config, species_names=species_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/tiled_particle_initialization_test.py b/tests/code_tests/tiled_particle_initialization_test.py new file mode 100644 index 00000000..522e91f8 --- /dev/null +++ b/tests/code_tests/tiled_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 PyPIC3D.tests.tiled_particle_fixtures import particle_species, to_tiled_particles + + +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_to_tiled_particles_packs_existing_particle_species_list(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 = particle_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 = to_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 = [ + particle_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"], + ), + particle_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 = to_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_to_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 = particle_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 = to_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/tiled_particle_refresh_test.py b/tests/code_tests/tiled_particle_refresh_test.py new file mode 100644 index 00000000..9046b045 --- /dev/null +++ b/tests/code_tests/tiled_particle_refresh_test.py @@ -0,0 +1,231 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 particle_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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 = to_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/tiled_pusher_test.py b/tests/code_tests/tiled_pusher_test.py new file mode 100644 index 00000000..4328276a --- /dev/null +++ b/tests/code_tests/tiled_pusher_test.py @@ -0,0 +1,448 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 = to_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 particle_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 particle_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 = to_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 particle_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 particle_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/tiled_refactor_contract_test.py b/tests/code_tests/tiled_refactor_contract_test.py new file mode 100644 index 00000000..b39758b9 --- /dev/null +++ b/tests/code_tests/tiled_refactor_contract_test.py @@ -0,0 +1,281 @@ +import inspect +import importlib +import importlib.util +import unittest + +import jax.numpy as jnp + +from PyPIC3D import evolve +from PyPIC3D.deposition import Esirkepov +from PyPIC3D.deposition import J_from_rhov +from PyPIC3D.deposition import rho +from PyPIC3D.boundary_conditions import boundaryconditions +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics import output_adapters +from PyPIC3D.initialization import build_tiled_array, initialize_fields +from PyPIC3D.solvers import electrostatic_yee +from PyPIC3D.solvers import first_order_yee +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid + + +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) + ntx = int(world["Nx"]) // tile_nx + nty = int(world["Ny"]) // tile_ny + ntz = int(world["Nz"]) // tile_nz + + 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) + + +class TestTiledRefactorContracts(unittest.TestCase): + def _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}, + "tile_shape": (2, 3, 2), + "guard_cells": 2, + } + vertex_grid, center_grid = build_yee_grid(world) + world["grids"] = {"vertex": vertex_grid, "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 _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) + return jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (ii + 2.0 * jj - kk)) + + def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): + world = self._world() + tile_shape = (2, 3, 2) + E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) + E_tiles = tuple(tile_scalar_field(component, world, tile_shape) for component in E) + + refreshed = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) + self.assertEqual(len(refreshed), 3) + for component in refreshed: + self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) + + def test_evolve_loops_read_runtime_controls_from_world(self): + electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) + electrostatic_signature = inspect.signature(evolve.time_loop_electrostatic) + for name in ("curl_func", "J_func", "tile_shape", "g", "current_deposition"): + self.assertNotIn(name, electrodynamic_signature.parameters) + for name in ("curl_func", "J_func", "tile_shape", "g"): + self.assertNotIn(name, electrostatic_signature.parameters) + + def test_tiled_yee_updates_read_tile_metadata_from_world(self): + update_E_signature = inspect.signature(first_order_yee.update_E) + update_B_signature = inspect.signature(first_order_yee.update_B) + for name in ("curl_func", "tile_shape", "g"): + self.assertNotIn(name, update_E_signature.parameters) + self.assertNotIn(name, update_B_signature.parameters) + self.assertFalse(hasattr(first_order_yee, "update_tiled_E")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_B")) + self.assertFalse(hasattr(first_order_yee, "tiled_grid_axes_from_world")) + self.assertFalse(hasattr(first_order_yee, "tile_grid_axes")) + self.assertFalse(hasattr(first_order_yee, "tile_vector_field")) + self.assertFalse(hasattr(first_order_yee, "stack_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "unstack_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_for_pml")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_for_pml")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells")) + self.assertFalse(hasattr(first_order_yee, "apply_tiled_conducting_bc")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells_periodic")) + self.assertFalse(hasattr(first_order_yee, "empty_tiled_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "empty_tiled_vector_field")) + self.assertFalse(hasattr(first_order_yee, "tile_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "assemble_tiled_scalar_field")) + self.assertFalse(hasattr(first_order_yee, "assemble_tiled_vector_field")) + + def test_global_assembly_lives_at_output_boundary(self): + self.assertTrue(hasattr(output_adapters, "assemble_tiled_scalar_field")) + self.assertTrue(hasattr(output_adapters, "assemble_tiled_vector_field")) + + source = inspect.getsource(output_adapters) + self.assertNotIn("from PyPIC3D.solvers.yee_tiled import", source) + + def test_ghost_cell_logic_lives_in_boundary_conditions_ghost_cells(self): + for name in ( + "update_tiled_ghost_cells", + "update_tiled_vector_ghost_cells", + "update_tiled_ghost_cells_for_pml", + "update_tiled_vector_ghost_cells_for_pml", + "fold_tiled_ghost_cells", + "fold_tiled_vector_ghost_cells", + "apply_tiled_conducting_bc", + ): + self.assertTrue(hasattr(ghost_cells, name)) + + for name in ( + "update_ghost_cells", + "fold_ghost_cells", + "apply_conducting_bc", + "apply_scalar_conducting_bc", + ): + self.assertFalse(hasattr(boundaryconditions, name)) + + def test_tiled_array_initialization_lives_in_initialization(self): + world = self._world() + world["tile_shape"] = (2, 3, 2) + world["guard_cells"] = 2 + + array = build_tiled_array(world, dtype=jnp.float32) + self.assertEqual(array.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(array.dtype, jnp.float32) + self.assertTrue(jnp.all(array == 0.0)) + + def test_initialize_fields_builds_tiled_field_state_from_world(self): + world = self._world() + world["tile_shape"] = (2, 3, 2) + world["guard_cells"] = 2 + + E, B, J, phi, rho_tiles = initialize_fields(world) + for vector_field in (E, B, J): + self.assertEqual(len(vector_field), 3) + for component in vector_field: + self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(component.dtype, jnp.float64) + self.assertTrue(jnp.all(component == 0.0)) + for scalar_field in (phi, rho_tiles): + self.assertEqual(scalar_field.shape, (4, 2, 2, 6, 7, 6)) + self.assertEqual(scalar_field.dtype, jnp.float64) + self.assertTrue(jnp.all(scalar_field == 0.0)) + + def test_grid_builders_live_in_utilities(self): + from PyPIC3D import utils + from PyPIC3D.utilities import grids + + self.assertFalse(hasattr(utils, "build_yee_grid")) + self.assertFalse(hasattr(utils, "build_collocated_grid")) + self.assertTrue(hasattr(grids, "build_yee_grid")) + self.assertTrue(hasattr(grids, "build_collocated_grid")) + self.assertFalse(hasattr(grids, "tile_grid_axes")) + self.assertTrue(hasattr(grids, "build_tiled_yee_grids")) + + def test_tiled_benchmarking_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) + + def test_deprecated_fdtd_errors_and_root_boris_modules_are_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.solvers.fdtd")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.errors")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.boris")) + self.assertIsNotNone(importlib.util.find_spec("PyPIC3D.pusher.boris")) + + def test_old_particle_species_surfaces_are_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.species_class")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.tiled_particle_diagnostics")) + + from PyPIC3D.particles import particle_initialization + + self.assertFalse(hasattr(particle_initialization, "to_tiled_particles")) + source = inspect.getsource(particle_initialization.load_particles_from_toml) + self.assertNotIn("particle_species", source) + + def test_particle_push_owns_tiled_pusher_contract(self): + particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") + + self.assertIsNone(importlib.util.find_spec("PyPIC3D.pusher.tiled_pusher")) + signature = inspect.signature(particle_push.particle_push) + for name in ("grid", "staggered_grid", "tile_shape", "g"): + self.assertNotIn(name, signature.parameters) + + source = inspect.getsource(particle_push.particle_push) + self.assertIn('world["grids"]["tiled_center_grid"]', source) + self.assertIn('world["grids"]["tiled_vertex_grid"]', source) + self.assertNotIn("tiled_grid_axes_from_world", source) + self.assertNotIn("tile_grid_axes", source) + self.assertNotIn("get_forward_position", source) + self.assertNotIn("set_velocity", source) + + def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): + source = inspect.getsource(evolve.time_loop_electrodynamic) + self.assertNotIn("assemble_tiled_scalar_field", source) + self.assertNotIn("assemble_tiled_vector_field", source) + self.assertNotIn("fields_for_output", source) + + def test_electrostatic_solver_uses_single_tile_arrays_directly(self): + source = inspect.getsource(electrostatic_yee.calculate_tiled_electrostatic_fields) + self.assertNotIn("assemble_tiled_scalar_field", source) + self.assertNotIn("tile_scalar_field", source) + self.assertIn("[0, 0, 0]", source) + + def test_flat_current_and_electrodynamic_references_are_removed(self): + self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) + self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_tiled")) + self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) + self.assertFalse(hasattr(evolve, "_time_loop_electrostatic_global_reference")) + + def test_direct_current_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(J_from_rhov.J_from_rhov) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + + def test_esirkepov_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(Esirkepov.Esirkepov_current) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + + def test_tiled_esirkepov_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.esirkepov_tiled")) + + def test_rho_api_reads_tile_metadata_from_world(self): + signature = inspect.signature(rho.compute_rho) + self.assertNotIn("tile_shape", signature.parameters) + self.assertNotIn("g", signature.parameters) + + def test_flat_rho_and_tiled_rho_module_are_removed(self): + self.assertFalse(hasattr(rho, "_compute_rho_flat")) + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.rho_tiled")) + + def test_current_methods_module_is_removed(self): + self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.current_methods")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/code_tests/tiled_rho_test.py b/tests/code_tests/tiled_rho_test.py new file mode 100644 index 00000000..410e88f9 --- /dev/null +++ b/tests/code_tests/tiled_rho_test.py @@ -0,0 +1,363 @@ +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 PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 = particle_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 = particle_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 = to_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.v1 = jnp.zeros_like(species.v1) + species.v2 = jnp.zeros_like(species.v2) + species.v3 = jnp.zeros_like(species.v3) + 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..060415da --- /dev/null +++ b/tests/code_tests/utils_test.py @@ -0,0 +1,499 @@ +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 PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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, +) +from PyPIC3D.tests.tiled_particle_fixtures import particle_species + +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} + 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="periodic", + ) + species.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) + + tiled_particles, species_config = to_tiled_particles([species], world, {"particle_tile_nx": 1, "particle_tile_ny": 1, "particle_tile_nz": 1}) + + _, _, 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 = particle_species( + name="beam electrons", + N_particles=2, + charge=-1.0, + mass=1.0, + weight=2.0, + T=0.0, + x1=jnp.array([-0.25, 0.25]), + x2=jnp.array([0.0, 0.0]), + x3=jnp.array([0.0, 0.0]), + v1=jnp.array([0.1, 0.2]), + v2=jnp.array([0.0, 0.0]), + v3=jnp.array([0.0, 0.0]), + xwind=1.0, + ywind=1.0, + zwind=1.0, + dx=0.25, + dy=0.5, + dz=1.0, + dt=0.1, + ) + 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 = to_tiled_particles([species], world, 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 = particle_species( + name="electrons", + N_particles=1, + charge=-1.0, + mass=1.0, + weight=1.0, + T=0.0, + x1=jnp.array([0.0]), + x2=jnp.array([0.0]), + x3=jnp.array([0.0]), + v1=jnp.array([0.2]), + v2=jnp.array([0.0]), + v3=jnp.array([0.0]), + xwind=1.0, + ywind=1.0, + zwind=1.0, + dx=0.25, + dy=0.5, + dz=1.0, + dt=0.1, + ) + 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 = to_tiled_particles([species], world, 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_tiled_test.py b/tests/code_tests/yee_tiled_test.py new file mode 100644 index 00000000..5d17deb0 --- /dev/null +++ b/tests/code_tests/yee_tiled_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/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/physics_tests/boris_test.py b/tests/physics_tests/boris_test.py new file mode 100644 index 00000000..b5e45de4 --- /dev/null +++ b/tests/physics_tests/boris_test.py @@ -0,0 +1,452 @@ +import unittest +import jax +import jax.numpy as jnp +import sys +import os +from functools import partial + + +from PyPIC3D.pusher.boris import boris_single_particle, interpolate_field_to_particles +from PyPIC3D.pusher.higuera_cary import higuera_cary_single_particle +from PyPIC3D.utils import mae, convergence_test + +jax.config.update("jax_enable_x64", True) + +class TestBorisMethods(unittest.TestCase): + + def setUp(self): + self.q = 1.0 + self.m = 1.0 + self.x = jnp.array([0.0]) + self.y = jnp.array([0.0]) + self.z = jnp.array([0.0]) + # single particle at the origin + + self.vx = jnp.array([1.0]) + self.vy = jnp.array([0.0]) + self.vz = jnp.array([0.0]) + # single particle moving in the x direction + + self.Ex = jnp.zeros( (5, 5, 5) ) + self.Ey = jnp.zeros( (5, 5, 5) ) + self.Ez = jnp.zeros( (5, 5, 5) ) + self.Bx = jnp.zeros( (5, 5, 5) ) + self.By = jnp.zeros( (5, 5, 5) ) + self.Bz = jnp.ones( (5, 5, 5) ) + # grid of 5x5x5 with a uniform magnetic field in the z direction + + self.grid = jnp.arange(-1/2, 1/2, (1/5)), jnp.arange(-1/2, 1/2, (1/5)), jnp.arange(-1/2, 1/2, (1/5)) + self.staggered_grid = jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)), jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)), jnp.arange(-1/2 + (1/5)/2, 1/2 + (1/5)/2, (1/5)) + self.dt = 0.1 + # grid and staggered grid for a 5x5x5 grid with a spacing of 1/5 and a timestep of 0.1 + + def test_boris(self): + + efield_atx = interpolate_field_to_particles(self.Ex, self.x, self.y, self.z, self.grid, shape_factor=1) + efield_aty = interpolate_field_to_particles(self.Ey, self.x, self.y, self.z, self.grid, shape_factor=1) + efield_atz = interpolate_field_to_particles(self.Ez, self.x, self.y, self.z, self.grid, shape_factor=1) + # calculate the electric field at the particle positions + bfield_atx = interpolate_field_to_particles(self.Bx, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) + bfield_aty = interpolate_field_to_particles(self.By, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) + bfield_atz = interpolate_field_to_particles(self.Bz, self.x, self.y, self.z, self.staggered_grid, shape_factor=1) + + boris_vmap = jax.vmap(boris_single_particle, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None)) + newvx, newvy, newvz = boris_vmap(self.vx, self.vy, self.vz, efield_atx, efield_aty, efield_atz, bfield_atx, bfield_aty, bfield_atz, self.q, self.m, self.dt, None) + + self.assertIsInstance(newvx, jnp.ndarray) + self.assertIsInstance(newvy, jnp.ndarray) + self.assertIsInstance(newvz, jnp.ndarray) + # make sure the velocities are jax arrays + self.assertTrue( jnp.allclose(jnp.abs(newvx), 1.0, rtol = 1e-2)) + # make sure the x velocity is unchanged + self.assertLess( newvy, 0.1 ) + # make sure the y velocity is unchanged + self.assertLess( newvz, 0.1 ) + # make sure the z velocity is 0.0 from the magnetic field + + def test_interpolate_field_to_particles_with_ghost_cells_on_reduced_axes(self): + field = jnp.full((6, 3, 3), 7.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, 6), + jnp.array([-1.0, 0.0, 1.0]), + jnp.array([-1.0, 0.0, 1.0]), + ) + + interpolated = interpolate_field_to_particles( + field, + x, + y, + z, + grid, + shape_factor=2, + ghost_cells=True, + ) + + 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): + vx, vy, vz = 1.0, 0.0, 0.0 + x, y, z = 0.0, 0.0, 0.0 + + q = 1.0 + m = 1.0 + E = jnp.array([0.0, 0.0, 0.0]) + B = jnp.array([0.0, 1.0, 0.0]) + dt = 0.1 + n_steps = 5000 + + xs = [] + ys = [] + zs = [] + + for i in range(n_steps): + vx, vy, vz = boris_single_particle(vx, vy, vz, E[0], E[1], E[2], B[0], B[1], B[2], q, m, dt, None) + x += vx * dt + y += vy * dt + z += vz * dt + + xs.append(x) + ys.append(y) + zs.append(z) + + + def measure_xz_radius(xs, zs): + """ + Measure the radius of the XZ cut by calculating the distance of each point + in the XZ plane from the origin and returning the average and maximum radius. + + Parameters: + xs (list): List of x-coordinates. + zs (list): List of z-coordinates. + + Returns: + tuple: A tuple containing the average radius and maximum radius. + """ + xs = jnp.array(xs) + zs = jnp.array(zs) + # Calculate the distance of each point from the origin in the XZ plane + radii = [jnp.sqrt(x**2 + z**2) for x, z in zip(xs, zs)] + + # Compute the average and maximum radius + avg_radius = jnp.mean(jnp.asarray(radii)) + + return avg_radius + + avg_radius = measure_xz_radius(xs, zs) + #print(f"Average radius: {avg_radius}") + self.assertTrue( jnp.isclose(avg_radius, 1.28, atol=0.5) ) + + def test_higuera_cary_zero_field_preserves_velocity(self): + vx, vy, vz = higuera_cary_single_particle( + 0.2, -0.1, 0.05, + 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # zero fields should not change the particle velocity + self.assertTrue(jnp.allclose(jnp.array([vx, vy, vz]), jnp.array([0.2, -0.1, 0.05]))) + + def test_higuera_cary_magnetic_field_preserves_speed(self): + vx0, vy0, vz0 = 0.2, 0.0, 0.0 + vx, vy, vz = higuera_cary_single_particle( + vx0, vy0, vz0, + 0.0, 0.0, 0.0, + 0.0, 0.0, 1.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # magnetic rotation should conserve speed + old_speed = jnp.sqrt(vx0**2 + vy0**2 + vz0**2) + new_speed = jnp.sqrt(vx**2 + vy**2 + vz**2) + self.assertTrue(jnp.allclose(new_speed, old_speed, rtol=1e-12, atol=1e-12)) + + def test_higuera_cary_electric_field_accelerates_along_field(self): + vx, vy, vz = higuera_cary_single_particle( + 0.0, 0.0, 0.0, + 1.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + 1.0, 1.0, 0.1, + {"C": 10.0}, + ) + # electric field in x should accelerate the particle in x only + self.assertGreater(vx, 0.0) + self.assertTrue(jnp.allclose(jnp.array([vy, vz]), jnp.array([0.0, 0.0]))) + + def test_higuera_cary_vectorizes_with_per_particle_charge_and_mass(self): + higuera_cary_vmap = jax.vmap( + higuera_cary_single_particle, + in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None), + ) + newvx, newvy, newvz = higuera_cary_vmap( + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([1.0, 1.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([0.0, 0.0]), + jnp.array([1.0, 2.0]), + jnp.array([1.0, 1.0]), + 0.1, + {"C": 10.0}, + ) + # per-particle charge and mass should work through the vectorized path + self.assertEqual(newvx.shape, (2,)) + self.assertGreater(newvx[1], newvx[0]) + self.assertTrue(jnp.allclose(newvy, jnp.zeros(2))) + self.assertTrue(jnp.allclose(newvz, jnp.zeros(2))) + + + def test_trilinear_convergence(self): + + # Run convergence tests for each interpolation method + ################### BASIC TRIGONOMETRIC TEST ############################# + slope = convergence_test(partial(interpolation_wave_test, shape_factor=1)) + # measure the convergence rate of the trilinear interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + ################### POLYNOMIAL TEST ####################################### + slope = convergence_test(partial(interpolation_polynomial_test, shape_factor=1)) + # measure the convergence rate of the trilinear interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + ################### HIGH-FREQUENCY TEST ################################### + slope = convergence_test(partial(interpolation_oscillatory_test, shape_factor=1)) + # measure the convergence rate of the trilinear interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + def test_quadratic_convergence(self): + + ################### BASIC TRIGONOMETRIC TEST ############################### + slope = convergence_test(partial(interpolation_wave_test, shape_factor=2)) + # measure the convergence rate of the second-order shape interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + ################### POLYNOMIAL TEST ######################################## + slope = convergence_test(partial(interpolation_polynomial_test, shape_factor=2)) + # measure the convergence rate of the second-order shape interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + ################### HIGH-FREQUENCY TEST #################################### + slope = convergence_test(partial(interpolation_oscillatory_test, shape_factor=2)) + # measure the convergence rate of the second-order shape interpolation + self.assertTrue(slope > 1.9) + # assert that the order of the error is at least 2nd order + ############################################################################ + + +def interpolation_wave_test(nx, shape_factor): + # Define a symmetric domain + x_wind = 2.0 * jnp.pi + y_wind = 2.0 * jnp.pi + z_wind = 2.0 * jnp.pi + + # Create uniform grid + x_grid = jnp.linspace(0, x_wind, nx) + y_grid = jnp.linspace(0, y_wind, nx) + z_grid = jnp.linspace(0, z_wind, nx) + + dx = x_wind / (nx - 1) + dy = y_wind / (nx - 1) + dz = z_wind / (nx - 1) + + # Create meshgrid for field evaluation + X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') + + # Define analytical test function: smooth trigonometric function + # f(x,y,z) = sin(x) * cos(y) * sin(z) + analytical_field = jnp.sin(X) * jnp.cos(Y) * jnp.sin(Z) + + grid = (x_grid, y_grid, z_grid) + + # Create test points that are offset from grid points + # This tests the interpolation accuracy between grid points + n_test = nx + x_test = jnp.linspace(dx/3, x_wind - dx/3, n_test) + y_test = jnp.linspace(dy/3, y_wind - dy/3, n_test) + z_test = jnp.linspace(dz/3, z_wind - dz/3, n_test) + + X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') + + # Flatten for vectorized interpolation + x_test_flat = X_test.flatten() + y_test_flat = Y_test.flatten() + z_test_flat = Z_test.flatten() + + # Compute analytical values at test points + analytical_values = jnp.sin(x_test_flat) * jnp.cos(y_test_flat) * jnp.sin(z_test_flat) + + # Interpolate values at test points + interpolated_values = interpolate_field_to_particles( + analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor + ) + + # Compute mean squared error + error = mae(interpolated_values, analytical_values) + + return error, dx + + +def interpolation_polynomial_test(nx, shape_factor): + + # Define domain + x_wind = 1.0 + y_wind = 1.0 + z_wind = 1.0 + + # Create uniform grid + x_grid = jnp.linspace(0, x_wind, nx) + y_grid = jnp.linspace(0, y_wind, nx) + z_grid = jnp.linspace(0, z_wind, nx) + + dx = x_wind / (nx - 1) + # create grid spacing + + # Create meshgrid + X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') + + # Use a polynomial test function: f(x,y,z) = x³y + y³z + z³x + # This has continuous derivatives and tests interpolation of curved surfaces + analytical_field = X**3 * Y + Y**3 * Z + Z**3 * X + + grid = (x_grid, y_grid, z_grid) + + # Test at mid-points between grid points (maximum interpolation error) + n_test = max(4, nx) # Ensure we have enough test points + x_test = jnp.linspace(dx/2, x_wind - dx/2, n_test) + y_test = jnp.linspace(dx/2, y_wind - dx/2, n_test) + z_test = jnp.linspace(dx/2, z_wind - dx/2, n_test) + + X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') + + # Flatten + x_test_flat = X_test.flatten() + y_test_flat = Y_test.flatten() + z_test_flat = Z_test.flatten() + + # Analytical values + analytical_values = (x_test_flat**3 * y_test_flat + + y_test_flat**3 * z_test_flat + + z_test_flat**3 * x_test_flat) + + # Interpolated values + interpolated_values = interpolate_field_to_particles( + analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor + ) + + # Compute error + error = mae(interpolated_values, analytical_values) + + return error, dx + + +def interpolation_oscillatory_test(nx, shape_factor): + + # Define domain + x_wind = jnp.pi + y_wind = jnp.pi + z_wind = jnp.pi + + # Create grid + x_grid = jnp.linspace(0, x_wind, nx) + y_grid = jnp.linspace(0, y_wind, nx) + z_grid = jnp.linspace(0, z_wind, nx) + + dx = x_wind / (nx - 1) + # create grid spacing + + # Create meshgrid + X, Y, Z = jnp.meshgrid(x_grid, y_grid, z_grid, indexing='ij') + + # High-frequency test function: f(x,y,z) = cos(2x) * sin(3y) * cos(4z) + analytical_field = jnp.cos(2*X) * jnp.sin(3*Y) * jnp.cos(4*Z) + + grid = (x_grid, y_grid, z_grid) + + # Test points slightly offset from grid + n_test = max(6, nx) + x_test = jnp.linspace(dx/4, x_wind - dx/4, n_test) + y_test = jnp.linspace(dx/4, y_wind - dx/4, n_test) + z_test = jnp.linspace(dx/4, z_wind - dx/4, n_test) + + X_test, Y_test, Z_test = jnp.meshgrid(x_test, y_test, z_test, indexing='ij') + + # Flatten + x_test_flat = X_test.flatten() + y_test_flat = Y_test.flatten() + z_test_flat = Z_test.flatten() + + # Analytical values + analytical_values = (jnp.cos(2*x_test_flat) * + jnp.sin(3*y_test_flat) * + jnp.cos(4*z_test_flat)) + + # Interpolated values + interpolated_values = interpolate_field_to_particles( + analytical_field, x_test_flat, y_test_flat, z_test_flat, grid, shape_factor + ) + + # Compute error + error = mae(interpolated_values, analytical_values) + + return error, dx + +if __name__ == '__main__': + unittest.main() diff --git a/tests/physics_tests/electrostatic_yee_test.py b/tests/physics_tests/electrostatic_yee_test.py new file mode 100644 index 00000000..3879d8ca --- /dev/null +++ b/tests/physics_tests/electrostatic_yee_test.py @@ -0,0 +1,164 @@ +import unittest + +import jax +import jax.numpy as jnp + +from PyPIC3D.tests.tiled_particle_fixtures import particle_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 = [ + 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", + 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)) + + +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..de4b6859 --- /dev/null +++ b/tests/physics_tests/pml_test.py @@ -0,0 +1,590 @@ +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) + + 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 + + +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/tiled_electrostatic_test.py b/tests/physics_tests/tiled_electrostatic_test.py new file mode 100644 index 00000000..68eb6863 --- /dev/null +++ b/tests/physics_tests/tiled_electrostatic_test.py @@ -0,0 +1,512 @@ +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_PERIODIC +from PyPIC3D.deposition.rho import compute_rho +from PyPIC3D.evolve import time_loop_electrostatic +from PyPIC3D.initialization import initialize_simulation +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from PyPIC3D.particles.particle_class import TiledParticles +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field +from PyPIC3D.utilities.grids import build_collocated_grid, 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 TestTiledElectrostatic(unittest.TestCase): + def _build_world(self, shape_factor=1): + world = { + "Nx": 8, + "Ny": 1, + "Nz": 1, + "dx": 0.5, + "dy": 1.0, + "dz": 1.0, + "dt": 0.01, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "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} + world["tile_shape"] = (2, 1, 1) + world = self._world_with_tiled_grids(world, world["tile_shape"]) + return world + + 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 _simulation_parameters(self): + return { + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + } + + 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 _one_tile_shape(self, world): + return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) + + def _empty_fields(self, world): + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + zero = jnp.zeros(shape) + E = (zero, zero, zero) + B = (zero, zero, zero) + J = (zero, zero, zero) + external_fields = (E, B) + return E, B, J, zero, zero, external_fields + + def _species(self, world): + return particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + def _neutral_species(self, world): + electrons = self._species(world) + ions = particle_species( + name="ions", + N_particles=4, + charge=1.0, + mass=4.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.25, -0.25, 0.75, 1.75]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([-0.03, 0.04, -0.02, 0.01]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + return [electrons, ions] + + def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): + current_slots = tiled_particles.active.shape[-1] + if current_slots >= min_slots: + return tiled_particles + + pad_width = [(0, 0)] * tiled_particles.x.ndim + pad_width[-2] = (0, min_slots - current_slots) + x = jnp.pad(tiled_particles.x, pad_width) + u = jnp.pad(tiled_particles.u, pad_width) + + scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim + scalar_pad_width[-1] = (0, min_slots - current_slots) + + return TiledParticles( + x=x, + u=u, + active=jnp.pad(tiled_particles.active, scalar_pad_width), + ) + + def _build_tiled_state(self, particles, fields, world, tile_shape, min_slots=None): + tiled_particles, species_config = to_tiled_particles( + particles, + world, + self._simulation_parameters_for_tile_shape(tile_shape), + ) + if min_slots is not None: + tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots) + + E, B, J, rho, phi, external_fields = fields + g = int(world["guard_cells"]) + tiled_fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + tile_vector_field(J, world, tile_shape, num_guard_cells=g), + tile_scalar_field(rho, world, tile_shape, num_guard_cells=g), + tile_scalar_field(phi, world, tile_shape, num_guard_cells=g), + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), + ), + None, + jnp.asarray(False), + ) + + return tiled_particles, species_config, tiled_fields + + def _long_two_stream_species(self, world): + x0 = jnp.linspace( + -0.5 * world["x_wind"], + 0.5 * world["x_wind"], + world["Nx"], + endpoint=False, + ) + 0.5 * world["dx"] + x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) + zeros = jnp.zeros_like(x0) + beam_speed = 0.08 + + species_kwargs = { + "N_particles": world["Nx"], + "mass": 1.0, + "weight": 1.0, + "T": 0.0, + "x1": x0, + "x2": zeros, + "x3": zeros, + "v2": zeros, + "v3": zeros, + "xwind": world["x_wind"], + "ywind": world["y_wind"], + "zwind": world["z_wind"], + "dx": world["dx"], + "dy": world["dy"], + "dz": world["dz"], + "dt": world["dt"], + } + return [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + + def _long_two_stream_state(self): + world = { + "Nx": 16, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.002, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "guard_cells": 2, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, + } + center_grid, vertex_grid = build_collocated_grid(world) + world["grids"] = {"center": center_grid, "vertex": vertex_grid} + + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = self._one_tile_shape(world) + world = self._world_with_tiled_grids(world, tile_shape) + Nt = 1500 + + E, B, J, rho, phi, external_fields = self._empty_fields(world) + fields = (E, B, J, rho, phi, external_fields) + tiled_particles, species_config, tiled_fields = self._build_tiled_state( + self._long_two_stream_species(world), + fields, + world, + tile_shape, + ) + + return tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt + + def test_scalar_tiles_round_trip_global_rho_and_phi(self): + world = self._build_world() + tile_shape = world["tile_shape"] + scalar = jnp.arange((world["Nx"] + 2) * (world["Ny"] + 2) * (world["Nz"] + 2), dtype=float) + scalar = scalar.reshape((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) + + scalar_tiles = tile_scalar_field(scalar, world, tile_shape) + scalar_from_tiles = assemble_tiled_scalar_field(scalar_tiles, world, tile_shape) + + self.assertTrue(jnp.allclose(scalar_from_tiles[1:-1, 1:-1, 1:-1], scalar[1:-1, 1:-1, 1:-1])) + + def test_tiled_rho_assembles_to_global_rho(self): + world = self._build_world(shape_factor=1) + constants = {"alpha": 1.0} + species = self._species(world) + tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) + _, _, _, rho, _, _ = self._empty_fields(world) + rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) + + tiled_rho = compute_rho( + tiled_particles, + species_config, + rho_tiles, + constants, + world, + ) + assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) + + reference_world = self._world_with_tiled_grids(world, self._one_tile_shape(world)) + reference_particles, reference_species_config = to_tiled_particles( + [self._species(reference_world)], + reference_world, + self._simulation_parameters_for_tile_shape(reference_world["tile_shape"]), + ) + reference_rho_tiles = tile_scalar_field(rho, reference_world, reference_world["tile_shape"], num_guard_cells=int(reference_world["guard_cells"])) + reference_rho_tiles = compute_rho( + reference_particles, + reference_species_config, + reference_rho_tiles, + constants, + reference_world, + ) + reference_rho = assemble_tiled_scalar_field( + reference_rho_tiles, + reference_world, + reference_world["tile_shape"], + num_guard_cells=int(reference_world["guard_cells"]), + ) + + self.assertTrue( + jnp.allclose( + assembled_rho[1:-1, 1:-1, 1:-1], + reference_rho[1:-1, 1:-1, 1:-1], + rtol=1.0e-12, + atol=1.0e-12, + ) + ) + + def test_tiled_electrostatic_step_uses_single_tile_state(self): + world = self._build_world(shape_factor=1) + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + E, B, J, rho, phi, external_fields = self._empty_fields(world) + fields = (E, B, J, rho, phi, external_fields) + tile_shape = self._one_tile_shape(world) + world = self._world_with_tiled_grids(world, tile_shape) + tiled_particles, species_config, tiled_fields = self._build_tiled_state( + self._neutral_species(world), + fields, + world, + tile_shape, + ) + tiled_particles, tiled_fields = time_loop_electrostatic( + tiled_particles, + species_config, + tiled_fields, + world, + constants, + solver="electrostatic", + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(B_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(J_tiles[0].shape[:3], (1, 1, 1)) + self.assertEqual(rho_tiles.shape[:3], (1, 1, 1)) + self.assertEqual(phi_tiles.shape[:3], (1, 1, 1)) + self.assertFalse(bool(tiled_fields[-1])) + self.assertTrue(jnp.all(jnp.isfinite(rho_tiles))) + self.assertTrue(jnp.all(jnp.isfinite(phi_tiles))) + + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", + "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step single-tile electrostatic check.", + ) + def test_long_two_stream_single_tile_electrostatic_remains_finite(self): + tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt = self._long_two_stream_state() + + tiled_loop = jax.jit( + time_loop_electrostatic, + static_argnames=("solver", "relativistic", "particle_pusher"), + ) + + for step in range(1, Nt + 1): + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + species_config, + tiled_fields, + world, + constants, + solver="electrostatic", + relativistic=False, + particle_pusher="boris", + ) + + E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields + self.assertEqual(tuple(world["tile_shape"]), tuple(tile_shape)) + self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) + self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") + for component in E_tiles + B_tiles + J_tiles: + self.assertTrue(jnp.all(jnp.isfinite(component)), f"step {step}: nonfinite vector field") + self.assertTrue(jnp.all(jnp.isfinite(rho_tiles)), f"step {step}: nonfinite rho") + self.assertTrue(jnp.all(jnp.isfinite(phi_tiles)), f"step {step}: nonfinite phi") + + def test_initialize_simulation_accepts_tiled_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": "tiled 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)) + self.assertIn("tiled_center_grid", world["grids"]) + self.assertIn("tiled_vertex_grid", world["grids"]) + self.assertEqual(int(world["guard_cells"]), 2) + self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (1, 1, 1, 12)) + self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (1, 1, 1, 5)) + 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/tiled_yee_integration_test.py b/tests/physics_tests/tiled_yee_integration_test.py new file mode 100644 index 00000000..2c544984 --- /dev/null +++ b/tests/physics_tests/tiled_yee_integration_test.py @@ -0,0 +1,1411 @@ +import os +import tempfile +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import toml + +from PyPIC3D.evolve import time_loop_electrodynamic +from PyPIC3D.initialization import initialize_simulation, validate_field_solver +from PyPIC3D.tests.tiled_particle_fixtures import particle_species +from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from PyPIC3D.particles.particle_class import TiledParticles +from PyPIC3D.boundary_conditions import ghost_cells +from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field +from PyPIC3D.solvers.first_order_yee import ( + update_B, + update_E, +) +from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid +from PyPIC3D.utils import compute_energy +from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC + + +jax.config.update("jax_enable_x64", True) + +ROUND_OFF_RTOL = 1.0e-11 +ROUND_OFF_ATOL = 1.0e-11 + + +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 TestTiledYeeIntegration(unittest.TestCase): + def _build_world(self, boundary_conditions=None): + if boundary_conditions is None: + boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} + world = { + "Nx": 8, + "Ny": 1, + "Nz": 1, + "dx": 0.5, + "dy": 1.0, + "dz": 1.0, + "dt": 0.01, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 1, + "guard_cells": 2, + "current_deposition": "direct", + "current_filter": "none", + "boundary_conditions": boundary_conditions, + } + 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): + 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_fields(self, 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)) + phi = jnp.zeros(shape) + rho = jnp.zeros(shape) + external_fields = (E, B) + return E, B, J, rho, phi, external_fields, None + + def _species(self, world): + return particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.zeros(4), + v3=jnp.zeros(4), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + + 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.lexsort((x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _long_two_stream_state(self, current_deposition="direct", current_filter="none"): + world = { + "Nx": 16, + "Ny": 1, + "Nz": 1, + "dx": 0.25, + "dy": 1.0, + "dz": 1.0, + "dt": 0.002, + "x_wind": 4.0, + "y_wind": 1.0, + "z_wind": 1.0, + "shape_factor": 2, + "guard_cells": 2, + "current_deposition": current_deposition, + "current_filter": current_filter, + "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, + "particle_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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (4, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) + Nt = 1500 + + x0 = jnp.linspace( + -0.5 * world["x_wind"], + 0.5 * world["x_wind"], + world["Nx"], + endpoint=False, + ) + 0.5 * world["dx"] + x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) + zeros = jnp.zeros_like(x0) + beam_speed = 0.08 + + species_kwargs = { + "N_particles": world["Nx"], + "mass": 1.0, + "weight": 1.0, + "T": 0.0, + "x1": x0, + "x2": zeros, + "x3": zeros, + "v2": zeros, + "v3": zeros, + "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 = [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + reference_species = [ + particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), + particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), + ] + + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) + tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) + g = int(tiled_world["guard_cells"]) + if current_deposition == "esirkepov": + J_template = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) + J_tiles = tuple(jnp.zeros_like(component) for component in J_template) + else: + J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) + + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=g), + tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=g), + J_tiles, + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=g), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + reference_species, + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + current_deposition, + ) + reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + tiled_world, + reference_tile_shape, + tile_shape, + step=0, + ) + + return ( + reference_particles, + reference_species_config, + reference_fields, + reference_tile_shape, + reference_world, + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + tile_shape, + Nt, + ) + + def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): + current_slots = tiled_particles.active.shape[-1] + if current_slots >= min_slots: + return tiled_particles + + pad_width = [(0, 0)] * tiled_particles.x.ndim + pad_width[-2] = (0, min_slots - current_slots) + x = jnp.pad(tiled_particles.x, pad_width) + u = jnp.pad(tiled_particles.u, pad_width) + + scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim + scalar_pad_width[-1] = (0, min_slots - current_slots) + active = jnp.pad(tiled_particles.active, scalar_pad_width) + + return TiledParticles( + x=x, + u=u, + active=active, + ) + + def _tiled_species_rows(self, tiled_particles, species_index): + active = tiled_particles.active[:, :, :, species_index, :] + x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] + order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) + return x[order], u[order] + + def _build_tiled_state(self, particles, fields, world, tile_shape, current_deposition="direct"): + simulation_parameters = { + "particle_tile_nx": tile_shape[0], + "particle_tile_ny": tile_shape[1], + "particle_tile_nz": tile_shape[2], + } + tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) + E, B, J, rho, phi, external_fields, pml_state = fields + g = int(world["guard_cells"]) + if current_deposition == "esirkepov": + J_template = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + J_tiles = tuple(jnp.zeros_like(component) for component in J_template) + else: + J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) + tiled_fields = ( + tile_vector_field(E, world, tile_shape, num_guard_cells=g), + tile_vector_field(B, world, tile_shape, num_guard_cells=g), + J_tiles, + rho, + phi, + ( + tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), + tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), + ), + pml_state, + jnp.asarray(False), + ) + return tiled_particles, species_config, tiled_fields + + def _jitted_electrodynamic_loop_with_static_world(self, world): + def loop_with_static_world( + particles, + species_config, + fields, + constants, + relativistic=True, + particle_pusher="boris", + ): + return time_loop_electrodynamic( + particles, + species_config, + fields, + world, + constants, + relativistic=relativistic, + particle_pusher=particle_pusher, + ) + + return jax.jit( + loop_with_static_world, + static_argnames=("relativistic", "particle_pusher"), + ) + + def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): + for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): + diff = jnp.max(jnp.abs(reference_component - tiled_component)) + self.assertTrue( + jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: {name}{component_name} max abs diff {diff}", + ) + + def _assert_tiled_state_matches_standard( + self, + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + step, + ): + if len(tiled_fields) > 7: + self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") + + species_count = reference_particles.active.shape[3] + for species_index in range(species_count): + reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) + tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) + self.assertTrue( + jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} position mismatch", + ) + self.assertTrue( + jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), + f"step {step}: species {species_index} velocity mismatch", + ) + + reference_E, reference_B, reference_J, *_ = reference_fields + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + num_guard_cells = int(world.get("guard_cells", 1)) + self._assert_vector_fields_close( + assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + step, + "E", + ) + self._assert_vector_fields_close( + assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + step, + "B", + ) + self._assert_vector_fields_close( + assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + step, + "J", + ) + + def _max_tiled_state_difference( + self, + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + world, + tile_shape, + ): + differences = {} + species_count = reference_particles.active.shape[3] + for species_index in range(species_count): + reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) + tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) + differences[f"x{species_index}"] = float(jnp.max(jnp.abs(tiled_x - reference_x))) + differences[f"u{species_index}"] = float(jnp.max(jnp.abs(tiled_u - reference_u))) + + reference_E, reference_B, reference_J, *_ = reference_fields + E_tiles, B_tiles, J_tiles, *_ = tiled_fields + num_guard_cells = int(world.get("guard_cells", 1)) + for name, reference_field, tiled_field in ( + ( + "E", + assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), + ( + "B", + assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), + ( + "J", + assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), + assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), + ), + ): + differences[name] = max(float(jnp.max(jnp.abs(tiled - reference))) for reference, tiled in zip(reference_field, tiled_field)) + + return differences + + def test_validate_field_solver_accepts_tiled_yee(self): + validate_field_solver("electrodynamic_yee") + + def test_tiled_yee_field_update_applies_standard_field_filter(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} + tile_shape = (2, 1, 1) + shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) + interior = (slice(1, -1), slice(1, -1), slice(1, -1)) + + x = jnp.arange(world["Nx"], dtype=jnp.float64)[:, None, None] + seed = jnp.zeros(shape) + Ex = _update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ey = _update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Ez = _update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bx = _update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + By = _update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + Bz = _update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) + E = (Ex, Ey, Ez) + B = (Bx, By, Bz) + J = (seed, seed, seed) + rho = seed + phi = seed + external_fields = ((seed, seed, seed), (seed, seed, seed)) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_with_tiled_grids(world, tile_shape) + reference_fields = ( + tile_vector_field(E, reference_world, reference_tile_shape), + tile_vector_field(B, reference_world, reference_tile_shape), + tile_vector_field(J, reference_world, reference_tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], reference_world, reference_tile_shape), + tile_vector_field(external_fields[1], reference_world, reference_tile_shape), + ), + None, + jnp.asarray(False), + ) + reference_E, reference_pml_state = update_E( + reference_fields[0], + reference_fields[1], + reference_fields[2], + reference_world, + constants, + ) + reference_B, reference_pml_state = update_B( + reference_E, + reference_fields[1], + reference_world, + constants, + reference_pml_state, + ) + reference_fields = (reference_E, reference_B, *reference_fields[2:]) + + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), + ), + None, + jnp.asarray(False), + ) + tiled_E, tiled_pml_state = update_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) + tiled_B, tiled_pml_state = update_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) + tiled_fields = (tiled_E, tiled_B, *tiled_fields[2:]) + + reference_E, reference_B, *_ = reference_fields + E_tiles, B_tiles, *_ = tiled_fields + self._assert_vector_fields_close( + assemble_tiled_vector_field(reference_E, reference_world, reference_tile_shape), + assemble_tiled_vector_field(E_tiles, tiled_world, tile_shape), + 1, + "E", + ) + self._assert_vector_fields_close( + assemble_tiled_vector_field(reference_B, reference_world, reference_tile_shape), + assemble_tiled_vector_field(B_tiles, tiled_world, tile_shape), + 1, + "B", + ) + + def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(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 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, + "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, + }, + } + config_path = os.path.join(tmpdir, "tiled_yee.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) + + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) + + def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(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 higuera cary 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, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "j_from_rhov", + "filter_j": "none", + "particle_pusher": "higuera_cary", + "relativistic": True, + }, + "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_higuera_cary.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["particle_pusher"], "higuera_cary") + + def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(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 digital filter 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, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "j_from_rhov", + "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, + }, + } + config_path = os.path.join(tmpdir, "tiled_yee_digital.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["filter_j"], "digital") + + def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(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 bilinear filter 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, + "particle_tile_nx": 2, + "particle_tile_ny": 1, + "particle_tile_nz": 1, + "current_calculation": "j_from_rhov", + "filter_j": "bilinear", + "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_bilinear.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(simulation_parameters["filter_j"], "bilinear") + + def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(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 conducting 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, + "x_bc": "conducting", + "y_bc": "conducting", + "z_bc": "conducting", + "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, + }, + } + config_path = os.path.join(tmpdir, "tiled_yee_conducting.toml") + with open(config_path, "w") as f: + toml.dump(config, f) + + loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) + + self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + self.assertIsInstance(particles, TiledParticles) + self.assertEqual(fields[0][0].ndim, 6) + self.assertEqual(world["boundary_conditions"], {"x": 1, "y": 1, "z": 1}) + + def test_tiled_yee_step_matches_standard_periodic_yee_step(self): + world = self._build_world() + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_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], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + ) + + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + reference_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + reference_E, reference_B, *_ = reference_fields + E_tiles, B_tiles, *_ = tiled_fields + reference_energy = compute_energy( + reference_particles, + reference_E, + reference_B, + world, + constants, + species_config=reference_species_config, + ) + tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants, species_config=species_config) + for reference, tiled in zip(reference_energy, tiled_energy): + self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) + + def test_tiled_yee_step_matches_standard_conducting_yee_step(self): + world = self._build_world( + boundary_conditions={"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING} + ) + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_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], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + ) + + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + reference_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + + def test_tiled_yee_step_uses_digital_filtered_current(self): + world = self._build_world() + world["current_filter"] = "digital" + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} + tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_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], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + ) + + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + reference_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + + def test_tiled_yee_step_uses_bilinear_filtered_current(self): + world = self._build_world() + world["current_filter"] = "bilinear" + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_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], + } + + species = self._species(world) + reference_species = self._species(world) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + ) + + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape), + tile_vector_field(B, tiled_world, tile_shape), + tile_vector_field(J, tiled_world, tile_shape), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape), + tile_vector_field(external_fields[1], tiled_world, tile_shape), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + reference_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + + def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_standard(self): + world = self._build_world() + world["shape_factor"] = 2 + world["guard_cells"] = 2 + world["current_filter"] = "bilinear" + constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} + tile_shape = (2, 1, 1) + reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) + reference_world = self._world_with_tiled_grids(world, reference_tile_shape) + tiled_world = self._world_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], + } + + species = particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.array([0.03, -0.04, 0.02, -0.01]), + v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + reference_species = particle_species( + name="electrons", + N_particles=4, + charge=-1.0, + mass=2.0, + weight=0.5, + T=1.0, + x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), + x2=jnp.zeros(4), + x3=jnp.zeros(4), + v1=jnp.array([0.10, -0.05, 0.07, -0.02]), + v2=jnp.array([0.03, -0.04, 0.02, -0.01]), + v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), + xwind=world["x_wind"], + ywind=world["y_wind"], + zwind=world["z_wind"], + dx=world["dx"], + dy=world["dy"], + dz=world["dz"], + dt=world["dt"], + ) + E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) + + reference_particles, reference_species_config, reference_fields = self._build_tiled_state( + [reference_species], + (E, B, J, rho, phi, external_fields, pml_state), + reference_world, + reference_tile_shape, + ) + + tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) + tiled_fields = ( + tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + rho, + phi, + ( + tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), + ), + None, + jnp.asarray(False), + ) + reference_particles, reference_fields = time_loop_electrodynamic( + reference_particles, + reference_species_config, + reference_fields, + reference_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = time_loop_electrodynamic( + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + relativistic=False, + particle_pusher="boris", + ) + + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + world, + reference_tile_shape, + tile_shape, + 1, + ) + + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_YEE") == "1", + "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", + ) + def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): + ( + reference_particles, + reference_species_config, + reference_fields, + reference_tile_shape, + reference_world, + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + tile_shape, + Nt, + ) = self._long_two_stream_state("direct", "none") + + reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) + + for t in range(Nt): + reference_particles, reference_fields = reference_loop( + reference_particles, + reference_species_config, + reference_fields, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + species_config, + tiled_fields, + constants, + relativistic=False, + particle_pusher="boris", + ) + + differences = self._max_tiled_state_difference( + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + tiled_world, + tile_shape, + ) + if t == 0: + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + tiled_world, + reference_tile_shape, + tile_shape, + t + 1, + ) + if any(value > ROUND_OFF_ATOL for value in differences.values()): + print(f"direct-current roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") + break + + @unittest.skipUnless( + os.environ.get("RUN_SLOW_TILED_YEE") == "1", + "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", + ) + def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): + ( + reference_particles, + reference_species_config, + reference_fields, + reference_tile_shape, + reference_world, + tiled_particles, + species_config, + tiled_fields, + tiled_world, + constants, + tile_shape, + Nt, + ) = self._long_two_stream_state("esirkepov", "none") + + reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) + tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) + + for t in range(Nt): + reference_particles, reference_fields = reference_loop( + reference_particles, + reference_species_config, + reference_fields, + constants, + relativistic=False, + particle_pusher="boris", + ) + tiled_particles, tiled_fields = tiled_loop( + tiled_particles, + species_config, + tiled_fields, + constants, + relativistic=False, + particle_pusher="boris", + ) + + differences = self._max_tiled_state_difference( + reference_particles, + reference_fields, + reference_tile_shape, + tiled_particles, + tiled_fields, + tiled_world, + tile_shape, + ) + if t == 0: + self._assert_tiled_state_matches_standard( + reference_particles, + reference_fields, + tiled_particles, + tiled_fields, + tiled_world, + reference_tile_shape, + tile_shape, + t + 1, + ) + if any(value > ROUND_OFF_ATOL for value in differences.values()): + print(f"Esirkepov roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") + break + + +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..074876d7 --- /dev/null +++ b/tests/physics_tests/yee_convergence_test.py @@ -0,0 +1,90 @@ +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 TestTiledYeeConvergence(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 _one_step_by_error(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._one_step_by_error(Nx, layout(Nx))) + + first_order = math.log(errors[0] / errors[1], 2.0) + second_order = math.log(errors[1] / errors[2], 2.0) + + self.assertGreater(first_order, 1.8) + self.assertGreater(second_order, 1.8) + + +if __name__ == "__main__": + unittest.main() From 9be8c7da1cebd705a311978826d1cb9570b870ef Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 9 Jul 2026 16:28:09 -0500 Subject: [PATCH 68/83] refactoring and commenting grid and stencil test --- tests/code_tests/grid_and_stencil_test.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/code_tests/grid_and_stencil_test.py b/tests/code_tests/grid_and_stencil_test.py index c31ab49f..d9b4ad2d 100644 --- a/tests/code_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() From 8dfb59c10db2b2598e801da54cf59f5d90fe3d0d Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 9 Jul 2026 16:42:43 -0500 Subject: [PATCH 69/83] commenting and reviewing initialization test --- tests/code_tests/initialization_test.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/code_tests/initialization_test.py b/tests/code_tests/initialization_test.py index 5f7f0f4c..a35e26e4 100644 --- a/tests/code_tests/initialization_test.py +++ b/tests/code_tests/initialization_test.py @@ -87,7 +87,7 @@ def test_initialize_simulation_returns_tiled_runtime_for_ordinary_electrodynamic loop, particles, fields, world, simulation_parameters, *_rest = initialize_simulation(toml.load(config_path)) - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) + 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)) @@ -99,6 +99,8 @@ def test_initialize_simulation_returns_tiled_runtime_for_ordinary_electrodynamic 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: @@ -146,6 +148,7 @@ def test_initialize_simulation_encodes_global_particle_boundary_conditions(self) 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: @@ -188,6 +191,7 @@ def test_initialize_simulation_uses_collocated_grid_for_electrostatic(self): 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: @@ -210,6 +214,7 @@ def test_initialize_simulation_rejects_unknown_solver(self): 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"): From 07d4a8b22b437b8aa418ac28405c59a4a2976700 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 9 Jul 2026 17:08:34 -0500 Subject: [PATCH 70/83] refactoring openPMD writing tests --- tests/code_tests/openpmd_test.py | 236 +++++++++++++++++++++---------- 1 file changed, 160 insertions(+), 76 deletions(-) diff --git a/tests/code_tests/openpmd_test.py b/tests/code_tests/openpmd_test.py index da2788fd..9abb08ea 100644 --- a/tests/code_tests/openpmd_test.py +++ b/tests/code_tests/openpmd_test.py @@ -1,17 +1,16 @@ import unittest +import math import threading import time from unittest.mock import patch -import jax 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.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles def _tile_axis_count(n_cells, cells_per_tile): @@ -30,46 +29,24 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): 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, + 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): @@ -192,14 +169,6 @@ def _world(): } -def _simulation_parameters(): - return { - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - } - - def _record_data(record): if len(record.chunks) <= 1: return record.data @@ -214,29 +183,86 @@ def _record_data(record): def _species(name, charge, mass, weight, x1): - world = _world() - return particle_species( - name=name, - N_particles=x1.shape[0], - charge=charge, - mass=mass, - weight=weight, - T=0.0, - x1=x1, - x2=jnp.zeros_like(x1), - x3=jnp.zeros_like(x1), - v1=jnp.ones_like(x1) * 0.1, - v2=jnp.zeros_like(x1), - v3=jnp.zeros_like(x1), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - active_mask=jnp.ones(x1.shape[0], dtype=bool), - dt=world["dt"], + 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): @@ -248,6 +274,66 @@ def test_openpmd_field_array_preserves_thin_y_axis(self): 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) @@ -411,8 +497,7 @@ def test_tiled_particle_snapshot_writes_same_records_as_synchronous_output(self) _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 = to_tiled_particles(species, world, _simulation_parameters()) - species_names = tuple(s.get_name() for s in species) + tiled_particles, species_config, species_names = _make_tiled_particles(species, world) constants = {"C": 10.0} sync_series = FakeSeries() @@ -590,8 +675,7 @@ def test_write_openpmd_initial_particles_flattens_tiled_particles_and_preserves_ _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 = to_tiled_particles(species, world, _simulation_parameters()) - species_names = tuple(s.get_name() for s in species) + tiled_particles, species_config, species_names = _make_tiled_particles(species, world) series = FakeSeries() with patch.object(openPMD.io, "Series", return_value=series): From af2974b4f160c0d869682973be96516681bf06d5 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Thu, 9 Jul 2026 18:11:06 -0500 Subject: [PATCH 71/83] refactoring and reviewing new esirkepov unittest --- ...ed_esirkepov_test.py => esirkepov_test.py} | 337 +++++++++++------- .../tiled_esirkepov_integration_test.py | 151 -------- tests/tiled_particle_fixtures.py | 280 --------------- 3 files changed, 210 insertions(+), 558 deletions(-) rename tests/code_tests/{tiled_esirkepov_test.py => esirkepov_test.py} (77%) delete mode 100644 tests/code_tests/tiled_esirkepov_integration_test.py delete mode 100644 tests/tiled_particle_fixtures.py diff --git a/tests/code_tests/tiled_esirkepov_test.py b/tests/code_tests/esirkepov_test.py similarity index 77% rename from tests/code_tests/tiled_esirkepov_test.py rename to tests/code_tests/esirkepov_test.py index 6cec6dce..e8b5948d 100644 --- a/tests/code_tests/tiled_esirkepov_test.py +++ b/tests/code_tests/esirkepov_test.py @@ -19,9 +19,7 @@ refresh_tiled_particle_tiles, update_tiled_particle_positions, ) -from tests.tiled_particle_fixtures import particle_species -from tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import TiledParticles +from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles from PyPIC3D.solvers.first_order_yee import ( update_E, ) @@ -46,52 +44,52 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=2): 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, + 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, @@ -140,52 +138,97 @@ 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(self, world, x1): - return particle_species( - name="electrons", - N_particles=3, - charge=-1.0, - mass=1.0, - weight=0.5, - T=1.0, - x1=x1, - x2=jnp.zeros_like(x1), - x3=jnp.zeros_like(x1), - v1=jnp.array([0.08, -0.06, 0.05]), - v2=jnp.zeros_like(x1), - v3=jnp.zeros_like(x1), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], + 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 _species_from_arrays(self, world, x, u, weight=0.5, active_mask=None): - return particle_species( - name="electrons", - N_particles=x.shape[0], - charge=-1.0, - mass=1.0, - weight=weight, - T=1.0, - x1=x[:, 0], - x2=x[:, 1], - x3=x[:, 2], - v1=u[:, 0], - v2=u[:, 1], - v3=u[:, 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=active_mask, + 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 ( @@ -197,6 +240,61 @@ def _tile_shape_for_world(self, world): 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) @@ -208,15 +306,9 @@ def _world_with_tiled_grids(self, world, tile_shape): world["grids"] = grids return world - def _assembled_esirkepov_current(self, world, old_species, constants, tile_shape): + def _assembled_esirkepov_current(self, world, tiled_particles, species_config, constants, tile_shape): world = self._world_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], - } g = int(world["guard_cells"]) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, @@ -234,11 +326,18 @@ def _assert_tiled_current_matches_reference(self, world, x_old, u, tile_shape=No if tile_shape is None: tile_shape = self._tile_shape_for_world(world) - old_species = self._species_from_arrays(world, x_old, u) - J_tiles, J_from_tiles = self._assembled_esirkepov_current(world, old_species, constants, tile_shape) + 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, - old_species, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, constants, self._one_tile_shape_for_world(world), ) @@ -361,15 +460,9 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) world = self._world_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], - } x_old = jnp.array([-1.10, -0.10, 1.05]) - old_species = self._species(world, x_old) + tiled_particles, species_config = self._one_dimensional_particles(world, x_old, tile_shape) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) g = int(world["guard_cells"]) _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( @@ -382,7 +475,8 @@ def test_tiled_esirkepov_matches_global_1d_periodic_current(self): J_from_tiles = assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=g) _, J_reference = self._assembled_esirkepov_current( world, - old_species, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, constants, self._one_tile_shape_for_world(world), ) @@ -398,15 +492,9 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre world["guard_cells"] = 2 world = self._world_with_tiled_grids(world, (2, 1, 1)) constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} - simulation_parameters = { - "particle_tile_nx": world["tile_shape"][0], - "particle_tile_ny": world["tile_shape"][1], - "particle_tile_nz": world["tile_shape"][2], - } x_old = jnp.array([-1.10, -0.10, 1.05]) - old_species = self._species(world, x_old) + tiled_particles, species_config = self._one_dimensional_particles(world, x_old, world["tile_shape"]) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) _, _, J_template, _, _ = initialize_fields(world) J_tiles = Esirkepov_current( tiled_particles, @@ -423,7 +511,8 @@ def test_public_Esirkepov_current_dispatches_tiled_particles_to_tile_local_curre ) _, J_reference = self._assembled_esirkepov_current( world, - old_species, + self._one_tile_particles_from_tiled(tiled_particles), + species_config, constants, self._one_tile_shape_for_world(world), ) @@ -442,13 +531,11 @@ def test_public_Esirkepov_current_rejects_tiled_filter_modes(self): tile_shape = (2, 1, 1) world = self._world_with_tiled_grids(world, tile_shape) constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - old_species = self._species(world, jnp.array([-1.10, -0.10, 1.05])) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + 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) @@ -539,11 +626,6 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): constants = {"C": 1.0, "eps": 1.0, "alpha": 1.0} tile_shape = (2, 1, 1) world = self._world_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], - } dx = world["dx"] x_old = jnp.array( [ @@ -559,8 +641,7 @@ def test_tiled_esirkepov_satisfies_tile_local_discrete_continuity(self): [-0.4 * dx / world["dt"], 0.0, 0.0], ] ) - old_species = self._species_from_arrays(world, x_old, u) - tiled_particles, species_config = to_tiled_particles([old_species], world, simulation_parameters) + 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) @@ -835,14 +916,16 @@ def test_tiled_yee_esirkepov_staging_uses_world_contract_not_function_identity(s particle_pusher=particle_pusher, ) - old_species = self._species_from_arrays( + 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, - old_species, + reference_particles, + reference_species_config, constants, self._one_tile_shape_for_world(world), ) diff --git a/tests/code_tests/tiled_esirkepov_integration_test.py b/tests/code_tests/tiled_esirkepov_integration_test.py deleted file mode 100644 index e5aec31a..00000000 --- a/tests/code_tests/tiled_esirkepov_integration_test.py +++ /dev/null @@ -1,151 +0,0 @@ -import os -import tempfile -import unittest -from contextlib import redirect_stdout -from io import StringIO -from unittest.mock import patch - -import jax -import numpy as np -import toml - -from PyPIC3D import __main__ as pypic_main -from PyPIC3D.__main__ import run_PyPIC3D - - -jax.config.update("jax_enable_x64", True) - - -class TiledEsirkepovIntegrationTest(unittest.TestCase): - def test_short_output_enabled_tiled_esirkepov_run_writes_openpmd_without_vtk(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 esirkepov output 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, - "plot_openpmd_fields": True, - "plot_openpmd_particles": True, - "plot_phasespace": True, - }, - "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, - }, - } - - _Nt, _plotting, _simulation, _plasma, _constants, _particles, fields, world, _species_config = run_PyPIC3D(config) - _E, _B, J, *_rest = fields - - self.assertEqual(J[0].shape[-3:], (6, 5, 5)) - self.assertEqual(tuple(int(width) for width in world["tile_shape"]), (2, 1, 1)) - for filename in ( - "total_energy.txt", - "electric_field_energy.txt", - "magnetic_field_energy.txt", - "kinetic_energy.txt", - "total_momentum.txt", - ): - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", filename))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "fields.h5"))) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "particles.h5"))) - self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "vector_field_slice"))) - self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "particles", "electrons.000000000.vtu"))) - self.assertFalse(os.path.exists(os.path.join(tmpdir, "data", "phase_space"))) - - def test_main_prints_final_energy_for_tiled_runtime(self): - with tempfile.TemporaryDirectory() as tmpdir: - x_path = os.path.join(tmpdir, "x.npy") - zeros_path = os.path.join(tmpdir, "zeros.npy") - config_path = os.path.join(tmpdir, "final_energy.toml") - np.save(x_path, np.array([0.0])) - np.save(zeros_path, np.zeros(1)) - - config = { - "simulation_parameters": { - "name": "final energy smoke", - "output_dir": tmpdir, - "solver": "electrodynamic_yee", - "Nx": 4, - "Ny": 1, - "Nz": 1, - "x_wind": 1.0, - "y_wind": 1.0, - "z_wind": 1.0, - "dt": 1.0e-10, - "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": 1, - "charge": -1.0, - "mass": 1.0, - "weight": 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, - }, - } - with open(config_path, "w") as f: - toml.dump(config, f) - - with patch("sys.argv", ["PyPIC3D", "--config", config_path]): - stdout = StringIO() - with redirect_stdout(stdout): - pypic_main.main() - - self.assertIn("Final Kinetic Energy", stdout.getvalue()) - self.assertTrue(os.path.exists(os.path.join(tmpdir, "data", "output.toml"))) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/tiled_particle_fixtures.py b/tests/tiled_particle_fixtures.py deleted file mode 100644 index d8f05cf6..00000000 --- a/tests/tiled_particle_fixtures.py +++ /dev/null @@ -1,280 +0,0 @@ -import math - -import jax.numpy as jnp -import numpy as np - -from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles - - -class particle_species: - """ - Test-only particle constructor for building tiled particle fixtures. - - Production code no longer has a flat species class. Several numerical tests - still need a compact way to name positions, velocities, masks, and per-species - constants before packing them into TiledParticles. - """ - - def __init__( - self, - name, - N_particles, - charge, - mass, - weight=1.0, - T=0.0, - x1=None, - x2=None, - x3=None, - v1=None, - v2=None, - v3=None, - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=1.0, - dy=1.0, - dz=1.0, - dt=1.0, - update_pos=True, - update_x=True, - update_y=True, - update_z=True, - update_v=True, - update_vx=True, - update_vy=True, - update_vz=True, - active_mask=None, - x_bc="periodic", - y_bc="periodic", - z_bc="periodic", - ): - self.name = name - self.N_particles = int(N_particles) - self.charge = charge - self.mass = mass - self.weight = weight - self.T = T - self.xwind = xwind - self.ywind = ywind - self.zwind = zwind - self.dx = dx - self.dy = dy - self.dz = dz - self.dt = dt - self.x_bc = x_bc - self.y_bc = y_bc - self.z_bc = z_bc - - self.update_pos = update_pos - self.update_x = update_pos and update_x - self.update_y = update_pos and update_y - self.update_z = update_pos and update_z - self.update_v = update_v - self.update_vx = update_v and update_vx - self.update_vy = update_v and update_vy - self.update_vz = update_v and update_vz - - self.x1 = jnp.asarray(x1 if x1 is not None else jnp.zeros(self.N_particles)) - self.x2 = jnp.asarray(x2 if x2 is not None else jnp.zeros(self.N_particles)) - self.x3 = jnp.asarray(x3 if x3 is not None else jnp.zeros(self.N_particles)) - self.v1 = jnp.asarray(v1 if v1 is not None else jnp.zeros(self.N_particles)) - self.v2 = jnp.asarray(v2 if v2 is not None else jnp.zeros(self.N_particles)) - self.v3 = jnp.asarray(v3 if v3 is not None else jnp.zeros(self.N_particles)) - - if active_mask is None: - self.active_mask = jnp.ones(self.N_particles, dtype=bool) - else: - self.active_mask = jnp.asarray(active_mask, dtype=bool) - - def get_name(self): - return self.name - - def get_charge(self): - return self.charge - - def get_mass(self): - return self.mass - - def get_weight(self): - return self.weight - - def get_temperature(self): - return self.T - - def get_number_of_particles(self): - return self.N_particles - - def get_velocity(self): - return self.v1, self.v2, self.v3 - - def get_active_mask(self): - return self.active_mask - - def get_forward_position(self): - return self.x1, self.x2, self.x3 - - def _half_step_axis(self, x, v, wind, bc): - x_diagnostic = x - v * self.dt / 2 - if bc == "periodic": - 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 get_position(self): - return ( - self._half_step_axis(self.x1, self.v1, self.xwind, self.x_bc), - self._half_step_axis(self.x2, self.v2, self.ywind, self.y_bc), - self._half_step_axis(self.x3, self.v3, self.zwind, self.z_bc), - ) - - def set_position(self, x1, x2, x3): - self.x1 = jnp.asarray(x1) - self.x2 = jnp.asarray(x2) - self.x3 = jnp.asarray(x3) - - def set_velocity(self, v1, v2, v3): - self.v1 = jnp.asarray(v1) - self.v2 = jnp.asarray(v2) - self.v3 = jnp.asarray(v3) - - def set_mass(self, mass): - self.mass = mass - - def set_weight(self, weight): - self.weight = weight - - def update_position(self): - self.x1 = jnp.where(self.update_x, self.x1 + self.v1 * self.dt, self.x1) - self.x2 = jnp.where(self.update_y, self.x2 + self.v2 * self.dt, self.x2) - self.x3 = jnp.where(self.update_z, self.x3 + self.v3 * self.dt, self.x3) - - def _apply_axis_bc(self, x, wind, bc): - half_wind = wind / 2 - if int(bc) == 0: - x = jnp.where(x > half_wind, x - wind, x) - x = jnp.where(x < -half_wind, x + wind, x) - return x, self.active_mask - if int(bc) == 2: - active = self.active_mask & (x >= -half_wind) & (x <= half_wind) - return x, active - return x, self.active_mask - - def boundary_conditions(self, world): - particle_bc = world.get("particle_boundary_conditions", {"x": 0, "y": 0, "z": 0}) - self.x1, self.active_mask = self._apply_axis_bc(self.x1, self.xwind, particle_bc.get("x", 0)) - self.x2, self.active_mask = self._apply_axis_bc(self.x2, self.ywind, particle_bc.get("y", 0)) - self.x3, self.active_mask = self._apply_axis_bc(self.x3, self.zwind, particle_bc.get("z", 0)) - - def kinetic_energy(self): - v2 = self.v1**2 + self.v2**2 + self.v3**2 - return jnp.sum(self.active_mask.astype(v2.dtype) * 0.5 * self.mass * self.weight * v2) - - def momentum(self): - vmag = jnp.sqrt(self.v1**2 + self.v2**2 + self.v3**2) - return jnp.sum(self.active_mask.astype(vmag.dtype) * self.mass * self.weight * vmag) - - -def _tile_shape(world, simulation_parameters): - if "tile_shape" in world: - return tuple(int(width) for width in world["tile_shape"]) - return ( - int(simulation_parameters.get("particle_tile_nx", world["Nx"])), - int(simulation_parameters.get("particle_tile_ny", world["Ny"])), - int(simulation_parameters.get("particle_tile_nz", world["Nz"])), - ) - - -def _tile_axis_count(n_cells, cells_per_tile): - return int(math.ceil(int(n_cells) / int(cells_per_tile))) - - -def _particle_tile_indices(x1, x2, x3, world, tile_shape): - tile_nx, tile_ny, tile_nz = tile_shape - Nx = int(world["Nx"]) - Ny = int(world["Ny"]) - Nz = 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 to_tiled_particles(particles, world, simulation_parameters): - tile_shape = _tile_shape(world, simulation_parameters) - 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) - - tile_counts = np.zeros((ntx, nty, ntz, len(particles)), dtype=int) - species_data = [] - - for species_index, species in enumerate(particles): - x1, x2, x3 = species.get_forward_position() - v1, v2, v3 = species.get_velocity() - active = species.get_active_mask() - tx, ty, tz = _particle_tile_indices(x1, x2, x3, world, tile_shape) - - for particle_index in range(species.get_number_of_particles()): - tile_counts[tx[particle_index], ty[particle_index], tz[particle_index], species_index] += 1 - - species_data.append((x1, x2, x3, v1, v2, v3, active, tx, ty, tz)) - - capacity_factor = float(simulation_parameters.get("particle_tile_capacity_factor", 1.0)) - max_particles_per_tile = int(math.ceil(max(1, tile_counts.max()) * capacity_factor)) - - x_tiles = np.zeros((ntx, nty, ntz, len(particles), max_particles_per_tile, 3), dtype=float) - u_tiles = np.zeros_like(x_tiles) - active_tiles = np.zeros((ntx, nty, ntz, len(particles), max_particles_per_tile), dtype=bool) - write_counts = np.zeros_like(tile_counts) - - for species_index, (x1, x2, x3, v1, v2, v3, active, tx, ty, tz) in enumerate(species_data): - for particle_index in range(particles[species_index].get_number_of_particles()): - tile = (tx[particle_index], ty[particle_index], tz[particle_index], species_index) - slot = write_counts[tile] - write_counts[tile] += 1 - x_tiles[tile + (slot, slice(None))] = [ - float(x1[particle_index]), - float(x2[particle_index]), - float(x3[particle_index]), - ] - u_tiles[tile + (slot, slice(None))] = [ - float(v1[particle_index]), - float(v2[particle_index]), - float(v3[particle_index]), - ] - active_tiles[tile + (slot,)] = bool(active[particle_index]) - - species_config = SpeciesConfig( - charge=jnp.asarray([species.charge for species in particles]), - mass=jnp.asarray([species.mass for species in particles]), - weight=jnp.asarray([species.weight for species in particles]), - update_x=jnp.asarray( - [[species.update_x, species.update_y, species.update_z] for species in particles], - dtype=bool, - ), - update_u=jnp.asarray( - [[species.update_vx, species.update_vy, species.update_vz] for species in particles], - dtype=bool, - ), - ) - - return ( - TiledParticles( - x=jnp.asarray(x_tiles), - u=jnp.asarray(u_tiles), - active=jnp.asarray(active_tiles), - ), - species_config, - ) From 456aad3a9bb62d83e58689d507d99d43f05a9371 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 09:47:47 -0500 Subject: [PATCH 72/83] updating code tests --- ...ities_test.py => fluid_quantities_test.py} | 85 +++++++------------ ...s_test.py => particle_diagnostics_test.py} | 37 ++++---- ...est.py => particle_initialization_test.py} | 20 ++--- ...fresh_test.py => particle_refresh_test.py} | 23 +++-- .../{tiled_pusher_test.py => pusher_test.py} | 15 ++-- ...ract_test.py => refactor_contract_test.py} | 0 .../{tiled_rho_test.py => rho_test.py} | 13 ++- tests/code_tests/utils_test.py | 76 +++++------------ .../{yee_tiled_test.py => yee_test.py} | 0 9 files changed, 106 insertions(+), 163 deletions(-) rename tests/code_tests/{tiled_fluid_quantities_test.py => fluid_quantities_test.py} (84%) rename tests/code_tests/{tiled_particle_diagnostics_test.py => particle_diagnostics_test.py} (77%) rename tests/code_tests/{tiled_particle_initialization_test.py => particle_initialization_test.py} (96%) rename tests/code_tests/{tiled_particle_refresh_test.py => particle_refresh_test.py} (88%) rename tests/code_tests/{tiled_pusher_test.py => pusher_test.py} (97%) rename tests/code_tests/{tiled_refactor_contract_test.py => refactor_contract_test.py} (100%) rename tests/code_tests/{tiled_rho_test.py => rho_test.py} (96%) rename tests/code_tests/{yee_tiled_test.py => yee_test.py} (100%) diff --git a/tests/code_tests/tiled_fluid_quantities_test.py b/tests/code_tests/fluid_quantities_test.py similarity index 84% rename from tests/code_tests/tiled_fluid_quantities_test.py rename to tests/code_tests/fluid_quantities_test.py index 83b1a1fa..28cd277e 100644 --- a/tests/code_tests/tiled_fluid_quantities_test.py +++ b/tests/code_tests/fluid_quantities_test.py @@ -16,8 +16,7 @@ compute_tiled_velocity_field_from_tiled_particles, ) from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from tests.initial_particles import build_tiled_particles, tiled_species from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -29,7 +28,6 @@ def _tile_axis_count(n_cells, cells_per_tile): 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) @@ -39,48 +37,29 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=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, + # 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): @@ -136,7 +115,7 @@ def _single_tile_world(self, world): ) def _assembled_mass_density(self, particles, world): - tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + 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, @@ -153,7 +132,7 @@ def _assembled_mass_density(self, particles, world): ) def _assembled_velocity_field(self, particles, world, direction): - tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + 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, @@ -171,7 +150,7 @@ def _assembled_velocity_field(self, particles, world, direction): ) def _assembled_pressure_field(self, particles, world, direction): - tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + 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, @@ -209,7 +188,7 @@ def _simulation_parameters(self): } def _particles(self, world): - electrons = particle_species( + electrons = tiled_species( name="electrons", N_particles=5, charge=-1.0, @@ -230,7 +209,7 @@ def _particles(self, world): dz=world["dz"], dt=world["dt"], ) - ions = particle_species( + ions = tiled_species( name="ions", N_particles=3, charge=2.0, @@ -257,7 +236,7 @@ def _particles(self, world): 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 = to_tiled_particles(particles, world, self._simulation_parameters()) + 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) @@ -266,7 +245,7 @@ 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 = to_tiled_particles(particles, world, self._simulation_parameters()) + 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( @@ -297,7 +276,7 @@ def test_compute_mass_density_prefers_tile_major_output_for_tiled_scalar_field(s world = self._build_world(shape_factor=2) world = self._world_with_tiled_grids(world) particles = self._particles(world) - tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + 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) @@ -323,7 +302,7 @@ 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 = to_tiled_particles(particles, world, self._simulation_parameters()) + 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( @@ -353,7 +332,7 @@ 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 = to_tiled_particles(particles, world, self._simulation_parameters()) + 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) diff --git a/tests/code_tests/tiled_particle_diagnostics_test.py b/tests/code_tests/particle_diagnostics_test.py similarity index 77% rename from tests/code_tests/tiled_particle_diagnostics_test.py rename to tests/code_tests/particle_diagnostics_test.py index 69ce6682..f6fe6f53 100644 --- a/tests/code_tests/tiled_particle_diagnostics_test.py +++ b/tests/code_tests/particle_diagnostics_test.py @@ -3,9 +3,8 @@ import jax import jax.numpy as jnp -from PyPIC3D.tests.tiled_particle_fixtures import particle_species from PyPIC3D.diagnostics.output_adapters import particles_for_output -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from tests.initial_particles import build_tiled_particles, species_names, tiled_species jax.config.update("jax_enable_x64", True) @@ -35,7 +34,7 @@ def _simulation_parameters(self): def _species(self): world = self._world() - ions = particle_species( + ions = tiled_species( name="ions", N_particles=4, charge=2.0, @@ -57,7 +56,7 @@ def _species(self): active_mask=jnp.array([True, False, True, True]), dt=world["dt"], ) - electrons = particle_species( + electrons = tiled_species( name="electrons", N_particles=3, charge=-1.0, @@ -83,29 +82,29 @@ def _species(self): def test_flatten_tiled_particles_matches_active_original_species(self): species_list = self._species() - tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + 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.get_active_mask() + active = original["active"] - original_x = jnp.stack(original.get_forward_position(), axis=-1)[active] - original_u = jnp.stack(original.get_velocity(), axis=-1)[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)))) + 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + 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) @@ -116,7 +115,7 @@ def test_inactive_tiled_slots_are_not_flattened(self): def test_absorbed_particles_do_not_appear_in_flattened_output(self): species_list = self._species() - tiled_particles, species_config = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) + 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) @@ -134,23 +133,23 @@ def test_flattened_diagnostic_positions_match_original_half_step_positions(self) "y": 0, "z": 0, } - tiled_particles, species_config = to_tiled_particles(species_list, world, self._simulation_parameters()) + 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.get_active_mask() - original_position = jnp.stack(original.get_position(), axis=-1)[active] + 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 = to_tiled_particles(species_list, self._world(), self._simulation_parameters()) - species_names = tuple(species.get_name() for species in species_list) + 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=species_names) + 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") diff --git a/tests/code_tests/tiled_particle_initialization_test.py b/tests/code_tests/particle_initialization_test.py similarity index 96% rename from tests/code_tests/tiled_particle_initialization_test.py rename to tests/code_tests/particle_initialization_test.py index 522e91f8..b624d855 100644 --- a/tests/code_tests/tiled_particle_initialization_test.py +++ b/tests/code_tests/particle_initialization_test.py @@ -9,7 +9,7 @@ from PyPIC3D.particles.particle_initialization import load_particles_from_toml from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles -from PyPIC3D.tests.tiled_particle_fixtures import particle_species, to_tiled_particles +from tests.initial_particles import build_tiled_particles, tiled_species jax.config.update("jax_enable_x64", True) @@ -21,7 +21,7 @@ def _write_array(self, tmpdir, name, values): np.save(path, np.asarray(values, dtype=float)) return path - def test_to_tiled_particles_packs_existing_particle_species_list(self): + def test_direct_tiled_particles_preserve_inactive_slots_and_metadata(self): world = { "Nx": 4, "Ny": 2, @@ -40,7 +40,7 @@ def test_to_tiled_particles_packs_existing_particle_species_list(self): "particle_tile_nz": 1, } - species = particle_species( + species = tiled_species( name="ions", N_particles=3, charge=2.0, @@ -65,7 +65,7 @@ def test_to_tiled_particles_packs_existing_particle_species_list(self): dt=world["dt"], ) - particles, species_config = to_tiled_particles([species], world, simulation_parameters) + particles, species_config = build_tiled_particles([species], world, simulation_parameters) self.assertIsInstance(particles, TiledParticles) self.assertIsInstance(species_config, SpeciesConfig) @@ -115,7 +115,7 @@ def test_species_metadata_is_not_slot_shaped(self): "particle_tile_nz": 1, } species = [ - particle_species( + tiled_species( name="electrons", N_particles=2, charge=-1.0, @@ -136,7 +136,7 @@ def test_species_metadata_is_not_slot_shaped(self): dz=world["dz"], dt=world["dt"], ), - particle_species( + tiled_species( name="ions", N_particles=2, charge=2.0, @@ -159,7 +159,7 @@ def test_species_metadata_is_not_slot_shaped(self): ), ] - particles, species_config = to_tiled_particles(species, world, simulation_parameters) + 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,)) @@ -168,7 +168,7 @@ def test_species_metadata_is_not_slot_shaped(self): self.assertEqual(species_config.update_x.shape, (2, 3)) self.assertEqual(species_config.update_u.shape, (2, 3)) - def test_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): + def test_direct_tiled_particles_can_allocate_inactive_capacity_headroom(self): world = { "Nx": 4, "Ny": 1, @@ -187,7 +187,7 @@ def test_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): "particle_tile_nz": 1, "particle_tile_capacity_factor": 3.0, } - species = particle_species( + species = tiled_species( name="ions", N_particles=3, charge=1.0, @@ -209,7 +209,7 @@ def test_to_tiled_particles_can_allocate_inactive_capacity_headroom(self): dt=world["dt"], ) - particles, species_config = to_tiled_particles([species], world, simulation_parameters) + 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) diff --git a/tests/code_tests/tiled_particle_refresh_test.py b/tests/code_tests/particle_refresh_test.py similarity index 88% rename from tests/code_tests/tiled_particle_refresh_test.py rename to tests/code_tests/particle_refresh_test.py index 9046b045..480ab32d 100644 --- a/tests/code_tests/tiled_particle_refresh_test.py +++ b/tests/code_tests/particle_refresh_test.py @@ -3,8 +3,7 @@ import jax import jax.numpy as jnp -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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, @@ -39,7 +38,7 @@ def _species(self, world, x1, v1, active_mask=None, update_x=True, x_bc="periodi if active_mask is None: active_mask = jnp.ones_like(jnp.asarray(x1), dtype=bool) n_particles = len(x1) - return particle_species( + return tiled_species( name="moving", N_particles=n_particles, charge=2.0, @@ -86,7 +85,7 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): v1=[0.25, 0.5, 0.75], active_mask=jnp.array([True, False, True]), ) - tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) + tiled_particles, species_config = build_tiled_particles([species], world, self._simulation_parameters()) moved = update_tiled_particle_positions(tiled_particles, species_config, world["dt"]) @@ -95,7 +94,7 @@ def test_update_tiled_particle_positions_respects_active_and_update_flags(self): 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 = to_tiled_particles([fixed_species], world, self._simulation_parameters()) + 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)) @@ -118,7 +117,7 @@ def test_tiles_need_repack_is_false_when_active_particles_stay_in_tile(self): 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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) @@ -135,7 +134,7 @@ def test_refresh_moves_particles_to_neighbor_tiles_with_static_shape(self): 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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)) @@ -148,7 +147,7 @@ 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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) @@ -166,7 +165,7 @@ 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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)) @@ -180,7 +179,7 @@ def test_refresh_absorbs_particle_from_global_boundary_condition(self): 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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)) @@ -196,7 +195,7 @@ def test_sparse_refresh_matches_compacting_refresh_for_neighbor_exchange(self): x1=[-1.5, -0.25, 0.25, 1.25], v1=[0.0, 0.0, 0.0, 0.0], ) - tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) + 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) @@ -217,7 +216,7 @@ def test_sparse_refresh_matches_compacting_refresh_for_neighbor_exchange(self): 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 = to_tiled_particles([species], world, self._simulation_parameters()) + 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)) diff --git a/tests/code_tests/tiled_pusher_test.py b/tests/code_tests/pusher_test.py similarity index 97% rename from tests/code_tests/tiled_pusher_test.py rename to tests/code_tests/pusher_test.py index 4328276a..e66ca0cd 100644 --- a/tests/code_tests/tiled_pusher_test.py +++ b/tests/code_tests/pusher_test.py @@ -4,8 +4,7 @@ import jax.numpy as jnp from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 @@ -116,7 +115,7 @@ def _simulation_parameters_for_tile_shape(self, tile_shape): } def _push_tiled_species(self, species, world, tile_shape, E, B, constants, relativistic=True, particle_pusher="boris"): - tiled_particles, species_config = to_tiled_particles( + tiled_particles, species_config = build_tiled_particles( [species], world, self._simulation_parameters_for_tile_shape(tile_shape), @@ -152,7 +151,7 @@ def _species(self, world, active_mask=None, update_vx=True, update_vy=True, upda if active_mask is None: active_mask = jnp.array([True, True, True, True]) - return particle_species( + return tiled_species( name="test particles", N_particles=4, charge=-1.0, @@ -233,7 +232,7 @@ def test_particle_push_matches_one_tile_higuera_cary(self): B = self._deterministic_vector_field(world, scale=0.05) def make_species(): - return particle_species( + return tiled_species( name="higuera cary particles", N_particles=4, charge=-1.0, @@ -303,7 +302,7 @@ def test_particle_push_respects_active_and_update_flags(self): update_vy=True, update_vz=False, ) - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = build_tiled_particles([species], world, simulation_parameters) pushed = particle_push( tiled_particles, @@ -327,7 +326,7 @@ def test_particle_push_matches_relativistic_one_tile_boris_on_reduced_axes(self) E = self._deterministic_vector_field(world, scale=0.1) B = self._deterministic_vector_field(world, scale=0.02) def make_species(): - return particle_species( + return tiled_species( name="one dimensional", N_particles=3, charge=1.0, @@ -390,7 +389,7 @@ def test_particle_push_matches_one_tile_boris_on_two_guard_reduced_axes(self): B = self._deterministic_vector_field(world, scale=0.02) def make_species(): - return particle_species( + return tiled_species( name="two guard one dimensional", N_particles=3, charge=1.0, diff --git a/tests/code_tests/tiled_refactor_contract_test.py b/tests/code_tests/refactor_contract_test.py similarity index 100% rename from tests/code_tests/tiled_refactor_contract_test.py rename to tests/code_tests/refactor_contract_test.py diff --git a/tests/code_tests/tiled_rho_test.py b/tests/code_tests/rho_test.py similarity index 96% rename from tests/code_tests/tiled_rho_test.py rename to tests/code_tests/rho_test.py index 410e88f9..8d3b1906 100644 --- a/tests/code_tests/tiled_rho_test.py +++ b/tests/code_tests/rho_test.py @@ -7,8 +7,7 @@ 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 PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +from tests.initial_particles import build_tiled_particles, tiled_species from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid @@ -129,7 +128,7 @@ def _one_tile_parameters(self, world): } def _particles(self, world): - electrons = particle_species( + electrons = tiled_species( name="electrons", N_particles=5, charge=-1.0, @@ -150,7 +149,7 @@ def _particles(self, world): dz=world["dz"], dt=world["dt"], ) - ions = particle_species( + ions = tiled_species( name="ions", N_particles=4, charge=2.0, @@ -175,7 +174,7 @@ def _particles(self, world): return [electrons, ions] def _tiled_with_noisy_inactive_slots(self, particles, world): - tiled_particles, species_config = to_tiled_particles(particles, world, self._simulation_parameters()) + 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) @@ -187,9 +186,7 @@ def _tiled_with_noisy_inactive_slots(self, particles, world): def _zero_species_velocities(self, particles): for species in particles: - species.v1 = jnp.zeros_like(species.v1) - species.v2 = jnp.zeros_like(species.v2) - species.v3 = jnp.zeros_like(species.v3) + species["u"] = jnp.zeros_like(species["u"]) return particles def _zero_tiled_velocities(self, tiled_particles): diff --git a/tests/code_tests/utils_test.py b/tests/code_tests/utils_test.py index 060415da..67247755 100644 --- a/tests/code_tests/utils_test.py +++ b/tests/code_tests/utils_test.py @@ -10,7 +10,7 @@ from PyPIC3D.boundary_conditions import ghost_cells from PyPIC3D.initialization import initialize_fields from PyPIC3D.particles.particle_class import SpeciesConfig, TiledParticles -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles +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 ( @@ -18,7 +18,6 @@ particle_sanity_check, load_external_fields_from_toml, add_external_fields, compute_energy, dump_parameters_to_toml, ) -from PyPIC3D.tests.tiled_particle_fixtures import particle_species jax.config.update("jax_enable_x64", True) @@ -292,30 +291,19 @@ def test_compute_energy_ignores_inactive_particles(self): } E, B, J, phi, rho = initialize_fields(world) 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="periodic", + 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.boundary_conditions({"particle_boundary_conditions": {"x": 2, "y": 0, "z": 0}}) - - tiled_particles, species_config = to_tiled_particles([species], world, {"particle_tile_nx": 1, "particle_tile_ny": 1, "particle_tile_nz": 1}) + 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) @@ -362,26 +350,17 @@ def test_package_does_not_export_vtk_diagnostics(self): self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.vtk")) def test_plot_positions_flattens_tiled_particles_and_preserves_species_names(self): - species = particle_species( + species = tiled_species( name="beam electrons", - N_particles=2, charge=-1.0, mass=1.0, weight=2.0, - T=0.0, x1=jnp.array([-0.25, 0.25]), x2=jnp.array([0.0, 0.0]), x3=jnp.array([0.0, 0.0]), - v1=jnp.array([0.1, 0.2]), - v2=jnp.array([0.0, 0.0]), - v3=jnp.array([0.0, 0.0]), - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=0.25, - dy=0.5, - dz=1.0, - dt=0.1, + u1=jnp.array([0.1, 0.2]), + u2=jnp.array([0.0, 0.0]), + u3=jnp.array([0.0, 0.0]), ) world = { "Nx": 4, @@ -402,7 +381,7 @@ def test_plot_positions_flattens_tiled_particles_and_preserves_species_names(sel "particle_tile_ny": 1, "particle_tile_nz": 1, } - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + tiled_particles, species_config = build_tiled_particles([species], world, simulation_parameters=simulation_parameters) class FakeFigure: def __init__(self): @@ -435,26 +414,17 @@ def write_html(self, *args, **kwargs): self.assertEqual(figure.trace_names, ["beam electrons"]) def test_particles_phase_space_flattens_tiled_particles(self): - species = particle_species( + species = tiled_species( name="electrons", - N_particles=1, charge=-1.0, mass=1.0, weight=1.0, - T=0.0, x1=jnp.array([0.0]), x2=jnp.array([0.0]), x3=jnp.array([0.0]), - v1=jnp.array([0.2]), - v2=jnp.array([0.0]), - v3=jnp.array([0.0]), - xwind=1.0, - ywind=1.0, - zwind=1.0, - dx=0.25, - dy=0.5, - dz=1.0, - dt=0.1, + u1=jnp.array([0.2]), + u2=jnp.array([0.0]), + u3=jnp.array([0.0]), ) world = { "Nx": 4, @@ -475,7 +445,7 @@ def test_particles_phase_space_flattens_tiled_particles(self): "particle_tile_ny": 1, "particle_tile_nz": 1, } - tiled_particles, species_config = to_tiled_particles([species], world, simulation_parameters) + 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")) diff --git a/tests/code_tests/yee_tiled_test.py b/tests/code_tests/yee_test.py similarity index 100% rename from tests/code_tests/yee_tiled_test.py rename to tests/code_tests/yee_test.py From e67519d52867ed2a23bb1851fec8408fcfbb3c38 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 09:53:14 -0500 Subject: [PATCH 73/83] testing initial particle and pml helpers --- tests/initial_particles.py | 204 ++++++++++++++++++++++++++++++++ tests/physics_tests/pml_test.py | 62 ++++------ 2 files changed, 226 insertions(+), 40 deletions(-) create mode 100644 tests/initial_particles.py 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/physics_tests/pml_test.py b/tests/physics_tests/pml_test.py index de4b6859..69a9ff27 100644 --- a/tests/physics_tests/pml_test.py +++ b/tests/physics_tests/pml_test.py @@ -42,47 +42,29 @@ def tile_scalar_field(field, world, tile_shape, num_guard_cells=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, + # 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): From 6815ea7502106caec206d32defccf1136d9dd689 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 09:57:51 -0500 Subject: [PATCH 74/83] updating and testing convergence of the yee grid --- .../tiled_yee_integration_test.py | 1411 ----------------- tests/physics_tests/yee_convergence_test.py | 8 +- 2 files changed, 5 insertions(+), 1414 deletions(-) delete mode 100644 tests/physics_tests/tiled_yee_integration_test.py diff --git a/tests/physics_tests/tiled_yee_integration_test.py b/tests/physics_tests/tiled_yee_integration_test.py deleted file mode 100644 index 2c544984..00000000 --- a/tests/physics_tests/tiled_yee_integration_test.py +++ /dev/null @@ -1,1411 +0,0 @@ -import os -import tempfile -import unittest - -import jax -import jax.numpy as jnp -import numpy as np -import toml - -from PyPIC3D.evolve import time_loop_electrodynamic -from PyPIC3D.initialization import initialize_simulation, validate_field_solver -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import TiledParticles -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics.output_adapters import assemble_tiled_vector_field -from PyPIC3D.solvers.first_order_yee import ( - update_B, - update_E, -) -from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid -from PyPIC3D.utils import compute_energy -from PyPIC3D.boundary_conditions.grid_and_stencil import BC_CONDUCTING, BC_PERIODIC - - -jax.config.update("jax_enable_x64", True) - -ROUND_OFF_RTOL = 1.0e-11 -ROUND_OFF_ATOL = 1.0e-11 - - -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 TestTiledYeeIntegration(unittest.TestCase): - def _build_world(self, boundary_conditions=None): - if boundary_conditions is None: - boundary_conditions = {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC} - world = { - "Nx": 8, - "Ny": 1, - "Nz": 1, - "dx": 0.5, - "dy": 1.0, - "dz": 1.0, - "dt": 0.01, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "guard_cells": 2, - "current_deposition": "direct", - "current_filter": "none", - "boundary_conditions": boundary_conditions, - } - 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): - 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_fields(self, 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)) - phi = jnp.zeros(shape) - rho = jnp.zeros(shape) - external_fields = (E, B) - return E, B, J, rho, phi, external_fields, None - - def _species(self, world): - return particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - 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.lexsort((x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - - def _long_two_stream_state(self, current_deposition="direct", current_filter="none"): - world = { - "Nx": 16, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.002, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 2, - "guard_cells": 2, - "current_deposition": current_deposition, - "current_filter": current_filter, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_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": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (4, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, tile_shape) - Nt = 1500 - - x0 = jnp.linspace( - -0.5 * world["x_wind"], - 0.5 * world["x_wind"], - world["Nx"], - endpoint=False, - ) + 0.5 * world["dx"] - x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) - zeros = jnp.zeros_like(x0) - beam_speed = 0.08 - - species_kwargs = { - "N_particles": world["Nx"], - "mass": 1.0, - "weight": 1.0, - "T": 0.0, - "x1": x0, - "x2": zeros, - "x3": zeros, - "v2": zeros, - "v3": zeros, - "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 = [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - reference_species = [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, tiled_world, simulation_parameters) - tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots=12) - g = int(tiled_world["guard_cells"]) - if current_deposition == "esirkepov": - J_template = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) - J_tiles = tuple(jnp.zeros_like(component) for component in J_template) - else: - J_tiles = tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=g) - - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=g), - tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=g), - J_tiles, - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=g), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - reference_species, - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - current_deposition, - ) - reference_particles = self._pad_tiled_particle_capacity(reference_particles, min_slots=12) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - step=0, - ) - - return ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) - - def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): - current_slots = tiled_particles.active.shape[-1] - if current_slots >= min_slots: - return tiled_particles - - pad_width = [(0, 0)] * tiled_particles.x.ndim - pad_width[-2] = (0, min_slots - current_slots) - x = jnp.pad(tiled_particles.x, pad_width) - u = jnp.pad(tiled_particles.u, pad_width) - - scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim - scalar_pad_width[-1] = (0, min_slots - current_slots) - active = jnp.pad(tiled_particles.active, scalar_pad_width) - - return TiledParticles( - x=x, - u=u, - active=active, - ) - - def _tiled_species_rows(self, tiled_particles, species_index): - active = tiled_particles.active[:, :, :, species_index, :] - x = tiled_particles.x[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - u = tiled_particles.u[:, :, :, species_index, :, :].reshape(-1, 3)[active.reshape(-1)] - order = jnp.lexsort((u[:, 0], x[:, 2], x[:, 1], x[:, 0])) - return x[order], u[order] - - def _build_tiled_state(self, particles, fields, world, tile_shape, current_deposition="direct"): - simulation_parameters = { - "particle_tile_nx": tile_shape[0], - "particle_tile_ny": tile_shape[1], - "particle_tile_nz": tile_shape[2], - } - tiled_particles, species_config = to_tiled_particles(particles, world, simulation_parameters) - E, B, J, rho, phi, external_fields, pml_state = fields - g = int(world["guard_cells"]) - if current_deposition == "esirkepov": - J_template = tile_vector_field(J, world, tile_shape, num_guard_cells=g) - J_tiles = tuple(jnp.zeros_like(component) for component in J_template) - else: - J_tiles = tile_vector_field(J, world, tile_shape, num_guard_cells=g) - tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=g), - tile_vector_field(B, world, tile_shape, num_guard_cells=g), - J_tiles, - rho, - phi, - ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), - ), - pml_state, - jnp.asarray(False), - ) - return tiled_particles, species_config, tiled_fields - - def _jitted_electrodynamic_loop_with_static_world(self, world): - def loop_with_static_world( - particles, - species_config, - fields, - constants, - relativistic=True, - particle_pusher="boris", - ): - return time_loop_electrodynamic( - particles, - species_config, - fields, - world, - constants, - relativistic=relativistic, - particle_pusher=particle_pusher, - ) - - return jax.jit( - loop_with_static_world, - static_argnames=("relativistic", "particle_pusher"), - ) - - def _assert_vector_fields_close(self, reference_field, tiled_field, step, name): - for component_name, reference_component, tiled_component in zip(("x", "y", "z"), reference_field, tiled_field): - diff = jnp.max(jnp.abs(reference_component - tiled_component)) - self.assertTrue( - jnp.allclose(tiled_component, reference_component, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: {name}{component_name} max abs diff {diff}", - ) - - def _assert_tiled_state_matches_standard( - self, - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - step, - ): - if len(tiled_fields) > 7: - self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - - species_count = reference_particles.active.shape[3] - for species_index in range(species_count): - reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) - tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) - self.assertTrue( - jnp.allclose(tiled_x, reference_x, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} position mismatch", - ) - self.assertTrue( - jnp.allclose(tiled_u, reference_u, rtol=ROUND_OFF_RTOL, atol=ROUND_OFF_ATOL), - f"step {step}: species {species_index} velocity mismatch", - ) - - reference_E, reference_B, reference_J, *_ = reference_fields - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - num_guard_cells = int(world.get("guard_cells", 1)) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "E", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "B", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - step, - "J", - ) - - def _max_tiled_state_difference( - self, - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - world, - tile_shape, - ): - differences = {} - species_count = reference_particles.active.shape[3] - for species_index in range(species_count): - reference_x, reference_u = self._tiled_species_rows(reference_particles, species_index) - tiled_x, tiled_u = self._tiled_species_rows(tiled_particles, species_index) - differences[f"x{species_index}"] = float(jnp.max(jnp.abs(tiled_x - reference_x))) - differences[f"u{species_index}"] = float(jnp.max(jnp.abs(tiled_u - reference_u))) - - reference_E, reference_B, reference_J, *_ = reference_fields - E_tiles, B_tiles, J_tiles, *_ = tiled_fields - num_guard_cells = int(world.get("guard_cells", 1)) - for name, reference_field, tiled_field in ( - ( - "E", - assemble_tiled_vector_field(reference_E, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(E_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ( - "B", - assemble_tiled_vector_field(reference_B, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(B_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ( - "J", - assemble_tiled_vector_field(reference_J, world, reference_tile_shape, num_guard_cells=num_guard_cells), - assemble_tiled_vector_field(J_tiles, world, tile_shape, num_guard_cells=num_guard_cells), - ), - ): - differences[name] = max(float(jnp.max(jnp.abs(tiled - reference))) for reference, tiled in zip(reference_field, tiled_field)) - - return differences - - def test_validate_field_solver_accepts_tiled_yee(self): - validate_field_solver("electrodynamic_yee") - - def test_tiled_yee_field_update_applies_standard_field_filter(self): - world = self._build_world() - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} - tile_shape = (2, 1, 1) - shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - interior = (slice(1, -1), slice(1, -1), slice(1, -1)) - - x = jnp.arange(world["Nx"], dtype=jnp.float64)[:, None, None] - seed = jnp.zeros(shape) - Ex = _update_ghost_cells(seed.at[interior].set(0.1 + 0.03 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ey = _update_ghost_cells(seed.at[interior].set(0.2 - 0.02 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Ez = _update_ghost_cells(seed.at[interior].set(0.05 * jnp.sin(x + 1.0)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bx = _update_ghost_cells(seed.at[interior].set(0.07 * jnp.cos(x + 0.5)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - By = _update_ghost_cells(seed.at[interior].set(0.04 + 0.01 * x), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - Bz = _update_ghost_cells(seed.at[interior].set(0.03 * jnp.sin(2.0 * x)), BC_PERIODIC, BC_PERIODIC, BC_PERIODIC) - E = (Ex, Ey, Ez) - B = (Bx, By, Bz) - J = (seed, seed, seed) - rho = seed - phi = seed - external_fields = ((seed, seed, seed), (seed, seed, seed)) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_with_tiled_grids(world, tile_shape) - reference_fields = ( - tile_vector_field(E, reference_world, reference_tile_shape), - tile_vector_field(B, reference_world, reference_tile_shape), - tile_vector_field(J, reference_world, reference_tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], reference_world, reference_tile_shape), - tile_vector_field(external_fields[1], reference_world, reference_tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_E, reference_pml_state = update_E( - reference_fields[0], - reference_fields[1], - reference_fields[2], - reference_world, - constants, - ) - reference_B, reference_pml_state = update_B( - reference_E, - reference_fields[1], - reference_world, - constants, - reference_pml_state, - ) - reference_fields = (reference_E, reference_B, *reference_fields[2:]) - - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - tiled_E, tiled_pml_state = update_E(tiled_fields[0], tiled_fields[1], tiled_fields[2], tiled_world, constants) - tiled_B, tiled_pml_state = update_B(tiled_E, tiled_fields[1], tiled_world, constants, tiled_pml_state) - tiled_fields = (tiled_E, tiled_B, *tiled_fields[2:]) - - reference_E, reference_B, *_ = reference_fields - E_tiles, B_tiles, *_ = tiled_fields - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_E, reference_world, reference_tile_shape), - assemble_tiled_vector_field(E_tiles, tiled_world, tile_shape), - 1, - "E", - ) - self._assert_vector_fields_close( - assemble_tiled_vector_field(reference_B, reference_world, reference_tile_shape), - assemble_tiled_vector_field(B_tiles, tiled_world, tile_shape), - 1, - "B", - ) - - def test_initialize_simulation_uses_tiled_particles_and_fields_for_tiled_yee(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 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, - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee.toml") - with open(config_path, "w") as f: - toml.dump(config, f) - - loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(tuple(world["tile_shape"]), (2, 1, 1)) - - def test_initialize_simulation_accepts_tiled_yee_higuera_cary_pusher(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 higuera cary 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "filter_j": "none", - "particle_pusher": "higuera_cary", - "relativistic": True, - }, - "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_higuera_cary.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["particle_pusher"], "higuera_cary") - - def test_initialize_simulation_accepts_tiled_yee_digital_current_filter(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 digital filter 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee_digital.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["filter_j"], "digital") - - def test_initialize_simulation_accepts_tiled_yee_bilinear_current_filter(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 bilinear filter 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, - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - "current_calculation": "j_from_rhov", - "filter_j": "bilinear", - "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_bilinear.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.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(simulation_parameters["filter_j"], "bilinear") - - def test_initialize_simulation_accepts_tiled_yee_conducting_field_boundaries(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 conducting 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, - "x_bc": "conducting", - "y_bc": "conducting", - "z_bc": "conducting", - "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, - }, - } - config_path = os.path.join(tmpdir, "tiled_yee_conducting.toml") - with open(config_path, "w") as f: - toml.dump(config, f) - - loop, particles, fields, world, *_rest = initialize_simulation(toml.load(config_path)) - - self.assertIs(loop.func if hasattr(loop, "func") else loop, time_loop_electrodynamic) - self.assertIsInstance(particles, TiledParticles) - self.assertEqual(fields[0][0].ndim, 6) - self.assertEqual(world["boundary_conditions"], {"x": 1, "y": 1, "z": 1}) - - def test_tiled_yee_step_matches_standard_periodic_yee_step(self): - world = self._build_world() - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - reference_E, reference_B, *_ = reference_fields - E_tiles, B_tiles, *_ = tiled_fields - reference_energy = compute_energy( - reference_particles, - reference_E, - reference_B, - world, - constants, - species_config=reference_species_config, - ) - tiled_energy = compute_energy(tiled_particles, E_tiles, B_tiles, world, constants, species_config=species_config) - for reference, tiled in zip(reference_energy, tiled_energy): - self.assertTrue(jnp.allclose(tiled, reference, rtol=1.0e-12, atol=1.0e-12)) - - def test_tiled_yee_step_matches_standard_conducting_yee_step(self): - world = self._build_world( - boundary_conditions={"x": BC_CONDUCTING, "y": BC_CONDUCTING, "z": BC_CONDUCTING} - ) - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_uses_digital_filtered_current(self): - world = self._build_world() - world["current_filter"] = "digital" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 0.6} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_uses_bilinear_filtered_current(self): - world = self._build_world() - world["current_filter"] = "bilinear" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = self._species(world) - reference_species = self._species(world) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape), - tile_vector_field(B, tiled_world, tile_shape), - tile_vector_field(J, tiled_world, tile_shape), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape), - tile_vector_field(external_fields[1], tiled_world, tile_shape), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - def test_tiled_yee_step_bilinear_quadratic_two_guards_reduced_axes_matches_standard(self): - world = self._build_world() - world["shape_factor"] = 2 - world["guard_cells"] = 2 - world["current_filter"] = "bilinear" - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = (2, 1, 1) - reference_tile_shape = (world["Nx"], world["Ny"], world["Nz"]) - reference_world = self._world_with_tiled_grids(world, reference_tile_shape) - tiled_world = self._world_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], - } - - species = particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.array([0.03, -0.04, 0.02, -0.01]), - v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - reference_species = particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.array([0.03, -0.04, 0.02, -0.01]), - v3=jnp.array([-0.02, 0.01, -0.03, 0.04]), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - E, B, J, rho, phi, external_fields, pml_state = self._empty_fields(world) - - reference_particles, reference_species_config, reference_fields = self._build_tiled_state( - [reference_species], - (E, B, J, rho, phi, external_fields, pml_state), - reference_world, - reference_tile_shape, - ) - - tiled_particles, species_config = to_tiled_particles([species], tiled_world, simulation_parameters) - tiled_fields = ( - tile_vector_field(E, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(B, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(J, tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - rho, - phi, - ( - tile_vector_field(external_fields[0], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - tile_vector_field(external_fields[1], tiled_world, tile_shape, num_guard_cells=int(world["guard_cells"])), - ), - None, - jnp.asarray(False), - ) - reference_particles, reference_fields = time_loop_electrodynamic( - reference_particles, - reference_species_config, - reference_fields, - reference_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = time_loop_electrodynamic( - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - relativistic=False, - particle_pusher="boris", - ) - - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - world, - reference_tile_shape, - tile_shape, - 1, - ) - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_YEE") == "1", - "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", - ) - def test_long_two_stream_tiled_yee_direct_current_roundoff_origin_probe(self): - ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) = self._long_two_stream_state("direct", "none") - - reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) - - for t in range(Nt): - reference_particles, reference_fields = reference_loop( - reference_particles, - reference_species_config, - reference_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - - differences = self._max_tiled_state_difference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - ) - if t == 0: - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - t + 1, - ) - if any(value > ROUND_OFF_ATOL for value in differences.values()): - print(f"direct-current roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") - break - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_YEE") == "1", - "Set RUN_SLOW_TILED_YEE=1 to run the 1500-step tiled Yee two-stream comparison.", - ) - def test_long_two_stream_tiled_yee_esirkepov_roundoff_origin_probe(self): - ( - reference_particles, - reference_species_config, - reference_fields, - reference_tile_shape, - reference_world, - tiled_particles, - species_config, - tiled_fields, - tiled_world, - constants, - tile_shape, - Nt, - ) = self._long_two_stream_state("esirkepov", "none") - - reference_loop = self._jitted_electrodynamic_loop_with_static_world(reference_world) - tiled_loop = self._jitted_electrodynamic_loop_with_static_world(tiled_world) - - for t in range(Nt): - reference_particles, reference_fields = reference_loop( - reference_particles, - reference_species_config, - reference_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - constants, - relativistic=False, - particle_pusher="boris", - ) - - differences = self._max_tiled_state_difference( - reference_particles, - reference_fields, - reference_tile_shape, - tiled_particles, - tiled_fields, - tiled_world, - tile_shape, - ) - if t == 0: - self._assert_tiled_state_matches_standard( - reference_particles, - reference_fields, - tiled_particles, - tiled_fields, - tiled_world, - reference_tile_shape, - tile_shape, - t + 1, - ) - if any(value > ROUND_OFF_ATOL for value in differences.values()): - print(f"Esirkepov roundoff probe first exceeded {ROUND_OFF_ATOL} at step {t + 1}: {differences}") - break - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/physics_tests/yee_convergence_test.py b/tests/physics_tests/yee_convergence_test.py index 074876d7..73510ccf 100644 --- a/tests/physics_tests/yee_convergence_test.py +++ b/tests/physics_tests/yee_convergence_test.py @@ -13,7 +13,7 @@ jax.config.update("jax_enable_x64", True) -class TestTiledYeeConvergence(unittest.TestCase): +class TestYeeConvergence(unittest.TestCase): def _world(self, Nx, tile_shape): x_wind = 2.0 * math.pi world = { @@ -45,7 +45,7 @@ def _world(self, Nx, tile_shape): world["grids"]["tiled_center_grid"] = tiled_center_grid return world - def _one_step_by_error(self, Nx, tile_shape): + 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 @@ -77,10 +77,12 @@ def test_public_tiled_update_B_converges_for_one_tile_and_multi_tile_layouts(sel for layout in layouts: errors = [] for Nx in (32, 64, 128): - errors.append(self._one_step_by_error(Nx, layout(Nx))) + 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) From 8ee5403f93dbe426eecfe27ed4ed7c59064a3b3e Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 10:38:29 -0500 Subject: [PATCH 75/83] refactoring electrostatic unit tests --- tests/physics_tests/electrostatic_yee_test.py | 79 ++- .../physics_tests/tiled_electrostatic_test.py | 512 ------------------ 2 files changed, 77 insertions(+), 514 deletions(-) delete mode 100644 tests/physics_tests/tiled_electrostatic_test.py diff --git a/tests/physics_tests/electrostatic_yee_test.py b/tests/physics_tests/electrostatic_yee_test.py index 3879d8ca..1b63754f 100644 --- a/tests/physics_tests/electrostatic_yee_test.py +++ b/tests/physics_tests/electrostatic_yee_test.py @@ -1,9 +1,16 @@ import unittest +import os +import tempfile import jax import jax.numpy as jnp +import numpy as np +import toml -from PyPIC3D.tests.tiled_particle_fixtures import particle_species +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, @@ -75,7 +82,7 @@ def setUp(self): self.initial_phi = jnp.zeros((self.Nx + 2 * self.g, self.Ny + 2 * self.g, self.Nz + 2 * self.g)) self.particles = [ - particle_species( + tiled_species( name="test", N_particles=1, charge=1.0, @@ -159,6 +166,74 @@ def test_electrostatic_field_solve_uses_local_centered_gradient(self): 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/tiled_electrostatic_test.py b/tests/physics_tests/tiled_electrostatic_test.py deleted file mode 100644 index 68eb6863..00000000 --- a/tests/physics_tests/tiled_electrostatic_test.py +++ /dev/null @@ -1,512 +0,0 @@ -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_PERIODIC -from PyPIC3D.deposition.rho import compute_rho -from PyPIC3D.evolve import time_loop_electrostatic -from PyPIC3D.initialization import initialize_simulation -from PyPIC3D.tests.tiled_particle_fixtures import particle_species -from PyPIC3D.tests.tiled_particle_fixtures import to_tiled_particles -from PyPIC3D.particles.particle_class import TiledParticles -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics.output_adapters import assemble_tiled_scalar_field -from PyPIC3D.utilities.grids import build_collocated_grid, 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 TestTiledElectrostatic(unittest.TestCase): - def _build_world(self, shape_factor=1): - world = { - "Nx": 8, - "Ny": 1, - "Nz": 1, - "dx": 0.5, - "dy": 1.0, - "dz": 1.0, - "dt": 0.01, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "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} - world["tile_shape"] = (2, 1, 1) - world = self._world_with_tiled_grids(world, world["tile_shape"]) - return world - - 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 _simulation_parameters(self): - return { - "particle_tile_nx": 2, - "particle_tile_ny": 1, - "particle_tile_nz": 1, - } - - 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 _one_tile_shape(self, world): - return (int(world["Nx"]), int(world["Ny"]), int(world["Nz"])) - - def _empty_fields(self, world): - shape = (world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2) - zero = jnp.zeros(shape) - E = (zero, zero, zero) - B = (zero, zero, zero) - J = (zero, zero, zero) - external_fields = (E, B) - return E, B, J, zero, zero, external_fields - - def _species(self, world): - return particle_species( - name="electrons", - N_particles=4, - charge=-1.0, - mass=2.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.5, -0.5, 0.5, 1.5]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([0.10, -0.05, 0.07, -0.02]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - - def _neutral_species(self, world): - electrons = self._species(world) - ions = particle_species( - name="ions", - N_particles=4, - charge=1.0, - mass=4.0, - weight=0.5, - T=1.0, - x1=jnp.array([-1.25, -0.25, 0.75, 1.75]), - x2=jnp.zeros(4), - x3=jnp.zeros(4), - v1=jnp.array([-0.03, 0.04, -0.02, 0.01]), - v2=jnp.zeros(4), - v3=jnp.zeros(4), - xwind=world["x_wind"], - ywind=world["y_wind"], - zwind=world["z_wind"], - dx=world["dx"], - dy=world["dy"], - dz=world["dz"], - dt=world["dt"], - ) - return [electrons, ions] - - def _pad_tiled_particle_capacity(self, tiled_particles, min_slots): - current_slots = tiled_particles.active.shape[-1] - if current_slots >= min_slots: - return tiled_particles - - pad_width = [(0, 0)] * tiled_particles.x.ndim - pad_width[-2] = (0, min_slots - current_slots) - x = jnp.pad(tiled_particles.x, pad_width) - u = jnp.pad(tiled_particles.u, pad_width) - - scalar_pad_width = [(0, 0)] * tiled_particles.active.ndim - scalar_pad_width[-1] = (0, min_slots - current_slots) - - return TiledParticles( - x=x, - u=u, - active=jnp.pad(tiled_particles.active, scalar_pad_width), - ) - - def _build_tiled_state(self, particles, fields, world, tile_shape, min_slots=None): - tiled_particles, species_config = to_tiled_particles( - particles, - world, - self._simulation_parameters_for_tile_shape(tile_shape), - ) - if min_slots is not None: - tiled_particles = self._pad_tiled_particle_capacity(tiled_particles, min_slots) - - E, B, J, rho, phi, external_fields = fields - g = int(world["guard_cells"]) - tiled_fields = ( - tile_vector_field(E, world, tile_shape, num_guard_cells=g), - tile_vector_field(B, world, tile_shape, num_guard_cells=g), - tile_vector_field(J, world, tile_shape, num_guard_cells=g), - tile_scalar_field(rho, world, tile_shape, num_guard_cells=g), - tile_scalar_field(phi, world, tile_shape, num_guard_cells=g), - ( - tile_vector_field(external_fields[0], world, tile_shape, num_guard_cells=g), - tile_vector_field(external_fields[1], world, tile_shape, num_guard_cells=g), - ), - None, - jnp.asarray(False), - ) - - return tiled_particles, species_config, tiled_fields - - def _long_two_stream_species(self, world): - x0 = jnp.linspace( - -0.5 * world["x_wind"], - 0.5 * world["x_wind"], - world["Nx"], - endpoint=False, - ) + 0.5 * world["dx"] - x0 = x0 + 0.02 * world["dx"] * jnp.sin(2.0 * jnp.pi * (x0 / world["x_wind"] + 0.5)) - zeros = jnp.zeros_like(x0) - beam_speed = 0.08 - - species_kwargs = { - "N_particles": world["Nx"], - "mass": 1.0, - "weight": 1.0, - "T": 0.0, - "x1": x0, - "x2": zeros, - "x3": zeros, - "v2": zeros, - "v3": zeros, - "xwind": world["x_wind"], - "ywind": world["y_wind"], - "zwind": world["z_wind"], - "dx": world["dx"], - "dy": world["dy"], - "dz": world["dz"], - "dt": world["dt"], - } - return [ - particle_species(name="electron_left", charge=-1.0, v1=-beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="electron_right", charge=-1.0, v1=beam_speed * jnp.ones_like(x0), **species_kwargs), - particle_species(name="ion_background", charge=2.0, v1=zeros, **species_kwargs), - ] - - def _long_two_stream_state(self): - world = { - "Nx": 16, - "Ny": 1, - "Nz": 1, - "dx": 0.25, - "dy": 1.0, - "dz": 1.0, - "dt": 0.002, - "x_wind": 4.0, - "y_wind": 1.0, - "z_wind": 1.0, - "shape_factor": 1, - "guard_cells": 2, - "boundary_conditions": {"x": BC_PERIODIC, "y": BC_PERIODIC, "z": BC_PERIODIC}, - "particle_boundary_conditions": {"x": 0, "y": 0, "z": 0}, - } - center_grid, vertex_grid = build_collocated_grid(world) - world["grids"] = {"center": center_grid, "vertex": vertex_grid} - - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - tile_shape = self._one_tile_shape(world) - world = self._world_with_tiled_grids(world, tile_shape) - Nt = 1500 - - E, B, J, rho, phi, external_fields = self._empty_fields(world) - fields = (E, B, J, rho, phi, external_fields) - tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._long_two_stream_species(world), - fields, - world, - tile_shape, - ) - - return tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt - - def test_scalar_tiles_round_trip_global_rho_and_phi(self): - world = self._build_world() - tile_shape = world["tile_shape"] - scalar = jnp.arange((world["Nx"] + 2) * (world["Ny"] + 2) * (world["Nz"] + 2), dtype=float) - scalar = scalar.reshape((world["Nx"] + 2, world["Ny"] + 2, world["Nz"] + 2)) - - scalar_tiles = tile_scalar_field(scalar, world, tile_shape) - scalar_from_tiles = assemble_tiled_scalar_field(scalar_tiles, world, tile_shape) - - self.assertTrue(jnp.allclose(scalar_from_tiles[1:-1, 1:-1, 1:-1], scalar[1:-1, 1:-1, 1:-1])) - - def test_tiled_rho_assembles_to_global_rho(self): - world = self._build_world(shape_factor=1) - constants = {"alpha": 1.0} - species = self._species(world) - tiled_particles, species_config = to_tiled_particles([species], world, self._simulation_parameters()) - _, _, _, rho, _, _ = self._empty_fields(world) - rho_tiles = tile_scalar_field(rho, world, world["tile_shape"]) - - tiled_rho = compute_rho( - tiled_particles, - species_config, - rho_tiles, - constants, - world, - ) - assembled_rho = assemble_tiled_scalar_field(tiled_rho, world, world["tile_shape"], num_guard_cells=int(world["guard_cells"])) - - reference_world = self._world_with_tiled_grids(world, self._one_tile_shape(world)) - reference_particles, reference_species_config = to_tiled_particles( - [self._species(reference_world)], - reference_world, - self._simulation_parameters_for_tile_shape(reference_world["tile_shape"]), - ) - reference_rho_tiles = tile_scalar_field(rho, reference_world, reference_world["tile_shape"], num_guard_cells=int(reference_world["guard_cells"])) - reference_rho_tiles = compute_rho( - reference_particles, - reference_species_config, - reference_rho_tiles, - constants, - reference_world, - ) - reference_rho = assemble_tiled_scalar_field( - reference_rho_tiles, - reference_world, - reference_world["tile_shape"], - num_guard_cells=int(reference_world["guard_cells"]), - ) - - self.assertTrue( - jnp.allclose( - assembled_rho[1:-1, 1:-1, 1:-1], - reference_rho[1:-1, 1:-1, 1:-1], - rtol=1.0e-12, - atol=1.0e-12, - ) - ) - - def test_tiled_electrostatic_step_uses_single_tile_state(self): - world = self._build_world(shape_factor=1) - constants = {"C": 10.0, "eps": 1.0, "mu": 1.0, "alpha": 1.0} - E, B, J, rho, phi, external_fields = self._empty_fields(world) - fields = (E, B, J, rho, phi, external_fields) - tile_shape = self._one_tile_shape(world) - world = self._world_with_tiled_grids(world, tile_shape) - tiled_particles, species_config, tiled_fields = self._build_tiled_state( - self._neutral_species(world), - fields, - world, - tile_shape, - ) - tiled_particles, tiled_fields = time_loop_electrostatic( - tiled_particles, - species_config, - tiled_fields, - world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", - ) - - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(B_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(J_tiles[0].shape[:3], (1, 1, 1)) - self.assertEqual(rho_tiles.shape[:3], (1, 1, 1)) - self.assertEqual(phi_tiles.shape[:3], (1, 1, 1)) - self.assertFalse(bool(tiled_fields[-1])) - self.assertTrue(jnp.all(jnp.isfinite(rho_tiles))) - self.assertTrue(jnp.all(jnp.isfinite(phi_tiles))) - - @unittest.skipUnless( - os.environ.get("RUN_SLOW_TILED_ELECTROSTATIC") == "1", - "Set RUN_SLOW_TILED_ELECTROSTATIC=1 to run the 1500-step single-tile electrostatic check.", - ) - def test_long_two_stream_single_tile_electrostatic_remains_finite(self): - tiled_particles, species_config, tiled_fields, world, constants, tile_shape, Nt = self._long_two_stream_state() - - tiled_loop = jax.jit( - time_loop_electrostatic, - static_argnames=("solver", "relativistic", "particle_pusher"), - ) - - for step in range(1, Nt + 1): - tiled_particles, tiled_fields = tiled_loop( - tiled_particles, - species_config, - tiled_fields, - world, - constants, - solver="electrostatic", - relativistic=False, - particle_pusher="boris", - ) - - E_tiles, B_tiles, J_tiles, rho_tiles, phi_tiles, *_ = tiled_fields - self.assertEqual(tuple(world["tile_shape"]), tuple(tile_shape)) - self.assertEqual(E_tiles[0].shape[:3], (1, 1, 1)) - self.assertFalse(bool(tiled_fields[-1]), f"step {step}: tiled particle capacity overflowed") - for component in E_tiles + B_tiles + J_tiles: - self.assertTrue(jnp.all(jnp.isfinite(component)), f"step {step}: nonfinite vector field") - self.assertTrue(jnp.all(jnp.isfinite(rho_tiles)), f"step {step}: nonfinite rho") - self.assertTrue(jnp.all(jnp.isfinite(phi_tiles)), f"step {step}: nonfinite phi") - - def test_initialize_simulation_accepts_tiled_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": "tiled 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)) - self.assertIn("tiled_center_grid", world["grids"]) - self.assertIn("tiled_vertex_grid", world["grids"]) - self.assertEqual(int(world["guard_cells"]), 2) - self.assertEqual(world["grids"]["tiled_center_grid"][0].shape, (1, 1, 1, 12)) - self.assertEqual(world["grids"]["tiled_center_grid"][1].shape, (1, 1, 1, 5)) - 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() From 93b8baf5007621ec7dc5c681b2e8dd0e2653d4c2 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 10:38:48 -0500 Subject: [PATCH 76/83] updating demo to include tiling parameters --- demos/two_stream/two_stream.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/demos/two_stream/two_stream.toml b/demos/two_stream/two_stream.toml index 29cf92c9..19d3ad5a 100644 --- a/demos/two_stream/two_stream.toml +++ b/demos/two_stream/two_stream.toml @@ -6,6 +6,10 @@ 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 From 9498fab97076d19c1ef0003b67e7807c5b604a6d Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 10:41:48 -0500 Subject: [PATCH 77/83] updating unit testing workflow for github to reflect unit test restructuring --- .github/workflows/unit_testing.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index f1fe3bdc..3837ad75 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -27,4 +27,5 @@ jobs: pip install . - name: Test with unittest run: | - python -m unittest tests/*.py + python -m unittest tests/code_tests/*.py + python -m unittest tests/physics_tests/*.py From 80079ce5cbf95639501ec806bce0d9f8f1d42386 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 11:43:57 -0500 Subject: [PATCH 78/83] updating unit testing workflows --- .github/workflows/coverage.yml | 4 ++-- .github/workflows/unit_testing.yml | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 971b4dd9..60078c60 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -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 tests/physics_tests/*.py coverage xml - name: Upload coverage report diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index 3837ad75..64e40b42 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -27,5 +27,4 @@ jobs: pip install . - name: Test with unittest run: | - python -m unittest tests/code_tests/*.py - python -m unittest tests/physics_tests/*.py + python -m unittest tests/code_tests/*.py tests/physics_tests/*.py From a879639563868f139c1c49e7a13c762c56d6fb1d Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 12:12:25 -0500 Subject: [PATCH 79/83] removing deprecated method --- tests/code_tests/refactor_contract_test.py | 281 --------------------- 1 file changed, 281 deletions(-) delete mode 100644 tests/code_tests/refactor_contract_test.py diff --git a/tests/code_tests/refactor_contract_test.py b/tests/code_tests/refactor_contract_test.py deleted file mode 100644 index b39758b9..00000000 --- a/tests/code_tests/refactor_contract_test.py +++ /dev/null @@ -1,281 +0,0 @@ -import inspect -import importlib -import importlib.util -import unittest - -import jax.numpy as jnp - -from PyPIC3D import evolve -from PyPIC3D.deposition import Esirkepov -from PyPIC3D.deposition import J_from_rhov -from PyPIC3D.deposition import rho -from PyPIC3D.boundary_conditions import boundaryconditions -from PyPIC3D.boundary_conditions import ghost_cells -from PyPIC3D.diagnostics import output_adapters -from PyPIC3D.initialization import build_tiled_array, initialize_fields -from PyPIC3D.solvers import electrostatic_yee -from PyPIC3D.solvers import first_order_yee -from PyPIC3D.utilities.grids import build_tiled_yee_grids, build_yee_grid - - -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) - ntx = int(world["Nx"]) // tile_nx - nty = int(world["Ny"]) // tile_ny - ntz = int(world["Nz"]) // tile_nz - - 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) - - -class TestTiledRefactorContracts(unittest.TestCase): - def _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}, - "tile_shape": (2, 3, 2), - "guard_cells": 2, - } - vertex_grid, center_grid = build_yee_grid(world) - world["grids"] = {"vertex": vertex_grid, "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 _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) - return jnp.zeros(shape).at[1:-1, 1:-1, 1:-1].set(scale * (ii + 2.0 * jj - kk)) - - def test_tiled_vector_ghost_refresh_accepts_tuple_field_state(self): - world = self._world() - tile_shape = (2, 3, 2) - E = tuple(self._field(world, scale) for scale in (1.0, -0.5, 0.25)) - E_tiles = tuple(tile_scalar_field(component, world, tile_shape) for component in E) - - refreshed = ghost_cells.update_tiled_vector_ghost_cells(E_tiles, world, num_guard_cells=2, tile_shape=tile_shape) - self.assertEqual(len(refreshed), 3) - for component in refreshed: - self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) - - def test_evolve_loops_read_runtime_controls_from_world(self): - electrodynamic_signature = inspect.signature(evolve.time_loop_electrodynamic) - electrostatic_signature = inspect.signature(evolve.time_loop_electrostatic) - for name in ("curl_func", "J_func", "tile_shape", "g", "current_deposition"): - self.assertNotIn(name, electrodynamic_signature.parameters) - for name in ("curl_func", "J_func", "tile_shape", "g"): - self.assertNotIn(name, electrostatic_signature.parameters) - - def test_tiled_yee_updates_read_tile_metadata_from_world(self): - update_E_signature = inspect.signature(first_order_yee.update_E) - update_B_signature = inspect.signature(first_order_yee.update_B) - for name in ("curl_func", "tile_shape", "g"): - self.assertNotIn(name, update_E_signature.parameters) - self.assertNotIn(name, update_B_signature.parameters) - self.assertFalse(hasattr(first_order_yee, "update_tiled_E")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_B")) - self.assertFalse(hasattr(first_order_yee, "tiled_grid_axes_from_world")) - self.assertFalse(hasattr(first_order_yee, "tile_grid_axes")) - self.assertFalse(hasattr(first_order_yee, "tile_vector_field")) - self.assertFalse(hasattr(first_order_yee, "stack_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "unstack_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_for_pml")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_for_pml")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells")) - self.assertFalse(hasattr(first_order_yee, "apply_tiled_conducting_bc")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "update_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "fold_tiled_vector_ghost_cells_periodic")) - self.assertFalse(hasattr(first_order_yee, "empty_tiled_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "empty_tiled_vector_field")) - self.assertFalse(hasattr(first_order_yee, "tile_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "assemble_tiled_scalar_field")) - self.assertFalse(hasattr(first_order_yee, "assemble_tiled_vector_field")) - - def test_global_assembly_lives_at_output_boundary(self): - self.assertTrue(hasattr(output_adapters, "assemble_tiled_scalar_field")) - self.assertTrue(hasattr(output_adapters, "assemble_tiled_vector_field")) - - source = inspect.getsource(output_adapters) - self.assertNotIn("from PyPIC3D.solvers.yee_tiled import", source) - - def test_ghost_cell_logic_lives_in_boundary_conditions_ghost_cells(self): - for name in ( - "update_tiled_ghost_cells", - "update_tiled_vector_ghost_cells", - "update_tiled_ghost_cells_for_pml", - "update_tiled_vector_ghost_cells_for_pml", - "fold_tiled_ghost_cells", - "fold_tiled_vector_ghost_cells", - "apply_tiled_conducting_bc", - ): - self.assertTrue(hasattr(ghost_cells, name)) - - for name in ( - "update_ghost_cells", - "fold_ghost_cells", - "apply_conducting_bc", - "apply_scalar_conducting_bc", - ): - self.assertFalse(hasattr(boundaryconditions, name)) - - def test_tiled_array_initialization_lives_in_initialization(self): - world = self._world() - world["tile_shape"] = (2, 3, 2) - world["guard_cells"] = 2 - - array = build_tiled_array(world, dtype=jnp.float32) - self.assertEqual(array.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(array.dtype, jnp.float32) - self.assertTrue(jnp.all(array == 0.0)) - - def test_initialize_fields_builds_tiled_field_state_from_world(self): - world = self._world() - world["tile_shape"] = (2, 3, 2) - world["guard_cells"] = 2 - - E, B, J, phi, rho_tiles = initialize_fields(world) - for vector_field in (E, B, J): - self.assertEqual(len(vector_field), 3) - for component in vector_field: - self.assertEqual(component.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(component.dtype, jnp.float64) - self.assertTrue(jnp.all(component == 0.0)) - for scalar_field in (phi, rho_tiles): - self.assertEqual(scalar_field.shape, (4, 2, 2, 6, 7, 6)) - self.assertEqual(scalar_field.dtype, jnp.float64) - self.assertTrue(jnp.all(scalar_field == 0.0)) - - def test_grid_builders_live_in_utilities(self): - from PyPIC3D import utils - from PyPIC3D.utilities import grids - - self.assertFalse(hasattr(utils, "build_yee_grid")) - self.assertFalse(hasattr(utils, "build_collocated_grid")) - self.assertTrue(hasattr(grids, "build_yee_grid")) - self.assertTrue(hasattr(grids, "build_collocated_grid")) - self.assertFalse(hasattr(grids, "tile_grid_axes")) - self.assertTrue(hasattr(grids, "build_tiled_yee_grids")) - - def test_tiled_benchmarking_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.diagnostics.tiled_benchmarking")) - - def test_deprecated_fdtd_errors_and_root_boris_modules_are_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.solvers.fdtd")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.errors")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.boris")) - self.assertIsNotNone(importlib.util.find_spec("PyPIC3D.pusher.boris")) - - def test_old_particle_species_surfaces_are_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.species_class")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.particles.tiled_particle_diagnostics")) - - from PyPIC3D.particles import particle_initialization - - self.assertFalse(hasattr(particle_initialization, "to_tiled_particles")) - source = inspect.getsource(particle_initialization.load_particles_from_toml) - self.assertNotIn("particle_species", source) - - def test_particle_push_owns_tiled_pusher_contract(self): - particle_push = importlib.import_module("PyPIC3D.pusher.particle_push") - - self.assertIsNone(importlib.util.find_spec("PyPIC3D.pusher.tiled_pusher")) - signature = inspect.signature(particle_push.particle_push) - for name in ("grid", "staggered_grid", "tile_shape", "g"): - self.assertNotIn(name, signature.parameters) - - source = inspect.getsource(particle_push.particle_push) - self.assertIn('world["grids"]["tiled_center_grid"]', source) - self.assertIn('world["grids"]["tiled_vertex_grid"]', source) - self.assertNotIn("tiled_grid_axes_from_world", source) - self.assertNotIn("tile_grid_axes", source) - self.assertNotIn("get_forward_position", source) - self.assertNotIn("set_velocity", source) - - def test_electrodynamic_hot_step_does_not_assemble_global_fields(self): - source = inspect.getsource(evolve.time_loop_electrodynamic) - self.assertNotIn("assemble_tiled_scalar_field", source) - self.assertNotIn("assemble_tiled_vector_field", source) - self.assertNotIn("fields_for_output", source) - - def test_electrostatic_solver_uses_single_tile_arrays_directly(self): - source = inspect.getsource(electrostatic_yee.calculate_tiled_electrostatic_fields) - self.assertNotIn("assemble_tiled_scalar_field", source) - self.assertNotIn("tile_scalar_field", source) - self.assertIn("[0, 0, 0]", source) - - def test_flat_current_and_electrodynamic_references_are_removed(self): - self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_flat")) - self.assertFalse(hasattr(J_from_rhov, "_J_from_rhov_tiled")) - self.assertFalse(hasattr(evolve, "_time_loop_electrodynamic_global_reference")) - self.assertFalse(hasattr(evolve, "_time_loop_electrostatic_global_reference")) - - def test_direct_current_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(J_from_rhov.J_from_rhov) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_esirkepov_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(Esirkepov.Esirkepov_current) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_tiled_esirkepov_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.esirkepov_tiled")) - - def test_rho_api_reads_tile_metadata_from_world(self): - signature = inspect.signature(rho.compute_rho) - self.assertNotIn("tile_shape", signature.parameters) - self.assertNotIn("g", signature.parameters) - - def test_flat_rho_and_tiled_rho_module_are_removed(self): - self.assertFalse(hasattr(rho, "_compute_rho_flat")) - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.rho_tiled")) - - def test_current_methods_module_is_removed(self): - self.assertIsNone(importlib.util.find_spec("PyPIC3D.deposition.current_methods")) - - -if __name__ == "__main__": - unittest.main() From b36f948b563343cbc5809f1c23fbecfa7bc2ea65 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 12:15:38 -0500 Subject: [PATCH 80/83] updating actions for faster testing --- .github/workflows/coverage.yml | 2 +- .github/workflows/unit_testing.yml | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 60078c60..6a44ee3f 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -26,7 +26,7 @@ jobs: - name: Run tests with coverage run: | - coverage run -m unittest tests/code_tests/*.py tests/physics_tests/*.py + 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 64e40b42..e7dd3937 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -24,7 +24,8 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip + pip install unittest-parallel pip install . - name: Test with unittest run: | - python -m unittest tests/code_tests/*.py tests/physics_tests/*.py + unittest-parallel -t . -s tests -p "*_test.py" From 0b64cd98994b7864120556aeed63da91d9ac1462 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 12:25:34 -0500 Subject: [PATCH 81/83] going back to unit test --- .github/workflows/unit_testing.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index e7dd3937..64e40b42 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -24,8 +24,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install unittest-parallel pip install . - name: Test with unittest run: | - unittest-parallel -t . -s tests -p "*_test.py" + python -m unittest tests/code_tests/*.py tests/physics_tests/*.py From efe8c32ce9f52c603f599270b0d5105f627284e6 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 12:26:08 -0500 Subject: [PATCH 82/83] just doing code tests for right now due to resource limits --- .github/workflows/unit_testing.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index 64e40b42..cec61c6e 100644 --- a/.github/workflows/unit_testing.yml +++ b/.github/workflows/unit_testing.yml @@ -27,4 +27,5 @@ jobs: pip install . - name: Test with unittest run: | - python -m unittest tests/code_tests/*.py tests/physics_tests/*.py + python -m unittest tests/code_tests/*.py +# tests/physics_tests/*.py From fbd0f6db4b7eefe1ebe90e9928c2aaea6c6752d2 Mon Sep 17 00:00:00 2001 From: Christopher Woolford Date: Tue, 14 Jul 2026 12:37:58 -0500 Subject: [PATCH 83/83] switching to self-hosted runners --- .github/workflows/coverage.yml | 2 +- .github/workflows/unit_testing.yml | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 6a44ee3f..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 diff --git a/.github/workflows/unit_testing.yml b/.github/workflows/unit_testing.yml index cec61c6e..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,5 +27,4 @@ jobs: pip install . - name: Test with unittest run: | - python -m unittest tests/code_tests/*.py -# tests/physics_tests/*.py + python -m unittest tests/code_tests/*.py tests/physics_tests/*.py