adept is a differentiable two-fluid code that was used to show how to incorporate kinetic effects into fluid codes using neural networks. It uses JAX and equinox to solve the plasma two-fluid equations in 1D as a proof of concept.
The manuscript that introduces and uses adept is
[arXiv] [IOP - Machine Learning: Science and Technology]
The code is [here]
tsadar performs parameter estimation of Thomson scattering spectra measured at the [Omega Laser Facility] via AD-powered gradient descent using JAX.
The code is [here]
plasmadisp calculates the complex roots to the electrostatic dispersion relation given by
$$ 1 + \frac{\omega_p^2}{k^2} \int du \frac{d_u g(u)}{\omega/k - u} = 0 $$
where g(u) is the 1D distribution function, here assumed to be a Maxwell-Boltzmann.