⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

Cet article présente HIcosmo, un cadre open-source basé sur JAX pour l'inférence cosmologique efficace qui exploite la différenciation automatique et l'accélération GPU pour atteindre des vitesses d'échantillonnage nettement plus rapides que les outils existants comme Cobaya, tout en prenant en charge une large gamme de vraisemblances de relevés de l'étape IV et de modèles cosmologiques.

Auteurs originaux : Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Publié 2026-06-29

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez l'univers comme un puzzle géant et complexe que les scientifiques tentent de résoudre. Depuis des décennies, ils utilisent un ensemble de règles standards (appelé le modèle $\Lambda$ CDM) pour comprendre comment l'univers a commencé, comment il s'étend et de quoi il est composé. Mais récemment, les pièces du puzzle ne s'emboîtent plus parfaitement. Les mesures de l'univers primordial (comme le Fond Diffus Cosmique Micro-onde) ne concordent pas avec les mesures de l'univers proche (comme les étoiles explosives appelées Supernovae). C'est ce qu'on appelle la « tension de Hubble ».

Pour résoudre cela, de nouveaux télescopes massifs sont en cours de construction pour collecter des milliards de points de données. Le problème ? Les programmes informatiques que les scientifiques utilisent pour résoudre ces énigmes sont trop lents et encombrants pour gérer ce nouveau déluge d'informations. C'est comme essayer de compter les grains de sable sur une plage à l'aide d'une seule calculatrice lente.

Cet article présente HIcosmo, un tout nouvel outil informatique ultra-rapide conçu pour résoudre ces puzzles cosmiques efficacement. Voici comment il fonctionne, expliqué simplement :

1. La « Calculatrice Intelligente » (JAX et la Différentiabilité)

Les outils traditionnels sont comme un chef qui doit mesurer manuellement chaque ingrédient, noter le résultat, puis calculer manuellement l'étape suivante. S'il veut savoir comment le changement d'un seul ingrédient affecte le goût final, il doit deviner et tester, ce qui prend un temps infini.

HIcosmo est différent. Il est construit sur une technologie appelée JAX, qui agit comme une « calculatrice intelligente » capable de comprendre l'intégralité de la recette d'un coup.

Différenciation Automatique : Au lieu de deviner, HIcosmo sait exactement comment le résultat change si vous modifiez un seul chiffre. C'est comme avoir un GPS qui vous indique instantanément le meilleur itinéraire sans avoir besoin de parcourir toutes les routes possibles pour voir laquelle est la plus rapide.
Compilation Juste-à-Temps (Just-in-Time) : La première fois que vous demandez à HIcosment d'effectuer un calcul, il prend une fraction de seconde pour « s'échauffer » et traduire les instructions en un code machine super-optimisé. Après cela, il fonctionne à la vitesse de l'éclair, contournant le « traducteur » lent (l'interprète Python) qui ralentit les outils plus anciens.

2. Le Gâteau à Trois Couches

Les auteurs ont construit HIcosmo comme un gâteau à trois couches, où chaque couche a un travail spécifique :

La Couche Inférieure (Le Modèle) : C'est la partie « théorie ». Elle prend les règles de l'univers (comme la quantité de matière noire ou d'énergie noire) et calcule à quoi l'univers devrait ressembler.
La Couche Médiane (La Vraisemblance) : C'est la partie « comparaison ». Elle prend la théorie de la couche inférieure et la compare aux données réelles provenant des télescopes (comme les données de supernovae Pantheon+ ou les cartes de galaxies DESI). Elle calcule à quel point la théorie est « fausse ».
La Couche Supérieure (L'Inférence) : C'est la partie « solveur ». Elle utilise les informations de la couche médiane pour trouver la meilleure réponse possible. Elle utilise une méthode intelligente appelée NUTS (No-U-Turn Sampler) qui utilise les gradients de la « calculatrice intelligente » pour foncer directement vers la bonne réponse, plutôt que de errer sans but comme le font les outils plus anciens.

3. Ce qu'il peut faire (et ce qu'il ne peut pas faire)

Ce qu'il fait :
HIcosmo est actuellement le maître de la cosmologie de fond (background cosmology). Voyez cela comme la mesure de la « vue d'ensemble » de l'univers : la vitesse à laquelle il s'étend, la distance des objets, et la forme globale de l'espace. Il gère les données de :

Étoiles explosives (Supernovae).
Regroupement de galaxies (BAO).
Le rayonnement fossile du Big Bang (priors de distance du CMB).
Délais temporels dans les lentilles gravitationnelles.

Il a été testé par rapport au logiciel de référence actuel (appelé Cobaya) et s'est révélé statistiquement identique dans ses réponses, mais beaucoup plus rapide.

Sur un processeur standard (CPU), il est environ 8,7 fois plus rapide que l'ancien outil.
Sur une carte graphique puissante (GPU), qui est idéale pour gérer de grandes quantités de données simultanément, il peut être jusqu'à 20 fois plus rapide à mesure que les données augmentent.

Ce qu'il ne fait pas :
L'article précise très clairement que HIcosmo ne gère pas actuellement les « détails fins » ou les « perturbations » de l'univers, telles que les ondulations complexes de la toile cosmique ou les spectres de puissance détaillés de l'univers primordial. Ceux-ci nécessitent un moteur très lourd et différent (appelé solveur de Boltzmann) que HIcosmo laisse à d'autres outils. Considérez HIcosmo comme l'expert sur la forme du ballon, tandis que d'autres outils sont nécessaires pour analyser la texture du caoutchouc.

4. Résultats en Conditions Réelles

Les auteurs ne se sont pas contentés de le construire ; ils l'ont utilisé pour résoudre de vrais problèmes :

Contraintes Conjointes : Ils ont combiné les données de supernovae, de relevés de galaxies et du CMB pour fixer avec une précision incroyable le taux d'expansion de l'univers ( $H_0$ ) et la densité de la matière ( $\Omega_m$ ).
Prévisions Futures : Ils ont utilisé HIcosmo pour prédire avec quelle précision les futurs télescopes radio (comme SKA, MeerKAT et CHIME) pourront mesurer « l'équation d'état » de l'énergie noire (une façon élégante de demander : « L'énergie noire est-elle constante, ou change-t-elle au fil du temps ? »). Ils ont découvert que ces nouveaux télescopes pourraient mesurer cela avec une précision extraordinaire.

L'essentiel

HIcosmo est un moteur moderne et ultra-rapide pour la cosmologie. Il remplace les méthodes manuelles et lentes du passé par un système automatisé et rationalisé, capable de gérer les flux massifs de données provenant de la prochaine génération de télescopes. Il ne change pas les lois de la physique, mais il change la rapidité et l'efficacité avec lesquelles nous pouvons les tester face à l'univers.

Résumé Technique : HIcosmo – Un cadre de calcul de l'inférence cosmologique différentiable basé sur JAX

1. Énoncé du problème

Les prochains relevés cosmologiques de l'étape IV (par exemple, Euclid, LSST, DESI, SKA) généreront des données observationnelles d'un volume sans précédent. Bien que l'inférence cosmologique au niveau du fond (impliquant les distances, les rapports de distances et les modules de distance) opère généralement dans des espaces de paramètres de faible dimension ( $d \sim 5\text{--}7$ ), le coût computationnel de l'inférence reste un goulot d'étranglement en raison du nombre colossal d'évaluations de vraisemblance requises par les méthodes de Monte Carlo par chaînes de Markov (MCMC).

Les pipelines traditionnels (par exemple, CosmoMC, Cobaya, MontePython) reposent sur des échantillonneurs de marche aléatoire de Metropolis-Hastings (MH). Ces méthodes souffrent de temps de mélange lents qui augmentent linéairement avec la dimensionnalité ( $d$ ) et nécessitent des approximations par différences finies pour les gradients, ce qui est coûteux en termes de calcul ( $O(d)$ évaluations supplémentaires) et sensible au choix de la taille du pas. De plus, les codes matures existants ne sont pas implémentés comme des programmes différentiables, ce qui empêche l'utilisation d'échantillonneurs basés sur le gradient comme l'échantillonneur No-U-Turn (NUTS) et la différenciation automatique (AD) pour les prévisions de Fisher. Bien que des outils basés sur JAX aient émergé pour les analyses au niveau des perturbations (par exemple, jax-cosmo, CosmoPower), il n'existait pas de cadre systématiquement validé et de bout en bout pour la cosmologie de fond différentiable intégrant les vraisemblances, l'échantillonnage NUTS et les prévisions de Fisher basées sur l'AD au sein d'une architecture unique.

2. Méthodologie et Architecture

Le document présente HIcosmo (High-performance Inference for Cosmology), un cadre open-source entièrement construit sur JAX. Le cadre est conçu autour du principe de la différentiabilité de bout en bout, garantissant que chaque opération, des paramètres cosmologiques au log-vraisemblance scalaire, est traçable par le système de différenciation automatique de JAX.

2.1 Architecture en couches

HIcosmo se compose de trois couches principales :

Couche de Modèle : Calcule les prédictions théoriques. Elle implémente quatre modèles de fond : $\Lambda$ CDM, $w$ CDM, $w_0w_a$ CDM (paramétrisation CPL) et un modèle d'énergie noire interactive ( $I\Lambda$ CDM). Crucialement, elle remplace les intégrateurs adaptatifs (qui sont non différentiables) par une quadrature de type trapèze sur grille fixe sur des grilles de redshift denses. Cela permet un calcul exact du gradient via la règle de la chaîne sans recourir aux différences finies.
Couche de Vraisemblance : Compare les prédictions du modèle avec les données observationnelles. Elle supporte :
- Supernovae de Type Ia : Pantheon+, DES-SN5YR, et Union3.
- Oscillations Acoustiques de Baryons (BAO) : DESI DR1/DR2, SDSS, BOSS, et 6dFGS.
- CMB : Priors de distance de Planck 2018 (compressés en $R$ , $l_A$ , $\omega_b$ ).
- $H_0$ locale : Échelle de distance SH0ES.
- Lentillage Fort : Délais temporels H0LiCOW.
- Toutes les vraisemblances sont implémentées en pur JAX, supportant la compilation Just-In-Time (JIT) et la vectorisation.
Couche d'Inférence : Consomme le log-vraisemblance et ses dérivées. Elle fournit :
- Échantillonnage NUTS : Implémenté via NumPyro, utilisant les gradients directement issus de JAX pour améliorer les temps de mélange (passant d'une échelle de $d$ à $d^{1/4}$ ).
- Prévision de Fisher : Calcule la matrice Hessienne via une différenciation automatique imbriquée, éliminant ainsi le besoin de sélection de la taille du pas et les erreurs de différences finies.

2.2 Détails d'implémentation

Différentiabilité : En utilisant une intégration à pas fixe et en précalculant les quantités indépendantes des paramètres (ex: matrices de covariance inverses), le cadre garantit que les gradients et les Hessiennes sont obtenus directement via l'AD en mode inverse.
Accélération Matérielle : Le cadre exploite l'exécution unifiée CPU/GPU/TPU de JAX. La compilation JIT optimise la vraisemblance en code machine, contournant le surcoût de l'interprète Python lors des exécutions MCMC.
Calibration : Pour garantir la précision physique, le cadre calibre les quantités dérivées (comme l'horizon sonore $r_d$ ) par rapport au solveur de Boltzmann CAMB, réduisant les résidus à environ $0,01\%$ .

3. Résultats Clés

3.1 Validation et Précision

HIcosmo a été systématiquement validé par rapport à des implémentations de référence (Cobaya) et des codes de vraisemblance officiels à travers les modèles $\Lambda$ CDM, $w$ CDM et $w_0w_a$ CDM.

Précision de la Vraisemblance : Les valeurs de $\chi^2$ concordent avec les codes de référence avec des différences absolues de $10^{-6}$ à $10^{-2}$ .
Contraintes de Paramètres : Les contraintes marginalisées de HIcosmo et Cobaya diffèrent de moins de 0,2 $\sigma$ dans toutes les analyses testées. Pour les analyses conjointes SN+BAO+CMB, l'accord est de 0,04 $\sigma$ .
Cohérence : Le cadre reproduit les contraintes établies, telles que $H_0 = 68,24 \pm 0,38$ km/s/Mpc et $\Omega_m = 0,307 \pm 0,005$ pour l'analyse conjointe, confirmant que les approximations différentiables n'introduisent pas de biais significatif.

3.2 Benchmarks de Performance

Accélération par rapport à Scipy : Pour les calculs de distance, HIcosmo atteint une accélération de 140 $\times$ par rapport à la quadrature adaptative de scipy après compilation JIT. Pour le pipeline complet de vraisemblance Pantheon+, il est 16 $\times$ plus rapide sur CPU.
Débit d'Échantillonnage : Comparé à Cobaya (utilisant des échantillonneurs MH) sur un CPU à 8 cœurs, HIcosmo (utilisant NUTS) atteint un débit de Taille d'Échantillon Effective (ESS) 8,7 $\times$ plus élevé (400 ESS/s contre 45,8 ESS/s). Ceci est dû à la fois à l'évaluation plus rapide de la vraisemblance et au meilleur mélange de NUTS ( $\tau \approx 1,2$ contre $\tau \approx 11,4$ ).
Accélération GPU : Bien que les évaluations uniques sur de petits jeux de données (Pantheon+) montrent des gains GPU modestes dus au surcoût, l'accélération GPU augmente avec la taille du jeu de données. Pour des jeux de données synthétiques avec $N_{SN} = 20\,000$ , le GPU atteint une accélération de 20,8 $\times$ par rapport au CPU.
Matrice de Fisher : Le calcul de la matrice de Fisher par AD est stable et ne nécessite aucun réglage de la taille du pas, évitant les erreurs d'arrondi de virgule flottante qui affectent les méthodes de différences finies à de petits pas.

3.3 Applications

Le document démontre l'utilité du cadre à travers :

Contraintes Conjointes Multi-sondes : Combiner Pantheon+, DESI BAO et les priors de Planck pour resserrer les contraintes sur $H_0$ et $\Omega_m$ .
Contraintes sur l'Énergie Noire : Dériver des contraintes sur $w_0$ et $w_a$ cohérentes avec les résultats de DESI 2024.
Prévisions de l'Imagerie de l'Intensité de la raie 21 cm : Réaliser des prévisions de Fisher pour six relevés (SKA1, MeerKAT, BINGO, Tianlai, CHIME). Le cadre prédit que SKA1 Band 1 pourrait contraindre l'équation d'état de l'énergie noire à $\sigma(w) \approx 0,046$ .

4. Signification et Revendications

Les auteurs affirment que HIcosmo comble une lacune critique dans le paysage de l'inférence cosmologique en fournissant le premier pipeline de cosmologie de fond, systématiquement validé, de bout en bout et natif de JAX.

Efficacité : Il démontre que l'échantillonnage basé sur le gradient (NUTS) combiné aux vraisemblances compilées JIT peut nettement surpasser le MCMC traditionnel par marche aléatoire en termes de vitesse de calcul et d'efficacité statistique (ESS).
Scalabilité : Le cadre est idéalement positionné pour gérer les jeux de données massifs des futurs relevés (Étape IV) en exploitant l'accélération GPU, où les gains de vitesse peuvent atteindre 20 $\times$ par rapport à l'exécution CPU.
Modularité et Extensibilité : En séparant les couches de modèle, de vraisemblance et d'inférence, le cadre permet l'ajout facile de nouveaux modèles ou jeux de données sans réécrire le code de dérivation.
Fondation pour l'Inférence au Niveau des Perturbations : Bien qu'actuellement limité à la cosmologie de fond (excluant les solveurs de Boltzmann complets), le document soutient que la validation réussie de bout en bout au niveau du fond établit une base robuste. L'architecture est conçue pour être étendue aux analyses de niveau de perturbation en remplaçant la couche de modèle par des émulateurs de Boltzmann différentiables (par exemple, CosmoPower-JAX), permettant ainsi l'inférence par gradient pour les problèmes de haute dimension de l'Étape IV.

En résumé, HIcosmo offre une alternative performante et reproductible aux outils existants, prouvant que la différenciation automatique et l'accélération matérielle moderne peuvent être appliquées efficacement aux tâches standards d'inférence cosmologique.