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

Anna Lena Schaible, Ufuk Çakır, Harald Mack, Aura Obreja, William H. Oliver, Nihat Oguz, Horea Cărămizaru, Tobias Buck

Keywords: astronomy, differentiable programming, gradient-based parameter estimation, JAX, IFU

TL;DR: We present \texttt{RUBIX}, a differentiable IFU forward model, and show the potential of gradient-based parameter estimation for galaxy observations.

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 25