Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

RUBIX introduces a fully differentiable, JAX-based forward model for galaxy IFU data cubes that links particle-scale inputs to observable spectra, enabling gradient-based parameter estimation and paving the way for variational and simulation-based inference. The authors validate gradient correctness against finite differences and demonstrate a successful gradient-driven recovery of input ages and metallicities in a controlled single-spaxel setup, all while leveraging GPU acceleration and a modular, configurable pipeline. While promising, they recognize limitations in scaling to realistic, multi-spaxel galaxies and discuss strategies to manage ill-posedness and uncertainty quantification through priors and variational methods. The work establishes a foundation for integrating differentiable forward models with ML workflows in IFU data analysis and outlines clear directions for future extensions to complex galaxies and inference techniques.

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.

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

TL;DR

RUBIX introduces a fully differentiable, JAX-based forward model for galaxy IFU data cubes that links particle-scale inputs to observable spectra, enabling gradient-based parameter estimation and paving the way for variational and simulation-based inference. The authors validate gradient correctness against finite differences and demonstrate a successful gradient-driven recovery of input ages and metallicities in a controlled single-spaxel setup, all while leveraging GPU acceleration and a modular, configurable pipeline. While promising, they recognize limitations in scaling to realistic, multi-spaxel galaxies and discuss strategies to manage ill-posedness and uncertainty quantification through priors and variational methods. The work establishes a foundation for integrating differentiable forward models with ML workflows in IFU data analysis and outlines clear directions for future extensions to complex galaxies and inference techniques.

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.

Paper Structure

This paper contains 10 sections, 3 figures.

Table of Contents

  1. Motivation
  2. The RUBIX software
  3. Results
  4. Gradient validation and the physical interpretation of the gradients
  5. Parameter estimation via gradient-based optimization
  6. Discussion and limitations
  7. Ill-posed inverse problems:
  8. Future applications:
  9. Uncertainty quantification via variational inference:
  10. Conclusion

Figures (3)

  • Figure 1: Different steps that happen inside the RUBIX pipeline. We start with a config file and input data, get a mock IFU cube and in between there are pure JAX functions.
  • Figure 2: First row shows the spectrum of a single spaxel calculated by RUBIX, which contains one stellar particles with age of 10 Gyrs and metallicity of ($1.4\times10^{-3}$). We calculate the gradient with respect to the stellar age in the second row and show the results for autograd (jax.jacfwd) and central finite difference. In the third row we show the gradient via autodiff and the finite difference with respect to the metallicity of the stellar particles.
  • Figure 3: Left: loss landscape and optimization trajectories. The true parameters are marked by the yellow star; initializations are shown as white circles. Right: loss history across iterations for three runs, showing rapid convergence.