jax-lensing

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:

Convergence posterior sampling

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:

Mass-mapping

Comparison between DLPosterior, DeepMass, Wiener Filter and KS93 methods.

Ground truth convergence	`DLPosterior` samples

`DeepMass`	`DLPosterior` mean

`Wiener Filter`	`Kaiser-Squires`

Install

jax-lensing is pure python and can be easily installed with pip:

$ cd jax-lensing
$ pip install .

Requirements

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]

Citation

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: