pith. sign in

arxiv: 2604.08463 · v1 · submitted 2026-04-09 · 🌌 astro-ph.CO

A GPU-Accelerated JAX Framework for Robust Parametric Component Separation and Clustering Optimization for CMB Polarization Satellites

Wassim Kabalan , Arianna Rizzieri , Wuhyun Sohn , Artem Basyrov , Alexandre Boucaud , Benjamin Beringue , Pierre Chanial , Ema Tsang King Sang
show 1 more author
Josquin Errard
This is my paper

Pith reviewed 2026-05-10 17:01 UTC · model grok-4.3

classification 🌌 astro-ph.CO
keywords CMB polarizationcomponent separationforeground SEDsK-means clusteringtensor-to-scalar ratioJAXLiteBIRDparametric modeling
0
0 comments X

The pith

An optimized K-means clustering configuration in a JAX pipeline reduces the 68% upper limit on the tensor-to-scalar ratio by about 30% in LiteBIRD-like CMB simulations.

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

The paper develops a JAX-based framework for parametric component separation of CMB polarization data that accounts for foregrounds with spatially varying spectral energy distributions. It groups pixels into clusters sharing similar spectral parameters using K-means and scans over thousands of such groupings to optimize the balance between model detail and contamination. The GPU implementation speeds up the spectral likelihood evaluations by up to 100 times compared to previous methods. Tests on LiteBIRD-like simulations show that the best clustering setup improves the upper limit on the tensor-to-scalar ratio r by roughly 30 percent over a standard multi-resolution approach, with similar statistical uncertainties. This is significant for efforts to detect primordial gravitational waves through CMB B-modes.

Core claim

The central claim is that a JAX-powered pipeline extending parametric component separation methods can efficiently optimize K-means clustering of pixels by shared foreground SEDs, leading to better separation performance and a 30 percent reduction in the 68 percent upper limit on r for LiteBIRD simulations while preserving competitive error bars.

What carries the argument

The vectorized, GPU-accelerated evaluation of the spectral likelihood across thousands of K-means pixel subset configurations within the parametric foreground model.

Load-bearing premise

The K-means algorithm on pixel subsets reliably identifies areas with comparable foreground spectral energy distributions without introducing biases or overfitting to the specific noise in the test simulations.

What would settle it

Applying the optimized clustering to an independent suite of LiteBIRD-like simulations generated with different foreground realizations and noise seeds, and checking whether the 30% improvement in the r upper limit is reproduced.

Figures

Figures reproduced from arXiv: 2604.08463 by Alexandre Boucaud, Arianna Rizzieri, Artem Basyrov, Benjamin Beringue, Ema Tsang King Sang, Josquin Errard, Pierre Chanial, Wassim Kabalan, Wuhyun Sohn.

Figure 2
Figure 2. Figure 2: Example of spherical K-means patching applied to a HEALPix￾formatted mask (Górski et al. 2005), showing patches sharing a common spectral parameter. Each color represents a distinct patch. GAL040) 1 . These masks divide the sky into regions with different levels of Galactic foreground contamination—typically distinguish￾ing low-, medium-, and high-foreground areas based on thresholding the emission detecte… view at source ↗
Figure 3
Figure 3. Figure 3: Overview of the end-to-end pipeline for adaptive component sep￾aration and tensor-to-scalar ratio estimation. The grid search (top, blue) enu￾merates all configurations K ∈ G (Section 2.3), distributed across devices via jax-grid-search (Kabalan 2025). For each configuration, the pipeline operates independently on three disjoint sky regions—high-latitude (hi-lat, GAL020), mid-latitude (mid-lat, GAL040 − 02… view at source ↗
Figure 5
Figure 5. Figure 5: Sky partitioning into three disjoint regions based on the Planck Galactic plane masks (Planck Collaboration 2020). The high-latitude region retains the cleanest 20% of the sky ( 𝑓sky = 0.2, GAL020, hereafter hi-lat), the mid-latitude region covers the next 20% (0.2 < 𝑓sky ≤ 0.4, hereafter mid-lat), and the low-latitude region the following 20% (0.4 < 𝑓sky ≤ 0.6, hereafter low-lat). The component separation… view at source ↗
Figure 6
Figure 6. Figure 6: Validation on simplified synthetic data. The 68% upper limit on 𝑟 (𝑟 + 𝜎(𝑟 )) is plotted as a function of the number of dust spectral-index patches 𝐾𝛽𝑑 . Green: synthetic sky with 𝐾𝛽𝑑 = 100, 𝐾𝑇𝑑 = 15, 𝐾𝛽𝑠 = 5; red: c1d0s0 sky (uniform SEDs). Dashed vertical lines mark the ground￾truth 𝐾𝛽𝑑 for each case. Note that the 𝑦-axis is truncated at 1.05 × 10−4 ; configurations with 𝑟 + 𝜎(𝑟 ) above this threshold ar… view at source ↗
Figure 7
Figure 7. Figure 7: Trade-off between recovered CMB spatial variance (Q+U, 𝑥-axis, in 𝜇K 2 ) and tensor-to-scalar ratio 𝑟 (𝑦-axis). Open circles show 𝑟; filled circles show 𝑟 + 𝜎(𝑟 ). Color indicates the total number of patches. Configurations with the lowest variance (left) tend to have high systematic bias, while the 68% upper bound exhibits a V-shaped envelope whose minimum reveals the optimal trade-off between bias and st… view at source ↗
Figure 8
Figure 8. Figure 8: Grid search results on the c1d1s1 sky. Left column: 68% upper limit on 𝑟 (𝑟 + 𝜎(𝑟 )) as a function of the number of patches for each spectral parameter, shown for the three Galactic regions (hi-lat, mid-lat, low-lat). The other two parameters are fixed at representative values (see text). Right column: resulting 𝐵-mode power spectra for the systematic, statistical and total residuals averaged over the numb… view at source ↗
Figure 9
Figure 9. Figure 9: Comparison of true input spectral parameter maps (d1s1, top) with recovered parameters from the optimal patch configuration (bottom), from a specific noise realisation. From left to right: dust spectral index 𝛽𝑑, dust temperature 𝑇𝑑, and synchrotron spectral index 𝛽𝑠. Grey regions are masked. Note that the colorbar ranges differ between the input and recovered maps to accommodate the broader range of recov… view at source ↗
Figure 10
Figure 10. Figure 10: Residual 𝐵-mode power spectra comparing the multi-resolution configuration of LiteBIRD Collaboration et al. (2023) (purple) with our op￾timized K-means configuration (“This work”, green), corresponding to the “All Combined” configuration of Fig. 8d. Line styles indicate total (𝐶res ℓ , dashed), systematic (𝐶 syst ℓ , solid), and statistical (𝐶stat ℓ , dotted) residuals. The gray band shows the primordial … view at source ↗
Figure 12
Figure 12. Figure 12: Comparison of patch geometries between the multi-resolution strategy of LiteBIRD Collaboration et al. (2023) and our optimized K-means clustering, summed over the three Galactic regions. From left to right: dust spectral index 𝛽𝑑, dust temperature 𝑇𝑑, and synchrotron spectral index 𝛽𝑠. Grey regions are masked. The K-means approach allocates significantly more patches to 𝑇𝑑 and 𝛽𝑠, which is associated with… view at source ↗
Figure 13
Figure 13. Figure 13: Residual 𝐵-mode power spectra for the uniform configuration (“low patches”, 𝐾𝛽𝑑 = 𝐾𝑇𝑑 = 𝐾𝛽𝑠 = 1) and an over-parameterized configuration (“high patches”), applied to input skies with 𝑟 = 0 and 𝑟 = 3 × 10−3 . Line styles indicate observed spectra (𝐶obs ℓ , dashed), systematic (𝐶 syst ℓ , solid), and statistical (𝐶stat ℓ , dotted) residuals. The grey band shows the primordial 𝐶 𝐵𝐵 ℓ for 𝑟 ∈ [10−3 , 4 × 10−3… view at source ↗
Figure 15
Figure 15. Figure 15: Sky partition obtained by binning the recovered spectral-parameter templates of the optimal K-means configuration into 𝑁bin = 100 equal￾width intervals per parameter (𝛽𝑑, 𝑇𝑑, 𝛽𝑠). Pixels sharing the same bin for a given parameter are assigned to the same pixel subset, yielding spatially disconnected regions. This is the coarsest configuration shown in [PITH_FULL_IMAGE:figures/full_fig_p015_15.png] view at source ↗
Figure 16
Figure 16. Figure 16: Statistical (dotted) and systematic (solid) 𝐵-mode residual power spectra for the optimal K-means configuration and three progressively coarser configurations obtained by binning the recovered spectral-parameter tem￾plates into 𝑁bin = 1000 and 100 equal-width intervals per parameter with the statistical residuals computed by averaging across 40 noise realisations using equation 14. As disconnected cluster… view at source ↗
read the original abstract

We present a novel, JAX-powered implementation of a parametric component-separation method for CMB polarization data, explicitly designed to handle spatially varying foreground Spectral Energy Distributions (SEDs). The approach models this variation across the sky by grouping sets of pixels that share common foreground spectral parameters, scanning over thousands of such configurations to evaluate the trade-off between model complexity and residual systematic contamination. Built within the FURAX framework -- a JAX-powered environment for CMB data analysis -- our pipeline extends the fgbuster parametric formalism. It enables fully vectorized, GPU-accelerated evaluation of the spectral likelihood, map reconstruction, and diagnostic metrics across tens of thousands of pixel subset configurations, noise realizations, and sky regions. Our implementation achieves up to $\sim 100\times$ speed-up over the scipy TNC optimizer used in fgbuster when running on GPUs, as well as giving more robust results. When applied to LiteBIRD-like simulations with spatially varying foreground SEDs, our optimized K-means configuration reduces the 68% upper limit on the tensor-to-scalar ratio $r$ by $\approx 30\%$ relative to a fixed, previously derived multi-resolution configuration, while maintaining competitive statistical uncertainties.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript presents a JAX-based, GPU-accelerated extension of the fgbuster parametric component-separation formalism for CMB polarization data. It groups pixels via K-means clustering to model spatially varying foreground SEDs, scans thousands of such groupings to balance model complexity against residuals, and reports up to 100x speed-up relative to scipy TNC. On LiteBIRD-like simulations the optimized configuration yields an approximately 30% tighter 68% upper limit on the tensor-to-scalar ratio r while preserving competitive statistical uncertainties.

Significance. If the reported improvement is robust, the framework would provide a practical, scalable tool for foreground cleaning in next-generation CMB experiments. The vectorized JAX implementation and explicit handling of spatially varying SEDs address a recognized limitation of fixed-resolution parametric methods; the claimed speed-up is a concrete engineering advance that could enable larger parameter scans or Monte-Carlo suites.

major comments (2)
  1. [abstract and results section describing LiteBIRD simulations] The central result (abstract and §4) selects the K-means configuration by minimizing the r upper limit on the identical LiteBIRD-like simulation suite used for final evaluation. No held-out noise realizations, cross-validation, or independent simulation set is described for the configuration choice; this leaves open the possibility that the 30% gain partly reflects noise-specific alignments rather than improved modeling of SED spatial variation.
  2. [methods section on K-means clustering] The weakest-assumption paragraph notes that K-means on pixel subsets is assumed to capture regions of similar foreground SEDs without introducing new biases. The manuscript provides no quantitative test (e.g., comparison of recovered SED parameters against input maps or residual power spectra binned by cluster) that would demonstrate the clustering step itself is unbiased at the level required for the r constraint.
minor comments (2)
  1. [abstract] The abstract states 'more robust results' without defining the metric; a brief clarification in the text would help readers interpret the comparison to fgbuster.
  2. [methods] Notation for the spectral likelihood and the precise definition of the 'multi-resolution configuration' baseline should be stated explicitly once in the methods to avoid ambiguity when comparing figures.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for their positive evaluation of the significance of our work and for the constructive major comments. We address each point below and describe the revisions we will implement to strengthen the manuscript.

read point-by-point responses

  1. Referee: [abstract and results section describing LiteBIRD simulations] The central result (abstract and §4) selects the K-means configuration by minimizing the r upper limit on the identical LiteBIRD-like simulation suite used for final evaluation. No held-out noise realizations, cross-validation, or independent simulation set is described for the configuration choice; this leaves open the possibility that the 30% gain partly reflects noise-specific alignments rather than improved modeling of SED spatial variation.

    Authors: We acknowledge that the configuration optimization was performed by selecting the K-means grouping that minimizes the r upper limit on the same LiteBIRD-like simulation suite used for the final reported results. While this is a standard approach in simulation studies aimed at identifying the best-performing model within a given framework, we agree that it does not fully exclude the possibility of alignment with specific noise realizations. In the revised manuscript we will add a cross-validation procedure: the available noise realizations will be partitioned into independent training and validation sets; the K-means configuration scan will be performed on the training set, and the selected configuration will then be evaluated on the held-out validation set. We will report the r constraints obtained on both sets to demonstrate that the reported improvement is robust and arises from improved modeling of spatially varying SEDs. revision: yes

  2. Referee: [methods section on K-means clustering] The weakest-assumption paragraph notes that K-means on pixel subsets is assumed to capture regions of similar foreground SEDs without introducing new biases. The manuscript provides no quantitative test (e.g., comparison of recovered SED parameters against input maps or residual power spectra binned by cluster) that would demonstrate the clustering step itself is unbiased at the level required for the r constraint.

    Authors: We agree that a direct quantitative validation of the clustering assumption would strengthen the paper. The current manuscript states the assumption in the weakest-assumption paragraph but does not include an explicit test against the input maps. In the revised version we will add, in the methods section, a dedicated validation subsection that (i) compares the recovered spectral parameters within each K-means cluster to the known input foreground SED maps and (ii) presents residual power spectra computed separately for pixels belonging to each cluster. These diagnostics will quantify any residual bias introduced by the clustering step at the angular scales relevant to the r constraint. revision: yes

Circularity Check

1 steps flagged

K-means config optimization on same LiteBIRD simulations forces the reported 30% r-limit improvement by selecting the minimum on evaluation data

specific steps
  1. fitted input called prediction [Abstract]
    "When applied to LiteBIRD-like simulations with spatially varying foreground SEDs, our optimized K-means configuration reduces the 68% upper limit on the tensor-to-scalar ratio r by ≈30% relative to a fixed, previously derived multi-resolution configuration, while maintaining competitive statistical uncertainties."

    The pipeline scans thousands of K-means configurations on the same LiteBIRD-like simulations (including their specific noise realizations) to evaluate trade-offs and selects the 'optimized' grouping. The 30% reduction is therefore the minimum achieved on the evaluation data itself rather than a prediction on independent data, making the improvement statistically forced by the selection procedure.

full rationale

The central scientific claim is the 30% reduction in the 68% upper limit on r achieved by the 'optimized K-means configuration'. This configuration is obtained by scanning thousands of pixel groupings on the identical LiteBIRD-like simulation suite used to compute the r upper limits. The reported improvement is therefore the minimized value on the data used for selection, with no indication of held-out noise realizations or cross-validation. This matches the fitted-input-called-prediction pattern: the headline result is statistically forced by the optimization procedure rather than representing an independent prediction. The comparison baseline is a fixed prior configuration, but the gain itself reduces to the selection step on the same inputs. No other circular steps (self-citations, ansatzes, or renamings) are present in the provided text; the JAX implementation and speed-up claims are independent engineering results.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claim rests on standard assumptions of parametric foreground modeling and the validity of the LiteBIRD-like simulations; no new physical entities are introduced.

axioms (1)
  • domain assumption Foreground spectral energy distributions can be adequately described by a small number of spatially varying parameters within each clustered pixel group
    Invoked when grouping pixels and scanning configurations; standard in CMB component separation but not proven for all sky regions.

pith-pipeline@v0.9.0 · 5551 in / 1299 out tokens · 50545 ms · 2026-05-10T17:01:33.184118+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

3 extracted references · 3 canonical work pages

  1. [1]

    release score

    Alonso D., Sanchez J., Slosar A., 2019, Monthly Notices of the Royal Astro- nomical Society, 484, 4127 BinghamE.,ChenJ.P.,JankowiakM.,ObermeyerF.,PradhanN.,Karaletsos T., Kingma D. P., Tran D., 2019, NumPyro: A Lightweight Library for ProbabilisticProgramminginJAX,https://github.com/pyro-ppl/ numpyro Bradbury J., et al., 2018, JAX: composable transformati...

    work page doi:10.5281/zenodo.17674778 2019

  2. [2]

    −gwith the feasible direction into the valid parameter space. For a parameter𝑖currently at a bound: score𝑖 =𝑝 𝑖 × (−𝑔int,𝑖 ),(A2) whereg int is the gradient in the internal parameter space, and𝑝𝑖 ∈ {−1,0,1}is the pivot vector.𝑝 𝑖 =−1indicates the parameter is at the lower bound,𝑝𝑖 =1at the upper bound, and𝑝𝑖 =0indicates it is free. WethenuseJAX’shardware-...

    work page 2026

  3. [3]

    chattering

    As disconnected clusters are merged, statistical residuals decrease while systematic residuals increase. The gray band shows the pri- mordial𝐶 𝐵𝐵 ℓ for𝑟∈ [10 −3 ,4×10 −3 ]; the solid gray line is the lensing contribution. most strongly pushed towards the feasible region by the gradient) and release them simultaneously (p𝑖 ←0). A3 Impact of the Top-K Fract...

    work page 1999