RUBIX: Differentiable forward modelling of galaxy spectral data cubes for gradient-based parameter estimation

TL;DR

RUBIX is a differentiable, JAX-based pipeline for modeling galaxy spectral data cubes, enabling gradient-based parameter estimation and paving the way for advanced machine learning applications in galaxy analysis.

Contribution

The paper introduces RUBIX, a fully differentiable, scalable pipeline for forward modeling of galaxy spectral data cubes using JAX, facilitating efficient parameter inference.

Findings

Validated gradients against finite differences.

Demonstrated gradient-based parameter estimation.

Showed feasibility of differentiable modeling for galaxy data.

Abstract

Although integral-field spectroscopy enables spatially resolved spectral studies of galaxies, bridging particle-based simulations to observations remains slow and non-differentiable. We present RUBIX, a JAX-based pipeline that models mock integral-field unit (IFU) cubes for galaxies end-to-end and calculates gradients with respect to particle inputs. Our implementation is purely functional, sharded, and differentiable throughout. We validate the gradients against central finite differences and demonstrate gradient-based parameter estimation on controlled setups. While current experiments are limited to basic test cases, they demonstrate the feasibility of differentiable forward modelling of IFU data. This paves the way for future work scaling up to realistic galaxy cubes and enabling machine learning workflows for IFU-based inference. The source code for the RUBIX software is publicly available under https://github.com/AstroAI-Lab/rubix.