a differentiable lensing simulator (preprint, accepted publication)
MADLens is a python package for producing non-Gaussian cosmic shear maps at arbitrary source redshifts. MADLens is designed to achieve high accuracy while keeping computational costs as low as possible. A MADLens simulation with only 256^3 particles produces convergence maps whose power agree with theoretical lensing power spectra up to scales of L=10000. MADlens is based on a highly parallelizable particle-mesh algorithm and employs a sub-evolution scheme in the lensing projection and a machine-learning inspired sharpening step to achieve these high accuracies.
MADLens is fully differentiable with respect to the initial conditions of the underlying particle-mesh simulations and a number of cosmological parameters. These properties allow MADLens to be used as a forward model in Bayesian inference algorithms that require optimization or derivative-aided sampling. Another use case for MADLens is the production of large, high resolution simulation sets as they are required for training novel deep-learning-based lensing analysis tools.
Download the code, then run
pip install -e .
which installs MADLens together with all its dependencies.
Alternatively, to create a conda environment with the required dependencies, run
conda env create -f MADLens.yml
Run the code with
python run_lightcone.py
Download the code.
For an interactive session, allocate the desired number of nodes, then run
source PrepareInteractiveJob.sh
then run
srun -n <number of tasks> -u python run_lightcone.py <flags>
For job submission, prepare a job script following PrepareInteractiveJob.sh
To display all parameters run
python run_lightcone.py --helpfull
to set a parameter, either change its default in run_lightcone.py or pass them as
python run_lightcone.py --parameter_name parameter_value
Here's a non-exhaustive list of parameters that can be set by the user:
| Parameter | Description | Typical Value(s) |
|---|---|---|
| BoxSize | side length of the simulation box | 128-1024 Mpc/h |
| Nmesh | resolution of the particle-mesh simulation | 64^3-512^3 |
| B | force resolution factor | 2 |
| Nsteps | number of steps in the FastPM simulation | 11-40 |
| N_maps | number of output maps | >1 |
| Nmesh2D | resolution of the convergence map | 256^2-2048^2 |
| BoxSize2D | size of the convergence map in degrees | 2.5-22 degrees |
| zs_source | list of source redshifts | 0.3-2.0 |
| Omega_m | total matter density | 0.32 |
| sigma_8 | amplitude of matter fluctuations | 0.82 |
| PGD | whether to use PGD enhancement or not | True/False |
| interpolation | whether to use the sub-evolution scheme | True/False |
We provide PGD parameters for a limited amount of configurations. Before using PGD, you have to dump those parameters into parameter files by running the notebook PGD_params.ipynb
See lightcone.py for examples of how to compute forward passes, vector-Jacobian and Jacobian-vector products.
For example,
kmaps, tape = model.compute(vout='kmaps', init=dict(rho=rho),return_tape=True)
evolves initial conditions rho into lensing maps kmaps (runs the forward model), while saving operations to the tape.
Once the tape has been created,
vjp = tape.get_vjp()
kmap_vjp = vjp.compute(init=dict(_kmaps=vector), vout='_rho')
computes the vector Jacobian product against a vector of the same shape as kmaps.
To use the MADLens version that support parameter derivatives, checkout the param_derivs branch.
Vanessa Böhm, Yu Feng, Max Lee, Biwei Dai
other packages MADlens relies on: