Score Matching on Large Geometric Graphs for Cosmology Generation

By Diana-Alexandra Onutu, Yue Zhao, Joaquin Vanschoren, Vlado Menkovski

This repository contains the code for the paper Score Matching on Large Geometric Graphs for Cosmology Generation, accepted at the International Conference on Discovery Science 2025 conference.

Overview

A score-based generative model for simulating galaxy clustering across different cosmologies. The model uses Equivariant Graph Neural Networks (EGNNs) and a topology-aware noise schedule to generate realistic halo catalogs conditioned on cosmological parameters. It scales to catalogs with up to ~600,000 halos while running over 6 orders of magnitude faster than traditional N-body simulations.

Key features:

• Physics-grounded prior: sampling starts from a uniform distribution, matching early-universe conditions, rather than Gaussian noise.
• E(3) equivariance: enforces Euclidean symmetries and periodic boundary conditions throughout training and inference.
• Topology-aware noise schedule: designed by tracking the percentage of preserved k-NN edges across noise levels, ensuring graph structure is maintained at low noise (σ_low = 0.001).
• Scalable: handles full Quijote simulation boxes via sub-box sampling during training.

---

Repository Structure

project_root/
│-- train.py                        # Single-GPU training
│-- train_ddp.py                    # Multi-GPU training (DistributedDataParallel)
│-- infer.py                        # Inference for standard experiments
│-- infer_denoising.py              # Inference with denoising visualization
│-- infer_large_graphs.py           # Inference for full-scale (~600k halo) catalogs
│-- profiling.py                    # Runtime and memory profiling
│
│-- models/
│   │-- egnn.py                     # Equivariant GNN (ScoreEGNN)
│   │-- gnn.py                      # Non-equivariant GNN (ScoreGNN) baseline
│   │-- egcl.py                     # Equivariant Graph Convolutional Layer
│   │-- gcl.py                      # Graph Convolutional Layer
│   │-- mlp.py                      # MLP building block
│
│-- utils/
│   │-- data.py                     # Dataset loading and preprocessing
│   │-- graph.py                    # k-NN graph construction with periodic BCs
│   │-- perturbation.py             # Forward noising process
│   │-- loss.py                     # Denoising score matching loss
│   │-- train.py                    # Training loop utilities
│   │-- evaluation.py               # 2PCF computation and metrics
│   │-- output.py                   # Logging and checkpoint utilities
│
│-- configs/
│   │-- train_egnn.yaml             # ScoreEGNN training (Ωm, σ8 conditioning)
│   │-- train_egnn_all.yaml         # ScoreEGNN training (all 5 parameters)
│   │-- train_egnn_large_graphs.yaml# ScoreEGNN training on full-scale boxes
│   │-- train_egnn_other_noise.yaml # Ablation: suboptimal noise schedule
│   │-- train_gnn_augm.yaml         # ScoreGNN training with augmentation
│   │-- infer_egnn.yaml             # Inference config (Ωm, σ8)
│   │-- infer_egnn_all.yaml         # Inference config (all 5 parameters)
│   │-- infer_egnn_large_graphs.yaml# Inference config for large graphs
│   │-- infer_egnn_other_noise.yaml # Inference config for noise ablation
│   │-- infer_egnn_quali.yaml       # Qualitative 2PCF experiment
│   │-- infer_gnn.yaml              # ScoreGNN inference config
│
│-- scripts/
│   │-- submit_train.sh             # Cluster job submission for training
│   │-- submit_train_ddp.sh         # Cluster job submission for DDP training
│   │-- submit_infer.sh             # Cluster job submission for inference
│   │-- submit_infer_denoising.sh   # Cluster job submission for denoising inference
│   │-- submit_infer_large_graphs.sh# Cluster job submission for large-graph inference

---

Dataset

This project uses the Quijote N-body simulation suite. The dataset consists of 2,000 halo catalogs at redshift z=0, with cosmological parameters (Ωm, Ωb, H₀, nₛ, σ₈) varied across simulations.

• Training set: 1,800 cosmologies
• Test set: 200 cosmologies
• Subset experiments: 5,000 most massive halos per simulation
• Large-graph experiments: full catalogs (~600,000 halos)

The raw data is not included in this repository. Follow the Quijote download instructions to obtain the halo catalogs, then update the data.cosmologies_dir paths in the relevant config files.

---

Installation

pip install -r requirements.txt

---

Training

The models were trained on NVIDIA H100 (SXM5) 94GiB GPUs on Snellius, the Dutch national supercomputer.

Single GPU:

python train.py --config configs/train_egnn.yaml --run_name my_run

Multi-GPU (DDP):

python train_ddp.py --config configs/train_egnn.yaml --run_name my_run

Or submit to a SLURM cluster:

bash scripts/submit_train_ddp.sh

Key config parameters (see configs/train_egnn.yaml):

Parameter	Default	Description
noise.L	100	Number of noise levels
noise.sigma_1	0.2	Highest noise level (σ_high)
noise.sigma_L	0.001	Lowest noise level (σ_low)
training.k	20	k-NN graph connectivity
data.cosm_param	["Omega_m", "sigma_8"]	Conditioning parameters
model.hidden_nf	128	Hidden feature dimensionality
model.n_layers	4	Number of EGCL layers

---

Inference

python infer.py --config configs/infer_egnn.yaml --run_name my_run

For large-scale (~600k halo) generation:

python infer_large_graphs.py --config configs/infer_egnn_large_graphs.yaml --run_name my_run

---

Acknowledgements

The project has been supported by the Dutch national e-infrastructure with the support of SURF (grant no. EINF-10541). We thank SURF (www.surf.nl) for the support in using the National Supercomputer Snellius.

---

Citation

@inproceedings{onuțu2025score,
  title={Score Matching on Large Geometric Graphs for Cosmology Generation},
  author={Onuțu, Diana-Alexandra and Zhao, Yue and Vanschoren, Joaquin and Menkovski, Vlado},
  booktitle={International Conference on Discovery Science},
  pages={17--31},
  year={2025},
  organization={Springer}
}