Pith. sign in

REVIEW 3 major objections 5 minor 2 cited by

Fast GPU-Powered and Auto-Differentiable Forward Modeling of IFU Data Cubes

T0 review · 3 major / 5 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read RUBIX, a JAX-based pipeline, makes mock IFU observations of galaxies hundreds of times faster on GPUs while keeping the entire forward model auto-differentiable.

desk verdict A fast, open-source JAX IFU forward model with an honest speed caveat, but the auto-differentiation claim is unverified and likely undermined by the discrete spaxel-assignment step. read the letter →

arxiv 2412.08265 v1 pith:BRJUSGRP submitted 2024-12-11 astro-ph.IM astro-ph.GAphysics.comp-phphysics.data-an

Ufuk Çak{i}r , Anna Lena Schaible , Tobias Buck This is my paper · ORCID
classification astro-ph.IMastro-ph.GAphysics.comp-phphysics.data-an
keywords IFUdatacubesforwardmodelingJAXGPUparallelizationauto-differentiationmockobservationscosmologicalsimulationsintegralfieldspectroscopy
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper introduces RUBIX, a JAX-based pipeline that converts particles from cosmological hydrodynamical simulations into integral-field-unit (IFU) data cubes, and claims it runs hundreds of times faster than existing CPU-only mock generators. The reported speedup is dramatic: a 6-million-particle MUSE-like cube takes 8.58 seconds on a single A100 GPU, compared with a quoted 1.4-hour runtime for the CPU-based GalCraft code. Because every step of the pipeline is expressed in JAX, the whole forward model is auto-differentiable, which the authors argue enables gradient-based optimization of galaxy model parameters and simulation-based inference. If these claims hold, mock IFU production ceases to be a bottleneck for building large statistical samples from simulations.

What carries the argument

The carrying mechanism is a vectorized, just-in-time-compiled JAX data-processing pipeline in which each galaxy particle is mapped to spectral contributions via vmap, operations are fused by XLA, and the datacube computation is parallelized across GPUs with pmap. Auto-differentiability follows directly from applying JAX's reverse-mode automatic differentiation to the compiled function. The performance comparison uses the GalCraft 1.4-hour runtime for a 6-million-particle MUSE field as the benchmark against which RUBIX's measured CPU and GPU times are reported.

What would settle it

Run GalCraft and RUBIX on the same machine with the same 6-million-particle input, instrument settings, and SSP library, timing both repeatedly, and compare JAX gradients to finite differences for one spaxel spectrum; a runtime ratio far below 600 or a gradient mismatch would falsify the paper's two central claims.

Watch

Extended reading notes

Core claim

The central claim is that a carefully vectorized, just-in-time-compiled JAX implementation of the IFU forward model—using vmap over particles, XLA operation fusion, and pmap across GPUs—reduces the execution time for a 6-million-particle MUSE observation from roughly 1.4 hours (the GalCraft CPU benchmark) to 123.57 seconds on a 24-core CPU and 8.58 seconds on a single NVIDIA A100 GPU, a 43.7-fold and approximately 600-fold improvement respectively. Because each pipeline stage (orientation, spaxel assignment, SSP lookup, Doppler shifting, wavelength resampling, spectral summation, PSF/LSF convolution, and noise addition) is implemented in JAX, automatic differentiation works end-to-end. The authors support the method with qualitative mock MUSE cubes from TNG50 subhalos and provide strong- and weak-scaling measurements, which show sub-ideal multi-GPU scaling due to communication overhead.

Load-bearing premise

The factor-of-600 speedup rests on the assumption that the quoted 1.4-hour GalCraft runtime on a 24-core 2.50 GHz CPU is a fair enough baseline that RUBIX's 123.57 seconds on a different 24-core CPU and 8.58 seconds on an A100 GPU can be described as 43.7x and roughly 600x faster.

Editorial extensions

If this is right

  • A 6-million-particle MUSE-like galaxy cube can be produced in 8.58 seconds on a single A100 GPU, making large mock-IFU catalogs computationally feasible.
  • End-to-end gradients of the mock cube with respect to input parameters are computable, enabling gradient-based optimization of galaxy model parameters and opening routes to simulation-based inference.
  • The CPU-only execution (123.57 seconds) is already 43.7 times faster than the quoted GalCraft benchmark, so the speedups do not depend on having GPU hardware.
  • Because the pipeline is modular and configured through JSON, users can swap SSP libraries, telescope configurations, PSF/LSF settings, and galaxy orientation without recompiling.
  • The authors report sub-ideal scaling beyond one GPU due to communication overhead and the jit-wrapping-pmap interaction, leaving headroom for further optimization.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If the end-to-end gradients are as clean in practice as in principle, one could use the mock-cube Jacobian for Hamiltonian Monte Carlo or variational inference over galaxy parameters, something CPU-only mock generators cannot do.
  • The authors' profiling points to the jit-of-pmap data-collection bottleneck; a fully sharded compilation strategy would likely recover near-ideal weak scaling and should be tested on more than eight GPUs.
  • A natural stress test is to run RUBIX on a 10^7-particle disk galaxy with dust and gas emission included; if runtime stays in the tens of seconds, the method can scale to survey-volume mock production.
  • A like-for-like benchmark with GalCraft on identical hardware, using the same input catalog and instrument configuration, would pin down how much of the 600x factor is algorithmic versus hardware-dependent.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

Summary. The paper presents RUBIX, a JAX-based forward-modeling code that generates mock integral-field-unit (IFU) datacubes from cosmological hydrodynamical simulations. The pipeline takes star particle data, orients the galaxy, assigns particles to spaxels, computes Doppler-shifted SSP spectra, resamples to the telescope wavelength grid, sums spectra per spaxel, and applies PSF/LSF convolution and noise. The central claims are: (1) the code is automatically parallelized across multiple GPUs and achieves a factor-of-600 speedup over the GalCraft code, and (2) JAX's auto-differentiation makes the entire pipeline differentiable, enabling gradient-based inference. The code is open source. Validation is limited to a qualitative visual inspection of two mock MUSE datacubes and scaling tests with up to 100,000 particles.

Significance. If the claims are sound, RUBIX would be a valuable community tool: it would lower the computational cost of producing large mock IFU samples for simulation-based inference, and its claimed end-to-end differentiability would open new avenues for gradient-based parameter optimization. The open-source release and use of JAX's vectorization are explicit strengths. However, the two headline claims are currently supported only by a qualitative check and by a benchmark against a published runtime on different hardware, respectively; the auto-differentiation claim in particular is neither demonstrated nor clearly reconciled with the discrete spaxel-assignment step. These issues are load-bearing for the methodological contribution.

major comments (3)
  1. [Abstract and Section 3] The claim that RUBIX provides 'gradient computations through the entire pipeline' is not substantiated by any gradient computation or numerical test in Sections 3–4. The pipeline includes a step that assigns particles to IFU spaxels, which in JAX is naturally implemented with integer indexing or scatter operations; gradients through such operations with respect to continuous inputs (e.g., galaxy distance and orientation) are zero almost everywhere unless a continuous/soft assignment is used. The authors must either demonstrate a working gradient (for example, by comparing a finite-difference Jacobian of the datacube with respect to distance and orientation to the JAX-computed gradient) or describe and verify the differentiable surrogate for the spaxel-assignment step. Without this, the promised gradient-based optimization and SBI applications do not follow.
  2. [Section 4, Figure 2] The abstract's 'factor of 600' compares a measured RUBIX GPU runtime (8.58 s on one A100) with a GalCraft runtime of 1.4 hours on a 24-core 2.50 GHz CPU quoted from Wang et al. (2024). The authors appropriately caution in the text that the hardware and methodology differ, but the abstract and the opening motivation present the number without this caveat. As stated, the comparison is not a controlled benchmark: there is no common input, no same-machine run of GalCraft, and no specification of the particle data and instrument configuration used for the RUBIX timing. To support the headline speedup, the authors should either perform a controlled comparison on identical hardware and input or rephrase the claim as '8.58 s on one A100 versus 1.4 h on a 24-core CPU from a published GalCraft run' and remove the unqualified factor of 600 from the abstract.
  3. [Section 4, Figure 1] The validation is purely qualitative: 'we can observe that RUBIX can reproduce the trends that we expect' (Section 4). The paper describes the code as 'fully tested', but no quantitative correctness tests are shown. A forward-modeling code should at minimum be checked against a simple analytic model (e.g., a point source or a constant-velocity stellar population) to verify flux normalization, Doppler shifting, and PSF/LSF convolution, or against a reference code on a small common input. Without such a test, the correctness of the pipeline, which is prerequisite for both the speed and differentiability claims to be useful, remains unverified.
minor comments (5)
  1. [Section 4, Figure 3(a)] Figure 3(a) is labeled 'Strong Scaling' but it plots average runtime versus number of particles with a fixed number of GPUs (8). Strong scaling normally holds the problem size fixed and increases the number of processors; this plot would be more accurately described as runtime scaling with particle number at fixed resources. The term 'strong scaling' should either be corrected or the plot replaced with a fixed-size speedup versus GPU count.
  2. [Abstract, Figure 2] The abstract states 'by a factor of 600', while Figure 2 reports '629.37x faster'. Please harmonize the numerical value or state it consistently as 'about 600'.
  3. [Section 4, Figure 3] The scaling experiments in Figure 3 use at most 100,000 particles, whereas the headline timing run uses 6 million particles; the scaling-efficiency conclusions are therefore presented as typical behavior at production scale without direct evidence. Please state this extrapolation explicitly.
  4. [Section 4] The paper would benefit from stating which specific simulation data (e.g., which IllustrisTNG subhalo, snapshot, and particle type) were used for the 6-million-particle timing run and from reporting the number of spaxels and wavelength bins, since the runtime depends strongly on these dimensions.
  5. [Section 3] The description of the parallelization is brief: the text mentions that only the datacube calculation is parallelized with pmap and that jit'ing a pmap function can cause inefficient data movement. A short paragraph explaining exactly which stages run on multiple GPUs and where data are gathered would help readers assess the scaling results.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: RUBIX's performance and capability claims rest on external benchmarks and code measurements, not on self-referential construction.

full rationale

This paper is a software and performance report, not a mathematical derivation or inference pipeline. Its central claims are (1) that RUBIX runs fast on GPUs and (2) that it is written in JAX and thus supports autodifferentiation. Both claims are empirical or architectural, not derived from fitting parameters to data or from a self-citation chain. The factor-of-600 speedup is an external comparison against a GalCraft runtime quoted from Wang et al. (2024); the authors explicitly caution that hardware and methodology differ, making this a benchmark-fairness concern rather than circularity. The autodifferentiation claim is asserted but not demonstrated with a gradient check; that is an unsupported claim or correctness risk, not a circular step, because no quantity is defined in terms of its own prediction. The paper does not invoke the authors' own prior results as load-bearing support, and no equation or construction reduces a predicted output to an input fitted quantity. Accordingly, the circularity score is 0.

Assumptions & free parameters 0 free parameters · 4 assumptions · 0 invented entities

The central claim is about software performance, so no new physical entities are introduced. The ledger records the domain assumptions about the forward-modeling physics that the tool relies on.

assumptions (4)
  • domain assumption An SSP library lookup accurately represents the spectra of stellar particles.
    Invoked in Section 3 ('Each star spectrum is calculated as a lookup from a simple stellar population (SSP) library'); the accuracy of the mock cubes depends on this.
  • domain assumption Doppler shifting and resampling onto the telescope wavelength grid correctly model observed spectra.
    Section 3 pipeline steps; any error in this mapping would propagate to the output cube, but it is a standard approximation.
  • domain assumption PSF and LSF convolution plus added noise adequately reproduce observational effects.
    Section 3 ('To simulate observational effects, we apply a point-spread function (PSF) and line-spread function (LSF) convolution and add realistic noise'); this is a simplification of real instrument response.
  • domain assumption IllustrisTNG simulation outputs provide a suitable ground truth for mock observations.
    Used in Section 4 for the example mocks; the realism of the mocks is limited by the simulation's physics.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Fast GPU-Powered and Auto-Differentiable Forward Modeling of IFU Data Cubes." pith.science (2026). https://pith.science/paper/BRJUSGRP

@misc{pith2026241208265,
  author       = {Pith},
  title        = {Pith review of: Fast GPU-Powered and Auto-Differentiable Forward Modeling of IFU Data Cubes},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/BRJUSGRP}},
  note         = {Machine review of arXiv:2412.08265}
}
read the original abstract

We present RUBIX, a fully tested, well-documented, and modular Open Source tool developed in JAX, designed to forward model IFU cubes of galaxies from cosmological hydrodynamical simulations. The code automatically parallelizes computations across multiple GPUs, demonstrating performance improvements over state-of-the-art codes by a factor of 600. This optimization reduces compute times from hours to only seconds. RUBIX leverages JAX's auto-differentiation capabilities to enable not only forward modeling but also gradient computations through the entire pipeline paving the way for new methodological approaches such as e.g. gradient-based optimization of astrophysics model parameters. RUBIX is open-source and available on GitHub: https://github.com/ufuk-cakir/rubix.

Figures

Figures reproduced from arXiv: 2412.08265 by the authors.

Figure 1
Figure 1. MUSE mock observations – for different Subhalos, the total flux in each pixel is shown as an image representation on the left. On the right, the spectra of three different spaxels are plotted. are illustrated in [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Speed comparison – the execution time of different codes are compared. Note that the y￾axis is logarithmic. Speed comparison The primary objective of this paper is to highlight the methodological improvements RUBIX provides for the forward modeling process. In [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Scaling plots – (a) Strong Scaling: Increasing particle size, while keeping number of GPUs fixed (8 NVIDIA A100 GPUs). (b) Average runtime of different RUBIX runs, where we proportionally increase particle size and number of available GPUs, such that the workload per GPU remains constant. (c) Scaling efficiency calculated as the ratio of the runtime with one GPU to the runtime with multiple GPUs. of GPUs is reached,… view at source ↗

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. End-to-end differentiable retrieval of molecular spectra using hydrodynamics, chemistry, and radiative transfer

    astro-ph.IM 2026-07 conditional novelty 6.0 of 10

    An end-to-end differentiable JAX pipeline couples 1D hydrodynamics, time-dependent chemistry, and radiative transfer, and recovers shock and rate parameters from synthetic HCO+ spectra.

  2. Star Formation Rates, Metallicities, and Stellar Masses on kpc-scales in TNG50

    astro-ph.GA 2025-01 conditional novelty 5.0 of 10

    In TNG50, the resolved star-forming main sequence has slope 0.30, much shallower than observed, while the resolved mass-metallicity relation matches data; both can be described by a leaky-box model that prefers low ne...

Reference graph

Works this paper leans on

10 extracted references · 10 canonical work pages · cited by 2 Pith papers

  1. [1]

    Realistic synthetic integral field spectroscopy with realsim- ifs

    Connor Bottrell and Maan H Hani. Realistic synthetic integral field spectroscopy with realsim- ifs. Monthly Notices of the Royal Astronomical Society , 514(2):2821–2838, June 2022

    work page 2022

  2. [2]

    Bershady, David R

    Kevin Bundy, Matthew A. Bershady, David R. Law, Renbin Yan, Niv Drory, Nicholas MacDon- ald, David A. Wake, Brian Cherinka, José R. Sánchez-Gallego, Anne-Marie Weijmans, Daniel Thomas, Christy Tremonti, Karen Masters, Lodovico Coccato, Aleksandar M. Diamond-Stanic, Alfonso Aragón-Salamanca, Vladimir Avila-Reese, Carles Badenes, Jésus Falcón-Barroso, Franc...

    work page 2015

  3. [3]

    Simspin: Kinematic analysis of simulated galaxies, 2023

    Katherine Harborne. Simspin: Kinematic analysis of simulated galaxies, 2023. Publications of the Astronomical Society of Australia, V olume 40, article id. e048, Oct 2023

    work page 2023

  4. [4]

    iMaNGA: mock MaNGA galaxies based on IllustrisTNG and MaStar SSPs – I

    Lorenza Nanni, Daniel Thomas, James Trayford, Claudia Maraston, Justus Neumann, David R Law, Lewis Hill, Annalisa Pillepich, Renbin Yan, Yanping Chen, and Dan Lazarz. iMaNGA: mock MaNGA galaxies based on IllustrisTNG and MaStar SSPs – I. Construction and analysis of the mock data cubes. Monthly Notices of the Royal Astronomical Society , 515(1):320–338, 0...

    work page 2022

  5. [5]

    First results from the tng50 simulation: galactic outflows driven by supernovae and black hole feedback

    Dylan Nelson, Annalisa Pillepich, V olker Springel, Rüdiger Pakmor, Rainer Weinberger, Shy Genel, Paul Torrey, Mark V ogelsberger, Federico Marinacci, and Lars Hernquist. First results from the tng50 simulation: galactic outflows driven by supernovae and black hole feedback. Monthly Notices of the Royal Astronomical Society , 490(3):3234–3261, August 2019

    work page 2019

  6. [6]

    First results from the tng50 simulation: the evolution of stellar and gaseous discs across cosmic time

    Annalisa Pillepich, Dylan Nelson, V olker Springel, Rüdiger Pakmor, Paul Torrey, Rainer Weinberger, Mark V ogelsberger, Federico Marinacci, Shy Genel, Arjen van der Wel, and Lars Hernquist. First results from the tng50 simulation: the evolution of stellar and gaseous discs across cosmic time. Monthly Notices of the Royal Astronomical Society , 490(3):3196...

    work page 2019

  7. [7]

    Knapen, Héctor Ibarra-Medel, Annalisa Pillepich, Sebastián F

    Regina Sarmiento, Marc Huertas-Company, Johan H. Knapen, Héctor Ibarra-Medel, Annalisa Pillepich, Sebastián F. Sánchez, and Alina Boecker. Mangia: 10 000 mock galaxies for stellar population analysis. Astronomy &; Astrophysics, 673:A23, April 2023

    work page 2023

  8. [8]

    Chemical enrichment and radial migration in the Galactic disc – the origin of the [αFe] double sequence

    Sanjib Sharma, Michael R Hayden, and Joss Bland-Hawthorn. Chemical enrichment and radial migration in the Galactic disc – the origin of the [αFe] double sequence. Monthly Notices of the Royal Astronomical Society, 507(4):5882–5901, 07 2021

    work page 2021

Show all 10 references
  1. [9]

    S. F. Sánchez, J. K. Barrera-Ballesteros, E. Lacerda, A. Mejía-Narvaez, A. Camps-Fariña, Gustavo Bruzual, C. Espinosa-Ponce, A. Rodríguez-Puebla, A. R. Calette, H. Ibarra-Medel, V . Avila-Reese, H. Hernandez-Toledo, M. A. Bershady, M. Cano-Diaz, and A. M. Munguia- Cordova. Sds...

    2022

  2. [10]

    Hayden, Jesse van de Sande, Joss Bland-Hawthorn, Sam Vaughan, Marie Martig, and Francesca Pinna

    Zixian Wang, Sanjib Sharma, Michael R. Hayden, Jesse van de Sande, Joss Bland-Hawthorn, Sam Vaughan, Marie Martig, and Francesca Pinna. Validating full-spectrum fitting with a synthetic integral-field spectroscopic observation of the Milky Way. MNRAS, 534(2):1175– 1204, Octobe...

    2024

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.