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

TL;DR

Ray-trax is a GPU-optimized, fully differentiable 3D ray tracer designed for efficient, time-dependent radiative transfer in turbulent astrophysical flows, enabling coupling with hydrodynamic simulations and inverse problem solving.

Contribution

It introduces Ray-trax, a novel differentiable ray tracing method that efficiently handles time-dependent radiative transfer directly on turbulent gas fields using JAX.

Findings

Validated against analytical solutions.

Demonstrated propagation in turbulent media.

Performed inverse optimization with gradients.

Abstract

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.