Train-AE

Generative modeling of galaxy images from the Euclid Q1 survey, developed during an internship at CosmoStat (CEA).

Overview

This repository trains a two-stage generative model on 64x64 galaxy postage stamps:

1. A convolutional autoencoder that reconstructs galaxy images, convolving its output with the instrument PSF (via jax-galsim) before comparing it to the observed, PSF-convolved image.
2. A normalizing flow fit on the frozen autoencoder's latent space, allowing new galaxy images to be sampled.

The training data is the Euclid Q1 dataset, hosted on Hugging Face (VincentB03/euclid-Q1-VF), each sample pairing a science image with its PSF stamp, noise map, and mask. Runs were executed on the Jean Zay supercomputer (IDRIS/CNRS).

The modeling code (autoencoder, flow, and building blocks in pshear/) builds on prior work by Benjamin Rémy, a former PhD student at the CosmoStat laboratory.

Repository structure

pshear/                 Core library
  galaxy.py               Galaxy autoencoder (PSF convolution via jax-galsim) and its loss
  utils.py                Checkpoint save/load, Weights & Biases checkpoint fetching
  nn/                      Autoencoder, flow, and shared network blocks (JAX/Equinox)

experiments/            Training and evaluation scripts
  train_test.py            Autoencoder training with the full PSF
  train_test_partial.py    Autoencoder training with a partial PSF
  train_test_parallel.py   Multi-device (pmap) variant of train_test.py
  train_test_partial_parallel.py  Multi-device (pmap) variant of train_test_partial.py
  train_flow.py            Fits the normalizing flow on the frozen autoencoder's latent space
  verification.py          PQMass-based check that generated samples match the real data distribution
  utils.py                 Shared paths and plotting helpers

test/
  test_requirements.py     Sanity check of the training environment (JAX/GPU, Hugging Face, Weights & Biases)

Full vs. partial PSF training

train_test.py and train_test_partial.py differ in how the PSF is handled during autoencoder training:

• train_test.py uses the full PSF. Because deconvolving with the full PSF is an ill-posed problem, the loss includes a total-variation regularization term to suppress the pixelization artifacts this introduces in the deconvolved output.
• train_test_partial.py uses a partial PSF, also provided in the dataset. This avoids the need for a non-physical regularization term altogether.

The _parallel variants (train_test_parallel.py, train_test_partial_parallel.py) are a multi-device (pmap) starting point for future work and were not used to produce the results of these experiments.

Environment

Dependencies are listed in requirements.txt. JAX itself is installed separately with the CUDA build matching the target cluster, e.g.:

pip install -U "jax[cuda12]"
pip install -r requirements.txt

Training runs are logged to Weights & Biases, and checkpoints/configs can be re-fetched from a W&B run for evaluation (see pshear.utils.fetch_wandb_checkpoint).

Verification

experiments/verification.py uses PQMass (the pqm package) to statistically test whether images generated by the flow follow the same distribution as real held-out data, both in the autoencoder's latent space and in image space, with a real-vs-real calibration test as a sanity check on the setup itself.

Pre-downloading checkpoints

download_wandb_weights.py fetches the AE and/or flow checkpoints of a W&B run into wandb_weights/, in the layout load_galaxy_autoencoder / load_flow expect:

wandb_weights/<run_id>/config.yaml                        run config
wandb_weights/<run_id>/epoch_<n>/model_checkpoint_<n>.eqx weights
wandb_weights/<run_id>/epoch_<n>/config.yaml              run config, un-wandbified

Run it on a login node (with network access); the resulting cache can then be reused as-is on a compute node with no network, where fetch_wandb_checkpoint skips the W&B API entirely. The run ids and epochs are set in the CONFIG block at the top of the file, and every value has a matching CLI flag for use in a job script:

python download_wandb_weights.py                          # use the CONFIG block
python download_wandb_weights.py --only flow
python download_wandb_weights.py --flow-run-id 4q23te9a --flow-epoch 420 --only flow
python download_wandb_weights.py --cache-dir /path/to/other/dir   # change the download destination

This is the mechanism experiments/verification.py relies on to load its models: it calls fetch_wandb_checkpoint to populate wandb_weights/, then load_galaxy_autoencoder / load_flow read the resulting epoch_<n>/ directory directly. The galaxy-morphometrics repo expects checkpoints in the same wandb_weights/<run_id>/epoch_<n>/ layout (its WandBGalaxyAutoencoder / WandBGalaxyFlow), so the cache produced here can be reused there as-is.

Because the destination is configurable with --cache-dir, you can point it straight at another repo's checkpoint directory and skip the copy step. For example, to make a checkpoint available to a galaxy-morphometrics checkout:

python download_wandb_weights.py --cache-dir /path/to/galaxy-morphometrics/wandb_weights

WandBGalaxyAutoencoder / WandBGalaxyFlow will then find the wandb_weights/<run_id>/epoch_<n>/ tree already populated and load from it without hitting the W&B API. Point verification.py at the same directory (its cache_dir argument to fetch_wandb_checkpoint) if you move the cache away from this repo's default.