Weak Lensing Uncertainty Challenge

This repository is implement from team THUML for NeurIPS 2025 Weak Lensing Uncertainty Challenge. We propose a Vision Transformer (ViT) architecture conditioned on handcrafted domain features from cosmology through Adaptive Layer Normalization (AdaLN), and an ensemble selected via greedy search over heterogeneous feature sets and network variants to achieve accurate and robust inference of $\Omega_m$ and $S_8$ from weak-lensing convergence maps.

Setup

First, run the following command to install the required Python libraries.

pip install -r requirements.txt

Pipeline

Our pipeline consists of three stages:

(1) Pre-extracting Handcrafted Features — computing handcrafted domain features including Power Spectrum, Peak Counts, Minkowski Functionals, etc. from each convergence map and saving.

(2) Training Multiple Model Variants — using different feature combinations and architectural choices to build a diverse set of candidate models.

(3) Greedy Ensemble Selection — iteratively selecting the model which yields the largest improvement in the validation score when combined with the current ensemble.

Pre-extracting Handcrafted Features

Run the following Python script to compute handcrafted domain features and save.

Note that argument data_dir should be set to your own directory where your data is stored. This directory contains at least noisy_kappa_train.npy noisy_kappa_val.npy noisy_kappa_test.npy, which correspond respectively to convergence maps for training, validating and testing.

python traditional_feature_extract.py --data_dir XXX

Training Multiple Model Variants

Run the following Bash script to train multiple model variants with different feature combinations and architectural choices.

bash main_1.sh / bash main_2.sh / bash main_3.sh

Note that argument data_dir should also be set to your own directory where your data is stored.

First, you can change argument traditional_f_pos to set feature combinations. Note that the IDs below and their corresponding features:

(0) Power Spectrum,

(1) Peak Counts,

(2) Minkowski Functionals,

(3) Lightweight Statistical Descriptors,

(4) Multiscale Moments,

(5) Multiscale Wavelet Statistics,

(6) Aperture Mass Statistics,

(7) 1-point PDF,

(8) Gradient and Curvature Statistics,

(9) Entropy and Information-Theoretic Features,

(10) Spatial Autocorrelation Features.

Second, you can change argument model to choose model architecture:

FAIR_model1: Standard architecture - ViT with AdaLN,

FAIR_model2: Standard architecture with a CNN front-end,

FAIR_model3: Standard architecture without convergence map input.

Greedy Ensemble Selection

Run the following Python script to conduct greedy search according to validation scores to achieve more effective ensemble selection.

The result including model selection will be saved in the experiment directory in section Training Multiple Model Variants.

python greedy_search.py