This repository contains the pipeline to train neural networks that learn informative summary statistics from the Dark Energy Survey Year 3 (DES Y3)-like weak lensing and galaxy clustering maps [Thomsen et al. 2025].
- Training Data: HEALPix maps in spherical (curved) geometry masked according to the DES Y3 survey footprint and stored as tensors in
.tfrecordformat as generated bymultiprobe-simulation-forward-model. - Architectures: Per default, we use the DeepSphere graph convolutional neural networks [Defferrard et al. 2020] from
deepsphere-cosmo-tf2. Alternative choices include 1D convolutional networks and vision transformers. - Loss Functions: The preferred loss function to train networks to implement a mapping from pixel space to low-dimensional, yet informative summary statistics is variational mutual information maximization. Alternative choices include the mean squared error, log-likelihood loss, and Fisher information maximization.
- HPC Distribution: Multi-GPU support (intra- and cross-node) for data parallel training on HPC clusters. Optimized for the Perlmutter A100 cluster at the National Energy Research Scientific Computing (NERSC) facility.
Requires Python >= 3.8, TensorFlow >= 2.0, TensorFlow-Probability, and Horovod.
Main dependencies:
multiprobe-simulation-forward-modelfor data loading and utilitiesdeepsphere-cosmo-tf2for graph convolutional neural networks on the pixelized sphere
Step 1: Install dependencies from GitHub
# Install multiprobe-simulation-forward-model (data loading)
pip install git+https://github.com/des-science/multiprobe-simulation-forward-model.git
# Install deepsphere-cosmo-tf2 (graph convolutional networks in TensorFlow 2)
pip install git+https://github.com/deepsphere/deepsphere-cosmo-tf2.gitStep 2: Install this package
On HPC clusters with pre-installed TensorFlow/Horovod (recommended):
pip install -e .On systems without TensorFlow/Horovod:
pip install -e .[tf]Use the first option when TensorFlow and Horovod are available via system modules (e.g., module load tensorflow horovod) to preserve optimized GPU/MPI configurations.
deep_lss/apps- Training and evaluation scripts forslurmsubmissiondeep_lss/models- Loss function-specific model classesdeep_lss/nets- Neural network implementationsdeep_lss/utils- Loss functions, multi-GPU distribution, and helper functions
Configuration files for network architecture, loss function, and optimizer hyperparameters, and analysis choices including the selected cosmological probe(s), parameters to be constrained, and smoothing scales.
Bash scripts for HPC slurm job scheduling.
- Forward modeling:
multiprobe-simulation-forward-model - Simulation-based inference:
multiprobe-simulation-inference
