---
title: 'Ray-trax: Differentiable Ray Tracing Method'
url: https://www.emergentmind.com/topics/ray-trax
type: topic
---

# Ray-trax: Differentiable Ray Tracing Method

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 [2511.09389].

## 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 [2511.09389].

## 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 [2511.09389].

## 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 [2511.09389].

## 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 [2511.09389].

## 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 [2511.09389].

## 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 [2511.09389].

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" [2503.20907], 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 [2511.09389].

---

**References**:  
• "Ray-trax: Fast, Time-Dependent, and Differentiable Ray Tracing for On-the-fly Radiative Transfer in Turbulent Astrophysical Flows" [2511.09389]  
• "Generalized Ray Tracing with Basis Functions for Tomographic Projections" [2503.20907]

Source: https://www.emergentmind.com/topics/ray-trax