RadJAX: Differentiable Radiative Transfer Engine

• RadJAX is a JAX-based LTE molecular line radiative transfer solver that transforms 3D disk models into synthetic ALMA spectral cubes.
• It integrates physics-constrained neural fields with embedded symmetry and hydrostatic constraints to derive temperature, density, and CO abundance profiles.
• Its GPU vectorization and automatic differentiation enable 3–4 orders of magnitude speedup over traditional CPU solvers for iterative data fitting.

Searching arXiv for the primary RadJAX paper and closely related JAX-based scientific solver papers for contextual support. Primary paper found: "Revealing Fine Structure in Protoplanetary Disks with Physics Constrained Neural Fields" (Levis et al., 3 Sep 2025). Also checking adjacent JAX-based differentiable radiative transfer and scientific-solver papers for context. RadJAX is a differentiable, GPU-accelerated molecular line radiative transfer engine introduced as the rendering core of a broader computational imaging framework for protoplanetary disks. In that framework, physics-constrained neural fields specify three-dimensional disk structure, RadJAX renders synthetic ALMA spectral-channel images from temperature, density, CO abundance, and velocity fields, and gradient-based optimization or MCMC is used to fit observed data cubes. The system was introduced in the context of ALMA CO observations of HD 163296, where it enabled reconstructions that recovered the vertical morphology of the CO-rich layer and revealed a pronounced narrowing and flattening of the emission surface beyond about $400$ au (Levis et al., 3 Sep 2025).

1. Definition and computational role

RadJAX is a lightweight LTE line radiative transfer solver written in JAX, designed for inverse problems rather than general-purpose astrophysical simulation. Its central function is to transform a three-dimensional physical disk model into a synthetic spectral image cube while remaining fully differentiable end-to-end. This differentiability allows gradients to propagate not only to a small set of parametric disk parameters but also to the weights of a high-dimensional neural field, which is the principal reason the method is useful for reconstruction rather than merely forward modeling (Levis et al., 3 Sep 2025).

The broader framework in which RadJAX appears has three layers. A physics-constrained neural field represents the disk temperature structure; derived physical rules then map that temperature field to density, CO abundance, and velocity; and RadJAX serves as the forward model that ray-traces the corresponding line emission. The rendered cube is compared to ALMA channel maps with a pixel-wise $\chi^2$-type objective during neural optimization, while the same renderer is also used inside MCMC for simpler parametric models. A persistent source of confusion is the tendency to identify the entire reconstruction system with RadJAX itself. More precisely, RadJAX is the renderer; the full framework also includes the neural representation, the embedded physical constraints, and the optimization or inference machinery (Levis et al., 3 Sep 2025).

This specialization defines both its capabilities and its tradeoffs. RadJAX sacrifices some of the flexibility of full-featured CPU solvers such as RADMC-3D—specifically adaptive mesh refinement, dust scattering, and multiple integration schemes—in exchange for GPU vectorization, automatic differentiation, and high-throughput repeated forward evaluation. This suggests that RadJAX is best understood as a solver optimized for iterative data fitting, not as a comprehensive replacement for general radiative transfer packages (Levis et al., 3 Sep 2025).

2. Integration with physics-constrained neural fields

In the nonparametric reconstruction mode, the learned quantity is the temperature field

$T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$

where $\theta$ are neural-network weights and $\gamma(r,|z|)$ is a sinusoidal positional encoding. The network uses positional encoding of degree $L=4$, has 4 fully connected layers with 64 hidden units each, uses ReLU activations, and applies a sigmoid output to constrain temperature to $[5,400]$ K. The field is continuous in space and explicitly constrained by azimuthal and midplane symmetry (Levis et al., 3 Sep 2025).

Temperature is the only learned field directly. The remaining quantities required by RadJAX are derived through embedded disk physics. Hydrostatic equilibrium determines $\rho(r,z)$ from $T(r,z)$; CO abundance is set to zero in freeze-out or photodissociated regions; and the velocity field follows height-dependent Keplerian rotation with pressure-gradient corrections. The surface-density normalization is imposed through

$$\Sigma(r) = \Sigma_0 \left(\frac{r}{r_c}\right)^{-\gamma} \exp\left[-\left(\frac{r}{r_c}\right)^{2 - \gamma}\right].$$

CO abundance is defined through a canonical abundance where gas-phase CO survives, modulated by freeze-out and photodissociation, with the vertical H$\chi^2$0 column above $\chi^2$1 entering through

$\chi^2$2

These derived fields—temperature, density, CO abundance, and velocity—constitute the physical input state that RadJAX ray-traces (Levis et al., 3 Sep 2025).

A key methodological point is that the neural field is flexible but not unconstrained. The main regularization is structural rather than penalty-based: network architecture, azimuthal and mirror symmetry, hydrostatic closure, CO survival rules, and Keplerian-like kinematics collectively define the admissible model class. A plausible implication is that RadJAX’s effectiveness depends as much on this upstream inductive bias as on the renderer itself, because the differentiable forward model alone does not determine identifiability in weakly constrained regions (Levis et al., 3 Sep 2025).

3. Radiative transfer formulation and rendering architecture

RadJAX implements standard LTE molecular line transfer along rays. The specific intensity at frequency $\chi^2$3 is written as

$\chi^2$4

with optical depth

$\chi^2$5

where $\chi^2$6 is emissivity and $\chi^2$7 is absorption or extinction. The usual background term $\chi^2$8 is omitted because ALMA calibration typically removes it, and volumetric continuum emission and absorption are neglected for the application considered (Levis et al., 3 Sep 2025).

At each point $\chi^2$9, the line emissivity and absorption coefficients are

$T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$0

$T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$1

with Einstein coefficients $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$2, LTE Boltzmann level populations, and a Gaussian line profile shifted by the local line-of-sight velocity. The source function is implied in the standard way as $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$3 where defined. The line width combines thermal and microturbulent broadening in quadrature; the supplied text notes that the typesetting of the explicit width formula is slightly corrupted but that the intended form is clearly thermal broadening plus turbulent broadening in quadrature (Levis et al., 3 Sep 2025).

The renderer uses backward ray tracing. Rays originate at image-plane pixels, are traced into the volume, and accumulate absorption and local emissivity along the path. The implementation operates on regularly spaced or log-spaced Cartesian or spherical grids rather than AMR grids. This design choice is central to its accelerator performance because it permits vectorized array programs over rays, channels, and grid cells. The paper does not present a novel transfer equation; its technical contribution lies in a differentiable, massively parallel implementation of standard LTE line transfer on regular or log-spaced grids (Levis et al., 3 Sep 2025).

4. Differentiability, optimization, and Bayesian inference

The practical significance of RadJAX is that every major stage of the forward model is implemented in JAX-compatible primitives and therefore participates in automatic differentiation: positional encoding, neural temperature evaluation, hydrostatic-equilibrium-derived density computation, CO masking by freeze-out and photodissociation rules, velocity calculation, ray integration, image synthesis, and loss evaluation. This permits end-to-end gradient propagation from an image-domain objective back to neural-field weights and, in principle, to any differentiable physical parameter upstream of the renderer (Levis et al., 3 Sep 2025).

For neural reconstruction, the rendered ALMA cube is compared to observations with a pixel-wise $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$4-type loss over channel mini-batches. Training uses ADAM with learning rate starting at $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$5 and decaying polynomially to $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$6, for 30,000 steps; each step uses a random batch of 50 slices from a 115-channel cube, with velocity sampling of about 40 m/s and training noise level of about 1.5% estimated from background pixels. The paper does not introduce explicit penalties such as total variation or weight decay; regularization is instead induced by architecture, symmetry, and embedded physical constraints (Levis et al., 3 Sep 2025).

For the 8-parameter parametric disk model, RadJAX is used inside MCMC rather than in backpropagation-based training. The setup uses the emcee ensemble sampler with 60 walkers, 100 warm-up iterations, re-centering around the chain median, 5000 additional iterations, and discarding the first 75% as burn-in. The stated computational consequence is that MCMC-based Bayesian inference for ALMA disk models becomes practical in a few hours on a single GPU rather than months on CPUs. This suggests that RadJAX’s differentiability is only part of its importance; the equally significant feature is the speed with which repeated forward evaluations become feasible (Levis et al., 3 Sep 2025).

5. Performance, validation, and empirical behavior

RadJAX’s acceleration is attributed to a combination of JAX implementation, just-in-time compilation, GPU-native vectorization, regular or log-spaced grids, specialization to LTE inverse problems, and automatic differentiation obtained directly from JAX rather than from custom adjoint derivations. The paper reports 3–4 orders of magnitude speedups, with claims of up to $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$7 relative to conventional CPU ray tracers. In one validation example, a spectral cube render that took about 20 minutes with RADMC-3D was reduced to about 100 ms with RadJAX. The paper also states that RadJAX maintains sub-second rendering even for ALMA-scale cubes where RADMC-3D may take up to an hour, and that speedups exceed $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$8 at high spectral resolution (Levis et al., 3 Sep 2025).

The benchmark configuration used a hydrodynamical simulation volume on a $T(r, |z|) = \mathrm{MLP}_{\theta}(\gamma(r, |z|)),$9 spherical grid, with RadJAX run on an NVIDIA A6000 GPU and RADMC-3D run on an Intel Xeon Gold 6348 CPU. Benchmarks were performed with fixed spatial resolution $\theta$0 while varying channel count, and with fixed 20-channel cubes while varying spatial resolution. The paper notes that at low spectral resolution, RADMC-3D runtime is partly bottlenecked by non-ray-tracing computations, which limits the interpretation of speedup factors in that regime (Levis et al., 3 Sep 2025).

Rendering accuracy was validated against RADMC-3D on hydrodynamical simulation outputs, with rendered pixel intensities and morphology reported to agree closely while maintaining roughly $\theta$1 acceleration. On the scientific model comparison, the neural field plus RadJAX reduced average pixel-averaged $\theta$2 by about 19% relative to the parametric model. Robustness was tested by training on only red-shifted or only blue-shifted channels and evaluating on the complementary half; under these withheld-data experiments, the neural model continued to outperform the parametric baseline, and the recovered thick-to-thin outer-disk transition persisted (Levis et al., 3 Sep 2025).

6. Scientific application, interpretive significance, and limitations

The principal application in (Levis et al., 3 Sep 2025) is ALMA $\theta$3 observations of the protoplanetary disk HD 163296, including a legacy dataset previously modeled by Flaherty et al. and a higher-resolution MAPS dataset. Within this application, the framework reconstructs the temperature field $\theta$4, derived H$\theta$5 density, CO-emitting layer morphology, and pressure-corrected Keplerian velocity field. The main scientific result is that the CO-emitting layer appears broadly flared or convex inside $\theta$6 au, then narrows and flattens outside $\theta$7 au, forming a nearly constant-height plateau in the outer disk (Levis et al., 3 Sep 2025).

This outer-disk flattening is presented as a feature missed by existing parametric approaches. The parametric model produces a “stair-case” separation artifact in image space, whereas the neural reconstruction recovers a smoother continuous emitting surface that matches ALMA channel maps more faithfully. Comparison with geometric surface extraction indicates broad agreement inside $\theta$8 au, but beyond that radius the geometric method shows a steep dip while the neural radiative transfer model yields a plateau; the paper suggests that geometric extraction becomes unreliable where the outer-disk emission becomes optically thin or near and far ridges merge (Levis et al., 3 Sep 2025).

The assumptions delimiting this result are explicit. RadJAX is LTE-only; it has no AMR, no dust scattering, no volumetric continuum transfer in the forward model, and it is limited to regular or log-spaced grids. The broader reconstruction additionally assumes azimuthal and mirror symmetry, CO abundance determined by threshold rules for freeze-out and photodissociation, and Keplerian-like velocity with pressure corrections. Weakly constrained regions can produce artifacts; the paper explicitly mentions isolated lower-surface CO “islands,” likely arising from underconstrained regions and the adopted CO-temperature thresholding in optically thick parts of the disk. For MCMC, the data model operates on CLEANed images rather than visibilities, so pixel correlations from the synthesized beam remain; for the higher-resolution MAPS data this is mitigated with pixel subsampling (Levis et al., 3 Sep 2025).

In a wider methodological context, RadJAX belongs to a growing class of JAX-native scientific solvers that emphasize accelerator execution and differentiability. Ray-trax is a GPU-oriented, fully differentiable 3D ray tracer for time-dependent emission–absorption transport on regular grids (Branca et al., 12 Nov 2025), JAXtronomy is a JAX reimplementation of lenstronomy built around JIT compilation and automatic differentiation (Huang et al., 1 Jun 2026), and gyaradax is a minimal JAX/CUDA solver for local flux-tube gyrokinetics with inverse-problem and sensitivity-analysis demonstrations (Galletti et al., 7 Apr 2026). The comparison is not one of direct equivalence: RadJAX is specifically an LTE molecular line renderer for disk inverse problems. A plausible implication is that its broader significance lies less in any new transfer physics than in exemplifying a general pattern in computational astrophysics: trusted forward models rewritten as JAX-native, differentiable, accelerator-oriented kernels that can be embedded directly in optimization and inference pipelines.