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

Radiative transfer is a key bottleneck in computational astrophysics: it is nonlocal, stiff, and tightly coupled to hydrodynamics. We introduce Ray-trax, a GPU-oriented, fully differentiable 3D ray tracer written in JAX that solves the time-dependent emission–absorption problem and runs directly on turbulent gas fields produced by hydrodynamic simulations. The method favors the widely used on-the-fly emission–absorption approximation, which is state of the art in many production hydro codes when scattering is isotropic. 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 preserving differentiability. We validate against analytical 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.