⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

Questo articolo introduce HIcosmo, un framework open-source basato su JAX per l'inferenza cosmologica efficiente che sfrutta la differiazione automatica e l'accelerazione su GPU per raggiungere velocità di campionamento significativamente superiori rispetto agli strumenti esistenti come Cobaya, supportando al contempo una vasta gamma di likelihood di survey Stage IV e modelli cosmologici.

Autori originali: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Pubblicato 2026-06-29

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immaginate l'universo come un enorme e complesso puzzle che gli scienziati stanno cercando di risolvere. Per decenni, hanno utilizzato un insieme standard di regole (chiamato modello $\Lambda$ CDM) per capire come sia iniziato l'universo, come si stia espandendo e di cosa sia fatto. Ma recentemente, i pezzi del puzzle non si incastrano perfettamente. Le misurazioni dell'universo primordiale (come la Radiazione Cosmica di Fondo) non concordano con le misurazioni dell'universo vicino (come le stelle esplodenti chiamate Supernovae). Questo è noto come la "Tensione di Hubble".

Per risolvere questo problema, vengono costruiti nuovi, massicci telescopi che raccoglieranno miliardi di punti dati. Il problema? I programmi informatici che gli scienziati usano per risolvere questi enigmi sono troppo lenti e macchinosi per gestire questo fiume di informazioni. Sono come cercare di contare i granelli di sabbia su una spiaggia usando una singola, lenta calcolatrice.

Questo articolo presenta HIcosmo, un nuovissimo strumento informatico super veloce progettato per risolvere questi enigmi cosmici in modo efficiente. Ecco come funziona, spiegato in modo semplice:

1. La "Calcolatrice Intelligente" (JAX e la Differenziabilità)

Gli strumenti tradizionali sono come uno chef che deve misurare manualmente ogni ingrediente, annotare il risultato e poi calcolare manualmente il passaggio successivo. Se hanno bisogno di sapere come la modifica di un singolo ingrediente influenzi il sapore finale, devono indovinare e controllare, il che richiede un tempo infinito.

HIcosmo è diverso. È costruito su una tecnologia chiamata JAX, che agisce come una "calcolatrice intelligente" che comprende l'intera ricetta in un colpo solo.

Differenziazione Automatica: Invece di indovinare, HIcosmo sa esattamente come cambia il risultato se si modifica un singolo numero. È come avere un GPS che ti dice istantaneamente il percorso migliore da seguire senza dover percorrere ogni strada possibile per vedere quale sia la più veloce.
Compilazione Just-in-Time: La prima volta che chiedete a HIcosmo di eseguire un calcolo, impiega una frazione di secondo per "scaldarsi" e tradurre le istruzioni in un codice macchina super ottimizzato. Dopo di che, gira a una velocità fulminea, bypassando il lento "traduttore" (l'interprete Python) che rallenta gli strumenti più vecchi.

2. La Torta a Tre Strati

Gli autori hanno costruito HIcosmo come una torta a tre strati, dove ogni strato ha un compito specifico:

Lo Strato Inferiore (Il Modello): Questa è la parte "teorica". Prende le regole dell'universo (come quanta materia oscura o energia oscura ci sia) e calcola come l'universo dovrebbe apparire.
Lo Strato Centrale (La Verosimiglianza/Likelihood): Questa è la parte di "confronto". Prende la teoria dallo strato inferiore e la confronta con i dati reali provenienti dai telescopi (come i dati delle supernovae Pantheon+ o le mappe delle galassie DESI). Calcola quanto la teoria sia "sbagliata".
Lo Strato Superiore (L'Inferenza): Questa è la parte del "solutore". Utilizza le informazioni dello strato centrale per trovare la risposta migliore. Utilizza un metodo intelligente chiamato NUTS (No-U-Turn Sampler) che usa i gradienti della "calcolatrice intelligente" per puntare direttamente alla risposta corretta, invece di vagare a vuoto come fanno gli strumenti più vecchi.

3. Cosa Può Fare (e Cosa Non Può Fare)

Cosa fa:
HIcosmo è attualmente il maestro della cosmologia di fondo (background cosmology). Pensate a questo come alla misurazione del "quadro generale" dell'universo: quanto velocemente si sta espandendo, quanto sono lontane le cose e la forma complessiva dello spazio. Gestisce dati provenienti da:

Stelle esplodenti (Supernovae).
Raggruppamento delle galassie (BAO).
L'eco residuo del Big Bang (prior di distanza della CMB).
Ritardi temporali nelle lenti gravitazionali.

È stato testato contro il software che è attualmente lo standard di riferimento (chiamato Cobaya) e ha dimostrato di essere statisticamente identico nelle sue risposte, ma molto più veloce.

Su un normale processore per computer (CPU), è circa 8,7 volte più veloce del vecchio strumento.
Su una potente scheda grafica (GPU), che è ottima per gestire enormi quantità di dati contemporaneamente, può essere fino a 20 volte più veloce man mano che i dati aumentano.

Cosa non fa:
L'articolo è molto chiaro: HIcosmo al momento non gestisce i "dettagli fini" o le "perturbazioni" dell'universo, come le complesse increspature della rete cosmica o gli spettri di potenza dettagliati dell'universo primordiale. Questi richiedono un motore diverso, molto più pesante (chiamato risolutore di Boltzmann), che HIcosmo lascia ad altri strumenti. Pensate a HIcosmo come all'esperto della forma del palloncino, mentre altri strumenti sono necessari per analizzare la trama della gomma.

4. Risultati nel Mondo Reale

Gli autori non si sono limitati a costruirlo; lo hanno usato per risolvere problemi reali:

Vincoli Congiunti: Hanno combinato i dati delle supernovae, dei sondaggi galattici e della CMB per fissare con incredibile precisione il tasso di espansione dell'universo ( $H_0$ ) e la densità della materia ( $\Omega_m$ ).
Previsioni Future: Hanno usato HIcosmo per prevedere quanto bene i futuri radiotelescopi (come SKA, MeerKAT e CHIME) saranno in grado di misurare l' "equazione di stato" dell'energia oscura (un modo elegante per chiedere: "L'energia oscura è costante o cambia nel tempo?"). Hanno scoperto che questi nuovi telescopi potrebbero misurarla con un'accuratezza straordinaria.

In Sintesi

HIcosmo è un motore moderno e ad alta velocità per la cosmologia. Sostituisce i metodi lenti e manuali del passato con un sistema automatico e snello, capace di gestire le enormi inondazioni di dati provenienti dalla prossima generazione di telescopi. Non cambia le leggi della fisica, ma cambia la velocità e l'efficienza con cui possiamo testarle contro l'universo.

Sintesi Tecnica: HIcosmo – Un Framework Differenziabile Basato su JAX per l'Inferenza Cosmologica

1. Definizione del Problema

Le imminenti indagini cosmologiche di Classe IV (ad esempio, Euclid, LSST, DESI, SKA) genereranno dati osservativi di un volume senza precedenti. Sebbene l'inferenza cosmologica a livello di background (che coinvolge distanze, rapporti di distanza e moduli di distanza) operi tipicamente in spazi di parametri a bassa dimensionalità ( $d \sim 5\text{--}7$ ), il costo computazionale dell'inferenza rimane un collo di bottiglia a causa dell'enorme numero di valutazioni della verosimiglianza (likelihood) richieste dai metodi Markov Chain Monte Carlo (MCMC).

Le pipeline tradizionali (ad esempio, CosmoMC, Cobaya, MontePython) si affidano a campionatori Metropolis–Hastings (MH) a cammino casuale. Questi metodi soffrono di tempi di miscelazione (mixing) lenti che scalano linearmente con la dimensionalità ( $d$ ) e richiedono approssimazioni alle differenze finite per i gradienti, che sono computazionalmente costose ( $O(d)$ valutazioni extra) e sensibili alla scelta della dimensione del passo (step-size). Inoltre, i codici esistenti già maturi non sono implementati come programmi differenziabili, impedendo l'uso di campionatori basati sul gradiente come il No-U-Turn Sampler (NUTS) e la differenziazione automatica (AD) per le previsioni di Fisher. Sebbene siano nati strumenti basati su JAX per analisi a livello di perturbazione (ad esempio, jax-cosmo, CosmoPower), non esisteva un framework sistematicamente validato e end-to-end differenziabile per la cosmologia a livello di background che integrasse verosimiglianze, campionamento NUTS e previsioni di Fisher basate su AD all'interno di un'unica architettura.

2. Metodologia e Architettura

Il documento presenta HIcosmo (High-performance Inference for Cosmology), un framework open-source costruito interamente su JAX. Il framework è progettato attorno al principio della differenziabilità end-to-end, garantendo che ogni operazione, dai parametri cosmologici alla log-verosimiglianza scalare, sia tracciabile dal sistema di differenziazione automatica di JAX.

2.1 Architettura a Livelli

HIcosmo consiste in tre livelli primari:

Livello del Modello (Model Layer): Calcola le predizioni teoriche. Implementa quattro modelli di background: $\Lambda$ CDM, $w$ CDM, $w_0w_a$ CDM (parametrizzazione CPL) e un modello di energia oscura interagente ( $I\Lambda$ CDM). Fondamentalmente, sostituisce gli integratori adattivi (non differenziabili) con la quadratura trapezoidale su griglia fissa su griglie di redshift dense. Ciò consente il calcolo esatto del gradiente tramite la regola della catena senza ricorrere alle differenze finite.
Livello di Verosimiglianza (Likelihood Layer): Confronta le predizioni del modello con i dati osservativi. Supporta:
- Supernovae di Tipo Ia: Pantheon+, DES-SN5YR, e Union3.
- Oscillazioni Acustiche dei Barioni (BAO): DESI DR1/DR2, SDSS, BOSS, e 6dFGS.
- CMB: Prior di distanza di Planck 2018 (compressi in $R$ , $l_A$ , $\omega_b$ ).
- $H_0$ Locale: Scala delle distanze SH0ES.
- Lente Gravitazionale Forte: Ritardi temporali H0LiCOW.
  Tutte le verosimiglianze sono implementate in puro JAX, supportando la compilazione Just-In-Time (JIT) e la vettorizzazione.
Livello di Inferenza (Inference Layer): Consuma la log-verosimiglianza e le sue derivate. Fornisce:
- Campionamento NUTS: Implementato tramite NumPyro, utilizzando i gradienti direttamente da JAX per migliorare i tempi di miscelazione (scalando come $d^{1/4}$ ).
- Previsioni di Fisher (Fisher Forecasting): Calcola la matrice Hessiana tramite differenziazione automatica nidificata, eliminando la necessità di selezione della dimensione del passo e gli errori delle differenze finite.

2.2 Dettagli di Implementazione

Differenziabilità: Utilizzando l'integrazione a passo fisso e precomputando quantità indipendenti dai parametri (ad esempio, matrici di covarianza inverse), il framework garantisce che i gradienti e le matrici Hessiane siano ottenuti direttamente tramite la differenziazione automatica in modalità inversa (reverse-mode AD).
Accelerazione Hardware: Il framework sfrutta l'esecuzione unificata CPU/GPU/TPU di JAX. La compilazione JIT ottimizza la verosimiglianza in codice macchina, bypassando l'overhead dell'interprete Python durante le corse MCMC.
Calibrazione: Per garantire l'accuratezza fisica, il framework calibra le quantità derivate (come l'orizzonte sonoro $r_d$ ) rispetto al solutore di Boltzmann CAMB, riducendo i residui a circa lo $0,01\%$ .

3. Risultati Chiave

3.1 Validazione e Accuratezza

HIcosmo è stato sistematicamente validato contro implementazioni di riferimento (Cobaya) e codici di verosimiglianza ufficiali attraverso modelli $\Lambda$ CDM, $w$ CDM e $w_0w_a$ CDM.

Precisione della Verosimiglianza: I valori di $\chi^2$ concordano con i codici di riferimento con differenze assolute da $10^{-6}$ a $10^{-2}$ .
Vincoli sui Parametri: I vincoli marginalizzati da HIcosmo e Cobaya differiscono per meno di 0,2 $\sigma$ in tutte le analisi testate. Per le analisi congiunte SN+BAO+CMB, l'accordo è entro 0,04 $\sigma$ .
Consistenza: Il framework riproduce i vincoli stabiliti, come $H_0 = 68,24 \pm 0,38$ km/s/Mpc e $\Omega_m = 0,307 \pm 0,005$ per l'analisi congiunta, confermando che le approssimazioni differenziabili non introducono bias significativi.

3.2 Benchmark delle Prestazioni

Velocità rispetto a Scipy: Per i calcoli delle distanze, HIcosmo ottiene un'accelerazione di 140 $\times$ rispetto alla quadratura adattiva di scipy dopo la compilazione JIT. Per l'intera pipeline di verosimiglianza Pantheon+, è 16 $\times$ più veloce su CPU.
Throughput di Campionamento: Rispetto a Cobaya (utilizzando campionatori MH) su una CPU a 8 core, HIcosmo (utilizzando NUTS) raggiunge un throughput di dimensione effettiva del campione (ESS) di 8,7 $\times$ superiore (400 ESS/s contro 45,8 ESS/s). Ciò è dovuto sia alla valutazione più veloce della verosimiglianza, sia alla migliore miscelazione di NUTS ( $\tau \approx 1,2$ contro $\tau \approx 11,4$ ).
Accelerazione GPU: Sebbene le singole valutazioni su piccoli dataset (Pantheon+) mostrino guadagni GPU modesti a causa dell'overhead, l'accelerazione GPU scala con la dimensione del dataset. Per dataset sintetici con $N_{SN} = 20.000$ , la GPU ottiene un'accelerazione di 20,8 $\times$ rispetto alla CPU.
Matrice di Fisher: Il calcolo della matrice di Fisher basato su AD è stabile e non richiede la sintonizzazione della dimensione del passo, evitando gli errori di arrotondamento in virgola mobile che affliggono i metodi a differenze finite a piccole dimensioni di passo.

3.3 Applicazioni

Il documento dimostra l'utilità del framework attraverso:

Vincoli Congiunti Multi-Sonda: Combinando Pantheon+, DESI BAO e i prior di Planck per stringere i vincoli su $H_0$ e $\Omega_m$ .
Vincoli sull'Energia Oscura: Derivando vincoli su $w_0$ e $w_a$ coerenti con i risultati DESI 2024.
Previsioni per la Mappatura dell'Intensità a 21 cm: Eseguendo previsioni di Fisher per sei survey (SKA1, MeerKAT, BINGO, Tianlai, CHIME). Il framework prevede che SKA1 Band 1 potrebbe vincolare l'equazione di stato dell'energia oscura a $\sigma(w) \approx 0,046$ .

4. Significato e Rivendicazioni

Gli autori affermano che HIcosmo colma una lacuna critica nel panorama dell'inferenza cosmologica fornendo la prima pipeline per la cosmologia di background, sistematicamente validata ed end-to-end differenziabile e nativa in JAX.

Efficienza: Dimostra che il campionamento basato sul gradiente (NUTS), combinato con le verosimiglianze compilate JIT, può superare significativamente il tradizionale MCMC a cammino casuale in termini di velocità computazionale ed efficienza statistica (ESS).
Scalabilità: Il framework è unicamente posizionato per gestire i massicci dataset delle future survey (Classe IV), sfruttando l'accelerazione GPU, dove i miglioramenti di velocità possono raggiungere le 20 $\times$ .
Modularità ed Estensibilità: Separando i livelli di modello, verosimiglianza e inferenza, il framework permette l'aggiunta di nuovi modelli o dataset senza dover riscrivere il codice delle derivate.
Fondazione per l'Inferenza a Livello di Perturbazione: Sebbene attualmente limitato alla cosmologia di background (escludendo i solutori di Boltzmann completi), il documento sostiene che la validazione end-to-end di successo a livello di background stabilisca una base robusta. L'architettura è progettata per essere estesa alle analisi a livello di perturbazione, sostituendo il livello del modello con emulatori di Boltzmann differenziabili (ad esempio, CosmoPower-JAX), abilitando così l'inferenza basata sul gradiente per problemi ad alta dimensionalità di Classe IV.

In sintesi, HIcosmo offre un'alternativa riproducibile e ad alte prestazioni agli strumenti esistenti, provando che la differenziazione automatica e l'accelerazione hardware moderna possono essere applicate efficacemente ai compiti standard di inferenza cosmologica.