Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Ma, Longyu, Jorge, Rogerio, Lu, Hongke, Tran, Aaron, Woolford, Christopher
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:https://arxiv.org/abs/2512.12160
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866912761879986176
author Ma, Longyu
Jorge, Rogerio
Lu, Hongke
Tran, Aaron
Woolford, Christopher
author_facet Ma, Longyu
Jorge, Rogerio
Lu, Hongke
Tran, Aaron
Woolford, Christopher
contents JAX-in-Cell is a fully electromagnetic, multispecies, and relativistic 1D3V Particle-in-Cell (PIC) framework implemented entirely in JAX. It provides a modern, Python-based alternative to traditional PIC frameworks. It leverages Just-In-Time compilation and automatic vectorization to achieve the performance of traditional compiled codes on CPUs, GPUs, and TPUs. The resulting framework bridges the gap between educational scripts and production codes, providing a testbed for differentiable physics and AI integration that enables end-to-end gradient-based optimization. The code solves the Vlasov-Maxwell system on a staggered Yee lattice with either periodic, reflective, or absorbing boundary conditions, allowing both an explicit Boris solver and an implicit Crank-Nicolson method via Picard iteration to ensure energy conservation. Here, we detail the numerical methods employed, validate against standard benchmarks, and showcase the use of its auto-differentiation capabilities.
format Preprint
id arxiv_https___arxiv_org_abs_2512_12160
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle JAX-in-Cell: A Differentiable Particle-in-Cell Code for Plasma Physics Applications
Ma, Longyu
Jorge, Rogerio
Lu, Hongke
Tran, Aaron
Woolford, Christopher
Plasma Physics
JAX-in-Cell is a fully electromagnetic, multispecies, and relativistic 1D3V Particle-in-Cell (PIC) framework implemented entirely in JAX. It provides a modern, Python-based alternative to traditional PIC frameworks. It leverages Just-In-Time compilation and automatic vectorization to achieve the performance of traditional compiled codes on CPUs, GPUs, and TPUs. The resulting framework bridges the gap between educational scripts and production codes, providing a testbed for differentiable physics and AI integration that enables end-to-end gradient-based optimization. The code solves the Vlasov-Maxwell system on a staggered Yee lattice with either periodic, reflective, or absorbing boundary conditions, allowing both an explicit Boris solver and an implicit Crank-Nicolson method via Picard iteration to ensure energy conservation. Here, we detail the numerical methods employed, validate against standard benchmarks, and showcase the use of its auto-differentiation capabilities.
title JAX-in-Cell: A Differentiable Particle-in-Cell Code for Plasma Physics Applications
topic Plasma Physics
url https://arxiv.org/abs/2512.12160