jaxincell

Submodules

Attributes

epsilon_0

mu_0

speed_of_light

elementary_charge

mass_electron

mass_proton

boltzmann_constant

_MAX_FILTER_PASSES

Classes

Simulation

Calling Simulation(parameters) will create a Simulation object with the provided parameters.

Functions

set_BC_single_particle(x_n, v_n, q, q_m, dx, grid, ...)

Applies boundary conditions (BCs) to a single particle's position and velocity.

set_BC_particles(xs_n, vs_n, qs, ms, q_ms, dx, grid, ...)

Applies boundary conditions to all particles in parallel.

set_BC_single_particle_positions(x_n, dx, grid, ...)

Applies boundary conditions to particle positions only (used for half-step updates).

set_BC_positions(xs_n, qs, dx, grid, box_size_x, ...)

Applies boundary conditions to particle positions for all particles during a half-step update.

diagnostics(output)

E_from_Gauss_1D_FFT(charge_density, dx)

Solve for the electric field E = -d(phi)/dx using FFT,

E_from_Poisson_1D_FFT(charge_density, dx)

Solve for the electric field E = -d(phi)/dx using FFT,

E_from_Gauss_1D_Cartesian(charge_density, dx)

Solve for the electric field at t=0 (E0) using the charge density distribution

curlE(E_field, B_field, dx, dt, field_BC_left, ...)

Compute the curl of the electric field, which is related to the time derivative of

curlB(B_field, E_field, dx, dt, field_BC_left, ...)

Compute the curl of the magnetic field, which is related to the time derivative of

field_update(E_fields, B_fields, dx, dt, j, ...)

Update the electric and magnetic fields based on Maxwell's equations

field_update1(E_fields, B_fields, dx, dt, j, ...)

field_update2(E_fields, B_fields, dx, dt, j, ...)

fields_to_particles_grid(x_n, field, dx, grid, ...)

This function retrieves the electric or magnetic field values at particle positions

fields_to_particles_periodic_CN(x_n, field, dx, grid_start)

Interpolates field to particle using Periodic BCs.

rotation(dt, B, vsub, q_m)

This function implements the Boris algorithm to rotate the particle velocity vector

boris_step(dt, xs_nplushalf, vs_n, q_ms, ...)

This function performs one step of the Boris algorithm for particle motion.

boris_step_relativistic(dt, xs_nplushalf, vs_n, q_s, ...)

Relativistic Boris pusher for N particles.

plot(output[, direction, threshold, save_mp4, fps, ...])

Production-ready plotting/animation for JAX-in-Cell outputs.

load_parameters(input_file)

Load parameters from a given .toml input file given the path to the file.

get_S2_weights_and_indices_periodic_CN(x, dx, ...)

Calculates weights and indices for Quadratic Spline (S2).

charge_density_BCs(particle_BC_left, ...)

Compute the charge contribution to the boundary points based on particle positions and boundary conditions.

single_particle_charge_density(x, q, dx, grid, ...)

Computes the charge density contribution of a single particle to the grid using a

calculate_charge_density(xs_n, qs, dx, grid, ...[, ...])

Computes the total charge density on the grid by summing contributions from all particles.

current_density(xs_nminushalf, xs_n, xs_nplushalf, ...)

Computes the current density j on the grid from particle motion.

current_density_periodic_CN(xs_n, vs_n, qs, dx, ...)

Deposits Current J using Periodic BCs.

_shift_with_bc_1d(x, shift, bc_left, bc_right)

Shift x by 'shift' cells along axis=0 with boundary conditions:

binomial_filter_3point(x[, alpha, stride, bc_left, ...])

3-point digital filter along axis 0 with BCs:

_repeat_filter(y, stride, passes, alpha[, bc_left, ...])

JAX-safe version of the 3-point digital filter with compensation:

filter_scalar_field(scalar_field[, passes, alpha, ...])

Apply a multi-pass 3-point binomial digital filter to a scalar field.

filter_vector_field(F[, passes, alpha, strides, ...])

Apply digital filter along the grid axis (axis=0) for each component.

Package Contents

jaxincell.set_BC_single_particle(x_n, v_n, q, q_m, dx, grid, box_size_x, box_size_y, box_size_z, BC_left, BC_right)

Applies boundary conditions (BCs) to a single particle’s position and velocity.

Parameters:

  • x_n (jnp.ndarray) – Particle position as a 1D array [x, y, z].

  • v_n (jnp.ndarray) – Particle velocity as a 1D array [vx, vy, vz].

  • q (float) – Particle charge.

  • q_m (float) – Charge-to-mass ratio of the particle.

  • dx (float) – Grid spacing.

  • grid (jnp.ndarray) – Discretized grid positions.

  • box_size_x (float) – Box dimensions in x, y, and z directions.

  • box_size_y (float) – Box dimensions in x, y, and z directions.

  • box_size_z (float) – Box dimensions in x, y, and z directions.

  • BC_left (int) – Boundary conditions for left and right boundaries in the x-direction. 0: Periodic 1: Reflective 2: Absorbing

  • BC_right (int) – Boundary conditions for left and right boundaries in the x-direction. 0: Periodic 1: Reflective 2: Absorbing

Returns:

Updated position (x_n), velocity (v_n), charge (q), and charge-to-mass ratio (q_m).

Return type:

tuple

jaxincell.set_BC_particles(xs_n, vs_n, qs, ms, q_ms, dx, grid, box_size_x, box_size_y, box_size_z, BC_left, BC_right)

Applies boundary conditions to all particles in parallel.

Parameters:

  • xs_n (jnp.ndarray) – Positions of all particles, shape (N, 3).

  • vs_n (jnp.ndarray) – Velocities of all particles, shape (N, 3).

  • qs (jnp.ndarray) – Charges of all particles, shape (N,).

  • ms (jnp.ndarray) – Masses of all particles, shape (N,).

  • q_ms (jnp.ndarray) – Charge-to-mass ratios of all particles, shape (N,).

  • parameters (Other) – Same as set_BCs.

Returns:

Updated positions, velocities, charges, masses, and charge-to-mass ratios for all particles.

Return type:

tuple

jaxincell.set_BC_single_particle_positions(x_n, dx, grid, box_size_x, box_size_y, box_size_z, BC_left, BC_right)

Applies boundary conditions to particle positions only (used for half-step updates).

Parameters:

  • x_n (jnp.ndarray) – Particle position as a 1D array [x, y, z].

  • parameters (Other) – Same as set_BCs.

Returns:

Updated particle position [x, y, z].

Return type:

jnp.ndarray

jaxincell.set_BC_positions(xs_n, qs, dx, grid, box_size_x, box_size_y, box_size_z, BC_left, BC_right)

Applies boundary conditions to particle positions for all particles during a half-step update.

Parameters:

  • xs_n (jnp.ndarray) – Positions of all particles, shape (N, 3).

  • qs (jnp.ndarray) – Charges of all particles, shape (N,).

  • parameters (Other) – Same as set_BCs.

Returns:

Updated positions of all particles, shape (N, 3).

Return type:

jnp.ndarray

jaxincell.epsilon_0 = 8.85418782e-12
jaxincell.mu_0 = 1.25663706e-06
jaxincell.speed_of_light = 299792458.0
jaxincell.elementary_charge = 1.60217663e-19
jaxincell.mass_electron = 9.10938371e-31
jaxincell.mass_proton = 1.67262193e-27
jaxincell.boltzmann_constant = 1.380649e-23
jaxincell.diagnostics(output)
jaxincell.E_from_Gauss_1D_FFT(charge_density, dx)

Solve for the electric field E = -d(phi)/dx using FFT, where phi is derived from the 1D Gauss’ law equation. Parameters: charge_density : 1D numpy array, source term (right-hand side of Poisson equation) dx : float, grid spacing in the x-direction Returns: E : 1D numpy array, electric field

jaxincell.E_from_Poisson_1D_FFT(charge_density, dx)

Solve for the electric field E = -d(phi)/dx using FFT, where phi is derived from the 1D Poisson equation. Parameters: charge_density : 1D numpy array, source term (right-hand side of Poisson equation) dx : float, grid spacing in the x-direction Returns: E : 1D numpy array, electric field

jaxincell.E_from_Gauss_1D_Cartesian(charge_density, dx)

Solve for the electric field at t=0 (E0) using the charge density distribution and applying Gauss’s law in a 1D system.

Parameters:

  • charge_density – 1D numpy array, source term (right-hand side of Gauss equation)

  • dx – float, grid spacing in the x-direction

Returns:

The electric field at each grid point due to the particles, shape (G,).

Return type:

array

jaxincell.curlE(E_field, B_field, dx, dt, field_BC_left, field_BC_right)

Compute the curl of the electric field, which is related to the time derivative of the magnetic field in Maxwell’s equations (Faraday’s law).

Parameters:

  • E_field (array) – Electric field at each grid point, shape (G, 3).

  • B_field (array) – Magnetic field at each grid point, shape (G, 3).

  • dx (float) – Grid spacing.

  • dt (float) – Time step.

  • field_BC_left (int) – Left boundary condition for fields (0: periodic, 1: reflective, 2: absorbing).

  • field_BC_right (int) – Right boundary condition for fields.

Returns:

The curl of the electric field, which is the source of the magnetic field.

Return type:

array

jaxincell.curlB(B_field, E_field, dx, dt, field_BC_left, field_BC_right)

Compute the curl of the magnetic field, which is related to the time derivative of the electric field in Maxwell’s equations (Ampère’s law with Maxwell correction).

Parameters:

  • B_field (array) – Magnetic field at each grid point, shape (G, 3).

  • E_field (array) – Electric field at each grid point, shape (G, 3).

  • dx (float) – Grid spacing.

  • dt (float) – Time step.

  • field_BC_left (int) – Left boundary condition for fields.

  • field_BC_right (int) – Right boundary condition for fields.

Returns:

The curl of the magnetic field, which is the source of the electric field.

Return type:

array

jaxincell.field_update(E_fields, B_fields, dx, dt, j, field_BC_left, field_BC_right)

Update the electric and magnetic fields based on Maxwell’s equations

Parameters:

  • E_fields (array) – Electric field at each grid point, shape (G, 3).

  • B_fields (array) – Magnetic field at each grid point, shape (G, 3).

  • dx (float) – Grid spacing.

  • dt (float) – Time step.

  • j (array) – Current density at each grid point, shape (G, 3).

  • field_BC_left (int) – Left boundary condition for fields.

  • field_BC_right (int) – Right boundary condition for fields.

Returns:

Updated electric and magnetic fields, each of shape (G, 3).

Return type:

tuple

jaxincell.field_update1(E_fields, B_fields, dx, dt, j, field_BC_left, field_BC_right)
jaxincell.field_update2(E_fields, B_fields, dx, dt, j, field_BC_left, field_BC_right)
jaxincell.fields_to_particles_grid(x_n, field, dx, grid, grid_start, field_BC_left, field_BC_right)

This function retrieves the electric or magnetic field values at particle positions using a field interpolation scheme. The function first adds ghost cells to the field array to handle boundary conditions, then interpolates the field based on the particle’s position in the grid.

Parameters:

  • x_n (array) – The position of particles at time step n, shape (N,).

  • field (array) – The field values at each grid point, shape (G,).

  • dx (float) – The spatial grid spacing.

  • grid (array) – The grid positions where the field is defined, shape (G,).

  • grid_start (float) – The starting position of the grid (usually the left boundary).

  • field_BC_left (int) – Boundary condition for the left side of the particle grid.

  • field_BC_right (int) – Boundary condition for the right side of the particle grid.

Returns:

The interpolated field values at the particle positions, shape (N,).

Return type:

array

jaxincell.fields_to_particles_periodic_CN(x_n, field, dx, grid_start)

Interpolates field to particle using Periodic BCs. :param field: The field array (size N). :param grid_start: Physical position of field[0].

jaxincell.rotation(dt, B, vsub, q_m)

This function implements the Boris algorithm to rotate the particle velocity vector in the magnetic field for one time step. This step is part of the numerical solution of the Lorentz force equation.

Parameters:

  • dt (float) – Time step for the simulation.

  • B (array) – Magnetic field at the particle’s position, shape (3,).

  • vsub (array) – The particle’s velocity before the rotation, shape (3,).

  • q_m (array) – The charge-to-mass ratio of the particle, shape (3,).

Returns:

The updated velocity after the rotation, shape (3,).

Return type:

array

jaxincell.boris_step(dt, xs_nplushalf, vs_n, q_ms, E_fields_at_x, B_fields_at_x)

This function performs one step of the Boris algorithm for particle motion. The particle velocity is updated using the electric and magnetic fields at its position, and the particle position is updated using the new velocity.

Parameters:

  • dt (float) – Time step for the simulation.

  • xs_nplushalf (array) – The particle positions at the half-time step n+1/2, shape (N, 3).

  • vs_n (array) – The particle velocities at time step n, shape (N, 3).

  • q_ms (array) – The charge-to-mass ratio of each particle, shape (N, 1).

  • E_fields_at_x (array) – The interpolated electric field values at the particle positions, shape (N, 3).

  • B_fields_at_x (array) – The magnetic field values at the particle positions, shape (N, 3).

Returns:

A tuple containing:

  • xs_nplus3_2 (array): The updated particle positions at time step n+3/2, shape (N, 3).

  • vs_nplus1 (array): The updated particle velocities at time step n+1, shape (N, 3).

Return type:

tuple

jaxincell.boris_step_relativistic(dt, xs_nplushalf, vs_n, q_s, m_s, E_fields_at_x, B_fields_at_x)

Relativistic Boris pusher for N particles.

Parameters:

  • dt – Time step

  • xs_nplushalf – Particle positions at t = n + 1/2, shape (N, 3)

  • vs_n – Velocities at time t = n, shape (N, 3)

  • q_s – Charges, shape (N,)

  • m_s – Masses, shape (N,)

  • E_fields_at_x – Electric fields at particle positions, shape (N, 3)

  • B_fields_at_x – Magnetic fields at particle positions, shape (N, 3)

  • c – Speed of light (default = 1.0 for normalized units)

Returns:

Updated positions at t = n + 3/2, shape (N, 3) vs_nplus1: Updated velocities at t = n + 1, shape (N, 3)

Return type:

xs_nplus3_2

jaxincell.plot(output, direction: str = 'x', threshold: float = 1e-12, save_mp4: str | None = None, fps: int = 30, dpi: int = 150, show: bool = True, animation_interval: int = 1, save_stride: int = 1, save_dpi: int | None = None, save_crf: int | None = None, save_preset: str | None = None, save_codec: str | None = None)

Production-ready plotting/animation for JAX-in-Cell outputs.

What you get:

  1. Heatmaps (x vs time): E, B (nonzero components), charge density. - IMPORTANT: heatmap color limits are fixed over the whole run using a robust

    percentile, so growth/decay in time is visible (no per-frame re-normalization).

  2. Instantaneous E(x,t) overlay: - For each plotted electric-field component, we draw a line on top of the heatmap. - The line is normalized by a single GLOBAL robust scale over the whole run,

    so amplitude growth is visible.

    • Overlay axes have no ticks (prevents clashes with colorbars).

  3. Distribution functions f(v,t) (LAB FRAME; NO drift centering): - Shown as clean line plots in their own subplot(s) (no current-density heatmap). - Solid: current frame - Dashed: initial (t=0), labeled “(initial)”

  4. Phase space (x vs v) for electrons and ions for each requested component: - Uses a robust velocity range per species per component, so ion dynamics remains visible. - Uses LogNorm on counts (with +1 internally) so low-density structure is visible.

Multi-species:

  • Uses diagnostics() legacy split if present (velocity_electrons/velocity_ions).

  • Otherwise combines output[“species”] by charge sign (q<0 as electrons, q>0 as ions).

class jaxincell.Simulation(parameters=None)

Calling Simulation(parameters) will create a Simulation object with the provided parameters. The parameters should be provided as a dictionary, but can also be provided as a path to a .toml file containing the parameters. The parameters will be cleaned and initialized using the appropriate cleaner functions for each parameter section, and the simulation state will be initialized based on the cleaned parameters.

The Simulation object will expose any differentiable parameters that were provided in input_parameters through Simulation_object.input_parameters. These input_parameters can then be passed to the simulation() or run() bound functions to run a simulation with these input_parameters exposed such that grads can be taken with respect to them. If simulation() or run() is called without input_parameters, it will use the parameters provided at initialization (which may include user-provided parameters and defaults) and will not overwrite any parameters with the input_parameters.

Example without input parameters:

sim = Simulation(parameters) simulation_output = sim.simulation() or simulation_output = sim.run()

Example with input parameters:

sim = Simulation(parameters) input_parameters = sim.input_parameters simulation_output = sim.simulation(input_parameters) or simulation_output = sim.run(input_parameters)

Additionally,

def scalar_objective(input_parameters):

simulation_output = sim.run(input_parameters) return some_scalar_function_of(simulation_output)

grad(scalar_objective)(input_parameters)

will provide the gradient of some scalar objective function of the simulation output with respect to the input_parameters.

simulation(input_parameters=None)

Exposed simulation call which doesn’t expose the hash values for each section to prevent unintentionally forcing recompiles or not recompiling when necessary.

input_parameters is a dictionary of differentiable parameters If simulation is called without input parameters, it will use the user input parameters or default parameters previously provided. input_parameters will overwrite any differentiable parameters previously provided. This is meant to make it simple to expose grads derivatives with respect to the input parameters.

run(input_parameters=None)
_simulation(input_parameters=None, domain_hash='', species_hash='', external_field_hash='', source_hash='', solver_hash='')

Run a plasma physics simulation using a Particle-In-Cell (PIC) method in JAX.

This function simulates the evolution of a plasma system by solving for particle motion (electrons and ions) and self-consistent electromagnetic fields on a grid. It uses the Boris algorithm for particle updates and a leapfrog scheme for field updates.

Parameters:

user_parametersdict

User-defined parameters for the simulation. These can include: - Physical parameters: box size, number of particles, thermal velocities. - Numerical parameters: grid resolution, time step size. - Boundary conditions for particles and fields. - Random seed for reproducibility.

Returns:

output : dict

assemble_output(simulation_output, input_parameters)
clean_and_initialize_parameters(parameters)
classify_and_sort_input_parameters(parameters)

Sort through input parameters to move parameters into their respective dictionaries to overwrite defaults and move differentiable parameters into a separate input_parameters dictionary. This input_parameters dictionary can then be accessed to use them as inputs to the simulation function without having to write multiple toml files for the differentiable inputs and the non-differentiable parameters.

build_hash_values()

Build the hashes for each of the parameter sections to help with determining when Jax needs to recompile the _simulation() function due to new parameters being passed.

reinitialize_simulation_state()

Reinitialize the simulation state based on the current parameter sections. This should be called whenever parameters are updated after initialization to ensure that the simulation state is consistent with the new parameters.

clean_runtime_input_parameters(input_parameters=None)

Clean the input_parameters provided to the simulation(…) or run(…) functions at runtime.

current_domain_state()
build_domain()
initialize_particles()
initialize_fields()
set_parameter_section(section_name, new_parameters)

Helper which is used by setters for the parameter sections to automatically clean parameters and reinitialize the state of the simulation including creating new hashes.

property domain_parameters
property species_parameters
property external_field_parameters
property source_parameters
property solver_parameters
property input_parameters
jaxincell.load_parameters(input_file)

Load parameters from a given .toml input file given the path to the file.

jaxincell.get_S2_weights_and_indices_periodic_CN(x, dx, grid_start, grid_size)

Calculates weights and indices for Quadratic Spline (S2). Applies Periodic Wrapping to indices immediately.

jaxincell.charge_density_BCs(particle_BC_left, particle_BC_right, position, dx, grid, charge)

Compute the charge contribution to the boundary points based on particle positions and boundary conditions.

Parameters:

  • particle_BC_left (int) – Boundary condition for the left edge (0: periodic, 1: reflective, 2: absorbing).

  • particle_BC_right (int) – Boundary condition for the right edge (0: periodic, 1: reflective, 2: absorbing).

  • position (float) – Position of the particle.

  • dx (float) – Grid spacing.

  • grid (array-like) – Grid points as a 1D array.

  • charge (float) – Charge of the particle.

Returns:

Charge contributions to the left and right boundaries.

Return type:

tuple

jaxincell.single_particle_charge_density(x, q, dx, grid, particle_BC_left, particle_BC_right)

Computes the charge density contribution of a single particle to the grid using a quadratic particle shape function.

Parameters:

  • x (float) – The particle position.

  • q (float) – The particle charge.

  • dx (float) – The grid spacing.

  • grid (array) – The grid points.

  • particle_BC_left (int) – Left boundary condition type (0: periodic, 1: reflective, 2: absorbing).

  • particle_BC_right (int) – Right boundary condition type (0: periodic, 1: reflective, 2: absorbing).

Returns:

The charge density contribution on the grid.

Return type:

array

jaxincell.calculate_charge_density(xs_n, qs, dx, grid, particle_BC_left, particle_BC_right, filter_passes=5, filter_alpha=0.5, filter_strides=(1, 2, 4), field_BC_left=0, field_BC_right=0)

Computes the total charge density on the grid by summing contributions from all particles.

Parameters:

  • xs_n (array) – Particle positions at the current timestep, shape (N, 1).

  • qs (array) – Particle charges, shape (N, 1).

  • dx (float) – The grid spacing.

  • grid (array) – The grid points.

  • particle_BC_left (int) – Left particle boundary condition type (0: periodic, 1: reflective, 2: absorbing).

  • particle_BC_right (int) – Right particle boundary condition type (0: periodic, 1: reflective, 2: absorbing).

  • filter_passes (int) – Number of digital filter passes to apply (default: 5). Internally capped at 17.

  • filter_alpha (float) – Filter strength parameter (default: 0.5). Controls the weight of the center point in the 3-point filter.

  • filter_strides (tuple) – Tuple of stride values for multi-scale filtering (default: (1, 2, 4)).

  • field_BC_left (int) – Left boundary condition for filtering (default: 0: periodic, 1: reflective, 2: absorbing).

  • field_BC_right (int) – Right boundary condition for filtering (default: 0: periodic, 1: reflective, 2: absorbing).

Returns:

Total charge density on the grid.

Return type:

array

jaxincell.current_density(xs_nminushalf, xs_n, xs_nplushalf, vs_n, qs, dx, dt, grid, grid_start, particle_BC_left, particle_BC_right, filter_passes=5, filter_alpha=0.5, filter_strides=(1, 2, 4), field_BC_left=0, field_BC_right=0)

Computes the current density j on the grid from particle motion.

Parameters:

  • xs_nminushalf (array) – Particle positions at the half timestep before the current one, shape (N, 1).

  • xs_n (array) – Particle positions at the current timestep, shape (N, 1).

  • xs_nplushalf (array) – Particle positions at the half timestep after the current one, shape (N, 1).

  • vs_n (array) – Particle velocities at the current timestep, shape (N, 3).

  • qs (array) – Particle charges, shape (N, 1).

  • dx (float) – The grid spacing.

  • dt (float) – The time step size.

  • grid (array) – The grid points.

  • grid_start (float) – The starting position of the grid.

  • particle_BC_left (int) – Left particle boundary condition type (0: periodic, 1: reflective, 2: absorbing).

  • particle_BC_right (int) – Right particle boundary condition type (0: periodic, 1: reflective, 2: absorbing).

  • filter_passes (int) – Number of digital filter passes to apply (default: 5). Internally capped at 17.

  • filter_alpha (float) – Filter strength parameter (default: 0.5). Controls the weight of the center point in the 3-point filter.

  • filter_strides (tuple) – Tuple of stride values for multi-scale filtering (default: (1, 2, 4)).

  • field_BC_left (int) – Left boundary condition for filtering (default: 0: periodic, 1: reflective, 2: absorbing).

  • field_BC_right (int) – Right boundary condition for filtering (default: 0: periodic, 1: reflective, 2: absorbing).

Returns:

Current density on the grid, shape (G, 3), where G is the number of grid points.

Return type:

array

jaxincell.current_density_periodic_CN(xs_n, vs_n, qs, dx, grid_start, grid_size)

Deposits Current J using Periodic BCs. Note: We removed xs_nminushalf/plus half arguments as we use the midpoint approximation (xs_n, vs_n) consistent with CN.

jaxincell._MAX_FILTER_PASSES = 16
jaxincell._shift_with_bc_1d(x, shift, bc_left, bc_right)

Shift x by ‘shift’ cells along axis=0 with boundary conditions:

bc = 0: periodic bc = 1: reflective (clamp to boundary cell) bc = 2: absorbing (outside domain -> 0)

Works for arrays with shape (G, …) – only axis 0 is shifted.

jaxincell.binomial_filter_3point(x, alpha=0.5, stride=1, bc_left=0, bc_right=0)

3-point digital filter along axis 0 with BCs:

x^f_j = α x_j + (1-α)/2 [ x_{j-stride} + x_{j+stride} ]

bc_left / bc_right:

0: periodic 1: reflective (clamp) 2: absorbing (outside -> 0)

jaxincell._repeat_filter(y, stride, passes, alpha, bc_left=0, bc_right=0)

JAX-safe version of the 3-point digital filter with compensation: - If passes <= 0: return y unchanged. - If passes > 0:

  • apply (passes - 1) regular binomial_filter_3point passes with alpha (up to a static maximum _MAX_FILTER_PASSES),

  • then a final compensation pass with comp_alpha = passes - alpha*(passes - 1).

This is fully jit- and grad-safe even when passes is a traced value.

Note: The number of regular filter passes is internally capped at _MAX_FILTER_PASSES (16). If passes > _MAX_FILTER_PASSES + 1, the function will only apply _MAX_FILTER_PASSES regular passes plus one compensation pass, which may not match the expected filtering behavior. For typical use cases (default is 5), this limit should not be reached.

jaxincell.filter_scalar_field(scalar_field, passes=5, alpha=0.5, strides=(1, 2, 4), bc_left=0, bc_right=0)

Apply a multi-pass 3-point binomial digital filter to a scalar field.

Parameters:

  • scalar_field – Input scalar field array to be filtered.

  • passes – Number of filter passes (default: 5). Note: internally capped at 17 (16 regular passes + 1 compensation pass).

  • alpha – Filter strength parameter (default: 0.5).

  • strides – Tuple/list of stride values for filtering (default: (1, 2, 4)).

  • bc_left – Boundary condition for the left side (default: 0). 0: periodic, 1: reflective, 2: absorbing.

  • bc_right – Boundary condition for the right side (default: 0). 0: periodic, 1: reflective, 2: absorbing.

Returns:

Filtered scalar field array.

jaxincell.filter_vector_field(F, passes=5, alpha=0.5, strides=(1, 2, 4), bc_left=0, bc_right=0)

Apply digital filter along the grid axis (axis=0) for each component. F has shape (G, C), typically (grid_points, 3) for a vector field.

Parameters:

  • F – Input vector field array, shape (G, C), typically (grid_points, 3).

  • passes – Number of filter passes (default: 5). Note: internally capped at 17 (16 regular passes + 1 compensation pass).

  • alpha – Filter strength parameter (default: 0.5).

  • strides – Tuple/list of stride values for filtering (default: (1, 2, 4)).

  • bc_left – Boundary condition for the left side (default: 0). 0: periodic, 1: reflective, 2: absorbing.

  • bc_right – Boundary condition for the right side (default: 0). 0: periodic, 1: reflective, 2: absorbing.

Returns:

Filtered vector field array with the same shape as input.