Bump-on-tail problem; arbitrary additional ion/electron populations; separate diagnostics(...) and simulation(...)

[aarontran]

First draft, review/comments welcome.

New features

• Bump-on-tail problem creates 4 wavelengths across box, clearly visible in e- phase space with 40,000 pseudoelectrons/population (i.e., 40k for bulk maxwellian, 40k for weak drifting maxwellian).

• Code edits allow user to specify arbitrary number of drifting Maxwellian populations via TOML array of tables. Usage is shown by [[species]] blocks within examples/bump-on-tail.toml.

• The user must call the diagnostics(...) function outside of simulation(...) to get value-added output fields that separate the ions and electrons, because (i) we cannot rely on static array index number_pseudoelectrons to distinguish ions/electrons within particle arrays, and (ii) I did not see a quick/easy way to sift particle species into separate arrays within a JAX jitted function, so the post-processing is split into its own non-jitted subroutine.

I have not verified the linear growth rate in E^2 energy for the bump-in-tail, for it will require more particles / testing to beat down E-field noise. Current problem takes ~ few minutes to run on my personal Mac laptop. I expect we will obtain multiple decades of clean E^2 linear growth with ~100x number of total number of particles (~100k+ particles/cell \approx 4 million total particles). Maybe one could try to estimate from the relaxation/phase-space mixing rate of the electrons' velocity distribution.

Backwards compatibility:

Two stream + weibel examples still run, but other example problems were not checked for compatibility with the new changes.

Possible improvements:

• dump particle weights in output, so that output distributions can be plotted correctly (in the below plot, it looks like the beam population has equal density as bulk, but the beam particle weight is ~3% that of the bulk population)
• tag species/population ID for each particle, or output separate files for each population

NOTE

Below plot shows a simulation wherein the bump-on-tail.toml file was manually edited to not initialize charge-neutralizing ions for the bump distribution, because the net charge does not matter for this particular problem, ions don't participate, the test runs faster w/o the ions. However, @luisHongkeLu 's simulations show that my edit may have messed up the display of ions vs. electrons -- the ion VDF plot should NOT show phase-space holes / multiple populations. The particle output code + plotting code may need to be fixed in the case where n_ion != n_electron. --ATr,2025sep16

[Copilot]

Pull Request Overview

This PR implements support for bump-on-tail instability simulations by adding arbitrary multiple ion/electron populations and refactoring the simulation output processing. The changes enable users to specify multiple drifting Maxwellian populations via TOML configuration files and separate the diagnostics computation from the main simulation function.

• Implements bump-on-tail problem with 4 wavelengths across the simulation box
• Adds support for arbitrary number of particle species via TOML array of tables configuration
• Separates diagnostics() function from simulation() to handle ion/electron segregation outside of JAX-jitted code

Reviewed Changes

Copilot reviewed 12 out of 20 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
jaxincell/_simulation.py	Core implementation adding multi-species support, new make_particles() function, and modified simulation output structure
jaxincell/_diagnostics.py	Updates diagnostics to segregate ions/electrons based on charge sign instead of array indices
jaxincell/_plot.py	Minor plotting improvements with symmetric colormap bounds and debug output
examples/*.py	Updates all example scripts to call diagnostics() externally after simulation
examples/bump-on-tail.toml	New configuration file demonstrating multi-species setup for bump-on-tail instability
examples/bump-on-tail.py	New example script implementing the bump-on-tail problem

[Copilot]

The assertion comment mentions 'brittle code' and hard-coded offsets, but the assertion only checks if rng_index >= 0. Consider adding proper validation for the upper bound or documenting the expected range more clearly.

Suggested change

	assert rng_index >= 0
	MAX_RNG_INDEX = 15 # Set this to the maximum number of populations expected
	assert 0 <= rng_index <= MAX_RNG_INDEX, f"rng_index ({rng_index}) must be in [0, {MAX_RNG_INDEX}]"

[aarontran]

Disagree with Copilot's suggestion. This PR doesn't limit the number of species/populations, so it doesn't make sense to hard-code an upper bound. The brittleness here is tied to hardcoding of RNG offsets in PRNGKey(seed+1), PRNGKey(seed+2), etc... like in lines 194–227 of _simulation.py.

To address the code brittleness, one could...

• use global state for RNG stream, so there is no need for hard-coded seed offsets,
• make data structure to enumerate distinct species, so that it's clear what block of RNG seeds is used for what species... (then safety checks could be inserted, to guard against user-customized init with arbitrary RNG calls colliding with existing RNG streams). But that's probably over-engineering at this point.

[luisHongkeLu]

I verified the new bump-on-tail example runs and the existing examples still work—nice! One suggestion: the current diagnostics(output) splits particles only by charge sign (ions vs. electrons). With the new [[species]] support, it would be great if diagnostics could also report per-species summaries.

My suggestion is to add a species_id: int32 during initialization in _simulation.py then in _diagonistics.py we can directly use species_id for the output.

[aarontran]

Thanks for the review & suggestions!

Species ID is added: default "electron" ID = 0, default "ion" ID = 1, user-added species ID=2, 3, 4, etc... in order of TOML file.

I made some fixes to diagnostic outputs to correctly handle variable particle weight. One possibly-confusing point: the same physical species (electron) can have pseudo-particles with different charges and masses, because of how weights are implemented internally.

--The plotting routine is also modified to allow plotting histograms for specific species only; example in bump-on-tail.py-- (edit: reverted this change to work with Rogerio's jittable version of diagnostics function)