A JAX package for gravitational lensing
This repository contains scripts and notebook to reproduce the results from the following paper:
Probabilistic Mass-Mapping with Neural Score Estimation, B. Remy, F. Lanusse, N. Jeffrey, J. Liu, J.-L. Starck, K. Osato, T. Schrabback, submitted to Astronomy and Astrophisics, 2021.
You will find all the scripts and instructions to reproduce traine models, sample maps and reproduce the paper figures here.
This repository also contains implementation of the following methods:
This package enables to sample high resolution convergence map from the posterior distribution in 10 GPU-minutes (on an Nvidia Tesla V100 GPU) in average. Have a look at the annealed Hamiltonian Monte Carlo sampling scheme bellow:
Comparison between
DLPosterior, DeepMass, Wiener Filter and KS93 methods.
| Ground truth convergence | DLPosterior samples |
|---|---|
 |
 |
DeepMass |
DLPosterior mean |
|---|---|
 |
 |
Wiener Filter |
Kaiser-Squires |
|---|---|
 |
 |
jax-lensing is pure python and can be easily installed with pip:
$ cd jax-lensing
$ pip install .
- Jax [==0.2.18]
- Haiku [==0.0.4]
- Astropy [==4.2]
- TensorFlow Probability [==0.13.0]
- TensorFlow Datasets [==4.3.0]
- TensorFlow [==2.5.0]
If you use jax-lensing in a scientific publication, we would appreciate citations to the following paper:
Probabilistic Mass-Mapping with Neural Score Estimation, B. Remy, F. Lanusse, N. Jeffrey, J. Liu, J.-L. Starck, K. Osato, T. Schrabback, submitted to Astronomy and Astrophisics, 2021.
The BibTeX citation is the following:
@Upcomming
Thanks goes to these wonderful people (emoji key):
 Benjamin Remy 🚇 |
 Francois Lanusse 🚇 |
 Niall Jeffrey |
 Jia Liu |
This project follows the all-contributors specification. Contributions of any kind welcome!