❯ jaxincell: 1D3V particle-in-cell code in JAX to perform simulation of plasmas

Table of Contents

• Overview
• Features
• Project Structure
  • Project Index
• Getting Started
  • Prerequisites
  • Installation
  • Usage
  • Testing
• Project Roadmap
• Contributing
• License
• Acknowledgments

---

Overview

JAX-in-Cell is an open-source project in Python that uses JAX to speedup simulations, leading to a simple to use, fast and concise code. It can be imported in a Python script using the jaxincell package, or run directly in the command line as jaxincell. To install it, use

pip install jaxincell

Alternatively, you can install the Python dependencies jax, jax_tqdm and matplotlib, and run an example script in the repository after downloading it as

git clone https://github.com/uwplasma/JAX-in-Cell
python examples/two-stream_instability.py

This allows JAX-in-Cell to be run without any installation.

An example without the use of an input toml file can be seen in the Weibel instability example

The project can be downloaded in its GitHub repository

---

Features

JAX-in-Cell can run in CPUs, GPUs and TPUs, has autodifferentiation and just-in-time compilation capabilities, is based on rigorous testing, uses CI/CD via GitHub actions and has detailed documentation.

Currently, it evolves particles using the non-relativisic Lorentz force $\mathbf F = q (\mathbf E + \mathbf v \times \mathbf B)$, and evolves the electric $\mathbf E$ and magnetic $\mathbf B$ field using Maxwell's equations.

Plenty of examples are provided in the examples folder, and the documentation can be found in Read the Docs.

---

Project Structure

└── JAX-in-Cell/
    ├── LICENSE
    ├── docs
    ├── examples
    │   ├── Landau_damping.py
    │   ├── Langmuir_wave.py
    │   ├── ion-acoustic_wave.py
	│   ├── two-stream_instability.py
	│   ├── auto-differentiability.py
	│   ├── scaling_energy_time.py
    │   └── optimize_two_stream_saturation.py
    ├── jaxincell
    │   ├── algorithms.py
    │   ├── boundary_conditions.py
    │   ├── constants.py
    │   ├── diagnostics.py
    │   ├── filters.py
    │   ├── fields.py
    │   ├── particles.py
    │   ├── plot.py
    │   ├── simulation.py
    │   └── sources.py
    ├── main.py
    └── tests
        └── test_simulation.py

---

Getting Started

Prerequisites

• Programming Language: Python

Besides Python, JAX-in-Cell has minimum requirements. These are stated in requirements.txt, and consist of the Python libraries jax, jax_tqdm and matplotlib.

Installation

Install JAX-in-Cell using one of the following methods:

Using PyPi:

1. Install JAX-in-Cell from anywhere in the terminal:

pip install jaxincell

Build from source:

1. Clone the JAX-in-Cell repository:

git clone https://github.com/uwplasma/JAX-in-Cell

2. Navigate to the project directory:

cd JAX-in-Cell

3. Install the project dependencies:

pip install -r /path/to/requirements.txt

4. Install JAX-in-Cell:

pip install -e .

Usage

To run a simple case of JAX-in-Cell, you can simply call jaxincell from the terminal

jaxincell

This runs JAX-in-Cell using standard input parameters of Landau damping. To change input parameters, use a TOML file similar to the example script present in the repository as

jaxincell examples/input.toml

Additionally, it can be run inside a script, as shown in the example script file

python examples/two-stream_instability.py

There, you can find most of the input parameters needed to run many test cases, as well as resolution parameters.

Parameters

JAX-in-Cell is highly configurable. Below is a list of the available parameters that can be defined in the TOML configuration file or the Python input dictionary.

Solver Parameters

These parameters control the numerical discretization, algorithm selection, and simulation resolution.

Click to expand full Solver Parameter Table

Parameter Key	Description
number_grid_points	Number of spatial grid cells.
total_steps	Total number of time steps to run.
number_pseudoelectrons	Total number of electron macroparticles.
number_pseudoparticles_species	List of particle counts for additional species.
Algorithms
time_evolution_algorithm	0: Explicit Boris pusher 1: Implicit Crank-Nicolson
field_solver	0: Electromagnetic (Curl_EB) 1: Electrostatic (Gauss FFT) 2: Electrostatic (Gauss Finite Diff) 3: Poisson (FFT)
Implicit Solver Settings
max_number_of_Picard_iterations_implicit_CN	Max Picard iterations per time step (only for algorithm 1).
number_of_particle_substeps_implicit_CN	Number of particle orbit substeps (only for algorithm 1).

Input Parameters

Click to expand full Parameter Table

Parameter Key	Description
length	Total length of the simulation box (meters).
grid_points_per_Debye_length	$\Delta x$ over Debye length.
timestep_over_spatialstep_times_c	CFL condition factor: $c \Delta t / \Delta x$.
Initialization
seed	Random seed for reproducibility.
random_positions_x	Randomize particle positions in x axis.
weight	Particle weight.
Species Properties
electron_charge_over_elementary_charge	Electron charge (normalized to $e$).
ion_charge_over_elementary_charge	Ion charge (normalized to $e$).
ion_mass_over_proton_mass	Ion mass (normalized to proton mass).
relativistic	Use relativistic Boris pusher.
vth_electrons_over_c_x,y,z	Electron thermal velocity in x,y,z.
ion_temperature_over_electron_temperature_x,y,z	Ratio $T_i/T_e$ in x,y,z.
electron_drift_speed_x	Electron drift speed (m/s) in X.
ion_drift_speed_x	Ion drift speed (m/s) in X.
velocity_plus_minus_electrons_x	Create counter-streaming electron populations in X.
Perturbations
amplitude_perturbation_x	Amplitude of density perturbation in X.
wavenumber_electrons_x	Mode number $k$ (factor of $2\pi/L$) for electrons in X.
External Fields
external_electric_field_amplitude	Amplitude of external E-field (cosine).
external_electric_field_wavenumber	Wavenumber of external E-field.
external_magnetic_field_amplitude	Amplitude of external B-field (cosine).
external_magnetic_field_wavenumber	Wavenumber of external B-field.
Boundary Conditions
particle_BC_left,right	Left,right particle boundary conditon (0: periodic, 1: reflective, 2: absorbing).
Numerics
filter_passes	Number of passes of the digital filter for charge/current density.
filter_alpha	Smoothing strength (0.0 to 1.0).
filter_strides	Multi-scale filtering strides.
tolerance_Picard_iterations_implicit_CN	Tolerance for implicit solver iterations.
print_info	Print simulation details to console.

Bump on Tail Domenstration

Here is a simple demonstration with electrons and ions moving in the x direction, including additional species with different weights (i.e., different real-particle number densities). A detailed list of parameters can be found in the bump-on-tail example.

Periodic Boundary Condition

bump_on_tail_explicit_periodic.mp4

Reflective Boundary Condition

bump_on_tail_explicit_reflective.mp4

Testing

Run the test suite using the following command:

pytest .

---

Project Roadmap

• Task 1: Run PIC simulation using several field solvers.
• Task 2: Finalize example scripts and their documentation.
• Task 3: Implement a relativistic equation of motion.
• Task 4: Implement collisions to allow the plasma to relax to a Maxwellian.
• Task 5: Implement guiding-center equations of motion.
• Task 6: Implement an implicit time-stepping algorithm.
• Task 7: Generalize JAX-in-Cell to 2D.

---

Contributing

• 💬 Join the Discussions: Share your insights, provide feedback, or ask questions.
• 🐛 Report Issues: Submit bugs found or log feature requests for the JAX-in-Cell project.
• 💡 Submit Pull Requests: Review open PRs, and submit your own PRs.

Contributing Guidelines

1. Fork the Repository: Start by forking the project repository to your github account.
2. Clone Locally: Clone the forked repository to your local machine using a git client.

   git clone https://github.com/uwplasma/JAX-in-Cell

3. Create a New Branch: Always work on a new branch, giving it a descriptive name.

   git checkout -b new-feature-x

4. Make Your Changes: Develop and test your changes locally.
5. Commit Your Changes: Commit with a clear message describing your updates.

   git commit -m 'Implemented new feature x.'

6. Push to github: Push the changes to your forked repository.

   git push origin new-feature-x

7. Submit a Pull Request: Create a PR against the original project repository. Clearly describe the changes and their motivations.
8. Review: Once your PR is reviewed and approved, it will be merged into the main branch. Congratulations on your contribution!

Contributor Graph

---

License

This project is protected under the MIT License. For more details, refer to the LICENSE file.

---

Acknowledgments

• This code was inspired by a previous implementation of a PIC code in JAX by Sean Lim here.
• We acknowledge the help of the whole UWPlasma plasma group.

---