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_partial_parallel.py  Multi-GPU (data-parallel) version 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)
  test_multi_gpu.py        Checks that JAX sees the node's GPUs and that they can all-reduce (NCCL)

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.

Multi-GPU training

experiments/train_partial_parallel.py is a working data-parallel version of train_test_partial.py: it trains on every GPU of a node (single process, Mesh + NamedSharding + shard_map, no pmap). Model, EMA and optimizer state are replicated, batch_size is the global batch split across GPUs, and gradients are averaged with pmean. Submit the job with one task for the whole node (e.g. --gres=gpu:4 --ntasks=1); run test/test_multi_gpu.py first to confirm JAX sees and can all-reduce across the GPUs.

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.