Efficiently emulating distribution functions in gigaparsec volumes for varying cosmological parameters

1. Problématique

Les simulations cosmologiques de type $N$-body sont essentielles pour comprendre la formation des structures de l'Univers, en particulier la fonction de masse des halos (HMF). Cependant, un compromis fondamental existe entre le volume simulé, la résolution et la couverture des paramètres cosmologiques :

• Volume vs Résolution : Les simulations périodiques de grande volume (Gigaparsecs) nécessaires pour capturer les halos massifs (amas de galaxies) et les statistiques d'ordre élevé souffrent d'une résolution limitée, empêchant l'étude des halos de faible masse. À l'inverse, les simulations hydrodynamiques haute résolution sont limitées à de petits volumes.
• Coût des paramètres : Pour contraindre les modèles cosmologiques avec la précision requise par les futurs relevés (Euclid, Roman, Rubin), il est nécessaire d'explorer un vaste espace de paramètres cosmologiques ($\Lambda$CDM). Le coût computationnel de l'exécution de milliers de simulations périodiques de grande volume pour couvrir cet espace est prohibitif.
• Limites des approches actuelles : Les fonctions d'ajustement analytiques deviennent des sources d'erreurs systématiques majeures à mesure que les erreurs statistiques des relevés descendent en dessous du pourcent. Les émulateurs existants, bien que précis, sont souvent entraînés sur des boîtes périodiques complètes, ce qui reste coûteux et limite la dynamique de masse accessible.

L'objectif de cet article est de proposer une méthode permettant d'émuler des fonctions de distribution (comme la HMF) dans des volumes effectifs de Gigaparsecs, pour une large gamme de paramètres cosmologiques, à une fraction du coût des simulations traditionnelles, tout en accédant à des masses de halos plus faibles.

2. Méthodologie

Les auteurs proposent un cadre novateur combinant des simulations « zoom » (ou régions ciblées) et des estimateurs de densité neuronaux (Normalizing Flows).

A. Données d'entraînement (Suite Quijote)

• Utilisation de la suite de simulations Quijote (BSQ - Big Sobol Sequence), composée de 32 768 simulations $N$-body de matière noire uniquement dans des boîtes périodiques de $(1000 \, h^{-1}\text{Mpc})^3$.
• Ces simulations varient cinq paramètres cosmologiques ($\Omega_m, \Omega_b, h, n_s, \sigma_8$) selon une séquence Sobol.
• Au lieu d'utiliser les boîtes complètes, les auteurs extraient aléatoirement de petites régions cubiques de volume $V \sim (50 \, h^{-1}\text{Mpc})^3$ à partir de ces simulations parents.
• Chaque région est caractérisée par son sur-densité locale ($\delta$) et les paramètres cosmologiques globaux ($\theta$) de la simulation parente.

B. Architecture de l'émulateur

• Modèle : Un Normalizing Flow conditionnel (flux de normalisation) est entraîné pour prédire la fonction de masse des halos ($n(M)$) dans une région donnée.
• Conditionnement : Le modèle est conditionné par :
  1. Les paramètres cosmologiques globaux $\theta$.
  2. Le percentile de sur-densité de la région ($p_S$) par rapport à la distribution globale de la simulation parente (plutôt que la valeur absolue de $\delta$, pour éviter de modéliser explicitement la PDF de $\delta$).
  3. La masse du halo.
• Décomposition : La prédiction est décomposée en une base globale (moyenne de toutes les régions d'entraînement) et un résidu appris par le réseau de neurones.
• Fonction de vraisemblance : La prédiction du nombre de halos suit une loi de Poisson. L'entraînement minimise une perte de vraisemblance négative (NLL) via Monte Carlo, en intégrant sur des échantillons latents.

C. Reconstruction de la distribution globale

• Pour obtenir la HMF globale d'un volume Gigaparsec, l'émulateur est intégré sur la distribution des sur-densités.
• Procédure d'inférence :
  1. On génère une réalisation d'un grand volume en « tuilant » virtuellement des régions indépendantes.
  2. On échantillonne des percentiles de sur-densité ($p_S$) uniformément.
  3. L'émulateur prédit la densité de nombre de halos pour chaque région conditionnée par $\theta$ et $p_S$.
  4. On somme les comptes de halos sur toutes les régions pour reconstruire la HMF globale, intégrant ainsi efficacement la variance environnementale sans avoir besoin de simuler le volume complet.

D. Implémentation

• Le modèle est développé en JAX, permettant une compilation JIT (Just-In-Time) pour des performances GPU optimisées et une différenciation automatique complète. Cela rend l'émulateur compatible avec des cadres d'inférence bayésienne basés sur le gradient (ex: BlackJax).

3. Contributions Clés

1. Efficacité computationnelle extrême : La méthode utilise seulement ~0,026 % du volume de simulation original pour reconstruire avec précision la distribution globale. Cela suggère que des suites de simulations « zoom » extraites de volumes parents à basse résolution peuvent remplacer des suites massives de simulations périodiques.
2. Extension de la dynamique de masse : En utilisant des régions ciblées, la méthode permet de résoudre des halos de masse beaucoup plus faible que ce qui est possible dans les simulations périodiques de Gigaparsecs à résolution équivalente, tout en conservant la capacité à prédire la queue de distribution des amas massifs.
3. Gestion de la sur-densité via les percentiles : L'utilisation du percentile de sur-densité comme variable d'entrée simplifie l'intégration pour la reconstruction globale et évite les erreurs hors-distribution aux extrêmes de densité.
4. Différentiabilité : Le modèle est entièrement différentiable, ouvrant la voie à une inférence directe des paramètres cosmologiques à partir des données observationnelles.

4. Résultats

• Précision au niveau des régions : L'émulateur reproduit la HMF locale avec une précision compatible avec le bruit de Poisson. Les résidus déviants signés montrent peu de biais systématique, et les probabilités de comptage nul sont bien calibrées.
• Reconstruction globale : La HMF globale reconstruite présente une erreur fractionnaire médiane d'environ 2 % pour les masses inférieures à $10^{14.5} \, M_\odot$. Pour 95 % des cosmologies testées, l'erreur reste inférieure à 5 % en dessous de cette masse.
• Comportement aux hautes masses : Aux masses très élevées (queue de la distribution), les erreurs convergent vers la limite du bruit de Poisson, ce qui est attendu étant donné la rareté statistique de ces objets.
• Dépendance aux paramètres : L'émulateur capture correctement les sensibilités de la HMF aux paramètres cosmologiques, notamment la forte dépendance de la normalisation à $\Omega_m$ et la sensibilité de la queue haute masse à $\sigma_8$ et $n_s$.
• Optimisation des ressources : L'analyse de sensibilité montre que la diversité des environnements (sur-densités) est aussi importante, voire plus, que la diversité des paramètres cosmologiques pour réduire l'erreur. Un volume de région de $(50\text{--}70 \, h^{-1}\text{Mpc})^3$ semble optimal pour supprimer la variance cosmique sans retour sur investissement décroissant.

5. Signification et Implications

Cette étude démontre qu'il est possible de surmonter le compromis volume-résolution-paramètres en utilisant une approche hybride de simulations « zoom » et d'apprentissage automatique.

• Pour les relevés futurs (Euclid, Roman, Rubin) : La méthode offre un moyen pratique de générer des prédictions théoriques précises et dépendantes des paramètres pour des sondages à grand champ, en incluant potentiellement des effets baryoniques via des simulations hydrodynamiques « zoom » (comme FLARES ou CAMELS), sans le coût prohibitif de simulations hydrodynamiques de Gigaparsecs.
• Étude des galaxies faibles : En permettant de modéliser des halos de masse beaucoup plus faible à grande échelle, cette approche ouvre la voie à l'étude de galaxies à raies d'émission et d'autres traceurs dans des environnements variés, ce qui est crucial pour la cosmologie de précision.
• Inférence bayésienne : La nature différentiable du modèle permet son intégration directe dans des pipelines d'inférence de paramètres cosmologiques, facilitant l'extraction de contraintes précises à partir des données observationnelles.

En résumé, ce travail propose un changement de paradigme : plutôt que de simuler des volumes gigantesques à basse résolution, on simule des régions ciblées à haute résolution et on utilise l'apprentissage automatique pour reconstruire la statistique globale, réduisant ainsi le coût computationnel d'un facteur de plusieurs milliers.