Ray-trax: Differentiable Ray Tracing Method

• Ray-trax is a differentiable ray tracing framework that simulates time-dependent emission–absorption radiative transfer using vectorized GPU acceleration.
• It employs operator splitting and on-the-fly approximations to efficiently compute radiative effects in turbulent, inhomogeneous media.
• Integrating with JAX, Ray-trax enables gradient-based optimization and inverse problems in astrophysical simulations with scalable performance.

Ray-trax refers to a class of ray-tracing methodologies and software frameworks targeting physically rigorous, high-performance simulation of radiative and electromagnetic transport in complex, often turbulent or inhomogeneous, environments. Originally developed for computational astrophysics, the Ray-trax formalism unifies fast, fully differentiable GPU-accelerated ray tracing with the time-dependent emission–absorption radiative transfer equation, while being composable in modern autodifferentiation frameworks (JAX). Ray-trax algorithms emphasize on-the-fly, vectorized integration of radiative effects in turbulent media and enable integration of gradient-based optimization, rendering them optimal for differentiable programming paradigms and inverse problems in astrophysics and beyond (Branca et al., 12 Nov 2025).

1. Time-Dependent Emission–Absorption Radiative Transfer Formalism

Ray-trax centers on the time-dependent, monochromatic emission–absorption equation for the specific intensity $I(\mathbf{x}, \hat{\mathbf{n}}, t)$,

$$\frac{1}{c}\frac{\partial I}{\partial t} + \hat{\mathbf{n}} \cdot \nabla I = -\kappa(\mathbf{x})I + j(\mathbf{x}),$$

where $c$ is the (effective) radiation speed, $\kappa(\mathbf{x})$ the opacity (absorption coefficient), and $j(\mathbf{x})$ the local emissivity. Ray-trax implements a semi-analytic discrete-ray marching algorithm with operator splitting in transport, absorption, and emission terms. Given a spatial step $\Delta s$ and substep index $k$, the update is

$$\begin{aligned} \mathbf{x}_{k+1} &= \mathbf{x}_k + \hat{\mathbf{n}}\Delta s,\ \tau_{k+1} &= \tau_k + \kappa(\mathbf{x}_k)\Delta s,\ I_{k+1} &= I_k\exp[-\kappa(\mathbf{x}_k)\Delta s] + j(\mathbf{x}_k)\exp[-\tau_k]\Delta s. \end{aligned}$$

The integration is performed in operator-split substeps, $N_s = \lceil c\Delta t / \Delta s \rceil$, with exact attenuation and first-order emission quadrature per segment (Branca et al., 12 Nov 2025).

2. On-the-Fly Emission–Absorption Approximation

Ray-trax follows the "on-the-fly emission–absorption" approximation used in production radiative hydrodynamics codes, such as RAMSES-RT and AREPO-RT, when scattering is isotropic. Instead of explicitly accounting for angular diffusion or a nonlocal scattering kernel, isotropic scattering is incorporated into the local absorption $\kappa$ and isotropic re-emission $j$. This approximation is justified when angular redistribution is dominated by local processes and only photo-heating and attenuation are required, making the method tractable for integration within hydrodynamics codes (Branca et al., 12 Nov 2025).

3. GPU-Oriented, JAX-Based Vectorization and Differentiation

Ray-trax is architected for accelerated, batch array programming. All ray and source loops are vectorized using JAX's vmap, enabling dense batching of rays, sources, and frequency bins as additional control-batch dimensions, without architectural or code changes. The underlying ray-marching loop is implemented as a single JIT-compiled lax.fori_loop for maximal device efficiency.

Spherical directions $\{\hat{\mathbf{n}}_m\}$ are sampled via a Fibonacci (golden-angle) sphere lattice, ensuring uniform solid-angle coverage and trivial vector slicing. The approach is fully agnostic to the number of frequency bins due to treating frequency as a batch axis. Multi-device (GPU) sharding is employed by splitting the set of rays ($N_\Omega$) across available accelerators, yielding parallel scalability with linear scaling in per-device memory and time (Branca et al., 12 Nov 2025).

4. End-to-End Differentiability and Inverse Problems

All Ray-trax primitives (e.g., trilinear interpolation for $j$, nearest-neighbor lookup for $\kappa$, grid deposition, lax.fori_loop) possess well-defined Jacobians in JAX, supporting forward-mode and reverse-mode automatic differentiation (jacfwd, jacrev). This facilitates gradient-based optimization of observables with respect to source properties, opacity parameters, or any upstream hydrodynamical parameters.

For example, with $j(\mathbf{x}) = A e_0(\mathbf{x})$ and mean-squared-error loss

$$\mathcal{L}(A) = \frac{1}{2}\sum_{\mathbf{x}}[I(\mathbf{x};A)-I_{\rm ref}(\mathbf{x})]^2,$$

the gradient is computed as

$$\frac{\partial\mathcal{L}}{\partial A} = \sum_{\mathbf{x}} [I(\mathbf{x};A)-I_{\rm ref}(\mathbf{x})] \frac{\partial I(\mathbf{x};A)}{\partial A},$$

compatible with any JAX-native optimizer (e.g., Optax). This property enables end-to-end coupling with differentiable hydro solvers (e.g., jf1uids), supporting full backpropagation from emission to dynamics (Branca et al., 12 Nov 2025).

5. Complexity, Memory Scaling, and Benchmarks

The work per time-step scales as

$$\mathcal{O}(N_{\rm src} \times N_\Omega \times N_s),$$

where $N_{\rm src}$ is the number of sources, $N_\Omega$ the ray count (directions), and $N_s$ the ray-marching substeps. If per-source, per-ray contributions are retained, transient memory scales as $\mathcal{O}(N_\Omega \times N_{\rm cells})$ or $\mathcal{O}(N_{\rm src} \times N_{\rm cells})$ if storing grids per source. Practical device sharding ensures that memory per device is

$$\mathcal{O}\left(\frac{N_\Omega}{N_{\rm dev}} \times N_{\rm cells}\right).$$

Benchmarks show that full ray-tracing from 2,097 sources, $4,096$ rays, and $128^3$ grid cells completes in $4.49$ s (single NVIDIA H200), with mean relative error $\lesssim 3.2\%$ against the analytic $1/r^2$-attenuated solution for constant opacity; turbulent media propagation is also demonstrated (Branca et al., 12 Nov 2025).

6. Representative Applications

Ray-trax is optimized for three principal astrophysical applications:

• Gradient-based inversion of source luminosity or opacity scaling, using one or more frequency bins from observations.
• Differentiable coupling to grid-based hydrodynamics (e.g., jf1uids), allowing errors from radiative observables to be propagated through the radiative transfer and hydro time-integration to subgrid model parameters.
• Production galaxy/star-formation prediction, enabling on-the-fly radiative substeps in turbulence-resolving simulation pipelines where photoionization/RTE is a bottleneck (Branca et al., 12 Nov 2025).

Ray-trax can also be integrated with higher-order spline/basis-function projection ray-tracing, as described in "Generalized Ray Tracing with Basis functions for Tomographic Projections" (Haouchat et al., 26 Mar 2025), where box-splines of degree 2 are recommended for optimal PSNR/SSIM and efficient GPU implementation.

7. Implementation, Validation, and Code Availability

The Ray-trax codebase (https://github.com/lorenzobranca/Ray-trax.git) emphasizes minimality and composability for JAX-based pipelines. Uniform direction sampling, vectorized batch control, and reproducible performance benchmarks are provided. Validation against analytical solutions (uniform media, point-source attenuation), turbulent astrophysical media, and code performance metrics are explicitly documented. The code base serves as a building block for any modern, differentiable astrophysical simulation seeking to couple radiative transfer tightly and efficiently into GPU-accelerated, autodifferentiable pipelines (Branca et al., 12 Nov 2025).

---

References:

• "Ray-trax: Fast, Time-Dependent, and Differentiable Ray Tracing for On-the-fly Radiative Transfer in Turbulent Astrophysical Flows" (Branca et al., 12 Nov 2025)
• "Generalized Ray Tracing with Basis Functions for Tomographic Projections" (Haouchat et al., 26 Mar 2025)