Fixed charge and conservation error in explicit + Faster boundary conditions

[rogeriojorge]

There was a bug in the current deposition formula, which would make the charge not be conserved in the explicit (Boris) method. By fixing it, energy is now conserved at a much better level, and charge conservation is also achieved.
In a few places, including boundary conditions, vmap was replaced by simple vectorization, making the code run faster.

Please try this new version to prepare for the merge.

@Longyu975 Implicit method does not appear to be conserving charge (it is not plotted in _plot.py if the charge error is too big), any way to fix that?

Last thing I want to make before JOSS submission is update README and docs to be more clear what all of the input parameters mean, and have some movies with different boundary conditions, such as the one on the bump on tail example now.

[Copilot]

Pull request overview

This PR fixes a charge conservation bug in the explicit (Boris) method and improves performance by replacing vmap with vectorization in boundary conditions. The fix significantly improves energy conservation and achieves proper charge conservation. Additionally, the code now enforces Gauss's law more robustly, removes the field_solver parameter (now always uses curl-based updates with Gauss enforcement), and enhances plotting capabilities with video export and improved diagnostics.

Key Changes:

• Fixed current deposition formula to conserve charge in explicit method via continuity equation
• Vectorized boundary condition code for better performance
• Enhanced Gauss law enforcement with periodic/non-periodic support
• Improved plotting with video export, distribution functions, and energy diagnostics

Reviewed changes

Copilot reviewed 20 out of 21 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
jaxincell/_sources.py	Rewrote current/charge deposition with new S2 functions and continuity-preserving current calculation
jaxincell/_fields.py	Replaced multiple Gauss solvers with single enforce_gauss_1d function using FFT or integration
jaxincell/_boundary_conditions.py	Vectorized particle BC functions, removed single-particle and ghost cell helpers
jaxincell/_particles.py	Updated field interpolation to use new S2 gather, simplified Boris step
jaxincell/_algorithms.py	Integrated new current deposition and Gauss enforcement, removed field_solver parameter
jaxincell/_simulation.py	Removed field_solver static argument, added Gauss parameters
jaxincell/_diagnostics.py	Added Gauss law error diagnostics (L∞ and L2 norms)
jaxincell/_plot.py	Major rewrite with video export, distribution functions, phase space improvements, and energy plots
tests/test_simulation.py	Added tests for load_parameters and initialize_simulation_parameters
tests/test_plots.py	Comprehensive new tests for plotting functionality
tests/test_boundary_conditions.py	Removed tests for deleted single-particle BC functions
examples/*.py	Updated examples to remove field_solver parameter and add video export demonstrations
examples/*.toml	Updated configuration files to remove field_solver and adjust parameters

Comments suppressed due to low confidence (1)

jaxincell/_fields.py:1

• The hardcoded epsilon 1e-15 should be a named constant or parameter. Consider using jnp.finfo(v.dtype).eps for dtype-aware precision, especially since the code elsewhere uses both float32 and float64.

from jax import jit, lax

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The rounding method has changed from jnp.round() to jnp.floor(x_norm + 0.5). While mathematically equivalent for most cases, this could produce different results for exact half-integers due to different tie-breaking rules. jnp.round() uses banker's rounding (round half to even), while floor(x + 0.5) always rounds half up. Consider documenting this behavior change or verifying it doesn't affect simulation results.

Suggested change

	k = jnp.floor(x_norm + 0.5).astype(jnp.int32)
	k = jnp.round(x_norm).astype(jnp.int32)

[Copilot]

The docstring should clarify the input/output shapes and data types. Add parameter descriptions for rho_prev, rho_next, dx, dt and return value description for Jx.

Suggested change

	Periodic solve for Jx enforcing discrete continuity with backward-diff div:
	(Jx - roll(Jx,1))/dx = -(rho_next - rho_prev)/dt
	Gauge: set k=0 mode of Jx to 0 (mean current is undetermined by continuity).
	Computes the current density Jx from the discrete continuity equation using a periodic
	backward-difference scheme and FFT-based solver.
	Solves for Jx such that:
	(Jx - roll(Jx, 1)) / dx = -(rho_next - rho_prev) / dt
	with periodic boundary conditions. The k=0 (mean) mode of Jx is set to zero (gauge choice).
	Parameters
	----------
	rho_prev : jax.numpy.ndarray, shape (n,), dtype float
	Charge density at the previous time step.
	rho_next : jax.numpy.ndarray, shape (n,), dtype float
	Charge density at the next time step.
	dx : float
	Grid spacing.
	dt : float
	Time step.
	Returns
	-------
	Jx : jax.numpy.ndarray, shape (n,), dtype float
	Current density array satisfying the discrete continuity equation with periodic boundary conditions.

[Copilot]

Extra space before comment marker. Should be # QuickTime instead of # QuickTime .

Suggested change

	tag_args = ["-tag:v", "hvc1"] # QuickTime compatibility for HEVC in MP4
	tag_args = ["-tag:v", "hvc1"] # QuickTime compatibility for HEVC in MP4

[Copilot]

The hardcoded epsilon 1e-12 for numerical stability should be a named constant or use a dtype-aware value like jnp.finfo(E_calculated.dtype).eps to handle both float32 and float64 correctly.

Suggested change

	delta_E = jnp.max(jnp.abs(E_calculated - E_guess)) / (jnp.max(jnp.abs(E_calculated)) + 1e-12)
	delta_E = jnp.max(jnp.abs(E_calculated - E_guess)) / (jnp.max(jnp.abs(E_calculated)) + jnp.finfo(E_calculated.dtype).eps)

[Longyu975]

I am skeptical that exact implicit charge conservation is feasible here. Chen's paper suggests a method similar to explicit schemes, but it requires particles to stop exactly at cell boundaries and update everything once to ensure energy conservation. In my view, implementing such adaptive sub-stepping is impossible because the amount of particle and also JAX enforces fixed substep sizes.

One note for Weibel, current energy is fine with the filter. But if we set the filter to 0, it breaks energy unless we change dt/c back to 0.5.

[Longyu975]

Another finding concerns the runtime comparison. CPU computations are now significantly faster, showing a speedup of roughly a factor of 4–6 compared to previous results. In contrast, the GPU implementation appears to have a lower bound of about 2 seconds. The current filter accounts for roughly a 1 second baseline, and the new ghost-cell implementation introduces an additional ~1 second overhead. Previous run time is about 0.2 seconds for comparison.

So, I was wondering will it would be better to use the previous version for JOSS submission and we could update these with better GPU and CPU for the later version.