RUBIX: Differentiable IFU Modeling

• RUBIX is a fully open-source, GPU-native, and differentiable forward-modeling framework that produces realistic integral-field spectroscopy datacubes from cosmological hydrodynamical simulations.
• It employs a modular JAX-based pipeline with jit, vmap, and pmap primitives to enable rapid, automatic parallel computation and gradient-based optimization.
• RUBIX achieves up to 600x performance improvements over traditional methods, reducing processing times from hours to seconds while supporting scalable multi-GPU execution.

RUBIX is a fully open-source, GPU-native, and auto-differentiable forward-modeling framework for integral-field spectroscopy datacubes, written end-to-end in JAX and designed to forward model IFU cubes of galaxies from cosmological hydrodynamical simulations. Its principal design goal is to take particle-level simulation outputs—positions, velocities, ages, metallicities, and related attributes—and produce realistic, instrument-specific mock IFU observations in seconds rather than hours, while making every step of the pipeline differentiable for gradient-based inference and optimization. In the reported benchmark scenario, it automatically parallelizes computations across multiple GPUs and demonstrates performance improvements over state-of-the-art codes by a factor of 600, reducing compute times from hours to only seconds (Çakır et al., 2024).

1. Definition and project scope

RUBIX is described as a fully tested, well-documented, and modular Open Source tool developed in JAX. It is intended to bridge cosmological hydrodynamical simulations and observationally realistic IFU data products by implementing the full forward model in a single differentiable software stack. The codebase is organized around instrument-specific mock observation generation rather than generic spectral synthesis alone, and the software is publicly available at https://github.com/ufuk-cakir/rubix (Çakır et al., 2024).

The framework operates on two principal inputs. The first is a JSON or Python-dict configuration specifying items such as SSP library choice, instrument parameters, cosmology, and galaxy orientation. The second is a pre-processed particle file containing star and gas particle attributes. This input design situates RUBIX as a simulation-to-observation layer: it consumes particle-based galaxy representations and emits mock IFU datacubes that are directly comparable to data products from integral-field instruments (Çakır et al., 2024).

A central point of clarification is that the current main release treats only stellar continuum. The paper explicitly presents radiative transfer and gas-emission components as subsequent modules rather than already completed parts of the main release. This distinction matters because the differentiable architecture is already in place, but the astrophysical scope of the released implementation is narrower than a full continuum-plus-nebular-plus-dust pipeline (Çakır et al., 2024).

2. Modular architecture and JAX execution model

RUBIX arranges the entire IFU forward model as a linear chain of pure functions, each of which performs one logical operation, such as orienting the galaxy, assigning particles to spaxels, computing spectra, convolving with PSF and LSF kernels, and adding noise. Each module is separately testable, can be jitted by XLA, can be vectorized via vmap on a single GPU or mapped across devices via pmap on multiple GPUs, and exposes internal parameters and intermediate outputs to JAX’s autodiff engine (Çakır et al., 2024).

The JAX primitives are central to the implementation. jit compiles each function, or the assembled pipeline, into optimized XLA kernels. vmap automatically vectorizes loops over particles or spaxels, eliminating Python-level loops. pmap shards data across multiple devices to compute in parallel. grad, jacrev, and jacfwd compute derivatives of outputs, such as synthetic spectra, with respect to any input parameter. In architectural terms, RUBIX is therefore not merely a GPU port of an existing forward model; it is a functional re-expression of the pipeline in a form compatible with compilation, device sharding, and end-to-end differentiation (Çakır et al., 2024).

This execution model has methodological implications. Because every stage is written as a pure JAX transformation, the forward model, the performance model, and the derivative structure are co-designed rather than retrofitted. That property is what enables the same code path to support both high-throughput mock-cube generation and gradient propagation through the full simulation-to-observation map.

3. Mathematical formulation of the forward model

The paper formulates the forward model at the level of spaxel spectra, instrumental response, and detector noise. If $i$ indexes spaxels $(x_i,y_i)$, $p$ indexes stellar or gas particles, and $\lambda$ denotes observer-frame wavelength, each particle $p$ has an intrinsic template spectrum $S_0(\lambda; t_p, Z_p)$ from an SSP library. At line-of-sight velocity $v_p$, that spectrum is Doppler-shifted as

$$S_p(\lambda) = S_0\!\left(\lambda /(1 + v_p/c); t_p, Z_p\right).$$

The contribution of particle $p$ to spaxel $i$ is weighted by a projected light-kernel $(x_i,y_i)$0, yielding the raw spaxel spectrum

$(x_i,y_i)$1

Instrumental effects are then introduced through sequential convolutions. The raw datacube $(x_i,y_i)$2 is convolved with a normalized two-dimensional PSF kernel $(x_i,y_i)$3,

$(x_i,y_i)$4

and each spaxel spectrum is convolved with an LSF kernel $(x_i,y_i)$5,

$(x_i,y_i)$6

The final simulated cube is obtained by mapping the convolved intensities onto the instrument wavelength grid, multiplying by the total system throughput $(x_i,y_i)$7, and adding a noise realization $(x_i,y_i)$8:

$(x_i,y_i)$9

The paper also states a prospective radiative-transfer extension. In that extension, the formal solution of

$p$0

would be evaluated along the line of sight, leading to

$p$1

The text explicitly notes that, in practice, this would be discretized over gas cells, and that this module lies beyond the current main release (Çakır et al., 2024).

4. GPU parallelization, benchmarking, and scaling

The reported performance gains come from replacing all loops over approximately $p$2–$p$3 particles or spaxels with vmap’d JAX kernels, sharding datacube assembly across devices with pmap for multi-GPU runs, and relying on XLA to fuse pointwise and small-matrix operations into large, efficient kernels. The benchmark scenario is a MUSE-like field of view with approximately $p$4 particles (Çakır et al., 2024).

Implementation	Runtime	Relative statement
GalCraft on 24-core CPU (2.5 GHz)	1.4 hours ($p$5 s)	Baseline
RUBIX on 24-core CPU (AMD Epyc 7452)	123.6 s	$p$6 faster
RUBIX on single NVIDIA A100 GPU	8.6 s	$p$7 faster than GalCraft
RUBIX on $p$8 A100 GPUs	$p$9 s	Weak-scaling tests reported

The scaling analysis is not presented as idealized. For strong scaling at fixed particle number, the timing follows $\lambda$0 with measured $\lambda$1, the deviation from linear scaling being attributed to communication overhead. For weak scaling, where $\lambda$2, the observed timing exhibits a slow rise rather than remaining constant, due to load imbalance and inter-GPU latency. The scaling efficiency

$\lambda$3

falls from approximately $\lambda$4 at $\lambda$5 to approximately $\lambda$6 at $\lambda$7. These results frame RUBIX as fast but not yet asymptotically optimal in multi-device scheduling (Çakır et al., 2024).

5. End-to-end differentiability and gradient-based estimation

Because every pipeline module is written in pure JAX, with no Python side-effects, the composed forward map

$\lambda$8

is itself jittable and differentiable. Here $\lambda$9 can include continuous parameters such as galaxy orientation angles, distance, dust-attenuation coefficients, instrument LSF width, and SSP ages. Reverse-mode autodiff then yields $p$0 for an arbitrary loss $p$1, including a $p$2 loss comparing $p$3 to data. The Jacobian

$p$4

can be recovered via jacrev for full-matrix studies or by vector-Jacobian products for optimization. The paper states that JAX records a computational graph of fused XLA kernels and that backpropagation on this graph yields exact, machine-precision gradients (Çakır et al., 2024).

The methodological significance is that forward modelling and inference are no longer disjoint stages. The paper explicitly identifies gradient-based optimization of astrophysics model parameters as an enabled use case, and it further proposes end-to-end Bayesian parameter estimation via gradient-based optimizers such as Adam and LBFGS rather than black-box samplers, integration with simulation-based inference suites such as Pyro and sbi, and hybrid “physics-decoder, neural-encoder” architectures trainable with RUBIX gradients (Çakır et al., 2024).

A later paper, "RUBIX: Differentiable forward modelling of galaxy spectral data cubes for gradient-based parameter estimation" (Schaible et al., 21 Nov 2025), provides explicit gradient validation and a controlled optimization example. In a toy setup of one particle and one spaxel, the maximum absolute difference between autodiff and central finite differences was reported as $p$5 for the age gradient and $p$6 for the metallicity gradient, and a one-particle parameter-estimation test using Adam with learning rate $p$7 converged in approximately $p$8 iterations (Schaible et al., 21 Nov 2025).

6. Scientific use cases, present limitations, and projected extensions

The paper’s experiments forward-modeled Subhalo 37 and 63874 from the TNG50-1 snapshot 99. A representative configuration specifies the SSP library "Mastar_CB19_SLOG_1_5", instrument type "MUSE", Gaussian PSF and LSF, orientation parameters with inclination $p$9 and position angle $S_0(\lambda; t_p, Z_p)$0, distance $S_0(\lambda; t_p, Z_p)$1 Mpc, and a noise specification with sky $S_0(\lambda; t_p, Z_p)$2 and read noise $S_0(\lambda; t_p, Z_p)$3. The pipeline call returns a datacube in which the main output is a 3D array flux[x,y,λ] in erg s⁻¹ cm⁻² Å⁻¹, with an optional gradient cube $S_0(\lambda; t_p, Z_p)$4 (Çakır et al., 2024).

The paper also states the current limitations. Strong and weak scaling are not yet optimal, and pmap plus jit interaction can cause collect/gather inefficiencies, discussed there in connection with JAX issue #2926. Gas-line emission, dust attenuation, and full radiative transfer are not yet implemented in pure JAX form. Memory fragmentation when handling very large particle sets can bottleneck GPU utilization. These constraints delimit the present operational envelope of the software more sharply than its architectural ambitions might suggest (Çakır et al., 2024).

The roadmap is correspondingly explicit. Planned extensions include detailed gas modeling with nebular emission lines, dust-attenuation modules fully jittable in JAX, multi-level radiative-transfer solvers such as ray-tracing and Monte Carlo in vectorized form, and improved parallel scheduling that moves more of the pipeline under pmap, reduces cross-device collect calls, and explores newer JAX parallel kernels. Taken together, these plans suggest that RUBIX is intended to develop from a fast stellar-continuum IFU forward model into a broader differentiable infrastructure for parameter inference, scientific machine learning, and real-time emulator construction in extragalactic astronomy (Çakır et al., 2024).