confocal PSF as a function of pinhole size

confocal_psf.py

import scipy.special
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import patches
from matplotlib.animation import FuncAnimation

WF_FWHM = 0.8436843684368438
WF_MAX = 2199.5516151207426
WF_INTEGRAL = 1300.2494235915415


def airy_disk(X: np.ndarray, corrected: bool = True):
    """Plot Airy disk at positions X, centered at center."""
    # not sure why this magic number is needed to put the first minimum at 1
    rho = np.pi * X * (1.2195 if corrected else 1)
    return (2 * scipy.special.j1(rho) / rho) ** 2


def tophat(X: np.ndarray, width: float):
    """Plot Airy disk at positions X, centered at center."""
    out = np.zeros_like(X)
    out[np.abs(X) < width] = 1
    return out


def confocal_psf(X, pinhole_au: float = 1, norm: bool = False):
    """Plot the PSF of a confocal microscope with a pinhole of diameter pinhole_au AU."""
    disk = airy_disk(X)
    pinhole = tophat(X, pinhole_au)
    psf = disk * np.convolve(disk, pinhole, mode="same")
    return psf / psf.max() if norm else psf


def fwhm(X, psf, norm_to=WF_FWHM):
    """Return the full width at half maximum of the PSF."""
    x0, x1 = np.where(np.diff(np.sign(psf - np.max(psf) / 2)))[0]
    return (X[x1] - X[x0]) / norm_to


def full_widths(X, pinholes):
    """Plot the full width at half maximum of the PSF for various pinhole sizes."""
    out = []
    for au in pinholes:
        psf = confocal_psf(X, au, norm=False)
        out.append(fwhm(X, psf))
    return np.array(out)


def integrals(X, pinholes):
    """Plot the integral (tot signal) of the PSF for various pinhole sizes."""
    out = []
    norm_to = np.trapz(confocal_psf(X, 100, norm=False), X)
    for au in pinholes:
        psf = confocal_psf(X, au, norm=False)
        out.append(np.trapz(psf, X) / norm_to)
    return np.array(out)


def pinhole_animation(
    x_range=2, steps=100, show_inset=True, show_unnormalized=True, show_pinhole=True
):
    fig, ax = plt.subplots()
    X = np.linspace(-x_range, x_range, 10000)
    pinholes = np.logspace(np.log10(2), np.log10(0.01), steps)

    ax.plot(
        X,
        confocal_psf(X, 100, True),
        alpha=0.2,
        color="black",
        linestyle="--",
        linewidth=0.5,
    )

    (line,) = ax.plot(X, confocal_psf(X, 1, True), alpha=0.6)
    if show_unnormalized:
        (line2,) = ax.plot(X, confocal_psf(X, 1, False) / WF_MAX)
        ax.legend(["(no pinhole)", "normalized", "unnormalized"])

    ax.set_title("Effective Confocal PSF (AU=2.00)")
    ax.set_xlabel("Position (normalized, AU)")
    ax.set_ylabel("Intensity")

    if show_pinhole:
        # pinhole rect
        rect = patches.Rectangle((-2, 0), 4, 1, facecolor="gray", alpha=0.1)
        ax.add_patch(rect)

    # inset
    if show_inset:
        axins = ax.inset_axes([0.1, 0.7, 0.25, 0.25])
        axins.plot(pinholes, full_widths(X, pinholes))
        axins.plot(pinholes, integrals(X, pinholes))
        axins.set_xlabel("pinhole (AU)", fontsize=6)
        # axins.set_ylabel("rel. FWHM", fontsize=6)
        axins.legend(["FWHM", "signal"], fontsize=6, loc="lower right")
        axins.set_ylim(0, 1.1)
        axins.tick_params(axis="both", which="major", labelsize=6)
        vline = axins.axvline(x=2, color="red", linestyle="--", linewidth=0.5)

    def animate(pinhole_au: float):
        """Plot the PSF of a confocal microscope with a pinhole of diameter pinhole_au AU."""
        line.set_ydata(confocal_psf(X, pinhole_au, True))
        if show_unnormalized:
            line2.set_ydata(confocal_psf(X, pinhole_au, False) / WF_MAX)
        if show_inset:
            vline.set_xdata(pinhole_au)

        if show_pinhole:
            rect.set_x(-pinhole_au)
            rect.set_width(pinhole_au * 2)
        ax.set_title(f"Confocal PSF (AU={pinhole_au:.2f})")
        fig.canvas.draw()

    anim = FuncAnimation(fig, animate, frames=pinholes, interval=steps)
    anim.save("confocal_psf.mp4")


if __name__ == "__main__":
    pinhole_animation()