Cosmos DESI Emulator

Notice: This project is the result of academic research currently in the experimental and paper-writing phase. The methodology, network architecture, and inference pipelines are in continuous development. As of now, the work does not present definitive conclusions or published final results.

Cosmological emulator based on Evidential Graph Physics-Informed Neural Networks (Graph E-PINN) for DESI DR1 data analysis.

Premise and Methodology

The main premise of this study is to investigate the capability of Physics-Informed Neural Networks (PINNs) in performing inference of underlying cosmological parameters — focusing on matter density ($\Omega_m$) and the amplitude of fluctuations ($\sigma_8$) — directly from the spatial distribution of galaxies. The central objective is to develop an analytical and computationally efficient alternative to traditional N-body simulations.

From a methodological standpoint, the pipeline integrates:

1. Spatial Graphs (kNN): Structuring the galaxy catalog into a k-Nearest Neighbors graph in comoving space, allowing for a mesh-free modeling free from traditional grid resolutions.
2. Physics as a Probabilistic Prior: The dynamics of collisionless phase space is governed by the Vlasov-Poisson system of equations. This physical behavior is not imposed as an isolated loss function penalty, but rather as a constraint in a Bayesian model.
3. Variational Inference and Uncertainty (SVI): Network training occurs within the NumPyro framework via Stochastic Variational Inference (SVI). Cosmological variables are modeled as latent variables. This Bayesian formulation combined with Normal-Inverse-Gamma distributions (Evidential Regression) allows for a strict decomposition between aleatoric uncertainty (intrinsic observational noise) and epistemic uncertainty (degree of model uncertainty).

Data Origin and Structure

The data used in this project comes from the Large Scale Structure (LSS) catalogs of the Dark Energy Spectroscopic Instrument (DESI) — Data Release 1 (DR1).

The emulator's pipeline natively manages the downloading and structural processing of the original files. The catalogs provide the astronomical coordinates of the mapped sources, fundamentally: Right Ascension (RA), Declination (DEC), and spectroscopic Redshift ($z$), along with their respective operational weights for completeness and systematics. The code has pre-configured support for ingesting and validating the main cosmological tracers surveyed by DESI, including Luminous Red Galaxies (LRG), Quasars (QSO), Emission Line Galaxies (ELG), and data from the Bright Galaxy Survey (BGS).

Dependencies and Installation

It is recommended to set up an isolated virtual environment for installing dependencies:

pip install -r requirements.txt

Note: For performance optimization, you need to install a JAX version with CUDA support compatible with your system. Consult the official JAX documentation for proper instructions for your execution environment.

Execution

The pipeline is orchestrated by the main.py script, which manages everything from data preparation and reading to model initialization and the iterative SVI process.

For structural validation of the code using a local simplified synthetic catalog:

python main.py --synthetic --n-galaxies 2000 --steps 1000

For inference using the LRG (Luminous Red Galaxies) catalog from the NGC region of DESI DR1:

python main.py --tracer LRG_NGC --max-rows 50000 --steps 2000