Differentiable Optics with dLux I: Deep calibration of Flat Field and Phase Retrieval with Automatic Differentiation

TL;DR

This paper introduces dLux, a Python package leveraging automatic differentiation for high-dimensional phase retrieval and detector calibration in astronomical imaging, enabling more accurate space telescope data analysis.

Contribution

The paper presents a novel GPU-accelerated, autodiff-based framework for simultaneous phase retrieval and detector calibration using high-dimensional optimization.

Findings

Gradient descent with autodiff effectively performs phase retrieval.

The framework scales to millions of parameters.

Enables improved calibration from science data alone.

Abstract

The sensitivity limits of space telescopes are imposed by uncalibrated errors in the point spread function, photon-noise, background light, and detector sensitivity. These are typically calibrated with specialized wavefront sensor hardware and with flat fields obtained on the ground or with calibration sources, but these leave vulnerabilities to residual time-varying or non-common path aberrations and variations in the detector conditions. It is therefore desirable to infer these from science data alone, facing the prohibitively high dimensional problems of phase retrieval and pixel-level calibration. We introduce a new Python package for physical optics simulation, dLux, which uses the machine learning framework JAX to achieve GPU acceleration and automatic differentiation (autodiff), and apply this to simulating astronomical imaging. In this first of a series of papers, we show that gradient descent enabled by autodiff can be used to simultaneously perform phase retrieval and calibration of detector sensitivity, scaling efficiently to inferring millions of parameters. This new framework enables high dimensional optimization and inference in data analysis and hardware design in astronomy and beyond, which we explore in subsequent papers in this series.