strong_lensing_simulation

This repo contains all our code for the imaging multi-regression task of parameter and uncertainty estimation for strong lensing systems.

Data generation / preparation using Lenstronomy

• This notebook is used to generate the images (data) and paramters (targets) as the dataset for the strong lensing systems.
• This notebook is used to process the dataset: data split and target normalization.

Source code for training models

• train_model.py is the main code to train models (ViT and ResNet).
• The src folder contains scripts for helper functions.

Train models

• This notebook is used to train a ViT model.
• The training_eccv folder contains the notebooks used to train models for our ECCV paper.
• The training_icml folder contains the notebooks used to train models for our ICML paper.

Model Prediction and visulization

• predict.py is the source code to make prediction using a trained model.
• This notebook uses predict.py to make predictions for our ECCV paper.
• visualization.py contains objects and functions for visulization.
• This notebook uses visualization.py to make figures for our ECCV paper.