-
The manuscript repeatedly claims an “end-to-end”, “fully differentiable”, “JAX-traceable” pipeline, but Sec. 3.2–3.3 indicate reliance on non-JAX FFTLog / `mcfit.TophatVar` NumPy code paths that are “not safe to JIT-trace”. As written, it is unclear which parts are actually inside the autodiff graph. This affects the scope of supported gradients—especially derivatives with respect to cosmological parameters if $\sigma(R,z)$ (and any downstream quantities depending on it) is computed outside JAX.
Recommendation: In Sec. 3.2–3.3, explicitly state what “differentiable” means in the public implementation. Add a compact “differentiability matrix” (table) listing which outputs (e.g. $C_\ell$, bandpowers, likelihood) are differentiable with respect to which parameter blocks (cosmology vs pressure/profile vs nuisance) and whether gradients pass through: (i) emulator calls, (ii) $\sigma(R,z)$/FFTLog, (iii) HMF/bias, (iv) interpolation steps, (v) halo integrals. If $\sigma(R,z)$ is precomputed and treated as a constant inside the closure at fixed cosmology, say so and constrain the “end-to-end” language accordingly in Sec. 1–3 and Sec. 6. If cosmology gradients are intended, describe (even briefly) the concrete plan (JAX-native FFTLog, custom JVP/VJP, or alternative $\sigma$ computation) and how this would affect timings.
-
The worked-example likelihood is under-specified and therefore hard to reproduce/interpret (Sec. 5.1, Sec. 5.3–5.5). The source of the 8 bandpowers, their exact $\ell$-binning/window functions, whether the data are real or synthetic, how the full $8\times 8$ covariance was obtained (and its correlation structure), and the priors / parameterization for $(P_0,\beta)$ are not fully stated. Because the posterior exhibits long tails, priors and parameter bounds materially affect results and sampler behavior.
Recommendation: Expand Sec. 5 with a precise likelihood specification: identify the dataset origin (or state it is synthetic), provide the 8 bandpowers with $\ell$-bin edges/centers and uncertainties (table or appendix), define or reference the bandpower window functions used, and describe covariance estimation (analytic vs simulations; diagonal vs correlated; any conditioning). Write the Gaussian likelihood explicitly in terms of the data vector and covariance. State priors (forms and bounds) and whether parameters are sampled in linear/log space, and clarify whether MAP includes priors (Sec. 5.3). If the full numerical vectors/matrices live in the repository, explicitly point to file paths and ensure they match the manuscript.
-
Scientific/technical validation against established reference calculations is currently too thin given the central reliance on (i) CosmoPower ede-v2 emulators (Sec. 3.1) and (ii) a new numerical/integration implementation. There is no in-paper end-to-end comparison of $C_\ell^{yy}$ (1h/2h/total) against a CLASS/CAMB-based pipeline (e.g. class_sz) over representative cosmologies, nor a discussion of how emulator errors propagate to downstream halo-model observables and inferred parameters.
Recommendation: Add a validation subsection (Sec. 3 or Sec. 5) comparing classy_szlite to a reference pipeline (class_sz/CLASS/CAMB-based) for multiple cosmologies within the ede-v2 training domain and at least one profile setting. Show fractional differences versus $\ell$ (and ideally bandpower-level differences) for 1h, 2h, and total. Summarize emulator accuracy relevant for tSZ inputs (Sec. 3.1) and briefly discuss error propagation to $C_\ell^{yy}$ and to degeneracies (e.g. $\sigma_8$–$P_0$). Also state how points near emulator boundaries are handled.
-
A central advertised use case is higher-dimensional joint cosmology+astrophysics inference (Sec. 1, Sec. 5.2, Sec. 6.1), yet the main demonstration is restricted to a 2-parameter $(P_0,\beta)$ inference at fixed cosmology (Sec. 5). Scaling claims for NUTS and the differentiable approach at $d\gtrsim 6$–30 are extrapolated rather than empirically shown in this work.
Recommendation: Add at least one modest higher-dimensional example using the same pipeline and NUTS (e.g. include $B$ and one additional GNFW shape parameter; or a small joint cosmology+profile run with $d\sim 6$–10, even with informative priors or a simplified/mock likelihood). Report wall time, gradient-evaluation counts, ESS, R-hat, divergences, and (ideally) ESS-per-gradient-evaluation. If infeasible, soften the language in Sec. 1, Sec. 5.2, Sec. 5.5, and Sec. 6.1 to clearly label higher-$d$ performance as an expectation rather than a demonstrated result.
-
Timing benchmarks and sampler-comparison methodology are not fully normalized or clearly documented (Sec. 2; Sec. 5.1–5.5; Sec. 6.2; Fig. 1). The RW-MH vs NUTS comparison mixes differing parallelism (e.g. RW-MH with “4 MPI walkers” vs NUTS chain execution assumptions), warmup/compilation inclusion is unclear, hardware details differ across locations (laptop CPU vs EPYC vs TPU), and the headline “$\sim 100\times$” statement appears numerically inconsistent in at least one place (Sec. 5.5).
Recommendation: For Sec. 5.1–5.5, state a clear protocol: (i) whether reported times include JAX/XLA compilation and warmup; (ii) whether NUTS chains are run sequentially or in parallel; (iii) for RW-MH, proposal tuning/adaptation, warmup, and whether “4 MPI walkers” implies 4 cores used concurrently (and whether times are wall-clock vs CPU-time). Report efficiency in units that factor out parallelism differences (e.g. ESS per forward-model evaluation, ESS per gradient evaluation) in addition to ESS/sec. Correct the inconsistent speedup arithmetic in Sec. 5.5 and reconcile the NUTS wall-time discrepancy between Table 1 and Fig. 4 caption by explicitly stating what differs (sample budget, hardware, warmup inclusion, parallelization). For Fig. 1 / Sec. 6.2, standardize time units, document thread settings (OpenMP/JAX/XLA), and provide reproducible benchmark scripts and variability/error bars.
-
Core mathematical definitions for the tSZ window function are not fully verifiable from the manuscript (Sec. 3.2; Sec. 4). In particular, $W_y(\ell,M,z)$ uses symbols that are not defined (e.g. $\ell_{500}$, $J_\ell$), the prefactor is dimensionally ambiguous, and consistency with Eq. (1)’s transform convention is difficult to audit.
Recommendation: Define $\ell_{500}$ explicitly (e.g. via $\ell_{500}=D_A(z)/r_{500}$ or equivalent) and provide an explicit definition of $J_\ell[\cdot]$ (including whether it contains the $r^2dr$ measure and any projection factors). Re-check and state the dimensional consistency so that Compton-$y$ is dimensionless, and ensure the $W_y$ definition is demonstrably consistent with Eq. (1) under the substitutions $r=x r_{500}$ and $k=k_\ell$ (Sec. 4).
-
The profile Fourier-transform integral is written to $r=\infty$ (Eq. (1), Sec. 3.2) while the inference explores outer slopes around $\beta\approx 2.7$ (Sec. 5.3–5.4), for which integrals like $\int r^2 P(r)dr$ can fail to converge without truncation/apodization. The manuscript does not state an $r_{\max}$, truncation scheme, or convergence/regularization strategy, so low-$k$/low-$\ell$ well-posedness is unclear from the PDF alone.
Recommendation: State explicitly whether the GNFW/pressure profile is truncated (and at what radius, e.g. multiple of $r_{500}$ or $r_{\rm vir}$), apodized, or otherwise regularized before applying Eq. (1). Document the effective $r$-grid used in the implementation and clarify under what conditions on $\beta$ the transform exists as $k\to 0$. If finite-$r$ limits are used in code, reflect that in the mathematical description (Eq. (1) and surrounding text in Sec. 3.2).
-
Key numerical settings controlling accuracy and runtime are not systematically documented (Sec. 3.1–3.3; Sec. 5; Sec. 6.2). Important details (mass/redshift grid ranges and resolution, FFTLog configuration, interpolation schemes, $\ell$ sampling, integration rules, and default configuration used to reproduce figures/timings) are scattered or implicit, which limits strict reproducibility and makes it harder to judge accuracy/runtime trade-offs.
Recommendation: Add a dedicated “Numerical configuration” subsection (e.g. Sec. 3.4 or a preamble to Sec. 5) that lists: (i) $z$ and $M$ grid definitions (ranges, spacing, sizes); (ii) FFTLog settings for $\sigma(R,z)$ and pressure transforms (grid sizes, bias/padding, extrapolation); (iii) $\ell$ arrays and bandpower windowing; (iv) integration strategy (vectorization, quadrature, any adaptivity); (v) runtime scaling with grid sizes. Point to a single default config file / parameter dictionary in the repository that reproduces the paper’s figures and timings.