2025/11/12 by Branca, Lorenzo, Rost, Rune, Buck, Tobias
#Astrophysics of Galaxies (astro-ph.GA) #Computational Engineering #FOS: Computer and information sciences #FOS: Physical sciences #Finance #Instrumentation and Methods for Astrophysics (astro-ph.IM) #Solar and Stellar Astrophysics (astro-ph.SR) #and Science (cs.CE)
paper · doi:10.48550/arxiv.2511.09389
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 O(Nsrc Ncells) while remaining highly efficient on accelerators.