SciPy 2024

Development of Caustics: a differentiable, GPU accelerated, gravitational lensing simulator

We present Caustics, a tool to accelerate the analysis of gravitational lensing systems for the next generation of astronomical data. Caustics will enable precision measurements of dark matter properties, the expansion rate of the Universe, lensed black holes, the first stars, and more. In this talk I will discuss the benefits and challenges of how we used PyTorch (a differentiable and GPU accelerated scientific python package) to allow for fast development without sacrificing numerical performance. I will detail our development process as well as how we encourage users of all skill levels to engage with our documentation/tools.


We present Caustics, a tool to accelerate the analysis of gravitational lensing systems for the next generation of astronomical data (https://caustics.readthedocs.io). In this talk I will discuss the development process and documentation of our lensing package using PyTorch as a numerical library; including its benefits and challenges. I will begin by explaining the basics of gravitational lensing for a broad audience (https://youtu.be/20-kR0cDPhA). With a basic foundation in place I can explain why GPUs are ideally suited for gravitational lensing raytracing simulations, and why a differentiable code gives us access to powerful new numerical techniques. The bulk of the talk will then be devoted to the package development, API design, and documentation to suit the scientific needs of our community; with a focus on lessons that apply generally to scientific python development.

While writing the package we aimed to meet the needs of users at all levels, and so we have developed an API which is an evolution from previous widely used packages. We include a yaml configuration interface, an object oriented interface, and a functional interface where each API builds on the one below it to allow for users to seamlessly transition from beginners to power users. The configuration file interface allows users to begin applying caustics to their own data within minutes, though it's flexibility is limited. The object oriented framework allows for complex modelling and advanced scientific analyses.
The functional interface allows for development of entirely new tools, while exploiting a robust and well tested core functionality.

Our documentation similarly meets the needs of all users, with units included directly in the parameter descriptions, links to relevant papers, and mathematical explanations of various operations. Scientific programming requires a level of rigour beyond many other applications and so we have built into our documentation as much precision as possible. There are also Jupyter notebooks directly integrated in our documentation (https://caustics.readthedocs.io/en/latest/tutorials/BasicIntroduction.html) which guide users to best practices and provide a skeleton from which to further develop their science. These notebooks can be a welcome mat for new users who want to learn about lensing; they are also convenient Lego pieces for advanced users to quickly develop new ideas.

Overall, our development approach has been to make a clean, complete, and reliable code to meet the needs of the gravitational lensing community. Our team has learned many lessons along the way which I would like to share in this talk to encourage a discussion on best practices for development of workhorse scientific python packages.

I have given numerous talks before, here is a selection of slides from last year:
https://docs.google.com/presentation/d/1kmUQx9JpceHQDL2g-d387REo4Fv6c6c7o6JDz3YnUYk/edit?usp=sharing
https://docs.google.com/presentation/d/1xhZphgw46wb7iKPY-2YyWkNuT1yb8H-RdcLClYIGRPs/edit?usp=sharing
https://drive.google.com/file/d/1_YwRY5GJzY5HBfJ5xy2bpgepfxDg2yI_/view?usp=sharing
https://docs.google.com/presentation/d/13COYyj7_FUXmirivS_m8cui-MgdozzANVy8Y7UjahUc/edit?usp=sharing