Tianao Li¹, Emma Alexander^2
1Tsinghua University, ²Northwestern University
Monthly Notices of the Royal Astronomical Society: Letters

Official code for Galaxy Image Deconvolution for Weak Gravitational Lensing with Unrolled Plug-and-Play ADMM.

To clone this project, run:

git clone https://github.com/Lukeli0425/Galaxy-Deconv.git

Create a virtual environment and download the required packages:

pip install -r requirements.txt

If you want to train the models with Shape Constraint, please install AlphaTransform.

Download our simulated galaxy dataset from Google Drive.

To train the models, run train.py and choose parameters and loss function for your training, for instance:

python train.py --model Unrolled_ADMM --n_iters 8 --n_epochs 50 --loss MultiScale --lr 1e-4

Test the algorithms from the perspectives of time and performance with test.py. Uncomment the methods to be tested in the code and specify the number of galaxies you want to use in the test dataset:

python test.py --n_gal 10000

Similarly, you can test the robustness of the algorithms to systematic errors in PSF with test_psf.py:

python test_psf.py --n_gal 10000

All the test results will be automatically saved in the results folder.

We saved our models trained on our LSST dataset (see saved_models). We also provide a tutorial for using the suggested model on your data, see tutorials/deconv.ipynb for details.

We simulated our dataset with the modular galaxy image simulation toolkit Galsim and the COSMOS Real Galaxy Dataset. To us our image simulation pipeline, one need to first download the COSMOS data here or download with Galsim:

galsim_download_cosmos -s 23.5

Run generate_data.py to simulate your own dataset under different settings (remember to change the path to your COSMOS data). Simulate your dataset for deconvolution task by running

python generate_data.py --task Deconv --n_train 40000

We provide the dataset generation code for denoising task. The denoising dataset is used in our ablation studies (see figures/ablation.ipynb) to train the plugin denoiser in ADMMNet. Simulate your dataset for denoising task by running

python generate_data.py --task Denoise --n_train 40000

We provide a detailed tutorial for image simulation (see tutorials/image_simulation.ipynb), where you can find out how to set up your simulations.

You can train the models on your dataset.

The figures folder holds the figures in this paper and the files that created them (see figures/README.md). To recreate the figures with your own results, you can use the given files and follow the instructions we provide.

@article{li2023galaxy,
  title={Galaxy image deconvolution for weak gravitational lensing with unrolled plug-and-play ADMM},
  author={Li, Tianao and Alexander, Emma},
  journal={Monthly Notices of the Royal Astronomical Society: Letters},
  volume={522},
  number={1},
  pages={L31--L35},
  year={2023},
  publisher={Oxford University Press}
}