text:
\section{Introduction}

\subsection{Context}

Each time an imaging instrument makes an observation, the image taken is affected by instrument-related distortions. We can encompass most of the distortions into the instrumental response or point spread function (PSF). Many science applications need to take into account the effects of the PSF for a quantitative analysis of images. The characterization of the instrument's response, also known as PSF modeling, is a subject of study in many fields, ranging from astronomical imaging to biomedical imaging. Cosmology provides an interesting application that poses very stringent requirements on the quality of the PSF model. This application is weak gravitational lensing (see \cite{kilbinger2015} or \cite{mandelbaum2018_bis3} for a review).

The light emitted from observed objects travels through the Universe towards our imaging instrument and is deflected by the gravitational potential of massive objects along the line of sight. The effect is similar to that of a magnifying glass and is therefore given the term gravitational lensing. In the weak lensing regime, it represents a very small change of shape of observed extended objects, usually galaxies. The precise measurement of this change in shape can help cosmologists answer questions that have puzzled them since the 1980s, for example, what is the nature of dark matter and its spatial distribution.
Weak gravitational lensing is based on the measurement of galaxy shapes with high accuracy. The instrument's PSF produces a systematic deformation of the images, thus obscuring the lensing-induced shape deformations.
It is therefore \textit{essential} to correct these images from the effects of the PSF.

The next generation of powerful imaging instruments is currently being built, such as the \textit{Euclid} space telescope \cite{laureijs2011}, the \textit{Roman} space telescope \cite{wfirst}, and the Vera C.~Rubin Observatory's Legacy Survey of Space and Time (LSST) \cite{LSST2009}. These instruments' wide field-of-view (FOV) will make the PSF vary significantly across the focal plane. The extremely tight instrumental knowledge required by weak-lensing applications requires capturing PSF morphological variations across the FOV into the image correction. The accuracy aimed by these future missions also demands taking into account the wavelength variations of the PSF. We assume that the PSF can be modeled locally by a convolution kernel. Then, the \textit{PSF field} refers to a continuous function that maps each point in the FOV and each wavelength to the corresponding convolution kernel. In astronomy, certain observations of unresolved sources like isolated stars provide a measurement of the PSF field at a given FOV position. PSF modeling consists in using these observations to build a model of the PSF field. Then, the PSF model can be used to infer the PSF at the positions of the objects of interest and to correct their images.

This problem has created much attention within the astronomer's community. Over the years, imaging instruments have become more and more powerful, while the accuracy requirements have become more and more stringent. This has motivated the development of a wide range of methodologies for PSF modeling. These include models based on Moffat functions \cite{moffat1988}, polynomial models \cite{piotrowski2013, bertin2011, miller2013, jarvis2020}, principal component analysis \cite{pca_psf2007, schrabback2010, gentile2013}, sparse matrix factorization \cite{ngole2015,ngole2016,schmitz2020}, specific basis functions \cite{refregier2003a,massey2005}, optimal transport \cite{ngole2017}, neural networks \cite{herbel2018,jia2020a,jia2020b,jia2020c}, and a parametric model of the telescope's optics \cite{krist1993, krist1995, krist2011}.

Future missions, like \textit{Euclid}, impose severe requirements to the PSF model.
The current state-of-the-art data-driven PSF model for \textit{Euclid} \cite{schmitz2020} is not capable of modeling wavelength variations and is far from requirements, at a factor between $100$ and $300$ for the PSF shape and a factor of $10^4$ for the PSF size. These facts triggered us to rethink the way we are building data-driven PSF models.

\subsection{Challenges in point spread function modeling for space missions}

The main difference compared to a ground-based telescope PSF is the effect of the atmosphere, which acts as a fast-changing stochastic filter. The PSF tends to resemble a low-pass filter as the exposure time increases. Therefore, it strongly affects the observations and changes over time. The improvements in the PSF stability and quality are one of the main reasons to send a telescope in space.

As already mentioned above, telescope observations in galaxy imaging surveys include unresolved objects, which are stars, and extended objects of interest, such as galaxies, at different positions in the FOV. The modeling of the PSF in the field of view (FOV) can be seen as an inverse problem. We can consider stars as samples of the PSF field in the FOV. These samples can be used to constrain our model, which we will later use to infer the PSF at target positions. PSF field modeling for space missions encompasses several challenges:
\begin{itemize}
    %

    \item[1.] \textit{The PSF varies spatially in the FOV}.
    The instruments are designed with a wide FOV and are composed of a large array of charge-coupled devices (CCDs). These dimensions make the PSF vary substantially across the focal plane. The lack of atmosphere and the high-quality optics used in telescopes makes the PSF close to the diffraction limit. As a consequence, the PSF has more complex shapes with higher spatial frequencies. The model needs to capture the different spatial frequency variations of the PSF from the stars in order to infer the PSF at target positions.
    %

    \item[2.] \textit{The observations are undersampled}.
    This is generally the case for space missions, and the model needs to super-resolve the output PSFs. This differs from the usual super-resolution (SR) task, as we do not have several low-resolution observations of the same object. In this case, we have several samples of the undersampled PSF field at \textit{different positions} in the FOV.
    %

    \item[3.] \textit{The PSF varies as a function of wavelength}.
    Also known as chromatic variations, they need to be included in the PSF model for most science goals. However, the modeling difficulty resides in some instruments' very broad passband. This means that each star observation is integrated with respect to its wavelength throughout the passband (e.g.~the \textit{Euclid} passband extends from $550$ to $900$nm \cite{cropper2016}). The PSF chromaticity will also be very likely important for the \textit{Roman} space telescope \cite{wfirst} or LSST \cite{LSST2009}.
\end{itemize}

In addition to the three major problems mentioned above, we need to also address the fact that during its journey in space, the telescope will change over time. Some reasons include strong temperature gradients or degradations of the instrument. Although it might seem intricate to capture these variations into the PSF model, in practice, one can build independent PSF models at each point in time. As a consequence, we have a different model for each time snapshot and we cope with the temporal variations. This will be the case for the present work where we analyze a single point in time. The only limitation that this approach imposes is the number of observed stars available to constraint the PSF model. The stability of a space mission allow using several exposures to increase the available information when building the PSF model.

\subsection{Current approaches for PSF modeling}
\label{sc:approaches_PSF_modeling}

There are two main approaches, or families of methods, for PSF modeling, as follows:

\paragraph{Parametric PSF models}
This approach consists in building a parametric model of the entire optical system that aims to be as close as possible to the actual telescope. Then, a few number of model parameters are fit to the star observations. This family of methods is capable of handling the chromatic variations of the PSF. This approach has been used for the Hubble Space Telescope (HST), with the Tiny Tim method \cite{krist1993, krist1995b, krist2011}. However, it was later shown that a simple data-driven model \cite{HST_PSF2017} outperformed Tiny Tim, exposing limitations of parametric modeling. Errors will arise if there is a mismatch between the parametric model and the ground truth. Furthermore, even if, ideally, there was no mismatch, the optimization of these models is a degenerate problem. It requires potentially expensive calibration information, usually in the form of out-of-focus observations, to break degeneracies. Some events such as launch vibrations and ice contaminations introduce significant variations into the model. These reasons prevent a complete ground characterization to be successful.

\paragraph{Data-driven PSF models}
This approach, also known as \textit{non-parametric PSF models}, only relies on the observed stars to build the model in pixel space. It is blind to the physics of the inverse problem. The models assume some regularity in the variation of the PSF field across the FOV. These methods differ in the way they exploit this regularity and handle the super-resolution, for example, \cite{liaudat2020, bertin2011, schmitz2020, jarvis2020, miller2013}. Data-driven methods can easily adapt to the current state of the optical system. However, they have difficulties in modeling complex PSF shapes. A limitation shared by all the data-driven methods is their sensitivity to the available number of stars to constrain the model. When the number of stars falls below some threshold, the model built is usually considered unusable. This family of methods has been widely used for modeling ground-based telescope PSFs. Nevertheless, they are not capable of successfully modeling the chromatic variations in addition to the spatial variations and the super-resolution.

\subsection{Our contribution}

In this paper, we propose a novel data-driven point spread function model coined WaveDiff. Our model can handle super-resolution and model spatial variations of wide-field telescopes. To the best of our knowledge, this is the first data-driven PSF model that can also capture the diffraction-induced chromatic variations of the PSF.

The proposed WaveDiff model represents a \textit{paradigm shift} with respect to current state-of-the-art \textit{data-driven} PSF models as it builds the model in the wavefront space. Although this is also the case for parametric models, WaveDiff remains data-driven. We can include the physics behind the PSF modeling inverse problem into the WaveDiff model using the \textit{differentiable optical forward model} that we have designed. The forward model is theoretically based on optics principles and practically implemented in a modern automatically differentiable framework that runs on GPU. The optical forward model represents a simplified version of the telescope's optical system and allows going from the wavefront to the pixel representation of the PSF. The forward model is fast and generic, as depends on a reduced number of parameters that are known for optical telescopes. It could be therefore easily adapted to model the PSF field of other telescopes. In addition, we propose a framework for training the WaveDiff model based on modern stochastic first-order optimization methods. Even though the WaveDiff model is built in the wavefront space, it is constrained only using noisy in-focus observations, without the need of using special calibration observations. The WaveDiff model can estimate a useful manifold in the wavefront space that is a good approximation of the underlying PSF field of the observations. We present three variations of the WaveDiff model that we refer to as its flavors.
This paper extends the work in \cite{liaudat2021,liaudat2021_b}.

The rest of the paper is organized as follows: \autoref{sc:psf_field_inv_problem} describes the inverse problem we are targeting, \autoref{sc:related_work} gives an overview of related work, \autoref{sc:wavefront_psf_model} presents the proposed model from the physical motivation to the practical implementation. In \autoref{sc:numerical_experiments} we describe the numerical experiments we have done to validate the WaveDiff model and benchmark it against state-of-the-art PSF models. We also include a description of the experiment setup and the presentation of the results. We discuss important points in \autoref{sc:discussions}, and finally conclude this study in \autoref{sc:conclusions}.

\subsection{Notation}
We use capital letters (e.g. $H$) for discrete functions where the codomain is $\mathbb{K}^{a \times b}$ for some $a,b \in \mathbb{N}$, and $\mathbb{K}$ is $\mathbb{R}$ or $\mathbb{C}$. We use calligraphic capital letters (e.g. $\mathcal{H}$) for continuous functions where the codomain is $\mathbb{K}$. The variables $(x,y)$ are assigned for the positions in the FOV, $[u,v]$ for positions in the focal plane, and $[\xi,\eta]$ for positions in the pupil plane.

\section{The PSF field and the inverse problem}
\label{sc:psf_field_inv_problem}

Let us define the PSF field for a particular image exposure as a function that has as inputs a position in the FOV and a specific wavelength, and outputs a monochromatic PSF. We denote this function by $I: \mathbb{R}^{2} \times \mathbb{R}_{+}  \to \mathbb{R}^{M \times M} $, where $M^2$ is the number of pixels of the desired super-resolved postage stamp. PSF modeling consists in building an estimator, $H$, of the PSF field, and then using it to output any monochromatic PSF at a set of target FOV positions $\{(x_j,y_j) \}_{j=1, \ldots, m_{\text{target}}}$.  The imaging instrument has a certain passband and several degradations that we need to take into account into our model in order to match our model with the observations.
A monochromatic PSF is a PSF image in a square postage stamp at a specific wavelength.

\subsection{The observation model}
We consider stars as point sources. Therefore, star observations provide samples of the PSF field at their corresponding positions. The spectral energy of the point source at each wavelength is not necessarily the same, and we have to take this into account for our observations. The spectral energy distribution (SED) gives us the normalized energy as a function of wavelength of a star at a FOV position $(x_i, y_i)$ or $\int_{\text{passband}} \text{SED}(x_i, y_i; \lambda)  \text{d}\lambda  = 1$. From now on, we assume the SEDs to be known. The following observational model,
\begin{equation}
    \bar{I}(x_i, y_i) = F_d \left\{ \int_{\text{passband}} \text{SED}(x_i, y_i; \lambda) \; I(x_i, y_i; \lambda) \; \text{d}\lambda \right\} + \mathbf{n}_{i}\;,
    \label{eq:poly_star}
\end{equation}
relates the star observation $\bar{I}(x_i, y_i) \in \mathbb{R}^{N \times N}$, where $N^2$ is the number of pixels in the observed postage stamp, at the FOV position $(x_i, y_i) \in \mathbb{R}^{2}$ to the objective PSF field $I(x_i, y_i; \lambda) \in \mathbb{R}^{M \times M}$, which is at a higher resolution with respect to the observations ($M>N$). The desired instrumental response $I$ is integrated in the instrument's passband weighted by the star's normalized spectral energy distribution $\text{SED}(x_i, y_i; \lambda) \in \mathbb{R}_{+}$. Then, it is degraded with the operator $F_d : \mathbb{R}^{M \times M} \to  \mathbb{R}^{N \times N}$,  which accounts for:
(i) a downsampling that affects the pixels in the stamp by a factor $D$ (meaning that $D N = M$),
(ii) a sub-pixel shift that depends on where the position of the centroid of the object is placed with respect to the pixel grid, and
(iii) other types of pixel-level degradations such as detector effects or guiding errors.
The observational noise is modeled by $\mathbf{n}_{i}$, which we assume to be white and Gaussian for simplicity, i.e. $\mathbf{n}_{i} \sim \mathcal{N}(0, \sigma_{i}^{2} \mathbf{I}_{N})$. We make a distinction between a monochromatic PSF $I(x_i, y_i; \lambda)$, that is evaluated at a single wavelength, and its polychromatic counterpart $\bar{I}(x_i, y_i)$, which has been integrated in wavelength over a passband. \autoref{fig:PSF_examples} presents examples of different representations of a PSF at a single FOV position.

\begin{figure}
    \begin{subfigure}{.24\textwidth}
      \centering
      \includegraphics[width=.9\linewidth]{figures/observation_example.pdf}
      \caption{}
      \label{fig:sub_a}
    \end{subfigure}%
    \begin{subfigure}{.24\textwidth}
      \centering
      \includegraphics[width=.9\linewidth]{figures/polyPSF_example.pdf}
      \caption{}
      \label{fig:sub_b}
    \end{subfigure}
    \begin{subfigure}{.24\textwidth}
      \centering
      \includegraphics[width=.9\linewidth]{figures/polyPSF_SR_example.pdf}
      \caption{}
      \label{fig:sub_c}
    \end{subfigure}
    \begin{subfigure}{.24\textwidth}
      \centering
      \includegraphics[width=.9\linewidth]{figures/opd_example.pdf}
      \caption{}
      \label{fig:sub_d}
    \end{subfigure}\\
    \begin{subfigure}{\textwidth}
      \centering
      \includegraphics[width=0.95\linewidth]{figures/chromatic_PSFs.pdf}
      \caption{}
      \label{fig:sub_e}
    \end{subfigure}
    \caption{\small Different representations of a space-like PSF at a single position in the field-of-view. The SED used corresponds to one of the templates from the numerical experiments described in \autoref{sc:experiment_set_up}. \textbf{(a)} Noisy polychromatic PSF at observation resolution, \textbf{(b)} Noiseless polychromatic observation at observation resolution, \textbf{(c)} High-resolution noiseless polychromatic observation, \textbf{(d)} Wavefront error map representing the aberrations in the optical system. Units are in $\mu$m, \textbf{(e)} High-resolution chromatic variations of the PSF at equally spaced wavelengths in the passband $[550, 900]$nm.}
    \label{fig:PSF_examples}
\end{figure}

\subsection{Inverse problem and regularizations}
\label{sc:inv_problem}

The PSF modeling problem consists in building a model $H: \mathbb{R}^{2} \times \mathbb{R}_{+}  \to \mathbb{R}^{M \times M}$  that best approximates $\{I(x_j, y_j; \lambda)\}_{j=1, \ldots, n_{\text{target}}}$ for all the target positions and wavelengths. To do so, we have as input a set of noisy degraded polychromatic observations $\{\bar{I}(x_i, y_i)\}_{i=1, \ldots, n_{\text{obs}}}$. The way we constrain our model with the observations is to minimize some type of reconstruction error between our model and the observations, for example, $\sum_{i=1}^{n_{\text{obs}}}\| \bar{I}(x_i, y_i) - \bar{H}(x_i, y_i) \|_{F}^{2}$, where $\| . \|_{F}$ denotes the Frobenius norm of a matrix. We need to consider the same degradations applied to $I(x_i, y_i; \lambda)$ in \autoref{eq:poly_star} to the model $H$ in order to match the observation. This inverse problem is ill-posed due to the degradations in $F_d$, the broad passband integration, and the noise. As a consequence, many PSF models can reproduce the observations.

To regularize the inverse problem, we can inject some type of prior knowledge we have about the PSF field. For example:

\begin{itemize}
\item \textit{PSF field regularity}: The most used regularization is to exploit the regularity of the PSF field in the FOV. This regularization imposes that two close-by positions in the FOV should have a close PSF representation. The regularity translates as a correlation in space (and time) of the PSFs. Most current PSF models rely on this fact to build a PSF model that can obtain a good generalization to target positions. Each PSF model imposes the regularity differently, e.g. \cite{liaudat2020, ngole2016, bertin2011, pca_psf2007}.

\item \textit{Positivity}: Given the way the PSF physically forms it should not contain negative values. The presence of negative values in the observations is due to the noise. Therefore, one can then consider the positivity of the PSF to regularize the problem.

\item \textit{Smoothness of the PSF}: One can consider that the PSFs are structured, and therefore the model should output smooth or piecewise-smooth images. One way to take this into account would be to use predefined parametric functions which are already smooth to build the PSF model. For example, the Moffat function \cite{moffat1988, li2016}. This choice can be effective in avoiding noise contamination to the model. However, it has severe difficulties to model PSF fields that are not well represented by the chosen base function. Another approach, proposed by \cite{ngole2016}, would be to use sparsity to enforce structure in the PSF image. In practice, this is used to denoise the model, so that the observed noise is not learned by the model. It can be achieved by enforcing a sparse representation of the PSF on some type of dictionary that can be built using, for example, isotropic undecimated wavelets, also known as starlets \cite{starck2011}. Although this technique proved to be flexible and adaptable to the observed data, the resulting PSFs are biased towards the properties of the dictionary. In the case of starlets, the resulting PSF would be biased towards isotropic features, smoothing out anisotropies of the PSF.

\item \textit{Manifold hypothesis}: From a general perspective, the manifold hypothesis \cite{fefferman2016} states that real-world data tends to lie in the vicinity of a low-dimensional manifold embedded in a higher-dimensional space. This hypothesis can be enforced in several ways depending on how the PSF model is built. For example, in \cite{ngole2016} the PSF model is built in a constrained matrix factorization framework. The PSF field is represented by a matrix factorization and the low dimensionality of the manifold is enforced by keeping a low rank in the factorization. The rank is left as a model hyperparameter.

\end{itemize}

Up to now, the properties mentioned were mild assumptions of the PSF modeling problem. One might notice that this prior knowledge is a consequence of the physics involved in the image formation of the telescope. For example, the CCD detectors measure the intensity of the electric field, which is naturally positive. Also, the smoothness and structure of the PSF come from the fact that it is the diffraction pattern of the incoming wavefront. Instead of building our PSF model over the assumptions mentioned above, we can take a step back and inject the physics of the optical system into our model. In practice, we account for the physics of the inverse problem with the addition of a differentiable optical forward model. The PSF modeling space is shifted from the pixels to the wavefront which is closer to the physics of the problem and a more natural space to model the PSF. With this change, we are not only considering the previously mentioned priors, but we are adding much more physical information to solve our inverse problem. This change is a fundamental difference in the way data-driven PSF models are built.

The PSFs generated using the optical forward model are naturally without negative pixel values as well as smooth. However, as already mentioned, current data-driven models can also account for the positivity and smoothness of the PSF. Then, what would be the advantage of building a wavefront-based model? Current data-driven models are only tackling consequences of the physical process involved in the inverse problem, for example, the positivity. The proposed wavefront-based model is incorporating the physics involved, the reconstructed pixel PSF is calculated as the intensity of the diffracted electric field in the focal plane. Furthermore, the optical forward model cannot reproduce observational noise which gives us a natural way of avoiding overfitting it without having to use any type of denoising technique.

It may seem that we are adding a complete model of the optical system as it is done in the parametric PSF models from \autoref{sc:approaches_PSF_modeling}. Nevertheless, we remain on a general framework that is data-driven as we will see in \autoref{sc:wavefront_psf_model}. We include in our model the diffraction phenomena which has been well studied, as well as general properties of the telescope like the focal length, the pupil diameter, the obscurations, and the pixel spacing. The optical forward model is end-to-end differentiable making use of an automatic differentiation framework. In addition, it uses the advances in stochastic first-order optimizations algorithms developed in recent years by the machine learning and deep learning community.

\section{Related work}
\label{sc:related_work}

\paragraph{Data-driven efforts to model chromatic variations}
Previous work attempted to solve the problem of modeling the chromatic variations from the data-driven side. An approach based on optimal transport \cite{schmitz2018} was developed in \cite{schmitz2019}. The chromatic variations were modeled as the displacement interpolation \cite{mccann1997}, or a weighted Wasserstein barycenter \cite{agueh2011}, between the monochromatic PSFs at the two extreme wavelengths. However, the results were not encouraging and the hypothesis taken by this approach might not hold for \textit{Euclid}. The authors in \cite{soulez2016} tried to model the propagation of the incoming wavefront from the stars through the different mirrors in the optical system. They propose to infer the aberrations introduced by each mirror by solving a phase retrieval type of problem. Finally, the proposed formulation is recast into a constrained optimization problem using the star observations in the FOV. However, it is a proof-of-concept article, as many difficulties remained unaddressed and results were only qualitative. None of these studies proved to be a satisfying solution.

\paragraph{Phase retrieval with automatic differentiation}
Estimating the wavefront of an in-focus observation falls in the category of phase retrieval problems. Recent works \cite{Wang2020, Wong2021}, based on a framework established by \cite{Jurling2014} are tackling the phase retrieval problem \cite{schechtman2015} using automatic differentiation. Their objective is solely to estimate the wavefront of a single image. By using a first-order optimization method over their forward model they are able to obtain a good estimate of the wavefront. In our PSF modeling problem, we do not have access to the PSF pixel image at the positions of interest. In addition, estimating the wavefront is not the final goal of the PSF model. We use the wavefront as an auxiliary product that is useful in our modeling purposes but not as an objective in itself. Therefore, these methods are not directly usable for our PSF modeling problem. Nevertheless, it illustrates the virtues and the utility of using automatic differentiation to tackle phase retrieval problems.

\paragraph{State-of-the-art data-driven PSF models}
Several state-of-the-art data-driven PSF models exist. Here we briefly review the three most relevant methods which we will later use for comparison. Although none of them is capable of modeling spectral variations, all of them can model spatial variations and handle super-resolution. The three methods are as follows:

\begin{itemize}
    \item \texttt{PSFEx}\footnote{\url{https://github.com/astromatic/psfex}} \cite{bertin2011} has been widely used in astronomy for weak-lensing surveys, e.g. \cite{zuntz2018}. It was designed to work together with the \texttt{SExtractor}\cite{bertin1996} software which builds catalogs from astronomical images and measures several properties of the observed stars. \texttt{PSFEx} models the variability of the PSF in the FOV as a function of these measured properties. It builds independent models for each CCD in the focal plane and works with polychromatic observations. It is not able to model the chromatic variations of the PSF field. The model is based on a matrix factorization scheme, where one matrix represents PSF features and the other matrix the feature weights. Each observed PSF is represented as a linear combination of PSF features. The feature weights are defined as a polynomial law of the selected measured properties. This choice allows having an easy interpolation framework for target positions. In practice, the properties that are generally used are both components of the PSF's FOV position. The PSF features are shared by all the observed PSFs and are learned in an optimization problem. The PSF reconstruction in the model's framework at a FOV position $(x_i,y_i)$ can be written as

    \begin{equation}
        \bar{H}^{\text{PSFEx}}(x_i,y_i) = F^{\text{PSFEx}} \left\{ \sum_{\substack{p,q > 0 \\ p+q \leq d}} x_{i}^{p} \, y_{i}^{q} \, S_{p,q} \,+\, S_{0} \right\} \,,
        \label{eq:psfex_model}
    \end{equation}

    where $S_{p,q}$ represents the learned PSF features, $S_{0}$ represents a first guess of the PSF, the polynomial law is defined to be of degree $d$, and $F^{\text{PSFEx}}$ represents the degradations considered to match the model with the observations. The first guess can be computed as a function of the median of all the observations. The PSF features are learned in an optimization problem that aims to minimize the reconstruction error between the \texttt{PSFEx} model and the observations which read

    \begin{equation}
        \min_{\substack{S_{p,q} \\ \forall p,q > 0 \,,\, p+q \leq d}} \left\{ \sum_{i=1}^{n_{\text{obs}}}  \left\| \frac{\bar{I}(x_i,y_i) - \bar{H}^{\text{PSFEx}}(x_i,y_i)}{\hat{\sigma}_{i}^{2}} \right\|_{F}^{2} + \sum_{\substack{p,q > 0 \\ p+q \leq d}} \left\| T_{p,q} S_{p,q} \right\|_{F}^{2} \right\} \,,
        \label{eq:psfex_optim}
    \end{equation}

    where $\hat{\sigma}_{i}^{2}$ represent the estimated per-pixel variances, $\bar{I}$ represents the noisy observations, and $\| \cdot \|_{F}$ the Frobenius norm of a matrix. The second term in \autoref{eq:psfex_optim} corresponds to a Tikhonov regularization, where $T_{p,q}$ represents some regularization weights to favor smoother PSF models. The PSF recovery at target positions is straightforward, one just needs to introduce new positions in the \autoref{eq:psfex_model} after having learned the PSF features $S_{p,q}$.

    \item \texttt{RCA}\footnote{\url{https://github.com/CosmoStat/rca}} \cite{ngole2016} is the state-of-the-art data-driven method designed for the \textit{Euclid} mission \cite{schmitz2020}. The model builds an independent model for each CCD, and, like \texttt{PSFEx}, is based on a matrix factorization scheme. However, there are two fundamental changes with respect to \texttt{PSFEx}. The first difference is that, in \texttt{RCA}, the feature weights are defined as a further matrix factorization. The feature weights are also learned from the data, and are constrained to be part of dictionary built with different spatial variations that are based on the harmonics of a fully connected graph. The graph is built using the star positions as the nodes, and using a function of the inverse of the distance between the stars to define the edge weights.
    The second difference corresponds to the regularizations used in the loss function and the optimization algorithms \cite{beck2009, condat2013}. \texttt{RCA} uses a positivity constraint, a denoising based on a sparsity constraint in the starlet \cite{starck2015} domain, which is a wavelet representation basis, and a constraint to learn the useful spatial variations from the graph-harmonics-based dictionary. See \cite{ngole2016} for more details.

    \item \texttt{MCCD}\footnote{\url{https://github.com/cosmostat/mccd}} \cite{liaudat2020} is a state-of-the-art data-driven method that was originally designed for the ground-based Canada-France Imaging Survey (CFIS)\footnote{\url{https://www.cfht.hawaii.edu/Science/CFIS/}} on the Canada-France-Hawaii Telescope (CFHT). The \texttt{MCCD} model extends the notions from the two methods described above and builds a single PSF model for the entire CCD mosaic in the focal plane. This choice allows for a more robust capture of spatial variations that extend over several CCDs. It is composed of global components that are shared between all CCDs, and local CCD-specific components that aim to correct for the imperfections of the global component. Both components are based on a matrix factorization framework, and the feature weights are defined as a combination of the polynomial strategy used in \texttt{PSFEx} and the graph-based strategy used in \texttt{RCA}. The model optimization is more complex, but it is also based on iterative schemes as the other two models. The constraints used in \texttt{MCCD} are motivated by the ones used in \texttt{RCA}. See \cite{liaudat2020} for more details.
\end{itemize}

\section{Data-driven wavefront PSF model}
\label{sc:wavefront_psf_model}

\subsection{Physical motivation and the wavefront error}
\label{sc:physical_motivation}

Current and future telescope optical systems require high optical performance, thus imposing very complex designs. To illustrate this, \autoref{fi:optics_a} presents an example of the optical system of the \textit{Euclid} space telescope from \cite{racca2016}. The PSF can be modeled with high accuracy and complexity using powerful commercial optic simulators. They depend on ray-tracing techniques to account for each optical surface. Although very accurate, these simulators are not practical for our problem due to their high complexity and are not suited for massive PSF computations as are required in survey-type missions. Nevertheless, we can consider a simpler optical system that is a good approximation of the actual system. It should be able to reproduce some effects of the real system while being more accessible to model. We use some ideas from parametric models \cite{krist2011}, and adopt a single converging lens system. Stars are well approximated by point sources at infinity and will be considered as plane waves when entering the system. An ideal optical system turns the plane waves into spherical waves, and the rays converge in a single point at the focal plane of the instrument. The exact point at which the rays converge depends on the angle of incidence of the incoming rays with respect to the optical axis.

\begin{figure}
    \centering
    \begin{subfigure}[b]{0.32\textwidth}
        \centering
        \includegraphics[width=\textwidth]{figures/euclid_optic_system.pdf}
        \caption{\label{fi:optics_a}}
    \end{subfigure}
    \begin{subfigure}[b]{0.65\textwidth}
        \centering
        \includegraphics[width=\textwidth]{figures/myWFEdiagram_v2.pdf}
        \caption{\label{fi:optics_b}}
    \end{subfigure}
    \caption{The figure in the left hand-side, (\subref{fi:optics_a}), shows a schematic functional view of the \textit{Euclid}'s payload module taken from \cite{racca2016}. The figure in the right-hand side, (\subref{fi:optics_b}), shows a one-dimensional illustration of the single-lens equivalent system were the wavefront error is visible.}
    \label{fi:wfe_explained}
\end{figure}

Due to imperfections and aberrations of the optical system, the outgoing wavefront is not exactly spherical, and the rays do not converge in a single point. The difference between the real aberrated wavefront and the ideal spherical wavefront is called the \textit{wavefront error} (WFE). \autoref{fi:optics_b} presents a simplified sketch of the WFE. The advantage of this simpler optical system is that we can take advantage of Fraunhofer's approximation of diffraction \cite[Chapter~4.3]{goodman2005} and its effect on the wavefront  \cite[Chapter~5.2]{goodman2005}. These approximations allow us to relate the electric field's propagation and diffraction between the pupil plane and the focal plane. The field at the focal plane is the two-dimensional Fourier transform of the field at the pupil plane. The sensors measure the intensity of the electric field at the focal plane. Therefore, we can relate the electric field at the pupil plane and the observed pixel PSF by taking the squared modulus of its Fourier transform. The monochromatic PSF at the focal plane can be written as
\begin{equation}
    \hat{\mathcal{H}}(x_i, y_i; \lambda)[u,v] \propto \left| \iint_{\mathbb{R}^{2}} \mathcal{U}_p(x_i, y_i; \lambda)[\xi, \eta] \exp\left[ \frac{- 2 \pi \text{i}}{\lambda f_{L}} (u \xi + v \eta) \right] \dd \xi \dd \eta \right|^{2} \;,
    \label{eq:mono_PSF_model_theory}
\end{equation}
where $f_L$ is the optical system's focal length, and $(x_i,y_i) \in \mathbb{R}^{2}$ is a FOV position that corresponds to a determined set of sky angular coordinates. We consider that for a specific FOV position and wavelength we have a given electric field at the pupil plane $\mathcal{U}_p(x_i, y_i; \lambda) : \mathbb{R}^{2} \to \mathbb{C}$. The two coordinates of the pupil plane that we integrate over are $[\xi, \eta]$, while the coordinates at the focal plane are $[u,v]$, where we observe our pixel PSF. We can rewrite the previous equation using the Fourier transform (FT) \cite{fourier1978} as
\begin{equation}
    \hat{\mathcal{H}}(x_i, y_i; \lambda)[u,v] \propto \left|  \text{FT}\left\{\mathcal{U}_p(x_i, y_i; \lambda) \right\} \left[ \frac{u}{\lambda f_L}, \frac{v}{\lambda f_L} \right] \right|^2 \;.
    \label{eq:mono_PSF_model_theory_2}
\end{equation}
Then, we need to specify the electric field at the pupil plane $\mathcal{U}_p$ that characterizes the aberrations of the system. The electric field writes
\begin{equation}
    \mathcal{U}_p(x_i, y_i; \lambda)[\xi, \eta] = \mathcal{P}(x_i, y_i; \lambda)[\xi, \eta] \exp \left[ \frac{2 \pi \text{i}}{\lambda} \phi_{\theta}(x_i, y_i)[\xi, \eta] \right] \;,
    \label{eq:elec_field_theory}
\end{equation}
where $\mathcal{P}(x_i, y_i; \lambda): \mathbb{R}^2 \to \mathbb{R}$ represents the obscurations encountered at the focal plane for a respective FOV position and wavelength, and $\phi_{\theta}(x_i, y_i): \mathbb{R}^2 \to \mathbb{R}$ represents the WFE or the aberrations of our optical system. It is common to have obscurations at the pupil due to the shape of the lenses and to the superposing of mirrors and supports. It is especially the case for telescopes due to occultation from a secondary mirror and its supporting structures. The obscurations are known and unlikely to change during the lifetime of the telescope, although their figure may vary with the FOV position. They significantly affect the shape of the PSF and are added to our model. For simplicity, we neglect the obscuration's $\mathcal{P}$ FOV variations. We focused on all-reflective-telescope designs and we neglect any wavelength-dependence of the WFE $\phi_{\theta}$. We note that such wavelength dependence could be added to our model, for example, to model a telescope with a diffractive element in the optical system.

To summarize, the WFE map is transformed into a complex electric field via \autoref{eq:elec_field_theory}. This field is propagated to the focal plane, and its intensity is computed to provide the pixel PSF with \autoref{eq:mono_PSF_model_theory_2}. This process of obtaining the pixel images from the WFE map is referred to as \textit{optical forward model}.

\subsection{Wavefront error PSF model}

In the following, we focus on how to use the WFE to model the PSF.
The core of the proposed framework is the WFE model, or WFE PSF model, $\Phi_{\theta}$, which is given by the following equation of the electric field at the pupil plane
\begin{equation}
    U_p(x_i, y_i; \lambda) = P \odot \exp\left[ \frac{2 \pi \text{i}}{\lambda} \; \Phi_{\theta}(x_i, y_i) \right]  \;,
    \label{eq:U_PSF_model}
\end{equation}
where $\odot$ is the Hadamard or element-wise product, $U_p(x_i, y_i; \lambda) \in \mathbb{C}^{K \times K}$, and $P, \Phi_{\theta}(x_i, y_i) \in \mathbb{R}^{K \times K}$. For the sake of simplicity, we use a fixed obscuration $P$ that does not depend on the FOV position and wavelength. The WFE PSF model depends on some parameters, represented with $\theta$, and its goal is to compute a WFE map at any position in the FOV. To conduct this task, we have to estimate, or \textit{learn}, the parameters $\theta$ of our model from the observed images using the full differentiable forward model that includes the degradations from the observational model in \autoref{eq:poly_star}. A schematic of the proposed framework is presented in \autoref{fi:psf_model_diagram}. We provide more detail of the forward model in \autoref{sc:optical_forward_model}.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/PSF_model_diagram_v5.pdf}
    \caption{A schematic of the proposed framework for data-driven wavefront PSF modeling.}
    \label{fi:psf_model_diagram}
\end{figure}

Our framework shifts a large part of the complexity from the PSF model into the forward model, which encodes all the diffraction phenomena and degradations. This change allows us to simplify the construction of the WFE PSF model block seen in \autoref{fi:psf_model_diagram}. This framework adds a physical meaning to the model as the WFE can be interpreted from a physical perspective. The WFE could even be measured directly on the telescope.

We next focus on a reliable generalization capability of the PSF to target positions, and the model's adaptation to observed data. For this reason we propose to use a weighted sum of wavefront features (or \textit{eigenWFE} if we draw a parallel with the notion of \textit{eigenPSF} \cite{liaudat2020, schmitz2020, ngole2016}). A wavefront feature is distributed across the FOV and, therefore, shared by all the PSFs at any FOV position. Then, the weight that goes with that wavefront feature will change depending on the FOV position. To facilitate the optimization of the model we consider two classes of wavefront features: fixed features that should provide a good representation basis of our WFE, and data-driven features that should be learned from the data and be able to more easily adapt to the observations. We can write our WFE PSF model as follows
\begin{equation}
    \Phi_{\theta}(x_i,y_i) = \underbrace{\sum_{l=1}^{n_{\text{Z}}} f^{\text{Z}}_l(x_i,y_i) \; S^{\text{Z}}_l}_{\Phi^{\text{Z}}(x_i,y_i)} + \underbrace{\sum_{k=1}^{n_{\text{DD}}} f^{\text{DD}}_k(x_i,y_i) \; S^{\text{DD}}_k}_{\Phi^{\text{DD}}(x_i,y_i)} \;,
    \label{eq:wfe_psf_model}
\end{equation}
where $f^{\text{Z}}_l, f^{\text{DD}}_k : \mathbb{R}^{2} \to \mathbb{R}$ and $S^{\text{Z}}_l, S^{\text{DD}}_{k} \in \mathbb{R}^{K \times K}$. The functions $f^{\text{Z}}_l, f^{\text{DD}}_k$ give the weights for the wavefront features $S^{\text{Z}}_l, S^{\text{DD}}_{k}$ that correspond to the fixed and the data-driven (DD) parts respectively. A fixed number of features is set for both parts, $n_{\text{Z}}$ and $n_{\text{DD}}$, for the fixed and the DD respectively. These hyperparameters allow to control the desired complexity of the model.

\subsubsection{Zernike polynomials as fixed features} Zernike polynomials \cite{noll1976} have been widely used in the optics community. These polynomials provide a good basis for modeling circular apertures as they are orthogonal on the unit disk. They are well suited to model the low-frequency variations of the WFE. \autoref{fi:zernike_example} shows the first $15$ Zernike polynomials that are function of the pupil coordinates $[\xi, \eta]$ on a circular aperture. The weights $f^{\text{Z}}_l$, for each Zernike polynomial $S^{\text{Z}}_l$, are defined as independent polynomial variations as a function of the FOV position, $(x_i, y_i)$. We include polynomials up to a maximum degree $d_{\text{Z}}$. For example, a maximum degree of $1$ gives us $f^{\text{Z}}_l(x_i, y_i) = c^{\text{Z}}_{l,0} + c^{\text{Z}}_{l,1} x_i + c^{\text{Z}}_{l,2} y_i$, where $\{ c^{\text{Z}}_{l,j} \}_{j=1, 2, 3}$ are the coefficients for the Zernike index $l$. The total number of parameters to estimate if we use $n_{\text{Z}}$ Zernike polynomials is $n_{\text{Z}} (d_{\text{Z}} + 1)(d_{\text{Z}} + 2)/2$. Note that the Zernike polynomial definition is in pupil plane coordinates $[\xi, \eta]$ (e.g.~the maps seen in \autoref{fi:zernike_example}). However, the weight accompanying each Zernike polynomial, $f^{\text{Z}}_l$, is a polynomial of FOV positions $(x_i, y_i)$.

\begin{figure}
    \centering
    \includegraphics[width=0.9\textwidth]{figures/zernike_maps_article.pdf}
    \caption{The first $15$ Zernike polynomials with their Noll's single-index prefaced by the hash symbol $\#$, and ordered vertically by their radial degree $n$ and horizontally by their azimuthal degree $m$.  \nblink{Zernike-example-plots}}
    \label{fi:zernike_example}
\end{figure}

\subsubsection{Data-driven features}
\label{sc:data_driven_features}
The data-driven features, represented by the $S^{\text{DD}}_{k}$ matrices, are non-parametric. We do not impose a specific structure on the features and these are entirely learned from the data. These DD features should adapt to capture variations that are not well modeled by the Zernike-based features and should also account for unexpected mismatches between our forward model and the ground truth model.

It is important to maintain a certain regularity in the PSF field and have a good generalization capability over the full FOV. The underlying hypothesis under the PSF field regularity is a consequence of the optical system that can be seen in \autoref{fi:optics_a}. The system is composed of mirrors of different sizes. The way the mirror's footprints propagate to the focal plane of the instrument imposes a regularity in the PSF field. The simplified model we propose in \autoref{fi:optics_b} cannot inherently account for this propagation, and it is necessary to impose the regularity in the $f_{k}^{\text{DD}}(x,y)$. The way we impose this regularity will define different flavors of the proposed WaveDiff model:
\begin{itemize}
    \item \textbf{WaveDiff-original:} the feature weights are defined as FOV position polynomials up to a degree $d_{\text{DD}}$. This is a constraining but effective choice and is the same one as for the fixed features.

    \item \textbf{WaveDiff-graph:} we want to add constraints to the feature weights so that we can capture smooth and localized spatial variations of the PSF field. To achieve this goal we follow the approach presented in \cite{ngole2016}, and already applied in \cite{schmitz2020, liaudat2020}. We aim to capture variations that occur at different spatial frequencies. Unfortunately, the fact that the PSF positions in the FOV are randomly distributed makes it hard to define the different types of variations. To tackle this issue, we define a fully-connected undirected weighted graph where each node corresponds to the position of a star. The weights are built as a parametric function of parameters $\Lambda$ of the distance between the different stars. Once the graph is defined, we compute its Laplacian matrix $L_{\text{p}, \Lambda} \in \mathbb{R}^{n_{\text{stars}} \times n_{\text{stars}}}$ \cite{chung1997}, that encodes all the weights and connections in the graph. This matrix is real and symmetric, and we can decompose it as $L_{\text{p}, \Lambda} = V_{\Lambda} \Sigma_{\Lambda} V_{\Lambda}^{\top}$, with $\Sigma_{\Lambda}$ being a diagonal matrix containing the eigenvalues. The eigenvectors of the Laplacian matrix, $V_{\Lambda}$, represent the graph harmonics or a basis to define the graph Fourier transform \cite{ricaud2019}.

    From these, a dictionary is built that is used to factorize the feature weights, as follows. In order to capture more spatial variations, we use $r$ different parameters, $\{\Lambda_1, \ldots,  \Lambda_r\}$ for the graph's weight function. These parameters define different graphs and subsequently different harmonics. We concatenate the graph-harmonics, or Laplacian eigenvectors, of all the graphs to build the graph dictionary that we define as $V = [V_{\Lambda_1}, \ldots, V_{\Lambda_r}]$.
    %
    We write $\mathbf{f}_{k}^{\text{DD}} = [f_{k}^{\text{DD}}(x_1,y_1),\, \ldots\,, f_{k}^{\text{DD}}(x_{n_{\text{stars}}},y_{n_{\text{stars}}})] \in \mathbb{R}^{1 \times n_{\text{stars}}}$, where the weight function of the DD feature $k$ is expressed as a vector of all star positions.
    Then, we factorize the feature weight vector as $\mathbf{f}_{k}^{\text{DD}} = \bm{\alpha}_k V^{\top}$, where $\bm{\alpha}_k \in \mathbb{R}^{1 \times r n_{\text{stars}}}$. The DD feature weights are build as a linear combination of graph-harmonics, that are represented as row vectors in $V^{\top}$. The vector $\bm{\alpha}_k$ that defines the weights of this combination.
    %
    We impose sparsity on the representation vector so that the learned feature can specialize on a specific, or a small subset of graph harmonics. The sparsity in the $\bm{\alpha}_k$ vector can be achieved by adding an $\ell_1$ constraint in the final loss function. We call it the \textit{graph constraint}. The optimization algorithm will be in charge of selecting the best graph-harmonics for a specific feature that allows obtaining a good representation of our observations.

    For a more detailed description of how we define the graph weights and the choice of the $\Lambda$ parameters see \cite[Sections~5.2,~5.5.3]{ngole2016}.

    \item \textbf{WaveDiff-polygraph:} the feature weights simultaneously use the two previously mentioned constraints. It contains two contributions one based on FOV position polynomials and the other based on the graph constraint.
\end{itemize}

\subsection{Practical optical forward model}
\label{sc:optical_forward_model}
The choice of the proposed model and Fourier optics allows computing the forward modeling of the wavefront to the pixels using a Fourier Transform, as described in \autoref{sc:physical_motivation}. Given the vast amount of data to process in upcoming missions, it is important to develop fast methods. This requirement motivates the use of the Fast Fourier Transform (\texttt{FFT}) algorithm. \autoref{eq:mono_PSF_model_theory_2} shows that the sampling of the wavefront space changes with wavelength ($\lambda$). This observation implies that a resampling of the space is required for each wavelength if we want to work with a fixed space $\mathbb{C}^{K \times K}$. It is highly impractical as we need to calculate several monochromatic PSFs, at different wavelengths, to be able to compute a single polychromatic PSF that is integrated in the instrument's passband. Furthermore, we have to take into account that the input and output matrices of the \texttt{FFT} algorithm are of the same dimension. We want to be Nyquist sampled at all wavelengths, to avoid any type of aliasing problem.

One way to accelerate the calculations is to fix the dimensions of the electric field at the pupil plane from \autoref{eq:U_PSF_model} as $\mathbb{C}^{K \times K}$, and then add a wavelength-dependent zero padding. By adding zeros in the obscured part, and fixing the non-obscured dimensions, we can control the sampling of the electrical field at the pupil plane. This practical method avoids a resampling of the electrical field for each wavelength and allows to be Nyquist sampled. We zero-pad the electrical field matrix from $\mathbb{C}^{K \times K}$ to $\mathbb{C}^{p(\lambda) \times p(\lambda)}$. The required dimensions for the input matrix of the \texttt{FFT} operation can be seen in the following equation
\begin{equation}
    p(\lambda) = \frac{K \; \lambda \; f_{\mathrm{L}}}{M_{\mathrm{D}} \; \Delta} \; Q \;,
    \label{eq:wavelength_dep_zero_pad}
\end{equation}
where $K$ is the fixed matrix dimension of the electrical field that is equal to the pupil diameter in pixels, $f_{\mathrm{L}}$ is the focal length, $M_{\mathrm{D}}$ is the length of the pupil diameter, $\Delta$ is the length of a squared-pixel in the focal plane, $\lambda$ is the wavelength, and $Q$ is a dimensionless oversampling factor. See \ref{sc:app_oversamp_formula} for the derivation of the \autoref{eq:wavelength_dep_zero_pad} were we follow \cite{schmidt2010}. In practice, we fix the telescope's parameters, we choose the dimension $K$ we will work with, and then we calculate the required $Q$ so that we respect the Nyquist criterion ($p(\lambda) \geq 2 K$) for every $\lambda$ in our passband. Following the variable zero padding, each computed monochromatic PSF will have different dimensions. This issue can be solved by an afterward crop to a fixed size of $\mathbb{R}^{M \times M}$. The energy contribution of the cropped pixels is negligible given that $M$ is sufficiently big. The propagation of the wavefront PSF to the oversampled pixel PSF can be resumed as follows
\begin{equation}
    \hat{H}(x_i, y_i; \lambda) \propto \underset{p(\lambda) \times p(\lambda) \to M \times M}{\text{crop}} \underbrace{\Bigg| \texttt{FFT}\bigg\{ \underset{K \times K \to p(\lambda) \times p(\lambda)}{\text{pad}} \Big( \overbrace{P \odot \exp\Big[ \frac{2 \pi i}{\lambda} \underbrace{\Phi_{\theta}(x_i, y_i)}_{\text{WFE PSF model }} }^{\text{Electric field}} \Big) \Big] \bigg\} \Bigg|^{2}}_{\text{Pixel representation}} .
    \label{eq:mono_PSF_model}
\end{equation}
\autoref{eq:mono_PSF_model} represents the optical system box in \autoref{fi:psf_model_diagram} that gives us access to the monochromatic oversampled pixel PSF. We still need to apply the degradations to the previous pixel PSF so that it can match our observations. This action resumes in applying the observational model seen in \autoref{eq:poly_star} to the monochromatic oversampled PSF. We first need to normalize the PSF so that it has a unit pixel sum, or $\sum\sum_{l,m = 1}^{M} \hat{H}(x_i, y_i; \lambda)[u_l,v_m] = 1$. We can then integrate over the instrument's passband and downsample to match the instrument's sampling, which is resumed as follows
\begin{equation}
    \hat{\bar{H}}(x_i, y_j) = D_\text{D} \left\{ \sum_{k=1}^{n_{\lambda}} \hat{H}(x_i, y_j; \lambda_k) \; \text{SED}(x_i, y_i; \lambda_k) \right\} \; ,
    \label{eq:spectral_approx}
\end{equation}
where we discretized the integral in \autoref{eq:poly_star} in $n_{\lambda}$ evenly spaced bins, and $D_\text{D}$ represents a downsampling by a factor $\text{D}$ that can include a further crop on the postage stamp. To match the observation sampling we need to use $\text{D} = Q$. The SED in \autoref{eq:spectral_approx} has been normalized so that for a given $n_{\lambda}$ and position $(x,y)$, it verifies $\sum_{k=1}^{n_{\lambda}} \text{SED}(x, y; \lambda_k) = 1$. As one can see from \autoref{eq:spectral_approx}, we are not including further detector effects for the sake of simplicity. We use an oversampling factor $Q \in \mathbb{N}$ as the implementation of the downsampling operator is simplified.

\subsection{Training}
\label{sc:training}
The parameters $\theta$ we need to estimate are contained in the WFE PSF model $\Phi_{\theta}$ from \autoref{eq:wfe_psf_model}. For the fixed features, we optimize with respect to $f_l^{\text{Z}}$, and for the data driven features with respect to $f_k^{\text{DD}}$ and $S_{k}^{\text{DD}}$. We assume white Gaussian noise as in \autoref{eq:poly_star} and propose to use the following loss function to optimize our model
\begin{equation}
    L(\theta) = \frac{1}{n_{\text{stars}}} \sum_{i=1}^{n_{\text{stars}}} \frac{1}{\hat{\sigma}_i} \left\| \hat{\bar{H}}(x_i, y_j; \Phi_{\theta}) - \bar{I}(x_i, y_i) \right\|^{2}_{F} + \mathcal{R}(\theta) \; ,
\end{equation}
where $\hat{\sigma}$ is an estimation of the observed star noise standard deviation. The estimation is carried out by using the median absolute deviation (MAD) \cite{starck2015}. We use the pixels from the outer region of the squared postage stamp, where we have masked the circular central region that contains most of the star's energy. We represent a set of regularizers on the model's parameters $\theta$ with $\mathcal{R}(\theta)$. Depending on which WaveDiff flavor we are using, the regularizer can take different forms. For the WaveDiff-graph model we use a $\mu \| \bm{\alpha}_k \|_{1}$ regularization to impose sparsity on the different $\bm{\alpha}_k$ vectors that were defined in \autoref{sc:data_driven_features}, and are used to build the DD feature weights. The hyperparameter $\mu$ allows us to control the strength of the regularization. We use the same regularization for the graph part in the WaveDiff-polygraph model. However, we do not use any regularization for the WaveDiff-original model, and we discuss this choice in \autoref{sc:discussions}.

The optimization algorithm used ins practice is Rectified Adam \cite{liu2020}, a stochastic gradient-based method widely used in the machine learning community. The data-driven features and the feature weight parameters are initialized using a uniform distribution with zero mean. The bounds of the distribution are kept close to zero, $-1$nm and $1$nm, so that the starting PSF model contains a small amount of aberrations, close to a perfect system, and almost no spatial variations. This choice corresponds to starting with an almost perfect system, close to a vanishing WFE. However, the small non-zero aberration values help to escape local minima at the start of the optimization algorithm. The optimization adds aberrations to the system until the observed PSF field is well represented. We have found that the best working optimization strategy is to iterate the optimization between the parameters of the parametric and the data-driven parts. This strategy means that we first fix the data-driven parameters, and optimize the parametric part. We then fix the parametric part and optimize the data-driven parameters. We repeat these iterations until convergence. We use different learning rates and number of epochs for each optimization sequence. Algorithm \ref{al:wavediff_training} presents a summary of the training procedure.

We note the versatility of the proposed PSF model framework using automatic differentiation. If prior information is available, it can be added as regularizations at the wavefront level or the pixel level. The framework can be adapted for other types of imaging applications, and other task-specific regularizers can be added to the loss. As the model is developed in the TensorFlow framework \cite{tensorflow2015}, making these adaptations is straightforward.

\begin{algorithm
    \caption{WaveDiff training procedure}
    \begin{algorithmic}[1]
    \label{al:wavediff_training}

    \STATE \underline{Hyperparameters}:
    \\ Learning rates $\{\eta_{m}^{\text{DD}} \}_{m=1, \ldots, m_{\max}}$, $\{\eta_{m}^{\text{Z}} \}_{m=1, \ldots, m_{\max}}$
    \\ Number of epochs $\{N_{\text{ep}, m}^{\text{DD}} \}_{m=1, \ldots, m_{\max}}$, $\{N_{\text{ep}, m}^{\text{Z}} \}_{m=1, \ldots, m_{\max}}$
    \\ Regularization strength $\mu$ \hfill (Graph flavors only)
    \bigskip

    \STATE \underline{Initialization}:
    \\ Data-driven features:  $S_{k}^{\text{DD}} \sim \mathcal{U}[-10^{-3}, 10^{-3}]^{K \times K}, \forall k$.
    \\ Feature weight parameters: $c^{\text{DD}}_{k,j}, c^{\text{Z}}_{l,j} \sim \mathcal{U}[-10^{-2}, 10^{-2}], \forall l,k,j.$
    \\ Estimate noise level, $\hat{\sigma}_i \leftarrow \text{MAD noise estimation} (\bar{I}(x_i, y_i))$
    \\ Graph constraint initialization following \cite[\S5.2,\S5.5.3]{ngole2016} \hfill (Graph flavors only)
    \\ Generate Zernike polynomial maps $S_{l}^{\text{Z}}, \forall l$ and telescope obscuration $P$

    \bigskip
    \STATE \underline{Alternate optimization}:
    \FOR{$m=1$ to $m_{\max}$}
        \STATE Optimize $L(\theta)$ over $c^{\text{Z}}_{l,j}, \;\forall l,j$ using $\eta_{m}^{\text{Z}}$, $N_{\text{ep}, m}^{\text{Z}}$ \hfill (parametric part)
        \STATE Optimize $L(\theta)$ over $c^{\text{DD}}_{k,j}, S_{k}^{\text{DD}}, \;\forall k,j$ using $\eta_{m}^{\text{DD}}$, $N_{\text{ep}, m}^{\text{DD}}$ \hfill (data-driven part)
    \ENDFOR

    \bigskip
    \STATE Return PSF model: $\Phi_{\theta}$

    \end{algorithmic}
\end{algorithm}

\subsection{PSF recovery}
\label{sc:inference}
PSF recovery consists in using the trained PSF model to estimate the PSF at other positions in the FOV where star observations are not available. The recovery requires to specify how the feature weight functions can extrapolate to different positions in the FOV. Then, it is possible to reuse the learned and fixed WFE features from \autoref{eq:wfe_psf_model} to calculate the WFE at the FOV position of interest.

In the case where the weight functions are defined as FOV polynomials, as in the parts involving Zernike polynomials or the WaveDiff-original, the procedure is straightforward. It consists in evaluating the polynomial on a new position obtaining the new feature weights.

In the case of the WaveDiff models with graph flavors, the PSF recovery is not direct. During the optimization, it is not only the DD features being learned but also their corresponding spatial variations that depend on the graph constraint, described in \autoref{sc:data_driven_features}. We want to be able to recover the PSF at any position in the FOV while respecting the learned spatial variations. To accomplish this, we consider the different weight functions $f_{k}^{\text{DD}}$, or the vector $\mathbf{f}_{k}^{\text{DD}}$, learned for each feature $S_{k}^{\text{DD}}$. These encompass the learned spatial variations that contain a selection of useful graph harmonics. We follow the strategy from \cite{liaudat2020} and use the learned $f_{k}^{\text{DD}}$ at different FOV positions to build a local interpolant. Then, we use it to interpolate the feature weights to the desired position. In practice, we use a Radial Basis Function (RBF) interpolation scheme. For a new position $(x_j, y_j)$, we select the closest $N_{\text{RBF}}$ training star positions $\{ (x_i, y_i) \}_{i=1, \ldots, N_{\text{RBF}}}$. The interpolant function of the weights corresponding to feature $k$ is built using the following equation
\begin{equation}
    g_k(x,y) = \sum_{i=1}^{N_{\text{RBF}}} \lambda_{k,i} \, K\left(\left\| [x - x_i \,,\, y - y_i]^{\top} \right\|_{2}\right),
    \label{eq:rbf_interpolation function}
\end{equation}
where $K: \mathbb{R}_{+} \to \mathbb{R}$ is the interpolation kernel or radial basis function, that takes as input the distance between the desired and the training position. The weights that define the interpolant $g_k$ are $\{ \lambda_{k,i} \}_{i=1, \ldots, N_{\text{RBF}}}$. We learn these weights by using a set of exact interpolation constraints in the training positions that writes $\{g_k(x_i,y_i) = f_{k}^{\text{DD}}(x_i,y_i) \}_{i=1, \ldots, N_{\text{RBF}}}$. We follow \cite[\S3.7]{liaudat2020} and use a thin plate kernel that is defined as $K(r) = r^2 \ln(r)$, and set $N_{\text{RBF}}$ to $20$. To resume, we have to build one interpolant function for each feature $k$ and each new position $(x_j, y_j)$. Even though it may seem that it is a costly procedure in terms of time and resources, the fact that the interpolation is done using $1$-D functions makes it fast and cheap.

\section{Numerical experiments}
\label{sc:numerical_experiments}
In this section, we start by defining the objectives of the numerical experiments. Then, we introduce the setup of the experiment and the data set we use to accomplish our goals. Finally, we present and discuss the results we find.

\subsection{Objectives}
To evaluate the novel framework and show its utility, we consider a scenario where unexpected physical phenomena introduced some complexity into the PSF field. This choice means that we are not specifically accounting for the increase in complexity in our model. Another interpretation is that we underestimated in our model the complexity of the PSF field. The practical way we set up the experiment is to simulate the WFE of a PSF field that is more complex than what the models can reproduce in the WFE space. In this setting, none of the models are able to have a zero reconstruction error over the ground truth WFE, even on a noiseless scenario. After having simulated the complex WFE field, we can generate the required pixel PSFs using the optical forward model. Note that the added complexity is represented in the spatial variations of the WFE field and not in a change of the optical forward model.

This chosen setup allows comparing different model performances and observing to which degree the models are capable of building an approximation of a complex WFE PSF field. The setup is also useful to study at which point the built approximation is good in terms of pixel PSF reconstruction error. We use this setting to study the following items:
\begin{itemize}
    \item[\textit{a)}] Pixel reconstruction performance at the observation resolution and at $3$ times the observation resolution (super-resolution) for polychromatic images.

    \item[\textit{b)}] Pixel reconstruction performance at the observation resolution as a function of wavelength for the best performing model of \textit{a)}.
    \item[\textit{c)}] Pixel reconstruction performance as a function of the total number of observed stars in FOV that are used to constrain the PSF models.
    \item[\textit{d)}] Error in the recovered WFE with respect to the ground truth WFE of the PSF field.
    \item[\textit{e)}] Errors using weak-lensing metrics.
\end{itemize}

\subsubsection{Reconstruction metrics}
The main metric used to compare the performance is the residual root mean square error (RMSE). As we will work with simulations, we have access to the ground truth PSF field. Therefore, we can calculate noiseless PSFs at any resolution and FOV position, which facilitates the performance comparison. We define the root mean square (RMS) of a matrix $A_i \in \mathbb{K}^{K \times K}$, or a vector in $\mathbb{K}^{K^2}$, as
\begin{equation}
    \text{RMS}(A_i) = \left( \frac{1}{K^2} \sum_{j=1}^{K^2} \left|A_{i;j}\right|^{2} \right)^{1/2} ,
\end{equation}
where we use a single index that corresponds to a flattened matrix. Let us consider a set of $n$ noiseless $K \times K$ dimensional images $\{ A_{i} \}_{i=1}^{n}$, and the corresponding reconstructions $\{ \hat{A}_{i} \}_{i=1}^{n}$. The main pixel reconstruction errors we use are the absolute error, $Err_{\text{abs}}$, and the relative error, $Err_{\text{rel}}$, that are defined as
\begin{equation}
    Err_{\text{abs}} = \frac{1}{n} \sum_{i=1}^{n} \text{RMS}(A_{i} - \hat{A}_{i}) \;,\quad Err_{\text{rel}} = \frac{1}{n} \sum_{i=1}^{n} \frac{\text{RMS}(A_{i} - \hat{A}_{i})}{\text{RMS}(A_{i})}  \times 100\% \;.
\end{equation}
A constant value over all the WFE matrix does not affect the pixel PSF.
Therefore, we remove the mean of each WFE matrix before computing the RMS of the residual in the WFE error calculation.
We only use the non-obscured elements of the WFE matrix to compute the RMS.

\subsubsection{Weak-lensing metrics}
Weak lensing sets the most stringent requirements on the errors of the PSF models. Cosmologists have developed formulations to relate the PSF errors expressed in terms of shape and size metrics \cite{massey2012} to the errors of the cosmological parameters of interest \cite{cropper2013}. Therefore, the PSF model requirements are usually set up in terms of these shape and size metrics. These are defined in terms of the moments of the polychromatic observation $\bar{I}[u, v]$, following \cite{hirata2003}, we define
\begin{align}
    \label{eq:psf_moments_1}
    \bar{\mu} & = \frac{\int \mu \; \bar{I}[u, v] \; w[u, v] \; \text{d}u \, \text{d}v }{\int \bar{I}[u, v] \; w[u, v] \; \text{d}u \, \text{d}v},  \\
    M_{\mu \nu} &= \frac{\int \bar{I}[u, v] \; (\mu - \bar{\mu}) \; (\nu - \bar{\nu}) \; w[u, v] \; \text{d}u \, \text{d}v}{\int \bar{I}[u, v] w[u, v] \; \text{d}u \, \text{d}v},
    \label{eq:psf_moments_2}
\end{align}
where $\mu , \nu \in \{u, v\}$ and $w[u, v]$ is a weight window that is used to avoid the divergence of the integrals due to noise. \autoref{eq:psf_moments_1} defines the first-order moments, while \autoref{eq:psf_moments_2} defines the second-order moments. The size metric is defined as
\begin{equation}
    R^{2} = T = M_{u u} + M_{v v},
\end{equation}
and the shape metrics, or ellipticities, are defined as
\begin{equation}
    e = e_1 + {\rm i} e_2 = \frac{(M_{u u} - M_{v v}) + {\rm i} \, 2 M_{u v}}{T} .
\end{equation}
The method used to estimate these metrics is the widely-used adaptive moment algorithm from \texttt{Galsim}'s HSM module\footnote{\url{https://github.com/GalSim-developers/GalSim}} \cite{hirata2003,mandelbaum2005}. The adaptive moment algorithm measurement provides $\sigma$ as size which relates to the above-defined size metric as $R^2 = 2\sigma^2$. The measurements are carried out on super-resolved polychromatic images.

To calculate the error metrics, we first have to measure the ellipticity and size parameters using the adaptive moments' algorithm over each image from the sets of ground truth images $\{ A_{i} \}_{i=1}^{n}$, and reconstructions $\{ \hat{A}_{i} \}_{i=1}^{n}$. Then, the measurements are used to build vectors as for example $\mathbf{e}_{1,2} = [e_{1,2; 1}, \ldots, e_{1,2; n}]$. Finally, the error metrics used are the RMSE over each metric that can be written as
\begin{equation}
    \text{RMSE}(\mathbf{e}_{1,2}) = \text{RMS}(\mathbf{e}_{1,2} - \hat{\mathbf{e}}_{1,2})\;, \quad \text{RMSE}(\mathbf{R}^{2}) = \frac{\text{RMS}(\mathbf{R}^{2} - \hat{\mathbf{R}}^{2})}{\langle \mathbf{R}^{2} \rangle},
\end{equation}
where the size metric is divided by its mean as it is usually done in the literature.

\subsection{Setup}
\label{sc:experiment_set_up}

To demonstrate our novel framework, we simulate a simplified FOV with $2000$ star observations for training and $400$ noiseless target stars for testing. All positions are randomly distributed in the FOV. To evaluate the impact of the number of stars in the PSF field on the performance, we build three subsets $\mathcal{S}_1$, $\mathcal{S}_2$, and $\mathcal{S}_3$ containing $200$, $500$, and $1000$ stars, respectively. To provide a fair comparison and be less sensitive to randomness, the stars are assigned to the subsets so that they verify $\mathcal{S}_1 \subset \mathcal{S}_2 \subset \mathcal{S}_3 \subset \mathcal{S}$, where $\mathcal{S}$ is the data set with the entire training stars. The observations have a variable signal-to-noise ratio (SNR) that is uniformly distributed in the range $[10, 110]$. The SNR definition we use is
\begin{equation}
    \text{SNR} = \frac{\left\| \bar{I} \right\|_{F}^{2}}{\sigma_{\bar{I}}^2 N^2}, \quad \left\| \bar{I} \right\|_{F}^{2} = \sum_{i,j = 1}^{N,N}  | \bar{I}_{i,j} |^{2},
\end{equation}
where $\bar{I} \in \mathbb{R}^{N \times N}$ is an observation corrupted with white Gaussian noise with standard deviation $\sigma_{\bar{I}}$, and the operation $\| \cdot \|_{F}^{2}$ denotes the Frobenius matrix norm.

We use the parameters of an optical model close to the \textit{Euclid}'s VIS instrument model \cite{laureijs2011, venancio2020}. We consider a broad passband of $[550, 900]$nm, a telescope's focal length of $f_{\text{L}}=24.5$m, a telescope pupil diameter of  $M_{\text{D}}=1.2$m, a squared pixel size of $\Delta=12\mu$m\footnote{This corresponds to $0.1$ arcsecond}, and an oversampling factor of $Q=3$. The obscuration used is fixed for every position in the FOV and is built as a superposition of simple geometrical shapes as it can be seen in \cite[Fig. 7]{venancio2020}. A moving average filter is applied to the obscuration in order to smooth the transitions between the obscured and illuminated parts to avoid aliasing. We set $n_{\lambda}=20$ wavelength bins for the approximation in \autoref{eq:spectral_approx}. The dimensions used are $K=256$ for the WFE space, $M=64$ for the high-resolution pixel space, and $N=32$ for the observation, or low resolution, pixel space.

To build polychromatic observation the stars need SED information. We randomly choose for each star one of $13$ stellar SED templates from \cite{pickles1998} following the work in \cite{kuntzer2016}. More information about the SEDs can be found in \ref{sc:stellar_SEDs}. The SED data is considered to be a perfectly-known input to the PSF model. The ground truth (GT) PSF field is built using $45$ Zernike polynomials with a $d_{\text{Z}}$ of $2$ for each mode. The polynomial coefficients are randomly chosen such that the total aberration at any position in the FOV is close to a nominal value. The total amount of aberrations can be characterized by the RMS value of the WFE map, and the nominal value is defined as $100$nm. No further detector effects are considered in this work. The instrument's focal plane is constituted of a squared matrix of squared CCD chips or CCD mosaic \cite{cropper2016}. This work does not consider any discontinuity between the different CCD chips and assumes that the CCDs are perfectly aligned in the focal plane and share the same properties. The diffraction phenomena are considered as the only sources of chromatic variation of the PSF field. \autoref{fi:dataset_example} presents four PSF examples from the simulated data set, where the observed stars and their corresponding WFE is shown.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/observed_stars.pdf}
    \includegraphics[width=\textwidth]{figures/observed_OPD.pdf}
    \caption{Example of PSFs from the simulated data set. The first row contains images of observed stars, while the second row the corresponding WFE map used to generate each one of them. \nblink{DD-features-dataset-examples}}
    \label{fi:dataset_example}
\end{figure}

The following models are compared:
\begin{itemize}
    \item[\textit{i)}] \textit{Zernike 15}: a model using the proposed differentiable framework, but without any data-driven features. Only the fixed features, Zernike polynomials, are being used. We use $n_{\text{Z}}=15$, and $d_{\text{Z}}=2$. In the context of this study, this model represents a parametric approach where its parameters were badly specified. As a consequence, the complexity of the model is far away from the ground truth PSF field.
    \item[\textit{ii)}] \textit{Zernike 40}: a model similar to \textit{i)}, with $d_{\text{Z}}=2$, but using $n_{\text{Z}}=40$. In the context of this study, this model represents a parametric model where its parameters were slightly misspecified. This model is close to the ground truth PSF field. However, it is still lacking complexity as the number of Zernike polynomials is lower than the ground truth.
    \item[\textit{iii)}] \textit{WaveDiff-original}: the model described in \autoref{sc:data_driven_features} using $n_{\text{Z}}=15$, $d_{\text{Z}}=2$, and $d_{\text{DD}}=5$ that corresponds to $n_{\text{DD}}=21$. This model uses the same badly specified parameters of the fixed features as the model \textit{i)}.
    \item[\textit{iv)}] \textit{WaveDiff-graph}: a similar model as \textit{iii)} with a badly specified number of Zernike polynomials, using $n_{\text{Z}}=15$ and $d_{\text{Z}}=2$. However, this model uses the previously described graph constraint for the data-driven features. In order to perform a fair comparison we use $n_{\text{DD}}=21$.
    \item[\textit{v)}] \textit{WaveDiff-polygraph}: a similar model as \textit{iii)} and \textit{iv)}, using $n_{\text{Z}}=15$ and $d_{\text{Z}}=2$. For the data-driven features, this model uses $d_{\text{DD}}=3$ that corresponds to $10$ features, and $10$ more features corresponding to the graph constraint.
    \item[\textit{vi)}] \texttt{PSFEx}: The choice of the model parameters is based on the numerical experiment done in \cite{schmitz2020}. The model is built independently on each CCD. The spatial variations of the PSF, \texttt{PSFVAR\_KEYS}, are a function of the position parameters, \texttt{XWIN\_IMAGE, YWIN\_IMAGE}. The sampling and the image dimension are adjusted for our experiments. We use a \texttt{PSF\_SAMPLING} of \texttt{1.} and \texttt{0.33}, and a \texttt{PSF\_SIZE} of \texttt{32,32} and \texttt{64,64}, for reconstructions at the observation resolution and at $3$ times the observation resolution, respectively. We use a \texttt{PSFVAR\_DEGREES} of \texttt{3}, the degree of the polynomial used for the spatial variation function, as we found that it gave better results than other values. See \autoref{sc:related_work} for a detailed description.
    \item[\textit{vii)}] \texttt{RCA}: the current state-of-the-art data-driven PSF model \cite{schmitz2020} which has been specially designed for the \textit{Euclid} mission. The model is build independently on each CCD. We verified that the parameters used in \cite{schmitz2020} were the best performing ones. Therefore, we use \texttt{4} features, or EigenPSFs, with a denoising parameter \texttt{K} of \texttt{3}. We adjust the upsampling parameter \texttt{D} to match the two different reconstruction resolutions desired that in our case is \texttt{1} and \texttt{0.33}.
    \item[\textit{viii)}] \texttt{MCCD}: a recently developed data-driven PSF model \cite{liaudat2020} that combines ideas from the two previous models. \texttt{PSFEx} and \texttt{RCA} build a PSF model per CCD chip, using fewer stars.
    The \texttt{MCCD} model is able to build a PSF model in the full FOV taking into account the geometry of the instrument.
    We compare models with different parameters and selected the one that was the best performing. The global polynomial degree is set to \texttt{6} with a global denoising parameter, \texttt{K\_GLOB}, set to \texttt{3}. For the local part, the hybrid model is being used with \texttt{4} graph features and a degree of \texttt{2} for the polynomial features. The local denoising parameter, \texttt{K\_LOC}, is also set to \texttt{3}.
\end{itemize}

Our setup for testing \texttt{PSFEx} and \texttt{RCA} is different from \cite{ngole2016} and \cite{schmitz2020}. Here, the observed PSF field has a spatially varying SNR, and we simulate more realistic polychromatic PSFs. The observed PSF positions are uniformly distributed in the FOV. The studies \cite{ngole2016, schmitz2020} used monochromatic PSFs sampled at $600$ nm, requiring a super-resolution factor of $2$ to be Nyquist sampled. However, in the current experiment, when working with polychromatic PSFs, we consider monochromatic PSFs from $550$ nm to $900$ nm, requiring an integer super-resolution factor of $3$ to be Nyquist sampled. As a consequence of the different experimental setup, the performance results of \texttt{PSFEx} and \texttt{RCA} are expected to be different compared to \cite{ngole2016, schmitz2020}.

Additional parameters of the models such as learning rates and the number of epochs can be found in \ref{sc:model_parameters}.

\subsubsection{Data set adaptation for \texttt{PSFEx}, \texttt{RCA}, and \texttt{MCCD}}

The \texttt{PSFEx} and \texttt{RCA} models build independent PSF models for each CCD chip. However, we have simulated the PSF field in the full FOV at once. Therefore, it is necessary to divide the FOV respecting the geometry of the instrument. In this case, we follow \textit{Euclid}'s  VIS instrument geometry \cite{cropper2016}. The focal plane consists of a $6\times6$ CCD matrix, where each CCD is a $4096$ pixel square. We divide the FOV into $36$ subsets, depending on the star positions. As the star positions are uniformly distributed over the FOV, the star density is approximately constant on each CCD. The input of the \texttt{RCA} model is shared with \texttt{MCCD}. The latter model continues with internal processing where all the CCDs are merged into a single data set with global coordinates. We have included the geometry and the coordinate transformations required by the \texttt{MCCD} model.

To create the PSF model with \texttt{PSFEx}, we need to process the simulations with \texttt{SExtractor} \cite{bertin2011}. For that, we create a full CCD image of $4096 \times 4096$ pixels with the postage stamps of the stars at their corresponding positions. Each star has its already defined varying SNR. Choosing the right parameters to make sure that \textit{all} the stars are detected and extracted by \texttt{SExtractor} can be tedious, in particular for low SNR stars. We build a second mock CCD image that we only use for the detection. In this new mock image with low-valued pixels, we place high-valued pixels on the positions of the stars thus forcing the detection of the stars. Then, the extraction of the stars is done using the first mock CCD image. Following this procedure, we make sure that all the stars are correctly detected and extracted by \texttt{SExtractor}. Finally, the star observations can be used to build PSF models with \texttt{PSFEx}.

\subsection{Results}
The results presented here were obtained from the set of testing stars that were not used for training the different PSF models. The WaveDiff models were trained using a Nvidia Tesla V100 GPU of 32 GB. The training time scales with the number of observations in the training set. However, the number of epochs used for the training also varies. The training of the WaveDiff-original model of the $2000$-star data set took $11.5$ hours, while the $200$-star data set took $8.2$ hours. Note that the hyperparameters, especially the number of epochs, were set trying to maximize the performance and not the speed.

\subsubsection{Results: (a), Polychromatic errors}
\label{sc:result_objective_a}

\autoref{tb:results_a} summarizes the pixel reconstruction performance results on polychromatic images. There is a significant improvement of WaveDiff-original and WaveDiff-polygraph, compared to \texttt{PSFEx}, \texttt{RCA}, and \texttt{MCCD}. These three comparison models do not use any type of forward model and are blind to the physics involved in the inverse problem we are tackling. Considering the reconstruction at the observation resolution (x$1$), the WaveDiff-original model reaches $6$ to almost $11$ times lower absolute errors. The errors are decreased by a factor between $34$ and $51$ when considering the super-resolution task (x$3$).

\begin{table}
    \centering
    \caption{Polychromatic test star reconstruction $Err_{\text{abs}}$ and $Err_{\text{rel}}$ at the observation resolution (x1) and at super resolution (x3). The presented results were obtained with models being trained on the largest training star set $\mathcal{S}$. \nblink{table-pixel-errors-metrics}}
        \begin{tabular}{lcc}
            \toprule
            &\multicolumn{2}{c}{$Err_{\text{abs}}$ [$\times 10^{-5}$] ($Err_{\text{rel}}$)}   \\
            \cmidrule(r){2-3}
            PSF model   & Resolution x1 & Resolution x3 \\
            \midrule
            \textit{i)} Zernike 15                   &  $77.5$ ($10.6\%$)    & $19.0$ ($12.8\%$)     \\
            \textit{ii)} Zernike 40                  &  $32.6$ ($4.4\%$)     & $8.7$ ($5.9\%$)      \\
            \textbf{\textit{iii)} WaveDiff-original} &  $\mathbf{6.4}$ ($\mathbf{0.86\%}$)     & $\mathbf{1.9}$ ($\mathbf{1.3\%}$)      \\
            \textit{iv)} WaveDiff-graph             &  $71.8$ ($9.8\%$)    & $19.0$ ($12.7\%$)      \\
            \textit{v)} WaveDiff-polygraph          &  $6.9$ ($0.94\%$)    & $1.9$ ($1.3\%$)      \\
            \textit{vi)} \texttt{PSFEx}             &  $69.2$ ($9.5\%$)    & $66.3$ ($43.0\%$)     \\
            \textit{vii)} \texttt{RCA}              &  $39.6$ ($5.4\%$)    & $85.3$ ($55.5\%$)     \\
            \textit{viii)} \texttt{MCCD}            &  $43.5$ ($6.0\%$)    & $97.7$ ($63.4\%$)     \\
            \bottomrule
        \end{tabular}
    \label{tb:results_a}
\end{table}

The Zernike $15$ model, or \textit{i)}, underperforms with respect to the WFE-based models at the observation resolution. It is even poorer than the state-of-the-art models, thus highlighting the lack of representation of a reduced number of Zernike polynomials. If we remove the data-driven features from the WaveDiff models we obtain a model that is equivalent with model \textit{i)}. This difference in the performance between model \textit{i)} and the WaveDiff models remark the importance of the data-driven features in the WFE-based PSF model. It also shows the effectiveness of the data-driven features in adapting to the observations and generalizing to the target positions.
Concerning the super-resolution task, the model \textit{i)} still outperforms the pixel-based models. There is a huge performance gap between the models using a differentiable forward model and the ones that are building their models directly in the pixel space. The super-resolution is taken into account naturally when using the optical forward model. This handling of the super-resolution is done in the image formation process that is encoded in the forward model with the subsequent downsampling to obtain a PSF at the observation resolution. The observed performance gap underlines the importance of adding prior physical information to the inverse problem to solve a hard task such as super-resolution.

Model \textit{ii)} is performing poorly with respect to WaveDiff-original and WaveDiff-polygraph, even though the number of Zernike polynomials, $d_{\text{Z}}=40$, is close to the ground truth. This small mismatch leads to large PSF reconstruction errors, which demonstrates the rigidity of parametric models. The WaveDiff-original model with only $15$ Zernike polynomials does not suffer from this restriction, showing the flexibility of the data-driven component of the model.

It is important to compare the performance of the different WaveDiff flavors. There is a remarkable difference, around an order of magnitude, between the WaveDiff-graph and the original and polygraph flavors. This fact shows that \textit{only} using the graph constraint for the weight function is not effective in building an appropriate WFE manifold for the PSF field, in contrast with the polynomial weight functions. \autoref{fi:eigen_psfs} presents some examples of learned data-driven WFE features for the WaveDiff-original model. The features show that some structure was learned indicating that we are estimating a WFE manifold. The WFE manifold learning is a consequence of the optimization through the differentiable optical model used in the proposed framework. Note that no constraint has been applied to the data-driven WFE features. We found that the polynomial weight functions, that are more constrained than their graph counterparts, allow learning more structure into the data-driven features.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/DD_features_example.pdf}
    \caption{Examples of learned data-driven eigenWFEs, or WFE features, of the WaveDiff-original model. These features correspond to polynomial feature weight functions. Units are in $\mu$m. \nblink{DD-features-dataset-examples}}
    \label{fi:eigen_psfs}
\end{figure}

\autoref{fi:visual_reconstruction_psf} presents the reconstruction of a test PSF done by the WaveDiff-original model. The first three rows show the excellent reconstruction quality of the WaveDiff-original model for the different pixel scenarios. The fourth row shows important errors in the WFE reconstruction, and nonetheless, the pixel reconstructions are very similar. This fact shows the richness of the WFE space, and how two very different WFE maps can recreate particularly similar pixel PSF images.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/PSF_reconstruction_example.pdf}
    \caption{Visual reconstruction of a test PSF with the WaveDiff-original model. The first column shows the ground truth noiseless PSF, the second column the reconstruction of the PSF model, and the third column the residual between the ground truth and the model.
    \textbf{First row:} Polychromatic pixel representation at the observation resolution. \textbf{Second row:} Polychromatic representation at $3$ times the observation resolution. \textbf{Third row:} Monochromatic pixel representation at $750$ nm and at $3$ times the observation resolution. \textbf{Fourth row:} WFE representation of the PSF, where units are in $\mu$m. We have removed the mean of the WFE and applied the obscurations.
    The first three rows show the absolute value of the residual while the last one just the residual.
    \nblink{DD-features-dataset-examples}}
    \label{fi:visual_reconstruction_psf}
\end{figure}

\subsubsection{Results: (b), Monochromatic errors}
\label{sc:result_objective_b}

\autoref{fi:result_b} presents the pixel reconstruction errors as a function of wavelength.
We only present the WaveDiff-original model as it is the best performing in the results from the objective \textit{(a)}. One can see that the errors are kept low over all the instrument's passband. The model successfully captures the chromatic variations of the ground truth PSF field. The mean relative error over the passband is $1.46\%$. It is important to highlight this result, as the data used to constrain the PSF model are noisy polychromatic PSFs. The loss function we use to optimize our model is built on the reconstruction of polychromatic PSFs. Estimating chromatic variations from polychromatic observations is a difficult ill-posed problem. There are several sets of monochromatic PSFs that can represent the same polychromatic PSF, as information is lost in the spectral integration of \autoref{eq:poly_star}. Ensuring that the monochromatic error is kept low is an essential step in the validation of the proposed PSF model. One can also notice the error bars in \autoref{fi:result_b}, that represent the standard deviation of the reconstruction error within the set of test stars. The variance in the error is kept considerably low with respect to $Err_{\text{abs}}$, the mean RMS error on the test set. A visual reconstruction example of a monochromatic PSF can be seen in the third row of \autoref{fi:visual_reconstruction_psf}.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/WaveDiff-original_monochromatic.pdf}
    \caption{Test star reconstruction error as a function of wavelength for the WaveDiff-original model at the observation resolution.
    \textbf{Left axis:} Absolute error, or $Err_{\text{abs}}$, and is plotted with a solid line. The error bars represent the variance over the RMS errors of the different test stars being used for the calculation of $Err_{\text{abs}}$.
    \textbf{Right axis:} Relative error, or $Err_{\text{rel}}$, and is plotted with a dashed line.
    \nblink{WaveDiff-original-monochromatic-rmse}}
    \label{fi:result_b}
\end{figure}

\subsubsection{Results: (c), Errors as a function of the total number of training stars}
\label{sc:result_objective_c}

The PSF modeling task becomes more arduous as we reduce the number of training stars used to constrain the PSF model. This fact is expected as there is less information available to build the PSF models. \autoref{fi:result_c_x1_res} presents the pixel reconstruction relative errors, $Err_{\text{rel}}$, as a function of the number of training stars. WaveDiff-original and WaveDiff-polygraph maintain excellent performance even with the smallest training set that has $10\%$ of the stars in $\mathcal{S}$. These results show the robustness of the proposed models with a low number of observed stars. Both models are able to properly handle a variable number of training stars with different SNR. WaveDiff-original and WaveDiff-polygraph are able to incorporate the information present in the stars as the reconstruction errors monotonically decrease when we augment the number of training stars. \texttt{PSFEx} shows a rather constant performance with the variation of the number of training stars. This result can be explained by the fact that it is the most constrained model. As a consequence, it has difficulties to incorporate more information into the model, but, in turn, it is more robust to a variation in the number of training stars. \texttt{RCA} exhibits rather similar behavior to \texttt{MCCD}, with the performance improving as the number of stars increases. \texttt{RCA} is not able to produce a result with the smallest data set $\mathcal{S}_{1}$. However, \texttt{MCCD} is still able to work with such a low number of training stars. The results from the largest data set, $\mathcal{S}$, in \autoref{fi:result_c_x1_res} are the same as the ones from the left column of \autoref{tb:results_a}.

We repeat the previous experiment with the reconstruction error being evaluated at super-resolution, and we obtain the results shown in \autoref{fi:result_c_x3_res}. The performance seen in the $2000$-star training data set, $\mathcal{S}$, coincides with the right column of \autoref{tb:results_a}. One can observe the same trend as before, with a huge performance gap between the WaveDiff family, which uses an optical forward model, and the state-of-the-art models. One can notice that now the errors of the WaveDiff-original and WaveDiff-polygraph models are not monotonically decreasing as we increase the number of stars. Nevertheless, the errors are always low, never going over the relative $6\%$ error. This result shows that, for a varying number of training stars, the proposed models are satisfactorily generalizing to test positions as well as capturing the complexity of the high-resolution PSF field. Between the state-of-the-art pixel-based models, \texttt{PSFEx} is the best one performing in the super-resolution reconstruction task. In the reconstruction at the observation resolution \texttt{PSFEx} is underperforming with respect to \texttt{RCA} and \texttt{MCCD}. One explanation of this behavior is that in our experimental setup we perform a super-resolution with a factor of $3$, which differs from the factor of $2$ used in \cite{ngole2016, schmitz2020} as already mentioned in \autoref{sc:experiment_set_up}. As the problem becomes harder, complex models tend to be degenerate if they are not well regularized. However, simpler models like \texttt{PSFEx} tend to be more robust and give consistent results, as they are regularized by their limited capacity. This fact illustrates the importance of validating the reconstruction results at both resolutions.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/AllModels_relative_PixRMSE_vs_NbStars_WF_SOTA.pdf}
    \caption{Polychromatic PSF relative reconstruction error, $Err_{\text{rel}}$, at observation resolution as a function of the total number of training stars in the FOV. The training data sets correspond, in ascending order, to $\mathcal{S}_1 \subset \mathcal{S}_2 \subset \mathcal{S}_3 \subset \mathcal{S}$, described in \autoref{sc:experiment_set_up}. The results are computed with respect to the reconstruction of test stars. The set of test stars is maintained throughout the models and the training sets allowing a truthful comparison. The result of \texttt{RCA} for the $\mathcal{S}_1$ ($200$ stars) data set is not shown as it is very poor.
    \nblink{WaveDiff-original-rmse-vs-star-nb}}
    \label{fi:result_c_x1_res}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/SR_AllModels_relative_PixRMSE_vs_NbStars_WF_SOTA.pdf}
    \caption{Polychromatic PSF relative reconstruction error, $Err_{\text{rel}}$, at three times (x$3$) the observation resolution as a function of the total number of training stars in the FOV. This experiment represents the super-resolution task.
    \nblink{WaveDiff-original-rmse-vs-star-nb}}
    \label{fi:result_c_x3_res}
\end{figure}

\subsubsection{Results: (d), WFE reconstruction errors}
\label{sc:result_objective_d}

The advantage of validating the PSF models with simulations is that we have access to the ground truth WFE of the PSF field. \autoref{tb:results_d_WFE} presents the performance results on the WFE recovery for the models based on WFE. It is interesting to observe that all the models obtain high errors in the WFE recovery. This result illustrates the difficulty of recovering the WFE from degraded in-focus observations. The ground truth WFE field is more complex than what the PSF models can represent, as mentioned in \autoref{sc:experiment_set_up}, so a zero WFE recovery error was not possible to achieve.

Even though WaveDiff-original and WaveDiff-polygraph achieve a very good pixel reconstruction performance, they have large errors in WFE recovery. These results indicate that the WFE manifold estimated by these models is far from the ground truth WFE manifold. However, the pixel PSF fields the different WFE manifolds represent are very close, and this is what matters to our PSF model. The WFE manifold estimated by our model in these experiments is an auxiliary product allowing our models to incorporate physical information, capture chromatic variations, and achieve good super-resolution performance. \autoref{fi:visual_reconstruction_psf} presents an example of two similar PSFs being generated by two very different WFEs. The WFE recovery can be found in the fourth row.

\begin{table}
    \centering
    \caption{Target stars WFE recovery errors with the PSF models being trained on the $\mathcal{S}$ data set. \nblink{table-WFE-error-metrics}}
        \begin{tabular}{lc}
            \toprule
            PSF model                 & WFE $Err_{\text{abs}}$ [nm] ($Err_{\text{rel}}$) \\
            \midrule
            \textit{i)} Zernike 15             &  $106$ ($136\%$)       \\
            \textit{ii)} Zernike 40   &  $102$  ($130\%$)        \\
            \textbf{\textit{iii)} WaveDiff-original}    &  $\mathbf{101}$ ($\mathbf{129\%}$)        \\
            \textit{iv)} WaveDiff-graph       &  $108$ ($138\%$)       \\
            \textit{v)} WaveDiff-polygraph    &  $111$ ($142\%$)       \\
            \bottomrule
        \end{tabular}
    \label{tb:results_d_WFE}
\end{table}

\subsubsection{Results: (e), Weak-lensing metrics}
\label{sc:result_objective_e}

\autoref{tb:results_e} presents the performance results of the different PSF models using weak-lensing metrics based on the shape and the size of the PSF. The results obtained are coherent with the pixel errors in the super-resolution column, see \autoref{tb:results_a}. This observation is expected as the weak-lensing metrics are computed on the polychromatic super-resolved reconstructed PSFs. The error reduction achieved by the WaveDiff-original model with respect to \texttt{RCA} constitutes a \textit{breakthrough in performance} for data-driven PSF modeling. The $e_1$ and $e_2$ errors are $20$ and $26$ times lower, and the size error, $R^{2}/\langle R^{2} \rangle $, is $276$ times lower. These results highlight the importance of the proposed model for current and future weak-lensing studies.

\begin{table}
    \centering
    \caption{Weak-lensing metrics of the different PSF models trained with the $\mathcal{S}$ data set.
    \nblink{table-shape-size-error-metrics}}
        \begin{tabular}{lccc}
            \toprule
            &\multicolumn{3}{c}{RMSE}   \\
            \cmidrule(r){2-4}
            PSF model   & $e_1$[$\times 10^{-2}$] & $e_2$[$\times 10^{-2}$] & $R^{2}/\langle R^{2} \rangle $[$\times 10^{-1}$] \\
            \midrule
            \textit{i)} Zernike 15                   &  $3.92$    & $2.69$  & $0.92$    \\
            \textit{ii)} Zernike 40                  &  $2.12$    & $1.38$  & $0.97$    \\
            \textbf{\textit{iii)} WaveDiff-original} &  $\mathbf{0.23}$     & $0.16$  & $\mathbf{0.13}$    \\
            \textit{iv)} WaveDiff-graph              &  $4.19$    & $1.64$  & $1.39$    \\
            \textit{v)} WaveDiff-polygraph           &  $0.24$    & $\mathbf{0.14}$   & $0.16$    \\
            \textit{vi)} \texttt{PSFEx}              &  $4.38$    & $4.25$  & $14.8$    \\
            \textit{vii)} \texttt{RCA}               &  $4.61$    & $4.17$  & $36.0$   \\
            \textit{viii)} \texttt{MCCD}             &  $9.79$    & $7.32$  & $47.4$   \\
            \bottomrule
        \end{tabular}
    \label{tb:results_e}
    \end{table}

\section{Discussion}
\label{sc:discussions}

In this section, we discuss several scientific choices, the utility of the present work, possible improvements, and extensions.

In \autoref{sc:training} we introduced a regularizer term in the loss function. One of the two best-performing models, WaveDiff-original, does not include any type of regularization. We tried using a $\nu \ell_2(\Phi_{\theta})$ regularization term that constrained the total amount of aberration that is represented by the energy of the WFE map. The physical motivation was to favor models that are less aberrated, as we saw that the WaveDiff-original model converged towards a WFE manifold with more aberration than the ground truth. This result is illustrated by the WFE reconstruction in the fourth row of \autoref{fi:visual_reconstruction_psf}. After tuning the $\nu$ hyperparameter accompanying the regularizer and running the experiments, we noticed that all the pixel performance results were degraded with respect to the original model. However, the WFE reconstruction error decreased considerably. Even though the estimated WFE manifold got closer to the ground truth WFE manifold of the PSF field, the pixel performance worsened.
This experience showcases that regularizers in the WFE space should be used with care. We identify two regimes of the proposed WFE-based PSF modeling. The first regime is aiming to reconstruct the \textit{ground truth} WFE manifold. A typical example would be parametric models. The second regime is aiming to reconstruct a \textit{useful} WFE manifold. In this case, \textit{useful} refers to a WFE manifold that has a similar pixel representation to the ground truth WFE manifold. The WaveDiff model we are proposing is an example of this regime. Introducing a regularizer on the WFE would make sense on the first regime. Particularly if \textit{a priori} information about the ground truth WFE manifold is available and could be used to build the regularizer. This situation could be the case if complementary measurements of the optical system are available. The second regime should be taken with more care as the WFE regularizer can bias the estimated useful WFE manifold. This bias represents a larger deviation in pixel space even if in the WFE space both manifolds become closer.

In the previous section, this work presented certain advantages of the proposed data-driven approach with respect to parametric PSF models. It was done in the comparison of the Zernike $15$ and Zernike $40$ models with the WaveDiff-original model. However, it does not mean that parametric and data-driven models cannot be complementary.
The WFE output of the parametric model can be used as input in the proposed framework. The model would shift from the first to the second regime mentioned above, trying to reconstruct the ground-truth WFE rather than a useful representation. The parametric model output could be introduced to the data-driven model using a WFE regularizer. The task of the data-driven optimization would be to estimate the difference between the ground truth WFE and the WFE estimated by the parametric model.

This work finds similar performances for the WaveDiff-original and WaveDiff-polygraph models. The results presented in \autoref{tb:results_a} showed a slightly better performance of the WaveDiff-original model. We therefore chose it to create the visual examples of \autoref{fi:eigen_psfs} and \autoref{fi:visual_reconstruction_psf}. It is however hard to conclude which model is better with the different metrics used. The structure imposed into the $\Phi_{\theta}$ function constrains the WFE manifold learned from the observations. Which $\Phi_{\theta}$ function allows learning a better WFE manifold depends on the ground-truth WFE which in turn depends on the telescope we are modeling. In the current setting, both models achieve great performances. A follow-up study would be to evaluate the models under high-fidelity optical simulations that generally involve costly ray-tracing techniques.

Even though the simulations used in this work are representative of a space telescope, there are several simplifications we have assumed. For example, we have considered that all the CCD chips are perfectly aligned in the focal plane, and not included any detector-level effects. These and other effects could be included in the proposed framework. Detector-level effects, such as guiding errors or charge transfer inefficiency, could be included in the optical forward model. An interesting study, but out of scope for this article, is to evaluate how do errors in the input SED data propagate to the estimated PSF model.

The proposed framework is not only applicable to space telescopes but could be used for ground-based telescopes. The next generation of ground telescopes being built, like the Vera C. Rubin observatory \cite{LSST2009}, pose serious challenges to the control of systematic errors. The PSF is one of the biggest contributors. Therefore, it becomes important to model the atmospheric chromaticity to build a chromatic PSF model like the one we are proposing. The atmosphere prevented the adoption of parametric models for ground-based telescopes. The fact that our model is data-driven and physically motivated makes it a good candidate to prove itself useful. We use the WFE as the modeling space, therefore it is possible to include WFE-based atmospheric models, as Kolmogorov phase screens \cite{jee2011}. One can then optimize over the atmospheric component, thus incorporating an atmospheric contribution in the PSF model.

New flavors of the WaveDiff PSF model can be developed by taking advantage of our framework and redefining the function $\Phi_{\theta}$ in \autoref{fi:psf_model_diagram}. Previous works that tried to leverage the power of modern neural networks into PSF models were blind to the physics of the inverse problem and built the model in pixel space \cite{jia2020a,jia2020b,jia2020c}. Our framework uses an end-to-end differentiable optical forward model, runs on GPU, and is entirely coded in TensorFlow \cite{tensorflow2015}. This construction makes it an interesting candidate to develop \textit{physically-motivated} neural network approaches to boost the PSF model.

\section{Conclusion}
\label{sc:conclusions}

This work introduces a novel data-driven modeling method of a point spread function field coined WaveDiff which represents a paradigm shift with respect to current state-of-the-art data-driven methods. We propose to include a differentiable physics-based forward model that allows changing the modeling space from pixels to the wavefront. This change transfers most of the modeling complexity into the forward model and simplifies the building of the point spread function model. As the entire optical forward model is differentiable we can use modern optimization techniques to estimate the model parameters in the wavefront space. The WaveDiff model is built using the TensorFlow framework and runs in GPU.
We have verified that our model can learn from in-focus observations a \textit{useful} wavefront-based manifold representation of the PSF field. We have not used special calibration data or wavefront information. The WaveDiff model is only constrained from noisy in-focus polychromatic under-sampled observations using the spectral energy distribution of each observed object. The proposed WaveDiff model is capable of super-resolving the PSF field, and is, for the best of our knowledge, the \textit{first data-driven PSF model} to effectively capture the wavelength dependence of the point spread function.

The proposed method was tested using a set of simulations of a space telescope which included a broad passband, different spectral energy distribution for the observed objects, variable signal-to-noise ratios, and the undersampling of the observations. We compared our WaveDiff model against three state-of-the-art data-driven point spread function models, including \texttt{PSFEx}, \texttt{MCCD}, and \texttt{RCA}, the current data-driven approach for the \textit{Euclid} satellite. The WaveDiff model represents a \textit{breakthrough in performance} obtaining $6$ and $44$ \textit{times} lower reconstruction error than \texttt{RCA} at observation resolution and super-resolution (x$3$), respectively. Our proposed model obtains $0.86\%$ and $1.3\%$ relative reconstruction error at x$1$ and x$3$ the observation resolution, respectively. We validated that the WaveDiff model is well capturing the chromatic variations of the point spread function with the monochromatic reconstruction error. The WaveDiff model achieves a mean relative error of $1.46\%$ over the broad instrument's passband. We have further tested the WaveDiff model varying the number of observed stars used to constrain the model. We have found that our model maintains good performance even when using the $10\%$ of the original observed stars.

The use of simulations for the validation of the model allows accessing and comparing the estimated and the ground truth wavefronts. We have found that the wavefront manifold estimated by the WaveDiff model and the ground truth are far from each other. However, and interestingly, both wavefront manifolds represent very similar pixel point spread function fields as was shown in the pixel errors above. These results support the idea of using the wavefront space as an auxiliary space to represent the pixel point spread function field. It has not necessarily to be the objective to find the ground truth wavefront manifold, but a manifold that is \textit{useful} for our objectives. These are defined in pixel space rather than in the wavefront space.

Finally, we have performed a comparison of the different point spread function models using weak-lensing metrics such as ellipticity and size. When compared to the \texttt{RCA} model, the WaveDiff model manages to reduce the ellipticity errors by a factor of $20$ and $26$ for each component, respectively. In addition, the size error obtained is $276$ times lower than the one obtained by \texttt{RCA}. These results highlight the utility and confirm the breakthrough in performance achieved by the proposed WaveDiff model.

\bigskip

In the spirit of reproducible and reusable research, all software, scripts, and analysis notebooks are publicly available at \url{https://github.com/tobias-liaudat/wf-psf}.

\section*{Acknowledgments}
This work was granted access to the HPC resources of IDRIS under the allocations 2021-AD011011554 and 2021-AD011012983 made by GENCI. This work has made use of the CANDIDE Cluster at the Institut d'Astrophysique de Paris and made possible by grants from the PNCG and the DIM-ACAV.
\textit{Software used in this work: Astropy \cite{Robitaille2013, PriceWhelan2018}, IPython \cite{Perez2007}, GalSim \cite{rowe2015}, Jupyter \cite{Kluyver2016}, Matplotlib \cite{Hunter2007}, Seaborn \cite{seaborn2021}, Numpy \cite{numpy2020}, TensorFlow \cite{tensorflow2015}.}

\section*{References}

redpajama_set_name: RedPajamaArXiv