Ray-trax: Fast, Time-Dependent and Differentiable Ray Tracing for On-the-Fly Radiative Transfer in Turbulent Astrophysical Flows

Lorenzo Branca, Rune Rost, Tobias Buck

Keywords: Radiative Transfer, Numerics, Astrophysics, Differentiable Programming

TL;DR: A new GPUs accelerated, time dependent and differentiable Ray Tracing code to compute Radiative Transfer in turbulent astrophysical environment.

Abstract:

Radiative transfer is a key bottleneck in computational astrophysics: it is nonlocal, stiff, and tightly coupled to hydrodynamics.
We introduce $\textbf{Ray-trax}$, a GPU-oriented, fully differentiable 3D ray tracer written in JAX that solves the $\textit{time-dependent}$ emission--absorption problem and runs directly on turbulent gas fields produced by hydrodynamical simulations.
The method favors the widely used on-the-fly $\textit{absorption--emission approximation}$, which is state-of-the-art in many production hydro codes when scattering is subdominant.
Ray-trax vectorizes across rays and sources, supports arbitrarily many frequency bins without architectural changes, and exposes end-to-end gradients, making it straightforward to couple with differentiable hydro solvers while $\textit{preserving differentiability}$.
We validate against analytic solutions, demonstrate propagation in turbulent media, and perform a simple inverse problem via gradient-based optimization.
In practice the memory footprint scales as $\mathcal{O}(N_{\text{src}}\,N_{\text{cells}})$ while remaining highly efficient on accelerators.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 7