https://github.com/proximafusion/raytrax
Raytrax: ECRH ray tracing in JAX
https://github.com/proximafusion/raytrax
diffrax fusion jax plasma-physics python pyvista
Last synced: about 2 months ago
JSON representation
Raytrax: ECRH ray tracing in JAX
- Host: GitHub
- URL: https://github.com/proximafusion/raytrax
- Owner: proximafusion
- License: mit
- Created: 2025-07-16T11:27:43.000Z (12 months ago)
- Default Branch: main
- Last Pushed: 2026-05-11T08:19:51.000Z (about 2 months ago)
- Last Synced: 2026-05-11T09:37:48.611Z (about 2 months ago)
- Topics: diffrax, fusion, jax, plasma-physics, python, pyvista
- Language: Jupyter Notebook
- Homepage: https://proximafusion.github.io/raytrax
- Size: 5.25 MB
- Stars: 13
- Watchers: 1
- Forks: 5
- Open Issues: 5
-
Metadata Files:
- Readme: README.md
- License: LICENSE.md
- Citation: CITATION.cff
- Codeowners: CODEOWNERS
Awesome Lists containing this project
- awesome-ml-in-plasma-physics - Raytrax - ECRH ray tracing for fusion plasmas in JAX with JIT compilation and automatic differentiation (Tools / Simulation and Modeling Frameworks)
README
# Raytrax
**An ECRH ray tracer for fusion plasmas, built on JAX.**
[](https://github.com/proximafusion/raytrax/actions/workflows/test.yml)
[](https://opensource.org/licenses/MIT)
[](https://www.python.org/)
[](https://proximafusion.github.io/raytrax/)
[](https://doi.org/10.5281/zenodo.18924104)
Raytrax simulates Electron Cyclotron Resonance Heating (ECRH) of magnetic confinement fusion plasmas. Powered by [JAX](https://docs.jax.dev), it features JIT-compiled ray tracing and automatic differentiation — making it well-suited for gradient-based beam optimization in fusion plant design.
> **Note:** Raytrax is in early development. Expect API changes and incomplete validation. For details [see the docs](https://proximafusion.github.io/raytrax/limitations).
## Usage
Load a magnetic equilibrium, define plasma profiles, and trace a beam:
```python
import jax.numpy as jnp
import raytrax
import vmecpp
vmec_wout = vmecpp.VmecWOut.from_wout_file("w7x.nc")
mag_conf = raytrax.MagneticConfiguration.from_vmec_wout(vmec_wout)
rho = jnp.linspace(0, 1, 200)
profiles = raytrax.RadialProfiles(
rho=rho,
electron_density=2.0 * (1 - rho**2),
electron_temperature=3.0 * (1 - rho**2),
)
beam = raytrax.Beam(
position=jnp.array([6.6, 0.0, 0.0]),
direction=jnp.array([-0.985, 0.0, -0.174]),
frequency=140e9,
mode="O",
power=1e6,
)
result = raytrax.trace(mag_conf, profiles, beam)
print(f"Optical depth τ = {result.optical_depth:.3f}")
print(f"Absorbed fraction = {result.absorbed_power_fraction:.1%}")
print(f"Deposition at ρ = {result.deposition_rho_mean:.2f} ± {result.deposition_rho_std:.2f}")
```
## Visualizations
### Beam trace on a poloidal cross-section
### 3D visualization
Raytrax can render flux surfaces and beam trajectories in 3D using [PyVista](https://pyvista.org):
```python
import pyvista as pv
from raytrax.plot.plot3d import plot_flux_surface_3d, plot_beam_profile_3d
plotter = pv.Plotter()
plot_flux_surface_3d(mag_conf, rho_value=1.0, plotter=plotter, opacity=0.25)
plot_beam_profile_3d(result.beam_profile, plotter=plotter, tube_radius=0.02)
plotter.export_html("scene.html") # interactive standalone HTML
```
> Try out a **[live interactive version](https://proximafusion.github.io/raytrax/3d-visualisation/)** in the docs.
## Jupyter Notebooks
Quick-start notebooks are in the [`notebooks/`](notebooks/) directory.
## Installation
```bash
python -m pip install raytrax
```
See the [documentation](https://proximafusion.github.io/raytrax/) for a full getting-started guide, theory background, and API reference.
## Citing Raytrax
If you use Raytrax in your research, please cite it via Zenodo: [doi:10.5281/zenodo.18924104](https://doi.org/10.5281/zenodo.18924104).
## Acknowledgements
The development of Raytrax is a collaboration between [Proxima Fusion](https://www.proximafusion.com) and the [Munich University of Applied Sciences (HM)](https://www.hm.edu) and was partially supported by the German Federal Ministry of Research, Technology and Space (BMFTR) under grant FPP-MC (13F1001B).
## License
Raytrax is released under the MIT License. See [LICENSE.md](LICENSE.md) for details.