Aberrated Image Recovery of Ultracold Atoms

Aberrated Image Recovery of Ultracold Atoms#

Overview#

In Project 01 you studied how a real imaging system distorts images of an ultracold atomic gas [1], [2]. Its finite resolution blurs the image, defocus and spherical aberration add tails when the gas is thicker than the depth of field (DOF), and the camera pixelizes and adds noise. In this project you will use machine learning to unfold these effects, and recover the true size of the gas from its aberrated image:

  • You will build a large training set with a fast simulator, validated against the simulation used in Project 01.

  • You will establish a conventional baseline: fitting a 2D Gaussian to each image.

  • You will train a convolutional neural network (CNN) to measure the transverse and axial size of the gas from the image alone.

  • You will test how robust the network is when the imaging system differs from the one it was trained on.

Data Sources#

Why a fast simulator? ImSim in ImageSimulation.ipynb convolves each of the \(N_z\) axial slices of the gas with its own defocused PSF, which takes about a second or more per image. That is too slow to make a training set of tens of thousands of images. However, the true density is a Gaussian, \(n(x,y,z) = G_{xy}(x,y)\,g(z)\). Because convolution is linear, the column-integrated image is therefore

\[ \sum_z \left[ G_{xy}\, g(z) \right] * \text{PSF}_z \;=\; G_{xy} * \underbrace{\sum_z g(z)\, \text{PSF}_z}_{\text{PSF}_\text{eff}}, \]

a single convolution with an effective PSF. The class below precomputes the stack of defocused PSFs once, using the same pupil-function model as PSF2D in ImageSimulation.ipynb. After that, each image takes about 15 ms instead of a second or more.

A larger field of view. In Project 01 the simulation grid extended only ±10 µm, which truncated the larger gases. The default geometry below uses a ±20.7 µm field of view (64 × 64 pixels of 0.65 µm in the object plane, as for the high-NA system with \(M = 20\)) and ±30 µm along the imaging axis. This is enough for the parameter ranges used in this project.

%matplotlib inline
import os
import tarfile
import time
import urllib.request

import numpy as np
import scipy as sp
import scipy.fft, scipy.signal, scipy.stats
import skimage.io, skimage.measure
import matplotlib.pyplot as plt

DATA_URL = "https://courses.physics.illinois.edu/phys498mlp/sp2025/data/projects/AberratedImages/data.tgz"
if not os.path.isdir("Heavy_tails"):
    urllib.request.urlretrieve(DATA_URL, "data.tgz")
    with tarfile.open("data.tgz") as tar:
        try:
            tar.extractall(".", filter="data")
        except TypeError:
            tar.extractall(".")

LAM = 0.780e-6   # imaging wavelength (m), Rb-87 D2 line
F = 31.3e-3      # focal length of the objective (m)
PS = 13e-6       # camera pixel size (m)


def psf2d(N, bound_p, defocus, Z04, D_l):
    """PSF = |FFT(pupil)|^2 with defocus (Zernike Z_02) and primary spherical aberration (Z_04),
    exactly as PSF2D in ImageSimulation.ipynb."""
    Z02 = ((1 / F) - (1 / (F + defocus))) * (D_l ** 2) / (8 * LAM)
    x, y = np.mgrid[-bound_p:bound_p:N*1j, -bound_p:bound_p:N*1j]
    r = np.hypot(x, y)
    phase = Z02 * (2 * r**2 - 1) + Z04 * (6 * r**4 - 6 * r**2 + 1)
    pupil = np.where(r <= 1, np.exp(2j * np.pi * phase), 0)
    return np.roll(np.abs(sp.fft.fft2(pupil)) ** 2, N // 2, axis=(0, 1))


class Imager:
    """Fast simulator of absorption images of a Gaussian gas.

    Precomputes the stack of defocused PSFs; each image then needs one 2D convolution with the
    effective PSF sum_z g(z) PSF_z. Physical lengths are in metres.
    """
    def __init__(self, NA=0.4, M=20, Z04=0.2, n_pix=64, oversample=6, N_z=121, bound_z=30e-6):
        self.NA, self.M, self.Z04, self.n_pix, self.b = NA, M, Z04, n_pix, oversample
        self.D_l = 2 * F * np.tan(np.arcsin(NA))
        N_xy = n_pix * oversample
        self.dr = PS / M / oversample                      # fine-grid spacing in the object plane
        self.bound_xy = self.dr * (N_xy - 1) / 2
        pupil_bound = LAM * F / (self.dr * self.D_l)
        assert pupil_bound >= 1, "grid too coarse to sample the PSF"
        self.x = np.linspace(-self.bound_xy, self.bound_xy, N_xy)
        self.z = np.linspace(-bound_z, bound_z, N_z)
        self.psf = np.stack([psf2d(N_xy, pupil_bound, dz, Z04, self.D_l) for dz in self.z]).astype(np.float32)

    def image(self, sigma_xy, sigma_z, N=1, mu=(0, 0, 0), rng=None):
        """Pixelized image of a gas with widths sigma_xy, sigma_z (m), normalized to N atoms.
        If rng is given, add the same per-pixel noise as ImSim (std = sqrt(pixel value))."""
        gx = sp.stats.norm.pdf(self.x, mu[0], sigma_xy)
        gy = sp.stats.norm.pdf(self.x, mu[1], sigma_xy)
        gz = sp.stats.norm.pdf(self.z, mu[2], sigma_z)
        psf_eff = np.tensordot(gz, self.psf, axes=(0, 0))
        col = sp.signal.fftconvolve(np.outer(gx, gy), psf_eff, mode="same")
        im = skimage.measure.block_reduce(col, (self.b, self.b))
        im = im * N / im.sum()
        if rng is not None:
            im = im + np.sqrt(np.clip(im, 0, None)) * rng.normal(size=im.shape)
        return im


def make_dataset(imager, n, rng, sigma_xy_range=(0.5e-6, 4e-6), sigma_z_range=(0.5e-6, 15e-6), N_range=(1e4, 2e5)):
    """Generate n noisy images with random sizes and atom numbers.
    sigma_xy is uniform, sigma_z and N are log-uniform. Returns images (n, n_pix, n_pix) and labels (n, 3) = (sigma_xy, sigma_z, N)."""
    X = np.empty((n, imager.n_pix, imager.n_pix), np.float32)
    Y = np.empty((n, 3), np.float32)
    for i in range(n):
        sxy = rng.uniform(*sigma_xy_range)
        sz = np.exp(rng.uniform(*np.log(sigma_z_range)))
        N = np.exp(rng.uniform(*np.log(N_range)))
        X[i] = imager.image(sxy, sz, N, rng=rng)
        Y[i] = (sxy * 1e6, sz * 1e6, N)                   # sizes stored in um
    return X, Y


t = time.time()
imager = Imager()
print(f"PSF stack: {time.time() - t:.1f} s; field of view ±{imager.bound_xy * 1e6:.1f} um; {imager.n_pix}x{imager.n_pix} pixels")
rng = np.random.default_rng(42)
X_demo, Y_demo = make_dataset(imager, 8, rng)
fig, axs = plt.subplots(1, 8, figsize=(16, 2.4))
for ax, im, (sxy, sz, N) in zip(axs, X_demo, Y_demo):
    ax.imshow(im); ax.set_title(f"$\\sigma_{{xy}}$={sxy:.1f}, $\\sigma_z$={sz:.1f}", fontsize=8); ax.axis("off")

Questions#

General instructions

  • Answer each question in a markdown cell directly below it (add code cells where needed). Your notebook should run from top to bottom without errors. Fix your random seeds.

  • Use PyTorch, as in the Convolutional and Recurrent Neural Networks, Deep Learning and GPU 101 lectures. Keep separate training, validation and test sets. Plot training and validation loss versus epoch, and report the number of trainable parameters, the training time and the hardware used.

Question 01#

Validating the simulator. Before trusting a fast simulator for training data, check it against the original.

  1. Explain in your own words why the effective-PSF method is exact for a Gaussian density, i.e. why the two expressions in the Data Sources section are equal. Would it still be exact for a non-Gaussian density \(n(x,y,z)\) that does not factorize as \(G_{xy}(x,y)\,g(z)\)?

  2. Verify the identity numerically. For one gas, e.g. \(\sigma_{xy} = 1\) µm and \(\sigma_z = 10\) µm, compute the image slice by slice: convolve \(G_{xy}\,g(z_i)\) with each PSF in imager.psf, sum over \(i\), and pixelize. Compare the result with imager.image. They should agree to numerical precision.

  3. Check the physics model against the Project 01 data. Run ImSim from ImageSimulation.ipynb (noise-free, atom_number=1) for one of the Heavy_tails images, e.g. \(\sigma = (1, 1, 20)\) µm, with the HiNA geometry defined in ImageSimulation.ipynb (about 20 s per image). Scale it to \(N = 10^5\) and plot the residuals of the data divided by the expected noise, \(\sqrt{\text{model}}\) (“pulls”). Are they consistent with pure noise (mean ≈ 0, standard deviation ≈ 1)? psf2d above uses the same pupil-function model as PSF2D, so this also validates the fast simulator’s physics.

  4. How much faster is the effective-PSF method than ImSim?

Question 02#

A conventional baseline. Generate a test set of a few thousand images with make_dataset. Fit each image with a 2D Gaussian (e.g. with scipy.optimize.curve_fit) and compare the fitted width with the true \(\sigma_{xy}\).

  • Plot the fitted width divided by the true width as a function of \(\sigma_{xy}\), and colour the points by \(\sigma_z\). Explain the bias you see, using the resolution and DOF from Project 01, Questions 03–04.

  • Can you deconvolve the bias analytically, e.g. by subtracting a PSF width in quadrature? Why does this fail for gases thicker than the DOF?

  • Can the Gaussian fit tell you anything about \(\sigma_z\)?

Question 03#

A CNN that measures the gas. Generate a training set of at least 20,000 images (about 5 minutes on a laptop CPU), plus validation and test sets. Train a CNN that takes an image as input and predicts \(\log\sigma_{xy}\) and \(\log\sigma_z\).

  • Describe your preprocessing. For example, normalizing each image by its sum removes the dependence on \(N\); is that what you want? Also describe your architecture and loss function.

  • Report the test-set accuracy for \(\sigma_{xy}\) and \(\sigma_z\) (e.g. the median relative error), and compare the \(\sigma_{xy}\) accuracy with the Gaussian-fit baseline of Question 02, especially for small gases.

  • Plot the relative error on \(\sigma_z\) as a function of the true \(\sigma_z\). Relate what you see to the DOF of the imaging system (\(\approx \lambda/\text{NA}^2 \approx 5\) µm). Why is the axial size of the gas only measurable in some range? How does the accuracy depend on the atom number, i.e. on the noise?

Question 04#

Robustness to a mis-modelled imaging system. In a real experiment, the aberrations of the imaging system are never known exactly. Generate test sets with an Imager that differs from the one used for training, while keeping your trained network fixed:

  • a different spherical aberration, e.g. Z04 = 0.1 and 0.3 instead of 0.2;

  • a gas whose centre is displaced along the imaging axis, e.g. mu=(0, 0, 3e-6), i.e. out of focus.

How do the predictions of \(\sigma_{xy}\) and \(\sigma_z\) change? Which kind of mis-modelling is most harmful, and why? How could you make the network robust to it? For example, consider training on a range of Z04 and \(\mu_z\) values, or predicting them as additional outputs.

Optional extensions

  • Uncertainty quantification (see the AI Explainability and Uncertainty Quantification lecture): have your network predict a mean and a variance for each output, trained with a Gaussian negative log-likelihood, or train an ensemble of networks. Are the predicted uncertainties calibrated on the test set?

  • Image deconvolution: train a CNN with an image as output (e.g. a small U-Net) to map the aberrated image to the true, un-aberrated column density. Compare it with the classical Richardson–Lucy algorithm [6], [7] (skimage.restoration.richardson_lucy), which needs the PSF as input. Refs. [3], [4] and [5] discuss neural-network approaches to deconvolution and PSF modelling.

  • Back to Project 01: compare your CNN with a small fully connected network that uses as input only the first \(k\) PCA (SVD) coefficients of each image (Project 01, Question 09). How large does \(k\) need to be to measure \(\sigma_z\)?

References#

[1] S. S. Kondov, W. R. McGehee, J. J. Zirbel, B. DeMarco, “Three-dimensional Anderson localization of ultracold matter,” Science 334, 66 (2011). doi:10.1126/science.1209019

[2] W. B. Cairncross et al., “Precision measurement of the electron’s electric dipole moment using trapped molecular ions,” Phys. Rev. Lett. 119, 153001 (2017). doi:10.1103/PhysRevLett.119.153001

[3] R. J. Steriti, M. A. Fiddy, “Blind deconvolution of images by use of neural networks,” Opt. Lett. 19, 575 (1994). doi:10.1364/OL.19.000575

[4] C. J. Schuler, Machine learning approaches to image deconvolution, Ph.D. thesis, Eberhard Karls Universität Tübingen (2015). https://publikationen.uni-tuebingen.de/xmlui/handle/10900/75986

[5] J. Herbel, T. Kacprzak, A. Amara, A. Refregier, A. Lucchi, “Fast point spread function modeling with deep learning,” J. Cosmol. Astropart. Phys. 2018 (07), 054 (2018). doi:10.1088/1475-7516/2018/07/054, arXiv:1801.07615

[6] W. H. Richardson, “Bayesian-based iterative method of image restoration,” J. Opt. Soc. Am. 62, 55 (1972). doi:10.1364/JOSA.62.000055

[7] L. B. Lucy, “An iterative technique for the rectification of observed distributions,” Astron. J. 79, 745 (1974). doi:10.1086/111605

[8] C. J. Foot, Atomic Physics (Oxford University Press, 2005); J. W. Goodman, Introduction to Fourier Optics, 3rd ed. (Roberts & Company, 2005).

Acknowledgements#

  • Initial version: Max Gold with some guidance from Mark Neubauer

  • Revised: 2026

© Copyright 2026