⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

Dit artikel introduceert HIcosmo, een open-source, op JAX gebaseerd framework voor efficiënte kosmologische inferentie dat gebruikmaakt van automatische differentiatie en GPU-acceleratie om aanzienlijk snellere sample-snelheden te bereiken dan bestaande tools zoals Cobaya, terwijl het een breed scala aan Stage IV survey-likelihoods en kosmologische modellen ondersteunt.

Oorspronkelijke auteurs: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Gepubliceerd 2026-06-29

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je het universum voor als een gigantische, complexe puzzel die wetenschappers proberen op te lossen. Decennia lang hebben ze een standaard set regels gebruikt (de $\Lambda$ CDM-modellen) om te begrijpen hoe het universum begon, hoe het uitdijt en waar het uit bestaat. Maar onlangs passen de stukjes van de puzzel niet meer perfect in elkaar. Metingen uit het vroege universum (zoals de kosmische achtergrondstraling) komen niet overeen met metingen uit het nabijgelegen universum (zoals exploderende sterren genaamd supernova's). Dit staat bekend als de "Hubble-spanning" (Hubble Tension).

Om dit op te lossen, worden er nieuwe, enorme telescopen gebouwd die miljarden datapunten zullen verzamelen. Het probleem? De computerprogramma's die wetenschappers gebruiken om deze puzzels op te lossen, zijn te traag en lomp om deze nieuwe vloedgolf aan informatie aan te kunnen. Ze zijn als proberen zandkorrels op een strand te tellen met een enkele, trage rekenmachine.

Dit artikel introduceert HIcosmo, een gloednieuwe, supersnelle computertool die ontworft is om deze kosmische puzzels efficiënt op te lossen. Hier is hoe het werkt, eenvoudig uitgelegd:

1. De "Slimme Rekenmachine" (JAX en Differentieerbaarheid)

Traditionele tools zijn als een chef die elke ingrediënt handmatig moet afwegen, het resultaat moet opschrijven en vervolgens handmatig de volgende stap moet berekenen. Als hij wil weten hoe het veranderen van één ingrediënt de uiteindelijke smaak beïnvloedt, moet hij gissen en controleren, wat eeuwen duurt.

HIcosmo is anders. Het is gebouwd op een technologie genaamd JAX, die fungeert als een "slimme rekenmachine" die het volledige recept in één keer begrijpt.

Automatische Differentiatie: In plaats van te gissen, weet HIcosmo precies hoe het resultaat verandert als je één getal aanpast. Het is alsoals een GPS die je direct de beste route vertelt zonder dat je elke mogelijke weg hoeft te rijden om te zien welke het snelst is.
Just-in-Time Compilatie: De eerste keer dat je HIcosmo vraagt een berekening uit te voeren, duurt het een fractie van een seconde om op te warmen en de instructies te vertalen naar super-geoptimaliseerde machinecode. Daarna werkt het met de snelheid van het licht, waarbij de trage "vertaler" (de Python-interpreter) die oudere tools vertraagt, wordt omzeild.

2. De Drielagen Taart

De auteurs hebben HIcosmo gebouwd als een drielagen taart, waarbij elke laag een specifieke taak heeft:

De Onderste Laag (Het Model): Dit is het "theorie"-gedeelte. Het neemt de regels van het universum (zoals hoeveel donkere materie of donkere energie er is) en berekent hoe het universum er zou moeten uitzien.
De Middelste Laag (De Likelihood): Dit is het "vergelijkings"-gedeelte. Het neemt de theorie van de onderste laag en vergelijkt deze met echte gegevens van telescopen (zoals de Pantheon+ supernova-data of DESI-sterrenstelselkaarten). Het berekent hoe "fout" de theorie is.
De Bovenste Laag (De Inferentie): Dit is het "oplossings"-gedeelte. Het gebruikt de informatie uit de middelste laag om het best mogelijke antwoord te vinden. Het maakt gebruik van een slimme methode genaamd NUTS (No-U-Turn Sampler), die de gradiënten van de "slimme rekenmachine" gebruikt om rechtstreeks naar het juiste antwoord te zoomen, in plaats van doelloos rond te dwalen zoals oudere tools doen.

3. Wat het kan (en wat het niet kan)

Wat het doet:
HIcosmo is momenteel de meester van de achtergrondkosmologie. Denk hierbij aan het meten van het "grote plaatje" van het universum: hoe snel het uitdijt, hoe ver dingen weg zijn en de algemene vorm van de ruimte. Het verwerkt data van:

Exploderende sterren (Supernova's).
Sterrenstelselclustering (BAO).
De nagloed van de oerknal (CMB afstandsprioriteiten).
Tijdvertragingen in gravitationele lenzen.

Het is getest tegen de huidige gouden standaard software (genaamd Cobaya) en bleek statistisch identiek in de antwoorden, maar veel sneller.

Op een standaard computerprocessor (CPU) is het ongeveer 8,7 keer sneller dan de oude tool.
Op een krachtige grafische kaart (GPU), die geweldig is in het tegelijkertijd verwerken van enorme hoeveelheden data, kan het tot wel 20 keer sneller zijn naarmate de data groter wordt.

Wat het niet doet:
De tekst is heel duidelijk dat HIcosmo momenteel de "fijne details" of "perturbaties" van het universum niet afhandelt, zoals de complexe rimpelingen in het kosmische web of de gedetailleerde vermogensspectra van het vroege universum. Hiervoor zijn andere, zeer zware motoren nodig (een zogenaamde Boltzmann-solver) die HIcosmo aan andere tools overlaat. Zie HIcosmo als de expert op de vorm van de ballon, terwijl andere tools nodig zijn om de textuur van het rubber te analyseren.

4. Resultaten in de Praktijk

De auteurs hebben het niet alleen gebouwd; ze hebben het gebruikt om echte problemen op te lossen:

Gecombineerde Constrainingen: Ze combineerden data van supernova's, sterrenstelselmetingen en de CMB om de expansiesnelheid van het universum ( $H_0$ ) en de materiedichtheid ( $\Omega_m$ ) met ongelooflijke precisie vast te leggen.
Toekomstvoorspellingen: Ze gebruikten HIcosmo om te voorspellen hoe goed toekomstige radiotelescopen (zoals SKA, MeerKAT en CHIME) de "toestandsvergelijking" van donkere energie kunnen meten (een chique manier om te vragen: "Is donkere energie constant, of verandert het in de loop van de tijd?"). Ze ontdekten dat deze nieuwe telescopen dit met verbazingwekkende nauwkeurigheid kunnen meten.

De Kernboodschap

HIcosmo is een moderne, hogesnelheidsmotor voor de kosmologie. Het vervangt de trage, handmatige methoden uit het verleden door een gestroomlijnd, automatisch systeem dat de enorme datastromen van de volgende generatie telescopen aan kan. Het verandert de natuurwetten niet, maar het verandert de snelheid en efficiëntie waarmee we ze kunnen testen tegen het universum.

Technische Samenvatting: HIcosmo – Een Differentieerbaar JAX-gebaseerd Framework voor Kosmologische Inferentie

1. Probleemstelling

De aanstaande Stage IV kosmologische surveys (bijv. Euclid, LSST, DESI, SKA) zullen observaties genereren van ongekende omvang. Hoewel inferentie op achtergrondniveau (betrekking hebbend op afstanden, afstandratio's en afstandmoduli) doorgaans opereert in laagdimensionale parameterruimtes ( $d \sim 5\text{--}7$ ), blijft de computationele kost van inferentie een bottleneck vanwege de enorme hoeveelheid vereiste likelihood-evaluaties door Markov Chain Monte Carlo (MCMC)-methoden.

Traditionele pipelines (bijv. CosmoMC, Cobaya, MontePython) vertrouwen op random-walk Metropolis–Hastings (MH) samplers. Deze methoden lijden onder trage mengtijden die lineair schalen met de dimensionaliteit ( $d$ ) en vereisen eindverschil-benaderingen voor gradiënten, wat computationeel duur is ( $O(d)$ extra evaluaties) en gevoelig is voor de keuze van de stapgrootte. Bovendien zijn bestaande volwassen codes niet geïmplementeerd als differentieerbare programma's, wat het gebruik van gradiënt-gebaseerde samplers zoals de No-U-Turn Sampler (NUTS) en automatische differentiatie (AD) voor Fisher-voorspellingen verhindert. Ho hoewel JAX-gebaseerde tools zijn opgekomen voor perturbatie-niveau analyses (bijv. jax-cosmo, CosmoPower), was er geen systematisch gevalideerd, end-to-end differentieerbaar framework voor achtergrond-niveau kosmologie dat likelihoods, NUTS-sampling en AD-gebaseerde Fisher-voorspellingen binnen één architectuur integreert.

2. Methodologie en Architectuur

Het artikel presenteert HIcosmo (High-performance Inference for Cosmology), een open-source framework dat volledig is gebouwd op JAX. Het framework is ontworpen rond het principe van end-to-end differentieerbaarheid, waarbij wordt gewaarborgd dat elke operatie van kosmologische parameters tot de scalaire log-likelihood traceerbaar is door JAX's systeem voor automatische differentiatie.

2.1 Gelaagde Architectuur

HIcosmo bestaat uit drie primaire lagen:

Model-laag: Berekent theoretische voorspellingen. Het implementeert vier achtergrondmodellen: $\Lambda$ CDM, $w$ CDM, $w_0w_a$ CDM (CPL-parametrisatie), en een interagerend donkere energie-model ( $I\Lambda$ CDM). Cruciaal is dat het adaptieve integratoren (die niet-differentieerbaar zijn) vervangt door fixed-grid trapezoidal quadrature op dichte roodverschuivingsroosters. Dit maakt exacte gradiëntberekening mogelijk via de kettingregel zonder eindverschillen.
Likelihood-laag: Vergelijkt modelvoorspellingen met observationele data. Het ondersteunt:
- Type Ia Supernovae: Pantheon+, DES-SN5YR, en Union3.
- Baryon Acoustic Oscillations (BAO): DESI DR1/DR2, SDSS, BOSS, en 6dFGS.
- CMB: Planck 2018 afstandspriors (gecomprimeerd in $R$ , $l_A$ , $\omega_b$ ).
- Lokale $H_0$ : SH0ES afstandsladder.
- Sterke Lenswerking: H0LiCOW tijdvertragingen.
- Alle likelihoods zijn geïmplementeerd in pure JAX, wat Just-In-Time (JIT) compilatie en vectorisatie ondersteunt.
Inferentie-laag: Consumeert de log-likelihood en de bijbehorende afgeleiden. Het biedt:
- NUTS Sampling: Geïmplementeerd via NumPyro, waarbij direct gebruik wordt gemaakt van gradiënten uit JAX om mengtijden te verbeteren (schaling als $d^{1/4}$ ).
- Fisher Forecasting: Berekent de Hessian-matrix via geneste automatische differentiatie, waardoor de noodzaak voor stapgrootte-selectie en eindverschil-fouten wordt geëlimineerd.

2.2 Implementatiedetails

Differentieerbaarheid: Door het gebruik van fixed-step integratie en het voorberekenen van parameter-onafhankelijke grootheden (zoals inverse covariantie-matrices), garandeert het framework dat gradiënten en Hessianen direct worden verkregen via reverse-mode AD.
Hardware Acceleratie: Het framework maakt gebruik van JAX's uniforme CPU/GPU/TPU executie. JIT-compilatie optimaliseert de likelihood tot machinecode, waardoor de Python-interpreter overhead tijdens MCMC-runs wordt omzeild.
Kalibratie: Om fysieke nauwkeurigheid te garanderen, kalibreert het framework afgeleide grootheden (zoals de geluidshorizon $r_d$ ) tegen de Boltzmann-solver CAMB, waarbij residuen worden teruggebracht tot $\sim 0.01\%$ .

3. Belangrijkste Resultaten

3.1 Validatie en Nauwkeurigheid

HIcosmo werd systematisch gevalideerd tegen referentie-implementaties (Cobaya) en officiële likelihood-codes over $\Lambda$ CDM, $w$ CDM, en $w_0w_a$ CDM modellen.

Likelihood Precisie: De $\chi^2$ -waarden komen overeen met referentiecodes tot absolute verschillen van $10^{-6}$ tot $10^{-2}$ .
Parameter Beperkingen: Gemarginaliseerde beperkingen van HIcosmo en Cobaya verschillen door minder dan 0.2 $\sigma$ in alle geteste analyses. Voor gezamenlijke SN+BAO+CMB analyses is de overeenkomst binnen 0.04 $\sigma$ .
Consistentie: Het framework reproduceert gevestigde beperkingen, zoals $H_0 = 68.24 \pm 0.38$ km/s/Mpc en $\Omega_m = 0.307 \pm 0.005$ voor de gezamenlijke analyse, wat bevestigt dat de differentieerbare benaderingen geen significante bias introduceren.

3.2 Prestatiebenchmarks

Versnelling ten opzichte van Scipy: Voor afstandsberekeningen bereikt HIcosmo een 140 $\times$ versnelling ten opzichte van scipy adaptieve kwadratuur na JIT-compilatie. Voor de volledige Pantheon+ likelihood pipeline is het 16 $\times$ sneller op CPU.
Sampling Doorvoer: Vergeleken met Cobaya (gebruikmakend van MH samplers) op een 8-core CPU, bereikt HIcosmo (gebruikmakend van NUTS) een 8.7 $\times$ hogere Effective Sample Size (ESS) doorvoer (400 ESS/s vs. 45.8 ESS/s). Dit komt door zowel de snellere likelihood-evaluatie als de superieure mixing van NUTS ( $\tau \approx 1.2$ vs. $\tau \approx 11.4$ ).
GPU Acceleratie: Hoewel single evaluaties op kleine datasets (Pantheon+) bescheiden GPU-winsten laten zien door overhead, schaalt de GPU-acceleratie met de datasetgrootte. Voor synthetische datasets met $N_{SN} = 20,000$ , bereikt de GPU een 20.8 $\times$ versnelling over de CPU.
Fisher Matrix: AD-gebaseerde Fisher-matrix berekening is stabiel en vereist geen stapgrootte-tuning, waardoor floating-point afrondingsfouten die eindverschil-methoden bij kleine stapgroottes teisteren, worden vermeden.

3.3 Toepassingen

Het artikel demonstreert de bruikbaarheid van het framework via:

Multi-probe Joint Constraints: Het combineren van Pantheon+, DESI BAO, en Planck-priors om beperkingen op $H_0$ en $\Omega_m$ aan te scherpen.
Donkere Energie Beperkingen: Het afleiden van beperkingen op $w_0$ en $w_a$ die consistent zijn met de DESI 2024 resultaten.
21 cm Intensity Mapping Voorspellingen: Het uitvoeren van Fisher-voorspellingen voor zes surveys (SKA1, MeerKAT, BINGO, Tianlai, CHIME). Het framework voorspelt dat SKA1 Band 1 de donkere energie toestandsvergelijking kan beperken tot $\sigma(w) \approx 0.046$ .

4. Betekenis en Claims

De auteurs claimen dat HIcosmo een kritieke kloof vult in het landschap van kosmologische inferentie door het eerste systematisch gevalideerde, end-to-end differentieerbare JAX-native pipeline voor achtergrond-kosmologie te bieden.

Efficiëntie: Het toont aan dat gradiënt-gebaseerde sampling (NUTS) gecombineerd met JIT-gecompileerde likelihoods aanzienlijk kan presteren ten opzichte van traditionele random-walk MCMC in termen van zowel computationele snelheid als statistische efficiëntie (ESS).
Schaalbaarheid: Het framework is uniek gepositioneerd om de enorme datasets van toekomstige surveys (Stage IV) aan te kunnen door gebruik te maken van GPU-acceleratie, waarbij versnellingen van 20 $\times$ over CPU-executie kunnen worden bereikt.
Modulariteit en Uitbreidbaarheid: Door de model-, likelihood- en inferentielagen te scheiden, maakt het framework het eenvoudig om nieuwe modellen of datasets toe te voegen zonder de afgeleide code te herschrijven.
Fundament voor Perturbatie-niveau Inferentie: Hoewel momenteel beperkt tot achtergrond-kosmologie (exclusief volledige Boltzmann-solvers), betogen de auteurs dat de succesvolle end-to-end validatie op achtergrondniveau een robuust fundament legt. De architectuur is ontworpen om uitgebreid te worden naar perturbatie-niveau analyses door de model-laag te vervangen door differentieerbare Boltzmann-emulators (bijv. CosmoPower-JAX), wat gradiënt-gebaseerde inferentie voor hoogdimensionale Stage IV problemen mogelijk maakt.

Samenvattend biedt HIcosmo een reproduceerbaar, hoogwaardig alternatief voor bestaande tools, waarmee wordt bewezen dat automatische differentiatie en moderne hardware-acceleratie effectief kunnen worden toegepast op standaard kosmologische inferentie-taken.