Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

TL;DR

Problem: Radiative transfer in astrophysical simulations is a bottleneck due to nonlocality, stiffness, and coupling to hydrodynamics. Approach: Ray-trax provides a time-dependent, differentiable, GPU-accelerated 3D ray tracer in JAX with a finite light-travel horizon, frequency-bin batching, and end-to-end gradients. Contributions: Analytical validation, turbulence tests, and gradient-based inversion are demonstrated; memory scales as $O(N_{ ext{src}} N_{ ext{cells}})$, and the implementation achieves high throughput via full vectorization and sharding. Significance: This compact RT building block enables on-the-fly RT integration with differentiable hydro, supporting physics-constrained learning and subgrid calibration in large-scale astrophysical simulations.

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.

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

TL;DR

Problem: Radiative transfer in astrophysical simulations is a bottleneck due to nonlocality, stiffness, and coupling to hydrodynamics. Approach: Ray-trax provides a time-dependent, differentiable, GPU-accelerated 3D ray tracer in JAX with a finite light-travel horizon, frequency-bin batching, and end-to-end gradients. Contributions: Analytical validation, turbulence tests, and gradient-based inversion are demonstrated; memory scales as , and the implementation achieves high throughput via full vectorization and sharding. Significance: This compact RT building block enables on-the-fly RT integration with differentiable hydro, supporting physics-constrained learning and subgrid calibration in large-scale astrophysical simulations.

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 while remaining highly efficient on accelerators.

Paper Structure

This paper contains 22 sections, 3 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Method and code structure
  3. Vectorization, memory, and speed.
  4. Frequency-bin agnosticism.
  5. Benchmarks in uniform media and turbulent fields
  6. Analytical validation.
  7. Turbulent snapshots from hydro.
  8. On-the-fly suitability.
  9. Differentiability: inversion and coupling
  10. Gradient-based inversion.
  11. Coupling to differentiable hydro.
  12. Discussion and conclusions
  13. Acknowledgments
  14. Notation
  15. Geometry and coordinates.
  16. ...and 7 more sections

Figures (3)

  • Figure 1: Analytical benchmark in a uniform medium with multiple point sources. Top: Specific intensity $I$ along a line of sight through the domain center, shown for both $\pm$ directions; the sightline intersects a second source in one of the directions. Bottom: Diagonal slice of $I$ across the volume (left to right: numerical solution, analytical reference, and $\log_{10}$ residuals).
  • Figure 2: Time-resolved propagation (from top left to bottom right) of specific intensity $I$ through a turbulent opacity field $\kappa(\mathbf{x})$ extracted from a hydrodynamic snapshot. The frequency bin is chosen to highlight shadowing. Time $t$ is in unit of total simulation time.
  • Figure 3: End-to-end differentiability enables straightforward parameter recovery.