jaxtronomy 0.1.5

Installation and Usage

GPU

To use JAXtronomy with an NVIDIA GPU on Linux, first install JAX with

pip install -U "jax[cuda13]"

For other GPUs or operating systems, installation is more complicated. See the JAX installation instructions for GPU for more details. Then, install JAXtronomy with

pip install jaxtronomy

By default, JAX will still use CPU for computations. To change this, run the following line of code immediately after importing JAX

jax.config.update("jax_platform_name", "gpu")

CPU

For standard CPU-only usage, JAXtronomy and JAX can both be installed with

pip install jaxtronomy

For computations parallelized across CPU cores, an environment variable needs to be set before importing JAX indicating the number CPU devices to use. For example, to use 16 CPU cores, this can be done with

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=16"

Default Behavior and Configuration Flags

By default, JAX uses 32 bit datatypes when performing computations. Users of JAXtronomy are highly recommended to overwrite this behavior and use 64 bit datatypes instead, since 32 bit simulations have not been well-tested. Typically, this can be done by setting an environment variable or a configuration flag before any JAX operations are performed. For details, see JAX documentation on configuration flags

Example notebook: An example notebook has been made available, which showcases the features and improvements in JAXtronomy.

Performance comparison between JAXtronomy and lenstronomy

We compare the runtimes between JAXtronomy and lenstronomy by timing 1,000 function executions. While lenstronomy is always run using CPU, JAXtronomy can be run using either CPU or GPU. These tests were run using an Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz, an NVIDIA A100 GPU, and JAX version 0.7.0. A performance comparison notebook has been made available for reproducibility.

LensModel ray-shooting

The table below shows how much faster JAXtronomy is compared to lenstronomy for different deflector profiles and different grid sizes. Some comparisons vary significantly with values of function arguments.

LightModel surface brightness

The table below shows how much faster JAXtronomy is compared to lenstronomy for different source profiles and different grid sizes.

FFT Pixel Kernel Convolution

Convolution runtimes vary significantly, depending on both grid size and kernel size.

• For a 60x60 grid, and kernel sizes ranging from 3 to 45, jaxtronomy on CPU is about 1.1x to 2.9x faster than lenstronomy, with no obvious correlation to kernel size.

• For a 60x60 grid, and kernel sizes ranging from 3 to 45, jaxtronomy on GPU is about 1.5x to 3.5x faster than lenstronomy, with JAX performing better with higher kernel sizes.

• For a 180x180 grid, and kernel sizes ranging from 9 to 135, jaxtronomy on CPU is about 0.7x to 2.5x as fast as lenstronomy, with no obvious correlation to kernel size.

• For a 180x180 grid, and kernel sizes ranging from 9 t0 135, jaxtronomy on GPU is about 10x to 20x as fast as lenstronomy, with JAX performing better with higher kernel sizes.

Related software packages

The following lensing software packages do use JAX-accelerated computing that in part were inspired or made use of lenstronomy functions:

• Herculens

• GIGA-lens

• PaltaX

Community guidelines

Contributing to jaxtronomy

The guidelines for contributing to JAXtronomy are the same as lenstronomy, which can be found here. In short,

• Fork the repository

• Write clean, well-documented code, following conventions

• Submit pull requests

Reporting issues, seeking support, and feature requests

• Submit a Github issue